From 61953318bf4fd16393c4a2a60bf398d791cd1f19 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 06:51:28 -0700 Subject: [PATCH 01/58] feat(mcp): share compatibility-aware result conversion across tool surfaces (#43089) * feat(mcp): share compatibility-aware result conversion across tool surfaces Adds one converter that turns text, JSON, SDK results, interim InputRequiredResult values and exceptions into a CallToolResult shaped for the negotiated MCP revision. Legacy revisions keep object-only structuredContent with a lossless text fallback for other JSON values, and reject interim results through failure accounting. Modern revisions pass arbitrary structuredContent and InputRequiredResult through without completed-success accounting or post-call hooks. OpenAPI tools keep the upstream body verbatim and gain structuredContent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): accept the wire compat argument in local-registry fakes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(mcp): read tagged outcome fields directly in the result converter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(mcp): keep the mutable-ok marker on the list literal it suppresses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(mcp): rerun the mcp-integration shard after a tcp cancellation timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): cover result conversion boundaries and explicit returns * fix(mcp): keep SSE connections on the legacy protocol --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 15 +- .../_experimental/mcp_server/contracts.py | 2 + .../mcp_server/mcp_server_manager.py | 65 +++-- .../mcp_server/openapi_to_mcp_generator.py | 9 +- .../_experimental/mcp_server/operations.py | 70 +++-- .../mcp_server/rest_endpoints.py | 4 +- .../mcp_server/result_conversion.py | 120 +++++++++ .../proxy/_experimental/mcp_server/server.py | 30 ++- .../_experimental/mcp_server/tool_outcome.py | 56 ++++ .../_experimental/mcp_server/tool_search.py | 4 +- .../mcp_server/test_mcp_hook_extra_headers.py | 13 +- .../test_mcp_max_concurrent_requests.py | 7 +- .../mcp_server/test_mcp_server.py | 185 ++++++++++++- .../mcp_server/test_mcp_server_manager.py | 41 ++- .../test_openapi_to_mcp_generator.py | 87 +++++-- .../mcp_server/test_openapi_tool_auth.py | 19 +- .../mcp_server/test_operations.py | 37 ++- .../mcp_server/test_rest_endpoints.py | 100 ++++++- .../mcp_server/test_result_conversion.py | 243 ++++++++++++++++++ 19 files changed, 985 insertions(+), 122 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/result_conversion.py create mode 100644 litellm/proxy/_experimental/mcp_server/tool_outcome.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_result_conversion.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 1206f9abcbd..01670be74c8 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -36,6 +36,7 @@ from mcp.types import ( REQUEST_TIMEOUT, GetPromptRequestParams, GetPromptResult, + InputRequiredResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult, @@ -44,7 +45,6 @@ from mcp.types import ( Prompt, ResourceTemplate, ServerNotification, - TextContent, ) from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult @@ -61,6 +61,7 @@ from litellm.constants import ( from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response +from litellm.proxy._experimental.mcp_server.result_conversion import error_text_result from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( MCPAuth, @@ -828,17 +829,15 @@ class MCPClient: @staticmethod def error_tool_result(exc: Exception) -> MCPCallToolResult: """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" - return MCPCallToolResult( - content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")], - is_error=True, - ) + return error_text_result(exc) async def call_tool( self, call_tool_request_params: MCPCallToolRequestParams, host_progress_callback: Callable | None = None, raise_on_error: bool = False, - ) -> MCPCallToolResult: + allow_input_required: bool = False, + ) -> MCPCallToolResult | InputRequiredResult: """ Call an MCP Tool. @@ -847,6 +846,9 @@ class MCPClient: ``isError=True`` result. The token-exchange (OBO) tool-call path uses this to detect an upstream 401 so it can re-mint the exchanged token and retry once; every other caller keeps the default and gets graceful ``isError`` degradation. + allow_input_required: When True, a 2026-07-28 upstream may answer with an interim + ``InputRequiredResult`` and it is returned as is. The SDK rejects it otherwise, so + callers only opt in when the downstream side can carry it. """ verbose_logger.info("MCP client calling tool '%s'", call_tool_request_params.name) @@ -869,6 +871,7 @@ class MCPClient: name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, progress_callback=on_progress, + allow_input_required=allow_input_required, ) try: diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index 1879e285789..a88d400282c 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -5,6 +5,7 @@ from datetime import datetime from types import MappingProxyType from typing import Final, Protocol +from litellm.proxy._experimental.mcp_server.tool_outcome import WireCompat from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -26,6 +27,7 @@ class OperationContext: raw_headers: Mapping[str, str] | None = field(default=None, repr=False) client_ip: str | None = None mcp_proxy_mode: bool = False + wire_compat: WireCompat = WireCompat.LEGACY def __post_init__(self) -> None: object.__setattr__(self, "_caller", copy_caller(self._caller)) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index be4df55ff58..24cae976174 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -44,6 +44,7 @@ from mcp.types import ( CallToolResult, GetPromptRequestParams, GetPromptResult, + InputRequiredResult, Prompt, ResourceTemplate, ) @@ -133,6 +134,12 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ServerSpec, TokenExchangeConfig, ) +from litellm.proxy._experimental.mcp_server.result_conversion import ( + WireCompat, + complete_call_tool_result, + handler_outcome, + to_gateway_tool, +) from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) @@ -5361,16 +5368,9 @@ class MCPServerManager: prefix: Final = get_server_prefix(server) for tool in tools: - tool_copy = tool.model_copy(deep=True) - - original_name = tool_copy.name + original_name = tool.name prefixed_name = add_server_prefix_to_name(original_name, prefix) - - name_to_use = prefixed_name if add_prefix else original_name - - # Preserve all tool fields including metadata/_meta by avoiding mutation - tool_copy.name = name_to_use - prefixed_tools.append(tool_copy) + prefixed_tools.append(to_gateway_tool(tool, prefixed_name if add_prefix else original_name)) # Register every known prefix form (alias, server_name, server_id, # short ID) so call_tool can resolve regardless of which form a @@ -5547,6 +5547,7 @@ class MCPServerManager: server: MCPServer, tool_name: str, arguments: _ToolArguments, + wire_compat: WireCompat = WireCompat.LEGACY, ) -> CallToolResult: """ Call an OpenAPI tool handler directly. @@ -5586,14 +5587,7 @@ class MCPServerManager: # Call the tool handler with the arguments # The handler is an async function that makes the HTTP request handler_result: Final = await tool.handler(**arguments) - - # Convert the handler result (string response) to CallToolResult format - result: Final = CallToolResult( - content=[TextContent(type="text", text=str(handler_result))], - is_error=False, - ) - - return result + return complete_call_tool_result(handler_outcome(handler_result), wire_compat) except MCPUpstreamAuthError: # The caller must re-authenticate upstream, so this keeps its type all the way to the @@ -5820,7 +5814,8 @@ class MCPServerManager: user_api_key_auth: UserAPIKeyAuth | None, raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, - ) -> CallToolResult: + allow_input_required: bool = False, + ) -> CallToolResult | InputRequiredResult: """Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry. The exchanged token is baked into the client at build time, so the retry invalidates the @@ -5830,7 +5825,10 @@ class MCPServerManager: """ try: return await client.call_tool( - call_tool_params, host_progress_callback=host_progress_callback, raise_on_error=True + call_tool_params, + host_progress_callback=host_progress_callback, + raise_on_error=True, + allow_input_required=allow_input_required, ) except Exception as exc: if _extract_upstream_auth_failure(exc) is None: @@ -5848,7 +5846,11 @@ class MCPServerManager: raw_headers=raw_headers, client_ip=client_ip, ) - return await retry_client.call_tool(call_tool_params, host_progress_callback=host_progress_callback) + return await retry_client.call_tool( + call_tool_params, + host_progress_callback=host_progress_callback, + allow_input_required=allow_input_required, + ) async def _call_regular_mcp_tool( self, @@ -5865,7 +5867,8 @@ class MCPServerManager: hook_extra_headers: dict[str, str] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, client_ip: str | None = None, - ) -> CallToolResult: + allow_input_required: bool = False, + ) -> CallToolResult | InputRequiredResult: """ Call a regular MCP tool using the MCP client. @@ -6036,6 +6039,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, raw_headers=raw_headers, client_ip=client_ip, + allow_input_required=allow_input_required, ) tool_call_coro = _obo_call_tool_limited() @@ -6049,7 +6053,11 @@ class MCPServerManager: async def _call_tool_via_client(client, params): async with self._limit_outbound_concurrency(mcp_server): if not relays_upstream_auth: - return await client.call_tool(params, host_progress_callback=host_progress_callback) + return await client.call_tool( + params, + host_progress_callback=host_progress_callback, + allow_input_required=allow_input_required, + ) # The client-forwarded modes carry the caller's own upstream token, so an upstream # 401 (expired/invalid token) is the caller's to resolve: relay it as # MCPUpstreamAuthError so single-server REST callers turn it into a 401 + @@ -6061,7 +6069,10 @@ class MCPServerManager: # the same isError degradation the default path produces. try: return await client.call_tool( - params, host_progress_callback=host_progress_callback, raise_on_error=True + params, + host_progress_callback=host_progress_callback, + raise_on_error=True, + allow_input_required=allow_input_required, ) except Exception as e: auth_info: Final = _extract_upstream_auth_failure(e) @@ -6114,7 +6125,7 @@ class MCPServerManager: result: Final = mcp_responses[result_index] self._remember_upstream_initialize_instructions(mcp_server, client) - return cast(CallToolResult, result) + return cast("CallToolResult | InputRequiredResult", result) def _resolve_mcp_server_for_tool_call( self, @@ -6318,7 +6329,8 @@ class MCPServerManager: litellm_logging_obj: "LiteLLMLoggingObj | None" = None, guardrail_context: Mapping[str, object] | None = None, client_ip: str | None = None, - ) -> CallToolResult: + wire_compat: WireCompat = WireCompat.LEGACY, + ) -> CallToolResult | InputRequiredResult: """ Call a tool with the given name and arguments @@ -6427,7 +6439,7 @@ class MCPServerManager: resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) try: async with self._limit_outbound_concurrency(mcp_server): - return await self._call_openapi_tool_handler(mcp_server, name, arguments) + return await self._call_openapi_tool_handler(mcp_server, name, arguments, wire_compat) finally: _request_auth_header.reset(auth_token) _request_extra_headers.reset(extra_token) @@ -6449,6 +6461,7 @@ class MCPServerManager: host_progress_callback=host_progress_callback, hook_extra_headers=hook_result.get("extra_headers"), user_api_key_auth=user_api_key_auth, + allow_input_required=wire_compat is WireCompat.MODERN, ) return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 1247ff1ac28..5b23695d06d 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -52,6 +52,7 @@ from litellm.llms.custom_httpx.http_handler import ( header_value, httpxSpecialProvider, ) +from litellm.proxy._experimental.mcp_server.tool_outcome import JsonResult, TextResult, parse_http_body from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -497,7 +498,7 @@ def create_tool_function( path_params, query_params, body_params = extract_parameters(operation) original_method: Final = method.lower() - async def tool_function(**kwargs: object) -> str: + async def tool_function(**kwargs: object) -> TextResult | JsonResult: """ Dynamically generated tool function. @@ -531,7 +532,7 @@ def create_tool_function( # Sanitize and encode path parameter to prevent traversal attacks safe_value = _sanitize_path_parameter_value(param_value, param_name) except ValueError as exc: - return "Invalid path parameter: " + str(exc) + return TextResult("Invalid path parameter: " + str(exc)) # Replace {param_name} or {{param_name}} in URL url = url.replace("{" + param_name + "}", safe_value) url = url.replace("{{" + param_name + "}}", safe_value) @@ -580,7 +581,7 @@ def create_tool_function( elif original_method == "patch": response = await client.patch(url, params=params, json=json_body, headers=effective_headers) else: - return f"Unsupported HTTP method: {original_method}" + return TextResult(f"Unsupported HTTP method: {original_method}") except MaskedHTTPStatusError as e: _raise_for_upstream_failure(e.response, upstream, relays_upstream_auth) raise @@ -588,7 +589,7 @@ def create_tool_function( _request_upstream_url.reset(url_token) _raise_for_upstream_failure(response, upstream, relays_upstream_auth) - return response.text + return parse_http_body(response.text) return tool_function diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index dcab43bdc76..ebd26e4bf87 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -17,6 +17,7 @@ from mcp.types import ( GetPromptRequest, GetPromptRequestParams, GetPromptResult, + InputRequiredResult, ListPromptsRequest, ListPromptsResult, ListResourcesRequest, @@ -85,6 +86,12 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_extra_headers, _request_resolved_auth_headers, ) +from litellm.proxy._experimental.mcp_server.result_conversion import ( + WireCompat, + complete_call_tool_result, + handler_outcome, + to_call_tool_result, +) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -1804,8 +1811,9 @@ async def execute_mcp_tool( host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, client_ip: str | None = None, + wire_compat: WireCompat = WireCompat.LEGACY, **kwargs: object, # kwargs-ok: preserves the existing REST and decorated logging call contract -) -> CallToolResult: +) -> CallToolResult | InputRequiredResult: context: Final = prepare_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -1813,6 +1821,7 @@ async def execute_mcp_tool( oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + wire_compat=wire_compat, ) operation: Final = AuthorizedToolCall( name=name, @@ -1839,8 +1848,9 @@ async def _execute_mcp_tool( host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, client_ip: str | None = None, + wire_compat: WireCompat = WireCompat.LEGACY, **kwargs: Any, -) -> CallToolResult: +) -> CallToolResult | InputRequiredResult: """ Execute MCP tool. @@ -2088,7 +2098,7 @@ async def _execute_mcp_tool( _extra_token: Final = _request_extra_headers.set(forwarded_headers) _resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) try: - response = await _handle_local_mcp_tool(name, arguments) + response = await _handle_local_mcp_tool(name, arguments, wire_compat) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) @@ -2112,6 +2122,7 @@ async def _execute_mcp_tool( litellm_logging_obj=litellm_logging_obj, guardrail_context=guardrail_context, host_progress_callback=host_progress_callback, + wire_compat=wire_compat, ) # Fall back to local tool registry with original name (legacy support) @@ -2169,10 +2180,13 @@ async def _execute_mcp_tool( if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args - response = await _handle_local_mcp_tool(original_tool_name, arguments) + response = await _handle_local_mcp_tool(original_tool_name, arguments, wire_compat) + converted: Final = to_call_tool_result(response, wire_compat) + if isinstance(converted, InputRequiredResult): + return converted return await _run_post_mcp_call_guardrails( - result=response, + result=converted, litellm_logging_obj=litellm_logging_obj, user_api_key_auth=user_api_key_auth, request_data=kwargs, @@ -2206,6 +2220,13 @@ async def _run_post_mcp_call_guardrails( ) +def suppress_completed_success_logging(logging_obj: LiteLLMLoggingObj) -> None: + """An interim ``InputRequiredResult`` is not a completed call, so the ``@client`` wrapper + on ``call_mcp_tool`` must not run the success handlers for it when the coroutine returns.""" + logging_obj.has_run_logging(event_type="sync_success") + logging_obj.has_run_logging(event_type="async_success") + + async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, result: CallToolResult, @@ -2322,10 +2343,14 @@ async def call_mcp_tool( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, client_ip: str | None = None, + wire_compat: WireCompat = WireCompat.LEGACY, **kwargs: Any, -) -> CallToolResult: +) -> CallToolResult | InputRequiredResult: """ Call a specific tool with the provided arguments (handles prefixed tool names). + + A modern ``InputRequiredResult`` is an interim answer, so it is returned as is and skips the + completed-call logging below. """ start_time: Final = datetime.now() litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) @@ -2376,12 +2401,17 @@ async def call_mcp_tool( oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + wire_compat=wire_compat, **kwargs, ) except Exception as e: await fire_mcp_tool_call_failure_logging(litellm_logging_obj, e, start_time, user_api_key_auth, kwargs) raise + if isinstance(response, InputRequiredResult): + if litellm_logging_obj: + suppress_completed_success_logging(litellm_logging_obj) + return response if litellm_logging_obj: response = await _fire_mcp_tool_call_logging( logging_obj=litellm_logging_obj, @@ -2547,7 +2577,8 @@ async def _handle_managed_mcp_tool( host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, client_ip: str | None = None, -) -> CallToolResult: + wire_compat: WireCompat = WireCompat.LEGACY, +) -> CallToolResult | InputRequiredResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj @@ -2566,12 +2597,15 @@ async def _handle_managed_mcp_tool( host_progress_callback=host_progress_callback, litellm_logging_obj=litellm_logging_obj, guardrail_context=guardrail_context, + wire_compat=wire_compat, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result -async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult: +async def _handle_local_mcp_tool( + name: str, arguments: dict[str, object], wire_compat: WireCompat = WireCompat.LEGACY +) -> CallToolResult: """Execute a local-registry tool and report whether it succeeded. Returns the result rather than bare content because the verdict is part of it: the content @@ -2604,10 +2638,7 @@ async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> Cal content=[TextContent(text=f"Error: {e}", type="text")], # mutable-ok: MCP result content is_error=True, ) - return CallToolResult( - content=[TextContent(text=str(result), type="text")], # mutable-ok: MCP result content - is_error=False, - ) + return complete_call_tool_result(handler_outcome(result), wire_compat) _MCP_CREDENTIAL_REQUEST_FIELDS: Final = frozenset( @@ -2694,7 +2725,7 @@ async def _execute_handle_list_tools( async def _execute_mcp_server_tool_call( context: OperationContext, params: CallToolRequestParams, host_progress_callback: ProgressCallback | None = None -) -> CallToolResult: +) -> CallToolResult | InputRequiredResult: from mcp.types import CallToolResult from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException @@ -2778,6 +2809,7 @@ async def _execute_mcp_server_tool_call( raw_headers=raw_headers, client_ip=_client_ip, host_progress_callback=host_progress_callback, + wire_compat=context.wire_compat, **data, # for logging ) except MCPMissingUserEnvVarsError as e: @@ -3032,6 +3064,7 @@ def prepare_context( raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, mcp_proxy_mode: bool = False, + wire_compat: WireCompat = WireCompat.LEGACY, ) -> OperationContext: return OperationContext( _caller=user_api_key_auth, @@ -3042,6 +3075,7 @@ def prepare_context( raw_headers=raw_headers, client_ip=client_ip, mcp_proxy_mode=mcp_proxy_mode, + wire_compat=wire_compat, ) @@ -3058,6 +3092,7 @@ GatewayOperation: TypeAlias = ( GatewayResult: TypeAlias = ( ListToolsResult | CallToolResult + | InputRequiredResult | ListPromptsResult | GetPromptResult | ListResourcesResult @@ -3071,13 +3106,17 @@ class GatewayOperations: self._host_progress_callback = host_progress_callback @overload - async def execute(self, operation: AuthorizedToolCall, context: OperationContext) -> CallToolResult: ... + async def execute( + self, operation: AuthorizedToolCall, context: OperationContext + ) -> CallToolResult | InputRequiredResult: ... @overload async def execute(self, operation: ListToolsRequest, context: OperationContext) -> ListToolsResult: ... @overload - async def execute(self, operation: CallToolRequest, context: OperationContext) -> CallToolResult: ... + async def execute( + self, operation: CallToolRequest, context: OperationContext + ) -> CallToolResult | InputRequiredResult: ... @overload async def execute(self, operation: ListPromptsRequest, context: OperationContext) -> ListPromptsResult: ... @@ -3113,6 +3152,7 @@ class GatewayOperations: client_ip=_client_ip, host_progress_callback=operation.host_progress_callback, guardrail_context=operation.guardrail_context, + wire_compat=context.wire_compat, **operation.logging_data, ) case ListToolsRequest(params=params): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5922285f643..7f519e2c0d9 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ) from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url +from litellm.proxy._experimental.mcp_server.result_conversion import WireCompat, complete_call_tool_result from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, build_effective_auth_contexts, @@ -1197,7 +1198,7 @@ if MCP_AVAILABLE: # Call execute_mcp_tool directly (permission checks already done) _tool_start_time: Final = datetime.now() - result: Final = await execute_mcp_tool( + executed: Final = await execute_mcp_tool( name=tool_name, arguments=tool_arguments, allowed_mcp_servers=allowed_mcp_servers, @@ -1212,6 +1213,7 @@ if MCP_AVAILABLE: guardrail_context=MCPRequestContext.resolve_guardrail_context(data), requested_server_id=canonical_server_id, ) + result: Final = complete_call_tool_result(executed, WireCompat.LEGACY) except Exception as e: request_data: Final = proxy_base_llm_response_processor.data await _safe_fire_mcp_tool_call_failure_logging( diff --git a/litellm/proxy/_experimental/mcp_server/result_conversion.py b/litellm/proxy/_experimental/mcp_server/result_conversion.py new file mode 100644 index 00000000000..52931fae116 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/result_conversion.py @@ -0,0 +1,120 @@ +"""Compatibility-aware conversion of upstream outcomes into MCP SDK results. + +Every gateway surface that turns a tool outcome (text, JSON, an SDK result, an +interim result, an exception) into the ``CallToolResult`` it sends downstream +goes through ``to_call_tool_result`` so the per-revision wire rules live in one +place. SDK 2.x serializes ``structuredContent`` as object-only on the handshake +revisions (``2024-11-05`` .. ``2025-11-25``) and admits any JSON value, plus +``input_required`` interim results, only on ``2026-07-28``. +""" + +from __future__ import annotations + +import json +from typing import Final, TypeAlias + +from mcp.types import CallToolResult, ContentBlock, InputRequiredResult, TextContent, Tool +from typing_extensions import ReadOnly, TypedDict, assert_never + +from litellm.proxy._experimental.mcp_server.tool_outcome import ( + JsonResult, + TextResult, + WireCompat, + handler_outcome, + parse_http_body, + wire_compat_for, +) + +__all__ = ( + "INPUT_REQUIRED_UNSUPPORTED_MESSAGE", + "JsonResult", + "TextResult", + "ToolOutcome", + "WireCompat", + "complete_call_tool_result", + "error_text_result", + "handler_outcome", + "parse_http_body", + "to_call_tool_result", + "to_gateway_tool", + "wire_compat_for", +) + +ToolOutcome: TypeAlias = TextResult | JsonResult | CallToolResult | InputRequiredResult | Exception + + +class _Downgraded(TypedDict): + structured_content: ReadOnly[None] + content: ReadOnly[list[ContentBlock]] # mutable-ok: SDK list field + + +class _Renamed(TypedDict): + name: ReadOnly[str] + + +INPUT_REQUIRED_UNSUPPORTED_MESSAGE: Final = ( + "Error: upstream tool returned an input_required interim result, which this MCP protocol revision cannot carry" +) + + +def error_text_result(exc: Exception) -> CallToolResult: + return CallToolResult( + content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")], # mutable-ok: SDK list field + is_error=True, + ) + + +def to_call_tool_result(outcome: ToolOutcome, compat: WireCompat) -> CallToolResult | InputRequiredResult: + match outcome: + case TextResult(): + return CallToolResult( + content=[TextContent(type="text", text=outcome.text)], # mutable-ok: SDK list field + is_error=False, + ) + case JsonResult(): + keep_structured: Final = compat is WireCompat.MODERN or isinstance(outcome.value, dict) + return CallToolResult( + content=[TextContent(type="text", text=outcome.original_text)], # mutable-ok: SDK list field + is_error=False, + structured_content=outcome.value if keep_structured else None, + ) + case CallToolResult(): + return _downgrade_structured_content(outcome) if compat is WireCompat.LEGACY else outcome + case InputRequiredResult(): + if compat is WireCompat.MODERN: + return outcome + return CallToolResult( + content=[TextContent(type="text", text=INPUT_REQUIRED_UNSUPPORTED_MESSAGE)], # mutable-ok: SDK + is_error=True, + ) + case Exception(): + return error_text_result(outcome) + return assert_never(outcome) + + +def complete_call_tool_result(outcome: ToolOutcome, compat: WireCompat) -> CallToolResult: + """``to_call_tool_result`` for callers that can never carry an interim result.""" + converted: Final = to_call_tool_result(outcome, compat) + if isinstance(converted, InputRequiredResult): + return CallToolResult( + content=[TextContent(type="text", text=INPUT_REQUIRED_UNSUPPORTED_MESSAGE)], # mutable-ok: SDK + is_error=True, + ) + return converted + + +def _downgrade_structured_content(result: CallToolResult) -> CallToolResult: + structured: Final = result.structured_content + if structured is None or isinstance(structured, dict): + return result + fallback: Final = TextContent(type="text", text=json.dumps(structured)) + update: Final[_Downgraded] = { + "structured_content": None, + "content": [*result.content, fallback], # mutable-ok: SDK list field + } + return result.model_copy(update=update) + + +def to_gateway_tool(tool: Tool, name: str) -> Tool: + update: Final[_Renamed] = {"name": name} + return tool.model_copy(deep=True, update=update) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6261f36983d..1bd31d971b0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -145,6 +145,7 @@ try: from mcp import ReadResourceResult, Resource from mcp.server import Server + from mcp.server.runner import serve_loop from mcp.server.session import ServerSession as _McpServerSession from mcp.types import ( BlobResourceContents, @@ -504,6 +505,7 @@ if MCP_AVAILABLE: _invalidate_byok_cred_cache, _mcp_session_id_from_headers, ) + from litellm.proxy._experimental.mcp_server.result_conversion import wire_compat_for try: from mcp.server.streamable_http_manager import StreamableHTTPSessionManager @@ -516,6 +518,7 @@ if MCP_AVAILABLE: GetPromptRequestParams, Implementation, InitializeRequest, + InputRequiredResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult, @@ -818,7 +821,15 @@ if MCP_AVAILABLE: client_ip, ) = await get_or_extract_auth_context() yield operations.prepare_context( - auth, token, servers, server_headers, oauth_headers, headers, client_ip, _mcp_proxy_mode.get() + auth, + token, + servers, + server_headers, + oauth_headers, + headers, + client_ip, + _mcp_proxy_mode.get(), + wire_compat_for(ctx.protocol_version), ) async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: @@ -875,7 +886,9 @@ if MCP_AVAILABLE: _dispatch_virtual_mcp_tool, ) - async def mcp_server_tool_call(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: + async def mcp_server_tool_call( + ctx: ServerRequestContext, params: CallToolRequestParams + ) -> CallToolResult | InputRequiredResult: async with _legacy_operation_context(ctx, trace=True) as context: return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( CallToolRequest(params=params), context @@ -2384,8 +2397,17 @@ if MCP_AVAILABLE: scoped_server_endpoint=scoped_server_endpoint, is_initialize=scope.get("method") == "GET", ): - async with sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream): - await server.run(read_stream, write_stream, server.create_initialization_options()) + async with ( + sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream), + server.lifespan(server) as lifespan_state, + ): + await serve_loop( + server, + read_stream, + write_stream, + lifespan_state=lifespan_state, + init_options=server.create_initialization_options(), + ) except MCPUpstreamAuthError as e: # Upstream delegated auth returned 401; surface it to the client so # standards-compliant MCP clients trigger the upstream OAuth flow. diff --git a/litellm/proxy/_experimental/mcp_server/tool_outcome.py b/litellm/proxy/_experimental/mcp_server/tool_outcome.py new file mode 100644 index 00000000000..ac712241b5b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/tool_outcome.py @@ -0,0 +1,56 @@ +"""SDK-free half of the result conversion boundary. + +``openapi_to_mcp_generator`` and ``contracts`` must import without the ``mcp`` +package installed, so the compatibility enum and the tagged handler outcomes +live here; ``result_conversion`` turns them into SDK results. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Final + +from mcp_types.version import MODERN_PROTOCOL_VERSIONS +from pydantic import JsonValue, TypeAdapter, ValidationError + +_JSON_VALUE: Final = TypeAdapter(JsonValue) + + +class WireCompat(str, Enum): + LEGACY = "legacy" + MODERN = "modern" + + +def wire_compat_for(protocol_version: str) -> WireCompat: + return WireCompat.MODERN if protocol_version in MODERN_PROTOCOL_VERSIONS else WireCompat.LEGACY + + +@dataclass(frozen=True, slots=True) +class TextResult: + text: str + + +@dataclass(frozen=True, slots=True) +class JsonResult: + value: JsonValue + original_text: str + + +def parse_http_body(body: str) -> TextResult | JsonResult: + if not body.strip(): + return TextResult(body) + try: + value: Final = _JSON_VALUE.validate_json(body) + except ValidationError: + return TextResult(body) + if value is None: + return TextResult(body) + return JsonResult(value=value, original_text=body) + + +def handler_outcome(value: object) -> TextResult | JsonResult: + """Normalize what a registered tool handler returned; OpenAPI handlers already return a tagged outcome.""" + if isinstance(value, (TextResult, JsonResult)): + return value + return TextResult(str(value)) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 71e46f8df25..9d117a1a1fa 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K +from litellm.proxy._experimental.mcp_server.result_conversion import WireCompat, complete_call_tool_result from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K from litellm.proxy.common_utils.semantic_text_index import ( Embedder, @@ -634,7 +635,7 @@ async def handle_mcp_tool_call( raise HTTPException(status_code=403, detail="User not allowed to call this tool.") - return await execute_mcp_tool( + result: Final = await execute_mcp_tool( name=tool_name, arguments=arguments, allowed_mcp_servers=allowed_mcp_servers, @@ -649,3 +650,4 @@ async def handle_mcp_tool_call( requested_server_id=requested_server_id, guardrail_context=guardrail_context, ) + return complete_call_tool_result(result, WireCompat.LEGACY) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 9659eb1cbc2..d39ec063538 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -17,6 +17,7 @@ from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from mcp.types import CallToolResult from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager from litellm.proxy._types import UserAPIKeyAuth @@ -409,7 +410,7 @@ class TestCallToolFlowsHookHeaders: manager, "_call_openapi_tool_handler", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=CallToolResult(content=[], isError=False), ): import litellm.proxy._experimental.mcp_server.mcp_server_manager as mgr_mod @@ -456,7 +457,7 @@ class TestCallToolFlowsHookHeaders: manager, "_call_openapi_tool_handler", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=CallToolResult(content=[], isError=False), ): proxy_logging = MagicMock(spec=ProxyLogging) @@ -1076,9 +1077,9 @@ class TestOpenApiByokCallTool: user_auth = UserAPIKeyAuth(user_id="default_user_id", api_key="sk-dashboard") captured_auth: dict[str, Optional[str]] = {} - async def fake_openapi_handler(_server, _name, _arguments): + async def fake_openapi_handler(_server, _name, _arguments, _wire_compat): captured_auth["value"] = _request_auth_header.get() - return MagicMock() + return CallToolResult(content=[], isError=False) with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): with patch( @@ -1316,9 +1317,9 @@ class TestOpenApiResolvedUpstreamAuth: user_auth = UserAPIKeyAuth(user_id="alice", api_key="sk-user") captured: Dict[str, Any] = {} - async def fake_openapi_handler(_server, _name, _arguments): + async def fake_openapi_handler(_server, _name, _arguments, _wire_compat): captured["resolved"] = _request_resolved_auth_headers.get() - return MagicMock() + return CallToolResult(content=[], isError=False) with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): with patch.object( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py index e11897b65c2..0dc7ac5ecd9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py @@ -2,6 +2,7 @@ import asyncio from typing import Dict, Optional import pytest +from mcp.types import CallToolResult, TextContent from unittest.mock import patch from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager @@ -46,7 +47,7 @@ def _make_server(server_id: str, max_concurrent_requests: Optional[int]) -> MCPS def _patch_client_with_tracker(manager: MCPServerManager, tracker: _ConcurrencyTracker): async def fake_create_mcp_client(server, **kwargs): class _ProbeClient: - async def call_tool(self, params, host_progress_callback=None): + async def call_tool(self, params, host_progress_callback=None, allow_input_required=False): tracker.enter(server.server_id) try: await asyncio.sleep(HOLD_SECONDS) @@ -145,11 +146,11 @@ async def test_openapi_backed_server_also_respects_the_cap(): server = _make_server("srv-openapi", max_concurrent_requests=2) server.spec_path = "/fake/openapi.json" - async def fake_openapi_handler(mcp_server, name, arguments): + async def fake_openapi_handler(mcp_server, name, arguments, wire_compat): tracker.enter(mcp_server.server_id) try: await asyncio.sleep(HOLD_SECONDS) - return "ok" + return CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) finally: tracker.exit(mcp_server.server_id) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 6887adf8283..97b242831a2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -12,8 +12,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from mcp import ReadResourceResult, Resource +from mcp.server.models import InitializationOptions from mcp.types import ( INVALID_REQUEST, + METHOD_NOT_FOUND, BlobResourceContents, CallToolResult, Prompt, @@ -21,7 +23,7 @@ from mcp.types import ( TextContent, TextResourceContents, ) -from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION, MODERN_PROTOCOL_VERSIONS from pydantic import TypeAdapter from starlette.types import Message, Receive, Scope, Send @@ -6565,8 +6567,11 @@ class TestGatewayCreateInitializationOptions: async def connect_sse(scope, receive, send): yield (None, None) - async def record_request(read_stream, write_stream, options): - captured["server_name"] = server.create_initialization_options().server_name + async def record_request( + serving_server: object, read_stream: object, write_stream: object, + *, lifespan_state: object, init_options: InitializationOptions, + ) -> None: + captured["server_name"] = init_options.server_name scope = { "type": "http", @@ -6612,8 +6617,8 @@ class TestGatewayCreateInitializationOptions: True, ), patch.object( - mcp_server.server, - "run", + mcp_server, + "serve_loop", side_effect=record_request, ), ): @@ -8021,7 +8026,7 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): ), patch( "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", - new=AsyncMock(return_value=[]), + new=AsyncMock(return_value=CallToolResult(content=[], is_error=False)), ), patch( "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", @@ -9255,6 +9260,152 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error(): proxy_logging_mock.post_call_failure_hook.assert_not_awaited() +def _interim_input_required_result(): + from mcp.types import InputRequiredResult + + return InputRequiredResult.model_validate( + { + "resultType": "input_required", + "inputRequests": { + "req-1": { + "method": "elicitation/create", + "params": {"message": "Pick one", "requestedSchema": {"type": "object", "properties": {}}}, + } + }, + "requestState": "state-1", + } + ) + + +@contextlib.contextmanager +def _managed_tool_returning(server, upstream_result, proxy_logging_mock): + from litellm.proxy._experimental.mcp_server.server import global_mcp_server_manager + + with ( + patch.object( + global_mcp_server_manager, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[server.server_id], + ), + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server), + patch.object(global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server), + patch.object(global_mcp_server_manager, "server_owning_tool_name_prefix", return_value=server), + patch( + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers_from_mcp_server_names", + new_callable=AsyncMock, + return_value=[server], + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._list_tools_before_first_call", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", + return_value=(None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._handle_managed_mcp_tool", + new_callable=AsyncMock, + return_value=upstream_result, + ) as managed_call, + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock), + ): + yield managed_call + + +@pytest.mark.asyncio +async def test_call_mcp_tool_legacy_interim_result_is_rejected_into_failure_accounting(): + """An upstream input_required interim on a legacy connection cannot be carried on the wire, so it + must come back as isError and go through the same failure accounting as any other errored call.""" + from mcp.types import CallToolResult + + from litellm.proxy._experimental.mcp_server.result_conversion import ( + INPUT_REQUIRED_UNSUPPORTED_MESSAGE, + WireCompat, + ) + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + from litellm.proxy._types import MCPTransport, UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="server-interim", + name="test_server", + alias="test_server", + server_name="test_server", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "test_server"}, + ) + proxy_logging_mock = _mock_mcp_proxy_logging() + logging_obj = _mock_mcp_logging_obj() + + with _managed_tool_returning(server, _interim_input_required_result(), proxy_logging_mock) as managed_call: + result = await call_mcp_tool( + name="test_server-any_tool", + arguments={"x": 1}, + user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + litellm_logging_obj=logging_obj, + wire_compat=WireCompat.LEGACY, + ) + + assert managed_call.await_args.kwargs["wire_compat"] is WireCompat.LEGACY + assert isinstance(result, CallToolResult) and result.is_error is True + assert result.content[0].text == INPUT_REQUIRED_UNSUPPORTED_MESSAGE + logging_obj.async_success_handler.assert_not_awaited() + logging_obj.async_failure_handler.assert_awaited_once() + assert str(logging_obj.async_failure_handler.await_args.args[0]) == INPUT_REQUIRED_UNSUPPORTED_MESSAGE + proxy_logging_mock.post_call_failure_hook.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_mcp_tool_modern_interim_result_passes_through_without_completed_accounting(): + """On a modern connection the interim result is returned with its fields intact and is neither + logged as a completed success nor run through the post-call guardrail and success hooks.""" + from mcp.types import InputRequiredResult + + from litellm.proxy._experimental.mcp_server.result_conversion import WireCompat + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + from litellm.proxy._types import MCPTransport, UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="server-interim", + name="test_server", + alias="test_server", + server_name="test_server", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "test_server"}, + ) + proxy_logging_mock = _mock_mcp_proxy_logging() + logging_obj = _mock_mcp_logging_obj() + interim = _interim_input_required_result() + + with _managed_tool_returning(server, interim, proxy_logging_mock) as managed_call: + result = await call_mcp_tool( + name="test_server-any_tool", + arguments={"x": 1}, + user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + litellm_logging_obj=logging_obj, + wire_compat=WireCompat.MODERN, + ) + + assert managed_call.await_args.kwargs["wire_compat"] is WireCompat.MODERN + assert isinstance(result, InputRequiredResult) + assert result.request_state == "state-1" + assert result.input_requests is not None and set(result.input_requests) == {"req-1"} + logging_obj.async_success_handler.assert_not_awaited() + logging_obj.async_failure_handler.assert_not_awaited() + logging_obj.async_post_mcp_tool_call_hook.assert_not_awaited() + proxy_logging_mock.post_mcp_call_hook.assert_not_awaited() + proxy_logging_mock.post_call_failure_hook.assert_not_awaited() + assert sorted(c.kwargs["event_type"] for c in logging_obj.has_run_logging.call_args_list) == [ + "async_success", + "sync_success", + ], "the @client wrapper would otherwise log the interim result as a completed success on return" + + @pytest.mark.asyncio async def test_aggregate_listing_reports_per_server_outcomes(): """A failed server must contribute a classified outcome, not just silently shrink the list: @@ -10371,7 +10522,10 @@ async def test_tool_listing_preserves_permission_denial_when_failure_logging_fai @pytest.mark.asyncio @pytest.mark.parametrize("prefix,suffix", (("", ""), ("/gateway", "/"))) -async def test_legacy_sse_mount_emits_message_endpoint(prefix: str, suffix: str) -> None: +@pytest.mark.parametrize("opening_protocol", (None, *MODERN_PROTOCOL_VERSIONS)) +async def test_legacy_sse_mount_emits_message_endpoint( + prefix: str, suffix: str, opening_protocol: str | None, +) -> None: from starlette.applications import Starlette from starlette.routing import Mount from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -10434,6 +10588,23 @@ async def test_legacy_sse_mount_emits_message_endpoint(prefix: str, suffix: str) await asyncio.wait_for(app(post_scope, requests.get, messages.put), 2) return (await messages.get())["status"] + if opening_protocol is not None: + discover: Final = json.dumps({ + "jsonrpc": "2.0", + "id": 0, + "method": "server/discover", + "params": {"_meta": { + "io.modelcontextprotocol/protocolVersion": opening_protocol, + "io.modelcontextprotocol/clientInfo": {"name": "modern-client", "version": "1"}, + "io.modelcontextprotocol/clientCapabilities": {}, + }}, + }).encode() + assert await post(discover) == 202 + discovered_frame: Final = (await asyncio.wait_for(outgoing.get(), 2))["body"].decode() + discovered: Final = json.loads(discovered_frame.split("data: ", 1)[1].splitlines()[0]) + assert discovered["id"] == 0 + assert discovered["error"]["code"] == METHOD_NOT_FOUND + initialization: Final = json.dumps( { "jsonrpc": "2.0", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index f3d37a858ca..64d94065674 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -39,6 +39,7 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl, TypeAdapter from litellm.constants import MCP_METADATA_TIMEOUT +from litellm.proxy._experimental.mcp_server.tool_outcome import TextResult from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _deserialize_json_dict, @@ -6730,7 +6731,7 @@ class TestMCPServerManager: # Create mock client that tracks call_tool usage mock_client = AsyncMock() - async def mock_call_tool(params, host_progress_callback=None): + async def mock_call_tool(params, host_progress_callback=None, allow_input_required=False): # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] @@ -10061,7 +10062,7 @@ class _RetryFakeClient: self._MCPClient = MCPClient self.attempts = 0 - async def call_tool(self, params, host_progress_callback=None, raise_on_error=False): + async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False): self.attempts += 1 if self._raises is not None: if raise_on_error: @@ -10273,7 +10274,7 @@ class TestOBOConcurrencyLimit: inflight = {"current": 0, "peak": 0} class _ConcurrencyRecordingClient: - async def call_tool(self, params, host_progress_callback=None, raise_on_error=False): + async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False): inflight["current"] += 1 inflight["peak"] = max(inflight["peak"], inflight["current"]) try: @@ -12250,6 +12251,38 @@ class TestOpenApiHandlerRelaysUpstreamAuth: assert result.is_error is True assert "upstream returned HTTP 503" in result.content[0].text + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("body", "compat", "expected_structured"), + [ + ('{"total": 1.10, "items": [ ]}', "legacy", {"total": 1.1, "items": []}), + ('{"total": 1.10, "items": [ ]}', "modern", {"total": 1.1, "items": []}), + ("[1, 2]", "legacy", None), + ("[1, 2]", "modern", [1, 2]), + ("plain text", "legacy", None), + ("plain text", "modern", None), + ], + ) + async def test_json_bodies_keep_verbatim_text_and_gain_structured_content(self, body, compat, expected_structured): + """The OpenAPI arm used to stringify the response; now the text block is the upstream body + byte for byte, exactly once, and JSON bodies carry structuredContent when the caller's revision admits it.""" + from litellm.proxy._experimental.mcp_server.tool_outcome import WireCompat, parse_http_body + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + manager = MCPServerManager() + + async def handler(**_kwargs): + return parse_http_body(body) + + tool = MagicMock() + tool.handler = handler + with patch.object(global_mcp_tool_registry, "get_tool", return_value=tool): + result = await manager._call_openapi_tool_handler(self._server(), "list_reports", {}, WireCompat(compat)) + + assert result.is_error is False + assert [block.text for block in result.content] == [body] + assert result.structured_content == expected_structured + class TestConfigServerIdPinning: """config.yaml servers may pin ``server_id`` so permission grants survive connection edits.""" @@ -14254,7 +14287,7 @@ class TestProtectedCredentialPreparation: caller_token: Final = _request_auth_header.set(caller) extra_token: Final = _request_extra_headers.set(forwarded) try: - assert await tool() == "authenticated" + assert await tool() == TextResult("authenticated") sent: Final = destination.calls.last.request.headers assert sent.get("x-api-key") == static.get("X-API-Key", (forwarded or {}).get("X-API-Key")) if caller: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index bd351f9106e..6b0211c3866 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -23,6 +23,7 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, _request_resolved_auth_headers, + _request_upstream_url, _resolve_param_list, _resolve_ref, build_input_schema, @@ -31,6 +32,7 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( get_base_url, resolve_operation_params, ) +from litellm.proxy._experimental.mcp_server.tool_outcome import JsonResult, TextResult from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, @@ -40,6 +42,43 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +async def test_unsupported_http_method_returns_text_without_sending_request( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function("/echo", "HEAD", {}, "https://upstream.example") + token: Final = _request_upstream_url.set("https://outer.example/request") + try: + assert await tool() == TextResult("Unsupported HTTP method: head") + assert len(respx_mock.calls) == 0 + assert _request_upstream_url.get() == "https://outer.example/request" + finally: + _request_upstream_url.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body,expected", [ + (' { "ok": true }\n', JsonResult({"ok": True}, ' { "ok": true }\n')), + (' [1, 2]\n', JsonResult([1, 2], ' [1, 2]\n')), + ('false', JsonResult(False, 'false')), + ('0', JsonResult(0, '0')), + ('""', JsonResult("", '""')), + ('null', TextResult('null')), + ('{"unfinished":', TextResult('{"unfinished":')), + ('', TextResult('')), +]) +async def test_http_response_preserves_body_and_classifies_json( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + body: str, expected: TextResult | JsonResult, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text=body) + assert await tool() == expected + assert destination.call_count == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("auth_type,value,accepted", [ (MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False), @@ -63,7 +102,7 @@ async def test_authorization_validates_credentials_before_http( caller_token: Final = _request_auth_header.set(value) try: if accepted: - assert await tool() == "authenticated" + assert await tool() == TextResult("authenticated") assert destination.call_count == 1 assert destination.calls.last.request.headers["authorization"] == value else: @@ -103,7 +142,7 @@ async def test_static_auth_validates_headers_after_existing_precedence( assert exc.value.status_code == 500 assert destination.call_count == 0 else: - assert await tool() == "authenticated" + assert await tool() == TextResult("authenticated") assert destination.call_count == 1 assert destination.calls.last.request.headers["authorization"] == expected finally: @@ -124,7 +163,7 @@ async def test_static_auth_uses_configured_custom_header( ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") if credential: - assert await tool() == "authenticated" + assert await tool() == TextResult("authenticated") assert destination.call_count == 1 assert destination.calls.last.request.headers["x-custom"] == credential else: @@ -144,7 +183,7 @@ async def test_static_auth_accepts_api_key_carried_by_static_header( ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") if credential: - assert await tool() == "authenticated" + assert await tool() == TextResult("authenticated") assert destination.calls.last.request.headers["apikey"] == credential assert "x-api-key" not in destination.calls.last.request.headers else: @@ -167,7 +206,7 @@ async def test_static_validation_preserves_no_auth_and_resolved_oauth( destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="echo") token: Final = _request_resolved_auth_headers.set(resolved) try: - assert await tool() == "echo" + assert await tool() == TextResult("echo") assert destination.call_count == 1 assert destination.calls.last.request.headers.get("authorization") == (resolved or {}).get("Authorization") finally: @@ -220,7 +259,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"repository-id": "test-repo"}) - assert result == '{"id": "123"}' + assert result == JsonResult({"id": "123"}, '{"id": "123"}') # Verify URL was constructed correctly call_args = async_client.get.call_args @@ -256,7 +295,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"2fa-code": "123456"}) - assert result == "verified" + assert result == TextResult("verified") # Verify query parameter was included call_args = async_client.post.call_args @@ -290,7 +329,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"user.name": "john.doe"}) - assert result == "found" + assert result == TextResult("found") call_args = async_client.get.call_args assert call_args[1]["params"]["user.name"] == "john.doe" @@ -323,7 +362,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"$filter": "name eq 'test'"}) - assert result == "[]" + assert result == JsonResult([], "[]") call_args = async_client.get.call_args assert call_args[1]["params"]["$filter"] == "name eq 'test'" @@ -356,7 +395,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"class": "premium"}) - assert result == "items" + assert result == TextResult("items") call_args = async_client.get.call_args assert call_args[1]["params"]["class"] == "premium" @@ -407,7 +446,7 @@ class TestCreateToolFunction: "$filter": "active", } ) - assert result == "success" + assert result == TextResult("success") @pytest.mark.asyncio async def test_request_body_parameter(self): @@ -440,7 +479,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"body": {"name": "test"}}) - assert result == "created" + assert result == TextResult("created") call_args = async_client.post.call_args assert call_args[1]["json"] == {"name": "test"} @@ -464,7 +503,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func() - assert result == "ok" + assert result == TextResult("ok") @pytest.mark.asyncio async def test_all_http_methods(self): @@ -497,7 +536,7 @@ class TestCreateToolFunction: mock_client.return_value = async_client result = await func(**{"repository-id": "test"}) - assert result == "success" + assert result == TextResult("success") def test_no_exec_usage(self): """Verify that create_tool_function does not use exec().""" @@ -614,7 +653,7 @@ class TestPathSecurity: response = await tool_function(**{"filename": "../admin"}) - assert "Invalid path parameter" in response + assert isinstance(response, TextResult) and "Invalid path parameter" in response.text @pytest.mark.asyncio async def test_should_encode_and_request_safe_path_parameters(self): @@ -643,7 +682,7 @@ class TestPathSecurity: response = await tool_function(**{"filename": "report 2024.json"}) - assert response == "dummy-response" + assert response == TextResult("dummy-response") # Verify URL was properly encoded call_args = async_client.get.call_args @@ -1181,7 +1220,7 @@ class TestRequestExtraHeaders: finally: _request_extra_headers.reset(token) - assert result == "ok" + assert result == TextResult("ok") call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert headers_sent.get("X-TOKEN") == "secret-value" @@ -1204,7 +1243,7 @@ class TestRequestExtraHeaders: result = await func() - assert result == "ok" + assert result == TextResult("ok") call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert headers_sent == {"X-Static": "static-value"} @@ -1232,7 +1271,7 @@ class TestRequestExtraHeaders: finally: _request_extra_headers.reset(token) - assert result == "created" + assert result == TextResult("created") call_args = async_client.post.call_args headers_sent = call_args[1]["headers"] assert headers_sent.get("X-Static") == "static-value" @@ -1260,7 +1299,7 @@ class TestRequestExtraHeaders: finally: _request_extra_headers.reset(token) - assert result == "ok" + assert result == TextResult("ok") call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert headers_sent.get("X-Tenant") == "operator-tenant" @@ -1288,7 +1327,7 @@ class TestRequestExtraHeaders: finally: _request_extra_headers.reset(token) - assert result == "ok" + assert result == TextResult("ok") call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert headers_sent.get("X-Tenant") == "operator-tenant" @@ -1320,7 +1359,7 @@ class TestRequestExtraHeaders: _request_auth_header.reset(auth_token) _request_extra_headers.reset(extra_token) - assert result == "secure-data" + assert result == TextResult("secure-data") call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert headers_sent.get("Authorization") == "Bearer byok-credential" @@ -1380,7 +1419,7 @@ class TestRequestExtraHeaders: _request_extra_headers.reset(extra_token) _request_resolved_auth_headers.reset(resolved_token) - assert result == "secure-data" + assert result == TextResult("secure-data") headers_sent = async_client.get.call_args[1]["headers"] authorization_values = [v for k, v in headers_sent.items() if k.lower() == "authorization"] assert authorization_values == ["Bearer resolved-oauth"] @@ -1436,7 +1475,7 @@ class TestUpstreamStatusIsClassified: async def test_success_still_returns_the_body(self): tool, client = self._tool(200, text='{"reports": []}') with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): - assert await tool() == '{"reports": []}' + assert await tool() == JsonResult({"reports": []}, '{"reports": []}') @pytest.mark.asyncio async def test_401_raises_the_reauth_signal_carrying_the_challenge(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index bb70f38285c..15d3b67e641 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -9,6 +9,7 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from mcp.types import CallToolResult from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, @@ -46,7 +47,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): fake_tool.name = "list_pets" pre_call = AsyncMock(return_value={}) - handle_local = AsyncMock(return_value=[]) + handle_local = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) with ( patch.object( @@ -131,7 +132,7 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): pre_call = AsyncMock( side_effect=HTTPException(status_code=403, detail="not allowed") ) - handle_local = AsyncMock(return_value=[]) + handle_local = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) with ( patch.object( @@ -191,7 +192,7 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): fake_tool.name = "list_pets" pre_call = AsyncMock(return_value={}) - handle_local = AsyncMock(return_value=[]) + handle_local = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) resolve_auth = MagicMock() # `_get_mcp_server_from_tool_name` returns None — no server context. @@ -275,9 +276,9 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): fake_tool.name = "get_values" captured: dict = {} - async def handle_local(_name, _arguments): + async def handle_local(_name, _arguments, _wire_compat): captured["resolved"] = _request_resolved_auth_headers.get() - return [] + return CallToolResult(content=[], is_error=False) with ( patch.object( @@ -603,13 +604,13 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc captured["resolver_credential"] = kwargs["mcp_auth_header"] return None, kwargs["forwarded_headers"] - async def capture_local(_name, _arguments): + async def capture_local(_name, _arguments, _wire_compat): captured["injected"] = _request_auth_header.get() - return [] + return CallToolResult(content=[], is_error=False) - async def capture_openapi_handler(_server, _name, _arguments): + async def capture_openapi_handler(_server, _name, _arguments, _wire_compat): captured["injected"] = _request_auth_header.get() - return [] + return CallToolResult(content=[], is_error=False) manager = mcp_operations.global_mcp_server_manager with ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py index 81877c38389..81f81045740 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -35,6 +35,7 @@ async def test_oauth_prefetch_failure_does_not_log_caller_or_exception_text(capl @pytest.mark.asyncio async def test_dispatch_uses_explicit_context_when_ambient_caller_differs(): from mcp.server.auth.middleware.auth_context import auth_context_var + from litellm.proxy._experimental.mcp_server.server import set_auth_context context = prepare_context( @@ -64,12 +65,13 @@ async def test_dispatch_uses_explicit_context_when_ambient_caller_differs(): @pytest.mark.asyncio async def test_legacy_adapter_cleans_context_after_cancelled_operation(): from types import SimpleNamespace + from litellm.proxy._experimental.mcp_server import server from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var previous_session = server.active_mcp_session_var.get() previous_request = active_mcp_request_ctx_var.get() - request = SimpleNamespace(session=object()) + request = SimpleNamespace(session=object(), protocol_version="2025-06-18") auth = (None, None, None, None, None, None, None) async def cancelled_operation(): @@ -90,6 +92,7 @@ async def test_legacy_adapter_cleans_context_after_cancelled_operation(): @pytest.mark.asyncio async def test_legacy_adapter_cleans_context_when_trace_setup_fails(): from types import SimpleNamespace + from litellm.proxy._experimental.mcp_server import server from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var @@ -111,6 +114,7 @@ async def test_legacy_adapter_cleans_context_when_trace_setup_fails(): @pytest.mark.asyncio async def test_prompt_sampling_receives_explicit_operation_caller_headers_and_ip(): from unittest.mock import MagicMock + from litellm.proxy._experimental.mcp_server import operations from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -201,9 +205,11 @@ def _catalog_case(method): @pytest.mark.parametrize("state", ["success", "denied", "upstream_failure", "scope_failure"]) async def test_native_catalog_operations_preserve_context_results_and_failure_policy(method, state): from types import SimpleNamespace + from fastapi import HTTPException from mcp.server.context import ServerRequestContext from mcp.types import PaginatedRequestParams + from litellm.proxy._experimental.mcp_server import operations, server from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -260,6 +266,7 @@ async def test_native_catalog_operations_preserve_context_results_and_failure_po async def test_explicit_proxy_context_rejects_catalog_operations_before_upstream_access(method): from mcp.shared.exceptions import MCPError from mcp.types import METHOD_NOT_FOUND + from litellm.proxy._experimental.mcp_server import operations operation, _, manager_method, _, _ = _catalog_case(method) @@ -276,6 +283,7 @@ async def test_explicit_proxy_context_rejects_catalog_operations_before_upstream @pytest.mark.parametrize("failure", ["missing_env", "pii", "guardrail", "unexpected"]) async def test_tool_operation_preserves_failure_messages_and_request_trace(failure): from mcp.types import CallToolRequest, CallToolRequestParams + from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy._experimental.mcp_server import operations from litellm.proxy._experimental.mcp_server.utils import MCPMissingUserEnvVarsError @@ -333,6 +341,7 @@ async def test_catalog_operation_preserves_empty_result_for_malformed_upstream_i @pytest.mark.parametrize("catalog_unavailable", [False, True]) async def test_tool_listing_returns_empty_result_without_dispatch_for_unavailable_catalog(catalog_unavailable): from mcp.types import ListToolsRequest + from litellm.proxy._experimental.mcp_server import operations allowed = AsyncMock( @@ -352,6 +361,7 @@ async def test_tool_listing_returns_empty_result_without_dispatch_for_unavailabl @pytest.mark.asyncio async def test_explicit_proxy_context_lists_builtin_tools_and_blocks_direct_tool_dispatch(): from mcp.types import CallToolRequest, CallToolRequestParams, ListToolsRequest + from litellm.proxy._experimental.mcp_server import operations context = prepare_context(mcp_proxy_mode=True) @@ -486,8 +496,10 @@ class TestChallengeMissingTokenExchangeSubject: @pytest.mark.asyncio async def test_execute_mcp_tool_challenges_missing_subject_before_cold_listing(): """On a cold catalog the challenge fires before any listing or tool resolution is attempted.""" - from fastapi import HTTPException from datetime import datetime, timezone + + from fastapi import HTTPException + from litellm.proxy._experimental.mcp_server import operations server = _server("te-exec", MCPAuth.oauth2_token_exchange) @@ -509,3 +521,24 @@ async def test_execute_mcp_tool_challenges_missing_subject_before_cold_listing() ) assert exc_info.value.status_code == 401 listing.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compat", ["legacy", "modern"]) +async def test_local_tool_json_array_is_converted_once_for_the_caller_revision(compat: str) -> None: + """The local-registry arm used to convert at MODERN and let the legacy downgrade append a second + text block; converting at the caller's revision keeps the upstream body exactly once.""" + from unittest.mock import MagicMock + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.tool_outcome import WireCompat, parse_http_body + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + body = '["a","b"]' + tool = MagicMock() + tool.handler = AsyncMock(return_value=parse_http_body(body)) + with patch.object(global_mcp_tool_registry, "get_tool", return_value=tool): + result = await operations._handle_local_mcp_tool("reports-list_tags", {}, WireCompat(compat)) + + assert [block.text for block in result.content] == [body] + assert result.structured_content == (["a", "b"] if compat == "modern" else None) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index e20d74ab60d..4da120cb26f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,3 @@ -from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import inspect import json @@ -7,12 +6,15 @@ from datetime import datetime from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock +from litellm.proxy._experimental.mcp_server import operations as mcp_operations + if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 from exceptiongroup import BaseExceptionGroup import httpx import pytest from fastapi import HTTPException +from mcp.types import CallToolResult, TextContent from starlette.requests import Request from litellm.constants import MCP_TOOL_LISTING_TIMEOUT @@ -29,6 +31,8 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +_OK_TOOL_RESULT: Final = CallToolResult(content=[TextContent(type="text", text='{"result": "ok"}')], is_error=False) + def _rendered_log_message(call): message = str(call.args[0]) @@ -1472,10 +1476,10 @@ class TestListToolsRestAPI: monkeypatch, ): """The REST tools/list path should include tools beyond the upstream first page.""" - import litellm.experimental_mcp_client.client as mcp_client_module from mcp.types import ListToolsResult, PaginatedRequestParams from mcp.types import Tool as MCPTool + import litellm.experimental_mcp_client.client as mcp_client_module from litellm.proxy._experimental.mcp_server.server import MCPServer from litellm.types.mcp import MCPTransport @@ -2470,7 +2474,7 @@ class TestCallToolRestAPI: async def fake_execute_mcp_tool(**kwargs): captured.update(kwargs) - return {"result": "ok"} + return _OK_TOOL_RESULT monkeypatch.setattr( rest_endpoints, @@ -2530,13 +2534,89 @@ class TestCallToolRestAPI: user_api_key_dict=UserAPIKeyAuth(), ) - assert result == {"result": "ok"} + assert result == _OK_TOOL_RESULT assert captured["name"] == "demo-tool" assert captured["arguments"] == {"foo": "bar"} assert captured["allowed_mcp_servers"] == [stub_server] assert captured["oauth2_headers"] is None fire_logging.assert_awaited_once() + @pytest.mark.parametrize( + ("structured", "expected_structured", "expected_texts"), + [ + ({"a": 1}, {"a": 1}, ['{"a": 1}']), + ([1, 2], None, ['{"a": 1}', "[1, 2]"]), + ], + ) + async def test_rest_keeps_its_serialization_shape_with_legacy_structured_admission( + self, monkeypatch, structured, expected_structured, expected_texts + ): + """REST has no negotiated revision, so it admits object structuredContent only and downgrades + anything else losslessly, while the response keeps the SDK model shape (resultType included) + rather than being run through the MCP legacy wire serializer.""" + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + auth_type = None + + stub_server = StubServer() + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + async def fake_execute_mcp_tool(**kwargs): + return CallToolResult( + content=[TextContent(type="text", text='{"a": 1}')], + structuredContent=structured, + isError=False, + ) + + async def fake_fire_logging(logging_obj, result, start_time, end_time, **kwargs): + return result + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers", fake_get_allowed_mcp_servers + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False) + monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False) + monkeypatch.setattr(rest_endpoints, "_fire_mcp_tool_call_logging", fake_fire_logging, raising=False) + + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}}, + ) + + result = await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth()) + + assert isinstance(result, CallToolResult) + dumped = result.model_dump(by_alias=True, mode="json", exclude_none=True) + assert dumped.get("structuredContent") == expected_structured + assert [block["text"] for block in dumped["content"]] == expected_texts + assert dumped["resultType"] == "complete" + assert dumped["isError"] is False + @pytest.mark.asyncio @pytest.mark.parametrize( ("auth_type", "per_user_oauth", "expected"), @@ -2580,7 +2660,7 @@ class TestCallToolRestAPI: async def fake_execute_mcp_tool(**kwargs): captured.update(kwargs) - return {"result": "ok"} + return _OK_TOOL_RESULT monkeypatch.setattr( rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers", fake_get_allowed_mcp_servers @@ -2607,7 +2687,7 @@ class TestCallToolRestAPI: result = await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth()) - assert result == {"result": "ok"} + assert result == _OK_TOOL_RESULT assert captured["oauth2_headers"] == expected assert captured["raw_headers"]["authorization"] == "Bearer user-subject-token" @@ -2637,7 +2717,7 @@ class TestCallToolRestAPI: return kwargs.get("data", {}) async def fake_execute_mcp_tool(**kwargs): - return {"content": [{"type": "text", "text": "jane@example.com"}]} + return CallToolResult(content=[TextContent(type="text", text="jane@example.com")], is_error=False) monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) monkeypatch.setattr( @@ -2659,7 +2739,7 @@ class TestCallToolRestAPI: ) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False) monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False) - masked_result = {"content": [{"type": "text", "text": ""}]} + masked_result = CallToolResult(content=[TextContent(type="text", text="")], is_error=False) monkeypatch.setattr( rest_endpoints, "_fire_mcp_tool_call_logging", @@ -2714,9 +2794,9 @@ class TestCallToolRestAPI: async def fake_execute_mcp_tool(**kwargs): captured.update(kwargs) - return {"result": "ok"} + return _OK_TOOL_RESULT - fire_logging = AsyncMock(return_value={"result": "ok"}) + fire_logging = AsyncMock(return_value=_OK_TOOL_RESULT) monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) monkeypatch.setattr( rest_endpoints.global_mcp_server_manager, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_result_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_result_conversion.py new file mode 100644 index 00000000000..d9b5063a811 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_result_conversion.py @@ -0,0 +1,243 @@ +import json +from typing import Final + +import pytest +from mcp.types import CallToolResult, ImageContent, InputRequiredResult, TextContent, Tool +from mcp_types.methods import serialize_server_result +from mcp_types.version import KNOWN_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS +from pydantic import JsonValue, ValidationError + +from litellm.proxy._experimental.mcp_server.result_conversion import ( + INPUT_REQUIRED_UNSUPPORTED_MESSAGE, + JsonResult, + TextResult, + WireCompat, + complete_call_tool_result, + error_text_result, + handler_outcome, + parse_http_body, + to_call_tool_result, + to_gateway_tool, + wire_compat_for, +) + +BOTH: Final = (WireCompat.LEGACY, WireCompat.MODERN) + + +def _interim() -> InputRequiredResult: + return InputRequiredResult.model_validate( + { + "resultType": "input_required", + "inputRequests": { + "req-1": { + "method": "elicitation/create", + "params": {"message": "Pick one", "requestedSchema": {"type": "object", "properties": {}}}, + } + }, + "requestState": "abc", + } + ) + + +def _wire(result: CallToolResult | InputRequiredResult, version: str) -> dict[str, object]: + return serialize_server_result( + "tools/call", version, result.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + + +class TestWireCompatFor: + def test_only_modern_revisions_map_to_modern(self): + for version in KNOWN_PROTOCOL_VERSIONS: + expected: Final = WireCompat.MODERN if version in MODERN_PROTOCOL_VERSIONS else WireCompat.LEGACY + assert wire_compat_for(version) is expected, version + assert wire_compat_for("1999-01-01") is WireCompat.LEGACY + + +class TestParseHttpBody: + @pytest.mark.parametrize("body", ["", " ", "{not json", "null"]) + def test_non_structured_bodies_stay_text(self, body: str): + assert parse_http_body(body) == TextResult(body) + + @pytest.mark.parametrize( + "body, value", + [ + ('{"a": 1}', {"a": 1}), + ("[1, 2]", [1, 2]), + ("1.10", 1.1), + ("true", True), + ('"hi"', "hi"), + ], + ) + def test_json_bodies_keep_original_text(self, body: str, value: object): + assert parse_http_body(body) == JsonResult(value=value, original_text=body) + + def test_handler_outcome_stringifies_unknown_values(self): + assert handler_outcome(42) == TextResult("42") + assert handler_outcome(TextResult("x")) == TextResult("x") + + +class TestTextAndJsonArms: + @pytest.mark.parametrize("compat", BOTH) + def test_text_result(self, compat: WireCompat): + result = to_call_tool_result(TextResult("plain"), compat) + assert isinstance(result, CallToolResult) + assert result.is_error is False + assert [c.text for c in result.content if isinstance(c, TextContent)] == ["plain"] + assert result.structured_content is None + + @pytest.mark.parametrize("compat", BOTH) + def test_json_object_is_structured_everywhere_and_text_is_verbatim(self, compat: WireCompat): + body: Final = '{"n": 1.10,\n"k": "v"}' + result = to_call_tool_result(parse_http_body(body), compat) + assert isinstance(result, CallToolResult) + assert result.structured_content == {"n": 1.1, "k": "v"} + assert [c.text for c in result.content if isinstance(c, TextContent)] == [body] + + @pytest.mark.parametrize("body", ["[1, 2]", "3", "true", '"s"']) + def test_non_object_json_is_structured_only_on_modern(self, body: str): + legacy = to_call_tool_result(parse_http_body(body), WireCompat.LEGACY) + modern = to_call_tool_result(parse_http_body(body), WireCompat.MODERN) + assert isinstance(legacy, CallToolResult) and isinstance(modern, CallToolResult) + assert legacy.structured_content is None + assert modern.structured_content == json.loads(body) + for result in (legacy, modern): + assert [c.text for c in result.content if isinstance(c, TextContent)] == [body] + + @pytest.mark.parametrize("compat", BOTH) + def test_json_null_keeps_text_and_claims_no_structured_field(self, compat: WireCompat): + result = to_call_tool_result(parse_http_body("null"), compat) + assert isinstance(result, CallToolResult) + assert [c.text for c in result.content if isinstance(c, TextContent)] == ["null"] + assert "structuredContent" not in _wire(result, "2026-07-28") + + +class TestSdkResultArm: + def _incoming(self, content: list[TextContent]) -> CallToolResult: + return CallToolResult(content=content, structured_content=[1, 2], meta={"trace": "t1"}, is_error=False) + + def test_modern_passes_through_the_same_object(self): + incoming = self._incoming([]) + assert to_call_tool_result(incoming, WireCompat.MODERN) is incoming + + def test_legacy_downgrade_with_empty_content_appends_json_text(self): + incoming = self._incoming([]) + result = to_call_tool_result(incoming, WireCompat.LEGACY) + assert isinstance(result, CallToolResult) + assert result.structured_content is None + assert [c.text for c in result.content if isinstance(c, TextContent)] == ["[1, 2]"] + assert result.meta == {"trace": "t1"} + assert incoming.structured_content == [1, 2] and incoming.content == [] + + def test_legacy_downgrade_keeps_unrelated_content_and_appends_json_text(self): + incoming = self._incoming([TextContent(type="text", text="Done")]) + result = to_call_tool_result(incoming, WireCompat.LEGACY) + assert isinstance(result, CallToolResult) + assert [c.text for c in result.content if isinstance(c, TextContent)] == ["Done", "[1, 2]"] + assert incoming.content == [TextContent(type="text", text="Done")] + assert incoming.structured_content == [1, 2] + + def test_legacy_keeps_object_structured_content(self): + incoming = CallToolResult(content=[], structured_content={"a": 1}, is_error=False) + assert to_call_tool_result(incoming, WireCompat.LEGACY) is incoming + + @pytest.mark.parametrize("value", [False, 0, "", []]) + def test_legacy_downgrade_preserves_falsy_values_and_non_text_blocks(self, value: JsonValue) -> None: + incoming: Final = CallToolResult( + content=[ + ImageContent(type="image", data="AA==", mime_type="image/png"), + TextContent(type="text", text="Done"), + ], + structured_content=value, + meta={"trace": "t1"}, + is_error=True, + ) + before: Final = incoming.model_dump(by_alias=True) + result: Final = to_call_tool_result(incoming, WireCompat.LEGACY) + assert isinstance(result, CallToolResult) + assert result.content == [*incoming.content, TextContent(type="text", text=json.dumps(value))] + assert result.structured_content is None + assert result.meta == incoming.meta + assert result.is_error is True + assert incoming.model_dump(by_alias=True) == before + + def test_is_error_survives_downgrade(self): + incoming = CallToolResult(content=[], structured_content=7, is_error=True) + result = to_call_tool_result(incoming, WireCompat.LEGACY) + assert isinstance(result, CallToolResult) and result.is_error is True + + +class TestInterimAndExceptionArms: + def test_modern_interim_passes_through(self): + interim = _interim() + assert to_call_tool_result(interim, WireCompat.MODERN) is interim + + def test_legacy_interim_becomes_error_result(self): + result = to_call_tool_result(_interim(), WireCompat.LEGACY) + assert isinstance(result, CallToolResult) + assert result.is_error is True + assert [c.text for c in result.content if isinstance(c, TextContent)] == [INPUT_REQUIRED_UNSUPPORTED_MESSAGE] + + def test_complete_call_tool_result_never_returns_interim(self): + result = complete_call_tool_result(_interim(), WireCompat.MODERN) + assert isinstance(result, CallToolResult) and result.is_error is True + + @pytest.mark.parametrize("compat", BOTH) + def test_exception_arm_matches_error_text_result(self, compat: WireCompat): + exc = ValueError("boom") + result = to_call_tool_result(exc, compat) + assert result == error_text_result(exc) + assert isinstance(result, CallToolResult) and result.is_error is True + assert [c.text for c in result.content if isinstance(c, TextContent)] == ["ValueError: boom"] + + +class TestSdkWireSerialization: + @pytest.mark.parametrize("version", KNOWN_PROTOCOL_VERSIONS) + def test_converted_results_serialize_on_their_negotiated_revision(self, version: str): + compat = wire_compat_for(version) + for body in ('{"a": 1}', "[1, 2]", "3", "null", "text"): + result = to_call_tool_result(parse_http_body(body), compat) + frame = _wire(result, version) + assert frame["content"] == [{"type": "text", "text": body}] + structured = json.loads(body) if body != "text" else None + expects_structured = structured is not None and ( + compat is WireCompat.MODERN or isinstance(structured, dict) + ) + assert ("structuredContent" in frame) is expects_structured, (version, body) + if expects_structured: + assert frame["structuredContent"] == structured + assert ("resultType" in frame) is (compat is WireCompat.MODERN), (version, body) + + @pytest.mark.parametrize("version", KNOWN_PROTOCOL_VERSIONS) + def test_downgraded_sdk_result_serializes_where_the_raw_one_would_not(self, version: str): + incoming = CallToolResult(content=[TextContent(type="text", text="Done")], structured_content=[1, 2]) + converted = to_call_tool_result(incoming, wire_compat_for(version)) + frame = _wire(converted, version) + if version in MODERN_PROTOCOL_VERSIONS: + assert frame["structuredContent"] == [1, 2] + return + with pytest.raises(ValidationError): + _wire(incoming, version) + assert "structuredContent" not in frame + assert frame["content"] == [{"type": "text", "text": "Done"}, {"type": "text", "text": "[1, 2]"}] + + def test_modern_interim_serializes_with_its_fields_intact(self): + frame = _wire(_interim(), "2026-07-28") + assert frame["resultType"] == "input_required" + assert frame["requestState"] == "abc" + assert frame["inputRequests"]["req-1"]["params"]["message"] == "Pick one" + + +class TestToGatewayTool: + def test_rename_is_a_deep_copy_that_keeps_every_other_field(self): + tool = Tool( + name="orig", + description="d", + inputSchema={"type": "object", "properties": {"q": {"type": "string"}}}, + _meta={"owner": "x"}, + ) + renamed = to_gateway_tool(tool, "srv-orig") + assert renamed.name == "srv-orig" + assert tool.name == "orig" + assert renamed.input_schema == tool.input_schema and renamed.input_schema is not tool.input_schema + assert renamed.meta == {"owner": "x"} + assert renamed.description == "d" From d33d36ce862459e144e814dcdab049a8bde3b06f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 07:56:46 -0700 Subject: [PATCH 02/58] fix(caching): propagate auth cache invalidation over Redis Cluster via a node-level pub/sub client (#43110) * fix(caching): give Redis Cluster clients a node-level pub/sub client so auth invalidation propagates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(caching): shorten pub/sub client docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(caching): noqa BLE001 on best-effort pubsub client close Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(caching): satisfy LIT002/LIT006 in pubsub client derivation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(auth): rename fake pubsub hook to init_pubsub_client Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): rename fake pubsub hook to init_pubsub_client Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(caching): drop class-level health ping patch from pub/sub client tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): close cluster pubsub pools and cover failure paths * fix(caching): clean up expired cluster pubsub clients safely * test(mcp): arm cancellation deadline after requests start --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .github/workflows/test-redis-compat.yml | 6 + litellm/caching/evicted_client_closer.py | 2 + litellm/caching/redis_cache.py | 62 ++++++- .../auth_cache_invalidation_pubsub.py | 12 -- .../proxy/common_utils/config_sync_pubsub.py | 29 +-- .../caching/test_evicted_client_closer.py | 29 +++ .../caching/test_redis_cluster_cache.py | 173 +++++++++++++++++- .../test_mcp_client.py | 6 +- .../mcp_server/test_byok_credential_cache.py | 2 +- .../proxy/auth/test_auth_checks.py | 4 +- .../test_auth_cache_invalidation_pubsub.py | 2 +- .../common_utils/test_config_sync_pubsub.py | 49 +++-- tests/test_litellm/proxy/test_proxy_server.py | 2 +- 13 files changed, 319 insertions(+), 59 deletions(-) diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml index 0b58cf9d486..25fb8f8bce3 100644 --- a/.github/workflows/test-redis-compat.yml +++ b/.github/workflows/test-redis-compat.yml @@ -8,9 +8,13 @@ on: paths: - "litellm/_redis.py" - "litellm/_redis_credential_provider.py" + - "litellm/caching/redis_cache.py" + - "litellm/caching/evicted_client_closer.py" - "tests/test_litellm/test_redis.py" - "tests/local_testing/test_caching.py" - "tests/test_litellm/caching/test_redis_connection_pool.py" + - "tests/test_litellm/caching/test_redis_cluster_cache.py" + - "tests/test_litellm/caching/test_evicted_client_closer.py" - ".github/workflows/test-redis-compat.yml" - "pyproject.toml" - "uv.lock" @@ -82,6 +86,8 @@ jobs: uv run --no-sync pytest \ tests/test_litellm/test_redis.py \ tests/test_litellm/caching/test_redis_connection_pool.py \ + tests/test_litellm/caching/test_redis_cluster_cache.py \ + tests/test_litellm/caching/test_evicted_client_closer.py \ tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \ tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \ --tb=short -vv \ diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py index eee7e2ea289..6e4635dd83a 100644 --- a/litellm/caching/evicted_client_closer.py +++ b/litellm/caching/evicted_client_closer.py @@ -136,6 +136,8 @@ def _has_connection_in_flight(client: object) -> bool: window as the only guard, exactly as it was before this check existed. """ try: + if getattr(getattr(client, "connection_pool", None), "_in_use_connections", None): + return True transport: Final = _transport_of(client) pooled_busy: Final = _pool_has_busy_connection(transport) if pooled_busy is not None: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 7cec84e0ebb..0b56c28f9b1 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -21,6 +21,7 @@ from collections.abc import Awaitable, Callable, Iterator, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast from pydantic import TypeAdapter @@ -290,6 +291,29 @@ def _opaque_kwarg_key(value: object) -> str: return f"{type(value).__name__}-{id(value)}" +_CLUSTER_ONLY_CONNECTION_KWARGS: Final[frozenset[str]] = frozenset({"response_callbacks"}) + + +def _cluster_node_pubsub_client( # pyright: ignore[reportUnknownParameterType] # redis generics + cluster: async_redis_cluster_client, # pyright: ignore[reportUnknownParameterType] # redis generics +) -> async_redis_client: + """Plain async client on one cluster node; classic PUBLISH/SUBSCRIBE is broadcast cluster-wide.""" + from redis.asyncio import ConnectionPool, Redis + + node: Final = cluster.get_default_node() or next(iter(cluster.nodes_manager.startup_nodes.values()), None) + if node is None: # pyright: ignore[reportUnnecessaryComparison] # get_default_node is None before cluster init + raise ValueError("cannot derive a pub/sub client: redis cluster has no default node and no startup nodes") + node_kwargs: Final = MappingProxyType( + { + key: value # pyright: ignore[reportAny] # connection_kwargs values are Any in redis stubs + for key, value in cluster.connection_kwargs.items() # pyright: ignore[reportAny] # connection_kwargs values are Any in redis stubs + if key not in _CLUSTER_ONLY_CONNECTION_KWARGS + } + ) + pool: Final = ConnectionPool(host=node.host, port=node.port, **node_kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # cluster kwargs validated by redis-py at runtime + return Redis.from_pool(pool) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics + + @functools.lru_cache(maxsize=1) def _redis_health_error_types() -> tuple[type, ...]: """Exception types that mean the Redis backend itself is unhealthy. @@ -738,6 +762,28 @@ class RedisCache(BaseCache): self.redis_async_client = redis_async_client return redis_async_client + def init_pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics + from redis.asyncio import RedisCluster + + from litellm import in_memory_llm_clients_cache + + client: Final = self.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics + if not isinstance(client, RedisCluster): + return client # pyright: ignore[reportUnknownVariableType] # redis generics + cache_key: Final = f"{self._get_async_client_cache_key()}-pubsub" + cached_client: Final = in_memory_llm_clients_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped in-memory client cache + key=cache_key + ) + if cached_client is not None: + return cast( # cast-ok: per-loop pub/sub client stored by this method # pyright: ignore[reportUnknownVariableType] # redis generics + async_redis_client, cached_client + ) + pubsub_client: Final = _cluster_node_pubsub_client( # pyright: ignore[reportUnknownVariableType] # redis generics + cluster=client + ) + in_memory_llm_clients_cache.set_cache(key=cache_key, value=pubsub_client, litellm_owned_client=True) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped in-memory client cache + return pubsub_client # pyright: ignore[reportUnknownVariableType] # redis generics + def _async_commands(self) -> _AsyncRedisCommands: return self.init_async_client() @@ -1785,7 +1831,21 @@ class RedisCache(BaseCache): self.redis_client.flushall() async def disconnect(self): - await self.async_redis_conn_pool.disconnect(inuse_connections=True) + from litellm import in_memory_llm_clients_cache + + if self.async_redis_conn_pool is not None: + await self.async_redis_conn_pool.disconnect(inuse_connections=True) + cached_pubsub_client: Final = cast( # cast-ok: only this module stores clients under this key # pyright: ignore[reportUnknownVariableType] # redis generics + async_redis_client | None, + in_memory_llm_clients_cache.get_cache( # pyright: ignore[reportUnknownMemberType] # untyped in-memory client cache + key=f"{self._get_async_client_cache_key()}-pubsub" + ), + ) + if cached_pubsub_client is not None: + try: + await cached_pubsub_client.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # redis stubs leave aclose unknown + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_logger.debug("Error closing cached pub/sub Redis client: %s", e) try: self.redis_client.close() except Exception as e: diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 11cb66d1a7f..3e09ad7157f 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -74,12 +74,6 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None: async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: str) -> None: try: client: Final = _pubsub_capable_client(redis_cache) - if client is None: - verbose_proxy_logger.debug( - "auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support", - cache_key, - ) - return async with _in_flight_publishes: await client.publish(auth_cache_invalidation_channel(redis_cache), message) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors @@ -185,12 +179,6 @@ class AuthCacheInvalidationSubscriber: while True: try: client = _pubsub_capable_client(self._redis_cache) - if client is None: - verbose_proxy_logger.warning( - "auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; " - "cross-worker eviction falls back to the local cache TTL" - ) - return pubsub = client.pubsub() try: await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache)) diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index b20c0d9c9a5..b4ebb5fa876 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -82,22 +82,13 @@ def config_sync_channel(redis_cache: "RedisCache") -> str: return f"{redis_cache.namespace}:{CONFIG_SYNC_CHANNEL}" -def _raw_async_client(redis_cache: "RedisCache") -> object: - return cast( # cast-ok: redis-py generics leave the client type partially unknown - object, - redis_cache.init_async_client(), # pyright: ignore[reportUnknownMemberType] # redis generics +def _pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient: + return cast( # cast-ok: protocol view of the pub/sub-capable async redis client + _ConfigSyncPubSubClient, + redis_cache.init_pubsub_client(), # pyright: ignore[reportUnknownMemberType] # redis generics ) -def _pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient | None: - from redis.asyncio import Redis - - client: Final = _raw_async_client(redis_cache) - if isinstance(client, Redis): - return cast(_ConfigSyncPubSubClient, client) # cast-ok: protocol view of the standalone redis client - return None - - @dataclass(frozen=True, slots=True) class _ConfigChangeMessage: object_type: str @@ -112,12 +103,6 @@ async def publish_config_change(redis_cache: "RedisCache | None", object_type: s return try: client: Final = _pubsub_capable_client(redis_cache) - if client is None: - verbose_proxy_logger.debug( - "config sync publish for %s skipped: cluster redis client has no pub/sub support", - object_type, - ) - return await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type)) except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e) @@ -238,12 +223,6 @@ class ConfigSyncSubscriber: while True: try: client = _pubsub_capable_client(self._redis_cache) - if client is None: - verbose_proxy_logger.warning( - "config sync subscriber disabled: cluster redis client has no pub/sub support; " - "interval polling remains the only sync mechanism" - ) - return pubsub = client.pubsub() try: await pubsub.subscribe(config_sync_channel(self._redis_cache)) diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/test_litellm/caching/test_evicted_client_closer.py index 939be5f3d6b..7679e276621 100644 --- a/tests/test_litellm/caching/test_evicted_client_closer.py +++ b/tests/test_litellm/caching/test_evicted_client_closer.py @@ -10,9 +10,11 @@ never closed, because litellm does not own its lifecycle. import asyncio import gc import weakref +from unittest.mock import AsyncMock import httpx import pytest +from redis.asyncio import ConnectionPool, Redis from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -72,6 +74,33 @@ def make_closer(clock: FakeClock, grace_seconds: float = 60.0) -> EvictedClientC return EvictedClientCloser(grace_seconds=grace_seconds, clock=clock) +@pytest.mark.asyncio +async def test_redis_client_is_closed_only_after_its_subscription_releases_the_connection(): + closer = EvictedClientCloser(grace_seconds=0) + pool = ConnectionPool() + client = Redis.from_pool(pool) + closed = asyncio.Event() + connection = AsyncMock() + connection.disconnect.side_effect = closed.set + pool._available_connections.append(connection) + borrowed = pool.get_available_connection() + closer.mark_owned(client) + closer.schedule(client) + + closer.reap() + await asyncio.sleep(0.05) + + connection.disconnect.assert_not_awaited() + assert closer.pending_count == 1 + + await pool.release(borrowed) + closer.reap() + await asyncio.wait_for(closed.wait(), timeout=1) + + connection.disconnect.assert_awaited_once() + assert closer.pending_count == 0 + + async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: """Serves a chunked body slowly, so a request stays on the wire long enough to observe.""" await reader.read(4096) diff --git a/tests/test_litellm/caching/test_redis_cluster_cache.py b/tests/test_litellm/caching/test_redis_cluster_cache.py index 0763b5110d5..ba1deabd0e1 100644 --- a/tests/test_litellm/caching/test_redis_cluster_cache.py +++ b/tests/test_litellm/caching/test_redis_cluster_cache.py @@ -1,13 +1,20 @@ +import asyncio from importlib import import_module import json -from unittest.mock import MagicMock, patch +import ssl +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from redis.asyncio import Redis, RedisCluster +from redis.asyncio.cluster import ClusterNode +from redis.asyncio.connection import SSLConnection from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.evicted_client_closer import EvictedClientCloser +from litellm.caching.llm_caching_handler import LLMClientCache @patch("litellm._redis.init_redis_cluster") @@ -175,3 +182,167 @@ def test_router_create_redis_cache_cluster_detection( with patch.object(RedisCache, "__init__", _mock_redis_cache_init): redis_cache = Router._create_redis_cache(cache_config) assert isinstance(redis_cache, expected_cache_type) + + +def _isolated_redis_cache(host: str) -> RedisCache: + """RedisCache whose sync client and pool are stubbed out.""" + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()), + ): + return RedisCache(host=host, port=6379) + + +def _cluster_for_pubsub(startup_node_host: str = "10.9.9.9") -> RedisCluster: + """Uninitialized RedisCluster carrying the connection kwargs a real one would.""" + return RedisCluster( + startup_nodes=[ClusterNode(host=startup_node_host, port=7000)], + password="cluster-secret", + socket_timeout=7.0, + ) + + +def test_init_pubsub_client_derives_a_node_client_for_cluster_backend() -> None: + """LIT-8543: a cluster-backed cache must return a pub/sub-capable client. + + The derived client pins a plain Redis connection pool to the cluster's + default node, inheriting the connection kwargs minus cluster-only keys. + """ + cache = _isolated_redis_cache("cluster-pubsub-default-node") + cluster = _cluster_for_pubsub() + node = ClusterNode(host="10.1.2.3", port=7001) + cluster.nodes_manager.default_node = node + cache.init_async_client = MagicMock(return_value=cluster) + + client = cache.init_pubsub_client() + + assert isinstance(client, Redis) and not isinstance(client, RedisCluster) + kwargs = client.connection_pool.connection_kwargs + assert kwargs["host"] == "10.1.2.3" + assert kwargs["port"] == 7001 + assert kwargs["password"] == "cluster-secret" + assert kwargs["socket_timeout"] == 7.0 + assert "response_callbacks" not in kwargs + + +def test_init_pubsub_client_falls_back_to_first_startup_node() -> None: + """Before cluster initialization there is no default node; the first + startup node is a valid pub/sub target.""" + cache = _isolated_redis_cache("cluster-pubsub-startup-fallback") + cluster = _cluster_for_pubsub(startup_node_host="10.8.8.8") + cache.init_async_client = MagicMock(return_value=cluster) + + client = cache.init_pubsub_client() + + assert isinstance(client, Redis) + assert client.connection_pool.connection_kwargs["host"] == "10.8.8.8" + + +def test_init_pubsub_client_returns_the_same_cached_client_on_repeat_calls() -> None: + cache = _isolated_redis_cache("cluster-pubsub-caching") + cluster = _cluster_for_pubsub() + cluster.nodes_manager.default_node = ClusterNode(host="10.1.2.3", port=7001) + cache.init_async_client = MagicMock(return_value=cluster) + + first = cache.init_pubsub_client() + second = cache.init_pubsub_client() + + assert first is second + + +def test_init_pubsub_client_returns_the_shared_async_client_for_standalone() -> None: + cache = _isolated_redis_cache("standalone-pubsub") + standalone = Redis() + cache.init_async_client = MagicMock(return_value=standalone) + + assert cache.init_pubsub_client() is standalone + + +def test_init_pubsub_client_preserves_tls_and_authentication() -> None: + cache = _isolated_redis_cache("cluster-pubsub-tls") + cluster = RedisCluster( + startup_nodes=[ClusterNode(host="redis.example.test", port=7000)], + ssl=True, + ssl_cert_reqs="required", + ssl_check_hostname=True, + username="pubsub-user", + password="test-password", + socket_connect_timeout=3.0, + socket_keepalive=True, + ) + cache.init_async_client = MagicMock(return_value=cluster) + + client = cache.init_pubsub_client() + connection = client.connection_pool.make_connection() + + assert isinstance(connection, SSLConnection) + assert connection.ssl_context.cert_reqs == ssl.CERT_REQUIRED + assert connection.ssl_context.check_hostname is True + assert connection.username == "pubsub-user" + assert connection.password == "test-password" + assert connection.socket_connect_timeout == 3.0 + assert connection.socket_keepalive is True + + +def test_init_pubsub_client_rejects_missing_nodes_and_can_retry() -> None: + cache = _isolated_redis_cache("cluster-pubsub-no-nodes") + cluster = _cluster_for_pubsub() + cluster.nodes_manager.startup_nodes = {} + cache.init_async_client = MagicMock(return_value=cluster) + + with pytest.raises(ValueError, match="no default node and no startup nodes"): + cache.init_pubsub_client() + + cluster.nodes_manager.default_node = ClusterNode(host="recovered.example.test", port=7001) + client = cache.init_pubsub_client() + + assert client.connection_pool.connection_kwargs["host"] == "recovered.example.test" + + +@pytest.mark.parametrize("close_fails", [False, True]) +@pytest.mark.parametrize("has_shared_pool", [False, True]) +def test_disconnect_closes_derived_pubsub_connections_even_when_pool_close_fails( + close_fails: bool, has_shared_pool: bool +) -> None: + cache = _isolated_redis_cache(f"cluster-pubsub-close-{close_fails}-{has_shared_pool}") + cache.async_redis_conn_pool = AsyncMock() if has_shared_pool else None + cache.init_async_client = MagicMock(return_value=_cluster_for_pubsub()) + + async def exercise() -> None: + client = cache.init_pubsub_client() + connection = AsyncMock() + connection.disconnect.side_effect = ConnectionError("connection close failed") if close_fails else None + client.connection_pool._available_connections.append(connection) + + await cache.disconnect() + + connection.disconnect.assert_awaited_once() + if has_shared_pool: + cache.async_redis_conn_pool.disconnect.assert_awaited_once_with(inuse_connections=True) + cache.redis_client.close.assert_called_once() + + asyncio.run(exercise()) + + +def test_expired_pubsub_client_closes_connections_after_eviction() -> None: + cache = _isolated_redis_cache("cluster-pubsub-expired") + cache.init_async_client = MagicMock(return_value=_cluster_for_pubsub()) + clients = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0)) + + async def exercise() -> None: + client = cache.init_pubsub_client() + closed = asyncio.Event() + connection = AsyncMock() + connection.disconnect.side_effect = closed.set + client.connection_pool._available_connections.append(connection) + cache_key = clients.update_cache_key_with_event_loop(f"{cache._get_async_client_cache_key()}-pubsub") + clients.ttl_dict[cache_key] = 0 + + replacement = cache.init_pubsub_client() + await asyncio.wait_for(closed.wait(), timeout=1) + + assert replacement is not client + connection.disconnect.assert_awaited_once() + + with patch("litellm.in_memory_llm_clients_cache", clients): + asyncio.run(exercise()) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index ae30b086c6e..368e34c455d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -2762,6 +2762,7 @@ async def test_cancellation_delivers_termination_over_tcp( cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool ) -> None: started: Final = asyncio.Event() + scope_ready: Final[asyncio.Future[anyio.CancelScope]] = asyncio.get_running_loop().create_future() terminations: Final[list[bytes]] = [] starts: Final[list[bytes]] = [] stop: Final = asyncio.Event() @@ -2858,13 +2859,16 @@ async def test_cancellation_delivers_termination_over_tcp( async def invoke(): if cancel_mode == "scope": - with anyio.fail_after(0.2): + with anyio.fail_after(None) as scope: + scope_ready.set_result(scope) return await calls() return await calls() try: task: Final = asyncio.create_task(invoke()) await asyncio.wait_for(started.wait(), 3) + if cancel_mode == "scope": + (await scope_ready).deadline = anyio.current_time() + 0.2 if cancel_mode == "task": task.cancel() expected_error: Final = ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_credential_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_credential_cache.py index 8ec5b8642bc..0ec4b431276 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_credential_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_credential_cache.py @@ -16,7 +16,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache class _FakeRedisCache: namespace = None - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return object() diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index b811d4453ca..e42a47a1091 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8998,7 +8998,7 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t def __init__(self) -> None: self.namespace = None - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return _RecordingRedisClient() local_spend_counter_cache = DualCache() @@ -9072,7 +9072,7 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_ def __init__(self) -> None: self.namespace = None - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return _RecordingRedisClient() local_spend_counter_cache = DualCache() diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 96770ee01c4..34d9741bcc0 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -87,7 +87,7 @@ class _FakeRedisCache: self._client = client self.namespace = namespace - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return self._client diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 0f64ef2b4ca..83ed3afa293 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -81,14 +81,25 @@ class _FailingPublishRedisClient(Redis): raise ConnectionError("redis down") -class _NotRedisClient: - def __init__(self) -> None: +class _ScriptedPubSubClient: + """Pub/sub-capable client that is not a redis.asyncio.Redis. + + Mirrors what RedisCache.init_pubsub_client returns for a cluster backend: + a node-level client exposing publish/pubsub without being an instance of + the standalone Redis class. + """ + + def __init__(self, pubsubs: Iterable["_QueuePubSub"]) -> None: + self._scripted_pubsubs = iter(pubsubs) self.published: List[Tuple[str, str]] = [] async def publish(self, channel: str, message: str) -> int: self.published.append((channel, message)) return 1 + def pubsub(self) -> "_QueuePubSub": + return next(self._scripted_pubsubs) + class _QueuePubSub: def __init__(self, initial_messages: Iterable[str] = ()) -> None: @@ -159,14 +170,14 @@ class _FakeRedisCache: self._client = client self.namespace = namespace - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return self._client class _ExplodingRedisCache: namespace: Optional[str] = None - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: raise ConnectionError("cannot connect") @@ -215,13 +226,17 @@ async def test_publish_swallows_client_init_errors() -> None: await publish_config_change(redis_cache=_ExplodingRedisCache(), object_type="litellm_proxymodeltable") -async def test_publish_skips_clients_without_pubsub_support() -> None: - client = _NotRedisClient() +async def test_publish_reaches_cluster_derived_pubsub_clients() -> None: + """LIT-8543: a cluster-backed cache returns a node-level client from + init_pubsub_client; publishes must go out on it instead of being skipped.""" + client = _ScriptedPubSubClient(pubsubs=[]) cache = _FakeRedisCache(client) await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") - assert client.published == [] + assert client.published == [ + (CONFIG_SYNC_CHANNEL, json.dumps({"object_type": "litellm_proxymodeltable"})) + ] async def test_subscriber_runs_injected_callbacks_in_order_on_message() -> None: @@ -558,20 +573,26 @@ async def test_stop_before_start_is_a_noop() -> None: await subscriber.stop() -async def test_subscriber_exits_without_callbacks_when_client_lacks_pubsub() -> None: - cache = _FakeRedisCache(_NotRedisClient()) +async def test_subscriber_subscribes_on_cluster_derived_pubsub_client() -> None: + """LIT-8543: the subscriber used to disable itself on cluster caches; now it + subscribes on the node-level client init_pubsub_client returns.""" + pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})]) + cache = _FakeRedisCache(_ScriptedPubSubClient(pubsubs=[pubsub])) resyncs: List[str] = [] + fired = asyncio.Event() subscriber = ConfigSyncSubscriber( redis_cache=cache, - resync_callbacks=(_recording_callback(resyncs, "resync", asyncio.Event()),), + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, ) subscriber.start() - task = subscriber._task - assert task is not None - await asyncio.wait_for(task, timeout=5) + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() - assert resyncs == [] + assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert resyncs == ["resync"] class _FakeTableActions: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 250556c9281..884a9c81500 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -15127,7 +15127,7 @@ async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached def __init__(self, client: object) -> None: self._client = client - def init_async_client(self) -> object: + def init_pubsub_client(self) -> object: return self._client byok_credential_cache.flush_cache() From c289d5d6fba3093d7dff15bb87fd9244a591998b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 09:16:17 -0700 Subject: [PATCH 03/58] feat(model-catalog): add Rust registry validation (#43136) * doc * feat(model-catalog): add Rust registry validation * test(model-catalog): ignore integration tests * docs(model-catalog): fix validation note grammar --------- Co-authored-by: Yujong Lee --- litellm-rust/Cargo.lock | 181 +++++- litellm-rust/crates/model-catalog/AGENTS.md | 6 + litellm-rust/crates/model-catalog/Cargo.toml | 8 +- litellm-rust/crates/model-catalog/README.md | 25 - .../crates/model-catalog/benches/catalog.rs | 21 - .../crates/model-catalog/src/capabilities.rs | 80 +++ .../crates/model-catalog/src/catalog.rs | 40 +- .../crates/model-catalog/src/error.rs | 11 +- .../crates/model-catalog/src/fallback.rs | 23 + litellm-rust/crates/model-catalog/src/lib.rs | 24 +- .../crates/model-catalog/src/model_info.rs | 587 ++++++------------ .../crates/model-catalog/src/pricing.rs | 97 +++ .../crates/model-catalog/src/schema.rs | 140 ++++- .../crates/model-catalog/src/validation.rs | 218 +++++++ .../crates/model-catalog/tests/catalog.rs | 82 ++- .../tests/registry_validation.rs | 76 +++ .../crates/model-catalog/tests/schema.rs | 121 ++++ .../crates/model-catalog/tests/spec_parity.rs | 121 ---- 18 files changed, 1227 insertions(+), 634 deletions(-) create mode 100644 litellm-rust/crates/model-catalog/AGENTS.md delete mode 100644 litellm-rust/crates/model-catalog/README.md delete mode 100644 litellm-rust/crates/model-catalog/benches/catalog.rs create mode 100644 litellm-rust/crates/model-catalog/src/capabilities.rs create mode 100644 litellm-rust/crates/model-catalog/src/fallback.rs create mode 100644 litellm-rust/crates/model-catalog/src/pricing.rs create mode 100644 litellm-rust/crates/model-catalog/src/validation.rs create mode 100644 litellm-rust/crates/model-catalog/tests/registry_validation.rs create mode 100644 litellm-rust/crates/model-catalog/tests/schema.rs delete mode 100644 litellm-rust/crates/model-catalog/tests/spec_parity.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c522bf205b4..5de32f5b622 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -897,6 +897,12 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "borrow-or-share" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0b364ead1874514c8c2855ab558056ebfeb775653e7ae45ff72f28f8f3166c" + [[package]] name = "bstr" version = "1.13.1" @@ -914,6 +920,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytecount" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e" + [[package]] name = "byteorder" version = "1.5.0" @@ -1540,6 +1552,15 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +[[package]] +name = "email_address" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e079f19b08ca6239f47f8ba8509c11cf3ea30095831f7fed61441475edd8c449" +dependencies = [ + "serde", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -1639,6 +1660,17 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "fluent-uri" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc74ac4d8359ae70623506d512209619e5cf8f347124910440dbc221714b328e" +dependencies = [ + "borrow-or-share", + "ref-cast", + "serde", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1660,6 +1692,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fraction" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e246562084dde8ebbcc943b261c406ce4f68e5032ec28029a251a47d6a295500" +dependencies = [ + "num", + "num-bigint 0.4.8", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1817,9 +1859,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 5.3.0", "wasip2", + "wasm-bindgen", ] [[package]] @@ -2660,6 +2704,59 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonschema" +version = "0.55.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b68339c3d874e48151d74ffe256d93a58cffa240983cb0967d3cbaea083a44fe" +dependencies = [ + "ahash", + "bytecount", + "data-encoding", + "email_address", + "fancy-regex 0.19.2", + "fraction", + "getrandom 0.3.4", + "itoa", + "jsonschema-regex", + "jsonschema-value", + "num-cmp", + "num-traits", + "percent-encoding", + "referencing", + "regex", + "serde", + "serde_json", + "strum", + "unicode-general-category", + "uuid-simd", +] + +[[package]] +name = "jsonschema-regex" +version = "0.55.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6307b5b51216ec9b941b52244c74043fa0b1d6b657b56199f57cb1416d3641c5" +dependencies = [ + "regex-syntax", +] + +[[package]] +name = "jsonschema-value" +version = "0.55.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0230ac05e09c6111e96c147b75c390579f5cbd45654b980c68ac60fe17b3f129" +dependencies = [ + "ahash", + "bytecount", + "fraction", + "getrandom 0.3.4", + "num-cmp", + "num-traits", + "serde_json", + "zmij", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -3126,14 +3223,14 @@ dependencies = [ name = "litellm-model-catalog" version = "0.1.0" dependencies = [ - "criterion", "indexmap 2.14.0", - "litellm-model-catalog", + "jsonschema", "rstest", "schemars 1.2.2", "serde", "serde_json", "thiserror 2.0.19", + "time", ] [[package]] @@ -3504,6 +3601,12 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "micromap" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a86d3146ed3995b5913c414f6664344b9617457320782e64f0bb44afd49d74" + [[package]] name = "mime" version = "0.3.17" @@ -3599,6 +3702,20 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint 0.4.8", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.8" @@ -3619,6 +3736,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-cmp" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63335b2e2c34fae2fb0aa2cecfd9f0832a1e24b3b32ecec612c3426d46dc8aaa" + [[package]] name = "num-complex" version = "0.4.6" @@ -3643,6 +3766,27 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint 0.4.8", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -4458,6 +4602,23 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "referencing" +version = "0.55.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a196a5b4a8a12f46b6353174df865a05d41a6055aff212ec30877492788618b6" +dependencies = [ + "ahash", + "fluent-uri", + "getrandom 0.3.4", + "hashbrown 0.17.1", + "itoa", + "micromap", + "parking_lot", + "percent-encoding", + "serde_json", +] + [[package]] name = "regex" version = "1.13.1" @@ -5941,6 +6102,12 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" +[[package]] +name = "unicode-general-category" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b993bddc193ae5bd0d623b49ec06ac3e9312875fdae725a975c51db1cc1677f" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -6021,6 +6188,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "uuid-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b082222b4f6619906941c17eb2297fff4c2fb96cb60164170522942a200bd8" +dependencies = [ + "outref", + "vsimd", +] + [[package]] name = "valuable" version = "0.1.1" diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md new file mode 100644 index 00000000000..9fcbcd57a76 --- /dev/null +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -0,0 +1,6 @@ +## Validation + +For `model_prices_and_context_window.json` validation, we should eventually: + +- Remove any schema file like `model_prices_and_context_window.schema.json` +- Stop skipping this crate's tests diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index ea75c6386d8..0b26e398ac8 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -14,12 +14,8 @@ schemars = { version = "1.0", optional = true } serde.workspace = true serde_json.workspace = true thiserror.workspace = true +time.workspace = true [dev-dependencies] -criterion.workspace = true +jsonschema = { version = "0.55.1", default-features = false } rstest.workspace = true -litellm-model-catalog = { path = ".", features = ["schema"] } - -[[bench]] -name = "catalog" -harness = false diff --git a/litellm-rust/crates/model-catalog/README.md b/litellm-rust/crates/model-catalog/README.md deleted file mode 100644 index 973f7190614..00000000000 --- a/litellm-rust/crates/model-catalog/README.md +++ /dev/null @@ -1,25 +0,0 @@ -# Model catalog - -`litellm-model-catalog` builds an immutable snapshot from caller supplied JSON bytes. It has no network, Python, registration, or refresh behavior. The caller supplies optional source, revision, and ETag provenance. Parse and validation are separate so small synthetic catalogs can use explicit integrity limits - -The parser treats `sample_spec` and `fallback_generalizations` as reserved top level metadata. `fallback_rules()` exposes the typed rule array when present; this crate does not execute regex generalizations. Model entries retain all JSON fields except `aliases`, including unknown fields. `field()` returns `None` for an absent key and a JSON null, false, or zero value for a present key. The returned values are borrowed, so callers cannot mutate the snapshot - -Each entry also deserializes into `ModelInfo`, a typed mirror of `model_prices_and_context_window.schema.json`'s `modelEntry` definition, reachable via `ModelEntry::info()`. All schema fields are optional on `ModelInfo`, including `litellm_provider` which the schema marks required, so small synthetic catalogs still parse. Unknown fields are not part of `ModelInfo`; they remain on `fields()`. Building with the `schema` feature adds `schemars` derives and exposes `model_entry_json_schema()` for emitting the entry's JSON Schema. Parse and validation failures are reported by the `Error` enum in `error.rs`, while catalog logic lives in `catalog.rs` - -The integration tests read the repository's catalog and schema files at test time, assert every entry round-trips through `ModelInfo`, and verify that the generated schema's properties match the repository schema - -Aliases point to their canonical entries. An alias that exactly matches any canonical key is skipped; the first canonical entry claiming an alias wins. Invalid alias lists and nonstring names are skipped and reported by `alias_issues()`. Exact lookup wins. For a case insensitive miss, the last key with the same lowercase spelling wins, following Python's lowercase map built after aliases are appended. This uses Rust Unicode lowercasing, which can differ from Python for unusual Unicode model IDs - -`validate()` counts canonical entries before alias expansion and excludes both reserved keys. It enforces an explicit minimum and backup shrink ratio, with Python defaults of 50 models and 0.5. Parsing rejects nonobject model entries and known fields with the wrong JSON type, but ignores unknown fields. It does not enforce every constraint in the JSON schema, calculate prices, resolve providers, or check provenance authenticity. The caller decides how to handle validation failures - -This snapshot does not represent Python's live mutable `litellm.model_cost`, nested dict and list mutation, or mutation of dicts previously returned by Python APIs. It has no bridge or runtime integration - -## Benchmarks - -`cargo bench -p litellm-model-catalog --bench catalog` measures parsing plus alias indexing and exact lookup. For a local Python baseline on the same fixture, use: - -```sh -python3 -m timeit -s 'import json, pathlib; body = pathlib.Path("../model_prices_and_context_window.json").read_bytes()' 'json.loads(body)' -``` - -Run these commands from `litellm-rust`. Python's command measures JSON loading only, without alias expansion or snapshot construction. The Rust benchmark does not include future Python object materialization, so these numbers are not an end to end runtime comparison diff --git a/litellm-rust/crates/model-catalog/benches/catalog.rs b/litellm-rust/crates/model-catalog/benches/catalog.rs deleted file mode 100644 index d1f51507c2b..00000000000 --- a/litellm-rust/crates/model-catalog/benches/catalog.rs +++ /dev/null @@ -1,21 +0,0 @@ -use criterion::{Criterion, criterion_group, criterion_main}; -use litellm_model_catalog::{Catalog, Provenance}; -use std::hint::black_box; - -fn benchmarks(c: &mut Criterion) { - let body = include_bytes!("../../../../model_prices_and_context_window.json"); - c.bench_function("parse_current_catalog", |b| { - b.iter(|| Catalog::parse(black_box(body), Provenance::default()).unwrap()) - }); - let catalog = Catalog::parse(body, Provenance::default()).unwrap(); - let key = catalog - .model_names() - .next() - .expect("catalog must have a benchmark key"); - c.bench_function("lookup_catalog_key", |b| { - b.iter(|| black_box(&catalog).lookup(black_box(key))) - }); -} - -criterion_group!(benches, benchmarks); -criterion_main!(benches); diff --git a/litellm-rust/crates/model-catalog/src/capabilities.rs b/litellm-rust/crates/model-catalog/src/capabilities.rs new file mode 100644 index 00000000000..66b5f1c5d2e --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/capabilities.rs @@ -0,0 +1,80 @@ +use serde::{Deserialize, Serialize}; + +/// Primary API surface / task type of the model. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum Mode { + AudioSpeech, + AudioTranscription, + Chat, + Completion, + Embedding, + Evaluation, + Guardrail, + ImageEdit, + ImageGeneration, + Moderation, + Ocr, + Realtime, + Rerank, + Responses, + Search, + VectorStore, + VideoGeneration, +} + +/// Reasoning effort level accepted or applied by the model. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum ReasoningEffort { + None, + Minimal, + Low, + Medium, + High, + Xhigh, + Max, +} + +/// Gemini audio generation API the model is served through. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum VertexAiAudioApi { + LyriaPredict, + LyriaInteractions, +} + +/// Audio container format the model can return. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum AudioFormat { + Mp3, + Wav, +} + +/// Input modality the model accepts. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum InputModality { + Text, + Image, + Audio, + Video, +} + +/// Output modality the model can produce. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum OutputModality { + Text, + Image, + Audio, + Video, + Code, +} diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index dc7564f9bee..0b113436a5b 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -1,5 +1,6 @@ use crate::error::Error; -use crate::model_info::{FallbackGeneralizations, FallbackRule, ModelInfo}; +use crate::fallback::{FallbackGeneralizations, FallbackRule}; +use crate::model_info::ModelInfo; use indexmap::IndexMap; use serde::Deserialize; use serde_json::{Map, Value}; @@ -14,19 +15,9 @@ pub struct Provenance { #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { - pub backup_model_count: usize, + pub reference_model_count: usize, pub min_model_count: usize, - pub min_backup_ratio: f64, -} - -impl IntegrityLimits { - pub fn python_defaults(backup_model_count: usize) -> Self { - Self { - backup_model_count, - min_model_count: 50, - min_backup_ratio: 0.5, - } - } + pub min_reference_ratio: f64, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -99,16 +90,11 @@ impl Catalog { } _ => {} } - let Value::Object(ref object) = value else { + let Value::Object(mut fields) = value else { return Err(Error::EntryNotObject { model: name }); }; - let info = ModelInfo::deserialize(object)?; - let Value::Object(mut fields) = value else { - unreachable!("value checked is_object above") - }; - if let Some(aliases) = fields.remove("aliases") - && !aliases.is_null() - { + let info = ModelInfo::deserialize(&fields)?; + if let Some(aliases) = fields.remove("aliases") { match aliases { Value::Array(names) => alias_lists.push((name.clone(), names)), _ => alias_issues.push(AliasIssue::InvalidList { @@ -161,7 +147,9 @@ impl Catalog { } pub fn validate(&self, limits: IntegrityLimits) -> Result<(), Error> { - if !limits.min_backup_ratio.is_finite() || !(0.0..=1.0).contains(&limits.min_backup_ratio) { + if !limits.min_reference_ratio.is_finite() + || !(0.0..=1.0).contains(&limits.min_reference_ratio) + { return Err(Error::InvalidRatio); } let actual = self.entries.len(); @@ -171,13 +159,13 @@ impl Catalog { minimum: limits.min_model_count, }); } - if limits.backup_model_count > 0 - && (actual as f64) < (limits.backup_model_count as f64) * limits.min_backup_ratio + if limits.reference_model_count > 0 + && (actual as f64) < (limits.reference_model_count as f64) * limits.min_reference_ratio { return Err(Error::Shrunk { actual, - backup: limits.backup_model_count, - ratio: limits.min_backup_ratio, + reference: limits.reference_model_count, + ratio: limits.min_reference_ratio, }); } Ok(()) diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 83617312fff..edb6ba1eb11 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -1,6 +1,5 @@ use thiserror::Error; -/// Failures from parsing or validating a catalog snapshot. #[derive(Debug, Error)] pub enum Error { /// The body is not valid JSON, or a model entry fails typed deserialization. @@ -15,14 +14,14 @@ pub enum Error { /// Canonical entry count is under the configured minimum. #[error("catalog has {actual} models, below minimum {minimum}")] BelowMinimum { actual: usize, minimum: usize }, - /// Canonical entry count is under the configured backup shrink ratio. - #[error("catalog has {actual} models, below {ratio} of backup count {backup}")] + /// Canonical entry count is under the configured reference ratio. + #[error("catalog has {actual} models, below {ratio} of reference count {reference}")] Shrunk { actual: usize, - backup: usize, + reference: usize, ratio: f64, }, - /// The configured minimum backup ratio is not finite or outside `[0, 1]`. - #[error("minimum backup ratio must be finite and between zero and one")] + /// The configured minimum reference ratio is not finite or outside `[0, 1]`. + #[error("minimum reference ratio must be finite and between zero and one")] InvalidRatio, } diff --git a/litellm-rust/crates/model-catalog/src/fallback.rs b/litellm-rust/crates/model-catalog/src/fallback.rs new file mode 100644 index 00000000000..62291a84929 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/fallback.rs @@ -0,0 +1,23 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::BTreeMap; + +/// One regex rule generalizing unknown model ids to known families. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct FallbackRule { + pub name: String, + pub pattern: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(flatten)] + pub extra: BTreeMap, +} + +/// Regex rules that generalize unknown model ids to known families; not a model entry. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct FallbackGeneralizations { + pub rules: Vec, +} diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index 9c942a5521c..066c2c83c6b 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -1,16 +1,20 @@ +mod capabilities; mod catalog; mod error; +mod fallback; mod model_info; +mod pricing; +mod validation; + +pub use capabilities::*; +pub use catalog::*; +pub use error::*; +pub use fallback::*; +pub use model_info::*; +pub use pricing::*; +pub use validation::*; + #[cfg(feature = "schema")] mod schema; - -pub use catalog::{AliasIssue, Catalog, IntegrityLimits, ModelEntry, ModelMatch, Provenance}; -pub use error::Error; -pub use model_info::{ - AudioFormat, FallbackGeneralizations, FallbackRule, InputModality, Mode, ModelInfo, - OffPeakPricing, OffPeakWindow, OutputModality, ReasoningEffort, SearchContextCostPerQuery, - TieredRate, UtcHours, VertexAiAudioApi, WebSearchBillingUnit, Weekday, -}; - #[cfg(feature = "schema")] -pub use schema::model_entry_json_schema; +pub use schema::*; diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 4a56e1112d1..361cb56e9b1 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -1,673 +1,482 @@ +use crate::capabilities::{ + AudioFormat, InputModality, Mode, OutputModality, ReasoningEffort, VertexAiAudioApi, +}; +use crate::pricing::{OffPeakPricing, SearchContextCostPerQuery, TieredRate, WebSearchBillingUnit}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::BTreeMap; -/// Primary API surface / task type of the model. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum Mode { - AudioSpeech, - AudioTranscription, - Chat, - Completion, - Embedding, - Evaluation, - Guardrail, - ImageEdit, - ImageGeneration, - Moderation, - Ocr, - Realtime, - Rerank, - Responses, - Search, - VectorStore, - VideoGeneration, -} - -/// Reasoning effort level accepted or applied by the model. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum ReasoningEffort { - None, - Minimal, - Low, - Medium, - High, - Xhigh, - Max, -} - -/// Gemini audio generation API the model is served through. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum VertexAiAudioApi { - LyriaPredict, - LyriaInteractions, -} - -/// Whether web search is billed per query or per prompt. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum WebSearchBillingUnit { - PerQuery, - PerPrompt, -} - -/// Audio container format the model can return. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum AudioFormat { - Mp3, - Wav, -} - -/// Input modality the model accepts. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum InputModality { - Text, - Image, - Audio, - Video, -} - -/// Output modality the model can produce. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum OutputModality { - Text, - Image, - Audio, - Video, - Code, -} - -/// UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(untagged)] -pub enum UtcHours { - Single(String), - Multiple(Vec), -} - -/// ISO-8601 weekday number (1 = Monday .. 7 = Sunday) or English day name. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(untagged)] -pub enum Weekday { - Number(u8), - Name(String), -} - -/// One off-peak window entry inside `windows`. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct OffPeakWindow { - pub hours_utc: UtcHours, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub weekdays: Option>, -} - -/// Rates that replace the same-named base fields inside the stated UTC windows. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct OffPeakPricing { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub hours_utc: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub windows: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub weekday_timezone: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_cost_per_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_cost_per_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_cost_per_reasoning_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost: Option, -} - -/// USD cost per web search query, keyed by search context size. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct SearchContextCostPerQuery { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub search_context_size_low: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub search_context_size_medium: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub search_context_size_high: Option, -} - -/// One tier of a context-length or result-count tiered rate. -#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct TieredRate { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub range: Option<[f64; 2]>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_results_range: Option<[f64; 2]>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_cost_per_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_cost_per_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_cost_per_reasoning_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_cost_per_query: Option, -} - -/// One regex rule generalizing unknown model ids to known families. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct FallbackRule { - pub name: String, - pub pattern: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(flatten)] - pub extra: BTreeMap, -} - -/// Regex rules that generalize unknown model ids to known families; not a model entry. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(deny_unknown_fields)] -pub struct FallbackGeneralizations { - pub rules: Vec, -} - /// Typed mirror of one catalog model entry. #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct ModelInfo { - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub annotation_cost_per_page: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub annotation_cost_per_page_batches: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub audio_transcription_config: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub bedrock_converse_supports_strict_tools: Option, /// Highest reasoning effort the Bedrock output_config accepts for this model. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub bedrock_output_config_effort_ceiling: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_audio_token_cost: Option, /// USD per token written to the provider's prompt cache. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_creation_input_token_cost_above_32k_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_128k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_1hr: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_200k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_256k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_272k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_272k_tokens_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_272k_tokens_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_272k_tokens_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_creation_input_token_cost_above_32k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_audio_token_cost: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_image_token_cost: Option, /// USD per prompt token served from the provider's prompt cache. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_input_token_cost_above_32k_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_128k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_200k_tokens: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_200k_tokens_priority: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_256k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_272k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_272k_tokens_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_272k_tokens_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_272k_tokens_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_read_input_token_cost_above_32k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_512k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub citation_cost_per_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub code_interpreter_cost_per_session: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub comment: Option, /// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub default_reasoning_effort: Option, /// Date the provider deprecates the model, YYYY-MM-DD. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub deprecation_date: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub gemini_audio_only_live: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub gemini_native_audio: Option, /// USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub google_maps_grounding_cost_per_query: Option, /// USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub guardrail_cost_per_unit: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_audio_per_second: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_audio_per_second_above_128k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_audio_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_audio_token_batches: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_audio_token_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_character: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_character_above_128k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_image: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_image_above_128k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_image_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_image_token_batches: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_pixel: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_query: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_request: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_second: Option, /// USD per prompt token. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token_above_32k_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_128k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_200k_tokens: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_200k_tokens_priority: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_256k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_272k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_272k_tokens_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_272k_tokens_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_272k_tokens_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_cost_per_token_above_32k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_512k_tokens: Option, /// USD per prompt token via the provider's batch API. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_batches: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_cache_hit: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_per_second: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_per_second_above_128k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_per_second_above_15s_interval: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_per_second_above_8s_interval: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_video_token_batches: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub input_dbu_cost_per_token: Option, /// LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub litellm_provider: Option, /// Maximum prompt/context tokens the model accepts. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub max_input_tokens: Option, /// Maximum tokens the model can generate in one response. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub max_output_tokens: Option, /// Legacy field: max output tokens if the provider specifies it, else max input tokens. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, /// Free-form notes about the entry (e.g. pricing derivation). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub metadata: Option>, /// Primary API surface / task type of the model. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub mode: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub ocr_cost_per_credit: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub ocr_cost_per_page: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub ocr_cost_per_page_batches: Option, /// Rates that replace the same-named base fields while the request falls inside the stated UTC windows. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub off_peak_pricing: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_audio_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_character: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_character_above_128k_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_1024: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_1536: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_512: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_pixel: Option, /// USD per reasoning/thinking token, when billed separately. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_reasoning_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_1080p: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_2k: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_480p: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_4k: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_720p: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_second_768p: Option, /// USD per generated token. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token_above_32k_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_128k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_200k_tokens: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_200k_tokens_priority: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_256k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_272k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_272k_tokens_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_272k_tokens_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_272k_tokens_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_cost_per_token_above_32k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_512k_tokens: Option, /// USD per generated token via the provider's batch API. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_batches: Option, /// Flex service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_flex: Option, /// Priority service-tier rate for the same-named base field. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_priority: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_video_per_second: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_video_token: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_dbu_cost_per_token: Option, /// Embedding dimension for embedding models. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub output_vector_size: Option, /// Smallest prefix the provider will actually cache; absent means the provider default applies. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub prompt_cache_min_tokens: Option, /// Provider-internal routing hints (e.g. bedrock_invocation_schema). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub provider_specific_entry: Option>, /// Exact reasoning_effort levels this deployment accepts; wins over supports_* flags. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub reasoning_effort_levels: Option>, /// Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub regional_endpoint_uplift_multiplier: Option, /// Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub regional_processing_uplift_multiplier_eu: Option, /// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub regional_processing_uplift_multiplier_us: Option, /// Provider default requests-per-minute limit. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub rpm: Option, /// USD cost per web search query, keyed by search context size. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub search_context_cost_per_query: Option, /// URL of the provider pricing/model page this entry was taken from. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub source: Option, /// Audio container formats the model can return. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supported_audio_formats: Option>, /// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supported_endpoints: Option>, /// Input modalities the model accepts. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supported_modalities: Option>, /// Output modalities the model can produce. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supported_output_modalities: Option>, /// Cloud regions the model is available in ('global' or region ids). - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supported_regions: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_adaptive_thinking: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_anthropic_compaction: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_anthropic_thinking_payload: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_assistant_prefill: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_audio_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_audio_output: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_computer_use: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_embedding_image_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_fast_mode: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_forced_tool_use: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_function_calling: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_image_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_image_size: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_legacy_thinking: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_low_reasoning_effort: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_max_reasoning_effort: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_mid_conversation_system: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_minimal_reasoning_effort: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_multimodal: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_native_streaming: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_native_structured_output: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_none_reasoning_effort: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_nova_canvas_image_edit: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_output_config: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_parallel_function_calling: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_parallel_tool_use_config: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_pdf_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_prompt_cache_breakpoint: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_prompt_caching: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_reasoning: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_response_schema: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_sampling_params: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_speed: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_system_messages: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_thinking_cache_preservation: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_tool_choice: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_tool_search: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_url_context: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_video_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_vision: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_web_search: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub supports_xhigh_reasoning_effort: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub thinking_always_on: Option, /// Context-length or result-count tiered rates; each tier's costs apply within its range. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub tiered_pricing: Option>, /// Provider default tokens-per-minute limit. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub tpm: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub use_openai_responses_path: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub uses_embed_content: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub vertex_ai_audio_api: Option, /// Whether web search is billed per query or per prompt. - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde(skip_serializing_if = "Option::is_none")] pub web_search_billing_unit: Option, } diff --git a/litellm-rust/crates/model-catalog/src/pricing.rs b/litellm-rust/crates/model-catalog/src/pricing.rs new file mode 100644 index 00000000000..b8ed652451c --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/pricing.rs @@ -0,0 +1,97 @@ +use serde::{Deserialize, Serialize}; + +/// Whether web search is billed per query or per prompt. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum WebSearchBillingUnit { + PerQuery, + PerPrompt, +} + +/// UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(untagged)] +pub enum UtcHours { + Single(String), + Multiple(Vec), +} + +/// ISO-8601 weekday number (1 = Monday .. 7 = Sunday) or English day name. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(untagged)] +pub enum Weekday { + Number(u8), + Name(String), +} + +/// One off-peak window entry inside `windows`. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct OffPeakWindow { + pub hours_utc: UtcHours, + #[serde(skip_serializing_if = "Option::is_none")] + pub weekdays: Option>, +} + +/// Rates that replace the same-named base fields inside the stated UTC windows. +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct OffPeakPricing { + #[serde(skip_serializing_if = "Option::is_none")] + pub hours_utc: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub windows: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub weekday_timezone: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_reasoning_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_input_token_cost: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_creation_input_token_cost: Option, +} + +/// USD cost per web search query, keyed by search context size. +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct SearchContextCostPerQuery { + #[serde(skip_serializing_if = "Option::is_none")] + pub search_context_size_low: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub search_context_size_medium: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub search_context_size_high: Option, +} + +/// One tier of a context-length or result-count tiered rate. +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct TieredRate { + #[serde(skip_serializing_if = "Option::is_none")] + pub range: Option<[f64; 2]>, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_results_range: Option<[f64; 2]>, + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_reasoning_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_input_token_cost: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_creation_input_token_cost: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_query: Option, +} diff --git a/litellm-rust/crates/model-catalog/src/schema.rs b/litellm-rust/crates/model-catalog/src/schema.rs index 82cfd6c0352..7de988b3ff3 100644 --- a/litellm-rust/crates/model-catalog/src/schema.rs +++ b/litellm-rust/crates/model-catalog/src/schema.rs @@ -1,7 +1,137 @@ -use crate::model_info::ModelInfo; +use schemars::Schema; +use serde_json::{Map, Value, json}; -/// JSON Schema for one catalog model entry, mirroring -/// `model_prices_and_context_window.schema.json`'s `modelEntry` definition. -pub fn model_entry_json_schema() -> schemars::Schema { - schemars::schema_for!(ModelInfo) +/// JSON Schema for one model entry, including registry validation constraints. +pub fn model_entry_json_schema() -> Schema { + let mut schema = serde_json::to_value(schemars::schema_for!(crate::ModelInfo)) + .expect("derived model schema serializes"); + remove_nullable_optional_fields(&mut schema); + decorate_model_entry(&mut schema); + Schema::from( + schema + .as_object() + .expect("derived schema is an object") + .clone(), + ) +} + +/// JSON Schema for the complete model prices registry document. +pub fn registry_json_schema() -> Schema { + let mut entry = model_entry_json_schema().as_value().clone(); + let mut definitions = take_definitions(&mut entry); + entry.as_object_mut().unwrap().remove("$schema"); + definitions.insert("modelEntry".into(), entry); + + let mut fallback = serde_json::to_value(schemars::schema_for!(crate::FallbackGeneralizations)) + .expect("derived fallback schema serializes"); + remove_nullable_optional_fields(&mut fallback); + definitions.extend(take_definitions(&mut fallback)); + fallback.as_object_mut().unwrap().remove("$schema"); + + let root = json!({ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "LiteLLM model prices and context window registry", + "type": "object", + "properties": { + "sample_spec": {"type": "object"}, + "fallback_generalizations": fallback + }, + "additionalProperties": {"$ref": "#/$defs/modelEntry"}, + "$defs": definitions + }); + Schema::from(root.as_object().unwrap().clone()) +} + +fn take_definitions(schema: &mut Value) -> Map { + schema + .as_object_mut() + .unwrap() + .remove("$defs") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default() +} + +fn remove_nullable_optional_fields(value: &mut Value) { + match value { + Value::Array(values) => values.iter_mut().for_each(remove_nullable_optional_fields), + Value::Object(map) => { + map.values_mut().for_each(remove_nullable_optional_fields); + if let Some(Value::Array(types)) = map.get_mut("type") { + types.retain(|value| value != "null"); + if types.len() == 1 { + let only = types[0].clone(); + map.insert("type".into(), only); + } + } + if let Some(Value::Array(branches)) = map.get_mut("anyOf") { + branches.retain(|branch| branch.get("type") != Some(&Value::String("null".into()))); + if branches.len() == 1 { + let only = branches[0] + .as_object() + .expect("schema branch is an object") + .clone(); + map.remove("anyOf"); + map.extend(only); + } + } + } + _ => {} + } +} + +fn decorate_model_entry(schema: &mut Value) { + let object = schema.as_object_mut().unwrap(); + object.insert("required".into(), json!(["litellm_provider"])); + object.insert("additionalProperties".into(), Value::Bool(true)); + let properties = object + .get_mut("properties") + .unwrap() + .as_object_mut() + .unwrap(); + properties.insert( + "aliases".into(), + json!({"type": "array", "items": {"type": "string"}}), + ); + properties.get_mut("deprecation_date").unwrap()["format"] = json!("date"); + properties.get_mut("deprecation_date").unwrap()["pattern"] = + json!(r"^\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"); + + properties.iter_mut().for_each(|(name, property)| { + if name.contains("cost") { + property["minimum"] = json!(0); + } else if name.contains("uplift_multiplier") { + property["minimum"] = json!(1); + } + }); + properties.get_mut("guardrail_cost_per_unit").unwrap()["additionalProperties"]["minimum"] = + json!(0); + + let definitions = object.get_mut("$defs").unwrap().as_object_mut().unwrap(); + for definition in ["OffPeakPricing", "TieredRate", "SearchContextCostPerQuery"] { + let properties = definitions[definition]["properties"] + .as_object_mut() + .unwrap(); + properties.iter_mut().for_each(|(name, property)| { + if name.contains("cost") || definition == "SearchContextCostPerQuery" { + property["minimum"] = json!(0); + } + }); + } + definitions["OffPeakPricing"]["anyOf"] = json!([ + {"required": ["hours_utc"]}, + {"required": ["windows"]} + ]); + definitions["OffPeakPricing"]["properties"]["windows"]["minItems"] = json!(1); + definitions["OffPeakWindow"]["properties"]["weekdays"]["minItems"] = json!(1); + definitions["TieredRate"]["properties"]["range"]["items"]["minimum"] = json!(0); + definitions["TieredRate"]["properties"]["max_results_range"]["items"]["minimum"] = json!(0); + definitions["Weekday"]["anyOf"][0]["minimum"] = json!(1); + definitions["Weekday"]["anyOf"][0]["maximum"] = json!(7); + definitions["Weekday"]["anyOf"][1]["pattern"] = json!( + r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" + ); + let window_pattern = json!(r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"); + definitions["UtcHours"]["anyOf"][0]["pattern"] = window_pattern.clone(); + definitions["UtcHours"]["anyOf"][1]["items"]["pattern"] = window_pattern; + definitions["UtcHours"]["anyOf"][1]["minItems"] = json!(1); } diff --git a/litellm-rust/crates/model-catalog/src/validation.rs b/litellm-rust/crates/model-catalog/src/validation.rs new file mode 100644 index 00000000000..73f08a865c1 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/validation.rs @@ -0,0 +1,218 @@ +use std::collections::BTreeSet; + +use serde_json::{Map, Value}; +use thiserror::Error; + +use crate::{AliasIssue, Catalog, ModelInfo, UtcHours, Weekday}; + +/// A registry entry violates the checked-in catalog contract. +#[derive(Debug, Error)] +pub enum RegistryValidationError { + #[error("{reason}")] + Entry { model: String, reason: String }, + #[error("alias issue: {0:?}")] + Alias(AliasIssue), +} + +/// Validate one registry entry without restricting the tolerant catalog reader. +pub fn validate_model_entry(model: &str, value: &Value) -> Result<(), RegistryValidationError> { + validate_entry_inner(model, value).map_err(|reason| RegistryValidationError::Entry { + model: model.to_owned(), + reason, + }) +} + +/// Check every model and alias in a parsed catalog against registry rules. +pub fn validate_registry(catalog: &Catalog) -> Result<(), RegistryValidationError> { + if let Some(issue) = catalog.alias_issues().first() { + return Err(RegistryValidationError::Alias(issue.clone())); + } + catalog.model_names().try_for_each(|name| { + let entry = catalog.lookup(name).expect("catalog name must resolve"); + validate_model_entry(name, &Value::Object(entry.entry.fields().clone())) + }) +} + +fn json_eq(left: &Value, right: &Value) -> bool { + match (left, right) { + (Value::Number(left), Value::Number(right)) => left.as_f64() == right.as_f64(), + (Value::Array(left), Value::Array(right)) => { + left.len() == right.len() && left.iter().zip(right).all(|(a, b)| json_eq(a, b)) + } + (Value::Object(left), Value::Object(right)) => { + left.len() == right.len() + && left + .iter() + .all(|(key, value)| right.get(key).is_some_and(|other| json_eq(value, other))) + } + _ => left == right, + } +} + +fn keys(value: &Map) -> BTreeSet { + value.keys().cloned().collect() +} + +fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> BTreeSet { + left.symmetric_difference(right).cloned().collect() +} + +fn validate_entry_inner(model_name: &str, value: &Value) -> Result<(), String> { + let object = value + .as_object() + .ok_or_else(|| format!("{model_name} must be an object"))?; + if let Some(aliases) = object.get("aliases") { + let names = aliases + .as_array() + .ok_or_else(|| format!("{model_name}.aliases must be an array"))?; + if names.iter().any(|name| !name.is_string()) { + return Err(format!("{model_name}.aliases must contain strings")); + } + } + let info: ModelInfo = + serde_json::from_value(value.clone()).map_err(|error| format!("{model_name}: {error}"))?; + if info.litellm_provider.is_none() { + return Err(format!("{model_name}.litellm_provider is required")); + } + validate_dates_and_windows(model_name, &info)?; + let serialized = serde_json::to_value(info).map_err(|error| error.to_string())?; + let mut expected = object.clone(); + expected.remove("aliases"); + if !json_eq(&Value::Object(expected.clone()), &serialized) { + let actual = serialized + .as_object() + .expect("ModelInfo serializes as an object"); + return Err(format!( + "{model_name} has an unknown field, null, or changed value: {:?}", + symmetric_difference(&keys(&expected), &keys(actual)) + )); + } + check_prices(model_name, value) +} + +fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), String> { + if let Some(date) = &info.deprecation_date { + let format = time::format_description::parse_borrowed::<2>("[year]-[month]-[day]").unwrap(); + time::Date::parse(date, &format) + .map_err(|error| format!("{model_name}.deprecation_date: {error}"))?; + } + let Some(pricing) = &info.off_peak_pricing else { + return Ok(()); + }; + if pricing.hours_utc.is_none() && pricing.windows.is_none() { + return Err(format!( + "{model_name}.off_peak_pricing needs hours or windows" + )); + } + if let Some(hours) = &pricing.hours_utc { + validate_hours(hours)?; + } + if let Some(windows) = &pricing.windows { + if windows.is_empty() { + return Err(format!("{model_name}.off_peak_pricing.windows is empty")); + } + windows.iter().try_for_each(|window| { + validate_hours(&window.hours_utc)?; + if let Some(days) = &window.weekdays + && (days.is_empty() || days.iter().any(|day| !valid_weekday(day))) + { + return Err(format!("{model_name}.off_peak_pricing.weekdays is invalid")); + } + Ok(()) + })?; + } + Ok(()) +} + +fn validate_hours(hours: &UtcHours) -> Result<(), String> { + let values = match hours { + UtcHours::Single(value) => std::slice::from_ref(value), + UtcHours::Multiple(values) => values.as_slice(), + }; + if values.is_empty() || values.iter().any(|value| !valid_utc_window(value)) { + return Err("off_peak_pricing.hours_utc is invalid".into()); + } + Ok(()) +} + +fn valid_utc_window(value: &str) -> bool { + let Some((start, end)) = value.split_once('-') else { + return false; + }; + [start, end].into_iter().all(|clock| { + let Some((hour, minute)) = clock.split_once(':') else { + return false; + }; + hour.len() == 2 + && minute.len() == 2 + && hour.parse::().is_ok_and(|hour| hour < 24) + && minute.parse::().is_ok_and(|minute| minute < 60) + }) +} + +fn valid_weekday(day: &Weekday) -> bool { + match day { + Weekday::Number(number) => (1..=7).contains(number), + Weekday::Name(name) => matches!( + name.to_ascii_lowercase().as_str(), + "mon" + | "monday" + | "tue" + | "tues" + | "tuesday" + | "wed" + | "wednesday" + | "thu" + | "thur" + | "thurs" + | "thursday" + | "fri" + | "friday" + | "sat" + | "saturday" + | "sun" + | "sunday" + ), + } +} + +fn check_prices(path: &str, value: &Value) -> Result<(), String> { + let Some(object) = value.as_object() else { + return Ok(()); + }; + object.iter().try_for_each(|(key, field)| { + let field_path = format!("{path}.{key}"); + if (key.contains("cost") + || path.ends_with(".guardrail_cost_per_unit") + || path.ends_with(".search_context_cost_per_query")) + && let Some(number) = field.as_f64() + && number < 0.0 + { + return Err(format!("{field_path} must be nonnegative")); + } + if key.contains("uplift_multiplier") + && let Some(number) = field.as_f64() + && number < 1.0 + { + return Err(format!("{field_path} must be at least one")); + } + if matches!(key.as_str(), "range" | "max_results_range") + && field.as_array().is_some_and(|values| { + values + .iter() + .any(|value| value.as_f64().is_some_and(|n| n < 0.0)) + }) + { + return Err(format!("{field_path} must be nonnegative")); + } + if matches!(key.as_str(), "metadata" | "provider_specific_entry") { + return Ok(()); + } + match field.as_array() { + Some(items) => items.iter().enumerate().try_for_each(|(index, item)| { + check_prices(&format!("{field_path}[{index}]"), item) + }), + None => check_prices(&field_path, field), + } + }) +} diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index bcadd38e908..d87de3c5b3b 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -43,6 +43,7 @@ fn fixture_catalog() -> Catalog { } #[rstest] +#[ignore] fn preserves_fields_and_metadata(fixture_catalog: Catalog) { let catalog = fixture_catalog; let entry = catalog.lookup("SHORT").unwrap(); @@ -68,6 +69,7 @@ fn preserves_fields_and_metadata(fixture_catalog: Catalog) { } #[rstest] +#[ignore] fn snapshot_does_not_borrow_source() { let mut source = ALPHA_FIXTURE.to_vec(); let catalog = Catalog::parse(&source, Provenance::default()).unwrap(); @@ -84,7 +86,8 @@ fn snapshot_does_not_borrow_source() { #[case("shared", "Second")] #[case("FIRST", "First")] #[case("sHaReD", "Second")] -fn alias_collisions_and_case_fallback_follow_python_order( +#[ignore] +fn alias_collisions_and_case_fallback_follow_entry_order( #[case] lookup: &str, #[case] expected: &str, ) { @@ -117,6 +120,36 @@ fn alias_collisions_and_case_fallback_follow_python_order( ); } +#[test] +#[ignore] +fn json_entry_order_controls_alias_ownership_and_case_fallback() { + let forward = Catalog::parse( + br#"{ + "Alpha":{"aliases":["shared"]}, + "Beta":{"aliases":["shared"]}, + "Foo":{}, + "fOO":{} + }"#, + Provenance::default(), + ) + .unwrap(); + let reversed = Catalog::parse( + br#"{ + "fOO":{}, + "Foo":{}, + "Beta":{"aliases":["shared"]}, + "Alpha":{"aliases":["shared"]} + }"#, + Provenance::default(), + ) + .unwrap(); + + assert_eq!(forward.lookup("shared").unwrap().canonical_key, "Alpha"); + assert_eq!(reversed.lookup("shared").unwrap().canonical_key, "Beta"); + assert_eq!(forward.lookup("foo").unwrap().canonical_key, "fOO"); + assert_eq!(reversed.lookup("foo").unwrap().canonical_key, "Foo"); +} + #[derive(Debug)] enum ValidationOutcome { Ok, @@ -128,36 +161,37 @@ enum ValidationOutcome { #[rstest] #[case( IntegrityLimits { - backup_model_count: 2, + reference_model_count: 2, min_model_count: 1, - min_backup_ratio: 0.5, + min_reference_ratio: 0.5, }, ValidationOutcome::Ok )] #[case( IntegrityLimits { - backup_model_count: 3, + reference_model_count: 3, min_model_count: 1, - min_backup_ratio: 0.5, + min_reference_ratio: 0.5, }, ValidationOutcome::Shrunk )] #[case( IntegrityLimits { - backup_model_count: 0, + reference_model_count: 0, min_model_count: 2, - min_backup_ratio: 0.5, + min_reference_ratio: 0.5, }, ValidationOutcome::BelowMinimum )] #[case( IntegrityLimits { - backup_model_count: 0, + reference_model_count: 0, min_model_count: 0, - min_backup_ratio: f64::NAN, + min_reference_ratio: f64::NAN, }, ValidationOutcome::InvalidRatio )] +#[ignore] fn integrity_uses_canonical_count_and_strict_shrink_boundary( #[case] limits: IntegrityLimits, #[case] expected: ValidationOutcome, @@ -195,6 +229,7 @@ enum MalformedOutcome { br#"{"fallback_generalizations":{},"a":{}}"#, MalformedOutcome::Json )] +#[ignore] fn malformed_input_and_aliases_have_typed_outcomes( #[case] body: &[u8], #[case] expected: MalformedOutcome, @@ -210,6 +245,7 @@ fn malformed_input_and_aliases_have_typed_outcomes( } #[rstest] +#[ignore] fn invalid_aliases_are_reported_not_fatal() { let catalog = Catalog::parse( br#"{"a":{"aliases":"bad"},"b":{"aliases":[9,"ok"]}}"#, @@ -228,7 +264,8 @@ fn invalid_aliases_are_reported_not_fatal() { } #[rstest] -fn parses_current_and_packaged_catalogs_without_pinning_counts( +#[ignore] +fn parses_current_and_packaged_catalogs_against_independent_baseline( current_catalog: Catalog, backup_catalog: Catalog, ) { @@ -236,18 +273,17 @@ fn parses_current_and_packaged_catalogs_without_pinning_counts( assert!(backup_catalog.model_count() > 0); assert!(current_catalog.sample_spec().is_some()); assert!(backup_catalog.sample_spec().is_some()); - assert!( - current_catalog - .validate(IntegrityLimits::python_defaults( - backup_catalog.model_count() - )) - .is_ok() - ); - for name in current_catalog.model_names() { + // Snapshot from 2026-09-23; the backup file mirrors the current file and cannot detect shrinkage. + const REFERENCE_MODEL_COUNT: usize = 4303; + current_catalog + .validate(IntegrityLimits { + reference_model_count: REFERENCE_MODEL_COUNT, + min_model_count: 50, + min_reference_ratio: 0.9, + }) + .unwrap(); + assert!(current_catalog.model_names().all(|name| { let entry = current_catalog.lookup(name).unwrap().entry; - assert_eq!( - entry.info().litellm_provider.is_some(), - entry.field("litellm_provider").is_some() - ); - } + entry.info().litellm_provider.is_some() == entry.field("litellm_provider").is_some() + })); } diff --git a/litellm-rust/crates/model-catalog/tests/registry_validation.rs b/litellm-rust/crates/model-catalog/tests/registry_validation.rs new file mode 100644 index 00000000000..8f555e1f884 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/registry_validation.rs @@ -0,0 +1,76 @@ +use std::path::{Path, PathBuf}; + +use litellm_model_catalog::{ + Catalog, FallbackGeneralizations, Provenance, validate_model_entry, validate_registry, +}; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value}; + +#[fixture] +fn repo_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..") +} + +#[rstest] +#[case("model_prices_and_context_window.json")] +#[case("litellm/model_prices_and_context_window_backup.json")] +#[ignore] +fn checked_in_registry_passes_strict_validation(repo_root: PathBuf, #[case] filename: &str) { + let body = std::fs::read(repo_root.join(filename)).unwrap(); + let catalog = Catalog::parse(&body, Provenance::default()).unwrap(); + validate_registry(&catalog).unwrap(); +} + +#[rstest] +#[ignore] +fn fallback_generalizations_are_typed(repo_root: PathBuf) { + let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap(); + let document: Map = serde_json::from_slice(&body).unwrap(); + let Some(raw_rules) = document.get("fallback_generalizations") else { + return; + }; + let _: FallbackGeneralizations = serde_json::from_value(raw_rules.clone()).unwrap(); + let catalog = Catalog::parse(&body, Provenance::default()).unwrap(); + assert!( + catalog + .fallback_rules() + .is_some_and(|rules| !rules.is_empty()) + ); +} + +#[rstest] +#[case::missing_provider(serde_json::json!({"mode": "chat"}), "litellm_provider")] +#[case::unknown_field(serde_json::json!({"litellm_provider": "test", "typo": true}), "unknown field")] +#[case::negative_price(serde_json::json!({"litellm_provider": "test", "input_cost_per_token": -1}), "nonnegative")] +#[case::negative_nested_price(serde_json::json!({"litellm_provider": "test", "guardrail_cost_per_unit": {"unit": -1}}), "nonnegative")] +#[case::invalid_mode(serde_json::json!({"litellm_provider": "test", "mode": "invalid"}), "unknown variant")] +#[case::invalid_date(serde_json::json!({"litellm_provider": "test", "deprecation_date": "2026-02-31"}), "deprecation_date")] +#[case::invalid_hours(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"hours_utc": "25:00-01:00"}}), "hours_utc")] +#[case::empty_windows(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}), "windows is empty")] +#[case::invalid_weekday(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}), "weekdays is invalid")] +#[case::invalid_aliases(serde_json::json!({"litellm_provider": "test", "aliases": ["good", 7]}), "aliases must contain strings")] +#[case::null_aliases(serde_json::json!({"litellm_provider": "test", "aliases": null}), "aliases must be an array")] +#[ignore] +fn registry_validation_rejects_malformed_entries(#[case] entry: Value, #[case] expected: &str) { + assert!( + validate_model_entry("test", &entry) + .unwrap_err() + .to_string() + .contains(expected) + ); +} + +#[test] +#[ignore] +fn checked_in_catalog_and_backup_match() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); + let current = std::fs::read(root.join("model_prices_and_context_window.json")).unwrap(); + let backup = + std::fs::read(root.join("litellm/model_prices_and_context_window_backup.json")).unwrap(); + assert_eq!(current, backup); + let catalog = Catalog::parse(¤t, Provenance::default()).unwrap(); + assert!( + catalog.alias_issues().is_empty(), + "invalid registry aliases" + ); +} diff --git a/litellm-rust/crates/model-catalog/tests/schema.rs b/litellm-rust/crates/model-catalog/tests/schema.rs new file mode 100644 index 00000000000..f9a028a78ab --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/schema.rs @@ -0,0 +1,121 @@ +#![cfg(feature = "schema")] + +use std::collections::BTreeSet; +use std::path::Path; + +use litellm_model_catalog::{model_entry_json_schema, registry_json_schema}; +use rstest::rstest; +use serde_json::{Value, json}; + +fn schema() -> Value { + serde_json::to_value(model_entry_json_schema()).expect("generated schema serializes") +} + +fn registry_validator() -> jsonschema::Validator { + let schema = serde_json::to_value(registry_json_schema()).unwrap(); + jsonschema::options() + .should_validate_formats(true) + .build(&schema) + .expect("generated registry schema is valid") +} + +#[rstest] +#[case("model_prices_and_context_window.json")] +#[case("litellm/model_prices_and_context_window_backup.json")] +#[ignore] +fn generated_registry_schema_validates_checked_in_catalog(#[case] path: &str) { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); + let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap(); + let validator = registry_validator(); + let errors: Vec<_> = validator + .iter_errors(&catalog) + .map(|error| error.to_string()) + .collect(); + assert!(errors.is_empty(), "{path}: {errors:?}"); +} + +#[rstest] +#[case(json!({"example": {"litellm_provider": "test"}}))] +#[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))] +#[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))] +#[ignore] +fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) { + assert!(registry_validator().is_valid(&document)); +} + +#[rstest] +#[case::missing_provider(json!({"mode": "chat"}))] +#[case::negative_cost(json!({"litellm_provider": "test", "input_cost_per_token": -1}))] +#[case::negative_guardrail_cost(json!({"litellm_provider": "test", "guardrail_cost_per_unit": {"unit": -1}}))] +#[case::negative_search_cost(json!({"litellm_provider": "test", "search_context_cost_per_query": {"search_context_size_low": -1}}))] +#[case::negative_tier_cost(json!({"litellm_provider": "test", "tiered_pricing": [{"input_cost_per_token": -1}]}))] +#[case::negative_tier_range(json!({"litellm_provider": "test", "tiered_pricing": [{"range": [-1, 2]}]}))] +#[case::low_uplift(json!({"litellm_provider": "test", "regional_endpoint_uplift_multiplier": 0.5}))] +#[case::nullable_cost(json!({"litellm_provider": "test", "input_cost_per_token": null}))] +#[case::invalid_mode(json!({"litellm_provider": "test", "mode": "telepathy"}))] +#[case::invalid_date(json!({"litellm_provider": "test", "deprecation_date": "2026-02-31"}))] +#[case::invalid_hours(json!({"litellm_provider": "test", "off_peak_pricing": {"hours_utc": "25:00-01:00"}}))] +#[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))] +#[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))] +#[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))] +#[case::non_object_model(json!(4))] +#[ignore] +fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) { + assert!(!registry_validator().is_valid(&json!({"example": entry}))); +} + +#[rstest] +#[case("model_prices_and_context_window.json")] +#[case("litellm/model_prices_and_context_window_backup.json")] +#[ignore] +fn generated_schema_covers_catalog_fields(#[case] path: &str) { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../.."); + let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap(); + let schema = schema(); + let properties = schema["properties"] + .as_object() + .expect("ModelInfo schema has properties"); + let fields: BTreeSet<&str> = catalog + .as_object() + .expect("catalog is an object") + .iter() + .filter(|(name, _)| *name != "sample_spec" && *name != "fallback_generalizations") + .flat_map(|(_, entry)| entry.as_object().expect("model entry is an object").keys()) + .map(String::as_str) + .filter(|name| *name != "aliases") + .collect(); + let missing: Vec<_> = fields + .into_iter() + .filter(|name| !properties.contains_key(*name)) + .collect(); + + assert!( + missing.is_empty(), + "{path}: fields missing from schema: {missing:?}" + ); +} + +#[rstest] +#[case("Mode", "chat")] +#[case("ReasoningEffort", "high")] +#[case("InputModality", "image")] +#[ignore] +fn generated_schema_includes_enum_values(#[case] definition: &str, #[case] value: &str) { + let schema = schema(); + let variants = schema["$defs"][definition]["enum"] + .as_array() + .expect("enum definition has variants"); + + assert!(variants.iter().any(|variant| variant == value)); +} + +#[test] +#[ignore] +fn generated_schema_includes_nested_pricing_types() { + let schema = schema(); + let definitions = schema["$defs"].as_object().expect("schema has definitions"); + + assert!(definitions.contains_key("OffPeakPricing")); + assert!(definitions.contains_key("TieredRate")); + assert!(definitions.contains_key("UtcHours")); +} diff --git a/litellm-rust/crates/model-catalog/tests/spec_parity.rs b/litellm-rust/crates/model-catalog/tests/spec_parity.rs deleted file mode 100644 index 7296d96f798..00000000000 --- a/litellm-rust/crates/model-catalog/tests/spec_parity.rs +++ /dev/null @@ -1,121 +0,0 @@ -use std::collections::{BTreeSet, HashSet}; -use std::path::{Path, PathBuf}; - -use indexmap::IndexMap; -use litellm_model_catalog::{ - Catalog, FallbackGeneralizations, ModelInfo, Provenance, model_entry_json_schema, -}; -use rstest::{fixture, rstest}; -use serde_json::{Map, Value}; - -#[fixture] -fn repo_root() -> PathBuf { - Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..") -} - -fn json_eq(left: &Value, right: &Value) -> bool { - match (left, right) { - (Value::Number(left), Value::Number(right)) => left.as_f64() == right.as_f64(), - (Value::Array(left), Value::Array(right)) => { - left.len() == right.len() && left.iter().zip(right).all(|(a, b)| json_eq(a, b)) - } - (Value::Object(left), Value::Object(right)) => { - left.len() == right.len() - && left - .iter() - .all(|(key, value)| right.get(key).is_some_and(|other| json_eq(value, other))) - } - _ => left == right, - } -} - -fn keys(value: &Map) -> BTreeSet { - value.keys().cloned().collect() -} - -fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> BTreeSet { - left.symmetric_difference(right).cloned().collect() -} - -#[rstest] -#[case("model_prices_and_context_window.json")] -#[case("litellm/model_prices_and_context_window_backup.json")] -fn every_entry_round_trips_through_model_info(repo_root: PathBuf, #[case] filename: &str) { - let body = std::fs::read(repo_root.join(filename)).unwrap(); - let document: IndexMap = serde_json::from_slice(&body).unwrap(); - for (model_name, value) in document { - if matches!( - model_name.as_str(), - "sample_spec" | "fallback_generalizations" - ) { - continue; - } - let object = value - .as_object() - .unwrap_or_else(|| panic!("{model_name} is not an object")); - let info: ModelInfo = serde_json::from_value(value.clone()) - .unwrap_or_else(|error| panic!("{model_name} does not deserialize: {error}")); - let serialized = serde_json::to_value(info).unwrap(); - let serialized_object = serialized - .as_object() - .unwrap_or_else(|| panic!("{model_name} did not serialize as an object")); - let mut expected = object.clone(); - expected.remove("aliases"); - let expected_keys = keys(&expected); - let serialized_keys = keys(serialized_object); - assert_eq!( - expected_keys, - serialized_keys, - "{model_name} key difference: {:?}", - symmetric_difference(&expected_keys, &serialized_keys) - ); - assert!( - json_eq(&Value::Object(expected), &serialized), - "{model_name} changed during ModelInfo round-trip" - ); - } -} - -#[rstest] -fn fallback_generalizations_are_typed(repo_root: PathBuf) { - let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap(); - let document: Map = serde_json::from_slice(&body).unwrap(); - let Some(raw_rules) = document.get("fallback_generalizations") else { - return; - }; - let _: FallbackGeneralizations = serde_json::from_value(raw_rules.clone()).unwrap(); - let catalog = Catalog::parse(&body, Provenance::default()).unwrap(); - assert!( - catalog - .fallback_rules() - .is_some_and(|rules| !rules.is_empty()) - ); -} - -#[rstest] -fn generated_schema_properties_match_repo_schema(repo_root: PathBuf) { - let body = - std::fs::read(repo_root.join("model_prices_and_context_window.schema.json")).unwrap(); - let document: Value = serde_json::from_slice(&body).unwrap(); - let repo_entry_properties = document["$defs"]["modelEntry"]["properties"] - .as_object() - .unwrap(); - let generated = serde_json::to_value(model_entry_json_schema()).unwrap(); - let generated_properties = generated["properties"].as_object().unwrap(); - let expected = keys(repo_entry_properties); - let actual = keys(generated_properties); - assert_eq!( - expected, - actual, - "modelEntry property difference: {:?}", - symmetric_difference(&expected, &actual) - ); - - let repo_root_properties = document["properties"].as_object().unwrap(); - let actual_root: HashSet = repo_root_properties.keys().cloned().collect(); - let expected_root: HashSet = ["sample_spec", "fallback_generalizations"] - .into_iter() - .map(str::to_owned) - .collect(); - assert_eq!(actual_root, expected_root); -} From 10413796c651aa4369a9db88c244dbb217ce6849 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:18:36 +0000 Subject: [PATCH 04/58] test(rust): reorganize core crate tests and split cache and OCR suites (#43177) * test(rust): group cache tests under cache/ and fold test_ocr.py into ocr/ The two failure cases in test_ocr.py duplicated the upstream-500 and timeout rows of PUBLIC_FAILURES, so only the file-input encoding case moves to ocr/test_requests.py Co-Authored-By: Claude Opus 5.5 * test(rust): split the response cache suite into one file per backend test_response_cache.py grew to 2400 lines. Each backend now has its own file, shared fixtures live in cache/conftest.py and shared helpers in support/cache.py. The helpers alias the private native test handles once, dropping the per-call reportPrivateUsage hits Co-Authored-By: Claude Opus 5.5 * split tokenizer test * test(core): consolidate route integration tests under tests/ with rstest and wiremock Moves the public-API OCR route tests out of src/ocr/route.rs and document.rs into tests/ocr/, split per provider plus lifecycle, machine, and document tests, merging the duplicated pairs. Messages, audio transcription, and chat completions share one wiremock-based upstream and recording secret source in tests/support, and gain table-driven cases for auth, routing, upstream errors, streaming, and declines. Tests of litellm-llms items move to that crate. Co-Authored-By: Claude Opus 5.5 * test(messages): keep the stream relay test independent of the stream head contents The stream head carries no headers on main, so the relay test asserts the open-then-deliver order and the relayed body instead of header hand-off. Co-Authored-By: Claude Sonnet 5 --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/chat_completions/handler.rs | 26 + .../core/src/chat_completions/prepare.rs | 244 -- .../crates/core/src/messages/common_utils.rs | 145 +- litellm-rust/crates/core/src/ocr/document.rs | 157 - litellm-rust/crates/core/src/ocr/mod.rs | 209 - litellm-rust/crates/core/src/ocr/prepare.rs | 205 +- litellm-rust/crates/core/src/ocr/route.rs | 3521 ----------------- .../crates/core/tests/audio_transcription.rs | 268 +- .../crates/core/tests/chat_completions.rs | 320 ++ litellm-rust/crates/core/tests/messages.rs | 471 --- .../crates/core/tests/messages/main.rs | 94 + .../crates/core/tests/messages/request.rs | 269 ++ .../crates/core/tests/messages/response.rs | 136 + .../crates/core/tests/messages/secrets.rs | 94 + .../crates/core/tests/messages/stream.rs | 163 + .../crates/core/tests/ocr/aws_textract.rs | 173 + .../crates/core/tests/ocr/azure_ai.rs | 270 ++ .../tests/ocr/azure_document_intelligence.rs | 441 +++ litellm-rust/crates/core/tests/ocr/cohere.rs | 42 + .../crates/core/tests/ocr/documents.rs | 182 + .../crates/core/tests/ocr/lifecycle.rs | 269 ++ litellm-rust/crates/core/tests/ocr/machine.rs | 284 ++ litellm-rust/crates/core/tests/ocr/main.rs | 125 + litellm-rust/crates/core/tests/ocr/mistral.rs | 248 ++ litellm-rust/crates/core/tests/ocr/reducto.rs | 321 ++ .../crates/core/tests/ocr/vertex_ai.rs | 184 + litellm-rust/crates/core/tests/support/mod.rs | 155 + .../llms/src/reducto/ocr/transformation.rs | 44 + litellm-rust/crates/llms/tests/ocr_handler.rs | 79 + tests/test_litellm_rust/AGENTS.md | 1 + tests/test_litellm_rust/cache/__init__.py | 1 + tests/test_litellm_rust/cache/conftest.py | 30 + .../cache/test_azure_blob.py | 173 + tests/test_litellm_rust/cache/test_disk.py | 117 + tests/test_litellm_rust/cache/test_facade.py | 397 ++ tests/test_litellm_rust/cache/test_gcs.py | 242 ++ .../cache/test_qdrant_semantic.py | 286 ++ tests/test_litellm_rust/cache/test_redis.py | 228 ++ .../cache/test_redis_semantic.py | 606 +++ tests/test_litellm_rust/cache/test_rollout.py | 264 ++ tests/test_litellm_rust/cache/test_s3.py | 187 + .../test_valkey_semantic.py} | 0 tests/test_litellm_rust/ocr/test_requests.py | 15 + tests/test_litellm_rust/support/cache.py | 40 + tests/test_litellm_rust/test_cache.py | 2397 ----------- tests/test_litellm_rust/test_ocr.py | 134 - tests/test_litellm_rust/tokenizer/__init__.py | 0 .../test_fast_count.py} | 61 - .../tokenizer/test_huggingface.py | 24 + .../tokenizer/test_tiktoken.py | 53 + 52 files changed, 7015 insertions(+), 7382 deletions(-) create mode 100644 litellm-rust/crates/core/tests/chat_completions.rs delete mode 100644 litellm-rust/crates/core/tests/messages.rs create mode 100644 litellm-rust/crates/core/tests/messages/main.rs create mode 100644 litellm-rust/crates/core/tests/messages/request.rs create mode 100644 litellm-rust/crates/core/tests/messages/response.rs create mode 100644 litellm-rust/crates/core/tests/messages/secrets.rs create mode 100644 litellm-rust/crates/core/tests/messages/stream.rs create mode 100644 litellm-rust/crates/core/tests/ocr/aws_textract.rs create mode 100644 litellm-rust/crates/core/tests/ocr/azure_ai.rs create mode 100644 litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs create mode 100644 litellm-rust/crates/core/tests/ocr/cohere.rs create mode 100644 litellm-rust/crates/core/tests/ocr/documents.rs create mode 100644 litellm-rust/crates/core/tests/ocr/lifecycle.rs create mode 100644 litellm-rust/crates/core/tests/ocr/machine.rs create mode 100644 litellm-rust/crates/core/tests/ocr/main.rs create mode 100644 litellm-rust/crates/core/tests/ocr/mistral.rs create mode 100644 litellm-rust/crates/core/tests/ocr/reducto.rs create mode 100644 litellm-rust/crates/core/tests/ocr/vertex_ai.rs create mode 100644 litellm-rust/crates/core/tests/support/mod.rs create mode 100644 litellm-rust/crates/llms/tests/ocr_handler.rs create mode 100644 tests/test_litellm_rust/AGENTS.md create mode 100644 tests/test_litellm_rust/cache/__init__.py create mode 100644 tests/test_litellm_rust/cache/conftest.py create mode 100644 tests/test_litellm_rust/cache/test_azure_blob.py create mode 100644 tests/test_litellm_rust/cache/test_disk.py create mode 100644 tests/test_litellm_rust/cache/test_facade.py create mode 100644 tests/test_litellm_rust/cache/test_gcs.py create mode 100644 tests/test_litellm_rust/cache/test_qdrant_semantic.py create mode 100644 tests/test_litellm_rust/cache/test_redis.py create mode 100644 tests/test_litellm_rust/cache/test_redis_semantic.py create mode 100644 tests/test_litellm_rust/cache/test_rollout.py create mode 100644 tests/test_litellm_rust/cache/test_s3.py rename tests/test_litellm_rust/{test_valkey_semantic_cache_native.py => cache/test_valkey_semantic.py} (100%) create mode 100644 tests/test_litellm_rust/support/cache.py delete mode 100644 tests/test_litellm_rust/test_cache.py delete mode 100644 tests/test_litellm_rust/test_ocr.py create mode 100644 tests/test_litellm_rust/tokenizer/__init__.py rename tests/test_litellm_rust/{test_tokenizer.py => tokenizer/test_fast_count.py} (51%) create mode 100644 tests/test_litellm_rust/tokenizer/test_huggingface.py create mode 100644 tests/test_litellm_rust/tokenizer/test_tiktoken.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5de32f5b622..2d96efa6077 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3092,6 +3092,7 @@ dependencies = [ "tokio-tungstenite", "url", "veil", + "wiremock", ] [[package]] diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index d7096cdd774..12410c187e2 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -40,3 +40,4 @@ litellm-auth-gcp.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true rstest_reuse.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index de926c715d5..2391ab83a60 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -85,3 +85,29 @@ pub(super) async fn outbound_request( other => other, }) } + +#[cfg(test)] +mod tests { + use super::{Error, as_response_error}; + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + for original in [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth(litellm_auth::Error::InvalidHeader), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), Error::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + let upstream = Error::Transport(litellm_http::transport::Error::Http { + status: 500, + body: "boom".to_string(), + }); + assert_eq!(as_response_error(upstream.clone()), upstream); + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index afea46221f5..b6425773964 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -736,248 +736,4 @@ mod tests { .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); } } - - mod round_trip { - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; - - use super::*; - use crate::chat_completions::chat_completions; - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") - { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - fn http_response(status: &str, body: &str) -> String { - format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - } - - /// Serve one request from a stub upstream and hand back what it received. - async fn serve_once( - status: &'static str, - body: &'static str, - ) -> (String, tokio::task::JoinHandle) { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let port = listener.local_addr().expect("addr").port(); - let handle = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts"); - let received = read_http_request(&mut socket).await; - socket - .write_all(http_response(status, body).as_bytes()) - .await - .expect("writes response"); - socket.flush().await.expect("flushes"); - received - }); - (format!("http://127.0.0.1:{port}/v1/messages"), handle) - } - - fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { - ChatCompletionsRequest { - model: "anthropic/claude-sonnet-4-5", - messages, - optional_params: match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: Some(api_base), - custom_llm_provider: None, - extra_headers: None, - timeout: Some(std::time::Duration::from_secs(10)), - } - } - - const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; - - #[tokio::test] - async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { - let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; - let response = chat_completions(call( - &api_base, - json!([ - {"role": "system", "content": "be terse"}, - {"role": "user", "content": "hi"} - ]), - json!({"max_tokens": 16}), - )) - .await - .expect("call succeeds"); - - let received = handle.await.expect("server task"); - let sent: Value = serde_json::from_str( - received - .split_once("\r\n\r\n") - .expect("request has a body") - .1, - ) - .expect("body is json"); - assert_eq!( - sent["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - sent["system"], - json!([{"type": "text", "text": "be terse"}]) - ); - assert_eq!(sent["max_tokens"], json!(16)); - assert!(received.to_lowercase().contains("x-api-key: sk-test")); - - assert_eq!( - response.choices[0].message.content.as_deref(), - Some("hello") - ); - assert_eq!(response.usage.total_tokens, 15); - } - - #[tokio::test] - async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { - // The provider was called and billed, so the host must not retry this - // on its own path. `MissingField` here would read as a pre-send - // decline and be retried; `InvalidResponse` cannot. - const NO_USAGE: &str = - r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; - let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { - const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; - let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn an_upstream_error_status_keeps_its_code() { - let (api_base, handle) = - serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("upstream rejects"); - handle.await.expect("server task"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "expected a 429, got {err:?}" - ); - } - - #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. - let port = { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - listener.local_addr().expect("has an address").port() - // Dropped here, so the port is closed and the connect is refused. - }; - let err = chat_completions(call( - &format!("http://127.0.0.1:{port}/v1/messages"), - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("nothing is listening"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Connect(_)) - ), - "expected a pre-send connect failure, got {err:?}" - ); - } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string() - })), - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); - } - } } diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 4327754ed05..d27b79bdc04 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -29,151 +29,10 @@ pub(super) fn string_headers( #[cfg(test)] mod tests { - use std::{sync::Arc, time::Duration}; - - use futures_util::future::BoxFuture; - use litellm_secrets::{SecretValue, source::SecretSource}; - use serde_json::{Value, json}; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; + use serde_json::json; use super::{messages_provider_config, string_headers, truncate_error_body}; - use crate::messages::{ - Error, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, - types::MessagesShaping, - }; - - struct RecordingSecrets { - values: Vec<(&'static str, String)>, - requested: std::sync::Mutex>, - } - - impl SecretSource for RecordingSecrets { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - Box::pin(async move { - self.requested.lock().unwrap().push(name.to_string()); - Ok(self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| SecretValue::new(value.clone()))) - }) - } - } - - fn secrets_call() -> MessagesCall { - let Value::Object(body) = json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }) else { - unreachable!("literal object") - }; - MessagesCall { - model: "claude-sonnet-4-5".into(), - body, - api_key: None, - api_base: None, - custom_llm_provider: Some("anthropic".into()), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - } - } - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - #[tokio::test] - async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - request - }); - let secrets = Arc::new(RecordingSecrets { - values: vec![ - ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), - ("ANTHROPIC_BASE_URL", format!("http://{addr}")), - ], - requested: std::sync::Mutex::new(Vec::new()), - }); - - let output = litellm_host::run::run( - messages_machine(secrets.clone()), - &LocalMessagesHost::new(secrets_call()), - ) - .await - .expect("messages request succeeds"); - - assert!(matches!(output, MessagesOutput::Message(_))); - let request = server.await.expect("server task completes"); - assert!( - request - .to_ascii_lowercase() - .contains("x-api-key: sk-from-manager"), - "{request}" - ); - let requested = secrets.requested.lock().unwrap().clone(); - assert_eq!( - requested, - messages_provider_config("anthropic") - .unwrap() - .secret_names() - .iter() - .map(ToString::to_string) - .collect::>() - ); - } + use crate::messages::Error; #[test] fn provider_config_resolves_anthropic_and_azure_ai() { diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index b78c09298de..c33ee053422 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -244,160 +244,3 @@ mod tests { } } } - -#[cfg(test)] -mod document_tests { - use litellm_host::event::WireRequest; - use litellm_llms::base_llm::ocr::error::Error; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, - request_body, wire_request_with_document, - }; - - #[derive(Clone, Copy, Debug)] - enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, - } - - impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } - } - - /// What the host does to the wire request in `before_send`. - #[derive(Clone, Copy, Debug)] - enum Host { - Detached, - ReplacesDocument, - } - - const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - - impl Host { - fn before_send(self, wire: WireRequest) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } - } - - struct Sent { - result: Result<(), Error>, - provider_body: Option, - } - - async fn send(route: Route, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = - mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document_type = route.document_type(); - let document = - json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = - LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - Sent { - result, - provider_body, - } - } - - fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) - } - - #[rstest] - #[case::azure_ai(Route::AzureAi)] - #[case::vertex_mistral(Route::VertexMistral)] - #[case::azure_cohere_parse(Route::AzureCohereParse)] - #[tokio::test] - async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::Detached, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); - } - - #[rstest] - #[tokio::test] - async fn document_replaced_by_the_host_reaches_the_provider( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - ) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::ReplacesDocument, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 270a402c9fa..2a0d20f69c9 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -7,212 +7,3 @@ pub mod provider_config; pub mod route; pub mod types; pub mod wire; - -#[cfg(test)] -pub(crate) mod test_support { - use std::sync::{Arc, Mutex}; - - use futures_util::future::BoxFuture; - use litellm_host::event::WireRequest; - use litellm_llms::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient}, - transformation::LiteLLMOcrResponse, - }; - use serde_json::{Value, json}; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, - }; - - use crate::ocr::{ - route::{LocalOcrHost, ocr_machine}, - types::LiteLLMOcrRequest, - wire::{OcrWireRequest, decode_request}, - }; - - /// Stands in for a host with no hooks registered: the wire request goes out unchanged - /// and response events go nowhere. - pub(crate) struct NoHooks; - - impl CallHooks for NoHooks { - fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(wire) }) - } - - fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(async { Ok(()) }) - } - } - - pub(crate) fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) - } - - pub(crate) async fn perform_ocr( - request: LiteLLMOcrRequest, - ) -> Result { - crate::ocr::client::perform(&ocr_client(), request).await - } - - pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { - litellm_host::run::run(ocr_machine(ocr_client()), &host).await - } - - pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - base, - json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - options, - ) - } - - pub(crate) fn wire_request_with_document( - model: &str, - base: &str, - document: Value, - options: Value, - ) -> LiteLLMOcrRequest { - decode_request(OcrWireRequest { - model: model.into(), - document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), - api_base: Some(base.into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap() - } - - pub(crate) fn resolved_request( - request: LiteLLMOcrRequest, - ) -> crate::ocr::types::ResolvedOcrRequest { - request - .map_document(crate::ocr::document::prepare_document) - .unwrap() - } - - pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { - let request = resolved_request(request); - let document = request.document.clone().with_source(source.into()); - request.with_document(document.into()) - } - - pub(crate) fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; - - /// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. - pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - loop { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut buffer = [0u8; 4096]; - let _ = socket.read(&mut buffer).await.unwrap(); - let head = format!( - "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - SERVED_DOCUMENT.len() - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(SERVED_DOCUMENT).await.unwrap(); - } - }); - (base, task) - } - - pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, - } - - impl MockResponse { - pub fn json(body: Value) -> Self { - Self { - status: 200, - headers: vec![], - body, - } - } - } - - pub(crate) async fn mock_server( - responses: Vec, - ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let requests = Arc::new(Mutex::new(Vec::new())); - let seen = requests.clone(); - let server_base = base.clone(); - let task = tokio::spawn(async move { - for response in responses { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0u8; 4096]; - let header_end = loop { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { - break index + 4; - } - }; - let length = String::from_utf8_lossy(&bytes[..header_end]) - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap_or(0); - while bytes.len() < header_end + length { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - } - seen.lock() - .unwrap() - .push(String::from_utf8_lossy(&bytes).into_owned()); - let body = serde_json::to_vec(&response.body).unwrap(); - let headers = response - .headers - .into_iter() - .map(|(name, value)| { - format!("{name}: {}\r\n", value.replace("{base}", &server_base)) - }) - .collect::(); - let head = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", - response.status, - body.len(), - headers - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(&body).await.unwrap(); - } - }); - (base, requests, task) - } - - pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { - request - .lines() - .take_while(|line| !line.is_empty()) - .find_map(|line| { - let (key, value) = line.split_once(':')?; - key.eq_ignore_ascii_case(name).then(|| value.trim()) - }) - } -} diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 37c0f18f659..18961ec96fa 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -70,20 +70,203 @@ pub(crate) fn prepare_request( } } -#[cfg(test)] -pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest { - prepare_request( - request, - true, - &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), - std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), - ) -} - #[cfg(test)] mod tests { + use std::time::Duration; + + use futures_util::future::BoxFuture; use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options}; - use serde_json::json; + use litellm_host::event::WireRequest; + use litellm_llms::{ + base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::{BaseOcrConfig, OcrResponseFormat}, + }, + cohere::ocr::transformation::CohereParseConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + use serde_json::{Value, json}; + + use super::*; + use crate::ocr::{ + document::prepare_document, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, + }; + + /// Stands in for a host with no hooks registered. + struct NoHooks; + + impl CallHooks for NoHooks { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(wire) }) + } + + fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + } + + fn client() -> OcrClient { + OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()) + } + + fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() + } + + fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest { + prepare_request( + request.map_document(prepare_document).unwrap(), + true, + &client(), + std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), + ) + } + + fn image(url: &str) -> Value { + json!({"type": "image_url", "image_url": url}) + } + + #[tokio::test] + async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() { + let request = request( + "cohere/parse", + "https://example.com", + image("https://example.com/original.png"), + json!({ + "output_format": "markdown", "timeout": 30, + "extra_body": { + "output_format": {"future": true}, + "document": {"type": "image_url", "image_url": "https://example.com/a.png", + "provider_options": {"nested": [false, 0, null]}} + } + }), + ); + + let http = CohereParseConfig + .prepare_request(&prepared(request), &client(), &NoHooks) + .await + .unwrap(); + + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "parse", "output_format": {"future": true}, + "document": {"type": "image_url", "image_url": "https://example.com/a.png", + "provider_options": {"nested": [false, 0, null]}} + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = request( + "cohere/parse", + "https://example.com", + image("https://example.com/a.png"), + json!({"output_format": null, "req_format": null}), + ); + assert_eq!( + request.response_format().unwrap(), + OcrResponseFormat::Litellm + ); + + let http = CohereParseConfig + .prepare_request(&prepared(request), &client(), &NoHooks) + .await + .unwrap(); + + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[tokio::test] + async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() { + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let document = + json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}); + let direct = prepared(request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + document.clone(), + options.clone(), + )); + let vertex = prepared(request( + "vertex_ai/mistral-ocr-maas", + "https://vertex.test", + document, + options, + )); + + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client(), &NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client(), &NoHooks) + .await + .unwrap(); + + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "preserved" + }) + ); + } + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } #[derive(serde::Deserialize)] struct KnownParams { diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index e3e57bd2d77..adb704a15d1 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -161,3524 +161,3 @@ impl litellm_host::host::Host for LocalOcrHost { Ok(()) } } - -#[cfg(test)] -mod aws_textract_tests { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; - use litellm_llms::base_llm::ocr::error::Error; - use serde_json::{Value, json}; - use time::{PrimitiveDateTime, format_description}; - - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, header, mock_server, perform_ocr_with, request_body, - wire_request_with_document, - }, - types::LiteLLMOcrRequest, - }; - - const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; - const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; - - fn textract_request(base: &str) -> LiteLLMOcrRequest { - textract_request_for("aws_textract/detect-document-text", base) - } - - fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - &format!("{base}/"), - json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), - json!({ - "aws_access_key_id": ACCESS_KEY_ID, - "aws_secret_access_key": SECRET_ACCESS_KEY, - "aws_region_name": "eu-west-1" - }), - ) - } - - fn textract_response() -> MockResponse { - MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] - })) - } - - /// Recomputes SigV4 over the bytes the server received, at the time the client claimed. - fn expected_authorization(url: &str, raw_request: &str) -> String { - let format = - format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") - .unwrap(); - let signed_at: SystemTime = - PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) - .unwrap() - .assume_utc() - .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] - .into_iter() - .map(|name| { - ( - name.to_string(), - header(raw_request, name).unwrap().to_string(), - ) - }) - .collect(); - let body = raw_request.split_once("\r\n\r\n").unwrap().1; - sign_post( - url, - body.as_bytes(), - &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, - ) - .unwrap()["Authorization"] - .clone() - } - - #[tokio::test] - async fn the_request_is_signed_for_textract_and_lines_become_the_page() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - - let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.DetectDocumentText") - ); - assert_eq!( - header(&raw, "content-type"), - Some("application/x-amz-json-1.1") - ); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "Invoice 12345"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); - } - - #[tokio::test] - async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { - assert!( - !wire - .headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), - "the hook ran after signing" - ); - wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - } - - #[tokio::test] - async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { - let (base, _, server) = mock_server(vec![MockResponse { - status: 400, - headers: vec![], - body: json!({ - "__type": "UnsupportedDocumentException", - "Message": "Request has unsupported document format" - }), - }]) - .await; - - let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap_err(); - server.await.unwrap(); - - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); - assert!( - body.contains("multi-page documents are not supported"), - "{body}" - ); - } - - #[tokio::test] - async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [ - {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, - {"Id": "t", "BlockType": "LAYOUT_TITLE", - "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} - ] - }))]) - .await; - let request = textract_request_for("aws_textract/analyze-document", &base); - - let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.AnalyzeDocument") - ); - assert_eq!( - request_body(&raw)["FeatureTypes"], - json!(["LAYOUT", "TABLES"]) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "# Quarterly Report"); - } -} - -#[cfg(test)] -mod azure_ai_tests { - use litellm_llms::base_llm::ocr::error::Error; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }; - - #[tokio::test] - async fn facade_executes_azure_mistral_with_prepared_auth() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"include_image_base64":true}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![( - "Authorization".into(), - "Bearer python-prepared-token".into(), - )]; - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "include_image_base64":true - }) - ); - } - - #[tokio::test] - async fn facade_acquires_supplied_entra_token_for_final_request() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"azure_ad_token":"rust-owned-token"}), - ); - request.credentials.api_key = None; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer rust-owned-token\r\n") - ); - } - - #[tokio::test] - async fn rejects_non_inline_body_after_guardrails() { - let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { - wire.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(wire) - }); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); - } - - mod transformation { - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, - }; - use rstest::rstest; - use serde_json::json; - - use super::*; - use crate::ocr::{ - test_support::{MockResponse, header, mock_server, perform_ocr}, - types::LiteLLMOcrRequest, - wire::decode_request, - }; - - #[derive(Debug)] - struct CountingToken { - token: fn(usize) -> String, - calls: AtomicUsize, - } - - impl CountingToken { - fn new(token: fn(usize) -> String) -> Arc { - Arc::new(Self { - token, - calls: AtomicUsize::new(0), - }) - } - - fn calls(&self) -> usize { - self.calls.load(Ordering::SeqCst) - } - } - - impl TokenProvider for CountingToken { - fn acquire(&self) -> TokenFuture<'_> { - let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - let token = SecretValue::new((self.token)(call)); - Box::pin(async move { - Ok(ResolvedCredential::AccessToken { - token, - expires_on: None, - }) - }) - } - } - - fn numbered_token(call: usize) -> String { - format!("callback-{call}") - } - - fn azure_request( - provider: &Arc, - api_base: Option<&str>, - api_key: Option<&str>, - extra_headers: Value, - optional_params: Value, - ) -> LiteLLMOcrRequest { - let wire = serde_json::from_value(json!({ - "model": "azure_ai/mistral-ocr-latest", - "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": null, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": 2.0 - })) - .unwrap(); - LiteLLMOcrRequest { - azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), - ..decode_request(wire).unwrap() - } - } - - fn ocr_page() -> MockResponse { - MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) - } - - #[tokio::test] - async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; - - for _ in 0..2 { - perform_ocr(azure_request( - &provider, - Some(&base), - None, - Value::Null, - json!({}), - )) - .await - .unwrap(); - } - server.await.unwrap(); - - assert_eq!(provider.calls(), 2); - let requests = seen.lock().unwrap(); - assert_eq!( - requests - .iter() - .map(|request| header(request, "authorization")) - .collect::>(), - [Some("Bearer callback-1"), Some("Bearer callback-2")] - ); - } - - #[rstest] - #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] - #[case::provider_beats_static_token( - None, - Value::Null, - json!({"azure_ad_token":"static-token"}), - "Bearer callback-1", - 1 - )] - #[case::header_wins_on_the_wire_but_provider_still_runs( - None, - json!({"Authorization":"Bearer override"}), - json!({}), - "Bearer override", - 1 - )] - #[tokio::test] - async fn credential_precedence( - #[case] api_key: Option<&str>, - #[case] extra_headers: Value, - #[case] optional_params: Value, - #[case] expected_authorization: &str, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - perform_ocr(azure_request( - &provider, - Some(&base), - api_key, - extra_headers, - optional_params, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(provider.calls(), expected_calls); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!( - header(&requests[0], "authorization"), - Some(expected_authorization) - ); - } - - #[rstest] - #[case::missing_api_base( - false, - json!({}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - environment_variable: "AZURE_AI_API_BASE", - })), - 0 - )] - #[case::unsupported_oidc_reference( - true, - json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), - 0 - )] - #[case::empty_provider_token_ignores_static_token( - true, - json!({"azure_ad_token":"static-token"}), - |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), - 1 - )] - #[tokio::test] - async fn credential_failures_send_no_provider_request( - #[case] with_api_base: bool, - #[case] optional_params: Value, - #[case] token: fn(usize) -> String, - #[case] expected: fn(&Error) -> bool, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - let error = perform_ocr(azure_request( - &provider, - with_api_base.then_some(base.as_str()), - None, - Value::Null, - optional_params, - )) - .await - .unwrap_err(); - server.abort(); - - assert!(expected(&error), "unexpected error: {error:?}"); - assert_eq!(provider.calls(), expected_calls); - assert!(seen.lock().unwrap().is_empty()); - } - } -} - -#[cfg(test)] -mod azure_document_intelligence_tests { - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, - }; - - fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) - } - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), - ); - request.document = serde_json::from_value::< - litellm_llms::base_llm::ocr::transformation::OcrDocument, - >(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf"}) - ); - } - - #[rstest] - #[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] - #[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] - #[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] - #[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] - #[case(json!({"features":"languages&pages=1"}), Error::Features)] - #[case(json!({"req_format":"azure"}), Error::RequestFormat)] - #[tokio::test] - async fn rejects_invalid_pages_features_and_format( - #[case] options: Value, - #[case] expected: Error, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: Some(base), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }); - let result = match result { - Ok(request) => perform_ocr(request).await, - Err(error) => Err(error), - }; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid options: {options}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); - } - - #[rstest] - #[case(json!({}))] - #[case(json!({"req_format":"litellm"}))] - #[tokio::test] - async fn missing_native_fields_keep_page_text_without_retaining_raw_response( - #[case] options: Value, - ) { - let operation = json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} - }); - let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; - let response = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - options, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "hello"); - assert_eq!(response.provider_native_response, None); - let serialized = response.into_json(); - assert_eq!(serialized.get("content"), Some(&Value::Null)); - assert_eq!(serialized.get("tables"), Some(&Value::Null)); - assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - let target = requests[0].split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - for field in ["pages", "features", "req_format"] { - assert_eq!(query_value(&url, field), None); - } - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); - } - - #[tokio::test] - async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - } - - #[tokio::test] - async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} - }))]) - .await; - let client = ocr_client().with_settings(OcrSettings { - document_intelligence_api_version: "2099-01-01".into(), - document_intelligence_dpi: 72, - ..OcrSettings::default() - }); - - let result = crate::ocr::client::perform( - &client, - wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - - let target = seen.lock().unwrap()[0] - .split_whitespace() - .nth(1) - .unwrap() - .to_string(); - assert_eq!( - query_value(&format!("{base}{target}"), "api-version").as_deref(), - Some("2099-01-01") - ); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":612,"height":792,"dpi":72}) - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { - return; - }; - match request_count.lock().unwrap().len() { - 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), - 2 => assert!(raw.body.contains("succeeded")), - count => panic!("unexpected callback after {count} requests"), - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[tokio::test] - async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); - } - - #[tokio::test] - async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); - } - - #[tokio::test] - async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); - } - - #[tokio::test] - async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } - } - - #[tokio::test] - async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } - } - - #[tokio::test] - async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - let client = ocr_client().with_settings(OcrSettings { - poll_timeout: std::time::Duration::from_millis(100), - ..OcrSettings::default() - }); - - let error = tokio::time::timeout( - std::time::Duration::from_secs(1), - crate::ocr::client::perform(&client, request), - ) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); - } - - #[tokio::test] - async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } - } - - mod transformation { - use std::sync::{Arc, Mutex}; - - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::transformation::OcrDocument; - use serde_json::{Value, json}; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }, - }; - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); - } - - #[tokio::test] - async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - "http://127.0.0.1:1", - options.clone(), - ); - let rejected = perform_ocr(request).await.is_err(); - assert!(rejected, "accepted {options}"); - } - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let request_count = seen.clone(); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed - .lock() - .unwrap() - .push((request_count.lock().unwrap().len(), raw.body.clone())); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *responses_received.lock().unwrap(), - [ - (1, r#"{"submitted":true}"#.to_string()), - (2, r#"{"status":"succeeded"}"#.to_string()), - ] - ); - } - } -} - -#[cfg(test)] -mod cohere_tests { - mod transformation { - use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, - }, - cohere::ocr::transformation::*, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - #[tokio::test] - async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({ - "output_format":"markdown", "timeout":30, - "extra_body":{ - "output_format": {"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - } - }), - ); - let request = request.with_document( - serde_json::from_value(json!({ - "type":"image_url","image_url":"https://example.com/original.png" - })) - .unwrap(), - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model":"parse", "output_format":{"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - }) - ); - } - - #[tokio::test] - async fn explicit_null_options_use_defaults_before_http() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({"output_format":null,"req_format":null}), - ); - let request = request.with_document( - serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png"}), - ) - .unwrap(), - ); - assert_eq!( - request.response_format().unwrap(), - OcrResponseFormat::Litellm - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!(body["output_format"], "markdown"); - assert!(body.get("req_format").is_none()); - } - - #[rstest] - #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] - #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] - #[tokio::test] - async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( - #[case] model: &str, - #[case] request_line: &str, - ) { - use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; - - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = crate::ocr::test_support::wire_request(model, &base, json!({})) - .with_document( - serde_json::from_value::( - json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), - ) - .unwrap() - .into(), - ); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with(request_line), "{}", requests[0]); - assert_eq!( - header(&requests[0], "authorization"), - Some("Bearer test-key") - ); - } - - #[rstest] - #[tokio::test] - async fn route_rejects_non_image_document_without_a_request( - #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, - ) { - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; - - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - - let error = perform_ocr(crate::ocr::test_support::wire_request( - model, - &base, - json!({}), - )) - .await - .unwrap_err(); - server.abort(); - - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); - assert!(seen.lock().unwrap().is_empty()); - } - } -} - -#[cfg(test)] -mod deepseek_tests { - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, - vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, - normalize_response as transform_ocr_response, - }, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - fn document() -> OcrDocument { - serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() - } - - #[rstest] - #[case("stream", json!(true))] - #[case("temperature", json!(0.1))] - #[case("max_tokens", json!(1024))] - #[case("top_p", json!(0.9))] - #[case("n", json!(2))] - #[case("stop", json!("done"))] - #[case("stop", json!(["done", "stop"]))] - fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); - } - - #[rstest] - #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] - #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] - fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); - } - - #[rstest] - #[case(json!("# hello"), "# hello")] - #[case(json!("{broken"), "{broken")] - #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] - #[case(json!({"pages":[]}), "")] - #[case(json!("[]"), "[]")] - #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] - #[case(json!({"pages":[{"markdown":"object"}]}), "object")] - fn response_codec_handles_text_json_and_objects( - #[case] content: Value, - #[case] expected: &str, - ) { - let structured = content - .as_object() - .is_some_and(|object| object.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if structured { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } - } - - #[test] - fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = transform_ocr_response("requested", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["index"], 2); - assert_eq!(result["pages"][0]["images"][0]["id"], "one"); - assert_eq!(result["model"], "provider-model"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - assert_eq!(result["document_annotation"]["language"], "en"); - assert_eq!(result["future"], "kept"); - } - - #[test] - fn response_codec_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } - } -} - -#[cfg(test)] -mod reducto_tests { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers( - #[case] model: &str, - #[values("application/pdf", "image/png")] mime_type: &str, - ) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let document = if mime_type.starts_with("image/") { - json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) - } else { - json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) - }; - let mut request = crate::ocr::types::LiteLLMOcrRequest { - document: serde_json::from_value::(document) - .unwrap() - .into(), - ..wire_request(&format!("reducto/{model}"), &base, json!({})) - }; - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; - assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); - assert!(multipart.contains("\r\n\r\nabc\r\n--")); - assert!(requests[1].starts_with("POST /parse ")); - let source_field = if model == "parse-legacy" { - "document_url" - } else { - "input" - }; - assert_eq!( - request_body(&requests[1]), - json!({source_field:"reducto://uploaded.pdf"}) - ); - for request in requests.iter() { - assert!( - request - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - } - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case(json!({"file_id":""}))] - #[case(json!({}))] - #[case(json!({"file_id":null}))] - #[tokio::test] - async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[tokio::test] - async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[rstest] - #[case("https://example.com/a.pdf", Error::ReductoSource)] - #[case("reducto://", Error::RequestField { path: "document file id".into() })] - #[case("data:application/pdf;base64", Error::InvalidDataUri)] - #[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network( - #[case] source: &str, - #[case] expected: Error, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - source, - ); - let result = perform_ocr(request).await; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid source: {source}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); - } - - #[test] - fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use litellm_llms::reducto::ocr::transformation::{ - ReductoResponse, normalize_response as transform_ocr_response, - }; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = transform_ocr_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = transform_ocr_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[tokio::test] - async fn native_format_retains_the_provider_response() { - let raw = json!({ - "result":{"chunks":[{"content":"native OCR response"}]}, - "usage":{"num_pages":1} - }); - let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages[0].markdown, "native OCR response"); - assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); - } - - #[tokio::test] - async fn unknown_model_reaches_parse_and_keeps_its_name() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[{"content":"future model response"}]} - }))]) - .await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/future-parse-model", &base, json!({})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.model, "future-parse-model"); - assert_eq!(response.pages[0].markdown, "future model response"); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!( - request_body(&requests[0]), - json!({"input":"reducto://ready.pdf"}) - ); - } - - #[tokio::test] - async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_before_send(|wire, _| { - assert_eq!( - wire.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(WireRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); - } - - mod transformation { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, - reducto::ocr::transformation::*, - }; - use rstest::rstest; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }, - }; - - #[tokio::test] - async fn v3_options_preserve_explicit_null() { - let overrides = - serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) - .unwrap(); - let params = ReductoParseV3Config - .map_ocr_params(&overrides, "parse-v3") - .unwrap(); - let client = crate::ocr::test_support::ocr_client(); - let connection = OcrConnection::default(); - let document = serde_json::from_value( - json!({"type":"document_url","document_url":"reducto://ready.pdf"}), - ) - .unwrap(); - let body = ReductoParseV3Config - .async_transform_ocr_request( - "parse-v3", - document, - ¶ms, - &[], - OcrRequestContext { - client: &client, - connection: &connection, - }, - ) - .await - .unwrap(); - assert_eq!( - serde_json::to_value(body).unwrap(), - json!({ - "input":"reducto://ready.pdf", "formatting":null, "settings":{} - }) - ); - let absent = ReductoParseV3Config - .map_ocr_params( - &litellm_core_utils::call_arguments::CallArguments::default(), - "parse-v3", - ) - .unwrap(); - assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case("https://example.com/a.pdf")] - #[case("reducto://")] - #[case("data:application/pdf;base64")] - #[case("data:application/pdf;base64,INVALID!")] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), - source, - ); - assert!(perform_ocr(request).await.is_err()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = - vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[rstest] - #[case("reducto/parse-v3")] - #[case("reducto/parse-legacy")] - #[tokio::test] - async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.transport.extra_headers = - vec![("authorization".into(), "Bearer original".into())]; - let host = LocalOcrHost::new(request).with_before_send(|wire, _| { - Ok(WireRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } - } - } -} - -#[cfg(test)] -mod vertex_ai_tests { - use litellm_auth::InputSource; - use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; - use serde_json::{Value, json}; - - use crate::ocr::test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, wire_request, - }; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[tokio::test] - async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); - } - - #[tokio::test] - async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let client = ocr_client().with_settings(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - crate::ocr::client::perform( - &client, - wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].starts_with( - "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - } - - #[tokio::test] - async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); - } - - #[tokio::test] - async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); - } - - #[tokio::test] - async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); - } - - #[tokio::test] - async fn adapters_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "ignored" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - for http in [&direct_http, &vertex_http] { - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "ignored" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } - - mod transformation { - - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::test_support::wire_request; - - #[rstest] - #[case::mistral(false)] - #[case::vertex(true)] - #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization( - #[case] use_vertex: bool, - ) { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - let http = if use_vertex { - &vertex_http - } else { - &direct_http - }; - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "preserved" - }) - ); - let payload = serde_json::to_vec( - &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), - ) - .unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &payload, Default::default()) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &payload, Default::default()) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } - } -} - -#[cfg(test)] -mod vertex_ai_deepseek_tests { - use litellm_auth::InputSource; - use serde_json::{Value, json}; - - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } - - #[test] - fn host_registration_selects_deepseek_without_affecting_mistral() { - assert!(crate::ocr::arguments::is_supported_request( - "deepseek-ocr-maas", - Some("vertex_ai") - )); - assert!(crate::ocr::arguments::is_supported_request( - "mistral-ocr-maas", - Some("vertex_ai") - )); - } - - #[tokio::test] - async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/deepseek-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); - } - - mod deepseek_transformation { - use serde_json::json; - - use super::*; - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = - crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } - } -} - -#[cfg(test)] -pub(crate) mod tests { - use std::sync::{Arc, Mutex}; - - use futures_util::future::BoxFuture; - use litellm_auth_gcp::VertexAuth; - use litellm_host::{ - event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp}, - machine::{HostFailure, Machine, MachineStep}, - }; - use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, - }; - use litellm_llms::base_llm::ocr::{ - error::Error as OcrError, - handler::OcrClient, - settings::OcrSettings, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, - }, - }; - use litellm_secrets::source::SecretSource; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::{LocalOcrHost, OcrOp, OcrProjection, ocr_machine}; - use crate::ocr::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, - }; - - struct RecordingSecretSource { - names: Arc>>, - values: &'static [(&'static str, &'static str)], - api_base: String, - } - - impl SecretSource for RecordingSecretSource { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> - { - self.names.lock().unwrap().push(name.to_owned()); - Box::pin(async move { - Ok(match name { - "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), - _ => self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), - } - .map(litellm_secrets::SecretValue::new)) - }) - } - } - - #[rstest] - #[case::mistral("mistral/model", json!({}))] - #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] - #[tokio::test] - async fn ocr_contract_upstream_error_preserves_status_body_and_headers( - #[case] model: &str, - #[case] options: Value, - ) { - let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); - let expected_body = serde_json::to_string(&payload).unwrap(); - let (base, seen, server) = mock_server(vec![MockResponse { - status: 422, - headers: vec![ - ("Retry-After", "17".into()), - ("X-Request-ID", "request-123".into()), - ("X-Future-Header", "retained".into()), - ], - body: payload, - }]) - .await; - let error = perform_ocr(wire_request(model, &base, options)) - .await - .unwrap_err(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - let OcrError::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; - assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) - ); - } - assert_eq!( - body.len(), - expected_body.len(), - "provider error body was truncated" - ); - assert_eq!(body, expected_body); - } - - #[test] - fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: serde_json::Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); - } - - #[tokio::test] - async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":"ignored" - }) - ); - } - - #[tokio::test] - async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = - mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.map(Value::Object), - Some(provider_response) - ); - } - - #[rstest] - #[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] - #[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] - #[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] - #[tokio::test] - async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( - #[case] secrets: &'static [(&'static str, &'static str)], - #[case] expected_key: &str, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: secrets, - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); - } - - #[tokio::test] - async fn mistral_ocr_resolves_provider_secrets_before_transformation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: &[("MISTRAL_API_KEY", "source-key")], - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); - } - - #[tokio::test] - async fn ocr_client_uses_the_injected_http_pool_configuration() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let settings = HttpSettings { - user_agent: Some("host-owned/1".into()), - ..HttpSettings::default() - }; - let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), - ) - .unwrap(); - crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); - } - - fn event_name(event: &CallEvent) -> &'static str { - match event { - CallEvent::Started { .. } => "started", - CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", - CallEvent::Succeeded { .. } => "success", - CallEvent::Failed { .. } => "failure", - } - } - - fn recording_host( - request: crate::ocr::types::LiteLLMOcrRequest, - events: Arc>>, - block: bool, - ) -> LocalOcrHost { - let before_send_events = events.clone(); - LocalOcrHost::new(request) - .with_before_send(move |wire, _| { - before_send_events.lock().unwrap().push("before_send"); - if block { - return Err(OcrError::InvalidRequest("blocked".into())); - } - Ok(wire) - }) - .with_observer(move |event| events.lock().unwrap().push(event_name(event))) - } - - #[tokio::test] - async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) - .with_before_send(|mut wire, _| { - wire.headers - .push(("x-core-callback".into(), "edited".into())); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); - } - - #[tokio::test] - async fn before_send_context_names_the_route_and_its_secrets() { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host = LocalOcrHost::new(wire_request( - "mistral/model", - &base, - json!({"pages": [0], "req_format": "native"}), - )) - .with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let (wire, context) = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.custom_llm_provider, "mistral"); - assert_eq!(context.model, "model"); - assert_eq!(wire.body["pages"], json!([0])); - assert!(context.secret_fields.is_empty()); - assert_eq!(context.optional_params["req_format"], "native"); - - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let request = wire_request( - "azure_ai/model", - &base, - json!({"client_secret": "shh", "tenant_id": "t"}), - ); - let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { - bytes: b"abc".as_slice().into(), - file_name: None, - mime_type: Some("application/pdf".into()), - }); - let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some(context.clone()); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let context = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.secret_fields, ["client_secret"]); - } - - #[tokio::test] - async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "response", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[tokio::test] - async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", "http://127.0.0.1:1", json!({})), - events.clone(), - true, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - } - - #[tokio::test] - async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - assert!(perform_ocr_with(host).await.is_err()); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - /// Drives the machine by hand, answering every op through `host` except `before_send`, - /// which `intercept` answers so a test can fail or cancel exactly there. - async fn drive_until( - client: OcrClient, - host: &LocalOcrHost, - mut intercept: impl FnMut(WireRequest) -> Result>, - ) -> ( - Result, - Vec<&'static str>, - crate::ocr::route::OcrMachine, - ) { - let mut machine = ocr_machine(client); - let mut ops = Vec::new(); - let outcome = loop { - let op = match machine.resume().await { - Ok(MachineStep::Host(op)) => op, - Ok(MachineStep::Complete(response)) => break Ok(response), - Err(error) => break Err(error), - }; - let answer = match op { - HostOp::Project(reply) => { - ops.push("Project"); - host.project() - .await - .map(|projection| reply.send(projection)) - .map_err(HostFailure::Error) - } - HostOp::Custom(op) => { - ops.push(match op { - OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken", - }); - host.custom_op(op).await.map_err(HostFailure::Error) - } - HostOp::BeforeSend { wire, reply, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| reply.send(wire)) - } - HostOp::Emit(event, reply) => { - let event = CallEvent::Machine(event); - ops.push(event_name(&event)); - host.emit(&event) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error) - } - }; - if let Err(failure) = answer { - break machine.interrupt(failure).await; - } - }; - (outcome, ops, machine) - } - - #[tokio::test] - async fn failed_before_send_does_not_replay_or_reach_transport() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Error(OcrError::InvalidRequest( - "before_send failed".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") - ); - assert_eq!(ops, ["Project", "BeforeSend"]); - assert!(machine.resume().await.is_err()); - } - - #[tokio::test] - async fn invalid_provider_response_emits_response_received_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed.lock().unwrap().push(raw.body.clone()); - } - }); - let error = perform_ocr_with(host).await.unwrap_err(); - server.await.unwrap(); - assert!(matches!(error, OcrError::ResponseField { .. })); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - *responses_received.lock().unwrap(), - [r#"{"pages":"invalid"}"#] - ); - } - - #[tokio::test] - async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; - server.await.unwrap(); - assert_eq!(outcome.unwrap().pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["Project", "BeforeSend", "response"]); - assert!(matches!( - machine.resume().await, - Err(OcrError::InvalidRequest(_)) - )); - } - - #[tokio::test] - async fn empty_byte_documents_fail_before_the_provider_is_called() { - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Bytes { - bytes: Default::default(), - file_name: None, - mime_type: None, - }, - ); - let response = perform_ocr_with(LocalOcrHost::new(request)).await; - assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); - assert!(seen.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn path_documents_are_read_by_core_without_a_host_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"path"}] - }))]) - .await; - let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); - std::fs::create_dir_all(&dir).unwrap(); - let path = dir.join("scan.png"); - std::fs::write(&path, b"abc").unwrap(); - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let (response, ops, _) = drive_until(ocr_client(), &LocalOcrHost::new(request), Ok).await; - server.await.unwrap(); - std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(ops, ["Project", "BeforeSend", "response"]); - assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); - - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let response = perform_ocr_with(LocalOcrHost::new(request)).await; - assert!(matches!( - response.unwrap_err(), - OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound - )); - assert!(seen.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Cancelled(OcrError::InvalidRequest( - "cancelled".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); - assert_eq!(ops, ["Project", "BeforeSend"]); - assert!(machine.resume().await.is_err()); - } - - #[tokio::test] - async fn resuming_before_answering_preserves_pending_operation() { - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let mut machine = ocr_machine(ocr_client()); - let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else { - panic!("expected the projection op first"); - }; - assert!(machine.resume().await.is_err()); - reply.send(OcrProjection { - request, - caller_token: false, - }); - assert!(matches!( - machine.resume().await, - Ok(MachineStep::Host(HostOp::BeforeSend { .. })) - )); - } - - async fn read_bounded_response( - response: Vec, - limit: usize, - ) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") - } - - #[tokio::test] - async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use litellm_llms::base_llm::ocr::error::Error; - - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(Error::TooLarge { limit: 8 }) - )); - } - } - - #[rstest] - #[case::declared("Content-Length: 1000000")] - #[case::chunked("Transfer-Encoding: chunked")] - #[tokio::test] - async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( - #[case] headers: &str, - ) { - let prefix = "x".repeat(4096); - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), prefix.len()) - .await - .unwrap_err(); - match error { - OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!(body, prefix); - } - error => panic!("unexpected error: {error}"), - } - } - - #[test] - fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.transport.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } - } - - #[derive(Debug)] - struct PendingToken { - entered: Arc, - dropped: Arc, - } - - struct TokenFutureDrop(Arc); - - impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } - } - - impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } - } - - #[tokio::test] - async fn interrupt_drops_provider_captures_before_returning() { - use std::sync::atomic::{AtomicBool, Ordering}; - - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = crate::ocr::types::LiteLLMOcrRequest { - transport: OcrTransportConfig { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.transport - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let host = LocalOcrHost::new(request); - let mut machine = ocr_machine(ocr_client()); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = machine.resume() => { - match step.unwrap() { - MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), - MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), - MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), - MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), - MachineStep::Complete(_) => panic!("pending provider completed"), - } - } - } - } - }) - .await - .unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = OcrError::InvalidRequest("cancelled".into()); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); - assert!( - dropped.load(Ordering::SeqCst), - "interrupt returned while provider captures were still alive" - ); - assert!( - matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); - } - - struct CallerTokenHost { - request: Mutex>, - trace: Mutex>, - } - - impl Host for CallerTokenHost { - async fn project(&self) -> Result { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrProjection { - request: self.request.lock().unwrap().take().unwrap(), - caller_token: true, - }) - } - - async fn custom_op(&self, op: OcrOp) -> Result<(), OcrError> { - match op { - OcrOp::AcquireAzureAdToken(reply) => { - self.trace.lock().unwrap().push("token".into()); - reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), - )); - Ok(()) - } - } - } - - async fn before_send( - &self, - wire: WireRequest, - _: &litellm_host::event::RequestContext, - ) -> Result { - let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); - let authorization = wire - .headers - .iter() - .find(|(name, _)| is_authorization(name)) - .map(|(_, value)| value.clone()) - .unwrap_or_default(); - self.trace - .lock() - .unwrap() - .push(format!("before_send:{authorization}")); - let headers = wire - .headers - .into_iter() - .map(|(name, value)| match is_authorization(&name) { - true => (name, "Bearer edited".to_string()), - false => (name, value), - }) - .collect(); - Ok(WireRequest { headers, ..wire }) - } - } - - #[tokio::test] - async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request("azure_ai/model", &base, json!({})); - request.credentials.api_key = None; - let host = CallerTokenHost { - request: Mutex::new(Some(request)), - trace: Mutex::new(Vec::new()), - }; - - litellm_host::run::run(ocr_machine(ocr_client()), &host) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!( - *host.trace.lock().unwrap(), - ["project", "token", "before_send:Bearer caller-token"] - ); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer edited\r\n") - ); - } - - #[tokio::test] - async fn interrupting_an_in_flight_provider_request_closes_its_connection() { - use tokio::io::AsyncReadExt; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let received = Arc::new(tokio::sync::Notify::new()); - let server_received = received.clone(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = Vec::new(); - let mut buffer = [0u8; 4096]; - while !request.windows(4).any(|window| window == b"\r\n\r\n") { - let read = socket.read(&mut buffer).await.unwrap(); - request.extend_from_slice(&buffer[..read]); - } - server_received.notify_one(); - loop { - if socket.read(&mut buffer).await.unwrap() == 0 { - break; - } - } - }); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let mut machine = ocr_machine(ocr_client()); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = received.notified() => break, - step = machine.resume() => { - match step.unwrap() { - MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), - MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), - MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), - MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), - MachineStep::Complete(_) => panic!("the stalled provider completed"), - } - } - } - } - }) - .await - .unwrap(); - - let cancelled = OcrError::InvalidRequest("cancelled".into()); - assert!( - machine - .interrupt(HostFailure::Cancelled(cancelled)) - .await - .is_err() - ); - tokio::time::timeout(std::time::Duration::from_secs(1), server) - .await - .expect("the provider connection stayed open after the interrupt") - .unwrap(); - } -} diff --git a/litellm-rust/crates/core/tests/audio_transcription.rs b/litellm-rust/crates/core/tests/audio_transcription.rs index aa5aef0149b..196f085a6c3 100644 --- a/litellm-rust/crates/core/tests/audio_transcription.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -1,50 +1,250 @@ -use std::{ - io::{Read, Write}, - net::TcpListener, - thread, +use litellm_core::audio_transcription::{ + Error, audio_transcription, types::AudioTranscriptionRequest, }; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; -use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; -use serde_json::{Map, json}; +mod support; +use support::*; -#[tokio::test] -async fn bedrock_request_is_signed_and_contains_audio() { - let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); - let address = listener.local_addr().expect("address"); - let server = thread::spawn(move || { - let (mut stream, _) = listener.accept().expect("connection"); - let mut request = Vec::new(); - let mut buffer = [0_u8; 16_384]; - let count = stream.read(&mut buffer).expect("request"); - request.extend_from_slice(&buffer[..count]); - let request = String::from_utf8_lossy(&request); - assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse")); - assert!(request.contains("authorization: AWS4-HMAC-SHA256")); - assert!(request.contains("x-amz-date:")); - assert!(request.contains("\"bytes\":\"AQI=\"")); - assert!(request.contains("Transcribe the audio. Respond with only the transcript.")); - let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}"; - stream.write_all(response).expect("response"); - }); +const MODEL: &str = "mistral.voxtral-mini-3b-2507"; - let optional_params = Map::from_iter([ +fn transcript_response(text: &str) -> ResponseTemplate { + json_response(json!({"output": {"message": {"content": [{"text": text}]}}})) +} + +fn aws_params(region: &str) -> Map { + Map::from_iter([ ("aws_access_key_id".to_string(), json!("access-key")), ("aws_secret_access_key".to_string(), json!("secret-key")), - ("aws_region_name".to_string(), json!("us-east-1")), - ]); - let api_base = format!("http://{address}"); - let response = audio_transcription(AudioTranscriptionRequest { - model: "mistral.voxtral-mini-3b-2507", + ("aws_region_name".to_string(), json!(region)), + ]) +} + +#[fixture] +fn request() -> AudioTranscriptionRequest<'static> { + AudioTranscriptionRequest { + model: MODEL, audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), api_key: None, - api_base: Some(&api_base), + api_base: None, custom_llm_provider: Some("bedrock"), extra_headers: None, - optional_params, + optional_params: aws_params("us-east-1"), timeout: None, + } +} + +#[rstest] +#[case::us_east_1("us-east-1")] +#[case::eu_west_1("eu-west-1")] +#[tokio::test] +async fn bedrock_converse_request_is_signed_for_the_requested_region( + request: AudioTranscriptionRequest<'static>, + #[case] region: &str, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + + let response = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + optional_params: aws_params(region), + ..request }) .await .expect("transcription"); + assert_eq!(response, json!({"text": "hello"})); - server.join().expect("server"); + let sent = only_request(&upstream).await; + assert_eq!(sent.method.as_str(), "POST"); + assert_eq!(sent.url.path(), format!("/model/{MODEL}/converse")); + let authorization = sent.header("authorization").expect("request is signed"); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256 Credential=access-key/"), + "{authorization}" + ); + assert!( + authorization.contains(&format!("/{region}/bedrock/aws4_request")), + "{authorization}" + ); + assert!(sent.header("x-amz-date").is_some()); + assert!(!sent.body_text().contains("secret-key")); +} + +#[rstest] +#[tokio::test] +async fn the_provider_can_come_from_the_model_prefix(request: AudioTranscriptionRequest<'static>) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + let model = format!("bedrock/{MODEL}"); + + audio_transcription(AudioTranscriptionRequest { + model: &model, + custom_llm_provider: None, + api_base: Some(&base), + ..request + }) + .await + .expect("transcription"); + + assert_eq!( + only_request(&upstream).await.url.path(), + format!("/model/{MODEL}/converse") + ); +} + +#[rstest] +#[tokio::test] +async fn audio_and_transcription_params_reach_the_converse_body( + request: AudioTranscriptionRequest<'static>, + #[values("wav", "mp3", "flac", "ogg")] format: &str, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + let optional_params = aws_params("us-east-1") + .into_iter() + .chain([ + ("language".to_string(), json!("fr")), + ("temperature".to_string(), json!(0.2)), + ]) + .collect(); + + audio_transcription(AudioTranscriptionRequest { + audio: json!({"data": "AQI=", "format": format}), + api_base: Some(&base), + optional_params, + ..request + }) + .await + .expect("transcription"); + + let body = only_request(&upstream).await.json(); + let content = &body["messages"][0]["content"]; + assert_eq!( + content[0], + json!({"audio": {"format": format, "source": {"bytes": "AQI="}}}) + ); + let instruction = content[1]["text"].as_str().expect("instruction text"); + assert!(instruction.contains("fr"), "{instruction}"); + assert_eq!(body["inferenceConfig"]["temperature"], 0.2); +} + +#[rstest] +#[case::unknown_format(json!({"data": "AQI=", "format": "aac"}))] +#[case::missing_data(json!({"format": "wav"}))] +#[case::not_an_object(json!("AQI="))] +#[tokio::test] +async fn invalid_audio_is_rejected_before_sending( + request: AudioTranscriptionRequest<'static>, + #[case] audio: Value, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + audio, + api_base: Some(&base), + ..request + }) + .await + .expect_err("invalid audio is rejected"); + + assert!( + matches!( + error, + Error::InvalidRequest(_) | Error::MissingField(_) | Error::InvalidType { .. } + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[case::unknown_provider(MODEL, Some("openai"), "openai")] +#[case::unresolvable_model( + "no-such-model", + None, + "unable to resolve custom_llm_provider for audio transcription request" +)] +#[tokio::test] +async fn unsupported_providers_are_rejected_before_sending( + request: AudioTranscriptionRequest<'static>, + #[case] model: &'static str, + #[case] provider: Option<&'static str>, + #[case] reported: &str, +) { + let error = audio_transcription(AudioTranscriptionRequest { + model, + custom_llm_provider: provider, + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("unsupported provider errors"); + + assert_eq!(error, Error::InvalidProvider(reported.into())); +} + +#[rstest] +#[tokio::test] +async fn a_non_string_extra_header_is_rejected(request: AudioTranscriptionRequest<'static>) { + let error = audio_transcription(AudioTranscriptionRequest { + extra_headers: Some(Map::from_iter([("x-count".to_string(), json!(3))])), + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("a non-string header is rejected"); + + assert!(matches!(error, Error::Headers(_)), "{error:?}"); +} + +#[rstest] +#[case::throttled(429)] +#[case::server_error(500)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_and_body( + request: AudioTranscriptionRequest<'static>, + #[case] status: u16, +) { + let upstream = + upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status, + body: "upstream said no".into() + }) + ); +} + +#[rstest] +#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] +#[case::no_output(json_response(json!({"unexpected": true})))] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response( + request: AudioTranscriptionRequest<'static>, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("an unreadable body fails"); + + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); } diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs new file mode 100644 index 00000000000..ae96509fe2e --- /dev/null +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -0,0 +1,320 @@ +use std::time::Duration; + +use litellm_core::chat_completions::{ + Error, chat_completions, chat_completions_decline_reason, types::ChatCompletionsRequest, +}; +use litellm_http::transport::Error as TransportError; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; + +mod support; +use support::*; + +const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn anthropic_response(body: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_raw(body, "application/json") +} + +fn hi() -> Value { + json!([{"role": "user", "content": "hi"}]) +} + +#[fixture] +fn request() -> ChatCompletionsRequest<'static> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages: hi(), + optional_params: object(json!({"max_tokens": 16})), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + timeout: Some(Duration::from_secs(10)), + } +} + +#[rstest] +#[tokio::test] +async fn anthropic_round_trip_translates_the_conversation_and_normalizes_the_response( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + let response = chat_completions(ChatCompletionsRequest { + messages: json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + api_base: Some(&base), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/v1/messages"); + assert_eq!(sent.header_values("x-api-key"), ["sk-test"]); + let body = sent.json(); + assert_eq!(body["model"], "claude-sonnet-4-5"); + assert_eq!( + body["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + body["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(body["max_tokens"], 16); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); +} + +#[rstest] +#[tokio::test] +async fn the_deployment_key_replaces_a_caller_supplied_x_api_key( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + extra_headers: Some(object( + json!({"x-api-key": "caller-key", "x-trace": "kept"}), + )), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!(sent.header_values("x-api-key"), ["sk-test"]); + assert_eq!(sent.header("x-trace"), Some("kept")); +} + +#[rstest] +#[tokio::test] +async fn bedrock_round_trip_is_signed_and_normalized(request: ChatCompletionsRequest<'static>) { + let upstream = upstream([json_response(json!({ + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15} + }))]) + .await; + let base = upstream.uri(); + + let response = chat_completions(ChatCompletionsRequest { + model: "bedrock/anthropic.claude-sonnet-4-5", + optional_params: object(json!({ + "aws_access_key_id": "access-key", + "aws_secret_access_key": "secret-key", + "aws_region_name": "eu-west-1" + })), + api_key: None, + api_base: Some(&base), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/model/anthropic.claude-sonnet-4-5/converse" + ); + let authorization = sent.header("authorization").expect("request is signed"); + assert!( + authorization.contains("/eu-west-1/bedrock/aws4_request"), + "{authorization}" + ); + assert_eq!( + sent.json()["messages"], + json!([{"role": "user", "content": [{"text": "hi"}]}]) + ); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); +} + +/// The provider already answered and billed these, so the host must not retry them on +/// its own path: they surface as `InvalidResponse`, never as a pre-send decline. +#[rstest] +#[case::missing_usage( + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"# +)] +#[case::tool_use_block(r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#)] +#[case::not_json("not json")] +#[tokio::test] +async fn a_response_it_cannot_normalize_is_reported_as_already_sent( + request: ChatCompletionsRequest<'static>, + #[case] body: &str, +) { + let upstream = upstream([anthropic_response(body)]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("response cannot be normalized"); + + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); +} + +#[rstest] +#[case::rate_limited(429)] +#[case::server_error(500)] +#[tokio::test] +async fn an_upstream_error_status_keeps_its_code_and_body( + request: ChatCompletionsRequest<'static>, + #[case] status: u16, +) { + let upstream = upstream([ResponseTemplate::new(status).set_body_string("slow down")]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("upstream rejects"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status, + body: "slow down".into() + }) + ); +} + +/// Nothing was sent, so nothing was billed and the host can still serve the request. +#[rstest] +#[tokio::test] +async fn a_connection_that_is_never_established_declines_instead_of_failing( + request: ChatCompletionsRequest<'static>, +) { + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("nothing is listening"); + + assert!( + matches!(error, Error::Transport(TransportError::Connect(_))), + "{error:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_timeout_after_sending_is_not_a_pre_send_decline( + request: ChatCompletionsRequest<'static>, +) { + let upstream = + upstream([anthropic_response(ANTHROPIC_MESSAGE).set_delay(Duration::from_secs(5))]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + timeout: Some(Duration::from_millis(100)), + ..request + }) + .await + .expect_err("the call times out"); + + assert!( + matches!(error, Error::Transport(TransportError::Network(_))), + "{error:?}" + ); +} + +#[rstest] +#[case::accepted("anthropic/claude-sonnet-4-5", None, hi(), json!({"max_tokens": 16}), None)] +#[case::accepted_bedrock("bedrock/anthropic.claude-sonnet-4-5", None, hi(), json!({}), None)] +#[case::unknown_provider( + "gpt-4o", + Some("openai"), + hi(), + json!({}), + Some("provider is not on the rust chat completions path") +)] +#[case::unreadable_messages( + "anthropic/claude-sonnet-4-5", + None, + json!("hi"), + json!({}), + Some("unreadable message list") +)] +#[case::empty_messages("anthropic/claude-sonnet-4-5", None, json!([]), json!({}), Some("empty message list"))] +#[case::streaming( + "anthropic/claude-sonnet-4-5", + None, + hi(), + json!({"stream": true}), + Some("streaming") +)] +#[case::unrecognized_param( + "anthropic/claude-sonnet-4-5", + None, + hi(), + json!({"not_a_param": 1}), + Some("unrecognized request parameter") +)] +#[case::opens_on_assistant_turn( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "assistant", "content": "hi"}]), + json!({}), + Some("conversation does not open on a user turn") +)] +fn decline_reason_names_why_the_core_would_not_serve_the_request( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] messages: Value, + #[case] params: Value, + #[case] reason: Option<&str>, +) { + assert_eq!( + chat_completions_decline_reason(model, provider, messages, &object(params)), + reason + ); +} + +/// A request the decline check accepts must not be declined by the call itself. +#[rstest] +#[tokio::test] +async fn a_declined_request_fails_the_call_before_sending( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + optional_params: object(json!({"stream": true})), + api_base: Some(&base), + ..request + }) + .await + .expect_err("streaming is declined"); + + assert_eq!(error, Error::Unsupported("streaming")); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/messages.rs b/litellm-rust/crates/core/tests/messages.rs deleted file mode 100644 index 18af8a7d619..00000000000 --- a/litellm-rust/crates/core/tests/messages.rs +++ /dev/null @@ -1,471 +0,0 @@ -use std::{sync::Arc, time::Duration}; - -use futures_util::future::BoxFuture; -use litellm_core::messages::{ - Error, messages, - route::{LocalMessagesHost, MessagesCall, messages_machine}, - types::{MessagesRequest, MessagesShaping}, -}; -use litellm_secrets::{SecretValue, source::SecretSource}; -use serde_json::{Map, Value, json}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, -}; - -struct RecordingSecrets { - values: Vec<(&'static str, String)>, - fails: bool, - requested: std::sync::Mutex>, -} - -impl RecordingSecrets { - fn new(values: Vec<(&'static str, String)>, fails: bool) -> Self { - Self { - values, - fails, - requested: std::sync::Mutex::new(Vec::new()), - } - } -} - -impl SecretSource for RecordingSecrets { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - Box::pin(async move { - self.requested.lock().unwrap().push(name.to_string()); - if self.fails { - return Err(litellm_secrets::Error::ManagedSecretMissing); - } - Ok(self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| SecretValue::new(value.clone()))) - }) - } -} - -fn secrets_call() -> MessagesCall { - let Value::Object(body) = json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }) else { - unreachable!("literal object") - }; - MessagesCall { - model: "claude-sonnet-4-5".into(), - body, - api_key: None, - api_base: None, - custom_llm_provider: Some("anthropic".into()), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - } -} - -#[tokio::test] -async fn route_surfaces_a_secret_manager_failure_before_the_call() { - let Err(error) = litellm_host::run::run( - messages_machine(Arc::new(RecordingSecrets::new(Vec::new(), true))), - &LocalMessagesHost::new(secrets_call()), - ) - .await - else { - panic!("a secret manager failure fails the call"); - }; - assert!( - matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), - "{error:?}" - ); -} - -async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") -} - -fn write_response(body: &str) -> String { - format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) -} - -#[tokio::test] -async fn messages_round_trip_builds_azure_request_and_passes_response_through() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let response = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 1024, - "messages": [{ - "role": "user", - "content": [{ - "type": "text", - "text": "hi", - "cache_control": {"type": "ephemeral", "scope": "global"} - }] - }] - }), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - assert_eq!(response.content[0]["text"], "hi"); - assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); - - let request = server.await.expect("server task completes"); - let (head, body) = request.split_once("\r\n\r\n").expect("has body"); - assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}"); - let head_lower = head.to_ascii_lowercase(); - assert!(head_lower.contains("x-api-key: sk-azure"), "{head}"); - assert!( - head_lower.contains("anthropic-version: 2023-06-01"), - "{head}" - ); - assert!( - head_lower.contains("content-type: application/json"), - "{head}" - ); - - let sent_body: Value = serde_json::from_str(body).expect("body is json"); - assert_eq!( - sent_body["messages"][0]["content"][0]["cache_control"], - json!({"type": "ephemeral"}) - ); -} - -#[tokio::test] -async fn messages_round_trip_builds_native_anthropic_request() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let response = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 1024, - "messages": [{"role": "user", "content": "hi"}] - }), - api_key: Some("sk-ant"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("anthropic"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - assert_eq!(response.content[0]["text"], "hi"); - assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); - - let request = server.await.expect("server task completes"); - let (head, _) = request.split_once("\r\n\r\n").expect("has body"); - assert!(head.starts_with("POST /v1/messages "), "{head}"); - let head_lower = head.to_ascii_lowercase(); - assert!(head_lower.contains("x-api-key: sk-ant"), "{head}"); - assert!( - head_lower.contains("anthropic-version: 2023-06-01"), - "{head}" - ); -} - -#[tokio::test] -async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "x-api-key".to_string(), - Value::String("from-python".to_string()), - ); - headers.insert( - "anthropic-beta".to_string(), - Value::String("token-efficient-tools-2025-02-19".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("rust-fallback-key"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - let api_key_count = head - .lines() - .filter(|line| line.starts_with("x-api-key:")) - .count(); - assert_eq!(api_key_count, 1, "{head}"); - assert!(head.contains("x-api-key: from-python"), "{head}"); - assert!( - head.contains("anthropic-beta: token-efficient-tools-2025-02-19"), - "{head}" - ); - assert!(!head.contains("rust-fallback-key"), "{head}"); -} - -#[tokio::test] -async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer entra-token".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: None, - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("entra id request succeeds without api key"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - assert!(head.contains("authorization: bearer entra-token"), "{head}"); - assert!(!head.contains("x-api-key"), "{head}"); -} - -#[tokio::test] -async fn messages_requires_auth_when_no_key_and_no_header() { - let err = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: None, - api_base: Some("http://127.0.0.1:1"), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_millis(50)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("missing auth errors"); - - assert!(matches!(err, Error::Auth(_))); -} - -#[tokio::test] -async fn messages_ignores_malformed_authorization_and_uses_api_key() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer ".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("falls back to api key"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - assert!(head.contains("x-api-key: sk-azure"), "{head}"); -} - -#[tokio::test] -async fn messages_maps_provider_error_status_to_http_error() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let _ = read_http_request(&mut socket).await; - let body = "unauthorized"; - let response = format!( - "HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - }); - - let err = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("provider error propagates"); - - assert!(matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 401, .. }) - )); -} - -#[tokio::test] -async fn messages_rejects_unsupported_provider() { - let err = messages(MessagesRequest { - model: "claude-3-5-sonnet", - body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}), - api_key: Some("sk"), - api_base: Some("http://127.0.0.1:1"), - custom_llm_provider: Some("openai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_millis(50)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("unsupported provider errors"); - - assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); -} diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs new file mode 100644 index 00000000000..4e549bae309 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -0,0 +1,94 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_core::messages::{ + Error, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + types::MessagesShaping, +}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use rstest::fixture; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; + +#[path = "../support/mod.rs"] +mod support; +use support::*; + +mod request; +mod response; +mod secrets; +mod stream; + +const MODEL: &str = "claude-sonnet-4-5"; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn message_body() -> Value { + json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": MODEL, + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 2} + }) +} + +fn message_response() -> ResponseTemplate { + json_response(message_body()) +} + +/// A non-streaming call with nothing that would authenticate or route it, so each test +/// states the provider, credentials, and base it depends on. +#[fixture] +fn call() -> MessagesCall { + MessagesCall { + model: MODEL.into(), + body: object(json!({ + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + })), + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +fn headers<'a>(pairs: impl IntoIterator) -> Option> { + Some( + pairs + .into_iter() + .map(|(name, value)| (name.to_string(), Value::from(value))) + .collect(), + ) +} + +async fn run_with( + secrets: Arc, + call: MessagesCall, +) -> Result { + litellm_host::run::run(messages_machine(secrets), &LocalMessagesHost::new(call)).await +} + +/// Runs the route with a secret source that knows nothing, so no environment leaks in. +async fn run(call: MessagesCall) -> Result { + run_with(Arc::new(RecordingSecrets::empty()), call).await +} + +async fn run_message(call: MessagesCall) -> AnthropicMessagesResponse { + match run(call).await.expect("messages call succeeds") { + MessagesOutput::Message(message) => *message, + MessagesOutput::Streamed => panic!("a non-streaming call returned a stream"), + } +} diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs new file mode 100644 index 00000000000..9353324d370 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -0,0 +1,269 @@ +use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::anthropic_key("anthropic", Some("sk-ant"), &[], ("x-api-key", "sk-ant"), &["authorization"])] +#[case::azure_key("azure_ai", Some("sk-azure"), &[], ("x-api-key", "sk-azure"), &["authorization"])] +#[case::caller_x_api_key_wins( + "azure_ai", + Some("rust-fallback-key"), + &[("x-api-key", "from-python")], + ("x-api-key", "from-python"), + &["authorization"] +)] +#[case::entra_bearer_without_key( + "azure_ai", + None, + &[("Authorization", "Bearer entra-token")], + ("authorization", "Bearer entra-token"), + &["x-api-key"] +)] +#[case::empty_bearer_falls_back_to_key( + "azure_ai", + Some("sk-azure"), + &[("Authorization", "Bearer ")], + ("x-api-key", "sk-azure"), + &[] +)] +#[case::anthropic_forwards_caller_authorization( + "anthropic", + Some("sk-ant"), + &[("Authorization", "Bearer caller")], + ("authorization", "Bearer caller"), + &["x-api-key"] +)] +#[case::anthropic_oauth_key_becomes_bearer( + "anthropic", + Some("sk-ant-oat01-token"), + &[], + ("authorization", "Bearer sk-ant-oat01-token"), + &["x-api-key"] +)] +#[tokio::test] +async fn credentials_become_exactly_one_auth_header( + call: MessagesCall, + #[case] provider: &str, + #[case] api_key: Option<&str>, + #[case] extra_headers: &[(&str, &str)], + #[case] expected: (&str, &str), + #[case] absent: &[&str], +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: api_key.map(Into::into), + api_base: Some(upstream.uri()), + extra_headers: headers(extra_headers.iter().copied()), + ..call + }) + .await; + + let request = only_request(&upstream).await; + let (name, value) = expected; + assert_eq!(request.header_values(name), [value]); + for name in absent { + assert_eq!(request.header(name), None, "{name} must not be sent"); + } +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn a_call_without_credentials_fails_before_sending( + call: MessagesCall, + #[case] provider: &str, +) { + let upstream = upstream([message_response()]).await; + + let error = run(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("a call without credentials fails"); + + assert!( + matches!( + error, + Error::Auth(litellm_auth::Error::MissingApiKey { .. }) + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[case::anthropic(MODEL, Some("anthropic"), "", "/v1/messages")] +#[case::anthropic_base_with_trailing_slash(MODEL, Some("anthropic"), "/", "/v1/messages")] +#[case::anthropic_base_with_the_messages_path( + MODEL, + Some("anthropic"), + "/v1/messages", + "/v1/messages" +)] +#[case::azure_ai(MODEL, Some("azure_ai"), "", "/anthropic/v1/messages")] +#[case::provider_from_model_prefix("anthropic/claude-sonnet-4-5", None, "", "/v1/messages")] +#[tokio::test] +async fn each_provider_posts_to_its_messages_endpoint( + call: MessagesCall, + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] base_suffix: &str, + #[case] path: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + model: model.into(), + custom_llm_provider: provider.map(Into::into), + api_key: Some("sk".into()), + api_base: Some(format!("{}{base_suffix}", upstream.uri())), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!(request.method.as_str(), "POST"); + assert_eq!(request.url.path(), path); + assert_eq!(request.json()["model"], MODEL); + assert_eq!(request.header("anthropic-version"), Some("2023-06-01")); + assert_eq!(request.header("content-type"), Some("application/json")); +} + +#[rstest] +#[case::unknown_provider(MODEL, Some("openai"), "openai")] +#[case::unresolvable_model( + "no-such-model", + None, + "unable to resolve custom_llm_provider for messages request" +)] +#[tokio::test] +async fn unsupported_providers_are_rejected_before_sending( + call: MessagesCall, + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] reported: &str, +) { + let error = run(MessagesCall { + model: model.into(), + custom_llm_provider: provider.map(Into::into), + api_key: Some("sk".into()), + api_base: Some(UNREACHABLE_BASE.into()), + ..call + }) + .await + .err() + .expect("unsupported provider errors"); + + assert_eq!(error, Error::InvalidProvider(reported.into())); +} + +#[rstest] +#[tokio::test] +async fn caller_headers_and_provider_scoped_headers_are_forwarded(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let scoped = |provider: &str, value: &str| ProviderSpecificHeader { + custom_llm_provider: provider.into(), + extra_headers: object(json!({"x-scoped": value})), + }; + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("anthropic-beta", "token-efficient-tools-2025-02-19")]), + provider_specific_header: Some(ProviderSpecificHeaders::Many(vec![ + scoped("bedrock", "other-provider"), + scoped("azure_ai, anthropic", "this-provider"), + ])), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!( + request.header("anthropic-beta"), + Some("token-efficient-tools-2025-02-19") + ); + assert_eq!(request.header_values("x-scoped"), ["this-provider"]); +} + +#[rstest] +#[tokio::test] +async fn azure_strips_the_cache_control_scope_anthropic_rejects(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + api_key: Some("sk-azure".into()), + api_base: Some(upstream.uri()), + body: object(json!({ + "model": MODEL, + "max_tokens": 16, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "hi", + "cache_control": {"type": "ephemeral", "scope": "global"} + }] + }] + })), + ..call + }) + .await; + + assert_eq!( + only_request(&upstream).await.json()["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); +} + +#[rstest] +#[tokio::test] +async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let mut body = call.body.clone(); + body.insert("temperature".into(), json!(0.5)); + body.insert("top_k".into(), json!(3)); + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + body, + shaping: MessagesShaping { + additional_drop_params: vec!["temperature".into()], + ..MessagesShaping::default() + }, + ..call + }) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("temperature"), None); + assert_eq!(sent["top_k"], 3); +} + +#[rstest] +#[case::anthropic_streams(MODEL, Some("anthropic"), true, true)] +#[case::anthropic_prefix_streams("anthropic/claude-sonnet-4-5", None, true, true)] +#[case::azure_without_stream(MODEL, Some("azure_ai"), false, true)] +#[case::azure_stream(MODEL, Some("azure_ai"), true, false)] +#[case::other_provider(MODEL, Some("openai"), false, false)] +#[case::unresolvable_model("no-such-model", None, false, false)] +fn supports_matches_what_the_route_can_serve( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] stream: bool, + #[case] supported: bool, +) { + assert_eq!( + litellm_core::messages::route::supports(model, provider, stream), + supported + ); +} diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs new file mode 100644 index 00000000000..38a18c415ba --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -0,0 +1,136 @@ +use litellm_core::messages::{messages, types::MessagesRequest}; +use litellm_http::transport::Error as TransportError; +use rstest::rstest; + +use super::*; + +#[rstest] +#[tokio::test] +async fn the_provider_message_is_returned(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + let message = run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(message.id, "msg_1"); + assert_eq!(message.content, [json!({"type": "text", "text": "hi"})]); + assert_eq!(message.stop_reason.as_deref(), Some("end_turn")); +} + +#[rstest] +#[case::bad_request(400)] +#[case::unauthorized(401)] +#[case::rate_limited(429)] +#[case::server_error(500)] +#[case::overloaded(529)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] status: u16) { + let upstream = + upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status, + body: "upstream said no".into() + }) + ); +} + +#[rstest] +#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] +#[case::not_a_message(json_response(json!({"unexpected": true})))] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response( + call: MessagesCall, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("an unreadable body fails"); + + assert!(error.is_response(), "{error:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { + let upstream = upstream([message_response().set_delay(Duration::from_secs(5))]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + timeout: Some(Duration::from_millis(100)), + ..call + }) + .await + .err() + .expect("the call times out"); + + assert!(matches!(error, Error::Transport(_)), "{error:?}"); +} + +fn facade_request(body: Value, api_base: &str) -> MessagesRequest<'_> { + MessagesRequest { + model: MODEL, + body, + api_key: Some("sk-ant"), + api_base: Some(api_base), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +#[tokio::test] +async fn the_facade_runs_the_route_in_process() { + let upstream = upstream([message_response()]).await; + let base = upstream.uri(); + + let message = messages(facade_request( + json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), + &base, + )) + .await + .expect("messages request succeeds"); + + assert_eq!(message.id, "msg_1"); + assert_eq!( + only_request(&upstream).await.header("x-api-key"), + Some("sk-ant") + ); +} + +#[tokio::test] +async fn the_facade_rejects_a_body_that_is_not_an_object() { + let error = messages(facade_request(json!([]), UNREACHABLE_BASE)) + .await + .expect_err("a non-object body is rejected"); + + assert_eq!( + error, + Error::InvalidRequest("messages body must be an object".into()) + ); +} diff --git a/litellm-rust/crates/core/tests/messages/secrets.rs b/litellm-rust/crates/core/tests/messages/secrets.rs new file mode 100644 index 00000000000..419b6d6c753 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/secrets.rs @@ -0,0 +1,94 @@ +use litellm_llms::{ + anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, + azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, +}; +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::anthropic("anthropic", &ANTHROPIC_MESSAGES_CONFIG, "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "/v1/messages")] +#[case::azure_ai("azure_ai", &AZURE_ANTHROPIC_MESSAGES_CONFIG, "AZURE_API_KEY", "AZURE_API_BASE", "/anthropic/v1/messages")] +#[tokio::test] +async fn the_credential_and_base_come_from_the_secret_source( + call: MessagesCall, + #[case] provider: &str, + #[case] config: &dyn BaseAnthropicMessagesConfig, + #[case] key_name: &str, + #[case] base_name: &str, + #[case] path: &str, +) { + let upstream = upstream([message_response()]).await; + let base = upstream.uri(); + let secrets = Arc::new(RecordingSecrets::new([ + (key_name, "sk-from-manager"), + (base_name, base.as_str()), + ])); + + let output = run_with( + secrets.clone(), + MessagesCall { + custom_llm_provider: Some(provider.into()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = only_request(&upstream).await; + assert_eq!(request.url.path(), path); + assert_eq!(request.header("x-api-key"), Some("sk-from-manager")); + assert_eq!(secrets.requested(), config.secret_names()); +} + +#[rstest] +#[tokio::test] +async fn call_arguments_win_over_the_secret_source(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let secrets = Arc::new(RecordingSecrets::new([ + ("ANTHROPIC_API_KEY", "sk-from-manager"), + ("ANTHROPIC_BASE_URL", UNREACHABLE_BASE), + ])); + + run_with( + secrets, + MessagesCall { + api_key: Some("sk-from-call".into()), + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + assert_eq!( + only_request(&upstream).await.header("x-api-key"), + Some("sk-from-call") + ); +} + +#[rstest] +#[tokio::test] +async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + let error = run_with( + Arc::new(RecordingSecrets::failing()), + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .err() + .expect("a secret manager failure fails the call"); + + assert!( + matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs new file mode 100644 index 00000000000..ea23a9e8e38 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -0,0 +1,163 @@ +use std::{convert::Infallible, sync::Mutex}; + +use bytes::Bytes; +use litellm_core::messages::route::Messages; +use litellm_host::host::{Demand, Host}; +use rstest::rstest; + +use super::*; + +const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + +enum Seen { + Open, + Deliver(Bytes), +} + +/// Projects like `LocalMessagesHost`, records every stream op in the order the route +/// performs it, and detaches after `detach_after` ops. +struct RecordingStreamHost { + call: LocalMessagesHost, + detach_after: usize, + seen: Mutex>, +} + +impl RecordingStreamHost { + fn new(call: MessagesCall, detach_after: usize) -> Self { + Self { + call: LocalMessagesHost::new(call), + detach_after, + seen: Mutex::new(Vec::new()), + } + } + + fn record(&self, op: Seen) -> Demand { + let mut seen = self.seen.lock().unwrap(); + seen.push(op); + match seen.len() < self.detach_after { + true => Demand::More, + false => Demand::Detached, + } + } +} + +impl Host for RecordingStreamHost { + async fn project(&self) -> Result { + self.call.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn open(&self, (): ()) -> Result { + Ok(self.record(Seen::Open)) + } + + async fn deliver(&self, chunk: Bytes) -> Result { + Ok(self.record(Seen::Deliver(chunk))) + } +} + +fn streaming(call: MessagesCall, api_base: String) -> MessagesCall { + let mut body = call.body.clone(); + body.insert("stream".into(), json!(true)); + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + body, + ..call + } +} + +fn sse_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream") +} + +async fn stream_through(host: &RecordingStreamHost) -> Result { + litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await +} + +#[rstest] +#[tokio::test] +async fn the_stream_opens_once_before_relaying_the_upstream_body(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + let outcome = stream_through(&host).await.expect("streamed call succeeds"); + + assert!(matches!(outcome, MessagesOutput::Streamed)); + let seen = host.seen.into_inner().unwrap(); + let [Seen::Open, chunks @ ..] = seen.as_slice() else { + panic!("the stream opens before any chunk is delivered"); + }; + let delivered: Vec = chunks + .iter() + .flat_map(|step| match step { + Seen::Deliver(chunk) => chunk.to_vec(), + Seen::Open => panic!("the stream opens exactly once"), + }) + .collect(); + assert_eq!(delivered, SSE_BODY.as_bytes()); +} + +#[rstest] +#[case::at_open(1)] +#[case::after_the_first_chunk(2)] +#[tokio::test] +async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] detach_after: usize) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), detach_after); + + let outcome = stream_through(&host) + .await + .expect("a detached stream still completes"); + + assert!(matches!(outcome, MessagesOutput::Streamed)); + assert_eq!(host.seen.into_inner().unwrap().len(), detach_after); +} + +#[rstest] +#[tokio::test] +async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + let error = stream_through(&host) + .await + .err() + .expect("upstream error propagates"); + + assert!( + matches!( + error, + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) + ), + "{error:?}" + ); + assert!(host.seen.into_inner().unwrap().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new( + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + ..streaming(call, upstream.uri()) + }, + usize::MAX, + ); + + let error = stream_through(&host) + .await + .err() + .expect("azure streaming is refused"); + + assert_eq!( + error, + Error::Unsupported("streaming messages for this provider") + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/aws_textract.rs b/litellm-rust/crates/core/tests/ocr/aws_textract.rs new file mode 100644 index 00000000000..790e16a95ec --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/aws_textract.rs @@ -0,0 +1,173 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use rstest::rstest; +use time::{PrimitiveDateTime, format_description}; +use wiremock::Request; + +use super::*; + +const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; +const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; +const DETECT: &str = "aws_textract/detect-document-text"; +const ANALYZE: &str = "aws_textract/analyze-document"; + +fn textract_request(model: &str, base: &str) -> LiteLLMOcrRequest { + ocr_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) +} + +fn textract_response() -> ResponseTemplate { + json_response(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) +} + +/// Recomputes SigV4 over the request the upstream received, at the time the client claimed. +fn expected_authorization(url: &str, sent: &Request) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(sent.header("x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| (name.to_string(), sent.header(name).unwrap().to_string())) + .collect(); + sign_post( + url, + &sent.body, + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() +} + +/// The recorded URL names wiremock's host, not the address the client signed for. +fn assert_signed(upstream: &MockServer, sent: &Request) { + let url = format!("{}/", upstream.uri()); + assert_eq!( + sent.header("authorization"), + Some(expected_authorization(&url, sent).as_str()) + ); +} + +#[tokio::test] +async fn detect_document_text_is_signed_and_lines_become_the_page() { + let upstream = upstream([textract_response()]).await; + + let response = perform_with(LocalOcrHost::new(textract_request(DETECT, &upstream.uri()))) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.header("x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + sent.header("content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!(sent.json(), json!({"Document": {"Bytes": "b3JpZ2luYWw="}})); + assert_signed(&upstream, &sent); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); +} + +#[tokio::test] +async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let upstream = upstream([textract_response()]).await; + let host = LocalOcrHost::new(textract_request(DETECT, &upstream.uri())).with_before_send( + |mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }, + ); + + perform_with(host).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.json(), json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}})); + assert_signed(&upstream, &sent); +} + +#[tokio::test] +async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let upstream = upstream([json_response(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + + let response = perform_with(LocalOcrHost::new(textract_request( + ANALYZE, + &upstream.uri(), + ))) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.header("x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!(sent.json()["FeatureTypes"], json!(["LAYOUT", "TABLES"])); + assert_signed(&upstream, &sent); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); +} + +#[rstest] +#[case::detect(DETECT)] +#[case::analyze(ANALYZE)] +#[tokio::test] +async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[case] model: &str) { + let upstream = upstream([status_response( + 400, + json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + )]) + .await; + + let error = perform_with(LocalOcrHost::new(textract_request(model, &upstream.uri()))) + .await + .unwrap_err(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/azure_ai.rs b/litellm-rust/crates/core/tests/ocr/azure_ai.rs new file mode 100644 index 00000000000..0d6ba024c15 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/azure_ai.rs @@ -0,0 +1,270 @@ +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use litellm_auth::{ + ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, +}; +use rstest::rstest; + +use super::*; + +#[derive(Debug)] +struct CountingToken { + token: fn(usize) -> String, + calls: AtomicUsize, +} + +impl CountingToken { + fn new(token: fn(usize) -> String) -> Arc { + Arc::new(Self { + token, + calls: AtomicUsize::new(0), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } +} + +impl TokenProvider for CountingToken { + fn acquire(&self) -> TokenFuture<'_> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let token = SecretValue::new((self.token)(call)); + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token, + expires_on: None, + }) + }) + } +} + +fn numbered_token(call: usize) -> String { + format!("callback-{call}") +} + +fn azure_request( + provider: &Arc, + api_base: Option<&str>, + api_key: Option<&str>, + extra_headers: Value, + optional_params: Value, +) -> LiteLLMOcrRequest { + let wire = serde_json::from_value(json!({ + "model": "azure_ai/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": null, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": 2.0 + })) + .unwrap(); + let mut request = decode_request(wire).unwrap(); + request.azure_ad_token_provider = Some(TokenProviderHandle::new(provider.clone())); + request +} + +fn ocr_page() -> ResponseTemplate { + json_response(json!({"pages": [{"index": 0, "markdown": "hello"}]})) +} + +#[tokio::test] +async fn mistral_on_azure_sends_the_prepared_bearer_and_the_mistral_body() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + let request = with_headers( + without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({"include_image_base64": true}), + )), + &[("Authorization", "Bearer python-prepared-token")], + ); + + let result = perform(request).await.unwrap(); + + assert_eq!(result.pages[0].markdown, "hello"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/providers/mistral/azure/ocr"); + assert_eq!( + sent.header("authorization"), + Some("Bearer python-prepared-token") + ); + assert_eq!( + sent.json(), + json!({ + "model": "model", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "include_image_base64": true + }) + ); +} + +#[tokio::test] +async fn a_static_entra_token_becomes_the_bearer() { + let upstream = upstream([pages_response()]).await; + let request = without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({"azure_ad_token": "rust-owned-token"}), + )); + + perform(request).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header("authorization"), + Some("Bearer rust-owned-token") + ); +} + +#[tokio::test] +async fn a_guardrail_that_swaps_in_a_remote_document_is_rejected() { + let host = LocalOcrHost::new(ocr_request("azure_ai/model", UNREACHABLE_BASE, json!({}))) + .with_before_send(|mut wire, _| { + wire.body["document"] = json!({ + "type": "document_url", + "document_url": "https://example.com/not-inline.pdf" + }); + Ok(wire) + }); + + let error = perform_with(host).await.unwrap_err(); + + assert!(error.to_string().contains("data URI"), "{error}"); +} + +#[tokio::test] +async fn the_token_provider_is_the_bearer_and_is_acquired_for_each_request() { + let provider = CountingToken::new(numbered_token); + let upstream = upstream([ocr_page(), ocr_page()]).await; + let base = upstream.uri(); + + for _ in 0..2 { + perform(azure_request( + &provider, + Some(&base), + None, + Value::Null, + json!({}), + )) + .await + .unwrap(); + } + + assert_eq!(provider.calls(), 2); + let authorizations: Vec = received(&upstream) + .await + .iter() + .map(|request| { + request + .header("authorization") + .unwrap_or_default() + .to_string() + }) + .collect(); + assert_eq!(authorizations, ["Bearer callback-1", "Bearer callback-2"]); +} + +#[rstest] +#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] +#[case::provider_beats_static_token( + None, + Value::Null, + json!({"azure_ad_token": "static-token"}), + "Bearer callback-1", + 1 +)] +#[case::header_wins_on_the_wire_but_provider_still_runs( + None, + json!({"Authorization": "Bearer override"}), + json!({}), + "Bearer override", + 1 +)] +#[tokio::test] +async fn credential_precedence( + #[case] api_key: Option<&str>, + #[case] extra_headers: Value, + #[case] optional_params: Value, + #[case] expected_authorization: &str, + #[case] expected_calls: usize, +) { + let provider = CountingToken::new(numbered_token); + let upstream = upstream([ocr_page()]).await; + + perform(azure_request( + &provider, + Some(&upstream.uri()), + api_key, + extra_headers, + optional_params, + )) + .await + .unwrap(); + + assert_eq!(provider.calls(), expected_calls); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + [expected_authorization] + ); +} + +#[rstest] +#[case::missing_api_base( + false, + json!({}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: "AZURE_AI_API_BASE", + })), + 0 +)] +#[case::unsupported_oidc_reference( + true, + json!({"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), + 0 +)] +#[case::empty_provider_token_ignores_static_token( + true, + json!({"azure_ad_token": "static-token"}), + |_| String::new(), + |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + 1 +)] +#[tokio::test] +async fn credential_failures_send_no_provider_request( + #[case] with_api_base: bool, + #[case] optional_params: Value, + #[case] token: fn(usize) -> String, + #[case] expected: fn(&Error) -> bool, + #[case] expected_calls: usize, +) { + let provider = CountingToken::new(token); + let upstream = upstream([ocr_page()]).await; + let base = upstream.uri(); + + let error = perform(azure_request( + &provider, + with_api_base.then_some(base.as_str()), + None, + Value::Null, + optional_params, + )) + .await + .unwrap_err(); + + assert!(expected(&error), "unexpected error: {error:?}"); + assert_eq!(provider.calls(), expected_calls); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs b/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs new file mode 100644 index 00000000000..1921176d158 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs @@ -0,0 +1,441 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_host::event::{CallEvent, MachineEvent}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use rstest::rstest; + +use super::*; + +const MODEL: &str = "azure_ai/doc-intelligence/prebuilt-read"; + +fn read_request(base: &str, options: Value) -> LiteLLMOcrRequest { + ocr_request(MODEL, base, options) +} + +#[tokio::test] +async fn pages_features_and_extra_options_map_to_the_analyze_call() { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": []} + }))]) + .await; + let request = read_request( + &upstream.uri(), + json!({ + "pages": [2, 0, 0, 1], + "features": ["keyValuePairs", "languages"], + "future_option": {"nested": null}, + "extra_body": {"provider_option": false} + }), + ) + .with_document( + document( + json!({"type": "document_url", "document_url": "https://example.com/document.pdf"}), + ) + .into(), + ); + + perform(request).await.unwrap(); + + let sent = only_request(&upstream).await; + assert!( + sent.url.path().ends_with("/prebuilt-read:analyze"), + "{}", + sent.url + ); + assert_eq!(sent.query("pages").as_deref(), Some("1,2,3")); + assert_eq!( + sent.query("features").as_deref(), + Some("keyValuePairs,languages") + ); + assert_eq!( + sent.json(), + json!({ + "urlSource": "https://example.com/document.pdf", + "future_option": {"nested": null}, + "provider_option": false + }) + ); +} + +#[rstest] +#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))] +#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))] +#[case(json!({"features": "languages&pages=1"}), Error::Features)] +#[case(json!({"req_format": "azure"}), Error::RequestFormat)] +#[tokio::test] +async fn invalid_pages_features_and_format_are_rejected_before_sending( + #[case] options: Value, + #[case] expected: Error, +) { + let upstream = upstream([json_response(json!({}))]).await; + + let result = match decode_request(wire( + MODEL, + &upstream.uri(), + json!({"type": "document_url", "document_url": "https://example.com/a.pdf"}), + options.clone(), + )) { + Ok(request) => perform(request).await, + Err(error) => Err(error), + }; + + assert!( + received(&upstream).await.is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[rstest] +#[case::no_options(json!({}))] +#[case::litellm_format(json!({"req_format": "litellm"}))] +#[tokio::test] +async fn an_inline_document_is_sent_as_base64_and_only_page_text_is_kept(#[case] options: Value) { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": [{"pageNumber": 1, "lines": [{"content": "hello"}]}]} + }))]) + .await; + + let response = perform(read_request(&upstream.uri(), options)) + .await + .unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + for field in ["content", "tables", "keyValuePairs"] { + assert_eq!(serialized.get(field), Some(&Value::Null), "{field}"); + } + let sent = only_request(&upstream).await; + for field in ["pages", "features", "req_format"] { + assert_eq!(sent.query(field), None, "{field}"); + } + assert_eq!(sent.json(), json!({"base64Source": "YWJj"})); +} + +#[tokio::test] +async fn native_format_normalizes_pages_and_keeps_the_provider_response() { + let operation = json!({ + "status": "succeeded", + "operationExtension": 42, + "analyzeResult": { + "content": "A\n\nB", + "tables": [{"cells": []}], + "keyValuePairs": [{"key": {"content": "A"}}], + "pages": [{ + "pageNumber": "2", + "width": "8.5", + "height": 11, + "unit": "inch", + "lines": [{"content": "A"}, {"content": null}, {"content": "B"}] + }] + } + }); + let upstream = upstream([json_response(operation.clone())]).await; + + let result = perform(read_request( + &upstream.uri(), + json!({"req_format": "native"}), + )) + .await + .unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width": 816, "height": 1056, "dpi": 96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells": []}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key": {"content": "A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); +} + +#[tokio::test] +async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": [{"pageNumber": 1, "width": 8.5, "height": 11, "unit": "inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = + litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({}))) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream) + .await + .query("api-version") + .as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width": 612, "height": 792, "dpi": 72}) + ); +} + +#[tokio::test] +async fn an_accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status": "succeeded", "analyzeResult": {"pages": []}}); + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "running"})).insert_header("Retry-After", "0"), + json_response(operation.clone()), + ], + ) + .await; + let request = with_headers( + read_request(&upstream.uri(), json!({"req_format": "native"})), + &[("X-Trace", "initial-only")], + ); + + let result = perform(request).await.unwrap(); + + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + let requests = received(&upstream).await; + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].header("x-trace"), Some("initial-only")); + for poll in &requests[1..] { + assert_eq!(poll.method.as_str(), "GET"); + assert_eq!(poll.url.path(), "/operation"); + assert_eq!(poll.header("x-trace"), None); + assert_eq!(poll.header("ocp-apim-subscription-key"), Some("test-key")); + } +} + +#[tokio::test] +async fn polling_forwards_bearer_credentials() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + let request = with_headers( + without_api_key(read_request(&upstream.uri(), json!({}))), + &[("Authorization", "Bearer token")], + ); + + perform(request).await.unwrap(); + + assert_eq!( + received(&upstream).await[1].header("authorization"), + Some("Bearer token") + ); +} + +#[tokio::test] +async fn response_received_fires_for_the_submission_and_the_completed_poll() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({"submitted": true})), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = + LocalOcrHost::new(read_request(&upstream.uri(), json!({}))).with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + perform_with(host).await.unwrap(); + + assert_eq!(received(&upstream).await.len(), 2); + assert_eq!( + *observed.lock().unwrap(), + [r#"{"submitted":true}"#, r#"{"status":"succeeded"}"#] + ); +} + +#[tokio::test] +async fn polling_does_not_follow_redirects() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + ResponseTemplate::new(302) + .insert_header("Location", format!("{}/redirected", upstream.uri())), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(received(&upstream).await.len(), 2); +} + +#[tokio::test] +async fn a_failed_operation_is_an_error() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "failed"})), + ], + ) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status failed"), "{error}"); +} + +#[tokio::test] +async fn the_polling_deadline_bounds_the_retry_delay() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "notStarted"})).insert_header("Retry-After", "9999"), + ], + ) + .await; + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: Duration::from_millis(100), + ..OcrSettings::default() + }); + + let error = tokio::time::timeout( + Duration::from_secs(1), + litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({}))), + ) + .await + .expect("the deadline cuts the retry delay short") + .unwrap_err(); + + assert!(error.to_string().contains("timed out"), "{error}"); +} + +#[rstest] +#[case::null_pages(json!({"pages": null}), "pages")] +#[case::null_page(json!({"pages": [null]}), "pages[0]")] +#[case::null_lines(json!({"pages": [{"lines": null}]}), "lines")] +#[case::bad_width(json!({"pages": [{"width": "bad"}]}), "width")] +#[tokio::test] +async fn malformed_provider_pages_report_the_response_path( + #[case] analysis: Value, + #[case] path: &str, +) { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": analysis + }))]) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains(path), "{error}"); +} + +#[rstest] +#[case::missing(None)] +#[case::relative(Some("/relative"))] +#[case::cross_origin(Some("http://example.com/operation"))] +#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))] +#[tokio::test] +async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) { + let response = location + .into_iter() + .fold(ResponseTemplate::new(202), |response, location| { + response.insert_header("Operation-Location", location) + }); + let upstream = upstream([response]).await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("operation-location"), "{error}"); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn the_model_id_is_percent_encoded() { + let upstream = upstream([json_response(json!({"status": "succeeded"}))]).await; + + perform(ocr_request( + "azure_ai/doc-intelligence/a ?#é", + &upstream.uri(), + json!({}), + )) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert!( + sent.url.path().ends_with("/a%20%3F%23%C3%A9:analyze"), + "{}", + sent.url + ); +} + +#[rstest] +#[case::dot("azure_ai/doc-intelligence/.")] +#[case::dot_dot("azure_ai/doc-intelligence/..")] +#[tokio::test] +async fn dot_segment_model_ids_are_rejected(#[case] model: &str) { + let error = perform(ocr_request(model, UNREACHABLE_BASE, json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("dot segment"), "{error}"); +} diff --git a/litellm-rust/crates/core/tests/ocr/cohere.rs b/litellm-rust/crates/core/tests/ocr/cohere.rs new file mode 100644 index 00000000000..007aa49a2fd --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/cohere.rs @@ -0,0 +1,42 @@ +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::cohere("cohere/parse-v5.0", "/v2/parse")] +#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "/providers/cohere/v2/parse")] +#[tokio::test] +async fn an_image_goes_to_the_parse_endpoint_with_the_bearer_key( + #[case] model: &str, + #[case] path: &str, +) { + let upstream = upstream([pages_response()]).await; + let request = ocr_request_with_document( + model, + &upstream.uri(), + json!({"type": "image_url", "image_url": "data:image/png;base64,YWJj"}), + json!({}), + ); + + perform(request).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.method.as_str(), "POST"); + assert_eq!(sent.url.path(), path); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); +} + +#[rstest] +#[tokio::test] +async fn a_non_image_document_is_rejected_before_sending( + #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, +) { + let upstream = upstream([pages_response()]).await; + + let error = perform(ocr_request(model, &upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/documents.rs b/litellm-rust/crates/core/tests/ocr/documents.rs new file mode 100644 index 00000000000..e29ff3e9ee9 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/documents.rs @@ -0,0 +1,182 @@ +use base64::Engine; +use litellm_core::ocr::types::OcrDocumentInput; +use litellm_host::event::WireRequest; +use rstest::rstest; +use wiremock::{Mock, matchers::any}; + +use super::*; + +const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Guardrail { + Detached, + ReplacesDocument, +} + +impl Guardrail { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::Detached | Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +/// Serves [`SERVED_DOCUMENT`] as `image/png` to every request. +async fn document_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_raw(SERVED_DOCUMENT, "image/png")) + .mount(&server) + .await; + server +} + +/// Sends a remote document through `route` and returns the document the provider saw. +async fn provider_document(route: Route, guardrail: Guardrail) -> Value { + let documents = document_server().await; + let upstream = upstream([pages_response()]).await; + let document_type = route.document_type(); + let request = ocr_request_with_document( + route.model(), + &upstream.uri(), + json!({"type": document_type, document_type: format!("{}/scan.png", documents.uri())}), + route.options(), + ); + let host = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(guardrail.before_send(wire))); + + perform_with(host).await.unwrap(); + + only_request(&upstream).await.json()["document"][document_type].clone() +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let expected = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ); + + assert_eq!( + provider_document(route, Guardrail::Detached).await, + expected + ); +} + +#[rstest] +#[tokio::test] +async fn a_document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + assert_eq!( + provider_document(route, Guardrail::ReplacesDocument).await, + REPLACED_DOCUMENT + ); +} + +#[tokio::test] +async fn an_empty_byte_document_fails_before_sending() { + let upstream = upstream([pages_response()]).await; + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Bytes { + bytes: Default::default(), + file_name: None, + mime_type: None, + }, + ); + + let error = perform(request).await.unwrap_err(); + + assert!(matches!(error, Error::EmptyFile), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn a_missing_path_document_fails_before_sending() { + let upstream = upstream([pages_response()]).await; + let path = + std::env::temp_dir().join(format!("litellm-ocr-missing-{}.png", rand::random::())); + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + + let error = perform(request).await.unwrap_err(); + + assert!( + matches!( + &error, + Error::FileRead { path: failed, source } + if *failed == path && source.kind() == std::io::ErrorKind::NotFound + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/lifecycle.rs b/litellm-rust/crates/core/tests/ocr/lifecycle.rs new file mode 100644 index 00000000000..65e64cce79b --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/lifecycle.rs @@ -0,0 +1,269 @@ +use std::sync::{Arc, Mutex}; + +use litellm_core::ocr::{ + route::{Ocr, OcrOp, OcrProjection, ocr_machine}, + types::OcrDocumentInput, +}; +use litellm_host::{ + event::{CallEvent, MachineEvent, RequestContext, WireRequest}, + host::Host, +}; +use rstest::rstest; + +use super::*; + +pub(crate) fn event_name(event: &CallEvent) -> &'static str { + match event { + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", + CallEvent::Succeeded { .. } => "success", + CallEvent::Failed { .. } => "failure", + } +} + +fn recording_host( + request: LiteLLMOcrRequest, + events: Arc>>, + block: bool, +) -> LocalOcrHost { + let before_send_events = events.clone(); + LocalOcrHost::new(request) + .with_before_send(move |wire, _| { + before_send_events.lock().unwrap().push("before_send"); + match block { + true => Err(Error::InvalidRequest("blocked".into())), + false => Ok(wire), + } + }) + .with_observer(move |event| events.lock().unwrap().push(event_name(event))) +} + +#[tokio::test] +async fn hooks_run_in_order_and_one_success_is_emitted() { + let upstream = upstream([pages_response()]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + false, + )) + .await + .unwrap(); + + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "response", "success"] + ); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() { + let upstream = upstream([pages_response()]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + let error = perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + true, + )) + .await + .unwrap_err(); + + assert!( + matches!(&error, Error::InvalidRequest(message) if message == "blocked"), + "{error:?}" + ); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn an_upstream_failure_emits_one_terminal_failure() { + let upstream = upstream([status_response(500, json!({"error": "failed"}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + let result = perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + false, + )) + .await; + + assert!(result.is_err()); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn an_invalid_provider_response_is_observed_before_normalization_fails() { + let upstream = upstream([json_response(json!({"pages": "invalid"}))]).await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + let error = perform_with(host).await.unwrap_err(); + + assert!(matches!(error, Error::ResponseField { .. }), "{error:?}"); + assert_eq!(*observed.lock().unwrap(), [r#"{"pages":"invalid"}"#]); +} + +#[tokio::test] +async fn headers_returned_by_before_send_are_sent() { + let upstream = upstream([pages_response()]).await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))) + .with_before_send(|mut wire, _| { + wire.headers + .push(("x-core-callback".into(), "edited".into())); + Ok(wire) + }); + + perform_with(host).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header("x-core-callback"), + Some("edited") + ); +} + +async fn before_send_context(request: LiteLLMOcrRequest) -> (WireRequest, RequestContext) { + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some((wire.clone(), context.clone())); + Ok(wire) + }); + perform_with(host).await.unwrap(); + let context = observed.lock().unwrap().take(); + context.expect("before_send ran") +} + +#[tokio::test] +async fn before_send_sees_the_route_its_params_and_the_body() { + let upstream = upstream([pages_response()]).await; + + let (wire, context) = before_send_context(ocr_request( + "mistral/model", + &upstream.uri(), + json!({"pages": [0], "req_format": "native"}), + )) + .await; + + assert_eq!(context.custom_llm_provider, "mistral"); + assert_eq!(context.model, "model"); + assert_eq!(context.optional_params["req_format"], "native"); + assert!(context.secret_fields.is_empty()); + assert_eq!(wire.body["pages"], json!([0])); +} + +#[rstest] +#[case::client_secret(json!({"client_secret": "shh", "tenant_id": "t"}), &["client_secret"])] +#[case::no_secrets(json!({"tenant_id": "t"}), &[])] +#[tokio::test] +async fn before_send_names_the_secret_params(#[case] options: Value, #[case] secrets: &[&str]) { + let upstream = upstream([pages_response()]).await; + let request = ocr_request("azure_ai/model", &upstream.uri(), options).with_document( + OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + }, + ); + + let (_, context) = before_send_context(request).await; + + assert_eq!(context.secret_fields, secrets); +} + +/// Hands the route a caller-owned Azure token and rewrites the bearer in `before_send`. +struct CallerTokenHost { + request: Mutex>, + trace: Mutex>, +} + +impl Host for CallerTokenHost { + async fn project(&self) -> Result { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrProjection { + request: self.request.lock().unwrap().take().unwrap(), + caller_token: true, + }) + } + + async fn custom_op(&self, op: OcrOp) -> Result<(), Error> { + match op { + OcrOp::AcquireAzureAdToken(reply) => { + self.trace.lock().unwrap().push("token".into()); + reply.send(litellm_auth::ResolvedCredential::Static( + litellm_auth::SecretValue::new("caller-token"), + )); + Ok(()) + } + } + } + + async fn before_send( + &self, + wire: WireRequest, + _: &RequestContext, + ) -> Result { + let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); + let authorization = wire + .headers + .iter() + .find(|(name, _)| is_authorization(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + self.trace + .lock() + .unwrap() + .push(format!("before_send:{authorization}")); + let headers = wire + .headers + .into_iter() + .map(|(name, value)| match is_authorization(&name) { + true => (name, "Bearer edited".to_string()), + false => (name, value), + }) + .collect(); + Ok(WireRequest { headers, ..wire }) + } +} + +#[tokio::test] +async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { + let upstream = upstream([pages_response()]).await; + let host = CallerTokenHost { + request: Mutex::new(Some(without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({}), + )))), + trace: Mutex::new(Vec::new()), + }; + + litellm_host::run::run(ocr_machine(ocr_client()), &host) + .await + .unwrap(); + + assert_eq!( + *host.trace.lock().unwrap(), + ["project", "token", "before_send:Bearer caller-token"] + ); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer edited"] + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/machine.rs b/litellm-rust/crates/core/tests/ocr/machine.rs new file mode 100644 index 00000000000..073ce67e74b --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/machine.rs @@ -0,0 +1,284 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +use litellm_core::ocr::{ + route::{OcrMachine, OcrOp, OcrProjection}, + types::OcrDocumentInput, +}; +use litellm_host::{ + event::{CallEvent, WireRequest}, + host::{Host, HostOp}, + machine::{HostFailure, Machine, MachineStep}, +}; +use litellm_llms::base_llm::ocr::transformation::OcrTransportConfig; +use rstest::rstest; +use tokio::{io::AsyncReadExt, net::TcpListener, sync::Notify}; + +use super::{lifecycle::event_name, *}; + +/// Drives the machine by hand, answering every op through `host` except `before_send`, +/// which `intercept` answers so a test can fail or cancel exactly there. +async fn drive_until( + host: &LocalOcrHost, + mut intercept: impl FnMut(WireRequest) -> Result>, +) -> ( + Result, + Vec<&'static str>, + OcrMachine, +) { + let mut machine = ocr_machine(ocr_client()); + let mut ops = Vec::new(); + let outcome = loop { + let op = match machine.resume().await { + Ok(MachineStep::Host(op)) => op, + Ok(MachineStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + }; + let answer = match op { + HostOp::Project(reply) => { + ops.push("Project"); + host.project() + .await + .map(|projection| reply.send(projection)) + .map_err(HostFailure::Error) + } + HostOp::Custom(op) => { + ops.push(match op { + OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken", + }); + host.custom_op(op).await.map_err(HostFailure::Error) + } + HostOp::BeforeSend { wire, reply, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| reply.send(wire)) + } + HostOp::Emit(event, reply) => { + let event = CallEvent::Machine(event); + ops.push(event_name(&event)); + host.emit(&event) + .await + .map(|()| reply.send(())) + .map_err(HostFailure::Error) + } + }; + if let Err(failure) = answer { + break machine.interrupt(failure).await; + } + }; + (outcome, ops, machine) +} + +/// Answers every op until `stop` fires, leaving the machine suspended mid-call. +async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, stop: &Notify) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + tokio::select! { + _ = stop.notified() => break, + step = machine.resume() => { + match step.unwrap() { + MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), + MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), + MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), + MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), + MachineStep::Complete(_) => panic!("the stalled call completed"), + } + } + } + } + }) + .await + .expect("the call reached the stall point"); +} + +#[tokio::test] +async fn a_hand_driven_machine_performs_the_same_call() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "native"}] + }))]) + .await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))); + + let (outcome, ops, mut machine) = drive_until(&host, Ok).await; + + assert_eq!(outcome.unwrap().pages[0].markdown, "native"); + assert_eq!(received(&upstream).await.len(), 1); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); + assert!(matches!( + machine.resume().await, + Err(Error::InvalidRequest(_)) + )); +} + +#[tokio::test] +async fn a_path_document_is_read_by_core_without_a_host_operation() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Path { + path, + mime_type: None, + }, + ); + + let (response, ops, _) = drive_until(&LocalOcrHost::new(request), Ok).await; + std::fs::remove_dir_all(&dir).unwrap(); + + assert_eq!(response.unwrap().pages[0].markdown, "path"); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); + assert_eq!( + only_request(&upstream).await.json()["document"]["image_url"], + "data:image/png;base64,YWJj" + ); +} + +#[rstest] +#[case::failed(HostFailure::Error(Error::InvalidRequest("before_send failed".into())), "before_send failed")] +#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")] +#[tokio::test] +async fn a_before_send_failure_ends_the_call_without_reaching_transport( + #[case] failure: HostFailure, + #[case] message: &str, +) { + let upstream = upstream([pages_response()]).await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))); + let failure = Arc::new(std::sync::Mutex::new(Some(failure))); + + let (outcome, ops, mut machine) = drive_until(&host, |_| { + Err(failure + .lock() + .unwrap() + .take() + .expect("before_send is asked once")) + }) + .await; + + assert!( + matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message), + "{outcome:?}" + ); + assert_eq!(ops, ["Project", "BeforeSend"]); + assert!(machine.resume().await.is_err()); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn resuming_before_answering_keeps_the_pending_operation() { + let request = ocr_request("mistral/model", UNREACHABLE_BASE, json!({})); + let mut machine = ocr_machine(ocr_client()); + let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else { + panic!("expected the projection op first"); + }; + + assert!(machine.resume().await.is_err()); + reply.send(OcrProjection { + request, + caller_token: false, + }); + assert!(matches!( + machine.resume().await, + Ok(MachineStep::Host(HostOp::BeforeSend { .. })) + )); +} + +#[derive(Debug)] +struct PendingToken { + entered: Arc, + dropped: Arc, +} + +struct TokenFutureDrop(Arc); + +impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } +} + +#[tokio::test] +async fn interrupt_drops_provider_captures_before_returning() { + let entered = Arc::new(Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let mut request = ocr_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + request.transport = OcrTransportConfig { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.transport + }; + request.azure_ad_token_provider = Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))); + let host = LocalOcrHost::new(request); + let mut machine = ocr_machine(ocr_client()); + + drive_until_notified(&mut machine, &host, &entered).await; + assert!(!dropped.load(Ordering::SeqCst)); + let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest( + "cancelled".into(), + ))); + + assert!( + dropped.load(Ordering::SeqCst), + "interrupt returned while provider captures were still alive" + ); + assert!( + matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled") + ); +} + +#[tokio::test] +async fn interrupting_an_in_flight_provider_request_closes_its_connection() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let received = Arc::new(Notify::new()); + let server_received = received.clone(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0u8; 4096]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut buffer).await.unwrap(); + request.extend_from_slice(&buffer[..read]); + } + server_received.notify_one(); + while socket.read(&mut buffer).await.unwrap() != 0 {} + }); + let host = LocalOcrHost::new(ocr_request("mistral/model", &base, json!({}))); + let mut machine = ocr_machine(ocr_client()); + + drive_until_notified(&mut machine, &host, &received).await; + let cancelled = Error::InvalidRequest("cancelled".into()); + + assert!( + machine + .interrupt(HostFailure::Cancelled(cancelled)) + .await + .is_err() + ); + tokio::time::timeout(Duration::from_secs(1), server) + .await + .expect("the provider connection stayed open after the interrupt") + .unwrap(); +} diff --git a/litellm-rust/crates/core/tests/ocr/main.rs b/litellm-rust/crates/core/tests/ocr/main.rs new file mode 100644 index 00000000000..1a915389b20 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/main.rs @@ -0,0 +1,125 @@ +use litellm_core::ocr::{ + document::prepare_document, + route::{LocalOcrHost, ocr_machine}, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, +}; +use litellm_llms::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{LiteLLMOcrResponse, OcrDocument}, +}; +use serde_json::{Map, Value, json}; +use wiremock::{MockServer, ResponseTemplate}; + +#[path = "../support/mod.rs"] +mod support; +use support::*; + +mod aws_textract; +mod azure_ai; +mod azure_document_intelligence; +mod cohere; +mod documents; +mod lifecycle; +mod machine; +mod mistral; +mod reducto; +mod vertex_ai; + +const INLINE_PDF: &str = "data:application/pdf;base64,YWJj"; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) +} + +async fn perform(request: LiteLLMOcrRequest) -> Result { + litellm_core::ocr::client::perform(&ocr_client(), request).await +} + +async fn perform_with(host: LocalOcrHost) -> Result { + litellm_host::run::run(ocr_machine(ocr_client()), &host).await +} + +fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequest { + OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: object(options), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + } +} + +/// A request for an inline PDF, authenticated with `test-key`. +fn ocr_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + ocr_request_with_document( + model, + base, + json!({"type": "document_url", "document_url": INLINE_PDF}), + options, + ) +} + +fn ocr_request_with_document( + model: &str, + base: &str, + document: Value, + options: Value, +) -> LiteLLMOcrRequest { + decode_request(wire(model, base, document, options)).expect("request decodes") +} + +fn document(value: Value) -> OcrDocument { + serde_json::from_value(value).expect("document parses") +} + +/// Points the request's resolved document at `source`, keeping its type. +fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let resolved = request + .map_document(prepare_document) + .expect("document resolves"); + let document = resolved.document.clone().with_source(source.into()); + resolved.with_document(document.into()) +} + +fn with_headers(request: LiteLLMOcrRequest, headers: &[(&str, &str)]) -> LiteLLMOcrRequest { + let mut request = request; + request.transport.extra_headers = headers + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + request +} + +fn without_api_key(request: LiteLLMOcrRequest) -> LiteLLMOcrRequest { + let mut request = request; + request.credentials.api_key = None; + request +} + +fn pages_response() -> ResponseTemplate { + json_response(json!({"pages": []})) +} + +/// An Azure Document Intelligence 202 whose operation lives on `server`. +fn accepted(server: &MockServer, body: Value) -> ResponseTemplate { + ResponseTemplate::new(202) + .insert_header("Operation-Location", format!("{}/operation", server.uri())) + .set_body_json(body) +} diff --git a/litellm-rust/crates/core/tests/ocr/mistral.rs b/litellm-rust/crates/core/tests/ocr/mistral.rs new file mode 100644 index 00000000000..f80e564b03f --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/mistral.rs @@ -0,0 +1,248 @@ +use std::sync::Arc; + +use litellm_auth_gcp::VertexAuth; +use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, +}; +use litellm_llms::{ + base_llm::ocr::{ + settings::OcrSettings, + transformation::{BaseOcrConfig, OCR_RESPONSE_MAX_BYTES}, + }, + mistral::ocr::transformation::MistralOcrConfig, +}; +use rstest::rstest; + +use super::*; + +#[tokio::test] +async fn direct_mistral_sends_one_request_with_every_option() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello", "custom": "preserved"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + + let result = perform(ocr_request( + "mistral/model", + &upstream.uri(), + json!({"pages": "0,2-4", "extract_header": true, "unknown": "ignored"}), + )) + .await + .unwrap(); + + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/v1/ocr"); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + assert_eq!( + sent.json(), + json!({ + "model": "model", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "pages": "0,2-4", + "extract_header": true, + "unknown": "ignored" + }) + ); +} + +#[rstest] +#[case::litellm_format(json!({}), false)] +#[case::native_format(json!({"req_format": "native"}), true)] +#[tokio::test] +async fn the_native_response_is_kept_only_when_requested( + #[case] options: Value, + #[case] kept: bool, +) { + let provider_response = json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1}, + "provider_only": "preserved" + }); + let upstream = upstream([json_response(provider_response.clone())]).await; + + let response = perform(ocr_request("mistral/model", &upstream.uri(), options)) + .await + .unwrap(); + + assert_eq!( + response.provider_native_response.map(Value::Object), + kept.then_some(provider_response) + ); +} + +#[rstest] +#[case::mistral("mistral/model", json!({}))] +#[case::vertex( + "vertex_ai/mistral-ocr-latest", + json!({"vertex_project": "test-project", "vertex_location": "us-central1"}) +)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_whole_body_and_headers( + #[case] model: &str, + #[case] options: Value, +) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let upstream = upstream([status_response(422, payload) + .insert_header("Retry-After", "17") + .insert_header("X-Request-ID", "request-123") + .insert_header("X-Future-Header", "retained")]) + .await; + + let error = perform(ocr_request(model, &upstream.uri(), options)) + .await + .unwrap_err(); + + assert_eq!(received(&upstream).await.len(), 1); + let Error::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value), + "{name} missing from {headers:?}" + ); + } + assert_eq!(body, expected_body); +} + +#[rstest] +#[case::mistral_prefix("mistral/model", None, true)] +#[case::unknown_provider("model", Some("unknown"), false)] +fn decoding_accepts_known_providers_and_rejects_unknown_ones( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] accepted: bool, +) { + let request = OcrWireRequest { + custom_llm_provider: provider.map(Into::into), + ..wire( + model, + "https://example.com", + json!({"type": "document_url", "document_url": "https://example.com/doc.pdf"}), + json!({"extract_header": true, "unknown": 42}), + ) + }; + + assert_eq!(decode_request(request).is_ok(), accepted); +} + +#[rstest] +#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] +#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] +#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] +#[tokio::test] +async fn missing_credentials_come_from_the_injected_secret_source( + #[case] secrets: &[(&str, &str)], + #[case] expected_key: &str, +) { + let upstream = upstream([pages_response()]).await; + let base = upstream.uri(); + let source = Arc::new(RecordingSecrets::new( + secrets + .iter() + .copied() + .chain([("MISTRAL_AZURE_API_BASE", base.as_str())]), + )); + let client = ocr_client().with_secrets(source.clone()); + let request = decode_request(OcrWireRequest { + api_key: None, + api_base: None, + ..wire( + "mistral/model", + &base, + json!({"type": "document_url", "document_url": INLINE_PDF}), + json!({}), + ) + }) + .unwrap(); + + litellm_core::ocr::client::perform(&client, request) + .await + .unwrap(); + + assert_eq!(source.requested(), MistralOcrConfig.secret_names()); + assert_eq!( + only_request(&upstream).await.header("authorization"), + Some(format!("Bearer {expected_key}").as_str()) + ); +} + +#[tokio::test] +async fn the_client_uses_the_injected_http_pool_configuration() { + let upstream = upstream([pages_response()]).await; + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&settings).config, + UrlPolicy::default(), + VertexAuth::default(), + OcrSettings::default(), + Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + ) + .unwrap(); + + litellm_core::ocr::client::perform( + &client, + ocr_request("mistral/model", &upstream.uri(), json!({})), + ) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream).await.header("user-agent"), + Some("host-owned/1") + ); +} + +#[test] +fn a_valid_response_limit_is_consumed_and_not_forwarded() { + let request = ocr_request( + "mistral/model", + UNREACHABLE_BASE, + json!({"max_response_bytes": 123}), + ); + + assert_eq!(request.transport.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); +} + +#[rstest] +#[case::zero(json!(0))] +#[case::negative(json!(-1))] +#[case::boolean(json!(true))] +#[case::string(json!("123"))] +#[case::fraction(json!(1.5))] +#[case::above_the_cap(json!(OCR_RESPONSE_MAX_BYTES + 1))] +#[case::null(Value::Null)] +fn an_invalid_response_limit_is_rejected(#[case] limit: Value) { + let Err(error) = decode_request(wire( + "mistral/model", + UNREACHABLE_BASE, + json!({"type": "document_url", "document_url": INLINE_PDF}), + json!({"max_response_bytes": limit}), + )) else { + panic!("invalid response limit {limit} accepted"); + }; + + assert!(error.to_string().contains("max_response_bytes"), "{error}"); +} diff --git a/litellm-rust/crates/core/tests/ocr/reducto.rs b/litellm-rust/crates/core/tests/ocr/reducto.rs new file mode 100644 index 00000000000..8ccab27e58d --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/reducto.rs @@ -0,0 +1,321 @@ +use std::sync::{Arc, Mutex}; + +use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; +use rstest::rstest; + +use super::*; + +fn upload_response() -> ResponseTemplate { + json_response(json!({"file_id": "reducto://uploaded.pdf"})) +} + +fn chunks_response(chunks: Value) -> ResponseTemplate { + json_response(json!({"result": {"chunks": chunks}})) +} + +fn source_field(model: &str) -> &'static str { + match model.ends_with("parse-legacy") { + true => "document_url", + false => "input", + } +} + +#[rstest] +#[case::v3( + "reducto/parse-v3", + json!({ + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + "reducto://already.pdf", + json!({ + "input": "reducto://already.pdf", + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + "future_ocr_option": true, + "provider_option": "value" + }) +)] +#[case::legacy( + "reducto/parse-legacy", + json!({ + "enhance": {"agentic": [{"type": "table"}]}, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url": "reducto://legacy.pdf", + "options": {"enhance": {"agentic": [{"type": "table"}]}}, + "future_ocr_option": true, + "provider_option": "value" + }) +)] +#[tokio::test] +async fn an_uploaded_document_is_parsed_with_mapped_options( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, +) { + let upstream = upstream([chunks_response(json!([]))]).await; + + perform(with_source( + ocr_request(model, &upstream.uri(), options), + source, + )) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), expected); +} + +#[rstest] +#[tokio::test] +async fn an_inline_document_is_uploaded_as_multipart_then_parsed( + #[values("parse-v3", "parse-legacy")] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, +) { + let upstream = upstream([ + upload_response(), + chunks_response(json!([{"content": "hello"}])), + ]) + .await; + let data_uri = format!("data:{mime_type};base64,YWJj"); + let document = match mime_type.starts_with("image/") { + true => json!({"type": "image_url", "image_url": data_uri}), + false => json!({"type": "document_url", "document_url": data_uri}), + }; + let request = with_headers( + ocr_request_with_document( + &format!("reducto/{model}"), + &upstream.uri(), + document, + json!({}), + ), + &[ + ("Content-Type", "application/json"), + ("X-Trace", "upload-test"), + ], + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.pages[0].markdown, "hello"); + let requests = received(&upstream).await; + let [upload, parse] = requests.as_slice() else { + panic!( + "expected an upload and a parse, got {} requests", + requests.len() + ); + }; + assert_eq!(upload.url.path(), "/upload"); + assert!( + upload + .header("content-type") + .is_some_and(|value| value.starts_with("multipart/form-data; boundary=")), + "{:?}", + upload.header("content-type") + ); + assert_eq!(upload.header("x-trace"), Some("upload-test")); + let multipart = upload.body_text(); + assert!( + multipart.contains(&format!("Content-Type: {mime_type}\r\n")), + "{multipart}" + ); + assert!(multipart.contains("\r\n\r\nabc\r\n--"), "{multipart}"); + assert_eq!(parse.url.path(), "/parse"); + assert_eq!( + parse.json(), + json!({source_field(model): "reducto://uploaded.pdf"}) + ); + for request in &requests { + assert_eq!(request.header("authorization"), Some("Bearer test-key")); + } +} + +#[tokio::test] +async fn response_received_fires_once_for_the_parse_response() { + let upstream = upstream([upload_response(), chunks_response(json!([]))]).await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + perform_with(host).await.unwrap(); + + assert_eq!(received(&upstream).await.len(), 2); + assert_eq!(*observed.lock().unwrap(), [r#"{"result":{"chunks":[]}}"#]); +} + +#[rstest] +#[case::empty_id(json_response(json!({"file_id": ""})))] +#[case::missing_id(json_response(json!({})))] +#[case::null_id(json_response(json!({"file_id": null})))] +#[case::upload_failure(status_response(503, json!({"error": "unavailable"})))] +#[tokio::test] +async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) { + let upstream = upstream([upload]).await; + + let result = perform(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))).await; + + assert!(result.is_err()); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[rstest] +#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)] +#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })] +#[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)] +#[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] +#[tokio::test] +async fn invalid_document_sources_are_rejected_before_sending( + #[case] source: &str, + #[case] expected: Error, +) { + let upstream = upstream([json_response(json!({}))]).await; + + let result = perform(with_source( + ocr_request("reducto/parse-v3", &upstream.uri(), json!({})), + source, + )) + .await; + + assert!( + received(&upstream).await.is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[tokio::test] +async fn a_forwarded_authorization_wins_and_the_native_response_is_omitted_by_default() { + let upstream = upstream([json_response( + json!({"job_id": "job-1", "result": {"chunks": []}}), + )]) + .await; + let request = with_headers( + with_source( + ocr_request("reducto/parse-v3", &upstream.uri(), json!({})), + "reducto://ready.pdf", + ), + &[("authorization", "Bearer existing")], + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.provider_native_response, None); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer existing"] + ); +} + +#[tokio::test] +async fn native_format_retains_the_provider_response() { + let raw = json!({ + "result": {"chunks": [{"content": "native OCR response"}]}, + "usage": {"num_pages": 1} + }); + let upstream = upstream([json_response(raw.clone())]).await; + + let response = perform(with_source( + ocr_request( + "reducto/parse-v3", + &upstream.uri(), + json!({"req_format": "native"}), + ), + "reducto://ready.pdf", + )) + .await + .unwrap(); + + assert_eq!(response.pages[0].markdown, "native OCR response"); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(raw) + ); +} + +#[tokio::test] +async fn an_unknown_model_reaches_parse_and_keeps_its_name() { + let upstream = upstream([chunks_response( + json!([{"content": "future model response"}]), + )]) + .await; + + let response = perform(with_source( + ocr_request("reducto/future-parse-model", &upstream.uri(), json!({})), + "reducto://ready.pdf", + )) + .await + .unwrap(); + + assert_eq!(response.model, "future-parse-model"); + assert_eq!(response.pages[0].markdown, "future model response"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), json!({"input": "reducto://ready.pdf"})); +} + +#[tokio::test] +async fn a_guardrail_can_replace_the_document_before_upload() { + let upstream = upstream([chunks_response(json!([]))]).await; + let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))) + .with_before_send(|wire, _| { + assert_eq!(wire.body["document_url"], INLINE_PDF); + Ok(WireRequest { + body: json!({"type": "document_url", "document_url": "reducto://guarded.pdf"}), + ..wire + }) + }); + + perform_with(host).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), json!({"input": "reducto://guarded.pdf"})); +} + +#[rstest] +#[tokio::test] +async fn guardrail_headers_reach_both_upload_and_parse( + #[values("reducto/parse-v3", "reducto/parse-legacy")] model: &str, +) { + let upstream = upstream([upload_response(), chunks_response(json!([]))]).await; + let request = with_headers( + ocr_request(model, &upstream.uri(), json!({})), + &[("authorization", "Bearer original")], + ); + let host = LocalOcrHost::new(request).with_before_send(|wire, _| { + Ok(WireRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..wire + }) + }); + + perform_with(host).await.unwrap(); + + let requests = received(&upstream).await; + let paths: Vec<&str> = requests.iter().map(|request| request.url.path()).collect(); + assert_eq!(paths, ["/upload", "/parse"]); + for request in &requests { + assert_eq!(request.header_values("authorization"), ["Bearer guarded"]); + } +} diff --git a/litellm-rust/crates/core/tests/ocr/vertex_ai.rs b/litellm-rust/crates/core/tests/ocr/vertex_ai.rs new file mode 100644 index 00000000000..f0b2488e494 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/vertex_ai.rs @@ -0,0 +1,184 @@ +use litellm_auth::{InputSource, Sourced}; +use litellm_core::ocr::arguments::is_supported_request; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use rstest::rstest; + +use super::*; + +#[tokio::test] +async fn mistral_is_served_at_the_resolved_project_and_location() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + + let response = perform(ocr_request( + "vertex_ai/mistral-ocr-maas", + &upstream.uri(), + json!({ + "vertex_project": "project-1", + "vertex_location": "europe-west4", + "extract_footer": true + }), + )) + .await + .unwrap(); + + assert_eq!(response.pages[0].markdown, "hello"); + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + assert_eq!( + sent.json(), + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "extract_footer": true + }) + ); +} + +#[tokio::test] +async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let upstream = upstream([pages_response()]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + litellm_core::ocr::client::perform( + &client, + ocr_request("vertex_ai/mistral-ocr-maas", &upstream.uri(), json!({})), + ) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream).await.url.path(), + "/v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); +} + +#[tokio::test] +async fn a_supplied_authorization_is_forwarded_without_a_static_token() { + let upstream = upstream([pages_response()]).await; + let request = with_headers( + without_api_key(ocr_request( + "vertex_ai/model", + &upstream.uri(), + json!({"vertex_project": "project-1"}), + )), + &[("authorization", "Bearer supplied")], + ); + + perform(request).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer supplied"] + ); +} + +#[tokio::test] +async fn invalid_credentials_fail_before_sending() { + let error = perform(ocr_request( + "vertex_ai/model", + UNREACHABLE_BASE, + json!({"vertex_credentials": true}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("vertex_credentials"), "{error}"); +} + +#[rstest] +#[tokio::test] +async fn a_request_controlled_api_base_is_rejected_before_vertex_auth( + #[values("vertex_ai/mistral-ocr-maas", "vertex_ai/deepseek-ocr-maas")] model: &str, +) { + let mut request = ocr_request( + model, + "https://caller.example", + json!({"vertex_project": "project-1"}), + ); + request.credentials.api_base = Some(Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform(request).await.unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint"), + "{error}" + ); +} + +#[tokio::test] +async fn deepseek_is_served_at_the_openai_compatible_endpoint() { + let upstream = upstream([json_response(json!({ + "choices": [{"message": {"content": "recognized"}}], + "usage": {"prompt_tokens": 1} + }))]) + .await; + let request = with_source( + ocr_request( + "vertex_ai/deepseek-ocr-maas", + &upstream.uri(), + json!({ + "vertex_project": "project-1", + "vertex_location": "europe-west4", + "temperature": 0.1, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + ), + "gs://bucket/document.pdf", + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions" + ); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + let body = sent.json(); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type": "image_url", "image_url": "gs://bucket/document.pdf"}) + ); +} + +#[rstest] +#[case::deepseek("deepseek-ocr-maas", Some("vertex_ai"), true)] +#[case::mistral("mistral-ocr-maas", Some("vertex_ai"), true)] +#[case::prefixed("vertex_ai/mistral-ocr-maas", None, true)] +#[case::unknown_provider("model", Some("unknown"), false)] +fn supported_requests_follow_the_registered_configs( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] supported: bool, +) { + assert_eq!(is_supported_request(model, provider), supported); +} diff --git a/litellm-rust/crates/core/tests/support/mod.rs b/litellm-rust/crates/core/tests/support/mod.rs new file mode 100644 index 00000000000..4d2fe0232d0 --- /dev/null +++ b/litellm-rust/crates/core/tests/support/mod.rs @@ -0,0 +1,155 @@ +//! Shared fixtures for route integration tests: a scripted upstream and a recording +//! secret source. + +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::sync::Mutex; + +use futures_util::future::BoxFuture; +use litellm_secrets::{SecretValue, source::SecretSource}; +use serde_json::Value; +use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; + +/// A port nothing listens on, for calls that must fail before any request is sent. +pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1"; + +/// Starts an upstream that answers its n-th request with the n-th response and 404s after. +pub async fn upstream(responses: impl IntoIterator) -> MockServer { + let server = MockServer::start().await; + respond_in_order(&server, responses).await; + server +} + +/// Scripts responses on a started server, for responses that need its address. +pub async fn respond_in_order( + server: &MockServer, + responses: impl IntoIterator, +) { + for response in responses { + Mock::given(any()) + .respond_with(response) + .up_to_n_times(1) + .mount(server) + .await; + } +} + +pub async fn received(server: &MockServer) -> Vec { + server + .received_requests() + .await + .expect("request recording is on") +} + +pub async fn only_request(server: &MockServer) -> Request { + let [request] = <[Request; 1]>::try_from(received(server).await) + .unwrap_or_else(|requests| panic!("expected one request, got {}", requests.len())); + request +} + +pub fn json_response(body: Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(body) +} + +pub fn status_response(status: u16, body: Value) -> ResponseTemplate { + ResponseTemplate::new(status).set_body_json(body) +} + +pub trait ReceivedRequest { + fn header(&self, name: &str) -> Option<&str>; + fn header_values(&self, name: &str) -> Vec<&str>; + fn json(&self) -> Value; + fn body_text(&self) -> String; + /// The path and query, as the request line carried them. + fn target(&self) -> String; + fn query(&self, name: &str) -> Option; +} + +impl ReceivedRequest for Request { + fn header(&self, name: &str) -> Option<&str> { + self.headers.get(name).and_then(|value| value.to_str().ok()) + } + + fn header_values(&self, name: &str) -> Vec<&str> { + self.headers + .get_all(name) + .iter() + .filter_map(|value| value.to_str().ok()) + .collect() + } + + fn json(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is json") + } + + fn body_text(&self) -> String { + String::from_utf8_lossy(&self.body).into_owned() + } + + fn target(&self) -> String { + match self.url.query() { + Some(query) => format!("{}?{query}", self.url.path()), + None => self.url.path().to_string(), + } + } + + fn query(&self, name: &str) -> Option { + self.url + .query_pairs() + .find_map(|(key, value)| (key == name).then(|| value.into_owned())) + } +} + +/// A secret source that answers from a fixed table and records every name it was asked for. +pub struct RecordingSecrets { + values: Vec<(String, String)>, + fails: bool, + requested: Mutex>, +} + +impl RecordingSecrets { + pub fn new<'a>(values: impl IntoIterator) -> Self { + Self { + values: values + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + fails: false, + requested: Mutex::new(Vec::new()), + } + } + + pub fn empty() -> Self { + Self::new([]) + } + + pub fn failing() -> Self { + Self { + fails: true, + ..Self::empty() + } + } + + pub fn requested(&self) -> Vec { + self.requested.lock().unwrap().clone() + } +} + +impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + if self.fails { + return Err(litellm_secrets::Error::ManagedSecretMissing); + } + Ok(self + .values + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } +} diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index f00259984ba..c3377536545 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -554,6 +554,50 @@ async fn upload_bytes_async( mod tests { use super::*; + #[tokio::test] + async fn v3_body_keeps_explicit_null_options_and_drops_unknown_ones() { + use crate::base_llm::ocr::{handler::OcrClient, transformation::OcrRequestContext}; + + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({"input":"reducto://ready.pdf", "formatting":null, "settings":{}}) + ); + let absent = ReductoParseV3Config + .map_ocr_params( + &litellm_core_utils::call_arguments::CallArguments::default(), + "parse-v3", + ) + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + #[test] fn options_preserve_null_and_select_the_provider_fields() { let overrides = serde_json::from_value(json!({ diff --git a/litellm-rust/crates/llms/tests/ocr_handler.rs b/litellm-rust/crates/llms/tests/ocr_handler.rs new file mode 100644 index 00000000000..6e46e6f76d4 --- /dev/null +++ b/litellm-rust/crates/llms/tests/ocr_handler.rs @@ -0,0 +1,79 @@ +use std::time::Duration; + +use litellm_llms::base_llm::ocr::{error::Error, handler::read_response_bytes}; +use rstest::rstest; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; + +/// Answers one request with raw `response` bytes and then holds the connection open, so a +/// read that waits for the rest of an oversized body hangs instead of passing. +async fn read_bounded(response: String, limit: usize) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(response.as_bytes()).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = + tokio::time::timeout(Duration::from_secs(2), read_response_bytes(response, limit)).await; + server.abort(); + result.expect("bounded reads must finish without waiting for the rest of an oversized body") +} + +#[rstest] +#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh")] +#[case::chunked( + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n" +)] +#[tokio::test] +async fn a_body_of_exactly_the_limit_is_read(#[case] response: &str) { + assert_eq!(read_bounded(response.into(), 8).await.unwrap(), "abcdefgh"); +} + +#[rstest] +#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n")] +#[case::chunked("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n")] +#[tokio::test] +async fn a_body_over_the_limit_is_rejected(#[case] response: &str) { + assert!(matches!( + read_bounded(response.into(), 8).await, + Err(Error::TooLarge { limit: 8 }) + )); +} + +#[rstest] +#[case::declared("Content-Length: 1000000")] +#[case::chunked("Transfer-Encoding: chunked")] +#[tokio::test] +async fn an_oversized_error_keeps_its_status_and_a_bounded_body_without_draining( + #[case] headers: &str, +) { + let prefix = "x".repeat(4096); + let body = match headers.starts_with("Transfer") { + true => format!("{:x}\r\n{prefix}\r\n", prefix.len()), + false => prefix.clone(), + }; + + let error = read_bounded( + format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"), + prefix.len(), + ) + .await + .unwrap_err(); + + let Error::Transport(litellm_http::transport::Error::Http { status, body }) = error else { + panic!("unexpected error: {error}"); + }; + assert_eq!(status, 429); + assert_eq!(body, prefix); +} diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md new file mode 100644 index 00000000000..d65ffd613aa --- /dev/null +++ b/tests/test_litellm_rust/AGENTS.md @@ -0,0 +1 @@ +This directory holds only the tests that cannot be written in the Rust code diff --git a/tests/test_litellm_rust/cache/__init__.py b/tests/test_litellm_rust/cache/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/cache/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/cache/conftest.py b/tests/test_litellm_rust/cache/conftest.py new file mode 100644 index 00000000000..07ccd0fde2b --- /dev/null +++ b/tests/test_litellm_rust/cache/conftest.py @@ -0,0 +1,30 @@ +import threading +from collections.abc import Generator +from typing import Final + +import fakeredis +import pytest + +from tests.test_litellm_rust.support.s3_stub import S3Stub + + +@pytest.fixture +def redis_url() -> Generator[str]: + server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + yield f"redis://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +@pytest.fixture +def s3_stub() -> Generator[S3Stub]: + stub: Final = S3Stub() + try: + yield stub + finally: + stub.close() diff --git a/tests/test_litellm_rust/cache/test_azure_blob.py b/tests/test_litellm_rust/cache/test_azure_blob.py new file mode 100644 index 00000000000..bbbab22baca --- /dev/null +++ b/tests/test_litellm_rust/cache/test_azure_blob.py @@ -0,0 +1,173 @@ +import asyncio +import json +import os +import time +import uuid +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, cast + +import pytest +from azure.storage.blob import ContainerClient + +from litellm.caching.azure_blob_cache import AzureBlobCache +from litellm.caching.caching import Cache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheLookup, + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + completion_kwargs, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def azure_blob_facade() -> Generator[Cache]: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + yield facade + finally: + backend.container_client.delete_container() + asyncio.run(backend.disconnect()) + + +def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + return CacheTestHandle.azure_blob( + backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), + backend.container_client.container_name, + ) + + +def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + handle: Final = azure_blob_handle(azure_blob_facade) + assert handle.backend == "azure-blob" + account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") + with pytest.raises(TypeError, match="containers must match"): + CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade( + azure_blob_facade + ) + handle._bind_facade(azure_blob_facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + + response: Final = { + "choices": [{"text": "caf\u00e9 \u2603"}], + "usage": {"total_tokens": 3}, + "flag": True, + "empty": None, + } + native.store({**request("sync"), "ttl_seconds": 0.001}, response) + native.store(request("sync"), {"choices": [{"text": "second"}]}) + time.sleep(0.01) + stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) + assert stored["response"] == response + assert isinstance(stored["timestamp"], float) + assert native.lookup(request("sync")) == response + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + backend.set_cache("python", {"timestamp": time.time(), "response": response}) + backend.set_cache("legacy", "bare legacy value") + backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) + assert native.lookup(request("python")) == response + assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") + assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { + "values": [response, None, None, response], + "missing_indices": [1, 2], + } + + with rebound(azure_blob_facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): + assert resolver.resolve().kind == "python_callback" + + def custom_get(*_args: object, **_kwargs: object) -> None: + return None + + with rebound(backend, "get_cache", custom_get): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + class CustomBlobCache(AzureBlobCache): + pass + + with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): + assert resolver.resolve().kind == "python_callback" + with pytest.raises(TypeError): + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + + +async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() + assert binding.kind == "native" + ping: Final = cast(dict[str, object], await binding.ping()) + assert ping["status"] == "success", ping + + await binding.async_store(request("async"), {"value": 1}) + await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) + time.sleep(0.01) + assert await binding.async_lookup(request("async")) == {"value": 2} + assert await backend.async_get_cache("async") == json.loads( + backend.container_client.download_blob("async").readall() + ) + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} + + await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) + assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { + "values": [{"value": 4}, None, {"value": 3}], + "missing_indices": [1], + } + await binding.async_flush() + assert [blob.name for blob in backend.container_client.list_blobs()] == [] + assert await binding.async_lookup(request("async")) is None + + +async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + assert_native_runtime(facade) + kwargs: Final = completion_kwargs("azure") + await facade.async_add_cache({"answer": "azure"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "azure"} + assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"} + finally: + backend.container_client.delete_container() + await backend.disconnect() diff --git a/tests/test_litellm_rust/cache/test_disk.py b/tests/test_litellm_rust/cache/test_disk.py new file mode 100644 index 00000000000..4f2907e6a09 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_disk.py @@ -0,0 +1,117 @@ +import asyncio +import json +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import diskcache +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.disk_cache import DiskCache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: + disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + disk_cache.disk_cache.set( + "sync", + {"timestamp": time.time(), "response": json.dumps(response)}, + ) + disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response})) + disk_cache.disk_cache.set("raw", json.dumps(response)) + disk_cache.disk_cache.set("invalid", "not a cache entry") + disk_cache.disk_cache.set( + "large", + {"timestamp": time.time(), "response": {"text": "x" * 70_000}}, + ) + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("large")) == {"text": "x" * 70_000} + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored_response: Final = disk_cache.get_cache("native") + assert isinstance(stored_response, dict) + assert stored_response["response"] == response + stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True) + assert stored is not None + assert time.time() < expire_time <= time.time() + 12.0 + await binding.async_store(request("no-ttl"), response) + _, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True) + assert no_expiry is None + + +async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None: + first: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + await first.async_store(request("persistent"), {"value": "persistent"}) + await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"}) + fresh: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + assert fresh.lookup(request("expiring")) == {"value": "expiring"} + await asyncio.sleep(0.4) + assert fresh.lookup(request("expiring")) is None + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + + +def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None: + facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + with pytest.raises(TypeError, match="directories must match"): + CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade) + handle: Final = CacheTestHandle.disk(str(tmp_path)) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + binding.store(request("native"), {"value": "native"}) + assert facade.get_cache(cache_key="native") == {"value": "native"} + + with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "native" + + class CustomDiskCache(DiskCache): + pass + + with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + + class CustomStore(diskcache.Cache): + pass + + custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + custom_facade.cache.disk_cache = CustomStore(str(tmp_path)) + with pytest.raises(TypeError, match="built-in diskcache store"): + CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade) + + +async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } diff --git a/tests/test_litellm_rust/cache/test_facade.py b/tests/test_litellm_rust/cache/test_facade.py new file mode 100644 index 00000000000..d99ea4e2baa --- /dev/null +++ b/tests/test_litellm_rust/cache/test_facade.py @@ -0,0 +1,397 @@ +import asyncio +import contextvars +import gc +import weakref +from types import SimpleNamespace +from typing import Final, cast + +import pytest + +import litellm +from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.rust_bridge import _native +from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheLookup, CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def test_existing_constructor_and_global_are_unchanged() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert type(facade.cache) is InMemoryCache + assert "_native_cache_handle" not in vars(facade) + assert resolve_response_cache(facade) is None + with rebound(litellm, "cache", facade): + resolver: Final = CacheTestResolver(litellm) + assert resolver.resolve().kind == "python_callback" + resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) + assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} + + +async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + + sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) + assert sync_request is not None + runtime.store(sync_request, {"answer": 1}) + assert runtime.lookup(sync_request) == {"answer": 1} + assert facade.cache.get_cache("sync") is None + + async_request: Final = runtime.request(facade, {"cache_key": "async"}) + assert async_request is not None + await runtime.async_store(async_request, {"answer": 2}) + assert await runtime.async_lookup(async_request) == {"answer": 2} + assert await facade.cache.async_get_cache("async") is None + + requests: Final = (sync_request, async_request) + expected: Final = { + "values": [{"answer": 1}, {"answer": 2}], + "missing_indices": [], + } + assert runtime.lookup_batch(requests) == expected + assert await runtime.async_lookup_batch(requests) == expected + + await runtime.async_flush() + assert runtime.lookup(sync_request) is None + assert await runtime.async_lookup(async_request) is None + + +async def test_inference_resolver_uses_the_configured_native_cache_directly() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + facade._native_cache = runtime + + selected: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert selected.kind == "native" + request: Final = runtime.request(facade, {"cache_key": "inference-native"}) + assert request is not None + await selected.async_store(request, {"answer": 42}) + assert await selected.async_lookup(request) == {"answer": 42} + assert await runtime.async_lookup(request) == {"answer": 42} + assert facade.cache.get_cache("inference-native") is None + + facade._native_cache = None + fallback: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert fallback.kind == "python_callback" + await fallback.async_store(None, {"answer": 7}, callback_kwargs={"cache_key": "inference-python"}) + assert facade.get_cache(cache_key="inference-python") == {"answer": 7} + assert facade.cache.get_cache("inference-python") is not None + + +async def test_inference_resolver_declines_a_native_runtime_whose_facade_changed() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + facade._native_cache = runtime + stale_request: Final = runtime.request(facade, {"cache_key": "stale-only"}) + assert stale_request is not None + await runtime.async_store(stale_request, {"answer": "stale"}) + + replacement: Final = InMemoryCache() + facade.cache = replacement + with pytest.raises(_native.RustBridgeDeclined): + _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert await runtime.async_lookup(stale_request) == {"answer": "stale"} + assert replacement.get_cache("stale-only") is None + assert replacement.get_cache("swapped-backend") is None + + +def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: + resolver: Final = CacheTestResolver(litellm) + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) + enabled: Final = litellm.cache + assert isinstance(enabled, Cache) + assert enabled.ttl == 30 + assert resolver.resolve().kind == "python_callback" + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + assert litellm.cache is enabled + + update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + updated: Final = litellm.cache + assert isinstance(updated, Cache) + assert updated is not enabled + assert updated.ttl == 60 + + disable_cache() + assert litellm.cache is None + assert resolver.resolve().kind == "disabled" + + +async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: + namespace: Final = SimpleNamespace(cache=CacheTestHandle.memory()) + resolver: Final = CacheTestResolver(namespace) + selected: Final = resolver.resolve() + assert selected.kind == "native" + selected.store(request(), {"answer": 1}) + assert await selected.async_lookup(request()) == {"answer": 1} + with rebound(namespace, "cache", CacheTestHandle.memory()): + replacement: Final = resolver.resolve() + await selected.async_store(request(), {"answer": 2}) + assert replacement.lookup(request()) is None + assert selected.lookup(request()) == {"answer": 2} + with rebound(namespace, "cache", None): + disabled: Final = resolver.resolve() + assert disabled.kind == "disabled" + assert disabled.lookup(None) is None + await disabled.async_store(None, object()) + assert await disabled.async_lookup(None) is None + assert selected.lookup(request()) == {"answer": 2} + + +async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: + context: Final = contextvars.ContextVar("cache_context", default="caller") + caller: Final = asyncio.current_task() + sentinel: Final = object() + failure: Final = RuntimeError("callback failed") + + class CustomCache: + async def async_get_cache(self, *, marker: object) -> object: + assert marker is sentinel + assert asyncio.current_task() is caller + context.set("callback") + return marker + + async def async_add_cache(self, response: object, *, marker: object) -> None: + assert response is sentinel + assert marker is sentinel + raise failure + + namespace: Final = SimpleNamespace(cache=CustomCache()) + binding: Final = CacheTestResolver(namespace).resolve() + assert binding.kind == "python_callback" + assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel + assert context.get() == "callback" + with pytest.raises(RuntimeError) as caught: + await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) + assert caught.value is failure + + +async def test_callback_cancellation_stays_in_the_callers_task() -> None: + entered: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class CustomCache: + async def async_get_cache(self) -> None: + entered.set() + try: + await asyncio.Future() + finally: + finished.set() + + binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + + async def lookup() -> object: + return await binding.async_lookup(None, callback_kwargs={}) + + task: Final = asyncio.create_task(lookup()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + + +def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle: Final = CacheTestHandle.memory() + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + native.store(request(), {"source": "native"}) + assert native.lookup(request()) == {"source": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="key") is None + sentinel: Final = object() + + def outer_override(**_kwargs: object) -> object: + return sentinel + + def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: + return {"source": "override"} + + with rebound(facade, "get_cache", outer_override): + fallback: Final = resolver.resolve() + assert fallback.kind == "python_callback" + assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache") + assert resolver.resolve().kind == "native" + with rebound(facade.cache, "get_cache", backend_override): + backend_fallback: Final = resolver.resolve() + assert backend_fallback.kind == "python_callback" + assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} + + +def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: + class CustomCache(Cache): + pass + + handle: Final = CacheTestHandle.memory() + with pytest.raises(TypeError): + handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + with rebound(facade, "cache", InMemoryCache()): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + + def custom_key(**_kwargs: object) -> str: + return "custom" + + with rebound(facade, "get_cache_key", custom_key): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache_key") + assert resolver.resolve().kind == "native" + + +def test_resolver_and_callback_cycles_can_be_collected() -> None: + class CustomCache: + pass + + def cyclic_reference() -> weakref.ReferenceType[CustomCache]: + callback: Final = CustomCache() + namespace: Final = SimpleNamespace(cache=callback) + binding: Final = CacheTestResolver(namespace).resolve() + setattr(callback, "binding", binding) + return weakref.ref(callback) + + reference: Final = cyclic_reference() + gc.collect() + assert reference() is None + + +def test_invalid_duration_and_request_shape_fail_before_storage() -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve() + for seconds in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): + binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) + assert binding.lookup(request()) is None + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): + CacheTestHandle.memory(ttl_seconds=-1) + + +async def test_memory_size_policy_is_applied_by_the_native_host() -> None: + handle: Final = CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + small: Final = {"answer": "ok"} + binding.store(request("small"), small) + assert await binding.async_lookup(request("small")) == small + await binding.async_store(request("large"), {"answer": "x" * 256}) + assert binding.lookup(request("large")) is None + assert binding.lookup(request("small")) == small + disabled: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory(capacity=0))).resolve() + await disabled.async_store(request(), small) + assert await disabled.async_lookup(request()) is None + + +async def test_native_batch_lookup_and_store_report_partial_hits() -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } + + +async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: + result: Final = object() + marker: Final = object() + + class CustomCache(Cache): + def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("sync", kwargs) + + async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("async", kwargs) + + async def async_add_cache_pipeline( + self, result: object, dynamic_cache_object: object = None, **kwargs: object + ) -> object: + return result, kwargs + + binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() + assert binding.kind == "python_callback" + requests: Final = [request("first"), request("second")] + kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] + + assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] + assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ + ("async", kwargs[0]), + ("async", kwargs[1]), + ] + with pytest.raises(ValueError, match="equal lengths"): + binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) + with pytest.raises(TypeError, match="callback_result"): + await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) + stored: Final = cast( + tuple[object, dict[str, object]], + await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), + ) + assert stored[0] is result + assert stored[1] == {"marker": marker} + + +async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: + async def ping() -> str: + return "pong" + + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache.cache.set_cache("key", "value") + binding: Final = CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + assert binding.kind == "python_callback" + + setattr(cache.cache, "ping", ping) + assert await binding.ping() == "pong" + await binding.async_flush() + assert cache.cache.get_cache("key") is None + + +def test_facade_registration_rejects_mismatched_capacity() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + with pytest.raises(TypeError, match="capacities must match"): + CacheTestHandle.memory(capacity=7)._bind_facade(facade) diff --git a/tests/test_litellm_rust/cache/test_gcs.py b/tests/test_litellm_rust/cache/test_gcs.py new file mode 100644 index 00000000000..bfc9ebbb4d7 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_gcs.py @@ -0,0 +1,242 @@ +import json +import time +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, cast + +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.gcs_cache import GCSCache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheLookup, CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.fake_gcs import FakeGcs +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None diff --git a/tests/test_litellm_rust/cache/test_qdrant_semantic.py b/tests/test_litellm_rust/cache/test_qdrant_semantic.py new file mode 100644 index 00000000000..160089c9002 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_qdrant_semantic.py @@ -0,0 +1,286 @@ +import hashlib +import http.server +import json +import math +import os +import threading +import time +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + request, + require_rust, +) + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def qdrant_request( + key: str, + messages: list[dict[str, object]], + **kwargs: object, +) -> dict[str, object]: + return {**request(key), "messages": messages, **kwargs} + + +def embedding_vector(text: str) -> list[float]: + raw: Final = hashlib.sha256(text.encode()).digest()[:8] + values: Final = [byte / 127.5 - 1 for byte in raw] + norm: Final = math.sqrt(sum(value * value for value in values)) + return [value / norm for value in values] + + +@pytest.fixture +def qdrant_url() -> str: + value: Final[str | None] = os.environ.get("QDRANT_URL") + if not value: + pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") + return value.rstrip("/") + + +@pytest.fixture +def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: + class EmbeddingHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: + length: Final = int(self.headers["Content-Length"]) + body: Final = json.loads(self.rfile.read(length)) + text: Final = body["input"] + response: Final = { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": embedding_vector(text), + } + ], + "model": body["model"], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + encoded: Final = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args: object) -> None: + return + + server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: + return Cache( + type=LiteLLMCacheType.QDRANT_SEMANTIC, + qdrant_api_base=qdrant_url, + qdrant_collection_name=collection_name, + similarity_threshold=0.99, + qdrant_semantic_cache_embedding_model="text-embedding-3-small", + qdrant_semantic_cache_vector_size=8, + ) + + +def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "shared prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + facade.cache.set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} + binding.store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = facade.cache.get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] + assert binding.lookup(qdrant_request("native-key", unrelated)) is None + assert facade.cache.get_cache("native-key", messages=unrelated) is None + assert binding.lookup(qdrant_request("different-key", messages)) is None + assert facade.cache.get_cache("different-key", messages=messages) is None + + +async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "async prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + await facade.cache.async_set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} + await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + + +async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + entries: Final = [ + qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), + qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), + ] + await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) + + assert binding.lookup(entries[0]) == {"id": "one"} + assert binding.lookup(entries[1]) == {"id": "two"} + assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { + "id": "one" + } + assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { + "id": "two" + } + + +async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( + qdrant_url: str, fake_embedding_endpoint: str +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "malformed prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + key: Final = "malformed-key" + response: Final = { + "points": [ + { + "id": str(uuid4()), + "vector": embedding_vector("malformed prompt"), + "payload": { + "litellm_cache_key": key, + "text": "malformed prompt", + "response": "not json", + }, + } + ] + } + facade.cache.sync_client.put( + url=f"{qdrant_url}/collections/{collection}/points", + headers=facade.cache.headers, + json=response, + ) + assert binding.lookup(qdrant_request(key, messages)) is None + with pytest.raises(RuntimeError, match="operation is not supported"): + binding.lookup_batch([qdrant_request(key, messages)]) + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.async_flush() + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.ping() + + +def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "persistent prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) + time.sleep(1.2) + assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} + python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "persistent"} + + +def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + facade.cache.qdrant_api_key = "rotated" + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + facade.cache.similarity_threshold = 0.5 + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + unsupported.cache.embedding_max_input_tokens = 100 + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unsupported) + unsupported.cache.embedding_max_input_tokens = None + unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" + with pytest.raises(TypeError, match="gRPC"): + handle._bind_facade(unsupported) + + +def test_qdrant_semantic_rust_required_rule_activates_natively( + qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + del fake_embedding_endpoint + require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC) + facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + assert_native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} + facade.add_cache({"answer": "qdrant"}, **kwargs) + assert facade.get_cache(**kwargs) == {"answer": "qdrant"} diff --git a/tests/test_litellm_rust/cache/test_redis.py b/tests/test_litellm_rust/cache/test_redis.py new file mode 100644 index 00000000000..dd88145ef21 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_redis.py @@ -0,0 +1,228 @@ +import json +import os +import time +from types import SimpleNamespace +from typing import Final +from urllib.parse import urlparse + +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import CacheRule +from litellm.rust_bridge.configuration import Rollout +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + completion_kwargs, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def cluster_nodes() -> tuple[tuple[str, int], ...]: + configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") + if not configured: + pytest.skip("LITELLM_TEST_REDIS_CLUSTER_NODES is not set") + return tuple((host, int(port)) for host, _, port in (node.partition(":") for node in configured.split(","))) + + +async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + namespace: Final = SimpleNamespace(cache=CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = CacheTestResolver(namespace).resolve() + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} + client.set("team:sync", str(envelope)) + client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + client.set("team:raw", json.dumps(response)) + client.set("team:invalid", "not a cache entry") + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("team:async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = client.get("team:native") + assert isinstance(stored, bytes) + assert json.loads(stored)["response"] == response + assert 0 < client.ttl("team:native") <= 12 + assert client.get("litellm-cache:team:native") is None + assert client.get("team:team:async") is None + client.close() + + +async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: + parsed: Final = urlparse(redis_url) + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache( + type=LiteLLMCacheType.REDIS, + host=parsed.hostname, + port=str(parsed.port), + redis_flush_size=2, + ) + with pytest.raises(TypeError, match="default TTLs must match"): + CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + with pytest.raises(TypeError, match="namespaces must match"): + CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(redis_url) + + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + + pool: Final = facade.cache.redis_client.connection_pool + with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + + await binding.async_store(request("first"), {"value": 1}) + assert client.get("first") is None + await binding.async_store(request("second"), {"value": 2}) + + assert client.get("first") is not None + assert client.get("second") is not None + await facade.cache.disconnect() + client.close() + + +async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( + cluster_nodes: tuple[tuple[str, int], ...], +) -> None: + startup_nodes: Final = [{"host": host, "port": port} for host, port in cluster_nodes] + url: Final = f"redis://{cluster_nodes[0][0]}:{cluster_nodes[0][1]}" + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache(type=LiteLLMCacheType.REDIS, redis_startup_nodes=startup_nodes, namespace="parity") + assert type(facade.cache) is RedisClusterCache + with pytest.raises(TypeError, match="types must match"): + CacheTestHandle.redis(url, namespace="parity")._bind_facade(facade) + CacheTestHandle.redis(url, namespace="parity", startup_nodes=list(cluster_nodes))._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + manager: Final = facade.cache.redis_client.nodes_manager + with rebound(manager, "connection_kwargs", {**manager.connection_kwargs, "db": 1}): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "startup_nodes": startup_nodes[:1]}): + assert resolver.resolve().kind == "python_callback" + binding: Final = resolver.resolve() + assert binding.kind == "native" + + client: Final = redis.RedisCluster(startup_nodes=[redis.cluster.ClusterNode(*node) for node in cluster_nodes]) + keys: Final = tuple(f"slot-{index}" for index in range(12)) + slots: Final = {client.keyslot(f"parity:{key}") for key in keys} + assert len(slots) > 1, slots + requests: Final = [request(key) for key in keys] + values: Final = [{"index": index} for index in range(len(keys))] + await binding.async_store_batch(requests, values) + client.set("parity:slot-3", "not a cache entry") + client.set("parity:slot-7", json.dumps({"timestamp": time.time(), "response": {"index": 7, "python": True}})) + + batch: Final = await binding.async_lookup_batch(requests) + assert batch == { + "values": [ + None if index == 3 else {"index": 7, "python": True} if index == 7 else value + for index, value in enumerate(values) + ], + "missing_indices": [3], + } + assert facade.cache.get_cache("parity:slot-0")["response"] == {"index": 0} + assert (await facade.cache.async_get_cache("parity:slot-11"))["response"] == {"index": 11} + assert facade.cache.redis_client.mget_nonatomic([f"parity:{key}" for key in keys[:2]]) == [ + client.get("parity:slot-0"), + client.get("parity:slot-1"), + ] + + await binding.async_store({**request("pinned"), "ttl_seconds": 12.0}, {"pinned": True}) + assert 0 < client.ttl("parity:pinned") <= 12 + client.set("unscoped", "stays") + + await binding.async_flush() + + remaining: Final = tuple( + sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) + ) + assert remaining == (), remaining + assert client.get("unscoped") == b"stays" + client.delete("unscoped") + client.close() + facade.cache.redis_client.close() + + +def redis_facade(redis_url: str, **settings: object) -> Cache: + parsed: Final = urlparse(redis_url) + return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port), **settings) + + +@pytest.mark.parametrize( + ("settings", "message"), + [ + pytest.param({"max_connections": 10}, "max_connections requires Python", id="pool-size"), + pytest.param({"socket_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="socket-timeout"), + pytest.param( + {"socket_connect_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="connect-timeout" + ), + pytest.param({"socket_keepalive": True}, "does not support socket_keepalive", id="keepalive"), + pytest.param({"health_check_interval": 5}, "does not support health_check_interval", id="health-check"), + pytest.param({"client_name": "litellm"}, "does not support client_name", id="client-name"), + pytest.param({"ssl": True}, "ssl_check_hostname=false require Python", id="tls-default-hostname-check"), + pytest.param({"ssl": True, "ssl_cert_reqs": "none"}, "ssl_cert_reqs=none", id="tls-without-verification"), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_ca_certs": "/ca.pem"}, + "does not support ssl_ca_certs", + id="tls-custom-ca", + ), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, + "does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile", + id="tls-client-certificate", + ), + ], +) +def test_redis_settings_the_native_client_cannot_honor_decline( + redis_url: str, monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + with pytest.raises(RuntimeError, match=f"declined the cache: native Redis.*{message}"): + redis_facade(redis_url, **settings) + + +def test_redis_verified_tls_activates_natively(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + assert_native_runtime(redis_facade(redis_url, ssl=True, ssl_check_hostname=True)) + + +async def test_redis_flush_size_buffers_native_facade_writes(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + facade: Final = redis_facade(redis_url, redis_flush_size=2, namespace="team") + assert_native_runtime(facade) + client: Final = redis.Redis.from_url(redis_url) + first: Final = completion_kwargs("first") + await facade.async_add_cache({"value": 1}, **first) + first_key: Final = facade.get_cache_key(**first) + assert first_key.startswith("team:") + assert client.get(first_key) is None + second: Final = completion_kwargs("second") + await facade.async_add_cache({"value": 2}, **second) + assert client.get(first_key) is not None + assert client.get(facade.get_cache_key(**second)) is not None + client.close() + + +def test_rust_with_fallback_keeps_python_when_the_native_client_declines( + redis_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + catalog, + "RULES", + (CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({LiteLLMCacheType.REDIS})),), + ) + assert redis_facade(redis_url, socket_timeout=1.0)._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor diff --git a/tests/test_litellm_rust/cache/test_redis_semantic.py b/tests/test_litellm_rust/cache/test_redis_semantic.py new file mode 100644 index 00000000000..279330d9060 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_redis_semantic.py @@ -0,0 +1,606 @@ +import asyncio +import contextvars +import hashlib +import json +import math +import os +from collections.abc import Callable, Generator +from contextlib import ExitStack +from types import SimpleNamespace +from typing import Final, cast +from uuid import uuid4 + +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.redis_semantic_cache import RedisSemanticCache +from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +PARAPHRASE_MARKER: Final = " (paraphrase)" + + +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" + + +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" + + +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) + == expected[key] + ), key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} + client.close() + + +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store(semantic_request("inline", "what is the capital of france"), response) + assert ( + await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) + + +async def test_redis_semantic_rust_required_rule_activates_natively( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding, monkeypatch: pytest.MonkeyPatch +) -> None: + del semantic_embedding + url, index = redis_stack + require_rust(monkeypatch, LiteLLMCacheType.REDIS_SEMANTIC) + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + assert_native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": semantic_messages("name a primary color")} + await facade.async_add_cache({"answer": "blue"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "blue"} diff --git a/tests/test_litellm_rust/cache/test_rollout.py b/tests/test_litellm_rust/cache/test_rollout.py new file mode 100644 index 00000000000..7f33e31599f --- /dev/null +++ b/tests/test_litellm_rust/cache/test_rollout.py @@ -0,0 +1,264 @@ +import asyncio +from collections.abc import Callable +from pathlib import Path +from types import SimpleNamespace +from typing import Final, TypeAlias, cast +from urllib.parse import urlparse +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime, resolve_response_cache +from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import EmbeddingResponse +from tests.test_litellm_rust.support.cache import assert_native_runtime, completion_kwargs, require_rust +from tests.test_litellm_rust.support.s3_stub import S3Stub + +pytestmark: Final = pytest.mark.requires_rust_extension + + +CacheFactory: TypeAlias = Callable[[], Cache] + + +@pytest.fixture +def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: + backend: Final = cast(LiteLLMCacheType, request.param) + match backend: + case LiteLLMCacheType.LOCAL: + return lambda: Cache(type=backend) + case LiteLLMCacheType.DISK: + return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) + case LiteLLMCacheType.REDIS: + parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) + return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) + case LiteLLMCacheType.S3: + stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) + return lambda: Cache( + type=backend, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + case LiteLLMCacheType.GCS: + return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") + case LiteLLMCacheType.REDIS_SEMANTIC: + return lambda: Cache( + type=backend, + redis_url="redis://127.0.0.1:6379", + similarity_threshold=0.8, + redis_semantic_cache_embedding_model="text-embedding-3-small", + ) + case LiteLLMCacheType.VALKEY_SEMANTIC: + return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) + case _: + raise AssertionError(f"no local factory for {backend}") + + +ROUND_TRIP_BACKENDS: Final = ( + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, +) + + +SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) + + +@pytest.mark.parametrize("backend", list(LiteLLMCacheType)) +def test_shipped_rules_keep_every_backend_on_python(backend: LiteLLMCacheType) -> None: + assert resolve_response_cache(cast(Cache, SimpleNamespace(type=backend))) is None + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_shipped_rules_construct_python_backed_facades(cache_factory: CacheFactory) -> None: + assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_rust_required_rule_activates_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + assert_native_runtime(cache_factory()) + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_facade_storage_calls_round_trip_through_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + assert_native_runtime(facade) + + sync_kwargs: Final = completion_kwargs("sync") + facade.add_cache({"answer": 1}, **sync_kwargs) + assert facade.get_cache(**sync_kwargs) == {"answer": 1} + + async_kwargs: Final = completion_kwargs("async") + await facade.async_add_cache({"answer": 2}, **async_kwargs) + assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} + assert facade.get_cache(**completion_kwargs("absent")) is None + + +async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.LOCAL) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert_native_runtime(facade) + kwargs: Final = completion_kwargs("memory") + facade.add_cache({"answer": 1}, **kwargs) + assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None + assert facade.get_cache(**kwargs) == {"answer": 1} + + +@pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) +async def test_native_and_python_facades_share_one_wire_format( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + python_facade: Final = cache_factory() + assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + native_facade: Final = cache_factory() + assert_native_runtime(native_facade) + + native_written: Final = completion_kwargs("native") + native_facade.add_cache({"writer": "native"}, **native_written) + assert python_facade.get_cache(**native_written) == {"writer": "native"} + + python_written: Final = completion_kwargs("python") + python_facade.add_cache({"writer": "python"}, **python_written) + assert native_facade.get_cache(**python_written) == {"writer": "python"} + + async_native: Final = completion_kwargs("async-native") + await native_facade.async_add_cache({"writer": "async-native"}, **async_native) + assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} + + async_python: Final = completion_kwargs("async-python") + await python_facade.async_add_cache({"writer": "async-python"}, **async_python) + assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_embedding_pipeline_stores_one_native_entry_per_input( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + assert_native_runtime(facade) + inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] + result: Final = EmbeddingResponse( + model="text-embedding-3-small", + data=[ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, + ], + ) + await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) + + keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] + assert len(set(keys)) == len(inputs) + for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): + cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) + assert isinstance(cached, dict) + assert cached["embedding"] == expected + assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None + + +@pytest.mark.parametrize( + ("backend", "settings", "message"), + [ + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, + "native Valkey semantic cache does not support TLS connections", + id="valkey-tls", + ), + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", + id="valkey-socket-timeout", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-tls", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-query", + ), + ], +) +def test_semantic_settings_the_native_client_cannot_honor_decline( + monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, backend) + with pytest.raises(RuntimeError, match=f"declined the cache: {message}"): + Cache(type=backend, **settings) + + +class _SemanticHit: + """A native semantic runtime that answers every lookup with one cached response.""" + + kind: Final = "native" + + def lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + +@pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +def test_native_semantic_hit_stamps_similarity_on_request_metadata( + semantic_type: LiteLLMCacheType, use_async: bool +) -> None: + """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the + facade copies it to the caller's metadata; the native path must report it the same way.""" + facade: Final = Cache() + facade.type = semantic_type + facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter + metadata: Final[dict[str, object]] = {} + kwargs: Final = { + "cache_key": "semantic-key", + "messages": [{"role": "user", "content": "hello"}], + "metadata": metadata, + } + + result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) + + assert result == {"answer": 42} + assert metadata["semantic-similarity"] == 0.97 diff --git a/tests/test_litellm_rust/cache/test_s3.py b/tests/test_litellm_rust/cache/test_s3.py new file mode 100644 index 00000000000..044bfc39f8d --- /dev/null +++ b/tests/test_litellm_rust/cache/test_s3.py @@ -0,0 +1,187 @@ +import json +import time +from datetime import datetime +from types import SimpleNamespace +from typing import Final, cast +from unittest.mock import Mock + +import boto3 +import botocore.config +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.s3_cache import S3Cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.s3_stub import S3Stub + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def python_s3(url: str) -> S3Cache: + return S3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + + +async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: + python_cache: Final = python_s3(s3_stub.url) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) + python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) + s3_stub.put_object("team/malformed", b"not a cache entry") + s3_stub.put_object( + "team/expired", + json.dumps({"timestamp": time.time(), "response": response}).encode(), + {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, + ) + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + ) + ).resolve() + + assert binding.lookup(request("sync:key")) == response + assert await binding.async_lookup(request("plain")) == response + assert binding.lookup(request("malformed")) is None + assert binding.lookup(request("expired")) is None + assert binding.lookup(request("absent")) is None + + binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) + await binding.async_store(request("no_ttl"), response) + stored: Final = s3_stub.objects["team/native/key"] + assert stored.headers["content-type"] == "application/json" + assert stored.headers["content-language"] == "en" + assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' + assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" + expires: Final = cast(datetime, s3_stub.expires("team/native/key")) + remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() + assert 60 < remaining <= 91 + no_ttl: Final = s3_stub.objects["team/no_ttl"] + assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" + assert "expires" not in no_ttl.headers + assert python_cache.get_cache("native:key")["response"] == response + + partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) + assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} + + +def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + handle: Final = CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + with pytest.raises(TypeError, match="buckets must match"): + CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) + with pytest.raises(TypeError, match="key prefixes must match"): + CacheTestHandle.s3( + "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" + )._bind_facade(facade) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + + handler: Final = Mock() + facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) + binding.store(request("native"), {"answer": 1}) + assert binding.lookup(request("native")) == {"answer": 1} + assert handler.call_count == 0 + assert "team/native" in s3_stub.objects + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + other_client: Final = boto3.client( + "s3", + region_name="us-east-1", + endpoint_url=s3_stub.url, + aws_access_key_id="key", + aws_secret_access_key="secret", + ) + with rebound(facade.cache, "s3_client", other_client): + assert resolver.resolve().kind == "python_callback" + + class CustomS3Cache(S3Cache): + pass + + subclassed: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + subclassed.cache = CustomS3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + with pytest.raises(TypeError): + handle._bind_facade(subclassed) + assert CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" + + +def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: + handle: Final = CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + unverified: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url="https://s3.example.test", + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_verify=False, + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unverified) + proxied: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(proxied) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/cache/test_valkey_semantic.py similarity index 100% rename from tests/test_litellm_rust/test_valkey_semantic_cache_native.py rename to tests/test_litellm_rust/cache/test_valkey_semantic.py diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 0d3b8ba472d..d09e60784fa 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -244,6 +244,21 @@ def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: R assert "invalid OCR request" in str(caught.value) +def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"}, + opaque_extension=object(), + ) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].body == { + "model": "mistral-ocr-latest", + "document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}, + } + + class TokenAbort(BaseException): pass diff --git a/tests/test_litellm_rust/support/cache.py b/tests/test_litellm_rust/support/cache.py new file mode 100644 index 00000000000..41eb4d25257 --- /dev/null +++ b/tests/test_litellm_rust/support/cache.py @@ -0,0 +1,40 @@ +from typing import Final, Protocol +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.rust_bridge import _native, catalog +from litellm.rust_bridge.catalog import CacheRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime +from litellm.types.caching import LiteLLMCacheType + +CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name + + +CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + + +class CacheLookup(Protocol): + def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... + + +def request(key: str = "key") -> dict[str, object]: + return {"key": {"preset": key}} + + +def require_rust(monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType) -> None: + monkeypatch.setattr(catalog, "RULES", (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({backend})),)) + + +def assert_native_runtime(facade: Cache) -> ResponseCacheRuntime: + runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + return runtime + + +def completion_kwargs(label: str) -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": f"{label} {uuid4().hex}"}]} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py deleted file mode 100644 index 96b3674fde3..00000000000 --- a/tests/test_litellm_rust/test_cache.py +++ /dev/null @@ -1,2397 +0,0 @@ -import asyncio -import contextvars -import gc -import hashlib -import http.server -import json -import math -import os -import threading -import time -import uuid -import weakref -from collections.abc import Callable, Generator -from contextlib import ExitStack -from datetime import datetime -from pathlib import Path -from types import SimpleNamespace -from typing import Final, Protocol, TypeAlias, cast -from unittest.mock import Mock -from urllib.parse import urlparse -from uuid import uuid4 - -import boto3 -import botocore.config -import diskcache -import fakeredis -import pytest -import redis -from azure.storage.blob import ContainerClient - -import litellm -from litellm.caching.azure_blob_cache import AzureBlobCache -from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache -from litellm.caching.disk_cache import DiskCache -from litellm.caching.gcs_cache import GCSCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.caching.redis_semantic_cache import RedisSemanticCache -from litellm.caching.s3_cache import S3Cache -from litellm.rust_bridge import _native, catalog -from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime, resolve_response_cache -from litellm.types.caching import LiteLLMCacheType -from litellm.types.llms.custom_llm import CustomLLMItem -from litellm.types.utils import EmbeddingResponse -from tests.test_litellm_rust.support.fake_gcs import FakeGcs -from tests.test_litellm_rust.support.isolation import rebound -from tests.test_litellm_rust.support.s3_stub import S3Stub - -_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name -_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name - -pytestmark: Final = pytest.mark.requires_rust_extension - - -class CacheLookup(Protocol): - def get_cache(self, **kwargs: object) -> object: ... - def flush_cache(self) -> object: ... - - -def request(key: str = "key") -> dict[str, object]: - return {"key": {"preset": key}} - - -def qdrant_request( - key: str, - messages: list[dict[str, object]], - **kwargs: object, -) -> dict[str, object]: - return {**request(key), "messages": messages, **kwargs} - - -def embedding_vector(text: str) -> list[float]: - raw: Final = hashlib.sha256(text.encode()).digest()[:8] - values: Final = [byte / 127.5 - 1 for byte in raw] - norm: Final = math.sqrt(sum(value * value for value in values)) - return [value / norm for value in values] - - -@pytest.fixture -def qdrant_url() -> str: - value: Final[str | None] = os.environ.get("QDRANT_URL") - if not value: - pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") - return value.rstrip("/") - - -@pytest.fixture -def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: - class EmbeddingHandler(http.server.BaseHTTPRequestHandler): - def do_POST(self) -> None: - length: Final = int(self.headers["Content-Length"]) - body: Final = json.loads(self.rfile.read(length)) - text: Final = body["input"] - response: Final = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": embedding_vector(text), - } - ], - "model": body["model"], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - encoded: Final = json.dumps(response).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(encoded))) - self.end_headers() - self.wfile.write(encoded) - - def log_message(self, *_args: object) -> None: - return - - server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) - worker: Final = threading.Thread(target=server.serve_forever, daemon=True) - worker.start() - monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") - monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - try: - yield f"http://127.0.0.1:{server.server_address[1]}" - finally: - server.shutdown() - server.server_close() - worker.join(timeout=5) - - -@pytest.fixture -def redis_url() -> Generator[str]: - server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") - worker: Final = threading.Thread(target=server.serve_forever, daemon=True) - worker.start() - try: - yield f"redis://127.0.0.1:{server.server_address[1]}" - finally: - server.shutdown() - server.server_close() - worker.join(timeout=5) - - -@pytest.fixture -def fake_gcs() -> Generator[FakeGcs]: - server: Final = FakeGcs() - try: - yield server - finally: - server.close() - - -@pytest.fixture -def azure_blob_facade() -> Generator[Cache]: - account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") - if account_url is None: - pytest.skip( - "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" - ) - facade: Final = Cache( - type=LiteLLMCacheType.AZURE_BLOB, - azure_account_url=account_url, - azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", - ) - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - try: - yield facade - finally: - backend.container_client.delete_container() - asyncio.run(backend.disconnect()) - - -def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - return _native._CacheTestHandle.azure_blob( - backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), - backend.container_client.container_name, - ) - - -@pytest.fixture -def cluster_nodes() -> tuple[tuple[str, int], ...]: - configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") - if not configured: - pytest.skip("LITELLM_TEST_REDIS_CLUSTER_NODES is not set") - return tuple((host, int(port)) for host, _, port in (node.partition(":") for node in configured.split(","))) - - -def test_existing_constructor_and_global_are_unchanged() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - assert type(facade.cache) is InMemoryCache - assert "_native_cache_handle" not in vars(facade) - assert resolve_response_cache(facade) is None - with rebound(litellm, "cache", facade): - resolver: Final = _CacheTestResolver(litellm) - assert resolver.resolve().kind == "python_callback" - resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) - assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} - - -async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - assert runtime.kind == "native" - - sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) - assert sync_request is not None - runtime.store(sync_request, {"answer": 1}) - assert runtime.lookup(sync_request) == {"answer": 1} - assert facade.cache.get_cache("sync") is None - - async_request: Final = runtime.request(facade, {"cache_key": "async"}) - assert async_request is not None - await runtime.async_store(async_request, {"answer": 2}) - assert await runtime.async_lookup(async_request) == {"answer": 2} - assert await facade.cache.async_get_cache("async") is None - - requests: Final = (sync_request, async_request) - expected: Final = { - "values": [{"answer": 1}, {"answer": 2}], - "missing_indices": [], - } - assert runtime.lookup_batch(requests) == expected - assert await runtime.async_lookup_batch(requests) == expected - - await runtime.async_flush() - assert runtime.lookup(sync_request) is None - assert await runtime.async_lookup(async_request) is None - - -async def test_inference_resolver_uses_the_configured_native_cache_directly() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - facade._native_cache = runtime - - selected: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert selected.kind == "native" - request: Final = runtime.request(facade, {"cache_key": "inference-native"}) - assert request is not None - await selected.async_store(request, {"answer": 42}) - assert await selected.async_lookup(request) == {"answer": 42} - assert await runtime.async_lookup(request) == {"answer": 42} - assert facade.cache.get_cache("inference-native") is None - - facade._native_cache = None - fallback: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert fallback.kind == "python_callback" - await fallback.async_store(None, {"answer": 7}, callback_kwargs={"cache_key": "inference-python"}) - assert facade.get_cache(cache_key="inference-python") == {"answer": 7} - assert facade.cache.get_cache("inference-python") is not None - - -async def test_inference_resolver_declines_a_native_runtime_whose_facade_changed() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - facade._native_cache = runtime - stale_request: Final = runtime.request(facade, {"cache_key": "stale-only"}) - assert stale_request is not None - await runtime.async_store(stale_request, {"answer": "stale"}) - - replacement: Final = InMemoryCache() - facade.cache = replacement - with pytest.raises(_native.RustBridgeDeclined): - _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert await runtime.async_lookup(stale_request) == {"answer": "stale"} - assert replacement.get_cache("stale-only") is None - assert replacement.get_cache("swapped-backend") is None - - -def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _CacheTestResolver(litellm) - - enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) - enabled: Final = litellm.cache - assert isinstance(enabled, Cache) - assert enabled.ttl == 30 - assert resolver.resolve().kind == "python_callback" - - enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) - assert litellm.cache is enabled - - update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) - updated: Final = litellm.cache - assert isinstance(updated, Cache) - assert updated is not enabled - assert updated.ttl == 60 - - disable_cache() - assert litellm.cache is None - assert resolver.resolve().kind == "disabled" - - -async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) - resolver: Final = _CacheTestResolver(namespace) - selected: Final = resolver.resolve() - assert selected.kind == "native" - selected.store(request(), {"answer": 1}) - assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _CacheTestHandle.memory()): - replacement: Final = resolver.resolve() - await selected.async_store(request(), {"answer": 2}) - assert replacement.lookup(request()) is None - assert selected.lookup(request()) == {"answer": 2} - with rebound(namespace, "cache", None): - disabled: Final = resolver.resolve() - assert disabled.kind == "disabled" - assert disabled.lookup(None) is None - await disabled.async_store(None, object()) - assert await disabled.async_lookup(None) is None - assert selected.lookup(request()) == {"answer": 2} - - -async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: - context: Final = contextvars.ContextVar("cache_context", default="caller") - caller: Final = asyncio.current_task() - sentinel: Final = object() - failure: Final = RuntimeError("callback failed") - - class CustomCache: - async def async_get_cache(self, *, marker: object) -> object: - assert marker is sentinel - assert asyncio.current_task() is caller - context.set("callback") - return marker - - async def async_add_cache(self, response: object, *, marker: object) -> None: - assert response is sentinel - assert marker is sentinel - raise failure - - namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _CacheTestResolver(namespace).resolve() - assert binding.kind == "python_callback" - assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel - assert context.get() == "callback" - with pytest.raises(RuntimeError) as caught: - await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) - assert caught.value is failure - - -async def test_callback_cancellation_stays_in_the_callers_task() -> None: - entered: Final = asyncio.Event() - finished: Final = asyncio.Event() - - class CustomCache: - async def async_get_cache(self) -> None: - entered.set() - try: - await asyncio.Future() - finally: - finished.set() - - binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() - - async def lookup() -> object: - return await binding.async_lookup(None, callback_kwargs={}) - - task: Final = asyncio.create_task(lookup()) - await entered.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert finished.is_set() - - -def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _CacheTestHandle.memory() - handle._bind_facade(facade) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - native: Final = resolver.resolve() - assert native.kind == "native" - native.store(request(), {"source": "native"}) - assert native.lookup(request()) == {"source": "native"} - assert cast(CacheLookup, facade).get_cache(cache_key="key") is None - sentinel: Final = object() - - def outer_override(**_kwargs: object) -> object: - return sentinel - - def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: - return {"source": "override"} - - with rebound(facade, "get_cache", outer_override): - fallback: Final = resolver.resolve() - assert fallback.kind == "python_callback" - assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel - assert resolver.resolve().kind == "python_callback" - delattr(facade, "get_cache") - assert resolver.resolve().kind == "native" - with rebound(facade.cache, "get_cache", backend_override): - backend_fallback: Final = resolver.resolve() - assert backend_fallback.kind == "python_callback" - assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} - - -def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: - class CustomCache(Cache): - pass - - handle: Final = _CacheTestHandle.memory() - with pytest.raises(TypeError): - handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle._bind_facade(facade) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - with rebound(facade, "cache", InMemoryCache()): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "semantic_cache_scope", "end_user"): - assert resolver.resolve().kind == "python_callback" - - def custom_key(**_kwargs: object) -> str: - return "custom" - - with rebound(facade, "get_cache_key", custom_key): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "python_callback" - delattr(facade, "get_cache_key") - assert resolver.resolve().kind == "native" - - -def test_resolver_and_callback_cycles_can_be_collected() -> None: - class CustomCache: - pass - - def cyclic_reference() -> weakref.ReferenceType[CustomCache]: - callback: Final = CustomCache() - namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _CacheTestResolver(namespace).resolve() - setattr(callback, "binding", binding) - return weakref.ref(callback) - - reference: Final = cyclic_reference() - gc.collect() - assert reference() is None - - -async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: - client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _CacheTestResolver(namespace).resolve() - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} - envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} - client.set("team:sync", str(envelope)) - client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) - client.set("team:raw", json.dumps(response)) - client.set("team:invalid", "not a cache entry") - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("team:async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored: Final = client.get("team:native") - assert isinstance(stored, bytes) - assert json.loads(stored)["response"] == response - assert 0 < client.ttl("team:native") <= 12 - assert client.get("litellm-cache:team:native") is None - assert client.get("team:team:async") is None - client.close() - - -def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() - for seconds in (-1.0, float("nan"), float("inf")): - with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) - assert binding.lookup(request()) is None - with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _CacheTestHandle.memory(ttl_seconds=-1) - - -async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() - small: Final = {"answer": "ok"} - binding.store(request("small"), small) - assert await binding.async_lookup(request("small")) == small - await binding.async_store(request("large"), {"answer": "x" * 256}) - assert binding.lookup(request("large")) is None - assert binding.lookup(request("small")) == small - disabled: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0))).resolve() - await disabled.async_store(request(), small) - assert await disabled.async_lookup(request()) is None - - -async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() - requests: Final = [request("hit"), request("miss"), request("disabled")] - requests[2]["controls"] = { - "supported_call_type": True, - "configured": True, - "native_backend": True, - "default_on": True, - "caching": False, - "no_cache": False, - "no_store": False, - "use_cache": False, - } - await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) - - partial: Final = await binding.async_lookup_batch(requests) - - assert partial == { - "values": [{"value": 1}, {"value": 2}, None], - "missing_indices": [2], - } - - -async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: - result: Final = object() - marker: Final = object() - - class CustomCache(Cache): - def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: - return ("sync", kwargs) - - async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: - return ("async", kwargs) - - async def async_add_cache_pipeline( - self, result: object, dynamic_cache_object: object = None, **kwargs: object - ) -> object: - return result, kwargs - - binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() - assert binding.kind == "python_callback" - requests: Final = [request("first"), request("second")] - kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] - - assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] - assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ - ("async", kwargs[0]), - ("async", kwargs[1]), - ] - with pytest.raises(ValueError, match="equal lengths"): - binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) - with pytest.raises(TypeError, match="callback_result"): - await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) - stored: Final = cast( - tuple[object, dict[str, object]], - await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), - ) - assert stored[0] is result - assert stored[1] == {"marker": marker} - - -async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: - async def ping() -> str: - return "pong" - - cache: Final = Cache(type=LiteLLMCacheType.LOCAL) - cache.cache.set_cache("key", "value") - binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() - assert binding.kind == "python_callback" - - setattr(cache.cache, "ping", ping) - assert await binding.ping() == "pong" - await binding.async_flush() - assert cache.cache.get_cache("key") is None - - -def test_facade_registration_rejects_mismatched_capacity() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - with pytest.raises(TypeError, match="capacities must match"): - _CacheTestHandle.memory(capacity=7)._bind_facade(facade) - - -def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: - backend: Final = azure_blob_facade.cache - assert isinstance(backend, AzureBlobCache) - handle: Final = azure_blob_handle(azure_blob_facade) - assert handle.backend == "azure-blob" - account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") - with pytest.raises(TypeError, match="containers must match"): - _native._CacheTestHandle.azure_blob( - account_url, f"{backend.container_client.container_name}-other" - )._bind_facade(azure_blob_facade) - handle._bind_facade(azure_blob_facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) - native: Final = resolver.resolve() - assert native.kind == "native" - - response: Final = { - "choices": [{"text": "caf\u00e9 \u2603"}], - "usage": {"total_tokens": 3}, - "flag": True, - "empty": None, - } - native.store({**request("sync"), "ttl_seconds": 0.001}, response) - native.store(request("sync"), {"choices": [{"text": "second"}]}) - time.sleep(0.01) - stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) - assert stored["response"] == response - assert isinstance(stored["timestamp"], float) - assert native.lookup(request("sync")) == response - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response - - backend.set_cache("python", {"timestamp": time.time(), "response": response}) - backend.set_cache("legacy", "bare legacy value") - backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) - assert native.lookup(request("python")) == response - assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") - assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { - "values": [response, None, None, response], - "missing_indices": [1, 2], - } - - with rebound(azure_blob_facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): - assert resolver.resolve().kind == "python_callback" - - def custom_get(*_args: object, **_kwargs: object) -> None: - return None - - with rebound(backend, "get_cache", custom_get): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "python_callback" - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response - - class CustomBlobCache(AzureBlobCache): - pass - - with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): - assert resolver.resolve().kind == "python_callback" - with pytest.raises(TypeError): - azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) - - -async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: - backend: Final = azure_blob_facade.cache - assert isinstance(backend, AzureBlobCache) - azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() - assert binding.kind == "native" - ping: Final = cast(dict[str, object], await binding.ping()) - assert ping["status"] == "success", ping - - await binding.async_store(request("async"), {"value": 1}) - await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) - time.sleep(0.01) - assert await binding.async_lookup(request("async")) == {"value": 2} - assert await backend.async_get_cache("async") == json.loads( - backend.container_client.download_blob("async").readall() - ) - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} - - await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) - assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { - "values": [{"value": 4}, None, {"value": 3}], - "missing_indices": [1], - } - await binding.async_flush() - assert [blob.name for blob in backend.container_client.list_blobs()] == [] - assert await binding.async_lookup(request("async")) is None - - -async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: - parsed: Final = urlparse(redis_url) - with rebound(litellm, "default_redis_ttl", 60): - facade: Final = Cache( - type=LiteLLMCacheType.REDIS, - host=parsed.hostname, - port=str(parsed.port), - redis_flush_size=2, - ) - with pytest.raises(TypeError, match="default TTLs must match"): - _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) - with pytest.raises(TypeError, match="namespaces must match"): - _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(redis_url) - - with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - - pool: Final = facade.cache.redis_client.connection_pool - with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - - await binding.async_store(request("first"), {"value": 1}) - assert client.get("first") is None - await binding.async_store(request("second"), {"value": 2}) - - assert client.get("first") is not None - assert client.get("second") is not None - await facade.cache.disconnect() - client.close() - - -async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: - disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} - disk_cache.disk_cache.set( - "sync", - {"timestamp": time.time(), "response": json.dumps(response)}, - ) - disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response})) - disk_cache.disk_cache.set("raw", json.dumps(response)) - disk_cache.disk_cache.set("invalid", "not a cache entry") - disk_cache.disk_cache.set( - "large", - {"timestamp": time.time(), "response": {"text": "x" * 70_000}}, - ) - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - assert binding.lookup(request("large")) == {"text": "x" * 70_000} - - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored_response: Final = disk_cache.get_cache("native") - assert isinstance(stored_response, dict) - assert stored_response["response"] == response - stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True) - assert stored is not None - assert time.time() < expire_time <= time.time() + 12.0 - await binding.async_store(request("no-ttl"), response) - _, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True) - assert no_expiry is None - - -async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None: - first: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - await first.async_store(request("persistent"), {"value": "persistent"}) - await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"}) - fresh: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - assert fresh.lookup(request("persistent")) == {"value": "persistent"} - assert fresh.lookup(request("expiring")) == {"value": "expiring"} - await asyncio.sleep(0.4) - assert fresh.lookup(request("expiring")) is None - assert fresh.lookup(request("persistent")) == {"value": "persistent"} - - -def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None: - facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) - with pytest.raises(TypeError, match="directories must match"): - _native._CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade) - handle: Final = _native._CacheTestHandle.disk(str(tmp_path)) - handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - binding.store(request("native"), {"value": "native"}) - assert facade.get_cache(cache_key="native") == {"value": "native"} - - with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "native" - - class CustomDiskCache(DiskCache): - pass - - with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))): - assert resolver.resolve().kind == "python_callback" - - class CustomStore(diskcache.Cache): - pass - - custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) - custom_facade.cache.disk_cache = CustomStore(str(tmp_path)) - with pytest.raises(TypeError, match="built-in diskcache store"): - _native._CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade) - - -async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None: - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - requests: Final = [request("hit"), request("miss"), request("disabled")] - requests[2]["controls"] = { - "supported_call_type": True, - "configured": True, - "native_backend": True, - "default_on": True, - "caching": False, - "no_cache": False, - "no_store": False, - "use_cache": False, - } - await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) - - partial: Final = await binding.async_lookup_batch(requests) - - assert partial == { - "values": [{"value": 1}, {"value": 2}, None], - "missing_indices": [2], - } - - -@pytest.fixture -def s3_stub() -> Generator[S3Stub]: - stub: Final = S3Stub() - try: - yield stub - finally: - stub.close() - - -def python_s3(url: str) -> S3Cache: - return S3Cache( - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - - -async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: - python_cache: Final = python_s3(s3_stub.url) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} - python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) - python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) - s3_stub.put_object("team/malformed", b"not a cache entry") - s3_stub.put_object( - "team/expired", - json.dumps({"timestamp": time.time(), "response": response}).encode(), - {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, - ) - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - ) - ).resolve() - - assert binding.lookup(request("sync:key")) == response - assert await binding.async_lookup(request("plain")) == response - assert binding.lookup(request("malformed")) is None - assert binding.lookup(request("expired")) is None - assert binding.lookup(request("absent")) is None - - binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) - await binding.async_store(request("no_ttl"), response) - stored: Final = s3_stub.objects["team/native/key"] - assert stored.headers["content-type"] == "application/json" - assert stored.headers["content-language"] == "en" - assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' - assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" - expires: Final = cast(datetime, s3_stub.expires("team/native/key")) - remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() - assert 60 < remaining <= 91 - no_ttl: Final = s3_stub.objects["team/no_ttl"] - assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" - assert "expires" not in no_ttl.headers - assert python_cache.get_cache("native:key")["response"] == response - - partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) - assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} - - -def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: - facade: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - handle: Final = _native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - with pytest.raises(TypeError, match="buckets must match"): - _native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) - with pytest.raises(TypeError, match="key prefixes must match"): - _native._CacheTestHandle.s3( - "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" - )._bind_facade(facade) - handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - - handler: Final = Mock() - facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) - binding.store(request("native"), {"answer": 1}) - assert binding.lookup(request("native")) == {"answer": 1} - assert handler.call_count == 0 - assert "team/native" in s3_stub.objects - - with rebound(facade.cache, "bucket_name", "other"): - assert resolver.resolve().kind == "python_callback" - other_client: Final = boto3.client( - "s3", - region_name="us-east-1", - endpoint_url=s3_stub.url, - aws_access_key_id="key", - aws_secret_access_key="secret", - ) - with rebound(facade.cache, "s3_client", other_client): - assert resolver.resolve().kind == "python_callback" - - class CustomS3Cache(S3Cache): - pass - - subclassed: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - subclassed.cache = CustomS3Cache( - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - with pytest.raises(TypeError): - handle._bind_facade(subclassed) - assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" - - -def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: - handle: Final = _native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - unverified: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url="https://s3.example.test", - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - s3_verify=False, - ) - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(unverified) - proxied: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), - ) - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(proxied) - - -async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} - fake_gcs.put( - "bucket", - "cache/sync", - json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), - ) - fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) - fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) - fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - assert binding.lookup(request("missing")) is None - - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored: Final = fake_gcs.objects[("bucket", "cache/native")] - stored_value: Final = cast(dict[str, object], json.loads(stored)) - assert stored_value["response"] == response - assert isinstance(stored_value["timestamp"], float) - upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") - assert upload.path == "/upload/storage/v1/b/bucket/o" - assert upload.query == "uploadType=media&name=cache%2Fnative" - assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" - assert upload.headers["Content-Type"] == "application/json" - upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" - assert "ttl" not in upload_text.lower() - assert "expiry" not in upload_text.lower() - download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) - assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" - assert download.query == "alt=media" - - binding.store(request("sync2"), response) - assert binding.lookup(request("sync2")) == response - assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" - assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" - assert GCSCache(bucket_name="bucket").key_prefix == "" - - -async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: - fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) - fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - requests: Final = [request("hit"), request("missing"), request("invalid")] - expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} - - assert await binding.async_lookup_batch(requests) == expected - assert binding.lookup_batch(requests) == expected - await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) - assert ("bucket", "cache/first") in fake_gcs.objects - assert ("bucket", "cache/second") in fake_gcs.objects - - -async def test_gcs_facade_binds_only_exact_matching_configuration( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") - facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - assert type(facade.cache) is GCSCache - - mismatched_bucket: Final = _native._CacheTestHandle.gcs( - "other", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="buckets must match"): - mismatched_bucket._bind_facade(facade) - mismatched_prefix: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="x", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="key prefixes must match"): - mismatched_prefix._bind_facade(facade) - mismatched_credentials: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - path_service_account="sa.json", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="credentials must match"): - mismatched_credentials._bind_facade(facade) - with pytest.raises(TypeError, match="types must match"): - _native._CacheTestHandle.memory()._bind_facade(facade) - - matching: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - matching._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - await binding.async_store(request("native"), {"value": "native"}) - assert await binding.async_lookup(request("native")) == {"value": "native"} - assert cast(CacheLookup, facade).get_cache(cache_key="native") is None - - with rebound(facade.cache, "bucket_name", "other"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "key_prefix", "x/"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "path_service_account", "sa.json"): - assert resolver.resolve().kind == "python_callback" - - def no_get_cache(*args: object, **kwargs: object) -> None: - return None - - with rebound(facade.cache, "get_cache", no_get_cache): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - - class CustomGcs(GCSCache): - pass - - with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): - assert resolver.resolve().kind == "python_callback" - custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): - with pytest.raises(TypeError, match="types must match"): - matching._bind_facade(custom_facade) - - missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) - with pytest.raises(TypeError, match="requires a configured bucket name"): - matching._bind_facade(missing_bucket) - - -async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - await binding.async_store(request("key"), {"value": "stored"}) - await binding.async_flush() - assert ("bucket", "cache/key") in fake_gcs.objects - assert await binding.async_lookup(request("key")) == {"value": "stored"} - with pytest.raises(NotImplementedError): - await binding.ping() - - facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - with pytest.raises(AttributeError): - await facade.ping() - assert cast(CacheLookup, facade.cache).flush_cache() is None - - -async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: - wrong_token: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token="wrong-token", - ) - ) - ).resolve() - with pytest.raises(RuntimeError): - wrong_token.lookup(request("missing")) - assert not fake_gcs.objects - - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - with pytest.raises(RuntimeError): - binding.lookup(request("server-error")) - assert binding.lookup(request("missing")) is None - - -async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( - cluster_nodes: tuple[tuple[str, int], ...], -) -> None: - startup_nodes: Final = [{"host": host, "port": port} for host, port in cluster_nodes] - url: Final = f"redis://{cluster_nodes[0][0]}:{cluster_nodes[0][1]}" - with rebound(litellm, "default_redis_ttl", 60): - facade: Final = Cache(type=LiteLLMCacheType.REDIS, redis_startup_nodes=startup_nodes, namespace="parity") - assert type(facade.cache) is RedisClusterCache - with pytest.raises(TypeError, match="types must match"): - _native._CacheTestHandle.redis(url, namespace="parity")._bind_facade(facade) - _native._CacheTestHandle.redis(url, namespace="parity", startup_nodes=list(cluster_nodes))._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - assert resolver.resolve().kind == "native" - - manager: Final = facade.cache.redis_client.nodes_manager - with rebound(manager, "connection_kwargs", {**manager.connection_kwargs, "db": 1}): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "startup_nodes": startup_nodes[:1]}): - assert resolver.resolve().kind == "python_callback" - binding: Final = resolver.resolve() - assert binding.kind == "native" - - client: Final = redis.RedisCluster(startup_nodes=[redis.cluster.ClusterNode(*node) for node in cluster_nodes]) - keys: Final = tuple(f"slot-{index}" for index in range(12)) - slots: Final = {client.keyslot(f"parity:{key}") for key in keys} - assert len(slots) > 1, slots - requests: Final = [request(key) for key in keys] - values: Final = [{"index": index} for index in range(len(keys))] - await binding.async_store_batch(requests, values) - client.set("parity:slot-3", "not a cache entry") - client.set("parity:slot-7", json.dumps({"timestamp": time.time(), "response": {"index": 7, "python": True}})) - - batch: Final = await binding.async_lookup_batch(requests) - assert batch == { - "values": [ - None if index == 3 else {"index": 7, "python": True} if index == 7 else value - for index, value in enumerate(values) - ], - "missing_indices": [3], - } - assert facade.cache.get_cache("parity:slot-0")["response"] == {"index": 0} - assert (await facade.cache.async_get_cache("parity:slot-11"))["response"] == {"index": 11} - assert facade.cache.redis_client.mget_nonatomic([f"parity:{key}" for key in keys[:2]]) == [ - client.get("parity:slot-0"), - client.get("parity:slot-1"), - ] - - await binding.async_store({**request("pinned"), "ttl_seconds": 12.0}, {"pinned": True}) - assert 0 < client.ttl("parity:pinned") <= 12 - client.set("unscoped", "stays") - - await binding.async_flush() - - remaining: Final = tuple( - sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) - ) - assert remaining == (), remaining - assert client.get("unscoped") == b"stays" - client.delete("unscoped") - client.close() - facade.cache.redis_client.close() - - -PARAPHRASE_MARKER: Final = " (paraphrase)" -SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" -SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" -SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") - - -def _normalized(vector: list[float]) -> list[float]: - norm: Final = math.sqrt(sum(component * component for component in vector)) - return [component / norm for component in vector] - - -def _base_embedding(prompt: str) -> list[float]: - digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() - return _normalized([float(digest[index] + 1) for index in range(8)]) - - -def _semantic_embedding(prompt: str) -> list[float]: - if PARAPHRASE_MARKER not in prompt: - return _base_embedding(prompt) - base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) - pivot: Final = min(range(8), key=lambda index: abs(base[index])) - direction: Final = _normalized( - [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] - ) - # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance - return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) - - -class DeterministicEmbedding(litellm.CustomLLM): - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - self.async_calls: list[dict[str, object]] = [] - self.entered = asyncio.Event() - self.gate: asyncio.Event | None = None - - def _respond( - self, - model: str, - input: object, - model_response: EmbeddingResponse, - ) -> EmbeddingResponse: - texts: Final = cast(list[object], input if isinstance(input, list) else [input]) - self.calls.append({"model": model, "input": texts}) - model_response.model = model - model_response.data = [ - {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} - for index, text in enumerate(texts) - ] - return model_response - - def embedding( - self, - model: str, - input: list[object], - model_response: EmbeddingResponse, - print_verbose: Callable[..., object], - logging_obj: object, - optional_params: dict[str, object], - api_key: object = None, - api_base: object = None, - timeout: object = None, - litellm_params: object = None, - ) -> EmbeddingResponse: - return self._respond(model, input, model_response) - - async def aembedding( - self, - model: str, - input: list[object], - model_response: EmbeddingResponse, - print_verbose: Callable[..., object], - logging_obj: object, - optional_params: dict[str, object], - api_key: object = None, - api_base: object = None, - timeout: object = None, - litellm_params: object = None, - ) -> EmbeddingResponse: - texts: Final = cast(list[object], input if isinstance(input, list) else [input]) - self.async_calls.append( - { - "model": model, - "input": texts, - "task": asyncio.current_task(), - "context": SEMANTIC_CONTEXT.get(), - } - ) - SEMANTIC_CONTEXT.set("written-in-aembedding") - self.entered.set() - if self.gate is not None: - await self.gate.wait() - return self._respond(model, input, model_response) - - -@pytest.fixture -def semantic_embedding() -> Generator[DeterministicEmbedding]: - handler: Final = DeterministicEmbedding() - with ExitStack() as stack: - stack.enter_context( - rebound( - litellm, - "custom_provider_map", - [ - *litellm.custom_provider_map, - cast( - CustomLLMItem, - {"provider": "semantic-test", "custom_handler": handler}, - ), - ], - ) - ) - stack.enter_context( - rebound( - litellm, - "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook - [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook - ) - ) - stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) - yield handler - - -@pytest.fixture -def redis_stack() -> Generator[tuple[str, str]]: - url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") - if url is None: - pytest.skip("LITELLM_REDIS_STACK_URL is not set") - index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" - yield url, index - client: Final = redis.Redis.from_url(url) - try: - client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown - except redis.RedisError: - pass - client.close() - - -def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: - return { - "key": {"preset": key}, - "messages": [{"role": "user", "content": prompt}], - **extra, - } - - -def semantic_messages(prompt: str) -> list[dict[str, object]]: - return [{"role": "user", "content": prompt}] - - -def semantic_entry_id(prompt: str, tag: str) -> str: - return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() - - -def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: - facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=similarity_threshold, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) - return facade - - -def test_redis_semantic_constructor_identity_and_provenance( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - backend: Final = cast(RedisSemanticCache, facade.cache) - assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" - assert type(backend) is RedisSemanticCache - assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config - assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config - assert backend.similarity_threshold == 0.8 - assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL - handle: Final = cast(object, getattr(facade, "_native_cache_handle")) - assert isinstance(handle, _CacheTestHandle) - assert handle.backend == "redis_semantic" - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - - -def test_redis_semantic_native_and_python_sync_entries_share_one_layout( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} - - binding.store(semantic_request("geo", "what is the capital of france"), response) - - native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" - stored: Final = client.hgetall(native_hash_key) - assert set(stored) == { - b"entry_id", - b"prompt", - b"response", - b"prompt_vector", - b"inserted_at", - b"updated_at", - b"litellm_cache_key", - }, stored - assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] - assert stored[b"prompt"] == b"what is the capital of france" - assert stored[b"litellm_cache_key"] == b"geo" - assert len(stored[b"prompt_vector"]) == 32 - decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) - assert decoded["response"] == response - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "geo", messages=semantic_messages("what is the capital of france") - ) - == decoded - ) - assert semantic_embedding.calls == [ - {"model": "deterministic", "input": ["what is the capital of france"]}, - {"model": "deterministic", "input": ["what is the capital of france"]}, - {"model": "deterministic", "input": ["dimension test"]}, - ] - - cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "math", - json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), - messages=semantic_messages("what is 6 times 7"), - ) - python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" - assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { - "timestamp": 1700000000.0, - "response": {"answer": 42}, - } - assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} - client.close() - - -async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) - hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" - decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) - python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "async", messages=semantic_messages("name a primary color") - ) - assert python_read == decoded - - await binding.async_store_batch( - [ - semantic_request("batch-one", "first batch prompt"), - semantic_request("batch-two", "second batch prompt"), - ], - [{"answer": 1}, {"answer": 2}], - ) - expected: Final = { - key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) - for key, prompt in ( - ("batch-one", "first batch prompt"), - ("batch-two", "second batch prompt"), - ) - } - for key, prompt in ( - ("batch-one", "first batch prompt"), - ("batch-two", "second batch prompt"), - ): - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - key, messages=semantic_messages(prompt) - ) - == expected[key] - ), key - - cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "async-python", - json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), - messages=semantic_messages("python written prompt"), - ) - assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} - client.close() - - -async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - caller: Final = asyncio.current_task() - SEMANTIC_CONTEXT.set("caller-sentinel") - response: Final = {"choices": [{"text": "paris"}]} - - await binding.async_store(semantic_request("inline", "what is the capital of france"), response) - assert ( - await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) - == response - ) - assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None - assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" - assert semantic_embedding.async_calls == [ - { - "model": "deterministic", - "input": ["what is the capital of france"], - "task": caller, - "context": "caller-sentinel", - }, - { - "model": "deterministic", - "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], - "task": caller, - "context": "written-in-aembedding", - }, - { - "model": "deterministic", - "input": ["python written prompt"], - "task": caller, - "context": "written-in-aembedding", - }, - ], semantic_embedding.async_calls - - -async def test_native_semantic_cancellation_during_embedding_skips_the_backend( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - semantic_embedding.gate = asyncio.Event() - - async def lookup() -> object: - return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) - - task: Final = asyncio.create_task(lookup()) - await semantic_embedding.entered.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - semantic_embedding.gate.set() - - assert len(semantic_embedding.async_calls) == 1 - assert ( - await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "cancel", messages=semantic_messages("cancelled prompt") - ) - is None - ) - - -def test_redis_semantic_similarity_tag_and_threshold_boundaries( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - - binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) - paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" - assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} - assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None - assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None - - strict: Final = semantic_facade(url, index, similarity_threshold=0.99) - strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() - assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None - assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} - - -def test_redis_semantic_ttl_is_written_only_when_requested( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) - expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" - assert 0 < client.ttl(expiring) <= 12 - - binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) - persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" - assert client.ttl(persistent) == -1 - - binding.store( - {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, - {"answer": 3}, - ) - fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" - assert client.ttl(fractional) == 2 - client.close() - - -def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) - hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" - client.hset(hash_key, "response", b"{not json") - assert binding.lookup(semantic_request("bad", "corrupt me")) is None - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "bad", messages=semantic_messages("corrupt me") - ) - is None - ) - client.close() - - -async def test_redis_semantic_unsupported_operations_raise_not_implemented( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - - with pytest.raises(NotImplementedError): - binding.lookup_batch([semantic_request("batch", "prompt one")]) - with pytest.raises(NotImplementedError): - await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) - with pytest.raises(NotImplementedError): - await binding.async_flush() - with pytest.raises(NotImplementedError): - await binding.ping() - - -def test_redis_semantic_requests_without_prompt_are_noops( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store(request("plain"), {"answer": 1}) - assert binding.lookup(request("plain")) is None - assert semantic_embedding.calls == [] - assert client.keys(f"{index}:*") == [] - client.close() - - -def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} - binding.store(scoped, {"answer": "kept"}) - hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" - assert client.hget(hash_key, "litellm_cache_key") == b"team-a" - assert binding.lookup(scoped) == {"answer": "kept"} - assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None - assert binding.lookup({**scoped, "scope": "team-b"}) is None - client.close() - - -def test_redis_semantic_configuration_drift_falls_back_to_python( - redis_stack: tuple[str, str], - semantic_embedding: DeterministicEmbedding, - monkeypatch: pytest.MonkeyPatch, -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - assert resolver.resolve().kind == "native" - - with rebound(facade.cache, "similarity_threshold", 0.5): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "semantic_cache_scope", "end_user"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "embedding_model", "other-model"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "_index_name", "other-index"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): - assert resolver.resolve().kind == "python_callback" - - def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: - return _semantic_embedding(prompt) - - monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) - assert resolver.resolve().kind == "python_callback" - - -def test_redis_semantic_handle_rejects_wrong_backends( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - - class CustomSemanticCache(RedisSemanticCache): - pass - - with pytest.raises(TypeError, match="built-in RedisSemanticCache"): - _CacheTestHandle.redis_semantic(object()) - with pytest.raises(TypeError, match="built-in RedisSemanticCache"): - _CacheTestHandle.redis_semantic( - CustomSemanticCache( - redis_url=url, - similarity_threshold=0.8, - embedding_model=SEMANTIC_EMBEDDING_MODEL, - index_name=f"{index}_subclass", - ) - ) - - facade: Final = semantic_facade(url, index) - with pytest.raises(TypeError, match="backend types must match"): - _CacheTestHandle.redis(url)._bind_facade(facade) - - subclassed_facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared - redis_url=url, - similarity_threshold=0.8, - embedding_model=SEMANTIC_EMBEDDING_MODEL, - index_name=index, - ) - with pytest.raises(TypeError): - _CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) - - replacement_facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - with pytest.raises(TypeError, match="must be the native embedder"): - _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) - - -def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: - return Cache( - type=LiteLLMCacheType.QDRANT_SEMANTIC, - qdrant_api_base=qdrant_url, - qdrant_collection_name=collection_name, - similarity_threshold=0.99, - qdrant_semantic_cache_embedding_model="text-embedding-3-small", - qdrant_semantic_cache_vector_size=8, - ) - - -def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "shared prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - facade.cache.set_cache( - "python-key", - {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, - messages=messages, - ) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} - binding.store(qdrant_request("native-key", messages), {"id": "native"}) - python_value: Final = facade.cache.get_cache("native-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "native"} - unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] - assert binding.lookup(qdrant_request("native-key", unrelated)) is None - assert facade.cache.get_cache("native-key", messages=unrelated) is None - assert binding.lookup(qdrant_request("different-key", messages)) is None - assert facade.cache.get_cache("different-key", messages=messages) is None - - -async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "async prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - await facade.cache.async_set_cache( - "python-key", - {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, - messages=messages, - ) - assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} - await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) - python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "native"} - - -async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - entries: Final = [ - qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), - qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), - ] - await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) - - assert binding.lookup(entries[0]) == {"id": "one"} - assert binding.lookup(entries[1]) == {"id": "two"} - assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { - "id": "one" - } - assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { - "id": "two" - } - - -async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "malformed prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - key: Final = "malformed-key" - response: Final = { - "points": [ - { - "id": str(uuid4()), - "vector": embedding_vector("malformed prompt"), - "payload": { - "litellm_cache_key": key, - "text": "malformed prompt", - "response": "not json", - }, - } - ] - } - facade.cache.sync_client.put( - url=f"{qdrant_url}/collections/{collection}/points", - headers=facade.cache.headers, - json=response, - ) - assert binding.lookup(qdrant_request(key, messages)) is None - with pytest.raises(RuntimeError, match="operation is not supported"): - binding.lookup_batch([qdrant_request(key, messages)]) - with pytest.raises(RuntimeError, match="operation is not supported"): - await binding.async_flush() - with pytest.raises(RuntimeError, match="operation is not supported"): - await binding.ping() - - -def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "persistent prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) - time.sleep(1.2) - assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} - python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "persistent"} - - -def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - facade.cache.qdrant_api_key = "rotated" - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - facade.cache.similarity_threshold = 0.5 - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") - unsupported.cache.embedding_max_input_tokens = 100 - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(unsupported) - unsupported.cache.embedding_max_input_tokens = None - unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" - with pytest.raises(TypeError, match="gRPC"): - handle._bind_facade(unsupported) - - -CacheFactory: TypeAlias = Callable[[], Cache] - - -def require_rust(monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType) -> None: - monkeypatch.setattr(catalog, "RULES", (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({backend})),)) - - -def native_runtime(facade: Cache) -> ResponseCacheRuntime: - runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - assert isinstance(runtime, ResponseCacheRuntime) - assert runtime.kind == "native" - return runtime - - -@pytest.fixture -def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: - backend: Final = cast(LiteLLMCacheType, request.param) - match backend: - case LiteLLMCacheType.LOCAL: - return lambda: Cache(type=backend) - case LiteLLMCacheType.DISK: - return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) - case LiteLLMCacheType.REDIS: - parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) - return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) - case LiteLLMCacheType.S3: - stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) - return lambda: Cache( - type=backend, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - case LiteLLMCacheType.GCS: - return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") - case LiteLLMCacheType.REDIS_SEMANTIC: - return lambda: Cache( - type=backend, - redis_url="redis://127.0.0.1:6379", - similarity_threshold=0.8, - redis_semantic_cache_embedding_model="text-embedding-3-small", - ) - case LiteLLMCacheType.VALKEY_SEMANTIC: - return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) - case _: - raise AssertionError(f"no local factory for {backend}") - - -ROUND_TRIP_BACKENDS: Final = ( - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, -) -SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) - - -def completion_kwargs(label: str) -> dict[str, object]: - return {"model": "gpt-4o", "messages": [{"role": "user", "content": f"{label} {uuid4().hex}"}]} - - -@pytest.mark.parametrize("backend", list(LiteLLMCacheType)) -def test_shipped_rules_keep_every_backend_on_python(backend: LiteLLMCacheType) -> None: - assert resolve_response_cache(cast(Cache, SimpleNamespace(type=backend))) is None - - -@pytest.mark.parametrize( - "cache_factory", - [ - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, - LiteLLMCacheType.GCS, - LiteLLMCacheType.REDIS_SEMANTIC, - LiteLLMCacheType.VALKEY_SEMANTIC, - ], - indirect=True, -) -def test_shipped_rules_construct_python_backed_facades(cache_factory: CacheFactory) -> None: - assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - - -@pytest.mark.parametrize( - "cache_factory", - [ - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, - LiteLLMCacheType.GCS, - LiteLLMCacheType.REDIS_SEMANTIC, - LiteLLMCacheType.VALKEY_SEMANTIC, - ], - indirect=True, -) -def test_rust_required_rule_activates_the_native_backend( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - native_runtime(cache_factory()) - - -@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) -async def test_facade_storage_calls_round_trip_through_the_native_backend( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - facade: Final = cache_factory() - native_runtime(facade) - - sync_kwargs: Final = completion_kwargs("sync") - facade.add_cache({"answer": 1}, **sync_kwargs) - assert facade.get_cache(**sync_kwargs) == {"answer": 1} - - async_kwargs: Final = completion_kwargs("async") - await facade.async_add_cache({"answer": 2}, **async_kwargs) - assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} - assert facade.get_cache(**completion_kwargs("absent")) is None - - -async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.LOCAL) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - native_runtime(facade) - kwargs: Final = completion_kwargs("memory") - facade.add_cache({"answer": 1}, **kwargs) - assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None - assert facade.get_cache(**kwargs) == {"answer": 1} - - -@pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) -async def test_native_and_python_facades_share_one_wire_format( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - python_facade: Final = cache_factory() - assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - native_facade: Final = cache_factory() - native_runtime(native_facade) - - native_written: Final = completion_kwargs("native") - native_facade.add_cache({"writer": "native"}, **native_written) - assert python_facade.get_cache(**native_written) == {"writer": "native"} - - python_written: Final = completion_kwargs("python") - python_facade.add_cache({"writer": "python"}, **python_written) - assert native_facade.get_cache(**python_written) == {"writer": "python"} - - async_native: Final = completion_kwargs("async-native") - await native_facade.async_add_cache({"writer": "async-native"}, **async_native) - assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} - - async_python: Final = completion_kwargs("async-python") - await python_facade.async_add_cache({"writer": "async-python"}, **async_python) - assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} - - -@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) -async def test_embedding_pipeline_stores_one_native_entry_per_input( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - facade: Final = cache_factory() - native_runtime(facade) - inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] - result: Final = EmbeddingResponse( - model="text-embedding-3-small", - data=[ - {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, - {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, - ], - ) - await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) - - keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] - assert len(set(keys)) == len(inputs) - for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): - cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) - assert isinstance(cached, dict) - assert cached["embedding"] == expected - assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None - - -def redis_facade(redis_url: str, **settings: object) -> Cache: - parsed: Final = urlparse(redis_url) - return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port), **settings) - - -@pytest.mark.parametrize( - ("settings", "message"), - [ - pytest.param({"max_connections": 10}, "max_connections requires Python", id="pool-size"), - pytest.param({"socket_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="socket-timeout"), - pytest.param( - {"socket_connect_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="connect-timeout" - ), - pytest.param({"socket_keepalive": True}, "does not support socket_keepalive", id="keepalive"), - pytest.param({"health_check_interval": 5}, "does not support health_check_interval", id="health-check"), - pytest.param({"client_name": "litellm"}, "does not support client_name", id="client-name"), - pytest.param({"ssl": True}, "ssl_check_hostname=false require Python", id="tls-default-hostname-check"), - pytest.param({"ssl": True, "ssl_cert_reqs": "none"}, "ssl_cert_reqs=none", id="tls-without-verification"), - pytest.param( - {"ssl": True, "ssl_check_hostname": True, "ssl_ca_certs": "/ca.pem"}, - "does not support ssl_ca_certs", - id="tls-custom-ca", - ), - pytest.param( - {"ssl": True, "ssl_check_hostname": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, - "does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile", - id="tls-client-certificate", - ), - ], -) -def test_redis_settings_the_native_client_cannot_honor_decline( - redis_url: str, monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], message: str -) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - with pytest.raises(RuntimeError, match=f"declined the cache: native Redis.*{message}"): - redis_facade(redis_url, **settings) - - -def test_redis_verified_tls_activates_natively(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - native_runtime(redis_facade(redis_url, ssl=True, ssl_check_hostname=True)) - - -async def test_redis_flush_size_buffers_native_facade_writes(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - facade: Final = redis_facade(redis_url, redis_flush_size=2, namespace="team") - native_runtime(facade) - client: Final = redis.Redis.from_url(redis_url) - first: Final = completion_kwargs("first") - await facade.async_add_cache({"value": 1}, **first) - first_key: Final = facade.get_cache_key(**first) - assert first_key.startswith("team:") - assert client.get(first_key) is None - second: Final = completion_kwargs("second") - await facade.async_add_cache({"value": 2}, **second) - assert client.get(first_key) is not None - assert client.get(facade.get_cache_key(**second)) is not None - client.close() - - -@pytest.mark.parametrize( - ("backend", "settings", "message"), - [ - pytest.param( - LiteLLMCacheType.VALKEY_SEMANTIC, - {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, - "native Valkey semantic cache does not support TLS connections", - id="valkey-tls", - ), - pytest.param( - LiteLLMCacheType.VALKEY_SEMANTIC, - {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, - "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", - id="valkey-socket-timeout", - ), - pytest.param( - LiteLLMCacheType.REDIS_SEMANTIC, - {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, - "native Redis semantic cache does not support TLS or query options in redis_url", - id="redis-semantic-tls", - ), - pytest.param( - LiteLLMCacheType.REDIS_SEMANTIC, - {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, - "native Redis semantic cache does not support TLS or query options in redis_url", - id="redis-semantic-query", - ), - ], -) -def test_semantic_settings_the_native_client_cannot_honor_decline( - monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str -) -> None: - require_rust(monkeypatch, backend) - with pytest.raises(RuntimeError, match=f"declined the cache: {message}"): - Cache(type=backend, **settings) - - -def test_rust_with_fallback_keeps_python_when_the_native_client_declines( - redis_url: str, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr( - catalog, - "RULES", - (CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({LiteLLMCacheType.REDIS})),), - ) - assert redis_facade(redis_url, socket_timeout=1.0)._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - - -def test_qdrant_semantic_rust_required_rule_activates_natively( - qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch -) -> None: - del fake_embedding_endpoint - require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC) - facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") - native_runtime(facade) - kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} - facade.add_cache({"answer": "qdrant"}, **kwargs) - assert facade.get_cache(**kwargs) == {"answer": "qdrant"} - - -async def test_redis_semantic_rust_required_rule_activates_natively( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding, monkeypatch: pytest.MonkeyPatch -) -> None: - del semantic_embedding - url, index = redis_stack - require_rust(monkeypatch, LiteLLMCacheType.REDIS_SEMANTIC) - facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - native_runtime(facade) - kwargs: Final = {"model": "gpt-4o", "messages": semantic_messages("name a primary color")} - await facade.async_add_cache({"answer": "blue"}, **kwargs) - assert await facade.async_get_cache(**kwargs) == {"answer": "blue"} - - -async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None: - account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") - if account_url is None: - pytest.skip( - "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" - ) - require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB) - facade: Final = Cache( - type=LiteLLMCacheType.AZURE_BLOB, - azure_account_url=account_url, - azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", - ) - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - try: - native_runtime(facade) - kwargs: Final = completion_kwargs("azure") - await facade.async_add_cache({"answer": "azure"}, **kwargs) - assert await facade.async_get_cache(**kwargs) == {"answer": "azure"} - assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"} - finally: - backend.container_client.delete_container() - await backend.disconnect() - - -class _SemanticHit: - """A native semantic runtime that answers every lookup with one cached response.""" - - kind: Final = "native" - - def lookup_semantic(self, request: object) -> tuple[object, float | None]: - return {"answer": 42}, 0.97 - - async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: - return {"answer": 42}, 0.97 - - -@pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) -@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) -def test_native_semantic_hit_stamps_similarity_on_request_metadata( - semantic_type: LiteLLMCacheType, use_async: bool -) -> None: - """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the - facade copies it to the caller's metadata; the native path must report it the same way.""" - facade: Final = Cache() - facade.type = semantic_type - facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter - metadata: Final[dict[str, object]] = {} - kwargs: Final = { - "cache_key": "semantic-key", - "messages": [{"role": "user", "content": "hello"}], - "metadata": metadata, - } - - result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) - - assert result == {"answer": 42} - assert metadata["semantic-similarity"] == 0.97 diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py deleted file mode 100644 index 2fbf9817a53..00000000000 --- a/tests/test_litellm_rust/test_ocr.py +++ /dev/null @@ -1,134 +0,0 @@ -import json -import threading -from collections.abc import Generator -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from io import BytesIO -from typing import Final - -import pytest - -import litellm - -pytestmark = pytest.mark.requires_rust_extension - - -@pytest.fixture -def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]]]: - requests: Final[list[dict[str, object]]] = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: - requests.append( - { - "headers": {name.lower(): value for name, value in self.headers.items()}, - "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), - } - ) - if self.headers.get("x-test-stall") == "true": - self.connection.settimeout(2) - try: - self.rfile.read(1) - except TimeoutError: - pass - return - if self.headers.get("User-Agent", "").startswith("python-httpx"): - self.send_response(418) - self.end_headers() - return - status = int(self.headers.get("x-test-status", "200")) - if status != 200: - body = b'{"error":"provider unavailable"}' - self.send_response(status) - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - return - response: Final = json.dumps( - { - "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], - "model": "mistral-ocr-latest", - "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, - } - ).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, format: str, *args: object) -> None: - pass - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread: Final = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) - thread.start() - try: - yield server, requests - finally: - server.shutdown() - server.server_close() - thread.join() - - -def test_native_lifecycle_core_encodes_python_file_input(ocr_server): - server, requests = ocr_server - litellm.rust(True) - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"}, - api_key="test-key", - api_base=f"http://127.0.0.1:{server.server_port}", - opaque_extension=object(), - ) - assert response.pages[0].markdown == "native OCR response" - assert requests[0]["body"]["document"] == {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} - assert "opaque_extension" not in requests[0]["body"] - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_native_ocr_failures_do_not_retry_on_python(ocr_server, asynchronous): - server, requests = ocr_server - arguments = { - "model": "mistral-ocr-latest", - "custom_llm_provider": "mistral", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "extra_headers": {"x-test-status": "503"}, - "num_retries": 0, - } - litellm.rust(True) - with pytest.raises(litellm.ServiceUnavailableError) as caught: - await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - assert caught.value.status_code == 503 - assert len(requests) == 1 - assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, asynchronous): - import asyncio - import time - - server, requests = ocr_server - litellm.rust(True) - arguments = { - "model": "mistral/mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "extra_headers": {"x-test-stall": "true"}, - "timeout": 0.1, - "num_retries": 0, - } - started = time.monotonic() - with pytest.raises(litellm.Timeout): - await asyncio.wait_for( - litellm.aocr(**arguments) if asynchronous else asyncio.to_thread(litellm.ocr, **arguments), - timeout=3, - ) - assert 0.09 <= time.monotonic() - started < 3 - assert len(requests) == 1 - assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") diff --git a/tests/test_litellm_rust/tokenizer/__init__.py b/tests/test_litellm_rust/tokenizer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/test_tokenizer.py b/tests/test_litellm_rust/tokenizer/test_fast_count.py similarity index 51% rename from tests/test_litellm_rust/test_tokenizer.py rename to tests/test_litellm_rust/tokenizer/test_fast_count.py index 98d5259b652..2902b79dca8 100644 --- a/tests/test_litellm_rust/test_tokenizer.py +++ b/tests/test_litellm_rust/tokenizer/test_fast_count.py @@ -12,67 +12,6 @@ from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOK pytestmark = pytest.mark.requires_rust_extension -def test_tiktoken_codec_round_trips_and_counts() -> None: - tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") - encoded: Final = tokenizer.encode("hello world") - - assert tokenizer.name == "cl100k_base" - assert tokenizer.count("hello world") == len(encoded) - assert tokenizer.decode(encoded) == "hello world" - - -def test_huggingface_codec_skips_special_tokens() -> None: - tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) - encoded: Final = tokenizer.encode("hello") - - assert "" in tokenizer.decode(encoded, skip_special_tokens=False) - assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" - - -def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None: - assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2" - assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base" - assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode( - "hi" - ) - - -def test_tiktoken_codec_exposes_its_vocabulary() -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") - - assert tokenizer.special_tokens() == reference._special_tokens - assert tokenizer.max_token_value() == reference.max_token_value - assert tokenizer.token_byte_values() == reference.token_byte_values() - assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello") - assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0) - with pytest.raises(KeyError): - tokenizer.encode_single_token(b"<|not-a-token|>") - - -def test_huggingface_codec_rejects_tiktoken_only_calls() -> None: - tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) - with pytest.raises(ValueError, match="requires a tiktoken encoding"): - tokenizer.token_byte_values() - with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"): - _native.Tokenizer.from_tiktoken("cl100k_base").get_vocab() - - -def test_unknown_tiktoken_encoding_raises_value_error() -> None: - with pytest.raises(ValueError, match="unsupported tokenizer"): - _native.Tokenizer.from_tiktoken("unknown-encoding") - - -def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name) - encoded: Final = reference.encode("🙂漢字") - - assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple( - reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1) - ) - - FAST_TEXTS: Final = ( "", "hello world <|endoftext|>", diff --git a/tests/test_litellm_rust/tokenizer/test_huggingface.py b/tests/test_litellm_rust/tokenizer/test_huggingface.py new file mode 100644 index 00000000000..05c5989c676 --- /dev/null +++ b/tests/test_litellm_rust/tokenizer/test_huggingface.py @@ -0,0 +1,24 @@ +from typing import Final + +import pytest + +from litellm.rust_bridge import _native +from litellm.utils import claude_json_str + +pytestmark = pytest.mark.requires_rust_extension + + +def test_huggingface_codec_skips_special_tokens() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + encoded: Final = tokenizer.encode("hello") + + assert "" in tokenizer.decode(encoded, skip_special_tokens=False) + assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" + + +def test_huggingface_codec_rejects_tiktoken_only_calls() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + with pytest.raises(ValueError, match="requires a tiktoken encoding"): + tokenizer.token_byte_values() + with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"): + _native.Tokenizer.from_tiktoken("cl100k_base").get_vocab() diff --git a/tests/test_litellm_rust/tokenizer/test_tiktoken.py b/tests/test_litellm_rust/tokenizer/test_tiktoken.py new file mode 100644 index 00000000000..c204d7eaf2f --- /dev/null +++ b/tests/test_litellm_rust/tokenizer/test_tiktoken.py @@ -0,0 +1,53 @@ +from typing import Final + +import pytest +import tiktoken + +from litellm.rust_bridge import _native + +pytestmark = pytest.mark.requires_rust_extension + + +def test_tiktoken_codec_round_trips_and_counts() -> None: + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + encoded: Final = tokenizer.encode("hello world") + + assert tokenizer.name == "cl100k_base" + assert tokenizer.count("hello world") == len(encoded) + assert tokenizer.decode(encoded) == "hello world" + + +def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None: + assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2" + assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base" + assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode( + "hi" + ) + + +def test_tiktoken_codec_exposes_its_vocabulary() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + + assert tokenizer.special_tokens() == reference._special_tokens + assert tokenizer.max_token_value() == reference.max_token_value + assert tokenizer.token_byte_values() == reference.token_byte_values() + assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello") + assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0) + with pytest.raises(KeyError): + tokenizer.encode_single_token(b"<|not-a-token|>") + + +def test_unknown_tiktoken_encoding_raises_value_error() -> None: + with pytest.raises(ValueError, match="unsupported tokenizer"): + _native.Tokenizer.from_tiktoken("unknown-encoding") + + +def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name) + encoded: Final = reference.encode("🙂漢字") + + assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple( + reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1) + ) From 1d039ed0095818beb522c90792796eef5ac28420 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 09:40:20 -0700 Subject: [PATCH 05/58] fix(proxy): give SpendLogToolIndex its own share of the cleanup budget and log a per-run summary (#41768) * fix(proxy): give SpendLogToolIndex its own share of the cleanup budget and log a per-run summary Resolves LIT-8090 starvation bug: _clean_spend_log_tables gave LiteLLM_SpendLogs and LiteLLM_SpendLogToolIndex one shared deadline, so a persistent SpendLogs backlog starved the index table of every delete batch. Split the group deadline between the two tables and emit one per-run summary line (WARNING when backlog remains). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rerun unit tests after an order-dependent allowlist flake Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 27 ++++-- .../proxy/test_spend_log_cleanup.py | 94 ++++++++++++++++++- 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 85e19fa8a32..c6f52bf074b 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -37,6 +37,7 @@ StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reac class TableCleanupResult: """Outcome of pruning one table, so the caller can report why a run ended.""" + table_name: str rows_deleted: int stop_reason: StopReason @@ -472,11 +473,11 @@ class SpendLogCleanup: from the last run that finished inside its budget. """ if time.monotonic() >= deadline: - return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) + return TableCleanupResult(table_name=table_name, rows_deleted=rows_deleted, stop_reason=stop_reason) remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline) if remaining is not None: SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining) - return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) + return TableCleanupResult(table_name=table_name, rows_deleted=rows_deleted, stop_reason=stop_reason) async def _delete_old_logs( self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float @@ -571,7 +572,9 @@ class SpendLogCleanup: ) verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped) - logs_result: Final = await self._delete_old_logs(prisma_client, cutoff_date, deadline) + logs_result: Final = await self._delete_old_logs( + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2) + ) verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted) index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline) @@ -638,6 +641,17 @@ class SpendLogCleanup: return "batch_cap_reached" return "completed" + @staticmethod + def _log_run_summary(outcome: RunOutcome, results: tuple[TableCleanupResult, ...], elapsed_seconds: float) -> None: + per_table: Final = ", ".join( + f"{result.table_name}: deleted={result.rows_deleted} stop_reason={result.stop_reason}" for result in results + ) + message: Final = "Spend log cleanup run finished: outcome=%s elapsed=%.1fs [%s]" + if outcome == "completed": + verbose_proxy_logger.info(message, outcome, elapsed_seconds, per_table) + return + verbose_proxy_logger.warning(message, outcome, elapsed_seconds, per_table) + async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ Main cleanup function. Deletes old spend logs in batches. @@ -724,9 +738,10 @@ class SpendLogCleanup: else () ) - SpendLogCleanupMetrics.record_run( - self._run_outcome(spend_log_results + session_results + health_check_results) - ) + results: Final = spend_log_results + session_results + health_check_results + outcome: Final = self._run_outcome(results) + SpendLogCleanupMetrics.record_run(outcome) + self._log_run_summary(outcome, results, time.monotonic() - run_started_at) except asyncio.CancelledError: verbose_proxy_logger.error( diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index e333da03950..72463e17c6b 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -3,6 +3,7 @@ Test cases for spend log cleanup functionality """ import asyncio +import logging import math import time from contextlib import asynccontextmanager @@ -1421,7 +1422,10 @@ def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(st results into one answer: a first-match-wins implementation would pass on whichever order happened to be written and fail on its mirror. """ - results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons) + results = tuple( + TableCleanupResult(table_name=f"t{i}", rows_deleted=0, stop_reason=reason) + for i, reason in enumerate(stop_reasons) + ) assert SpendLogCleanup._run_outcome(results) == expected @@ -1545,3 +1549,91 @@ async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): (error_call,) = mock_logger.error.call_args_list rendered = error_call[0][0] % error_call[0][1:] assert "(rows_deleted=100, batches=1)" in rendered + + +@pytest.mark.asyncio +async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): + """ + Both spend-log tables share one run budget. Before the fix the spend-log + loop ran against the whole deadline, so a backlog that outlasted the budget + meant LiteLLM_SpendLogToolIndex never received a single delete batch, run + after run. The index table must still get its own share of the budget. + """ + mock_prisma_client = MagicMock() + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=1000) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_max_batches": 500, + "maximum_spend_logs_cleanup_run_budget": "1s", + } + ) + cleaner.pod_lock_manager = None + + started_at = time.monotonic() + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + elapsed = time.monotonic() - started_at + + tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list] + assert tables.count("LiteLLM_SpendLogs") > 0 + assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog" + assert elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s" + + +@pytest.mark.asyncio +async def test_run_that_leaves_backlog_logs_a_warning_summary_naming_each_table(caplog): + """ + Operators running at warning or error level saw nothing when a run stopped + with expired rows still present. A run that ends on a bound must emit one + WARNING line that names every table, its rows deleted and its stop reason. + """ + mock_prisma_client = MagicMock() + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=1000) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "maximum_spend_logs_cleanup_max_batches": 2, + } + ) + cleaner.pod_lock_manager = None + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()] + assert len(summaries) == 1 + summary = summaries[0] + assert summary.levelno == logging.WARNING + message = summary.getMessage() + assert "outcome=batch_cap_reached" in message + assert "LiteLLM_SpendLogs: deleted=2000 stop_reason=batch_cap_reached" in message + assert "LiteLLM_SpendLogToolIndex: deleted=2000 stop_reason=batch_cap_reached" in message + + +@pytest.mark.asyncio +async def test_run_that_drains_every_table_logs_the_summary_at_info_not_warning(caplog): + """A healthy run must not page anyone: the summary stays at INFO.""" + mock_prisma_client = MagicMock() + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()] + assert len(summaries) == 1 + assert summaries[0].levelno == logging.INFO + assert "outcome=completed" in summaries[0].getMessage() From 081f73f021620bce438fb86a71435ec34a707d54 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:03:15 -0700 Subject: [PATCH 06/58] feat(rust): hand upstream response headers to the native Messages stream (#43178) * ci: drop the ocr_testing job now that tests/ocr_tests is gone Co-Authored-By: Claude Opus 5.5 * test(ocr): restore the live OCR matrix and the ocr_testing job The public litellm.ocr / aocr / Router interface is unchanged by the Rust migration, so the live provider matrix still applies. Drops the stale VCR skip list for the deleted test_rust_bridge.py. Co-Authored-By: Claude Opus 5.5 * test(messages): show streamed upstream headers never reach the native stream The Python handler puts the upstream response headers on the stream's _hidden_params before the first chunk so the proxy can forward them as llm_provider-* headers. The native route drops them, and this test fails on the Rust path while passing on Python. Co-Authored-By: Claude Fable 5.1 * feat(messages): hand upstream response headers to the native stream before its first chunk The Messages route fills MessagesStreamHead from the upstream response and yields it on Open. The Python driver converts it through the protocol host and hands it to Stream and SyncStream as their _hidden_params, so a streamed native call carries additional_headers the same way the Python handler does and the proxy can forward them as llm_provider-* headers. The relay contract lives in the core crate test, the hand-off in the host-python driver test, and the header projection in the route host test, so the recording-server test that showed the gap is dropped. Co-Authored-By: Claude Fable 5.1 * wip --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Opus 5.5 --- .../crates/core/src/messages/handler.rs | 4 +- .../crates/core/src/messages/route.rs | 38 +- .../crates/core/tests/messages/host.rs | 210 +++++++++ .../crates/core/tests/messages/main.rs | 1 + .../crates/core/tests/messages/request.rs | 436 +++++++++++++++++- .../crates/core/tests/messages/response.rs | 87 +++- .../crates/core/tests/messages/secrets.rs | 124 ++++- .../crates/core/tests/messages/stream.rs | 135 +++++- .../crates/host-python/src/adapter.rs | 7 + litellm-rust/crates/host-python/src/driver.rs | 187 +++++++- litellm-rust/crates/host-python/src/handle.rs | 8 +- .../python-bridge/src/routes/messages/host.rs | 9 +- .../python-bridge/src/routes/messages/mod.rs | 15 +- .../python-bridge/src/routes/ocr/host.rs | 4 + litellm/messages/dispatch.py | 11 +- litellm/rust_bridge/catalog.py | 1 + litellm/rust_bridge/lifecycle.py | 14 +- litellm/rust_bridge/messages/route_host.py | 9 + .../rust_bridge/messages/test_route_host.py | 12 + .../messages/test_callbacks.py | 57 ++- 20 files changed, 1274 insertions(+), 95 deletions(-) create mode 100644 litellm-rust/crates/core/tests/messages/host.rs diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index fe7e8bb4b80..de1a5f476ed 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -17,8 +17,10 @@ pub(super) async fn send( body: &Value, timeout: Option, ) -> Result { + let encoded = serde_json::to_vec(body) + .map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?; let builder = headers.iter().fold( - http_client().post(url).json(body), + http_client().post(url).body(encoded), |builder, (key, value)| builder.header(key, value), ); let builder = match timeout { diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index fc1a9b63252..40aff185e81 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -6,7 +6,6 @@ use std::{ use bytes::Bytes; use litellm_auth::SecretValue; -use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, @@ -22,7 +21,6 @@ use serde_json::{Map, Value}; use super::{ Error, - common_utils::messages_provider_config, handler::{decode_response, network, provider_error, send}, prepare::{prepare_provider_request, resolve_provider}, types::{MessagesRequest, MessagesShaping}, @@ -54,6 +52,11 @@ pub enum MessagesOutput { Streamed, } +/// The upstream response as the caller sees it at stream hand-off, before any chunk. +pub struct MessagesStreamHead { + pub headers: Vec<(String, String)>, +} + pub struct Messages; impl Protocol for Messages { @@ -62,7 +65,7 @@ impl Protocol for Messages { type Projection = MessagesCall; type Op = Infallible; type Chunk = Bytes; - type StreamHead = (); + type StreamHead = MessagesStreamHead; } impl From for Error { @@ -77,19 +80,6 @@ impl From for Error { pub type MessagesHost = HostChannel; pub type MessagesMachine = CallMachine; -/// Whether this route serves the request, decided before any callback runs so a host -/// can still run its own path. -pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool { - let provider = get_custom_llm_provider(model, custom_llm_provider) - .map(|resolved| resolved.custom_llm_provider) - .or(custom_llm_provider); - match provider { - Some(ANTHROPIC_MESSAGES_PROVIDER) => true, - Some(provider) => !stream && messages_provider_config(provider).is_some(), - None => false, - } -} - /// The in-process host for a request already in hand. It answers projection once and /// observes nothing. pub struct LocalMessagesHost { @@ -152,8 +142,11 @@ async fn execute( model: request.model.clone(), custom_llm_provider: request.provider.clone(), optional_params: Value::Object( - call.body - .iter() + request + .body + .as_object() + .into_iter() + .flatten() .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) .map(|(name, value)| (name.clone(), value.clone())) .collect(), @@ -193,7 +186,14 @@ async fn relay( host: &MessagesHost, mut response: reqwest::Response, ) -> Result { - if host.open(()).await? == Demand::Detached { + let head = MessagesStreamHead { + headers: response + .headers() + .iter() + .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) + .collect(), + }; + if host.open(head).await? == Demand::Detached { return Ok(MessagesOutput::Streamed); } while let Some(chunk) = response.chunk().await.map_err(network)? { diff --git a/litellm-rust/crates/core/tests/messages/host.rs b/litellm-rust/crates/core/tests/messages/host.rs new file mode 100644 index 00000000000..ca2aece5ebd --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/host.rs @@ -0,0 +1,210 @@ +use std::{convert::Infallible, sync::Mutex}; + +use litellm_core::messages::route::Messages; +use litellm_host::{ + event::{CallEvent, MachineEvent, RequestContext, WireRequest}, + host::Host, +}; +use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities; +use rstest::rstest; + +use super::*; + +type Rewrite = Box Result + Send + Sync>; + +/// Projects like `LocalMessagesHost`, answers `before_send` through `rewrite`, and keeps +/// every event the driver emits. +struct RecordingHost { + call: LocalMessagesHost, + rewrite: Rewrite, + events: Mutex>, + optional_params: Mutex>, +} + +impl RecordingHost { + fn new(call: MessagesCall, rewrite: Rewrite) -> Self { + Self { + call: LocalMessagesHost::new(call), + rewrite, + events: Mutex::new(Vec::new()), + optional_params: Mutex::new(Vec::new()), + } + } + + fn passthrough(call: MessagesCall) -> Self { + Self::new(call, Box::new(Ok)) + } + + fn raw_responses(&self) -> Vec { + self.events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + CallEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + Some(raw.body.clone()) + } + _ => None, + }) + .collect() + } +} + +impl Host for RecordingHost { + async fn project(&self) -> Result { + self.call.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn before_send( + &self, + wire: WireRequest, + context: &RequestContext, + ) -> Result { + self.optional_params + .lock() + .unwrap() + .push(context.optional_params.clone()); + (self.rewrite)(wire) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), Error> { + self.events.lock().unwrap().push(event.clone()); + Ok(()) + } +} + +async fn run_through(host: &RecordingHost) -> Result { + litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await +} + +fn authenticated(call: MessagesCall, api_base: String) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + ..call + } +} + +#[rstest] +#[tokio::test] +async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::new( + authenticated(call, upstream.uri()), + Box::new(|wire| { + let mut body = wire.body; + body["system"] = json!("added by the host"); + Ok(WireRequest { + headers: wire + .headers + .into_iter() + .chain([("x-host".to_string(), "seen".to_string())]) + .collect(), + body, + ..wire + }) + }), + ); + + run_through(&host).await.expect("messages call succeeds"); + + let request = only_request(&upstream).await; + assert_eq!(request.json()["system"], "added by the host"); + assert_eq!(request.header("x-host"), Some("seen")); + assert_eq!(request.header("x-api-key"), Some("sk-ant")); +} + +#[rstest] +#[tokio::test] +async fn a_before_send_failure_never_sends(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::new( + authenticated(call, upstream.uri()), + Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))), + ); + + let error = run_through(&host) + .await + .err() + .expect("the host failure fails the call"); + + assert_eq!(error, Error::InvalidRequest("vetoed by the host".into())); + assert!(received(&upstream).await.is_empty()); + assert!(host.raw_responses().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall) { + let raw = message_body(); + let upstream = upstream([json_response(raw.clone())]).await; + let host = RecordingHost::passthrough(authenticated(call, upstream.uri())); + + let output = run_through(&host).await.expect("messages call succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let [emitted] = <[String; 1]>::try_from(host.raw_responses()) + .unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len())); + assert_eq!(serde_json::from_str::(&emitted).unwrap(), raw); +} + +#[rstest] +#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))] +#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))] +#[tokio::test] +async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( + call: MessagesCall, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + let mut body = call.body.clone(); + body.insert("stream".into(), json!(true)); + let host = + RecordingHost::passthrough(authenticated(MessagesCall { body, ..call }, upstream.uri())); + + let _ = run_through(&host).await; + + assert_eq!(received(&upstream).await.len(), 1); + assert!(host.raw_responses().is_empty()); +} + +/// Python logs `optional_params` as what it is about to send, so a dropped param must +/// not resurface in callbacks. +#[rstest] +#[tokio::test] +async fn the_request_context_carries_the_shaped_params_without_model_or_messages( + call: MessagesCall, +) { + let upstream = upstream([message_response()]).await; + let body: Map = call + .body + .clone() + .into_iter() + .chain([("temperature".to_string(), json!(0.2))]) + .collect(); + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + body, + shaping: MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + }, + drop_params: true, + ..MessagesShaping::default() + }, + ..call + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len())); + assert_eq!(optional_params, json!({"max_tokens": 16})); +} diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 4e549bae309..21ee678ced3 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -14,6 +14,7 @@ use wiremock::ResponseTemplate; mod support; use support::*; +mod host; mod request; mod response; mod secrets; diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 9353324d370..2927356b773 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -1,3 +1,7 @@ +use litellm_llms::anthropic::common_utils::{ + ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities, + SupportedEffortTiers, beta, +}; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use rstest::rstest; @@ -132,8 +136,8 @@ async fn each_provider_posts_to_its_messages_endpoint( assert_eq!(request.method.as_str(), "POST"); assert_eq!(request.url.path(), path); assert_eq!(request.json()["model"], MODEL); - assert_eq!(request.header("anthropic-version"), Some("2023-06-01")); - assert_eq!(request.header("content-type"), Some("application/json")); + assert_eq!(request.header_values("anthropic-version"), ["2023-06-01"]); + assert_eq!(request.header_values("content-type"), ["application/json"]); } #[rstest] @@ -249,21 +253,423 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) assert_eq!(sent["top_k"], 3); } +fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall { + let body: Map = call.body.into_iter().chain(object(fields)).collect(); + MessagesCall { body, ..call } +} + +fn sent_betas(request: &wiremock::Request) -> Vec { + let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta")) + .unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}")); + header + .split(',') + .map(str::trim) + .map(str::to_string) + .collect() +} + #[rstest] -#[case::anthropic_streams(MODEL, Some("anthropic"), true, true)] -#[case::anthropic_prefix_streams("anthropic/claude-sonnet-4-5", None, true, true)] -#[case::azure_without_stream(MODEL, Some("azure_ai"), false, true)] -#[case::azure_stream(MODEL, Some("azure_ai"), true, false)] -#[case::other_provider(MODEL, Some("openai"), false, false)] -#[case::unresolvable_model("no-such-model", None, false, false)] -fn supports_matches_what_the_route_can_serve( - #[case] model: &str, - #[case] provider: Option<&str>, - #[case] stream: bool, - #[case] supported: bool, +#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] +#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] +#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] +#[case::context_management_edits( + json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}), + &[beta::CONTEXT_MANAGEMENT_2025_06_27] +)] +#[case::per_message_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &[beta::PER_TURN_CONTROL_2026_07_01] +)] +#[case::advisor_tool( + json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}), + &[beta::ADVISOR_TOOL_2026_03_01] +)] +#[case::several_features_at_once( + json!({"speed": "fast", "output_format": {"type": "json_schema"}}), + &[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01] +)] +#[tokio::test] +async fn feature_betas_join_the_callers_betas_in_one_sorted_header( + call: MessagesCall, + #[case] fields: Value, + #[case] features: &[&str], ) { + let upstream = upstream([message_response()]).await; + let capabilities = AnthropicModelCapabilities { + supports_speed: true, + ..AnthropicModelCapabilities::default() + }; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..call + }, + fields, + )) + .await; + + let sent = sent_betas(&only_request(&upstream).await); + let mut expected: Vec = features + .iter() + .map(|feature| feature.to_string()) + .chain(["caller-beta-2025-01-01".to_string()]) + .collect(); + expected.sort(); + assert_eq!(sent, expected); +} + +#[rstest] +#[tokio::test] +async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + api_key: Some("sk-ant-oat01-token".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + let request = only_request(&upstream).await; assert_eq!( - litellm_core::messages::route::supports(model, provider, stream), - supported + request.header("anthropic-dangerous-direct-browser-access"), + Some("true") + ); + assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]); + assert_eq!(request.header("x-api-key"), None); +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([ + ("Anthropic-Version", "2024-01-01"), + ("Content-Type", "application/json; charset=utf-8"), + ]), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]); + assert_eq!( + request.header_values("content-type"), + ["application/json; charset=utf-8"] ); } + +fn sampling_removed() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + } +} + +#[rstest] +#[case::sampling_params(sampling_removed(), json!({"temperature": 0.2, "top_p": 0.9, "top_k": 5}), &["temperature", "top_p", "top_k"], "temperature=0.2")] +#[case::speed(AnthropicModelCapabilities::default(), json!({"speed": "fast"}), &["speed"], "speed='fast'")] +#[tokio::test] +async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_it( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] dropped: &[&str], + #[case] rejected_as: &str, +) { + let upstream = upstream([message_response(), message_response()]).await; + let shaped = |drop_params: bool| { + with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: capabilities.clone(), + drop_params, + ..MessagesShaping::default() + }, + body: call.body.clone(), + custom_llm_provider: call.custom_llm_provider.clone(), + extra_headers: None, + provider_specific_header: None, + model: call.model.clone(), + timeout: call.timeout, + }, + fields.clone(), + ) + }; + + let error = run(shaped(false)) + .await + .err() + .expect("an unsupported param is rejected without drop_params"); + assert!( + matches!(&error, Error::InvalidRequest(message) if message.contains(rejected_as)), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); + + run_message(shaped(true)).await; + let sent = only_request(&upstream).await.json(); + for name in dropped { + assert_eq!(sent.get(*name), None, "{name} must be dropped"); + } + assert_eq!(sent["max_tokens"], 16); +} + +#[rstest] +#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))] +#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))] +#[tokio::test] +async fn reasoning_auto_summary_marks_active_thinking_on_the_wire( + call: MessagesCall, + #[case] thinking: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + ..AnthropicModelCapabilities::default() + }, + reasoning_auto_summary: true, + ..MessagesShaping::default() + }, + ..call + }, + json!({"thinking": thinking}), + )) + .await; + + assert_eq!(only_request(&upstream).await.json()["thinking"], expected); +} + +#[rstest] +#[case::reasoning_effort_on_an_adaptive_model( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { high: true, ..SupportedEffortTiers::default() }, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "high"}), + json!({"thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) +)] +#[case::reasoning_effort_on_a_legacy_model_caps_the_budget_below_max_tokens( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "high"}), + json!({"thinking": {"type": "enabled", "budget_tokens": 2999}}) +)] +#[case::adaptive_payload_on_a_legacy_model_becomes_a_capped_budget( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, "temperature": 0}), + json!({"thinking": {"type": "enabled", "budget_tokens": 2999}}) +)] +#[case::adaptive_payload_on_a_model_without_reasoning_is_dropped( + AnthropicModelCapabilities::default(), + json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}), + json!({}) +)] +#[tokio::test] +async fn reasoning_is_translated_by_the_model_capabilities( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..call + }, + [("max_tokens".to_string(), json!(3000))] + .into_iter() + .chain(object(fields)) + .collect(), + )) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("reasoning_effort"), None); + assert_eq!(sent.get("temperature"), None); + let reasoning: Map = ["thinking", "output_config"] + .into_iter() + .filter_map(|name| Some((name.to_string(), sent.get(name)?.clone()))) + .collect(); + assert_eq!(Value::Object(reasoning), expected); +} + +#[rstest] +#[case::empty_text_blocks( + json!([{"role": "assistant", "content": [{"type": "text", "text": " "}, {"type": "text", "text": "kept"}]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}]) +)] +#[case::provider_specific_fields( + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept", "provider_specific_fields": {"x": 1}}]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}]) +)] +#[case::unencrypted_web_search_results_become_text( + json!([{"role": "assistant", "content": [{ + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "title": "T", "url": "https://e.x", "page_age": null}] + }]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "Web search results:\n\nTitle: T\nURL: https://e.x"}]}]) +)] +#[tokio::test] +async fn replayed_history_is_cleaned_before_sending( + call: MessagesCall, + #[case] history: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({"messages": history}), + )) + .await; + + assert_eq!(only_request(&upstream).await.json()["messages"], expected); +} + +#[rstest] +#[tokio::test] +async fn metadata_is_reduced_to_the_user_id(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({"metadata": {"user_id": "u-1", "trace_id": "internal", "tags": ["a"]}}), + )) + .await; + + assert_eq!( + only_request(&upstream).await.json()["metadata"], + json!({"user_id": "u-1"}) + ); +} + +#[rstest] +#[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))] +#[case::missing_max_tokens(json!({"max_tokens": null}))] +#[tokio::test] +async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fields: Value) { + let upstream = upstream([message_response()]).await; + + let error = run(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + fields, + )) + .await + .err() + .expect("the request is rejected"); + + assert!(error.is_request(), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[tokio::test] +async fn azure_folds_system_role_messages_into_the_system_prompt(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + api_key: Some("sk-azure".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({ + "system": "top level", + "messages": [ + {"role": "system", "content": "from a message"}, + {"role": "user", "content": "hi"} + ] + }), + )) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!( + sent["system"], + json!([ + {"type": "text", "text": "top level"}, + {"type": "text", "text": "from a message"} + ]) + ); + assert_eq!(sent["messages"], json!([{"role": "user", "content": "hi"}])); +} + +#[rstest] +#[case::bare_model(MODEL, MODEL)] +#[case::one_prefix("anthropic/claude-sonnet-4-5", MODEL)] +#[case::doubled_prefix_loses_one_segment( + "anthropic/anthropic/claude-sonnet-4-5", + "anthropic/claude-sonnet-4-5" +)] +#[tokio::test] +async fn the_provider_prefix_is_stripped_exactly_once( + call: MessagesCall, + #[case] model: &str, + #[case] sent_model: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + model: model.into(), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(only_request(&upstream).await.json()["model"], sent_model); +} diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index 38a18c415ba..133b7d2b162 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -5,11 +5,14 @@ use rstest::rstest; use super::*; #[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] #[tokio::test] -async fn the_provider_message_is_returned(call: MessagesCall) { +async fn the_provider_message_is_returned(call: MessagesCall, #[case] provider: &str) { let upstream = upstream([message_response()]).await; let message = run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), api_key: Some("sk".into()), api_base: Some(upstream.uri()), ..call @@ -21,6 +24,88 @@ async fn the_provider_message_is_returned(call: MessagesCall) { assert_eq!(message.stop_reason.as_deref(), Some("end_turn")); } +/// A refusal and fields the route does not model come back exactly as the provider sent +/// them, since the Python side returns the raw message and the router decides what to do. +#[rstest] +#[tokio::test] +async fn the_message_passes_through_losslessly(call: MessagesCall) { + let upstream_body = json!({ + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, + {"type": "text", "text": "no", "citations": [{"type": "web_search_result_location", "url": "https://e.x"}]} + ], + "stop_reason": "refusal", + "stop_sequence": null, + "stop_details": {"type": "safeguard", "safeguard_types": ["dangerous_tool_use"]}, + "container": {"id": "container_1", "expires_at": "2026-01-01T00:00:00Z"}, + "context_management": {"applied_edits": []}, + "usage": {"input_tokens": 1, "output_tokens": 2, "server_tool_use": {"web_search_requests": 1}}, + "unknown_future_field": {"nested": true} + }); + let upstream = upstream([json_response(upstream_body.clone())]).await; + + let message = run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(message.stop_reason.as_deref(), Some("refusal")); + assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body); +} + +#[rstest] +#[tokio::test] +async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) { + let envelope = + json!({"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}); + let upstream = upstream([status_response(400, envelope.clone())]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + let Error::Transport(TransportError::Http { status, body }) = error else { + panic!("{error:?}"); + }; + assert_eq!(status, 400); + assert_eq!(serde_json::from_str::(&body).unwrap(), envelope); +} + +#[rstest] +#[tokio::test] +async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) { + let long = "x".repeat(600); + let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status: 500, + body: format!("{}... (truncated)", &long[..256]) + }) + ); +} + #[rstest] #[case::bad_request(400)] #[case::unauthorized(401)] diff --git a/litellm-rust/crates/core/tests/messages/secrets.rs b/litellm-rust/crates/core/tests/messages/secrets.rs index 419b6d6c753..55e510d00d3 100644 --- a/litellm-rust/crates/core/tests/messages/secrets.rs +++ b/litellm-rust/crates/core/tests/messages/secrets.rs @@ -1,23 +1,30 @@ -use litellm_llms::{ - anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, - azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, - base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, -}; use rstest::rstest; use super::*; #[rstest] -#[case::anthropic("anthropic", &ANTHROPIC_MESSAGES_CONFIG, "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "/v1/messages")] -#[case::azure_ai("azure_ai", &AZURE_ANTHROPIC_MESSAGES_CONFIG, "AZURE_API_KEY", "AZURE_API_BASE", "/anthropic/v1/messages")] +#[case::anthropic( + "anthropic", + "ANTHROPIC_API_KEY", + "ANTHROPIC_BASE_URL", + "/v1/messages", + &["ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"] +)] +#[case::azure_ai( + "azure_ai", + "AZURE_API_KEY", + "AZURE_API_BASE", + "/anthropic/v1/messages", + &["AZURE_API_KEY", "AZURE_API_BASE"] +)] #[tokio::test] async fn the_credential_and_base_come_from_the_secret_source( call: MessagesCall, #[case] provider: &str, - #[case] config: &dyn BaseAnthropicMessagesConfig, #[case] key_name: &str, #[case] base_name: &str, #[case] path: &str, + #[case] looked_up: &[&str], ) { let upstream = upstream([message_response()]).await; let base = upstream.uri(); @@ -40,7 +47,7 @@ async fn the_credential_and_base_come_from_the_secret_source( let request = only_request(&upstream).await; assert_eq!(request.url.path(), path); assert_eq!(request.header("x-api-key"), Some("sk-from-manager")); - assert_eq!(secrets.requested(), config.secret_names()); + assert_eq!(secrets.requested(), looked_up); } #[rstest] @@ -92,3 +99,102 @@ async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCa ); assert!(received(&upstream).await.is_empty()); } + +#[derive(Clone, Copy)] +enum Base { + Upstream, + Unreachable, + Blank, + Absent, +} + +fn base_value(base: Base, upstream: &str) -> Option { + match base { + Base::Upstream => Some(upstream.to_string()), + Base::Unreachable => Some(UNREACHABLE_BASE.to_string()), + Base::Blank => Some(" ".to_string()), + Base::Absent => None, + } +} + +#[rstest] +#[case::api_base_beats_base_url(Base::Upstream, Base::Unreachable)] +#[case::blank_api_base_falls_through_to_base_url(Base::Blank, Base::Upstream)] +#[case::base_url_alone(Base::Absent, Base::Upstream)] +#[tokio::test] +async fn the_anthropic_base_env_precedence_picks_the_upstream( + call: MessagesCall, + #[case] api_base: Base, + #[case] base_url: Base, +) { + let upstream = upstream([message_response()]).await; + let uri = upstream.uri(); + let values: Vec<(&str, &str)> = [ + ("ANTHROPIC_API_KEY", Some("sk-env".to_string())), + ("ANTHROPIC_API_BASE", base_value(api_base, &uri)), + ("ANTHROPIC_BASE_URL", base_value(base_url, &uri)), + ] + .iter() + .filter_map(|(name, value)| Some((*name, value.as_deref()?))) + .map(|(name, value)| (name, Box::leak(value.to_string().into_boxed_str()) as &str)) + .collect(); + + run_with(Arc::new(RecordingSecrets::new(values)), call) + .await + .expect("messages call reaches the upstream the precedence picks"); + + assert_eq!(only_request(&upstream).await.url.path(), "/v1/messages"); +} + +#[rstest] +#[case::auth_token_alone( + &[("ANTHROPIC_AUTH_TOKEN", "tok")], + ("authorization", "Bearer tok"), + "x-api-key" +)] +#[case::api_key_beats_the_auth_token( + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "tok")], + ("x-api-key", "sk-env"), + "authorization" +)] +#[tokio::test] +async fn the_auth_token_env_is_a_bearer_only_without_a_key( + call: MessagesCall, + #[case] values: &[(&str, &str)], + #[case] expected: (&str, &str), + #[case] absent: &str, +) { + let upstream = upstream([message_response()]).await; + + run_with( + Arc::new(RecordingSecrets::new(values.iter().copied())), + MessagesCall { + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + let request = only_request(&upstream).await; + let (name, value) = expected; + assert_eq!(request.header_values(name), [value]); + assert_eq!(request.header(absent), None); +} + +#[rstest] +#[tokio::test] +async fn azure_without_a_base_anywhere_fails_before_sending(call: MessagesCall) { + let error = run_with( + Arc::new(RecordingSecrets::new([("AZURE_API_KEY", "sk-azure")])), + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + ..call + }, + ) + .await + .err() + .expect("azure needs a base"); + + assert_eq!(error, Error::Auth(litellm_auth::Error::MissingAzureApiBase)); +} diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index ea23a9e8e38..c4be3127d66 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -1,16 +1,25 @@ use std::{convert::Infallible, sync::Mutex}; use bytes::Bytes; -use litellm_core::messages::route::Messages; +use litellm_core::messages::route::{Messages, MessagesStreamHead}; use litellm_host::host::{Demand, Host}; use rstest::rstest; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; use super::*; +const UPSTREAM_HEADERS: [(&str, &str); 2] = [ + ("request-id", "req_upstream_123"), + ("anthropic-ratelimit-requests-remaining", "41"), +]; + const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; enum Seen { - Open, + Open(Vec<(String, String)>), Deliver(Bytes), } @@ -50,8 +59,8 @@ impl Host for RecordingStreamHost { match op {} } - async fn open(&self, (): ()) -> Result { - Ok(self.record(Seen::Open)) + async fn open(&self, head: MessagesStreamHead) -> Result { + Ok(self.record(Seen::Open(head.headers))) } async fn deliver(&self, chunk: Bytes) -> Result { @@ -71,7 +80,10 @@ fn streaming(call: MessagesCall, api_base: String) -> MessagesCall { } fn sse_response() -> ResponseTemplate { - ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream") + UPSTREAM_HEADERS.iter().fold( + ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream"), + |response, (name, value)| response.insert_header(*name, *value), + ) } async fn stream_through(host: &RecordingStreamHost) -> Result { @@ -80,7 +92,7 @@ async fn stream_through(host: &RecordingStreamHost) -> Result = headers + .iter() + .filter(|(name, _)| { + UPSTREAM_HEADERS + .iter() + .any(|(upstream, _)| upstream == name) + }) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect(); + assert_eq!(surfaced, UPSTREAM_HEADERS); let delivered: Vec = chunks .iter() .flat_map(|step| match step { Seen::Deliver(chunk) => chunk.to_vec(), - Seen::Open => panic!("the stream opens exactly once"), + Seen::Open(_) => panic!("the stream opens exactly once"), }) .collect(); assert_eq!(delivered, SSE_BODY.as_bytes()); @@ -118,9 +140,18 @@ async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] det } #[rstest] +#[case::text_body(ResponseTemplate::new(429).set_body_string("slow down"), "slow down")] +#[case::json_envelope( + status_response(429, json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}})), + r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"# +)] #[tokio::test] -async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: MessagesCall) { - let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; +async fn an_upstream_error_fails_the_call_without_opening_the_stream( + call: MessagesCall, + #[case] response: ResponseTemplate, + #[case] body: &str, +) { + let upstream = upstream([response]).await; let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); let error = stream_through(&host) @@ -128,16 +159,88 @@ async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: Messa .err() .expect("upstream error propagates"); - assert!( - matches!( - error, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "{error:?}" + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status: 429, + body: body.into() + }) ); assert!(host.seen.into_inner().unwrap().is_empty()); } +/// The native route relays bytes as they are. Python's synthetic `api_error` for a stream +/// that never reaches `message_stop` lives in its SSE wrapper, above this route. +#[rstest] +#[tokio::test] +async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: MessagesCall) { + const INCOMPLETE: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n"; + let upstream = + upstream([ResponseTemplate::new(200).set_body_raw(INCOMPLETE, "text/event-stream")]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + stream_through(&host).await.expect("streamed call succeeds"); + + let delivered: Vec = host + .seen + .into_inner() + .unwrap() + .iter() + .flat_map(|step| match step { + Seen::Deliver(chunk) => chunk.to_vec(), + Seen::Open(_) => Vec::new(), + }) + .collect(); + assert_eq!(delivered, INCOMPLETE.as_bytes()); +} + +/// Serves one SSE chunk and then holds the connection open without ever finishing. +async fn stalling_upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n\ + 1f\r\nevent: message_start\ndata: {}\n\n\r\n", + ) + .await + .unwrap(); + std::future::pending::<()>().await; + }); + base +} + +#[rstest] +#[tokio::test] +async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { + let base = stalling_upstream().await; + let host = RecordingStreamHost::new( + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + usize::MAX, + ); + + let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host)) + .await + .expect("the stalled stream gives up within the timeout") + .err() + .expect("a stalled body fails the call"); + + assert!(matches!(error, Error::Transport(_)), "{error:?}"); + let seen = host.seen.into_inner().unwrap(); + assert!( + matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"), + "the chunk before the stall reached the caller, saw {} ops", + seen.len() + ); +} + #[rstest] #[tokio::test] async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) { diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 87481aa89b7..7f07475bc4c 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -134,6 +134,13 @@ pub trait ProtocolHost: Send + Sync { response: ::Response, ) -> PyResult>; + /// What the stream carries at hand-off, as the caller's stream receives it. + fn head( + &mut self, + py: Python<'_>, + head: ::StreamHead, + ) -> PyResult>; + /// One streamed chunk as the caller receives it. fn chunk( &mut self, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index aaa0752522b..372af2843bd 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -134,10 +134,10 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Open => py + ExecutionStep::Open(head) => py .import("litellm.rust_bridge.lifecycle")? .getattr("SyncStream")? - .call1((Py::new(py, Execution::suspended(driver))?,)) + .call1((Py::new(py, Execution::suspended(driver))?, head)) .map(Bound::unbind), ExecutionStep::Await(_) | ExecutionStep::Yield(_) => { Err(PyRuntimeError::new_err("sync call suspended")) @@ -312,7 +312,7 @@ where Ok(_) => return Err(missing_state()), Err(error) => Err(error), }, - HostOp::Open(_, reply) => return self.opened(py, reply).map(Next::Return), + HostOp::Open(head, reply) => return self.opened(py, head, reply).map(Next::Return), HostOp::Deliver(chunk, reply) => { return self.delivered(py, chunk, reply).map(Next::Return); } @@ -340,12 +340,21 @@ where } } - fn opened(&mut self, py: Python<'_>, reply: Reply) -> PyResult { + fn opened( + &mut self, + py: Python<'_>, + head: as Protocol>::StreamHead, + reply: Reply, + ) -> PyResult { self.stage = Stage::Streaming; + let head = match self.host.head(py, head) { + Ok(head) => head, + Err(error) => return self.interrupt(py, error), + }; match self.adapter.opened(py) { Ok(()) => { self.pending = Some(Pending::Consumer(reply)); - Ok(ExecutionStep::Open) + Ok(ExecutionStep::Open(head)) } Err(error) => self.interrupt(py, error), } @@ -699,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .map(|answer| reply.send(answer)) } + fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult> { + match head {} + } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } @@ -945,6 +958,163 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + struct Streaming; + + impl Protocol for Streaming { + type Response = (); + type Error = Error; + type Projection = (); + type Op = std::convert::Infallible; + type Chunk = &'static str; + type StreamHead = Vec<(&'static str, &'static str)>; + } + + struct StreamingHost; + + impl ProtocolHost for StreamingHost { + type Protocol = Streaming; + type Failure = Classified; + + fn project( + &mut self, + _: Python<'_>, + _: &Bound<'_, PyDict>, + ) -> Result<(), InvokeError> { + Ok(()) + } + + fn invoke( + &mut self, + _: Python<'_>, + op: std::convert::Infallible, + ) -> Result<(), InvokeError> { + match op {} + } + + fn head( + &mut self, + py: Python<'_>, + head: Vec<(&'static str, &'static str)>, + ) -> PyResult> { + let headers = PyDict::new(py); + for (name, value) in head { + headers.set_item(name, value)?; + } + let hidden = PyDict::new(py); + hidden.set_item("additional_headers", headers)?; + Ok(hidden.into_any().unbind()) + } + + fn chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> { + Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind()) + } + + fn complete(&mut self, py: Python<'_>, (): ()) -> PyResult> { + Ok(py.None()) + } + + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + Ok(Classified(error.0)) + } + + fn host_error(error: &PyErr) -> Error { + Error(error.to_string()) + } + + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } + } + + fn streaming_machine() -> CallMachine { + CallMachine::new(|host| { + Box::pin(async move { + host.project().await?; + if host.open(vec![("request-id", "req_1")]).await? == Demand::Detached { + return Ok(()); + } + for chunk in ["first", "second"] { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(()) + }) + }) + } + + /// Drives a `Stream` (async) or `SyncStream` to completion from a sync test. + fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec { + if !asynchronous { + return stream + .try_iter() + .unwrap() + .map(|chunk| chunk.unwrap().extract().unwrap()) + .collect(); + } + std::iter::from_fn(|| { + let stop = stream + .call_method0("__anext__") + .unwrap() + .call_method1("send", (py.None(),)) + .unwrap_err(); + if stop.is_instance_of::(py) { + return None; + } + assert!(stop.is_instance_of::(py)); + Some(stop.value(py).getattr("value").unwrap().extract().unwrap()) + }) + .collect() + } + + #[test] + fn a_stream_carries_its_head_as_hidden_params_before_the_first_chunk() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + install_lifecycle_module(py); + for asynchronous in [false, true] { + let log = Log::default(); + let adapter = SyntheticAdapter { + log: Log(log.0.clone()), + script: AdapterScript::Plain, + }; + let handed = run_call( + py, + streaming_machine(), + StreamingHost, + Box::new(adapter), + PyDict::new(py).unbind(), + asynchronous, + ) + .unwrap(); + let stream = if asynchronous { + let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err(); + stop.value(py).getattr("value").unwrap() + } else { + handed.into_bound(py) + }; + let hidden: std::collections::HashMap< + String, + std::collections::HashMap, + > = stream.getattr("_hidden_params").unwrap().extract().unwrap(); + assert_eq!( + hidden["additional_headers"], + std::collections::HashMap::from([( + "request-id".to_string(), + "req_1".to_string() + )]) + ); + assert_eq!(log.entries(), ["started", "begin", "opened"]); + assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]); + } + }); + } + fn failing_machine() -> CallMachine { CallMachine::new(|host| { Box::pin(async move { @@ -1202,6 +1372,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) -> Result<(), InvokeError> { Err(missing_state().into()) } + fn head( + &mut self, + _: Python<'_>, + head: std::convert::Infallible, + ) -> PyResult> { + match head {} + } fn chunk( &mut self, _: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 10abbadbda5..24adfd404d7 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -8,9 +8,9 @@ use pyo3::prelude::*; pub enum ExecutionStep { Return(Py), Await(Py), - /// The call streams: the caller gets a stream over this execution, which stays - /// suspended until the stream asks for a chunk. - Open, + /// The call streams: the caller gets a stream over this execution carrying this head, + /// and the execution stays suspended until the stream asks for a chunk. + Open(Py), Yield(Py), } @@ -75,7 +75,7 @@ impl Execution { let step = body.resume(result)?; let (tag, value, suspended) = match step { ExecutionStep::Await(value) => ("Await", value, true), - ExecutionStep::Open => ("Open", py.None(), true), + ExecutionStep::Open(head) => ("Open", head, true), ExecutionStep::Yield(value) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 6de4e1320e1..a253f4f5670 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,7 @@ use std::convert::Infallible; use bytes::Bytes; use litellm_core::messages::{ Error, - route::{Messages, MessagesCall, MessagesOutput}, + route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead}, types::MessagesShaping, }; use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; @@ -238,6 +238,13 @@ impl ProtocolHost for MessagesPythonHost { } } + fn head(&mut self, py: Python<'_>, head: MessagesStreamHead) -> PyResult> { + py.import(ROUTE_HOST_MODULE)? + .getattr("stream_hidden_params")? + .call1((to_py(py, &head.headers)?,)) + .map(Bound::unbind) + } + fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { Ok(PyBytes::new(py, &chunk).into_any().unbind()) } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 65040f31684..dae8623979a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -4,14 +4,12 @@ use host::MessagesPythonHost; use litellm_callbacks_legacy_python::{ LegacySurface, PassThroughStream, PublicCall, run_legacy_call, }; -use litellm_core::messages::route::{messages_machine, supports}; +use litellm_core::messages::route::messages_machine; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::errors::RustBridgeDeclined; - const SURFACE: LegacySurface = LegacySurface { call_type: "anthropic_messages", input_description: "Messages", @@ -28,17 +26,6 @@ fn run_messages( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - let model: String = request.getattr("model")?.extract()?; - let provider: Option = request.getattr("custom_llm_provider")?.extract()?; - let stream = request - .getattr("stream")? - .extract::>()? - .unwrap_or(false); - if !supports(&model, provider.as_deref(), stream) { - return Err(RustBridgeDeclined::new_err( - "the Rust Messages route does not serve this provider", - )); - } let secrets = crate::secrets::source(py)?; run_legacy_call( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 5a3806e61e3..dc01ced15a0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -117,6 +117,10 @@ impl ProtocolHost for OcrPythonHost { .map(Bound::unbind) } + fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult> { + match head {} + } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index a0c791a136c..13a030e7ebe 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -3,6 +3,8 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Itera from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.experimental_pass_through.messages import handler as main from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook @@ -71,10 +73,17 @@ def _public_request( ) +def _resolved_provider(request: LiteLLMMessagesRequest) -> str | None: + try: + return get_llm_provider(request.model, request.custom_llm_provider)[1] + except BadRequestError: + return request.custom_llm_provider + + def _context(request: LiteLLMMessagesRequest) -> RouteContext: return RouteContext( Route.MESSAGES, - provider=request.custom_llm_provider, + provider=_resolved_provider(request), model=request.model, delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED, ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index d9834adc7e8..6e455817194 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -109,6 +109,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index 2f243e8c212..7a6485a5f2c 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Iterator +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping from dataclasses import dataclass from typing import Final, Protocol @@ -17,7 +17,7 @@ class Complete: @dataclass(frozen=True, slots=True) class Open: - value: None + value: Mapping[str, object] | None @dataclass(frozen=True, slots=True) @@ -68,7 +68,7 @@ async def drive(execution: Execution) -> object: step: Final = await _settle(execution, execution.start()) if isinstance(step, Open): handed_off = True - return Stream(execution) + return Stream(execution, step.value) return step.value finally: if not handed_off: @@ -78,10 +78,10 @@ async def drive(execution: Execution) -> object: class Stream(AsyncIterator[object]): """A streamed native call: each read resumes the execution until its next chunk.""" - def __init__(self, execution: Execution) -> None: + def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None: self._execution: Final = execution self._done = False - self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place + self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it def __aiter__(self) -> Stream: return self @@ -115,10 +115,10 @@ class Stream(AsyncIterator[object]): class SyncStream(Iterator[object]): """The sync form of `Stream`; its execution never suspends on an awaitable.""" - def __init__(self, execution: Execution) -> None: + def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None: self._execution: Final = execution self._done = False - self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place + self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it def __iter__(self) -> SyncStream: return self diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index d49d7b75a6f..0a23989a59c 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from typing import Final, cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict +import httpx from pydantic import TypeAdapter, ValidationError import litellm @@ -53,6 +54,14 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: ) +def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + anthropic_messages_stream_hidden_params, + ) + + return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) + + def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: return request.kwargs diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/test_litellm/rust_bridge/messages/test_route_host.py index f47333a45d9..c5a442e0709 100644 --- a/tests/test_litellm/rust_bridge/messages/test_route_host.py +++ b/tests/test_litellm/rust_bridge/messages/test_route_host.py @@ -110,3 +110,15 @@ def test_native_request_rejections_map_to_the_public_400() -> None: assert "does not support top_k=5" in mapped.message assert mapped.model == "claude-sonnet-5" assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError) + + +def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None: + hidden: Final = route_host.stream_hidden_params( + (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41")) + ) + + additional: Final = hidden["additional_headers"] + assert isinstance(additional, dict) + assert additional["llm_provider-request-id"] == "req_upstream_123" + assert additional["x-ratelimit-remaining-requests"] == "41" + assert "request-id" not in additional diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index 19043780eb6..dc66852d214 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -127,7 +127,7 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on **arguments(messages_server, stream=True, callbacks=[recorder]) ) assert isinstance(stream, AsyncIterator) - assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} + assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" first: Final = await anext(stream) await drain_logging() assert "async_log_success_event" not in recorder.names @@ -171,7 +171,7 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) assert isinstance(stream, Iterator) - assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} + assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" assert b"".join(stream) == sse_payload() assert_served_natively(messages_server) @@ -186,3 +186,56 @@ def test_native_sync_messages_returns_the_provider_message(messages_server: Reco assert_served_natively(messages_server) assert response["content"] == MESSAGES_RESPONSE["content"] assert len(recorder.wait_for("log_success_event")) == 1 + + +@pytest.mark.asyncio +async def test_native_messages_pre_call_sees_the_shaped_optional_params( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + await litellm.anthropic.messages.acreate( + **arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True) + ) + + sent: Final = messages_server.requests[0].body + assert not {"temperature", "top_k"} & sent.keys() + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + optional_params: Final = pre_call["optional_params"] + assert isinstance(optional_params, dict) + assert not {"model", "messages", "temperature", "top_k"} & optional_params.keys() + assert optional_params["max_tokens"] == sent["max_tokens"] + + +@pytest.mark.asyncio +async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None: + class Broken(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + raise RuntimeError("logger exploded") + + response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()])) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + + +@pytest.mark.asyncio +async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + async for _ in stream: + pass + + success: Final = await recorder.wait_for_async("async_log_success_event") + usage: Final = success[0].response.usage + assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"] + assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] + assert success[0].response.choices[0].message.content == "Hello from native Messages" From 9c10e0f985787875ddb36818959b0667f41c9cf6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:27:30 +0000 Subject: [PATCH 07/58] feat(testkit): agent clients for Claude Code, Codex and opencode (#43181) * feat(testkit): install and configure Claude Code, Codex and opencode against a gateway Co-Authored-By: Claude Sonnet 5 * refactor(testkit): derive targets from target-lexicon and split agents behind a trait Co-Authored-By: Claude Sonnet 5 * refactor(testkit): group sources into agent and install folders Co-Authored-By: Claude Sonnet 5 * feat(testkit): split agents into install, configure and drive with semver-aware launch Co-Authored-By: Claude Sonnet 5 * style(testkit): drop a needless borrow Co-Authored-By: Claude Sonnet 5 --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Sonnet 5 --- litellm-rust/Cargo.lock | 152 +++++++++- litellm-rust/Cargo.toml | 6 + litellm-rust/crates/testkit/Cargo.toml | 32 +++ .../crates/testkit/src/agent/claude.rs | 181 ++++++++++++ .../crates/testkit/src/agent/codex.rs | 174 ++++++++++++ .../crates/testkit/src/agent/configure.rs | 69 +++++ .../crates/testkit/src/agent/drive.rs | 57 ++++ .../crates/testkit/src/agent/install.rs | 17 ++ litellm-rust/crates/testkit/src/agent/mod.rs | 20 ++ .../crates/testkit/src/agent/opencode.rs | 187 +++++++++++++ litellm-rust/crates/testkit/src/error.rs | 56 ++++ .../crates/testkit/src/install/archive.rs | 52 ++++ .../crates/testkit/src/install/fetch.rs | 55 ++++ .../crates/testkit/src/install/mod.rs | 118 ++++++++ .../crates/testkit/src/install/release.rs | 65 +++++ litellm-rust/crates/testkit/src/lib.rs | 15 + litellm-rust/crates/testkit/src/session.rs | 76 +++++ litellm-rust/crates/testkit/src/target.rs | 69 +++++ .../crates/testkit/tests/configure.rs | 133 +++++++++ litellm-rust/crates/testkit/tests/install.rs | 262 ++++++++++++++++++ litellm-rust/crates/testkit/tests/live.rs | 133 +++++++++ litellm-rust/crates/testkit/tests/session.rs | 155 +++++++++++ .../crates/testkit/tests/support/mod.rs | 70 +++++ 23 files changed, 2151 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/testkit/Cargo.toml create mode 100644 litellm-rust/crates/testkit/src/agent/claude.rs create mode 100644 litellm-rust/crates/testkit/src/agent/codex.rs create mode 100644 litellm-rust/crates/testkit/src/agent/configure.rs create mode 100644 litellm-rust/crates/testkit/src/agent/drive.rs create mode 100644 litellm-rust/crates/testkit/src/agent/install.rs create mode 100644 litellm-rust/crates/testkit/src/agent/mod.rs create mode 100644 litellm-rust/crates/testkit/src/agent/opencode.rs create mode 100644 litellm-rust/crates/testkit/src/error.rs create mode 100644 litellm-rust/crates/testkit/src/install/archive.rs create mode 100644 litellm-rust/crates/testkit/src/install/fetch.rs create mode 100644 litellm-rust/crates/testkit/src/install/mod.rs create mode 100644 litellm-rust/crates/testkit/src/install/release.rs create mode 100644 litellm-rust/crates/testkit/src/lib.rs create mode 100644 litellm-rust/crates/testkit/src/session.rs create mode 100644 litellm-rust/crates/testkit/src/target.rs create mode 100644 litellm-rust/crates/testkit/tests/configure.rs create mode 100644 litellm-rust/crates/testkit/tests/install.rs create mode 100644 litellm-rust/crates/testkit/tests/live.rs create mode 100644 litellm-rust/crates/testkit/tests/session.rs create mode 100644 litellm-rust/crates/testkit/tests/support/mod.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2d96efa6077..e7d911f5fd9 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -73,6 +73,15 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + [[package]] name = "arc-swap" version = "1.9.2" @@ -1470,6 +1479,17 @@ dependencies = [ "serde_core", ] +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "derive_builder" version = "0.20.2" @@ -1643,6 +1663,16 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -3453,6 +3483,27 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-testkit" +version = "0.1.0" +dependencies = [ + "flate2", + "futures-util", + "reqwest 0.12.28", + "rstest", + "semver", + "serde", + "serde_json", + "sha2 0.10.9", + "tar", + "target-lexicon", + "tempfile", + "thiserror 2.0.19", + "tokio", + "toml", + "zip", +] + [[package]] name = "litellm-token-counter" version = "0.1.0" @@ -5206,6 +5257,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -5537,6 +5597,17 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", + "xattr", +] + [[package]] name = "target-lexicon" version = "0.13.5" @@ -5806,6 +5877,30 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap 2.14.0", + "serde_core", + "serde_spanned", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -5822,9 +5917,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" dependencies = [ "indexmap 2.14.0", - "toml_datetime", + "toml_datetime 1.1.1+spec-1.1.0", "toml_parser", - "winnow", + "winnow 1.0.4", ] [[package]] @@ -5833,9 +5928,15 @@ version = "1.1.3+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" dependencies = [ - "winnow", + "winnow 1.0.4", ] +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + [[package]] name = "tonic" version = "0.14.6" @@ -6597,6 +6698,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" + [[package]] name = "winnow" version = "1.0.4" @@ -6659,6 +6766,16 @@ dependencies = [ "time", ] +[[package]] +name = "xattr" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" +dependencies = [ + "libc", + "rustix", +] + [[package]] name = "xmlparser" version = "0.13.6" @@ -6784,6 +6901,23 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zip" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "flate2", + "indexmap 2.14.0", + "memchr", + "thiserror 2.0.19", + "zopfli", +] + [[package]] name = "zlib-rs" version = "0.6.7" @@ -6795,3 +6929,15 @@ name = "zmij" version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 0c7236e807e..022e8f13311 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -81,6 +81,12 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" +flate2 = "1" +semver = "1" +tar = "0.4" +target-lexicon = "0.13.5" +tempfile = "3" +zip = { version = "2", default-features = false, features = ["deflate"] } moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" diff --git a/litellm-rust/crates/testkit/Cargo.toml b/litellm-rust/crates/testkit/Cargo.toml new file mode 100644 index 00000000000..98a36a1e87f --- /dev/null +++ b/litellm-rust/crates/testkit/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "litellm-testkit" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +publish = false + +[dependencies] +flate2.workspace = true +reqwest.workspace = true +serde.workspace = true +semver.workspace = true +serde_json.workspace = true +sha2.workspace = true +tar.workspace = true +target-lexicon.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["fs", "process"] } +zip.workspace = true + +[dev-dependencies] +flate2.workspace = true +rstest.workspace = true +sha2.workspace = true +tar.workspace = true +target-lexicon.workspace = true +futures-util.workspace = true +tempfile.workspace = true +tokio.workspace = true +toml = "0.9" +zip.workspace = true diff --git a/litellm-rust/crates/testkit/src/agent/claude.rs b/litellm-rust/crates/testkit/src/agent/claude.rs new file mode 100644 index 00000000000..6870fff6bf8 --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/claude.rs @@ -0,0 +1,181 @@ +use std::collections::BTreeMap; +use std::path::Path; + +use semver::Version; +use serde::Deserialize; + +use super::{ + Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines, + path_string, +}; +use crate::install::release::parse; +use crate::install::{Packaging, Release}; +use crate::{Error, Fetch, Target}; + +const RELEASES: &str = "https://downloads.claude.ai/claude-code-releases"; + +pub struct ClaudeCode; + +#[derive(Deserialize)] +struct Manifest { + platforms: BTreeMap, +} + +#[derive(Deserialize)] +struct Platform { + checksum: String, +} + +impl Install for ClaudeCode { + fn binary(&self) -> &'static str { + "claude" + } + + async fn release( + &self, + fetch: &impl Fetch, + version: &Version, + target: Target, + ) -> Result { + let manifest_url = format!("{RELEASES}/{version}/manifest.json"); + let manifest: Manifest = parse(&manifest_url, &fetch.get(&manifest_url).await?)?; + let key = format!( + "{}-{}{}", + target.os_name(), + target.arch_name(), + target.musl_suffix() + ); + let platform = manifest + .platforms + .get(&key) + .ok_or_else(|| Error::AssetNotFound(key.clone()))?; + Ok(Release { + url: format!("{RELEASES}/{version}/{key}/claude"), + asset: key, + sha256: platform.checksum.clone(), + packaging: Packaging::Bare, + }) + } +} + +impl Configure for ClaudeCode { + fn configure( + &self, + _version: &Version, + settings: &Settings, + home: &Path, + ) -> Result { + if settings.wire != Wire::Messages { + return Err(Error::UnsupportedWire { + agent: "claude", + wire: settings.wire, + }); + } + Ok(LaunchSpec { + env: env([ + ("HOME", path_string(home)), + ("CLAUDE_CONFIG_DIR", path_string(&home.join(".claude"))), + ("ANTHROPIC_BASE_URL", settings.base_url.clone()), + ("ANTHROPIC_AUTH_TOKEN", settings.api_key.clone()), + ("ANTHROPIC_MODEL", settings.model.clone()), + ("DISABLE_AUTOUPDATER", "1".to_owned()), + ]), + files: BTreeMap::new(), + }) + } +} + +#[derive(Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum Event { + Assistant { + message: AssistantMessage, + }, + Result(Finished), + #[serde(other)] + Other, +} + +#[derive(Deserialize)] +struct AssistantMessage { + content: Vec, +} + +#[derive(Deserialize)] +struct Block { + #[serde(rename = "type")] + kind: String, + name: Option, +} + +#[derive(Deserialize)] +struct Finished { + is_error: bool, + result: Option, + usage: Option, +} + +#[derive(Deserialize)] +struct TokenUsage { + input_tokens: u64, + output_tokens: u64, +} + +impl Drive for ClaudeCode { + fn args(&self, _version: &Version, settings: &Settings, prompt: &Prompt) -> Vec { + let base = [ + "-p", + &prompt.text, + "--output-format", + "stream-json", + "--verbose", + "--model", + &settings.model, + ]; + let tools = ["--allowedTools", "Bash,Read,Write,Edit"]; + base.into_iter() + .chain(tools.into_iter().filter(|_| prompt.allow_tools)) + .map(str::to_owned) + .collect() + } + + fn parse(&self, _version: &Version, stdout: &str) -> Outcome { + let events: Vec = json_lines(stdout).collect(); + let tool_calls = events + .iter() + .filter_map(|event| match event { + Event::Assistant { message } => Some(&message.content), + _ => None, + }) + .flatten() + .filter(|block| block.kind == "tool_use") + .filter_map(|block| block.name.clone()) + .collect(); + let finished = events.into_iter().find_map(|event| match event { + Event::Result(finished) => Some(finished), + _ => None, + }); + let Some(finished) = finished else { + return Outcome { + tool_calls, + ..Outcome::default() + }; + }; + let result = finished.result.unwrap_or_default(); + let (text, errors) = if finished.is_error { + (String::new(), vec![result]) + } else { + (result, Vec::new()) + }; + Outcome { + text, + tool_calls, + usage: finished.usage.map_or_else(Usage::default, |usage| Usage { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + }), + errors, + exit_code: None, + } + } +} diff --git a/litellm-rust/crates/testkit/src/agent/codex.rs b/litellm-rust/crates/testkit/src/agent/codex.rs new file mode 100644 index 00000000000..6749e81a471 --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/codex.rs @@ -0,0 +1,174 @@ +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use semver::Version; +use serde::Deserialize; + +use super::{ + Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines, + path_string, quoted, v1, +}; +use crate::install::release::github_release; +use crate::install::{Packaging, Release}; +use crate::target::{Arch, Os}; +use crate::{Error, Fetch, Target}; + +const RELEASES: &str = "https://api.github.com/repos/openai/codex/releases/tags"; + +pub struct Codex; + +fn triple(target: Target) -> String { + let arch = match target.arch { + Arch::Aarch64 => "aarch64", + Arch::X86_64 => "x86_64", + }; + match target.os { + Os::Macos => format!("{arch}-apple-darwin"), + Os::Linux => format!("{arch}-unknown-linux-musl"), + } +} + +impl Install for Codex { + fn binary(&self) -> &'static str { + "codex" + } + + async fn release( + &self, + fetch: &impl Fetch, + version: &Version, + target: Target, + ) -> Result { + let triple = triple(target); + github_release( + fetch, + RELEASES, + &format!("rust-v{version}"), + &format!("codex-{triple}.tar.gz"), + Packaging::TarGz { + member: format!("codex-{triple}"), + }, + ) + .await + } +} + +impl Configure for Codex { + fn configure( + &self, + _version: &Version, + settings: &Settings, + home: &Path, + ) -> Result { + if settings.wire != Wire::Responses { + return Err(Error::UnsupportedWire { + agent: "codex", + wire: settings.wire, + }); + } + let config = format!( + "model = {model}\nmodel_provider = \"litellm\"\n\n[model_providers.litellm]\nname = \"LiteLLM\"\nbase_url = {base_url}\nenv_key = \"LITELLM_API_KEY\"\nwire_api = \"responses\"\n", + model = quoted(&settings.model), + base_url = quoted(&v1(settings)), + ); + Ok(LaunchSpec { + env: env([ + ("HOME", path_string(home)), + ("CODEX_HOME", path_string(&home.join(".codex"))), + ("LITELLM_API_KEY", settings.api_key.clone()), + ]), + files: BTreeMap::from([(PathBuf::from(".codex/config.toml"), config)]), + }) + } +} + +#[derive(Deserialize)] +enum EventKind { + #[serde(rename = "item.completed")] + ItemCompleted, + #[serde(rename = "turn.completed")] + TurnCompleted, + #[serde(rename = "turn.failed")] + TurnFailed, + #[serde(other)] + Other, +} + +#[derive(Deserialize)] +struct Event { + #[serde(rename = "type")] + kind: EventKind, + item: Option, + usage: Option, + error: Option, +} + +#[derive(Deserialize)] +struct Item { + #[serde(rename = "type")] + kind: String, + text: Option, +} + +#[derive(Deserialize)] +struct TokenUsage { + input_tokens: u64, + output_tokens: u64, +} + +#[derive(Deserialize)] +struct Failure { + message: String, +} + +const NON_TOOL_ITEMS: [&str; 3] = ["agent_message", "reasoning", "error"]; + +impl Drive for Codex { + fn args(&self, _version: &Version, _settings: &Settings, prompt: &Prompt) -> Vec { + let sandbox = ["--sandbox", "workspace-write"]; + ["exec", "--json", "--skip-git-repo-check"] + .into_iter() + .chain(sandbox.into_iter().filter(|_| prompt.allow_tools)) + .chain([prompt.text.as_str()]) + .map(str::to_owned) + .collect() + } + + fn parse(&self, _version: &Version, stdout: &str) -> Outcome { + let events: Vec = json_lines(stdout).collect(); + let items: Vec<&Item> = events + .iter() + .filter(|event| matches!(event.kind, EventKind::ItemCompleted)) + .filter_map(|event| event.item.as_ref()) + .collect(); + Outcome { + text: items + .iter() + .rev() + .find(|item| item.kind == "agent_message") + .and_then(|item| item.text.clone()) + .unwrap_or_default(), + tool_calls: items + .iter() + .filter(|item| !NON_TOOL_ITEMS.contains(&item.kind.as_str())) + .map(|item| item.kind.clone()) + .collect(), + usage: events + .iter() + .filter(|event| matches!(event.kind, EventKind::TurnCompleted)) + .filter_map(|event| event.usage.as_ref()) + .map(|usage| Usage { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + }) + .fold(Usage::default(), |total, turn| total + turn), + errors: events + .iter() + .filter(|event| matches!(event.kind, EventKind::TurnFailed)) + .filter_map(|event| event.error.as_ref()) + .map(|failure| failure.message.clone()) + .collect(), + exit_code: None, + } + } +} diff --git a/litellm-rust/crates/testkit/src/agent/configure.rs b/litellm-rust/crates/testkit/src/agent/configure.rs new file mode 100644 index 00000000000..a095aacc92f --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/configure.rs @@ -0,0 +1,69 @@ +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use semver::Version; + +use crate::Error; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Wire { + ChatCompletions, + Messages, + Responses, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Settings { + pub base_url: String, + pub api_key: String, + pub model: String, + pub wire: Wire, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LaunchSpec { + pub env: BTreeMap, + pub files: BTreeMap, +} + +impl LaunchSpec { + pub fn write_files(&self, home: &Path) -> std::io::Result<()> { + self.files.iter().try_for_each(|(relative, contents)| { + let path = home.join(relative); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, contents) + }) + } +} + +pub trait Configure { + fn configure( + &self, + version: &Version, + settings: &Settings, + home: &Path, + ) -> Result; +} + +pub(crate) fn env( + pairs: impl IntoIterator, +) -> BTreeMap { + pairs + .into_iter() + .map(|(key, value)| (key.to_owned(), value)) + .collect() +} + +pub(crate) fn path_string(path: &Path) -> String { + path.to_string_lossy().into_owned() +} + +pub(crate) fn quoted(value: &str) -> String { + serde_json::Value::from(value).to_string() +} + +pub(crate) fn v1(settings: &Settings) -> String { + format!("{}/v1", settings.base_url.trim_end_matches('/')) +} diff --git a/litellm-rust/crates/testkit/src/agent/drive.rs b/litellm-rust/crates/testkit/src/agent/drive.rs new file mode 100644 index 00000000000..2c238843ed7 --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/drive.rs @@ -0,0 +1,57 @@ +use std::ops::Add; + +use semver::Version; + +use crate::Settings; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Prompt { + pub text: String, + pub allow_tools: bool, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Usage { + pub input_tokens: u64, + pub output_tokens: u64, +} + +impl Add for Usage { + type Output = Self; + + fn add(self, other: Self) -> Self { + Self { + input_tokens: self.input_tokens + other.input_tokens, + output_tokens: self.output_tokens + other.output_tokens, + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Outcome { + pub text: String, + pub tool_calls: Vec, + pub usage: Usage, + pub errors: Vec, + pub exit_code: Option, +} + +impl Outcome { + pub fn succeeded(&self) -> bool { + self.exit_code == Some(0) && self.errors.is_empty() + } +} + +pub trait Drive { + fn args(&self, version: &Version, settings: &Settings, prompt: &Prompt) -> Vec; + + fn parse(&self, version: &Version, stdout: &str) -> Outcome; +} + +pub(crate) fn json_lines<'a, T: serde::de::DeserializeOwned + 'a>( + stdout: &'a str, +) -> impl Iterator + 'a { + stdout + .lines() + .filter_map(|line| serde_json::from_str(line).ok()) +} diff --git a/litellm-rust/crates/testkit/src/agent/install.rs b/litellm-rust/crates/testkit/src/agent/install.rs new file mode 100644 index 00000000000..3f1a0f8950d --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/install.rs @@ -0,0 +1,17 @@ +use std::future::Future; + +use semver::Version; + +use crate::install::Release; +use crate::{Error, Fetch, Target}; + +pub trait Install: Sync { + fn binary(&self) -> &'static str; + + fn release( + &self, + fetch: &impl Fetch, + version: &Version, + target: Target, + ) -> impl Future> + Send; +} diff --git a/litellm-rust/crates/testkit/src/agent/mod.rs b/litellm-rust/crates/testkit/src/agent/mod.rs new file mode 100644 index 00000000000..03a70583f66 --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/mod.rs @@ -0,0 +1,20 @@ +mod claude; +mod codex; +mod configure; +mod drive; +mod install; +mod opencode; + +pub use claude::ClaudeCode; +pub use codex::Codex; +pub use configure::{Configure, LaunchSpec, Settings, Wire}; +pub use drive::{Drive, Outcome, Prompt, Usage}; +pub use install::Install; +pub use opencode::Opencode; + +pub(crate) use configure::{env, path_string, quoted, v1}; +pub(crate) use drive::json_lines; + +pub trait Agent: Install + Configure + Drive {} + +impl Agent for T {} diff --git a/litellm-rust/crates/testkit/src/agent/opencode.rs b/litellm-rust/crates/testkit/src/agent/opencode.rs new file mode 100644 index 00000000000..a1b2fa01eef --- /dev/null +++ b/litellm-rust/crates/testkit/src/agent/opencode.rs @@ -0,0 +1,187 @@ +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use semver::Version; +use serde::Deserialize; + +use super::{ + Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines, + path_string, v1, +}; +use crate::install::release::github_release; +use crate::install::{Packaging, Release}; +use crate::target::Os; +use crate::{Error, Fetch, Target}; + +const RELEASES: &str = "https://api.github.com/repos/sst/opencode/releases/tags"; + +pub struct Opencode; + +impl Install for Opencode { + fn binary(&self) -> &'static str { + "opencode" + } + + async fn release( + &self, + fetch: &impl Fetch, + version: &Version, + target: Target, + ) -> Result { + let stem = format!( + "opencode-{}-{}{}", + target.os_name(), + target.arch_name(), + target.musl_suffix() + ); + let member = "opencode".to_owned(); + let (asset, packaging) = match target.os { + Os::Macos => (format!("{stem}.zip"), Packaging::Zip { member }), + Os::Linux => (format!("{stem}.tar.gz"), Packaging::TarGz { member }), + }; + github_release(fetch, RELEASES, &format!("v{version}"), &asset, packaging).await + } +} + +impl Configure for Opencode { + fn configure( + &self, + _version: &Version, + settings: &Settings, + home: &Path, + ) -> Result { + let npm = match settings.wire { + Wire::ChatCompletions => "@ai-sdk/openai-compatible", + Wire::Responses => "@ai-sdk/openai", + Wire::Messages => "@ai-sdk/anthropic", + }; + let config = serde_json::json!({ + "$schema": "https://opencode.ai/config.json", + "model": format!("litellm/{}", settings.model), + "provider": { + "litellm": { + "npm": npm, + "name": "LiteLLM", + "options": { "baseURL": v1(settings), "apiKey": settings.api_key }, + "models": { settings.model.clone(): { "name": settings.model } }, + } + }, + }); + Ok(LaunchSpec { + env: env([ + ("HOME", path_string(home)), + ("XDG_CONFIG_HOME", path_string(&home.join(".config"))), + ("XDG_DATA_HOME", path_string(&home.join(".local/share"))), + ("OPENCODE_DISABLE_AUTOUPDATE", "true".to_owned()), + ]), + files: BTreeMap::from([( + PathBuf::from(".config/opencode/opencode.json"), + config.to_string(), + )]), + }) + } +} + +#[derive(Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum Event { + Text { + part: TextPart, + }, + ToolUse { + part: ToolPart, + }, + StepFinish { + part: StepFinish, + }, + Error { + error: Failure, + }, + #[serde(other)] + Other, +} + +#[derive(Deserialize)] +struct TextPart { + text: String, +} + +#[derive(Deserialize)] +struct ToolPart { + tool: String, +} + +#[derive(Deserialize)] +struct StepFinish { + tokens: Tokens, +} + +#[derive(Deserialize)] +struct Tokens { + input: u64, + output: u64, +} + +#[derive(Deserialize)] +struct Failure { + name: String, + data: Option, +} + +#[derive(Deserialize)] +struct FailureData { + message: Option, +} + +impl Drive for Opencode { + fn args(&self, _version: &Version, _settings: &Settings, prompt: &Prompt) -> Vec { + ["run", "--format", "json", &prompt.text] + .map(str::to_owned) + .to_vec() + } + + fn parse(&self, _version: &Version, stdout: &str) -> Outcome { + let events: Vec = json_lines(stdout).collect(); + Outcome { + text: events + .iter() + .rev() + .find_map(|event| match event { + Event::Text { part } => Some(part.text.clone()), + _ => None, + }) + .unwrap_or_default(), + tool_calls: events + .iter() + .filter_map(|event| match event { + Event::ToolUse { part } => Some(part.tool.clone()), + _ => None, + }) + .collect(), + usage: events + .iter() + .filter_map(|event| match event { + Event::StepFinish { part } => Some(Usage { + input_tokens: part.tokens.input, + output_tokens: part.tokens.output, + }), + _ => None, + }) + .fold(Usage::default(), |total, step| total + step), + errors: events + .iter() + .filter_map(|event| match event { + Event::Error { error } => Some( + error + .data + .as_ref() + .and_then(|data| data.message.clone()) + .unwrap_or_else(|| error.name.clone()), + ), + _ => None, + }) + .collect(), + exit_code: None, + } + } +} diff --git a/litellm-rust/crates/testkit/src/error.rs b/litellm-rust/crates/testkit/src/error.rs new file mode 100644 index 00000000000..03520d827f1 --- /dev/null +++ b/litellm-rust/crates/testkit/src/error.rs @@ -0,0 +1,56 @@ +use std::io; +use std::path::PathBuf; + +use thiserror::Error; + +use crate::Wire; + +#[derive(Debug, Error)] +pub enum Error { + #[error("unsupported target {0}")] + UnsupportedTarget(String), + #[error("{0} is not a plain x.y.z release version")] + InvalidVersion(String), + #[error("request to {url} failed")] + Request { + url: String, + #[source] + source: reqwest::Error, + }, + #[error("{url} answered with status {status}")] + Status { url: String, status: u16 }, + #[error("release metadata at {url} is malformed")] + Metadata { + url: String, + #[source] + source: serde_json::Error, + }, + #[error("release has no asset named {0}")] + AssetNotFound(String), + #[error("release publishes no sha256 for {0}")] + MissingChecksum(String), + #[error("sha256 mismatch for {asset}: expected {expected}, got {actual}")] + ChecksumMismatch { + asset: String, + expected: String, + actual: String, + }, + #[error("archive does not contain {0}")] + ArchiveMemberNotFound(String), + #[error("archive is unreadable")] + Archive(#[source] io::Error), + #[error("zip archive is unreadable")] + Zip(#[from] zip::result::ZipError), + #[error("{binary} reports version '{reported}', expected {expected}")] + VersionMismatch { + binary: PathBuf, + expected: String, + reported: String, + }, + #[error("{agent} cannot talk to the gateway over {wire:?}")] + UnsupportedWire { agent: &'static str, wire: Wire }, + #[error("agent did not finish within {0:?}")] + Timeout(std::time::Duration), + #[error("io failure")] + Io(#[from] io::Error), +} diff --git a/litellm-rust/crates/testkit/src/install/archive.rs b/litellm-rust/crates/testkit/src/install/archive.rs new file mode 100644 index 00000000000..c8d08f66f47 --- /dev/null +++ b/litellm-rust/crates/testkit/src/install/archive.rs @@ -0,0 +1,52 @@ +use std::io::{Cursor, Read}; + +use flate2::read::GzDecoder; +use sha2::{Digest, Sha256}; + +use super::release::Packaging; +use crate::Error; + +pub(crate) fn verify_sha256(asset: &str, expected: &str, bytes: &[u8]) -> Result<(), Error> { + let actual = format!("{:x}", Sha256::digest(bytes)); + if actual.eq_ignore_ascii_case(expected) { + return Ok(()); + } + Err(Error::ChecksumMismatch { + asset: asset.to_owned(), + expected: expected.to_owned(), + actual, + }) +} + +pub(crate) fn extract_binary(packaging: &Packaging, bytes: &[u8]) -> Result, Error> { + match packaging { + Packaging::Bare => Ok(bytes.to_vec()), + Packaging::TarGz { member } => extract_tar_gz(member, bytes), + Packaging::Zip { member } => extract_zip(member, bytes), + } +} + +fn extract_tar_gz(member: &str, bytes: &[u8]) -> Result, Error> { + let mut archive = tar::Archive::new(GzDecoder::new(bytes)); + for entry in archive.entries().map_err(Error::Archive)? { + let mut entry = entry.map_err(Error::Archive)?; + let path = entry.path().map_err(Error::Archive)?; + if path.file_name().is_some_and(|name| name == member) { + let mut binary = Vec::new(); + entry.read_to_end(&mut binary).map_err(Error::Archive)?; + return Ok(binary); + } + } + Err(Error::ArchiveMemberNotFound(member.to_owned())) +} + +fn extract_zip(member: &str, bytes: &[u8]) -> Result, Error> { + let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?; + let mut file = archive.by_name(member).map_err(|error| match error { + zip::result::ZipError::FileNotFound => Error::ArchiveMemberNotFound(member.to_owned()), + other => Error::Zip(other), + })?; + let mut binary = Vec::new(); + file.read_to_end(&mut binary).map_err(Error::Archive)?; + Ok(binary) +} diff --git a/litellm-rust/crates/testkit/src/install/fetch.rs b/litellm-rust/crates/testkit/src/install/fetch.rs new file mode 100644 index 00000000000..73008f7a0da --- /dev/null +++ b/litellm-rust/crates/testkit/src/install/fetch.rs @@ -0,0 +1,55 @@ +use std::future::Future; + +use crate::Error; + +pub trait Fetch: Sync { + fn get(&self, url: &str) -> impl Future, Error>> + Send; +} + +pub struct HttpFetch { + client: reqwest::Client, + github_token: Option, +} + +impl HttpFetch { + pub fn new(github_token: Option) -> Self { + Self { + client: reqwest::Client::new(), + github_token, + } + } + + pub fn from_env() -> Self { + Self::new(std::env::var("GITHUB_TOKEN").ok()) + } +} + +impl Fetch for HttpFetch { + async fn get(&self, url: &str) -> Result, Error> { + let request = self + .client + .get(url) + .header("user-agent", "litellm-testkit") + .header("accept", "application/json, application/octet-stream"); + let request = match ( + &self.github_token, + url.starts_with("https://api.github.com/"), + ) { + (Some(token), true) => request.bearer_auth(token), + _ => request, + }; + let request_error = |source| Error::Request { + url: url.to_owned(), + source, + }; + let response = request.send().await.map_err(request_error)?; + let status = response.status(); + if !status.is_success() { + return Err(Error::Status { + url: url.to_owned(), + status: status.as_u16(), + }); + } + Ok(response.bytes().await.map_err(request_error)?.to_vec()) + } +} diff --git a/litellm-rust/crates/testkit/src/install/mod.rs b/litellm-rust/crates/testkit/src/install/mod.rs new file mode 100644 index 00000000000..1104bcec102 --- /dev/null +++ b/litellm-rust/crates/testkit/src/install/mod.rs @@ -0,0 +1,118 @@ +mod archive; +mod fetch; +pub(crate) mod release; + +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::sync::atomic::{AtomicU64, Ordering}; + +use semver::Version; +use tokio::fs; +use tokio::process::Command; + +use crate::{Error, Install, Target}; +use archive::{extract_binary, verify_sha256}; + +static STAGING_COUNTER: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Installed { + pub version: Version, + pub binary: PathBuf, +} + +pub struct Installer { + fetch: F, + cache_root: PathBuf, + target: Target, +} + +impl Installer { + pub fn new(fetch: F, cache_root: impl Into, target: Target) -> Self { + Self { + fetch, + cache_root: cache_root.into(), + target, + } + } + + pub async fn install( + &self, + agent: &impl Install, + version: &Version, + ) -> Result { + validate_release(version)?; + let dir = self + .cache_root + .join(agent.binary()) + .join(version.to_string()); + let binary = dir.join(agent.binary()); + let installed = Installed { + version: version.clone(), + binary: binary.clone(), + }; + if fs::try_exists(&binary).await? && probe_version(&binary, version).await.is_ok() { + return Ok(installed); + } + + let release = agent.release(&self.fetch, version, self.target).await?; + let archive = self.fetch.get(&release.url).await?; + verify_sha256(&release.asset, &release.sha256, &archive)?; + let contents = extract_binary(&release.packaging, &archive)?; + + fs::create_dir_all(&dir).await?; + let staging = dir.join(format!( + ".{}.{}.{}.partial", + agent.binary(), + std::process::id(), + STAGING_COUNTER.fetch_add(1, Ordering::Relaxed) + )); + fs::write(&staging, contents).await?; + fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)).await?; + fs::rename(&staging, &binary).await?; + + match probe_version(&binary, version).await { + Ok(()) => Ok(installed), + Err(error) => { + fs::remove_file(&binary).await?; + Err(error) + } + } + } +} + +fn validate_release(version: &Version) -> Result<(), Error> { + if version.pre.is_empty() && version.build.is_empty() { + return Ok(()); + } + Err(Error::InvalidVersion(version.to_string())) +} + +async fn probe_version(binary: &Path, expected: &Version) -> Result<(), Error> { + let home = std::env::temp_dir(); + let output = Command::new(binary) + .arg("--version") + .env_clear() + .env("HOME", home) + .env("DISABLE_AUTOUPDATER", "1") + .stdin(Stdio::null()) + .output() + .await?; + let stdout = String::from_utf8_lossy(&output.stdout); + if stdout + .split_whitespace() + .filter_map(|token| Version::parse(token).ok()) + .any(|reported| &reported == expected) + { + return Ok(()); + } + Err(Error::VersionMismatch { + binary: binary.to_owned(), + expected: expected.to_string(), + reported: stdout.trim().to_owned(), + }) +} + +pub use fetch::{Fetch, HttpFetch}; +pub use release::{Packaging, Release}; diff --git a/litellm-rust/crates/testkit/src/install/release.rs b/litellm-rust/crates/testkit/src/install/release.rs new file mode 100644 index 00000000000..a21b9a14f1b --- /dev/null +++ b/litellm-rust/crates/testkit/src/install/release.rs @@ -0,0 +1,65 @@ +use serde::Deserialize; + +use crate::{Error, Fetch}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Packaging { + Bare, + TarGz { member: String }, + Zip { member: String }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Release { + pub asset: String, + pub url: String, + pub sha256: String, + pub packaging: Packaging, +} + +#[derive(Deserialize)] +struct GithubRelease { + assets: Vec, +} + +#[derive(Deserialize)] +struct GithubAsset { + name: String, + digest: Option, + browser_download_url: String, +} + +pub(crate) async fn github_release( + fetch: &impl Fetch, + releases_url: &str, + tag: &str, + asset_name: &str, + packaging: Packaging, +) -> Result { + let url = format!("{releases_url}/{tag}"); + let release: GithubRelease = parse(&url, &fetch.get(&url).await?)?; + let asset = release + .assets + .into_iter() + .find(|asset| asset.name == asset_name) + .ok_or_else(|| Error::AssetNotFound(asset_name.to_owned()))?; + let sha256 = asset + .digest + .as_deref() + .and_then(|digest| digest.strip_prefix("sha256:")) + .ok_or_else(|| Error::MissingChecksum(asset_name.to_owned()))? + .to_owned(); + Ok(Release { + asset: asset.name, + url: asset.browser_download_url, + sha256, + packaging, + }) +} + +pub(crate) fn parse Deserialize<'de>>(url: &str, body: &[u8]) -> Result { + serde_json::from_slice(body).map_err(|source| Error::Metadata { + url: url.to_owned(), + source, + }) +} diff --git a/litellm-rust/crates/testkit/src/lib.rs b/litellm-rust/crates/testkit/src/lib.rs new file mode 100644 index 00000000000..9ea6123a176 --- /dev/null +++ b/litellm-rust/crates/testkit/src/lib.rs @@ -0,0 +1,15 @@ +mod agent; +mod error; +mod install; +mod session; +mod target; + +pub use agent::{ + Agent, ClaudeCode, Codex, Configure, Drive, Install, LaunchSpec, Opencode, Outcome, Prompt, + Settings, Usage, Wire, +}; +pub use error::Error; +pub use install::{Fetch, HttpFetch, Installed, Installer, Packaging, Release}; +pub use semver::Version; +pub use session::Session; +pub use target::{Arch, Os, Target}; diff --git a/litellm-rust/crates/testkit/src/session.rs b/litellm-rust/crates/testkit/src/session.rs new file mode 100644 index 00000000000..6b06e756cca --- /dev/null +++ b/litellm-rust/crates/testkit/src/session.rs @@ -0,0 +1,76 @@ +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::process::Stdio; +use std::time::Duration; + +use semver::Version; +use tokio::process::Command; +use tokio::time::timeout; + +use crate::{Configure, Drive, Error, Installed, Outcome, Prompt, Settings}; + +const STDERR_LIMIT_CHARS: usize = 2000; + +pub struct Session { + binary: PathBuf, + home: PathBuf, + version: Version, + settings: Settings, + env: BTreeMap, +} + +impl Session { + pub fn prepare( + agent: &impl Configure, + installed: &Installed, + settings: Settings, + home: impl Into, + ) -> Result { + let home = home.into(); + let spec = agent.configure(&installed.version, &settings, &home)?; + spec.write_files(&home)?; + Ok(Self { + binary: installed.binary.clone(), + home, + version: installed.version.clone(), + settings, + env: spec.env, + }) + } + + pub async fn run( + &self, + agent: &impl Drive, + prompt: &Prompt, + limit: Duration, + ) -> Result { + let child = Command::new(&self.binary) + .args(agent.args(&self.version, &self.settings, prompt)) + .env_clear() + .env("PATH", "/usr/bin:/bin") + .envs(&self.env) + .current_dir(&self.home) + .stdin(Stdio::null()) + .kill_on_drop(true) + .output(); + let output = timeout(limit, child) + .await + .map_err(|_| Error::Timeout(limit))??; + let parsed = agent.parse(&self.version, &String::from_utf8_lossy(&output.stdout)); + let failed_silently = !output.status.success() && parsed.errors.is_empty(); + Ok(Outcome { + errors: if failed_silently { + vec![ + String::from_utf8_lossy(&output.stderr) + .chars() + .take(STDERR_LIMIT_CHARS) + .collect(), + ] + } else { + parsed.errors + }, + exit_code: output.status.code(), + ..parsed + }) + } +} diff --git a/litellm-rust/crates/testkit/src/target.rs b/litellm-rust/crates/testkit/src/target.rs new file mode 100644 index 00000000000..a9d4b012d52 --- /dev/null +++ b/litellm-rust/crates/testkit/src/target.rs @@ -0,0 +1,69 @@ +use target_lexicon::{Architecture, Environment, OperatingSystem, Triple}; + +use crate::Error; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Os { + Macos, + Linux, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Arch { + Aarch64, + X86_64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Target { + pub os: Os, + pub arch: Arch, + pub musl: bool, +} + +impl Target { + pub fn host() -> Result { + Self::try_from(&Triple::host()) + } + + pub(crate) const fn os_name(self) -> &'static str { + match self.os { + Os::Macos => "darwin", + Os::Linux => "linux", + } + } + + pub(crate) const fn arch_name(self) -> &'static str { + match self.arch { + Arch::Aarch64 => "arm64", + Arch::X86_64 => "x64", + } + } + + pub(crate) const fn musl_suffix(self) -> &'static str { + if self.musl { "-musl" } else { "" } + } +} + +impl TryFrom<&Triple> for Target { + type Error = Error; + + fn try_from(triple: &Triple) -> Result { + let unsupported = || Error::UnsupportedTarget(triple.to_string()); + let os = match triple.operating_system { + OperatingSystem::Darwin(_) | OperatingSystem::MacOSX(_) => Os::Macos, + OperatingSystem::Linux => Os::Linux, + _ => return Err(unsupported()), + }; + let arch = match triple.architecture { + Architecture::Aarch64(_) => Arch::Aarch64, + Architecture::X86_64 => Arch::X86_64, + _ => return Err(unsupported()), + }; + Ok(Self { + os, + arch, + musl: triple.environment == Environment::Musl, + }) + } +} diff --git a/litellm-rust/crates/testkit/tests/configure.rs b/litellm-rust/crates/testkit/tests/configure.rs new file mode 100644 index 00000000000..ca3587c3474 --- /dev/null +++ b/litellm-rust/crates/testkit/tests/configure.rs @@ -0,0 +1,133 @@ +use std::path::Path; + +use litellm_testkit::{ClaudeCode, Codex, Configure, Error, Opencode, Settings, Version, Wire}; +use rstest::rstest; + +fn settings(wire: Wire) -> Settings { + Settings { + base_url: "http://localhost:4000/".to_owned(), + api_key: "sk-test \"quoted\"".to_owned(), + model: "some-model".to_owned(), + wire, + } +} + +fn version() -> Version { + Version::new(1, 2, 3) +} + +#[rstest] +#[case(&ClaudeCode, Wire::Messages)] +#[case(&Codex, Wire::Responses)] +#[case(&Opencode, Wire::ChatCompletions)] +fn every_agent_runs_inside_the_given_home(#[case] agent: &impl Configure, #[case] wire: Wire) { + let home = Path::new("/scratch/home"); + + let spec = agent.configure(&version(), &settings(wire), home).unwrap(); + + assert_eq!(spec.env["HOME"], "/scratch/home"); + assert!( + spec.env + .iter() + .filter(|(key, _)| key.ends_with("_HOME") || key.as_str() == "CLAUDE_CONFIG_DIR") + .all(|(_, value)| value.starts_with("/scratch/home")) + ); + assert!(spec.files.keys().all(|path| path.is_relative())); +} + +#[rstest] +#[case::claude_code(&ClaudeCode, &[Wire::ChatCompletions, Wire::Responses])] +#[case::codex(&Codex, &[Wire::ChatCompletions, Wire::Messages])] +fn wires_an_agent_cannot_speak_are_refused( + #[case] agent: &impl Configure, + #[case] refused: &[Wire], +) { + refused.iter().for_each(|wire| { + let result = agent.configure(&version(), &settings(*wire), Path::new("/h")); + + assert!(matches!(result, Err(Error::UnsupportedWire { wire: got, .. }) if got == *wire)); + }); +} + +#[test] +fn claude_code_points_at_the_gateway_root_with_the_key_and_model() { + let spec = ClaudeCode + .configure(&version(), &settings(Wire::Messages), Path::new("/h")) + .unwrap(); + + assert_eq!(spec.env["ANTHROPIC_BASE_URL"], "http://localhost:4000/"); + assert_eq!(spec.env["ANTHROPIC_AUTH_TOKEN"], "sk-test \"quoted\""); + assert_eq!(spec.env["ANTHROPIC_MODEL"], "some-model"); +} + +#[test] +fn codex_config_is_valid_toml_routing_the_responses_api_to_the_gateway() { + let dir = tempfile::tempdir().unwrap(); + let spec = Codex + .configure(&version(), &settings(Wire::Responses), dir.path()) + .unwrap(); + spec.write_files(dir.path()).unwrap(); + + let config: toml::Table = + toml::from_str(&std::fs::read_to_string(dir.path().join(".codex/config.toml")).unwrap()) + .unwrap(); + let provider = &config["model_providers"]["litellm"]; + + assert_eq!(config["model"].as_str(), Some("some-model")); + assert_eq!(config["model_provider"].as_str(), Some("litellm")); + assert_eq!( + provider["base_url"].as_str(), + Some("http://localhost:4000/v1") + ); + assert_eq!(provider["wire_api"].as_str(), Some("responses")); + let key_var = provider["env_key"].as_str().unwrap(); + assert_eq!(spec.env[key_var], "sk-test \"quoted\""); +} + +#[rstest] +#[case(Wire::ChatCompletions)] +#[case(Wire::Responses)] +#[case(Wire::Messages)] +fn opencode_config_is_valid_json_registering_the_gateway_model(#[case] wire: Wire) { + let dir = tempfile::tempdir().unwrap(); + let spec = Opencode + .configure(&version(), &settings(wire), dir.path()) + .unwrap(); + spec.write_files(dir.path()).unwrap(); + + let config: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(dir.path().join(".config/opencode/opencode.json")).unwrap(), + ) + .unwrap(); + let provider = &config["provider"]["litellm"]; + + assert_eq!(config["model"], "litellm/some-model"); + assert_eq!(provider["options"]["baseURL"], "http://localhost:4000/v1"); + assert_eq!(provider["options"]["apiKey"], "sk-test \"quoted\""); + assert!(provider["models"]["some-model"].is_object()); +} + +#[test] +fn opencode_uses_a_different_provider_package_for_every_wire() { + let package = |wire| { + let dir = tempfile::tempdir().unwrap(); + let spec = Opencode + .configure(&version(), &settings(wire), dir.path()) + .unwrap(); + let config: serde_json::Value = + serde_json::from_str(spec.files.values().next().unwrap()).unwrap(); + config["provider"]["litellm"]["npm"] + .as_str() + .unwrap() + .to_owned() + }; + let packages = [Wire::ChatCompletions, Wire::Responses, Wire::Messages].map(package); + + assert_eq!( + packages + .iter() + .collect::>() + .len(), + packages.len() + ); +} diff --git a/litellm-rust/crates/testkit/tests/install.rs b/litellm-rust/crates/testkit/tests/install.rs new file mode 100644 index 00000000000..7edaf27de6c --- /dev/null +++ b/litellm-rust/crates/testkit/tests/install.rs @@ -0,0 +1,262 @@ +mod support; + +use std::str::FromStr; + +use litellm_testkit::{ClaudeCode, Codex, Error, Installer, Opencode, Target, Version}; +use rstest::rstest; +use serde_json::json; +use support::{FakeFetch, script_printing, sha256, tar_gz, zip_archive}; +use target_lexicon::Triple; + +fn target(triple: &str) -> Target { + Target::try_from(&Triple::from_str(triple).unwrap()).unwrap() +} + +fn linux() -> Target { + target("x86_64-unknown-linux-gnu") +} +fn version() -> Version { + Version::new(9, 8, 7) +} + +fn github_release(asset: &str, download_url: &str, digest: Option) -> Vec { + json!({ + "assets": [ + { "name": "unrelated.txt", "digest": "sha256:00", "browser_download_url": "https://example.test/unrelated" }, + { "name": asset, "digest": digest, "browser_download_url": download_url }, + ] + }) + .to_string() + .into_bytes() +} + +fn claude_routes(binary: &[u8], checksum: &str) -> Vec<(String, Vec)> { + let base = "https://downloads.claude.ai/claude-code-releases/9.8.7"; + let manifest = json!({ "platforms": { "linux-x64": { "checksum": checksum } } }); + vec![ + ( + format!("{base}/manifest.json"), + manifest.to_string().into_bytes(), + ), + (format!("{base}/linux-x64/claude"), binary.to_vec()), + ] +} + +fn codex_routes(archive: Vec, digest: Option) -> Vec<(String, Vec)> { + let release = github_release( + "codex-x86_64-unknown-linux-musl.tar.gz", + "https://example.test/codex.tar.gz", + digest, + ); + vec![ + ( + "https://api.github.com/repos/openai/codex/releases/tags/rust-v9.8.7".to_owned(), + release, + ), + ("https://example.test/codex.tar.gz".to_owned(), archive), + ] +} + +#[tokio::test] +async fn claude_bare_binary_is_installed_and_runnable() { + let binary = script_printing("9.8.7 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(&binary))); + let cache = tempfile::tempdir().unwrap(); + + let installed = Installer::new(&fetch, cache.path(), linux()) + .install(&ClaudeCode, &version()) + .await + .unwrap(); + + assert_eq!(installed.binary, cache.path().join("claude/9.8.7/claude")); + assert_eq!(std::fs::read(&installed.binary).unwrap(), binary); +} + +#[tokio::test] +async fn codex_binary_is_extracted_from_the_tarball_under_its_own_name() { + let binary = script_printing("codex-cli 9.8.7"); + let archive = tar_gz("codex-x86_64-unknown-linux-musl", &binary); + let fetch = FakeFetch::new(codex_routes( + archive.clone(), + Some(format!("sha256:{}", sha256(&archive))), + )); + let cache = tempfile::tempdir().unwrap(); + + let installed = Installer::new(&fetch, cache.path(), linux()) + .install(&Codex, &version()) + .await + .unwrap(); + + assert_eq!(std::fs::read(&installed.binary).unwrap(), binary); + assert_eq!(installed.binary, cache.path().join("codex/9.8.7/codex")); +} + +#[tokio::test] +async fn opencode_binary_is_extracted_from_the_darwin_zip() { + let binary = script_printing("9.8.7"); + let archive = zip_archive("opencode", &binary); + let release = github_release( + "opencode-darwin-arm64.zip", + "https://example.test/opencode.zip", + Some(format!("sha256:{}", sha256(&archive))), + ); + let fetch = FakeFetch::new([ + ( + "https://api.github.com/repos/sst/opencode/releases/tags/v9.8.7".to_owned(), + release, + ), + ("https://example.test/opencode.zip".to_owned(), archive), + ]); + let cache = tempfile::tempdir().unwrap(); + + let installed = Installer::new(&fetch, cache.path(), target("aarch64-apple-darwin")) + .install(&Opencode, &version()) + .await + .unwrap(); + + assert_eq!(std::fs::read(&installed.binary).unwrap(), binary); +} + +#[tokio::test] +async fn tampered_download_is_rejected_and_nothing_is_left_behind() { + let binary = script_printing("9.8.7 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(b"what the vendor signed"))); + let cache = tempfile::tempdir().unwrap(); + + let result = Installer::new(&fetch, cache.path(), linux()) + .install(&ClaudeCode, &version()) + .await; + + assert!(matches!(result, Err(Error::ChecksumMismatch { .. }))); + assert!(!cache.path().join("claude/9.8.7").exists()); +} + +#[tokio::test] +async fn github_asset_without_a_digest_is_refused() { + let archive = tar_gz( + "codex-x86_64-unknown-linux-musl", + &script_printing("codex-cli 9.8.7"), + ); + let fetch = FakeFetch::new(codex_routes(archive, None)); + let cache = tempfile::tempdir().unwrap(); + + let result = Installer::new(&fetch, cache.path(), linux()) + .install(&Codex, &version()) + .await; + + assert!(matches!(result, Err(Error::MissingChecksum(_)))); +} + +#[tokio::test] +async fn binary_reporting_a_different_version_is_removed() { + let binary = script_printing("1.0.0 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(&binary))); + let cache = tempfile::tempdir().unwrap(); + + let result = Installer::new(&fetch, cache.path(), linux()) + .install(&ClaudeCode, &version()) + .await; + + assert!(matches!(result, Err(Error::VersionMismatch { .. }))); + assert!(!cache.path().join("claude/9.8.7/claude").exists()); +} + +#[tokio::test] +async fn second_install_reuses_the_cached_binary_without_downloading() { + let binary = script_printing("9.8.7 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(&binary))); + let cache = tempfile::tempdir().unwrap(); + let installer = Installer::new(&fetch, cache.path(), linux()); + + let first = installer.install(&ClaudeCode, &version()).await.unwrap(); + let calls_after_first = fetch.calls(); + let second = installer.install(&ClaudeCode, &version()).await.unwrap(); + + assert_eq!(first, second); + assert_eq!(fetch.calls(), calls_after_first); +} + +#[tokio::test] +async fn corrupted_cache_entry_is_replaced_by_a_fresh_download() { + let binary = script_printing("9.8.7 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(&binary))); + let cache = tempfile::tempdir().unwrap(); + let installer = Installer::new(&fetch, cache.path(), linux()); + let installed = installer.install(&ClaudeCode, &version()).await.unwrap(); + std::fs::write(&installed.binary, script_printing("0.0.1")).unwrap(); + + installer.install(&ClaudeCode, &version()).await.unwrap(); + + assert_eq!(std::fs::read(&installed.binary).unwrap(), binary); +} + +#[rstest] +#[case("9.8.7-beta.1")] +#[case("9.8.7+build.5")] +#[tokio::test] +async fn pre_releases_never_reach_the_network_or_the_filesystem(#[case] version: &str) { + let fetch = FakeFetch::new([]); + let cache = tempfile::tempdir().unwrap(); + + let result = Installer::new(&fetch, cache.path(), linux()) + .install(&ClaudeCode, &Version::parse(version).unwrap()) + .await; + + assert!(matches!(result, Err(Error::InvalidVersion(_)))); + assert_eq!(fetch.calls(), 0); + assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 0); +} + +#[tokio::test] +async fn musl_linux_picks_the_musl_claude_build() { + let binary = script_printing("9.8.7 (Claude Code)"); + let base = "https://downloads.claude.ai/claude-code-releases/9.8.7"; + let manifest = json!({ "platforms": { + "linux-x64": { "checksum": sha256(b"glibc build") }, + "linux-x64-musl": { "checksum": sha256(&binary) }, + } }); + let fetch = FakeFetch::new([ + ( + format!("{base}/manifest.json"), + manifest.to_string().into_bytes(), + ), + (format!("{base}/linux-x64-musl/claude"), binary.clone()), + ]); + let cache = tempfile::tempdir().unwrap(); + + let installed = Installer::new(&fetch, cache.path(), target("x86_64-unknown-linux-musl")) + .install(&ClaudeCode, &version()) + .await + .unwrap(); + + assert_eq!(std::fs::read(&installed.binary).unwrap(), binary); +} + +#[rstest] +#[case("x86_64-pc-windows-msvc")] +#[case("riscv64gc-unknown-linux-gnu")] +#[case("wasm32-unknown-unknown")] +fn targets_no_agent_ships_for_are_rejected(#[case] triple: &str) { + let result = Target::try_from(&Triple::from_str(triple).unwrap()); + + assert!(matches!(result, Err(Error::UnsupportedTarget(_)))); +} + +#[tokio::test] +async fn concurrent_installs_of_the_same_version_both_succeed() { + let binary = script_printing("9.8.7 (Claude Code)"); + let fetch = FakeFetch::new(claude_routes(&binary, &sha256(&binary))); + let cache = tempfile::tempdir().unwrap(); + let installer = Installer::new(&fetch, cache.path(), linux()); + + let wanted = version(); + let installs = + futures_util::future::join_all((0..8).map(|_| installer.install(&ClaudeCode, &wanted))) + .await; + + assert!(installs.iter().all(Result::is_ok)); + assert_eq!( + std::fs::read(&installs[0].as_ref().unwrap().binary).unwrap(), + binary + ); +} diff --git a/litellm-rust/crates/testkit/tests/live.rs b/litellm-rust/crates/testkit/tests/live.rs new file mode 100644 index 00000000000..b805596879a --- /dev/null +++ b/litellm-rust/crates/testkit/tests/live.rs @@ -0,0 +1,133 @@ +//! Drives the real agents through a real gateway. Run with `cargo test -p litellm-testkit --test live -- --ignored` +//! after exporting `TESTKIT_GATEWAY_URL`, `TESTKIT_GATEWAY_KEY`, one `TESTKIT_MODEL_` per wire +//! (`MESSAGES`, `RESPONSES`, `CHAT_COMPLETIONS`) and one `TESTKIT__VERSION` per agent +//! (`CLAUDE`, `CODEX`, `OPENCODE`). `TESTKIT_CACHE_DIR` and `GITHUB_TOKEN` are optional. + +use std::path::PathBuf; +use std::time::Duration; + +use litellm_testkit::{ + Agent, ClaudeCode, Codex, HttpFetch, Installer, Opencode, Outcome, Prompt, Session, Settings, + Target, Version, Wire, +}; +use rstest::rstest; + +const LIMIT: Duration = Duration::from_secs(180); + +fn required(name: &str) -> String { + std::env::var(name).unwrap_or_else(|_| panic!("{name} must be set to run the live tests")) +} + +fn model_var(wire: Wire) -> &'static str { + match wire { + Wire::Messages => "TESTKIT_MODEL_MESSAGES", + Wire::Responses => "TESTKIT_MODEL_RESPONSES", + Wire::ChatCompletions => "TESTKIT_MODEL_CHAT_COMPLETIONS", + } +} + +async fn drive( + agent: &impl Agent, + version_var: &str, + wire: Wire, + model: Option<&str>, + prompt: Prompt, +) -> Outcome { + let cache = std::env::var("TESTKIT_CACHE_DIR") + .map(PathBuf::from) + .unwrap_or_else(|_| std::env::temp_dir().join("litellm-testkit-cache")); + let installer = Installer::new(HttpFetch::from_env(), cache, Target::host().unwrap()); + let installed = installer + .install(agent, &Version::parse(&required(version_var)).unwrap()) + .await + .unwrap(); + let settings = Settings { + base_url: required("TESTKIT_GATEWAY_URL"), + api_key: required("TESTKIT_GATEWAY_KEY"), + model: model.map_or_else(|| required(model_var(wire)), str::to_owned), + wire, + }; + let home = tempfile::tempdir().unwrap(); + let session = Session::prepare(agent, &installed, settings, home.path()).unwrap(); + session.run(agent, &prompt, LIMIT).await.unwrap() +} + +fn text_prompt() -> Prompt { + Prompt { + text: "Reply with the single word: pong".to_owned(), + allow_tools: false, + } +} + +fn tool_prompt() -> Prompt { + Prompt { + text: "Run the shell command 'echo tool-ok' and reply with exactly its output.".to_owned(), + allow_tools: true, + } +} + +#[rstest] +#[case::claude_messages(&ClaudeCode, "TESTKIT_CLAUDE_VERSION", Wire::Messages)] +#[case::codex_responses(&Codex, "TESTKIT_CODEX_VERSION", Wire::Responses)] +#[case::opencode_chat(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::ChatCompletions)] +#[case::opencode_responses(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Responses)] +#[case::opencode_messages(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Messages)] +#[ignore = "needs a live gateway, see the module docs"] +#[tokio::test] +async fn plain_prompt_gets_an_answer_and_token_usage( + #[case] agent: &impl Agent, + #[case] version_var: &str, + #[case] wire: Wire, +) { + let outcome = drive(agent, version_var, wire, None, text_prompt()).await; + + assert!(outcome.succeeded(), "{outcome:?}"); + assert!(outcome.text.to_lowercase().contains("pong"), "{outcome:?}"); + assert!(outcome.usage.output_tokens > 0, "{outcome:?}"); +} + +#[rstest] +#[case::claude_messages(&ClaudeCode, "TESTKIT_CLAUDE_VERSION", Wire::Messages)] +#[case::codex_responses(&Codex, "TESTKIT_CODEX_VERSION", Wire::Responses)] +#[case::opencode_chat(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::ChatCompletions)] +#[case::opencode_responses(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Responses)] +#[case::opencode_messages(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Messages)] +#[ignore = "needs a live gateway, see the module docs"] +#[tokio::test] +async fn tool_use_is_reported_and_its_result_reaches_the_answer( + #[case] agent: &impl Agent, + #[case] version_var: &str, + #[case] wire: Wire, +) { + let outcome = drive(agent, version_var, wire, None, tool_prompt()).await; + + assert!(outcome.succeeded(), "{outcome:?}"); + assert!(!outcome.tool_calls.is_empty(), "{outcome:?}"); + assert!(outcome.text.contains("tool-ok"), "{outcome:?}"); +} + +#[rstest] +#[case::claude_messages(&ClaudeCode, "TESTKIT_CLAUDE_VERSION", Wire::Messages)] +#[case::codex_responses(&Codex, "TESTKIT_CODEX_VERSION", Wire::Responses)] +#[case::opencode_chat(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::ChatCompletions)] +#[case::opencode_responses(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Responses)] +#[case::opencode_messages(&Opencode, "TESTKIT_OPENCODE_VERSION", Wire::Messages)] +#[ignore = "needs a live gateway, see the module docs"] +#[tokio::test] +async fn model_the_gateway_rejects_is_reported_as_an_error( + #[case] agent: &impl Agent, + #[case] version_var: &str, + #[case] wire: Wire, +) { + let outcome = drive( + agent, + version_var, + wire, + Some("testkit-no-such-model"), + text_prompt(), + ) + .await; + + assert!(!outcome.succeeded(), "{outcome:?}"); + assert!(!outcome.errors.is_empty(), "{outcome:?}"); +} diff --git a/litellm-rust/crates/testkit/tests/session.rs b/litellm-rust/crates/testkit/tests/session.rs new file mode 100644 index 00000000000..cd5e0dcbc71 --- /dev/null +++ b/litellm-rust/crates/testkit/tests/session.rs @@ -0,0 +1,155 @@ +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use litellm_testkit::{ + Configure, Drive, Error, Installed, LaunchSpec, Outcome, Prompt, Session, Settings, Version, + Wire, +}; + +struct Scripted; + +impl Configure for Scripted { + fn configure( + &self, + version: &Version, + _settings: &Settings, + home: &Path, + ) -> Result { + Ok(LaunchSpec { + env: [ + ("AGENT_HOME".to_owned(), home.to_string_lossy().into_owned()), + ("AGENT_SAW_VERSION".to_owned(), version.to_string()), + ] + .into(), + files: [( + PathBuf::from("conf/agent.toml"), + "configured = true\n".to_owned(), + )] + .into(), + }) + } +} + +impl Drive for Scripted { + fn args(&self, _version: &Version, _settings: &Settings, prompt: &Prompt) -> Vec { + vec!["--prompt".to_owned(), prompt.text.clone()] + } + + fn parse(&self, _version: &Version, stdout: &str) -> Outcome { + Outcome { + text: stdout.to_owned(), + ..Outcome::default() + } + } +} + +fn settings() -> Settings { + Settings { + base_url: "http://gateway.test".to_owned(), + api_key: "sk-test".to_owned(), + model: "some-model".to_owned(), + wire: Wire::Messages, + } +} + +fn prompt(text: &str) -> Prompt { + Prompt { + text: text.to_owned(), + allow_tools: false, + } +} + +fn session(script: &str) -> (Session, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let binary = dir.path().join("agent"); + std::fs::write(&binary, format!("#!/bin/sh\n{script}\n")).unwrap(); + std::fs::set_permissions(&binary, std::fs::Permissions::from_mode(0o755)).unwrap(); + let home = dir.path().join("home"); + std::fs::create_dir(&home).unwrap(); + let installed = Installed { + version: Version::new(4, 5, 6), + binary, + }; + ( + Session::prepare(&Scripted, &installed, settings(), home).unwrap(), + dir, + ) +} + +const LIMIT: Duration = Duration::from_secs(20); + +#[tokio::test] +async fn prepare_writes_the_config_files_under_home() { + let (_session, dir) = session("true"); + + let written = std::fs::read_to_string(dir.path().join("home/conf/agent.toml")).unwrap(); + + assert_eq!(written, "configured = true\n"); +} + +#[tokio::test] +async fn configure_and_drive_are_given_the_installed_version() { + let (session, _dir) = session("echo \"$AGENT_SAW_VERSION\""); + + let outcome = session.run(&Scripted, &prompt("hi"), LIMIT).await.unwrap(); + + assert_eq!(outcome.text.trim(), "4.5.6"); +} + +#[tokio::test] +async fn agent_runs_in_home_with_only_its_own_environment() { + let (session, dir) = session("pwd -P; env"); + + let outcome = session.run(&Scripted, &prompt("hi"), LIMIT).await.unwrap(); + + let home = dir.path().join("home").canonicalize().unwrap(); + assert_eq!(outcome.text.lines().next().unwrap(), home.to_string_lossy()); + assert!(outcome.text.contains("AGENT_HOME=")); + assert!( + !outcome.text.contains("CARGO_"), + "test runner environment leaked into the agent" + ); +} + +#[tokio::test] +async fn prompt_reaches_the_agent_as_one_untouched_argument() { + let (session, _dir) = session("printf '%s|' \"$@\""); + let text = "two spaces; $(echo injected) 'quoted'"; + + let outcome = session.run(&Scripted, &prompt(text), LIMIT).await.unwrap(); + + assert_eq!(outcome.text, format!("--prompt|{text}|")); +} + +#[tokio::test] +async fn clean_exit_is_a_success() { + let (session, _dir) = session("echo done"); + + let outcome = session.run(&Scripted, &prompt("hi"), LIMIT).await.unwrap(); + + assert_eq!(outcome.exit_code, Some(0)); + assert!(outcome.succeeded()); +} + +#[tokio::test] +async fn failing_exit_without_a_parsed_error_reports_stderr() { + let (session, _dir) = session("echo boom >&2; exit 3"); + + let outcome = session.run(&Scripted, &prompt("hi"), LIMIT).await.unwrap(); + + assert_eq!(outcome.exit_code, Some(3)); + assert!(!outcome.succeeded()); + assert_eq!(outcome.errors, ["boom\n"]); +} + +#[tokio::test] +async fn agent_that_outlives_the_limit_is_stopped() { + let (session, _dir) = session("sleep 30"); + + let result = session + .run(&Scripted, &prompt("hi"), Duration::from_millis(200)) + .await; + + assert!(matches!(result, Err(Error::Timeout(_)))); +} diff --git a/litellm-rust/crates/testkit/tests/support/mod.rs b/litellm-rust/crates/testkit/tests/support/mod.rs new file mode 100644 index 00000000000..f4a13759941 --- /dev/null +++ b/litellm-rust/crates/testkit/tests/support/mod.rs @@ -0,0 +1,70 @@ +use std::collections::HashMap; +use std::io::Write; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_testkit::{Error, Fetch}; +use sha2::{Digest, Sha256}; + +pub struct FakeFetch { + routes: HashMap>, + calls: AtomicUsize, +} + +impl FakeFetch { + pub fn new(routes: impl IntoIterator)>) -> Self { + Self { + routes: routes.into_iter().collect(), + calls: AtomicUsize::new(0), + } + } + + pub fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } +} + +impl Fetch for FakeFetch { + async fn get(&self, url: &str) -> Result, Error> { + self.calls.fetch_add(1, Ordering::SeqCst); + self.routes.get(url).cloned().ok_or_else(|| Error::Status { + url: url.to_owned(), + status: 404, + }) + } +} + +impl Fetch for &FakeFetch { + async fn get(&self, url: &str) -> Result, Error> { + (*self).get(url).await + } +} + +pub fn sha256(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) +} + +pub fn script_printing(output: &str) -> Vec { + format!("#!/bin/sh\necho '{output}'\n").into_bytes() +} + +pub fn tar_gz(member: &str, contents: &[u8]) -> Vec { + let mut builder = tar::Builder::new(Vec::new()); + let mut header = tar::Header::new_gnu(); + header.set_size(contents.len() as u64); + header.set_mode(0o755); + header.set_cksum(); + builder.append_data(&mut header, member, contents).unwrap(); + let tarball = builder.into_inner().unwrap(); + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(&tarball).unwrap(); + encoder.finish().unwrap() +} + +pub fn zip_archive(member: &str, contents: &[u8]) -> Vec { + let mut writer = zip::ZipWriter::new(std::io::Cursor::new(Vec::new())); + writer + .start_file(member, zip::write::SimpleFileOptions::default()) + .unwrap(); + writer.write_all(contents).unwrap(); + writer.finish().unwrap().into_inner() +} From 694783ebbeca0e15031c537ccf2427573af1460f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 11:11:51 -0700 Subject: [PATCH 08/58] ci: run migrated unit selections on every event in legacy GHA shards (#43182) * ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/_test-unit-base.yml | 18 ++++++------- .github/workflows/test-unit-proxy-db.yml | 33 ++++++++++++------------ .github/workflows/test-unit.yml | 20 +++++++------- 3 files changed, 36 insertions(+), 35 deletions(-) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index ef1dc53b4a6..fac0d766535 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -13,12 +13,13 @@ on: have its path existence-checked like any other token. required: true type: string - fork-flag: + unit-flag: description: >- Codecov flag of the `.circleci/tests.yml` job that now owns part of - this shard. CircleCI does not run on pull requests from forks, so on - those events this shard also runs the files - `.circleci/scripts/unit_selection.sh` lists for the flag. + this shard. The shard also runs the files + `.circleci/scripts/unit_selection.sh` lists for the flag, on every + event, because the CircleCI pipeline is manual-only while the tests + migrate. required: false type: string default: "" @@ -175,8 +176,7 @@ jobs: timeout-minutes: ${{ inputs.timeout-minutes }} env: TEST_PATH: ${{ inputs.test-path }} - FORK_FLAG: ${{ inputs.fork-flag }} - IS_FORK: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository }} + UNIT_FLAG: ${{ inputs.unit-flag }} MAX_FAILURES: ${{ inputs.max-failures }} WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} @@ -186,11 +186,11 @@ jobs: run: | echo "has-coverage=false" >> "$GITHUB_OUTPUT" selection="${TEST_PATH}" - if [ "${IS_FORK}" = "true" ] && [ -n "${FORK_FLAG}" ]; then - selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${FORK_FLAG}" | tr '\n' ' ')" + if [ -n "${UNIT_FLAG}" ]; then + selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${UNIT_FLAG}" | tr '\n' ' ')" fi if [ -z "${selection// /}" ]; then - echo "shard selection is empty on this event (CircleCI flag ${FORK_FLAG:-none} owns it); nothing to run" + echo "shard selection is empty; nothing to run" exit 0 fi pytest_args=() diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 86b385d91a7..da4477b6947 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -22,9 +22,10 @@ concurrency: # # `.circleci/tests.yml` runs each group's files on same-repo events under the # `proxy-db-` Codecov flag; `.circleci/scripts/unit_selection.sh` holds -# the file lists. CircleCI does not build pull requests from forks, so `fork-flag` -# makes the shard run that list there. `test-path` keeps the files that still -# reach real providers and never left tests/proxy_unit_tests. +# the file lists. That pipeline is manual-only while the tests migrate, so +# `unit-flag` makes the shard run that list on every event. `test-path` keeps +# the files that still reach real providers and never left +# tests/proxy_unit_tests. # # Design targets: # * Every shard runs in <= 7 minutes of wall-clock on the default runner. @@ -78,7 +79,7 @@ jobs: # Must run serially — event-loop conflict with the logging worker. - test-group: key-generation test-path: "" - fork-flag: proxy-db-key-generation + unit-flag: proxy-db-key-generation workers: 0 dist: loadscope timeout: 20 @@ -86,13 +87,13 @@ jobs: # ---- auth: split into 2 shards ---- - test-group: auth-checks test-path: "" - fork-flag: proxy-db-auth-checks + unit-flag: proxy-db-auth-checks workers: 4 dist: loadscope timeout: 15 - test-group: jwt-and-keys test-path: "" - fork-flag: proxy-db-jwt-and-keys + unit-flag: proxy-db-jwt-and-keys workers: 4 dist: loadscope timeout: 15 @@ -100,7 +101,7 @@ jobs: # ---- test_proxy_utils.py, single shard, worksteal distribution ---- - test-group: proxy-utils test-path: "" - fork-flag: proxy-db-proxy-utils + unit-flag: proxy-db-proxy-utils workers: 4 dist: worksteal timeout: 15 @@ -108,13 +109,13 @@ jobs: # ---- proxy server: split into 2 shards ---- - test-group: proxy-server-core test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py" - fork-flag: proxy-db-proxy-server-core + unit-flag: proxy-db-proxy-server-core workers: 4 dist: loadscope timeout: 15 - test-group: proxy-runtime test-path: "" - fork-flag: proxy-db-proxy-runtime + unit-flag: proxy-db-proxy-runtime workers: 4 dist: loadscope timeout: 15 @@ -122,20 +123,20 @@ jobs: # ---- logging: split into 2 shards ---- - test-group: custom-logging test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py" - fork-flag: proxy-db-custom-logging + unit-flag: proxy-db-custom-logging workers: 4 dist: loadscope timeout: 15 - test-group: logging-misc test-path: "" - fork-flag: proxy-db-logging-misc + unit-flag: proxy-db-logging-misc workers: 4 dist: loadscope timeout: 15 - test-group: db-and-spend test-path: "" - fork-flag: proxy-db-db-and-spend + unit-flag: proxy-db-db-and-spend workers: 4 dist: loadscope timeout: 15 @@ -143,27 +144,27 @@ jobs: # ---- guardrails + budget + hooks: split into 2 ---- - test-group: guardrails-hooks test-path: "" - fork-flag: proxy-db-guardrails-hooks + unit-flag: proxy-db-guardrails-hooks workers: 4 dist: loadscope timeout: 15 - test-group: budgets test-path: "" - fork-flag: proxy-db-budgets + unit-flag: proxy-db-budgets workers: 4 dist: loadscope timeout: 15 - test-group: endpoints-and-responses test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py" - fork-flag: proxy-db-endpoints-and-responses + unit-flag: proxy-db-endpoints-and-responses workers: 4 dist: loadscope timeout: 15 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} - fork-flag: ${{ matrix.fork-flag }} + unit-flag: ${{ matrix.unit-flag }} workers: ${{ matrix.workers }} reruns: 2 timeout-minutes: ${{ matrix.timeout }} diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 126a6e26e6f..a60d230d05f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -36,9 +36,9 @@ concurrency: # Folding it in here is a follow-up, together with generalising that guard into # assert_ci_coverage.py. # -# `fork-flag` names the `.circleci/tests.yml` job that now runs part of the -# shard under the same Codecov flag. CircleCI does not build pull requests from -# forks, so the shard still runs those files there and skips them elsewhere. +# `unit-flag` names the `.circleci/tests.yml` job that now runs part of the +# shard under the same Codecov flag. That pipeline is manual-only while the +# tests migrate, so the shard also runs those files on every event. jobs: unit: name: ${{ matrix.shard }} @@ -53,7 +53,7 @@ jobs: - shard: mcp-integration artifact-name: mcp-integration test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" - fork-flag: mcp-integration + unit-flag: mcp-integration workers: 2 reruns: 0 timeout-minutes: 20 @@ -73,7 +73,7 @@ jobs: tests/test_litellm/google_genai tests/test_litellm/router_utils tests/test_litellm/router_strategy - fork-flag: enterprise-routing + unit-flag: enterprise-routing workers: 2 reruns: 2 timeout-minutes: 20 @@ -205,7 +205,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py - fork-flag: proxy-infra + unit-flag: proxy-infra workers: 4 reruns: 2 timeout-minutes: 20 @@ -214,7 +214,7 @@ jobs: - shard: caching-local artifact-name: caching-local test-path: "" - fork-flag: caching-local + unit-flag: caching-local workers: 2 reruns: 2 timeout-minutes: 20 @@ -223,7 +223,7 @@ jobs: - shard: proxy-extras artifact-name: proxy-extras test-path: "" - fork-flag: proxy-extras + unit-flag: proxy-extras workers: 2 reruns: 2 timeout-minutes: 20 @@ -232,7 +232,7 @@ jobs: - shard: enterprise-package artifact-name: enterprise-package test-path: "" - fork-flag: enterprise-package + unit-flag: enterprise-package workers: 4 reruns: 2 timeout-minutes: 20 @@ -251,7 +251,7 @@ jobs: uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} - fork-flag: ${{ matrix.fork-flag || '' }} + unit-flag: ${{ matrix.unit-flag || '' }} workers: ${{ matrix.workers }} reruns: ${{ matrix.reruns }} timeout-minutes: ${{ matrix.timeout-minutes }} From f6882246d4a86be4a5666f70c166802cf029d746 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 11:30:43 -0700 Subject: [PATCH 09/58] test: move tests/test_litellm root and small trees into tests/unit (#43186) * ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/classify_changes.sh | 4 +- .circleci/scripts/unit_selection.sh | 23 + .circleci/tests.yml | 8 + .github/merge-smoke-tests.json | 4 +- .github/workflows/test-redis-compat.yml | 4 +- .github/workflows/test-unit.yml | 20 +- Makefile | 4 +- tests/_vcr_conftest_common.py | 2 +- .../code_qa_check_tests.py | 13 +- .../router_code_coverage.py | 2 +- tests/llm_translation/test_skills_api.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 387 -- .../chat_completions/test_dispatch.py | 117 - tests/test_litellm/conftest.py | 9 + .../test_litellm_responses_bridge.py | 72 - tests/test_litellm/messages/__init__.py | 0 tests/test_litellm/messages/test_dispatch.py | 155 - tests/test_litellm/rag/__init__.py | 0 tests/test_litellm/rag/ingestion/__init__.py | 0 tests/test_litellm/rerank_api/__init__.py | 0 .../test_router_tag_routing.py | 2 +- tests/test_litellm/test_compression.py | 608 --- tests/test_litellm/test_main.py | 4078 ---------------- tests/test_litellm/types/__init__.py | 0 tests/test_litellm/types/proxy/__init__.py | 0 .../types/proxy/policy_engine/__init__.py | 0 tests/test_litellm/vector_stores/__init__.py | 0 tests/test_litellm/videos/__init__.py | 0 tests/unit/batches/test_batch_utils.py | 345 ++ tests/unit/chat_completions/test_dispatch.py | 99 + .../__init__.py | 0 ...itellm_responses_transformation_handler.py | 0 ...responses_transformation_transformation.py | 0 tests/unit/conftest.py | 182 +- .../providers => unit/containers}/__init__.py | 0 .../test_azure_container_transformation.py | 0 .../containers/test_container_api.py | 0 .../containers/test_container_handler_url.py | 0 .../containers/test_container_integration.py | 0 .../test_container_proxy_ownership.py | 0 .../test_container_regional_api_base.py | 0 .../test_container_transformation.py | 0 .../containers/test_container_utils.py | 0 .../containers/test_endpoint_factory.py | 0 .../embeddings}/__init__.py | 0 .../embeddings/test_dispatch.py | 0 .../experimental_mcp_client}/__init__.py | 0 .../test_mcp_client.py | 0 .../experimental_mcp_client/test_tools.py | 0 .../batches => unit/files}/__init__.py | 0 .../{test_litellm => unit}/files/test_main.py | 0 .../fixtures}/__init__.py | 0 .../fixtures/together_ai_sync}/__init__.py | 0 .../fixtures/together_ai_sync/deprecations.md | 0 .../together_ai_sync/models_serverless.json | 0 .../google_genai}/__init__.py | 0 .../google_genai/test_google_genai_adapter.py | 0 .../test_google_genai_adapter_fixes.py | 0 .../google_genai/test_google_genai_handler.py | 86 - .../google_genai/test_google_genai_main.py | 0 .../test_google_genai_streaming_iterator.py | 0 .../test_google_genai_transformation.py | 0 .../endpoints => unit/images}/__init__.py | 0 .../images/test_image_edit_extra_params.py | 0 .../images/test_image_edit_utils.py | 0 .../test_image_generation_extra_headers.py | 0 .../speech => unit/interactions}/__init__.py | 0 .../interactions/test_agents_http_handler.py | 0 .../test_agents_main_and_utils.py | 0 .../test_background_cost_polling.py | 0 ...test_gemini_interactions_transformation.py | 0 .../test_interactions_streaming_iterator.py | 0 .../test_litellm_responses_bridge.py | 80 + .../interactions/test_openapi_compliance.py | 2 +- tests/unit/messages/test_dispatch.py | 136 + tests/{test_litellm => unit}/rag/test_main.py | 0 .../rerank_api}/__init__.py | 0 .../rerank_api/test_main.py | 0 .../test_a2a_registry_lookup.py | 0 .../test_acompletion_session_reuse_e2e.py | 0 .../test_add_deployment_no_master_key.py | 0 .../test_aembedding_session_reuse_e2e.py | 0 .../test_anthropic_beta_headers_filtering.py | 0 .../test_anthropic_skills_transformation.py | 0 .../test_assert_ci_coverage.py | 0 .../test_assert_workflow_dir_hygiene.py | 0 .../test_audio_transcription_rust_bridge.py | 0 ...to_update_price_and_context_window_file.py | 0 ...st_azure_ad_token_credential_resolution.py | 0 .../test_azure_ai_grok_4_3_model_metadata.py | 0 .../test_azure_ai_grok_4_6_model_metadata.py | 0 .../test_baseten_glm_5_3_model_metadata.py | 0 ...t_batch_completion_models_all_responses.py | 0 ..._bedrock_marengo_embed_3_model_metadata.py | 0 .../test_budget_ratchet_check.py | 0 .../test_chat_ui_responses_session.py | 0 .../test_check_licenses.py | 0 .../test_check_mcp_operation_boundary.py | 0 .../test_check_migrations_no_data_rewrites.py | 0 .../test_check_py310_typing_imports.py | 0 .../test_check_test_quality.py | 0 .../test_check_type_discipline.py | 0 .../test_circleci_path_filter.py | 0 .../test_circleci_rust_toolchain.py | 0 .../test_claude_fable_5_config.py | 0 .../test_claude_opus_4_6_config.py | 0 .../test_claude_opus_4_8_config.py | 0 .../test_claude_opus_5_config.py | 0 .../test_claude_sonnet_5_config.py | 0 ...st_cloudflare_workers_ai_model_metadata.py | 0 .../test_completion_timeout_resolution.py | 0 .../test_component_entrypoint.py | 0 tests/unit/test_compression.py | 649 +++ .../test_conftest_isolation.py | 0 .../{test_litellm => unit}/test_constants.py | 0 .../test_container_router.py | 0 .../test_cost_calculation_log_level.py | 0 .../test_cost_calculator.py | 0 .../test_cost_map_guard.py | 0 .../test_count_tokens_public_api.py | 0 .../test_dashscope_image_generation.py | 2 +- .../test_daybreak_model_metadata.py | 0 .../test_deepseek_model_metadata.py | 0 .../test_default_branch.py | 0 .../test_detect_changes.py | 0 .../test_dockerfile_apk_repository.py | 0 .../test_dockerfile_bedrock_realtime_extra.py | 0 .../test_dockerfile_non_root.py | 0 .../test_drop_params_env_var.py | 0 .../test_e2e_egress_sentinel.py | 0 .../test_eager_tiktoken_load.py | 0 .../test_env_key_doc_gate.py | 0 .../test_exception_exports.py | 0 .../test_exception_header_preservation.py | 0 ...est_exception_mapping_request_attribute.py | 0 .../test_filter_out_litellm_params.py | 0 .../test_fireworks_serverless_model_costs.py | 0 .../test_gate_slot_lock.py | 0 ...est_gemini_3_1_flash_lite_image_pricing.py | 0 .../test_gemini_tts_native_audio_pricing.py | 0 .../test_get_blog_posts.py | 0 .../{test_litellm => unit}/test_git_hooks.py | 0 .../test_gpt_5_4_model_metadata.py | 0 .../test_gpt_5_5_model_metadata.py | 0 .../test_gpt_image_cost_calculator.py | 0 .../test_gpt_realtime_mode.py | 0 .../test_groq_streaming_encoding.py | 0 .../test_guardrail_exception_status_codes.py | 0 .../test_lazy_imports.py | 0 .../test_lint_workflow_diff_gates.py | 0 .../test_litellm_params_reserved_keys.py | 0 tests/{test_litellm => unit}/test_logging.py | 0 .../test_lowest_latency_zero_tokens.py | 0 tests/unit/test_main.py | 4124 +++++++++++++++++ .../test_main_module_header.py | 0 .../test_mistral_medium_3_5_model_metadata.py | 0 .../test_mistral_small_4_0_model_metadata.py | 0 ...test_mistral_zai_glm_5_2_model_metadata.py | 0 .../test_model_block_unblock.py | 0 .../test_model_cost_aliases.py | 0 .../test_model_param_helper.py | 0 .../test_model_prices_schema.py | 0 .../test_model_response_normalization.py | 0 .../test_muse_spark_1_1_model_metadata.py | 0 .../test_muse_spark_1_2_model_metadata.py | 0 .../test_muse_spark_1_3_model_metadata.py | 0 .../test_mutation_report.py | 0 .../test_nested_drop_params.py | 0 .../test_non_chat_routes_open_llm_spans.py | 0 ...penai_embedding_encoding_format_default.py | 0 ...penai_service_tier_long_context_pricing.py | 0 .../test_pre_commit_lint.py | 0 .../test_prisma_generate_if_needed.py | 0 .../test_process_helpers.py | 0 .../test_project_alias_tracking.py | 0 .../test_project_tags_pydantic.py | 0 .../{test_litellm => unit}/test_proxy_auth.py | 0 .../test_rag_openai_ingestion.py | 0 .../test_rate_limit_error_unification.py | 0 .../test_read_rc_version.py | 0 .../test_redact_string_in_error_paths.py | 0 tests/{test_litellm => unit}/test_redis.py | 0 .../test_redis_credential_provider.py | 0 .../test_register_model_custom_pricing.py | 0 ...st_register_model_zero_cost_persistence.py | 0 .../test_replicate_model_key_format.py | 0 .../test_responses_api_bridge_non_stream.py | 0 .../test_responses_id_security.py | 60 +- ...responses_streaming_container_ownership.py | 0 .../test_retrieve_batch_bedrock_dispatch.py | 0 .../test_router}/test_router.py | 0 .../test_router_block_helpers.py | 0 .../test_router_exception_redaction.py | 0 .../test_router_google_genai.py | 0 .../test_router_model_cost_isolation.py | 0 .../test_router_order_fallback.py | 0 .../test_router_per_deployment_num_retries.py | 0 .../test_router_redis_init.py | 0 .../test_router_retry_backoff_headers.py | 0 .../test_router_retry_non_retryable_errors.py | 0 .../test_router_retry_policy_update.py | 0 .../test_router_silent_experiment.py | 82 +- ...test_router_streaming_fallback_metadata.py | 0 .../test_router_weighted_failover.py | 0 .../test_ruff_strict_gate.py | 0 .../test_sambanova_model_metadata.py | 0 .../test_secret_redaction.py | 0 .../test_select_ui_test_scope.py | 0 .../test_service_logger.py | 0 .../test_setup_wizard.py | 0 .../test_shared_session_integration.py | 0 .../test_ssl_verify_unit.py | 35 - .../test_stream_chunk_builder_annotations.py | 0 .../test_stream_chunk_builder_citations.py | 0 .../test_stream_chunk_builder_images.py | 0 .../test_streaming_connection_cleanup.py | 0 .../test_sync_together_ai_models.py | 0 .../test_system_message_format_bug.py | 0 .../test_test_quality_gate.py | 0 .../test_thinking_enabled.py | 0 .../test_together_ai_model_metadata.py | 0 .../test_type_check_gate.py | 0 .../test_type_discipline_gate.py | 0 .../test_typesafe_model_metadata.py | 0 .../test_unit_shard_missing_paths.py | 1 + .../test_unit_shard_per_test_timeout.py | 0 tests/{test_litellm => unit}/test_utils.py | 0 .../test_utils_module_docstring.py | 0 .../test_uuid_helper.py | 0 .../test_vcr_safe_body_matcher.py | 8 - ...tex_ai_xai_grok_prompt_caching_metadata.py | 0 .../test_video_generation.py | 0 .../test_with_dashboard_node.py | 0 .../test_xai_grok_4_3_model_metadata.py | 0 .../test_xai_responses_auto_routing.py | 0 .../types/test_completion.py | 2 +- .../test_guardrails_case_normalization.py | 0 .../{test_litellm => unit}/types/test_mcp.py | 0 .../types/test_presidio_entity_expansion.py | 0 .../test_prometheus_label_value_sanitize.py | 0 .../types/test_prometheus_latency_buckets.py | 0 .../types/test_router.py | 0 .../types/test_types_utils.py | 0 .../types/test_uk_pii_entities.py | 0 .../files => unit/vector_stores}/__init__.py | 0 .../vector_stores/test_main.py | 0 ...test_vector_store_create_provider_logic.py | 0 .../test_vector_store_registry.py | 0 248 files changed, 5722 insertions(+), 5685 deletions(-) delete mode 100644 tests/test_litellm/batches/test_batch_utils.py delete mode 100644 tests/test_litellm/chat_completions/test_dispatch.py delete mode 100644 tests/test_litellm/messages/__init__.py delete mode 100644 tests/test_litellm/messages/test_dispatch.py delete mode 100644 tests/test_litellm/rag/__init__.py delete mode 100644 tests/test_litellm/rag/ingestion/__init__.py delete mode 100644 tests/test_litellm/rerank_api/__init__.py delete mode 100644 tests/test_litellm/types/__init__.py delete mode 100644 tests/test_litellm/types/proxy/__init__.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/__init__.py delete mode 100644 tests/test_litellm/vector_stores/__init__.py delete mode 100644 tests/test_litellm/videos/__init__.py rename tests/{test_litellm/a2a_protocol => unit/completion_extras/litellm_responses_transformation}/__init__.py (100%) rename tests/{test_litellm => unit}/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py (100%) rename tests/{test_litellm => unit}/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py (100%) rename tests/{test_litellm/a2a_protocol/providers => unit/containers}/__init__.py (100%) rename tests/{test_litellm => unit}/containers/test_azure_container_transformation.py (100%) rename tests/{test_litellm => unit}/containers/test_container_api.py (100%) rename tests/{test_litellm => unit}/containers/test_container_handler_url.py (100%) rename tests/{test_litellm => unit}/containers/test_container_integration.py (100%) rename tests/{test_litellm => unit}/containers/test_container_proxy_ownership.py (100%) rename tests/{test_litellm => unit}/containers/test_container_regional_api_base.py (100%) rename tests/{test_litellm => unit}/containers/test_container_transformation.py (100%) rename tests/{test_litellm => unit}/containers/test_container_utils.py (100%) rename tests/{test_litellm => unit}/containers/test_endpoint_factory.py (100%) rename tests/{test_litellm/a2a_protocol/providers/bedrock_agentcore => unit/embeddings}/__init__.py (100%) rename tests/{test_litellm => unit}/embeddings/test_dispatch.py (100%) rename tests/{test_litellm/a2a_protocol/providers/pydantic_ai_agents => unit/experimental_mcp_client}/__init__.py (100%) rename tests/{test_litellm => unit}/experimental_mcp_client/test_mcp_client.py (100%) rename tests/{test_litellm => unit}/experimental_mcp_client/test_tools.py (100%) rename tests/{test_litellm/batches => unit/files}/__init__.py (100%) rename tests/{test_litellm => unit}/files/test_main.py (100%) rename tests/{test_litellm/chat_completions => unit/fixtures}/__init__.py (100%) rename tests/{test_litellm/completion_extras => unit/fixtures/together_ai_sync}/__init__.py (100%) rename tests/{test_litellm => unit}/fixtures/together_ai_sync/deprecations.md (100%) rename tests/{test_litellm => unit}/fixtures/together_ai_sync/models_serverless.json (100%) rename tests/{test_litellm/containers => unit/google_genai}/__init__.py (100%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_adapter.py (100%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_adapter_fixes.py (100%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_handler.py (76%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_main.py (100%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/google_genai/test_google_genai_transformation.py (100%) rename tests/{test_litellm/endpoints => unit/images}/__init__.py (100%) rename tests/{test_litellm => unit}/images/test_image_edit_extra_params.py (100%) rename tests/{test_litellm => unit}/images/test_image_edit_utils.py (100%) rename tests/{test_litellm => unit}/images/test_image_generation_extra_headers.py (100%) rename tests/{test_litellm/endpoints/speech => unit/interactions}/__init__.py (100%) rename tests/{test_litellm => unit}/interactions/test_agents_http_handler.py (100%) rename tests/{test_litellm => unit}/interactions/test_agents_main_and_utils.py (100%) rename tests/{test_litellm => unit}/interactions/test_background_cost_polling.py (100%) rename tests/{test_litellm => unit}/interactions/test_gemini_interactions_transformation.py (100%) rename tests/{test_litellm => unit}/interactions/test_interactions_streaming_iterator.py (100%) create mode 100644 tests/unit/interactions/test_litellm_responses_bridge.py rename tests/{test_litellm => unit}/interactions/test_openapi_compliance.py (99%) rename tests/{test_litellm => unit}/rag/test_main.py (100%) rename tests/{test_litellm/endpoints/speech/speech_to_completion_bridge => unit/rerank_api}/__init__.py (100%) rename tests/{test_litellm => unit}/rerank_api/test_main.py (100%) rename tests/{test_litellm => unit}/test_a2a_registry_lookup.py (100%) rename tests/{test_litellm => unit}/test_acompletion_session_reuse_e2e.py (100%) rename tests/{test_litellm => unit}/test_add_deployment_no_master_key.py (100%) rename tests/{test_litellm => unit}/test_aembedding_session_reuse_e2e.py (100%) rename tests/{test_litellm => unit}/test_anthropic_beta_headers_filtering.py (100%) rename tests/{test_litellm => unit}/test_anthropic_skills_transformation.py (100%) rename tests/{test_litellm => unit}/test_assert_ci_coverage.py (100%) rename tests/{test_litellm => unit}/test_assert_workflow_dir_hygiene.py (100%) rename tests/{test_litellm => unit}/test_audio_transcription_rust_bridge.py (100%) rename tests/{test_litellm => unit}/test_auto_update_price_and_context_window_file.py (100%) rename tests/{test_litellm => unit}/test_azure_ad_token_credential_resolution.py (100%) rename tests/{test_litellm => unit}/test_azure_ai_grok_4_3_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_azure_ai_grok_4_6_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_baseten_glm_5_3_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_batch_completion_models_all_responses.py (100%) rename tests/{test_litellm => unit}/test_bedrock_marengo_embed_3_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_budget_ratchet_check.py (100%) rename tests/{test_litellm => unit}/test_chat_ui_responses_session.py (100%) rename tests/{test_litellm => unit}/test_check_licenses.py (100%) rename tests/{test_litellm => unit}/test_check_mcp_operation_boundary.py (100%) rename tests/{test_litellm => unit}/test_check_migrations_no_data_rewrites.py (100%) rename tests/{test_litellm => unit}/test_check_py310_typing_imports.py (100%) rename tests/{test_litellm => unit}/test_check_test_quality.py (100%) rename tests/{test_litellm => unit}/test_check_type_discipline.py (100%) rename tests/{test_litellm => unit}/test_circleci_path_filter.py (100%) rename tests/{test_litellm => unit}/test_circleci_rust_toolchain.py (100%) rename tests/{test_litellm => unit}/test_claude_fable_5_config.py (100%) rename tests/{test_litellm => unit}/test_claude_opus_4_6_config.py (100%) rename tests/{test_litellm => unit}/test_claude_opus_4_8_config.py (100%) rename tests/{test_litellm => unit}/test_claude_opus_5_config.py (100%) rename tests/{test_litellm => unit}/test_claude_sonnet_5_config.py (100%) rename tests/{test_litellm => unit}/test_cloudflare_workers_ai_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_completion_timeout_resolution.py (100%) rename tests/{test_litellm => unit}/test_component_entrypoint.py (100%) create mode 100644 tests/unit/test_compression.py rename tests/{test_litellm => unit}/test_conftest_isolation.py (100%) rename tests/{test_litellm => unit}/test_constants.py (100%) rename tests/{test_litellm => unit}/test_container_router.py (100%) rename tests/{test_litellm => unit}/test_cost_calculation_log_level.py (100%) rename tests/{test_litellm => unit}/test_cost_calculator.py (100%) rename tests/{test_litellm => unit}/test_cost_map_guard.py (100%) rename tests/{test_litellm => unit}/test_count_tokens_public_api.py (100%) rename tests/{test_litellm => unit}/test_dashscope_image_generation.py (99%) rename tests/{test_litellm => unit}/test_daybreak_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_deepseek_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_default_branch.py (100%) rename tests/{test_litellm => unit}/test_detect_changes.py (100%) rename tests/{test_litellm => unit}/test_dockerfile_apk_repository.py (100%) rename tests/{test_litellm => unit}/test_dockerfile_bedrock_realtime_extra.py (100%) rename tests/{test_litellm => unit}/test_dockerfile_non_root.py (100%) rename tests/{test_litellm => unit}/test_drop_params_env_var.py (100%) rename tests/{test_litellm => unit}/test_e2e_egress_sentinel.py (100%) rename tests/{test_litellm => unit}/test_eager_tiktoken_load.py (100%) rename tests/{test_litellm => unit}/test_env_key_doc_gate.py (100%) rename tests/{test_litellm => unit}/test_exception_exports.py (100%) rename tests/{test_litellm => unit}/test_exception_header_preservation.py (100%) rename tests/{test_litellm => unit}/test_exception_mapping_request_attribute.py (100%) rename tests/{test_litellm => unit}/test_filter_out_litellm_params.py (100%) rename tests/{test_litellm => unit}/test_fireworks_serverless_model_costs.py (100%) rename tests/{test_litellm => unit}/test_gate_slot_lock.py (100%) rename tests/{test_litellm => unit}/test_gemini_3_1_flash_lite_image_pricing.py (100%) rename tests/{test_litellm => unit}/test_gemini_tts_native_audio_pricing.py (100%) rename tests/{test_litellm => unit}/test_get_blog_posts.py (100%) rename tests/{test_litellm => unit}/test_git_hooks.py (100%) rename tests/{test_litellm => unit}/test_gpt_5_4_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_gpt_5_5_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_gpt_image_cost_calculator.py (100%) rename tests/{test_litellm => unit}/test_gpt_realtime_mode.py (100%) rename tests/{test_litellm => unit}/test_groq_streaming_encoding.py (100%) rename tests/{test_litellm => unit}/test_guardrail_exception_status_codes.py (100%) rename tests/{test_litellm => unit}/test_lazy_imports.py (100%) rename tests/{test_litellm => unit}/test_lint_workflow_diff_gates.py (100%) rename tests/{test_litellm => unit}/test_litellm_params_reserved_keys.py (100%) rename tests/{test_litellm => unit}/test_logging.py (100%) rename tests/{test_litellm => unit}/test_lowest_latency_zero_tokens.py (100%) create mode 100644 tests/unit/test_main.py rename tests/{test_litellm => unit}/test_main_module_header.py (100%) rename tests/{test_litellm => unit}/test_mistral_medium_3_5_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_mistral_small_4_0_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_mistral_zai_glm_5_2_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_model_block_unblock.py (100%) rename tests/{test_litellm => unit}/test_model_cost_aliases.py (100%) rename tests/{test_litellm => unit}/test_model_param_helper.py (100%) rename tests/{test_litellm => unit}/test_model_prices_schema.py (100%) rename tests/{test_litellm => unit}/test_model_response_normalization.py (100%) rename tests/{test_litellm => unit}/test_muse_spark_1_1_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_muse_spark_1_2_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_muse_spark_1_3_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_mutation_report.py (100%) rename tests/{test_litellm => unit}/test_nested_drop_params.py (100%) rename tests/{test_litellm => unit}/test_non_chat_routes_open_llm_spans.py (100%) rename tests/{test_litellm => unit}/test_openai_embedding_encoding_format_default.py (100%) rename tests/{test_litellm => unit}/test_openai_service_tier_long_context_pricing.py (100%) rename tests/{test_litellm => unit}/test_pre_commit_lint.py (100%) rename tests/{test_litellm => unit}/test_prisma_generate_if_needed.py (100%) rename tests/{test_litellm => unit}/test_process_helpers.py (100%) rename tests/{test_litellm => unit}/test_project_alias_tracking.py (100%) rename tests/{test_litellm => unit}/test_project_tags_pydantic.py (100%) rename tests/{test_litellm => unit}/test_proxy_auth.py (100%) rename tests/{test_litellm => unit}/test_rag_openai_ingestion.py (100%) rename tests/{test_litellm => unit}/test_rate_limit_error_unification.py (100%) rename tests/{test_litellm => unit}/test_read_rc_version.py (100%) rename tests/{test_litellm => unit}/test_redact_string_in_error_paths.py (100%) rename tests/{test_litellm => unit}/test_redis.py (100%) rename tests/{test_litellm => unit}/test_redis_credential_provider.py (100%) rename tests/{test_litellm => unit}/test_register_model_custom_pricing.py (100%) rename tests/{test_litellm => unit}/test_register_model_zero_cost_persistence.py (100%) rename tests/{test_litellm => unit}/test_replicate_model_key_format.py (100%) rename tests/{test_litellm => unit}/test_responses_api_bridge_non_stream.py (100%) rename tests/{test_litellm => unit}/test_responses_id_security.py (94%) rename tests/{test_litellm => unit}/test_responses_streaming_container_ownership.py (100%) rename tests/{test_litellm => unit}/test_retrieve_batch_bedrock_dispatch.py (100%) rename tests/{test_litellm => unit/test_router}/test_router.py (100%) rename tests/{test_litellm => unit}/test_router_block_helpers.py (100%) rename tests/{test_litellm => unit}/test_router_exception_redaction.py (100%) rename tests/{test_litellm => unit}/test_router_google_genai.py (100%) rename tests/{test_litellm => unit}/test_router_model_cost_isolation.py (100%) rename tests/{test_litellm => unit}/test_router_order_fallback.py (100%) rename tests/{test_litellm => unit}/test_router_per_deployment_num_retries.py (100%) rename tests/{test_litellm => unit}/test_router_redis_init.py (100%) rename tests/{test_litellm => unit}/test_router_retry_backoff_headers.py (100%) rename tests/{test_litellm => unit}/test_router_retry_non_retryable_errors.py (100%) rename tests/{test_litellm => unit}/test_router_retry_policy_update.py (100%) rename tests/{test_litellm => unit}/test_router_silent_experiment.py (92%) rename tests/{test_litellm => unit}/test_router_streaming_fallback_metadata.py (100%) rename tests/{test_litellm => unit}/test_router_weighted_failover.py (100%) rename tests/{test_litellm => unit}/test_ruff_strict_gate.py (100%) rename tests/{test_litellm => unit}/test_sambanova_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_secret_redaction.py (100%) rename tests/{test_litellm => unit}/test_select_ui_test_scope.py (100%) rename tests/{test_litellm => unit}/test_service_logger.py (100%) rename tests/{test_litellm => unit}/test_setup_wizard.py (100%) rename tests/{test_litellm => unit}/test_shared_session_integration.py (100%) rename tests/{test_litellm => unit}/test_ssl_verify_unit.py (83%) rename tests/{test_litellm => unit}/test_stream_chunk_builder_annotations.py (100%) rename tests/{test_litellm => unit}/test_stream_chunk_builder_citations.py (100%) rename tests/{test_litellm => unit}/test_stream_chunk_builder_images.py (100%) rename tests/{test_litellm => unit}/test_streaming_connection_cleanup.py (100%) rename tests/{test_litellm => unit}/test_sync_together_ai_models.py (100%) rename tests/{test_litellm => unit}/test_system_message_format_bug.py (100%) rename tests/{test_litellm => unit}/test_test_quality_gate.py (100%) rename tests/{test_litellm => unit}/test_thinking_enabled.py (100%) rename tests/{test_litellm => unit}/test_together_ai_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_type_check_gate.py (100%) rename tests/{test_litellm => unit}/test_type_discipline_gate.py (100%) rename tests/{test_litellm => unit}/test_typesafe_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_unit_shard_missing_paths.py (97%) rename tests/{test_litellm => unit}/test_unit_shard_per_test_timeout.py (100%) rename tests/{test_litellm => unit}/test_utils.py (100%) rename tests/{test_litellm => unit}/test_utils_module_docstring.py (100%) rename tests/{test_litellm => unit}/test_uuid_helper.py (100%) rename tests/{test_litellm => unit}/test_vcr_safe_body_matcher.py (98%) rename tests/{test_litellm => unit}/test_vertex_ai_xai_grok_prompt_caching_metadata.py (100%) rename tests/{test_litellm => unit}/test_video_generation.py (100%) rename tests/{test_litellm => unit}/test_with_dashboard_node.py (100%) rename tests/{test_litellm => unit}/test_xai_grok_4_3_model_metadata.py (100%) rename tests/{test_litellm => unit}/test_xai_responses_auto_routing.py (100%) rename tests/{test_litellm => unit}/types/test_completion.py (99%) rename tests/{test_litellm => unit}/types/test_guardrails_case_normalization.py (100%) rename tests/{test_litellm => unit}/types/test_mcp.py (100%) rename tests/{test_litellm => unit}/types/test_presidio_entity_expansion.py (100%) rename tests/{test_litellm => unit}/types/test_prometheus_label_value_sanitize.py (100%) rename tests/{test_litellm => unit}/types/test_prometheus_latency_buckets.py (100%) rename tests/{test_litellm => unit}/types/test_router.py (100%) rename tests/{test_litellm => unit}/types/test_types_utils.py (100%) rename tests/{test_litellm => unit}/types/test_uk_pii_entities.py (100%) rename tests/{test_litellm/files => unit/vector_stores}/__init__.py (100%) rename tests/{test_litellm => unit}/vector_stores/test_main.py (100%) rename tests/{test_litellm => unit}/vector_stores/test_vector_store_create_provider_logic.py (100%) rename tests/{test_litellm => unit}/vector_stores/test_vector_store_registry.py (100%) diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index ad265a5e39f..8c2ac019b99 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -14,12 +14,12 @@ while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in *.md | *.mdx) : ;; - pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py) + pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py) has_mcp_dependencies=true ;; esac case "$file" in tests/e2e/*/*.py) : ;; - tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock) + tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock) has_provider_harness=true ;; esac case "$file" in diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index f2ee7550df3..5ce8b6c84ba 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -8,6 +8,7 @@ legacy_flags=( enterprise-package enterprise-routing mcp-integration + misc proxy-db-auth-checks proxy-db-budgets proxy-db-custom-logging @@ -22,6 +23,7 @@ legacy_flags=( proxy-db-proxy-utils proxy-extras proxy-infra + responses-caching-types ) legacy_paths() { @@ -36,6 +38,7 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; enterprise-routing) + echo tests/unit/google_genai echo tests/unit/enterprise/enterprise_callbacks/send_emails echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py @@ -48,9 +51,28 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; mcp-integration) + echo tests/unit/experimental_mcp_client echo tests/unit/proxy/_experimental/mcp_server echo tests/unit/responses/mcp echo tests/mcp_tests/test_proxy_mcp_e2e.py ;; + misc) + find tests/unit -maxdepth 1 -name 'test_*.py' + echo tests/unit/test_router + echo tests/unit/a2a_protocol + echo tests/unit/batches + echo tests/unit/chat_completions + echo tests/unit/completion_extras + echo tests/unit/containers + echo tests/unit/embeddings + echo tests/unit/endpoints + echo tests/unit/files + echo tests/unit/images + echo tests/unit/interactions + echo tests/unit/messages + echo tests/unit/rag + echo tests/unit/rerank_api + echo tests/unit/vector_stores + echo tests/unit/videos ;; proxy-db-auth-checks) echo tests/unit/proxy/auth/test_auth_checks.py echo tests/unit/proxy/auth/test_user_api_key_auth.py @@ -113,6 +135,7 @@ legacy_paths() { proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; proxy-extras) echo tests/unit/litellm_proxy_extras ;; proxy-infra) echo tests/unit/gateway ;; + responses-caching-types) echo tests/unit/types ;; *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; esac } diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 264d7695a94..10ee19f146a 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -341,6 +341,7 @@ workflows: flag: - enterprise-package - proxy-infra + - responses-caching-types - proxy-db-auth-checks - proxy-db-jwt-and-keys - proxy-db-proxy-server-core @@ -353,6 +354,13 @@ workflows: - proxy-db-endpoints-and-responses base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-misc + flag: misc + shards: 2 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-proxy-db-proxy-utils flag: proxy-db-proxy-utils diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json index 6088953b7eb..8ed7b917460 100644 --- a/.github/merge-smoke-tests.json +++ b/.github/merge-smoke-tests.json @@ -5,8 +5,8 @@ "CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", "MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key", "MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]", - "COST-EXPLICIT": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", - "COST-ZERO": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero", + "COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", + "COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero", "LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on", "LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off", "CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger", diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml index 25fb8f8bce3..2f5ce4d441a 100644 --- a/.github/workflows/test-redis-compat.yml +++ b/.github/workflows/test-redis-compat.yml @@ -10,7 +10,7 @@ on: - "litellm/_redis_credential_provider.py" - "litellm/caching/redis_cache.py" - "litellm/caching/evicted_client_closer.py" - - "tests/test_litellm/test_redis.py" + - "tests/unit/test_redis.py" - "tests/local_testing/test_caching.py" - "tests/test_litellm/caching/test_redis_connection_pool.py" - "tests/test_litellm/caching/test_redis_cluster_cache.py" @@ -84,7 +84,7 @@ jobs: run: | redis-server --version uv run --no-sync pytest \ - tests/test_litellm/test_redis.py \ + tests/unit/test_redis.py \ tests/test_litellm/caching/test_redis_connection_pool.py \ tests/test_litellm/caching/test_redis_cluster_cache.py \ tests/test_litellm/caching/test_evicted_client_closer.py \ diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a60d230d05f..91b54f4ee70 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -52,7 +52,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" + test-path: "tests/mcp_tests" unit-flag: mcp-integration workers: 2 reruns: 0 @@ -70,7 +70,6 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing test-path: >- - tests/test_litellm/google_genai tests/test_litellm/router_utils tests/test_litellm/router_strategy unit-flag: enterprise-routing @@ -106,26 +105,13 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/test_litellm/batches tests/test_litellm/secret_managers - tests/test_litellm/a2a_protocol - tests/test_litellm/chat_completions - tests/test_litellm/completion_extras - tests/test_litellm/containers - tests/test_litellm/endpoints - tests/test_litellm/files - tests/test_litellm/images tests/test_litellm/interactions - tests/test_litellm/messages - tests/test_litellm/embeddings tests/test_litellm/ocr tests/test_litellm/passthrough - tests/test_litellm/rag - tests/test_litellm/rerank_api tests/test_litellm/rust_bridge - tests/test_litellm/vector_stores - tests/test_litellm/videos tests/test_litellm/test_*.py + unit-flag: misc workers: 2 reruns: 2 timeout-minutes: 20 @@ -243,7 +229,7 @@ jobs: test-path: >- tests/test_litellm/responses tests/test_litellm/caching - tests/test_litellm/types + unit-flag: responses-caching-types workers: 2 reruns: 2 timeout-minutes: 20 diff --git a/Makefile b/Makefile index 28daf589a23..62e6ae53275 100644 --- a/Makefile +++ b/Makefile @@ -332,10 +332,10 @@ test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps - $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 # Proxy unit tests (tests/unit/proxy split alphabetically) test-proxy-unit-a: install-test-deps diff --git a/tests/_vcr_conftest_common.py b/tests/_vcr_conftest_common.py index ab046674eb6..3adc671021b 100644 --- a/tests/_vcr_conftest_common.py +++ b/tests/_vcr_conftest_common.py @@ -52,7 +52,7 @@ from tests._vcr_redis_persister import ( # network call entirely, so skip tests record nothing (NOOP) and passing tests # stop carrying a volatile github episode. This matches the established idiom in # the unit-test suite, which sets the same flag (see e.g. -# tests/test_litellm/test_cost_calculator.py). ``setdefault`` so an explicit +# tests/unit/test_cost_calculator.py). ``setdefault`` so an explicit # override still wins. os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/code_coverage_tests/code_qa_check_tests.py b/tests/code_coverage_tests/code_qa_check_tests.py index 025f836511c..6c620a02522 100644 --- a/tests/code_coverage_tests/code_qa_check_tests.py +++ b/tests/code_coverage_tests/code_qa_check_tests.py @@ -13,15 +13,16 @@ def check_for_litellm_module_deletion(base_dir): del sys.modules[module] """ problematic_files = [] - test_dir = os.path.join(base_dir, "test_litellm") + candidate_dirs = [os.path.join(base_dir, name) for name in ("test_litellm", "unit")] + test_dirs = [test_dir for test_dir in candidate_dirs if os.path.exists(test_dir)] - if not os.path.exists(test_dir): - print(f"Warning: Directory {test_dir} does not exist.") + if not test_dirs: + print(f"Warning: None of {candidate_dirs} exist.") return [] - print(f"Checking directory: {test_dir}") + print(f"Checking directories: {test_dirs}") - for root, _, files in os.walk(test_dir): + for root, _, files in (entry for test_dir in test_dirs for entry in os.walk(test_dir)): for file in files: if file.endswith(".py"): file_path = os.path.join(root, file) @@ -173,7 +174,7 @@ def main(): f"This can cause import issues and test failures. Files: {problematic_files}" ) else: - print("✓ No litellm module deletion patterns found in test_litellm directory.") + print("✓ No litellm module deletion patterns found in tests/test_litellm or tests/unit.") if __name__ == "__main__": diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7332a533872..06e5b020836 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -31,7 +31,7 @@ def get_all_functions_called_in_tests(base_dir): specifically in files containing the word 'router'. """ called_functions = set() - test_dirs = ["local_testing", "router_unit_tests", "test_litellm"] + test_dirs = ["local_testing", "router_unit_tests", "test_litellm", "unit"] for test_dir in test_dirs: dir_path = os.path.join(base_dir, test_dir) diff --git a/tests/llm_translation/test_skills_api.py b/tests/llm_translation/test_skills_api.py index aeab5f0da3e..d21e7376ea7 100644 --- a/tests/llm_translation/test_skills_api.py +++ b/tests/llm_translation/test_skills_api.py @@ -277,4 +277,4 @@ class BaseSkillsAPITest(ABC): # # Transformation logic (URL construction, headers, request/response parsing) is # covered by unit tests in: -# tests/test_litellm/test_anthropic_skills_transformation.py +# tests/unit/test_anthropic_skills_transformation.py diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py deleted file mode 100644 index 0b2bfe9d266..00000000000 --- a/tests/test_litellm/batches/test_batch_utils.py +++ /dev/null @@ -1,387 +0,0 @@ -import json - -import pytest - -import litellm -import litellm.batches.batch_utils as bu -from litellm.types.llms.openai import Batch - -GROUNDED_USAGE_METADATA = { - "promptTokenCount": 19, - "candidatesTokenCount": 59, - "thoughtsTokenCount": 406, - "toolUsePromptTokenCount": 73, - "totalTokenCount": 557, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], - "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], - "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], - "trafficType": "ON_DEMAND", -} -PASSTHROUGH_OUTPUT_URI = ( - "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" - "predictions.jsonl" -) -UNGROUNDED_USAGE_METADATA = { - "promptTokenCount": 20, - "candidatesTokenCount": 48, - "thoughtsTokenCount": 195, - "toolUsePromptTokenCount": 73, - "totalTokenCount": 336, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], - "trafficType": "ON_DEMAND", -} - - -def _batch(output_file_id: str) -> Batch: - return Batch( - id="b", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="f", - object="batch", - status="completed", - output_file_id=output_file_id, - ) - - -def _vertex_jsonl(rows: list[dict]) -> bytes: - return "\n".join(json.dumps(row) for row in rows).encode() - - -def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict: - return { - "id": f"batch_req_{custom_id}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": custom_id, - "body": { - "id": f"chatcmpl-{custom_id}", - "object": "chat.completion", - "model": model, - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - }, - }, - }, - "error": None, - } - - -def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): - candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} - grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} - response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} - return { - "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, - "status": "", - "response": {**response, **({"modelVersion": model_version} if model_version else {})}, - "processed_time": "2026-09-23T19:02:00.000+00:00", - } - - -def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: - import litellm.cost_calculator as cc - - calls: list = [] - - def _calc(**kw): - calls.append(kw) - return (prompt_cost, completion_cost) - - monkeypatch.setattr(cc, "batch_cost_calculator", _calc) - return calls - - -def test_vertex_native_cost_bills_embedding_rows(monkeypatch): - monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) - rows = [ - { - "key": "id_1", - "status": "", - "request": {"content": {"parts": [{"text": "hello world"}]}}, - "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, - }, - { - "key": "id_2", - "status": "", - "request": {"content": {"parts": [{"text": "hello"}]}}, - "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, - }, - {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") - - assert (result.successful_requests, result.failed_requests) == (2, 1) - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) - assert result.cost == pytest.approx(5 * 1e-7) - assert result.models == ["gemini-embedding-2"] - - -@pytest.mark.asyncio -async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), - ] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert result.cost == pytest.approx(1.5) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.models == ["gemini-2.5-flash"] - assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} - - -@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) - monkeypatch.setattr( - bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") - ) - _capture_cost_calls(monkeypatch) - rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert result.successful_requests == 1 - - -@pytest.mark.asyncio -async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): - monkeypatch.setattr( - bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") - ) - _capture_cost_calls(monkeypatch) - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - custom_llm_provider="openai", - ) - - assert result.successful_requests == 0 - - -@pytest.mark.asyncio -async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) - raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] - - async def fake_fetch(batch, custom_llm_provider, litellm_params=None): - return _vertex_jsonl(raw_rows) - - monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - - result = await bu._handle_completed_batch( - _batch(PASSTHROUGH_OUTPUT_URI), - custom_llm_provider="vertex_ai", - model_name="gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert result.cost == pytest.approx(1.0) - assert result.usage.total_tokens == 557 - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) - ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) - - result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") - - grounded_usage, ungrounded_usage = (call["usage"] for call in calls) - assert grounded_usage.prompt_tokens == 19 - assert grounded_usage.completion_tokens == 59 + 406 - assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 - assert ungrounded_usage.prompt_tokens == 20 + 73 - assert ungrounded_usage.completion_tokens == 48 + 195 - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( - 19 + 93, - 465 + 243, - 557 + 336, - ) - - -def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) - - assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] - assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] - assert result.cost == pytest.approx(1.5) - assert result.successful_requests == 3 - assert result.usage.total_tokens == 557 + 336 + 336 - - -def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): - _capture_cost_calls(monkeypatch) - rows = [ - {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, - {"request": {"contents": []}, "response": {"candidates": []}}, - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert (result.successful_requests, result.failed_requests) == (1, 2) - assert result.usage.total_tokens == 557 - - -def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert result.models == ["gemini-2.5-flash"] - assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) - assert calls == [] - - -def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - - bu.calculate_vertex_ai_batch_cost_and_usage( - [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - "gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -@pytest.mark.asyncio -async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - deployment_model_info = {"input_cost_per_token_batches": 1e-6} - - await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - custom_llm_provider="vertex_ai", - model_name="gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert [call["model"] for call in calls] == ["gemini-2.5-flash"] - assert result.models == ["gemini-2.5-flash"] - - -def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [ - {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert (result.successful_requests, result.failed_requests) == (1, 1) - assert result.usage.total_tokens == 557 - assert len(calls) == 1 - - -@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) -def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) - - assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] - assert result.cost == pytest.approx(1.5) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.usage.total_tokens == 557 + 336 - - -def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") - without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) - - twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) - both = bu.calculate_vertex_ai_batch_cost_and_usage( - [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info - ) - - assert twin.cost > 0 - assert both.cost == pytest.approx(2 * twin.cost) - assert (both.successful_requests, both.failed_requests) == (2, 0) - - -def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): - import litellm.cost_calculator as cc - - def _calc(**kw): - if kw["model"] == "gemini-unpriced": - raise ValueError("no pricing") - return (0.5, 0.25) - - monkeypatch.setattr(cc, "batch_cost_calculator", _calc) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) - - assert result.cost == pytest.approx(0.75) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.usage.total_tokens == 557 + 336 - assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] - - -@pytest.mark.asyncio -async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch) - rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert calls == [] - assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/test_litellm/chat_completions/test_dispatch.py deleted file mode 100644 index ddb6e827309..00000000000 --- a/tests/test_litellm/chat_completions/test_dispatch.py +++ /dev/null @@ -1,117 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest - -import litellm -from litellm.chat_completions import dispatch -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules -from litellm.rust_bridge.chat_completions.entrypoints import ( - LiteLLMChatCompletionsRequest, - NativeAcompletion, - NativeCompletion, -) -from litellm.rust_bridge.configuration import Rollout -from litellm.types.utils import ModelResponse - -MESSAGES: Final = [{"role": "user", "content": "hi"}] - - -@pytest.mark.asyncio -async def test_public_completion_calls_keep_the_python_result() -> None: - sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - - assert isinstance(sync_response, ModelResponse) - assert isinstance(async_response, ModelResponse) - assert sync_response.choices[0].message.content == "ok" - assert async_response.choices[0].message.content == "ok" - - -def test_sync_completion_request_projects_public_arguments() -> None: - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) - expected: Final = ModelResponse() - - def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - assert request.model == "test-model" - assert request.messages == MESSAGES - assert request.custom_llm_provider == "openai" - assert request.stream is True - return expected - - binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {"custom_llm_provider": "openai", "stream": True}, - python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -@pytest.mark.asyncio -async def test_async_completion_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = ModelResponse() - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),) - - async def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_internal_acompletion_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) - expected: Final = ModelResponse() - - def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - pytest.fail("acompletion's inner completion call must stay on Python") - - binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {"custom_llm_provider": "openai", "acompletion": True}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index beca10d5555..f8c7d5273d1 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -14,6 +14,7 @@ from pathlib import Path from types import SimpleNamespace import httpx import pytest +from pytest_socket import _remove_restrictions import asyncio @@ -509,6 +510,14 @@ def setup_and_teardown(): print(f"[conftest] Module teardown complete (worker: {worker_id or 'master'})") +def pytest_collectstart(): + _remove_restrictions() + + +def pytest_runtest_setup(): + _remove_restrictions() + + def pytest_collection_modifyitems(config, items): """ Customize test collection order. diff --git a/tests/test_litellm/interactions/test_litellm_responses_bridge.py b/tests/test_litellm/interactions/test_litellm_responses_bridge.py index 8400f2c4840..17e7f9fc4ff 100644 --- a/tests/test_litellm/interactions/test_litellm_responses_bridge.py +++ b/tests/test_litellm/interactions/test_litellm_responses_bridge.py @@ -7,10 +7,6 @@ the litellm_responses bridge provider, which calls litellm.responses() internall import os -from litellm.interactions.litellm_responses_transformation.transformation import ( - LiteLLMResponsesInteractionsConfig, -) -from litellm.types.interactions import Turn from tests.test_litellm.interactions.base_interactions_test import ( BaseInteractionsTest, ) @@ -30,71 +26,3 @@ class TestLiteLLMResponsesBridge(BaseInteractionsTest): def get_api_key(self) -> str: """Return the OpenAI API key from environment.""" return os.getenv("OPENAI_API_KEY", "") - - -class TestBridgeInputTransformation: - """Regression tests for translating Interactions input into Responses API input. - - The bridge used to pass Google content parts through raw ({"type": "text"}), - which the Responses API rejects with a 400, and it dropped the role encoded - in step types and in the legacy "model" turn role. - """ - - def test_step_input_maps_roles_and_content_types(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [ - {"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]}, - {"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]}, - {"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]}, - ] - ) - assert transformed == [ - {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, - {"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]}, - ] - - def test_legacy_turn_input_maps_model_role_to_assistant(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [ - {"role": "user", "content": [{"type": "text", "text": "I like apples."}]}, - {"role": "model", "content": [{"type": "text", "text": "I like oranges."}]}, - ] - ) - assert transformed == [ - {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, - ] - - def test_turn_pydantic_model_with_string_content(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [Turn(role="model", content="I like oranges.")] - ) - assert transformed == [ - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]} - ] - - def test_string_input_passes_through(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello") - assert transformed == "Hello" - - def test_content_list_input_becomes_single_user_message(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [{"type": "text", "text": "Hello"}, "world"] - ) - assert transformed == [ - { - "role": "user", - "content": [ - {"type": "input_text", "text": "Hello"}, - {"type": "input_text", "text": "world"}, - ], - } - ] - - def test_non_text_content_passes_through_unchanged(self): - image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"} - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [{"type": "user_input", "content": [image_part]}] - ) - assert transformed == [{"role": "user", "content": [image_part]}] diff --git a/tests/test_litellm/messages/__init__.py b/tests/test_litellm/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/test_litellm/messages/test_dispatch.py deleted file mode 100644 index 4da060f809a..00000000000 --- a/tests/test_litellm/messages/test_dispatch.py +++ /dev/null @@ -1,155 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest -from pydantic import TypeAdapter - -import litellm -from litellm.messages import dispatch -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.messages.entrypoints import ( - LiteLLMMessagesRequest, - NativeAmessages, - NativeMessages, -) -from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse - -MESSAGES: Final = [{"role": "user", "content": "hi"}] - - -@pytest.mark.asyncio -async def test_public_anthropic_messages_keeps_the_python_result() -> None: - response: Final = await litellm.anthropic_messages( - model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok" - ) - - assert isinstance(response, dict) - content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", [])) - assert content[0]["text"] == "ok" - - -def test_sync_messages_request_projects_public_arguments() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - assert request.model == "claude-test" - assert request.messages == MESSAGES - assert request.max_tokens == 10 - assert request.custom_llm_provider == "anthropic" - return expected - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - { - "model": "claude-test", - "messages": MESSAGES, - "max_tokens": 10, - "custom_llm_provider": "anthropic", - }, - python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_messages_binding_error_delegates_unchanged_to_python() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - pytest.fail("a call without max_tokens cannot project a request and must stay on Python") - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - {"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -@pytest.mark.asyncio -async def test_async_messages_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = AnthropicMessagesResponse(model="claude-test") - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),) - - async def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - {"model": "claude-test", "messages": MESSAGES, "max_tokens": 10}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_internal_is_async_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - pytest.fail("anthropic_messages' inner handler call must stay on Python") - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - { - "model": "claude-test", - "messages": MESSAGES, - "max_tokens": 10, - "custom_llm_provider": "anthropic", - "is_async": True, - }, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected diff --git a/tests/test_litellm/rag/__init__.py b/tests/test_litellm/rag/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/rag/ingestion/__init__.py b/tests/test_litellm/rag/ingestion/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/rerank_api/__init__.py b/tests/test_litellm/rerank_api/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 16c641b8d29..e4b8860a7a6 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -2823,7 +2823,7 @@ def test_update_router_config_schema_includes_tag_routing_prefix(): # UpdateRouterConfig before calling update_settings; a field missing here # causes model_dump(exclude_none=True) to silently drop it before # update_settings is ever called -- the same bug shape LIT-3152 fixed for - # retry_policy (see tests/test_litellm/test_router_retry_policy_update.py). + # retry_policy (see tests/unit/test_router_retry_policy_update.py). from litellm.types.router import UpdateRouterConfig config = UpdateRouterConfig(tag_routing_prefix="route:") diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py index 4fbcd4ed30d..997778d0a1b 100644 --- a/tests/test_litellm/test_compression.py +++ b/tests/test_litellm/test_compression.py @@ -3,20 +3,13 @@ Unit tests for litellm.compress(). """ import os -import importlib import pytest import litellm -from litellm.compression.scoring.bm25 import bm25_score_messages -from litellm.compression.scoring.embedding_scorer import embedding_score_messages -from litellm.compression.content_detection import detect_content_type -from litellm.compression.message_stubbing import extract_key, stub_message -from litellm.compression.retrieval_tool import build_retrieval_tool from litellm.types.utils import CallTypes CALL_TYPE = CallTypes.completion -ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- @@ -24,420 +17,26 @@ ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- -def test_bm25_relevance_ranking(): - query = "Fix the authentication bug in the login handler" - messages = [ - { - "role": "user", - "content": "def login_handler(): authentication check bug fix", - }, - {"role": "user", "content": "def render_template(name): css styling layout"}, - {"role": "user", "content": "def verify(): authentication token bug handler"}, - ] - scores = bm25_score_messages(query, messages) - # Messages sharing query terms should score higher than unrelated ones - assert scores[0] > scores[1] - assert scores[2] > scores[1] - - -def test_bm25_empty_query(): - scores = bm25_score_messages("", [{"role": "user", "content": "hello"}]) - assert scores == [0.0] - - -def test_bm25_empty_messages(): - scores = bm25_score_messages("query", []) - assert scores == [] - - -def test_bm25_empty_content(): - scores = bm25_score_messages("query", [{"role": "user", "content": ""}]) - assert scores == [0.0] - - # --------------------------------------------------------------------------- # Content detection # --------------------------------------------------------------------------- -def test_detect_code(): - code = """ -import os -from pathlib import Path - -def main(): - class Foo: - pass - return Foo() -""" - assert detect_content_type(code) == "code" - - -def test_detect_json(): - assert detect_content_type('{"key": "value", "num": 42}') == "json" - assert detect_content_type("[1, 2, 3]") == "json" - - -def test_detect_text(): - assert detect_content_type("This is a plain text paragraph about dogs.") == "text" - - -def test_detect_empty(): - assert detect_content_type("") == "text" - - # --------------------------------------------------------------------------- # Message stubbing # --------------------------------------------------------------------------- -def test_extract_key_with_filename(): - msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"} - used: set = set() - key = extract_key(msg, fallback_index=0, used_keys=used) - assert key == "auth.py" - - -def test_extract_key_fallback(): - msg = {"role": "user", "content": "Some random content without a filename"} - used: set = set() - key = extract_key(msg, fallback_index=5, used_keys=used) - assert key == "message_5" - - -def test_extract_key_duplicates(): - used: set = set() - msg = {"role": "user", "content": "# auth.py\ncode here"} - k1 = extract_key(msg, fallback_index=0, used_keys=used) - k2 = extract_key(msg, fallback_index=1, used_keys=used) - assert k1 == "auth.py" - assert k2 == "auth.py_2" - - -def test_stub_message(): - msg = {"role": "user", "content": "line1\nline2\nline3"} - stubbed = stub_message(msg, "test_key") - assert stubbed["role"] == "user" - assert "test_key" in stubbed["content"] - assert "litellm_content_retrieve" in stubbed["content"] - assert "3 lines" in stubbed["content"] - - # --------------------------------------------------------------------------- # Retrieval tool # --------------------------------------------------------------------------- -def test_retrieval_tool_schema(): - tool = build_retrieval_tool(["auth.py", "utils.py"]) - assert tool["type"] == "function" - assert tool["function"]["name"] == "litellm_content_retrieve" - assert "key" in tool["function"]["parameters"]["properties"] - assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [ - "auth.py", - "utils.py", - ] - assert tool["function"]["parameters"]["required"] == ["key"] - - -def test_retrieval_tool_description_lists_keys(): - tool = build_retrieval_tool(["foo.py", "bar.js"]) - desc = tool["function"]["description"] - assert "foo.py" in desc - assert "bar.js" in desc - - # --------------------------------------------------------------------------- # compress() — end-to-end # --------------------------------------------------------------------------- -def test_compress_below_trigger_passthrough(): - messages = [{"role": "user", "content": "hello"}] - result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) - assert result["messages"] == messages - assert result["cache"] == {} - assert result["tools"] == [] - assert result["compression_ratio"] == 0.0 - assert result["original_tokens"] == result["compressed_tokens"] - - -def test_compress_above_trigger(): - big_messages = [ - {"role": "system", "content": "You are a coding assistant."}, - { - "role": "user", - "content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000, - }, - { - "role": "user", - "content": "# utils.py\n" + "def helper():\n pass\n" * 2000, - }, - { - "role": "user", - "content": "# readme.md\n" + "This is documentation. " * 2000, - }, - {"role": "user", "content": "Fix the bug in auth.py"}, - ] - - result = litellm.compress( - big_messages, - model="gpt-4o", - call_type=CALL_TYPE, - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] < result["original_tokens"] - assert result["compression_ratio"] > 0 - assert len(result["cache"]) > 0 - assert len(result["tools"]) == 1 - assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" - - -def test_compress_anthropic_list_content_is_boundary_stable(): - messages = [ - {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, - { - "role": "user", - "content": [ - {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, - { - "type": "image_url", - "image_url": {"url": "https://example.com/a.png"}, - }, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "# b.py\n" + "beta " * 2000}, - { - "type": "image_url", - "image_url": {"url": "https://example.com/b.png"}, - }, - ], - }, - { - "role": "user", - "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], - }, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] < result["original_tokens"] - assert len(result["messages"]) == len(messages) - assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] - assert len(result["cache"]) > 0 - assert len(result["tools"]) == 1 - assert result["tools"][0]["type"] == "custom" - assert result["tools"][0]["name"] == "litellm_content_retrieve" - assert "input_schema" in result["tools"][0] - - -def test_compress_preserves_system_message(): - messages = [ - {"role": "system", "content": "System prompt. " * 500}, - {"role": "user", "content": "Large file content. " * 5000}, - {"role": "user", "content": "Fix the bug"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - assert result["messages"][0]["role"] == "system" - assert "System prompt" in result["messages"][0]["content"] - - -def test_compress_preserves_last_user_message(): - messages = [ - {"role": "user", "content": "Big context " * 5000}, - {"role": "user", "content": "Fix the bug in auth.py"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - last_user = [m for m in result["messages"] if m["role"] == "user"][-1] - assert "Fix the bug in auth.py" in last_user["content"] - - -def test_compress_preserves_last_assistant_message(): - messages = [ - {"role": "user", "content": "Big context " * 5000}, - {"role": "assistant", "content": "I'll help with that. " * 2000}, - {"role": "user", "content": "Now fix the bug"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] - assert len(assistant_msgs) >= 1 - # The last assistant message should be preserved (not stubbed) - last_assistant = assistant_msgs[-1] - assert "I'll help with that" in last_assistant["content"] - - -def test_cache_keys_match_stubs(): - messages = [ - {"role": "user", "content": "# auth.py\n" + "code " * 5000}, - {"role": "user", "content": "Fix it"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - if result["tools"]: - tool_desc = result["tools"][0]["function"]["description"] - for key in result["cache"]: - assert key in tool_desc - - -def test_compress_default_target(): - """compression_target defaults to compression_trigger // 2.""" - messages = [ - {"role": "user", "content": "content " * 5000}, - {"role": "user", "content": "query"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 - ) - # Should have compressed — target = 1000 - assert result["compressed_tokens"] <= result["original_tokens"] - - -def test_compress_nested_tool_result_extracts_text_only(): - messages = [ - {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, - { - "role": "user", - "content": [ - {"type": "text", "text": "prefix"}, - { - "type": "tool_result", - "tool_use_id": "toolu_1", - "content": [ - {"type": "text", "text": "nested text fragment"}, - { - "type": "image_url", - "image_url": { - "url": "https://example.com/secret-tool.png", - }, - }, - ], - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/top.png"}, - }, - {"type": "text", "text": " " + ("irrelevant " * 3000)}, - ], - }, - { - "role": "user", - "content": [{"type": "text", "text": "final query that must remain"}], - }, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=500, - compression_target=100, - ) - - cached_text = " ".join(result["cache"].values()) - assert "nested text fragment" in cached_text - assert "https://example.com/secret-tool.png" not in cached_text - assert "https://example.com/top.png" not in cached_text - - -def test_compress_default_call_type_is_completion(): - result = litellm.compress( - messages=[ - {"role": "user", "content": "Large context " * 4000}, - {"role": "user", "content": "query"}, - ], - model="gpt-4o", - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] <= result["original_tokens"] - assert isinstance(result["tools"], list) - - -def test_compress_forwards_embedding_model_params(monkeypatch): - captured = {} - - def fake_embedding_score_messages( - query, messages, model, cache=None, embedding_model_params=None - ): - captured["query"] = query - captured["model"] = model - captured["embedding_model_params"] = embedding_model_params - return [0.0] * len(messages) - - monkeypatch.setattr( - "litellm.compression.scoring.embedding_scorer.embedding_score_messages", - fake_embedding_score_messages, - ) - - result = litellm.compress( - messages=[ - {"role": "user", "content": "Authentication code " * 2000}, - {"role": "user", "content": "Fix auth"}, - ], - model="gpt-4o", - call_type=CALL_TYPE, - compression_trigger=1000, - embedding_model="text-embedding-3-small", - embedding_model_params={"api_base": "https://example-embeddings.test"}, - ) - - assert result["compressed_tokens"] <= result["original_tokens"] - assert captured["model"] == "text-embedding-3-small" - assert captured["embedding_model_params"] == { - "api_base": "https://example-embeddings.test" - } - - -def test_embedding_scorer_forwards_embedding_model_params(monkeypatch): - captured = {} - - class _MockResponse: - data = [ - {"embedding": [1.0, 0.0]}, - {"embedding": [1.0, 0.0]}, - {"embedding": [0.0, 1.0]}, - ] - - def fake_embedding(**kwargs): - captured.update(kwargs) - return _MockResponse() - - monkeypatch.setattr(litellm, "embedding", fake_embedding) - - scores = embedding_score_messages( - query="auth", - messages=[ - {"role": "user", "content": "auth code"}, - {"role": "user", "content": "cooking recipe"}, - ], - model="text-embedding-3-small", - embedding_model_params={"api_base": "https://example-embeddings.test"}, - ) - - assert len(scores) == 2 - assert captured["model"] == "text-embedding-3-small" - assert captured["api_base"] == "https://example-embeddings.test" - - # --------------------------------------------------------------------------- # Embedding scorer — integration test (skipped without API key) # --------------------------------------------------------------------------- @@ -458,210 +57,3 @@ def test_embedding_scorer(): ) assert result["compression_ratio"] > 0 assert len(result["cache"]) > 0 - - -@pytest.mark.parametrize( - "final_user_message, expected_content", - [ - ("How to cook?", "Unrelated cooking recipes "), - ("Fix auth", "Authentication code "), - ], -) -def test_simple_compression(final_user_message, expected_content): - messages = [ - {"role": "user", "content": "Authentication code " * 2000}, - {"role": "user", "content": "Unrelated cooking recipes " * 2000}, - {"role": "user", "content": final_user_message}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - if expected_content == "Unrelated cooking recipes ": - assert "Unrelated cooking recipes " in result["messages"][1]["content"] - assert "Authentication code " not in result["messages"][0]["content"] - elif expected_content == "Authentication code ": - assert "Authentication code " in result["messages"][0]["content"] - assert "Unrelated cooking recipes " not in result["messages"][1]["content"] - else: - raise ValueError(f"Unexpected expected_content: {expected_content}") - - -def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): - compress_module = importlib.import_module("litellm.compression.compress") - - def fake_bm25_score_messages(query, messages): - assert "final query" in query - assert len(messages) == 5 - # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) - return [0.95, 0.01, 0.02, 0.8, 1.0] - - def fake_token_counter(model, messages=None, text=None): - if messages is not None: - return 1000 - if text is None: - return 0 - if "final query" in text: - return 50 - if "assistant_tail" in text: - return 20 - if "other_blob" in text: - return 220 - if "tool_payload_relevant" in text: - return 200 - if text == "": - return 1 - return 10 - - monkeypatch.setattr( - compress_module, "bm25_score_messages", fake_bm25_score_messages - ) - monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) - - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_drop", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_drop", - "content": [{"type": "text", "text": "tool_payload_relevant"}], - } - ], - }, - {"role": "assistant", "content": "assistant_tail"}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) - assert len(result["messages"]) == 3 - assert result["messages"][0]["role"] == "user" - assert "other_blob" in result["messages"][0]["content"] - assert result["messages"][1]["content"] == "assistant_tail" - assert result["messages"][2]["content"] == "final query" - assert result["cache"] == {} - - -def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): - compress_module = importlib.import_module("litellm.compression.compress") - - def fake_bm25_score_messages(query, messages): - assert "final query" in query - assert len(messages) == 5 - # Prefer the tool exchange span over idx=0 - return [0.05, 0.01, 0.92, 0.8, 1.0] - - def fake_token_counter(model, messages=None, text=None): - if messages is not None: - return 1000 - if text is None: - return 0 - if "final query" in text: - return 50 - if "assistant_tail" in text: - return 20 - if "other_blob" in text: - return 220 - if "tool_payload_relevant" in text: - return 200 - if text == "": - return 1 - return 10 - - monkeypatch.setattr( - compress_module, "bm25_score_messages", fake_bm25_score_messages - ) - monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) - - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_keep", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_keep", - "content": [{"type": "text", "text": "tool_payload_relevant"}], - } - ], - }, - {"role": "assistant", "content": "assistant_tail"}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - assert len(result["messages"]) == 5 - assert result["messages"][1]["role"] == "assistant" - assert result["messages"][2]["role"] == "user" - # idx=0 should be compressed instead - assert "litellm_content_retrieve" in result["messages"][0]["content"] - assert len(result["cache"]) == 1 - - -def test_compress_anthropic_malformed_tool_sequence_passes_through(): - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_broken", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - assert result["messages"] == messages - assert result["cache"] == {} - assert result["tools"] == [] - assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 227fb48bb08..78728d6fd58 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1,31 +1,12 @@ -import asyncio -import base64 -from datetime import datetime -import contextlib -import copy import json -import logging import os -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Final -import httpx import pytest -import respx -from fastapi.testclient import TestClient -import urllib.parse -from importlib import import_module from unittest.mock import MagicMock, patch import litellm -from litellm import main as litellm_main -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage async def _async_fake_bedrock_image_details(image_url): @@ -61,111 +42,6 @@ def add_api_keys_to_env(monkeypatch): monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) -@pytest.fixture -def openai_api_response(): - mock_response_data = { - "id": "chatcmpl-B0W3vmiM78Xkgx7kI7dr7PC949DMS", - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "content": "", - "refusal": None, - "role": "assistant", - "audio": None, - "function_call": None, - "tool_calls": None, - }, - } - ], - "created": 1739462947, - "model": "gpt-4o-mini-2024-07-18", - "object": "chat.completion", - "service_tier": "default", - "system_fingerprint": "fp_bd83329f63", - "usage": { - "completion_tokens": 1, - "prompt_tokens": 121, - "total_tokens": 122, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, - }, - } - - return mock_response_data - - -def test_completion_missing_role(openai_api_response): - from openai import OpenAI - - from litellm.types.utils import ModelResponse - - client = OpenAI(api_key="test_api_key") - - mock_raw_response = MagicMock() - mock_raw_response.headers = { - "x-request-id": "123", - "openai-organization": "org-123", - "x-ratelimit-limit-requests": "100", - "x-ratelimit-remaining-requests": "99", - } - mock_raw_response.parse.return_value = ModelResponse(**openai_api_response) - - print(f"openai_api_response: {openai_api_response}") - - with patch.object( - client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) - ) as mock_create: - litellm.completion( - model="gpt-4o-mini", - messages=[ - {"role": "user", "content": "Hey"}, - { - "content": "", - "tool_calls": [ - { - "id": "call_m0vFJjQmTH1McvaHBPR2YFwY", - "function": { - "arguments": '{"input": "dksjsdkjdhskdjshdskhjkhlk"}', - "name": "tool_name", - }, - "type": "function", - "index": 0, - }, - { - "id": "call_Vw6RaqV2n5aaANXEdp5pYxo2", - "function": { - "arguments": '{"input": "jkljlkjlkjlkjlk"}', - "name": "tool_name", - }, - "type": "function", - "index": 1, - }, - { - "id": "call_hBIKwldUEGlNh6NlSXil62K4", - "function": { - "arguments": '{"input": "jkjlkjlkjlkj;lj"}', - "name": "tool_name", - }, - "type": "function", - "index": 2, - }, - ], - }, - ], - client=client, - ) - - mock_create.assert_called_once() - - @pytest.mark.parametrize( "model", [ @@ -277,210 +153,6 @@ async def test_url_with_format_param(model, sync_mode, monkeypatch): assert "jpeg" not in json_str -@pytest.mark.parametrize("model", ["gpt-4o-mini"]) -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_url_with_format_param_openai(model, sync_mode): - from openai import AsyncOpenAI, OpenAI - - from litellm import acompletion, completion - - if sync_mode: - client = OpenAI() - else: - client = AsyncOpenAI() - - args = { - "model": model, - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", - "format": "image/png", - }, - }, - {"type": "text", "text": "Describe this image"}, - ], - } - ], - } - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - try: - if sync_mode: - response = completion(**args, client=client) - else: - response = await acompletion(**args, client=client) - print(response) - except Exception as e: - print(e) - - mock_client.assert_called() - - print(mock_client.call_args.kwargs) - - json_str = json.dumps(mock_client.call_args.kwargs) - - assert "format" not in json_str - - -def test_bedrock_latency_optimized_inference(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - client = HTTPHandler() - with patch.object(client, "post") as mock_post: - try: - response = litellm.completion( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[{"role": "user", "content": "Hello, how are you?"}], - performanceConfig={"latency": "optimized"}, - client=client, - ) - except Exception as e: - print(e) - - mock_post.assert_called_once() - json_data = json.loads(mock_post.call_args.kwargs["data"]) - assert json_data["performanceConfig"]["latency"] == "optimized" - - -@pytest.mark.parametrize( - ("custom_llm_provider", "model", "expected"), - [ - ("anthropic", "claude-sonnet-5", True), - ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), - ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), - ("bedrock", "us.amazon.nova-2-lite-v1:0", False), - ("vertex_ai", "claude-sonnet-5", True), - ("vertex_ai", "gemini-3.8-flash", False), - ("azure_ai", "claude-sonnet-4-6", True), - ("azure_ai", "gpt-5.6", False), - ("openai", "gpt-5.6", False), - ("gemini", "gemini-3.8-flash", False), - ], -) -def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): - assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected - - -@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) -def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): - tools = [ - {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, - "opaque_tool", - ] - - cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) - - assert cleaned == [ - {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, - "opaque_tool", - ] - assert tools[0][key] is True - assert tools[0]["function"][key] is True - - -def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): - api_base: Final = "http://localhost:12346/v1" - mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( - return_value=httpx.Response(status_code=200, json=openai_api_response) - ) - - litellm.completion( - model="openai/gpt-5.6", - messages=[{"role": "user", "content": "Write the file"}], - tools=[ - { - "type": "function", - "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, - "eager_input_streaming": True, - } - ], - api_base=api_base, - api_key="fake_openai_api_key", - ) - - assert mock_route.called - sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] - assert "eager_input_streaming" not in sent_tool - assert sent_tool["function"]["name"] == "write_file" - - -def test_custom_provider_with_extra_headers(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - headers={"X-Custom-Header": "custom-value"}, - api_base="https://example.com/api/v1", - ) - - mock_post.assert_called_once() - assert mock_post.call_args[1]["headers"]["X-Custom-Header"] == "custom-value" - - -def test_custom_provider_with_extra_body(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - extra_body={ - "X-Custom-BodyValue": "custom-value", - "X-Custom-BodyValue2": "custom-value2", - }, - api_base="https://example.com/api/v1", - ) - mock_post.assert_called_once() - - assert mock_post.call_args[1]["json"]["X-Custom-BodyValue"] == "custom-value" - assert mock_post.call_args[1]["json"] == { - "model": "custom", - "params": { - "prompt": ["Hello, how are you?"], - "max_tokens": None, - "temperature": None, - "top_p": None, - "top_k": None, - }, - "X-Custom-BodyValue": "custom-value", - "X-Custom-BodyValue2": "custom-value2", - } - - # test that extra_body is not passed if not provided - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - api_base="https://example.com/api/v1", - ) - mock_post.assert_called_once() - assert mock_post.call_args[1]["json"] == { - "model": "custom", - "params": { - "prompt": ["Hello, how are you?"], - "max_tokens": None, - "temperature": None, - "top_p": None, - "top_k": None, - }, - } - - @pytest.fixture(autouse=True) def set_openrouter_api_key(): original_api_key = os.environ.get("OPENROUTER_API_KEY") @@ -490,3753 +162,3 @@ def set_openrouter_api_key(): os.environ["OPENROUTER_API_KEY"] = original_api_key else: del os.environ["OPENROUTER_API_KEY"] - - -@pytest.mark.asyncio -async def test_extra_body_with_fallback( - respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch -): - """ - test regression for https://github.com/BerriAI/litellm/issues/8425. - - This was perhaps a wider issue with the acompletion function not passing kwargs such as extra_body correctly when fallbacks are specified. - """ - - # Save original state to restore after test - original_disable_aiohttp = litellm.disable_aiohttp_transport - - try: - # since this uses respx, we need to set use_aiohttp_transport to False - # Set both the global variable and environment variable to ensure it takes effect - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - # Flush cache to ensure no stale aiohttp clients are used - litellm.in_memory_llm_clients_cache.flush_cache() - - # Set up test parameters - model = "openrouter/deepseek/deepseek-chat" - messages = [{"role": "user", "content": "Hello, world!"}] - extra_body = { - "provider": { - "order": ["DeepSeek"], - "allow_fallbacks": False, - "require_parameters": True, - } - } - fallbacks = [{"model": "openrouter/google/gemini-flash-1.5-8b"}] - - # Set up mock to respond to any POST request to the OpenRouter endpoint - # This ensures it works for both primary and fallback models - mock_route = respx_mock.post("https://openrouter.ai/api/v1/chat/completions") - mock_route.return_value = httpx.Response( - 200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - - response = await litellm.acompletion( - model=model, - messages=messages, - extra_body=extra_body, - fallbacks=fallbacks, - api_key="fake-openrouter-api-key", - ) - - # Verify the response - assert response is not None - assert ( - len(respx_mock.calls) > 0 - ), "Mock was not called - check if aiohttp transport is properly disabled" - - # Get the request from the mock - request: httpx.Request = respx_mock.calls[0].request - request_body = request.read() - request_body = json.loads(request_body) - - # Verify basic parameters - assert request_body["model"] == "deepseek/deepseek-chat" - assert request_body["messages"] == messages - - # Verify the extra_body parameters remain under the provider key - assert request_body["provider"]["order"] == ["DeepSeek"] - assert request_body["provider"]["allow_fallbacks"] is False - assert request_body["provider"]["require_parameters"] is True - finally: - # Restore original state to prevent test pollution - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -@pytest.mark.parametrize("env_base", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_openai_env_base( - respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch -): - "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" - # Ensure aiohttp transport is disabled to use httpx which respx can mock - litellm.disable_aiohttp_transport = True - - expected_base_url = "http://localhost:12345/v1" - - # Assign the environment variable based on env_base, and use a fake API key. - monkeypatch.setenv(env_base, expected_base_url) - monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") - - model = "gpt-4o" - messages = [{"role": "user", "content": "Hello, how are you?"}] - - # Configure respx mock to intercept the request - mock_route = respx_mock.post( - url__regex=r"http://localhost:12345/v1/chat/completions.*" - ).mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - ) - - try: - response = await litellm.acompletion(model=model, messages=messages) - - # verify we had a response - assert response.choices[0].message.content == "Hello from mocked response!" - - # Verify the mock was called - assert ( - mock_route.called - ), "Mock route was not called - request may have bypassed respx" - finally: - # Clean up to avoid affecting other tests - litellm.disable_aiohttp_transport = False - - -def build_database_url(username, password, host, dbname): - username_enc = urllib.parse.quote_plus(username) - password_enc = urllib.parse.quote_plus(password) - dbname_enc = urllib.parse.quote_plus(dbname) - return f"postgresql://{username_enc}:{password_enc}@{host}/{dbname_enc}" - - -def test_build_database_url(): - url = build_database_url("user@name", "p@ss:word", "localhost", "db/name") - assert url == "postgresql://user%40name:p%40ss%3Aword@localhost/db%2Fname" - - -def test_bedrock_llama(): - litellm._turn_on_debug() - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "bedrock/invoke/us.meta.llama4-scout-17b-instruct-v1:0" - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": [ - {"role": "user", "content": "hi"}, - ], - }, - ) - print(request) - - assert ( - request["raw_request_body"]["prompt"] - == "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" - ) - - -def _mocked_openai_chat_response(model: str) -> httpx.Response: - return httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - - -def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter): - """Regression for #33952: return_raw_request must transform without contacting the provider. - - Previously return_raw_request invoked the real endpoint with a fake key and relied on the - provider rejecting it, which sent an unintended inference request and (in the async proxy - route) blocked the event loop on provider I/O. - """ - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "gpt-4o" - route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": [{"role": "user", "content": "hi"}], - }, - ) - - assert route.call_count == 0 - assert request.get("error") is None - assert request["raw_request_body"]["model"] == model - assert request["raw_request_body"]["messages"] == [ - {"role": "user", "content": "hi"} - ] - - -def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter): - """Regression test: completion() must forward the verbosity param to the provider request body.""" - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "gpt-5.2" - messages = [{"role": "user", "content": "hi"}] - respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": messages, - "verbosity": "high", - }, - ) - - assert request["raw_request_body"]["verbosity"] == "high" - assert request["raw_request_body"]["model"] == model - assert request["raw_request_body"]["messages"] == messages - - -@pytest.mark.asyncio -async def test_acompletion_forwards_verbosity_to_provider_request( - respx_mock: respx.MockRouter, monkeypatch -): - """Regression test: acompletion() must forward the verbosity param to the provider request body.""" - original_disable_aiohttp = litellm.disable_aiohttp_transport - try: - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - - model = "gpt-5.2" - messages = [{"role": "user", "content": "hi"}] - mock_route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - response = await litellm.acompletion( - model=model, - messages=messages, - verbosity="low", - api_key="fake-openai-api-key", - ) - - assert response.choices[0].message.content == "Hello from mocked response!" - assert mock_route.called - request_body = json.loads(respx_mock.calls[0].request.read()) - assert request_body["verbosity"] == "low" - assert request_body["model"] == model - assert request_body["messages"] == messages - finally: - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -def test_responses_api_bridge_check_strips_responses_prefix(): - """Test that responses_api_bridge_check strips 'responses/' prefix and sets mode.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - - model_info, model = responses_api_bridge_check( - model="responses/gpt-4-responses", - custom_llm_provider="openai", - ) - - assert model == "gpt-4-responses" - assert model_info["mode"] == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_pro(): - """Test that gpt-5.4-pro routes through responses API bridge, not chat completions. - - Regression test for https://github.com/BerriAI/litellm/issues/23014 - gpt-5.4-pro is a responses-only model and must not be sent to /v1/chat/completions. - """ - from litellm.main import responses_api_bridge_check - - for model_name in ["gpt-5.4-pro", "gpt-5.4-pro-2026-03-05"]: - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="openai", - ) - assert ( - model_info.get("mode") == "responses" - ), f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" - - -def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_responses(): - """gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="xhigh", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-6-astra", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - ) - - assert model == "gpt-6-astra" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): - """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.5-pro", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="xhigh", - ) - - assert model == "gpt-5.5-pro" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to_responses(): - """Azure gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="high", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables - reasoning by default for gpt-5.4+, and Chat Completions rejects function tools - whenever reasoning is on. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning - by default for gpt-5.4+, and Chat Completions rejects function tools whenever - reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "model_name, expected_mode", - [ - pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), - pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), - ], -) -def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( - monkeypatch, model_name, expected_mode -): - """ - gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit - reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with - "Function tools with reasoning_effort are not supported for gpt-5.6-sol in - /v1/chat/completions". - - Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The - gate parses the version and drops any suffix, so the family members bridge - identically and only the boundary distinguishes behaviour. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == model_name - assert model_info.get("mode") == expected_mode - - -def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): - """ - Explicit reasoning_effort "none" is OpenAI's documented escape hatch that keeps - function tools servable on Chat Completions; the bridge must not fire. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="none", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_reasoning_none_with_summary_still_routes_to_responses(): - """A reasoning summary is Responses-only regardless of effort value.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - reasoning_effort="none", - reasoning_summary="detailed", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_custom_tools_only_stays_chat(): - """ - Chat Completions serves custom (grammar) tools natively with reasoning on; only - FUNCTION tools trigger the OpenAI rejection. Custom-only requests must stay on chat - so responses keep the native custom tool_call shape instead of the bridge's - function-shaped mapping. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "custom", "custom": {"name": "ApplyPatch", "description": "V4A patch"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_gpt_5_4_mixed_function_and_custom_tools_routes_to_responses(): - """One function tool in the mix is enough to make chat unservable with reasoning on.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[ - {"type": "custom", "custom": {"name": "ApplyPatch"}}, - {"type": "function", "function": {"name": "shell"}}, - ], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_flat_function_tool_routes_to_responses(): - """Responses-style flat function tool defs still count as function tools.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "name": "shell", "parameters": {"type": "object"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "custom_llm_provider, model_name, api_base", - [ - pytest.param("openai", "gpt-5.6", None, id="openai"), - pytest.param("azure_ai", "gpt-6-astra", "https://myproject.services.ai.azure.com", id="azure-ai-foundry"), - ], -) -def test_responses_api_bridge_check_function_tool_without_body_stays_chat( - monkeypatch, custom_llm_provider, model_name, api_base -): - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider=custom_llm_provider, - tools=[{"type": "function"}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_dict_effort_none_stays_chat(): - """The escape hatch must honor litellm's dict form: {"effort": "none"} means reasoning off.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "none"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_dict_effort_active_routes_to_responses(): - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "low"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_dict_effort_none_with_summary_routes_to_responses(): - """A summary inside the dict form is Responses-only even when effort is none.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "none", "summary": "concise"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize("blank_api_base", [None, "", " ", "\t"]) -def test_responses_api_bridge_check_blank_api_base_is_default_openai(blank_api_base): - """ - A blank api_base (None, empty, or whitespace) resolves to the default OpenAI - endpoint downstream, which enforces the reasoning+tools constraint, so gpt-5.4+ - function-tool requests with unset reasoning_effort must still auto-bridge. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=blank_api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_custom_api_base_with_unset_effort_stays_chat(): - """ - Chat-only OpenAI-compatible backends registered under the openai provider with a - custom api_base and gpt-5.4+ model names serve tools-without-reasoning fine and - have no /responses route; the unset-effort arm must not reroute them. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base="http://vllm.internal:8000/v1", - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_custom_api_base_via_global_with_unset_effort_stays_chat(monkeypatch): - """ - A custom base set through the litellm.api_base global (not the call arg) is resolved the - same way the chat handler resolves it, so the unset-effort arm must not reroute a chat-only - backend to a /responses route it lacks. Regression guard: the gate previously inspected only - the call-level api_base and bridged these requests. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", "http://vllm.internal:8000/v1") - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -@pytest.mark.parametrize("env_var", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) -def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_stays_chat(monkeypatch, env_var): - """ - A custom base set via OPENAI_BASE_URL/OPENAI_API_BASE env is resolved identically to the chat - handler, so the unset-effort arm leaves the request on chat instead of bridging it. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", None) - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setenv(env_var, "http://vllm.internal:8000/v1") - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -@pytest.mark.parametrize( - "api_base", - [ - "https://southcentralus.privatelink.api.openai.com/v1", - "https://privatelink.corp.api.openai.com/v1", - "https://api.openai.com:443/v1", - "https://api.openai.com/v1/", - "HTTPS://API.OPENAI.COM/v1", - ], -) -def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): - """ - A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a - port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects - function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as - it does for the literal default URL. Regression guard for GH #39353. - """ - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "api_base", - [ - "https://api.openai.com.evil.example/v1", - "https://notapi.openai.com/v1", - "https://gateway.example/v1?upstream=api.openai.com", - "https://openai.internal.example/api.openai.com/v1", - ], -) -def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): - """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): - """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", None) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): - """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="high", - api_base="http://vllm.internal:8000/v1", - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): - """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base="https://myresource.openai.azure.com", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -_FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" -_FOUNDRY_FUNCTION_TOOL: Final = ({"type": "function", "function": {"name": "get_weather"}},) - - -@pytest.mark.parametrize( - "model_name, api_base, reasoning_effort", - [ - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, None, id="gpt-6-unset-effort"), - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "low", id="gpt-6-explicit-effort"), - pytest.param("gpt-6-astra", "https://myresource.openai.azure.com", None, id="gpt-6-azure-openai-host"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "low", id="gpt-5.6-explicit-effort"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, {"effort": "high"}, id="gpt-5.6-explicit-effort-dict"), - ], -) -def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_responses( - model_name, api_base, reasoning_effort -): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="azure_ai", - tools=_FOUNDRY_FUNCTION_TOOL, - reasoning_effort=reasoning_effort, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "model_name, api_base, reasoning_effort", - [ - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="explicit-none-stays-chat"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, None, id="gpt-5.6-unset-effort-stays-chat"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "none", id="gpt-5.6-explicit-none-stays-chat"), - pytest.param("gpt-5.5", _FOUNDRY_API_BASE, "high", id="gpt-5.5-explicit-effort-stays-chat"), - pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, None, id="gpt-5.4-mini-unset-effort-stays-chat"), - pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, "low", id="gpt-5.4-mini-explicit-effort-stays-chat"), - pytest.param("gpt-6-astra", "https://myproject.models.ai.azure.com", None, id="serverless-host-stays-chat"), - pytest.param("Mistral-large-2411", _FOUNDRY_API_BASE, None, id="non-gpt-5-model-stays-chat"), - pytest.param("claude-opus-4-1", _FOUNDRY_API_BASE, None, id="claude-on-foundry-stays-chat"), - ], -) -def test_responses_api_bridge_check_azure_ai_without_foundry_responses_route_stays_chat( - model_name, api_base, reasoning_effort -): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="azure_ai", - tools=_FOUNDRY_FUNCTION_TOOL, - reasoning_effort=reasoning_effort, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_older_gpt_5_tools_without_reasoning_stays_chat(): - """Pre-5.4 GPT-5 names keep the old boundary: tools alone never bridge.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.1", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.1" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_gpt_5_4_reasoning_summary_without_tools_routes_to_responses(): - """gpt-5.4+ with reasoning_effort + reasoningSummary but no tools should bridge (AI SDK).""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=None, - reasoning_effort="medium", - reasoning_summary="auto", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_reasoning_summary_routes_to_responses(): - """Bare ``gpt-5`` with reasoning_effort + reasoningSummary should bridge (not 5.4+).""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5", - custom_llm_provider="openai", - tools=None, - reasoning_effort="medium", - reasoning_summary="auto", - ) - - assert model == "gpt-5" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_tools_without_summary_stays_chat(): - """gpt-5 with tools + reasoning_effort but no summary should stay on chat.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="medium", - reasoning_summary=None, - ) - - assert model == "gpt-5" - assert model_info.get("mode") != "responses" - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( - mock_responses_completion, -): - """When routed to Responses, preserve reasoning_effort summary dict.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5.4", - messages=[{"role": "user", "content": "What is the capital of France?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_capital", - "description": "Get the capital of a country", - "parameters": { - "type": "object", - "properties": {"country": {"type": "string"}}, - }, - }, - } - ], - reasoning_effort={"effort": "xhigh", "summary": "detailed"}, - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == { - "effort": "xhigh", - "summary": "detailed", - } - - -@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) -def test_responses_bridge_preserves_reasoning_effort_with_drop_params( - reasoning_effort, - restore_model_registry, - respx_mock: respx.MockRouter, - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - response_body: Final = { - "id": "resp_test", - "object": "response", - "created_at": 1734366691, - "status": "completed", - "model": "test-responses-bridge", - "output": [ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "Done.", "annotations": []}], - } - ], - "parallel_tool_calls": True, - "usage": { - "input_tokens": 1, - "output_tokens": 1, - "total_tokens": 2, - "output_tokens_details": {"reasoning_tokens": 0}, - }, - "error": None, - "incomplete_details": None, - "instructions": None, - "metadata": None, - "temperature": None, - "tool_choice": "auto", - "tools": [], - "top_p": None, - "max_output_tokens": None, - "previous_response_id": None, - "reasoning": None, - "truncation": None, - "user": None, - } - response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) - model: Final = "perplexity/test-responses-bridge" - litellm.register_model( - { - model: { - "litellm_provider": "perplexity", - "mode": "responses", - "supports_reasoning": False, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - } - }, - persist_across_reloads=False, - ) - - litellm.completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - reasoning_effort=reasoning_effort, - drop_params=True, - api_key="fake-key", - api_base="https://api.perplexity.ai", - ) - - request_body: Final = json.loads(response_route.calls[0].request.content) - assert request_body["reasoning"] == {"effort": "high"} - - -_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY: Final = { - "id": "resp_foundry", - "object": "response", - "created_at": 1789852145, - "status": "completed", - "model": "gpt-6-astra", - "output": [ - { - "id": "fc_1", - "type": "function_call", - "status": "completed", - "arguments": '{"city":"Paris"}', - "call_id": "call_1", - "name": "get_weather", - } - ], - "parallel_tool_calls": True, - "usage": { - "input_tokens": 53, - "output_tokens": 18, - "total_tokens": 71, - "output_tokens_details": {"reasoning_tokens": 0}, - }, - "error": None, - "incomplete_details": None, - "instructions": None, - "metadata": {}, - "temperature": 1.0, - "tool_choice": "auto", - "tools": [], - "top_p": 1.0, - "max_output_tokens": 200, - "previous_response_id": None, - "reasoning": {"effort": "medium", "summary": None}, - "truncation": "disabled", - "user": None, -} - - -def test_completion_bridges_azure_ai_foundry_gpt_5_4_plus_function_tools_to_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - responses_route: Final = respx_mock.post(f"{_FOUNDRY_API_BASE}/openai/v1/responses").respond( - json=_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY - ) - - response: Final = litellm.completion( - model="azure_ai/gpt-6-astra", - messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather for a city", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, - }, - } - ], - max_tokens=200, - api_base=_FOUNDRY_API_BASE, - api_key="fake-foundry-key", - ) - - assert [str(call.request.url) for call in respx_mock.calls] == [f"{_FOUNDRY_API_BASE}/openai/v1/responses"] - request: Final = responses_route.calls[0].request - request_body: Final = json.loads(request.content) - assert request_body["tools"][0]["type"] == "function" - assert request_body["tools"][0]["name"] == "get_weather" - assert request.headers["api-key"] == "fake-foundry-key" - assert response.choices[0].finish_reason == "tool_calls" - assert response.choices[0].message.tool_calls[0].function.name == "get_weather" - - -@pytest.mark.parametrize( - "model, model_info, expected_model_param, expected_base_model_param", - [ - ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None), - ( - "gemini/gemini-3.1-pro", - {"base_model": "gemini-3.1-pro-preview"}, - "gemini-3.1-pro", - "gemini-3.1-pro-preview", - ), - ], -) -def test_completion_optional_params_base_model( - model: str, - model_info: dict | None, - expected_model_param: str, - expected_base_model_param: str | None, -): - """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model`` - (an additive capability hint), without overwriting ``model`` with the label. - - Regression for #29618: overwriting ``model`` with a friendly ``base_model`` - label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``.""" - with patch("litellm.main.get_optional_params") as mock_get_optional_params: - mock_get_optional_params.return_value = MagicMock() - - import litellm - - kwargs = { - "model": model, - "messages": [{"role": "user", "content": "What is the capital of France?"}], - "api_key": "fake-key", - "mock_response": "Hey, how's it going?", - } - if model_info is not None: - kwargs["model_info"] = model_info - - litellm.completion(**kwargs) - - assert mock_get_optional_params.called is True - call_kwargs = mock_get_optional_params.call_args.kwargs - assert call_kwargs["model"] == expected_model_param - assert call_kwargs["base_model"] == expected_base_model_param - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_4_responses_bridge_merges_reasoning_summary_kwarg_without_tools( - mock_responses_completion, -): - """reasoningSummary without tools should route and merge into reasoning_effort dict.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5.4", - messages=[{"role": "user", "content": "ok"}], - reasoning_effort="medium", - reasoningSummary="auto", - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == { - "effort": "medium", - "summary": "auto", - } - assert "reasoningSummary" not in optional_params - assert "reasoning_summary" not in optional_params - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_responses_bridge_preserves_reasoning_summary_without_effort( - mock_responses_completion, -): - """Reasoning summary should survive responses routing even without effort.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "ok"}], - reasoningSummary="auto", - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == {"summary": "auto"} - assert "reasoningSummary" not in optional_params - assert "reasoning_summary" not in optional_params - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_responses_bridge_tools_and_reasoning_summary( - mock_responses_completion, -): - """Bare gpt-5 with tools + reasoningSummary should bridge (OpenCode-style).""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5", - messages=[{"role": "user", "content": "ok"}], - tools=[ - { - "type": "function", - "function": { - "name": "apply_patch", - "parameters": {"type": "object", "properties": {}}, - }, - } - ], - tool_choice="auto", - reasoning_effort="medium", - reasoningSummary="auto", - stream=True, - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params.get("reasoning_effort") == { - "effort": "medium", - "summary": "auto", - } - - -def test_responses_api_bridge_check_handles_exception(): - """Test that responses_api_bridge_check handles exceptions and still processes responses/ models.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.side_effect = Exception("Model not found") - - model_info, model = responses_api_bridge_check( - model="responses/custom-model", custom_llm_provider="custom" - ) - - assert model == "custom-model" - assert model_info["mode"] == "responses" - - -def test_responses_api_bridge_check_global_flag_routes_openai(): - """When route_all_chat_openai_to_responses is True, any OpenAI model routes to responses.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="openai", - ) - - assert model == "gpt-4o" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_global_flag_does_not_affect_azure(): - """route_all_chat_openai_to_responses should not affect Azure models.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="azure", - ) - - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_global_flag_default_false(): - """By default, route_all_chat_openai_to_responses is False and doesn't affect routing.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", False): - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="openai", - ) - - assert model_info.get("mode") != "responses" - - -@pytest.mark.asyncio -async def test_async_mock_delay(): - """Use asyncio await for mock delay on acompletion""" - import time - - from litellm import acompletion - - start_time = time.time() - result = await acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - mock_delay=0.01, - mock_response="Hello world", - ) - end_time = time.time() - delay = end_time - start_time - assert delay >= 0.01 - - -def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): - from litellm import stream_chunk_builder - from litellm.types.utils import ( - ChatCompletionDeltaToolCall, - Delta, - Function, - ModelResponseStream, - StreamingChoices, - ) - - def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: - return ModelResponseStream( - id="chatcmpl-multi-choice", - created=1751934860, - model="gpt-4.1-mini", - object="chat.completion.chunk", - choices=choices, - ) - - chunks = [ - chunk( - [ - StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), - StreamingChoices( - index=1, - delta=Delta( - role="assistant", - tool_calls=[ - ChatCompletionDeltaToolCall( - id="call_1", - index=0, - type="function", - function=Function(name="lookup_fruit", arguments='{"fruit":'), - ) - ], - ), - ), - ] - ), - chunk( - [ - StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), - StreamingChoices( - index=1, - delta=Delta( - tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] - ), - finish_reason="tool_calls", - ), - ] - ), - ] - - response = stream_chunk_builder(chunks=chunks) - - tool_calls = response.choices[0].message.tool_calls - assert tool_calls is not None - assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ - ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') - ] - - -def test_stream_chunk_builder_thinking_blocks(): - from litellm import stream_chunk_builder - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - - chunks = [ - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="I need to summar", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "I need to summar", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "I need to summar", - "signature": None, - } - ] - }, - content="", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="ize the previous agent's thinking process into a", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "ize the previous agent's thinking process into a", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "ize the previous agent's thinking process into a", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" short description. Based on the input data provide", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " short description. Based on the input data provide", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " short description. Based on the input data provide", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="d, it seems the agent was planning to refine their search", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "d, it seems the agent was planning to refine their search", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "d, it seems the agent was planning to refine their search", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" to focus more on technical aspects of home automation and home", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " to focus more on technical aspects of home automation and home", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " to focus more on technical aspects of home automation and home", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" energy system management.\n\nI'll create a brief", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " energy system management.\n\nI'll create a brief", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " energy system management.\n\nI'll create a brief", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" summary of what the agent was doing.", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " summary of what the agent was doing.", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " summary of what the agent was doing.", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "", - "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "", - "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='{"a', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='gent_doing"', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content=': "Re', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="searching", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content=" technic", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="al aspect", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="s of home au", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='tomation"}', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason="tool_calls", - index=0, - delta=Delta( - provider_specific_fields=None, - content=None, - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - ), - ] - - response = stream_chunk_builder(chunks=chunks) - print(response) - - assert response is not None - assert response.choices[0].message.content is not None - assert response.choices[0].message.thinking_blocks is not None - - -from litellm.llms.openai.openai import OpenAIChatCompletion - - -def throw_retryable_error(*_, **__): - raise RuntimeError("BOOM") - - -@pytest.mark.asyncio -async def test_retrying() -> None: - litellm.num_retries = 10 - with ( - patch.object( - OpenAIChatCompletion, - "make_openai_chat_completion_request", - side_effect=throw_retryable_error, - ) as mock_request, - pytest.raises(litellm.InternalServerError, match="LiteLLM Retried: 10 times"), - ): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}], - ) - - -def test_anthropic_disable_url_suffix_env_var(): - """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/messages suffix.""" - import os - from unittest.mock import MagicMock, patch - - from litellm import completion - - # Test with environment variable disabled (default behavior) - with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): - actual_api_base = None - - with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - mock_anthropic.completion = capture_completion - - # This should append /v1/messages - completion( - model="anthropic/claude-3-sonnet", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base has /v1/messages appended - assert actual_api_base.endswith("/v1/messages") - assert actual_api_base == "https://api.example.com/v1/messages" - - # Test with environment variable enabled - with patch.dict( - os.environ, - { - "ANTHROPIC_API_BASE": "https://api.example.com/custom/path", - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", - }, - ): - actual_api_base = None - - with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - mock_anthropic.completion = capture_completion - - # This should NOT append /v1/messages - completion( - model="anthropic/claude-3-sonnet", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base does not have /v1/messages appended - assert actual_api_base == "https://api.example.com/custom/path" - assert not actual_api_base.endswith("/v1/messages") - - -def test_anthropic_text_disable_url_suffix_env_var(): - """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/complete suffix for anthropic_text.""" - import os - from unittest.mock import MagicMock, patch - - from litellm import completion - - # Test with environment variable disabled (default behavior) - with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): - actual_api_base = None - - with patch("litellm.main.base_llm_http_handler") as mock_handler: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - return MagicMock() - - mock_handler.completion = capture_completion - - # This should append /v1/complete - completion( - model="anthropic_text/claude-instant-1", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base has /v1/complete appended - assert actual_api_base.endswith("/v1/complete") - assert actual_api_base == "https://api.example.com/v1/complete" - - # Test with environment variable enabled - with patch.dict( - os.environ, - { - "ANTHROPIC_API_BASE": "https://api.example.com/custom/complete", - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", - }, - ): - actual_api_base = None - - with patch("litellm.main.base_llm_http_handler") as mock_handler: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - return MagicMock() - - mock_handler.completion = capture_completion - - # This should NOT append /v1/complete - completion( - model="anthropic_text/claude-instant-1", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base does not have /v1/complete appended - assert actual_api_base == "https://api.example.com/custom/complete" - assert not actual_api_base.endswith("/v1/complete") - - -def test_image_edit_merges_headers_and_extra_headers(): - from litellm.images.main import base_llm_http_handler - - combined_headers = { - "x-test-header-one": "value-1", - "x-test-header-two": "value-2", - } - - mock_image_edit_config = MagicMock() - mock_image_edit_config.get_supported_openai_params.return_value = set() - mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict( - kwargs["image_edit_optional_params"] - ) - - with ( - patch( - "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", - return_value=mock_image_edit_config, - ) as mock_config, - patch.object( - base_llm_http_handler, - "image_edit_handler", - return_value="ok", - ) as mock_handler, - ): - response = litellm.image_edit( - image=MagicMock(name="image"), - prompt="test", - model="azure/gpt-image-1", - headers={"x-test-header-one": "value-1"}, - extra_headers={ - "x-test-header-two": "value-2", - }, - ) - - assert response == "ok" - mock_config.assert_called_once() - - handler_kwargs = mock_handler.call_args.kwargs - assert handler_kwargs["extra_headers"] == combined_headers - assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] - - -@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) -@pytest.mark.parametrize("input_tokens", (51234, 0)) -def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int): - response = litellm.completion( - model="anthropic/claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - api_key="mock", - **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}}, - ) - - assert response.usage.prompt_tokens == input_tokens - assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens - - -def test_mock_completion_usage_falls_back_to_default_without_admission_count(): - response = litellm.completion( - model="anthropic/claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - api_key="mock", - metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, - ) - - assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT - - -_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT: Final = { - "model_name": "azure-ai-custom-priced", - "litellm_params": { - "model": "azure_ai/gpt-5.6", - "api_key": "mock", - "api_base": "https://example.services.ai.azure.com", - "mock_response": "ok", - "input_cost_per_token": 3e-6, - "output_cost_per_token": 7e-6, - "cache_read_input_token_cost": 1e-7, - "cache_creation_input_token_cost": 5e-7, - }, - "model_info": {"id": "azure-ai-custom-priced-deployment-id"}, -} - - -def _expected_custom_price(response: litellm.ModelResponse) -> float: - params: Final = _AZURE_AI_CUSTOM_PRICED_DEPLOYMENT["litellm_params"] - return ( - response.usage.prompt_tokens * params["input_cost_per_token"] - + response.usage.completion_tokens * params["output_cost_per_token"] - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("use_async", (False, True)) -async def test_mock_completion_prices_azure_ai_router_deployment_with_custom_pricing(use_async: bool): - router: Final = litellm.Router(model_list=[_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT]) - messages: Final = [{"role": "user", "content": "hello"}] - - response: Final = ( - await router.acompletion(model="azure-ai-custom-priced", messages=messages) - if use_async - else router.completion(model="azure-ai-custom-priced", messages=messages) - ) - - assert response._hidden_params["response_cost"] == pytest.approx(_expected_custom_price(response)) - assert response._hidden_params["custom_llm_provider"] == "azure_ai" - - -@pytest.mark.parametrize( - ("model", "expected_provider"), - (("anthropic/claude-sonnet-5", "anthropic"), ("no-such-provider-model", None)), -) -def test_mock_completion_infers_provider_when_called_directly_without_one(model: str, expected_provider: str | None): - response: Final = litellm.mock_completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - ) - - assert response.choices[0].message.content == "ok" - assert response._hidden_params.get("custom_llm_provider") == expected_provider - - -_ADMISSION_INPUT_TOKENS: Final = 51234 - - -def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata - return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} - - -_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) -_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] -_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" - - -def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: - return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] - - -def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: - return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] - - -@pytest.mark.parametrize("n", (None, 2)) -def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - n=n, - stream_options={"include_usage": True}, - metadata=_ADMISSION_METADATA, - ) - ) - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT - assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens - assert _prompt_token_counter_calls(token_counter) == [] - assert all(chunk.choices for chunk in chunks[:-1]) - assert {chunk.id for chunk in chunks} == {chunks[0].id} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("n", (None, 2)) -async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( - n: int | None, -): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - n=n, - stream_options={"include_usage": True}, - litellm_metadata=_ADMISSION_METADATA, - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens - assert _prompt_token_counter_calls(token_counter) == [] - assert all(chunk.choices for chunk in chunks[:-1]) - assert {chunk.id for chunk in chunks} == {chunks[0].id} - - -def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - metadata=_ADMISSION_METADATA, - ) - ) - - assert _client_usage_chunks(chunks) == [] - assert all(len(chunk.choices) == 1 for chunk in chunks) - assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={}, - metadata=_ADMISSION_METADATA, - ) - ) - - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" - assert _client_usage_chunks(chunks) == [] - assert _prompt_token_counter_calls(token_counter) == [] - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={}, - litellm_metadata=_ADMISSION_METADATA, - ) - chunks: Final = [chunk async for chunk in response] - - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" - assert _client_usage_chunks(chunks) == [] - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): - expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, - ) - ) - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == expected_prompt_tokens - assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens - assert len(_prompt_token_counter_calls(token_counter)) >= 1 - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): - expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == expected_prompt_tokens - assert len(_prompt_token_counter_calls(token_counter)) >= 1 - - -def _usage_triple(usage: Usage) -> tuple[int, int, int]: - return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) - - -@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) -def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): - metadata: Final = _admission_metadata(input_tokens) - non_stream: Final = litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - metadata=metadata, - ) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata=metadata, - ) - ) - - assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) - assert non_stream.usage.prompt_tokens == input_tokens - assert _prompt_token_counter_calls(token_counter) == [] - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=[{"role": "user", "content": ""}], - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - litellm_metadata=_admission_metadata(0), - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): - metadata: Final = _admission_metadata(0) - non_stream: Final = litellm.text_completion( - model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata - ) - chunks: Final = list( - litellm.text_completion( - model="openai/gpt-5.4-mini", - prompt="", - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata=metadata, - ) - ) - - stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) - assert len(stream_usages) == 1 - assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) - assert non_stream.usage.prompt_tokens == 0 - - -def test_mock_completion_stream_with_model_response(): - """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" - from litellm import completion - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a ModelResponse object - mock_model_response = ModelResponse( - id="chatcmpl-test-123", - created=1234567890, - model="gpt-4o-mini", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="This is a test response", - role="assistant", - ), - ) - ], - usage=Usage( - prompt_tokens=10, - completion_tokens=20, - total_tokens=30, - ), - ) - - # Call completion with stream=True and mock_response as ModelResponse - response = completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - mock_response=mock_model_response, - ) - - # Verify that the response is a stream - assert response is not None - - # Collect all chunks from the stream - chunks = [] - for chunk in response: - chunks.append(chunk) - print(f"Chunk: {chunk}") - - # Verify we got chunks - assert len(chunks) > 0 - - # Verify the content is streamed correctly - accumulated_content = "" - for chunk in chunks: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - accumulated_content += chunk.choices[0].delta.content - - assert "This is a test response" in accumulated_content or len(chunks) > 0 - - -@pytest.mark.asyncio -async def test_async_mock_completion_stream_with_model_response(): - """Test that async mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" - from litellm import acompletion - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a ModelResponse object - mock_model_response = ModelResponse( - id="chatcmpl-test-456", - created=1234567890, - model="gpt-4o-mini", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="This is an async test response", - role="assistant", - ), - ) - ], - usage=Usage( - prompt_tokens=15, - completion_tokens=25, - total_tokens=40, - ), - ) - - # Call acompletion with stream=True and mock_response as ModelResponse - response = await acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello async"}], - stream=True, - mock_response=mock_model_response, - ) - - # Verify that the response is a stream - assert response is not None - - # Collect all chunks from the stream - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"Async Chunk: {chunk}") - - # Verify we got chunks - assert len(chunks) > 0 - - # Verify the content is streamed correctly - accumulated_content = "" - for chunk in chunks: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - accumulated_content += chunk.choices[0].delta.content - - assert "This is an async test response" in accumulated_content or len(chunks) > 0 - - -class TestCallTypesOCR: - """Test that OCR call types are properly defined in CallTypes enum. - - Fixes https://github.com/BerriAI/litellm/issues/17381 - """ - - def test_ocr_call_type_exists(self): - """Test that CallTypes.ocr exists and has correct value.""" - from litellm.types.utils import CallTypes - - assert hasattr(CallTypes, "ocr") - assert CallTypes.ocr.value == "ocr" - - def test_aocr_call_type_exists(self): - """Test that CallTypes.aocr exists and has correct value.""" - from litellm.types.utils import CallTypes - - assert hasattr(CallTypes, "aocr") - assert CallTypes.aocr.value == "aocr" - - def test_ocr_call_type_from_string(self): - """Test that CallTypes can be constructed from 'ocr' string.""" - from litellm.types.utils import CallTypes - - call_type = CallTypes("ocr") - assert call_type == CallTypes.ocr - - def test_aocr_call_type_from_string(self): - """Test that CallTypes can be constructed from 'aocr' string. - - This is the actual use case that was failing - the OCR endpoint - uses route_type='aocr' and guardrails try to instantiate - CallTypes('aocr'). - """ - from litellm.types.utils import CallTypes - - call_type = CallTypes("aocr") - assert call_type == CallTypes.aocr - - -def test_stream_chunk_builder_text_completion_combines_text_and_usage(): - from litellm.main import stream_chunk_builder_text_completion - from litellm.types.utils import TextCompletionResponse - - chunks = [ - TextCompletionResponse( - id="cmpl-1", - object="text_completion", - created=1, - model="gpt-3.5-turbo-instruct", - choices=[{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}], - ), - TextCompletionResponse( - id="cmpl-1", - object="text_completion", - created=1, - model="gpt-3.5-turbo-instruct", - choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], - ), - ] - - response = stream_chunk_builder_text_completion( - chunks=chunks, messages=[{"role": "user", "content": "say hello"}] - ) - - assert response.choices[0].text == "Hello world" - assert response.choices[0].finish_reason == "stop" - assert response.usage.prompt_tokens > 0 - assert response.usage.completion_tokens > 0 - assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens - - -def test_completion_forwards_store_and_prompt_cache_key_to_openai(): - """ - Regression test for https://github.com/BerriAI/litellm/issues/33184 - - store and prompt_cache_key are documented OpenAI chat completion params that - were accepted as supported but silently dropped before the provider request - was built, because they were not named parameters of completion() and - get_optional_params() the way safety_identifier is. - """ - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - store=False, - prompt_cache_key="test-cache-key", - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert request_body["store"] is False - assert request_body["prompt_cache_key"] == "test-cache-key" - - -@pytest.mark.asyncio -async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): - """ - Async variant of the store/prompt_cache_key forwarding regression test for - https://github.com/BerriAI/litellm/issues/33184 - """ - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - await litellm.acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - store=False, - prompt_cache_key="test-cache-key", - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert request_body["store"] is False - assert request_body["prompt_cache_key"] == "test-cache-key" - - -def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): - """ - When store and prompt_cache_key are not passed, they must not appear in the - outbound request body (guards against always forwarding None defaults). - """ - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert "store" not in request_body - assert "prompt_cache_key" not in request_body - - -def test_completion_forwards_store_and_prompt_cache_key_to_mcp_gateway(): - """ - Regression test for the MCP gateway early-return in completion(): store and - prompt_cache_key are named params, so they no longer travel via **kwargs and - must be forwarded explicitly like safety_identifier and service_tier. - """ - with patch.object( - import_module("litellm.responses.mcp.chat_completions_handler"), "acompletion_with_mcp" - ) as mock_mcp: - result = litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy"}], - store=False, - prompt_cache_key="test-cache-key", - ) - - result.close() - mock_mcp.assert_called_once() - call_kwargs = mock_mcp.call_args.kwargs - assert call_kwargs["store"] is False - assert call_kwargs["prompt_cache_key"] == "test-cache-key" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "aws_credential_kwargs", - [ - { - "aws_session_name": "litellm-gcp", - "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", - "aws_web_identity_token": "oidc/google/108963886734710037768", - }, - { - "aws_access_key_id": "AKIASTATICKEYFORTEST", - "aws_secret_access_key": "static-secret-key", - "aws_session_token": "static-session-token", - }, - ], - ids=["web_identity", "static_keys"], -) -async def test_acompletion_forwards_aws_credentials_through_responses_bridge( - respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict -): - from botocore.credentials import Credentials - - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - original_disable_aiohttp = litellm.disable_aiohttp_transport - try: - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) - monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) - - get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) - monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) - - respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( - json={ - "id": "resp_123", - "object": "response", - "created_at": 1760144904, - "status": "completed", - "model": "openai.gpt-5.4", - "output": [ - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "ok", "annotations": []}], - } - ], - } - ) - - response = await litellm.acompletion( - model="bedrock_mantle/openai.gpt-5.4", - messages=[{"role": "user", "content": "hi"}], - api_base="https://bedrock-mantle.us-east-2.api.aws/v1", - aws_region_name="us-east-2", - num_retries=0, - **aws_credential_kwargs, - ) - - assert response.choices[0].message.content == "ok" - credential_kwargs = get_credentials_mock.call_args.kwargs - assert credential_kwargs["aws_region_name"] == "us-east-2" - for key, value in aws_credential_kwargs.items(): - assert credential_kwargs[key] == value - authorization = respx_mock.calls.last.request.headers["Authorization"] - assert authorization.startswith("AWS4-HMAC-SHA256") - assert "fake-key" in authorization - finally: - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -_GEMINI_RESPONSE_BODY = { - "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}], - "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, -} - - -def _gemini_client_returning_a_reply(): - """An injected HTTP client whose post() answers like generativelanguage does.""" - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - client = HTTPHandler() - request = httpx.Request("POST", "https://generativelanguage.googleapis.com/") - post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request)) - return client, post - - -@pytest.fixture -def restore_model_registry(): - """litellm.model_cost and the provider name sets are module-global. - - register_model merges into the existing entry in place, hence the deep copy. - """ - model_cost = copy.deepcopy(litellm.model_cost) - openai_models = set(litellm.open_ai_chat_completion_models) - yield - litellm.model_cost.clear() - litellm.model_cost.update(model_cost) - litellm.open_ai_chat_completion_models.clear() - litellm.open_ai_chat_completion_models.update(openai_models) - - -def test_openai_model_name_does_not_outrank_explicit_provider(): - """`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler. - - completion() checks `model in litellm.open_ai_chat_completion_models` ahead of - the gemini branch, so the call used to reach the OpenAI handler carrying - VertexGeminiConfig, whose transform_request raises NotImplementedError. - """ - assert "gpt-4o" in litellm.open_ai_chat_completion_models - client, post = _gemini_client_returning_a_reply() - - with patch.object(client, "post", new=post): - response = litellm.completion( - model="gemini/gpt-4o", - messages=[{"role": "user", "content": "hello"}], - api_key="test-api-key", - client=client, - ) - - assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] - assert "models/gpt-4o" in post.call_args.kwargs["url"] - assert response.choices[0].message.content == "hello" - - -def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry): - """register_model is the other way into the same failure. - - An entry claiming litellm_provider "openai" adds its name to - open_ai_chat_completion_models, so one mislabelled price reroutes every later - call to that model in the process. - """ - litellm.register_model( - { - "gemini-2.5-pro": { - "litellm_provider": "openai", - "mode": "chat", - "input_cost_per_token": 1e-06, - "output_cost_per_token": 4e-06, - } - } - ) - assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models - client, post = _gemini_client_returning_a_reply() - - with patch.object(client, "post", new=post): - response = litellm.completion( - model="gemini/gemini-2.5-pro", - messages=[{"role": "user", "content": "hello"}], - api_key="test-api-key", - client=client, - ) - - assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] - assert response.choices[0].message.content == "hello" - - -def test_openai_model_without_a_provider_still_routes_to_openai(): - from openai import OpenAI - - client = OpenAI(api_key="fake-key") - raw_response = client.chat.completions.with_raw_response - with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hello"}], - client=client, - ) - - mock_create.assert_called() - - -def _openai_chat_create_kwargs(client, **completion_kwargs): - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - with contextlib.suppress(Exception): - litellm.completion( - messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], - cache_control_injection_points=[{"location": "message", "role": "system"}], - client=client, - **completion_kwargs, - ) - - mock_client.assert_called_once() - return mock_client.call_args.kwargs - - -@pytest.fixture -def _no_openai_api_base_override(monkeypatch): - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("_no_openai_api_base_override") -async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: - with contextlib.suppress(Exception): - await litellm.acompletion( - model="gpt-5.6", - messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], - cache_control_injection_points=[{"location": "message", "role": "system"}], - client=client, - base_url="http://127.0.0.1:9/v1", - ) - - mock_create.assert_called_once() - request_body = mock_create.call_args.kwargs - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") - - assert request_body["messages"][0]["content"] == [ - {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} - ] - assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "explicit"} - - -_SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" - - -def _scoped_headers_for_oauth_request(): - from litellm.types.utils import ProviderSpecificHeader - - return [ - ProviderSpecificHeader( - custom_llm_provider="anthropic,bedrock,vertex_ai", - extra_headers={"anthropic-version": "2023-06-01"}, - ), - ProviderSpecificHeader( - custom_llm_provider="anthropic", - extra_headers={"authorization": _SUBSCRIPTION_OAUTH_CREDENTIAL}, - ), - ] - - -def _run_anthropic_hop_with_shared_headers(shared_headers): - litellm.completion( - model="anthropic/claude-3-5-sonnet-20240620", - messages=[{"role": "user", "content": "Say OK"}], - extra_headers=shared_headers, - provider_specific_header=_scoped_headers_for_oauth_request(), - api_key="sk-fake-anthropic-key", - mock_response="OK", - ) - - -def test_completion_does_not_mutate_caller_supplied_headers(): - shared_headers = {"x-tenant": "acme"} - - _run_anthropic_hop_with_shared_headers(shared_headers) - - assert shared_headers == {"x-tenant": "acme"} - - -def test_anthropic_oauth_credential_does_not_persist_into_next_provider_hop(): - shared_headers = {"x-tenant": "acme"} - - _run_anthropic_hop_with_shared_headers(shared_headers) - - leaked = [name for name, value in shared_headers.items() if value == _SUBSCRIPTION_OAUTH_CREDENTIAL] - assert leaked == [] - assert "anthropic-version" not in shared_headers - - -STREAM_COST_MODEL = "gpt-4o" -STREAMED_USAGE = {"prompt_tokens": 137, "completion_tokens": 42, "total_tokens": 179} - - -def _text_chunk(content, finish_reason=None, usage=None): - chunk = { - "id": "chatcmpl-stream-cost", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": STREAM_COST_MODEL, - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": content}, - "finish_reason": finish_reason, - } - ], - } - if usage is not None: - chunk["usage"] = usage - return chunk - - -def _priced_at(prompt_tokens, completion_tokens): - prices = litellm.model_cost[STREAM_COST_MODEL] - return ( - prompt_tokens * prices["input_cost_per_token"] - + completion_tokens * prices["output_cost_per_token"] - ) - - -@pytest.fixture -def local_cost_map(monkeypatch): - """The prices these tests assert are the checked-in ones. Setting the environment - variable alone does not reload the map, so pin the map itself. - - Prices are read through two separate lru_caches, so pinning ``model_cost`` is not - enough on its own: an entry warmed against the network-fetched map keeps its old - prices and billing reads those while the assertions read the pinned map. - ``_invalidate_model_cost_lowercase_map`` clears both caches, where - ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out - so entries never leak across tests in either direction.""" - from litellm.utils import _invalidate_model_cost_lowercase_map - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() - yield - _invalidate_model_cost_lowercase_map() - - -def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), - ], - messages=[{"role": "user", "content": "hi"}], - ) - - assert rebuilt.choices[0].message.content == "Hello there" - assert rebuilt.usage.prompt_tokens == STREAMED_USAGE["prompt_tokens"] - assert rebuilt.usage.completion_tokens == STREAMED_USAGE["completion_tokens"] - - cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) - - assert cost == pytest.approx(_priced_at(137, 42)) - - -def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), - ], - messages=[{"role": "user", "content": "hi"}], - ) - whole = litellm.ModelResponse( - id="chatcmpl-stream-cost", - model=STREAM_COST_MODEL, - object="chat.completion", - created=1700000000, - choices=[ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello there"}, - "finish_reason": "stop", - } - ], - usage=STREAMED_USAGE, - ) - - assert litellm.completion_cost( - completion_response=rebuilt, model=STREAM_COST_MODEL - ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) - - -def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop"), - ], - messages=[{"role": "user", "content": "hi"}], - ) - - assert rebuilt.usage.prompt_tokens > 0 - assert rebuilt.usage.completion_tokens > 0 - - cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) - - assert cost > 0 - assert cost == pytest.approx( - _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) - ) - - -@pytest.mark.asyncio -async def test_acompletion_resolves_provider_from_api_base(): - response = await litellm.acompletion( - model="deepseek-chat", - api_base="https://api.deepseek.com/v1", - api_key="fake-key", - messages=[{"role": "user", "content": "hi"}], - mock_response="resolved", - ) - - assert response.choices[0].message.content == "resolved" - - -@dataclass(frozen=True, slots=True) -class _RecordedSpeechSuccess: - call_type: str | None - spend_metadata: Mapping[str, object] - response_cost: float | None - logged_response_cost: float | None - - -def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess: - call_type: Final = payload.get("call_type") - response_cost: Final = payload.get("response_cost") - logging_payload: Final = payload.get("standard_logging_object") - logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None - return _RecordedSpeechSuccess( - call_type=call_type if isinstance(call_type, str) else None, - spend_metadata=get_litellm_metadata_from_kwargs(payload), - response_cost=response_cost if isinstance(response_cost, float) else None, - logged_response_cost=logged_cost if isinstance(logged_cost, float) else None, - ) - - -class _SuccessEventRecorder(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events - - async def async_log_success_event( - self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object - ) -> None: - self.events.append(_record_speech_success(kwargs)) - - -async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: str) -> _RecordedSpeechSuccess: - for _ in range(100): - if (event := next((e for e in recorder.events if e.call_type == call_type), None)) is not None: - return event - await asyncio.sleep(0.05) - pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}") - - -def _gemini_tts_generate_content_response() -> dict[str, object]: - return { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "audio/L16;codec=pcm;rate=24000", - "data": base64.b64encode(b"pcm-audio-bytes").decode(), - } - } - ], - "role": "model", - }, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 5, - "candidatesTokenCount": 60, - "totalTokenCount": 65, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], - "candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": 60}], - }, - "modelVersion": "gemini-2.5-flash-preview-tts", - } - - -@pytest.mark.asyncio -async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.delenv("GEMINI_API_KEY", raising=False) - monkeypatch.delenv("GOOGLE_API_KEY", raising=False) - recorder: Final = _SuccessEventRecorder() - monkeypatch.setattr(litellm, "callbacks", [recorder]) - mock_route: Final = respx_mock.post( - url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" - ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) - - await litellm.aspeech( - model="gemini/gemini-2.5-flash-preview-tts", - input="spend tracking check", - voice="Kore", - api_key="fake-gemini-key", - metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"}, - ) - - assert mock_route.called - assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key" - speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech") - assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key" - assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1" - expected_prompt_cost, expected_completion_cost = litellm.cost_per_token( - model="gemini/gemini-2.5-flash-preview-tts", - usage_object=Usage(prompt_tokens=5, completion_tokens=60, total_tokens=65), - ) - expected_cost: Final = expected_prompt_cost + expected_completion_cost - assert expected_cost > 0 - assert speech_event.response_cost == pytest.approx(expected_cost) - assert speech_event.logged_response_cost == pytest.approx(expected_cost) - - -def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream: - return ModelResponseStream( - id="chatcmpl-cost", - created=1724900000, - model=model, - object="chat.completion.chunk", - choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content, role="assistant"))], - ) - - -def test_stream_chunk_builder_sets_hidden_response_cost_for_known_model(): - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) - expected_cost: Final = prompt_cost + completion_cost - assert expected_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) - - -def test_stream_chunk_builder_unknown_model_leaves_response_cost_unset(): - chunks: Final = [ - _stream_builder_text_chunk("totally-unknown-model-xyz", "Hello "), - _stream_builder_text_chunk("totally-unknown-model-xyz", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response._hidden_params.get("response_cost") is None - assert response.choices[0].message.content == "Hello world." - - -def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): - chunks: Final = [ - _stream_builder_text_chunk("claude-opus-5", "Hello "), - _stream_builder_text_chunk("claude-opus-5", "world.", finish_reason="stop"), - ] - for chunk in chunks: - chunk._hidden_params = {"custom_llm_provider": "openai"} - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response._hidden_params["custom_llm_provider"] == "openai" - prompt_cost, completion_cost = litellm.cost_per_token(model="claude-opus-5", usage_object=response.usage) - expected_cost: Final = prompt_cost + completion_cost - assert expected_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) - - -def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: - logging_obj: Final = LiteLLMLogging( - model=model, - messages=[{"role": "user", "content": "hi"}], - stream=True, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-call-id", - function_id="test-function-id", - ) - logging_obj.update_environment_variables( - model=model, - user=None, - optional_params={}, - litellm_params={"custom_llm_provider": custom_llm_provider}, - custom_llm_provider=custom_llm_provider, - ) - return logging_obj - - -def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() - ) - - assert response is not None - usage_cost: Final = getattr(response.usage, "cost", None) - assert usage_cost is not None - assert usage_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) - - -def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): - import time as time_module - - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj: Final = LiteLLMLogging( - model="us.anthropic.claude-opus-5", - messages=[{"role": "user", "content": "hi"}], - stream=True, - call_type="completion", - start_time=time_module.time(), - litellm_call_id="stream-builder-alias-unpriceable", - function_id="1", - ) - logging_obj.model_call_details["custom_llm_provider"] = "bedrock" - logging_obj.optional_params = {} - usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") - usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) - chunks: Final = [ - _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj - ) - - assert response is not None - assert getattr(response.usage, "cost", None) is None - assert response._hidden_params.get("response_cost") is None - - -def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): - usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") - usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() - ) - - assert response is not None - assert getattr(response.usage, "cost", None) == pytest.approx(0.5) - assert response._hidden_params["response_cost"] == pytest.approx(0.5) - - -def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): - from openai.types.completion_usage import CompletionUsage - - usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") - usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) - assert type(usage_chunk.usage) is CompletionUsage - chunks: Final = [ - _stream_builder_text_chunk("mantle-claude", "Hello "), - _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response.usage.prompt_tokens == 20 - assert response.usage.completion_tokens == 60 - assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) - assert response._hidden_params["response_cost"] == pytest.approx(0.000704) - - -def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) - usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") - usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) - chunks: Final = [ - _stream_builder_text_chunk("grok-4", "Hello "), - _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), - usage_chunk, - ] - logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj - ) - - assert response is not None - assert getattr(response.usage, "cost", None) == pytest.approx(0.42) - assert response._hidden_params.get("response_cost") is None - assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) - - -def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") - audio_bytes: Final = b"ID3-fake-mp3-bytes" - mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( - return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) - ) - - response: Final = litellm.speech( - model="mistral/voxtral-mini-tts-2603", - input="hello from litellm", - voice="en_paul_neutral", - response_format="wav", - speed=2, - instructions="sound cheerful", - ) - - assert mock_route.called - request_body: Final = json.loads(mock_route.calls.last.request.content) - assert request_body == { - "model": "voxtral-mini-tts-2603", - "input": "hello from litellm", - "voice_id": "en_paul_neutral", - "response_format": "wav", - } - assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" - assert response.content == audio_bytes - - -def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") - audio_bytes: Final = b"ID3-gateway-bytes" - gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( - return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) - ) - - response: Final = litellm.speech( - model="mistral/voxtral-mini-tts-2603", - input="hello from litellm", - voice="en_paul_neutral", - api_base="https://mistral.gateway.internal", - ) - - assert gateway_route.called - assert response.content == audio_bytes - - -FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" - - -def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( - respx_mock: respx.MockRouter, -): - route: Final = respx_mock.post( - url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" - ).mock(return_value=httpx.Response(200, json={"text": "hello"})) - - response: Final = litellm.transcription( - model="azure_ai/whisper-1", - file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), - api_base=FOUNDRY_HOST, - api_key="fake-key", - ) - - assert route.called - assert response.text == "hello" - - -def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( - respx_mock: respx.MockRouter, -): - route: Final = respx_mock.post( - url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" - ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) - - response: Final = litellm.speech( - model="azure_ai/tts-1", - input="hello", - voice="alloy", - api_base=FOUNDRY_HOST, - api_key="fake-key", - ) - - assert route.called - assert response.content == b"mp3-bytes" - - -FORWARDED_CLIENT_HEADERS: Final = {"x-forwarded-for": "10.0.0.1", "x-amzn-trace-id": "Root=1-lit7694"} - - -def _chat_completion_json() -> Mapping[str, object]: - return { - "id": "chatcmpl-lit7694", - "object": "chat.completion", - "created": 1, - "model": "gpt-5.4", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - - -def _chat_completion_sse() -> bytes: - chunk: Final = { - "id": "chatcmpl-lit7694", - "object": "chat.completion.chunk", - "created": 1, - "model": "gpt-5.4", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - } - return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() - - -@pytest.mark.parametrize("stream", [False, True]) -def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_of_the_body( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool -): - monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, content=_chat_completion_sse(), headers={"content-type": "text/event-stream"}) - if stream - else httpx.Response(200, json=_chat_completion_json()) - ) - - response: Final = litellm.responses( - model="openai/gpt-5.4", - input="Reply with the single word ok", - stream=stream, - use_chat_completions_api=True, - headers=dict(FORWARDED_CLIENT_HEADERS), - api_key="sk-test", - ) - if stream: - list(response) - - assert route.called - request: Final = route.calls.last.request - body: Final = json.loads(request.content) - assert "extra_headers" not in body - assert body["model"] == "gpt-5.4" - assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS - - -@pytest.mark.parametrize("http2_on", [True, False]) -def test_aiohttp_openai_warns_only_when_http2_enabled( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, http2_on: bool -): - from litellm.main import base_llm_aiohttp_handler - - monkeypatch.setattr(litellm, "http2", http2_on) - monkeypatch.delenv("LITELLM_HTTP2", raising=False) - - handler_completion: Final = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(base_llm_aiohttp_handler, "completion", handler_completion) - - with caplog.at_level(logging.WARNING, logger="LiteLLM"): - litellm.completion( - model="aiohttp_openai/gpt-4o", - messages=[{"role": "user", "content": "hi"}], - api_key="sk-test", - ) - - assert handler_completion.called - warned: Final = "aiohttp_openai/ always uses aiohttp" in caplog.text - assert warned is http2_on - - -@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) -def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): - with pytest.raises(litellm.BadRequestError) as exc_info: - litellm.completion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "Which fruit is red?"}], - tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], - tool_choice=tool_choice, - api_key="sk-unused", - ) - assert exc_info.value.status_code == 400 - assert f"tool_choice={tool_choice}" in str(exc_info.value) diff --git a/tests/test_litellm/types/__init__.py b/tests/test_litellm/types/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/types/proxy/__init__.py b/tests/test_litellm/types/proxy/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/types/proxy/policy_engine/__init__.py b/tests/test_litellm/types/proxy/policy_engine/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/vector_stores/__init__.py b/tests/test_litellm/vector_stores/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/videos/__init__.py b/tests/test_litellm/videos/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index d1572f4a7c9..dd95addac40 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2072,3 +2072,348 @@ def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch): ) assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2) assert result.usage.total_tokens == 15 + + +GROUNDED_USAGE_METADATA = { + "promptTokenCount": 19, + "candidatesTokenCount": 59, + "thoughtsTokenCount": 406, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 557, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], + "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], + "trafficType": "ON_DEMAND", +} + + +PASSTHROUGH_OUTPUT_URI = ( + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" + "predictions.jsonl" +) + + +UNGROUNDED_USAGE_METADATA = { + "promptTokenCount": 20, + "candidatesTokenCount": 48, + "thoughtsTokenCount": 195, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 336, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], + "trafficType": "ON_DEMAND", +} + + +def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): + candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} + response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} + return { + "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, + "status": "", + "response": {**response, **({"modelVersion": model_version} if model_version else {})}, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + + +def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: + import litellm.cost_calculator as cc + + calls: list = [] + + def _calc(**kw): + calls.append(kw) + return (prompt_cost, completion_cost) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + return calls + + +def test_vertex_native_cost_bills_embedding_rows(monkeypatch): + monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) + rows = [ + { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, + }, + { + "key": "id_2", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, + }, + {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") + + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) + assert result.cost == pytest.approx(5 * 1e-7) + assert result.models == ["gemini-embedding-2"] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), + ] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.models == ["gemini-2.5-flash"] + assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} + + +@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) + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.successful_requests == 1 + + +@pytest.mark.asyncio +async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="openai", + ) + + assert result.successful_requests == 0 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(raw_rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + result = await bu._handle_completed_batch( + _batch(PASSTHROUGH_OUTPUT_URI), + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert result.cost == pytest.approx(1.0) + assert result.usage.total_tokens == 557 + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) + ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) + + result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") + + grounded_usage, ungrounded_usage = (call["usage"] for call in calls) + assert grounded_usage.prompt_tokens == 19 + assert grounded_usage.completion_tokens == 59 + 406 + assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 + assert ungrounded_usage.prompt_tokens == 20 + 73 + assert ungrounded_usage.completion_tokens == 48 + 195 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 19 + 93, + 465 + 243, + 557 + 336, + ) + + +def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.cost == pytest.approx(1.5) + assert result.successful_requests == 3 + assert result.usage.total_tokens == 557 + 336 + 336 + + +def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): + _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, + {"request": {"contents": []}, "response": {"candidates": []}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 2) + assert result.usage.total_tokens == 557 + + +def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert result.models == ["gemini-2.5-flash"] + assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) + assert calls == [] + + +def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + bu.calculate_vertex_ai_batch_cost_and_usage( + [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + "gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6} + + await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert [call["model"] for call in calls] == ["gemini-2.5-flash"] + assert result.models == ["gemini-2.5-flash"] + + +def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 1) + assert result.usage.total_tokens == 557 + assert len(calls) == 1 + + +@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) +def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + + +def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") + without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) + + twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) + both = bu.calculate_vertex_ai_batch_cost_and_usage( + [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info + ) + + assert twin.cost > 0 + assert both.cost == pytest.approx(2 * twin.cost) + assert (both.successful_requests, both.failed_requests) == (2, 0) + + +def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): + import litellm.cost_calculator as cc + + def _calc(**kw): + if kw["model"] == "gemini-unpriced": + raise ValueError("no pricing") + return (0.5, 0.25) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert result.cost == pytest.approx(0.75) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] + + +@pytest.mark.asyncio +async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert calls == [] + assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py index 2807ed7f8f7..40b1c0ef019 100644 --- a/tests/unit/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -20,6 +20,8 @@ from litellm.rust_bridge.chat_completions.entrypoints import ( ) from litellm.rust_bridge.configuration import Rollout from litellm.types.utils import ModelResponse +from litellm.chat_completions import dispatch +from litellm.rust_bridge.catalog import Rules MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () @@ -256,3 +258,100 @@ async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.Mo NATIVE_ACOMPLETION.reset() assert result is expected assert [request.model for request in captured] == ["gpt-4o"] + + +@pytest.mark.asyncio +async def test_public_completion_calls_keep_the_python_result() -> None: + sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok") + async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok") + + assert isinstance(sync_response, ModelResponse) + assert isinstance(async_response, ModelResponse) + assert sync_response.choices[0].message.content == "ok" + assert async_response.choices[0].message.content == "ok" + + +def test_sync_completion_request_projects_public_arguments() -> None: + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) + expected: Final = ModelResponse() + + def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + assert request.model == "test-model" + assert request.messages == MESSAGES + assert request.custom_llm_provider == "openai" + assert request.stream is True + return expected + + binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {"custom_llm_provider": "openai", "stream": True}, + python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +@pytest.mark.asyncio +async def test_async_completion_falls_back_after_native_declines() -> None: + from litellm.rust_bridge.bindings import native_exception_types + + native_types: Final = native_exception_types() + if native_types is None: + pytest.skip("native bridge is unavailable") + declined, _ = native_types + expected: Final = ModelResponse() + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),) + + async def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + raise declined("unsupported") + + async def python(*args: object, **kwargs: object) -> ModelResponse: + return expected + + binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None) + binding.override(native) + response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_internal_acompletion_marker_bypasses_native() -> None: + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) + expected: Final = ModelResponse() + + def python(*args: object, **kwargs: object) -> ModelResponse: + return expected + + def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + pytest.fail("acompletion's inner completion call must stay on Python") + + binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {"custom_llm_provider": "openai", "acompletion": True}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected diff --git a/tests/test_litellm/a2a_protocol/__init__.py b/tests/unit/completion_extras/litellm_responses_transformation/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/__init__.py rename to tests/unit/completion_extras/litellm_responses_transformation/__init__.py diff --git a/tests/test_litellm/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 similarity index 100% rename from tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py rename to tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py diff --git a/tests/test_litellm/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 similarity index 100% rename from tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py rename to tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 202ecb80d7b..653b2c9914a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,7 +1,11 @@ +import asyncio +import importlib import os -from collections.abc import Iterator +from collections.abc import Coroutine, Iterator +from pathlib import Path from typing import Final +import boto3 import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -10,6 +14,14 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency +from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency + image_handling as image_handling_module, +) +from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency + close_litellm_async_clients, +) +from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( @@ -20,6 +32,63 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( "AZURE_USERNAME", "AZURE_PASSWORD", ) +AMBIENT_AWS_ENV_VARS: Final = ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_SESSION_TOKEN", + "AWS_ROLE_ARN", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", +) +MODULES_WITH_AWS_AUTH_HANDLERS: Final = ( + "litellm.main", + "litellm.files.main", + "litellm.rerank_api.main", + "litellm.realtime_api.main", +) +CALLBACK_LISTS: Final = ( + "callbacks", + "success_callback", + "failure_callback", + "input_callback", + "_async_success_callback", + "_async_failure_callback", + "_async_input_callback", +) +RESET_TO_NONE_GLOBALS: Final = ("model_fallbacks", "cache") +RESTORED_GLOBALS: Final = ( + "disable_aiohttp_transport", + "force_ipv4", + "drop_params", + "secret_manager_client", + "_key_management_system", + "_key_management_settings", + "api_base", + "num_retries", + "modify_params", + "ssl_verify", + "credential_list", + "model_group_settings", + "default_internal_user_params", + "default_team_params", + "prometheus_emit_stream_label", + "vector_store_registry", + "model_cost", + "cost_margin_config", + "cost_discount_config", + "disable_hf_tokenizer_download", + "disable_copilot_system_to_assistant", + "cohere_models", + "anthropic_models", + "token_counter", + "initialized_langfuse_clients", +) +MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient") +SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client") def _allow_loopback_only() -> None: @@ -29,11 +98,116 @@ def _allow_loopback_only() -> None: _allow_loopback_only() +def pytest_collectstart() -> None: + _allow_loopback_only() + + @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() +def _run_coroutine_if_needed(result: object) -> None: + if not asyncio.iscoroutine(result): + return + coroutine: Final[Coroutine[object, object, object]] = result + try: + asyncio.run(coroutine) + except RuntimeError: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + coroutine.close() + return + loop.create_task(coroutine) + + +def _close_handler_if_needed(handler: object) -> None: + close: Final = getattr(handler, "close", None) + if not callable(close): + return + _run_coroutine_if_needed(close()) + + +def _reset_aws_auth_caches() -> None: + modules: Final = tuple(importlib.import_module(name) for name in MODULES_WITH_AWS_AUTH_HANDLERS) + flushes: Final = ( + getattr(getattr(getattr(module, attr_name), "iam_cache", None), "flush_cache", None) + for module in modules + for attr_name in dir(module) + ) + for flush in filter(callable, flushes): + flush() + boto3.DEFAULT_SESSION = None + + +def _flush_client_caches() -> None: + litellm.in_memory_llm_clients_cache.flush_cache() + image_handling_module.in_memory_cache.flush_cache() + _reset_aws_auth_caches() + + +@pytest.fixture(scope="session") +def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + aws_dir: Final = tmp_path_factory.mktemp("aws-config") + credentials: Final = aws_dir / "credentials" + config: Final = aws_dir / "config" + credentials.write_text("", encoding="utf-8") + config.write_text("", encoding="utf-8") + return credentials, config + + +@pytest.fixture(autouse=True) +def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]: + credentials, config = isolated_aws_config_files + with pytest.MonkeyPatch.context() as environment: + environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials)) + environment.setenv("AWS_CONFIG_FILE", str(config)) + environment.setenv("AWS_EC2_METADATA_DISABLED", "true") + for name in AMBIENT_AWS_ENV_VARS: + environment.delenv(name, raising=False) + environment.delenv("PROXY_BASE_URL", raising=False) + environment.setenv("LITELLM_CLI_DISABLE_KEYRING", "1") + yield + + +@pytest.fixture(autouse=True) +def isolate_litellm_globals() -> Iterator[None]: + original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} + original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} + original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} + original_loggers: Final = { + logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters)) + for logger in ALL_LOGGERS + } + original_tool_policy_registry: Final = tool_registry_writer_module._tool_policy_registry + _flush_client_caches() + for name in CALLBACK_LISTS: + setattr(litellm, name, []) + for name in RESET_TO_NONE_GLOBALS: + setattr(litellm, name, None) + for name in MODULE_LEVEL_CLIENTS: + litellm.__dict__.pop(name, None) + tool_registry_writer_module._tool_policy_registry = None + yield + _flush_client_caches() + leaked_clients: Final = tuple(litellm.__dict__.pop(name, None) for name in MODULE_LEVEL_CLIENTS) + for name, client in zip(MODULE_LEVEL_CLIENTS, leaked_clients): + if client is not original_clients.get(name): + _close_handler_if_needed(client) + litellm.__dict__.update(original_clients) + for name, value in (original_callbacks | original_reset | original_restored).items(): + setattr(litellm, name, value) + for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): + logger.setLevel(level) + logger.disabled = disabled + logger.propagate = propagate + logger.handlers = handlers + logger.filters = filters + tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry + + @pytest.fixture(autouse=True) def isolate_router_model_cost_state() -> Iterator[None]: original_live_routers: Final = frozenset(litellm_router_module._live_routers) @@ -41,6 +215,7 @@ def isolate_router_model_cost_state() -> Iterator[None]: model_key: dict(model_value) for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + litellm_utils_module._invalidate_model_cost_lowercase_map() yield for router in tuple(litellm_router_module._live_routers): litellm_router_module._live_routers.discard(router) @@ -68,4 +243,9 @@ def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: def pytest_sessionfinish() -> None: + for name in MODULE_LEVEL_CLIENTS: + _close_handler_if_needed(litellm.__dict__.pop(name, None)) + for name in SESSION_CLIENTS: + _close_handler_if_needed(getattr(litellm, name, None)) + _run_coroutine_if_needed(close_litellm_async_clients()) enable_socket() diff --git a/tests/test_litellm/a2a_protocol/providers/__init__.py b/tests/unit/containers/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/__init__.py rename to tests/unit/containers/__init__.py diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/unit/containers/test_azure_container_transformation.py similarity index 100% rename from tests/test_litellm/containers/test_azure_container_transformation.py rename to tests/unit/containers/test_azure_container_transformation.py diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/unit/containers/test_container_api.py similarity index 100% rename from tests/test_litellm/containers/test_container_api.py rename to tests/unit/containers/test_container_api.py diff --git a/tests/test_litellm/containers/test_container_handler_url.py b/tests/unit/containers/test_container_handler_url.py similarity index 100% rename from tests/test_litellm/containers/test_container_handler_url.py rename to tests/unit/containers/test_container_handler_url.py diff --git a/tests/test_litellm/containers/test_container_integration.py b/tests/unit/containers/test_container_integration.py similarity index 100% rename from tests/test_litellm/containers/test_container_integration.py rename to tests/unit/containers/test_container_integration.py diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/unit/containers/test_container_proxy_ownership.py similarity index 100% rename from tests/test_litellm/containers/test_container_proxy_ownership.py rename to tests/unit/containers/test_container_proxy_ownership.py diff --git a/tests/test_litellm/containers/test_container_regional_api_base.py b/tests/unit/containers/test_container_regional_api_base.py similarity index 100% rename from tests/test_litellm/containers/test_container_regional_api_base.py rename to tests/unit/containers/test_container_regional_api_base.py diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/unit/containers/test_container_transformation.py similarity index 100% rename from tests/test_litellm/containers/test_container_transformation.py rename to tests/unit/containers/test_container_transformation.py diff --git a/tests/test_litellm/containers/test_container_utils.py b/tests/unit/containers/test_container_utils.py similarity index 100% rename from tests/test_litellm/containers/test_container_utils.py rename to tests/unit/containers/test_container_utils.py diff --git a/tests/test_litellm/containers/test_endpoint_factory.py b/tests/unit/containers/test_endpoint_factory.py similarity index 100% rename from tests/test_litellm/containers/test_endpoint_factory.py rename to tests/unit/containers/test_endpoint_factory.py diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/__init__.py b/tests/unit/embeddings/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/__init__.py rename to tests/unit/embeddings/__init__.py diff --git a/tests/test_litellm/embeddings/test_dispatch.py b/tests/unit/embeddings/test_dispatch.py similarity index 100% rename from tests/test_litellm/embeddings/test_dispatch.py rename to tests/unit/embeddings/test_dispatch.py diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/tests/unit/experimental_mcp_client/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py rename to tests/unit/experimental_mcp_client/__init__.py diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py similarity index 100% rename from tests/test_litellm/experimental_mcp_client/test_mcp_client.py rename to tests/unit/experimental_mcp_client/test_mcp_client.py diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/unit/experimental_mcp_client/test_tools.py similarity index 100% rename from tests/test_litellm/experimental_mcp_client/test_tools.py rename to tests/unit/experimental_mcp_client/test_tools.py diff --git a/tests/test_litellm/batches/__init__.py b/tests/unit/files/__init__.py similarity index 100% rename from tests/test_litellm/batches/__init__.py rename to tests/unit/files/__init__.py diff --git a/tests/test_litellm/files/test_main.py b/tests/unit/files/test_main.py similarity index 100% rename from tests/test_litellm/files/test_main.py rename to tests/unit/files/test_main.py diff --git a/tests/test_litellm/chat_completions/__init__.py b/tests/unit/fixtures/__init__.py similarity index 100% rename from tests/test_litellm/chat_completions/__init__.py rename to tests/unit/fixtures/__init__.py diff --git a/tests/test_litellm/completion_extras/__init__.py b/tests/unit/fixtures/together_ai_sync/__init__.py similarity index 100% rename from tests/test_litellm/completion_extras/__init__.py rename to tests/unit/fixtures/together_ai_sync/__init__.py diff --git a/tests/test_litellm/fixtures/together_ai_sync/deprecations.md b/tests/unit/fixtures/together_ai_sync/deprecations.md similarity index 100% rename from tests/test_litellm/fixtures/together_ai_sync/deprecations.md rename to tests/unit/fixtures/together_ai_sync/deprecations.md diff --git a/tests/test_litellm/fixtures/together_ai_sync/models_serverless.json b/tests/unit/fixtures/together_ai_sync/models_serverless.json similarity index 100% rename from tests/test_litellm/fixtures/together_ai_sync/models_serverless.json rename to tests/unit/fixtures/together_ai_sync/models_serverless.json diff --git a/tests/test_litellm/containers/__init__.py b/tests/unit/google_genai/__init__.py similarity index 100% rename from tests/test_litellm/containers/__init__.py rename to tests/unit/google_genai/__init__.py diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_adapter.py rename to tests/unit/google_genai/test_google_genai_adapter.py diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py b/tests/unit/google_genai/test_google_genai_adapter_fixes.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py rename to tests/unit/google_genai/test_google_genai_adapter_fixes.py diff --git a/tests/test_litellm/google_genai/test_google_genai_handler.py b/tests/unit/google_genai/test_google_genai_handler.py similarity index 76% rename from tests/test_litellm/google_genai/test_google_genai_handler.py rename to tests/unit/google_genai/test_google_genai_handler.py index bf037c59854..5361d91718d 100644 --- a/tests/test_litellm/google_genai/test_google_genai_handler.py +++ b/tests/unit/google_genai/test_google_genai_handler.py @@ -2,99 +2,13 @@ """ Test to verify the Google GenAI generate_content handler functionality """ -import json from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter -from litellm.types.utils import ModelResponse - - -def test_non_stream_response_when_stream_requested_sync(): - """ - Test that when a non-stream response is returned but streaming was requested, - the sync handler correctly transforms it to generate_content format. - """ - from litellm.types.utils import Choices - - # Mock a non-stream response (ModelResponse with valid choices) - mock_response = ModelResponse( - id="test-123", - choices=[ - Choices( - index=0, - message={"role": "assistant", "content": "Hello, world!"}, - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - ) - - # Create an instance of the adapter - adapter = GoogleGenAIAdapter() - - # Test the adapter's translate_completion_to_generate_content method directly - result = adapter.translate_completion_to_generate_content(mock_response) - - # Verify the result is a valid Google GenAI format response - assert "candidates" in result - assert isinstance(result["candidates"], list) - assert len(result["candidates"]) > 0 - candidate = result["candidates"][0] - assert "content" in candidate - assert "parts" in candidate["content"] - assert isinstance(candidate["content"]["parts"], list) - assert len(candidate["content"]["parts"]) > 0 - assert "text" in candidate["content"]["parts"][0] - assert candidate["content"]["parts"][0]["text"] == "Hello, world!" - - -@pytest.mark.asyncio -async def test_non_stream_response_when_stream_requested_async(): - """ - Test that when a non-stream response is returned but streaming was requested, - the async handler correctly transforms it to generate_content format. - """ - from litellm.types.utils import Choices - - # Mock a non-stream response (ModelResponse with valid choices) - mock_response = ModelResponse( - id="test-123", - choices=[ - Choices( - index=0, - message={"role": "assistant", "content": "Hello, world!"}, - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - ) - - # Create an instance of the adapter - adapter = GoogleGenAIAdapter() - - # Test the adapter's translate_completion_to_generate_content method directly - result = adapter.translate_completion_to_generate_content(mock_response) - - # Verify the result is a valid Google GenAI format response - assert "candidates" in result - assert isinstance(result["candidates"], list) - assert len(result["candidates"]) > 0 - candidate = result["candidates"][0] - assert "content" in candidate - assert "parts" in candidate["content"] - assert isinstance(candidate["content"]["parts"], list) - assert len(candidate["content"]["parts"]) > 0 - assert "text" in candidate["content"]["parts"][0] - assert candidate["content"]["parts"][0]["text"] == "Hello, world!" def test_stream_response_when_stream_requested_sync(): diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_main.py rename to tests/unit/google_genai/test_google_genai_main.py diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py rename to tests/unit/google_genai/test_google_genai_streaming_iterator.py diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/unit/google_genai/test_google_genai_transformation.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_transformation.py rename to tests/unit/google_genai/test_google_genai_transformation.py diff --git a/tests/test_litellm/endpoints/__init__.py b/tests/unit/images/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/__init__.py rename to tests/unit/images/__init__.py diff --git a/tests/test_litellm/images/test_image_edit_extra_params.py b/tests/unit/images/test_image_edit_extra_params.py similarity index 100% rename from tests/test_litellm/images/test_image_edit_extra_params.py rename to tests/unit/images/test_image_edit_extra_params.py diff --git a/tests/test_litellm/images/test_image_edit_utils.py b/tests/unit/images/test_image_edit_utils.py similarity index 100% rename from tests/test_litellm/images/test_image_edit_utils.py rename to tests/unit/images/test_image_edit_utils.py diff --git a/tests/test_litellm/images/test_image_generation_extra_headers.py b/tests/unit/images/test_image_generation_extra_headers.py similarity index 100% rename from tests/test_litellm/images/test_image_generation_extra_headers.py rename to tests/unit/images/test_image_generation_extra_headers.py diff --git a/tests/test_litellm/endpoints/speech/__init__.py b/tests/unit/interactions/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/speech/__init__.py rename to tests/unit/interactions/__init__.py diff --git a/tests/test_litellm/interactions/test_agents_http_handler.py b/tests/unit/interactions/test_agents_http_handler.py similarity index 100% rename from tests/test_litellm/interactions/test_agents_http_handler.py rename to tests/unit/interactions/test_agents_http_handler.py diff --git a/tests/test_litellm/interactions/test_agents_main_and_utils.py b/tests/unit/interactions/test_agents_main_and_utils.py similarity index 100% rename from tests/test_litellm/interactions/test_agents_main_and_utils.py rename to tests/unit/interactions/test_agents_main_and_utils.py diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py similarity index 100% rename from tests/test_litellm/interactions/test_background_cost_polling.py rename to tests/unit/interactions/test_background_cost_polling.py diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/unit/interactions/test_gemini_interactions_transformation.py similarity index 100% rename from tests/test_litellm/interactions/test_gemini_interactions_transformation.py rename to tests/unit/interactions/test_gemini_interactions_transformation.py diff --git a/tests/test_litellm/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py similarity index 100% rename from tests/test_litellm/interactions/test_interactions_streaming_iterator.py rename to tests/unit/interactions/test_interactions_streaming_iterator.py diff --git a/tests/unit/interactions/test_litellm_responses_bridge.py b/tests/unit/interactions/test_litellm_responses_bridge.py new file mode 100644 index 00000000000..3abd0a6ca98 --- /dev/null +++ b/tests/unit/interactions/test_litellm_responses_bridge.py @@ -0,0 +1,80 @@ +""" +Tests for LiteLLM Responses bridge provider. + +Inherits from BaseInteractionsTest to run the same test suite against +the litellm_responses bridge provider, which calls litellm.responses() internally. +""" + + +from litellm.interactions.litellm_responses_transformation.transformation import ( + LiteLLMResponsesInteractionsConfig, +) +from litellm.types.interactions import Turn + + +class TestBridgeInputTransformation: + """Regression tests for translating Interactions input into Responses API input. + + The bridge used to pass Google content parts through raw ({"type": "text"}), + which the Responses API rejects with a 400, and it dropped the role encoded + in step types and in the legacy "model" turn role. + """ + + def test_step_input_maps_roles_and_content_types(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [ + {"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]}, + {"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]}, + {"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]}, + ] + ) + assert transformed == [ + {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, + {"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]}, + ] + + def test_legacy_turn_input_maps_model_role_to_assistant(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [ + {"role": "user", "content": [{"type": "text", "text": "I like apples."}]}, + {"role": "model", "content": [{"type": "text", "text": "I like oranges."}]}, + ] + ) + assert transformed == [ + {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, + ] + + def test_turn_pydantic_model_with_string_content(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [Turn(role="model", content="I like oranges.")] + ) + assert transformed == [ + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]} + ] + + def test_string_input_passes_through(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello") + assert transformed == "Hello" + + def test_content_list_input_becomes_single_user_message(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [{"type": "text", "text": "Hello"}, "world"] + ) + assert transformed == [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Hello"}, + {"type": "input_text", "text": "world"}, + ], + } + ] + + def test_non_text_content_passes_through_unchanged(self): + image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"} + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [{"type": "user_input", "content": [image_part]}] + ) + assert transformed == [{"role": "user", "content": [image_part]}] diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py similarity index 99% rename from tests/test_litellm/interactions/test_openapi_compliance.py rename to tests/unit/interactions/test_openapi_compliance.py index 2665f8703a6..d3f1183cea6 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -4,7 +4,7 @@ OpenAPI compliance tests for Google Interactions API. Validates that our SDK requests/responses match the OpenAPI spec at: https://ai.google.dev/static/api/interactions.openapi.json -Run with: pytest tests/test_litellm/interactions/test_openapi_compliance.py -v +Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v """ import json diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index 88ef849f0e2..3d5059b200f 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -22,6 +22,8 @@ from litellm.rust_bridge.messages.entrypoints import ( NativeMessages, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from pydantic import TypeAdapter +from litellm.messages import dispatch MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () @@ -284,3 +286,137 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon NATIVE_AMESSAGES.reset() assert result is expected assert [request.model for request in captured] == ["claude-sonnet-4-5"] + + +@pytest.mark.asyncio +async def test_public_anthropic_messages_keeps_the_python_result() -> None: + response: Final = await litellm.anthropic_messages( + model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok" + ) + + assert isinstance(response, dict) + content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", [])) + assert content[0]["text"] == "ok" + + +def test_sync_messages_request_projects_public_arguments() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + assert request.model == "claude-test" + assert request.messages == MESSAGES + assert request.max_tokens == 10 + assert request.custom_llm_provider == "anthropic" + return expected + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + { + "model": "claude-test", + "messages": MESSAGES, + "max_tokens": 10, + "custom_llm_provider": "anthropic", + }, + python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_messages_binding_error_delegates_unchanged_to_python() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + pytest.fail("a call without max_tokens cannot project a request and must stay on Python") + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + {"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +@pytest.mark.asyncio +async def test_async_messages_falls_back_after_native_declines() -> None: + from litellm.rust_bridge.bindings import native_exception_types + + native_types: Final = native_exception_types() + if native_types is None: + pytest.skip("native bridge is unavailable") + declined, _ = native_types + expected: Final = AnthropicMessagesResponse(model="claude-test") + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),) + + async def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + raise declined("unsupported") + + async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None) + binding.override(native) + response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + {"model": "claude-test", "messages": MESSAGES, "max_tokens": 10}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_internal_is_async_marker_bypasses_native() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + pytest.fail("anthropic_messages' inner handler call must stay on Python") + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + { + "model": "claude-test", + "messages": MESSAGES, + "max_tokens": 10, + "custom_llm_provider": "anthropic", + "is_async": True, + }, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected diff --git a/tests/test_litellm/rag/test_main.py b/tests/unit/rag/test_main.py similarity index 100% rename from tests/test_litellm/rag/test_main.py rename to tests/unit/rag/test_main.py diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/unit/rerank_api/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py rename to tests/unit/rerank_api/__init__.py diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/unit/rerank_api/test_main.py similarity index 100% rename from tests/test_litellm/rerank_api/test_main.py rename to tests/unit/rerank_api/test_main.py diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/unit/test_a2a_registry_lookup.py similarity index 100% rename from tests/test_litellm/test_a2a_registry_lookup.py rename to tests/unit/test_a2a_registry_lookup.py diff --git a/tests/test_litellm/test_acompletion_session_reuse_e2e.py b/tests/unit/test_acompletion_session_reuse_e2e.py similarity index 100% rename from tests/test_litellm/test_acompletion_session_reuse_e2e.py rename to tests/unit/test_acompletion_session_reuse_e2e.py diff --git a/tests/test_litellm/test_add_deployment_no_master_key.py b/tests/unit/test_add_deployment_no_master_key.py similarity index 100% rename from tests/test_litellm/test_add_deployment_no_master_key.py rename to tests/unit/test_add_deployment_no_master_key.py diff --git a/tests/test_litellm/test_aembedding_session_reuse_e2e.py b/tests/unit/test_aembedding_session_reuse_e2e.py similarity index 100% rename from tests/test_litellm/test_aembedding_session_reuse_e2e.py rename to tests/unit/test_aembedding_session_reuse_e2e.py diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py similarity index 100% rename from tests/test_litellm/test_anthropic_beta_headers_filtering.py rename to tests/unit/test_anthropic_beta_headers_filtering.py diff --git a/tests/test_litellm/test_anthropic_skills_transformation.py b/tests/unit/test_anthropic_skills_transformation.py similarity index 100% rename from tests/test_litellm/test_anthropic_skills_transformation.py rename to tests/unit/test_anthropic_skills_transformation.py diff --git a/tests/test_litellm/test_assert_ci_coverage.py b/tests/unit/test_assert_ci_coverage.py similarity index 100% rename from tests/test_litellm/test_assert_ci_coverage.py rename to tests/unit/test_assert_ci_coverage.py diff --git a/tests/test_litellm/test_assert_workflow_dir_hygiene.py b/tests/unit/test_assert_workflow_dir_hygiene.py similarity index 100% rename from tests/test_litellm/test_assert_workflow_dir_hygiene.py rename to tests/unit/test_assert_workflow_dir_hygiene.py diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/unit/test_audio_transcription_rust_bridge.py similarity index 100% rename from tests/test_litellm/test_audio_transcription_rust_bridge.py rename to tests/unit/test_audio_transcription_rust_bridge.py diff --git a/tests/test_litellm/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py similarity index 100% rename from tests/test_litellm/test_auto_update_price_and_context_window_file.py rename to tests/unit/test_auto_update_price_and_context_window_file.py diff --git a/tests/test_litellm/test_azure_ad_token_credential_resolution.py b/tests/unit/test_azure_ad_token_credential_resolution.py similarity index 100% rename from tests/test_litellm/test_azure_ad_token_credential_resolution.py rename to tests/unit/test_azure_ad_token_credential_resolution.py diff --git a/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py b/tests/unit/test_azure_ai_grok_4_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py rename to tests/unit/test_azure_ai_grok_4_3_model_metadata.py diff --git a/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py b/tests/unit/test_azure_ai_grok_4_6_model_metadata.py similarity index 100% rename from tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py rename to tests/unit/test_azure_ai_grok_4_6_model_metadata.py diff --git a/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py b/tests/unit/test_baseten_glm_5_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_baseten_glm_5_3_model_metadata.py rename to tests/unit/test_baseten_glm_5_3_model_metadata.py diff --git a/tests/test_litellm/test_batch_completion_models_all_responses.py b/tests/unit/test_batch_completion_models_all_responses.py similarity index 100% rename from tests/test_litellm/test_batch_completion_models_all_responses.py rename to tests/unit/test_batch_completion_models_all_responses.py diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/unit/test_bedrock_marengo_embed_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py rename to tests/unit/test_bedrock_marengo_embed_3_model_metadata.py diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/unit/test_budget_ratchet_check.py similarity index 100% rename from tests/test_litellm/test_budget_ratchet_check.py rename to tests/unit/test_budget_ratchet_check.py diff --git a/tests/test_litellm/test_chat_ui_responses_session.py b/tests/unit/test_chat_ui_responses_session.py similarity index 100% rename from tests/test_litellm/test_chat_ui_responses_session.py rename to tests/unit/test_chat_ui_responses_session.py diff --git a/tests/test_litellm/test_check_licenses.py b/tests/unit/test_check_licenses.py similarity index 100% rename from tests/test_litellm/test_check_licenses.py rename to tests/unit/test_check_licenses.py diff --git a/tests/test_litellm/test_check_mcp_operation_boundary.py b/tests/unit/test_check_mcp_operation_boundary.py similarity index 100% rename from tests/test_litellm/test_check_mcp_operation_boundary.py rename to tests/unit/test_check_mcp_operation_boundary.py diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/unit/test_check_migrations_no_data_rewrites.py similarity index 100% rename from tests/test_litellm/test_check_migrations_no_data_rewrites.py rename to tests/unit/test_check_migrations_no_data_rewrites.py diff --git a/tests/test_litellm/test_check_py310_typing_imports.py b/tests/unit/test_check_py310_typing_imports.py similarity index 100% rename from tests/test_litellm/test_check_py310_typing_imports.py rename to tests/unit/test_check_py310_typing_imports.py diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/unit/test_check_test_quality.py similarity index 100% rename from tests/test_litellm/test_check_test_quality.py rename to tests/unit/test_check_test_quality.py diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/unit/test_check_type_discipline.py similarity index 100% rename from tests/test_litellm/test_check_type_discipline.py rename to tests/unit/test_check_type_discipline.py diff --git a/tests/test_litellm/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py similarity index 100% rename from tests/test_litellm/test_circleci_path_filter.py rename to tests/unit/test_circleci_path_filter.py diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/unit/test_circleci_rust_toolchain.py similarity index 100% rename from tests/test_litellm/test_circleci_rust_toolchain.py rename to tests/unit/test_circleci_rust_toolchain.py diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/unit/test_claude_fable_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_fable_5_config.py rename to tests/unit/test_claude_fable_5_config.py diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/unit/test_claude_opus_4_6_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_4_6_config.py rename to tests/unit/test_claude_opus_4_6_config.py diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/unit/test_claude_opus_4_8_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_4_8_config.py rename to tests/unit/test_claude_opus_4_8_config.py diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/unit/test_claude_opus_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_5_config.py rename to tests/unit/test_claude_opus_5_config.py diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/unit/test_claude_sonnet_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_sonnet_5_config.py rename to tests/unit/test_claude_sonnet_5_config.py diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/unit/test_cloudflare_workers_ai_model_metadata.py similarity index 100% rename from tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py rename to tests/unit/test_cloudflare_workers_ai_model_metadata.py diff --git a/tests/test_litellm/test_completion_timeout_resolution.py b/tests/unit/test_completion_timeout_resolution.py similarity index 100% rename from tests/test_litellm/test_completion_timeout_resolution.py rename to tests/unit/test_completion_timeout_resolution.py diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/unit/test_component_entrypoint.py similarity index 100% rename from tests/test_litellm/test_component_entrypoint.py rename to tests/unit/test_component_entrypoint.py diff --git a/tests/unit/test_compression.py b/tests/unit/test_compression.py new file mode 100644 index 00000000000..be718f03963 --- /dev/null +++ b/tests/unit/test_compression.py @@ -0,0 +1,649 @@ +""" +Unit tests for litellm.compress(). +""" + +import importlib + +import pytest + +import litellm +from litellm.compression.scoring.bm25 import bm25_score_messages +from litellm.compression.scoring.embedding_scorer import embedding_score_messages +from litellm.compression.content_detection import detect_content_type +from litellm.compression.message_stubbing import extract_key, stub_message +from litellm.compression.retrieval_tool import build_retrieval_tool +from litellm.types.utils import CallTypes + +CALL_TYPE = CallTypes.completion +ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages + + +# --------------------------------------------------------------------------- +# BM25 scorer +# --------------------------------------------------------------------------- + + +def test_bm25_relevance_ranking(): + query = "Fix the authentication bug in the login handler" + messages = [ + { + "role": "user", + "content": "def login_handler(): authentication check bug fix", + }, + {"role": "user", "content": "def render_template(name): css styling layout"}, + {"role": "user", "content": "def verify(): authentication token bug handler"}, + ] + scores = bm25_score_messages(query, messages) + # Messages sharing query terms should score higher than unrelated ones + assert scores[0] > scores[1] + assert scores[2] > scores[1] + + +def test_bm25_empty_query(): + scores = bm25_score_messages("", [{"role": "user", "content": "hello"}]) + assert scores == [0.0] + + +def test_bm25_empty_messages(): + scores = bm25_score_messages("query", []) + assert scores == [] + + +def test_bm25_empty_content(): + scores = bm25_score_messages("query", [{"role": "user", "content": ""}]) + assert scores == [0.0] + + +# --------------------------------------------------------------------------- +# Content detection +# --------------------------------------------------------------------------- + + +def test_detect_code(): + code = """ +import os +from pathlib import Path + +def main(): + class Foo: + pass + return Foo() +""" + assert detect_content_type(code) == "code" + + +def test_detect_json(): + assert detect_content_type('{"key": "value", "num": 42}') == "json" + assert detect_content_type("[1, 2, 3]") == "json" + + +def test_detect_text(): + assert detect_content_type("This is a plain text paragraph about dogs.") == "text" + + +def test_detect_empty(): + assert detect_content_type("") == "text" + + +# --------------------------------------------------------------------------- +# Message stubbing +# --------------------------------------------------------------------------- + + +def test_extract_key_with_filename(): + msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"} + used: set = set() + key = extract_key(msg, fallback_index=0, used_keys=used) + assert key == "auth.py" + + +def test_extract_key_fallback(): + msg = {"role": "user", "content": "Some random content without a filename"} + used: set = set() + key = extract_key(msg, fallback_index=5, used_keys=used) + assert key == "message_5" + + +def test_extract_key_duplicates(): + used: set = set() + msg = {"role": "user", "content": "# auth.py\ncode here"} + k1 = extract_key(msg, fallback_index=0, used_keys=used) + k2 = extract_key(msg, fallback_index=1, used_keys=used) + assert k1 == "auth.py" + assert k2 == "auth.py_2" + + +def test_stub_message(): + msg = {"role": "user", "content": "line1\nline2\nline3"} + stubbed = stub_message(msg, "test_key") + assert stubbed["role"] == "user" + assert "test_key" in stubbed["content"] + assert "litellm_content_retrieve" in stubbed["content"] + assert "3 lines" in stubbed["content"] + + +# --------------------------------------------------------------------------- +# Retrieval tool +# --------------------------------------------------------------------------- + + +def test_retrieval_tool_schema(): + tool = build_retrieval_tool(["auth.py", "utils.py"]) + assert tool["type"] == "function" + assert tool["function"]["name"] == "litellm_content_retrieve" + assert "key" in tool["function"]["parameters"]["properties"] + assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [ + "auth.py", + "utils.py", + ] + assert tool["function"]["parameters"]["required"] == ["key"] + + +def test_retrieval_tool_description_lists_keys(): + tool = build_retrieval_tool(["foo.py", "bar.js"]) + desc = tool["function"]["description"] + assert "foo.py" in desc + assert "bar.js" in desc + + +# --------------------------------------------------------------------------- +# compress() — end-to-end +# --------------------------------------------------------------------------- + + +def test_compress_below_trigger_passthrough(): + messages = [{"role": "user", "content": "hello"}] + result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_ratio"] == 0.0 + assert result["original_tokens"] == result["compressed_tokens"] + + +def test_compress_above_trigger(): + big_messages = [ + {"role": "system", "content": "You are a coding assistant."}, + { + "role": "user", + "content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000, + }, + { + "role": "user", + "content": "# utils.py\n" + "def helper():\n pass\n" * 2000, + }, + { + "role": "user", + "content": "# readme.md\n" + "This is documentation. " * 2000, + }, + {"role": "user", "content": "Fix the bug in auth.py"}, + ] + + result = litellm.compress( + big_messages, + model="gpt-4o", + call_type=CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert result["compression_ratio"] > 0 + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" + + +def test_compress_anthropic_list_content_is_boundary_stable(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/a.png"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "# b.py\n" + "beta " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/b.png"}, + }, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert len(result["messages"]) == len(messages) + assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "custom" + assert result["tools"][0]["name"] == "litellm_content_retrieve" + assert "input_schema" in result["tools"][0] + + +def test_compress_preserves_system_message(): + messages = [ + {"role": "system", "content": "System prompt. " * 500}, + {"role": "user", "content": "Large file content. " * 5000}, + {"role": "user", "content": "Fix the bug"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + assert result["messages"][0]["role"] == "system" + assert "System prompt" in result["messages"][0]["content"] + + +def test_compress_preserves_last_user_message(): + messages = [ + {"role": "user", "content": "Big context " * 5000}, + {"role": "user", "content": "Fix the bug in auth.py"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + last_user = [m for m in result["messages"] if m["role"] == "user"][-1] + assert "Fix the bug in auth.py" in last_user["content"] + + +def test_compress_preserves_last_assistant_message(): + messages = [ + {"role": "user", "content": "Big context " * 5000}, + {"role": "assistant", "content": "I'll help with that. " * 2000}, + {"role": "user", "content": "Now fix the bug"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] + assert len(assistant_msgs) >= 1 + # The last assistant message should be preserved (not stubbed) + last_assistant = assistant_msgs[-1] + assert "I'll help with that" in last_assistant["content"] + + +def test_cache_keys_match_stubs(): + messages = [ + {"role": "user", "content": "# auth.py\n" + "code " * 5000}, + {"role": "user", "content": "Fix it"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + if result["tools"]: + tool_desc = result["tools"][0]["function"]["description"] + for key in result["cache"]: + assert key in tool_desc + + +def test_compress_default_target(): + """compression_target defaults to compression_trigger // 2.""" + messages = [ + {"role": "user", "content": "content " * 5000}, + {"role": "user", "content": "query"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 + ) + # Should have compressed — target = 1000 + assert result["compressed_tokens"] <= result["original_tokens"] + + +def test_compress_nested_tool_result_extracts_text_only(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "prefix"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "nested text fragment"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/secret-tool.png", + }, + }, + ], + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/top.png"}, + }, + {"type": "text", "text": " " + ("irrelevant " * 3000)}, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "final query that must remain"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=500, + compression_target=100, + ) + + cached_text = " ".join(result["cache"].values()) + assert "nested text fragment" in cached_text + assert "https://example.com/secret-tool.png" not in cached_text + assert "https://example.com/top.png" not in cached_text + + +def test_compress_default_call_type_is_completion(): + result = litellm.compress( + messages=[ + {"role": "user", "content": "Large context " * 4000}, + {"role": "user", "content": "query"}, + ], + model="gpt-4o", + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert isinstance(result["tools"], list) + + +def test_compress_forwards_embedding_model_params(monkeypatch): + captured = {} + + def fake_embedding_score_messages( + query, messages, model, cache=None, embedding_model_params=None + ): + captured["query"] = query + captured["model"] = model + captured["embedding_model_params"] = embedding_model_params + return [0.0] * len(messages) + + monkeypatch.setattr( + "litellm.compression.scoring.embedding_scorer.embedding_score_messages", + fake_embedding_score_messages, + ) + + result = litellm.compress( + messages=[ + {"role": "user", "content": "Authentication code " * 2000}, + {"role": "user", "content": "Fix auth"}, + ], + model="gpt-4o", + call_type=CALL_TYPE, + compression_trigger=1000, + embedding_model="text-embedding-3-small", + embedding_model_params={"api_base": "https://example-embeddings.test"}, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert captured["model"] == "text-embedding-3-small" + assert captured["embedding_model_params"] == { + "api_base": "https://example-embeddings.test" + } + + +def test_embedding_scorer_forwards_embedding_model_params(monkeypatch): + captured = {} + + class _MockResponse: + data = [ + {"embedding": [1.0, 0.0]}, + {"embedding": [1.0, 0.0]}, + {"embedding": [0.0, 1.0]}, + ] + + def fake_embedding(**kwargs): + captured.update(kwargs) + return _MockResponse() + + monkeypatch.setattr(litellm, "embedding", fake_embedding) + + scores = embedding_score_messages( + query="auth", + messages=[ + {"role": "user", "content": "auth code"}, + {"role": "user", "content": "cooking recipe"}, + ], + model="text-embedding-3-small", + embedding_model_params={"api_base": "https://example-embeddings.test"}, + ) + + assert len(scores) == 2 + assert captured["model"] == "text-embedding-3-small" + assert captured["api_base"] == "https://example-embeddings.test" + + +# --------------------------------------------------------------------------- +# Embedding scorer — integration test (skipped without API key) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "final_user_message, expected_content", + [ + ("How to cook?", "Unrelated cooking recipes "), + ("Fix auth", "Authentication code "), + ], +) +def test_simple_compression(final_user_message, expected_content): + messages = [ + {"role": "user", "content": "Authentication code " * 2000}, + {"role": "user", "content": "Unrelated cooking recipes " * 2000}, + {"role": "user", "content": final_user_message}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + if expected_content == "Unrelated cooking recipes ": + assert "Unrelated cooking recipes " in result["messages"][1]["content"] + assert "Authentication code " not in result["messages"][0]["content"] + elif expected_content == "Authentication code ": + assert "Authentication code " in result["messages"][0]["content"] + assert "Unrelated cooking recipes " not in result["messages"][1]["content"] + else: + raise ValueError(f"Unexpected expected_content: {expected_content}") + + +def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) + return [0.95, 0.01, 0.02, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_drop", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_drop", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) + assert len(result["messages"]) == 3 + assert result["messages"][0]["role"] == "user" + assert "other_blob" in result["messages"][0]["content"] + assert result["messages"][1]["content"] == "assistant_tail" + assert result["messages"][2]["content"] == "final query" + assert result["cache"] == {} + + +def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer the tool exchange span over idx=0 + return [0.05, 0.01, 0.92, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_keep", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_keep", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert len(result["messages"]) == 5 + assert result["messages"][1]["role"] == "assistant" + assert result["messages"][2]["role"] == "user" + # idx=0 should be compressed instead + assert "litellm_content_retrieve" in result["messages"][0]["content"] + assert len(result["cache"]) == 1 + + +def test_compress_anthropic_malformed_tool_sequence_passes_through(): + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_broken", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py similarity index 100% rename from tests/test_litellm/test_conftest_isolation.py rename to tests/unit/test_conftest_isolation.py diff --git a/tests/test_litellm/test_constants.py b/tests/unit/test_constants.py similarity index 100% rename from tests/test_litellm/test_constants.py rename to tests/unit/test_constants.py diff --git a/tests/test_litellm/test_container_router.py b/tests/unit/test_container_router.py similarity index 100% rename from tests/test_litellm/test_container_router.py rename to tests/unit/test_container_router.py diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/unit/test_cost_calculation_log_level.py similarity index 100% rename from tests/test_litellm/test_cost_calculation_log_level.py rename to tests/unit/test_cost_calculation_log_level.py diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/unit/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/test_cost_calculator.py rename to tests/unit/test_cost_calculator.py diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py similarity index 100% rename from tests/test_litellm/test_cost_map_guard.py rename to tests/unit/test_cost_map_guard.py diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py similarity index 100% rename from tests/test_litellm/test_count_tokens_public_api.py rename to tests/unit/test_count_tokens_public_api.py diff --git a/tests/test_litellm/test_dashscope_image_generation.py b/tests/unit/test_dashscope_image_generation.py similarity index 99% rename from tests/test_litellm/test_dashscope_image_generation.py rename to tests/unit/test_dashscope_image_generation.py index 1dd0b322623..6f91fe9a0e0 100644 --- a/tests/test_litellm/test_dashscope_image_generation.py +++ b/tests/unit/test_dashscope_image_generation.py @@ -2,7 +2,7 @@ Unit tests for DashScope image generation support (qwen-image-2.0, qwen-image-2.0-pro, qwen-image-3.0, qwen-image-3.0-pro). -Run in docker: pytest tests/test_litellm/test_dashscope_image_generation.py -v +Run in docker: pytest tests/unit/test_dashscope_image_generation.py -v """ from unittest.mock import MagicMock, patch diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/unit/test_daybreak_model_metadata.py similarity index 100% rename from tests/test_litellm/test_daybreak_model_metadata.py rename to tests/unit/test_daybreak_model_metadata.py diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/unit/test_deepseek_model_metadata.py similarity index 100% rename from tests/test_litellm/test_deepseek_model_metadata.py rename to tests/unit/test_deepseek_model_metadata.py diff --git a/tests/test_litellm/test_default_branch.py b/tests/unit/test_default_branch.py similarity index 100% rename from tests/test_litellm/test_default_branch.py rename to tests/unit/test_default_branch.py diff --git a/tests/test_litellm/test_detect_changes.py b/tests/unit/test_detect_changes.py similarity index 100% rename from tests/test_litellm/test_detect_changes.py rename to tests/unit/test_detect_changes.py diff --git a/tests/test_litellm/test_dockerfile_apk_repository.py b/tests/unit/test_dockerfile_apk_repository.py similarity index 100% rename from tests/test_litellm/test_dockerfile_apk_repository.py rename to tests/unit/test_dockerfile_apk_repository.py diff --git a/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py b/tests/unit/test_dockerfile_bedrock_realtime_extra.py similarity index 100% rename from tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py rename to tests/unit/test_dockerfile_bedrock_realtime_extra.py diff --git a/tests/test_litellm/test_dockerfile_non_root.py b/tests/unit/test_dockerfile_non_root.py similarity index 100% rename from tests/test_litellm/test_dockerfile_non_root.py rename to tests/unit/test_dockerfile_non_root.py diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/unit/test_drop_params_env_var.py similarity index 100% rename from tests/test_litellm/test_drop_params_env_var.py rename to tests/unit/test_drop_params_env_var.py diff --git a/tests/test_litellm/test_e2e_egress_sentinel.py b/tests/unit/test_e2e_egress_sentinel.py similarity index 100% rename from tests/test_litellm/test_e2e_egress_sentinel.py rename to tests/unit/test_e2e_egress_sentinel.py diff --git a/tests/test_litellm/test_eager_tiktoken_load.py b/tests/unit/test_eager_tiktoken_load.py similarity index 100% rename from tests/test_litellm/test_eager_tiktoken_load.py rename to tests/unit/test_eager_tiktoken_load.py diff --git a/tests/test_litellm/test_env_key_doc_gate.py b/tests/unit/test_env_key_doc_gate.py similarity index 100% rename from tests/test_litellm/test_env_key_doc_gate.py rename to tests/unit/test_env_key_doc_gate.py diff --git a/tests/test_litellm/test_exception_exports.py b/tests/unit/test_exception_exports.py similarity index 100% rename from tests/test_litellm/test_exception_exports.py rename to tests/unit/test_exception_exports.py diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py similarity index 100% rename from tests/test_litellm/test_exception_header_preservation.py rename to tests/unit/test_exception_header_preservation.py diff --git a/tests/test_litellm/test_exception_mapping_request_attribute.py b/tests/unit/test_exception_mapping_request_attribute.py similarity index 100% rename from tests/test_litellm/test_exception_mapping_request_attribute.py rename to tests/unit/test_exception_mapping_request_attribute.py diff --git a/tests/test_litellm/test_filter_out_litellm_params.py b/tests/unit/test_filter_out_litellm_params.py similarity index 100% rename from tests/test_litellm/test_filter_out_litellm_params.py rename to tests/unit/test_filter_out_litellm_params.py diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/unit/test_fireworks_serverless_model_costs.py similarity index 100% rename from tests/test_litellm/test_fireworks_serverless_model_costs.py rename to tests/unit/test_fireworks_serverless_model_costs.py diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/unit/test_gate_slot_lock.py similarity index 100% rename from tests/test_litellm/test_gate_slot_lock.py rename to tests/unit/test_gate_slot_lock.py diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/unit/test_gemini_3_1_flash_lite_image_pricing.py similarity index 100% rename from tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py rename to tests/unit/test_gemini_3_1_flash_lite_image_pricing.py diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/unit/test_gemini_tts_native_audio_pricing.py similarity index 100% rename from tests/test_litellm/test_gemini_tts_native_audio_pricing.py rename to tests/unit/test_gemini_tts_native_audio_pricing.py diff --git a/tests/test_litellm/test_get_blog_posts.py b/tests/unit/test_get_blog_posts.py similarity index 100% rename from tests/test_litellm/test_get_blog_posts.py rename to tests/unit/test_get_blog_posts.py diff --git a/tests/test_litellm/test_git_hooks.py b/tests/unit/test_git_hooks.py similarity index 100% rename from tests/test_litellm/test_git_hooks.py rename to tests/unit/test_git_hooks.py diff --git a/tests/test_litellm/test_gpt_5_4_model_metadata.py b/tests/unit/test_gpt_5_4_model_metadata.py similarity index 100% rename from tests/test_litellm/test_gpt_5_4_model_metadata.py rename to tests/unit/test_gpt_5_4_model_metadata.py diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/unit/test_gpt_5_5_model_metadata.py similarity index 100% rename from tests/test_litellm/test_gpt_5_5_model_metadata.py rename to tests/unit/test_gpt_5_5_model_metadata.py diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/unit/test_gpt_image_cost_calculator.py similarity index 100% rename from tests/test_litellm/test_gpt_image_cost_calculator.py rename to tests/unit/test_gpt_image_cost_calculator.py diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/unit/test_gpt_realtime_mode.py similarity index 100% rename from tests/test_litellm/test_gpt_realtime_mode.py rename to tests/unit/test_gpt_realtime_mode.py diff --git a/tests/test_litellm/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py similarity index 100% rename from tests/test_litellm/test_groq_streaming_encoding.py rename to tests/unit/test_groq_streaming_encoding.py diff --git a/tests/test_litellm/test_guardrail_exception_status_codes.py b/tests/unit/test_guardrail_exception_status_codes.py similarity index 100% rename from tests/test_litellm/test_guardrail_exception_status_codes.py rename to tests/unit/test_guardrail_exception_status_codes.py diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/unit/test_lazy_imports.py similarity index 100% rename from tests/test_litellm/test_lazy_imports.py rename to tests/unit/test_lazy_imports.py diff --git a/tests/test_litellm/test_lint_workflow_diff_gates.py b/tests/unit/test_lint_workflow_diff_gates.py similarity index 100% rename from tests/test_litellm/test_lint_workflow_diff_gates.py rename to tests/unit/test_lint_workflow_diff_gates.py diff --git a/tests/test_litellm/test_litellm_params_reserved_keys.py b/tests/unit/test_litellm_params_reserved_keys.py similarity index 100% rename from tests/test_litellm/test_litellm_params_reserved_keys.py rename to tests/unit/test_litellm_params_reserved_keys.py diff --git a/tests/test_litellm/test_logging.py b/tests/unit/test_logging.py similarity index 100% rename from tests/test_litellm/test_logging.py rename to tests/unit/test_logging.py diff --git a/tests/test_litellm/test_lowest_latency_zero_tokens.py b/tests/unit/test_lowest_latency_zero_tokens.py similarity index 100% rename from tests/test_litellm/test_lowest_latency_zero_tokens.py rename to tests/unit/test_lowest_latency_zero_tokens.py diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py new file mode 100644 index 00000000000..effc038f85b --- /dev/null +++ b/tests/unit/test_main.py @@ -0,0 +1,4124 @@ +import asyncio +import base64 +from datetime import datetime +import contextlib +import copy +import json +import logging +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +import respx + + +import urllib.parse +from importlib import import_module +from unittest.mock import MagicMock, patch + +import litellm +from litellm import main as litellm_main +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +@pytest.fixture(autouse=True) +def clear_client_cache(): + """ + Clear the HTTP client cache before each test to ensure mocks are used. + This prevents cached real clients from being reused across tests. + """ + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is not None: + cache.flush_cache() + yield + if cache is not None: + cache.flush_cache() + + +@pytest.fixture(autouse=True) +def add_api_keys_to_env(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-api03-1234567890") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-api03-1234567890") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "my-fake-aws-access-key-id") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "my-fake-aws-secret-access-key") + monkeypatch.setenv("AWS_REGION", "us-east-1") + # Keep these transformation tests on the simple access-key path. A leaked + # session token or role/web-identity env var pushes Bedrock auth down a + # different branch and fails before the mocked HTTP client is exercised. + monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + + +@pytest.fixture +def openai_api_response(): + mock_response_data = { + "id": "chatcmpl-B0W3vmiM78Xkgx7kI7dr7PC949DMS", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "", + "refusal": None, + "role": "assistant", + "audio": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + "created": 1739462947, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion", + "service_tier": "default", + "system_fingerprint": "fp_bd83329f63", + "usage": { + "completion_tokens": 1, + "prompt_tokens": 121, + "total_tokens": 122, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, + }, + } + + return mock_response_data + + +def test_completion_missing_role(openai_api_response): + from openai import OpenAI + + from litellm.types.utils import ModelResponse + + client = OpenAI(api_key="test_api_key") + + mock_raw_response = MagicMock() + mock_raw_response.headers = { + "x-request-id": "123", + "openai-organization": "org-123", + "x-ratelimit-limit-requests": "100", + "x-ratelimit-remaining-requests": "99", + } + mock_raw_response.parse.return_value = ModelResponse(**openai_api_response) + + print(f"openai_api_response: {openai_api_response}") + + with patch.object( + client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) + ) as mock_create: + litellm.completion( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "Hey"}, + { + "content": "", + "tool_calls": [ + { + "id": "call_m0vFJjQmTH1McvaHBPR2YFwY", + "function": { + "arguments": '{"input": "dksjsdkjdhskdjshdskhjkhlk"}', + "name": "tool_name", + }, + "type": "function", + "index": 0, + }, + { + "id": "call_Vw6RaqV2n5aaANXEdp5pYxo2", + "function": { + "arguments": '{"input": "jkljlkjlkjlkjlk"}', + "name": "tool_name", + }, + "type": "function", + "index": 1, + }, + { + "id": "call_hBIKwldUEGlNh6NlSXil62K4", + "function": { + "arguments": '{"input": "jkjlkjlkjlkj;lj"}', + "name": "tool_name", + }, + "type": "function", + "index": 2, + }, + ], + }, + ], + client=client, + ) + + mock_create.assert_called_once() + + +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_url_with_format_param_openai(model, sync_mode): + from openai import AsyncOpenAI, OpenAI + + from litellm import acompletion, completion + + if sync_mode: + client = OpenAI() + else: + client = AsyncOpenAI() + + args = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", + "format": "image/png", + }, + }, + {"type": "text", "text": "Describe this image"}, + ], + } + ], + } + with patch.object( + client.chat.completions.with_raw_response, "create" + ) as mock_client: + try: + if sync_mode: + response = completion(**args, client=client) + else: + response = await acompletion(**args, client=client) + print(response) + except Exception as e: + print(e) + + mock_client.assert_called() + + print(mock_client.call_args.kwargs) + + json_str = json.dumps(mock_client.call_args.kwargs) + + assert "format" not in json_str + + +def test_bedrock_latency_optimized_inference(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + with patch.object(client, "post") as mock_post: + try: + response = litellm.completion( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello, how are you?"}], + performanceConfig={"latency": "optimized"}, + client=client, + ) + except Exception as e: + print(e) + + mock_post.assert_called_once() + json_data = json.loads(mock_post.call_args.kwargs["data"]) + assert json_data["performanceConfig"]["latency"] == "optimized" + + +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected + + +@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) +def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): + tools = [ + {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, + "opaque_tool", + ] + + cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) + + assert cleaned == [ + {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, + "opaque_tool", + ] + assert tools[0][key] is True + assert tools[0]["function"][key] is True + + +def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + + litellm.completion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Write the file"}], + tools=[ + { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, + "eager_input_streaming": True, + } + ], + api_base=api_base, + api_key="fake_openai_api_key", + ) + + assert mock_route.called + sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] + assert "eager_input_streaming" not in sent_tool + assert sent_tool["function"]["name"] == "write_file" + + +def test_custom_provider_with_extra_headers(): + + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + headers={"X-Custom-Header": "custom-value"}, + api_base="https://example.com/api/v1", + ) + + mock_post.assert_called_once() + assert mock_post.call_args[1]["headers"]["X-Custom-Header"] == "custom-value" + + +def test_custom_provider_with_extra_body(): + + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + extra_body={ + "X-Custom-BodyValue": "custom-value", + "X-Custom-BodyValue2": "custom-value2", + }, + api_base="https://example.com/api/v1", + ) + mock_post.assert_called_once() + + assert mock_post.call_args[1]["json"]["X-Custom-BodyValue"] == "custom-value" + assert mock_post.call_args[1]["json"] == { + "model": "custom", + "params": { + "prompt": ["Hello, how are you?"], + "max_tokens": None, + "temperature": None, + "top_p": None, + "top_k": None, + }, + "X-Custom-BodyValue": "custom-value", + "X-Custom-BodyValue2": "custom-value2", + } + + # test that extra_body is not passed if not provided + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + api_base="https://example.com/api/v1", + ) + mock_post.assert_called_once() + assert mock_post.call_args[1]["json"] == { + "model": "custom", + "params": { + "prompt": ["Hello, how are you?"], + "max_tokens": None, + "temperature": None, + "top_p": None, + "top_k": None, + }, + } + + +@pytest.fixture(autouse=True) +def set_openrouter_api_key(): + original_api_key = os.environ.get("OPENROUTER_API_KEY") + os.environ["OPENROUTER_API_KEY"] = "fake-key-for-testing" + yield + if original_api_key is not None: + os.environ["OPENROUTER_API_KEY"] = original_api_key + else: + del os.environ["OPENROUTER_API_KEY"] + + +@pytest.mark.asyncio +async def test_extra_body_with_fallback( + respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch +): + """ + test regression for https://github.com/BerriAI/litellm/issues/8425. + + This was perhaps a wider issue with the acompletion function not passing kwargs such as extra_body correctly when fallbacks are specified. + """ + + # Save original state to restore after test + original_disable_aiohttp = litellm.disable_aiohttp_transport + + try: + # since this uses respx, we need to set use_aiohttp_transport to False + # Set both the global variable and environment variable to ensure it takes effect + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + # Flush cache to ensure no stale aiohttp clients are used + litellm.in_memory_llm_clients_cache.flush_cache() + + # Set up test parameters + model = "openrouter/deepseek/deepseek-chat" + messages = [{"role": "user", "content": "Hello, world!"}] + extra_body = { + "provider": { + "order": ["DeepSeek"], + "allow_fallbacks": False, + "require_parameters": True, + } + } + fallbacks = [{"model": "openrouter/google/gemini-flash-1.5-8b"}] + + # Set up mock to respond to any POST request to the OpenRouter endpoint + # This ensures it works for both primary and fallback models + mock_route = respx_mock.post("https://openrouter.ai/api/v1/chat/completions") + mock_route.return_value = httpx.Response( + 200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + + response = await litellm.acompletion( + model=model, + messages=messages, + extra_body=extra_body, + fallbacks=fallbacks, + api_key="fake-openrouter-api-key", + ) + + # Verify the response + assert response is not None + assert ( + len(respx_mock.calls) > 0 + ), "Mock was not called - check if aiohttp transport is properly disabled" + + # Get the request from the mock + request: httpx.Request = respx_mock.calls[0].request + request_body = request.read() + request_body = json.loads(request_body) + + # Verify basic parameters + assert request_body["model"] == "deepseek/deepseek-chat" + assert request_body["messages"] == messages + + # Verify the extra_body parameters remain under the provider key + assert request_body["provider"]["order"] == ["DeepSeek"] + assert request_body["provider"]["allow_fallbacks"] is False + assert request_body["provider"]["require_parameters"] is True + finally: + # Restore original state to prevent test pollution + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.parametrize("env_base", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_openai_env_base( + respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch +): + "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" + # Ensure aiohttp transport is disabled to use httpx which respx can mock + litellm.disable_aiohttp_transport = True + + expected_base_url = "http://localhost:12345/v1" + + # Assign the environment variable based on env_base, and use a fake API key. + monkeypatch.setenv(env_base, expected_base_url) + monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") + + model = "gpt-4o" + messages = [{"role": "user", "content": "Hello, how are you?"}] + + # Configure respx mock to intercept the request + mock_route = respx_mock.post( + url__regex=r"http://localhost:12345/v1/chat/completions.*" + ).mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + ) + + try: + response = await litellm.acompletion(model=model, messages=messages) + + # verify we had a response + assert response.choices[0].message.content == "Hello from mocked response!" + + # Verify the mock was called + assert ( + mock_route.called + ), "Mock route was not called - request may have bypassed respx" + finally: + # Clean up to avoid affecting other tests + litellm.disable_aiohttp_transport = False + + +def build_database_url(username, password, host, dbname): + username_enc = urllib.parse.quote_plus(username) + password_enc = urllib.parse.quote_plus(password) + dbname_enc = urllib.parse.quote_plus(dbname) + return f"postgresql://{username_enc}:{password_enc}@{host}/{dbname_enc}" + + +def test_build_database_url(): + url = build_database_url("user@name", "p@ss:word", "localhost", "db/name") + assert url == "postgresql://user%40name:p%40ss%3Aword@localhost/db%2Fname" + + +def test_bedrock_llama(): + litellm._turn_on_debug() + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "bedrock/invoke/us.meta.llama4-scout-17b-instruct-v1:0" + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": [ + {"role": "user", "content": "hi"}, + ], + }, + ) + print(request) + + assert ( + request["raw_request_body"]["prompt"] + == "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + ) + + +def _mocked_openai_chat_response(model: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + + +def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter): + """Regression for #33952: return_raw_request must transform without contacting the provider. + + Previously return_raw_request invoked the real endpoint with a fake key and relied on the + provider rejecting it, which sent an unintended inference request and (in the async proxy + route) blocked the event loop on provider I/O. + """ + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "gpt-4o" + route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert route.call_count == 0 + assert request.get("error") is None + assert request["raw_request_body"]["model"] == model + assert request["raw_request_body"]["messages"] == [ + {"role": "user", "content": "hi"} + ] + + +def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter): + """Regression test: completion() must forward the verbosity param to the provider request body.""" + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "gpt-5.2" + messages = [{"role": "user", "content": "hi"}] + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": messages, + "verbosity": "high", + }, + ) + + assert request["raw_request_body"]["verbosity"] == "high" + assert request["raw_request_body"]["model"] == model + assert request["raw_request_body"]["messages"] == messages + + +@pytest.mark.asyncio +async def test_acompletion_forwards_verbosity_to_provider_request( + respx_mock: respx.MockRouter, monkeypatch +): + """Regression test: acompletion() must forward the verbosity param to the provider request body.""" + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + + model = "gpt-5.2" + messages = [{"role": "user", "content": "hi"}] + mock_route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + response = await litellm.acompletion( + model=model, + messages=messages, + verbosity="low", + api_key="fake-openai-api-key", + ) + + assert response.choices[0].message.content == "Hello from mocked response!" + assert mock_route.called + request_body = json.loads(respx_mock.calls[0].request.read()) + assert request_body["verbosity"] == "low" + assert request_body["model"] == model + assert request_body["messages"] == messages + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +def test_responses_api_bridge_check_strips_responses_prefix(): + """Test that responses_api_bridge_check strips 'responses/' prefix and sets mode.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + + model_info, model = responses_api_bridge_check( + model="responses/gpt-4-responses", + custom_llm_provider="openai", + ) + + assert model == "gpt-4-responses" + assert model_info["mode"] == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_pro(): + """Test that gpt-5.4-pro routes through responses API bridge, not chat completions. + + Regression test for https://github.com/BerriAI/litellm/issues/23014 + gpt-5.4-pro is a responses-only model and must not be sent to /v1/chat/completions. + """ + from litellm.main import responses_api_bridge_check + + for model_name in ["gpt-5.4-pro", "gpt-5.4-pro-2026-03-05"]: + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="openai", + ) + assert ( + model_info.get("mode") == "responses" + ), f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" + + +def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_responses(): + """gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="xhigh", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) + + assert model == "gpt-6-astra" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): + """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.5-pro", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="xhigh", + ) + + assert model == "gpt-5.5-pro" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to_responses(): + """Azure gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="high", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): + """ + Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables + reasoning by default for gpt-5.4+, and Chat Completions rejects function tools + whenever reasoning is on. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): + """ + gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning + by default for gpt-5.4+, and Chat Completions rejects function tools whenever + reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "model_name, expected_mode", + [ + pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), + pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), + ], +) +def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( + monkeypatch, model_name, expected_mode +): + """ + gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit + reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with + "Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions". + + Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The + gate parses the version and drops any suffix, so the family members bridge + identically and only the boundary distinguishes behaviour. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == model_name + assert model_info.get("mode") == expected_mode + + +def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): + """ + Explicit reasoning_effort "none" is OpenAI's documented escape hatch that keeps + function tools servable on Chat Completions; the bridge must not fire. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="none", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_reasoning_none_with_summary_still_routes_to_responses(): + """A reasoning summary is Responses-only regardless of effort value.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + reasoning_effort="none", + reasoning_summary="detailed", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_custom_tools_only_stays_chat(): + """ + Chat Completions serves custom (grammar) tools natively with reasoning on; only + FUNCTION tools trigger the OpenAI rejection. Custom-only requests must stay on chat + so responses keep the native custom tool_call shape instead of the bridge's + function-shaped mapping. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "custom", "custom": {"name": "ApplyPatch", "description": "V4A patch"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_gpt_5_4_mixed_function_and_custom_tools_routes_to_responses(): + """One function tool in the mix is enough to make chat unservable with reasoning on.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[ + {"type": "custom", "custom": {"name": "ApplyPatch"}}, + {"type": "function", "function": {"name": "shell"}}, + ], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_flat_function_tool_routes_to_responses(): + """Responses-style flat function tool defs still count as function tools.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "name": "shell", "parameters": {"type": "object"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "custom_llm_provider, model_name, api_base", + [ + pytest.param("openai", "gpt-5.6", None, id="openai"), + pytest.param("azure_ai", "gpt-6-astra", "https://myproject.services.ai.azure.com", id="azure-ai-foundry"), + ], +) +def test_responses_api_bridge_check_function_tool_without_body_stays_chat( + monkeypatch, custom_llm_provider, model_name, api_base +): + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider=custom_llm_provider, + tools=[{"type": "function"}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_dict_effort_none_stays_chat(): + """The escape hatch must honor litellm's dict form: {"effort": "none"} means reasoning off.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "none"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_dict_effort_active_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "low"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_dict_effort_none_with_summary_routes_to_responses(): + """A summary inside the dict form is Responses-only even when effort is none.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "none", "summary": "concise"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize("blank_api_base", [None, "", " ", "\t"]) +def test_responses_api_bridge_check_blank_api_base_is_default_openai(blank_api_base): + """ + A blank api_base (None, empty, or whitespace) resolves to the default OpenAI + endpoint downstream, which enforces the reasoning+tools constraint, so gpt-5.4+ + function-tool requests with unset reasoning_effort must still auto-bridge. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=blank_api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_custom_api_base_with_unset_effort_stays_chat(): + """ + Chat-only OpenAI-compatible backends registered under the openai provider with a + custom api_base and gpt-5.4+ model names serve tools-without-reasoning fine and + have no /responses route; the unset-effort arm must not reroute them. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base="http://vllm.internal:8000/v1", + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_custom_api_base_via_global_with_unset_effort_stays_chat(monkeypatch): + """ + A custom base set through the litellm.api_base global (not the call arg) is resolved the + same way the chat handler resolves it, so the unset-effort arm must not reroute a chat-only + backend to a /responses route it lacks. Regression guard: the gate previously inspected only + the call-level api_base and bridged these requests. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", "http://vllm.internal:8000/v1") + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +@pytest.mark.parametrize("env_var", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) +def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_stays_chat(monkeypatch, env_var): + """ + A custom base set via OPENAI_BASE_URL/OPENAI_API_BASE env is resolved identically to the chat + handler, so the unset-effort arm leaves the request on chat instead of bridging it. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv(env_var, "http://vllm.internal:8000/v1") + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://southcentralus.privatelink.api.openai.com/v1", + "https://privatelink.corp.api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://api.openai.com/v1/", + "HTTPS://API.OPENAI.COM/v1", + ], +) +def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): + """ + A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a + port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects + function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as + it does for the literal default URL. Regression guard for GH #39353. + """ + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): + """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): + """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): + """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="high", + api_base="http://vllm.internal:8000/v1", + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): + """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base="https://myresource.openai.azure.com", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +_FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" +_FOUNDRY_FUNCTION_TOOL: Final = ({"type": "function", "function": {"name": "get_weather"}},) + + +@pytest.mark.parametrize( + "model_name, api_base, reasoning_effort", + [ + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, None, id="gpt-6-unset-effort"), + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "low", id="gpt-6-explicit-effort"), + pytest.param("gpt-6-astra", "https://myresource.openai.azure.com", None, id="gpt-6-azure-openai-host"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "low", id="gpt-5.6-explicit-effort"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, {"effort": "high"}, id="gpt-5.6-explicit-effort-dict"), + ], +) +def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_responses( + model_name, api_base, reasoning_effort +): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="azure_ai", + tools=_FOUNDRY_FUNCTION_TOOL, + reasoning_effort=reasoning_effort, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "model_name, api_base, reasoning_effort", + [ + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="explicit-none-stays-chat"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, None, id="gpt-5.6-unset-effort-stays-chat"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "none", id="gpt-5.6-explicit-none-stays-chat"), + pytest.param("gpt-5.5", _FOUNDRY_API_BASE, "high", id="gpt-5.5-explicit-effort-stays-chat"), + pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, None, id="gpt-5.4-mini-unset-effort-stays-chat"), + pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, "low", id="gpt-5.4-mini-explicit-effort-stays-chat"), + pytest.param("gpt-6-astra", "https://myproject.models.ai.azure.com", None, id="serverless-host-stays-chat"), + pytest.param("Mistral-large-2411", _FOUNDRY_API_BASE, None, id="non-gpt-5-model-stays-chat"), + pytest.param("claude-opus-4-1", _FOUNDRY_API_BASE, None, id="claude-on-foundry-stays-chat"), + ], +) +def test_responses_api_bridge_check_azure_ai_without_foundry_responses_route_stays_chat( + model_name, api_base, reasoning_effort +): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="azure_ai", + tools=_FOUNDRY_FUNCTION_TOOL, + reasoning_effort=reasoning_effort, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_older_gpt_5_tools_without_reasoning_stays_chat(): + """Pre-5.4 GPT-5 names keep the old boundary: tools alone never bridge.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.1", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.1" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_gpt_5_4_reasoning_summary_without_tools_routes_to_responses(): + """gpt-5.4+ with reasoning_effort + reasoningSummary but no tools should bridge (AI SDK).""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=None, + reasoning_effort="medium", + reasoning_summary="auto", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_reasoning_summary_routes_to_responses(): + """Bare ``gpt-5`` with reasoning_effort + reasoningSummary should bridge (not 5.4+).""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5", + custom_llm_provider="openai", + tools=None, + reasoning_effort="medium", + reasoning_summary="auto", + ) + + assert model == "gpt-5" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_tools_without_summary_stays_chat(): + """gpt-5 with tools + reasoning_effort but no summary should stay on chat.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="medium", + reasoning_summary=None, + ) + + assert model == "gpt-5" + assert model_info.get("mode") != "responses" + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( + mock_responses_completion, +): + """When routed to Responses, preserve reasoning_effort summary dict.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5.4", + messages=[{"role": "user", "content": "What is the capital of France?"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_capital", + "description": "Get the capital of a country", + "parameters": { + "type": "object", + "properties": {"country": {"type": "string"}}, + }, + }, + } + ], + reasoning_effort={"effort": "xhigh", "summary": "detailed"}, + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == { + "effort": "xhigh", + "summary": "detailed", + } + + +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": False, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} + + +_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY: Final = { + "id": "resp_foundry", + "object": "response", + "created_at": 1789852145, + "status": "completed", + "model": "gpt-6-astra", + "output": [ + { + "id": "fc_1", + "type": "function_call", + "status": "completed", + "arguments": '{"city":"Paris"}', + "call_id": "call_1", + "name": "get_weather", + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 53, + "output_tokens": 18, + "total_tokens": 71, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "max_output_tokens": 200, + "previous_response_id": None, + "reasoning": {"effort": "medium", "summary": None}, + "truncation": "disabled", + "user": None, +} + + +def test_completion_bridges_azure_ai_foundry_gpt_5_4_plus_function_tools_to_responses( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + responses_route: Final = respx_mock.post(f"{_FOUNDRY_API_BASE}/openai/v1/responses").respond( + json=_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY + ) + + response: Final = litellm.completion( + model="azure_ai/gpt-6-astra", + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + } + ], + max_tokens=200, + api_base=_FOUNDRY_API_BASE, + api_key="fake-foundry-key", + ) + + assert [str(call.request.url) for call in respx_mock.calls] == [f"{_FOUNDRY_API_BASE}/openai/v1/responses"] + request: Final = responses_route.calls[0].request + request_body: Final = json.loads(request.content) + assert request_body["tools"][0]["type"] == "function" + assert request_body["tools"][0]["name"] == "get_weather" + assert request.headers["api-key"] == "fake-foundry-key" + assert response.choices[0].finish_reason == "tool_calls" + assert response.choices[0].message.tool_calls[0].function.name == "get_weather" + + +@pytest.mark.parametrize( + "model, model_info, expected_model_param, expected_base_model_param", + [ + ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None), + ( + "gemini/gemini-3.1-pro", + {"base_model": "gemini-3.1-pro-preview"}, + "gemini-3.1-pro", + "gemini-3.1-pro-preview", + ), + ], +) +def test_completion_optional_params_base_model( + model: str, + model_info: dict | None, + expected_model_param: str, + expected_base_model_param: str | None, +): + """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model`` + (an additive capability hint), without overwriting ``model`` with the label. + + Regression for #29618: overwriting ``model`` with a friendly ``base_model`` + label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``.""" + with patch("litellm.main.get_optional_params") as mock_get_optional_params: + mock_get_optional_params.return_value = MagicMock() + + import litellm + + kwargs = { + "model": model, + "messages": [{"role": "user", "content": "What is the capital of France?"}], + "api_key": "fake-key", + "mock_response": "Hey, how's it going?", + } + if model_info is not None: + kwargs["model_info"] = model_info + + litellm.completion(**kwargs) + + assert mock_get_optional_params.called is True + call_kwargs = mock_get_optional_params.call_args.kwargs + assert call_kwargs["model"] == expected_model_param + assert call_kwargs["base_model"] == expected_base_model_param + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_4_responses_bridge_merges_reasoning_summary_kwarg_without_tools( + mock_responses_completion, +): + """reasoningSummary without tools should route and merge into reasoning_effort dict.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5.4", + messages=[{"role": "user", "content": "ok"}], + reasoning_effort="medium", + reasoningSummary="auto", + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == { + "effort": "medium", + "summary": "auto", + } + assert "reasoningSummary" not in optional_params + assert "reasoning_summary" not in optional_params + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_responses_bridge_preserves_reasoning_summary_without_effort( + mock_responses_completion, +): + """Reasoning summary should survive responses routing even without effort.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "ok"}], + reasoningSummary="auto", + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == {"summary": "auto"} + assert "reasoningSummary" not in optional_params + assert "reasoning_summary" not in optional_params + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_responses_bridge_tools_and_reasoning_summary( + mock_responses_completion, +): + """Bare gpt-5 with tools + reasoningSummary should bridge (OpenCode-style).""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5", + messages=[{"role": "user", "content": "ok"}], + tools=[ + { + "type": "function", + "function": { + "name": "apply_patch", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice="auto", + reasoning_effort="medium", + reasoningSummary="auto", + stream=True, + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params.get("reasoning_effort") == { + "effort": "medium", + "summary": "auto", + } + + +def test_responses_api_bridge_check_handles_exception(): + """Test that responses_api_bridge_check handles exceptions and still processes responses/ models.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.side_effect = Exception("Model not found") + + model_info, model = responses_api_bridge_check( + model="responses/custom-model", custom_llm_provider="custom" + ) + + assert model == "custom-model" + assert model_info["mode"] == "responses" + + +def test_responses_api_bridge_check_global_flag_routes_openai(): + """When route_all_chat_openai_to_responses is True, any OpenAI model routes to responses.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="openai", + ) + + assert model == "gpt-4o" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_global_flag_does_not_affect_azure(): + """route_all_chat_openai_to_responses should not affect Azure models.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="azure", + ) + + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_global_flag_default_false(): + """By default, route_all_chat_openai_to_responses is False and doesn't affect routing.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", False): + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="openai", + ) + + assert model_info.get("mode") != "responses" + + +@pytest.mark.asyncio +async def test_async_mock_delay(): + """Use asyncio await for mock delay on acompletion""" + import time + + from litellm import acompletion + + start_time = time.time() + result = await acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_delay=0.01, + mock_response="Hello world", + ) + end_time = time.time() + delay = end_time - start_time + assert delay >= 0.01 + + +def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): + from litellm import stream_chunk_builder + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-multi-choice", + created=1751934860, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=choices, + ) + + chunks = [ + chunk( + [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), + StreamingChoices( + index=1, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + index=0, + type="function", + function=Function(name="lookup_fruit", arguments='{"fruit":'), + ) + ], + ), + ), + ] + ), + chunk( + [ + StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), + StreamingChoices( + index=1, + delta=Delta( + tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] + ), + finish_reason="tool_calls", + ), + ] + ), + ] + + response = stream_chunk_builder(chunks=chunks) + + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ + ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') + ] + + +def test_stream_chunk_builder_thinking_blocks(): + from litellm import stream_chunk_builder + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + chunks = [ + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="I need to summar", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "I need to summar", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "I need to summar", + "signature": None, + } + ] + }, + content="", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="ize the previous agent's thinking process into a", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "ize the previous agent's thinking process into a", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "ize the previous agent's thinking process into a", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" short description. Based on the input data provide", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " short description. Based on the input data provide", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " short description. Based on the input data provide", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="d, it seems the agent was planning to refine their search", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "d, it seems the agent was planning to refine their search", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "d, it seems the agent was planning to refine their search", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" to focus more on technical aspects of home automation and home", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " to focus more on technical aspects of home automation and home", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " to focus more on technical aspects of home automation and home", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" energy system management.\n\nI'll create a brief", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " energy system management.\n\nI'll create a brief", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " energy system management.\n\nI'll create a brief", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" summary of what the agent was doing.", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " summary of what the agent was doing.", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " summary of what the agent was doing.", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "", + "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "", + "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='{"a', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='gent_doing"', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content=': "Re', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="searching", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content=" technic", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="al aspect", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="s of home au", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='tomation"}', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="tool_calls", + index=0, + delta=Delta( + provider_specific_fields=None, + content=None, + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + ), + ] + + response = stream_chunk_builder(chunks=chunks) + print(response) + + assert response is not None + assert response.choices[0].message.content is not None + assert response.choices[0].message.thinking_blocks is not None + + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +def throw_retryable_error(*_, **__): + raise RuntimeError("BOOM") + + +@pytest.mark.asyncio +async def test_retrying() -> None: + litellm.num_retries = 10 + with ( + patch.object( + OpenAIChatCompletion, + "make_openai_chat_completion_request", + side_effect=throw_retryable_error, + ) as mock_request, + pytest.raises(litellm.InternalServerError, match="LiteLLM Retried: 10 times"), + ): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + ) + + +def test_anthropic_disable_url_suffix_env_var(): + """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/messages suffix.""" + import os + from unittest.mock import MagicMock, patch + + from litellm import completion + + # Test with environment variable disabled (default behavior) + with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): + actual_api_base = None + + with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + mock_anthropic.completion = capture_completion + + # This should append /v1/messages + completion( + model="anthropic/claude-3-sonnet", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base has /v1/messages appended + assert actual_api_base.endswith("/v1/messages") + assert actual_api_base == "https://api.example.com/v1/messages" + + # Test with environment variable enabled + with patch.dict( + os.environ, + { + "ANTHROPIC_API_BASE": "https://api.example.com/custom/path", + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", + }, + ): + actual_api_base = None + + with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + mock_anthropic.completion = capture_completion + + # This should NOT append /v1/messages + completion( + model="anthropic/claude-3-sonnet", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base does not have /v1/messages appended + assert actual_api_base == "https://api.example.com/custom/path" + assert not actual_api_base.endswith("/v1/messages") + + +def test_anthropic_text_disable_url_suffix_env_var(): + """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/complete suffix for anthropic_text.""" + import os + from unittest.mock import MagicMock, patch + + from litellm import completion + + # Test with environment variable disabled (default behavior) + with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): + actual_api_base = None + + with patch("litellm.main.base_llm_http_handler") as mock_handler: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + return MagicMock() + + mock_handler.completion = capture_completion + + # This should append /v1/complete + completion( + model="anthropic_text/claude-instant-1", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base has /v1/complete appended + assert actual_api_base.endswith("/v1/complete") + assert actual_api_base == "https://api.example.com/v1/complete" + + # Test with environment variable enabled + with patch.dict( + os.environ, + { + "ANTHROPIC_API_BASE": "https://api.example.com/custom/complete", + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", + }, + ): + actual_api_base = None + + with patch("litellm.main.base_llm_http_handler") as mock_handler: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + return MagicMock() + + mock_handler.completion = capture_completion + + # This should NOT append /v1/complete + completion( + model="anthropic_text/claude-instant-1", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base does not have /v1/complete appended + assert actual_api_base == "https://api.example.com/custom/complete" + assert not actual_api_base.endswith("/v1/complete") + + +def test_image_edit_merges_headers_and_extra_headers(): + from litellm.images.main import base_llm_http_handler + + combined_headers = { + "x-test-header-one": "value-1", + "x-test-header-two": "value-2", + } + + mock_image_edit_config = MagicMock() + mock_image_edit_config.get_supported_openai_params.return_value = set() + mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict( + kwargs["image_edit_optional_params"] + ) + + with ( + patch( + "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", + return_value=mock_image_edit_config, + ) as mock_config, + patch.object( + base_llm_http_handler, + "image_edit_handler", + return_value="ok", + ) as mock_handler, + ): + response = litellm.image_edit( + image=MagicMock(name="image"), + prompt="test", + model="azure/gpt-image-1", + headers={"x-test-header-one": "value-1"}, + extra_headers={ + "x-test-header-two": "value-2", + }, + ) + + assert response == "ok" + mock_config.assert_called_once() + + handler_kwargs = mock_handler.call_args.kwargs + assert handler_kwargs["extra_headers"] == combined_headers + assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] + + +@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) +@pytest.mark.parametrize("input_tokens", (51234, 0)) +def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}}, + ) + + assert response.usage.prompt_tokens == input_tokens + assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens + + +def test_mock_completion_usage_falls_back_to_default_without_admission_count(): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + + assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + + +_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT: Final = { + "model_name": "azure-ai-custom-priced", + "litellm_params": { + "model": "azure_ai/gpt-5.6", + "api_key": "mock", + "api_base": "https://example.services.ai.azure.com", + "mock_response": "ok", + "input_cost_per_token": 3e-6, + "output_cost_per_token": 7e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 5e-7, + }, + "model_info": {"id": "azure-ai-custom-priced-deployment-id"}, +} + + +def _expected_custom_price(response: litellm.ModelResponse) -> float: + params: Final = _AZURE_AI_CUSTOM_PRICED_DEPLOYMENT["litellm_params"] + return ( + response.usage.prompt_tokens * params["input_cost_per_token"] + + response.usage.completion_tokens * params["output_cost_per_token"] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", (False, True)) +async def test_mock_completion_prices_azure_ai_router_deployment_with_custom_pricing(use_async: bool): + router: Final = litellm.Router(model_list=[_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT]) + messages: Final = [{"role": "user", "content": "hello"}] + + response: Final = ( + await router.acompletion(model="azure-ai-custom-priced", messages=messages) + if use_async + else router.completion(model="azure-ai-custom-priced", messages=messages) + ) + + assert response._hidden_params["response_cost"] == pytest.approx(_expected_custom_price(response)) + assert response._hidden_params["custom_llm_provider"] == "azure_ai" + + +@pytest.mark.parametrize( + ("model", "expected_provider"), + (("anthropic/claude-sonnet-5", "anthropic"), ("no-such-provider-model", None)), +) +def test_mock_completion_infers_provider_when_called_directly_without_one(model: str, expected_provider: str | None): + response: Final = litellm.mock_completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + assert response._hidden_params.get("custom_llm_provider") == expected_provider + + +_ADMISSION_INPUT_TOKENS: Final = 51234 + + +def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata + return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} + + +_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) +_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] +_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" + + +def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: + return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] + + +def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: + return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + +@pytest.mark.parametrize("n", (None, 2)) +def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", (None, 2)) +async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( + n: int | None, +): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + metadata=_ADMISSION_METADATA, + ) + ) + + assert _client_usage_chunks(chunks) == [] + assert all(len(chunk.choices) == 1 for chunk in chunks) + assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + metadata=_ADMISSION_METADATA, + ) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +def _usage_triple(usage: Usage) -> tuple[int, int, int]: + return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) + + +@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) +def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): + metadata: Final = _admission_metadata(input_tokens) + non_stream: Final = litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + metadata=metadata, + ) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) + assert non_stream.usage.prompt_tokens == input_tokens + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=[{"role": "user", "content": ""}], + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + litellm_metadata=_admission_metadata(0), + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): + metadata: Final = _admission_metadata(0) + non_stream: Final = litellm.text_completion( + model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata + ) + chunks: Final = list( + litellm.text_completion( + model="openai/gpt-5.4-mini", + prompt="", + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) + assert len(stream_usages) == 1 + assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) + assert non_stream.usage.prompt_tokens == 0 + + +def test_mock_completion_stream_with_model_response(): + """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" + from litellm import completion + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # Create a ModelResponse object + mock_model_response = ModelResponse( + id="chatcmpl-test-123", + created=1234567890, + model="gpt-4o-mini", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="This is a test response", + role="assistant", + ), + ) + ], + usage=Usage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + ), + ) + + # Call completion with stream=True and mock_response as ModelResponse + response = completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + mock_response=mock_model_response, + ) + + # Verify that the response is a stream + assert response is not None + + # Collect all chunks from the stream + chunks = [] + for chunk in response: + chunks.append(chunk) + print(f"Chunk: {chunk}") + + # Verify we got chunks + assert len(chunks) > 0 + + # Verify the content is streamed correctly + accumulated_content = "" + for chunk in chunks: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): + accumulated_content += chunk.choices[0].delta.content + + assert "This is a test response" in accumulated_content or len(chunks) > 0 + + +@pytest.mark.asyncio +async def test_async_mock_completion_stream_with_model_response(): + """Test that async mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" + from litellm import acompletion + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # Create a ModelResponse object + mock_model_response = ModelResponse( + id="chatcmpl-test-456", + created=1234567890, + model="gpt-4o-mini", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="This is an async test response", + role="assistant", + ), + ) + ], + usage=Usage( + prompt_tokens=15, + completion_tokens=25, + total_tokens=40, + ), + ) + + # Call acompletion with stream=True and mock_response as ModelResponse + response = await acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello async"}], + stream=True, + mock_response=mock_model_response, + ) + + # Verify that the response is a stream + assert response is not None + + # Collect all chunks from the stream + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"Async Chunk: {chunk}") + + # Verify we got chunks + assert len(chunks) > 0 + + # Verify the content is streamed correctly + accumulated_content = "" + for chunk in chunks: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): + accumulated_content += chunk.choices[0].delta.content + + assert "This is an async test response" in accumulated_content or len(chunks) > 0 + + +class TestCallTypesOCR: + """Test that OCR call types are properly defined in CallTypes enum. + + Fixes https://github.com/BerriAI/litellm/issues/17381 + """ + + def test_ocr_call_type_exists(self): + """Test that CallTypes.ocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "ocr") + assert CallTypes.ocr.value == "ocr" + + def test_aocr_call_type_exists(self): + """Test that CallTypes.aocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "aocr") + assert CallTypes.aocr.value == "aocr" + + def test_ocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'ocr' string.""" + from litellm.types.utils import CallTypes + + call_type = CallTypes("ocr") + assert call_type == CallTypes.ocr + + def test_aocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'aocr' string. + + This is the actual use case that was failing - the OCR endpoint + uses route_type='aocr' and guardrails try to instantiate + CallTypes('aocr'). + """ + from litellm.types.utils import CallTypes + + call_type = CallTypes("aocr") + assert call_type == CallTypes.aocr + + +def test_stream_chunk_builder_text_completion_combines_text_and_usage(): + from litellm.main import stream_chunk_builder_text_completion + from litellm.types.utils import TextCompletionResponse + + chunks = [ + TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="gpt-3.5-turbo-instruct", + choices=[{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}], + ), + TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="gpt-3.5-turbo-instruct", + choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], + ), + ] + + response = stream_chunk_builder_text_completion( + chunks=chunks, messages=[{"role": "user", "content": "say hello"}] + ) + + assert response.choices[0].text == "Hello world" + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +def test_completion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/33184 + + store and prompt_cache_key are documented OpenAI chat completion params that + were accepted as supported but silently dropped before the provider request + was built, because they were not named parameters of completion() and + get_optional_params() the way safety_identifier is. + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Async variant of the store/prompt_cache_key forwarding regression test for + https://github.com/BerriAI/litellm/issues/33184 + """ + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): + """ + When store and prompt_cache_key are not passed, they must not appear in the + outbound request body (guards against always forwarding None defaults). + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert "store" not in request_body + assert "prompt_cache_key" not in request_body + + +def test_completion_forwards_store_and_prompt_cache_key_to_mcp_gateway(): + """ + Regression test for the MCP gateway early-return in completion(): store and + prompt_cache_key are named params, so they no longer travel via **kwargs and + must be forwarded explicitly like safety_identifier and service_tier. + """ + with patch.object( + import_module("litellm.responses.mcp.chat_completions_handler"), "acompletion_with_mcp" + ) as mock_mcp: + result = litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy"}], + store=False, + prompt_cache_key="test-cache-key", + ) + + result.close() + mock_mcp.assert_called_once() + call_kwargs = mock_mcp.call_args.kwargs + assert call_kwargs["store"] is False + assert call_kwargs["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "aws_credential_kwargs", + [ + { + "aws_session_name": "litellm-gcp", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", + "aws_web_identity_token": "oidc/google/108963886734710037768", + }, + { + "aws_access_key_id": "AKIASTATICKEYFORTEST", + "aws_secret_access_key": "static-secret-key", + "aws_session_token": "static-session-token", + }, + ], + ids=["web_identity", "static_keys"], +) +async def test_acompletion_forwards_aws_credentials_through_responses_bridge( + respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict +): + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) + monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) + + respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( + json={ + "id": "resp_123", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "openai.gpt-5.4", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ) + + response = await litellm.acompletion( + model="bedrock_mantle/openai.gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + aws_region_name="us-east-2", + num_retries=0, + **aws_credential_kwargs, + ) + + assert response.choices[0].message.content == "ok" + credential_kwargs = get_credentials_mock.call_args.kwargs + assert credential_kwargs["aws_region_name"] == "us-east-2" + for key, value in aws_credential_kwargs.items(): + assert credential_kwargs[key] == value + authorization = respx_mock.calls.last.request.headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "fake-key" in authorization + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +_GEMINI_RESPONSE_BODY = { + "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, +} + + +def _gemini_client_returning_a_reply(): + """An injected HTTP client whose post() answers like generativelanguage does.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + request = httpx.Request("POST", "https://generativelanguage.googleapis.com/") + post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request)) + return client, post + + +@pytest.fixture +def restore_model_registry(): + """litellm.model_cost and the provider name sets are module-global. + + register_model merges into the existing entry in place, hence the deep copy. + """ + model_cost = copy.deepcopy(litellm.model_cost) + openai_models = set(litellm.open_ai_chat_completion_models) + yield + litellm.model_cost.clear() + litellm.model_cost.update(model_cost) + litellm.open_ai_chat_completion_models.clear() + litellm.open_ai_chat_completion_models.update(openai_models) + + +def test_openai_model_name_does_not_outrank_explicit_provider(): + """`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler. + + completion() checks `model in litellm.open_ai_chat_completion_models` ahead of + the gemini branch, so the call used to reach the OpenAI handler carrying + VertexGeminiConfig, whose transform_request raises NotImplementedError. + """ + assert "gpt-4o" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gpt-4o", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert "models/gpt-4o" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry): + """register_model is the other way into the same failure. + + An entry claiming litellm_provider "openai" adds its name to + open_ai_chat_completion_models, so one mislabelled price reroutes every later + call to that model in the process. + """ + litellm.register_model( + { + "gemini-2.5-pro": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + } + } + ) + assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gemini-2.5-pro", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_openai_model_without_a_provider_still_routes_to_openai(): + from openai import OpenAI + + client = OpenAI(api_key="fake-key") + raw_response = client.chat.completions.with_raw_response + with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + mock_create.assert_called() + + +def _openai_chat_create_kwargs(client, **completion_kwargs): + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with contextlib.suppress(Exception): + litellm.completion( + messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], + cache_control_injection_points=[{"location": "message", "role": "system"}], + client=client, + **completion_kwargs, + ) + + mock_client.assert_called_once() + return mock_client.call_args.kwargs + + +@pytest.fixture +def _no_openai_api_base_override(monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_no_openai_api_base_override") +async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with contextlib.suppress(Exception): + await litellm.acompletion( + model="gpt-5.6", + messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], + cache_control_injection_points=[{"location": "message", "role": "system"}], + client=client, + base_url="http://127.0.0.1:9/v1", + ) + + mock_create.assert_called_once() + request_body = mock_create.call_args.kwargs + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") + + assert request_body["messages"][0]["content"] == [ + {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "explicit"} + + +_SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" + + +def _scoped_headers_for_oauth_request(): + from litellm.types.utils import ProviderSpecificHeader + + return [ + ProviderSpecificHeader( + custom_llm_provider="anthropic,bedrock,vertex_ai", + extra_headers={"anthropic-version": "2023-06-01"}, + ), + ProviderSpecificHeader( + custom_llm_provider="anthropic", + extra_headers={"authorization": _SUBSCRIPTION_OAUTH_CREDENTIAL}, + ), + ] + + +def _run_anthropic_hop_with_shared_headers(shared_headers): + litellm.completion( + model="anthropic/claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Say OK"}], + extra_headers=shared_headers, + provider_specific_header=_scoped_headers_for_oauth_request(), + api_key="sk-fake-anthropic-key", + mock_response="OK", + ) + + +def test_completion_does_not_mutate_caller_supplied_headers(): + shared_headers = {"x-tenant": "acme"} + + _run_anthropic_hop_with_shared_headers(shared_headers) + + assert shared_headers == {"x-tenant": "acme"} + + +def test_anthropic_oauth_credential_does_not_persist_into_next_provider_hop(): + shared_headers = {"x-tenant": "acme"} + + _run_anthropic_hop_with_shared_headers(shared_headers) + + leaked = [name for name, value in shared_headers.items() if value == _SUBSCRIPTION_OAUTH_CREDENTIAL] + assert leaked == [] + assert "anthropic-version" not in shared_headers + + +STREAM_COST_MODEL = "gpt-4o" +STREAMED_USAGE = {"prompt_tokens": 137, "completion_tokens": 42, "total_tokens": 179} + + +def _text_chunk(content, finish_reason=None, usage=None): + chunk = { + "id": "chatcmpl-stream-cost", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": STREAM_COST_MODEL, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": content}, + "finish_reason": finish_reason, + } + ], + } + if usage is not None: + chunk["usage"] = usage + return chunk + + +def _priced_at(prompt_tokens, completion_tokens): + prices = litellm.model_cost[STREAM_COST_MODEL] + return ( + prompt_tokens * prices["input_cost_per_token"] + + completion_tokens * prices["output_cost_per_token"] + ) + + +@pytest.fixture +def local_cost_map(monkeypatch): + """The prices these tests assert are the checked-in ones. Setting the environment + variable alone does not reload the map, so pin the map itself. + + Prices are read through two separate lru_caches, so pinning ``model_cost`` is not + enough on its own: an entry warmed against the network-fetched map keeps its old + prices and billing reads those while the assertions read the pinned map. + ``_invalidate_model_cost_lowercase_map`` clears both caches, where + ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out + so entries never leak across tests in either direction.""" + from litellm.utils import _invalidate_model_cost_lowercase_map + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + yield + _invalidate_model_cost_lowercase_map() + + +def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.choices[0].message.content == "Hello there" + assert rebuilt.usage.prompt_tokens == STREAMED_USAGE["prompt_tokens"] + assert rebuilt.usage.completion_tokens == STREAMED_USAGE["completion_tokens"] + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost == pytest.approx(_priced_at(137, 42)) + + +def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + whole = litellm.ModelResponse( + id="chatcmpl-stream-cost", + model=STREAM_COST_MODEL, + object="chat.completion", + created=1700000000, + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello there"}, + "finish_reason": "stop", + } + ], + usage=STREAMED_USAGE, + ) + + assert litellm.completion_cost( + completion_response=rebuilt, model=STREAM_COST_MODEL + ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) + + +def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop"), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.usage.prompt_tokens > 0 + assert rebuilt.usage.completion_tokens > 0 + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost > 0 + assert cost == pytest.approx( + _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) + ) + + +@pytest.mark.asyncio +async def test_acompletion_resolves_provider_from_api_base(): + response = await litellm.acompletion( + model="deepseek-chat", + api_base="https://api.deepseek.com/v1", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + mock_response="resolved", + ) + + assert response.choices[0].message.content == "resolved" + + +@dataclass(frozen=True, slots=True) +class _RecordedSpeechSuccess: + call_type: str | None + spend_metadata: Mapping[str, object] + response_cost: float | None + logged_response_cost: float | None + + +def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess: + call_type: Final = payload.get("call_type") + response_cost: Final = payload.get("response_cost") + logging_payload: Final = payload.get("standard_logging_object") + logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None + return _RecordedSpeechSuccess( + call_type=call_type if isinstance(call_type, str) else None, + spend_metadata=get_litellm_metadata_from_kwargs(payload), + response_cost=response_cost if isinstance(response_cost, float) else None, + logged_response_cost=logged_cost if isinstance(logged_cost, float) else None, + ) + + +class _SuccessEventRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events.append(_record_speech_success(kwargs)) + + +async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: str) -> _RecordedSpeechSuccess: + for _ in range(100): + if (event := next((e for e in recorder.events if e.call_type == call_type), None)) is not None: + return event + await asyncio.sleep(0.05) + pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}") + + +def _gemini_tts_generate_content_response() -> dict[str, object]: + return { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "audio/L16;codec=pcm;rate=24000", + "data": base64.b64encode(b"pcm-audio-bytes").decode(), + } + } + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 60, + "totalTokenCount": 65, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], + "candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": 60}], + }, + "modelVersion": "gemini-2.5-flash-preview-tts", + } + + +@pytest.mark.asyncio +async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + recorder: Final = _SuccessEventRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + mock_route: Final = respx_mock.post( + url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" + ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) + + await litellm.aspeech( + model="gemini/gemini-2.5-flash-preview-tts", + input="spend tracking check", + voice="Kore", + api_key="fake-gemini-key", + metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"}, + ) + + assert mock_route.called + assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key" + speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech") + assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key" + assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1" + expected_prompt_cost, expected_completion_cost = litellm.cost_per_token( + model="gemini/gemini-2.5-flash-preview-tts", + usage_object=Usage(prompt_tokens=5, completion_tokens=60, total_tokens=65), + ) + expected_cost: Final = expected_prompt_cost + expected_completion_cost + assert expected_cost > 0 + assert speech_event.response_cost == pytest.approx(expected_cost) + assert speech_event.logged_response_cost == pytest.approx(expected_cost) + + +def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-cost", + created=1724900000, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content, role="assistant"))], + ) + + +def test_stream_chunk_builder_sets_hidden_response_cost_for_known_model(): + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) + expected_cost: Final = prompt_cost + completion_cost + assert expected_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + +def test_stream_chunk_builder_unknown_model_leaves_response_cost_unset(): + chunks: Final = [ + _stream_builder_text_chunk("totally-unknown-model-xyz", "Hello "), + _stream_builder_text_chunk("totally-unknown-model-xyz", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response._hidden_params.get("response_cost") is None + assert response.choices[0].message.content == "Hello world." + + +def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): + chunks: Final = [ + _stream_builder_text_chunk("claude-opus-5", "Hello "), + _stream_builder_text_chunk("claude-opus-5", "world.", finish_reason="stop"), + ] + for chunk in chunks: + chunk._hidden_params = {"custom_llm_provider": "openai"} + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response._hidden_params["custom_llm_provider"] == "openai" + prompt_cost, completion_cost = litellm.cost_per_token(model="claude-opus-5", usage_object=response.usage) + expected_cost: Final = prompt_cost + completion_cost + assert expected_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + +def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: + logging_obj: Final = LiteLLMLogging( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.update_environment_variables( + model=model, + user=None, + optional_params={}, + litellm_params={"custom_llm_provider": custom_llm_provider}, + custom_llm_provider=custom_llm_provider, + ) + return logging_obj + + +def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + usage_cost: Final = getattr(response.usage, "cost", None) + assert usage_cost is not None + assert usage_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) + + +def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj: Final = LiteLLMLogging( + model="us.anthropic.claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-unpriceable", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response._hidden_params.get("response_cost") is None + + +def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + assert response._hidden_params["response_cost"] == pytest.approx(0.5) + + +def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): + from openai.types.completion_usage import CompletionUsage + + usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) + assert type(usage_chunk.usage) is CompletionUsage + chunks: Final = [ + _stream_builder_text_chunk("mantle-claude", "Hello "), + _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response.usage.prompt_tokens == 20 + assert response.usage.completion_tokens == 60 + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + assert response._hidden_params["response_cost"] == pytest.approx(0.000704) + + +def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") + usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) + chunks: Final = [ + _stream_builder_text_chunk("grok-4", "Hello "), + _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), + usage_chunk, + ] + logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.42) + assert response._hidden_params.get("response_cost") is None + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) + + +def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-fake-mp3-bytes" + mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + response_format="wav", + speed=2, + instructions="sound cheerful", + ) + + assert mock_route.called + request_body: Final = json.loads(mock_route.calls.last.request.content) + assert request_body == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" + assert response.content == audio_bytes + + +def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-gateway-bytes" + gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + api_base="https://mistral.gateway.internal", + ) + + assert gateway_route.called + assert response.content == audio_bytes + + +FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" + + +def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" + ).mock(return_value=httpx.Response(200, json={"text": "hello"})) + + response: Final = litellm.transcription( + model="azure_ai/whisper-1", + file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.text == "hello" + + +def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" + ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) + + response: Final = litellm.speech( + model="azure_ai/tts-1", + input="hello", + voice="alloy", + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.content == b"mp3-bytes" + + +FORWARDED_CLIENT_HEADERS: Final = {"x-forwarded-for": "10.0.0.1", "x-amzn-trace-id": "Root=1-lit7694"} + + +def _chat_completion_json() -> Mapping[str, object]: + return { + "id": "chatcmpl-lit7694", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +def _chat_completion_sse() -> bytes: + chunk: Final = { + "id": "chatcmpl-lit7694", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + } + return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() + + +@pytest.mark.parametrize("stream", [False, True]) +def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_of_the_body( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, content=_chat_completion_sse(), headers={"content-type": "text/event-stream"}) + if stream + else httpx.Response(200, json=_chat_completion_json()) + ) + + response: Final = litellm.responses( + model="openai/gpt-5.4", + input="Reply with the single word ok", + stream=stream, + use_chat_completions_api=True, + headers=dict(FORWARDED_CLIENT_HEADERS), + api_key="sk-test", + ) + if stream: + list(response) + + assert route.called + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert "extra_headers" not in body + assert body["model"] == "gpt-5.4" + assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS + + +@pytest.mark.parametrize("http2_on", [True, False]) +def test_aiohttp_openai_warns_only_when_http2_enabled( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, http2_on: bool +): + from litellm.main import base_llm_aiohttp_handler + + monkeypatch.setattr(litellm, "http2", http2_on) + monkeypatch.delenv("LITELLM_HTTP2", raising=False) + + handler_completion: Final = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(base_llm_aiohttp_handler, "completion", handler_completion) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + litellm.completion( + model="aiohttp_openai/gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + ) + + assert handler_completion.called + warned: Final = "aiohttp_openai/ always uses aiohttp" in caplog.text + assert warned is http2_on + + +@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Which fruit is red?"}], + tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert f"tool_choice={tool_choice}" in str(exc_info.value) diff --git a/tests/test_litellm/test_main_module_header.py b/tests/unit/test_main_module_header.py similarity index 100% rename from tests/test_litellm/test_main_module_header.py rename to tests/unit/test_main_module_header.py diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/unit/test_mistral_medium_3_5_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_medium_3_5_model_metadata.py rename to tests/unit/test_mistral_medium_3_5_model_metadata.py diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/unit/test_mistral_small_4_0_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_small_4_0_model_metadata.py rename to tests/unit/test_mistral_small_4_0_model_metadata.py diff --git a/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py b/tests/unit/test_mistral_zai_glm_5_2_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py rename to tests/unit/test_mistral_zai_glm_5_2_model_metadata.py diff --git a/tests/test_litellm/test_model_block_unblock.py b/tests/unit/test_model_block_unblock.py similarity index 100% rename from tests/test_litellm/test_model_block_unblock.py rename to tests/unit/test_model_block_unblock.py diff --git a/tests/test_litellm/test_model_cost_aliases.py b/tests/unit/test_model_cost_aliases.py similarity index 100% rename from tests/test_litellm/test_model_cost_aliases.py rename to tests/unit/test_model_cost_aliases.py diff --git a/tests/test_litellm/test_model_param_helper.py b/tests/unit/test_model_param_helper.py similarity index 100% rename from tests/test_litellm/test_model_param_helper.py rename to tests/unit/test_model_param_helper.py diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py similarity index 100% rename from tests/test_litellm/test_model_prices_schema.py rename to tests/unit/test_model_prices_schema.py diff --git a/tests/test_litellm/test_model_response_normalization.py b/tests/unit/test_model_response_normalization.py similarity index 100% rename from tests/test_litellm/test_model_response_normalization.py rename to tests/unit/test_model_response_normalization.py diff --git a/tests/test_litellm/test_muse_spark_1_1_model_metadata.py b/tests/unit/test_muse_spark_1_1_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_1_model_metadata.py rename to tests/unit/test_muse_spark_1_1_model_metadata.py diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/unit/test_muse_spark_1_2_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_2_model_metadata.py rename to tests/unit/test_muse_spark_1_2_model_metadata.py diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/unit/test_muse_spark_1_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_3_model_metadata.py rename to tests/unit/test_muse_spark_1_3_model_metadata.py diff --git a/tests/test_litellm/test_mutation_report.py b/tests/unit/test_mutation_report.py similarity index 100% rename from tests/test_litellm/test_mutation_report.py rename to tests/unit/test_mutation_report.py diff --git a/tests/test_litellm/test_nested_drop_params.py b/tests/unit/test_nested_drop_params.py similarity index 100% rename from tests/test_litellm/test_nested_drop_params.py rename to tests/unit/test_nested_drop_params.py diff --git a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py b/tests/unit/test_non_chat_routes_open_llm_spans.py similarity index 100% rename from tests/test_litellm/test_non_chat_routes_open_llm_spans.py rename to tests/unit/test_non_chat_routes_open_llm_spans.py diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/unit/test_openai_embedding_encoding_format_default.py similarity index 100% rename from tests/test_litellm/test_openai_embedding_encoding_format_default.py rename to tests/unit/test_openai_embedding_encoding_format_default.py diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/unit/test_openai_service_tier_long_context_pricing.py similarity index 100% rename from tests/test_litellm/test_openai_service_tier_long_context_pricing.py rename to tests/unit/test_openai_service_tier_long_context_pricing.py diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/unit/test_pre_commit_lint.py similarity index 100% rename from tests/test_litellm/test_pre_commit_lint.py rename to tests/unit/test_pre_commit_lint.py diff --git a/tests/test_litellm/test_prisma_generate_if_needed.py b/tests/unit/test_prisma_generate_if_needed.py similarity index 100% rename from tests/test_litellm/test_prisma_generate_if_needed.py rename to tests/unit/test_prisma_generate_if_needed.py diff --git a/tests/test_litellm/test_process_helpers.py b/tests/unit/test_process_helpers.py similarity index 100% rename from tests/test_litellm/test_process_helpers.py rename to tests/unit/test_process_helpers.py diff --git a/tests/test_litellm/test_project_alias_tracking.py b/tests/unit/test_project_alias_tracking.py similarity index 100% rename from tests/test_litellm/test_project_alias_tracking.py rename to tests/unit/test_project_alias_tracking.py diff --git a/tests/test_litellm/test_project_tags_pydantic.py b/tests/unit/test_project_tags_pydantic.py similarity index 100% rename from tests/test_litellm/test_project_tags_pydantic.py rename to tests/unit/test_project_tags_pydantic.py diff --git a/tests/test_litellm/test_proxy_auth.py b/tests/unit/test_proxy_auth.py similarity index 100% rename from tests/test_litellm/test_proxy_auth.py rename to tests/unit/test_proxy_auth.py diff --git a/tests/test_litellm/test_rag_openai_ingestion.py b/tests/unit/test_rag_openai_ingestion.py similarity index 100% rename from tests/test_litellm/test_rag_openai_ingestion.py rename to tests/unit/test_rag_openai_ingestion.py diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/unit/test_rate_limit_error_unification.py similarity index 100% rename from tests/test_litellm/test_rate_limit_error_unification.py rename to tests/unit/test_rate_limit_error_unification.py diff --git a/tests/test_litellm/test_read_rc_version.py b/tests/unit/test_read_rc_version.py similarity index 100% rename from tests/test_litellm/test_read_rc_version.py rename to tests/unit/test_read_rc_version.py diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/unit/test_redact_string_in_error_paths.py similarity index 100% rename from tests/test_litellm/test_redact_string_in_error_paths.py rename to tests/unit/test_redact_string_in_error_paths.py diff --git a/tests/test_litellm/test_redis.py b/tests/unit/test_redis.py similarity index 100% rename from tests/test_litellm/test_redis.py rename to tests/unit/test_redis.py diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/unit/test_redis_credential_provider.py similarity index 100% rename from tests/test_litellm/test_redis_credential_provider.py rename to tests/unit/test_redis_credential_provider.py diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py similarity index 100% rename from tests/test_litellm/test_register_model_custom_pricing.py rename to tests/unit/test_register_model_custom_pricing.py diff --git a/tests/test_litellm/test_register_model_zero_cost_persistence.py b/tests/unit/test_register_model_zero_cost_persistence.py similarity index 100% rename from tests/test_litellm/test_register_model_zero_cost_persistence.py rename to tests/unit/test_register_model_zero_cost_persistence.py diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/unit/test_replicate_model_key_format.py similarity index 100% rename from tests/test_litellm/test_replicate_model_key_format.py rename to tests/unit/test_replicate_model_key_format.py diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/unit/test_responses_api_bridge_non_stream.py similarity index 100% rename from tests/test_litellm/test_responses_api_bridge_non_stream.py rename to tests/unit/test_responses_api_bridge_non_stream.py diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/unit/test_responses_id_security.py similarity index 94% rename from tests/test_litellm/test_responses_id_security.py rename to tests/unit/test_responses_id_security.py index a6081670172..704a52fc202 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -4,7 +4,7 @@ Tests for ResponsesIDSecurity hook. Tests the security hook that prevents user B from seeing response from user A. """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException @@ -113,63 +113,6 @@ class TestDecryptResponseId: assert team_id is None -class TestEncryptResponseId: - """Test _encrypt_response_id function""" - - @pytest.mark.skip( - reason="Flaky on CI; disabling temporarily until responses_id_security is fixed" - ) - def test_encrypt_response_id_success( - self, responses_id_security, mock_user_api_key_dict - ): - """Test encrypting a response ID with user information""" - mock_response = ResponsesAPIResponse( - id="resp_123", created_at=1234567890, output=[], status="completed" - ) - - with patch( - "litellm.proxy.hooks.responses_id_security.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "encrypted_base64_value" - - with patch.object( - responses_id_security, "_get_signing_key", return_value="test-key" - ): - result = responses_id_security._encrypt_response_id( - mock_response, mock_user_api_key_dict - ) - - assert result.id == "resp_encrypted_base64_value" - assert result.id.startswith("resp_") - mock_encrypt.assert_called_once() - - @pytest.mark.skip( - reason="Flaky on CI; disabling temporarily until responses_id_security is fixed" - ) - def test_encrypt_response_id_maintains_prefix( - self, responses_id_security, mock_user_api_key_dict - ): - """Test that encrypted response ID maintains 'resp_' prefix""" - mock_response = ResponsesAPIResponse( - id="resp_456", created_at=1234567890, output=[], status="in_progress" - ) - - with patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", - return_value="test-salt-key", - ): - with patch.object( - responses_id_security, "_get_signing_key", return_value="test-key" - ): - result = responses_id_security._encrypt_response_id( - mock_response, mock_user_api_key_dict - ) - - assert result.id.startswith("resp_") - # The encrypted ID should be different from the original - assert result.id != "resp_456" - - class TestCheckUserAccessToResponseId: """Test check_user_access_to_response_id function""" @@ -857,7 +800,6 @@ class TestAsyncPostCallSuccessHook: assert result == mock_response - _FABRICATED_PROVIDER_RESPONSE_ID = "resp_fabricatedprovideridaaaaaaaaaaaaaaaa" _FABRICATED_UNMANAGED_ID = "resp_fabricatedunmanagedidbbbbbbbbbbbbbbbb" _UNIT_TEST_SALT_KEY = "lit6837-unit-test-salt-key" diff --git a/tests/test_litellm/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py similarity index 100% rename from tests/test_litellm/test_responses_streaming_container_ownership.py rename to tests/unit/test_responses_streaming_container_ownership.py diff --git a/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py b/tests/unit/test_retrieve_batch_bedrock_dispatch.py similarity index 100% rename from tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py rename to tests/unit/test_retrieve_batch_bedrock_dispatch.py diff --git a/tests/test_litellm/test_router.py b/tests/unit/test_router/test_router.py similarity index 100% rename from tests/test_litellm/test_router.py rename to tests/unit/test_router/test_router.py diff --git a/tests/test_litellm/test_router_block_helpers.py b/tests/unit/test_router_block_helpers.py similarity index 100% rename from tests/test_litellm/test_router_block_helpers.py rename to tests/unit/test_router_block_helpers.py diff --git a/tests/test_litellm/test_router_exception_redaction.py b/tests/unit/test_router_exception_redaction.py similarity index 100% rename from tests/test_litellm/test_router_exception_redaction.py rename to tests/unit/test_router_exception_redaction.py diff --git a/tests/test_litellm/test_router_google_genai.py b/tests/unit/test_router_google_genai.py similarity index 100% rename from tests/test_litellm/test_router_google_genai.py rename to tests/unit/test_router_google_genai.py diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py similarity index 100% rename from tests/test_litellm/test_router_model_cost_isolation.py rename to tests/unit/test_router_model_cost_isolation.py diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py similarity index 100% rename from tests/test_litellm/test_router_order_fallback.py rename to tests/unit/test_router_order_fallback.py diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py similarity index 100% rename from tests/test_litellm/test_router_per_deployment_num_retries.py rename to tests/unit/test_router_per_deployment_num_retries.py diff --git a/tests/test_litellm/test_router_redis_init.py b/tests/unit/test_router_redis_init.py similarity index 100% rename from tests/test_litellm/test_router_redis_init.py rename to tests/unit/test_router_redis_init.py diff --git a/tests/test_litellm/test_router_retry_backoff_headers.py b/tests/unit/test_router_retry_backoff_headers.py similarity index 100% rename from tests/test_litellm/test_router_retry_backoff_headers.py rename to tests/unit/test_router_retry_backoff_headers.py diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/unit/test_router_retry_non_retryable_errors.py similarity index 100% rename from tests/test_litellm/test_router_retry_non_retryable_errors.py rename to tests/unit/test_router_retry_non_retryable_errors.py diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py similarity index 100% rename from tests/test_litellm/test_router_retry_policy_update.py rename to tests/unit/test_router_retry_policy_update.py diff --git a/tests/test_litellm/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py similarity index 92% rename from tests/test_litellm/test_router_silent_experiment.py rename to tests/unit/test_router_silent_experiment.py index d62962da275..ab65e09e133 100644 --- a/tests/test_litellm/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -388,47 +388,6 @@ async def test_shadow_of_a_shadow_is_not_launched(recording_logger): assert model_groups == ["shadow-a"] -def test_silent_experiment_completion_direct(): - """ - Test _silent_experiment_completion directly (for router code coverage). - Mocks router.completion to avoid real API call. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): - router._silent_experiment_completion( - silent_model="gpt-3.5-turbo", - messages=messages, - ) - - -@pytest.mark.asyncio -async def test_silent_experiment_acompletion_direct(): - """ - Test _silent_experiment_acompletion directly (for router code coverage). - Mocks router.acompletion to avoid real API call. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): - await router._silent_experiment_acompletion( - silent_model="gpt-3.5-turbo", - messages=messages, - ) - - @pytest.mark.asyncio async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger): router = Router(model_list=_streaming_model_list(None)) @@ -602,3 +561,44 @@ def test_router_silent_experiment_completion(): assert silent_call[1]["model"] == "openai/gpt-4" # Verify model_group is set to the silent model name for correct metric attribution assert silent_call[1]["metadata"]["model_group"] == "silent-model" + + +SILENT_EXPERIMENT_RUNNERS: Final = ( + pytest.param(lambda router, **kwargs: router._silent_experiment_completion(**kwargs), id="sync"), + pytest.param(lambda router, **kwargs: asyncio.run(router._silent_experiment_acompletion(**kwargs)), id="async"), +) + + +@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) +def test_silent_experiment_sends_shadow_request_attributed_to_the_silent_model(run_silent_experiment): + router = Router(model_list=_streaming_model_list(["shadow-a"])) + primary_metadata: Final = {"model_group": "primary-model"} + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: + run_silent_experiment( + router, + silent_model="shadow-a", + messages=[{"role": "user", "content": "hi"}], + metadata=primary_metadata, + ) + + acompletion.assert_awaited_once() + shadow_call: Final = acompletion.await_args.kwargs + assert shadow_call["model"] == "shadow-a" + assert shadow_call["messages"] == [{"role": "user", "content": "hi"}] + assert shadow_call["metadata"]["model_group"] == "shadow-a" + assert shadow_call["metadata"]["is_silent_experiment"] is True + assert primary_metadata == {"model_group": "primary-model"} + + +@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) +def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment): + router = Router(model_list=_streaming_model_list(["shadow-a"])) + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: + run_silent_experiment( + router, + silent_model="shadow-a", + messages=[{"role": "user", "content": "hi"}], + metadata={"is_silent_experiment": True}, + ) + + acompletion.assert_not_awaited() diff --git a/tests/test_litellm/test_router_streaming_fallback_metadata.py b/tests/unit/test_router_streaming_fallback_metadata.py similarity index 100% rename from tests/test_litellm/test_router_streaming_fallback_metadata.py rename to tests/unit/test_router_streaming_fallback_metadata.py diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py similarity index 100% rename from tests/test_litellm/test_router_weighted_failover.py rename to tests/unit/test_router_weighted_failover.py diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/unit/test_ruff_strict_gate.py similarity index 100% rename from tests/test_litellm/test_ruff_strict_gate.py rename to tests/unit/test_ruff_strict_gate.py diff --git a/tests/test_litellm/test_sambanova_model_metadata.py b/tests/unit/test_sambanova_model_metadata.py similarity index 100% rename from tests/test_litellm/test_sambanova_model_metadata.py rename to tests/unit/test_sambanova_model_metadata.py diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/unit/test_secret_redaction.py similarity index 100% rename from tests/test_litellm/test_secret_redaction.py rename to tests/unit/test_secret_redaction.py diff --git a/tests/test_litellm/test_select_ui_test_scope.py b/tests/unit/test_select_ui_test_scope.py similarity index 100% rename from tests/test_litellm/test_select_ui_test_scope.py rename to tests/unit/test_select_ui_test_scope.py diff --git a/tests/test_litellm/test_service_logger.py b/tests/unit/test_service_logger.py similarity index 100% rename from tests/test_litellm/test_service_logger.py rename to tests/unit/test_service_logger.py diff --git a/tests/test_litellm/test_setup_wizard.py b/tests/unit/test_setup_wizard.py similarity index 100% rename from tests/test_litellm/test_setup_wizard.py rename to tests/unit/test_setup_wizard.py diff --git a/tests/test_litellm/test_shared_session_integration.py b/tests/unit/test_shared_session_integration.py similarity index 100% rename from tests/test_litellm/test_shared_session_integration.py rename to tests/unit/test_shared_session_integration.py diff --git a/tests/test_litellm/test_ssl_verify_unit.py b/tests/unit/test_ssl_verify_unit.py similarity index 83% rename from tests/test_litellm/test_ssl_verify_unit.py rename to tests/unit/test_ssl_verify_unit.py index c39362c01a2..f47cdf3e6cd 100644 --- a/tests/test_litellm/test_ssl_verify_unit.py +++ b/tests/unit/test_ssl_verify_unit.py @@ -50,41 +50,6 @@ class TestBaseAWSLLMSSLVerify: # Result depends on environment, just verify it doesn't crash assert result is not None or result is None # Can be None, True, False, or path - @patch("boto3.client") - def test_get_credentials_propagates_ssl_verify(self, mock_boto_client): - """Test that get_credentials propagates ssl_verify to boto3 clients.""" - base_llm = BaseAWSLLM() - - # Mock the boto3 client - mock_sts_client = Mock() - mock_sts_client.assume_role.return_value = { - "Credentials": { - "AccessKeyId": "test_key", - "SecretAccessKey": "test_secret", - "SessionToken": "test_token", - "Expiration": "2026-01-20T00:00:00Z", - } - } - mock_boto_client.return_value = mock_sts_client - - # Call get_credentials with ssl_verify parameter - cert_path = "/path/to/cert.pem" - try: - base_llm.get_credentials( - aws_access_key_id="test_key", - aws_secret_access_key="test_secret", - aws_region_name="us-east-1", - ssl_verify=cert_path, - ) - except Exception: - # May fail due to missing credentials, but we're checking the call - pass - - # Verify boto3.client was called with verify parameter - # Note: This test verifies the parameter is accepted, actual propagation - # is tested in integration tests - assert True # If we got here without error, parameter was accepted - class TestAimGuardrailSSLVerify: """Test SSL verification parameter handling in AimGuardrail.""" diff --git a/tests/test_litellm/test_stream_chunk_builder_annotations.py b/tests/unit/test_stream_chunk_builder_annotations.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_annotations.py rename to tests/unit/test_stream_chunk_builder_annotations.py diff --git a/tests/test_litellm/test_stream_chunk_builder_citations.py b/tests/unit/test_stream_chunk_builder_citations.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_citations.py rename to tests/unit/test_stream_chunk_builder_citations.py diff --git a/tests/test_litellm/test_stream_chunk_builder_images.py b/tests/unit/test_stream_chunk_builder_images.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_images.py rename to tests/unit/test_stream_chunk_builder_images.py diff --git a/tests/test_litellm/test_streaming_connection_cleanup.py b/tests/unit/test_streaming_connection_cleanup.py similarity index 100% rename from tests/test_litellm/test_streaming_connection_cleanup.py rename to tests/unit/test_streaming_connection_cleanup.py diff --git a/tests/test_litellm/test_sync_together_ai_models.py b/tests/unit/test_sync_together_ai_models.py similarity index 100% rename from tests/test_litellm/test_sync_together_ai_models.py rename to tests/unit/test_sync_together_ai_models.py diff --git a/tests/test_litellm/test_system_message_format_bug.py b/tests/unit/test_system_message_format_bug.py similarity index 100% rename from tests/test_litellm/test_system_message_format_bug.py rename to tests/unit/test_system_message_format_bug.py diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/unit/test_test_quality_gate.py similarity index 100% rename from tests/test_litellm/test_test_quality_gate.py rename to tests/unit/test_test_quality_gate.py diff --git a/tests/test_litellm/test_thinking_enabled.py b/tests/unit/test_thinking_enabled.py similarity index 100% rename from tests/test_litellm/test_thinking_enabled.py rename to tests/unit/test_thinking_enabled.py diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/unit/test_together_ai_model_metadata.py similarity index 100% rename from tests/test_litellm/test_together_ai_model_metadata.py rename to tests/unit/test_together_ai_model_metadata.py diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/unit/test_type_check_gate.py similarity index 100% rename from tests/test_litellm/test_type_check_gate.py rename to tests/unit/test_type_check_gate.py diff --git a/tests/test_litellm/test_type_discipline_gate.py b/tests/unit/test_type_discipline_gate.py similarity index 100% rename from tests/test_litellm/test_type_discipline_gate.py rename to tests/unit/test_type_discipline_gate.py diff --git a/tests/test_litellm/test_typesafe_model_metadata.py b/tests/unit/test_typesafe_model_metadata.py similarity index 100% rename from tests/test_litellm/test_typesafe_model_metadata.py rename to tests/unit/test_typesafe_model_metadata.py diff --git a/tests/test_litellm/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py similarity index 97% rename from tests/test_litellm/test_unit_shard_missing_paths.py rename to tests/unit/test_unit_shard_missing_paths.py index b91c2cff764..4fa9c5bd3c1 100644 --- a/tests/test_litellm/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -36,6 +36,7 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl **os.environ, **_SHARD_ENV, "PATH": f"{shim_dir}{os.pathsep}{os.environ['PATH']}", + "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, "WORKERS": workers, }, diff --git a/tests/test_litellm/test_unit_shard_per_test_timeout.py b/tests/unit/test_unit_shard_per_test_timeout.py similarity index 100% rename from tests/test_litellm/test_unit_shard_per_test_timeout.py rename to tests/unit/test_unit_shard_per_test_timeout.py diff --git a/tests/test_litellm/test_utils.py b/tests/unit/test_utils.py similarity index 100% rename from tests/test_litellm/test_utils.py rename to tests/unit/test_utils.py diff --git a/tests/test_litellm/test_utils_module_docstring.py b/tests/unit/test_utils_module_docstring.py similarity index 100% rename from tests/test_litellm/test_utils_module_docstring.py rename to tests/unit/test_utils_module_docstring.py diff --git a/tests/test_litellm/test_uuid_helper.py b/tests/unit/test_uuid_helper.py similarity index 100% rename from tests/test_litellm/test_uuid_helper.py rename to tests/unit/test_uuid_helper.py diff --git a/tests/test_litellm/test_vcr_safe_body_matcher.py b/tests/unit/test_vcr_safe_body_matcher.py similarity index 98% rename from tests/test_litellm/test_vcr_safe_body_matcher.py rename to tests/unit/test_vcr_safe_body_matcher.py index 712ecf09911..cf4e4a1c276 100644 --- a/tests/test_litellm/test_vcr_safe_body_matcher.py +++ b/tests/unit/test_vcr_safe_body_matcher.py @@ -52,14 +52,6 @@ def test_safe_body_matcher_accepts_str_bytes_equivalent(): _safe_body_matcher(_req("hello"), _req(b"hello")) -def test_safe_body_matcher_handles_jsonl_without_crashing(): - jsonl = ( - b'{"recordId": "request-1", "modelInput": {}}\n' - b'{"recordId": "request-2", "modelInput": {}}\n' - ) - _safe_body_matcher(_req(jsonl), _req(jsonl)) - - def test_safe_body_matcher_rejects_different_jsonl_bodies(): a = b'{"recordId": "request-1"}\n{"recordId": "request-2"}\n' b = b'{"recordId": "request-1"}\n{"recordId": "request-3"}\n' diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/unit/test_vertex_ai_xai_grok_prompt_caching_metadata.py similarity index 100% rename from tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py rename to tests/unit/test_vertex_ai_xai_grok_prompt_caching_metadata.py diff --git a/tests/test_litellm/test_video_generation.py b/tests/unit/test_video_generation.py similarity index 100% rename from tests/test_litellm/test_video_generation.py rename to tests/unit/test_video_generation.py diff --git a/tests/test_litellm/test_with_dashboard_node.py b/tests/unit/test_with_dashboard_node.py similarity index 100% rename from tests/test_litellm/test_with_dashboard_node.py rename to tests/unit/test_with_dashboard_node.py diff --git a/tests/test_litellm/test_xai_grok_4_3_model_metadata.py b/tests/unit/test_xai_grok_4_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_xai_grok_4_3_model_metadata.py rename to tests/unit/test_xai_grok_4_3_model_metadata.py diff --git a/tests/test_litellm/test_xai_responses_auto_routing.py b/tests/unit/test_xai_responses_auto_routing.py similarity index 100% rename from tests/test_litellm/test_xai_responses_auto_routing.py rename to tests/unit/test_xai_responses_auto_routing.py diff --git a/tests/test_litellm/types/test_completion.py b/tests/unit/types/test_completion.py similarity index 99% rename from tests/test_litellm/types/test_completion.py rename to tests/unit/types/test_completion.py index cd51913c5dd..4971a0c7e0a 100644 --- a/tests/test_litellm/types/test_completion.py +++ b/tests/unit/types/test_completion.py @@ -5,7 +5,7 @@ This test suite validates the CompletionRequest model and its compatibility with OpenAI ChatCompletion API message formats. Usage: - pytest tests/test_litellm/types/test_completion.py -v + pytest tests/unit/types/test_completion.py -v """ import dataclasses diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/unit/types/test_guardrails_case_normalization.py similarity index 100% rename from tests/test_litellm/types/test_guardrails_case_normalization.py rename to tests/unit/types/test_guardrails_case_normalization.py diff --git a/tests/test_litellm/types/test_mcp.py b/tests/unit/types/test_mcp.py similarity index 100% rename from tests/test_litellm/types/test_mcp.py rename to tests/unit/types/test_mcp.py diff --git a/tests/test_litellm/types/test_presidio_entity_expansion.py b/tests/unit/types/test_presidio_entity_expansion.py similarity index 100% rename from tests/test_litellm/types/test_presidio_entity_expansion.py rename to tests/unit/types/test_presidio_entity_expansion.py diff --git a/tests/test_litellm/types/test_prometheus_label_value_sanitize.py b/tests/unit/types/test_prometheus_label_value_sanitize.py similarity index 100% rename from tests/test_litellm/types/test_prometheus_label_value_sanitize.py rename to tests/unit/types/test_prometheus_label_value_sanitize.py diff --git a/tests/test_litellm/types/test_prometheus_latency_buckets.py b/tests/unit/types/test_prometheus_latency_buckets.py similarity index 100% rename from tests/test_litellm/types/test_prometheus_latency_buckets.py rename to tests/unit/types/test_prometheus_latency_buckets.py diff --git a/tests/test_litellm/types/test_router.py b/tests/unit/types/test_router.py similarity index 100% rename from tests/test_litellm/types/test_router.py rename to tests/unit/types/test_router.py diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/unit/types/test_types_utils.py similarity index 100% rename from tests/test_litellm/types/test_types_utils.py rename to tests/unit/types/test_types_utils.py diff --git a/tests/test_litellm/types/test_uk_pii_entities.py b/tests/unit/types/test_uk_pii_entities.py similarity index 100% rename from tests/test_litellm/types/test_uk_pii_entities.py rename to tests/unit/types/test_uk_pii_entities.py diff --git a/tests/test_litellm/files/__init__.py b/tests/unit/vector_stores/__init__.py similarity index 100% rename from tests/test_litellm/files/__init__.py rename to tests/unit/vector_stores/__init__.py diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py similarity index 100% rename from tests/test_litellm/vector_stores/test_main.py rename to tests/unit/vector_stores/test_main.py diff --git a/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py b/tests/unit/vector_stores/test_vector_store_create_provider_logic.py similarity index 100% rename from tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py rename to tests/unit/vector_stores/test_vector_store_create_provider_logic.py diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py similarity index 100% rename from tests/test_litellm/vector_stores/test_vector_store_registry.py rename to tests/unit/vector_stores/test_vector_store_registry.py From 88fd15315c9961cfd6770754889a7a1eda6ce333 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Fri, 25 Sep 2026 20:58:17 +0200 Subject: [PATCH 10/58] fix(sso): gate /sso/debug routes behind ENABLE_SSO_DEBUG, off by default (#43150) /sso/debug/login and /sso/debug/callback are diagnostic pages that had no off switch. They cannot carry a bearer credential because the IdP redirects a bare browser to the callback, so the gate is an explicit opt-in flag rather than key auth: both routes return 404 unless ENABLE_SSO_DEBUG is set to a truthy value. --- litellm/proxy/management_endpoints/ui_sso.py | 11 ++++++ .../proxy/management_endpoints/test_ui_sso.py | 37 +++++++++++++++++-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7859c678c07..618b200a14c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -4618,6 +4618,13 @@ class GoogleSSOHandler: return result or {} +def _raise_if_sso_debug_disabled() -> None: + """The debug routes run the browser-redirect SSO flow, so they cannot carry a + bearer credential; an explicit opt-in flag is the only way to gate them.""" + if get_secret_bool("ENABLE_SSO_DEBUG") is not True: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not Found") + + @router.get("/sso/debug/login", tags=["experimental"], include_in_schema=False) async def debug_sso_login(request: Request): """ @@ -4625,6 +4632,8 @@ async def debug_sso_login(request: Request): PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" Example: """ + _raise_if_sso_debug_disabled() + from litellm.proxy.proxy_server import premium_user microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None) @@ -4670,6 +4679,8 @@ async def debug_sso_callback(request: Request): """ Returns the OpenID object returned by the SSO provider """ + _raise_if_sso_debug_disabled() + import json from fastapi.responses import HTMLResponse diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 1230c548281..21c0f565486 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -8029,6 +8029,37 @@ class TestPKCEStateCookieBinding: assert result is not None +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_sso_debug_value", [None, "false", "0"]) +async def test_sso_debug_routes_return_404_unless_explicitly_enabled(enable_sso_debug_value): + """ + /sso/debug/login and /sso/debug/callback must 404 unless ENABLE_SSO_DEBUG is + explicitly set to a truthy value. + """ + from litellm.proxy.management_endpoints.ui_sso import debug_sso_callback, debug_sso_login + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.example.com/" + mock_request.cookies = {} + mock_request.query_params = {} + + env = {"GENERIC_CLIENT_ID": "test_client_id"} + if enable_sso_debug_value is not None: + env["ENABLE_SSO_DEBUG"] = enable_sso_debug_value + + with patch.dict(os.environ, env, clear=False): + if enable_sso_debug_value is None: + os.environ.pop("ENABLE_SSO_DEBUG", None) + + with pytest.raises(HTTPException) as login_exc: + await debug_sso_login(mock_request) + with pytest.raises(HTTPException) as callback_exc: + await debug_sso_callback(mock_request) + + assert login_exc.value.status_code == 404 + assert callback_exc.value.status_code == 404 + + @pytest.mark.asyncio async def test_debug_sso_callback_renders_full_jwt_claims(): """ @@ -8080,7 +8111,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims(): with ( patch.dict( os.environ, - {"GENERIC_CLIENT_ID": "test_client_id"}, + {"GENERIC_CLIENT_ID": "test_client_id", "ENABLE_SSO_DEBUG": "true"}, clear=False, ), patch( @@ -8165,7 +8196,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): with ( patch.dict( os.environ, - {"MICROSOFT_CLIENT_ID": "test_microsoft_id"}, + {"MICROSOFT_CLIENT_ID": "test_microsoft_id", "ENABLE_SSO_DEBUG": "true"}, clear=False, ), patch.object( @@ -8213,7 +8244,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) return parsed stack = [ - patch.dict(os.environ, provider_env, clear=False), + patch.dict(os.environ, {**provider_env, "ENABLE_SSO_DEBUG": "true"}, clear=False), patch( # test-quality-ok: endpoint test stubs the upstream generic IdP boundary "litellm.proxy.management_endpoints.ui_sso.get_generic_sso_response", side_effect=fake_generic ), From 3c93ea1697a41aa432a7b69aba5d24767391c67a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 12:18:32 -0700 Subject: [PATCH 11/58] refactor(framer): replace Framer trait with tokio-util codecs (#43193) * refactor(framer): replace Framer trait with tokio-util codecs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(framer): port SSE and AWS event stream framing to codecs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust): drop clone on Copy capabilities in messages request test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust): use field init shorthand in messages request test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 16 +- .../crates/core/tests/messages/request.rs | 2 +- litellm-rust/crates/framer/Cargo.toml | 5 +- .../crates/framer/src/aws_event_stream.rs | 95 ++++---- litellm-rust/crates/framer/src/error.rs | 24 +- litellm-rust/crates/framer/src/framed.rs | 21 ++ litellm-rust/crates/framer/src/lib.rs | 4 +- litellm-rust/crates/framer/src/sse.rs | 189 +++++++++++++--- .../crates/framer/tests/aws_event_stream.rs | 209 ++++++++++++------ litellm-rust/crates/framer/tests/chaining.rs | 74 +++++-- litellm-rust/crates/framer/tests/sse.rs | 188 ++++++++++++---- .../crates/framer/tests/support/mod.rs | 68 ++++-- .../messages/streaming_iterator.rs | 37 ++-- 13 files changed, 657 insertions(+), 275 deletions(-) create mode 100644 litellm-rust/crates/framer/src/framed.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e7d911f5fd9..3677d1d654f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3166,10 +3166,11 @@ dependencies = [ "aws-smithy-types", "bytes", "futures-util", + "proptest", "rstest", - "sse-stream", "thiserror 2.0.19", "tokio", + "tokio-util", ] [[package]] @@ -5468,19 +5469,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "sse-stream" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4" -dependencies = [ - "bytes", - "futures-util", - "http-body 1.1.0", - "http-body-util", - "pin-project-lite", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 2927356b773..d37910d4ac4 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -398,7 +398,7 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i api_key: Some("sk".into()), api_base: Some(upstream.uri()), shaping: MessagesShaping { - capabilities: capabilities.clone(), + capabilities, drop_params, ..MessagesShaping::default() }, diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index 62bfcc7da3d..e11f2c02a97 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -8,16 +8,17 @@ repository.workspace = true [features] default = ["aws", "sse"] aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"] -sse = ["dep:sse-stream"] +sse = [] [dependencies] aws-smithy-eventstream = { version = "=0.61.4", optional = true } aws-smithy-types = { version = "1.6.1", optional = true } bytes = "1" futures-util.workspace = true -sse-stream = { version = "=0.2.6", optional = true } thiserror.workspace = true +tokio-util = { version = "0.7", features = ["codec", "io"] } [dev-dependencies] +proptest.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/framer/src/aws_event_stream.rs b/litellm-rust/crates/framer/src/aws_event_stream.rs index efd7adeb64b..405ec2d5ad2 100644 --- a/litellm-rust/crates/framer/src/aws_event_stream.rs +++ b/litellm-rust/crates/framer/src/aws_event_stream.rs @@ -1,66 +1,47 @@ -use bytes::{Buf, Bytes, BytesMut}; -use futures_util::{Stream, StreamExt}; +use aws_smithy_eventstream::frame::{read_message_from, write_message_to}; +pub use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; +use bytes::BytesMut; +use tokio_util::codec::{Decoder, Encoder}; -use aws_smithy_eventstream::frame::read_message_from; -use aws_smithy_types::event_stream::Header; - -use crate::{Error, Framer}; +use crate::EventStreamError; +const MIN_FRAME_BYTES: usize = 16; const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; -#[derive(Clone, Debug, PartialEq)] -pub struct AwsEventStreamFrame { - pub headers: Vec
, - pub payload: Bytes, -} - #[derive(Clone, Copy, Debug, Default)] -pub struct AwsEventStreamFramer; +pub struct AwsEventStreamCodec; -impl Framer for AwsEventStreamFramer { - type Frame = AwsEventStreamFrame; +impl Decoder for AwsEventStreamCodec { + type Item = Message; + type Error = EventStreamError; - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: Buf + Send, - E: std::error::Error + Send + Sync + 'static, - { - futures_util::stream::try_unfold( - (Box::pin(input), BytesMut::new()), - |(mut input, mut buffer)| async move { - loop { - if buffer.len() >= 4 { - let length = (&buffer[..4]).get_u32() as usize; - if !(16..=MAX_FRAME_BYTES).contains(&length) { - return Err(Error::InvalidLength(length)); - } - if buffer.len() >= length { - let raw = buffer.split_to(length).freeze(); - let message = read_message_from(raw)?; - let frame = AwsEventStreamFrame { - headers: message.headers().to_vec(), - payload: message.payload().clone(), - }; - return Ok(Some((frame, (input, buffer)))); - } - } - match input.next().await { - Some(Ok(mut chunk)) => { - while chunk.has_remaining() { - let bytes = chunk.chunk(); - buffer.extend_from_slice(bytes); - let length = bytes.len(); - chunk.advance(length); - } - } - Some(Err(error)) => return Err(Error::Body(Box::new(error))), - None if buffer.is_empty() => return Ok(None), - None => return Err(Error::Truncated), - } - } - }, - ) - .fuse() + fn decode(&mut self, src: &mut BytesMut) -> Result, EventStreamError> { + let Some(prefix) = src.first_chunk::<4>() else { + return Ok(None); + }; + let length = u32::from_be_bytes(*prefix) as usize; + if !(MIN_FRAME_BYTES..=MAX_FRAME_BYTES).contains(&length) { + return Err(EventStreamError::InvalidLength(length)); + } + if src.len() < length { + return Ok(None); + } + Ok(Some(read_message_from(src.split_to(length).freeze())?)) + } + + fn decode_eof(&mut self, src: &mut BytesMut) -> Result, EventStreamError> { + match self.decode(src)? { + Some(message) => Ok(Some(message)), + None if src.is_empty() => Ok(None), + None => Err(EventStreamError::Truncated), + } + } +} + +impl Encoder for AwsEventStreamCodec { + type Error = EventStreamError; + + fn encode(&mut self, message: Message, dst: &mut BytesMut) -> Result<(), EventStreamError> { + Ok(write_message_to(&message, dst)?) } } diff --git a/litellm-rust/crates/framer/src/error.rs b/litellm-rust/crates/framer/src/error.rs index b1f7ed96c5a..879d7557671 100644 --- a/litellm-rust/crates/framer/src/error.rs +++ b/litellm-rust/crates/framer/src/error.rs @@ -1,17 +1,21 @@ +#[cfg(feature = "sse")] #[derive(Debug, thiserror::Error)] -pub enum Error { - #[cfg(feature = "sse")] - #[error("SSE framing failed: {0}")] - Sse(#[from] sse_stream::Error), - #[cfg(feature = "aws")] - #[error("AWS EventStream framing failed: {0}")] - Aws(#[from] aws_smithy_eventstream::error::Error), +pub enum SseError { #[error("body stream failed: {0}")] - Body(#[source] Box), - #[cfg(feature = "aws")] + Body(#[from] std::io::Error), + #[error("SSE field is not UTF-8: {0}")] + InvalidUtf8(#[from] std::str::Utf8Error), +} + +#[cfg(feature = "aws")] +#[derive(Debug, thiserror::Error)] +pub enum EventStreamError { + #[error("body stream failed: {0}")] + Body(#[from] std::io::Error), #[error("invalid AWS EventStream frame length: {0}")] InvalidLength(usize), - #[cfg(feature = "aws")] #[error("truncated AWS EventStream frame")] Truncated, + #[error("malformed AWS EventStream frame: {0}")] + Malformed(#[from] aws_smithy_eventstream::error::Error), } diff --git a/litellm-rust/crates/framer/src/framed.rs b/litellm-rust/crates/framer/src/framed.rs new file mode 100644 index 00000000000..7a19dd40e13 --- /dev/null +++ b/litellm-rust/crates/framer/src/framed.rs @@ -0,0 +1,21 @@ +use std::io; + +use bytes::Buf; +use futures_util::{Stream, StreamExt, TryStreamExt}; +use tokio_util::{ + codec::{Decoder, FramedRead}, + io::StreamReader, +}; + +pub fn frames( + input: S, + codec: D, +) -> impl Stream> + Send +where + S: Stream> + Send, + B: Buf + Send, + E: std::error::Error + Send + Sync + 'static, + D: Decoder + Send, +{ + FramedRead::new(StreamReader::new(input.map_err(io::Error::other)), codec).fuse() +} diff --git a/litellm-rust/crates/framer/src/lib.rs b/litellm-rust/crates/framer/src/lib.rs index 552de419984..223f2f64120 100644 --- a/litellm-rust/crates/framer/src/lib.rs +++ b/litellm-rust/crates/framer/src/lib.rs @@ -1,8 +1,8 @@ mod error; -mod framer; +mod framed; pub use error::*; -pub use framer::*; +pub use framed::frames; #[cfg(feature = "aws")] pub mod aws_event_stream; diff --git a/litellm-rust/crates/framer/src/sse.rs b/litellm-rust/crates/framer/src/sse.rs index 79659f6ce13..6fee1cfab7f 100644 --- a/litellm-rust/crates/framer/src/sse.rs +++ b/litellm-rust/crates/framer/src/sse.rs @@ -1,43 +1,170 @@ -use futures_util::{Stream, StreamExt}; +use std::str; -use crate::{Error, Framer}; +use bytes::{Buf, BufMut, BytesMut}; +use tokio_util::codec::{Decoder, Encoder}; -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct SseFrame { +use crate::SseError; + +const BOM: &[u8] = b"\xEF\xBB\xBF"; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SseEvent { pub event: Option, - pub data: Option, + pub data: String, pub id: Option, pub retry: Option, } #[derive(Clone, Copy, Debug, Default)] -pub struct SseFramer; +pub struct SseCodec { + past_bom: bool, +} -impl Framer for SseFramer { - type Frame = SseFrame; +impl Decoder for SseCodec { + type Item = SseEvent; + type Error = SseError; - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: bytes::Buf + Send, - E: std::error::Error + Send + Sync + 'static, - { - let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input)); - futures_util::stream::try_unfold(frames, |mut frames| async move { - let Some(frame) = frames.next().await else { - return Ok(None); - }; - let frame = frame?; - Ok(Some(( - SseFrame { - event: frame.event, - data: frame.data, - id: frame.id, - retry: frame.retry, - }, - frames, - ))) - }) - .fuse() + fn decode(&mut self, src: &mut BytesMut) -> Result, SseError> { + if !self.skip_bom(src) { + return Ok(None); + } + while let Some(end) = block_end(src) { + let block = src.split_to(end); + let pending = lines(&block) + .map(|(line, _)| line) + .take_while(|line| !line.is_empty()) + .try_fold(Pending::default(), Pending::apply)?; + if let Some(event) = pending.dispatch() { + return Ok(Some(event)); + } + } + Ok(None) + } + + fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result, SseError> { + Ok(None) + } +} + +impl SseCodec { + fn skip_bom(&mut self, src: &mut BytesMut) -> bool { + if self.past_bom { + return true; + } + if src.starts_with(BOM) { + src.advance(BOM.len()); + } else if BOM.starts_with(src) { + return false; + } + self.past_bom = true; + true + } +} + +fn block_end(bytes: &[u8]) -> Option { + lines(bytes) + .find(|(line, _)| line.is_empty()) + .map(|(_, end)| end) +} + +fn lines(bytes: &[u8]) -> impl Iterator { + let mut cursor: usize = 0; + std::iter::from_fn(move || { + let rest = &bytes[cursor..]; + let end = rest.iter().position(|byte| matches!(byte, b'\n' | b'\r'))?; + cursor += end + terminator_len(&rest[end..]); + Some((&rest[..end], cursor)) + }) +} + +fn terminator_len(terminated: &[u8]) -> usize { + match terminated { + [b'\r', b'\n', ..] => 2, + _ => 1, + } +} + +#[derive(Default)] +struct Pending { + event: Option, + data: Option, + id: Option, + retry: Option, +} + +impl Pending { + fn apply(self, line: &[u8]) -> Result { + let (name, value) = split_field(line); + Ok(match name { + b"event" => Self { + event: Some(str::from_utf8(value)?.to_owned()), + ..self + }, + b"data" => Self { + data: Some(append_data(self.data, str::from_utf8(value)?)), + ..self + }, + b"id" if !value.contains(&0) => Self { + id: Some(str::from_utf8(value)?.to_owned()), + ..self + }, + b"retry" => Self { + retry: parse_retry(value).or(self.retry), + ..self + }, + _ => self, + }) + } + + fn dispatch(self) -> Option { + Some(SseEvent { + event: self.event, + data: self.data?, + id: self.id, + retry: self.retry, + }) + } +} + +fn split_field(line: &[u8]) -> (&[u8], &[u8]) { + let Some(colon) = line.iter().position(|byte| *byte == b':') else { + return (line, &[]); + }; + let value = &line[colon + 1..]; + (&line[..colon], value.strip_prefix(b" ").unwrap_or(value)) +} + +fn append_data(buffer: Option, line: &str) -> String { + match buffer { + Some(existing) => format!("{existing}\n{line}"), + None => line.to_owned(), + } +} + +fn parse_retry(value: &[u8]) -> Option { + if !value.iter().all(u8::is_ascii_digit) { + return None; + } + str::from_utf8(value).ok()?.parse().ok() +} + +impl Encoder for SseCodec { + type Error = SseError; + + fn encode(&mut self, event: SseEvent, dst: &mut BytesMut) -> Result<(), SseError> { + if let Some(name) = event.event { + dst.put_slice(format!("event: {name}\n").as_bytes()); + } + for line in event.data.split('\n') { + dst.put_slice(format!("data: {line}\n").as_bytes()); + } + if let Some(id) = event.id { + dst.put_slice(format!("id: {id}\n").as_bytes()); + } + if let Some(retry) = event.retry { + dst.put_slice(format!("retry: {retry}\n").as_bytes()); + } + dst.put_u8(b'\n'); + Ok(()) } } diff --git a/litellm-rust/crates/framer/tests/aws_event_stream.rs b/litellm-rust/crates/framer/tests/aws_event_stream.rs index c90a15a2b0e..d16caa39948 100644 --- a/litellm-rust/crates/framer/tests/aws_event_stream.rs +++ b/litellm-rust/crates/framer/tests/aws_event_stream.rs @@ -4,89 +4,174 @@ mod support; use std::io; -use futures_util::TryStreamExt; -use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}; -use litellm_framing::{Error, Framer}; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream}; +use litellm_framing::{ + EventStreamError, + aws_event_stream::{AwsEventStreamCodec, Header, HeaderValue, Message}, + frames, +}; +use proptest::prelude::*; use rstest::{fixture, rstest}; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -use support::encode; - -async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result, Error> { - AwsEventStreamFramer - .frame(futures_util::stream::iter( - bytes.chunks(chunk_size).map(Ok::<_, io::Error>), - )) +async fn collect(pieces: Vec) -> Result, EventStreamError> { + frames(input(pieces), AwsEventStreamCodec) .try_collect() .await } -#[fixture] -fn two_frames() -> Vec { - [encode(b"\xff\x00"), encode(b"second")].concat() +fn message(payload: &[u8]) -> Message { + Message::new(Bytes::copy_from_slice(payload)) + .add_header(Header::new( + ":event-type", + HeaderValue::String("payload".into()), + )) + .add_header(Header::new("sequence", HeaderValue::Int32(7))) } #[fixture] fn payload_frame() -> Vec { - encode(b"payload") + encode_all(AwsEventStreamCodec, [message(b"payload")]) +} + +fn header_value() -> impl Strategy { + prop_oneof![ + "[a-z]{0,8}".prop_map(|text| HeaderValue::String(text.into())), + any::().prop_map(HeaderValue::Int32), + any::().prop_map(HeaderValue::Bool), + proptest::collection::vec(any::(), 0..8) + .prop_map(|bytes| HeaderValue::ByteArray(bytes.into())), + ] +} + +fn arbitrary_message() -> impl Strategy { + ( + proptest::collection::vec(("[a-z:-]{1,12}", header_value()), 0..3), + proptest::collection::vec(any::(), 0..32), + ) + .prop_map(|(headers, payload)| { + headers.into_iter().fold( + Message::new(Bytes::from(payload)), + |message, (name, value)| message.add_header(Header::new(name, value)), + ) + }) +} + +proptest! { + #[test] + fn any_messages_survive_a_round_trip_through_any_cuts( + messages in proptest::collection::vec(arbitrary_message(), 1..4), + cuts in proptest::collection::vec(0_usize..512, 0..4), + ) { + let wire = encode_all(AwsEventStreamCodec, messages.clone()); + let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap(); + prop_assert_eq!(decoded, messages); + } } #[rstest] -#[case(1)] -#[case(3)] -#[case(12)] -#[case(usize::MAX)] +#[case::prelude_crc(8)] +#[case::message_crc(usize::MAX)] #[tokio::test] -async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads( - two_frames: Vec, - #[case] chunk_size: usize, -) { - let chunk_size = chunk_size.min(two_frames.len()); - let frames = collect_aws(&two_frames, chunk_size).await.unwrap(); - assert_eq!(frames.len(), 2); - assert_eq!(frames[0].payload, &b"\xff\x00"[..]); - assert_eq!(frames[1].payload, "second"); - assert_eq!( - frames[0].headers[0].value().as_string().unwrap().as_str(), - "payload" - ); - assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7)); -} - -#[rstest] -#[case(8)] -#[case(0)] -#[tokio::test] -async fn rejects_corrupt_crcs(payload_frame: Vec, #[case] index: usize) { - let corrupt_index = if index == 0 { - payload_frame.len() - 1 - } else { - index - }; +async fn a_corrupt_crc_is_malformed(payload_frame: Vec, #[case] index: usize) { let mut corrupt = payload_frame; - corrupt[corrupt_index] ^= 1; - assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_)))); -} - -#[rstest] -#[case(0_u32)] -#[case(15)] -#[case(u32::MAX)] -#[tokio::test] -async fn rejects_invalid_lengths(#[case] length: u32) { + let flipped = index.min(corrupt.len() - 1); + corrupt[flipped] ^= 1; assert!(matches!( - collect_aws(&length.to_be_bytes(), 1).await, - Err(Error::InvalidLength(_)) + collect(every(&corrupt, 3)).await, + Err(EventStreamError::Malformed(_)) )); } #[rstest] -#[case(1)] -#[case(3)] -#[case(5)] +#[case::zero(0)] +#[case::below_minimum(15)] +#[case::above_maximum(16 * 1024 * 1024 + 1)] +#[case::u32_max(u32::MAX)] #[tokio::test] -async fn rejects_truncation(payload_frame: Vec, #[case] end: usize) { +async fn a_length_outside_the_frame_bounds_fails_before_buffering(#[case] length: u32) { assert!(matches!( - collect_aws(&payload_frame[..end], 1).await, - Err(Error::Truncated) + collect(every(&length.to_be_bytes(), 1)).await, + Err(EventStreamError::InvalidLength(seen)) if seen == length as usize )); } + +#[rstest] +#[case::before_the_length(1)] +#[case::inside_the_prelude(5)] +#[case::one_byte_short(usize::MAX)] +#[tokio::test] +async fn eof_inside_a_frame_is_truncation(payload_frame: Vec, #[case] end: usize) { + let end = end.min(payload_frame.len() - 1); + assert!(matches!( + collect(every(&payload_frame[..end], 1)).await, + Err(EventStreamError::Truncated) + )); +} + +const FRAME_OVERHEAD_BYTES: usize = 16; +const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; + +#[tokio::test] +async fn a_frame_at_exactly_the_maximum_length_decodes() { + let largest = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES]); + let wire = encode_all(AwsEventStreamCodec, [largest.clone()]); + assert_eq!(wire.len(), MAX_FRAME_BYTES); + assert_eq!(collect(every(&wire, 1 << 20)).await.unwrap(), vec![largest]); +} + +#[tokio::test] +async fn a_frame_one_byte_over_the_maximum_length_is_rejected_by_its_prelude() { + let oversized = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES + 1]); + let wire = encode_all(AwsEventStreamCodec, [oversized]); + assert!(matches!( + collect(every(&wire[..4], 1)).await, + Err(EventStreamError::InvalidLength(length)) if length == MAX_FRAME_BYTES + 1 + )); +} + +#[tokio::test] +async fn an_empty_body_yields_nothing() { + assert_eq!(collect(vec![]).await.unwrap(), vec![]); +} + +#[tokio::test] +async fn a_complete_frame_precedes_a_truncated_following_frame() { + let wire = encode_all(AwsEventStreamCodec, [message(b"first"), message(b"second")]); + let mut messages = Box::pin(frames( + input(every(&wire[..wire.len() - 1], 3)), + AwsEventStreamCodec, + )); + + assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first")); + assert!(matches!( + messages.next().await, + Some(Err(EventStreamError::Truncated)) + )); + assert!(messages.next().await.is_none()); +} + +#[tokio::test] +async fn a_body_error_after_a_complete_frame_preserves_its_cause() { + let first = encode_all(AwsEventStreamCodec, [message(b"first")]); + let mut messages = Box::pin(frames( + stream::iter([ + Ok(cut_at(&first, [5])[0].clone()), + Ok(cut_at(&first, [5])[1].clone()), + Ok(Bytes::from_static(b"\0\0\0")), + Err(io::Error::new(io::ErrorKind::ConnectionReset, "reset")), + ]), + AwsEventStreamCodec, + )); + + assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first")); + let Some(Err(EventStreamError::Body(body))) = messages.next().await else { + panic!("the body error surfaces"); + }; + assert_eq!( + body_cause::(&body).unwrap().kind(), + io::ErrorKind::ConnectionReset + ); + assert!(messages.next().await.is_none()); +} diff --git a/litellm-rust/crates/framer/tests/chaining.rs b/litellm-rust/crates/framer/tests/chaining.rs index afd24a90704..81884d58ba1 100644 --- a/litellm-rust/crates/framer/tests/chaining.rs +++ b/litellm-rust/crates/framer/tests/chaining.rs @@ -2,28 +2,64 @@ mod support; -use std::io; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt}; +use litellm_framing::{ + EventStreamError, SseError, + aws_event_stream::{AwsEventStreamCodec, Message}, + frames, + sse::{SseCodec, SseEvent}, +}; +use proptest::prelude::*; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -use futures_util::TryStreamExt; -use litellm_framing::Framer; -use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}; -use litellm_framing::sse::SseFramer; +fn delta(data: &str) -> SseEvent { + SseEvent { + event: Some("delta".into()), + data: data.into(), + id: Some("7".into()), + retry: None, + } +} -use support::encode; +fn envelopes(payloads: Vec) -> Vec { + encode_all(AwsEventStreamCodec, payloads.into_iter().map(Message::new)) +} + +proptest! { + #[test] + fn an_sse_event_cut_anywhere_across_envelopes_is_reassembled(cut in 0_usize..64, chunk in 1_usize..8) { + let sse = encode_all(SseCodec::default(), [delta("hello")]); + let wire = envelopes(cut_at(&sse, [cut.min(sse.len())])); + let events = runtime().block_on(async { + let payloads = frames(input(every(&wire, chunk)), AwsEventStreamCodec) + .map_ok(|message| message.payload().clone()); + frames(payloads, SseCodec::default()).try_collect::>().await + }) + .unwrap(); + prop_assert_eq!(events, vec![delta("hello")]); + } +} #[tokio::test] -async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() { - let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat(); - let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter( - bytes.chunks(3).map(Ok::<_, io::Error>), +async fn a_truncated_envelope_after_an_sse_event_keeps_the_event_and_its_cause() { + let complete = encode_all(SseCodec::default(), [delta("complete")]); + let incomplete = encode_all(SseCodec::default(), [delta("incomplete")]); + let wire = envelopes(vec![complete.into(), incomplete.into()]); + let payloads = frames( + input(every(&wire[..wire.len() - 1], 3)), + AwsEventStreamCodec, + ) + .map_ok(|message| message.payload().clone()); + let mut events = Box::pin(frames(payloads, SseCodec::default())); + + assert_eq!(events.next().await.unwrap().unwrap(), delta("complete")); + let Some(Err(SseError::Body(body))) = events.next().await else { + panic!("the envelope error surfaces through the SSE layer"); + }; + assert!(matches!( + body_cause::(&body), + Some(EventStreamError::Truncated) )); - let frames = SseFramer - .frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload)) - .try_collect::>() - .await - .unwrap(); - assert_eq!(frames.len(), 1); - assert_eq!(frames[0].event.as_deref(), Some("delta")); - assert_eq!(frames[0].data.as_deref(), Some("hello")); - assert_eq!(frames[0].id.as_deref(), Some("7")); + assert!(events.next().await.is_none()); } diff --git a/litellm-rust/crates/framer/tests/sse.rs b/litellm-rust/crates/framer/tests/sse.rs index 66339dfbfd2..2fa064653a6 100644 --- a/litellm-rust/crates/framer/tests/sse.rs +++ b/litellm-rust/crates/framer/tests/sse.rs @@ -1,67 +1,169 @@ #![cfg(feature = "sse")] +mod support; + use std::io; -use futures_util::{StreamExt, TryStreamExt}; -use litellm_framing::sse::{SseFrame, SseFramer}; -use litellm_framing::{Error, Framer}; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream}; +use litellm_framing::{ + SseError, frames, + sse::{SseCodec, SseEvent}, +}; +use proptest::prelude::*; use rstest::rstest; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -async fn collect_sse(chunks: &[&[u8]]) -> Result, Error> { - SseFramer - .frame(futures_util::stream::iter( - chunks.iter().copied().map(Ok::<_, io::Error>), - )) +async fn collect(pieces: Vec) -> Result, SseError> { + frames(input(pieces), SseCodec::default()) .try_collect() .await } +fn event(name: Option<&str>, data: &str) -> SseEvent { + SseEvent { + event: name.map(str::to_owned), + data: data.to_owned(), + id: None, + retry: None, + } +} + +fn sse_event() -> impl Strategy { + ( + proptest::option::of("[^\r\n\0]{0,8}"), + "[^\r\0]{0,16}", + proptest::option::of("[^\r\n\0]{0,8}"), + proptest::option::of(any::()), + ) + .prop_map(|(event, data, id, retry)| SseEvent { + event, + data, + id, + retry, + }) +} + +fn terminators() -> impl Strategy { + prop_oneof![Just(&b"\n"[..]), Just(&b"\r\n"[..]), Just(&b"\r"[..])] +} + +proptest! { + #[test] + fn any_events_survive_a_round_trip_through_any_terminator_and_any_cuts( + events in proptest::collection::vec(sse_event(), 1..4), + terminator in terminators(), + cuts in proptest::collection::vec(0_usize..256, 0..4), + bom in any::(), + ) { + let lf_wire = encode_all(SseCodec::default(), events.clone()); + let body: Vec = lf_wire + .iter() + .flat_map(|byte| if *byte == b'\n' { terminator.to_vec() } else { vec![*byte] }) + .collect(); + let wire = if bom { [&b"\xEF\xBB\xBF"[..], &body].concat() } else { body }; + let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap(); + prop_assert_eq!(decoded, events); + } +} + #[rstest] -#[case( - &[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]], - vec![ - SseFrame { - event: Some("delta".into()), - data: Some("€\nnext".into()), - id: Some("7".into()), - retry: Some(10), - }, - SseFrame { - event: None, - data: Some("[DONE]".into()), - id: None, - retry: None, - }, - ] -)] +#[case::comment(b":ping\ndata: x\n\n")] +#[case::unknown_field(b"vendor: 1\ndata: x\n\n")] +#[case::field_without_colon(b"garbage\ndata: x\n\n")] +#[case::retry_with_non_digits(b"retry: soon\ndata: x\n\n")] +#[case::retry_with_a_sign(b"retry: +5\ndata: x\n\n")] +#[case::retry_without_a_value(b"retry:\ndata: x\n\n")] +#[case::id_with_nul(b"id: a\0b\ndata: x\n\n")] #[tokio::test] -async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel( - #[case] chunks: &[&[u8]], - #[case] expected: Vec, +async fn lines_the_spec_ignores_do_not_change_the_event(#[case] wire: &[u8]) { + assert_eq!( + collect(every(wire, 1)).await.unwrap(), + vec![event(None, "x")] + ); +} + +#[rstest] +#[case::no_data_at_all(b"event: ping\nid: 1\n\ndata: x\n\n", vec![event(None, "x")])] +#[case::empty_data_field(b"data:\n\n", vec![event(None, "")])] +#[case::one_leading_space_stripped(b"data: x\n\n", vec![event(None, " x")])] +#[case::multiline_data(b"data: a\ndata: b\ndata:\n\n", vec![event(None, "a\nb\n")])] +#[case::last_event_name_wins(b"event: a\nevent: b\ndata: x\n\n", vec![event(Some("b"), "x")])] +#[case::last_retry_wins(b"retry: 1\nretry: 2\ndata: x\n\n", vec![SseEvent { retry: Some(2), ..event(None, "x") }])] +#[case::split_utf8_across_lines_is_not_joined(b"data: \xe2\x82\xac\ndata: \xe2\x82\xac\n\n", vec![event(None, "€\n€")])] +#[tokio::test] +async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected: Vec) { + assert_eq!(collect(every(wire, 1)).await.unwrap(), expected); +} + +#[rstest] +#[case::unterminated_single(b"data: partial\n", vec![])] +#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])] +#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])] +#[case::lone_cr_line_then_eof(b"data: x\r", vec![])] +#[tokio::test] +async fn eof_dispatches_only_terminated_events( + #[case] wire: &[u8], + #[case] expected: Vec, ) { - assert_eq!(collect_sse(chunks).await.unwrap(), expected); + assert_eq!( + collect(vec![Bytes::copy_from_slice(wire)]).await.unwrap(), + expected + ); +} + +#[rstest] +#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])] +#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])] +#[tokio::test] +async fn a_crlf_split_across_chunks_is_one_terminator(#[case] pieces: Vec<&[u8]>) { + let pieces = pieces.into_iter().map(Bytes::copy_from_slice).collect(); + assert_eq!(collect(pieces).await.unwrap(), vec![event(None, "a\nb")]); } #[tokio::test] -async fn eof_does_not_dispatch_an_unterminated_frame() { - assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty()); +async fn a_bom_is_stripped_only_at_the_start_of_the_stream() { + let wire = b"\xEF\xBB\xBFdata: a\n\n\xEF\xBB\xBFdata: b\ndata: c\n\n"; + let decoded = collect(every(wire, 2)).await.unwrap(); + assert_eq!(decoded, vec![event(None, "a"), event(None, "c")]); +} + +#[tokio::test] +async fn invalid_utf8_in_a_field_fails_after_earlier_events_and_terminates() { + let mut events = Box::pin(frames( + input(every(b"data: ok\n\ndata: \xff\n\n", 3)), + SseCodec::default(), + )); + + assert_eq!(events.next().await.unwrap().unwrap(), event(None, "ok")); + assert!(matches!( + events.next().await, + Some(Err(SseError::InvalidUtf8(_))) + )); + assert!(events.next().await.is_none()); } #[rstest] #[case(io::ErrorKind::ConnectionReset)] #[case(io::ErrorKind::UnexpectedEof)] #[tokio::test] -async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) { - let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([ - Err(io::Error::new(kind, "reset")), - Ok(&b"data: later\n\n"[..]), - ]))); - let error = frames.next().await.unwrap().unwrap_err(); - assert!(matches!( - error, - Error::Sse(sse_stream::Error::Body(ref cause)) - if cause.downcast_ref::().unwrap().kind() == kind +async fn a_body_error_keeps_earlier_events_and_its_cause_then_terminates( + #[case] kind: io::ErrorKind, +) { + let mut events = Box::pin(frames( + stream::iter([ + Ok(&b"data: first\n\ndata: partial"[..]), + Err(io::Error::new(kind, "reset")), + Ok(&b"\n\n"[..]), + ]), + SseCodec::default(), )); - assert!(frames.next().await.is_none()); - assert!(frames.next().await.is_none()); + + assert_eq!(events.next().await.unwrap().unwrap(), event(None, "first")); + let Some(Err(SseError::Body(body))) = events.next().await else { + panic!("the body error surfaces"); + }; + assert_eq!(body_cause::(&body).unwrap().kind(), kind); + assert!(events.next().await.is_none()); + assert!(events.next().await.is_none()); } diff --git a/litellm-rust/crates/framer/tests/support/mod.rs b/litellm-rust/crates/framer/tests/support/mod.rs index 9db305af073..9ff67aef149 100644 --- a/litellm-rust/crates/framer/tests/support/mod.rs +++ b/litellm-rust/crates/framer/tests/support/mod.rs @@ -1,15 +1,57 @@ -use aws_smithy_eventstream::frame::write_message_to; -use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; -use bytes::Bytes; +#![allow(dead_code)] -pub fn encode(payload: &'static [u8]) -> Vec { - let message = Message::new(Bytes::from_static(payload)) - .add_header(Header::new( - ":event-type", - HeaderValue::String("payload".into()), - )) - .add_header(Header::new("sequence", HeaderValue::Int32(7))); - let mut bytes = Vec::new(); - write_message_to(&message, &mut bytes).unwrap(); - bytes +use std::{error::Error, io}; + +use bytes::{Bytes, BytesMut}; +use futures_util::{Stream, stream}; +use tokio_util::codec::Encoder; + +pub fn encode_all(mut codec: C, items: impl IntoIterator) -> Vec +where + C: Encoder, + C::Error: std::fmt::Debug, +{ + let mut wire = BytesMut::new(); + for item in items { + codec.encode(item, &mut wire).unwrap(); + } + wire.to_vec() +} + +pub fn cut_at(bytes: &[u8], offsets: impl IntoIterator) -> Vec { + let mut sorted: Vec = offsets + .into_iter() + .filter(|offset| *offset <= bytes.len()) + .collect(); + sorted.sort_unstable(); + sorted.dedup(); + let bounds = std::iter::once(0) + .chain(sorted) + .chain(std::iter::once(bytes.len())) + .collect::>(); + bounds + .windows(2) + .map(|pair| Bytes::copy_from_slice(&bytes[pair[0]..pair[1]])) + .collect() +} + +pub fn every(bytes: &[u8], size: usize) -> Vec { + bytes + .chunks(size.max(1)) + .map(Bytes::copy_from_slice) + .collect() +} + +pub fn input(pieces: Vec) -> impl Stream> + Send { + stream::iter(pieces.into_iter().map(Ok)) +} + +pub fn body_cause(body: &io::Error) -> Option<&T> { + body.get_ref()?.downcast_ref::() +} + +pub fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() } diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs index 35e7d5820b0..3f1b7ed9bcc 100644 --- a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs @@ -2,9 +2,9 @@ use base64::Engine; use bytes::Buf; use futures_util::{Stream, StreamExt}; use litellm_framing::{ - Framer, - aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}, - sse::{SseFrame, SseFramer}, + aws_event_stream::{AwsEventStreamCodec, Message}, + frames, + sse::{SseCodec, SseEvent}, }; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -13,8 +13,6 @@ use serde_json::{Map, Value}; pub enum Error { #[error("stream framing failed: {0}")] StreamFraming(String), - #[error("Anthropic SSE frame has no data")] - MissingStreamData, #[error("Anthropic stream event is invalid: {0}")] InvalidStreamEvent(String), #[error("Bedrock event payload is invalid: {0}")] @@ -165,15 +163,14 @@ struct BedrockChunkPayload { bytes: String, } -pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result { - let data = frame.data.ok_or(Error::MissingStreamData)?; - serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string())) +pub fn decode_anthropic_sse_frame(event: SseEvent) -> Result { + serde_json::from_str(&event.data).map_err(|error| Error::InvalidStreamEvent(error.to_string())) } pub fn decode_bedrock_anthropic_frame( - frame: AwsEventStreamFrame, + message: Message, ) -> Result { - let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload) + let payload: BedrockChunkPayload = serde_json::from_slice(message.payload()) .map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?; let event = base64::engine::general_purpose::STANDARD .decode(payload.bytes) @@ -189,9 +186,8 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - SseFramer.frame(input).map(|frame| { - let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?; - decode_anthropic_sse_frame(frame) + frames(input, SseCodec::default()).map(|event| { + decode_anthropic_sse_frame(event.map_err(|error| Error::StreamFraming(error.to_string()))?) }) } @@ -203,9 +199,10 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - AwsEventStreamFramer.frame(input).map(|frame| { - let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?; - decode_bedrock_anthropic_frame(frame) + frames(input, AwsEventStreamCodec).map(|message| { + decode_bedrock_anthropic_frame( + message.map_err(|error| Error::StreamFraming(error.to_string()))?, + ) }) } @@ -247,12 +244,10 @@ mod tests { #[test] fn decodes_citations_delta_events() { - let event = decode_anthropic_sse_frame(SseFrame { + let event = decode_anthropic_sse_frame(SseEvent { event: Some("content_block_delta".into()), - data: Some( - r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"# - .into(), - ), + data: r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"# + .into(), id: None, retry: None, }) From 636eb4c396c194235d375db6b021e90f7a53099c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 12:22:48 -0700 Subject: [PATCH 12/58] fix(anthropic): surface Responses bridge stream failures as Anthropic error events (#43126) * fix(anthropic): surface Responses bridge stream failures as Anthropic error events The /v1/messages Responses bridge logged every upstream failure and ended the SSE stream as if it had completed, so a rate limit, a provider 500, a dropped connection, or a read timeout reached the client as HTTP 200 with a lone message_start and no error event. Map response.failed and any raised upstream exception to a redacted Anthropic error frame, stop pulling upstream after it, and never fabricate end_turn or message_stop after a failure. * fix(anthropic): close a Responses bridge stream that ends without a terminal event with an error event Normalize the failure status behind the error type to an int or digit string within 400..599, narrow the response.failed event through pydantic, reuse the native Messages path's incomplete-stream message for a clean upstream EOF, and cover the pydantic event, the unwrapped fallback error, and the EOF cases --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../messages/streaming_iterator.py | 6 +- .../messages/utils.py | 6 + .../responses_adapters/streaming_iterator.py | 114 +++++++++++- litellm/responses/streaming_iterator.py | 5 + ...t_responses_adapters_streaming_iterator.py | 176 +++++++++++++++++- 5 files changed, 291 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 0bd46382fef..5550590d0c0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP +from litellm.llms.anthropic.experimental_pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -28,11 +29,6 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains -INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( - "Provider stream ended before emitting a message_stop event; " - "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." -) - def _is_message_stop_chunk(chunk: object) -> bool: if isinstance(chunk, dict): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 89105c00428..fe8ac2cd7a2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -15,6 +15,12 @@ if TYPE_CHECKING: from litellm.exceptions import ContentPolicyViolationError +INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( + "Provider stream ended before emitting a message_stop event; " + "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." +) + + def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None: """ Return the ``stop_details`` of an Anthropic Messages response refused by a diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f753e87fee3..59ccde872fc 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -2,20 +2,25 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio import json -import traceback from collections import deque from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final +from pydantic import BaseModel, ConfigDict, field_validator + from litellm import verbose_logger +from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + INCOMPLETE_STREAM_ERROR_MESSAGE, refusal_stop_details, responses_output_refusal_text, ) +from litellm.responses.streaming_iterator import stream_error_status_and_message from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from .transformation import ( @@ -27,6 +32,72 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +class _UpstreamFailure(BaseModel): + model_config = ConfigDict(frozen=True) + + status_code: int | None = None + message: str | None = None + + @field_validator("status_code", mode="before") + @classmethod + def http_error_status_or_none(cls, value: object) -> int | None: + candidate: Final = ( + value + if isinstance(value, int) and not isinstance(value, bool) + else int(value) + if isinstance(value, str) and value.isdecimal() + else None + ) + return candidate if candidate is not None and 400 <= candidate <= 599 else None + + @field_validator("message", mode="before") + @classmethod + def str_or_none(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + + +class _FailedResponse(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + error: object | None = None + + +class _FailedResponseEvent(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + response: _FailedResponse | None = None + + +def _original_failure(exception: Exception) -> Exception: + failure = exception # rebind-ok: walks the MidStreamFallbackError chain down to the provider failure + while isinstance(failure, MidStreamFallbackError) and failure.original_exception is not None: + failure = failure.original_exception + return failure + + +def _failure_status_and_message(exception: Exception) -> tuple[int, str]: + original: Final = _original_failure(exception) + failure: Final = _UpstreamFailure.model_validate( + {"status_code": getattr(original, "status_code", None), "message": getattr(original, "message", None)} + ) + status_code: Final = failure.status_code if failure.status_code is not None else 500 + message: Final = failure.message or str(original) or INCOMPLETE_STREAM_ERROR_MESSAGE + return status_code, message + + +def _anthropic_error_chunk(status_code: int, message: str) -> dict[str, object]: + from litellm.anthropic_interface.exceptions.exception_mapping_utils import ( + AnthropicExceptionMapping, + ) + + return dict( + AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status_code, + raw_message=redact_internal_details_from_client_message(message), + ) + ) + + class AnthropicResponsesStreamWrapper: """ Wraps a Responses API streaming iterator and re-emits events in Anthropic SSE format. @@ -40,6 +111,7 @@ class AnthropicResponsesStreamWrapper: response.function_call_arguments.delta -> content_block_delta (input_json_delta) response.output_item.done -> content_block_delta (signature_delta) + content_block_stop response.completed -> message_delta + message_stop + response.failed -> error (the stream ends without message_stop) """ def __init__( @@ -60,6 +132,7 @@ class AnthropicResponsesStreamWrapper: self._pending_tool_ids: dict[str, str] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False + self._stream_failed = False self._chunk_queue: deque[dict[str, object]] = deque() self._refusal_text: str = "" self._sync_responses_iterator: Iterator[object] | None = None @@ -293,10 +366,23 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.failed": + failed: Final = _FailedResponseEvent.model_validate(event) + status_code, message = stream_error_status_and_message( + failed.response.error if failed.response is not None else None + ) + verbose_logger.error( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s", + self.model, + status_code, + message, + ) + self._fail_stream(status_code, message) + return + # ---- response completed -> message_delta + message_stop ---- if event_type in ( "response.completed", - "response.failed", "response.incomplete", ): response_obj: Final = getattr(event, "response", None) or ( @@ -350,21 +436,24 @@ class AnthropicResponsesStreamWrapper: self._sent_message_stop = True return + def _fail_stream(self, status_code: int, message: str) -> None: + self._stream_failed = True + self._chunk_queue.append(_anthropic_error_chunk(status_code, message)) + def __aiter__(self) -> "AnthropicResponsesStreamWrapper": return self async def __anext__(self) -> dict[str, object]: - # Return any queued chunks first if self._chunk_queue: return self._chunk_queue.popleft() + if self._stream_failed: + raise StopAsyncIteration - # Emit message_start if not yet done (fallback if response.created wasn't fired) if not self._sent_message_start: self._sent_message_start = True self._chunk_queue.append(self._make_message_start()) return self._chunk_queue.popleft() - # Consume the upstream stream try: if hasattr(self.responses_stream, "__aiter__"): async for event in self.responses_stream: @@ -382,10 +471,19 @@ class AnthropicResponsesStreamWrapper: return self._chunk_queue.popleft() except StopAsyncIteration: pass - except Exception as e: - verbose_logger.error("AnthropicResponsesStreamWrapper error: %s\n%s", e, traceback.format_exc()) + except Exception as e: # noqa: BLE001 # every upstream failure becomes a client error event + verbose_logger.exception( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model + ) + self._fail_stream(*_failure_status_and_message(e)) + + if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed: + verbose_logger.error( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s ended without a terminal event", + self.model, + ) + self._fail_stream(500, INCOMPLETE_STREAM_ERROR_MESSAGE) - # Drain any remaining queued chunks if self._chunk_queue: return self._chunk_queue.popleft() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 70f2a7db6da..fdc702af005 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -230,6 +230,11 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500) +def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: + message, error_type, error_code = _error_event_fields(error_obj) + return _status_code_for_error_fields(error_type, error_code), message + + def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index bfe2d6b7cea..392ecc2bcdd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -4,18 +4,25 @@ Tests for AnthropicResponsesStreamWrapper """ import asyncio +import json import os import sys from types import SimpleNamespace +import pytest + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) +from litellm.llms.anthropic.experimental_pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) +from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse def _process_all(events: list) -> list: @@ -132,6 +139,7 @@ class TestReasoningItemWithoutSummaryText: {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.completed"}, ] def test_reasoning_without_summary_emits_no_thinking_block(self): @@ -144,6 +152,8 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 0), ("content_block_delta", 0), ("content_block_stop", 0), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == {"type": "text", "text": ""} @@ -166,6 +176,8 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 1), ("content_block_delta", 1), ("content_block_stop", 1), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" @@ -215,6 +227,8 @@ class TestEncryptedReasoningIsStreamedForReplay: ("content_block_start", 1), ("content_block_delta", 1), ("content_block_stop", 1), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == { "type": "redacted_thinking", @@ -234,9 +248,7 @@ class TestEncryptedReasoningIsStreamedForReplay: ] chunks = _process_all(events) - thinking = "".join( - c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta" - ) + thinking = "".join(c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta") assert thinking == "First.\n\nSecond." assert [c["type"] for c in chunks].count("content_block_start") == 1 @@ -283,6 +295,7 @@ class TestToolUseBlockClosedExactlyOnce: "type": "response.output_item.done", "item": {"type": "message", "id": "chatcmpl-123", "status": "completed"}, }, + {"type": "response.completed"}, ] def test_one_content_block_stop_per_content_block_start(self): @@ -302,6 +315,8 @@ class TestToolUseBlockClosedExactlyOnce: ("content_block_delta", 0), ("content_block_delta", 0), ("content_block_stop", 0), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == { "type": "tool_use", @@ -452,3 +467,158 @@ class TestRefusalStreamEvents: message_delta = next(c for c in chunks if c["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "max_tokens" assert "stop_details" not in message_delta["delta"] + + +def _collect(stream) -> list: + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m") + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + +class TestUpstreamFailureEndsStreamWithErrorEvent: + """A provider failure must reach the Anthropic client as an ``error`` event that + ends the stream, never as a fabricated ``end_turn`` or a silent close.""" + + def test_response_failed_event_emits_error_event_and_stops_pulling_upstream(self): + failed = SimpleNamespace( + status="failed", + output=[], + usage=None, + error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."}, + ) + + async def _gen(): + yield {"type": "response.created"} + yield {"type": "response.failed", "response": failed} + raise AssertionError("upstream was pulled again after the failure") + + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=_gen(), model="m") + return [frame async for frame in wrapper.async_anthropic_sse_wrapper()] + + frames = asyncio.run(_run()) + assert [frame.split(b"\n", 1)[0] for frame in frames] == [b"event: message_start", b"event: error"] + error_payload = json.loads(frames[1].split(b"data: ", 1)[1]) + assert error_payload["type"] == "error" + assert error_payload["error"] == { + "type": "rate_limit_error", + "message": "Rate limit reached for gpt-5.5, try again in 20s.", + } + + def test_raised_mid_stream_fallback_error_is_unwrapped_to_the_provider_failure(self): + rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + wrapped = MidStreamFallbackError( + message=str(rate_limit), + model="m", + llm_provider="openai", + original_exception=rate_limit, + is_pre_first_chunk=True, + ) + + async def _gen(): + yield {"type": "response.created"} + raise wrapped + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == {"type": "rate_limit_error", "message": rate_limit.message} + + def test_sync_upstream_transport_error_after_content_becomes_api_error_event(self): + def _events(): + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise ConnectionResetError("Response payload is not completed") + + chunks = _collect(_events()) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"} + + def test_error_event_message_is_redacted_before_it_reaches_the_client(self): + async def _gen(): + yield {"type": "response.created"} + raise RuntimeError("upstream failed with key sk-proj-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJ") + + chunks = _collect(_gen()) + assert chunks[-1]["type"] == "error" + assert "sk-proj-" not in chunks[-1]["error"]["message"] + assert chunks[-1]["error"]["message"].startswith("upstream failed with key") + + @pytest.mark.parametrize( + ("raised", "expected_error"), + [ + ( + MidStreamFallbackError(message="boom", model="m", llm_provider="openai"), + {"type": "api_error", "message": "litellm.MidStreamFallbackError: boom"}, + ), + ( + type("StringStatusError", (Exception,), {"status_code": "429"})("throttled"), + {"type": "rate_limit_error", "message": "throttled"}, + ), + ( + type("NonErrorStatusError", (Exception,), {"status_code": 200})("odd status"), + {"type": "api_error", "message": "odd status"}, + ), + ], + ids=["mid-stream-fallback-without-original", "digit-string-status", "status-outside-4xx-5xx"], + ) + def test_raised_failure_status_is_normalized_into_the_error_type(self, raised, expected_error): + async def _gen(): + yield {"type": "response.created"} + raise raised + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == expected_error + + def test_pydantic_response_failed_event_is_mapped_like_a_dict_event(self): + failed = ResponsesAPIResponse( + id="resp_1", + created_at=1, + error={"code": "server_error", "message": "The server had an error while processing your request."}, + status="failed", + output=[], + model="m", + object="response", + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + async def _gen(): + yield {"type": "response.created"} + yield ResponseFailedEvent(type="response.failed", response=failed) + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == { + "type": "api_error", + "message": "The server had an error while processing your request.", + } + + def test_upstream_ending_without_a_terminal_event_is_an_error_not_a_silent_close(self): + async def _gen(): + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert chunks[-1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} + + def test_sync_upstream_ending_before_any_event_is_an_error_not_a_silent_close(self): + chunks = _collect(iter(())) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} From 5e6dc89ba169167fedd64e48171e5c0152a43687 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 12:43:23 -0700 Subject: [PATCH 13/58] test: move tests/test_litellm/llms into tests/unit/llms (#43191) * ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: move tests/test_litellm/llms into tests/unit/llms Rename-only. Moves the provider tests and the fine-tuning fixtures they load, mirroring the old paths. Follow-up commits merge, split and wire them. * test: merge, split and prune the moved llms tests Merges the Databricks chat transformation tests into the existing unit file, keeps the tests that need real keys or the network in tests/test_litellm, deletes the audited tests a stronger unit test already covers, and points imports at tests.unit.llms. * ci: run the moved llms tests under their legacy flags The Vertex AI and All Other Providers shards keep their legacy test-path for the retained files and add the llm-vertex-ai and llm-other-providers unit selections. CircleCI gets matching unit jobs. * test: make the tests/unit/llms directories packages Adds __init__.py to the moved dirs and drops the legacy ones whose directories no longer hold tests. * test: drop script runners and path hacks the llms split left dangling The __main__ runners in the split openai_like files and the Databricks e2e runner called tests that now live in the other half of the split or were deleted. The retained legacy halves also no longer need sys.path edits. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. * test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path The Databricks e2e file is a manual script whose main() calls the tests that were pruned, so pruning them broke the documented run. It is back to its main version. The SageMaker Nova docstring now points at the file's real location in tests/local_testing. * test: keep the job's UNIT_FLAG out of the shard-script tests --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/unit_selection.sh | 4 + .circleci/tests.yml | 15 + .github/merge-smoke-tests.json | 6 +- .github/workflows/test-unit.yml | 2 + Makefile | 2 +- tests/llm_translation/test_bedrock_gpt_oss.py | 2 +- tests/local_testing/test_function_calling.py | 2 +- .../test_handler_gc_does_not_close_client.py | 2 +- .../test_sagemaker_nova_integration.py | 4 +- .../test_bing_grounding_search.py | 2 +- tests/search_tests/test_nimble_search.py | 2 +- .../integrations/test_helicone.py | 2 +- .../chat/test_cometapi_chat_transformation.py | 165 - .../test_databricks_chat_transformation.py | 79 - .../test_deepinfra_rerank_integration.py | 433 --- .../llms/gemini/files/__init__.py | 1 - .../llms/gemini/videos/__init__.py | 1 - tests/test_litellm/llms/manus/__init__.py | 1 - .../llms/manus/responses/__init__.py | 1 - tests/test_litellm/llms/minimax/__init__.py | 1 - .../llms/minimax/chat/__init__.py | 1 - .../llms/minimax/messages/__init__.py | 1 - ...tral_audio_transcription_transformation.py | 191 -- .../llms/openai_like/test_json_providers.py | 363 +-- .../llms/openai_like/test_xiaomi_mimo.py | 103 +- ...loud_audio_transcription_transformation.py | 58 - .../test_ovhcloud_chat_transformation.py | 237 -- tests/test_litellm/llms/reducto/__init__.py | 1 - .../test_litellm/llms/s3_vectors/__init__.py | 1 - .../llms/s3_vectors/vector_stores/__init__.py | 1 - tests/test_litellm/llms/soniox/__init__.py | 1 - .../test_vertex_ai_gemini_transformation.py | 2733 +---------------- .../llms/vertex_ai/image_edit/__init__.py | 1 - ...rtex_ai_image_generation_transformation.py | 624 +--- .../vertex_ai/vertex_gemma_models/__init__.py | 1 - .../llms/vertex_ai/videos/__init__.py | 3 - .../test_bedrock_guardrails.py | 2 +- .../test_bedrock_invoke_guardrail_checks.py | 2 +- .../test_llm_pass_through_endpoints.py | 2 +- tests/unit/conftest.py | 50 + .../expected_fine_tuning_api}/__init__.py | 0 .../azure_cancel_expected_output.json | 0 .../azure_cancel_raw_response.json | 0 .../azure_cancel_request.json | 0 .../azure_create_expected_output.json | 0 .../azure_create_raw_response.json | 0 .../azure_create_request.json | 0 .../azure_list_raw_response.json | 0 .../azure_list_request.json | 0 .../batches => unit/llms/aiml}/__init__.py | 0 .../llms/aiml/image_generation}/__init__.py | 0 ...st_aiml_image_generation_transformation.py | 0 .../anthropic/batches/test_transformation.py | 2 +- .../llms/anthropic/chat}/__init__.py | 0 .../llms/anthropic/chat/conftest.py | 0 .../chat/guardrail_translation}/__init__.py | 0 .../test_anthropic_guardrail_handler.py | 0 .../chat/test_anthropic_chat_handler.py | 0 .../test_anthropic_chat_transformation.py | 0 ...est_code_interpreter_results_extraction.py | 0 .../adapters}/__init__.py | 0 ...al_pass_through_adapters_transformation.py | 0 .../test_handler_output_config_passthrough.py | 0 .../adapters/test_handler_prompt_cache_key.py | 0 ..._handler_reasoning_effort_normalization.py | 0 .../test_streaming_iterator_combined_chunk.py | 0 .../test_streaming_iterator_compaction.py | 0 .../test_streaming_iterator_empty_choices.py | 0 .../test_streaming_iterator_first_delta.py | 0 .../test_streaming_iterator_message_id.py | 0 ...est_streaming_iterator_mid_stream_error.py | 0 .../test_streaming_iterator_stop_reason.py | 0 .../test_streaming_iterator_tool_args.py | 0 .../context_management}/__init__.py | 0 .../test_clear_tool_uses.py | 0 .../context_management/test_compact.py | 0 .../context_management/test_dispatcher.py | 0 .../messages}/__init__.py | 0 .../messages/test_advisor_integration.py | 0 .../test_agentic_streaming_iterator.py | 0 ...erimental_pass_through_messages_handler.py | 0 .../test_anthropic_messages_effort.py | 0 ..._anthropic_messages_encrypted_reasoning.py | 0 ...est_anthropic_messages_per_turn_control.py | 0 .../messages/test_anthropic_messages_speed.py | 0 ...t_anthropic_messages_structured_outputs.py | 0 .../test_content_after_stop_reason.py | 0 .../messages/test_mcp_handler.py | 0 .../messages/test_mid_conversation_system.py | 0 .../messages/test_parallel_tool_calls.py | 0 .../test_reasoning_auto_summary_messages.py | 0 .../test_reasoning_effort_translation.py | 0 .../test_request_optional_param_utils.py | 0 .../messages/test_response_cache.py | 0 .../messages/test_sse_wrapper.py | 0 .../messages/test_streaming_iterator.py | 0 .../responses_adapters}/__init__.py | 0 .../test_responses_adapters_handler.py | 0 ...t_responses_adapters_streaming_iterator.py | 0 .../test_responses_adapters_transformation.py | 0 .../anthropic/test_anthropic_common_utils.py | 0 ...t_anthropic_count_tokens_transformation.py | 0 .../test_anthropic_files_and_batches.py | 0 .../test_anthropic_output_format_filter.py | 0 .../test_anthropic_prompt_cache_prediction.py | 0 .../test_anthropic_reasoning_effort.py | 0 .../anthropic/test_anthropic_schema_filter.py | 0 .../test_anthropic_structured_output.py | 0 .../anthropic/test_azure_ai_cache_pricing.py | 0 .../test_cost_calculation_dict_safety.py | 0 .../llms/anthropic/test_count_tokens_oauth.py | 0 .../anthropic/test_message_sanitization.py | 0 .../llms/azure/batches}/__init__.py | 0 .../llms/azure/batches/test_handler.py | 0 .../llms/azure/chat}/__init__.py | 0 .../chat/test_azure_base_model_routing.py | 0 .../test_azure_chat_gpt_transformation.py | 0 ...test_azure_chat_o_series_transformation.py | 0 .../chat/test_azure_gpt5_transformation.py | 0 .../llms/azure/realtime/test_handler.py | 0 .../llms/azure/test_audio_transcriptions.py | 0 .../llms/azure/test_azure.py | 0 .../llms/azure/test_azure_common_utils.py | 0 .../llms/azure/test_azure_cost_calculation.py | 0 .../llms/azure/test_azure_embedding.py | 0 .../azure/test_azure_exception_mapping.py | 0 .../llms/azure/test_azure_fine_tuning_api.py | 0 .../test_azure_speech_audio_transcription.py | 0 .../llms/azure/videos}/__init__.py | 0 .../videos/test_azure_video_transformation.py | 0 .../llms/azure_ai/claude}/__init__.py | 0 ...e_anthropic_count_tokens_transformation.py | 0 .../claude/test_azure_anthropic_handler.py | 0 ...azure_anthropic_messages_transformation.py | 0 .../test_azure_anthropic_provider_routing.py | 0 .../test_azure_anthropic_transformation.py | 0 .../test_main_azure_anthropic_timeout.py | 0 .../azure_ai/image_generation}/__init__.py | 0 .../test_azure_ai_flux2_image_generation.py | 0 .../test_mai_image_generation.py | 0 .../azure_ai/test_azure_ai_agents_handler.py | 0 .../azure_ai/test_azure_ai_cost_calculator.py | 0 .../llms/azure_ai/test_azure_ai_entra_auth.py | 0 ...azure_ai_foundry_catalog_model_metadata.py | 0 .../test_azure_ai_fw_models_metadata.py | 0 .../test_azure_ai_kimi_k26_metadata.py | 0 .../batches/base_batches_config_test.py | 0 .../llms/base_llm/files}/__init__.py | 0 .../files/test_azure_blob_storage_backend.py | 0 .../files/test_litellm_db_storage_backend.py | 0 .../files/test_storage_backend_factory.py | 0 .../llms/base_llm/responses}/__init__.py | 0 .../base_llm/responses/test_codex_compat.py | 0 .../base_llm/responses/test_transformation.py | 0 .../llms/base_llm/search}/__init__.py | 0 .../search/test_base_search_transformation.py | 0 .../base_llm/test_base_managed_resource.py | 0 .../llms/base_llm/test_base_model_iterator.py | 0 .../test_managed_resource_isolation.py | 0 .../base_llm/test_managed_resources_utils.py | 0 .../llms/bedrock/batches}/__init__.py | 0 .../test_batch_metadata_sanitization.py | 0 .../llms/bedrock/batches/test_handler.py | 0 .../bedrock/batches/test_transformation.py | 2 +- .../chat/test_bedrock_converse_handler.py | 2 +- .../chat/test_converse_transformation.py | 0 .../test_converse_transformation_nova_2.py | 0 .../llms/bedrock/chat/test_invoke_handler.py | 0 .../llms/bedrock/chat/test_mistral_config.py | 0 .../llms/bedrock/chat/test_service_tier.py | 0 .../chat/test_streaming_choice_index.py | 0 .../llms/bedrock/chat/test_writer_palmyra.py | 0 .../test_bedrock_count_tokens_handler.py | 2 +- .../llms/bedrock/embed}/__init__.py | 0 .../test_bedrock_async_invoke_embedding.py | 2 +- .../bedrock/embed/test_bedrock_embedding.py | 2 +- .../llms/bedrock/embed/test_embedding.py | 0 ...est_twelvelabs_marengo_3_transformation.py | 0 .../llms/bedrock/event_loop_probe.py | 0 .../llms/bedrock/messages}/__init__.py | 0 .../invoke_transformations}/__init__.py | 0 .../test_anthropic_claude3_transformation.py | 0 .../llms/bedrock/rerank/transformation.py | 0 .../llms/bedrock/responses}/__init__.py | 0 .../test_bedrock_openai_responses.py | 0 .../llms/bedrock/search}/__init__.py | 0 .../test_agentcore_search_transformation.py | 0 .../bedrock/test_anthropic_beta_support.py | 0 .../llms/bedrock/test_base_aws_llm.py | 2 +- .../llms/bedrock/test_bedrock_common_utils.py | 0 .../llms/bedrock/test_bedrock_ssl_verify.py | 0 .../bedrock/test_claude_platform_provider.py | 0 .../test_converse_context_management.py | 0 ..._cross_region_inference_profile_mapping.py | 0 .../llms/bedrock/test_mantle.py | 0 .../llms/bedrock/test_nova_imported_models.py | 0 .../llms/bedrock/test_request_metadata.py | 0 .../test_web_identity_session_policy.py | 0 ..._bedrock_mantle_messages_transformation.py | 0 ...bedrock_mantle_responses_transformation.py | 0 .../test_bedrock_mantle_transformation.py | 2 +- .../llms/cometapi}/__init__.py | 0 .../llms/cometapi/chat}/__init__.py | 0 .../chat/test_cometapi_chat_transformation.py | 183 ++ .../llms/compactifai}/__init__.py | 0 .../llms/compactifai/test_compactifai.py | 50 - .../llms/custom_httpx}/__init__.py | 0 .../test_aiohttp_cleanup_closed.py | 0 .../llms/custom_httpx/test_aiohttp_handler.py | 0 .../custom_httpx/test_aiohttp_so_keepalive.py | 0 .../custom_httpx/test_aiohttp_transport.py | 0 .../llms/custom_httpx/test_asgi_handler.py | 0 .../custom_httpx/test_async_client_cleanup.py | 0 .../custom_httpx/test_container_handler.py | 0 .../test_credential_leak_prevention.py | 0 .../custom_httpx/test_gemini_session_leak.py | 0 .../llms/custom_httpx/test_http_handler.py | 0 .../custom_httpx/test_llm_http_handler.py | 2 +- .../llms/custom_httpx/test_mock_transport.py | 0 .../llms/dashscope}/__init__.py | 0 .../test_dashscope_chat_transformation.py | 0 .../test_dashscope_cost_calculator.py | 0 ...test_dashscope_embedding_transformation.py | 0 .../test_dashscope_rerank_transformation.py | 0 .../llms/dashscope/test_qwen_brand_aliases.py | 0 .../test_databricks_chat_transformation.py | 75 + .../test_databricks_common_utils.py | 0 .../test_databricks_cost_calculator.py | 0 .../test_databricks_partner_integration.py | 0 .../test_databricks_streaming_utils.py | 0 .../llms/deepgram}/__init__.py | 0 .../deepgram/audio_transcription}/__init__.py | 0 ...gram_audio_transcription_transformation.py | 0 .../deepgram/test_deepgram_common_utils.py | 0 .../test_deepgram_mock_transcription.py | 0 .../llms/deepinfra}/__init__.py | 0 .../test_deepinfra_chat_transformation.py | 0 .../llms/deepinfra/test_deepinfra_rerank.py | 0 .../test_deepinfra_rerank_integration.py | 159 + .../test_deepinfra_rerank_transformation.py | 0 .../llms/edenai}/__init__.py | 0 .../edenai/audio_transcription}/__init__.py | 0 ...enai_audio_transcription_transformation.py | 0 .../llms/edenai/chat}/__init__.py | 0 .../chat/test_edenai_chat_transformation.py | 0 .../llms/edenai/conftest.py | 0 .../llms/edenai}/embedding/__init__.py | 0 .../test_edenai_embedding_transformation.py | 0 .../llms/edenai/image_generation}/__init__.py | 0 ..._edenai_image_generation_transformation.py | 0 .../llms/edenai}/messages/__init__.py | 0 ...denai_anthropic_messages_transformation.py | 0 .../llms/edenai/responses}/__init__.py | 0 .../test_edenai_responses_transformation.py | 0 .../llms/edenai/test_edenai_common_utils.py | 0 .../llms/edenai/text_to_speech}/__init__.py | 0 ...st_edenai_text_to_speech_transformation.py | 0 .../llms/edenai/videos}/__init__.py | 0 .../test_edenai_video_transformation.py | 0 .../chat => unit/llms/fal_ai}/__init__.py | 0 .../llms/fal_ai/chat}/__init__.py | 0 .../chat/test_fal_ai_chat_transformation.py | 0 .../llms/fal_ai/image_edit}/__init__.py | 0 ...t_fal_ai_flux_lora_depth_transformation.py | 0 .../test_fal_ai_image_edit_transformation.py | 0 .../llms/fal_ai/image_generation}/__init__.py | 0 .../test_fal_ai_flux_dev_transformation.py | 0 .../test_fal_ai_gpt_image_2_transformation.py | 0 .../test_fal_ai_nano_banana_transformation.py | 0 .../llms/fal_ai/test_cost_calculator.py | 0 .../llms/fal_ai/videos}/__init__.py | 0 .../test_fal_ai_video_transformation.py | 0 .../llms/featherless_ai}/__init__.py | 0 .../llms/featherless_ai/chat}/__init__.py | 0 .../test_featherless_chat_transformation.py | 0 .../llms/fireworks_ai/completion}/__init__.py | 0 ..._fireworks_ai_completion_transformation.py | 0 ...works_ai_text_completion_transformation.py | 0 .../responses => unit/llms/gdc}/__init__.py | 0 .../llms/gdc/chat}/__init__.py | 0 .../gdc/chat/test_gdc_chat_transformation.py | 0 .../llms/gemini/test_cost_calculator.py | 0 .../llms/gemini/test_gemini_client_setup.py | 0 .../llms/gemini/test_gemini_common_utils.py | 0 ..._gemini_image_generation_transformation.py | 0 .../llms/gemini/test_gemini_tts.py | 0 .../test_github_copilot_authenticator.py | 0 .../test_github_copilot_transformation.py | 0 .../llms/heroku}/__init__.py | 0 .../heroku/test_heroku_chat_transformation.py | 0 .../llms/huggingface/embedding}/__init__.py | 0 .../test_huggingface_embedding_handler.py | 0 .../llms/langflow/test_langflow_a2a.py | 0 .../llms/lemonade}/__init__.py | 0 .../llms/lemonade/test_lemonade.py | 0 .../llms/lm_studio}/__init__.py | 0 .../test_lm_studio_chat_transformation.py | 0 .../mistral/audio_transcription}/__init__.py | 0 ...tral_audio_transcription_transformation.py | 195 ++ .../test_mistral_chat_transformation.py | 0 .../llms/mistral/test_mistral_completion.py | 0 .../llms/modelscope/chat}/__init__.py | 0 .../test_modelscope_chat_transformation.py | 0 .../tencent => unit/llms/nadir}/__init__.py | 0 .../llms/nadir/test_nadir.py | 0 .../chat => unit/llms/nebius}/__init__.py | 0 .../nebius/test_nebius_chat_transformation.py | 0 .../test_nebius_embedding_transformation.py | 0 .../llms/oci/rerank}/__init__.py | 0 .../llms/oci/test_oci_common_utils.py | 0 .../llms/oci/test_oci_coverage_boost.py | 0 .../llms/ollama}/__init__.py | 0 .../ollama/test_ollama_chat_transformation.py | 0 .../test_ollama_completion_transformation.py | 0 .../llms/ollama/test_ollama_embedding.py | 0 .../llms/ollama/test_ollama_model_info.py | 0 .../llms/openai/realtime/README.md | 0 .../llms/openai/realtime}/__init__.py | 0 .../realtime/test_openai_realtime_handler.py | 0 .../realtime/test_transcription_sessions.py | 0 .../llms/openai/responses}/__init__.py | 0 ...test_openai_count_tokens_transformation.py | 0 .../test_openai_responses_data_residency.py | 0 ...test_openai_responses_guardrail_handler.py | 0 ...t_openai_responses_guardrail_tool_merge.py | 0 .../test_openai_responses_transformation.py | 0 .../llms/openai/test_cost_calculation.py | 0 .../llms/openai/test_data_residency.py | 0 .../llms/openai/test_gpt5_transformation.py | 0 .../llms/openai/test_is_model_gpt_5_model.py | 0 .../openai/test_o_series_transformation.py | 0 .../llms/openai/test_openai.py | 0 .../llms/openai/test_openai_common_utils.py | 0 .../llms/openai/test_openai_empty_response.py | 0 .../test_openai_file_content_streaming.py | 0 .../test_openai_image_edit_transformation.py | 0 .../openai/test_openai_workload_identity.py | 0 .../llms/openai/test_organization_costs.py | 0 .../test_use_chat_completions_api_no_leak.py | 0 .../test_openai_transcriptions_handler.py | 0 .../llms/openai_like/responses}/__init__.py | 0 .../responses/test_openai_like_responses.py | 0 .../openai_like/test_abliteration_provider.py | 0 .../openai_like/test_assemblyai_provider.py | 0 .../llms/openai_like/test_charity_engine.py | 0 .../openai_like/test_cognition_provider.py | 0 .../llms/openai_like/test_dynamic_config.py | 3 - .../openai_like/test_empiriolabs_provider.py | 0 .../llms/openai_like/test_json_providers.py | 317 ++ .../openai_like/test_libertai_provider.py | 0 .../llms/openai_like/test_meta_provider.py | 0 .../llms/openai_like/test_model_info.py | 0 .../openai_like/test_pinstripes_provider.py | 25 - .../test_provider_affinity_forwarding.py | 0 .../llms/openai_like/test_scx_ai_provider.py | 0 .../openai_like/test_tensormesh_provider.py | 0 .../unit/llms/openai_like/test_xiaomi_mimo.py | 84 + .../files => unit/llms/ovhcloud}/__init__.py | 0 ...loud_audio_transcription_transformation.py | 58 + .../test_ovhcloud_chat_transformation.py | 250 ++ ...test_ovhcloud_embeddings_transformation.py | 0 .../llms/pass_through}/__init__.py | 0 .../guardrail_translation}/__init__.py | 0 .../llms/perplexity/test_perplexity.py | 0 .../test_perplexity_cost_calculator.py | 0 .../perplexity/test_perplexity_integration.py | 0 .../llms/pg_vector}/__init__.py | 0 .../llms/pg_vector/vector_stores}/__init__.py | 0 .../test_pg_vector_transformation.py | 0 .../llms/reducto}/__init__.py | 0 .../llms/reducto/conftest.py | 0 .../llms/reducto/test_cost.py | 0 .../llms/reducto/test_model_info.py | 0 .../llms/reducto/test_parse_legacy.py | 0 .../llms/reducto/test_parse_v3.py | 0 .../llms/reducto/test_upload.py | 0 .../qwen => unit/llms/sagemaker}/__init__.py | 0 .../sagemaker/test_sagemaker_chat_handler.py | 0 .../test_sagemaker_chat_transformation.py | 0 .../sagemaker/test_sagemaker_common_utils.py | 0 .../test_sagemaker_completion_handler.py | 0 ...est_sagemaker_embedding_role_assumption.py | 0 .../test_sagemaker_embedding_voyage.py | 0 .../test_sagemaker_nova_transformation.py | 0 .../llms/sambanova}/__init__.py | 0 ...ests_sambanova_embedding_transformation.py | 0 .../llms/sap/chat}/__init__.py | 0 .../llms/sap/chat/test_sap_chat_calls.py | 0 .../chat/test_sap_langchain_strict_param.py | 0 .../llms/sap/chat/test_sap_response_format.py | 0 .../llms/sap/chat/test_sap_tool_parameters.py | 0 .../llms/sap/chat/test_sap_transformation.py | 0 .../llms/sap/embed}/__init__.py | 0 .../embed/test_sap_embed_transformation.py | 0 .../llms/sap/embed/test_sap_embedding.py | 0 .../llms/snowflake/chat}/__init__.py | 0 .../test_snowflake_chat_transformation.py | 0 .../llms/snowflake/embedding}/__init__.py | 0 .../embedding/test_snowflake_embedding.py | 0 .../test_snowflake_native_endpoints.py | 2 +- .../soniox/audio_transcription/__init__.py | 0 ...test_soniox_audio_transcription_handler.py | 0 ...niox_audio_transcription_transformation.py | 0 .../llms/test_cache_control_and_reasoning.py | 0 .../llms/test_file_content_block.py | 0 .../llms/test_file_search_responses.py | 0 .../llms/test_lifecycle_fix.py | 0 .../llms/test_polling_url_origin_match.py | 0 .../llms/test_predibase_transformation.py | 0 tests/unit/llms/tinyfish/__init__.py | 0 .../llms/tinyfish/test_tinyfish_search.py | 0 .../test_vercel_ai_gateway.py | 0 .../vertex_ai/audio_transcription/__init__.py | 0 ...x_ai_audio_transcription_transformation.py | 0 ...tex_ai_gemini_transcribe_transformation.py | 0 .../test_vertex_ai_realtime_backend.py | 0 .../test_vertex_ai_realtime_transformation.py | 0 tests/unit/llms/vertex_ai/batches/__init__.py | 0 .../llms/vertex_ai/batches/test_handler.py | 0 .../vertex_ai/batches/test_transformation.py | 0 .../vertex_ai/files/test_transformation.py | 0 tests/unit/llms/vertex_ai/gemini/__init__.py | 0 .../gemini/test_context_circulation.py | 0 .../test_function_call_args_serialization.py | 0 .../test_gemini_image_url_missing_field.py | 0 ...emini_streaming_tool_call_finish_reason.py | 0 .../gemini/test_grounding_requests.py | 0 .../test_thought_signature_in_tool_call_id.py | 0 ...st_tool_call_followed_by_text_assistant.py | 0 .../vertex_ai/gemini/test_transformation.py | 0 .../test_vertex_ai_gemini_transformation.py | 2729 ++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 38 - .../test_vertex_gemini_unbound_local_error.py | 0 .../vertex_ai/image_generation/__init__.py | 0 ...tex_ai_image_generation_cost_calculator.py | 0 ...rtex_ai_image_generation_transformation.py | 637 ++++ tests/unit/llms/vertex_ai/rerank/__init__.py | 0 .../test_vertex_ai_rerank_integration.py | 0 .../test_vertex_ai_rerank_transformation.py | 0 .../test_vertex_ai_rerank_userlabels_e2e.py | 0 .../llms/vertex_ai/test_bge_embedding.py | 0 .../test_bge_response_transformation.py | 0 .../vertex_ai/test_gemini_batch_embeddings.py | 0 .../vertex_ai/test_gemini_empty_properties.py | 0 .../test_gemini_header_forwarding.py | 0 .../llms/vertex_ai/test_http_status_201.py | 0 .../llms/vertex_ai/test_vertex.py | 42 - .../test_vertex_ai_batch_transformation.py | 0 .../vertex_ai/test_vertex_ai_common_utils.py | 0 .../test_vertex_ai_psc_endpoint_support.py | 0 ...x_ai_search_vector_store_transformation.py | 0 .../test_vertex_gemini_gcs_uri_mime.py | 0 .../test_vertex_global_url_support.py | 0 .../vertex_ai/test_vertex_image_generation.py | 0 .../llms/vertex_ai/test_vertex_llm_base.py | 0 .../test_vertex_model_garden_openapi.py | 0 ...test_vertex_passthrough_logging_handler.py | 0 .../anthropic/__init__.py | 0 ..._vertex_ai_anthropic_image_url_handling.py | 0 ...artner_models_anthropic_messages_config.py | 0 ...partner_models_anthropic_transformation.py | 0 .../gemma/__init__.py | 0 .../test_vertex_ai_gemma_global_endpoint.py | 0 .../gpt_oss/__init__.py | 0 .../test_vertex_ai_gpt_oss_transformation.py | 0 .../vertex_ai_partner_models/qwen/__init__.py | 0 .../test_vertex_ai_qwen_global_endpoint.py | 0 .../test_partner_models_credential_reuse.py | 0 .../llms/volcengine/embedding/__init__.py | 0 .../llms/volcengine/test_volcengine.py | 0 tests/unit/llms/wandb/__init__.py | 0 .../wandb/test_wandb_chat_transformation.py | 0 ..._xai_audio_transcription_transformation.py | 0 .../llms/xai/test_xai_chat_transformation.py | 0 .../llms/xai/test_xai_cost_calculator.py | 0 .../llms/xai/test_xai_key_fallback.py | 0 .../llms/xai/test_xai_model_registry.py | 0 .../llms/xai/test_xai_oauth.py | 0 tests/unit/test_unit_shard_missing_paths.py | 1 + 479 files changed, 4789 insertions(+), 5180 deletions(-) delete mode 100644 tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py delete mode 100644 tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py delete mode 100644 tests/test_litellm/llms/gemini/files/__init__.py delete mode 100644 tests/test_litellm/llms/gemini/videos/__init__.py delete mode 100644 tests/test_litellm/llms/manus/__init__.py delete mode 100644 tests/test_litellm/llms/manus/responses/__init__.py delete mode 100644 tests/test_litellm/llms/minimax/__init__.py delete mode 100644 tests/test_litellm/llms/minimax/chat/__init__.py delete mode 100644 tests/test_litellm/llms/minimax/messages/__init__.py delete mode 100644 tests/test_litellm/llms/reducto/__init__.py delete mode 100644 tests/test_litellm/llms/s3_vectors/__init__.py delete mode 100644 tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py delete mode 100644 tests/test_litellm/llms/soniox/__init__.py delete mode 100644 tests/test_litellm/llms/vertex_ai/image_edit/__init__.py delete mode 100644 tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py delete mode 100644 tests/test_litellm/llms/vertex_ai/videos/__init__.py rename tests/{test_litellm/llms/anthropic => unit/expected_fine_tuning_api}/__init__.py (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_cancel_expected_output.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_cancel_raw_response.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_cancel_request.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_create_expected_output.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_create_raw_response.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_create_request.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_list_raw_response.json (100%) rename tests/{test_litellm => unit}/expected_fine_tuning_api/azure_list_request.json (100%) rename tests/{test_litellm/llms/anthropic/batches => unit/llms/aiml}/__init__.py (100%) rename tests/{test_litellm/llms/anthropic/experimental_pass_through/context_management => unit/llms/aiml/image_generation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/aiml/image_generation/test_aiml_image_generation_transformation.py (100%) rename tests/{test_litellm/llms/anthropic/experimental_pass_through/responses_adapters => unit/llms/anthropic/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/chat/conftest.py (100%) rename tests/{test_litellm/llms/anthropic/files => unit/llms/anthropic/chat/guardrail_translation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/chat/test_anthropic_chat_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/chat/test_anthropic_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/chat/test_code_interpreter_results_extraction.py (100%) rename tests/{test_litellm/llms/azure/batches => unit/llms/anthropic/experimental_pass_through/adapters}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py (100%) rename tests/{test_litellm/llms/azure/vector_stores => unit/llms/anthropic/experimental_pass_through/context_management}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/context_management/test_compact.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py (100%) rename tests/{test_litellm/llms/base_llm => unit/llms/anthropic/experimental_pass_through/messages}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_response_cache.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py (100%) rename tests/{test_litellm/llms/base_llm/batches => unit/llms/anthropic/experimental_pass_through/responses_adapters}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_count_tokens_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_files_and_batches.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_output_format_filter.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_prompt_cache_prediction.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_reasoning_effort.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_schema_filter.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_anthropic_structured_output.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_azure_ai_cache_pricing.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_cost_calculation_dict_safety.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_count_tokens_oauth.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/test_message_sanitization.py (100%) rename tests/{test_litellm/llms/base_llm/files => unit/llms/azure/batches}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/azure/batches/test_handler.py (100%) rename tests/{test_litellm/llms/base_llm/realtime => unit/llms/azure/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/azure/chat/test_azure_base_model_routing.py (100%) rename tests/{test_litellm => unit}/llms/azure/chat/test_azure_chat_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/chat/test_azure_chat_o_series_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/chat/test_azure_gpt5_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/realtime/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_audio_transcriptions.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_cost_calculation.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_embedding.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_exception_mapping.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_fine_tuning_api.py (100%) rename tests/{test_litellm => unit}/llms/azure/test_azure_speech_audio_transcription.py (100%) rename tests/{test_litellm/llms/bedrock => unit/llms/azure/videos}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/azure/videos/test_azure_video_transformation.py (100%) rename tests/{test_litellm/llms/bedrock/batches => unit/llms/azure_ai/claude}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_azure_anthropic_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_azure_anthropic_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py (100%) rename tests/{test_litellm/llms/bedrock/chat/agentcore => unit/llms/azure_ai/image_generation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_generation/test_mai_image_generation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_agents_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_entra_auth.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_fw_models_metadata.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/batches/base_batches_config_test.py (100%) rename tests/{test_litellm/llms/bedrock/passthrough/guardrail_translation => unit/llms/base_llm/files}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/files/test_azure_blob_storage_backend.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/files/test_litellm_db_storage_backend.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/files/test_storage_backend_factory.py (100%) rename tests/{test_litellm/llms/black_forest_labs => unit/llms/base_llm/responses}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/responses/test_codex_compat.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/responses/test_transformation.py (100%) rename tests/{test_litellm/llms/black_forest_labs/image_edit => unit/llms/base_llm/search}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/search/test_base_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/test_base_managed_resource.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/test_base_model_iterator.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/test_managed_resource_isolation.py (100%) rename tests/{test_litellm => unit}/llms/base_llm/test_managed_resources_utils.py (100%) rename tests/{test_litellm/llms/black_forest_labs/image_generation => unit/llms/bedrock/batches}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/batches/test_batch_metadata_sanitization.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/batches/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/batches/test_transformation.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_bedrock_converse_handler.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_converse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_converse_transformation_nova_2.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_invoke_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_mistral_config.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_service_tier.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_streaming_choice_index.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/chat/test_writer_palmyra.py (100%) rename tests/{test_litellm/llms/cerebras => unit/llms/bedrock/embed}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/embed/test_bedrock_embedding.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/embed/test_embedding.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/event_loop_probe.py (100%) rename tests/{test_litellm/llms/chatgpt => unit/llms/bedrock/messages}/__init__.py (100%) rename tests/{test_litellm/llms/chatgpt/chat => unit/llms/bedrock/messages/invoke_transformations}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/rerank/transformation.py (100%) rename tests/{test_litellm/llms/crusoe => unit/llms/bedrock/responses}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/responses/test_bedrock_openai_responses.py (100%) rename tests/{test_litellm/llms/databricks/chat => unit/llms/bedrock/search}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/search/test_agentcore_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_anthropic_beta_support.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_base_aws_llm.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/test_bedrock_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_bedrock_ssl_verify.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_claude_platform_provider.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_converse_context_management.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_cross_region_inference_profile_mapping.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_mantle.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_nova_imported_models.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_request_metadata.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/test_web_identity_session_policy.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/test_bedrock_mantle_transformation.py (99%) rename tests/{test_litellm/llms/databricks/responses => unit/llms/cometapi}/__init__.py (100%) rename tests/{test_litellm/llms/deepseek => unit/llms/cometapi/chat}/__init__.py (100%) create mode 100644 tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py rename tests/{test_litellm/llms/deepseek/chat => unit/llms/compactifai}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/compactifai/test_compactifai.py (84%) rename tests/{test_litellm/llms/deepseek/messages => unit/llms/custom_httpx}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_aiohttp_cleanup_closed.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_aiohttp_handler.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_aiohttp_so_keepalive.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_aiohttp_transport.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_asgi_handler.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_async_client_cleanup.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_container_handler.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_credential_leak_prevention.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_gemini_session_leak.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_http_handler.py (100%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_llm_http_handler.py (99%) rename tests/{test_litellm => unit}/llms/custom_httpx/test_mock_transport.py (100%) rename tests/{test_litellm/llms/gemini => unit/llms/dashscope}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/dashscope/test_dashscope_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/dashscope/test_dashscope_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/dashscope/test_dashscope_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/dashscope/test_dashscope_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/dashscope/test_qwen_brand_aliases.py (100%) rename tests/{test_litellm => unit}/llms/databricks/test_databricks_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/databricks/test_databricks_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/databricks/test_databricks_partner_integration.py (100%) rename tests/{test_litellm => unit}/llms/databricks/test_databricks_streaming_utils.py (100%) rename tests/{test_litellm/llms/gemini/audio_transcription => unit/llms/deepgram}/__init__.py (100%) rename tests/{test_litellm/llms/gemini/google_genai => unit/llms/deepgram/audio_transcription}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepgram/test_deepgram_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/deepgram/test_deepgram_mock_transcription.py (100%) rename tests/{test_litellm/llms/gemini/google_genai/guardrail_translation => unit/llms/deepinfra}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/deepinfra/test_deepinfra_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/deepinfra/test_deepinfra_rerank.py (100%) create mode 100644 tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py rename tests/{test_litellm => unit}/llms/deepinfra/test_deepinfra_rerank_transformation.py (100%) rename tests/{test_litellm/llms/gemini/image_edit => unit/llms/edenai}/__init__.py (100%) rename tests/{test_litellm/llms/gemini/realtime => unit/llms/edenai/audio_transcription}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py (100%) rename tests/{test_litellm/llms/gigachat => unit/llms/edenai/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/chat/test_edenai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/edenai/conftest.py (100%) rename tests/{test_litellm/llms/gigachat => unit/llms/edenai}/embedding/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/embedding/test_edenai_embedding_transformation.py (100%) rename tests/{test_litellm/llms/gigachat/passthrough => unit/llms/edenai/image_generation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/image_generation/test_edenai_image_generation_transformation.py (100%) rename tests/{test_litellm/llms/github_copilot => unit/llms/edenai}/messages/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py (100%) rename tests/{test_litellm/llms/gradient_ai => unit/llms/edenai/responses}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/responses/test_edenai_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/edenai/test_edenai_common_utils.py (100%) rename tests/{test_litellm/llms/gradient_ai/chat => unit/llms/edenai/text_to_speech}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py (100%) rename tests/{test_litellm/llms/groq => unit/llms/edenai/videos}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/edenai/videos/test_edenai_video_transformation.py (100%) rename tests/{test_litellm/llms/groq/chat => unit/llms/fal_ai}/__init__.py (100%) rename tests/{test_litellm/llms/huggingface => unit/llms/fal_ai/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/chat/test_fal_ai_chat_transformation.py (100%) rename tests/{test_litellm/llms/inception => unit/llms/fal_ai/image_edit}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py (100%) rename tests/{test_litellm/llms/mistral/batches => unit/llms/fal_ai/image_generation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/test_cost_calculator.py (100%) rename tests/{test_litellm/llms/mistral/files => unit/llms/fal_ai/videos}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/fal_ai/videos/test_fal_ai_video_transformation.py (100%) rename tests/{test_litellm/llms/nvidia_riva => unit/llms/featherless_ai}/__init__.py (100%) rename tests/{test_litellm/llms/oci/rerank => unit/llms/featherless_ai/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/featherless_ai/chat/test_featherless_chat_transformation.py (100%) rename tests/{test_litellm/llms/ocr => unit/llms/fireworks_ai/completion}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py (100%) rename tests/{test_litellm/llms/openai_like/responses => unit/llms/gdc}/__init__.py (100%) rename tests/{test_litellm/llms/parallel_ai => unit/llms/gdc/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/gdc/chat/test_gdc_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/test_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/gemini/test_gemini_client_setup.py (100%) rename tests/{test_litellm => unit}/llms/gemini/test_gemini_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/gemini/test_gemini_image_generation_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/test_gemini_tts.py (100%) rename tests/{test_litellm => unit}/llms/github_copilot/test_github_copilot_authenticator.py (100%) rename tests/{test_litellm => unit}/llms/github_copilot/test_github_copilot_transformation.py (100%) rename tests/{test_litellm/llms/pass_through => unit/llms/heroku}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/heroku/test_heroku_chat_transformation.py (100%) rename tests/{test_litellm/llms/pass_through/guardrail_translation => unit/llms/huggingface/embedding}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/huggingface/embedding/test_huggingface_embedding_handler.py (100%) rename tests/{test_litellm => unit}/llms/langflow/test_langflow_a2a.py (100%) rename tests/{test_litellm/llms/perplexity => unit/llms/lemonade}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/lemonade/test_lemonade.py (100%) rename tests/{test_litellm/llms/perplexity/embedding => unit/llms/lm_studio}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/lm_studio/test_lm_studio_chat_transformation.py (100%) rename tests/{test_litellm/llms/stability => unit/llms/mistral/audio_transcription}/__init__.py (100%) create mode 100644 tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py rename tests/{test_litellm => unit}/llms/mistral/test_mistral_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/mistral/test_mistral_completion.py (100%) rename tests/{test_litellm/llms/stability/image_generation => unit/llms/modelscope/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/modelscope/chat/test_modelscope_chat_transformation.py (100%) rename tests/{test_litellm/llms/tencent => unit/llms/nadir}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/nadir/test_nadir.py (100%) rename tests/{test_litellm/llms/tencent/chat => unit/llms/nebius}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/nebius/test_nebius_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/nebius/test_nebius_embedding_transformation.py (100%) rename tests/{test_litellm/llms/tencent/messages => unit/llms/oci/rerank}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/oci/test_oci_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/oci/test_oci_coverage_boost.py (100%) rename tests/{test_litellm/llms/vercel_ai_gateway/embedding => unit/llms/ollama}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/ollama/test_ollama_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/ollama/test_ollama_completion_transformation.py (100%) rename tests/{test_litellm => unit}/llms/ollama/test_ollama_embedding.py (100%) rename tests/{test_litellm => unit}/llms/ollama/test_ollama_model_info.py (100%) rename tests/{test_litellm => unit}/llms/openai/realtime/README.md (100%) rename tests/{test_litellm/llms/vertex_ai/agent_engine => unit/llms/openai/realtime}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/openai/realtime/test_openai_realtime_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/realtime/test_transcription_sessions.py (100%) rename tests/{test_litellm/llms/vertex_ai/audio_transcription => unit/llms/openai/responses}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/openai/responses/test_openai_count_tokens_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/responses/test_openai_responses_data_residency.py (100%) rename tests/{test_litellm => unit}/llms/openai/responses/test_openai_responses_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py (100%) rename tests/{test_litellm => unit}/llms/openai/responses/test_openai_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_cost_calculation.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_data_residency.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_gpt5_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_is_model_gpt_5_model.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_o_series_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai_empty_response.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai_file_content_streaming.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_openai_workload_identity.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_organization_costs.py (100%) rename tests/{test_litellm => unit}/llms/openai/test_use_chat_completions_api_no_leak.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_openai_transcriptions_handler.py (100%) rename tests/{test_litellm/llms/vertex_ai/batches => unit/llms/openai_like/responses}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/responses/test_openai_like_responses.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_abliteration_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_assemblyai_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_charity_engine.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_cognition_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_dynamic_config.py (96%) rename tests/{test_litellm => unit}/llms/openai_like/test_empiriolabs_provider.py (100%) create mode 100644 tests/unit/llms/openai_like/test_json_providers.py rename tests/{test_litellm => unit}/llms/openai_like/test_libertai_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_meta_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_model_info.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_pinstripes_provider.py (68%) rename tests/{test_litellm => unit}/llms/openai_like/test_provider_affinity_forwarding.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_scx_ai_provider.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/test_tensormesh_provider.py (100%) create mode 100644 tests/unit/llms/openai_like/test_xiaomi_mimo.py rename tests/{test_litellm/llms/vertex_ai/files => unit/llms/ovhcloud}/__init__.py (100%) create mode 100644 tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py create mode 100644 tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py rename tests/{test_litellm => unit}/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py (100%) rename tests/{test_litellm/llms/vertex_ai/gemini_embeddings => unit/llms/pass_through}/__init__.py (100%) rename tests/{test_litellm/llms/vertex_ai/text_to_speech => unit/llms/pass_through/guardrail_translation}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/test_perplexity.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/test_perplexity_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/test_perplexity_integration.py (100%) rename tests/{test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens => unit/llms/pg_vector}/__init__.py (100%) rename tests/{test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma => unit/llms/pg_vector/vector_stores}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/pg_vector/vector_stores/test_pg_vector_transformation.py (100%) rename tests/{test_litellm/llms/azure/realtime => unit/llms/reducto}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/reducto/conftest.py (100%) rename tests/{test_litellm => unit}/llms/reducto/test_cost.py (100%) rename tests/{test_litellm => unit}/llms/reducto/test_model_info.py (100%) rename tests/{test_litellm => unit}/llms/reducto/test_parse_legacy.py (100%) rename tests/{test_litellm => unit}/llms/reducto/test_parse_v3.py (100%) rename tests/{test_litellm => unit}/llms/reducto/test_upload.py (100%) rename tests/{test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen => unit/llms/sagemaker}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_chat_handler.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_completion_handler.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_embedding_role_assumption.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_embedding_voyage.py (100%) rename tests/{test_litellm => unit}/llms/sagemaker/test_sagemaker_nova_transformation.py (100%) rename tests/{test_litellm/llms/voyage/rerank => unit/llms/sambanova}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/sambanova/tests_sambanova_embedding_transformation.py (100%) rename tests/{test_litellm/llms/watsonx => unit/llms/sap/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/sap/chat/test_sap_chat_calls.py (100%) rename tests/{test_litellm => unit}/llms/sap/chat/test_sap_langchain_strict_param.py (100%) rename tests/{test_litellm => unit}/llms/sap/chat/test_sap_response_format.py (100%) rename tests/{test_litellm => unit}/llms/sap/chat/test_sap_tool_parameters.py (100%) rename tests/{test_litellm => unit}/llms/sap/chat/test_sap_transformation.py (100%) rename tests/{test_litellm/llms/watsonx/audio_transcription => unit/llms/sap/embed}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/sap/embed/test_sap_embed_transformation.py (100%) rename tests/{test_litellm => unit}/llms/sap/embed/test_sap_embedding.py (100%) rename tests/{test_litellm/llms/watsonx/rerank => unit/llms/snowflake/chat}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/snowflake/chat/test_snowflake_chat_transformation.py (100%) rename tests/{test_litellm/llms/you_com => unit/llms/snowflake/embedding}/__init__.py (100%) rename tests/{test_litellm => unit}/llms/snowflake/embedding/test_snowflake_embedding.py (100%) rename tests/{test_litellm => unit}/llms/soniox/audio_transcription/__init__.py (100%) rename tests/{test_litellm => unit}/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py (100%) rename tests/{test_litellm => unit}/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/test_cache_control_and_reasoning.py (100%) rename tests/{test_litellm => unit}/llms/test_file_content_block.py (100%) rename tests/{test_litellm => unit}/llms/test_file_search_responses.py (100%) rename tests/{test_litellm => unit}/llms/test_lifecycle_fix.py (100%) rename tests/{test_litellm => unit}/llms/test_polling_url_origin_match.py (100%) rename tests/{test_litellm => unit}/llms/test_predibase_transformation.py (100%) create mode 100644 tests/unit/llms/tinyfish/__init__.py rename tests/{test_litellm => unit}/llms/tinyfish/test_tinyfish_search.py (100%) rename tests/{test_litellm => unit}/llms/vercel_ai_gateway/test_vercel_ai_gateway.py (100%) create mode 100644 tests/unit/llms/vertex_ai/audio_transcription/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/batches/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/batches/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/batches/test_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/files/test_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/gemini/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_context_circulation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_function_call_args_serialization.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_grounding_requests.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py (99%) rename tests/{test_litellm => unit}/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py (100%) create mode 100644 tests/unit/llms/vertex_ai/image_generation/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py (100%) create mode 100644 tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py create mode 100644 tests/unit/llms/vertex_ai/rerank/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_bge_embedding.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_bge_response_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_gemini_batch_embeddings.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_gemini_empty_properties.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_gemini_header_forwarding.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_http_status_201.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex.py (97%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_ai_batch_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_ai_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_global_url_support.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_image_generation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_llm_base.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_model_garden_openapi.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/test_vertex_passthrough_logging_handler.py (100%) create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py (100%) create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py (100%) rename tests/{test_litellm => unit}/llms/volcengine/embedding/__init__.py (100%) rename tests/{test_litellm => unit}/llms/volcengine/test_volcengine.py (100%) create mode 100644 tests/unit/llms/wandb/__init__.py rename tests/{test_litellm => unit}/llms/wandb/test_wandb_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_key_fallback.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_model_registry.py (100%) rename tests/{test_litellm => unit}/llms/xai/test_xai_oauth.py (100%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 5ce8b6c84ba..e9e5dd3d66b 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,8 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + llm-other-providers + llm-vertex-ai mcp-integration misc proxy-db-auth-checks @@ -50,6 +52,8 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + llm-other-providers) find tests/unit/llms -name 'test_*.py' -not -path 'tests/unit/llms/vertex_ai/*' ;; + llm-vertex-ai) echo tests/unit/llms/vertex_ai ;; mcp-integration) echo tests/unit/experimental_mcp_client echo tests/unit/proxy/_experimental/mcp_server diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 10ee19f146a..994d67da64d 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -354,6 +354,21 @@ workflows: - proxy-db-endpoints-and-responses base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-llm-vertex-ai + flag: llm-vertex-ai + shards: 2 + workers: 1 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-llm-other-providers + flag: llm-other-providers + shards: 3 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-misc flag: misc diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json index 8ed7b917460..a563424c230 100644 --- a/.github/merge-smoke-tests.json +++ b/.github/merge-smoke-tests.json @@ -1,8 +1,8 @@ { "cases": { - "CHAT-JSON": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport", - "CHAT-TEXT-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport", - "CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", + "CHAT-JSON": "tests/unit/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport", + "CHAT-TEXT-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport", + "CHAT-TOOL-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", "MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key", "MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]", "COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 91b54f4ee70..2fa05879350 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -89,6 +89,7 @@ jobs: - shard: Vertex AI artifact-name: llm-vertex-ai test-path: "tests/test_litellm/llms/vertex_ai" + unit-flag: llm-vertex-ai workers: 1 reruns: 2 timeout-minutes: 20 @@ -97,6 +98,7 @@ jobs: - shard: All Other Providers artifact-name: llm-other-providers test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai" + unit-flag: llm-other-providers workers: 2 reruns: 2 timeout-minutes: 20 diff --git a/Makefile b/Makefile index 62e6ae53275..e86047b1987 100644 --- a/Makefile +++ b/Makefile @@ -314,7 +314,7 @@ test-unit: install-test-deps # Matrix test targets (matching CI workflow groups) test-unit-llms: install-test-deps - $(UV_RUN) pytest tests/test_litellm/llms --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/llms --tb=short -vv -n 4 --durations=20 test-unit-proxy-guardrails: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/management_helpers --tb=short -vv -n 4 --durations=20 diff --git a/tests/llm_translation/test_bedrock_gpt_oss.py b/tests/llm_translation/test_bedrock_gpt_oss.py index 4af81ee81f7..b264c16601f 100644 --- a/tests/llm_translation/test_bedrock_gpt_oss.py +++ b/tests/llm_translation/test_bedrock_gpt_oss.py @@ -22,7 +22,7 @@ class TestBedrockGPTOSS(BaseLLMChatTest): """Bedrock GPT-OSS intermittently emits truncated toolUse.input deltas on the live endpoint, which makes the inherited live integration test flaky. The accumulation side is covered deterministically by - tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py::test_transform_tool_calls_index; + tests/unit/llms/bedrock/chat/test_invoke_handler.py::test_transform_tool_calls_index; the GPT-OSS-specific request-body transformation is covered by test_function_calling_request_body_gpt_oss below. """ diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 3a5e2209f1e..2d79f8a6af6 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -324,7 +324,7 @@ def test_parallel_function_call_anthropic_error_msg(model, messages): Anthropic (and Bedrock Invoke via ``AnthropicConfig.transform_request``) inject a dummy tool so CLIs work with ``modify_params`` left off. Bedrock Converse's no-raise behavior is covered offline in - ``tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py`` + ``tests/unit/llms/bedrock/chat/test_converse_transformation.py`` (see #24158, #27138), which needs no live credentials. """ # Force modify_params off as a clean baseline: it exercises the Anthropic diff --git a/tests/local_testing/test_handler_gc_does_not_close_client.py b/tests/local_testing/test_handler_gc_does_not_close_client.py index 1a6ab1b1827..63c5694dd89 100644 --- a/tests/local_testing/test_handler_gc_does_not_close_client.py +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -17,7 +17,7 @@ body can still arrive, released once the caller is done with the response. Nothing here re-tests the shapes ``_handler_may_close_client`` covers -- a borrowed ``handler.client``, a caller-supplied client, an evicted-but-held -client. Those are pinned in ``tests/test_litellm/llms/custom_httpx/ +client. Those are pinned in ``tests/unit/llms/custom_httpx/ test_http_handler.py``. What is uncovered there is the in-flight response, so no test here may keep the client in a local: that inflates the very refcount under test, and the test then passes on a broken handler. They hold weak references diff --git a/tests/local_testing/test_sagemaker_nova_integration.py b/tests/local_testing/test_sagemaker_nova_integration.py index beeb1fa2db3..95f28fe9892 100644 --- a/tests/local_testing/test_sagemaker_nova_integration.py +++ b/tests/local_testing/test_sagemaker_nova_integration.py @@ -4,7 +4,7 @@ Integration tests for SageMaker Nova provider. These tests require a live SageMaker Nova endpoint and AWS credentials. They are skipped by default — run manually with: - pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py -v --no-header -rN + pytest tests/local_testing/test_sagemaker_nova_integration.py -v --no-header -rN Prerequisites: export AWS_PROFILE= # or set AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY @@ -251,7 +251,7 @@ class TestSagemakerNova2LiteIntegration: Run with: export SAGEMAKER_NOVA2_LITE_ENDPOINT= - pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v + pytest tests/local_testing/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v """ def test_should_accept_reasoning_effort_low(self): diff --git a/tests/search_tests/test_bing_grounding_search.py b/tests/search_tests/test_bing_grounding_search.py index 3d1737477a1..f532158e462 100644 --- a/tests/search_tests/test_bing_grounding_search.py +++ b/tests/search_tests/test_bing_grounding_search.py @@ -85,7 +85,7 @@ class TestBingGroundingSearch(BaseSearchTest): class TestBingGroundingSearchTransformation: """ Full-stack tests through `litellm.search` / `litellm.asearch` with the HTTP layer mocked. - Transformation details are unit-tested in tests/test_litellm/llms/azure/search/. + Transformation details are unit-tested in tests/unit/llms/azure/search/. """ @pytest.fixture(autouse=True) diff --git a/tests/search_tests/test_nimble_search.py b/tests/search_tests/test_nimble_search.py index df432f8ae84..3426fc712f4 100644 --- a/tests/search_tests/test_nimble_search.py +++ b/tests/search_tests/test_nimble_search.py @@ -58,7 +58,7 @@ class TestNimbleSearch(BaseSearchTest): class TestNimbleSearchTransformation: """ Full-stack tests through `litellm.search` / `litellm.asearch` with the HTTP layer mocked. - Transformation details are unit-tested in tests/test_litellm/llms/nimble/search/. + Transformation details are unit-tested in tests/unit/llms/nimble/search/. """ @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/integrations/test_helicone.py b/tests/test_litellm/integrations/test_helicone.py index 64960de050a..99cb1380dd7 100644 --- a/tests/test_litellm/integrations/test_helicone.py +++ b/tests/test_litellm/integrations/test_helicone.py @@ -13,7 +13,7 @@ def _claude_mapping(messages, response_obj): def test_claude_mapping_serializes_custom_tool_calls(monkeypatch): """ Stub the anthropic module unconditionally: the SDK may be absent (it lives in the - proxy-runtime extra), and the tests/test_litellm/llms/anthropic test package can + proxy-runtime extra), and the tests/unit/llms/anthropic test package can shadow it on sys.path, so an import probe proves nothing about the real SDK. """ stub = types.ModuleType("anthropic") diff --git a/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py index 7a69b676667..f692259db2e 100644 --- a/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py +++ b/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py @@ -9,171 +9,6 @@ import os import pytest -from litellm.llms.cometapi.chat.transformation import ( - CometAPIChatCompletionStreamingHandler, - CometAPIConfig, -) -from litellm.llms.cometapi.common_utils import CometAPIException - - -class TestCometAPIChatCompletionStreamingHandler: - def test_chunk_parser_successful(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test input chunk - chunk = { - "id": "test_id", - "created": 1234567890, - "model": "gpt-3.5-turbo", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - "choices": [ - {"delta": {"content": "test content", "reasoning": "test reasoning"}} - ], - } - - # Parse chunk - result = handler.chunk_parser(chunk) - - # Verify response - assert result.id == "test_id" - assert result.object == "chat.completion.chunk" - assert result.created == 1234567890 - assert result.model == "gpt-3.5-turbo" - assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] - assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] - assert result.usage.total_tokens == chunk["usage"]["total_tokens"] - assert len(result.choices) == 1 - assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" - - def test_chunk_parser_error_response(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test error chunk - error_chunk = { - "error": { - "message": "test error", - "code": 400, - } - } - - # Verify error handling - with pytest.raises(CometAPIException) as exc_info: - handler.chunk_parser(error_chunk) - - assert "CometAPI Error: test error" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - def test_chunk_parser_key_error(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test invalid chunk missing required fields - invalid_chunk = {"incomplete": "data"} - - # Verify KeyError handling - with pytest.raises(CometAPIException) as exc_info: - handler.chunk_parser(invalid_chunk) - - assert "KeyError" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - -class TestCometAPIConfig: - def test_transform_request_basic(self): - """Test basic request transformation""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["model"] == "cometapi/gpt-3.5-turbo" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_transform_request_with_extra_body(self): - """Test request transformation with extra_body parameters""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-4", - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={"extra_body": {"custom_param": "custom_value"}}, - litellm_params={}, - headers={}, - ) - - # Validate that extra_body parameters are merged into the request - assert transformed_request["custom_param"] == "custom_value" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_cache_control_flag_removal(self): - """Test cache control flag removal from messages""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-3.5-turbo", - messages=[ - { - "role": "user", - "content": "Hello, world!", - "cache_control": {"type": "ephemeral"}, - } - ], - optional_params={}, - litellm_params={}, - headers={}, - ) - - # CometAPI should remove cache_control flags by default - assert transformed_request["messages"][0].get("cache_control") is None - - def test_map_openai_params(self): - """Test OpenAI parameter mapping""" - config = CometAPIConfig() - - non_default_params = { - "temperature": 0.7, - "max_tokens": 100, - "top_p": 0.9, - } - - mapped_params = config.map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model="cometapi/gpt-3.5-turbo", - drop_params=False, - ) - - assert mapped_params["temperature"] == 0.7 - assert mapped_params["max_tokens"] == 100 - assert mapped_params["top_p"] == 0.9 - - def test_get_error_class(self): - """Test error class creation""" - config = CometAPIConfig() - - error = config.get_error_class( - error_message="Test error", - status_code=400, - headers={"Content-Type": "application/json"}, - ) - - assert isinstance(error, CometAPIException) - assert error.message == "Test error" - assert error.status_code == 400 # Integration test example (requires real API key) diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py deleted file mode 100644 index a3391a2c585..00000000000 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ /dev/null @@ -1,79 +0,0 @@ -import json -from typing import Final - -import httpx -import respx - -import litellm - - -def test_completion_merges_leading_system_and_developer_messages_for_chat_template_models( - respx_mock: respx.MockRouter, -): - upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "my-custom-model", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, - }, - ) - ) - - response: Final = litellm.completion( - model="databricks/my-custom-model", - messages=[ - {"role": "system", "content": "You are terse."}, - {"role": "developer", "content": "Skills: none."}, - {"role": "user", "content": "Hello"}, - ], - api_base="https://example.databricks.test/serving-endpoints", - api_key="fake-databricks-api-key", - num_retries=0, - ) - - assert upstream.call_count == 1 - request_body: Final = json.loads(upstream.calls[0].request.read()) - assert request_body["messages"] == [ - {"role": "system", "content": "You are terse.\n\nSkills: none."}, - {"role": "user", "content": "Hello"}, - ] - assert response.choices[0].message.content == "Answer" - - -def test_completion_merges_system_messages_when_one_has_empty_content(respx_mock: respx.MockRouter): - upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "my-custom-model", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, - }, - ) - ) - - litellm.completion( - model="databricks/my-custom-model", - messages=[ - {"role": "system", "content": "You are terse."}, - {"role": "system", "content": ""}, - {"role": "user", "content": "Hello"}, - ], - api_base="https://example.databricks.test/serving-endpoints", - api_key="fake-databricks-api-key", - num_retries=0, - ) - - request_body: Final = json.loads(upstream.calls[0].request.read()) - assert request_body["messages"] == [ - {"role": "system", "content": "You are terse."}, - {"role": "user", "content": "Hello"}, - ] diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py deleted file mode 100644 index 5b013681864..00000000000 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py +++ /dev/null @@ -1,433 +0,0 @@ -""" -Integration tests for DeepInfra rerank functionality. -Tests the full rerank flow following the repository patterns. -""" - -import asyncio -import json -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -import litellm - - -def assert_response_shape(response, custom_llm_provider): - """Helper function to validate response structure specific to DeepInfra.""" - assert hasattr(response, "id") - assert hasattr(response, "results") - assert hasattr(response, "meta") - assert isinstance(response.results, list) - - for result in response.results: - assert "index" in result - assert "relevance_score" in result - assert isinstance(result["index"], int) - assert isinstance(result["relevance_score"], (int, float)) - - # Check meta structure - assert "tokens" in response.meta - assert "billed_units" in response.meta - assert "input_tokens" in response.meta["tokens"] - assert "total_tokens" in response.meta["billed_units"] - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_basic_rerank_deepinfra(mock_sync_post, mock_async_post, sync_mode): - """Test basic DeepInfra rerank functionality.""" - # Mock response data that matches DeepInfra API format - mock_response_data = { - "scores": [0.9, 0.1], - "input_tokens": 25, - "request_id": "deepinfra-request-123", - "inference_status": { - "status": "success", - "runtime_ms": 150, - "cost": 0.0001, - "tokens_generated": 0, - "tokens_input": 25, - }, - } - - def return_val(): - return mock_response_data - - api_key = "test_deepinfra_api_key" - api_base = "https://api.deepinfra.com" - - if sync_mode: - # Create mock response object for sync - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_sync_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - top_n=2, - custom_llm_provider="deepinfra", - api_key=api_key, - api_base=api_base, - ) - mock_sync_post.assert_called_once() - else: - # Create mock response object for async - mock_response = AsyncMock() - - def return_val(): - return mock_response_data - - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_async_post.return_value = mock_response - - response = asyncio.run( - litellm.arerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - top_n=2, - custom_llm_provider="deepinfra", - api_key=api_key, - api_base=api_base, - ) - ) - mock_async_post.assert_called_once() - - # Verify response structure - assert response.id == "deepinfra-request-123" - assert response.results is not None - assert len(response.results) == 2 - assert response.results[0]["index"] == 0 - assert response.results[0]["relevance_score"] == 0.9 - assert response.results[1]["index"] == 1 - assert response.results[1]["relevance_score"] == 0.1 - - # Verify metadata - assert response.meta["tokens"]["input_tokens"] == 25 - assert response.meta["billed_units"]["total_tokens"] == 25 - - # Verify hidden params specific to DeepInfra - assert response._hidden_params["status"] == "success" - assert response._hidden_params["runtime_ms"] == 150 - assert response._hidden_params["cost"] == 0.0001 - # Note: The model name is processed and the 'deepinfra/' prefix is removed - assert response._hidden_params["model"] == "Qwen/Qwen3-Reranker-0.6B" - - assert_response_shape(response, custom_llm_provider="deepinfra") - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_queries_param( - mock_sync_post, mock_async_post, sync_mode -): - """Test DeepInfra rerank with multiple queries parameter.""" - mock_response_data = { - "scores": [0.8, 0.6, 0.2], - "input_tokens": 35, - "request_id": "deepinfra-multi-query-123", - "inference_status": {"status": "success", "runtime_ms": 200}, - } - - def return_val(): - return mock_response_data - - if sync_mode: - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_sync_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-4B", - query="hello", - documents=["hello", "world", "test"], - queries=["hello", "hi there"], # DeepInfra specific param - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - mock_sync_post.assert_called_once() - # Verify that queries parameter was passed in request - call_data = json.loads(mock_sync_post.call_args.kwargs["data"]) - assert "queries" in call_data - assert call_data["queries"] == ["hello", "hi there"] - else: - mock_response = AsyncMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_async_post.return_value = mock_response - - response = asyncio.run( - litellm.arerank( - model="deepinfra/Qwen/Qwen3-Reranker-4B", - query="hello", - documents=["hello", "world", "test"], - queries=["hello", "hi there"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - ) - - mock_async_post.assert_called_once() - call_data = json.loads(mock_async_post.call_args.kwargs["data"]) - assert "queries" in call_data - assert call_data["queries"] == ["hello", "hi there"] - - assert response.results is not None - assert len(response.results) == 3 - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_service_tier(mock_post): - """Test DeepInfra rerank with service_tier parameter.""" - mock_response_data = { - "scores": [0.95, 0.75], - "input_tokens": 30, - "request_id": "deepinfra-premium-123", - } - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-8B", - query="premium search", - documents=["doc1", "doc2"], - service_tier="premium", # DeepInfra specific param - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - mock_post.assert_called_once() - - # Verify URL - call_url = mock_post.call_args.kwargs["url"] - assert "api.deepinfra.com/inference/Qwen/Qwen3-Reranker-8B" in call_url - - # Verify request contains service_tier - call_data = json.loads(mock_post.call_args.kwargs["data"]) - assert call_data["service_tier"] == "premium" - - assert response.results is not None - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_env_vars(mock_post, monkeypatch): - """Test DeepInfra rerank with environment variable configuration.""" - monkeypatch.setenv("DEEPINFRA_API_KEY", "env_test_key") - monkeypatch.setenv("DEEPINFRA_API_BASE", "https://custom-deepinfra.com") - - mock_response_data = { - "scores": [0.88, 0.22], - "input_tokens": 28, - "request_id": "env-test-123", - } - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - ) - - mock_post.assert_called_once() - - # Verify headers contain env API key - headers = mock_post.call_args.kwargs.get("headers", {}) - assert "Bearer env_test_key" in headers.get("Authorization", "") - - assert response.results is not None - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_error_handling(mock_post): - """Test DeepInfra rerank error handling.""" - error_response = {"detail": {"error": "Invalid API key"}} - - def return_val(): - return error_response - - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.json = return_val - mock_response.text = json.dumps(error_response) - mock_response.headers = {"content-type": "application/json"} - mock_post.return_value = mock_response - - # The current implementation handles errors gracefully, so we expect a successful response - # with the error information in the hidden params - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="invalid_key", - api_base="https://api.deepinfra.com", - ) - - # Verify that the response contains error information - assert ( - response._hidden_params["status"] == "unknown" - ) # Default status when error occurs - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_defaults_api_base_when_missing(mock_post, monkeypatch): - """With no api_base anywhere, the call still goes out against DeepInfra's own base.""" - monkeypatch.delenv("DEEPINFRA_API_BASE", raising=False) - - mock_response = MagicMock() - mock_response.json = lambda: {"scores": [0.9, 0.1], "input_tokens": 20} - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="test_key", - # api_base is intentionally missing - ) - - assert "api.deepinfra.com" in mock_post.call_args.kwargs["url"] - assert [result["relevance_score"] for result in response.results] == [0.9, 0.1] - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_request_format(mock_post): - """Test that the request is properly formatted for DeepInfra API.""" - mock_response_data = {"scores": [0.9, 0.1], "input_tokens": 20} - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="test query", - documents=["doc1", "doc2"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - instruction="custom instruction", - webhook="https://webhook.example.com", - ) - - mock_post.assert_called_once() - - # Verify URL format - call_url = mock_post.call_args.kwargs["url"] - assert call_url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B" - - # Verify headers - headers = mock_post.call_args.kwargs["headers"] - assert headers["Authorization"] == "Bearer test_key" - assert headers["accept"] == "application/json" - assert headers["content-type"] == "application/json" - - # Verify request body format - request_data = json.loads(mock_post.call_args.kwargs["data"]) - assert request_data["queries"] == [ - "test query", - "test query", - ] # DeepInfra requires queries to match documents length - assert request_data["documents"] == ["doc1", "doc2"] - assert request_data["instruction"] == "custom instruction" - assert request_data["webhook"] == "https://webhook.example.com" - - assert response.results is not None - - -def test_deepinfra_rerank_models(): - """Test that DeepInfra Qwen rerank models are recognized.""" - # These should not raise errors during model validation - models = [ - "deepinfra/Qwen/Qwen3-Reranker-0.6B", - "deepinfra/Qwen/Qwen3-Reranker-4B", - "deepinfra/Qwen/Qwen3-Reranker-8B", - ] - - for model in models: - resolved_model, provider, _, api_base = litellm.get_llm_provider(model=model) - assert provider == "deepinfra" - assert resolved_model == model.removeprefix("deepinfra/") - assert api_base == "https://api.deepinfra.com/v1/openai" - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_minimal_response(mock_post): - """Test handling of minimal DeepInfra response.""" - # Minimal response with just scores - mock_response_data = {"scores": [0.7, 0.3]} - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - # Should handle minimal response gracefully - assert response.results is not None - assert len(response.results) == 2 - assert response.results[0]["relevance_score"] == 0.7 - assert response.results[1]["relevance_score"] == 0.3 - - # Should have default values for missing fields - assert response.meta["tokens"]["input_tokens"] == 0 # Default when missing - assert response._hidden_params["status"] == "unknown" # Default when missing diff --git a/tests/test_litellm/llms/gemini/files/__init__.py b/tests/test_litellm/llms/gemini/files/__init__.py deleted file mode 100644 index f48fe7dbe2b..00000000000 --- a/tests/test_litellm/llms/gemini/files/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for Gemini files functionality""" diff --git a/tests/test_litellm/llms/gemini/videos/__init__.py b/tests/test_litellm/llms/gemini/videos/__init__.py deleted file mode 100644 index e0780c08321..00000000000 --- a/tests/test_litellm/llms/gemini/videos/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Gemini Video Generation Tests diff --git a/tests/test_litellm/llms/manus/__init__.py b/tests/test_litellm/llms/manus/__init__.py deleted file mode 100644 index c9121a7b2a4..00000000000 --- a/tests/test_litellm/llms/manus/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Manus provider tests diff --git a/tests/test_litellm/llms/manus/responses/__init__.py b/tests/test_litellm/llms/manus/responses/__init__.py deleted file mode 100644 index ea7ebb64d55..00000000000 --- a/tests/test_litellm/llms/manus/responses/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Manus Responses API tests diff --git a/tests/test_litellm/llms/minimax/__init__.py b/tests/test_litellm/llms/minimax/__init__.py deleted file mode 100644 index 451f542f4ad..00000000000 --- a/tests/test_litellm/llms/minimax/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax tests diff --git a/tests/test_litellm/llms/minimax/chat/__init__.py b/tests/test_litellm/llms/minimax/chat/__init__.py deleted file mode 100644 index 4a7916ae6cf..00000000000 --- a/tests/test_litellm/llms/minimax/chat/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax chat tests diff --git a/tests/test_litellm/llms/minimax/messages/__init__.py b/tests/test_litellm/llms/minimax/messages/__init__.py deleted file mode 100644 index de5a80602ea..00000000000 --- a/tests/test_litellm/llms/minimax/messages/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax messages tests diff --git a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py index d1eb6241ceb..db77eabba23 100644 --- a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py @@ -1,19 +1,9 @@ import os from typing import Dict -from unittest.mock import MagicMock -import httpx import litellm import pytest -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) -from litellm.llms.mistral.audio_transcription.transformation import ( - MistralAudioTranscriptionConfig, -) -from litellm.types.utils import TranscriptionResponse -from litellm.utils import ProviderConfigManager from tests.llm_translation.base_audio_transcription_unit_tests import ( BaseLLMAudioTranscriptionTest, ) @@ -37,184 +27,3 @@ class TestMistralAudioTranscription(BaseLLMAudioTranscriptionTest): "Async audio transcription test for Mistral is skipped in this suite; " "async test plugins (e.g. pytest-asyncio/anyio) are not configured here." ) - - -def test_mistral_audio_transcription_config_installed(): - """Ensure Mistral audio transcription config is registered with ProviderConfigManager.""" - config = ProviderConfigManager.get_provider_audio_transcription_config( - model="mistral/voxtral-mini-latest", - provider=litellm.LlmProviders.MISTRAL, - ) - assert config is not None - assert isinstance(config, BaseAudioTranscriptionConfig) - assert isinstance(config, MistralAudioTranscriptionConfig) - - -def test_mistral_audio_transcription_get_complete_url(): - config = MistralAudioTranscriptionConfig() - url = config.get_complete_url( - api_base=None, - api_key="fake-key", - model="voxtral-mini-latest", - optional_params={}, - litellm_params={}, - ) - assert url == "https://api.mistral.ai/v1/audio/transcriptions" - - -def test_mistral_audio_transcription_get_complete_url_custom_base(): - config = MistralAudioTranscriptionConfig() - url = config.get_complete_url( - api_base="https://custom.api.example.com/v1/", - api_key="fake-key", - model="voxtral-mini-latest", - optional_params={}, - litellm_params={}, - ) - assert url == "https://custom.api.example.com/v1/audio/transcriptions" - - -def test_mistral_audio_transcription_validate_environment(): - config = MistralAudioTranscriptionConfig() - headers = config.validate_environment( - headers={}, - model="voxtral-mini-latest", - messages=[], - optional_params={}, - litellm_params={}, - api_key="test-key-123", - ) - assert headers["Authorization"] == "Bearer test-key-123" - assert headers["accept"] == "application/json" - - -def test_mistral_audio_transcription_supported_params(): - config = MistralAudioTranscriptionConfig() - params = config.get_supported_openai_params("voxtral-mini-latest") - assert "language" in params - assert "temperature" in params - assert "response_format" in params - assert "timestamp_granularities" in params - - -def test_mistral_audio_transcription_request_transform(): - config = MistralAudioTranscriptionConfig() - - wav_path = os.path.join( - os.path.dirname(__file__), - "../../../../..", - "tests", - "llm_translation", - "gettysburg.wav", - ) - audio_file = open(wav_path, "rb") - - result = config.transform_audio_transcription_request( - model="voxtral-mini-latest", - audio_file=audio_file, - optional_params={"language": "en", "temperature": 0.0}, - litellm_params={}, - ) - - audio_file.close() - - assert isinstance(result.data, dict) - assert result.data["model"] == "voxtral-mini-latest" - assert result.data["language"] == "en" - assert result.data["temperature"] == 0.0 - assert result.files is not None - assert "file" in result.files - - -def test_mistral_audio_transcription_request_with_diarize(): - """Test that Mistral-specific params like diarize are passed through.""" - config = MistralAudioTranscriptionConfig() - - wav_path = os.path.join( - os.path.dirname(__file__), - "../../../../..", - "tests", - "llm_translation", - "gettysburg.wav", - ) - audio_file = open(wav_path, "rb") - - result = config.transform_audio_transcription_request( - model="voxtral-mini-latest", - audio_file=audio_file, - optional_params={"diarize": True}, - litellm_params={}, - ) - - audio_file.close() - - assert isinstance(result.data, dict) - assert result.data["diarize"] == "true" - - -def test_mistral_audio_transcription_response_transform(): - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = {"text": "Four score and seven years ago..."} - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "Four score and seven years ago..." - - -def test_mistral_audio_transcription_response_transform_diarized(): - """Test that diarized responses preserve segments and language.""" - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = { - "model": "voxtral-mini-latest", - "text": "Hello, how are you? I am fine.", - "language": None, - "segments": [ - { - "text": "Hello, how are you?", - "start": 0.3, - "end": 2.1, - "speaker_id": "speaker_1", - "type": "transcription_segment", - }, - { - "text": "I am fine.", - "start": 2.5, - "end": 3.8, - "speaker_id": "speaker_2", - "type": "transcription_segment", - }, - ], - "usage": { - "prompt_audio_seconds": 4, - "prompt_tokens": 5, - "total_tokens": 50, - "completion_tokens": 20, - }, - } - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "Hello, how are you? I am fine." - assert response["segments"] is not None - assert len(response["segments"]) == 2 - assert response["segments"][0]["speaker_id"] == "speaker_1" - assert response["segments"][1]["speaker_id"] == "speaker_2" - assert response["language"] is None - - -def test_mistral_audio_transcription_response_transform_empty(): - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = {} - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "" diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index d84cc8d3237..55703063fae 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -3,321 +3,12 @@ Tests for JSON-based provider configuration system. """ import os -import sys -from unittest.mock import patch -try: - import pytest -except ImportError: - # pytest not available, will run as standalone script - pytest = None - -# Add workspace to path -workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) -sys.path.insert(0, workspace_path) +import pytest import litellm -class TestJSONProviderLoader: - """Test JSON provider loading and configuration""" - - def test_load_json_providers(self): - """Test that JSON providers load correctly""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - # Verify publicai is loaded - assert JSONProviderRegistry.exists("publicai") - - # Get publicai config - publicai = JSONProviderRegistry.get("publicai") - assert publicai is not None - assert publicai.base_url == "https://api.publicai.co/v1" - assert publicai.api_key_env == "PUBLICAI_API_KEY" - assert publicai.api_base_env == "PUBLICAI_API_BASE" - assert publicai.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_dynamic_config_generation(self): - """Test dynamic config class creation""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Test API info resolution - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://api.publicai.co/v1" - - # Test with custom base - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.api.com", "test-key" - ) - assert api_base == "https://custom.api.com" - assert api_key == "test-key" - - def test_parameter_mapping(self): - """Test parameter mapping works""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Test parameter mapping - optional_params = {} - non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} - result = config.map_openai_params( - non_default_params, optional_params, "gpt-4", False - ) - - # max_completion_tokens should be mapped to max_tokens - assert "max_tokens" in result - assert result["max_tokens"] == 100 - assert "max_completion_tokens" not in result - - # temperature should be passed through - assert result["temperature"] == 0.7 - - def test_supported_params(self): - """Test that config returns supported params""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Get supported params - supported = config.get_supported_openai_params("gpt-4") - - # Should have standard OpenAI params - assert isinstance(supported, list) - assert len(supported) > 0 - - def test_tool_params_excluded_when_function_calling_not_supported(self): - """Test that tool-related params are excluded for models that don't support - function calling. Regression test for https://github.com/BerriAI/litellm/issues/21125 - """ - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Mock supports_function_calling to return False - with patch("litellm.utils.supports_function_calling", return_value=False): - supported = config.get_supported_openai_params("some-model-without-fc") - - tool_params = [ - "tools", - "tool_choice", - "function_call", - "functions", - "parallel_tool_calls", - ] - for param in tool_params: - assert ( - param not in supported - ), f"'{param}' should not be in supported params when function calling is not supported" - - # Non-tool params should still be present - assert "temperature" in supported - assert "max_tokens" in supported - assert "stop" in supported - - def test_tool_params_included_when_function_calling_supported(self): - """Test that tool-related params are included for models that support function calling.""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Mock supports_function_calling to return True - with patch("litellm.utils.supports_function_calling", return_value=True): - supported = config.get_supported_openai_params("some-model-with-fc") - - assert "tools" in supported - assert "tool_choice" in supported - - def test_provider_resolution(self): - """Test that provider resolution finds JSON providers""" - from litellm.litellm_core_utils.get_llm_provider_logic import ( - get_llm_provider, - ) - - model, provider, api_key, api_base = get_llm_provider( - model="publicai/gpt-4", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "gpt-4" - assert provider == "publicai" - assert api_base == "https://api.publicai.co/v1" - - def test_provider_config_manager(self): - """Test that ProviderConfigManager returns JSON-based configs""" - from litellm import LlmProviders - from litellm.utils import ProviderConfigManager - - config = ProviderConfigManager.get_provider_chat_config( - model="gpt-4", provider=LlmProviders.PUBLICAI - ) - - assert config is not None - assert config.custom_llm_provider == "publicai" - - -class TestPinstripes: - """Tests for Pinstripes JSON-configured provider""" - - def test_pinstripes_json_config_exists(self): - """Test that pinstripes is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - assert JSONProviderRegistry.exists("pinstripes") - - pinstripes = JSONProviderRegistry.get("pinstripes") - assert pinstripes is not None - assert pinstripes.base_url == "https://pinstripes.io/v1" - assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" - assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_pinstripes_provider_resolution(self): - """Test that provider resolution finds pinstripes and returns the default base URL""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="pinstripes/ps/glm-4.5-air", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "ps/glm-4.5-air" - assert provider == "pinstripes" - assert api_base == "https://pinstripes.io/v1" - - def test_pinstripes_dynamic_config(self): - """Test dynamic config class creation for pinstripes""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("pinstripes") - config_class = create_config_class(provider) - config = config_class() - - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://pinstripes.io/v1" - - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.pinstripes.io/v1", "test-key" - ) - assert api_base == "https://custom.pinstripes.io/v1" - assert api_key == "test-key" - - def test_pinstripes_parameter_mapping(self): - """Test that max_completion_tokens is mapped to max_tokens for pinstripes""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("pinstripes") - config_class = create_config_class(provider) - config = config_class() - - optional_params = {} - non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} - result = config.map_openai_params( - non_default_params, optional_params, "ps/glm-4.5-air", False - ) - - assert "max_tokens" in result - assert result["max_tokens"] == 100 - assert "max_completion_tokens" not in result - assert result["temperature"] == 0.7 - - -class TestDarkbloom: - def test_darkbloom_json_config_exists(self): - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - darkbloom = JSONProviderRegistry.get("darkbloom") - assert darkbloom is not None - assert darkbloom.base_url == "https://api.darkbloom.dev/v1" - assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" - assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" - assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_darkbloom_provider_resolution(self): - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="darkbloom/gemma-4-26b", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "gemma-4-26b" - assert provider == "darkbloom" - assert api_key is None - assert api_base == "https://api.darkbloom.dev/v1" - - def test_darkbloom_dynamic_config(self): - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("darkbloom") - config_class = create_config_class(provider) - config = config_class() - - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://api.darkbloom.dev/v1" - - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.darkbloom.dev/v1", "test-key" - ) - assert api_base == "https://custom.darkbloom.dev/v1" - assert api_key == "test-key" - - def test_darkbloom_complete_url_appends_endpoint(self): - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("darkbloom") - config_class = create_config_class(provider) - config = config_class() - - url = config.get_complete_url( - api_base="https://api.darkbloom.dev/v1", - api_key="test-key", - model="darkbloom/gemma-4-26b", - optional_params={}, - litellm_params={}, - stream=True, - ) - - assert url == "https://api.darkbloom.dev/v1/chat/completions" - - def test_darkbloom_provider_config_manager(self): - from litellm import LlmProviders - from litellm.utils import ProviderConfigManager - - config = ProviderConfigManager.get_provider_chat_config( - model="gemma-4-26b", provider=LlmProviders.DARKBLOOM - ) - - assert config is not None - assert config.custom_llm_provider == "darkbloom" - - class TestPublicAIIntegration: """Integration tests for PublicAI provider""" @@ -457,55 +148,3 @@ class TestPublicAIIntegration: pytest.fail(f"Content list conversion test failed: {str(e)}") else: raise - - -if __name__ == "__main__": - # Run basic tests - print("Testing JSON Provider System...") - - test_loader = TestJSONProviderLoader() - print("\n1. Testing JSON provider loading...") - test_loader.test_load_json_providers() - print(" ✓ JSON providers loaded") - - print("\n2. Testing dynamic config generation...") - test_loader.test_dynamic_config_generation() - print(" ✓ Dynamic config works") - - print("\n3. Testing parameter mapping...") - test_loader.test_parameter_mapping() - print(" ✓ Parameter mapping works") - - print("\n4. Testing excluded params...") - test_loader.test_excluded_params() - print(" ✓ Excluded params work") - - print("\n5. Testing provider resolution...") - test_loader.test_provider_resolution() - print(" ✓ Provider resolution works") - - print("\n6. Testing provider config manager...") - test_loader.test_provider_config_manager() - print(" ✓ Config manager works") - - print("\n" + "=" * 50) - print("PublicAI Integration Tests...") - print("=" * 50) - - test_integration = TestPublicAIIntegration() - - print("\n7. Testing basic completion...") - test_integration.test_publicai_completion_basic() - - print("\n8. Testing streaming...") - test_integration.test_publicai_completion_with_streaming() - - print("\n9. Testing parameter mapping...") - test_integration.test_publicai_parameter_mapping() - - print("\n10. Testing content list conversion...") - test_integration.test_publicai_content_list_conversion() - - print("\n" + "=" * 50) - print("✓ All tests passed!") - print("=" * 50) diff --git a/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py b/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py index 8104fb12943..580994f60b8 100644 --- a/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py +++ b/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py @@ -4,86 +4,12 @@ Related to issue #18794 """ import os -import sys -from unittest.mock import MagicMock, patch -try: - import pytest -except ImportError: - pytest = None - -# Add workspace to path -workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) -sys.path.insert(0, workspace_path) +import pytest import litellm -class TestXiaomiMiMoProviderConfig: - """Test Xiaomi MiMo provider configuration""" - - def test_xiaomi_mimo_in_provider_list(self): - """Test that xiaomi_mimo is in the provider list (fixes #18794)""" - from litellm import LlmProviders - - # Verify xiaomi_mimo is in the enum - assert hasattr(LlmProviders, "XIAOMI_MIMO") - assert LlmProviders.XIAOMI_MIMO.value == "xiaomi_mimo" - - # Verify it's in the provider list - assert "xiaomi_mimo" in litellm.provider_list - - def test_xiaomi_mimo_json_config_exists(self): - """Test that xiaomi_mimo is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - # Verify xiaomi_mimo is loaded - assert JSONProviderRegistry.exists("xiaomi_mimo") - - # Get xiaomi_mimo config - xiaomi_mimo = JSONProviderRegistry.get("xiaomi_mimo") - assert xiaomi_mimo is not None - assert xiaomi_mimo.base_url == "https://api.xiaomimimo.com/v1" - assert xiaomi_mimo.api_key_env == "XIAOMI_MIMO_API_KEY" - assert xiaomi_mimo.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_xiaomi_mimo_provider_resolution(self): - """Test that provider resolution finds xiaomi_mimo""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="xiaomi_mimo/mimo-v2-flash", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "mimo-v2-flash" - assert provider == "xiaomi_mimo" - assert api_base == "https://api.xiaomimimo.com/v1" - - def test_xiaomi_mimo_router_config(self): - """Test that xiaomi_mimo can be used in Router configuration (fixes #18794)""" - from litellm import Router - - # This should not raise "Unsupported provider - xiaomi_mimo" - router = Router( - model_list=[ - { - "model_name": "mimo-v2-flash", - "litellm_params": { - "model": "xiaomi_mimo/mimo-v2-flash", - "api_key": "test-key", - }, - } - ] - ) - - # Verify the deployment was created successfully - assert len(router.model_list) == 1 - assert router.model_list[0]["model_name"] == "mimo-v2-flash" - - class TestXiaomiMiMoIntegration: """Integration tests for Xiaomi MiMo provider""" @@ -128,30 +54,3 @@ class TestXiaomiMiMoIntegration: pytest.fail(f"Xiaomi MiMo completion failed: {str(e)}") else: raise - - -if __name__ == "__main__": - # Run basic tests - print("Testing Xiaomi MiMo Provider...") - - test_config = TestXiaomiMiMoProviderConfig() - - print("\n1. Testing provider in list...") - test_config.test_xiaomi_mimo_in_provider_list() - print(" ✓ xiaomi_mimo in provider list") - - print("\n2. Testing JSON config...") - test_config.test_xiaomi_mimo_json_config_exists() - print(" ✓ xiaomi_mimo JSON config loaded") - - print("\n3. Testing provider resolution...") - test_config.test_xiaomi_mimo_provider_resolution() - print(" ✓ Provider resolution works") - - print("\n4. Testing router configuration...") - test_config.test_xiaomi_mimo_router_config() - print(" ✓ Router configuration works (issue #18794 fixed)") - - print("\n" + "=" * 50) - print("✓ All configuration tests passed!") - print("=" * 50) diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py index c8751fb2d95..8cc46dc98d0 100644 --- a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py @@ -54,61 +54,3 @@ def test_ovhcloud_audio_transcription_config_installed(): assert config is not None assert isinstance(config, BaseAudioTranscriptionConfig) - - - -class TestOVHCloudDurationFieldMigration: - """Tests for OVHCloud duration -> seconds field migration.""" - - def test_seconds_field_mapped_to_duration(self): - """New `seconds` field should be normalized to `duration`.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello world", - "seconds": 3.14, - } - - result = config.transform_audio_transcription_response(mock_response) - - assert result.text == "Hello world" - assert result._hidden_params["duration"] == 3.14 - - def test_legacy_duration_field_still_works(self): - """Legacy `duration` field should still be accepted.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello world", - "duration": 2.71, - } - - result = config.transform_audio_transcription_response(mock_response) - - assert result.text == "Hello world" - assert result._hidden_params["duration"] == 2.71 - - - - def test_seconds_zero_mapped_to_duration(self): - """seconds=0.0 must not be treated as falsy and lost.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = {"text": "silence", "seconds": 0.0} - result = config.transform_audio_transcription_response(mock_response) - assert result._hidden_params["duration"] == 0.0 \ No newline at end of file diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py index 057ab9ede9a..34954587ed0 100644 --- a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py +++ b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py @@ -6,174 +6,12 @@ import os import pytest -from litellm.llms.ovhcloud.utils import OVHCloudException -from litellm.utils import get_optional_params -from litellm.llms.ovhcloud.chat.transformation import ( - OVHCloudChatCompletionStreamingHandler, - OVHCloudChatConfig, -) -config = OVHCloudChatConfig() model = "ovhcloud/Mistral-7B-Instruct-v0.3" -class TestOvhCloudChatCompletionStreamingHandler: - def test_chunk_parser_successful(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - chunk = { - "id": "test_id", - "created": 1234567890, - "model": "gpt-oss-20b", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - "choices": [ - {"delta": {"content": "test content", "reasoning": "test reasoning"}} - ], - } - - result = handler.chunk_parser(chunk) - - assert result.id == "test_id" - assert result.object == "chat.completion.chunk" - assert result.created == 1234567890 - assert result.model == "gpt-oss-20b" - assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] - assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] - assert result.usage.total_tokens == chunk["usage"]["total_tokens"] - assert len(result.choices) == 1 - assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" - - def test_chunk_parser_error_response(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - error_chunk = { - "error": { - "message": "test error", - "code": 400, - } - } - - with pytest.raises(OVHCloudException) as exc_info: - handler.chunk_parser(error_chunk) - - assert "OVHCloud Error: test error" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - def test_chunk_parser_key_error(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - invalid_chunk = {"incomplete": "data"} - - with pytest.raises(OVHCloudException) as exc_info: - handler.chunk_parser(invalid_chunk) - - assert "KeyError" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - -class TestOVHCloudConfig: - def test_transform_request_basic(self): - """Test basic request transformation""" - transformed_request = config.transform_request( - model, - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["model"] == model - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_transform_request_with_extra_body(self): - """Test request transformation with extra_body parameters""" - transformed_request = config.transform_request( - model, - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={"extra_body": {"custom_param": "custom_value"}}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["custom_param"] == "custom_value" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_map_openai_params(self): - """Test OpenAI parameter mapping""" - non_default_params = { - "temperature": 0.7, - "max_tokens": 100, - "top_p": 0.9, - } - - mapped_params = config.map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - drop_params=False, - ) - - assert mapped_params["temperature"] == 0.7 - assert mapped_params["max_tokens"] == 100 - assert mapped_params["top_p"] == 0.9 - - def test_get_error_class(self): - """Test error class creation""" - error = config.get_error_class( - error_message="Test error", - status_code=400, - headers={"Content-Type": "application/json"}, - ) - - assert isinstance(error, OVHCloudException) - assert error.message == "Test error" - assert error.status_code == 400 - - @pytest.mark.parametrize( - "model", - [ - "Meta-Llama-3_3-70B-Instruct", - "Meta-Llama-3_1-70B-Instruct", - "Mixtral-8x7B-Instruct-v0.1", - "gpt-oss-120b", - "some-model-not-in-the-cost-map", - ], - ) - def test_tools_not_filtered_by_static_model_map(self, model): - """ - OVHCloud AI Endpoints are OpenAI-compatible; tools/tool_choice must pass - through for any model. The server is responsible for rejecting unsupported - tool calls — LiteLLM must not strip them based on a stale static catalog. - """ - - params = get_optional_params( - model=model, - custom_llm_provider="ovhcloud", - tools=[ - { - "type": "function", - "function": {"name": "x", "parameters": {}}, - } - ], - tool_choice="auto", - ) - - assert "tools" in params - assert "tool_choice" in params - - def test_ovhcloud_integration(): from litellm import completion @@ -285,78 +123,3 @@ def test_ovhcloud_with_custom_base_url(): if __name__ == "__main__": pytest.main([__file__, "-v"]) - - -class TestOVHCloudReasoningFieldMigration: - """Tests for OVHCloud reasoning_content -> reasoning field migration.""" - - def test_streaming_new_reasoning_field(self): - """New `reasoning` field should be mapped to `reasoning_content`.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "role": "assistant", - "reasoning": "Let me think...", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "Let me think..." - - def test_streaming_legacy_reasoning_content_unchanged(self): - """Legacy `reasoning_content` field should pass through untouched.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "role": "assistant", - "reasoning_content": "Already correct field.", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "Already correct field." - - def test_streaming_both_fields_legacy_wins(self): - """When both fields present, existing `reasoning_content` is not overwritten.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "reasoning": "new field", - "reasoning_content": "legacy field", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "legacy field" - - diff --git a/tests/test_litellm/llms/reducto/__init__.py b/tests/test_litellm/llms/reducto/__init__.py deleted file mode 100644 index 8b137891791..00000000000 --- a/tests/test_litellm/llms/reducto/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/test_litellm/llms/s3_vectors/__init__.py b/tests/test_litellm/llms/s3_vectors/__init__.py deleted file mode 100644 index d4b0c4d8550..00000000000 --- a/tests/test_litellm/llms/s3_vectors/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# S3 Vectors tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py deleted file mode 100644 index 231735c1de7..00000000000 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# S3 Vectors vector store tests diff --git a/tests/test_litellm/llms/soniox/__init__.py b/tests/test_litellm/llms/soniox/__init__.py deleted file mode 100644 index b2cd496d66a..00000000000 --- a/tests/test_litellm/llms/soniox/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Soniox provider tests.""" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 4679b978f78..d3a7ba7a1bd 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,1681 +1,13 @@ -import base64 - import pytest from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) -from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - _transform_request_body, - check_if_part_exists_in_parts, - _get_highest_media_resolution, - _extract_max_media_resolution_from_messages, -) from litellm.types.llms.vertex_ai import BlobType -from litellm.types.utils import Message - - -def test_check_if_part_exists_in_parts(): - parts = [ - {"text": "Hello", "thought": True}, - {"text": "World", "thought": False}, - ] - part = {"text": "Hello", "thought": True} - new_part = {"text": "Hello World", "thought": True} - assert check_if_part_exists_in_parts(parts, part) - assert not check_if_part_exists_in_parts(parts, new_part, ["thought"]) - assert check_if_part_exists_in_parts(parts, new_part, ["text"]) - - -def test_check_if_part_exists_in_parts_camel_case_snake_case(): - """Test that function handles both camelCase and snake_case key variations""" - # Test snake_case to camelCase matching - parts_with_snake_case = [ - { - "function_call": { - "name": "get_current_weather", - "args": {"location": "San Francisco, CA"}, - } - }, - {"text": "Some other content"}, - ] - - part_with_camel_case = { - "functionCall": { - "name": "get_current_weather", - "args": {"location": "San Francisco, CA"}, - } - } - - # Should find match between function_call and functionCall - assert check_if_part_exists_in_parts(parts_with_snake_case, part_with_camel_case) - - # Test camelCase to snake_case matching - parts_with_camel_case = [ - {"functionCall": {"name": "calculate_sum", "args": {"a": 1, "b": 2}}} - ] - - part_with_snake_case = { - "function_call": {"name": "calculate_sum", "args": {"a": 1, "b": 2}} - } - - # Should find match between functionCall and function_call - assert check_if_part_exists_in_parts(parts_with_camel_case, part_with_snake_case) - - # Test no match when values differ - part_with_different_values = { - "function_call": {"name": "different_function", "args": {"x": 5}} - } - - assert not check_if_part_exists_in_parts( - parts_with_snake_case, part_with_different_values - ) - - # Test multiple keys with mixed casing - parts_mixed = [ - { - "function_call": {"name": "test"}, - "thoughtSignature": "reasoning", - "text": "content", - } - ] - - part_mixed_casing = { - "functionCall": {"name": "test"}, - "thought_signature": "reasoning", - "text": "content", - } - - assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) - - -def test_cached_content_respects_modify_params_for_cache_incompatible_fields(): - """Regression: cachedContent drops system/tools/toolConfig only when modify_params=True.""" - import litellm - - cache_name = "projects/p/locations/us-central1/cachedContents/abc123" - messages = [ - {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "hi"}, - ] - optional_params = { - "tools": [ - { - "functionDeclarations": [ - {"name": "get_weather", "description": "Get weather"}, - ] - } - ], - "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, - } - - original_modify_params = litellm.modify_params - try: - # With modify_params=False (default), keep fields even with cachedContent. - litellm.modify_params = False - result = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=cache_name, - ) - assert result.get("cachedContent") == cache_name - assert "system_instruction" in result - assert "tools" in result - assert "toolConfig" in result - assert "contents" in result - - # With modify_params=True, drop cache-incompatible fields. - litellm.modify_params = True - result_modify_true = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=cache_name, - ) - assert result_modify_true.get("cachedContent") == cache_name - assert "system_instruction" not in result_modify_true - assert "tools" not in result_modify_true - assert "toolConfig" not in result_modify_true - assert "contents" in result_modify_true - - # Without cache, fields are always included. - result_no_cache = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - assert "system_instruction" in result_no_cache - assert "tools" in result_no_cache - assert "toolConfig" in result_no_cache - finally: - litellm.modify_params = original_modify_params - - -# Tests for issue #14556: Labels field provider-aware filtering -def test_google_genai_excludes_labels(): - """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"labels": {"project": "test", "team": "ai"}} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="gemini", - litellm_params=litellm_params, - cached_content=None, - ) - - # Google GenAI/AI Studio should NOT include labels - assert "labels" not in result - assert "contents" in result - - -def test_vertex_ai_includes_labels(): - """Test that Vertex AI endpoints include labels when custom_llm_provider='vertex_ai'""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"labels": {"project": "test", "team": "ai"}} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - # Vertex AI SHOULD include labels - assert "labels" in result - assert result["labels"] == {"project": "test", "team": "ai"} - - -def test_service_tier_forwarded_to_vertex_ai(): - """Test that service_tier in optional_params is mapped to serviceTier in request body.""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"service_tier": "flex"} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - assert "serviceTier" in result - assert result["serviceTier"] == "flex" - - -def test_extra_body_cache_not_forwarded_to_vertex_ai(): - """ - 'cache' inside extra_body is a LiteLLM-internal proxy caching control. - It must NOT be forwarded to the Vertex AI request body. - - Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." - Vertex AI enforces a strict JSON schema and rejects any unknown field. - """ - messages = [{"role": "user", "content": "test"}] - optional_params = { - "extra_body": { - "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal - "some_vertex_param": "value", # legitimate provider extra - }, - } - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - # 'cache' must be stripped — Vertex AI has no such field - assert "cache" not in result, ( - "extra_body.cache must not be forwarded to Vertex AI. " - 'Vertex AI rejects it with 400: Unknown name "cache": Cannot find field.' - ) - - # Other legitimate extra_body keys should still pass through - assert "some_vertex_param" in result - assert result["some_vertex_param"] == "value" - - # Core request fields must be present - assert "contents" in result - - -def test_extra_body_tags_not_forwarded_to_vertex_ai(): - """ - 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. - It must NOT be forwarded to the Vertex AI request body. - Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" - """ - messages = [{"role": "user", "content": "test"}] - optional_params = { - "extra_body": { - "tags": ["user:alice", "env:prod"], - "custom_param": "allowed", - }, - } - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - assert "tags" not in result - assert "custom_param" in result - assert result["custom_param"] == "allowed" - - -def test_extra_body_google_maps_rewrites_json_response_format(): - messages = [{"role": "user", "content": "test"}] - optional_params = { - "response_mime_type": "application/json", - "response_schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - "extra_body": { - "tools": [{"googleMaps": {}}], - }, - } - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - - generation_config = result["generationConfig"] - assert "response_mime_type" not in generation_config - assert generation_config["responseFormat"] == { - "text": { - "mimeType": "APPLICATION_JSON", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - } - } - - -def test_extra_body_generation_config_cannot_restore_google_maps_json_mime_type(): - messages = [{"role": "user", "content": "test"}] - optional_params = { - "tools": [{"googleMaps": {}}], - "response_mime_type": "application/json", - "extra_body": { - "generationConfig": { - "response_mime_type": "application/json", - "response_json_schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - }, - }, - } - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - - generation_config = result["generationConfig"] - assert "response_mime_type" not in generation_config - assert "response_json_schema" not in generation_config - assert generation_config["responseFormat"] == { - "text": { - "mimeType": "APPLICATION_JSON", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - } - } - - -def test_metadata_to_labels_vertex_only(): - """Test that metadata->labels conversion only happens for Vertex AI""" - messages = [{"role": "user", "content": "test"}] - optional_params = {} - litellm_params = { - "metadata": { - "requester_metadata": {"user": "john_doe", "project": "test-project"} - } - } - - # Google GenAI/AI Studio should not include labels from metadata - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params.copy(), - custom_llm_provider="gemini", - litellm_params=litellm_params.copy(), - cached_content=None, - ) - assert "labels" not in result - - # Vertex AI should include labels from metadata - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params.copy(), - custom_llm_provider="vertex_ai", - litellm_params=litellm_params.copy(), - cached_content=None, - ) - assert "labels" in result - assert result["labels"] == {"user": "john_doe", "project": "test-project"} - - -def test_empty_content_handling(): - """Test that empty content strings are properly handled in Gemini message transformation""" - # Test with empty content in user message - messages = [{"content": "", "role": "user"}] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify that the content was properly transformed - assert len(contents) == 1 - assert contents[0]["role"] == "user" - assert len(contents[0]["parts"]) == 1 - assert "text" in contents[0]["parts"][0] - assert contents[0]["parts"][0]["text"] == "" - - -def test_thought_signature_extraction_from_response(): - """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - # Test case: Single function call with thought signature - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, - ) - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - # Verify thought signature is stored in provider_specific_fields - assert tools is not None - assert len(tools) == 1 - assert "provider_specific_fields" in tools[0] - assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature - - -def test_thought_signature_parallel_function_calls(): - """Test that only the first function call in parallel calls has thought signature""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Parallel function calls - only first has signature - parts_parallel = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, # First FC has signature - ), - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "London"}, - }, - # Second FC has no signature (parallel call) - ), - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_parallel, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - # Verify only first tool call has thought signature - assert tools is not None - assert len(tools) == 2 - assert "provider_specific_fields" in tools[0] - assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature - # Second tool call should not have thought signature - assert "provider_specific_fields" not in tools[ - 1 - ] or "thought_signature" not in tools[1].get("provider_specific_fields", {}) - - -def test_thought_signature_preservation_in_conversion(): - """Test that thought signatures are preserved when converting assistant messages back to Gemini format""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Assistant message with tool calls containing thought signatures - assistant_message = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": test_signature, - }, - }, - { - "id": "call_def456", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "London"}', - }, - "index": 1, - # No thought signature for parallel call - }, - ], - } - - gemini_parts = convert_to_gemini_tool_call_invoke(assistant_message) - - # Verify thought signature is preserved in first function call part - assert len(gemini_parts) == 2 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - assert gemini_parts[0]["thoughtSignature"] == test_signature - - # Verify second function call part does not have thought signature - assert "function_call" in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[1] - - -def test_thought_signature_sequential_function_calls(): - """Test that each sequential function call preserves its own thought signature""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - signature_1 = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - signature_2 = "DifferentSignatureForSecondCall1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ" - - # Sequential function calls - each has its own signature - # This simulates a multi-step conversation where each step has a signature - assistant_message_step1 = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_step1", - "type": "function", - "function": { - "name": "check_flight", - "arguments": '{"flight": "AA100"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": signature_1, - }, - }, - ], - } - - assistant_message_step2 = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_step2", - "type": "function", - "function": { - "name": "book_taxi", - "arguments": '{"destination": "airport"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": signature_2, - }, - }, - ], - } - - gemini_parts_step1 = convert_to_gemini_tool_call_invoke(assistant_message_step1) - gemini_parts_step2 = convert_to_gemini_tool_call_invoke(assistant_message_step2) - - # Verify each step preserves its own signature - assert len(gemini_parts_step1) == 1 - assert gemini_parts_step1[0]["thoughtSignature"] == signature_1 - - assert len(gemini_parts_step2) == 1 - assert gemini_parts_step2[0]["thoughtSignature"] == signature_2 - - -def test_thought_signature_with_function_call_mode(): - """Test thought signature extraction in function_call mode (is_function_call=True)""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_weather", - "args": {"location": "Tokyo"}, - }, - thoughtSignature=test_signature, - ) - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=True, - ) - - # Verify thought signature is stored in function's provider_specific_fields - assert function is not None - # Function should be dict-like (TypedDict or dict) - assert hasattr(function, "__getitem__") or isinstance(function, dict) - assert "provider_specific_fields" in function - assert function["provider_specific_fields"]["thought_signature"] == test_signature - assert tools is None - - -def test_dummy_signature_added_for_gemini_3_conversation_history(): - """Test that dummy signatures are added when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3.""" - import base64 - - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Simulate conversation history from gemini-2.5-flash (no thought signature) - assistant_message_from_older_model = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - # No provider_specific_fields - older model doesn't provide signatures - }, - ], - } - - # Convert to Gemini format for gemini-3-pro-preview (should add dummy signature) - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_from_older_model, model="gemini-3-pro-preview" - ) - - # Verify dummy signature is added - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - - # Verify it's the expected dummy signature (base64 encoded "skip_thought_signature_validator") - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) - assert gemini_parts[0]["thoughtSignature"] == expected_dummy - - -def test_dummy_signature_not_added_for_gemini_2_5(): - """Test that dummy signatures are NOT added when target model is not gemini-3.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Simulate conversation history from gemini-2.5-flash (no thought signature) - assistant_message = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - # No provider_specific_fields - }, - ], - } - - # Convert to Gemini format for gemini-2.5-flash (should NOT add dummy signature) - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message, model="gemini-2.5-flash" - ) - - # Verify no dummy signature is added for non-gemini-3 models - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" not in gemini_parts[0] - - -def test_dummy_signature_not_added_when_signature_exists(): - """Test that dummy signatures are NOT added when a real signature already exists.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - real_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Assistant message with existing thought signature - assistant_message_with_signature = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - "provider_specific_fields": { - "thought_signature": real_signature, - }, - }, - "index": 0, - }, - ], - } - - # Convert to Gemini format for gemini-3-pro-preview - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_with_signature, model="gemini-3-pro-preview" - ) - - # Verify real signature is preserved, not replaced with dummy - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - assert gemini_parts[0]["thoughtSignature"] == real_signature - - -def test_dummy_signature_with_function_call_mode(): - """Test that dummy signatures are added for function_call mode when converting to gemini-3.""" - import base64 - - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Assistant message with function_call (not tool_calls) and no signature - assistant_message_function_call = { - "role": "assistant", - "content": None, - "function_call": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - # No provider_specific_fields - }, - } - - # Convert to Gemini format for gemini-3-pro-preview - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_function_call, model="gemini-3-pro-preview" - ) - - # Verify dummy signature is added - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - - # Verify it's the expected dummy signature - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) - assert gemini_parts[0]["thoughtSignature"] == expected_dummy - - -def _parallel_tool_calls(*signatures): - return [ - { - "id": f"call_{idx}", - "type": "function", - "function": { - "name": f"tool_{idx}", - "arguments": '{"location": "Paris"}', - **( - {"provider_specific_fields": {"thought_signature": signature}} - if signature is not None - else {} - ), - }, - "index": idx, - } - for idx, signature in enumerate(signatures) - ] - - -def _parallel_tool_calls_signed_via_id(*signatures): - """Parallel tool calls in the shape LiteLLM actually hands back to clients. - - The signature rides in the tool call id behind __thought__, which is what an - OpenAI-format client echoes back on the next turn. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - _encode_tool_call_id_with_signature, - ) - - return [ - { - "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), - "type": "function", - "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, - "index": idx, - } - for idx, signature in enumerate(signatures) - ] - - -REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" -PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" -) - - -def test_dummy_signature_only_on_first_parallel_tool_call(): - """Google documents the placeholder as a last resort that degrades quality, so an unsigned - parallel turn replayed to gemini-3 gets a budget of exactly one.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None, None), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): - """Gemini signs only the first of N parallel function calls, so a faithful replay has - nothing to attach to the siblings.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_real_signature_on_later_parallel_tool_call_is_preserved(): - """Clients may reorder or drop calls, so a signature that lands on a non-first call is - still the model's own and must survive the round trip.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - - -def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): - """Non-gemini-3 models never get a placeholder signature, on any call.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None), - }, - model="gemini-2.5-flash", - ) - - assert len(gemini_parts) == 2 - assert all("thoughtSignature" not in part for part in gemini_parts) - - -def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): - """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): - """A signature on the tool call itself, rather than on its function, behaves the same way.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - tool_calls = _parallel_tool_calls(None, None) - tool_calls[0]["provider_specific_fields"] = { - "thought_signature": REAL_THOUGHT_SIGNATURE - } - - gemini_parts = convert_to_gemini_tool_call_invoke( - {"role": "assistant", "content": None, "tool_calls": tool_calls}, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): - """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not - consume the one placeholder slot and leave the real first function call bare.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - tool_calls = [ - {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} - ] + _parallel_tool_calls(None, None) - - gemini_parts = convert_to_gemini_tool_call_invoke( - {"role": "assistant", "content": None, "tool_calls": tool_calls}, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_no_placeholder_when_model_is_unknown(): - """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None), - }, - ) - - assert len(gemini_parts) == 2 - assert all("thoughtSignature" not in part for part in gemini_parts) - - -def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): - """Older models still receive a real signature that a client replays, and still get no placeholder.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), - }, - model="gemini-2.5-flash", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_parallel_tool_call_history_replayed_through_full_message_conversion(): - """End to end through the message-history converter, the path a real /chat/completions replay takes.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - messages = [ - {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - ] - - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-3-pro-preview" - ) - - model_parts = contents[1]["parts"] - assert len(model_parts) == 3 - assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in model_parts[1] - assert "thoughtSignature" not in model_parts[2] - - -@pytest.mark.parametrize( - "model", - ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], -) -def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): - """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. - - Fabricating the placeholder alongside a real signature is what produced empty text responses - on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. - """ - import json - - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - messages = [ - {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages, model=model) - - model_parts = contents[1]["parts"] - assert len(model_parts) == 3 - assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in model_parts[1] - assert "thoughtSignature" not in model_parts[2] - assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) - - -@pytest.mark.parametrize( - "model", - [ - "gemini-3-pro-preview", - "gemini-3-flash-preview", - "gemini-3.1-pro-preview", - "gemini-3.5-flash", - "gemini-3.6-flash", - "gemini-3.7-flash", - "gemini-3.8-flash", - "vertex_ai/gemini-3.5-flash", - "vertex_ai/gemini-3.7-flash", - "vertex_ai/gemini-3.8-flash", - "gemini/gemini-3.5-flash", - "gemini/gemini-3.7-flash", - "gemini/gemini-3.8-flash", - ], -) -def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): - """The gemini-3 gate is a substring match, so every family member and prefix form has to - land on the same one-placeholder budget rather than only the versions we happened to try.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None, None), - }, - model=model, - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): - """Text-part and function-call signatures are collected by separate code paths, so scoping the - placeholder must not disturb a real signature that arrived on the text part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Checking all three cities.", - "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, - "tool_calls": _parallel_tool_calls(None, None, None), - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-3-pro-preview" - )[0]["parts"] - - assert parts[0]["text"] == "Checking all three cities." - assert parts[0]["thoughtSignature"] == "real_25_signature" - assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in parts[2] - assert "thoughtSignature" not in parts[3] - - -# Tests for media_resolution (detail parameter) handling - Issue #17084 -class TestMediaResolution: - """Tests for media_resolution handling in Gemini 2.x models""" - - def test_get_highest_media_resolution_high_wins(self): - """Test that 'high' resolution takes precedence over 'low'""" - assert _get_highest_media_resolution("low", "high") == "high" - assert _get_highest_media_resolution("high", "low") == "high" - assert _get_highest_media_resolution(None, "high") == "high" - assert _get_highest_media_resolution("high", None) == "high" - - def test_get_highest_media_resolution_low_over_none(self): - """Test that 'low' resolution takes precedence over None""" - assert _get_highest_media_resolution(None, "low") == "low" - assert _get_highest_media_resolution("low", None) == "low" - - def test_get_highest_media_resolution_same_values(self): - """Test handling of same resolution values""" - assert _get_highest_media_resolution("high", "high") == "high" - assert _get_highest_media_resolution("low", "low") == "low" - assert _get_highest_media_resolution(None, None) is None - - def test_get_highest_media_resolution_medium(self): - """Test that 'medium' resolution is correctly ranked between 'low' and 'high'""" - assert _get_highest_media_resolution("low", "medium") == "medium" - assert _get_highest_media_resolution("medium", "low") == "medium" - assert _get_highest_media_resolution("medium", "high") == "high" - assert _get_highest_media_resolution("high", "medium") == "high" - assert _get_highest_media_resolution(None, "medium") == "medium" - assert _get_highest_media_resolution("medium", None) == "medium" - - def test_get_highest_media_resolution_ultra_high(self): - """Test that 'ultra_high' resolution takes precedence over all others""" - assert _get_highest_media_resolution("high", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("ultra_high", "high") == "ultra_high" - assert _get_highest_media_resolution("medium", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("low", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution(None, "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("ultra_high", None) == "ultra_high" - - def test_extract_max_media_resolution_single_image_high(self): - """Test extraction of media resolution from single image with detail=high""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_single_image_low(self): - """Test extraction of media resolution from single image with detail=low""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "low" - - def test_extract_max_media_resolution_no_detail(self): - """Test extraction when no detail parameter is provided""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,abc123"}, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) is None - - def test_extract_max_media_resolution_multiple_images_mixed(self): - """Test that highest resolution is returned when multiple images have different details""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Compare these images"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,def456", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_text_only(self): - """Test extraction from messages with no images""" - messages = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well!"}, - ] - assert _extract_max_media_resolution_from_messages(messages) is None - - def test_transform_request_body_gemini_2x_adds_media_resolution(self): - """Test that media_resolution is added to generationConfig for Gemini 2.x models""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - assert "generationConfig" in result - assert "mediaResolution" in result["generationConfig"] - assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_HIGH" - - def test_transform_request_body_gemini_2x_low_resolution(self): - """Test that low media_resolution is correctly added for Gemini 2.x""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "low", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - assert "generationConfig" in result - assert "mediaResolution" in result["generationConfig"] - assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_LOW" - - def test_transform_request_body_gemini_3_no_global_media_resolution(self): - """Test that Gemini 3 models don't add media_resolution to generationConfig (they use per-part)""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-3-pro-preview", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # Gemini 3 should NOT have mediaResolution in generationConfig - # (it's handled per-part in the content transformation) - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - def test_transform_request_body_no_detail_no_media_resolution(self): - """Test that no mediaResolution is added when detail is not specified""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # When no detail is specified, mediaResolution should not be in generationConfig - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - def test_extract_max_media_resolution_file_type_with_detail(self): - """Test that detail is extracted from file content type, not just image_url""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is in this file?"}, - { - "type": "file", - "file": { - "url": "data:image/png;base64,abc123", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_mixed_image_and_file(self): - """Test that highest detail is returned across both image_url and file types""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Compare these"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - { - "type": "file", - "file": { - "url": "data:image/png;base64,def456", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_transform_request_body_gemini_1x_no_media_resolution(self): - """Test that Gemini 1.x models don't get mediaResolution in generationConfig""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-1.5-pro", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # Gemini 1.x should NOT have mediaResolution (not supported) - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - -# Tests for VideoMetadata support across all Gemini models (Issue #25474) -class TestVideoMetadataAllGeminiModels: - """Tests that video_metadata (fps, start_offset, end_offset) works for all Gemini models""" - - def _make_video_messages(self, video_metadata: dict) -> list: - return [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Analyze this video"}, - { - "type": "file", - "file": { - "file_id": "gs://bucket/video.mp4", - "format": "video/mp4", - "video_metadata": video_metadata, - }, - }, - ], - } - ] - - def _get_file_part(self, contents: list) -> dict: - for part in contents[0]["parts"]: - if "file_data" in part: - return part - raise AssertionError("No file part found in contents") - - def test_video_metadata_fps_gemini_2_5_flash(self): - """Gemini 2.5 Flash: fps in video_metadata should be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 5}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 5 - - def test_video_metadata_fps_gemini_2_5_pro(self): - """Gemini 2.5 Pro: fps in video_metadata should be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 10}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-pro" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 10 - - def test_video_metadata_offsets_gemini_2_5_flash(self): - """Gemini 2.5 Flash: start_offset/end_offset converted to camelCase (Issue #25474)""" - messages = self._make_video_messages( - {"start_offset": "5s", "end_offset": "30s"} - ) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - vm = file_part["video_metadata"] - assert vm["startOffset"] == "5s" - assert vm["endOffset"] == "30s" - - def test_video_metadata_all_fields_gemini_2_5_flash(self): - """Gemini 2.5 Flash: all video_metadata fields forwarded correctly (Issue #25474)""" - messages = self._make_video_messages( - {"fps": 5, "start_offset": "10s", "end_offset": "60s"} - ) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - vm = file_part["video_metadata"] - assert vm["fps"] == 5 - assert vm["startOffset"] == "10s" - assert vm["endOffset"] == "60s" - - def test_video_metadata_gemini_1_5_pro(self): - """Gemini 1.5 Pro: video_metadata should also be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 2}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-1.5-pro" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 2 - - -def test_convert_tool_response_with_base64_image(): - """Test tool response with base64 data URI image.""" - # Create a small test image (1x1 red pixel PNG) - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create tool message with image - tool_message = { - "role": "tool", - "tool_call_id": "call_test123", - "content": [ - { - "type": "text", - "text": '{"url": "https://example.com", "status": "success"}', - }, - {"type": "input_image", "image_url": image_data_uri}, - ], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_test123", - "function": {"name": "click_at", "arguments": '{"x": 100, "y": 200}'}, - } - ] - } - - # Convert tool response with nested multimodal functionResponse.parts. - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - result_part = result[0] - assert "function_response" in result_part - assert "inline_data" not in result_part - function_response = result_part["function_response"] - assert function_response["name"] == "click_at" - assert "response" in function_response - # Verify JSON response is parsed correctly - assert "url" in function_response["response"] - assert function_response["response"]["url"] == "https://example.com" - - # Check inline_data is nested under functionResponse.parts. - assert "parts" in function_response - assert len(function_response["parts"]) == 1 - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "image/png" - assert inline_data["data"] == test_image_base64 - - -def test_gemini_history_nests_multimodal_tool_response_parts(): - """Full history conversion should not emit sibling inline_data tool result parts.""" - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - messages = [ - {"role": "user", "content": "Get me an image"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_get_image", - "type": "function", - "function": {"name": "get_image", "arguments": "{}"}, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_get_image", - "content": [ - {"type": "text", "text": '{"image_ref": "inline"}'}, - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": test_image_base64, - }, - }, - ], - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - tool_response_parts = contents[-1]["parts"] - assert len(tool_response_parts) == 1 - assert "inline_data" not in tool_response_parts[0] - function_response = tool_response_parts[0]["function_response"] - assert function_response["parts"] == [ - { - "inline_data": { - "data": test_image_base64, - "mime_type": "image/png", - } - } - ] def test_convert_tool_response_with_url_image(): """Test tool response with HTTP URL image (will download and convert).""" - import pytest - # Use a publicly accessible test image URL test_image_url = "https://via.placeholder.com/1x1.png" @@ -1701,13 +33,9 @@ def test_convert_tool_response_with_url_image(): } try: - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) + result = convert_to_gemini_tool_call_result(tool_message, last_message_with_tool_calls) - assert isinstance( - result, list - ), "Should return a parts list when media is present" + assert isinstance(result, list), "Should return a parts list when media is present" assert len(result) == 1, "Should return one function_response part" result_part = result[0] assert "function_response" in result_part @@ -1724,1060 +52,3 @@ def test_convert_tool_response_with_url_image(): except Exception as e: # Skip test if URL download fails (no internet connection, etc.) pytest.skip(f"Failed to download image from URL: {e}") - - -def test_convert_tool_response_text_only(): - """Test tool response with only text (no image).""" - tool_message = { - "role": "tool", - "tool_call_id": "call_test789", - "content": [ - {"type": "text", "text": '{"status": "completed", "result": "success"}'} - ], - } - - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_test789", - "function": {"name": "wait_5_seconds", "arguments": "{}"}, - } - ] - } - - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Should be a single part (no list) when no image - assert not isinstance(result, list), "Should return single part when no image" - - # Check function_response exists - assert "function_response" in result - function_response = result["function_response"] - assert function_response["name"] == "wait_5_seconds" - # Verify JSON response is parsed correctly - assert "status" in function_response["response"] - assert function_response["response"]["status"] == "completed" - - # Check inline_data does NOT exist (no image provided) - assert "inline_data" not in result - - -def test_file_data_field_order(): - """ - Test that file_data fields are in the correct order (mime_type before file_uri). - - The Gemini API is sensitive to field order in the file_data object. - This test verifies that mime_type comes before file_uri in both: - 1. Dictionary key order - 2. JSON serialization - - Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. - """ - import json - - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - # Test with HTTPS URL and explicit format (audio file) - file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" - format = "audio/mpeg" - - result = _process_gemini_media(image_url=file_url, format=format) - - # Verify the result has file_data - assert "file_data" in result - file_data = result["file_data"] - - # Verify both fields are present - assert "mime_type" in file_data - assert "file_uri" in file_data - assert file_data["mime_type"] == "audio/mpeg" - assert file_data["file_uri"] == file_url - - # Verify field order by checking dictionary keys - # In Python 3.7+, dict maintains insertion order - file_data_keys = list(file_data.keys()) - assert file_data_keys.index("mime_type") < file_data_keys.index( - "file_uri" - ), "mime_type must come before file_uri in the file_data dict" - - # Also verify by serializing to JSON string - json_str = json.dumps(file_data) - mime_type_pos = json_str.find('"mime_type"') - file_uri_pos = json_str.find('"file_uri"') - assert ( - mime_type_pos < file_uri_pos - ), "mime_type must appear before file_uri in JSON serialization" - - -def test_file_data_field_order_gcs_urls(): - """Test that GCS URLs also maintain correct field order.""" - import json - - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - # Test with GCS URL - gcs_url = "gs://bucket/audio.mp3" - - result = _process_gemini_media(image_url=gcs_url) - - # Verify the result has file_data - assert "file_data" in result - file_data = result["file_data"] - - # Verify both fields are present - assert "mime_type" in file_data - assert "file_uri" in file_data - - # Verify field order - file_data_keys = list(file_data.keys()) - assert file_data_keys.index("mime_type") < file_data_keys.index( - "file_uri" - ), "mime_type must come before file_uri in the file_data dict" - - -def test_gemini_files_api_uri_without_format(): - """ - Test that Gemini Files API URIs work WITHOUT an explicit format/mime_type. - - When a user uploads a file via the Gemini Files API and then references it - by URI (https://generativelanguage.googleapis.com/v1beta/files/...), - the file is already on Google's servers. These URLs return 403 when - fetched directly, so _process_gemini_media must NOT try to resolve the - MIME type via HTTP. Instead it should pass the URI through as file_data - and let the Gemini API resolve the type from its stored metadata. - - Related issue: https://github.com/BerriAI/litellm/issues/24907 - """ - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - file_url = "https://generativelanguage.googleapis.com/v1beta/files/37eh7rsw1vfe" - - # Should NOT raise — previously this hit the generic https:// handler - # which called _get_image_mime_type_from_url() and got a 403. - result = _process_gemini_media(image_url=file_url) - - assert "file_data" in result - file_data = result["file_data"] - assert file_data["file_uri"] == file_url - # When no format is provided, mime_type should be absent so the - # Gemini API infers it from the stored file metadata. - assert "mime_type" not in file_data - - -def test_gemini_files_api_uri_with_format(): - """ - Test that Gemini Files API URIs correctly forward an explicit format. - - Related issue: https://github.com/BerriAI/litellm/issues/24907 - """ - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - file_url = "https://generativelanguage.googleapis.com/v1beta/files/n1vhxa28lyaw" - - result = _process_gemini_media(image_url=file_url, format="text/plain") - - assert "file_data" in result - file_data = result["file_data"] - assert file_data["file_uri"] == file_url - assert file_data["mime_type"] == "text/plain" - - -def test_extract_file_data_with_path_object(): - """ - Test that filename is correctly extracted from Path objects for MIME type detection. - - When uploading files using Path objects (e.g., Path("speech.mp3")), the filename - must be extracted to enable proper MIME type detection. Without this, files get - uploaded with 'application/octet-stream' instead of the correct MIME type. - - Related issue: Files uploaded with wrong MIME type cause Gemini API to reject - requests where the specified format doesn't match the uploaded file's MIME type. - """ - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - # Create a temporary MP3 file - with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp: - tmp.write(b"fake mp3 content") - tmp_path = tmp.name - - try: - # Test with Path object - path_obj = Path(tmp_path) - extracted = extract_file_data(path_obj) - - # Verify filename was extracted - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".mp3") - - # Verify MIME type was correctly detected - assert ( - extracted["content_type"] == "audio/mpeg" - ), f"Expected 'audio/mpeg' but got '{extracted['content_type']}'" - - # Verify content was read - assert extracted["content"] == b"fake mp3 content" - - finally: - # Clean up temporary file - os.unlink(tmp_path) - - -def test_extract_file_data_with_pathlib_path(): - """Test that filename is correctly extracted from pathlib.Path inputs. - Bare str paths are rejected — when this runs in a proxy request handler - the value is attacker-controlled and opening it as a path is an LFI.""" - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp.write(b"fake wav content") - tmp_path = Path(tmp.name) - - try: - extracted = extract_file_data(tmp_path) - - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".wav") - assert extracted["content_type"] in [ - "audio/wav", - "audio/x-wav", - ], f"Expected 'audio/wav' or 'audio/x-wav' but got '{extracted['content_type']}'" - assert extracted["content"] == b"fake wav content" - finally: - os.unlink(str(tmp_path)) - - -def test_extract_file_data_with_tuple_format(): - """Test that tuple format (with explicit content_type) still works correctly.""" - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - # Test with tuple format: (filename, content, content_type) - filename = "test_audio.mp3" - content = b"test audio content" - content_type = "audio/mpeg" - - extracted = extract_file_data((filename, content, content_type)) - - # Verify all fields are correct - assert extracted["filename"] == filename - assert extracted["content"] == content - assert extracted["content_type"] == content_type - - -def test_extract_file_data_fallback_to_octet_stream(): - """Unknown file types fall back to application/octet-stream.""" - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp: - tmp.write(b"unknown content") - tmp_path = Path(tmp.name) - - try: - extracted = extract_file_data(tmp_path) - - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".xyz123") - assert ( - extracted["content_type"] == "application/octet-stream" - ), f"Expected 'application/octet-stream' for unknown type, got '{extracted['content_type']}'" - finally: - os.unlink(str(tmp_path)) - - -def test_convert_tool_response_with_pdf_file(): - """Test tool response with PDF file content using file_data field.""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with file - tool_message = { - "role": "tool", - "tool_call_id": "call_pdf_test", - "content": [ - {"type": "text", "text": '{"status": "success", "pages": 1}'}, - {"type": "file", "file_data": file_data_uri}, - ], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_pdf_test", - "function": { - "name": "analyze_document", - "arguments": '{"path": "/tmp/doc.pdf"}', - }, - } - ] - } - - # Convert tool response with nested multimodal functionResponse.parts. - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - result_part = result[0] - assert "function_response" in result_part - assert "inline_data" not in result_part - function_response = result_part["function_response"] - assert function_response["name"] == "analyze_document" - assert "response" in function_response - # Verify JSON response is parsed correctly - assert "status" in function_response["response"] - assert function_response["response"]["status"] == "success" - - # Check inline_data is nested under functionResponse.parts. - assert "parts" in function_response - assert len(function_response["parts"]) == 1 - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "application/pdf" - assert inline_data["data"] == test_pdf_base64 - - -def test_convert_tool_response_with_input_file_type(): - """Test tool response with input_file content type (Responses API format).""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with input_file type - tool_message = { - "role": "tool", - "tool_call_id": "call_input_file_test", - "content": [{"type": "input_file", "file_data": file_data_uri}], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_input_file_test", - "function": {"name": "read_file", "arguments": "{}"}, - } - ] - } - - # Convert tool response - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Check inline_data is nested under functionResponse.parts. - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - function_response = result[0]["function_response"] - assert ( - function_response["parts"][0]["inline_data"]["mime_type"] == "application/pdf" - ) - - -def test_convert_tool_response_with_nested_file_object(): - """Test tool response with file content using nested file object format.""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with nested file object (OpenAI Agents SDK format) - tool_message = { - "role": "tool", - "tool_call_id": "call_nested_test", - "content": [{"type": "file", "file": {"file_data": file_data_uri}}], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_nested_test", - "function": {"name": "process_document", "arguments": "{}"}, - } - ] - } - - # Convert tool response - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Check inline_data is nested under functionResponse.parts. - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - function_response = result[0]["function_response"] - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "application/pdf" - assert inline_data["data"] == test_pdf_base64 - - -def test_assistant_message_with_images_field(): - """ - Test that assistant messages with images field are properly converted to Gemini format. - - This handles the case where an assistant message contains generated images in the - `images` field (e.g., from image generation models like gemini-2.5-flash-image). - The images should be converted to inline_data parts in the Gemini format. - """ - # Create a small test image (1x1 red pixel PNG) - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create messages with assistant message containing images field - messages = [ - { - "role": "user", - "content": "Generate an image of a banana wearing a costume that says LiteLLM", - }, - { - "role": "assistant", - "content": "Here's your banana in a LiteLLM costume!", - "images": [ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - }, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify structure - assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}" - - # Verify user message - assert contents[0]["role"] == "user" - assert len(contents[0]["parts"]) == 1 - assert ( - contents[0]["parts"][0]["text"] - == "Generate an image of a banana wearing a costume that says LiteLLM" - ) - - # Verify assistant message - assert contents[1]["role"] == "model" - assert ( - len(contents[1]["parts"]) == 2 - ), f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}" - - # Find text part and inline_data part - text_part = None - inline_data_part = None - for part in contents[1]["parts"]: - if "text" in part: - text_part = part - elif "inline_data" in part: - inline_data_part = part - - # Verify text part - assert text_part is not None, "Missing text part in assistant message" - assert text_part["text"] == "Here's your banana in a LiteLLM costume!" - - # Verify inline_data part (image) - assert inline_data_part is not None, "Missing inline_data part in assistant message" - inline_data: BlobType = inline_data_part["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "image/png" - assert inline_data["data"] == test_image_base64 - - -def test_assistant_message_with_multiple_images(): - """Test that assistant messages with multiple images are properly converted.""" - # Create two test images - test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" - image1_data_uri = f"data:image/png;base64,{test_image1_base64}" - image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}" - - messages = [ - {"role": "user", "content": "Generate two images"}, - { - "role": "assistant", - "content": "Here are your images:", - "images": [ - { - "image_url": {"url": image1_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - }, - { - "image_url": {"url": image2_data_uri, "detail": "high"}, - "index": 1, - "type": "image_url", - }, - ], - }, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify assistant message has 3 parts (1 text + 2 images) - assert contents[1]["role"] == "model" - assert ( - len(contents[1]["parts"]) == 3 - ), f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}" - - # Count inline_data parts - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert ( - len(inline_data_parts) == 2 - ), f"Expected 2 inline_data parts, got {len(inline_data_parts)}" - - # Verify first image - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64 - - # Verify second image - assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg" - assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64 - - -def test_assistant_message_with_images_using_message_object(): - """Test that Message objects with images field are properly converted.""" - # Create a small test image - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create messages using Message object (as returned by LiteLLM) - user_message = {"role": "user", "content": "Generate an image"} - - assistant_message = Message( - content="Here's your image!", - role="assistant", - tool_calls=None, - function_call=None, - images=[ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - ) - - messages = [user_message, assistant_message] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify assistant message has both text and image - assert contents[1]["role"] == "model" - assert len(contents[1]["parts"]) == 2 - - # Verify image was converted - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert len(inline_data_parts) == 1 - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64 - - -def test_assistant_message_with_images_in_conversation_history(): - """ - Test multi-turn conversation where assistant message with images is in history. - - This simulates the real use case where: - 1. User asks for image generation - 2. Assistant generates image (with images field) - 3. User asks follow-up question about the image - """ - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - messages = [ - {"role": "user", "content": "Generate an image of a cat"}, - { - "role": "assistant", - "content": "Here's a cat image:", - "images": [ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - }, - {"role": "user", "content": "Can you make it more colorful?"}, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify structure: user -> model (with image) -> user - assert len(contents) == 3 - assert contents[0]["role"] == "user" - assert contents[1]["role"] == "model" - assert contents[2]["role"] == "user" - - # Verify assistant message has image in history - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert len(inline_data_parts) == 1 - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - - -def test_function_response_has_user_role(): - """ - Test that function response ContentType blocks include role="user". - - Gemini API only accepts two roles: "user" and "model". Function responses - must be sent with role="user". Previously, LiteLLM omitted the role field - entirely, causing 400 errors from the Gemini API. - - Fixes: https://github.com/BerriAI/litellm/issues/22003 - Fixes: https://github.com/BerriAI/litellm/issues/20690 - """ - messages = [ - {"role": "user", "content": "What is the weather in Berlin?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Berlin"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_abc123", - "content": '{"temperature": "15°C", "condition": "Cloudy"}', - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Expect: user -> model (functionCall) -> user (functionResponse) - assert len(contents) == 3 - - assert contents[0]["role"] == "user" - assert contents[1]["role"] == "model" - assert "function_call" in contents[1]["parts"][0] - - # The critical assertion: function response must have role="user" - assert contents[2]["role"] == "user" - assert "function_response" in contents[2]["parts"][0] - - -def test_multi_turn_function_calling_roles(): - """ - Test a full multi-turn function calling conversation produces correct roles. - - Simulates: user asks → model calls tool → tool responds → model answers → user asks again. - Every content block must have an explicit role of "user" or "model". - - Fixes: https://github.com/BerriAI/litellm/issues/22003 - """ - messages = [ - {"role": "user", "content": "What is the weather in Berlin?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_001", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Berlin"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_001", - "content": '{"temperature": "15°C"}', - }, - { - "role": "assistant", - "content": "The weather in Berlin is 15°C.", - }, - {"role": "user", "content": "And in Paris?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_002", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Paris"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_002", - "content": '{"temperature": "18°C"}', - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Every content block must have a valid role - for i, content in enumerate(contents): - assert "role" in content, f"Content block {i} missing 'role' field" - assert content["role"] in ( - "user", - "model", - ), f"Content block {i} has invalid role: {content.get('role')}" - - # Verify the function response blocks specifically have role="user" - for i, content in enumerate(contents): - for part in content["parts"]: - if "function_response" in part: - assert ( - content["role"] == "user" - ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" - - -def test_gemini_thought_signature_preservation_real_response(): - """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - real_candidate = { - "content": { - "parts": [ - { - "text": "I will explain and then list files.", - "thoughtSignature": "mock_signature_from_text_part", - }, - { - "functionCall": { - "name": "list_files", - "args": {}, - } - }, - ] - } - } - - parts = real_candidate["content"]["parts"] - - content, reasoning_content = ( - VertexGeminiConfig().get_assistant_content_message(parts=parts) - ) - thought_signatures = ( - VertexGeminiConfig()._extract_thought_signatures_from_parts( - parts=parts - ) - ) - functions, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - msg: dict = {"role": "assistant"} - if content is not None: - msg["content"] = content - if tools: - msg["tool_calls"] = tools - if functions is not None: - msg["function_call"] = functions - if thought_signatures is not None: - msg["provider_specific_fields"] = { - "thought_signatures": thought_signatures - } - - converted_real = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted_real) == 1 - assert "parts" in converted_real[0] - parts_out = converted_real[0]["parts"] - assert len(parts_out) == 2 - assert "text" in parts_out[0] - assert ( - parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" - ) - assert "function_call" in parts_out[1] - assert "thoughtSignature" not in parts_out[1] - - -def test_gemini_thought_signature_deduplication_assumed_response(): - """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - pr_assumed_msg = { - "role": "assistant", - "content": "I will list the directory.", - "provider_specific_fields": { - "thought_signatures": ["mock_signature_63k"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": { - "thought_signature": "mock_signature_63k" - }, - } - ], - } - - converted_pr = _gemini_convert_messages_with_history( - messages=[pr_assumed_msg], - model="gemini-2.5-pro", - ) - - assert len(converted_pr) == 1 - assert "parts" in converted_pr[0] - parts_out = converted_pr[0]["parts"] - assert len(parts_out) == 2 - assert "text" in parts_out[0] - assert "thoughtSignature" not in parts_out[0] - assert "function_call" in parts_out[1] - assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" - - -def test_gemini_thought_signature_pure_text(): - """Test that thought signatures are preserved on the text part for responses with no tool calls.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Hello, I am a model.", - "provider_specific_fields": { - "thought_signatures": ["pure_text_signature"] - }, - } - - converted = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted) == 1 - assert "parts" in converted[0] - parts_out = converted[0]["parts"] - assert len(parts_out) == 1 - assert "text" in parts_out[0] - assert parts_out[0]["thoughtSignature"] == "pure_text_signature" - - -def test_gemini_thought_signature_pure_tool_call(): - """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": None, - "provider_specific_fields": { - "thought_signatures": ["pure_tool_signature"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": { - "thought_signature": "pure_tool_signature" - }, - } - ], - } - - converted = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted) == 1 - assert "parts" in converted[0] - parts_out = converted[0]["parts"] - assert len(parts_out) == 1 - assert "function_call" in parts_out[0] - assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" - - -def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): - """A text-part signature that differs from the tool-call signature must stay on the text part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Some analysis.", - "provider_specific_fields": { - "thought_signatures": ["text_signature", "tool_signature"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": {"thought_signature": "tool_signature"}, - } - ], - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-2.5-pro" - )[0]["parts"] - - assert parts[0]["text"] == "Some analysis." - assert parts[0]["thoughtSignature"] == "text_signature" - assert "function_call" in parts[1] - assert parts[1]["thoughtSignature"] == "tool_signature" - - -def test_gemini_25_text_signature_survives_replay_to_gemini_3(): - """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real - text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, - ) - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "I will list the directory.", - "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - } - ], - } - - parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ - 0 - ]["parts"] - - assert parts[0]["text"] == "I will list the directory." - assert parts[0]["thoughtSignature"] == "real_25_signature" - assert "function_call" in parts[1] - assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() - - -def test_gemini_function_call_signature_round_trip_no_duplicate(): - """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and - re-serialized sends the signature exactly once, on the function-call part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - response_parts = [ - {"text": "I will calculate the result for you."}, - { - "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, - "thoughtSignature": "signature_from_function_call", - }, - ] - - config = VertexGeminiConfig() - content, _ = config.get_assistant_content_message(parts=response_parts) - thought_signatures = config._extract_thought_signatures_from_parts( - parts=response_parts - ) - _, tools, _ = VertexGeminiConfig._transform_parts( - parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False - ) - - msg = { - "role": "assistant", - "content": content, - "tool_calls": tools, - "provider_specific_fields": {"thought_signatures": thought_signatures}, - } - - parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ - 0 - ]["parts"] - - signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] - assert signatures == ["signature_from_function_call"] - assert "thoughtSignature" not in parts[0] - assert "function_call" in parts[1] - - -def test_gemini_server_side_tool_signature_not_duplicated_on_text(): - """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "The weather in Buenos Aires is sunny.", - "provider_specific_fields": { - "thought_signatures": ["server_side_signature"], - "server_side_tool_invocations": [ - { - "tool_type": "GOOGLE_SEARCH_WEB", - "id": "abc123", - "args": {"queries": ["weather Buenos Aires"]}, - "response": {"weather": "Sunny"}, - "thought_signature": "server_side_signature", - } - ], - }, - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-2.5-pro" - )[0]["parts"] - - text_part = next(p for p in parts if "text" in p) - assert "thoughtSignature" not in text_part - tool_call_part = next(p for p in parts if "toolCall" in p) - assert tool_call_part["thoughtSignature"] == "server_side_signature" diff --git a/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py b/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py deleted file mode 100644 index 50135ba1f92..00000000000 --- a/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Vertex AI Image Edit Tests diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 54607cc5284..aeba9f0fa3c 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -1,13 +1,9 @@ import os -from unittest.mock import MagicMock, patch +from unittest.mock import patch -import httpx import pytest -from litellm.llms.vertex_ai.image_generation import ( - get_vertex_ai_image_generation_config, -) from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( VertexAIGeminiImageGenerationConfig, ) @@ -16,588 +12,6 @@ from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ) -class TestVertexAIGeminiImageGenerationConfig: - def setup_method(self): - """Set up test fixtures""" - self.config = VertexAIGeminiImageGenerationConfig() - - def test_get_supported_openai_params(self): - """Test get_supported_openai_params returns correct params""" - supported = self.config.get_supported_openai_params("gemini-2.5-flash-image") - assert "n" in supported - assert "size" in supported - - def test_map_openai_params_n(self): - """Test mapping n parameter to candidate_count""" - non_default_params = {"n": 3} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("candidate_count") == 3 - - def test_map_openai_params_size(self): - """Test mapping size parameter to aspectRatio""" - non_default_params = {"size": "1024x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("aspectRatio") == "1:1" - - def test_map_openai_params_size_16_9(self): - """Test mapping 16:9 size""" - non_default_params = {"size": "1792x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("aspectRatio") == "16:9" - - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" - assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" - assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default - - def test_get_supported_openai_params_includes_native_gemini_params(self): - """Test that native Gemini imageConfig params are supported""" - supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") - assert "aspectRatio" in supported - assert "aspect_ratio" in supported - assert "imageSize" in supported - assert "image_size" in supported - assert "imageConfig" in supported - - def test_map_openai_params_aspect_ratio_camel_case(self): - """Test mapping native aspectRatio parameter""" - result = self.config.map_openai_params({"aspectRatio": "9:16"}, {}, "gemini-3-pro-image-preview", False) - assert result["aspectRatio"] == "9:16" - - def test_map_openai_params_aspect_ratio_snake_case(self): - """Test mapping native aspect_ratio parameter""" - result = self.config.map_openai_params({"aspect_ratio": "16:9"}, {}, "gemini-3-pro-image-preview", False) - assert result["aspectRatio"] == "16:9" - - def test_map_openai_params_image_size_camel_case(self): - """Test mapping native imageSize parameter""" - result = self.config.map_openai_params({"imageSize": "4K"}, {}, "gemini-3-pro-image-preview", False) - assert result["imageSize"] == "4K" - - def test_map_openai_params_image_size_snake_case(self): - """Test mapping native image_size parameter""" - result = self.config.map_openai_params({"image_size": "2K"}, {}, "gemini-3-pro-image-preview", False) - assert result["imageSize"] == "2K" - - def test_map_openai_params_image_config_dict_stored_whole(self): - """imageConfig dict is stored as-is so all fields survive""" - result = self.config.map_openai_params( - {"imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}}, - {}, - "gemini-3.1-flash-image", - False, - ) - assert result["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "2K"} - - def test_map_openai_params_image_config_all_fields(self): - """All ImageConfig fields (personGeneration, imageOutputOptions) pass through""" - payload = { - "imageConfig": { - "aspectRatio": "9:16", - "imageSize": "4K", - "personGeneration": "DONT_ALLOW", - "imageOutputOptions": { - "mimeType": "image/jpeg", - "compressionQuality": 80, - }, - } - } - result = self.config.map_openai_params(payload, {}, "gemini-3.1-flash-image", False) - assert result["imageConfig"] == payload["imageConfig"] - - def test_map_openai_params_image_config_non_dict_warns_and_drops(self): - """Non-dict imageConfig is dropped with a warning, not silently discarded""" - with patch("litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation.verbose_logger") as mock_log: - result = self.config.map_openai_params( - {"imageConfig": "bad-string-value"}, {}, "gemini-3.1-flash-image", False - ) - assert "imageConfig" not in result - mock_log.warning.assert_called_once() - - def test_transform_image_generation_request_from_image_config(self): - """Full imageConfig dict is forwarded verbatim into generationConfig""" - full_config = { - "aspectRatio": "16:9", - "imageSize": "2K", - "personGeneration": "DONT_ALLOW", - "imageOutputOptions": {"mimeType": "image/jpeg", "compressionQuality": 85}, - } - mapped = self.config.map_openai_params( - {"imageConfig": full_config}, - {}, - "gemini-3.1-flash-image", - False, - ) - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image", - prompt="A nano banana on a desk", - optional_params=mapped, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"] == full_config - - def test_transform_image_generation_flat_params_override_image_config(self): - """Explicit flat params win over the same key inside imageConfig""" - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image", - prompt="A nano banana", - optional_params={ - "imageConfig": {"aspectRatio": "1:1", "personGeneration": "DONT_ALLOW"}, - "aspectRatio": "16:9", # should win - }, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert request["generationConfig"]["imageConfig"]["personGeneration"] == "DONT_ALLOW" - - def test_transform_image_generation_request_basic(self): - """Test basic request transformation""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={}, - litellm_params={}, - headers={}, - ) - assert "contents" in request - assert "generationConfig" in request - assert request["generationConfig"]["responseModalities"] == ["IMAGE"] - assert request["contents"][0]["parts"][0]["text"] == "A nano banana" - - def test_transform_image_generation_request_with_aspect_ratio(self): - """Test request transformation with aspectRatio""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"aspectRatio": "16:9"}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - - def test_transform_image_generation_request_with_image_size(self): - """Test request transformation with imageSize (Gemini 3 Pro)""" - request = self.config.transform_image_generation_request( - model="gemini-3-pro-image-preview", - prompt="A nano banana", - optional_params={"imageSize": "4K"}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K" - - def test_map_openai_params_web_search_options(self): - """Test web_search_options maps to googleSearch tool""" - result = self.config.map_openai_params({"web_search_options": {}}, {}, "gemini-3.1-flash-image-preview", False) - assert result["tools"] == [{"googleSearch": {}}] - - def test_transform_image_generation_request_with_web_search_tools(self): - """Test request transformation includes googleSearch tools""" - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image-preview", - prompt="Generate an image of the latest iPhone", - optional_params={"tools": [{"googleSearch": {}}]}, - litellm_params={}, - headers={}, - ) - assert request["tools"] == [{"googleSearch": {}}] - - def test_transform_image_generation_request_forwards_tool_config(self): - """Test request transformation forwards toolConfig side-effects from tool mapping""" - mapped = self.config.map_openai_params( - {"tools": [{"googleMaps": {"latitude": 37.7, "longitude": -122.4}}]}, - {}, - "gemini-3.1-flash-image-preview", - False, - ) - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image-preview", - prompt="Generate an image of a coffee shop nearby", - optional_params=mapped, - litellm_params={}, - headers={}, - ) - assert request["tools"] == [{"googleMaps": {}}] - assert request["toolConfig"] == {"retrievalConfig": {"latLng": {"latitude": 37.7, "longitude": -122.4}}} - - def test_transform_image_generation_request_with_candidate_count(self): - """Test request transformation with candidate_count""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"candidate_count": 2}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["candidateCount"] == 2 - - def test_transform_image_generation_request_with_n(self): - """Test request transformation with n parameter""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"n": 2}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["candidateCount"] == 2 - - def test_transform_image_generation_response(self): - """Test response transformation""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - } - } - ] - } - } - ], - "usageMetadata": { - "promptTokenCount": 93, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 54, - }, - { - "modality": "IMAGE", - "tokenCount": 39, - }, - ], - "candidatesTokenCount": 17, - "totalTokenCount": 110, - }, - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].url is None - assert result.usage.input_tokens == 93 - assert result.usage.input_tokens_details.text_tokens == 54 - assert result.usage.input_tokens_details.image_tokens == 39 - assert result.usage.output_tokens == 17 - assert result.usage.total_tokens == 110 - - def test_transform_image_generation_response_multiple_images(self): - """Test response transformation with multiple images""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "image1", - } - }, - { - "inlineData": { - "mimeType": "image/png", - "data": "image2", - } - }, - ] - } - } - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 2 - assert result.data[0].b64_json == "image1" - assert result.data[1].b64_json == "image2" - - def test_transform_image_generation_response_signature(self): - """Test response transformation includes thoughtSignature for Gemini 3 Pro""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - }, - "thoughtSignature": "test_signature_abc123", - } - ] - } - } - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-3-pro-image-preview", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" - - def test_transform_image_generation_response_tracks_web_search_requests(self): - """Grounding queries are carried onto usage so search spend can be billed""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - } - } - ] - }, - "groundingMetadata": {"webSearchQueries": ["eiffel tower", "paris skyline"]}, - } - ], - "usageMetadata": { - "promptTokenCount": 93, - "candidatesTokenCount": 17, - "totalTokenCount": 110, - }, - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=ImageResponse(), - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert result.usage.web_search_requests == 2 - - -class TestVertexAIImagenImageGenerationConfig: - def setup_method(self): - """Set up test fixtures""" - self.config = VertexAIImagenImageGenerationConfig() - - def test_get_supported_openai_params(self): - """Test get_supported_openai_params returns correct params""" - supported = self.config.get_supported_openai_params("imagegeneration@006") - assert "n" in supported - assert "size" in supported - - def test_map_openai_params_n(self): - """Test mapping n parameter to sampleCount""" - non_default_params = {"n": 3} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) - assert result.get("sampleCount") == 3 - - def test_map_openai_params_size(self): - """Test mapping size parameter to aspectRatio""" - non_default_params = {"size": "1024x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) - assert result.get("aspectRatio") == "1:1" - - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default - - def test_transform_image_generation_request_basic(self): - """Test basic request transformation""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={}, - litellm_params={}, - headers={}, - ) - assert "instances" in request - assert "parameters" in request - assert request["instances"][0]["prompt"] == "A cat" - assert request["parameters"]["sampleCount"] == 1 - - def test_transform_image_generation_request_with_params(self): - """Test request transformation with parameters""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={"sampleCount": 2, "aspectRatio": "16:9"}, - litellm_params={}, - headers={}, - ) - assert request["parameters"]["sampleCount"] == 2 - assert request["parameters"]["aspectRatio"] == "16:9" - - def test_transform_image_generation_request_labels_from_metadata(self): - """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={}, - litellm_params={"metadata": {"requester_metadata": {"team": "platform", "env": "prod"}}}, - headers={}, - ) - assert request["labels"] == {"team": "platform", "env": "prod"} - assert "labels" not in request["parameters"] - - def test_transform_image_generation_response(self): - """Test response transformation""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "base64_encoded_image_data"}]} - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="imagegeneration@006", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].url is None - - def test_transform_image_generation_response_multiple_images(self): - """Test response transformation with multiple images""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "predictions": [ - {"bytesBase64Encoded": "image1"}, - {"bytesBase64Encoded": "image2"}, - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="imagegeneration@006", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 2 - assert result.data[0].b64_json == "image1" - assert result.data[1].b64_json == "image2" - - -class TestGetVertexAIImageGenerationConfig: - """Test the router function that selects the correct config""" - - def test_get_gemini_model_config(self): - """Test that Gemini models return Gemini config""" - config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("vertex_ai/gemini-2.5-flash-image") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - def test_get_imagen_model_config(self): - """Test that Imagen models return Imagen config""" - config = get_vertex_ai_image_generation_config("imagegeneration@006") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("vertex_ai/imagegeneration@006") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - def test_get_non_gemini_model_config(self): - """Test that non-Gemini models default to Imagen config""" - config = get_vertex_ai_image_generation_config("some-other-model") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - class TestVertexAIImageGenerationIntegration: """Integration tests for Vertex AI image generation""" @@ -642,39 +56,3 @@ class TestVertexAIImageGenerationIntegration: litellm_params={}, ) assert "Authorization" in headers - - def test_gemini_get_complete_url(self): - """Test Gemini config URL generation""" - config = VertexAIGeminiImageGenerationConfig() - url = config.get_complete_url( - api_base=None, - api_key=None, - model="gemini-2.5-flash-image", - optional_params={}, - litellm_params={ - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - ) - assert "test-project" in url - assert "us-central1" in url - assert "gemini-2.5-flash-image" in url - assert "generateContent" in url - - def test_imagen_get_complete_url(self): - """Test Imagen config URL generation""" - config = VertexAIImagenImageGenerationConfig() - url = config.get_complete_url( - api_base=None, - api_key=None, - model="imagegeneration@006", - optional_params={}, - litellm_params={ - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - ) - assert "test-project" in url - assert "us-central1" in url - assert "imagegeneration@006" in url - assert "predict" in url diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py deleted file mode 100644 index 8b41c5ab3f8..00000000000 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for Vertex AI Gemma-AI models""" diff --git a/tests/test_litellm/llms/vertex_ai/videos/__init__.py b/tests/test_litellm/llms/vertex_ai/videos/__init__.py deleted file mode 100644 index f29c2a16fd5..00000000000 --- a/tests/test_litellm/llms/vertex_ai/videos/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Tests for Vertex AI video generation. -""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 1d3d7a452b6..45ad336368c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -30,7 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index bbc8fd539a3..5169d4c9ec6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -23,7 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 353ffadfa46..227921d6150 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,7 +24,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 653b2c9914a..ecea4723bf4 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,11 +1,14 @@ import asyncio +import base64 import importlib import os from collections.abc import Coroutine, Iterator +from dataclasses import dataclass, field from pathlib import Path from typing import Final import boto3 +import httpx import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -15,9 +18,12 @@ import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency +from litellm.anthropic_beta_headers_manager import reload_beta_headers_config # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory_module # noqa: E402 # same import-time dependency from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency image_handling as image_handling_module, ) +from litellm.llms.gemini.chat import transformation as gemini_chat_transformation_module # noqa: E402 # same import-time dependency from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency close_litellm_async_clients, ) @@ -89,6 +95,9 @@ RESTORED_GLOBALS: Final = ( ) MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient") SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client") +ONE_PIXEL_PNG: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) def _allow_loopback_only() -> None: @@ -236,6 +245,47 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def local_beta_headers_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + reload_beta_headers_config() + yield + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + reload_beta_headers_config() + + +@dataclass(slots=True) +class AsyncOnlyImageFetch: + fetched: list[str] = field(default_factory=list) # mutable-ok: tests assert on the URLs fetched, in order + base64_png: str = base64.b64encode(ONE_PIXEL_PNG).decode() + data_url: str = "data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode() + + +@pytest.fixture +def async_only_image_fetch(monkeypatch: pytest.MonkeyPatch) -> AsyncOnlyImageFetch: + fetch: Final = AsyncOnlyImageFetch() + + def forbid_sync_fetch(client: object, url: str, **kwargs: object) -> httpx.Response: + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client: object, url: str, **kwargs: object) -> httpx.Response: + fetch.fetched.append(url) + return httpx.Response( + 200, content=ONE_PIXEL_PNG, headers={"content-type": "image/png"}, request=httpx.Request("GET", url) + ) + + def forbid_sync_convert(url: str, *args: object, **kwargs: object) -> str: + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling_module, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling_module, "async_safe_get", serve_png) + for module in (image_handling_module, prompt_factory_module, gemini_chat_transformation_module): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + @pytest.fixture def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS: diff --git a/tests/test_litellm/llms/anthropic/__init__.py b/tests/unit/expected_fine_tuning_api/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/__init__.py rename to tests/unit/expected_fine_tuning_api/__init__.py diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json b/tests/unit/expected_fine_tuning_api/azure_cancel_expected_output.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_expected_output.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_cancel_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json b/tests/unit/expected_fine_tuning_api/azure_cancel_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_request.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json b/tests/unit/expected_fine_tuning_api/azure_create_expected_output.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json rename to tests/unit/expected_fine_tuning_api/azure_create_expected_output.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_create_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_create_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json b/tests/unit/expected_fine_tuning_api/azure_create_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_request.json rename to tests/unit/expected_fine_tuning_api/azure_create_request.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_list_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_list_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json b/tests/unit/expected_fine_tuning_api/azure_list_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_list_request.json rename to tests/unit/expected_fine_tuning_api/azure_list_request.json diff --git a/tests/test_litellm/llms/anthropic/batches/__init__.py b/tests/unit/llms/aiml/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/__init__.py rename to tests/unit/llms/aiml/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/tests/unit/llms/aiml/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py rename to tests/unit/llms/aiml/image_generation/__init__.py diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py rename to tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py diff --git a/tests/unit/llms/anthropic/batches/test_transformation.py b/tests/unit/llms/anthropic/batches/test_transformation.py index eacd2c9d03b..419fc7740eb 100644 --- a/tests/unit/llms/anthropic/batches/test_transformation.py +++ b/tests/unit/llms/anthropic/batches/test_transformation.py @@ -616,7 +616,7 @@ def test_transform_response_reraises_unexpected_error(config): # automatically. See base_batches_config_test.py. # --------------------------------------------------------------------------- # -from tests.test_litellm.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 +from tests.unit.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 BatchesConfigContractTests, ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py b/tests/unit/llms/anthropic/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py rename to tests/unit/llms/anthropic/chat/__init__.py diff --git a/tests/test_litellm/llms/anthropic/chat/conftest.py b/tests/unit/llms/anthropic/chat/conftest.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/conftest.py rename to tests/unit/llms/anthropic/chat/conftest.py diff --git a/tests/test_litellm/llms/anthropic/files/__init__.py b/tests/unit/llms/anthropic/chat/guardrail_translation/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/files/__init__.py rename to tests/unit/llms/anthropic/chat/guardrail_translation/__init__.py diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py rename to tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py rename to tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py rename to tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py b/tests/unit/llms/anthropic/chat/test_code_interpreter_results_extraction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py rename to tests/unit/llms/anthropic/chat/test_code_interpreter_results_extraction.py diff --git a/tests/test_litellm/llms/azure/batches/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/batches/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py diff --git a/tests/test_litellm/llms/azure/vector_stores/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/vector_stores/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py diff --git a/tests/test_litellm/llms/base_llm/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py diff --git a/tests/test_litellm/llms/base_llm/batches/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/batches/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py rename to tests/unit/llms/anthropic/test_anthropic_common_utils.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py rename to tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/unit/llms/anthropic/test_anthropic_files_and_batches.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py rename to tests/unit/llms/anthropic/test_anthropic_files_and_batches.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py rename to tests/unit/llms/anthropic/test_anthropic_output_format_filter.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py rename to tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py rename to tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py rename to tests/unit/llms/anthropic/test_anthropic_schema_filter.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py b/tests/unit/llms/anthropic/test_anthropic_structured_output.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py rename to tests/unit/llms/anthropic/test_anthropic_structured_output.py diff --git a/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py b/tests/unit/llms/anthropic/test_azure_ai_cache_pricing.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py rename to tests/unit/llms/anthropic/test_azure_ai_cache_pricing.py diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/unit/llms/anthropic/test_cost_calculation_dict_safety.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py rename to tests/unit/llms/anthropic/test_cost_calculation_dict_safety.py diff --git a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py rename to tests/unit/llms/anthropic/test_count_tokens_oauth.py diff --git a/tests/test_litellm/llms/anthropic/test_message_sanitization.py b/tests/unit/llms/anthropic/test_message_sanitization.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_message_sanitization.py rename to tests/unit/llms/anthropic/test_message_sanitization.py diff --git a/tests/test_litellm/llms/base_llm/files/__init__.py b/tests/unit/llms/azure/batches/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/__init__.py rename to tests/unit/llms/azure/batches/__init__.py diff --git a/tests/test_litellm/llms/azure/batches/test_handler.py b/tests/unit/llms/azure/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/azure/batches/test_handler.py rename to tests/unit/llms/azure/batches/test_handler.py diff --git a/tests/test_litellm/llms/base_llm/realtime/__init__.py b/tests/unit/llms/azure/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/realtime/__init__.py rename to tests/unit/llms/azure/chat/__init__.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py rename to tests/unit/llms/azure/chat/test_azure_base_model_routing.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py rename to tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py rename to tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py rename to tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/test_handler.py b/tests/unit/llms/azure/realtime/test_handler.py similarity index 100% rename from tests/test_litellm/llms/azure/realtime/test_handler.py rename to tests/unit/llms/azure/realtime/test_handler.py diff --git a/tests/test_litellm/llms/azure/test_audio_transcriptions.py b/tests/unit/llms/azure/test_audio_transcriptions.py similarity index 100% rename from tests/test_litellm/llms/azure/test_audio_transcriptions.py rename to tests/unit/llms/azure/test_audio_transcriptions.py diff --git a/tests/test_litellm/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure.py rename to tests/unit/llms/azure/test_azure.py diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_common_utils.py rename to tests/unit/llms/azure/test_azure_common_utils.py diff --git a/tests/test_litellm/llms/azure/test_azure_cost_calculation.py b/tests/unit/llms/azure/test_azure_cost_calculation.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_cost_calculation.py rename to tests/unit/llms/azure/test_azure_cost_calculation.py diff --git a/tests/test_litellm/llms/azure/test_azure_embedding.py b/tests/unit/llms/azure/test_azure_embedding.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_embedding.py rename to tests/unit/llms/azure/test_azure_embedding.py diff --git a/tests/test_litellm/llms/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_exception_mapping.py rename to tests/unit/llms/azure/test_azure_exception_mapping.py diff --git a/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py b/tests/unit/llms/azure/test_azure_fine_tuning_api.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py rename to tests/unit/llms/azure/test_azure_fine_tuning_api.py diff --git a/tests/test_litellm/llms/azure/test_azure_speech_audio_transcription.py b/tests/unit/llms/azure/test_azure_speech_audio_transcription.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_speech_audio_transcription.py rename to tests/unit/llms/azure/test_azure_speech_audio_transcription.py diff --git a/tests/test_litellm/llms/bedrock/__init__.py b/tests/unit/llms/azure/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/__init__.py rename to tests/unit/llms/azure/videos/__init__.py diff --git a/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py b/tests/unit/llms/azure/videos/test_azure_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py rename to tests/unit/llms/azure/videos/test_azure_video_transformation.py diff --git a/tests/test_litellm/llms/bedrock/batches/__init__.py b/tests/unit/llms/azure_ai/claude/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/__init__.py rename to tests/unit/llms/azure_ai/claude/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_handler.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_handler.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py b/tests/unit/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py rename to tests/unit/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/__init__.py b/tests/unit/llms/azure_ai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/agentcore/__init__.py rename to tests/unit/llms/azure_ai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py rename to tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py rename to tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py b/tests/unit/llms/azure_ai/test_azure_ai_agents_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py rename to tests/unit/llms/azure_ai/test_azure_ai_agents_handler.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/unit/llms/azure_ai/test_azure_ai_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py rename to tests/unit/llms/azure_ai/test_azure_ai_cost_calculator.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py rename to tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py diff --git a/tests/test_litellm/llms/base_llm/batches/base_batches_config_test.py b/tests/unit/llms/base_llm/batches/base_batches_config_test.py similarity index 100% rename from tests/test_litellm/llms/base_llm/batches/base_batches_config_test.py rename to tests/unit/llms/base_llm/batches/base_batches_config_test.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/__init__.py b/tests/unit/llms/base_llm/files/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/__init__.py rename to tests/unit/llms/base_llm/files/__init__.py diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/unit/llms/base_llm/files/test_azure_blob_storage_backend.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py rename to tests/unit/llms/base_llm/files/test_azure_blob_storage_backend.py diff --git a/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py b/tests/unit/llms/base_llm/files/test_litellm_db_storage_backend.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py rename to tests/unit/llms/base_llm/files/test_litellm_db_storage_backend.py diff --git a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py b/tests/unit/llms/base_llm/files/test_storage_backend_factory.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py rename to tests/unit/llms/base_llm/files/test_storage_backend_factory.py diff --git a/tests/test_litellm/llms/black_forest_labs/__init__.py b/tests/unit/llms/base_llm/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/__init__.py rename to tests/unit/llms/base_llm/responses/__init__.py diff --git a/tests/test_litellm/llms/base_llm/responses/test_codex_compat.py b/tests/unit/llms/base_llm/responses/test_codex_compat.py similarity index 100% rename from tests/test_litellm/llms/base_llm/responses/test_codex_compat.py rename to tests/unit/llms/base_llm/responses/test_codex_compat.py diff --git a/tests/test_litellm/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/responses/test_transformation.py rename to tests/unit/llms/base_llm/responses/test_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py b/tests/unit/llms/base_llm/search/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py rename to tests/unit/llms/base_llm/search/__init__.py diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/unit/llms/base_llm/search/test_base_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py rename to tests/unit/llms/base_llm/search/test_base_search_transformation.py diff --git a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py b/tests/unit/llms/base_llm/test_base_managed_resource.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_base_managed_resource.py rename to tests/unit/llms/base_llm/test_base_managed_resource.py diff --git a/tests/test_litellm/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_base_model_iterator.py rename to tests/unit/llms/base_llm/test_base_model_iterator.py diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/unit/llms/base_llm/test_managed_resource_isolation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py rename to tests/unit/llms/base_llm/test_managed_resource_isolation.py diff --git a/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py b/tests/unit/llms/base_llm/test_managed_resources_utils.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_managed_resources_utils.py rename to tests/unit/llms/base_llm/test_managed_resources_utils.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py b/tests/unit/llms/bedrock/batches/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py rename to tests/unit/llms/bedrock/batches/__init__.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py b/tests/unit/llms/bedrock/batches/test_batch_metadata_sanitization.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py rename to tests/unit/llms/bedrock/batches/test_batch_metadata_sanitization.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/test_handler.py rename to tests/unit/llms/bedrock/batches/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/batches/test_transformation.py rename to tests/unit/llms/bedrock/batches/test_transformation.py index 347c459a369..5e987239c54 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -878,7 +878,7 @@ def test_validate_environment_passes_headers_through(config): # Shared BaseBatchesConfig contract suite. # --------------------------------------------------------------------------- # -from tests.test_litellm.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 +from tests.unit.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 BatchesConfigContractTests, ) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py similarity index 99% rename from tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py rename to tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 67ffe7570a1..08bcac33a35 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -20,7 +20,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import configuration from litellm.types.utils import ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe RESOLVED_CREDENTIALS = Credentials( access_key="AKIARESOLVED", diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py rename to tests/unit/llms/bedrock/chat/test_converse_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py rename to tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py rename to tests/unit/llms/bedrock/chat/test_invoke_handler.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_mistral_config.py b/tests/unit/llms/bedrock/chat/test_mistral_config.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_mistral_config.py rename to tests/unit/llms/bedrock/chat/test_mistral_config.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_service_tier.py b/tests/unit/llms/bedrock/chat/test_service_tier.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_service_tier.py rename to tests/unit/llms/bedrock/chat/test_service_tier.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py b/tests/unit/llms/bedrock/chat/test_streaming_choice_index.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py rename to tests/unit/llms/bedrock/chat/test_streaming_choice_index.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_writer_palmyra.py b/tests/unit/llms/bedrock/chat/test_writer_palmyra.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_writer_palmyra.py rename to tests/unit/llms/bedrock/chat/test_writer_palmyra.py diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 3622ce7f212..d67724f261d 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -7,7 +7,7 @@ from botocore.credentials import RefreshableCredentials from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe class _ProbedCountTokensHandler(BedrockCountTokensHandler): diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/unit/llms/bedrock/embed/__init__.py similarity index 100% rename from tests/test_litellm/llms/cerebras/__init__.py rename to tests/unit/llms/bedrock/embed/__init__.py diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py similarity index 99% rename from tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py rename to tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 18f4b0f6ced..fbcbd0aaea6 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -9,7 +9,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py similarity index 99% rename from tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py rename to tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index e5a460e2f1a..ad21cadaa4b 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -11,7 +11,7 @@ import litellm from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.bedrock.embed.embedding import BedrockEmbedding -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} diff --git a/tests/test_litellm/llms/bedrock/embed/test_embedding.py b/tests/unit/llms/bedrock/embed/test_embedding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/embed/test_embedding.py rename to tests/unit/llms/bedrock/embed/test_embedding.py diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/unit/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py rename to tests/unit/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/unit/llms/bedrock/event_loop_probe.py similarity index 100% rename from tests/test_litellm/llms/bedrock/event_loop_probe.py rename to tests/unit/llms/bedrock/event_loop_probe.py diff --git a/tests/test_litellm/llms/chatgpt/__init__.py b/tests/unit/llms/bedrock/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/__init__.py rename to tests/unit/llms/bedrock/messages/__init__.py diff --git a/tests/test_litellm/llms/chatgpt/chat/__init__.py b/tests/unit/llms/bedrock/messages/invoke_transformations/__init__.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/chat/__init__.py rename to tests/unit/llms/bedrock/messages/invoke_transformations/__init__.py diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/transformation.py b/tests/unit/llms/bedrock/rerank/transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/transformation.py rename to tests/unit/llms/bedrock/rerank/transformation.py diff --git a/tests/test_litellm/llms/crusoe/__init__.py b/tests/unit/llms/bedrock/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/crusoe/__init__.py rename to tests/unit/llms/bedrock/responses/__init__.py diff --git a/tests/test_litellm/llms/bedrock/responses/test_bedrock_openai_responses.py b/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py similarity index 100% rename from tests/test_litellm/llms/bedrock/responses/test_bedrock_openai_responses.py rename to tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py diff --git a/tests/test_litellm/llms/databricks/chat/__init__.py b/tests/unit/llms/bedrock/search/__init__.py similarity index 100% rename from tests/test_litellm/llms/databricks/chat/__init__.py rename to tests/unit/llms/bedrock/search/__init__.py diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py rename to tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/unit/llms/bedrock/test_anthropic_beta_support.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py rename to tests/unit/llms/bedrock/test_anthropic_beta_support.py diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py similarity index 99% rename from tests/test_litellm/llms/bedrock/test_base_aws_llm.py rename to tests/unit/llms/bedrock/test_base_aws_llm.py index 6b9450afed4..db144ab6d56 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -28,7 +28,7 @@ from litellm.llms.bedrock.base_aws_llm import ( run_aws_signing, sign_request_off_loop_if_aws, ) -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py rename to tests/unit/llms/bedrock/test_bedrock_common_utils.py diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py b/tests/unit/llms/bedrock/test_bedrock_ssl_verify.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py rename to tests/unit/llms/bedrock/test_bedrock_ssl_verify.py diff --git a/tests/test_litellm/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_claude_platform_provider.py rename to tests/unit/llms/bedrock/test_claude_platform_provider.py diff --git a/tests/test_litellm/llms/bedrock/test_converse_context_management.py b/tests/unit/llms/bedrock/test_converse_context_management.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_converse_context_management.py rename to tests/unit/llms/bedrock/test_converse_context_management.py diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/unit/llms/bedrock/test_cross_region_inference_profile_mapping.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py rename to tests/unit/llms/bedrock/test_cross_region_inference_profile_mapping.py diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/unit/llms/bedrock/test_mantle.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_mantle.py rename to tests/unit/llms/bedrock/test_mantle.py diff --git a/tests/test_litellm/llms/bedrock/test_nova_imported_models.py b/tests/unit/llms/bedrock/test_nova_imported_models.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_nova_imported_models.py rename to tests/unit/llms/bedrock/test_nova_imported_models.py diff --git a/tests/test_litellm/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_request_metadata.py rename to tests/unit/llms/bedrock/test_request_metadata.py diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py rename to tests/unit/llms/bedrock/test_web_identity_session_policy.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 0cc3963358f..4bf3dd11fa1 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -19,7 +19,7 @@ import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture diff --git a/tests/test_litellm/llms/databricks/responses/__init__.py b/tests/unit/llms/cometapi/__init__.py similarity index 100% rename from tests/test_litellm/llms/databricks/responses/__init__.py rename to tests/unit/llms/cometapi/__init__.py diff --git a/tests/test_litellm/llms/deepseek/__init__.py b/tests/unit/llms/cometapi/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/__init__.py rename to tests/unit/llms/cometapi/chat/__init__.py diff --git a/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py new file mode 100644 index 00000000000..607648dd6c9 --- /dev/null +++ b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py @@ -0,0 +1,183 @@ +""" +Unit tests for CometAPI Chat Configuration + +Tests the CometAPIChatConfig class methods using mocks +""" + + +import pytest + + +from litellm.llms.cometapi.chat.transformation import ( + CometAPIChatCompletionStreamingHandler, + CometAPIConfig, +) +from litellm.llms.cometapi.common_utils import CometAPIException + + +class TestCometAPIChatCompletionStreamingHandler: + def test_chunk_parser_successful(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test input chunk + chunk = { + "id": "test_id", + "created": 1234567890, + "model": "gpt-3.5-turbo", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "choices": [ + {"delta": {"content": "test content", "reasoning": "test reasoning"}} + ], + } + + # Parse chunk + result = handler.chunk_parser(chunk) + + # Verify response + assert result.id == "test_id" + assert result.object == "chat.completion.chunk" + assert result.created == 1234567890 + assert result.model == "gpt-3.5-turbo" + assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] + assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] + assert result.usage.total_tokens == chunk["usage"]["total_tokens"] + assert len(result.choices) == 1 + assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" + + def test_chunk_parser_error_response(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test error chunk + error_chunk = { + "error": { + "message": "test error", + "code": 400, + } + } + + # Verify error handling + with pytest.raises(CometAPIException) as exc_info: + handler.chunk_parser(error_chunk) + + assert "CometAPI Error: test error" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + def test_chunk_parser_key_error(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test invalid chunk missing required fields + invalid_chunk = {"incomplete": "data"} + + # Verify KeyError handling + with pytest.raises(CometAPIException) as exc_info: + handler.chunk_parser(invalid_chunk) + + assert "KeyError" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + +class TestCometAPIConfig: + def test_transform_request_basic(self): + """Test basic request transformation""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["model"] == "cometapi/gpt-3.5-turbo" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_transform_request_with_extra_body(self): + """Test request transformation with extra_body parameters""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-4", + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={"extra_body": {"custom_param": "custom_value"}}, + litellm_params={}, + headers={}, + ) + + # Validate that extra_body parameters are merged into the request + assert transformed_request["custom_param"] == "custom_value" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_cache_control_flag_removal(self): + """Test cache control flag removal from messages""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-3.5-turbo", + messages=[ + { + "role": "user", + "content": "Hello, world!", + "cache_control": {"type": "ephemeral"}, + } + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + # CometAPI should remove cache_control flags by default + assert transformed_request["messages"][0].get("cache_control") is None + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + config = CometAPIConfig() + + non_default_params = { + "temperature": 0.7, + "max_tokens": 100, + "top_p": 0.9, + } + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model="cometapi/gpt-3.5-turbo", + drop_params=False, + ) + + assert mapped_params["temperature"] == 0.7 + assert mapped_params["max_tokens"] == 100 + assert mapped_params["top_p"] == 0.9 + + def test_get_error_class(self): + """Test error class creation""" + config = CometAPIConfig() + + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, CometAPIException) + assert error.message == "Test error" + assert error.status_code == 400 + + +# Integration test example (requires real API key) + + +if __name__ == "__main__": + # Quick test runner + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/deepseek/chat/__init__.py b/tests/unit/llms/compactifai/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/chat/__init__.py rename to tests/unit/llms/compactifai/__init__.py diff --git a/tests/test_litellm/llms/compactifai/test_compactifai.py b/tests/unit/llms/compactifai/test_compactifai.py similarity index 84% rename from tests/test_litellm/llms/compactifai/test_compactifai.py rename to tests/unit/llms/compactifai/test_compactifai.py index fd31049731a..1367c703fda 100644 --- a/tests/test_litellm/llms/compactifai/test_compactifai.py +++ b/tests/unit/llms/compactifai/test_compactifai.py @@ -104,56 +104,6 @@ def test_compactifai_completion_streaming(respx_mock): assert chunks[0].choices[0].delta.content == "Hello" -@pytest.mark.respx() -def test_compactifai_models_endpoint(respx_mock): - """Test CompactifAI models listing""" - litellm.disable_aiohttp_transport = True - - mock_response = { - "object": "list", - "data": [ - { - "id": "cai-llama-3-1-8b-slim", - "object": "model", - "created": 1677610602, - "owned_by": "compactifai", - }, - { - "id": "mistral-7b-compressed", - "object": "model", - "created": 1677610602, - "owned_by": "compactifai", - }, - ], - } - - respx_mock.post("https://api.compactif.ai/v1/chat/completions").respond( - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "cai-llama-3-1-8b-slim", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}, - }, - status_code=200, - ) - - # This would be tested if litellm had a models() function - # For now, we'll test that the provider is properly configured - response = litellm.completion( - model="compactifai/cai-llama-3-1-8b-slim", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - @pytest.mark.respx() def test_compactifai_authentication_error(respx_mock): """Test CompactifAI authentication error handling""" diff --git a/tests/test_litellm/llms/deepseek/messages/__init__.py b/tests/unit/llms/custom_httpx/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/messages/__init__.py rename to tests/unit/llms/custom_httpx/__init__.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/unit/llms/custom_httpx/test_aiohttp_cleanup_closed.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py rename to tests/unit/llms/custom_httpx/test_aiohttp_cleanup_closed.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py rename to tests/unit/llms/custom_httpx/test_aiohttp_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/unit/llms/custom_httpx/test_aiohttp_so_keepalive.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py rename to tests/unit/llms/custom_httpx/test_aiohttp_so_keepalive.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py rename to tests/unit/llms/custom_httpx/test_aiohttp_transport.py diff --git a/tests/test_litellm/llms/custom_httpx/test_asgi_handler.py b/tests/unit/llms/custom_httpx/test_asgi_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_asgi_handler.py rename to tests/unit/llms/custom_httpx/test_asgi_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_async_client_cleanup.py b/tests/unit/llms/custom_httpx/test_async_client_cleanup.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_async_client_cleanup.py rename to tests/unit/llms/custom_httpx/test_async_client_cleanup.py diff --git a/tests/test_litellm/llms/custom_httpx/test_container_handler.py b/tests/unit/llms/custom_httpx/test_container_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_container_handler.py rename to tests/unit/llms/custom_httpx/test_container_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/unit/llms/custom_httpx/test_credential_leak_prevention.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py rename to tests/unit/llms/custom_httpx/test_credential_leak_prevention.py diff --git a/tests/test_litellm/llms/custom_httpx/test_gemini_session_leak.py b/tests/unit/llms/custom_httpx/test_gemini_session_leak.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_gemini_session_leak.py rename to tests/unit/llms/custom_httpx/test_gemini_session_leak.py diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_http_handler.py rename to tests/unit/llms/custom_httpx/test_http_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py similarity index 99% rename from tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py rename to tests/unit/llms/custom_httpx/test_llm_http_handler.py index 0350ca74904..399e4dbf206 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -45,7 +45,7 @@ from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" diff --git a/tests/test_litellm/llms/custom_httpx/test_mock_transport.py b/tests/unit/llms/custom_httpx/test_mock_transport.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_mock_transport.py rename to tests/unit/llms/custom_httpx/test_mock_transport.py diff --git a/tests/test_litellm/llms/gemini/__init__.py b/tests/unit/llms/dashscope/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/__init__.py rename to tests/unit/llms/dashscope/__init__.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_chat_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/unit/llms/dashscope/test_dashscope_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py rename to tests/unit/llms/dashscope/test_dashscope_cost_calculator.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py rename to tests/unit/llms/dashscope/test_qwen_brand_aliases.py diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 52bb89fed5a..9cd17bd3580 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -16,6 +16,9 @@ from litellm.llms.databricks.chat.transformation import ( DatabricksConfig, _sanitize_empty_content, ) +from typing import Final +import httpx +import respx @pytest.fixture() @@ -808,3 +811,75 @@ def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> assert parsed.choices[0].delta.reasoning_content == "We need answer" assert parsed.choices[0].delta.content is None + + +def test_completion_merges_leading_system_and_developer_messages_for_chat_template_models( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "developer", "content": "Skills: none."}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse.\n\nSkills: none."}, + {"role": "user", "content": "Hello"}, + ] + assert response.choices[0].message.content == "Answer" + + +def test_completion_merges_system_messages_when_one_has_empty_content(respx_mock: respx.MockRouter): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": ""}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hello"}, + ] diff --git a/tests/test_litellm/llms/databricks/test_databricks_common_utils.py b/tests/unit/llms/databricks/test_databricks_common_utils.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_common_utils.py rename to tests/unit/llms/databricks/test_databricks_common_utils.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py rename to tests/unit/llms/databricks/test_databricks_cost_calculator.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/unit/llms/databricks/test_databricks_partner_integration.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_partner_integration.py rename to tests/unit/llms/databricks/test_databricks_partner_integration.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py b/tests/unit/llms/databricks/test_databricks_streaming_utils.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py rename to tests/unit/llms/databricks/test_databricks_streaming_utils.py diff --git a/tests/test_litellm/llms/gemini/audio_transcription/__init__.py b/tests/unit/llms/deepgram/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/audio_transcription/__init__.py rename to tests/unit/llms/deepgram/__init__.py diff --git a/tests/test_litellm/llms/gemini/google_genai/__init__.py b/tests/unit/llms/deepgram/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/__init__.py rename to tests/unit/llms/deepgram/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py rename to tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/unit/llms/deepgram/test_deepgram_common_utils.py similarity index 100% rename from tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py rename to tests/unit/llms/deepgram/test_deepgram_common_utils.py diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_mock_transcription.py b/tests/unit/llms/deepgram/test_deepgram_mock_transcription.py similarity index 100% rename from tests/test_litellm/llms/deepgram/test_deepgram_mock_transcription.py rename to tests/unit/llms/deepgram/test_deepgram_mock_transcription.py diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/unit/llms/deepinfra/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py rename to tests/unit/llms/deepinfra/__init__.py diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py rename to tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py rename to tests/unit/llms/deepinfra/test_deepinfra_rerank.py diff --git a/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py new file mode 100644 index 00000000000..8a2a1d09cb6 --- /dev/null +++ b/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py @@ -0,0 +1,159 @@ +""" +Integration tests for DeepInfra rerank functionality. +Tests the full rerank flow following the repository patterns. +""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_with_queries_param( + mock_sync_post, mock_async_post, sync_mode +): + """Test DeepInfra rerank with multiple queries parameter.""" + mock_response_data = { + "scores": [0.8, 0.6, 0.2], + "input_tokens": 35, + "request_id": "deepinfra-multi-query-123", + "inference_status": {"status": "success", "runtime_ms": 200}, + } + + def return_val(): + return mock_response_data + + if sync_mode: + mock_response = MagicMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_sync_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-4B", + query="hello", + documents=["hello", "world", "test"], + queries=["hello", "hi there"], # DeepInfra specific param + custom_llm_provider="deepinfra", + api_key="test_key", + api_base="https://api.deepinfra.com", + ) + + mock_sync_post.assert_called_once() + # Verify that queries parameter was passed in request + call_data = json.loads(mock_sync_post.call_args.kwargs["data"]) + assert "queries" in call_data + assert call_data["queries"] == ["hello", "hi there"] + else: + mock_response = AsyncMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_async_post.return_value = mock_response + + response = asyncio.run( + litellm.arerank( + model="deepinfra/Qwen/Qwen3-Reranker-4B", + query="hello", + documents=["hello", "world", "test"], + queries=["hello", "hi there"], + custom_llm_provider="deepinfra", + api_key="test_key", + api_base="https://api.deepinfra.com", + ) + ) + + mock_async_post.assert_called_once() + call_data = json.loads(mock_async_post.call_args.kwargs["data"]) + assert "queries" in call_data + assert call_data["queries"] == ["hello", "hi there"] + + assert response.results is not None + assert len(response.results) == 3 + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_with_env_vars(mock_post, monkeypatch): + """Test DeepInfra rerank with environment variable configuration.""" + monkeypatch.setenv("DEEPINFRA_API_KEY", "env_test_key") + monkeypatch.setenv("DEEPINFRA_API_BASE", "https://custom-deepinfra.com") + + mock_response_data = { + "scores": [0.88, 0.22], + "input_tokens": 28, + "request_id": "env-test-123", + } + + def return_val(): + return mock_response_data + + mock_response = MagicMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-0.6B", + query="hello", + documents=["hello", "world"], + custom_llm_provider="deepinfra", + ) + + mock_post.assert_called_once() + + # Verify headers contain env API key + headers = mock_post.call_args.kwargs.get("headers", {}) + assert "Bearer env_test_key" in headers.get("Authorization", "") + + assert response.results is not None + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_defaults_api_base_when_missing(mock_post, monkeypatch): + """With no api_base anywhere, the call still goes out against DeepInfra's own base.""" + monkeypatch.delenv("DEEPINFRA_API_BASE", raising=False) + + mock_response = MagicMock() + mock_response.json = lambda: {"scores": [0.9, 0.1], "input_tokens": 20} + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-0.6B", + query="hello", + documents=["hello", "world"], + custom_llm_provider="deepinfra", + api_key="test_key", + # api_base is intentionally missing + ) + + assert "api.deepinfra.com" in mock_post.call_args.kwargs["url"] + assert [result["relevance_score"] for result in response.results] == [0.9, 0.1] + + +def test_deepinfra_rerank_models(): + """Test that DeepInfra Qwen rerank models are recognized.""" + # These should not raise errors during model validation + models = [ + "deepinfra/Qwen/Qwen3-Reranker-0.6B", + "deepinfra/Qwen/Qwen3-Reranker-4B", + "deepinfra/Qwen/Qwen3-Reranker-8B", + ] + + for model in models: + resolved_model, provider, _, api_base = litellm.get_llm_provider(model=model) + assert provider == "deepinfra" + assert resolved_model == model.removeprefix("deepinfra/") + assert api_base == "https://api.deepinfra.com/v1/openai" diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py rename to tests/unit/llms/deepinfra/test_deepinfra_rerank_transformation.py diff --git a/tests/test_litellm/llms/gemini/image_edit/__init__.py b/tests/unit/llms/edenai/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/image_edit/__init__.py rename to tests/unit/llms/edenai/__init__.py diff --git a/tests/test_litellm/llms/gemini/realtime/__init__.py b/tests/unit/llms/edenai/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/realtime/__init__.py rename to tests/unit/llms/edenai/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py b/tests/unit/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py rename to tests/unit/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/gigachat/__init__.py b/tests/unit/llms/edenai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/__init__.py rename to tests/unit/llms/edenai/chat/__init__.py diff --git a/tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py b/tests/unit/llms/edenai/chat/test_edenai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py rename to tests/unit/llms/edenai/chat/test_edenai_chat_transformation.py diff --git a/tests/test_litellm/llms/edenai/conftest.py b/tests/unit/llms/edenai/conftest.py similarity index 100% rename from tests/test_litellm/llms/edenai/conftest.py rename to tests/unit/llms/edenai/conftest.py diff --git a/tests/test_litellm/llms/gigachat/embedding/__init__.py b/tests/unit/llms/edenai/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/embedding/__init__.py rename to tests/unit/llms/edenai/embedding/__init__.py diff --git a/tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py rename to tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py diff --git a/tests/test_litellm/llms/gigachat/passthrough/__init__.py b/tests/unit/llms/edenai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/passthrough/__init__.py rename to tests/unit/llms/edenai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py rename to tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py diff --git a/tests/test_litellm/llms/github_copilot/messages/__init__.py b/tests/unit/llms/edenai/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/messages/__init__.py rename to tests/unit/llms/edenai/messages/__init__.py diff --git a/tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py b/tests/unit/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py rename to tests/unit/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/gradient_ai/__init__.py b/tests/unit/llms/edenai/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/__init__.py rename to tests/unit/llms/edenai/responses/__init__.py diff --git a/tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py b/tests/unit/llms/edenai/responses/test_edenai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py rename to tests/unit/llms/edenai/responses/test_edenai_responses_transformation.py diff --git a/tests/test_litellm/llms/edenai/test_edenai_common_utils.py b/tests/unit/llms/edenai/test_edenai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/edenai/test_edenai_common_utils.py rename to tests/unit/llms/edenai/test_edenai_common_utils.py diff --git a/tests/test_litellm/llms/gradient_ai/chat/__init__.py b/tests/unit/llms/edenai/text_to_speech/__init__.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/chat/__init__.py rename to tests/unit/llms/edenai/text_to_speech/__init__.py diff --git a/tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py b/tests/unit/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py rename to tests/unit/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/groq/__init__.py b/tests/unit/llms/edenai/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/groq/__init__.py rename to tests/unit/llms/edenai/videos/__init__.py diff --git a/tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py b/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py rename to tests/unit/llms/edenai/videos/test_edenai_video_transformation.py diff --git a/tests/test_litellm/llms/groq/chat/__init__.py b/tests/unit/llms/fal_ai/__init__.py similarity index 100% rename from tests/test_litellm/llms/groq/chat/__init__.py rename to tests/unit/llms/fal_ai/__init__.py diff --git a/tests/test_litellm/llms/huggingface/__init__.py b/tests/unit/llms/fal_ai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/huggingface/__init__.py rename to tests/unit/llms/fal_ai/chat/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py rename to tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/inception/__init__.py b/tests/unit/llms/fal_ai/image_edit/__init__.py similarity index 100% rename from tests/test_litellm/llms/inception/__init__.py rename to tests/unit/llms/fal_ai/image_edit/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py rename to tests/unit/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py rename to tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/mistral/batches/__init__.py b/tests/unit/llms/fal_ai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/mistral/batches/__init__.py rename to tests/unit/llms/fal_ai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/unit/llms/fal_ai/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/test_cost_calculator.py rename to tests/unit/llms/fal_ai/test_cost_calculator.py diff --git a/tests/test_litellm/llms/mistral/files/__init__.py b/tests/unit/llms/fal_ai/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/mistral/files/__init__.py rename to tests/unit/llms/fal_ai/videos/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py rename to tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py diff --git a/tests/test_litellm/llms/nvidia_riva/__init__.py b/tests/unit/llms/featherless_ai/__init__.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/__init__.py rename to tests/unit/llms/featherless_ai/__init__.py diff --git a/tests/test_litellm/llms/oci/rerank/__init__.py b/tests/unit/llms/featherless_ai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/oci/rerank/__init__.py rename to tests/unit/llms/featherless_ai/chat/__init__.py diff --git a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py b/tests/unit/llms/featherless_ai/chat/test_featherless_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py rename to tests/unit/llms/featherless_ai/chat/test_featherless_chat_transformation.py diff --git a/tests/test_litellm/llms/ocr/__init__.py b/tests/unit/llms/fireworks_ai/completion/__init__.py similarity index 100% rename from tests/test_litellm/llms/ocr/__init__.py rename to tests/unit/llms/fireworks_ai/completion/__init__.py diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py rename to tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py rename to tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py diff --git a/tests/test_litellm/llms/openai_like/responses/__init__.py b/tests/unit/llms/gdc/__init__.py similarity index 100% rename from tests/test_litellm/llms/openai_like/responses/__init__.py rename to tests/unit/llms/gdc/__init__.py diff --git a/tests/test_litellm/llms/parallel_ai/__init__.py b/tests/unit/llms/gdc/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/__init__.py rename to tests/unit/llms/gdc/chat/__init__.py diff --git a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py b/tests/unit/llms/gdc/chat/test_gdc_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py rename to tests/unit/llms/gdc/chat/test_gdc_chat_transformation.py diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/unit/llms/gemini/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_cost_calculator.py rename to tests/unit/llms/gemini/test_cost_calculator.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py b/tests/unit/llms/gemini/test_gemini_client_setup.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_client_setup.py rename to tests/unit/llms/gemini/test_gemini_client_setup.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_common_utils.py rename to tests/unit/llms/gemini/test_gemini_common_utils.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py rename to tests/unit/llms/gemini/test_gemini_image_generation_transformation.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_tts.py b/tests/unit/llms/gemini/test_gemini_tts.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_tts.py rename to tests/unit/llms/gemini/test_gemini_tts.py diff --git a/tests/test_litellm/llms/github_copilot/test_github_copilot_authenticator.py b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/test_github_copilot_authenticator.py rename to tests/unit/llms/github_copilot/test_github_copilot_authenticator.py diff --git a/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py rename to tests/unit/llms/github_copilot/test_github_copilot_transformation.py diff --git a/tests/test_litellm/llms/pass_through/__init__.py b/tests/unit/llms/heroku/__init__.py similarity index 100% rename from tests/test_litellm/llms/pass_through/__init__.py rename to tests/unit/llms/heroku/__init__.py diff --git a/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py b/tests/unit/llms/heroku/test_heroku_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py rename to tests/unit/llms/heroku/test_heroku_chat_transformation.py diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/__init__.py b/tests/unit/llms/huggingface/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/pass_through/guardrail_translation/__init__.py rename to tests/unit/llms/huggingface/embedding/__init__.py diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py similarity index 100% rename from tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py rename to tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py diff --git a/tests/test_litellm/llms/langflow/test_langflow_a2a.py b/tests/unit/llms/langflow/test_langflow_a2a.py similarity index 100% rename from tests/test_litellm/llms/langflow/test_langflow_a2a.py rename to tests/unit/llms/langflow/test_langflow_a2a.py diff --git a/tests/test_litellm/llms/perplexity/__init__.py b/tests/unit/llms/lemonade/__init__.py similarity index 100% rename from tests/test_litellm/llms/perplexity/__init__.py rename to tests/unit/llms/lemonade/__init__.py diff --git a/tests/test_litellm/llms/lemonade/test_lemonade.py b/tests/unit/llms/lemonade/test_lemonade.py similarity index 100% rename from tests/test_litellm/llms/lemonade/test_lemonade.py rename to tests/unit/llms/lemonade/test_lemonade.py diff --git a/tests/test_litellm/llms/perplexity/embedding/__init__.py b/tests/unit/llms/lm_studio/__init__.py similarity index 100% rename from tests/test_litellm/llms/perplexity/embedding/__init__.py rename to tests/unit/llms/lm_studio/__init__.py diff --git a/tests/test_litellm/llms/lm_studio/test_lm_studio_chat_transformation.py b/tests/unit/llms/lm_studio/test_lm_studio_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/lm_studio/test_lm_studio_chat_transformation.py rename to tests/unit/llms/lm_studio/test_lm_studio_chat_transformation.py diff --git a/tests/test_litellm/llms/stability/__init__.py b/tests/unit/llms/mistral/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/stability/__init__.py rename to tests/unit/llms/mistral/audio_transcription/__init__.py diff --git a/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py new file mode 100644 index 00000000000..68875ff6d32 --- /dev/null +++ b/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py @@ -0,0 +1,195 @@ +import os +from unittest.mock import MagicMock + +import httpx +import litellm + +from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig, +) +from litellm.llms.mistral.audio_transcription.transformation import ( + MistralAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse +from litellm.utils import ProviderConfigManager + + +def test_mistral_audio_transcription_config_installed(): + """Ensure Mistral audio transcription config is registered with ProviderConfigManager.""" + config = ProviderConfigManager.get_provider_audio_transcription_config( + model="mistral/voxtral-mini-latest", + provider=litellm.LlmProviders.MISTRAL, + ) + assert config is not None + assert isinstance(config, BaseAudioTranscriptionConfig) + assert isinstance(config, MistralAudioTranscriptionConfig) + + +def test_mistral_audio_transcription_get_complete_url(): + config = MistralAudioTranscriptionConfig() + url = config.get_complete_url( + api_base=None, + api_key="fake-key", + model="voxtral-mini-latest", + optional_params={}, + litellm_params={}, + ) + assert url == "https://api.mistral.ai/v1/audio/transcriptions" + + +def test_mistral_audio_transcription_get_complete_url_custom_base(): + config = MistralAudioTranscriptionConfig() + url = config.get_complete_url( + api_base="https://custom.api.example.com/v1/", + api_key="fake-key", + model="voxtral-mini-latest", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.example.com/v1/audio/transcriptions" + + +def test_mistral_audio_transcription_validate_environment(): + config = MistralAudioTranscriptionConfig() + headers = config.validate_environment( + headers={}, + model="voxtral-mini-latest", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key-123", + ) + assert headers["Authorization"] == "Bearer test-key-123" + assert headers["accept"] == "application/json" + + +def test_mistral_audio_transcription_supported_params(): + config = MistralAudioTranscriptionConfig() + params = config.get_supported_openai_params("voxtral-mini-latest") + assert "language" in params + assert "temperature" in params + assert "response_format" in params + assert "timestamp_granularities" in params + + +def test_mistral_audio_transcription_request_transform(): + config = MistralAudioTranscriptionConfig() + + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + audio_file = open(wav_path, "rb") + + result = config.transform_audio_transcription_request( + model="voxtral-mini-latest", + audio_file=audio_file, + optional_params={"language": "en", "temperature": 0.0}, + litellm_params={}, + ) + + audio_file.close() + + assert isinstance(result.data, dict) + assert result.data["model"] == "voxtral-mini-latest" + assert result.data["language"] == "en" + assert result.data["temperature"] == 0.0 + assert result.files is not None + assert "file" in result.files + + +def test_mistral_audio_transcription_request_with_diarize(): + """Test that Mistral-specific params like diarize are passed through.""" + config = MistralAudioTranscriptionConfig() + + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + audio_file = open(wav_path, "rb") + + result = config.transform_audio_transcription_request( + model="voxtral-mini-latest", + audio_file=audio_file, + optional_params={"diarize": True}, + litellm_params={}, + ) + + audio_file.close() + + assert isinstance(result.data, dict) + assert result.data["diarize"] == "true" + + +def test_mistral_audio_transcription_response_transform(): + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = {"text": "Four score and seven years ago..."} + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago..." + + +def test_mistral_audio_transcription_response_transform_diarized(): + """Test that diarized responses preserve segments and language.""" + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = { + "model": "voxtral-mini-latest", + "text": "Hello, how are you? I am fine.", + "language": None, + "segments": [ + { + "text": "Hello, how are you?", + "start": 0.3, + "end": 2.1, + "speaker_id": "speaker_1", + "type": "transcription_segment", + }, + { + "text": "I am fine.", + "start": 2.5, + "end": 3.8, + "speaker_id": "speaker_2", + "type": "transcription_segment", + }, + ], + "usage": { + "prompt_audio_seconds": 4, + "prompt_tokens": 5, + "total_tokens": 50, + "completion_tokens": 20, + }, + } + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Hello, how are you? I am fine." + assert response["segments"] is not None + assert len(response["segments"]) == 2 + assert response["segments"][0]["speaker_id"] == "speaker_1" + assert response["segments"][1]["speaker_id"] == "speaker_2" + assert response["language"] is None + + +def test_mistral_audio_transcription_response_transform_empty(): + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = {} + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "" diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py rename to tests/unit/llms/mistral/test_mistral_chat_transformation.py diff --git a/tests/test_litellm/llms/mistral/test_mistral_completion.py b/tests/unit/llms/mistral/test_mistral_completion.py similarity index 100% rename from tests/test_litellm/llms/mistral/test_mistral_completion.py rename to tests/unit/llms/mistral/test_mistral_completion.py diff --git a/tests/test_litellm/llms/stability/image_generation/__init__.py b/tests/unit/llms/modelscope/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/stability/image_generation/__init__.py rename to tests/unit/llms/modelscope/chat/__init__.py diff --git a/tests/test_litellm/llms/modelscope/chat/test_modelscope_chat_transformation.py b/tests/unit/llms/modelscope/chat/test_modelscope_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/modelscope/chat/test_modelscope_chat_transformation.py rename to tests/unit/llms/modelscope/chat/test_modelscope_chat_transformation.py diff --git a/tests/test_litellm/llms/tencent/__init__.py b/tests/unit/llms/nadir/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/__init__.py rename to tests/unit/llms/nadir/__init__.py diff --git a/tests/test_litellm/llms/nadir/test_nadir.py b/tests/unit/llms/nadir/test_nadir.py similarity index 100% rename from tests/test_litellm/llms/nadir/test_nadir.py rename to tests/unit/llms/nadir/test_nadir.py diff --git a/tests/test_litellm/llms/tencent/chat/__init__.py b/tests/unit/llms/nebius/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/chat/__init__.py rename to tests/unit/llms/nebius/__init__.py diff --git a/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py b/tests/unit/llms/nebius/test_nebius_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py rename to tests/unit/llms/nebius/test_nebius_chat_transformation.py diff --git a/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py b/tests/unit/llms/nebius/test_nebius_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py rename to tests/unit/llms/nebius/test_nebius_embedding_transformation.py diff --git a/tests/test_litellm/llms/tencent/messages/__init__.py b/tests/unit/llms/oci/rerank/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/messages/__init__.py rename to tests/unit/llms/oci/rerank/__init__.py diff --git a/tests/test_litellm/llms/oci/test_oci_common_utils.py b/tests/unit/llms/oci/test_oci_common_utils.py similarity index 100% rename from tests/test_litellm/llms/oci/test_oci_common_utils.py rename to tests/unit/llms/oci/test_oci_common_utils.py diff --git a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py b/tests/unit/llms/oci/test_oci_coverage_boost.py similarity index 100% rename from tests/test_litellm/llms/oci/test_oci_coverage_boost.py rename to tests/unit/llms/oci/test_oci_coverage_boost.py diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py b/tests/unit/llms/ollama/__init__.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py rename to tests/unit/llms/ollama/__init__.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py rename to tests/unit/llms/ollama/test_ollama_chat_transformation.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py rename to tests/unit/llms/ollama/test_ollama_completion_transformation.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_embedding.py rename to tests/unit/llms/ollama/test_ollama_embedding.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/unit/llms/ollama/test_ollama_model_info.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_model_info.py rename to tests/unit/llms/ollama/test_ollama_model_info.py diff --git a/tests/test_litellm/llms/openai/realtime/README.md b/tests/unit/llms/openai/realtime/README.md similarity index 100% rename from tests/test_litellm/llms/openai/realtime/README.md rename to tests/unit/llms/openai/realtime/README.md diff --git a/tests/test_litellm/llms/vertex_ai/agent_engine/__init__.py b/tests/unit/llms/openai/realtime/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/agent_engine/__init__.py rename to tests/unit/llms/openai/realtime/__init__.py diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py rename to tests/unit/llms/openai/realtime/test_openai_realtime_handler.py diff --git a/tests/test_litellm/llms/openai/realtime/test_transcription_sessions.py b/tests/unit/llms/openai/realtime/test_transcription_sessions.py similarity index 100% rename from tests/test_litellm/llms/openai/realtime/test_transcription_sessions.py rename to tests/unit/llms/openai/realtime/test_transcription_sessions.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/__init__.py b/tests/unit/llms/openai/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/__init__.py rename to tests/unit/llms/openai/responses/__init__.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py rename to tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py b/tests/unit/llms/openai/responses/test_openai_responses_data_residency.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py rename to tests/unit/llms/openai/responses/test_openai_responses_data_residency.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py rename to tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py rename to tests/unit/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py rename to tests/unit/llms/openai/responses/test_openai_responses_transformation.py diff --git a/tests/test_litellm/llms/openai/test_cost_calculation.py b/tests/unit/llms/openai/test_cost_calculation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_cost_calculation.py rename to tests/unit/llms/openai/test_cost_calculation.py diff --git a/tests/test_litellm/llms/openai/test_data_residency.py b/tests/unit/llms/openai/test_data_residency.py similarity index 100% rename from tests/test_litellm/llms/openai/test_data_residency.py rename to tests/unit/llms/openai/test_data_residency.py diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/unit/llms/openai/test_gpt5_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_gpt5_transformation.py rename to tests/unit/llms/openai/test_gpt5_transformation.py diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/unit/llms/openai/test_is_model_gpt_5_model.py similarity index 100% rename from tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py rename to tests/unit/llms/openai/test_is_model_gpt_5_model.py diff --git a/tests/test_litellm/llms/openai/test_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_o_series_transformation.py rename to tests/unit/llms/openai/test_o_series_transformation.py diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai.py rename to tests/unit/llms/openai/test_openai.py diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_common_utils.py rename to tests/unit/llms/openai/test_openai_common_utils.py diff --git a/tests/test_litellm/llms/openai/test_openai_empty_response.py b/tests/unit/llms/openai/test_openai_empty_response.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_empty_response.py rename to tests/unit/llms/openai/test_openai_empty_response.py diff --git a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py b/tests/unit/llms/openai/test_openai_file_content_streaming.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_file_content_streaming.py rename to tests/unit/llms/openai/test_openai_file_content_streaming.py diff --git a/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py b/tests/unit/llms/openai/test_openai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py rename to tests/unit/llms/openai/test_openai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_workload_identity.py rename to tests/unit/llms/openai/test_openai_workload_identity.py diff --git a/tests/test_litellm/llms/openai/test_organization_costs.py b/tests/unit/llms/openai/test_organization_costs.py similarity index 100% rename from tests/test_litellm/llms/openai/test_organization_costs.py rename to tests/unit/llms/openai/test_organization_costs.py diff --git a/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py similarity index 100% rename from tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py rename to tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py b/tests/unit/llms/openai/transcriptions/test_openai_transcriptions_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py rename to tests/unit/llms/openai/transcriptions/test_openai_transcriptions_handler.py diff --git a/tests/test_litellm/llms/vertex_ai/batches/__init__.py b/tests/unit/llms/openai_like/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/__init__.py rename to tests/unit/llms/openai_like/responses/__init__.py diff --git a/tests/test_litellm/llms/openai_like/responses/test_openai_like_responses.py b/tests/unit/llms/openai_like/responses/test_openai_like_responses.py similarity index 100% rename from tests/test_litellm/llms/openai_like/responses/test_openai_like_responses.py rename to tests/unit/llms/openai_like/responses/test_openai_like_responses.py diff --git a/tests/test_litellm/llms/openai_like/test_abliteration_provider.py b/tests/unit/llms/openai_like/test_abliteration_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_abliteration_provider.py rename to tests/unit/llms/openai_like/test_abliteration_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_assemblyai_provider.py b/tests/unit/llms/openai_like/test_assemblyai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_assemblyai_provider.py rename to tests/unit/llms/openai_like/test_assemblyai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_charity_engine.py b/tests/unit/llms/openai_like/test_charity_engine.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_charity_engine.py rename to tests/unit/llms/openai_like/test_charity_engine.py diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/unit/llms/openai_like/test_cognition_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_cognition_provider.py rename to tests/unit/llms/openai_like/test_cognition_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_dynamic_config.py b/tests/unit/llms/openai_like/test_dynamic_config.py similarity index 96% rename from tests/test_litellm/llms/openai_like/test_dynamic_config.py rename to tests/unit/llms/openai_like/test_dynamic_config.py index 55e1a1679de..de70f98c3f1 100644 --- a/tests/test_litellm/llms/openai_like/test_dynamic_config.py +++ b/tests/unit/llms/openai_like/test_dynamic_config.py @@ -20,9 +20,6 @@ def _isolate_generated_class_cache(): class TestClassCaching: - def test_same_slug_returns_the_identical_class_object(self): - provider = _provider("cache_same_slug") - assert create_responses_config_class(provider) is create_responses_config_class(provider) def test_cache_is_keyed_on_slug_not_on_the_provider_instance(self): first = create_responses_config_class(_provider("cache_by_slug")) diff --git a/tests/test_litellm/llms/openai_like/test_empiriolabs_provider.py b/tests/unit/llms/openai_like/test_empiriolabs_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_empiriolabs_provider.py rename to tests/unit/llms/openai_like/test_empiriolabs_provider.py diff --git a/tests/unit/llms/openai_like/test_json_providers.py b/tests/unit/llms/openai_like/test_json_providers.py new file mode 100644 index 00000000000..a56108ca9ac --- /dev/null +++ b/tests/unit/llms/openai_like/test_json_providers.py @@ -0,0 +1,317 @@ +""" +Tests for JSON-based provider configuration system. +""" + +import os +import sys +from unittest.mock import patch + +try: + import pytest +except ImportError: + # pytest not available, will run as standalone script + pytest = None + +# Add workspace to path +workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +sys.path.insert(0, workspace_path) + + + +class TestJSONProviderLoader: + """Test JSON provider loading and configuration""" + + def test_load_json_providers(self): + """Test that JSON providers load correctly""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + # Verify publicai is loaded + assert JSONProviderRegistry.exists("publicai") + + # Get publicai config + publicai = JSONProviderRegistry.get("publicai") + assert publicai is not None + assert publicai.base_url == "https://api.publicai.co/v1" + assert publicai.api_key_env == "PUBLICAI_API_KEY" + assert publicai.api_base_env == "PUBLICAI_API_BASE" + assert publicai.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_dynamic_config_generation(self): + """Test dynamic config class creation""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Test API info resolution + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.publicai.co/v1" + + # Test with custom base + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.api.com", "test-key" + ) + assert api_base == "https://custom.api.com" + assert api_key == "test-key" + + def test_parameter_mapping(self): + """Test parameter mapping works""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Test parameter mapping + optional_params = {} + non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} + result = config.map_openai_params( + non_default_params, optional_params, "gpt-4", False + ) + + # max_completion_tokens should be mapped to max_tokens + assert "max_tokens" in result + assert result["max_tokens"] == 100 + assert "max_completion_tokens" not in result + + # temperature should be passed through + assert result["temperature"] == 0.7 + + def test_supported_params(self): + """Test that config returns supported params""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Get supported params + supported = config.get_supported_openai_params("gpt-4") + + # Should have standard OpenAI params + assert isinstance(supported, list) + assert len(supported) > 0 + + def test_tool_params_excluded_when_function_calling_not_supported(self): + """Test that tool-related params are excluded for models that don't support + function calling. Regression test for https://github.com/BerriAI/litellm/issues/21125 + """ + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Mock supports_function_calling to return False + with patch("litellm.utils.supports_function_calling", return_value=False): + supported = config.get_supported_openai_params("some-model-without-fc") + + tool_params = [ + "tools", + "tool_choice", + "function_call", + "functions", + "parallel_tool_calls", + ] + for param in tool_params: + assert ( + param not in supported + ), f"'{param}' should not be in supported params when function calling is not supported" + + # Non-tool params should still be present + assert "temperature" in supported + assert "max_tokens" in supported + assert "stop" in supported + + def test_tool_params_included_when_function_calling_supported(self): + """Test that tool-related params are included for models that support function calling.""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Mock supports_function_calling to return True + with patch("litellm.utils.supports_function_calling", return_value=True): + supported = config.get_supported_openai_params("some-model-with-fc") + + assert "tools" in supported + assert "tool_choice" in supported + + def test_provider_resolution(self): + """Test that provider resolution finds JSON providers""" + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + + model, provider, api_key, api_base = get_llm_provider( + model="publicai/gpt-4", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gpt-4" + assert provider == "publicai" + assert api_base == "https://api.publicai.co/v1" + + def test_provider_config_manager(self): + """Test that ProviderConfigManager returns JSON-based configs""" + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gpt-4", provider=LlmProviders.PUBLICAI + ) + + assert config is not None + assert config.custom_llm_provider == "publicai" + + +class TestPinstripes: + """Tests for Pinstripes JSON-configured provider""" + + def test_pinstripes_json_config_exists(self): + """Test that pinstripes is configured in providers.json""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("pinstripes") + + pinstripes = JSONProviderRegistry.get("pinstripes") + assert pinstripes is not None + assert pinstripes.base_url == "https://pinstripes.io/v1" + assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" + assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_pinstripes_provider_resolution(self): + """Test that provider resolution finds pinstripes and returns the default base URL""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="pinstripes/ps/glm-4.5-air", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "ps/glm-4.5-air" + assert provider == "pinstripes" + assert api_base == "https://pinstripes.io/v1" + + def test_pinstripes_dynamic_config(self): + """Test dynamic config class creation for pinstripes""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("pinstripes") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://pinstripes.io/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.pinstripes.io/v1", "test-key" + ) + assert api_base == "https://custom.pinstripes.io/v1" + assert api_key == "test-key" + + def test_pinstripes_parameter_mapping(self): + """Test that max_completion_tokens is mapped to max_tokens for pinstripes""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("pinstripes") + config_class = create_config_class(provider) + config = config_class() + + optional_params = {} + non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} + result = config.map_openai_params( + non_default_params, optional_params, "ps/glm-4.5-air", False + ) + + assert "max_tokens" in result + assert result["max_tokens"] == 100 + assert "max_completion_tokens" not in result + assert result["temperature"] == 0.7 + + +class TestDarkbloom: + def test_darkbloom_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + darkbloom = JSONProviderRegistry.get("darkbloom") + assert darkbloom is not None + assert darkbloom.base_url == "https://api.darkbloom.dev/v1" + assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" + assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" + assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_darkbloom_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="darkbloom/gemma-4-26b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gemma-4-26b" + assert provider == "darkbloom" + assert api_key is None + assert api_base == "https://api.darkbloom.dev/v1" + + def test_darkbloom_dynamic_config(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.darkbloom.dev/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.darkbloom.dev/v1", "test-key" + ) + assert api_base == "https://custom.darkbloom.dev/v1" + assert api_key == "test-key" + + def test_darkbloom_complete_url_appends_endpoint(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + url = config.get_complete_url( + api_base="https://api.darkbloom.dev/v1", + api_key="test-key", + model="darkbloom/gemma-4-26b", + optional_params={}, + litellm_params={}, + stream=True, + ) + + assert url == "https://api.darkbloom.dev/v1/chat/completions" + + def test_darkbloom_provider_config_manager(self): + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gemma-4-26b", provider=LlmProviders.DARKBLOOM + ) + + assert config is not None + assert config.custom_llm_provider == "darkbloom" diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/unit/llms/openai_like/test_libertai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_libertai_provider.py rename to tests/unit/llms/openai_like/test_libertai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/unit/llms/openai_like/test_meta_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_meta_provider.py rename to tests/unit/llms/openai_like/test_meta_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_model_info.py b/tests/unit/llms/openai_like/test_model_info.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_model_info.py rename to tests/unit/llms/openai_like/test_model_info.py diff --git a/tests/test_litellm/llms/openai_like/test_pinstripes_provider.py b/tests/unit/llms/openai_like/test_pinstripes_provider.py similarity index 68% rename from tests/test_litellm/llms/openai_like/test_pinstripes_provider.py rename to tests/unit/llms/openai_like/test_pinstripes_provider.py index 70bb786b2e6..e7a2dfb92dc 100644 --- a/tests/test_litellm/llms/openai_like/test_pinstripes_provider.py +++ b/tests/unit/llms/openai_like/test_pinstripes_provider.py @@ -16,17 +16,6 @@ class TestPinstripeProviderConfig: assert LlmProviders.PINSTRIPES.value == "pinstripes" assert "pinstripes" in litellm.provider_list - def test_pinstripes_json_config_exists(self): - """Test that pinstripes is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - assert JSONProviderRegistry.exists("pinstripes") - - pinstripes = JSONProviderRegistry.get("pinstripes") - assert pinstripes is not None - assert pinstripes.base_url == "https://pinstripes.io/v1" - assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" - assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" def test_pinstripes_in_openai_compatible_providers(self): """Test that pinstripes is in the openai_compatible_providers list""" @@ -34,20 +23,6 @@ class TestPinstripeProviderConfig: assert "pinstripes" in openai_compatible_providers - def test_pinstripes_provider_resolution(self): - """Test that provider resolution finds pinstripes and returns the default base URL""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="pinstripes/ps/glm-4.5-air", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "ps/glm-4.5-air" - assert provider == "pinstripes" - assert api_base == "https://pinstripes.io/v1" def test_pinstripes_api_base_override(self): """Test that an explicit api_base / api_key overrides the default""" diff --git a/tests/test_litellm/llms/openai_like/test_provider_affinity_forwarding.py b/tests/unit/llms/openai_like/test_provider_affinity_forwarding.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_provider_affinity_forwarding.py rename to tests/unit/llms/openai_like/test_provider_affinity_forwarding.py diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/unit/llms/openai_like/test_scx_ai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_scx_ai_provider.py rename to tests/unit/llms/openai_like/test_scx_ai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py b/tests/unit/llms/openai_like/test_tensormesh_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_tensormesh_provider.py rename to tests/unit/llms/openai_like/test_tensormesh_provider.py diff --git a/tests/unit/llms/openai_like/test_xiaomi_mimo.py b/tests/unit/llms/openai_like/test_xiaomi_mimo.py new file mode 100644 index 00000000000..a642cc91f90 --- /dev/null +++ b/tests/unit/llms/openai_like/test_xiaomi_mimo.py @@ -0,0 +1,84 @@ +""" +Tests for Xiaomi MiMo provider configuration and integration. +Related to issue #18794 +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +try: + import pytest +except ImportError: + pytest = None + +# Add workspace to path +workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +sys.path.insert(0, workspace_path) + +import litellm + + +class TestXiaomiMiMoProviderConfig: + """Test Xiaomi MiMo provider configuration""" + + def test_xiaomi_mimo_in_provider_list(self): + """Test that xiaomi_mimo is in the provider list (fixes #18794)""" + from litellm import LlmProviders + + # Verify xiaomi_mimo is in the enum + assert hasattr(LlmProviders, "XIAOMI_MIMO") + assert LlmProviders.XIAOMI_MIMO.value == "xiaomi_mimo" + + # Verify it's in the provider list + assert "xiaomi_mimo" in litellm.provider_list + + def test_xiaomi_mimo_json_config_exists(self): + """Test that xiaomi_mimo is configured in providers.json""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + # Verify xiaomi_mimo is loaded + assert JSONProviderRegistry.exists("xiaomi_mimo") + + # Get xiaomi_mimo config + xiaomi_mimo = JSONProviderRegistry.get("xiaomi_mimo") + assert xiaomi_mimo is not None + assert xiaomi_mimo.base_url == "https://api.xiaomimimo.com/v1" + assert xiaomi_mimo.api_key_env == "XIAOMI_MIMO_API_KEY" + assert xiaomi_mimo.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_xiaomi_mimo_provider_resolution(self): + """Test that provider resolution finds xiaomi_mimo""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="xiaomi_mimo/mimo-v2-flash", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "mimo-v2-flash" + assert provider == "xiaomi_mimo" + assert api_base == "https://api.xiaomimimo.com/v1" + + def test_xiaomi_mimo_router_config(self): + """Test that xiaomi_mimo can be used in Router configuration (fixes #18794)""" + from litellm import Router + + # This should not raise "Unsupported provider - xiaomi_mimo" + router = Router( + model_list=[ + { + "model_name": "mimo-v2-flash", + "litellm_params": { + "model": "xiaomi_mimo/mimo-v2-flash", + "api_key": "test-key", + }, + } + ] + ) + + # Verify the deployment was created successfully + assert len(router.model_list) == 1 + assert router.model_list[0]["model_name"] == "mimo-v2-flash" diff --git a/tests/test_litellm/llms/vertex_ai/files/__init__.py b/tests/unit/llms/ovhcloud/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/__init__.py rename to tests/unit/llms/ovhcloud/__init__.py diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py new file mode 100644 index 00000000000..87e54dfba9b --- /dev/null +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py @@ -0,0 +1,58 @@ + + + + +class TestOVHCloudDurationFieldMigration: + """Tests for OVHCloud duration -> seconds field migration.""" + + def test_seconds_field_mapped_to_duration(self): + """New `seconds` field should be normalized to `duration`.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello world", + "seconds": 3.14, + } + + result = config.transform_audio_transcription_response(mock_response) + + assert result.text == "Hello world" + assert result._hidden_params["duration"] == 3.14 + + def test_legacy_duration_field_still_works(self): + """Legacy `duration` field should still be accepted.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello world", + "duration": 2.71, + } + + result = config.transform_audio_transcription_response(mock_response) + + assert result.text == "Hello world" + assert result._hidden_params["duration"] == 2.71 + + + def test_seconds_zero_mapped_to_duration(self): + """seconds=0.0 must not be treated as falsy and lost.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = {"text": "silence", "seconds": 0.0} + result = config.transform_audio_transcription_response(mock_response) + assert result._hidden_params["duration"] == 0.0 diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py new file mode 100644 index 00000000000..c2bc4ee4a4c --- /dev/null +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py @@ -0,0 +1,250 @@ +""" +Unit tests for OVHCloud AI Endpoints chat integration. +""" + + +import pytest + +from litellm.llms.ovhcloud.utils import OVHCloudException +from litellm.utils import get_optional_params + + +from litellm.llms.ovhcloud.chat.transformation import ( + OVHCloudChatCompletionStreamingHandler, + OVHCloudChatConfig, +) + +config = OVHCloudChatConfig() +model = "ovhcloud/Mistral-7B-Instruct-v0.3" + + +class TestOvhCloudChatCompletionStreamingHandler: + def test_chunk_parser_successful(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + chunk = { + "id": "test_id", + "created": 1234567890, + "model": "gpt-oss-20b", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "choices": [ + {"delta": {"content": "test content", "reasoning": "test reasoning"}} + ], + } + + result = handler.chunk_parser(chunk) + + assert result.id == "test_id" + assert result.object == "chat.completion.chunk" + assert result.created == 1234567890 + assert result.model == "gpt-oss-20b" + assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] + assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] + assert result.usage.total_tokens == chunk["usage"]["total_tokens"] + assert len(result.choices) == 1 + assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" + + def test_chunk_parser_error_response(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + error_chunk = { + "error": { + "message": "test error", + "code": 400, + } + } + + with pytest.raises(OVHCloudException) as exc_info: + handler.chunk_parser(error_chunk) + + assert "OVHCloud Error: test error" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + def test_chunk_parser_key_error(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + invalid_chunk = {"incomplete": "data"} + + with pytest.raises(OVHCloudException) as exc_info: + handler.chunk_parser(invalid_chunk) + + assert "KeyError" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + +class TestOVHCloudConfig: + def test_transform_request_basic(self): + """Test basic request transformation""" + transformed_request = config.transform_request( + model, + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["model"] == model + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_transform_request_with_extra_body(self): + """Test request transformation with extra_body parameters""" + transformed_request = config.transform_request( + model, + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={"extra_body": {"custom_param": "custom_value"}}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["custom_param"] == "custom_value" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + non_default_params = { + "temperature": 0.7, + "max_tokens": 100, + "top_p": 0.9, + } + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=False, + ) + + assert mapped_params["temperature"] == 0.7 + assert mapped_params["max_tokens"] == 100 + assert mapped_params["top_p"] == 0.9 + + def test_get_error_class(self): + """Test error class creation""" + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, OVHCloudException) + assert error.message == "Test error" + assert error.status_code == 400 + + @pytest.mark.parametrize( + "model", + [ + "Meta-Llama-3_3-70B-Instruct", + "Meta-Llama-3_1-70B-Instruct", + "Mixtral-8x7B-Instruct-v0.1", + "gpt-oss-120b", + "some-model-not-in-the-cost-map", + ], + ) + def test_tools_not_filtered_by_static_model_map(self, model): + """ + OVHCloud AI Endpoints are OpenAI-compatible; tools/tool_choice must pass + through for any model. The server is responsible for rejecting unsupported + tool calls — LiteLLM must not strip them based on a stale static catalog. + """ + + params = get_optional_params( + model=model, + custom_llm_provider="ovhcloud", + tools=[ + { + "type": "function", + "function": {"name": "x", "parameters": {}}, + } + ], + tool_choice="auto", + ) + + assert "tools" in params + assert "tool_choice" in params + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +class TestOVHCloudReasoningFieldMigration: + """Tests for OVHCloud reasoning_content -> reasoning field migration.""" + + def test_streaming_new_reasoning_field(self): + """New `reasoning` field should be mapped to `reasoning_content`.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "role": "assistant", + "reasoning": "Let me think...", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "Let me think..." + + def test_streaming_legacy_reasoning_content_unchanged(self): + """Legacy `reasoning_content` field should pass through untouched.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "role": "assistant", + "reasoning_content": "Already correct field.", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "Already correct field." + + def test_streaming_both_fields_legacy_wins(self): + """When both fields present, existing `reasoning_content` is not overwritten.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "reasoning": "new field", + "reasoning_content": "legacy field", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "legacy field" diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py similarity index 100% rename from tests/test_litellm/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py rename to tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/unit/llms/pass_through/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py rename to tests/unit/llms/pass_through/__init__.py diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/__init__.py b/tests/unit/llms/pass_through/guardrail_translation/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/text_to_speech/__init__.py rename to tests/unit/llms/pass_through/guardrail_translation/__init__.py diff --git a/tests/test_litellm/llms/perplexity/test_perplexity.py b/tests/unit/llms/perplexity/test_perplexity.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity.py rename to tests/unit/llms/perplexity/test_perplexity.py diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py rename to tests/unit/llms/perplexity/test_perplexity_cost_calculator.py diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py b/tests/unit/llms/perplexity/test_perplexity_integration.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity_integration.py rename to tests/unit/llms/perplexity/test_perplexity_integration.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py b/tests/unit/llms/pg_vector/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py rename to tests/unit/llms/pg_vector/__init__.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py b/tests/unit/llms/pg_vector/vector_stores/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py rename to tests/unit/llms/pg_vector/vector_stores/__init__.py diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/unit/llms/pg_vector/vector_stores/test_pg_vector_transformation.py similarity index 100% rename from tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py rename to tests/unit/llms/pg_vector/vector_stores/test_pg_vector_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/__init__.py b/tests/unit/llms/reducto/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/realtime/__init__.py rename to tests/unit/llms/reducto/__init__.py diff --git a/tests/test_litellm/llms/reducto/conftest.py b/tests/unit/llms/reducto/conftest.py similarity index 100% rename from tests/test_litellm/llms/reducto/conftest.py rename to tests/unit/llms/reducto/conftest.py diff --git a/tests/test_litellm/llms/reducto/test_cost.py b/tests/unit/llms/reducto/test_cost.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_cost.py rename to tests/unit/llms/reducto/test_cost.py diff --git a/tests/test_litellm/llms/reducto/test_model_info.py b/tests/unit/llms/reducto/test_model_info.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_model_info.py rename to tests/unit/llms/reducto/test_model_info.py diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/unit/llms/reducto/test_parse_legacy.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_parse_legacy.py rename to tests/unit/llms/reducto/test_parse_legacy.py diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/unit/llms/reducto/test_parse_v3.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_parse_v3.py rename to tests/unit/llms/reducto/test_parse_v3.py diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/unit/llms/reducto/test_upload.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_upload.py rename to tests/unit/llms/reducto/test_upload.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py b/tests/unit/llms/sagemaker/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py rename to tests/unit/llms/sagemaker/__init__.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py rename to tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py rename to tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py rename to tests/unit/llms/sagemaker/test_sagemaker_common_utils.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py rename to tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_role_assumption.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py rename to tests/unit/llms/sagemaker/test_sagemaker_embedding_role_assumption.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py rename to tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py rename to tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py diff --git a/tests/test_litellm/llms/voyage/rerank/__init__.py b/tests/unit/llms/sambanova/__init__.py similarity index 100% rename from tests/test_litellm/llms/voyage/rerank/__init__.py rename to tests/unit/llms/sambanova/__init__.py diff --git a/tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py b/tests/unit/llms/sambanova/tests_sambanova_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py rename to tests/unit/llms/sambanova/tests_sambanova_embedding_transformation.py diff --git a/tests/test_litellm/llms/watsonx/__init__.py b/tests/unit/llms/sap/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/__init__.py rename to tests/unit/llms/sap/chat/__init__.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_chat_calls.py b/tests/unit/llms/sap/chat/test_sap_chat_calls.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_chat_calls.py rename to tests/unit/llms/sap/chat/test_sap_chat_calls.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_langchain_strict_param.py b/tests/unit/llms/sap/chat/test_sap_langchain_strict_param.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_langchain_strict_param.py rename to tests/unit/llms/sap/chat/test_sap_langchain_strict_param.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_response_format.py b/tests/unit/llms/sap/chat/test_sap_response_format.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_response_format.py rename to tests/unit/llms/sap/chat/test_sap_response_format.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_tool_parameters.py b/tests/unit/llms/sap/chat/test_sap_tool_parameters.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_tool_parameters.py rename to tests/unit/llms/sap/chat/test_sap_tool_parameters.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_transformation.py b/tests/unit/llms/sap/chat/test_sap_transformation.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_transformation.py rename to tests/unit/llms/sap/chat/test_sap_transformation.py diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/__init__.py b/tests/unit/llms/sap/embed/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/audio_transcription/__init__.py rename to tests/unit/llms/sap/embed/__init__.py diff --git a/tests/test_litellm/llms/sap/embed/test_sap_embed_transformation.py b/tests/unit/llms/sap/embed/test_sap_embed_transformation.py similarity index 100% rename from tests/test_litellm/llms/sap/embed/test_sap_embed_transformation.py rename to tests/unit/llms/sap/embed/test_sap_embed_transformation.py diff --git a/tests/test_litellm/llms/sap/embed/test_sap_embedding.py b/tests/unit/llms/sap/embed/test_sap_embedding.py similarity index 100% rename from tests/test_litellm/llms/sap/embed/test_sap_embedding.py rename to tests/unit/llms/sap/embed/test_sap_embedding.py diff --git a/tests/test_litellm/llms/watsonx/rerank/__init__.py b/tests/unit/llms/snowflake/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/rerank/__init__.py rename to tests/unit/llms/snowflake/chat/__init__.py diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py rename to tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py diff --git a/tests/test_litellm/llms/you_com/__init__.py b/tests/unit/llms/snowflake/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/you_com/__init__.py rename to tests/unit/llms/snowflake/embedding/__init__.py diff --git a/tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py b/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py similarity index 100% rename from tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py rename to tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py diff --git a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py index 7970f7771fc..344b8e5573d 100644 --- a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py @@ -7,7 +7,7 @@ Covers: - Claude models → /messages (Anthropic format) Run: - pytest tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py -v + pytest tests/unit/llms/snowflake/test_snowflake_native_endpoints.py -v """ import json diff --git a/tests/test_litellm/llms/soniox/audio_transcription/__init__.py b/tests/unit/llms/soniox/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/__init__.py rename to tests/unit/llms/soniox/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py rename to tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py rename to tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/test_cache_control_and_reasoning.py b/tests/unit/llms/test_cache_control_and_reasoning.py similarity index 100% rename from tests/test_litellm/llms/test_cache_control_and_reasoning.py rename to tests/unit/llms/test_cache_control_and_reasoning.py diff --git a/tests/test_litellm/llms/test_file_content_block.py b/tests/unit/llms/test_file_content_block.py similarity index 100% rename from tests/test_litellm/llms/test_file_content_block.py rename to tests/unit/llms/test_file_content_block.py diff --git a/tests/test_litellm/llms/test_file_search_responses.py b/tests/unit/llms/test_file_search_responses.py similarity index 100% rename from tests/test_litellm/llms/test_file_search_responses.py rename to tests/unit/llms/test_file_search_responses.py diff --git a/tests/test_litellm/llms/test_lifecycle_fix.py b/tests/unit/llms/test_lifecycle_fix.py similarity index 100% rename from tests/test_litellm/llms/test_lifecycle_fix.py rename to tests/unit/llms/test_lifecycle_fix.py diff --git a/tests/test_litellm/llms/test_polling_url_origin_match.py b/tests/unit/llms/test_polling_url_origin_match.py similarity index 100% rename from tests/test_litellm/llms/test_polling_url_origin_match.py rename to tests/unit/llms/test_polling_url_origin_match.py diff --git a/tests/test_litellm/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py similarity index 100% rename from tests/test_litellm/llms/test_predibase_transformation.py rename to tests/unit/llms/test_predibase_transformation.py diff --git a/tests/unit/llms/tinyfish/__init__.py b/tests/unit/llms/tinyfish/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/unit/llms/tinyfish/test_tinyfish_search.py similarity index 100% rename from tests/test_litellm/llms/tinyfish/test_tinyfish_search.py rename to tests/unit/llms/tinyfish/test_tinyfish_search.py diff --git a/tests/test_litellm/llms/vercel_ai_gateway/test_vercel_ai_gateway.py b/tests/unit/llms/vercel_ai_gateway/test_vercel_ai_gateway.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/test_vercel_ai_gateway.py rename to tests/unit/llms/vercel_ai_gateway/test_vercel_ai_gateway.py diff --git a/tests/unit/llms/vertex_ai/audio_transcription/__init__.py b/tests/unit/llms/vertex_ai/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py diff --git a/tests/unit/llms/vertex_ai/batches/__init__.py b/tests/unit/llms/vertex_ai/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/test_handler.py rename to tests/unit/llms/vertex_ai/batches/test_handler.py diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/unit/llms/vertex_ai/batches/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/test_transformation.py rename to tests/unit/llms/vertex_ai/batches/test_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_transformation.py b/tests/unit/llms/vertex_ai/files/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_transformation.py rename to tests/unit/llms/vertex_ai/files/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/gemini/__init__.py b/tests/unit/llms/vertex_ai/gemini/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py b/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py rename to tests/unit/llms/vertex_ai/gemini/test_context_circulation.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py b/tests/unit/llms/vertex_ai/gemini/test_function_call_args_serialization.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py rename to tests/unit/llms/vertex_ai/gemini/test_function_call_args_serialization.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py b/tests/unit/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py rename to tests/unit/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py b/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py rename to tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_grounding_requests.py b/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_grounding_requests.py rename to tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py rename to tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py b/tests/unit/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py rename to tests/unit/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py rename to tests/unit/llms/vertex_ai/gemini/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py new file mode 100644 index 00000000000..4f23ac1773a --- /dev/null +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -0,0 +1,2729 @@ +import base64 + +import pytest + +from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_result, +) +from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + _transform_request_body, + check_if_part_exists_in_parts, + _get_highest_media_resolution, + _extract_max_media_resolution_from_messages, +) +from litellm.types.llms.vertex_ai import BlobType +from litellm.types.utils import Message + + +def test_check_if_part_exists_in_parts(): + parts = [ + {"text": "Hello", "thought": True}, + {"text": "World", "thought": False}, + ] + part = {"text": "Hello", "thought": True} + new_part = {"text": "Hello World", "thought": True} + assert check_if_part_exists_in_parts(parts, part) + assert not check_if_part_exists_in_parts(parts, new_part, ["thought"]) + assert check_if_part_exists_in_parts(parts, new_part, ["text"]) + + +def test_check_if_part_exists_in_parts_camel_case_snake_case(): + """Test that function handles both camelCase and snake_case key variations""" + # Test snake_case to camelCase matching + parts_with_snake_case = [ + { + "function_call": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + }, + {"text": "Some other content"}, + ] + + part_with_camel_case = { + "functionCall": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + } + + # Should find match between function_call and functionCall + assert check_if_part_exists_in_parts(parts_with_snake_case, part_with_camel_case) + + # Test camelCase to snake_case matching + parts_with_camel_case = [ + {"functionCall": {"name": "calculate_sum", "args": {"a": 1, "b": 2}}} + ] + + part_with_snake_case = { + "function_call": {"name": "calculate_sum", "args": {"a": 1, "b": 2}} + } + + # Should find match between functionCall and function_call + assert check_if_part_exists_in_parts(parts_with_camel_case, part_with_snake_case) + + # Test no match when values differ + part_with_different_values = { + "function_call": {"name": "different_function", "args": {"x": 5}} + } + + assert not check_if_part_exists_in_parts( + parts_with_snake_case, part_with_different_values + ) + + # Test multiple keys with mixed casing + parts_mixed = [ + { + "function_call": {"name": "test"}, + "thoughtSignature": "reasoning", + "text": "content", + } + ] + + part_mixed_casing = { + "functionCall": {"name": "test"}, + "thought_signature": "reasoning", + "text": "content", + } + + assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) + + +def test_cached_content_respects_modify_params_for_cache_incompatible_fields(): + """Regression: cachedContent drops system/tools/toolConfig only when modify_params=True.""" + import litellm + + cache_name = "projects/p/locations/us-central1/cachedContents/abc123" + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "hi"}, + ] + optional_params = { + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": "Get weather"}, + ] + } + ], + "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, + } + + original_modify_params = litellm.modify_params + try: + # With modify_params=False (default), keep fields even with cachedContent. + litellm.modify_params = False + result = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result.get("cachedContent") == cache_name + assert "system_instruction" in result + assert "tools" in result + assert "toolConfig" in result + assert "contents" in result + + # With modify_params=True, drop cache-incompatible fields. + litellm.modify_params = True + result_modify_true = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result_modify_true.get("cachedContent") == cache_name + assert "system_instruction" not in result_modify_true + assert "tools" not in result_modify_true + assert "toolConfig" not in result_modify_true + assert "contents" in result_modify_true + + # Without cache, fields are always included. + result_no_cache = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + assert "system_instruction" in result_no_cache + assert "tools" in result_no_cache + assert "toolConfig" in result_no_cache + finally: + litellm.modify_params = original_modify_params + + +# Tests for issue #14556: Labels field provider-aware filtering +def test_google_genai_excludes_labels(): + """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="gemini", + litellm_params=litellm_params, + cached_content=None, + ) + + # Google GenAI/AI Studio should NOT include labels + assert "labels" not in result + assert "contents" in result + + +def test_vertex_ai_includes_labels(): + """Test that Vertex AI endpoints include labels when custom_llm_provider='vertex_ai'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # Vertex AI SHOULD include labels + assert "labels" in result + assert result["labels"] == {"project": "test", "team": "ai"} + + +def test_service_tier_forwarded_to_vertex_ai(): + """Test that service_tier in optional_params is mapped to serviceTier in request body.""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"service_tier": "flex"} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "serviceTier" in result + assert result["serviceTier"] == "flex" + + +def test_extra_body_cache_not_forwarded_to_vertex_ai(): + """ + 'cache' inside extra_body is a LiteLLM-internal proxy caching control. + It must NOT be forwarded to the Vertex AI request body. + + Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." + Vertex AI enforces a strict JSON schema and rejects any unknown field. + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal + "some_vertex_param": "value", # legitimate provider extra + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # 'cache' must be stripped — Vertex AI has no such field + assert "cache" not in result, ( + "extra_body.cache must not be forwarded to Vertex AI. " + 'Vertex AI rejects it with 400: Unknown name "cache": Cannot find field.' + ) + + # Other legitimate extra_body keys should still pass through + assert "some_vertex_param" in result + assert result["some_vertex_param"] == "value" + + # Core request fields must be present + assert "contents" in result + + +def test_extra_body_tags_not_forwarded_to_vertex_ai(): + """ + 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. + It must NOT be forwarded to the Vertex AI request body. + Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "tags": ["user:alice", "env:prod"], + "custom_param": "allowed", + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "tags" not in result + assert "custom_param" in result + assert result["custom_param"] == "allowed" + + +def test_extra_body_google_maps_rewrites_json_response_format(): + messages = [{"role": "user", "content": "test"}] + optional_params = { + "response_mime_type": "application/json", + "response_schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + "extra_body": { + "tools": [{"googleMaps": {}}], + }, + } + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + generation_config = result["generationConfig"] + assert "response_mime_type" not in generation_config + assert generation_config["responseFormat"] == { + "text": { + "mimeType": "APPLICATION_JSON", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } + } + + +def test_extra_body_generation_config_cannot_restore_google_maps_json_mime_type(): + messages = [{"role": "user", "content": "test"}] + optional_params = { + "tools": [{"googleMaps": {}}], + "response_mime_type": "application/json", + "extra_body": { + "generationConfig": { + "response_mime_type": "application/json", + "response_json_schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + }, + }, + } + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + generation_config = result["generationConfig"] + assert "response_mime_type" not in generation_config + assert "response_json_schema" not in generation_config + assert generation_config["responseFormat"] == { + "text": { + "mimeType": "APPLICATION_JSON", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } + } + + +def test_metadata_to_labels_vertex_only(): + """Test that metadata->labels conversion only happens for Vertex AI""" + messages = [{"role": "user", "content": "test"}] + optional_params = {} + litellm_params = { + "metadata": { + "requester_metadata": {"user": "john_doe", "project": "test-project"} + } + } + + # Google GenAI/AI Studio should not include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="gemini", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" not in result + + # Vertex AI should include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="vertex_ai", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" in result + assert result["labels"] == {"user": "john_doe", "project": "test-project"} + + +def test_empty_content_handling(): + """Test that empty content strings are properly handled in Gemini message transformation""" + # Test with empty content in user message + messages = [{"content": "", "role": "user"}] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify that the content was properly transformed + assert len(contents) == 1 + assert contents[0]["role"] == "user" + assert len(contents[0]["parts"]) == 1 + assert "text" in contents[0]["parts"][0] + assert contents[0]["parts"][0]["text"] == "" + + +def test_thought_signature_extraction_from_response(): + """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + # Test case: Single function call with thought signature + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, + ) + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + # Verify thought signature is stored in provider_specific_fields + assert tools is not None + assert len(tools) == 1 + assert "provider_specific_fields" in tools[0] + assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature + + +def test_thought_signature_parallel_function_calls(): + """Test that only the first function call in parallel calls has thought signature""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Parallel function calls - only first has signature + parts_parallel = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, # First FC has signature + ), + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "London"}, + }, + # Second FC has no signature (parallel call) + ), + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_parallel, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + # Verify only first tool call has thought signature + assert tools is not None + assert len(tools) == 2 + assert "provider_specific_fields" in tools[0] + assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature + # Second tool call should not have thought signature + assert "provider_specific_fields" not in tools[ + 1 + ] or "thought_signature" not in tools[1].get("provider_specific_fields", {}) + + +def test_thought_signature_preservation_in_conversion(): + """Test that thought signatures are preserved when converting assistant messages back to Gemini format""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Assistant message with tool calls containing thought signatures + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": test_signature, + }, + }, + { + "id": "call_def456", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "London"}', + }, + "index": 1, + # No thought signature for parallel call + }, + ], + } + + gemini_parts = convert_to_gemini_tool_call_invoke(assistant_message) + + # Verify thought signature is preserved in first function call part + assert len(gemini_parts) == 2 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + assert gemini_parts[0]["thoughtSignature"] == test_signature + + # Verify second function call part does not have thought signature + assert "function_call" in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[1] + + +def test_thought_signature_sequential_function_calls(): + """Test that each sequential function call preserves its own thought signature""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + signature_1 = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + signature_2 = "DifferentSignatureForSecondCall1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ" + + # Sequential function calls - each has its own signature + # This simulates a multi-step conversation where each step has a signature + assistant_message_step1 = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_step1", + "type": "function", + "function": { + "name": "check_flight", + "arguments": '{"flight": "AA100"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": signature_1, + }, + }, + ], + } + + assistant_message_step2 = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_step2", + "type": "function", + "function": { + "name": "book_taxi", + "arguments": '{"destination": "airport"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": signature_2, + }, + }, + ], + } + + gemini_parts_step1 = convert_to_gemini_tool_call_invoke(assistant_message_step1) + gemini_parts_step2 = convert_to_gemini_tool_call_invoke(assistant_message_step2) + + # Verify each step preserves its own signature + assert len(gemini_parts_step1) == 1 + assert gemini_parts_step1[0]["thoughtSignature"] == signature_1 + + assert len(gemini_parts_step2) == 1 + assert gemini_parts_step2[0]["thoughtSignature"] == signature_2 + + +def test_thought_signature_with_function_call_mode(): + """Test thought signature extraction in function_call mode (is_function_call=True)""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_weather", + "args": {"location": "Tokyo"}, + }, + thoughtSignature=test_signature, + ) + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=True, + ) + + # Verify thought signature is stored in function's provider_specific_fields + assert function is not None + # Function should be dict-like (TypedDict or dict) + assert hasattr(function, "__getitem__") or isinstance(function, dict) + assert "provider_specific_fields" in function + assert function["provider_specific_fields"]["thought_signature"] == test_signature + assert tools is None + + +def test_dummy_signature_added_for_gemini_3_conversation_history(): + """Test that dummy signatures are added when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3.""" + import base64 + + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Simulate conversation history from gemini-2.5-flash (no thought signature) + assistant_message_from_older_model = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + # No provider_specific_fields - older model doesn't provide signatures + }, + ], + } + + # Convert to Gemini format for gemini-3-pro-preview (should add dummy signature) + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_from_older_model, model="gemini-3-pro-preview" + ) + + # Verify dummy signature is added + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + + # Verify it's the expected dummy signature (base64 encoded "skip_thought_signature_validator") + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + + +def test_dummy_signature_not_added_for_gemini_2_5(): + """Test that dummy signatures are NOT added when target model is not gemini-3.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Simulate conversation history from gemini-2.5-flash (no thought signature) + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + # No provider_specific_fields + }, + ], + } + + # Convert to Gemini format for gemini-2.5-flash (should NOT add dummy signature) + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message, model="gemini-2.5-flash" + ) + + # Verify no dummy signature is added for non-gemini-3 models + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" not in gemini_parts[0] + + +def test_dummy_signature_not_added_when_signature_exists(): + """Test that dummy signatures are NOT added when a real signature already exists.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + real_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Assistant message with existing thought signature + assistant_message_with_signature = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + "provider_specific_fields": { + "thought_signature": real_signature, + }, + }, + "index": 0, + }, + ], + } + + # Convert to Gemini format for gemini-3-pro-preview + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_with_signature, model="gemini-3-pro-preview" + ) + + # Verify real signature is preserved, not replaced with dummy + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + assert gemini_parts[0]["thoughtSignature"] == real_signature + + +def test_dummy_signature_with_function_call_mode(): + """Test that dummy signatures are added for function_call mode when converting to gemini-3.""" + import base64 + + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Assistant message with function_call (not tool_calls) and no signature + assistant_message_function_call = { + "role": "assistant", + "content": None, + "function_call": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + # No provider_specific_fields + }, + } + + # Convert to Gemini format for gemini-3-pro-preview + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_function_call, model="gemini-3-pro-preview" + ) + + # Verify dummy signature is added + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + + # Verify it's the expected dummy signature + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + + +def _parallel_tool_calls(*signatures): + return [ + { + "id": f"call_{idx}", + "type": "function", + "function": { + "name": f"tool_{idx}", + "arguments": '{"location": "Paris"}', + **( + {"provider_specific_fields": {"thought_signature": signature}} + if signature is not None + else {} + ), + }, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +def _parallel_tool_calls_signed_via_id(*signatures): + """Parallel tool calls in the shape LiteLLM actually hands back to clients. + + The signature rides in the tool call id behind __thought__, which is what an + OpenAI-format client echoes back on the next turn. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _encode_tool_call_id_with_signature, + ) + + return [ + { + "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), + "type": "function", + "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" +PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" +) + + +def test_dummy_signature_only_on_first_parallel_tool_call(): + """Google documents the placeholder as a last resort that degrades quality, so an unsigned + parallel turn replayed to gemini-3 gets a budget of exactly one.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): + """Gemini signs only the first of N parallel function calls, so a faithful replay has + nothing to attach to the siblings.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_later_parallel_tool_call_is_preserved(): + """Clients may reorder or drop calls, so a signature that lands on a non-first call is + still the model's own and must survive the round trip.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + + +def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): + """Non-gemini-3 models never get a placeholder signature, on any call.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): + """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): + """A signature on the tool call itself, rather than on its function, behaves the same way.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = _parallel_tool_calls(None, None) + tool_calls[0]["provider_specific_fields"] = { + "thought_signature": REAL_THOUGHT_SIGNATURE + } + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): + """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not + consume the one placeholder slot and leave the real first function call bare.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = [ + {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} + ] + _parallel_tool_calls(None, None) + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_no_placeholder_when_model_is_unknown(): + """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): + """Older models still receive a real signature that a client replays, and still get no placeholder.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_parallel_tool_call_history_replayed_through_full_message_conversion(): + """End to end through the message-history converter, the path a real /chat/completions replay takes.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + + +@pytest.mark.parametrize( + "model", + ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], +) +def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): + """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. + + Fabricating the placeholder alongside a real signature is what produced empty text responses + on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages, model=model) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) + + +@pytest.mark.parametrize( + "model", + [ + "gemini-3-pro-preview", + "gemini-3-flash-preview", + "gemini-3.1-pro-preview", + "gemini-3.5-flash", + "gemini-3.6-flash", + "gemini-3.7-flash", + "gemini-3.8-flash", + "vertex_ai/gemini-3.5-flash", + "vertex_ai/gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", + "gemini/gemini-3.5-flash", + "gemini/gemini-3.7-flash", + "gemini/gemini-3.8-flash", + ], +) +def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): + """The gemini-3 gate is a substring match, so every family member and prefix form has to + land on the same one-placeholder budget rather than only the versions we happened to try.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model=model, + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): + """Text-part and function-call signatures are collected by separate code paths, so scoping the + placeholder must not disturb a real signature that arrived on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Checking all three cities.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": _parallel_tool_calls(None, None, None), + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-3-pro-preview" + )[0]["parts"] + + assert parts[0]["text"] == "Checking all three cities." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in parts[2] + assert "thoughtSignature" not in parts[3] + + +# Tests for media_resolution (detail parameter) handling - Issue #17084 +class TestMediaResolution: + """Tests for media_resolution handling in Gemini 2.x models""" + + def test_get_highest_media_resolution_high_wins(self): + """Test that 'high' resolution takes precedence over 'low'""" + assert _get_highest_media_resolution("low", "high") == "high" + assert _get_highest_media_resolution("high", "low") == "high" + assert _get_highest_media_resolution(None, "high") == "high" + assert _get_highest_media_resolution("high", None) == "high" + + def test_get_highest_media_resolution_low_over_none(self): + """Test that 'low' resolution takes precedence over None""" + assert _get_highest_media_resolution(None, "low") == "low" + assert _get_highest_media_resolution("low", None) == "low" + + def test_get_highest_media_resolution_same_values(self): + """Test handling of same resolution values""" + assert _get_highest_media_resolution("high", "high") == "high" + assert _get_highest_media_resolution("low", "low") == "low" + assert _get_highest_media_resolution(None, None) is None + + def test_get_highest_media_resolution_medium(self): + """Test that 'medium' resolution is correctly ranked between 'low' and 'high'""" + assert _get_highest_media_resolution("low", "medium") == "medium" + assert _get_highest_media_resolution("medium", "low") == "medium" + assert _get_highest_media_resolution("medium", "high") == "high" + assert _get_highest_media_resolution("high", "medium") == "high" + assert _get_highest_media_resolution(None, "medium") == "medium" + assert _get_highest_media_resolution("medium", None) == "medium" + + def test_get_highest_media_resolution_ultra_high(self): + """Test that 'ultra_high' resolution takes precedence over all others""" + assert _get_highest_media_resolution("high", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("ultra_high", "high") == "ultra_high" + assert _get_highest_media_resolution("medium", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("low", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution(None, "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("ultra_high", None) == "ultra_high" + + def test_extract_max_media_resolution_single_image_high(self): + """Test extraction of media resolution from single image with detail=high""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_single_image_low(self): + """Test extraction of media resolution from single image with detail=low""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "low" + + def test_extract_max_media_resolution_no_detail(self): + """Test extraction when no detail parameter is provided""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc123"}, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) is None + + def test_extract_max_media_resolution_multiple_images_mixed(self): + """Test that highest resolution is returned when multiple images have different details""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare these images"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,def456", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_text_only(self): + """Test extraction from messages with no images""" + messages = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well!"}, + ] + assert _extract_max_media_resolution_from_messages(messages) is None + + def test_transform_request_body_gemini_2x_adds_media_resolution(self): + """Test that media_resolution is added to generationConfig for Gemini 2.x models""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert "generationConfig" in result + assert "mediaResolution" in result["generationConfig"] + assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_HIGH" + + def test_transform_request_body_gemini_2x_low_resolution(self): + """Test that low media_resolution is correctly added for Gemini 2.x""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "low", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert "generationConfig" in result + assert "mediaResolution" in result["generationConfig"] + assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_LOW" + + def test_transform_request_body_gemini_3_no_global_media_resolution(self): + """Test that Gemini 3 models don't add media_resolution to generationConfig (they use per-part)""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-3-pro-preview", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # Gemini 3 should NOT have mediaResolution in generationConfig + # (it's handled per-part in the content transformation) + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + def test_transform_request_body_no_detail_no_media_resolution(self): + """Test that no mediaResolution is added when detail is not specified""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # When no detail is specified, mediaResolution should not be in generationConfig + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + def test_extract_max_media_resolution_file_type_with_detail(self): + """Test that detail is extracted from file content type, not just image_url""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this file?"}, + { + "type": "file", + "file": { + "url": "data:image/png;base64,abc123", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_mixed_image_and_file(self): + """Test that highest detail is returned across both image_url and file types""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare these"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + { + "type": "file", + "file": { + "url": "data:image/png;base64,def456", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_transform_request_body_gemini_1x_no_media_resolution(self): + """Test that Gemini 1.x models don't get mediaResolution in generationConfig""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-1.5-pro", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # Gemini 1.x should NOT have mediaResolution (not supported) + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + +# Tests for VideoMetadata support across all Gemini models (Issue #25474) +class TestVideoMetadataAllGeminiModels: + """Tests that video_metadata (fps, start_offset, end_offset) works for all Gemini models""" + + def _make_video_messages(self, video_metadata: dict) -> list: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": video_metadata, + }, + }, + ], + } + ] + + def _get_file_part(self, contents: list) -> dict: + for part in contents[0]["parts"]: + if "file_data" in part: + return part + raise AssertionError("No file part found in contents") + + def test_video_metadata_fps_gemini_2_5_flash(self): + """Gemini 2.5 Flash: fps in video_metadata should be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 5}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 5 + + def test_video_metadata_fps_gemini_2_5_pro(self): + """Gemini 2.5 Pro: fps in video_metadata should be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 10}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-pro" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 10 + + def test_video_metadata_offsets_gemini_2_5_flash(self): + """Gemini 2.5 Flash: start_offset/end_offset converted to camelCase (Issue #25474)""" + messages = self._make_video_messages( + {"start_offset": "5s", "end_offset": "30s"} + ) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + vm = file_part["video_metadata"] + assert vm["startOffset"] == "5s" + assert vm["endOffset"] == "30s" + + def test_video_metadata_all_fields_gemini_2_5_flash(self): + """Gemini 2.5 Flash: all video_metadata fields forwarded correctly (Issue #25474)""" + messages = self._make_video_messages( + {"fps": 5, "start_offset": "10s", "end_offset": "60s"} + ) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + vm = file_part["video_metadata"] + assert vm["fps"] == 5 + assert vm["startOffset"] == "10s" + assert vm["endOffset"] == "60s" + + def test_video_metadata_gemini_1_5_pro(self): + """Gemini 1.5 Pro: video_metadata should also be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 2}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-1.5-pro" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 2 + + +def test_convert_tool_response_with_base64_image(): + """Test tool response with base64 data URI image.""" + # Create a small test image (1x1 red pixel PNG) + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create tool message with image + tool_message = { + "role": "tool", + "tool_call_id": "call_test123", + "content": [ + { + "type": "text", + "text": '{"url": "https://example.com", "status": "success"}', + }, + {"type": "input_image", "image_url": image_data_uri}, + ], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_test123", + "function": {"name": "click_at", "arguments": '{"x": 100, "y": 200}'}, + } + ] + } + + # Convert tool response with nested multimodal functionResponse.parts. + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + result_part = result[0] + assert "function_response" in result_part + assert "inline_data" not in result_part + function_response = result_part["function_response"] + assert function_response["name"] == "click_at" + assert "response" in function_response + # Verify JSON response is parsed correctly + assert "url" in function_response["response"] + assert function_response["response"]["url"] == "https://example.com" + + # Check inline_data is nested under functionResponse.parts. + assert "parts" in function_response + assert len(function_response["parts"]) == 1 + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "image/png" + assert inline_data["data"] == test_image_base64 + + +def test_gemini_history_nests_multimodal_tool_response_parts(): + """Full history conversion should not emit sibling inline_data tool result parts.""" + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + messages = [ + {"role": "user", "content": "Get me an image"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_get_image", + "type": "function", + "function": {"name": "get_image", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_get_image", + "content": [ + {"type": "text", "text": '{"image_ref": "inline"}'}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": test_image_base64, + }, + }, + ], + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + tool_response_parts = contents[-1]["parts"] + assert len(tool_response_parts) == 1 + assert "inline_data" not in tool_response_parts[0] + function_response = tool_response_parts[0]["function_response"] + assert function_response["parts"] == [ + { + "inline_data": { + "data": test_image_base64, + "mime_type": "image/png", + } + } + ] + + +def test_convert_tool_response_text_only(): + """Test tool response with only text (no image).""" + tool_message = { + "role": "tool", + "tool_call_id": "call_test789", + "content": [ + {"type": "text", "text": '{"status": "completed", "result": "success"}'} + ], + } + + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_test789", + "function": {"name": "wait_5_seconds", "arguments": "{}"}, + } + ] + } + + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Should be a single part (no list) when no image + assert not isinstance(result, list), "Should return single part when no image" + + # Check function_response exists + assert "function_response" in result + function_response = result["function_response"] + assert function_response["name"] == "wait_5_seconds" + # Verify JSON response is parsed correctly + assert "status" in function_response["response"] + assert function_response["response"]["status"] == "completed" + + # Check inline_data does NOT exist (no image provided) + assert "inline_data" not in result + + +def test_file_data_field_order(): + """ + Test that file_data fields are in the correct order (mime_type before file_uri). + + The Gemini API is sensitive to field order in the file_data object. + This test verifies that mime_type comes before file_uri in both: + 1. Dictionary key order + 2. JSON serialization + + Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + # Test with HTTPS URL and explicit format (audio file) + file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" + format = "audio/mpeg" + + result = _process_gemini_media(image_url=file_url, format=format) + + # Verify the result has file_data + assert "file_data" in result + file_data = result["file_data"] + + # Verify both fields are present + assert "mime_type" in file_data + assert "file_uri" in file_data + assert file_data["mime_type"] == "audio/mpeg" + assert file_data["file_uri"] == file_url + + # Verify field order by checking dictionary keys + # In Python 3.7+, dict maintains insertion order + file_data_keys = list(file_data.keys()) + assert file_data_keys.index("mime_type") < file_data_keys.index( + "file_uri" + ), "mime_type must come before file_uri in the file_data dict" + + # Also verify by serializing to JSON string + json_str = json.dumps(file_data) + mime_type_pos = json_str.find('"mime_type"') + file_uri_pos = json_str.find('"file_uri"') + assert ( + mime_type_pos < file_uri_pos + ), "mime_type must appear before file_uri in JSON serialization" + + +def test_file_data_field_order_gcs_urls(): + """Test that GCS URLs also maintain correct field order.""" + import json + + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + # Test with GCS URL + gcs_url = "gs://bucket/audio.mp3" + + result = _process_gemini_media(image_url=gcs_url) + + # Verify the result has file_data + assert "file_data" in result + file_data = result["file_data"] + + # Verify both fields are present + assert "mime_type" in file_data + assert "file_uri" in file_data + + # Verify field order + file_data_keys = list(file_data.keys()) + assert file_data_keys.index("mime_type") < file_data_keys.index( + "file_uri" + ), "mime_type must come before file_uri in the file_data dict" + + +def test_gemini_files_api_uri_without_format(): + """ + Test that Gemini Files API URIs work WITHOUT an explicit format/mime_type. + + When a user uploads a file via the Gemini Files API and then references it + by URI (https://generativelanguage.googleapis.com/v1beta/files/...), + the file is already on Google's servers. These URLs return 403 when + fetched directly, so _process_gemini_media must NOT try to resolve the + MIME type via HTTP. Instead it should pass the URI through as file_data + and let the Gemini API resolve the type from its stored metadata. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/37eh7rsw1vfe" + + # Should NOT raise — previously this hit the generic https:// handler + # which called _get_image_mime_type_from_url() and got a 403. + result = _process_gemini_media(image_url=file_url) + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + # When no format is provided, mime_type should be absent so the + # Gemini API infers it from the stored file metadata. + assert "mime_type" not in file_data + + +def test_gemini_files_api_uri_with_format(): + """ + Test that Gemini Files API URIs correctly forward an explicit format. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/n1vhxa28lyaw" + + result = _process_gemini_media(image_url=file_url, format="text/plain") + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + assert file_data["mime_type"] == "text/plain" + + +def test_extract_file_data_with_path_object(): + """ + Test that filename is correctly extracted from Path objects for MIME type detection. + + When uploading files using Path objects (e.g., Path("speech.mp3")), the filename + must be extracted to enable proper MIME type detection. Without this, files get + uploaded with 'application/octet-stream' instead of the correct MIME type. + + Related issue: Files uploaded with wrong MIME type cause Gemini API to reject + requests where the specified format doesn't match the uploaded file's MIME type. + """ + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + # Create a temporary MP3 file + with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp: + tmp.write(b"fake mp3 content") + tmp_path = tmp.name + + try: + # Test with Path object + path_obj = Path(tmp_path) + extracted = extract_file_data(path_obj) + + # Verify filename was extracted + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".mp3") + + # Verify MIME type was correctly detected + assert ( + extracted["content_type"] == "audio/mpeg" + ), f"Expected 'audio/mpeg' but got '{extracted['content_type']}'" + + # Verify content was read + assert extracted["content"] == b"fake mp3 content" + + finally: + # Clean up temporary file + os.unlink(tmp_path) + + +def test_extract_file_data_with_pathlib_path(): + """Test that filename is correctly extracted from pathlib.Path inputs. + Bare str paths are rejected — when this runs in a proxy request handler + the value is attacker-controlled and opening it as a path is an LFI.""" + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(b"fake wav content") + tmp_path = Path(tmp.name) + + try: + extracted = extract_file_data(tmp_path) + + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".wav") + assert extracted["content_type"] in [ + "audio/wav", + "audio/x-wav", + ], f"Expected 'audio/wav' or 'audio/x-wav' but got '{extracted['content_type']}'" + assert extracted["content"] == b"fake wav content" + finally: + os.unlink(str(tmp_path)) + + +def test_extract_file_data_with_tuple_format(): + """Test that tuple format (with explicit content_type) still works correctly.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + # Test with tuple format: (filename, content, content_type) + filename = "test_audio.mp3" + content = b"test audio content" + content_type = "audio/mpeg" + + extracted = extract_file_data((filename, content, content_type)) + + # Verify all fields are correct + assert extracted["filename"] == filename + assert extracted["content"] == content + assert extracted["content_type"] == content_type + + +def test_extract_file_data_fallback_to_octet_stream(): + """Unknown file types fall back to application/octet-stream.""" + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp: + tmp.write(b"unknown content") + tmp_path = Path(tmp.name) + + try: + extracted = extract_file_data(tmp_path) + + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".xyz123") + assert ( + extracted["content_type"] == "application/octet-stream" + ), f"Expected 'application/octet-stream' for unknown type, got '{extracted['content_type']}'" + finally: + os.unlink(str(tmp_path)) + + +def test_convert_tool_response_with_pdf_file(): + """Test tool response with PDF file content using file_data field.""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with file + tool_message = { + "role": "tool", + "tool_call_id": "call_pdf_test", + "content": [ + {"type": "text", "text": '{"status": "success", "pages": 1}'}, + {"type": "file", "file_data": file_data_uri}, + ], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_pdf_test", + "function": { + "name": "analyze_document", + "arguments": '{"path": "/tmp/doc.pdf"}', + }, + } + ] + } + + # Convert tool response with nested multimodal functionResponse.parts. + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + result_part = result[0] + assert "function_response" in result_part + assert "inline_data" not in result_part + function_response = result_part["function_response"] + assert function_response["name"] == "analyze_document" + assert "response" in function_response + # Verify JSON response is parsed correctly + assert "status" in function_response["response"] + assert function_response["response"]["status"] == "success" + + # Check inline_data is nested under functionResponse.parts. + assert "parts" in function_response + assert len(function_response["parts"]) == 1 + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "application/pdf" + assert inline_data["data"] == test_pdf_base64 + + +def test_convert_tool_response_with_input_file_type(): + """Test tool response with input_file content type (Responses API format).""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with input_file type + tool_message = { + "role": "tool", + "tool_call_id": "call_input_file_test", + "content": [{"type": "input_file", "file_data": file_data_uri}], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_input_file_test", + "function": {"name": "read_file", "arguments": "{}"}, + } + ] + } + + # Convert tool response + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Check inline_data is nested under functionResponse.parts. + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + function_response = result[0]["function_response"] + assert ( + function_response["parts"][0]["inline_data"]["mime_type"] == "application/pdf" + ) + + +def test_convert_tool_response_with_nested_file_object(): + """Test tool response with file content using nested file object format.""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with nested file object (OpenAI Agents SDK format) + tool_message = { + "role": "tool", + "tool_call_id": "call_nested_test", + "content": [{"type": "file", "file": {"file_data": file_data_uri}}], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_nested_test", + "function": {"name": "process_document", "arguments": "{}"}, + } + ] + } + + # Convert tool response + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Check inline_data is nested under functionResponse.parts. + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + function_response = result[0]["function_response"] + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "application/pdf" + assert inline_data["data"] == test_pdf_base64 + + +def test_assistant_message_with_images_field(): + """ + Test that assistant messages with images field are properly converted to Gemini format. + + This handles the case where an assistant message contains generated images in the + `images` field (e.g., from image generation models like gemini-2.5-flash-image). + The images should be converted to inline_data parts in the Gemini format. + """ + # Create a small test image (1x1 red pixel PNG) + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages with assistant message containing images field + messages = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a costume that says LiteLLM", + }, + { + "role": "assistant", + "content": "Here's your banana in a LiteLLM costume!", + "images": [ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + }, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure + assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}" + + # Verify user message + assert contents[0]["role"] == "user" + assert len(contents[0]["parts"]) == 1 + assert ( + contents[0]["parts"][0]["text"] + == "Generate an image of a banana wearing a costume that says LiteLLM" + ) + + # Verify assistant message + assert contents[1]["role"] == "model" + assert ( + len(contents[1]["parts"]) == 2 + ), f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}" + + # Find text part and inline_data part + text_part = None + inline_data_part = None + for part in contents[1]["parts"]: + if "text" in part: + text_part = part + elif "inline_data" in part: + inline_data_part = part + + # Verify text part + assert text_part is not None, "Missing text part in assistant message" + assert text_part["text"] == "Here's your banana in a LiteLLM costume!" + + # Verify inline_data part (image) + assert inline_data_part is not None, "Missing inline_data part in assistant message" + inline_data: BlobType = inline_data_part["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "image/png" + assert inline_data["data"] == test_image_base64 + + +def test_assistant_message_with_multiple_images(): + """Test that assistant messages with multiple images are properly converted.""" + # Create two test images + test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" + image1_data_uri = f"data:image/png;base64,{test_image1_base64}" + image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}" + + messages = [ + {"role": "user", "content": "Generate two images"}, + { + "role": "assistant", + "content": "Here are your images:", + "images": [ + { + "image_url": {"url": image1_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + }, + { + "image_url": {"url": image2_data_uri, "detail": "high"}, + "index": 1, + "type": "image_url", + }, + ], + }, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has 3 parts (1 text + 2 images) + assert contents[1]["role"] == "model" + assert ( + len(contents[1]["parts"]) == 3 + ), f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}" + + # Count inline_data parts + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert ( + len(inline_data_parts) == 2 + ), f"Expected 2 inline_data parts, got {len(inline_data_parts)}" + + # Verify first image + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64 + + # Verify second image + assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg" + assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64 + + +def test_assistant_message_with_images_using_message_object(): + """Test that Message objects with images field are properly converted.""" + # Create a small test image + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages using Message object (as returned by LiteLLM) + user_message = {"role": "user", "content": "Generate an image"} + + assistant_message = Message( + content="Here's your image!", + role="assistant", + tool_calls=None, + function_call=None, + images=[ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + ) + + messages = [user_message, assistant_message] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has both text and image + assert contents[1]["role"] == "model" + assert len(contents[1]["parts"]) == 2 + + # Verify image was converted + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64 + + +def test_assistant_message_with_images_in_conversation_history(): + """ + Test multi-turn conversation where assistant message with images is in history. + + This simulates the real use case where: + 1. User asks for image generation + 2. Assistant generates image (with images field) + 3. User asks follow-up question about the image + """ + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + messages = [ + {"role": "user", "content": "Generate an image of a cat"}, + { + "role": "assistant", + "content": "Here's a cat image:", + "images": [ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + }, + {"role": "user", "content": "Can you make it more colorful?"}, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure: user -> model (with image) -> user + assert len(contents) == 3 + assert contents[0]["role"] == "user" + assert contents[1]["role"] == "model" + assert contents[2]["role"] == "user" + + # Verify assistant message has image in history + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + + +def test_function_response_has_user_role(): + """ + Test that function response ContentType blocks include role="user". + + Gemini API only accepts two roles: "user" and "model". Function responses + must be sent with role="user". Previously, LiteLLM omitted the role field + entirely, causing 400 errors from the Gemini API. + + Fixes: https://github.com/BerriAI/litellm/issues/22003 + Fixes: https://github.com/BerriAI/litellm/issues/20690 + """ + messages = [ + {"role": "user", "content": "What is the weather in Berlin?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Berlin"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": '{"temperature": "15°C", "condition": "Cloudy"}', + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Expect: user -> model (functionCall) -> user (functionResponse) + assert len(contents) == 3 + + assert contents[0]["role"] == "user" + assert contents[1]["role"] == "model" + assert "function_call" in contents[1]["parts"][0] + + # The critical assertion: function response must have role="user" + assert contents[2]["role"] == "user" + assert "function_response" in contents[2]["parts"][0] + + +def test_multi_turn_function_calling_roles(): + """ + Test a full multi-turn function calling conversation produces correct roles. + + Simulates: user asks → model calls tool → tool responds → model answers → user asks again. + Every content block must have an explicit role of "user" or "model". + + Fixes: https://github.com/BerriAI/litellm/issues/22003 + """ + messages = [ + {"role": "user", "content": "What is the weather in Berlin?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_001", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Berlin"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_001", + "content": '{"temperature": "15°C"}', + }, + { + "role": "assistant", + "content": "The weather in Berlin is 15°C.", + }, + {"role": "user", "content": "And in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_002", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_002", + "content": '{"temperature": "18°C"}', + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Every content block must have a valid role + for i, content in enumerate(contents): + assert "role" in content, f"Content block {i} missing 'role' field" + assert content["role"] in ( + "user", + "model", + ), f"Content block {i} has invalid role: {content.get('role')}" + + # Verify the function response blocks specifically have role="user" + for i, content in enumerate(contents): + for part in content["parts"]: + if "function_response" in part: + assert ( + content["role"] == "user" + ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" + + +def test_gemini_thought_signature_preservation_real_response(): + """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + real_candidate = { + "content": { + "parts": [ + { + "text": "I will explain and then list files.", + "thoughtSignature": "mock_signature_from_text_part", + }, + { + "functionCall": { + "name": "list_files", + "args": {}, + } + }, + ] + } + } + + parts = real_candidate["content"]["parts"] + + content, reasoning_content = ( + VertexGeminiConfig().get_assistant_content_message(parts=parts) + ) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=parts + ) + ) + functions, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + msg: dict = {"role": "assistant"} + if content is not None: + msg["content"] = content + if tools: + msg["tool_calls"] = tools + if functions is not None: + msg["function_call"] = functions + if thought_signatures is not None: + msg["provider_specific_fields"] = { + "thought_signatures": thought_signatures + } + + converted_real = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted_real) == 1 + assert "parts" in converted_real[0] + parts_out = converted_real[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert ( + parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" + ) + assert "function_call" in parts_out[1] + assert "thoughtSignature" not in parts_out[1] + + +def test_gemini_thought_signature_deduplication_assumed_response(): + """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + pr_assumed_msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": { + "thought_signatures": ["mock_signature_63k"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "mock_signature_63k" + }, + } + ], + } + + converted_pr = _gemini_convert_messages_with_history( + messages=[pr_assumed_msg], + model="gemini-2.5-pro", + ) + + assert len(converted_pr) == 1 + assert "parts" in converted_pr[0] + parts_out = converted_pr[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert "thoughtSignature" not in parts_out[0] + assert "function_call" in parts_out[1] + assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" + + +def test_gemini_thought_signature_pure_text(): + """Test that thought signatures are preserved on the text part for responses with no tool calls.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Hello, I am a model.", + "provider_specific_fields": { + "thought_signatures": ["pure_text_signature"] + }, + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "text" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_text_signature" + + +def test_gemini_thought_signature_pure_tool_call(): + """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "thought_signatures": ["pure_tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "pure_tool_signature" + }, + } + ], + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "function_call" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" + + +def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): + """A text-part signature that differs from the tool-call signature must stay on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Some analysis.", + "provider_specific_fields": { + "thought_signatures": ["text_signature", "tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": {"thought_signature": "tool_signature"}, + } + ], + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + assert parts[0]["text"] == "Some analysis." + assert parts[0]["thoughtSignature"] == "text_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == "tool_signature" + + +def test_gemini_25_text_signature_survives_replay_to_gemini_3(): + """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real + text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + } + ], + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + assert parts[0]["text"] == "I will list the directory." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + + +def test_gemini_function_call_signature_round_trip_no_duplicate(): + """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and + re-serialized sends the signature exactly once, on the function-call part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + response_parts = [ + {"text": "I will calculate the result for you."}, + { + "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, + "thoughtSignature": "signature_from_function_call", + }, + ] + + config = VertexGeminiConfig() + content, _ = config.get_assistant_content_message(parts=response_parts) + thought_signatures = config._extract_thought_signatures_from_parts( + parts=response_parts + ) + _, tools, _ = VertexGeminiConfig._transform_parts( + parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False + ) + + msg = { + "role": "assistant", + "content": content, + "tool_calls": tools, + "provider_specific_fields": {"thought_signatures": thought_signatures}, + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] + assert signatures == ["signature_from_function_call"] + assert "thoughtSignature" not in parts[0] + assert "function_call" in parts[1] + + +def test_gemini_server_side_tool_signature_not_duplicated_on_text(): + """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "The weather in Buenos Aires is sunny.", + "provider_specific_fields": { + "thought_signatures": ["server_side_signature"], + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny"}, + "thought_signature": "server_side_signature", + } + ], + }, + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + text_part = next(p for p in parts if "text" in p) + assert "thoughtSignature" not in text_part + tool_call_part = next(p for p in parts if "toolCall" in p) + assert tool_call_part["thoughtSignature"] == "server_side_signature" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py similarity index 99% rename from tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py rename to tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 88ba7fc37d9..739744336a1 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1894,42 +1894,6 @@ def test_vertex_ai_tool_call_id_format(): ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" -def test_vertex_ai_code_line_length(): - """ - Test that the specific code line generating tool call IDs is within character limit. - - This is a meta-test to ensure the code change meets the 40-character requirement. - """ - import inspect - - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - # Get the source code of the _transform_parts method - source_lines = inspect.getsource(VertexGeminiConfig._transform_parts).split("\n") - - # Find the line that generates the ID - id_line = None - for line in source_lines: - if '"id": f"call_' in line and "uuid.uuid4().hex[:28]" in line: - id_line = line.strip() # Remove indentation for length check - break - - assert id_line is not None, "Could not find the ID generation line in source code" - - # Check that the line is 40 characters or less (excluding indentation) - line_length = len(id_line) - assert ( - line_length <= 40 - ), f"ID generation line is {line_length} characters, should be ≤40: {id_line}" - - # Verify it contains the expected UUID format - assert ( - "uuid.uuid4().hex[:28]" in id_line - ), f"Line should contain shortened UUID format: {id_line}" - - def test_vertex_ai_map_google_maps_tool_simple(): """ Test googleMaps tool transformation without location data. @@ -2530,8 +2494,6 @@ def test_fine_tuned_endpoint_and_gemma_get_no_gemini_3_default_temperature(model assert "temperature" not in mapped - - def _tool_call_messages(tool_call_id: str): return [ {"role": "user", "content": "hi"}, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py rename to tests/unit/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py diff --git a/tests/unit/llms/vertex_ai/image_generation/__init__.py b/tests/unit/llms/vertex_ai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py rename to tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py new file mode 100644 index 00000000000..a72a570c2a2 --- /dev/null +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -0,0 +1,637 @@ +from unittest.mock import MagicMock, patch + +import httpx + + +from litellm.llms.vertex_ai.image_generation import ( + get_vertex_ai_image_generation_config, +) +from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( + VertexAIGeminiImageGenerationConfig, +) +from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ( + VertexAIImagenImageGenerationConfig, +) + + +class TestVertexAIGeminiImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIGeminiImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("gemini-2.5-flash-image") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to candidate_count""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("candidate_count") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("aspectRatio") == "1:1" + + def test_map_openai_params_size_16_9(self): + """Test mapping 16:9 size""" + non_default_params = {"size": "1792x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("aspectRatio") == "16:9" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" + assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" + assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_get_supported_openai_params_includes_native_gemini_params(self): + """Test that native Gemini imageConfig params are supported""" + supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") + assert "aspectRatio" in supported + assert "aspect_ratio" in supported + assert "imageSize" in supported + assert "image_size" in supported + assert "imageConfig" in supported + + def test_map_openai_params_aspect_ratio_camel_case(self): + """Test mapping native aspectRatio parameter""" + result = self.config.map_openai_params({"aspectRatio": "9:16"}, {}, "gemini-3-pro-image-preview", False) + assert result["aspectRatio"] == "9:16" + + def test_map_openai_params_aspect_ratio_snake_case(self): + """Test mapping native aspect_ratio parameter""" + result = self.config.map_openai_params({"aspect_ratio": "16:9"}, {}, "gemini-3-pro-image-preview", False) + assert result["aspectRatio"] == "16:9" + + def test_map_openai_params_image_size_camel_case(self): + """Test mapping native imageSize parameter""" + result = self.config.map_openai_params({"imageSize": "4K"}, {}, "gemini-3-pro-image-preview", False) + assert result["imageSize"] == "4K" + + def test_map_openai_params_image_size_snake_case(self): + """Test mapping native image_size parameter""" + result = self.config.map_openai_params({"image_size": "2K"}, {}, "gemini-3-pro-image-preview", False) + assert result["imageSize"] == "2K" + + def test_map_openai_params_image_config_dict_stored_whole(self): + """imageConfig dict is stored as-is so all fields survive""" + result = self.config.map_openai_params( + {"imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}}, + {}, + "gemini-3.1-flash-image", + False, + ) + assert result["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "2K"} + + def test_map_openai_params_image_config_all_fields(self): + """All ImageConfig fields (personGeneration, imageOutputOptions) pass through""" + payload = { + "imageConfig": { + "aspectRatio": "9:16", + "imageSize": "4K", + "personGeneration": "DONT_ALLOW", + "imageOutputOptions": { + "mimeType": "image/jpeg", + "compressionQuality": 80, + }, + } + } + result = self.config.map_openai_params(payload, {}, "gemini-3.1-flash-image", False) + assert result["imageConfig"] == payload["imageConfig"] + + def test_map_openai_params_image_config_non_dict_warns_and_drops(self): + """Non-dict imageConfig is dropped with a warning, not silently discarded""" + with patch("litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation.verbose_logger") as mock_log: + result = self.config.map_openai_params( + {"imageConfig": "bad-string-value"}, {}, "gemini-3.1-flash-image", False + ) + assert "imageConfig" not in result + mock_log.warning.assert_called_once() + + def test_transform_image_generation_request_from_image_config(self): + """Full imageConfig dict is forwarded verbatim into generationConfig""" + full_config = { + "aspectRatio": "16:9", + "imageSize": "2K", + "personGeneration": "DONT_ALLOW", + "imageOutputOptions": {"mimeType": "image/jpeg", "compressionQuality": 85}, + } + mapped = self.config.map_openai_params( + {"imageConfig": full_config}, + {}, + "gemini-3.1-flash-image", + False, + ) + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image", + prompt="A nano banana on a desk", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"] == full_config + + def test_transform_image_generation_flat_params_override_image_config(self): + """Explicit flat params win over the same key inside imageConfig""" + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image", + prompt="A nano banana", + optional_params={ + "imageConfig": {"aspectRatio": "1:1", "personGeneration": "DONT_ALLOW"}, + "aspectRatio": "16:9", # should win + }, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" + assert request["generationConfig"]["imageConfig"]["personGeneration"] == "DONT_ALLOW" + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "contents" in request + assert "generationConfig" in request + assert request["generationConfig"]["responseModalities"] == ["IMAGE"] + assert request["contents"][0]["parts"][0]["text"] == "A nano banana" + + def test_transform_image_generation_request_with_aspect_ratio(self): + """Test request transformation with aspectRatio""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_request_with_image_size(self): + """Test request transformation with imageSize (Gemini 3 Pro)""" + request = self.config.transform_image_generation_request( + model="gemini-3-pro-image-preview", + prompt="A nano banana", + optional_params={"imageSize": "4K"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + def test_map_openai_params_web_search_options(self): + """Test web_search_options maps to googleSearch tool""" + result = self.config.map_openai_params({"web_search_options": {}}, {}, "gemini-3.1-flash-image-preview", False) + assert result["tools"] == [{"googleSearch": {}}] + + def test_transform_image_generation_request_with_web_search_tools(self): + """Test request transformation includes googleSearch tools""" + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image-preview", + prompt="Generate an image of the latest iPhone", + optional_params={"tools": [{"googleSearch": {}}]}, + litellm_params={}, + headers={}, + ) + assert request["tools"] == [{"googleSearch": {}}] + + def test_transform_image_generation_request_forwards_tool_config(self): + """Test request transformation forwards toolConfig side-effects from tool mapping""" + mapped = self.config.map_openai_params( + {"tools": [{"googleMaps": {"latitude": 37.7, "longitude": -122.4}}]}, + {}, + "gemini-3.1-flash-image-preview", + False, + ) + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image-preview", + prompt="Generate an image of a coffee shop nearby", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + assert request["tools"] == [{"googleMaps": {}}] + assert request["toolConfig"] == {"retrievalConfig": {"latLng": {"latitude": 37.7, "longitude": -122.4}}} + + def test_transform_image_generation_request_with_candidate_count(self): + """Test request transformation with candidate_count""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"candidate_count": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_request_with_n(self): + """Test request transformation with n parameter""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"n": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 93, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 54, + }, + { + "modality": "IMAGE", + "tokenCount": 39, + }, + ], + "candidatesTokenCount": 17, + "totalTokenCount": 110, + }, + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + assert result.usage.input_tokens == 93 + assert result.usage.input_tokens_details.text_tokens == 54 + assert result.usage.input_tokens_details.image_tokens == 39 + assert result.usage.output_tokens == 17 + assert result.usage.total_tokens == 110 + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "image1", + } + }, + { + "inlineData": { + "mimeType": "image/png", + "data": "image2", + } + }, + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + def test_transform_image_generation_response_signature(self): + """Test response transformation includes thoughtSignature for Gemini 3 Pro""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + }, + "thoughtSignature": "test_signature_abc123", + } + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-3-pro-image-preview", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" + + def test_transform_image_generation_response_tracks_web_search_requests(self): + """Grounding queries are carried onto usage so search spend can be billed""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + }, + "groundingMetadata": {"webSearchQueries": ["eiffel tower", "paris skyline"]}, + } + ], + "usageMetadata": { + "promptTokenCount": 93, + "candidatesTokenCount": 17, + "totalTokenCount": 110, + }, + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.usage.web_search_requests == 2 + + +class TestVertexAIImagenImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIImagenImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("imagegeneration@006") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to sampleCount""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) + assert result.get("sampleCount") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) + assert result.get("aspectRatio") == "1:1" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "instances" in request + assert "parameters" in request + assert request["instances"][0]["prompt"] == "A cat" + assert request["parameters"]["sampleCount"] == 1 + + def test_transform_image_generation_request_with_params(self): + """Test request transformation with parameters""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={"sampleCount": 2, "aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["parameters"]["sampleCount"] == 2 + assert request["parameters"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_request_labels_from_metadata(self): + """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={"metadata": {"requester_metadata": {"team": "platform", "env": "prod"}}}, + headers={}, + ) + assert request["labels"] == {"team": "platform", "env": "prod"} + assert "labels" not in request["parameters"] + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "base64_encoded_image_data"}]} + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + {"bytesBase64Encoded": "image1"}, + {"bytesBase64Encoded": "image2"}, + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + +class TestGetVertexAIImageGenerationConfig: + """Test the router function that selects the correct config""" + + def test_get_gemini_model_config(self): + """Test that Gemini models return Gemini config""" + config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("vertex_ai/gemini-2.5-flash-image") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + def test_get_imagen_model_config(self): + """Test that Imagen models return Imagen config""" + config = get_vertex_ai_image_generation_config("imagegeneration@006") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("vertex_ai/imagegeneration@006") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + def test_get_non_gemini_model_config(self): + """Test that non-Gemini models default to Imagen config""" + config = get_vertex_ai_image_generation_config("some-other-model") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + +class TestVertexAIImageGenerationIntegration: + """Integration tests for Vertex AI image generation""" + + + def test_gemini_get_complete_url(self): + """Test Gemini config URL generation""" + config = VertexAIGeminiImageGenerationConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-2.5-flash-image", + optional_params={}, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + ) + assert "test-project" in url + assert "us-central1" in url + assert "gemini-2.5-flash-image" in url + assert "generateContent" in url + + def test_imagen_get_complete_url(self): + """Test Imagen config URL generation""" + config = VertexAIImagenImageGenerationConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="imagegeneration@006", + optional_params={}, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + ) + assert "test-project" in url + assert "us-central1" in url + assert "imagegeneration@006" in url + assert "predict" in url diff --git a/tests/unit/llms/vertex_ai/rerank/__init__.py b/tests/unit/llms/vertex_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/unit/llms/vertex_ai/test_bge_embedding.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_bge_embedding.py rename to tests/unit/llms/vertex_ai/test_bge_embedding.py diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py b/tests/unit/llms/vertex_ai/test_bge_response_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py rename to tests/unit/llms/vertex_ai/test_bge_response_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py rename to tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py b/tests/unit/llms/vertex_ai/test_gemini_empty_properties.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py rename to tests/unit/llms/vertex_ai/test_gemini_empty_properties.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py b/tests/unit/llms/vertex_ai/test_gemini_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py rename to tests/unit/llms/vertex_ai/test_gemini_header_forwarding.py diff --git a/tests/test_litellm/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_http_status_201.py rename to tests/unit/llms/vertex_ai/test_http_status_201.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/unit/llms/vertex_ai/test_vertex.py similarity index 97% rename from tests/test_litellm/llms/vertex_ai/test_vertex.py rename to tests/unit/llms/vertex_ai/test_vertex.py index e3007bac7f3..ab8bf123ab2 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/unit/llms/vertex_ai/test_vertex.py @@ -1193,7 +1193,6 @@ def test_logprobs(): def test_process_gemini_media(): """Test the _process_gemini_media function for different image sources""" - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import FileDataType # Test GCS URI @@ -1271,7 +1270,6 @@ def test_process_gemini_media(): assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." - def test_get_image_mime_type_from_url(): """Test the _get_image_mime_type_from_url function for different image URLs""" from litellm.llms.vertex_ai.gemini.transformation import ( @@ -1372,46 +1370,6 @@ def encoded_images(): return [encode_image_to_base64(path) for path in image_paths] -@pytest.fixture -def mock_convert_url_to_base64(): - with patch( - "litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64", - ) as mock: - # Setup the mock to return a valid image object - mock.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - yield mock - - -@pytest.fixture -def mock_blob(): - return Mock(spec=BlobType) - - -@pytest.mark.parametrize( - "http_url", - [ - "http://img1.etsystatic.com/260/0/7813604/il_fullxfull.4226713999_q86e.jpg", - "http://example.com/image.jpg", - "http://subdomain.domain.com/path/to/image.png", - ], -) -def test_process_gemini_media_http_url( - http_url: str, mock_convert_url_to_base64: Mock, mock_blob: Mock -) -> None: - """ - Test that _process_gemini_media correctly handles HTTP URLs. - - Args: - http_url: Test HTTP URL - mock_convert_to_anthropic: Mocked convert_to_anthropic_image_obj function - mock_blob: Mocked BlobType instance - - Vertex AI supports image urls. Ensure no network requests are made. - """ - expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - mock_convert_url_to_base64.return_value = expected_image_data - # Act - result = _process_gemini_media(http_url) # assert result["file_data"]["file_uri"] == http_url diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/unit/llms/vertex_ai/test_vertex_ai_batch_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_batch_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py b/tests/unit/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/unit/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py b/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py rename to tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_global_url_support.py b/tests/unit/llms/vertex_ai/test_vertex_global_url_support.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_global_url_support.py rename to tests/unit/llms/vertex_ai/test_vertex_global_url_support.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_image_generation.py b/tests/unit/llms/vertex_ai/test_vertex_image_generation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_image_generation.py rename to tests/unit/llms/vertex_ai/test_vertex_image_generation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py rename to tests/unit/llms/vertex_ai/test_vertex_llm_base.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py rename to tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/unit/llms/vertex_ai/test_vertex_passthrough_logging_handler.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py rename to tests/unit/llms/vertex_ai/test_vertex_passthrough_logging_handler.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py diff --git a/tests/test_litellm/llms/volcengine/embedding/__init__.py b/tests/unit/llms/volcengine/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/volcengine/embedding/__init__.py rename to tests/unit/llms/volcengine/embedding/__init__.py diff --git a/tests/test_litellm/llms/volcengine/test_volcengine.py b/tests/unit/llms/volcengine/test_volcengine.py similarity index 100% rename from tests/test_litellm/llms/volcengine/test_volcengine.py rename to tests/unit/llms/volcengine/test_volcengine.py diff --git a/tests/unit/llms/wandb/__init__.py b/tests/unit/llms/wandb/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py b/tests/unit/llms/wandb/test_wandb_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py rename to tests/unit/llms/wandb/test_wandb_chat_transformation.py diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/unit/llms/xai/test_xai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py rename to tests/unit/llms/xai/test_xai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_chat_transformation.py rename to tests/unit/llms/xai/test_xai_chat_transformation.py diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/unit/llms/xai/test_xai_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_cost_calculator.py rename to tests/unit/llms/xai/test_xai_cost_calculator.py diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/unit/llms/xai/test_xai_key_fallback.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_key_fallback.py rename to tests/unit/llms/xai/test_xai_key_fallback.py diff --git a/tests/test_litellm/llms/xai/test_xai_model_registry.py b/tests/unit/llms/xai/test_xai_model_registry.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_model_registry.py rename to tests/unit/llms/xai/test_xai_model_registry.py diff --git a/tests/test_litellm/llms/xai/test_xai_oauth.py b/tests/unit/llms/xai/test_xai_oauth.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_oauth.py rename to tests/unit/llms/xai/test_xai_oauth.py diff --git a/tests/unit/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py index 4fa9c5bd3c1..e464402c9d8 100644 --- a/tests/unit/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -39,6 +39,7 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, "WORKERS": workers, + "UNIT_FLAG": "", }, capture_output=True, text=True, From cf491d1df91afa50527d0253ac960a8bf81ff678 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 12:57:07 -0700 Subject: [PATCH 14/58] test: move tests/test_litellm integrations and secret_managers into tests/unit (#43194) * ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: move tests/test_litellm/llms into tests/unit/llms Rename-only. Moves the provider tests and the fine-tuning fixtures they load, mirroring the old paths. Follow-up commits merge, split and wire them. * test: merge, split and prune the moved llms tests Merges the Databricks chat transformation tests into the existing unit file, keeps the tests that need real keys or the network in tests/test_litellm, deletes the audited tests a stronger unit test already covers, and points imports at tests.unit.llms. * ci: run the moved llms tests under their legacy flags The Vertex AI and All Other Providers shards keep their legacy test-path for the retained files and add the llm-vertex-ai and llm-other-providers unit selections. CircleCI gets matching unit jobs. * test: make the tests/unit/llms directories packages Adds __init__.py to the moved dirs and drops the legacy ones whose directories no longer hold tests. * test: drop script runners and path hacks the llms split left dangling The __main__ runners in the split openai_like files and the Databricks e2e runner called tests that now live in the other half of the split or were deleted. The retained legacy halves also no longer need sys.path edits. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. * test: move tests/test_litellm integrations and secret_managers into tests/unit Rename-only. Mirrors the old paths, including the directory conftests and the prompt and JSON fixtures. Follow-up commits prune and wire them. * test: prune and repoint the moved integrations tests Deletes the 7 audited tests a stronger test in the same tree already covers, imports the TLS sink helpers from their new conftest path, and restores os.environ after each integrations test. Some presets write OTEL_EXPORTER_OTLP_HEADERS straight into os.environ, and without the legacy tree's test ordering that header leaked into the AgentOps tests. * ci: run the moved integrations tests under their legacy flag The integrations GHA shard and a new CircleCI job run the integrations unit selection. secret_managers joins the misc selection. * docs: point integrations and secret_managers references at tests/unit * test: make the moved integrations directories packages * test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path The Databricks e2e file is a manual script whose main() calls the tests that were pruned, so pruning them broke the documented run. It is back to its main version. The SageMaker Nova docstring now points at the file's real location in tests/local_testing. * test: keep the job's UNIT_FLAG out of the shard-script tests --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/unit_selection.sh | 3 + .circleci/tests.yml | 7 ++ .github/workflows/test-unit.yml | 4 +- Makefile | 4 +- .../dashboard_all_metrics/readme.md | 2 +- .../tests/secret_manager/support.rs | 2 +- litellm-rust/crates/secrets/PARITY.md | 20 +++--- litellm/integrations/levo/README.md | 2 +- .../integrations/levo/__init__.py | 1 - .../integrations/SlackAlerting}/__init__.py | 0 .../SlackAlerting/test_budget_alert_types.py | 0 .../test_hanging_request_check.py | 0 .../test_model_deprecation_alert.py | 0 .../SlackAlerting/test_ms_teams.py | 0 .../SlackAlerting/test_slack_alerting.py | 0 .../test_slack_alerting_digest.py | 0 .../test_slack_alerting_utils.py | 0 .../SlackAlerting/test_user_spend_alerts.py | 0 .../integrations/arize}/__init__.py | 0 .../integrations/arize/test_arize.py | 0 .../arize/test_arize_health_check.py | 0 .../arize/test_arize_otel_coexistence.py | 0 .../integrations/arize/test_arize_phoenix.py | 0 .../integrations/arize/test_arize_utils.py | 0 .../integrations/azure_storage}/__init__.py | 0 .../azure_storage/test_azure_storage.py | 0 tests/unit/integrations/bitbucket/__init__.py | 0 .../bitbucket/test_bitbucket_integration.py | 0 .../test_bitbucket_prompt_manager.py | 31 --------- tests/unit/integrations/cloudzero/__init__.py | 0 .../integrations/cloudzero/test_cloudzero.py | 0 .../cloudzero/test_cloudzero_database.py | 0 .../cloudzero/test_cz_stream_api.py | 0 .../cloudzero/test_dry_run_endpoint.py | 0 .../integrations/cloudzero/test_transform.py | 0 .../code_interpreter_interception/__init__.py | 0 .../test_handler.py | 0 .../integrations/conftest.py | 9 +++ tests/unit/integrations/datadog/__init__.py | 0 .../datadog/test_datadog_cost_management.py | 0 .../datadog/test_datadog_llm_obs.py | 0 .../datadog/test_datadog_llm_obs_agent.py | 0 .../datadog/test_datadog_logger_batching.py | 0 .../datadog/test_datadog_metrics.py | 0 .../datadog/test_datadog_tags_regression.py | 0 .../datadog/test_datadog_team_handler.py | 0 tests/unit/integrations/dotprompt/__init__.py | 0 .../integrations/dotprompt/chat_prompt.prompt | 0 .../dotprompt/chat_prompt.v1.prompt | 0 .../dotprompt/chat_prompt.v2.prompt | 0 .../dotprompt/coding_assistant.prompt | 0 .../dotprompt/sample_prompt.prompt | 0 .../dotprompt/test_prompt_manager.py | 0 tests/unit/integrations/focus/__init__.py | 0 .../integrations/focus/test_csv_serializer.py | 0 .../focus/test_destination_factory.py | 0 .../integrations/focus/test_focus_database.py | 0 .../focus/test_focus_gcs_destination.py | 0 .../focus/test_focus_transformer.py | 0 .../focus/test_mavvrik_destination.py | 0 .../integrations/focus/test_s3_destination.py | 0 .../integrations/focus/test_transformer.py | 0 .../focus/test_vantage_destination.py | 0 tests/unit/integrations/gitlab/__init__.py | 0 .../integrations/gitlab/test_gitlab_client.py | 0 .../gitlab/test_gitlab_integration.py | 0 .../gitlab/test_gitlab_prompt_manager.py | 0 tests/unit/integrations/langfuse/__init__.py | 0 .../langfuse/test_gemini_cached_tokens.py | 0 .../test_langfuse_prompt_management.py | 0 .../langfuse/test_langfuse_sdk.py | 0 tests/unit/integrations/newrelic/__init__.py | 0 .../integrations/newrelic/test_newrelic.py | 0 .../newrelic/test_newrelic_metrics.py | 0 .../newrelic/test_newrelic_team_handler.py | 0 .../integrations/open_telemetry/__init__.py | 0 .../integrations/open_telemetry/_helpers.py | 0 .../integrations/open_telemetry/conftest.py | 0 .../open_telemetry/data/__init__.py | 0 .../open_telemetry/data/captured_kwargs.json | 0 .../data/captured_response.json | 0 .../test_otel_admin_endpoints.py | 0 .../test_otel_exception_handler.py | 0 .../test_otel_passthrough_endpoints.py | 0 .../test_otel_unified_endpoints.py | 0 .../test_passthrough_parent_span.py | 0 tests/unit/integrations/otel/__init__.py | 0 .../integrations/otel/test_db_endpoint.py | 0 .../integrations/otel/test_langfuse_logger.py | 0 .../integrations/otel/test_otel_v2_baggage.py | 0 .../otel/test_otel_v2_components.py | 2 +- ..._v2_config_baggage_parenting_guardrails.py | 0 .../otel/test_otel_v2_destinations.py | 0 .../integrations/otel/test_otel_v2_dynamic.py | 0 .../integrations/otel/test_otel_v2_emitter.py | 0 .../integrations/otel/test_otel_v2_logger.py | 8 --- .../integrations/otel/test_otel_v2_metrics.py | 0 .../integrations/otel/test_otel_v2_mount.py | 0 .../otel/test_otel_v2_multibackend.py | 0 .../integrations/otel/test_otel_v2_presets.py | 0 .../otel/test_otel_v2_sources_of_truth.py | 0 .../otel/test_otel_v2_vendor_mappers.py | 0 .../integrations/otel/test_runtime.py | 0 .../integrations/rubrik_test_helpers.py | 0 .../integrations/test_agentops.py | 0 .../test_anthropic_cache_control_hook.py | 0 .../integrations/test_athina.py | 0 .../integrations/test_azure_sentinel.py | 0 .../integrations/test_braintrust_logging.py | 0 .../integrations/test_braintrust_span_name.py | 0 .../integrations/test_custom_guardrail.py | 0 .../test_custom_guardrail_recursion.py | 0 .../test_custom_prompt_management.py | 0 .../integrations/test_deepeval.py | 0 .../integrations/test_galileo.py | 0 .../test_guardrail_logging_sync.py | 0 .../integrations/test_helicone.py | 0 .../integrations/test_langfuse.py | 0 .../integrations/test_langfuse_otel.py | 0 .../integrations/test_langsmith_init.py | 0 .../integrations/test_lunary.py | 0 .../integrations/test_mlflow.py | 0 .../integrations/test_openmeter.py | 0 .../integrations/test_opentelemetry.py | 64 +------------------ .../test_opentelemetry_dynamic_imports.py | 0 .../integrations/test_opik_utils.py | 0 .../test_otel_guardrail_violation_spans.py | 0 .../test_otel_team_attributes_matrix.py | 0 .../test_prometheus_api_promql_escape.py | 0 .../test_prometheus_budget_metric_guard.py | 0 ...st_prometheus_budget_metrics_db_lookups.py | 0 .../test_prometheus_budget_metrics_timeout.py | 0 .../test_prometheus_cache_metrics.py | 2 +- .../test_prometheus_caller_identity.py | 0 .../test_prometheus_carried_budget_state.py | 0 .../test_prometheus_client_ip_user_agent.py | 0 ...prometheus_custom_metadata_label_counts.py | 0 ...ometheus_deployment_state_proxy_rejects.py | 0 .../test_prometheus_end_user_cardinality.py | 0 ..._prometheus_input_sequence_length_label.py | 0 .../test_prometheus_invalid_key_filtering.py | 0 .../integrations/test_prometheus_labels.py | 0 .../test_prometheus_mcp_tool_metrics.py | 2 +- ...est_prometheus_media_generation_metrics.py | 0 ...test_prometheus_metric_name_consistency.py | 0 .../test_prometheus_metrics_endpoint.py | 0 .../test_prometheus_missing_metrics.py | 0 .../test_prometheus_none_metadata.py | 0 ...est_prometheus_overhead_with_guardrails.py | 0 ...test_prometheus_queue_guardrail_metrics.py | 0 .../test_prometheus_rate_limit_labels.py | 0 ...etheus_remaining_tokens_router_fallback.py | 0 ..._prometheus_requested_model_cardinality.py | 0 .../test_prometheus_service_tier_label.py | 2 +- .../integrations/test_prometheus_services.py | 0 .../test_prometheus_spend_capture_rate.py | 0 .../test_prometheus_spend_logs_metadata.py | 0 .../test_prometheus_stream_label.py | 0 .../test_prometheus_token_detail_metrics.py | 2 +- .../test_prometheus_user_team_metrics.py | 21 ------ .../test_prometheus_zero_cost_metric.py | 0 .../integrations/test_prompt_manager_ssti.py | 0 .../test_responses_background_cost.py | 0 .../integrations/test_rubrik.py | 2 +- .../integrations/test_s3.py | 0 .../integrations/test_s3_v2.py | 0 .../integrations/test_shadow_eval_logger.py | 0 .../integrations/test_weave_otel.py | 0 .../websearch_interception/__init__.py | 0 .../test_websearch_agentic_loop_cap.py | 0 .../test_websearch_chat_completion.py | 0 .../test_websearch_interception_handler.py | 0 .../test_websearch_interception_thinking.py | 0 .../test_websearch_native_blocks.py | 0 .../test_websearch_responses.py | 0 .../test_websearch_rich_query_shape.py | 0 .../test_websearch_short_circuit.py | 0 .../test_websearch_streaming_wrap.py | 0 .../test_websearch_thinking_constraint.py | 0 tests/unit/secret_managers/__init__.py | 0 .../hashicorp_vault_parity.json | 0 .../test_aws_secret_manager_replication.py | 0 .../test_aws_secret_manager_rotation.py | 0 .../test_aws_secret_manager_v2.py | 0 .../test_base_secret_manager.py | 0 .../test_custom_secret_manager.py | 0 .../test_cyberark_secret_manager.py | 0 .../test_get_azure_ad_token_provider.py | 0 .../test_hashicorp_secret_manager.py | 0 .../test_secret_manager_handler.py | 0 .../test_secret_managers_main.py | 0 191 files changed, 43 insertions(+), 147 deletions(-) delete mode 100644 tests/test_litellm/integrations/levo/__init__.py rename tests/{test_litellm/integrations/code_interpreter_interception => unit/integrations/SlackAlerting}/__init__.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_budget_alert_types.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_hanging_request_check.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_model_deprecation_alert.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_ms_teams.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_slack_alerting.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_slack_alerting_digest.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_slack_alerting_utils.py (100%) rename tests/{test_litellm => unit}/integrations/SlackAlerting/test_user_spend_alerts.py (100%) rename tests/{test_litellm/integrations/gitlab => unit/integrations/arize}/__init__.py (100%) rename tests/{test_litellm => unit}/integrations/arize/test_arize.py (100%) rename tests/{test_litellm => unit}/integrations/arize/test_arize_health_check.py (100%) rename tests/{test_litellm => unit}/integrations/arize/test_arize_otel_coexistence.py (100%) rename tests/{test_litellm => unit}/integrations/arize/test_arize_phoenix.py (100%) rename tests/{test_litellm => unit}/integrations/arize/test_arize_utils.py (100%) rename tests/{test_litellm/integrations/open_telemetry => unit/integrations/azure_storage}/__init__.py (100%) rename tests/{test_litellm => unit}/integrations/azure_storage/test_azure_storage.py (100%) create mode 100644 tests/unit/integrations/bitbucket/__init__.py rename tests/{test_litellm => unit}/integrations/bitbucket/test_bitbucket_integration.py (100%) rename tests/{test_litellm => unit}/integrations/bitbucket/test_bitbucket_prompt_manager.py (93%) create mode 100644 tests/unit/integrations/cloudzero/__init__.py rename tests/{test_litellm => unit}/integrations/cloudzero/test_cloudzero.py (100%) rename tests/{test_litellm => unit}/integrations/cloudzero/test_cloudzero_database.py (100%) rename tests/{test_litellm => unit}/integrations/cloudzero/test_cz_stream_api.py (100%) rename tests/{test_litellm => unit}/integrations/cloudzero/test_dry_run_endpoint.py (100%) rename tests/{test_litellm => unit}/integrations/cloudzero/test_transform.py (100%) create mode 100644 tests/unit/integrations/code_interpreter_interception/__init__.py rename tests/{test_litellm => unit}/integrations/code_interpreter_interception/test_handler.py (100%) rename tests/{test_litellm => unit}/integrations/conftest.py (94%) create mode 100644 tests/unit/integrations/datadog/__init__.py rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_cost_management.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_llm_obs.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_llm_obs_agent.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_logger_batching.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_tags_regression.py (100%) rename tests/{test_litellm => unit}/integrations/datadog/test_datadog_team_handler.py (100%) create mode 100644 tests/unit/integrations/dotprompt/__init__.py rename tests/{test_litellm => unit}/integrations/dotprompt/chat_prompt.prompt (100%) rename tests/{test_litellm => unit}/integrations/dotprompt/chat_prompt.v1.prompt (100%) rename tests/{test_litellm => unit}/integrations/dotprompt/chat_prompt.v2.prompt (100%) rename tests/{test_litellm => unit}/integrations/dotprompt/coding_assistant.prompt (100%) rename tests/{test_litellm => unit}/integrations/dotprompt/sample_prompt.prompt (100%) rename tests/{test_litellm => unit}/integrations/dotprompt/test_prompt_manager.py (100%) create mode 100644 tests/unit/integrations/focus/__init__.py rename tests/{test_litellm => unit}/integrations/focus/test_csv_serializer.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_destination_factory.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_focus_database.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_focus_gcs_destination.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_focus_transformer.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_mavvrik_destination.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_s3_destination.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_transformer.py (100%) rename tests/{test_litellm => unit}/integrations/focus/test_vantage_destination.py (100%) create mode 100644 tests/unit/integrations/gitlab/__init__.py rename tests/{test_litellm => unit}/integrations/gitlab/test_gitlab_client.py (100%) rename tests/{test_litellm => unit}/integrations/gitlab/test_gitlab_integration.py (100%) rename tests/{test_litellm => unit}/integrations/gitlab/test_gitlab_prompt_manager.py (100%) create mode 100644 tests/unit/integrations/langfuse/__init__.py rename tests/{test_litellm => unit}/integrations/langfuse/test_gemini_cached_tokens.py (100%) rename tests/{test_litellm => unit}/integrations/langfuse/test_langfuse_prompt_management.py (100%) rename tests/{test_litellm => unit}/integrations/langfuse/test_langfuse_sdk.py (100%) create mode 100644 tests/unit/integrations/newrelic/__init__.py rename tests/{test_litellm => unit}/integrations/newrelic/test_newrelic.py (100%) rename tests/{test_litellm => unit}/integrations/newrelic/test_newrelic_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/newrelic/test_newrelic_team_handler.py (100%) create mode 100644 tests/unit/integrations/open_telemetry/__init__.py rename tests/{test_litellm => unit}/integrations/open_telemetry/_helpers.py (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/conftest.py (100%) create mode 100644 tests/unit/integrations/open_telemetry/data/__init__.py rename tests/{test_litellm => unit}/integrations/open_telemetry/data/captured_kwargs.json (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/data/captured_response.json (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/test_otel_admin_endpoints.py (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/test_otel_exception_handler.py (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/test_otel_passthrough_endpoints.py (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/test_otel_unified_endpoints.py (100%) rename tests/{test_litellm => unit}/integrations/open_telemetry/test_passthrough_parent_span.py (100%) create mode 100644 tests/unit/integrations/otel/__init__.py rename tests/{test_litellm => unit}/integrations/otel/test_db_endpoint.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_langfuse_logger.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_baggage.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_components.py (99%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_destinations.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_dynamic.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_emitter.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_logger.py (99%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_mount.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_multibackend.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_presets.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_sources_of_truth.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_otel_v2_vendor_mappers.py (100%) rename tests/{test_litellm => unit}/integrations/otel/test_runtime.py (100%) rename tests/{test_litellm => unit}/integrations/rubrik_test_helpers.py (100%) rename tests/{test_litellm => unit}/integrations/test_agentops.py (100%) rename tests/{test_litellm => unit}/integrations/test_anthropic_cache_control_hook.py (100%) rename tests/{test_litellm => unit}/integrations/test_athina.py (100%) rename tests/{test_litellm => unit}/integrations/test_azure_sentinel.py (100%) rename tests/{test_litellm => unit}/integrations/test_braintrust_logging.py (100%) rename tests/{test_litellm => unit}/integrations/test_braintrust_span_name.py (100%) rename tests/{test_litellm => unit}/integrations/test_custom_guardrail.py (100%) rename tests/{test_litellm => unit}/integrations/test_custom_guardrail_recursion.py (100%) rename tests/{test_litellm => unit}/integrations/test_custom_prompt_management.py (100%) rename tests/{test_litellm => unit}/integrations/test_deepeval.py (100%) rename tests/{test_litellm => unit}/integrations/test_galileo.py (100%) rename tests/{test_litellm => unit}/integrations/test_guardrail_logging_sync.py (100%) rename tests/{test_litellm => unit}/integrations/test_helicone.py (100%) rename tests/{test_litellm => unit}/integrations/test_langfuse.py (100%) rename tests/{test_litellm => unit}/integrations/test_langfuse_otel.py (100%) rename tests/{test_litellm => unit}/integrations/test_langsmith_init.py (100%) rename tests/{test_litellm => unit}/integrations/test_lunary.py (100%) rename tests/{test_litellm => unit}/integrations/test_mlflow.py (100%) rename tests/{test_litellm => unit}/integrations/test_openmeter.py (100%) rename tests/{test_litellm => unit}/integrations/test_opentelemetry.py (99%) rename tests/{test_litellm => unit}/integrations/test_opentelemetry_dynamic_imports.py (100%) rename tests/{test_litellm => unit}/integrations/test_opik_utils.py (100%) rename tests/{test_litellm => unit}/integrations/test_otel_guardrail_violation_spans.py (100%) rename tests/{test_litellm => unit}/integrations/test_otel_team_attributes_matrix.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_api_promql_escape.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_budget_metric_guard.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_budget_metrics_db_lookups.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_budget_metrics_timeout.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_cache_metrics.py (99%) rename tests/{test_litellm => unit}/integrations/test_prometheus_caller_identity.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_carried_budget_state.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_client_ip_user_agent.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_custom_metadata_label_counts.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_deployment_state_proxy_rejects.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_end_user_cardinality.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_input_sequence_length_label.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_invalid_key_filtering.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_labels.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_mcp_tool_metrics.py (99%) rename tests/{test_litellm => unit}/integrations/test_prometheus_media_generation_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_metric_name_consistency.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_metrics_endpoint.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_missing_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_none_metadata.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_overhead_with_guardrails.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_queue_guardrail_metrics.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_rate_limit_labels.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_remaining_tokens_router_fallback.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_requested_model_cardinality.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_service_tier_label.py (98%) rename tests/{test_litellm => unit}/integrations/test_prometheus_services.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_spend_capture_rate.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_spend_logs_metadata.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_stream_label.py (100%) rename tests/{test_litellm => unit}/integrations/test_prometheus_token_detail_metrics.py (99%) rename tests/{test_litellm => unit}/integrations/test_prometheus_user_team_metrics.py (98%) rename tests/{test_litellm => unit}/integrations/test_prometheus_zero_cost_metric.py (100%) rename tests/{test_litellm => unit}/integrations/test_prompt_manager_ssti.py (100%) rename tests/{test_litellm => unit}/integrations/test_responses_background_cost.py (100%) rename tests/{test_litellm => unit}/integrations/test_rubrik.py (99%) rename tests/{test_litellm => unit}/integrations/test_s3.py (100%) rename tests/{test_litellm => unit}/integrations/test_s3_v2.py (100%) rename tests/{test_litellm => unit}/integrations/test_shadow_eval_logger.py (100%) rename tests/{test_litellm => unit}/integrations/test_weave_otel.py (100%) create mode 100644 tests/unit/integrations/websearch_interception/__init__.py rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_agentic_loop_cap.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_chat_completion.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_interception_handler.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_interception_thinking.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_native_blocks.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_responses.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_rich_query_shape.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_short_circuit.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_streaming_wrap.py (100%) rename tests/{test_litellm => unit}/integrations/websearch_interception/test_websearch_thinking_constraint.py (100%) create mode 100644 tests/unit/secret_managers/__init__.py rename tests/{test_litellm => unit}/secret_managers/hashicorp_vault_parity.json (100%) rename tests/{test_litellm => unit}/secret_managers/test_aws_secret_manager_replication.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_aws_secret_manager_rotation.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_aws_secret_manager_v2.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_base_secret_manager.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_custom_secret_manager.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_cyberark_secret_manager.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_get_azure_ad_token_provider.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_hashicorp_secret_manager.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_secret_manager_handler.py (100%) rename tests/{test_litellm => unit}/secret_managers/test_secret_managers_main.py (100%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index e9e5dd3d66b..d56e29fb627 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,7 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + integrations llm-other-providers llm-vertex-ai mcp-integration @@ -52,6 +53,7 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + integrations) echo tests/unit/integrations ;; llm-other-providers) find tests/unit/llms -name 'test_*.py' -not -path 'tests/unit/llms/vertex_ai/*' ;; llm-vertex-ai) echo tests/unit/llms/vertex_ai ;; mcp-integration) @@ -75,6 +77,7 @@ legacy_paths() { echo tests/unit/messages echo tests/unit/rag echo tests/unit/rerank_api + echo tests/unit/secret_managers echo tests/unit/vector_stores echo tests/unit/videos ;; proxy-db-auth-checks) diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 994d67da64d..41e9f11cefa 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -369,6 +369,13 @@ workflows: reruns: 2 base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-integrations + flag: integrations + shards: 2 + reruns: 3 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-misc flag: misc diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2fa05879350..4dca8075440 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -80,7 +80,8 @@ jobs: - shard: integrations artifact-name: integrations - test-path: "tests/test_litellm/integrations" + test-path: "" + unit-flag: integrations workers: 2 reruns: 3 timeout-minutes: 20 @@ -107,7 +108,6 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/test_litellm/secret_managers tests/test_litellm/interactions tests/test_litellm/ocr tests/test_litellm/passthrough diff --git a/Makefile b/Makefile index e86047b1987..f27525b58ff 100644 --- a/Makefile +++ b/Makefile @@ -326,13 +326,13 @@ test-unit-proxy-misc: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps - $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20 test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md index 70562b18aa6..a3869213be2 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md @@ -2,7 +2,7 @@ Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about -Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard +Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs index 30e46f93248..e1797784b84 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs @@ -118,7 +118,7 @@ pub(super) struct ParityCase { pub(super) fn parity_cases() -> Vec { serde_json::from_str(include_str!(concat!( env!("CARGO_MANIFEST_DIR"), - "/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json" + "/../../../tests/unit/secret_managers/hashicorp_vault_parity.json" ))) .unwrap() } diff --git a/litellm-rust/crates/secrets/PARITY.md b/litellm-rust/crates/secrets/PARITY.md index aeed4ba4b83..9439da93d65 100644 --- a/litellm-rust/crates/secrets/PARITY.md +++ b/litellm-rust/crates/secrets/PARITY.md @@ -32,7 +32,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo `_SecretManagerRuntime` is a private implementation detail, not a replacement SDK class. Its async methods return Futures; public `async def` methods retain lazy coroutine creation and `asyncio.create_task` support. Passing the same names and arguments is insufficient to claim parity until the remaining return-value, error, cache and configuration differences above are closed -## [tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py) +## [tests/unit/secret_managers/test_aws_secret_manager_replication.py](../../../tests/unit/secret_managers/test_aws_secret_manager_replication.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -48,7 +48,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_replicate_secret_http_error_raises` | [direct_replication_returns_response_or_service_error](../secrets-aws/tests/secret_manager/writes.rs) | | `test_replicate_secret_timeout_raises` | [write_and_replication_timeouts_remain_errors](../secrets-aws/tests/secret_manager/writes.rs) | -## [tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py) +## [tests/unit/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/unit/secret_managers/test_aws_secret_manager_rotation.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -59,7 +59,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_write_secret_to_name_inside_recovery_window_restores_and_stores_new_value` | [recovery_window_alias_is_restored_updated_and_tagged](../secrets-aws/tests/secret_manager/writes.rs) | | `test_write_secret_to_live_existing_name_still_fails_without_overwriting` | [create_failure_does_not_overwrite_an_alias_without_a_deletion_date](../secrets-aws/tests/secret_manager/writes.rs) | -## [tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py) +## [tests/unit/secret_managers/test_aws_secret_manager_v2.py](../../../tests/unit/secret_managers/test_aws_secret_manager_v2.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -70,14 +70,14 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) | | `test_prepare_request_env_bedrock_runtime_endpoint_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_base_secret_manager.py](../../../tests/test_litellm/secret_managers/test_base_secret_manager.py) +## [tests/unit/secret_managers/test_base_secret_manager.py](../../../tests/unit/secret_managers/test_base_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | | `test_raise_if_unsafe_secret_name_rejects_traversal_and_line_breaks` | [names_reject_path_traversal_and_control_characters](../secrets-types/tests/rotation.rs) | | `test_raise_if_unsafe_secret_name_allows_legitimate_aliases` | [names_allow_safe_values](../secrets-types/tests/rotation.rs) | -## [tests/test_litellm/secret_managers/test_custom_secret_manager.py](../../../tests/test_litellm/secret_managers/test_custom_secret_manager.py) +## [tests/unit/secret_managers/test_custom_secret_manager.py](../../../tests/unit/secret_managers/test_custom_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -89,7 +89,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_custom_secret_manager_integration_with_litellm` | [manager_strings_are_coerced_like_literal_eval](../secrets/tests/resolution.rs) | | `test_minimal_custom_secret_manager` | Exercises the Python example subclass itself or Python default methods. Caller-authored Python implementations remain Python callbacks; resolver integration is covered by `manager_strings_are_coerced_like_literal_eval` | -## [tests/test_litellm/secret_managers/test_cyberark_secret_manager.py](../../../tests/test_litellm/secret_managers/test_cyberark_secret_manager.py) +## [tests/unit/secret_managers/test_cyberark_secret_manager.py](../../../tests/unit/secret_managers/test_cyberark_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -97,7 +97,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_async_write_matches_parity_fixture` | [writes_match_python_parity_fixture](../secrets-cyberark/tests/secret_manager/writes.rs) | | `test_missing_credentials_raise_value_error` | [new_validates_credentials_before_license_and_configuration](../secrets-cyberark/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py) +## [tests/unit/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/unit/secret_managers/test_get_azure_ad_token_provider.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -115,7 +115,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_get_azure_ad_token_provider_prefers_workload_identity_over_managed_identity` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate | | `test_get_azure_ad_token_provider_defaults_to_default_azure_credential` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate | -## [tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py](../../../tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py) +## [tests/unit/secret_managers/test_hashicorp_secret_manager.py](../../../tests/unit/secret_managers/test_hashicorp_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -130,13 +130,13 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_tls_login_uses_login_namespace` | [tls_login_posts_the_role_and_uses_the_client_identity](../secrets-hashicorp/tests/secret_manager/configuration.rs) | | `test_configuration_matches_native_parity_fixture` | [configuration_matches_python_parity_fixture](../secrets-hashicorp/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_secret_manager_handler.py](../../../tests/test_litellm/secret_managers/test_secret_manager_handler.py) +## [tests/unit/secret_managers/test_secret_manager_handler.py](../../../tests/unit/secret_managers/test_secret_manager_handler.py) | Python test | Rust coverage or boundary | | --- | --- | | `test_azure_key_vault_matches_rust_parity_fixture` | [parity_fixture_matches_python_backend_contract](../secrets-azure/tests/key_vault.rs) | -## [tests/test_litellm/secret_managers/test_secret_managers_main.py](../../../tests/test_litellm/secret_managers/test_secret_managers_main.py) +## [tests/unit/secret_managers/test_secret_managers_main.py](../../../tests/unit/secret_managers/test_secret_managers_main.py) | Python test | Rust coverage or boundary | | --- | --- | diff --git a/litellm/integrations/levo/README.md b/litellm/integrations/levo/README.md index 5296acb7ff4..1fbd202d9a5 100644 --- a/litellm/integrations/levo/README.md +++ b/litellm/integrations/levo/README.md @@ -92,7 +92,7 @@ litellm/integrations/levo/ ## Testing -See the test files in `tests/test_litellm/integrations/levo/`: +See the test files in `tests/unit/integrations/levo/`: - `test_levo.py`: Unit tests for configuration - `test_levo_integration.py`: Integration tests for callback registration diff --git a/tests/test_litellm/integrations/levo/__init__.py b/tests/test_litellm/integrations/levo/__init__.py deleted file mode 100644 index 1560e78b7b9..00000000000 --- a/tests/test_litellm/integrations/levo/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Levo integration tests diff --git a/tests/test_litellm/integrations/code_interpreter_interception/__init__.py b/tests/unit/integrations/SlackAlerting/__init__.py similarity index 100% rename from tests/test_litellm/integrations/code_interpreter_interception/__init__.py rename to tests/unit/integrations/SlackAlerting/__init__.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py rename to tests/unit/integrations/SlackAlerting/test_budget_alert_types.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py b/tests/unit/integrations/SlackAlerting/test_hanging_request_check.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py rename to tests/unit/integrations/SlackAlerting/test_hanging_request_check.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_model_deprecation_alert.py b/tests/unit/integrations/SlackAlerting/test_model_deprecation_alert.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_model_deprecation_alert.py rename to tests/unit/integrations/SlackAlerting/test_model_deprecation_alert.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py b/tests/unit/integrations/SlackAlerting/test_ms_teams.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py rename to tests/unit/integrations/SlackAlerting/test_ms_teams.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting_digest.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting_digest.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting_utils.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting_utils.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py b/tests/unit/integrations/SlackAlerting/test_user_spend_alerts.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py rename to tests/unit/integrations/SlackAlerting/test_user_spend_alerts.py diff --git a/tests/test_litellm/integrations/gitlab/__init__.py b/tests/unit/integrations/arize/__init__.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/__init__.py rename to tests/unit/integrations/arize/__init__.py diff --git a/tests/test_litellm/integrations/arize/test_arize.py b/tests/unit/integrations/arize/test_arize.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize.py rename to tests/unit/integrations/arize/test_arize.py diff --git a/tests/test_litellm/integrations/arize/test_arize_health_check.py b/tests/unit/integrations/arize/test_arize_health_check.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_health_check.py rename to tests/unit/integrations/arize/test_arize_health_check.py diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/unit/integrations/arize/test_arize_otel_coexistence.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py rename to tests/unit/integrations/arize/test_arize_otel_coexistence.py diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_phoenix.py rename to tests/unit/integrations/arize/test_arize_phoenix.py diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_utils.py rename to tests/unit/integrations/arize/test_arize_utils.py diff --git a/tests/test_litellm/integrations/open_telemetry/__init__.py b/tests/unit/integrations/azure_storage/__init__.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/__init__.py rename to tests/unit/integrations/azure_storage/__init__.py diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py similarity index 100% rename from tests/test_litellm/integrations/azure_storage/test_azure_storage.py rename to tests/unit/integrations/azure_storage/test_azure_storage.py diff --git a/tests/unit/integrations/bitbucket/__init__.py b/tests/unit/integrations/bitbucket/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/unit/integrations/bitbucket/test_bitbucket_integration.py similarity index 100% rename from tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py rename to tests/unit/integrations/bitbucket/test_bitbucket_integration.py diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py similarity index 93% rename from tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py rename to tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py index d6668bf9ad8..a1a88653ee6 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -307,37 +307,6 @@ def test_bitbucket_prompt_manager_render_template_not_found(): manager.prompt_manager.render_template("nonexistent", {"some": "variable"}) -@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") -def test_bitbucket_prompt_manager_integration(mock_client_class): - """Test BitBucketPromptManager integration with BitBucketClient.""" - # Mock the BitBucket client - mock_client = MagicMock() - mock_client.get_file_content.return_value = """--- -model: gpt-4 -temperature: 0.7 ---- -Hello {{name}}!""" - mock_client_class.return_value = mock_client - - config = { - "workspace": "test-workspace", - "repository": "test-repo", - "access_token": "test-token", - } - - manager = BitBucketPromptManager(config, prompt_id="test_prompt") - - # Should have loaded the prompt - assert "test_prompt" in manager.prompt_manager.prompts - template = manager.prompt_manager.prompts["test_prompt"] - assert template.model == "gpt-4" - assert template.temperature == 0.7 - - # Test rendering - rendered = manager.prompt_manager.render_template("test_prompt", {"name": "World"}) - assert rendered == "Hello World!" - - def test_bitbucket_prompt_manager_parse_prompt_to_messages(): """Test parsing prompt content into messages.""" config = { diff --git a/tests/unit/integrations/cloudzero/__init__.py b/tests/unit/integrations/cloudzero/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/unit/integrations/cloudzero/test_cloudzero.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cloudzero.py rename to tests/unit/integrations/cloudzero/test_cloudzero.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/unit/integrations/cloudzero/test_cloudzero_database.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py rename to tests/unit/integrations/cloudzero/test_cloudzero_database.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/unit/integrations/cloudzero/test_cz_stream_api.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py rename to tests/unit/integrations/cloudzero/test_cz_stream_api.py diff --git a/tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py b/tests/unit/integrations/cloudzero/test_dry_run_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py rename to tests/unit/integrations/cloudzero/test_dry_run_endpoint.py diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/unit/integrations/cloudzero/test_transform.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_transform.py rename to tests/unit/integrations/cloudzero/test_transform.py diff --git a/tests/unit/integrations/code_interpreter_interception/__init__.py b/tests/unit/integrations/code_interpreter_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py similarity index 100% rename from tests/test_litellm/integrations/code_interpreter_interception/test_handler.py rename to tests/unit/integrations/code_interpreter_interception/test_handler.py diff --git a/tests/test_litellm/integrations/conftest.py b/tests/unit/integrations/conftest.py similarity index 94% rename from tests/test_litellm/integrations/conftest.py rename to tests/unit/integrations/conftest.py index adc8e36e0af..48a01ed8d48 100644 --- a/tests/test_litellm/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -1,6 +1,7 @@ import functools import http.server import ipaddress +import os import queue import ssl import threading @@ -74,6 +75,14 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]: return certificate_path, key_path +@pytest.fixture(autouse=True) +def restore_process_environment() -> Iterator[None]: + original: Final = dict(os.environ) + yield + os.environ.clear() + os.environ.update(original) + + @pytest.fixture def tls_sink(tmp_path: Path) -> Iterator[TlsSink]: certificate_path, key_path = write_self_signed_cert(tmp_path, "sink") diff --git a/tests/unit/integrations/datadog/__init__.py b/tests/unit/integrations/datadog/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/unit/integrations/datadog/test_datadog_cost_management.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_cost_management.py rename to tests/unit/integrations/datadog/test_datadog_cost_management.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/unit/integrations/datadog/test_datadog_llm_obs.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py rename to tests/unit/integrations/datadog/test_datadog_llm_obs.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py b/tests/unit/integrations/datadog/test_datadog_llm_obs_agent.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py rename to tests/unit/integrations/datadog/test_datadog_llm_obs_agent.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py b/tests/unit/integrations/datadog/test_datadog_logger_batching.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py rename to tests/unit/integrations/datadog/test_datadog_logger_batching.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/unit/integrations/datadog/test_datadog_metrics.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_metrics.py rename to tests/unit/integrations/datadog/test_datadog_metrics.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py b/tests/unit/integrations/datadog/test_datadog_tags_regression.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py rename to tests/unit/integrations/datadog/test_datadog_tags_regression.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/unit/integrations/datadog/test_datadog_team_handler.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_team_handler.py rename to tests/unit/integrations/datadog/test_datadog_team_handler.py diff --git a/tests/unit/integrations/dotprompt/__init__.py b/tests/unit/integrations/dotprompt/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.prompt b/tests/unit/integrations/dotprompt/chat_prompt.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.prompt diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.v1.prompt b/tests/unit/integrations/dotprompt/chat_prompt.v1.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.v1.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.v1.prompt diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.v2.prompt b/tests/unit/integrations/dotprompt/chat_prompt.v2.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.v2.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.v2.prompt diff --git a/tests/test_litellm/integrations/dotprompt/coding_assistant.prompt b/tests/unit/integrations/dotprompt/coding_assistant.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/coding_assistant.prompt rename to tests/unit/integrations/dotprompt/coding_assistant.prompt diff --git a/tests/test_litellm/integrations/dotprompt/sample_prompt.prompt b/tests/unit/integrations/dotprompt/sample_prompt.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/sample_prompt.prompt rename to tests/unit/integrations/dotprompt/sample_prompt.prompt diff --git a/tests/test_litellm/integrations/dotprompt/test_prompt_manager.py b/tests/unit/integrations/dotprompt/test_prompt_manager.py similarity index 100% rename from tests/test_litellm/integrations/dotprompt/test_prompt_manager.py rename to tests/unit/integrations/dotprompt/test_prompt_manager.py diff --git a/tests/unit/integrations/focus/__init__.py b/tests/unit/integrations/focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/focus/test_csv_serializer.py b/tests/unit/integrations/focus/test_csv_serializer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_csv_serializer.py rename to tests/unit/integrations/focus/test_csv_serializer.py diff --git a/tests/test_litellm/integrations/focus/test_destination_factory.py b/tests/unit/integrations/focus/test_destination_factory.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_destination_factory.py rename to tests/unit/integrations/focus/test_destination_factory.py diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/unit/integrations/focus/test_focus_database.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_database.py rename to tests/unit/integrations/focus/test_focus_database.py diff --git a/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py b/tests/unit/integrations/focus/test_focus_gcs_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_gcs_destination.py rename to tests/unit/integrations/focus/test_focus_gcs_destination.py diff --git a/tests/test_litellm/integrations/focus/test_focus_transformer.py b/tests/unit/integrations/focus/test_focus_transformer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_transformer.py rename to tests/unit/integrations/focus/test_focus_transformer.py diff --git a/tests/test_litellm/integrations/focus/test_mavvrik_destination.py b/tests/unit/integrations/focus/test_mavvrik_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_mavvrik_destination.py rename to tests/unit/integrations/focus/test_mavvrik_destination.py diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/unit/integrations/focus/test_s3_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_s3_destination.py rename to tests/unit/integrations/focus/test_s3_destination.py diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/unit/integrations/focus/test_transformer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_transformer.py rename to tests/unit/integrations/focus/test_transformer.py diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/unit/integrations/focus/test_vantage_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_vantage_destination.py rename to tests/unit/integrations/focus/test_vantage_destination.py diff --git a/tests/unit/integrations/gitlab/__init__.py b/tests/unit/integrations/gitlab/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py b/tests/unit/integrations/gitlab/test_gitlab_client.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_client.py rename to tests/unit/integrations/gitlab/test_gitlab_client.py diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_integration.py b/tests/unit/integrations/gitlab/test_gitlab_integration.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_integration.py rename to tests/unit/integrations/gitlab/test_gitlab_integration.py diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py b/tests/unit/integrations/gitlab/test_gitlab_prompt_manager.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py rename to tests/unit/integrations/gitlab/test_gitlab_prompt_manager.py diff --git a/tests/unit/integrations/langfuse/__init__.py b/tests/unit/integrations/langfuse/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/langfuse/test_gemini_cached_tokens.py b/tests/unit/integrations/langfuse/test_gemini_cached_tokens.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_gemini_cached_tokens.py rename to tests/unit/integrations/langfuse/test_gemini_cached_tokens.py diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py b/tests/unit/integrations/langfuse/test_langfuse_prompt_management.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py rename to tests/unit/integrations/langfuse/test_langfuse_prompt_management.py diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/unit/integrations/langfuse/test_langfuse_sdk.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py rename to tests/unit/integrations/langfuse/test_langfuse_sdk.py diff --git a/tests/unit/integrations/newrelic/__init__.py b/tests/unit/integrations/newrelic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic.py b/tests/unit/integrations/newrelic/test_newrelic.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic.py rename to tests/unit/integrations/newrelic/test_newrelic.py diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py b/tests/unit/integrations/newrelic/test_newrelic_metrics.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py rename to tests/unit/integrations/newrelic/test_newrelic_metrics.py diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py b/tests/unit/integrations/newrelic/test_newrelic_team_handler.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py rename to tests/unit/integrations/newrelic/test_newrelic_team_handler.py diff --git a/tests/unit/integrations/open_telemetry/__init__.py b/tests/unit/integrations/open_telemetry/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/open_telemetry/_helpers.py b/tests/unit/integrations/open_telemetry/_helpers.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/_helpers.py rename to tests/unit/integrations/open_telemetry/_helpers.py diff --git a/tests/test_litellm/integrations/open_telemetry/conftest.py b/tests/unit/integrations/open_telemetry/conftest.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/conftest.py rename to tests/unit/integrations/open_telemetry/conftest.py diff --git a/tests/unit/integrations/open_telemetry/data/__init__.py b/tests/unit/integrations/open_telemetry/data/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json b/tests/unit/integrations/open_telemetry/data/captured_kwargs.json similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json rename to tests/unit/integrations/open_telemetry/data/captured_kwargs.json diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json b/tests/unit/integrations/open_telemetry/data/captured_response.json similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/data/captured_response.json rename to tests/unit/integrations/open_telemetry/data/captured_response.json diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_exception_handler.py b/tests/unit/integrations/open_telemetry/test_otel_exception_handler.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_exception_handler.py rename to tests/unit/integrations/open_telemetry/test_otel_exception_handler.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_passthrough_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_passthrough_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_passthrough_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_passthrough_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_unified_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_unified_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_unified_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_unified_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py b/tests/unit/integrations/open_telemetry/test_passthrough_parent_span.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py rename to tests/unit/integrations/open_telemetry/test_passthrough_parent_span.py diff --git a/tests/unit/integrations/otel/__init__.py b/tests/unit/integrations/otel/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/otel/test_db_endpoint.py b/tests/unit/integrations/otel/test_db_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_db_endpoint.py rename to tests/unit/integrations/otel/test_db_endpoint.py diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/unit/integrations/otel/test_langfuse_logger.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_langfuse_logger.py rename to tests/unit/integrations/otel/test_langfuse_logger.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_baggage.py b/tests/unit/integrations/otel/test_otel_v2_baggage.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_baggage.py rename to tests/unit/integrations/otel/test_otel_v2_baggage.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py similarity index 99% rename from tests/test_litellm/integrations/otel/test_otel_v2_components.py rename to tests/unit/integrations/otel/test_otel_v2_components.py index 07705e17d9a..fd10210c5ba 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -42,7 +42,7 @@ from opentelemetry.trace.propagation.tracecontext import ( # noqa: E402 ) import litellm # noqa: E402 -from conftest import TlsSink # noqa: E402 +from tests.unit.integrations.conftest import TlsSink # noqa: E402 from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.model.config import OpenTelemetryV2Config # noqa: E402 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py b/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py rename to tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_destinations.py rename to tests/unit/integrations/otel/test_otel_v2_destinations.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/unit/integrations/otel/test_otel_v2_dynamic.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py rename to tests/unit/integrations/otel/test_otel_v2_dynamic.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_emitter.py rename to tests/unit/integrations/otel/test_otel_v2_emitter.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py similarity index 99% rename from tests/test_litellm/integrations/otel/test_otel_v2_logger.py rename to tests/unit/integrations/otel/test_otel_v2_logger.py index 00c1343f72e..d478c670e58 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -2945,14 +2945,6 @@ def test_success_without_pre_call_emits_deferred_span(): assert spans[0].end_time == 101_500_000_000 -def test_no_carrier_and_no_payload_is_noop(): - logger, exporter = _logger() - asyncio.run( - logger.async_log_success_event({"litellm_params": {}}, None, None, None) - ) - assert exporter.get_finished_spans() == () - - def test_second_close_for_same_call_does_not_duplicate_span(): """Success and failure can both fire on one logging object for the same call id. The first close pops the carrier and finishes the boundary span; the diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py b/tests/unit/integrations/otel/test_otel_v2_metrics.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_metrics.py rename to tests/unit/integrations/otel/test_otel_v2_metrics.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_mount.py b/tests/unit/integrations/otel/test_otel_v2_mount.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_mount.py rename to tests/unit/integrations/otel/test_otel_v2_mount.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py b/tests/unit/integrations/otel/test_otel_v2_multibackend.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py rename to tests/unit/integrations/otel/test_otel_v2_multibackend.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_presets.py b/tests/unit/integrations/otel/test_otel_v2_presets.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_presets.py rename to tests/unit/integrations/otel/test_otel_v2_presets.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py rename to tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py rename to tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py diff --git a/tests/test_litellm/integrations/otel/test_runtime.py b/tests/unit/integrations/otel/test_runtime.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_runtime.py rename to tests/unit/integrations/otel/test_runtime.py diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/unit/integrations/rubrik_test_helpers.py similarity index 100% rename from tests/test_litellm/integrations/rubrik_test_helpers.py rename to tests/unit/integrations/rubrik_test_helpers.py diff --git a/tests/test_litellm/integrations/test_agentops.py b/tests/unit/integrations/test_agentops.py similarity index 100% rename from tests/test_litellm/integrations/test_agentops.py rename to tests/unit/integrations/test_agentops.py diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py similarity index 100% rename from tests/test_litellm/integrations/test_anthropic_cache_control_hook.py rename to tests/unit/integrations/test_anthropic_cache_control_hook.py diff --git a/tests/test_litellm/integrations/test_athina.py b/tests/unit/integrations/test_athina.py similarity index 100% rename from tests/test_litellm/integrations/test_athina.py rename to tests/unit/integrations/test_athina.py diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/unit/integrations/test_azure_sentinel.py similarity index 100% rename from tests/test_litellm/integrations/test_azure_sentinel.py rename to tests/unit/integrations/test_azure_sentinel.py diff --git a/tests/test_litellm/integrations/test_braintrust_logging.py b/tests/unit/integrations/test_braintrust_logging.py similarity index 100% rename from tests/test_litellm/integrations/test_braintrust_logging.py rename to tests/unit/integrations/test_braintrust_logging.py diff --git a/tests/test_litellm/integrations/test_braintrust_span_name.py b/tests/unit/integrations/test_braintrust_span_name.py similarity index 100% rename from tests/test_litellm/integrations/test_braintrust_span_name.py rename to tests/unit/integrations/test_braintrust_span_name.py diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_guardrail.py rename to tests/unit/integrations/test_custom_guardrail.py diff --git a/tests/test_litellm/integrations/test_custom_guardrail_recursion.py b/tests/unit/integrations/test_custom_guardrail_recursion.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_guardrail_recursion.py rename to tests/unit/integrations/test_custom_guardrail_recursion.py diff --git a/tests/test_litellm/integrations/test_custom_prompt_management.py b/tests/unit/integrations/test_custom_prompt_management.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_prompt_management.py rename to tests/unit/integrations/test_custom_prompt_management.py diff --git a/tests/test_litellm/integrations/test_deepeval.py b/tests/unit/integrations/test_deepeval.py similarity index 100% rename from tests/test_litellm/integrations/test_deepeval.py rename to tests/unit/integrations/test_deepeval.py diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/unit/integrations/test_galileo.py similarity index 100% rename from tests/test_litellm/integrations/test_galileo.py rename to tests/unit/integrations/test_galileo.py diff --git a/tests/test_litellm/integrations/test_guardrail_logging_sync.py b/tests/unit/integrations/test_guardrail_logging_sync.py similarity index 100% rename from tests/test_litellm/integrations/test_guardrail_logging_sync.py rename to tests/unit/integrations/test_guardrail_logging_sync.py diff --git a/tests/test_litellm/integrations/test_helicone.py b/tests/unit/integrations/test_helicone.py similarity index 100% rename from tests/test_litellm/integrations/test_helicone.py rename to tests/unit/integrations/test_helicone.py diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/unit/integrations/test_langfuse.py similarity index 100% rename from tests/test_litellm/integrations/test_langfuse.py rename to tests/unit/integrations/test_langfuse.py diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py similarity index 100% rename from tests/test_litellm/integrations/test_langfuse_otel.py rename to tests/unit/integrations/test_langfuse_otel.py diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/unit/integrations/test_langsmith_init.py similarity index 100% rename from tests/test_litellm/integrations/test_langsmith_init.py rename to tests/unit/integrations/test_langsmith_init.py diff --git a/tests/test_litellm/integrations/test_lunary.py b/tests/unit/integrations/test_lunary.py similarity index 100% rename from tests/test_litellm/integrations/test_lunary.py rename to tests/unit/integrations/test_lunary.py diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/unit/integrations/test_mlflow.py similarity index 100% rename from tests/test_litellm/integrations/test_mlflow.py rename to tests/unit/integrations/test_mlflow.py diff --git a/tests/test_litellm/integrations/test_openmeter.py b/tests/unit/integrations/test_openmeter.py similarity index 100% rename from tests/test_litellm/integrations/test_openmeter.py rename to tests/unit/integrations/test_openmeter.py diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py similarity index 99% rename from tests/test_litellm/integrations/test_opentelemetry.py rename to tests/unit/integrations/test_opentelemetry.py index 974961f2eb5..52eeec31e71 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -33,7 +33,7 @@ from parameterized import parameterized import requests -from conftest import TlsSink, write_self_signed_cert +from tests.unit.integrations.conftest import TlsSink, write_self_signed_cert import litellm from litellm.integrations import opentelemetry as otel_module from litellm.integrations.opentelemetry import ( @@ -1244,64 +1244,6 @@ class TestOpenTelemetry(unittest.TestCase): time.sleep(self.POLL_INTERVAL) return [] - @patch("litellm.integrations.opentelemetry.datetime") - def test_create_guardrail_span_with_valid_info(self, mock_datetime): - # Setup - otel = OpenTelemetry() - otel.tracer = MagicMock() - mock_span = MagicMock() - otel.tracer.start_span.return_value = mock_span - - # Create guardrail information - guardrail_info = { - "guardrail_name": "test_guardrail", - "guardrail_mode": "input", - "masked_entity_count": {"CREDIT_CARD": 2}, - "guardrail_response": "filtered_content", - "start_time": 1609459200.0, - "end_time": 1609459201.0, - } - - # Create a kwargs dict with standard_logging_object containing guardrail information - kwargs = { - "standard_logging_object": {"guardrail_information": [guardrail_info]} - } - - # Call the method - otel._create_guardrail_span(kwargs=kwargs, context=None) - - # Assertions - otel.tracer.start_span.assert_called_once() - - # print all calls to mock_span.set_attribute - print("Calls to mock_span.set_attribute:") - for call in mock_span.set_attribute.call_args_list: - print(call) - - # Check that the span has the correct attributes set - mock_span.set_attribute.assert_any_call("guardrail_name", "test_guardrail") - mock_span.set_attribute.assert_any_call("guardrail_mode", "input") - mock_span.set_attribute.assert_any_call( - "guardrail_response", safe_dumps("filtered_content") - ) - mock_span.set_attribute.assert_any_call( - "masked_entity_count", safe_dumps({"CREDIT_CARD": 2}) - ) - - # Verify that the span was ended - mock_span.end.assert_called_once() - - def test_create_guardrail_span_with_no_info(self): - # Setup - otel = OpenTelemetry() - otel.tracer = MagicMock() - - # Test with no guardrail information - kwargs = {"standard_logging_object": {}} - otel._create_guardrail_span(kwargs=kwargs, context=None) - - # Verify that start_span was never called - otel.tracer.start_span.assert_not_called() def test_get_tracer_to_use_for_request_with_dynamic_headers(self): """Test that get_tracer_to_use_for_request returns a dynamic tracer when dynamic headers are present.""" @@ -5461,10 +5403,6 @@ class TestOpenTelemetryPreprocessingDuration(unittest.TestCase): ) assert "litellm.preprocessing.duration_ms" not in self._attr(span, exp) - def test_none_span_is_noop(self): - OpenTelemetry().set_preprocessing_duration_attribute( - None, {"first_api_call_start_time": datetime(2026, 1, 1)} - ) def test_non_dict_container_is_noop(self): otel = OpenTelemetry() diff --git a/tests/test_litellm/integrations/test_opentelemetry_dynamic_imports.py b/tests/unit/integrations/test_opentelemetry_dynamic_imports.py similarity index 100% rename from tests/test_litellm/integrations/test_opentelemetry_dynamic_imports.py rename to tests/unit/integrations/test_opentelemetry_dynamic_imports.py diff --git a/tests/test_litellm/integrations/test_opik_utils.py b/tests/unit/integrations/test_opik_utils.py similarity index 100% rename from tests/test_litellm/integrations/test_opik_utils.py rename to tests/unit/integrations/test_opik_utils.py diff --git a/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py b/tests/unit/integrations/test_otel_guardrail_violation_spans.py similarity index 100% rename from tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py rename to tests/unit/integrations/test_otel_guardrail_violation_spans.py diff --git a/tests/test_litellm/integrations/test_otel_team_attributes_matrix.py b/tests/unit/integrations/test_otel_team_attributes_matrix.py similarity index 100% rename from tests/test_litellm/integrations/test_otel_team_attributes_matrix.py rename to tests/unit/integrations/test_otel_team_attributes_matrix.py diff --git a/tests/test_litellm/integrations/test_prometheus_api_promql_escape.py b/tests/unit/integrations/test_prometheus_api_promql_escape.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_api_promql_escape.py rename to tests/unit/integrations/test_prometheus_api_promql_escape.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py b/tests/unit/integrations/test_prometheus_budget_metric_guard.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py rename to tests/unit/integrations/test_prometheus_budget_metric_guard.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metrics_db_lookups.py b/tests/unit/integrations/test_prometheus_budget_metrics_db_lookups.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metrics_db_lookups.py rename to tests/unit/integrations/test_prometheus_budget_metrics_db_lookups.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metrics_timeout.py b/tests/unit/integrations/test_prometheus_budget_metrics_timeout.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metrics_timeout.py rename to tests/unit/integrations/test_prometheus_budget_metrics_timeout.py diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/unit/integrations/test_prometheus_cache_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_cache_metrics.py rename to tests/unit/integrations/test_prometheus_cache_metrics.py index aa031bb813b..21f13f0ec5c 100644 --- a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py +++ b/tests/unit/integrations/test_prometheus_cache_metrics.py @@ -1,7 +1,7 @@ """ Unit tests for cache Prometheus metrics. -Run with: uv run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v +Run with: uv run pytest tests/unit/integrations/test_prometheus_cache_metrics.py -v """ import pytest diff --git a/tests/test_litellm/integrations/test_prometheus_caller_identity.py b/tests/unit/integrations/test_prometheus_caller_identity.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_caller_identity.py rename to tests/unit/integrations/test_prometheus_caller_identity.py diff --git a/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_carried_budget_state.py rename to tests/unit/integrations/test_prometheus_carried_budget_state.py diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/unit/integrations/test_prometheus_client_ip_user_agent.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py rename to tests/unit/integrations/test_prometheus_client_ip_user_agent.py diff --git a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py b/tests/unit/integrations/test_prometheus_custom_metadata_label_counts.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py rename to tests/unit/integrations/test_prometheus_custom_metadata_label_counts.py diff --git a/tests/test_litellm/integrations/test_prometheus_deployment_state_proxy_rejects.py b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_deployment_state_proxy_rejects.py rename to tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py diff --git a/tests/test_litellm/integrations/test_prometheus_end_user_cardinality.py b/tests/unit/integrations/test_prometheus_end_user_cardinality.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_end_user_cardinality.py rename to tests/unit/integrations/test_prometheus_end_user_cardinality.py diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/unit/integrations/test_prometheus_input_sequence_length_label.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py rename to tests/unit/integrations/test_prometheus_input_sequence_length_label.py diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py rename to tests/unit/integrations/test_prometheus_invalid_key_filtering.py diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_labels.py rename to tests/unit/integrations/test_prometheus_labels.py diff --git a/tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py b/tests/unit/integrations/test_prometheus_mcp_tool_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py rename to tests/unit/integrations/test_prometheus_mcp_tool_metrics.py index 22c36f00ca9..da5a0b35e9d 100644 --- a/tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py +++ b/tests/unit/integrations/test_prometheus_mcp_tool_metrics.py @@ -5,7 +5,7 @@ These metrics expose ``mcp_tool_call_metadata`` in Prometheus so Grafana dashboards can break down MCP usage by server and tool name. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py -v + uv run pytest tests/unit/integrations/test_prometheus_mcp_tool_metrics.py -v """ from typing import get_args diff --git a/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py b/tests/unit/integrations/test_prometheus_media_generation_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py rename to tests/unit/integrations/test_prometheus_media_generation_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/unit/integrations/test_prometheus_metric_name_consistency.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py rename to tests/unit/integrations/test_prometheus_metric_name_consistency.py diff --git a/tests/test_litellm/integrations/test_prometheus_metrics_endpoint.py b/tests/unit/integrations/test_prometheus_metrics_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_metrics_endpoint.py rename to tests/unit/integrations/test_prometheus_metrics_endpoint.py diff --git a/tests/test_litellm/integrations/test_prometheus_missing_metrics.py b/tests/unit/integrations/test_prometheus_missing_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_missing_metrics.py rename to tests/unit/integrations/test_prometheus_missing_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_none_metadata.py b/tests/unit/integrations/test_prometheus_none_metadata.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_none_metadata.py rename to tests/unit/integrations/test_prometheus_none_metadata.py diff --git a/tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py b/tests/unit/integrations/test_prometheus_overhead_with_guardrails.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py rename to tests/unit/integrations/test_prometheus_overhead_with_guardrails.py diff --git a/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py b/tests/unit/integrations/test_prometheus_queue_guardrail_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py rename to tests/unit/integrations/test_prometheus_queue_guardrail_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/unit/integrations/test_prometheus_rate_limit_labels.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py rename to tests/unit/integrations/test_prometheus_rate_limit_labels.py diff --git a/tests/test_litellm/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_remaining_tokens_router_fallback.py rename to tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py diff --git a/tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py b/tests/unit/integrations/test_prometheus_requested_model_cardinality.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py rename to tests/unit/integrations/test_prometheus_requested_model_cardinality.py diff --git a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py b/tests/unit/integrations/test_prometheus_service_tier_label.py similarity index 98% rename from tests/test_litellm/integrations/test_prometheus_service_tier_label.py rename to tests/unit/integrations/test_prometheus_service_tier_label.py index b2212c4ff41..8b8131b5af2 100644 --- a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py +++ b/tests/unit/integrations/test_prometheus_service_tier_label.py @@ -6,7 +6,7 @@ between the tier a provider served and the tier a caller requested, and the end-to-end emit wiring through async_log_success_event. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_service_tier_label.py -v + uv run pytest tests/unit/integrations/test_prometheus_service_tier_label.py -v """ import datetime diff --git a/tests/test_litellm/integrations/test_prometheus_services.py b/tests/unit/integrations/test_prometheus_services.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_services.py rename to tests/unit/integrations/test_prometheus_services.py diff --git a/tests/test_litellm/integrations/test_prometheus_spend_capture_rate.py b/tests/unit/integrations/test_prometheus_spend_capture_rate.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_spend_capture_rate.py rename to tests/unit/integrations/test_prometheus_spend_capture_rate.py diff --git a/tests/test_litellm/integrations/test_prometheus_spend_logs_metadata.py b/tests/unit/integrations/test_prometheus_spend_logs_metadata.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_spend_logs_metadata.py rename to tests/unit/integrations/test_prometheus_spend_logs_metadata.py diff --git a/tests/test_litellm/integrations/test_prometheus_stream_label.py b/tests/unit/integrations/test_prometheus_stream_label.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_stream_label.py rename to tests/unit/integrations/test_prometheus_stream_label.py diff --git a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py b/tests/unit/integrations/test_prometheus_token_detail_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py rename to tests/unit/integrations/test_prometheus_token_detail_metrics.py index 5e3846d6fa2..03d67316080 100644 --- a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py +++ b/tests/unit/integrations/test_prometheus_token_detail_metrics.py @@ -6,7 +6,7 @@ from the Usage object that providers report. They are sparse — only incremented when the underlying detail is populated and > 0. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py -v + uv run pytest tests/unit/integrations/test_prometheus_token_detail_metrics.py -v """ from typing import get_args diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py similarity index 98% rename from tests/test_litellm/integrations/test_prometheus_user_team_metrics.py rename to tests/unit/integrations/test_prometheus_user_team_metrics.py index 0fc91748af2..ab0fb67b52f 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -102,27 +102,6 @@ class TestPrometheusUserTeamCountMetrics: f"litellm_teams_count_metric should accept value {value}: {e}" ) - def test_user_count_metric_with_zero(self, prometheus_logger): - """Test that user count metric handles zero users""" - metric = prometheus_logger.litellm_total_users_metric - - # Should handle zero gracefully - try: - metric.set(0) - assert True - except Exception as e: - pytest.fail(f"litellm_total_users_metric should handle zero: {e}") - - def test_team_count_metric_with_zero(self, prometheus_logger): - """Test that team count metric handles zero teams""" - metric = prometheus_logger.litellm_teams_count_metric - - # Should handle zero gracefully - try: - metric.set(0) - assert True - except Exception as e: - pytest.fail(f"litellm_teams_count_metric should handle zero: {e}") def test_metrics_can_be_updated_multiple_times(self, prometheus_logger): """Test that metrics can be updated multiple times (simulating refresh cycle)""" diff --git a/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py b/tests/unit/integrations/test_prometheus_zero_cost_metric.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py rename to tests/unit/integrations/test_prometheus_zero_cost_metric.py diff --git a/tests/test_litellm/integrations/test_prompt_manager_ssti.py b/tests/unit/integrations/test_prompt_manager_ssti.py similarity index 100% rename from tests/test_litellm/integrations/test_prompt_manager_ssti.py rename to tests/unit/integrations/test_prompt_manager_ssti.py diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/unit/integrations/test_responses_background_cost.py similarity index 100% rename from tests/test_litellm/integrations/test_responses_background_cost.py rename to tests/unit/integrations/test_responses_background_cost.py diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/unit/integrations/test_rubrik.py similarity index 99% rename from tests/test_litellm/integrations/test_rubrik.py rename to tests/unit/integrations/test_rubrik.py index 4a2ee487c65..f3fea292bde 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/unit/integrations/test_rubrik.py @@ -19,7 +19,7 @@ from litellm.integrations.rubrik import ( ) from litellm.proxy._types import UserAPIKeyAuth -from tests.test_litellm.integrations.rubrik_test_helpers import ( +from tests.unit.integrations.rubrik_test_helpers import ( make_inputs_with_tools, make_tool_call_dict, ) diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/unit/integrations/test_s3.py similarity index 100% rename from tests/test_litellm/integrations/test_s3.py rename to tests/unit/integrations/test_s3.py diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py similarity index 100% rename from tests/test_litellm/integrations/test_s3_v2.py rename to tests/unit/integrations/test_s3_v2.py diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/unit/integrations/test_shadow_eval_logger.py similarity index 100% rename from tests/test_litellm/integrations/test_shadow_eval_logger.py rename to tests/unit/integrations/test_shadow_eval_logger.py diff --git a/tests/test_litellm/integrations/test_weave_otel.py b/tests/unit/integrations/test_weave_otel.py similarity index 100% rename from tests/test_litellm/integrations/test_weave_otel.py rename to tests/unit/integrations/test_weave_otel.py diff --git a/tests/unit/integrations/websearch_interception/__init__.py b/tests/unit/integrations/websearch_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py b/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py rename to tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py rename to tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py rename to tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_thinking.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_thinking.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_interception_thinking.py rename to tests/unit/integrations/websearch_interception/test_websearch_interception_thinking.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py rename to tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_responses.py rename to tests/unit/integrations/websearch_interception/test_websearch_responses.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py rename to tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py rename to tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py b/tests/unit/integrations/websearch_interception/test_websearch_streaming_wrap.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py rename to tests/unit/integrations/websearch_interception/test_websearch_streaming_wrap.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py rename to tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py diff --git a/tests/unit/secret_managers/__init__.py b/tests/unit/secret_managers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/unit/secret_managers/hashicorp_vault_parity.json similarity index 100% rename from tests/test_litellm/secret_managers/hashicorp_vault_parity.json rename to tests/unit/secret_managers/hashicorp_vault_parity.json diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py b/tests/unit/secret_managers/test_aws_secret_manager_replication.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py rename to tests/unit/secret_managers/test_aws_secret_manager_replication.py diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py b/tests/unit/secret_managers/test_aws_secret_manager_rotation.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py rename to tests/unit/secret_managers/test_aws_secret_manager_rotation.py diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/unit/secret_managers/test_aws_secret_manager_v2.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py rename to tests/unit/secret_managers/test_aws_secret_manager_v2.py diff --git a/tests/test_litellm/secret_managers/test_base_secret_manager.py b/tests/unit/secret_managers/test_base_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_base_secret_manager.py rename to tests/unit/secret_managers/test_base_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_custom_secret_manager.py b/tests/unit/secret_managers/test_custom_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_custom_secret_manager.py rename to tests/unit/secret_managers/test_custom_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_cyberark_secret_manager.py b/tests/unit/secret_managers/test_cyberark_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_cyberark_secret_manager.py rename to tests/unit/secret_managers/test_cyberark_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py similarity index 100% rename from tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py rename to tests/unit/secret_managers/test_get_azure_ad_token_provider.py diff --git a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py b/tests/unit/secret_managers/test_hashicorp_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py rename to tests/unit/secret_managers/test_hashicorp_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_secret_manager_handler.py b/tests/unit/secret_managers/test_secret_manager_handler.py similarity index 100% rename from tests/test_litellm/secret_managers/test_secret_manager_handler.py rename to tests/unit/secret_managers/test_secret_manager_handler.py diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py similarity index 100% rename from tests/test_litellm/secret_managers/test_secret_managers_main.py rename to tests/unit/secret_managers/test_secret_managers_main.py From 6b7688869e098778356d84d226427f09a0714b06 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 25 Sep 2026 20:13:52 +0000 Subject: [PATCH 15/58] feat(mcp): configure protocol versions and capability discovery (#43169) * feat(mcp): configure protocol versions and capability discovery * test(mcp): return SDK initialization result in REST pagination fixture * test(mcp): arm cancellation deadlines after TCP calls start * fix(mcp): avoid serialized discovery and listing spend logs * test(mcp): isolate default protocol header policy * fix(mcp): honor preview protocol pins and refresh migrated tests * fix(mcp): preserve edited protocol pins in saved OAuth previews * fix(mcp): retain saved protocol pins when previews omit versions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 48 +++++- .../_experimental/mcp_server/capabilities.py | 145 ++++++++++++++++++ .../_experimental/mcp_server/contracts.py | 1 + .../mcp_server/mcp_server_manager.py | 15 ++ .../_experimental/mcp_server/operations.py | 75 ++++++++- .../mcp_server/rest_endpoints.py | 12 +- .../proxy/_experimental/mcp_server/server.py | 20 ++- litellm/proxy/_types.py | 6 + litellm/proxy/proxy_server.py | 5 + litellm/types/mcp.py | 16 +- .../types/mcp_server/mcp_server_manager.py | 23 ++- .../mcp/test_mcp_protocol_errors.py | 39 +++++ tests/integration/mcp/test_mcp_transports.py | 40 +++++ .../mcp_server/test_capabilities.py | 106 +++++++++++++ .../mcp_server/test_mcp_server.py | 32 +++- .../mcp_server/test_mcp_server_manager.py | 26 +++- .../mcp_server/test_operations.py | 123 +++++++++++++++ .../mcp_server/test_rest_endpoints.py | 81 +++++++++- .../proxy/proxy_server/test_proxy_config.py | 16 ++ tests/test_litellm/proxy/test__types.py | 18 +++ .../test_mcp_client.py | 95 +++++++++--- .../mcp_server/test_mcp_client_unit.py | 29 +++- .../mcp_server/test_mcp_server.py | 4 +- tests/unit/test_unit_shard_missing_paths.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 25 files changed, 939 insertions(+), 42 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/capabilities.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 01670be74c8..4e3b92edc89 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -10,6 +10,7 @@ import os from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Sequence from contextlib import AbstractAsyncContextManager from functools import partial +from importlib.metadata import version from types import MappingProxyType from typing import Final, TypeAlias, TypeVar, cast @@ -34,8 +35,16 @@ _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams] from mcp.types import ( METHOD_NOT_FOUND, REQUEST_TIMEOUT, + ClientCapabilities, + ElicitationCapability, + FormElicitationCapability, GetPromptRequestParams, GetPromptResult, + Implementation, + InitializedNotification, + InitializeRequest, + InitializeRequestParams, + InitializeResult, InputRequiredResult, ListPromptsResult, ListResourcesResult, @@ -44,12 +53,14 @@ from mcp.types import ( PaginatedResult, Prompt, ResourceTemplate, + SamplingCapability, ServerNotification, + UrlElicitationCapability, ) from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -64,11 +75,13 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_er from litellm.proxy._experimental.mcp_server.result_conversion import error_text_result from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( + MCP_LEGACY_VERSIONS, MCPAuth, MCPAuthType, MCPStdioConfig, MCPTransport, MCPTransportType, + MCPUpstreamProtocol, credential_redirect_hook, has_header, without_header, @@ -386,7 +399,9 @@ class MCPClient: sampling_callback: Callable | None = None, elicitation_callback: Callable | None = None, logging_callback: Callable | None = None, + protocol_version: MCPUpstreamProtocol = "auto", ): + self.protocol_version: MCPUpstreamProtocol = TypeAdapter(MCPUpstreamProtocol).validate_python(protocol_version) self.server_url: str = server_url self.transport_type: MCPTransport = transport_type self.auth_type: MCPAuthType = auth_type @@ -525,6 +540,35 @@ class MCPClient: return safe_env + async def _initialize_session(self, session: ClientSession) -> InitializeResult: + if self.protocol_version == "auto": + automatic: Final = await session.initialize() + if automatic.protocol_version not in MCP_LEGACY_VERSIONS: + raise MCPError(code=-32022, message="Upstream selected an unsupported MCP protocol version") + return automatic + result: Final = await session.send_request( + InitializeRequest( + params=InitializeRequestParams( + protocol_version=self.protocol_version, + client_info=Implementation(name="litellm", version=version("litellm")), + capabilities=ClientCapabilities( + sampling=SamplingCapability() if self._sampling_callback is not None else None, + elicitation=ElicitationCapability( + form=FormElicitationCapability(), url=UrlElicitationCapability() + ) + if self._elicitation_callback is not None + else None, + ), + ) + ), + InitializeResult, + ) + if result.protocol_version != self.protocol_version: + raise MCPError(code=-32022, message="Upstream did not accept the configured MCP protocol version") + session.adopt(result) + await session.send_notification(InitializedNotification()) + return result + async def _execute_session_operation( self, transport_ctx: _TransportContext, @@ -579,7 +623,7 @@ class MCPClient: ) session: Final = await session_ctx.__aenter__() try: - init_result: Final = await session.initialize() + init_result: Final = await self._initialize_session(session) instructions: Final = getattr(init_result, "instructions", None) self._last_initialize_instructions = ( instructions.strip() or None if isinstance(instructions, str) else None diff --git a/litellm/proxy/_experimental/mcp_server/capabilities.py b/litellm/proxy/_experimental/mcp_server/capabilities.py new file mode 100644 index 00000000000..bfd00327eb4 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/capabilities.py @@ -0,0 +1,145 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from itertools import product +from types import MappingProxyType +from typing import Final, Literal + +from mcp.server.context import CallNext, HandlerResult, ServerRequestContext +from mcp.shared.exceptions import MCPError +from mcp.types import DiscoverResult, InitializeRequestParams, InitializeResult, ServerCapabilities +from mcp_types.methods import CLIENT_REQUESTS +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION +from pydantic import TypeAdapter + +from litellm.types.mcp import MCP_LEGACY_VERSIONS, MCPAdvertisedVersions, MCPLegacyVersion, MCPSpecVersion, MCPTransport + +GATEWAY_OPERATIONS: Final = frozenset( + { + "tools/list", + "tools/call", + "prompts/list", + "prompts/get", + "resources/list", + "resources/read", + "resources/templates/list", + } +) + + +@dataclass(frozen=True, slots=True) +class RevisionSupport: + transports: frozenset[MCPTransport] + operations: frozenset[str] + results: frozenset[Literal["complete", "input_required"]] + extensions: frozenset[str] + completed: bool + + +REVISION_SUPPORT: Final[Mapping[str, RevisionSupport]] = MappingProxyType( + { + version.value: RevisionSupport( + transports=frozenset(MCPTransport) + if version.value in HANDSHAKE_PROTOCOL_VERSIONS + else frozenset({MCPTransport.http, MCPTransport.stdio}), + operations=frozenset(method for method in GATEWAY_OPERATIONS if (method, version.value) in CLIENT_REQUESTS), + results=frozenset({"complete"}) + if version.value in HANDSHAKE_PROTOCOL_VERSIONS + else frozenset({"complete", "input_required"}), + extensions=frozenset(), + completed=version.value in HANDSHAKE_PROTOCOL_VERSIONS, + ) + for version in MCPSpecVersion + } +) +_COMPLETED_REVISIONS: Final = tuple(version for version, support in REVISION_SUPPORT.items() if support.completed) +TRANSLATION_PAIRS: Final = frozenset(product(_COMPLETED_REVISIONS, repeat=2)) +_ADVERTISED_VERSIONS: Final[TypeAdapter[tuple[MCPLegacyVersion, ...]]] = TypeAdapter(MCPAdvertisedVersions) + + +def configured_versions() -> tuple[str, ...]: + from litellm.proxy.proxy_server import general_settings_view + + configured: Final = general_settings_view().get("mcp_advertised_versions") + return _ADVERTISED_VERSIONS.validate_python(MCP_LEGACY_VERSIONS if configured is None else configured) + + +def build_discovery( + *, + configured: tuple[str, ...], + revision: str, + transport: MCPTransport, + authorized_operations: frozenset[str], + upstream_versions: frozenset[str], + capabilities: ServerCapabilities, + client_extensions: frozenset[str] = frozenset(), + upstream_extensions: frozenset[str] = frozenset(), + instructions: str | None = None, +) -> DiscoverResult: + supported: Final = tuple( + version + for version, support in REVISION_SUPPORT.items() + if version in configured and support.completed and transport in support.transports + ) + revision_support: Final = REVISION_SUPPORT.get(revision) + operations: Final[frozenset[str]] = ( + authorized_operations & revision_support.operations + if revision in supported + and revision_support is not None + and any((revision, upstream) in TRANSLATION_PAIRS for upstream in upstream_versions) + else frozenset() + ) + extensions: Final[frozenset[str]] = ( + revision_support.extensions & client_extensions & upstream_extensions + if operations and revision_support is not None + else frozenset() + ) + caller_capabilities: Final = capabilities.model_copy(deep=True) + return DiscoverResult( + supported_versions=list(supported), + capabilities=ServerCapabilities( + tools=caller_capabilities.tools if {"tools/list", "tools/call"} <= operations else None, + prompts=caller_capabilities.prompts if {"prompts/list", "prompts/get"} <= operations else None, + resources=caller_capabilities.resources if {"resources/list", "resources/read"} <= operations else None, + extensions={ + key: value for key, value in (caller_capabilities.extensions or {}).items() if key in extensions + } + or None, + ), + instructions=instructions, + cache_scope="private", + ttl_ms=0, + ) + + +class GatewayVersionPolicy: + def __init__(self, versions: Callable[[], tuple[str, ...]] = configured_versions) -> None: + self._versions = versions + + async def __call__(self, ctx: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult: + versions: Final = self._versions() + requested: Final = ( + InitializeRequestParams.model_validate(ctx.params or {}).protocol_version + if ctx.method == "initialize" + else ctx.protocol_version + ) + negotiated: Final = ( + (requested if requested in HANDSHAKE_PROTOCOL_VERSIONS else LATEST_HANDSHAKE_VERSION) + if ctx.method == "initialize" + else requested + ) + if negotiated not in versions: + raise MCPError(code=-32022, message="Unsupported MCP protocol version", data={"supported": list(versions)}) + result: Final = await call_next(ctx) + if ctx.method != "initialize": + return result + initialized: Final = InitializeResult.model_validate(result) + discovery: Final = build_discovery( + configured=versions, + revision=initialized.protocol_version, + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=initialized.capabilities, + instructions=initialized.instructions, + ) + return initialized.model_copy(update={"capabilities": discovery.capabilities}) diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index a88d400282c..a0a08dc08ce 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -28,6 +28,7 @@ class OperationContext: client_ip: str | None = None mcp_proxy_mode: bool = False wire_compat: WireCompat = WireCompat.LEGACY + protocol_version: str | None = None def __post_init__(self) -> None: object.__setattr__(self, "_caller", copy_caller(self._caller)) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 24cae976174..6baa695433c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -193,6 +193,7 @@ from litellm.types.mcp import ( MCPAuth, MCPStdioConfig, MCPTokenEndpointAuthMethod, + MCPUpstreamProtocol, has_header, without_header, ) @@ -340,6 +341,7 @@ class MCPServerConfig(TypedDict, total=False): whatever the admin wrote, and each read applies its own default.""" server_id: ReadOnly[str] + protocol_version: ReadOnly[MCPUpstreamProtocol] alias: str description: str mcp_info: MCPInfo @@ -2549,6 +2551,9 @@ class MCPServerManager: new_server = MCPServer( server_id=server_id, name=name_for_prefix, + protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( + server_config.get("protocol_version", mcp_info.get("protocol_version", "auto")) + ), alias=alias, server_name=server_name, spec_path=server_config.get("spec_path", None), @@ -3109,6 +3114,9 @@ class MCPServerManager: new_server: Final = MCPServer( server_id=mcp_server.server_id, name=name_for_prefix, + protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( + _mcp_info.get("protocol_version", "auto") + ), alias=getattr(mcp_server, "alias", None), server_name=getattr(mcp_server, "server_name", None), url=mcp_server.url, @@ -4145,6 +4153,7 @@ class MCPServerManager: cred_provider: UpstreamCredentialProvider | None = None, raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, + protocol_version_override: MCPUpstreamProtocol | None = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -4168,6 +4177,9 @@ class MCPServerManager: """ record_auth_resolution(server.server_id, AuthResolution.unresolved) resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + protocol_version: Final = ( + protocol_version_override if protocol_version_override is not None else resolved_server.protocol_version + ) transport: Final = resolved_server.transport or MCPTransport.sse spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) provider: Final = cred_provider or self._cred_provider @@ -4249,6 +4261,7 @@ class MCPServerManager: return MCPClient( server_url="", # Not used for stdio transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, auth_value=auth_value, timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), @@ -4281,6 +4294,7 @@ class MCPServerManager: MCPClient( server_url=server_url, transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, timeout=( resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT @@ -4324,6 +4338,7 @@ class MCPServerManager: MCPClient( server_url=server_url, transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, auth_value=auth_value, auth_header_name=auth_header_name, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index ebd26e4bf87..a19246b6e90 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -14,6 +14,8 @@ from mcp.types import ( CallToolRequest, CallToolRequestParams, CallToolResult, + DiscoverRequest, + DiscoverResult, GetPromptRequest, GetPromptRequestParams, GetPromptResult, @@ -28,10 +30,14 @@ from mcp.types import ( ListToolsResult, PaginatedRequestParams, Prompt, + PromptsCapability, ReadResourceRequest, ReadResourceRequestParams, + ResourcesCapability, ResourceTemplate, + ServerCapabilities, TextContent, + ToolsCapability, ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter @@ -51,6 +57,11 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) +from litellm.proxy._experimental.mcp_server.capabilities import ( + GATEWAY_OPERATIONS, + build_discovery, + configured_versions, +) from litellm.proxy._experimental.mcp_server.contracts import ( AuthorizedToolCall, OperationContext, @@ -122,7 +133,9 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.types.mcp import ( DEFAULT_CREDENTIAL_HEADER, + MCP_LEGACY_VERSIONS, MCPAuth, + MCPTransport, without_header, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer @@ -2657,7 +2670,11 @@ class _McpDeniedDetail(TypedDict): async def _execute_handle_list_tools( - context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None + context: OperationContext, + params: PaginatedRequestParams, + host_progress_callback: ProgressCallback | None = None, + *, + log_list_tools_to_spendlogs: bool = True, ) -> ListToolsResult: try: ( @@ -2700,7 +2717,7 @@ async def _execute_handle_list_tools( mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, - log_list_tools_to_spendlogs=True, + log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source="mcp_protocol", client_ip=_client_ip, ) @@ -3065,6 +3082,7 @@ def prepare_context( client_ip: str | None = None, mcp_proxy_mode: bool = False, wire_compat: WireCompat = WireCompat.LEGACY, + protocol_version: str | None = None, ) -> OperationContext: return OperationContext( _caller=user_api_key_auth, @@ -3076,11 +3094,13 @@ def prepare_context( client_ip=client_ip, mcp_proxy_mode=mcp_proxy_mode, wire_compat=wire_compat, + protocol_version=protocol_version, ) GatewayOperation: TypeAlias = ( AuthorizedToolCall + | DiscoverRequest | ListToolsRequest | CallToolRequest | ListPromptsRequest @@ -3090,7 +3110,8 @@ GatewayOperation: TypeAlias = ( | ReadResourceRequest ) GatewayResult: TypeAlias = ( - ListToolsResult + DiscoverResult + | ListToolsResult | CallToolResult | InputRequiredResult | ListPromptsResult @@ -3105,6 +3126,9 @@ class GatewayOperations: def __init__(self, host_progress_callback: ProgressCallback | None = None) -> None: self._host_progress_callback = host_progress_callback + @overload + async def execute(self, operation: DiscoverRequest, context: OperationContext) -> DiscoverResult: ... + @overload async def execute( self, operation: AuthorizedToolCall, context: OperationContext @@ -3137,6 +3161,51 @@ class GatewayOperations: async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: + case DiscoverRequest(): + listings: Final = ( + () + if context.mcp_proxy_mode + else (ListPromptsRequest(), ListResourcesRequest(), ListResourceTemplatesRequest()) + ) + tasks: Final = ( + asyncio.create_task( + _execute_handle_list_tools( + context, + PaginatedRequestParams(), + self._host_progress_callback, + log_list_tools_to_spendlogs=False, + ) + ), + *(asyncio.create_task(self.execute(listing, context)) for listing in listings), + ) + try: + results: Final = await asyncio.gather(*tasks) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return build_discovery( + configured=configured_versions(), + revision=context.protocol_version or "2025-11-25", + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=frozenset(MCP_LEGACY_VERSIONS), + capabilities=ServerCapabilities( + tools=ToolsCapability() + if any(isinstance(result, ListToolsResult) and result.tools for result in results) + else None, + prompts=PromptsCapability() + if any(isinstance(result, ListPromptsResult) and result.prompts for result in results) + else None, + resources=ResourcesCapability() + if any( + (isinstance(result, ListResourcesResult) and result.resources) + or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates) + for result in results + ) + else None, + ), + ) case AuthorizedToolCall(): auth, token, _servers, server_headers, oauth_headers, headers, _client_ip = context.legacy_auth() return await _execute_mcp_tool( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 7f519e2c0d9..02694f110b1 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1375,7 +1375,16 @@ if MCP_AVAILABLE: and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) else None ) - return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers) + preview_request: Final = ( + request.model_copy( + update={"mcp_info": {**(request.mcp_info or {}), "protocol_version": saved_server.protocol_version}} + ) + if saved_server is not None and "protocol_version" not in (request.mcp_info or {}) + else request + ) + return _StagedServerTest( + request=preview_request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers + ) async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: with anyio.move_on_after(deadline): @@ -1512,6 +1521,7 @@ if MCP_AVAILABLE: extra_headers=merged_headers, stdio_env=stdio_env, cred_provider=preview_cred_provider, + protocol_version_override=server_model.protocol_version, ) return await operation(client) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1bd31d971b0..555aebc7434 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -125,12 +125,14 @@ def unsupported_protocol_version(scope: Scope) -> str | None: ``HANDSHAKE_PROTOCOL_VERSIONS`` to the modern single-exchange path, which bypasses litellm's session/auth model, so the ASGI entry rejects it. """ + from litellm.proxy._experimental.mcp_server.capabilities import configured_versions + headers: Final[Iterable[tuple[bytes, bytes]]] = scope.get("headers") or () values: Final = tuple( raw.decode("latin-1").strip() for key, raw in headers if key.lower() == _MCP_PROTOCOL_VERSION_HEADER ) for value in values: - if value and value not in HANDSHAKE_PROTOCOL_VERSIONS: + if value and value not in configured_versions(): return value return None @@ -149,7 +151,10 @@ try: from mcp.server.session import ServerSession as _McpServerSession from mcp.types import ( BlobResourceContents, + DiscoverRequest, + DiscoverResult, GetPromptResult, + RequestParams, ResourceTemplate, TextResourceContents, ) @@ -526,11 +531,11 @@ if MCP_AVAILABLE: PaginatedRequestParams, ReadResourceRequestParams, ) - from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) + from litellm.proxy._experimental.mcp_server.capabilities import GatewayVersionPolicy, configured_versions from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, global_mcp_server_manager, @@ -585,6 +590,7 @@ if MCP_AVAILABLE: name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) + server.middleware.append(GatewayVersionPolicy()) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/sse/messages") @@ -830,6 +836,7 @@ if MCP_AVAILABLE: client_ip, _mcp_proxy_mode.get(), wire_compat_for(ctx.protocol_version), + ctx.protocol_version, ) async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: @@ -948,6 +955,11 @@ if MCP_AVAILABLE: ReadResourceRequest(params=params), context ) + async def discover(ctx: ServerRequestContext, params: RequestParams) -> DiscoverResult: + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations().execute(DiscoverRequest(params=params), context) + + server.add_request_handler("server/discover", RequestParams, discover) server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools) server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call) server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) @@ -1954,7 +1966,7 @@ if MCP_AVAILABLE: reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: - supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) + supported: Final = ", ".join(configured_versions()) await JSONResponse( status_code=400, content={ # mutable-ok: JSON-RPC error payload @@ -2299,7 +2311,7 @@ if MCP_AVAILABLE: reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: - supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) + supported: Final = ", ".join(configured_versions()) await JSONResponse( status_code=400, content={ # mutable-ok: JSON-RPC error payload diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3da30070b9e..89fa0058644 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -37,6 +37,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.mcp import ( + MCPAdvertisedVersions, MCPAllowedClient, MCPAuth, MCPAuthType, @@ -2998,6 +2999,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", ) + mcp_advertised_versions: MCPAdvertisedVersions | None = Field( + None, + description="MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. " + "Modern protocol serving and Apps/Tasks remain disabled.", + ) mcp_allowed_clients: list[MCPAllowedClient] | None = Field( None, description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ed4ea347c2e..61304f0d919 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6342,6 +6342,11 @@ class ProxyConfig: if general_settings is None: general_settings = {} + if general_settings.get("mcp_advertised_versions") is not None: + from litellm.types.mcp import MCPAdvertisedVersions + + TypeAdapter(MCPAdvertisedVersions).validate_python(general_settings["mcp_advertised_versions"]) + if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None: warn_login_counters_are_per_worker(os.getenv("NUM_WORKERS", "1")) if declared_proxy_ranges(general_settings) is None: diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 83e719810d5..fec5e84c8df 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -4,7 +4,7 @@ import enum import re from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal from urllib.parse import urlsplit import httpx @@ -34,6 +34,8 @@ class MCPSpecVersion(str, enum.Enum): nov_2024 = "2024-11-05" mar_2025 = "2025-03-26" jun_2025 = "2025-06-18" + nov_2025 = "2025-11-25" + jul_2026 = "2026-07-28" class MCPAuth(str, enum.Enum): @@ -59,7 +61,17 @@ DEFAULT_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:access_tok # MCP Literals MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio] -MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025] +MCPLegacyVersion = Literal["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25"] +MCP_LEGACY_VERSIONS: Final[tuple[MCPLegacyVersion, ...]] = ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25") +MCPUpstreamProtocol = MCPLegacyVersion | Literal["auto"] +MCPAdvertisedVersions = Annotated[tuple[MCPLegacyVersion, ...], Field(min_length=1)] +MCPSpecVersionType = Literal[ + MCPSpecVersion.nov_2024, + MCPSpecVersion.mar_2025, + MCPSpecVersion.jun_2025, + MCPSpecVersion.nov_2025, + MCPSpecVersion.jul_2026, +] MCPAuthType = ( Literal[ MCPAuth.none, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cb32299b143..c3b106c11d5 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,7 @@ from datetime import datetime -from typing import Any, Final, Literal +from typing import Annotated, Any, Final, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, TypeAdapter, field_validator, model_validator from typing_extensions import Self from litellm.types.mcp import ( @@ -10,11 +10,19 @@ from litellm.types.mcp import ( MCPAuthType, MCPTokenEndpointAuthMethod, MCPTransportType, + MCPUpstreamProtocol, normalize_upstream_header_name, ) + # MCPInfo now allows arbitrary additional fields for custom metadata -MCPInfo = dict[str, Any] +def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, object]: + if "protocol_version" in value: + TypeAdapter(MCPUpstreamProtocol).validate_python(value["protocol_version"]) + return value + + +MCPInfo = Annotated[dict[str, Any], AfterValidator(_validate_mcp_protocol_metadata)] class MCPOAuthMetadata(BaseModel): @@ -66,6 +74,7 @@ class MCPServer(BaseModel): server_name: str | None = None url: str | None = None transport: MCPTransportType + protocol_version: MCPUpstreamProtocol = "auto" spec_path: str | None = None auth_type: MCPAuthType | None = None authentication_token: str | None = None @@ -246,6 +255,14 @@ class MCPServer(BaseModel): """ return self.oauth2_flow == "client_credentials" + @model_validator(mode="after") + def resolve_protocol_version(self) -> Self: + if "protocol_version" not in self.model_fields_set and self.mcp_info is not None: + self.protocol_version = TypeAdapter(MCPUpstreamProtocol).validate_python( + self.mcp_info.get("protocol_version", "auto") + ) + return self + @model_validator(mode="after") def validate_identity_binding_mode(self) -> Self: binding: Final = self.oauth_identity_binding diff --git a/tests/integration/mcp/test_mcp_protocol_errors.py b/tests/integration/mcp/test_mcp_protocol_errors.py index bb06d8c6068..257e90e9d0e 100644 --- a/tests/integration/mcp/test_mcp_protocol_errors.py +++ b/tests/integration/mcp/test_mcp_protocol_errors.py @@ -84,3 +84,42 @@ def test_jsonrpc_error_and_malformed_tool_result_remain_errors(gateway: Gateway) control: Final = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5}) assert control.status_code == 200 and control.json()["isError"] is False, control.text assert control.json()["content"][0]["text"] == "8" + + +@pytest.mark.parametrize("ingress", ("http", "sse")) +def test_configured_revision_blocks_unadvertised_handshake_and_keeps_allowed_control( + gateway: Gateway, tmp_path, ingress: str +) -> None: + import asyncio + from pathlib import Path + + import yaml + from integration._support.mcp import mcp_peer + from integration._support.process import owned_proxy + from litellm.experimental_mcp_client.client import MCPClient + from litellm.types.mcp import MCPTransport + from mcp import MCPError + from mcp.types import CallToolRequestParams + + with mcp_peer() as upstream, gateway.scenario() as scenario: + alias: Final = "restricted" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, upstream, alias) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"]["mcp_advertised_versions"] = ["2024-11-05"] + config_path: Final = tmp_path / "restricted.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {"DISABLE_SCHEMA_UPDATE": "true"}, config=config_path) as restricted: + endpoint: Final = str(restricted.client.base_url).rstrip("/") + ("/mcp/sse" if ingress == "sse" else "/mcp") + headers: Final = {"Authorization": f"Bearer {key}", "x-mcp-servers": identity} + denied: Final = MCPClient(server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version="2025-11-25", extra_headers=headers) + allowed: Final = MCPClient(server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version="2024-11-05", extra_headers=headers) + + async def exercise() -> None: + with pytest.raises(MCPError, match="Unsupported MCP protocol version"): + await denied.list_tools(raise_on_error=True) + assert f"{alias}-add" in tuple(tool.name for tool in await allowed.list_tools(raise_on_error=True)) + result: Final = await allowed.call_tool(CallToolRequestParams(name=f"{alias}-add", arguments={"a": 2, "b": 5})) + assert result.is_error is False and result.content[0].text == "7" + + asyncio.run(exercise()) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 1862f11d07e..16004cdd501 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -155,3 +155,43 @@ def test_server_initiated_sampling_and_elicitation_surface_as_errors_not_success assert outcome.error is not None, outcome.raw assert outcome.text is None or not outcome.text.startswith(("sampled:", "elicited:")), outcome.raw assert len(tool_calls(peer.drain())) == 1 + + +@pytest.mark.parametrize("downstream", ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25")) +@pytest.mark.parametrize("upstream", ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25")) +@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio")) +@pytest.mark.parametrize("ingress", ("http", "sse")) +def test_pinned_revision_pairs_list_and_call_through_gateway( + gateway: Gateway, downstream: str, upstream: str, peer_kind: PeerKind, ingress: str +) -> None: + import asyncio + + from mcp.types import CallToolRequestParams + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.types.mcp import MCPTransport + + with peer_of(peer_kind) as peer, gateway.scenario() as scenario: + alias: Final = "versions" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias, mcp_info={"protocol_version": upstream}) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + endpoint: Final = str(gateway.client.base_url).rstrip("/") + ("/mcp/sse" if ingress == "sse" else "/mcp") + client: Final = MCPClient( + server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version=downstream, + extra_headers={"Authorization": f"Bearer {key}", "x-mcp-servers": identity}, timeout=15, + ) + + async def exercise() -> None: + tools: Final = await client.list_tools(raise_on_error=True) + assert f"{alias}-add" in tuple(tool.name for tool in tools) + result: Final = await client.call_tool(CallToolRequestParams(name=f"{alias}-add", arguments={"a": 3, "b": 4})) + assert result.is_error is False + assert result.content[0].text == "7" + + peer.drain() + asyncio.run(exercise()) + observed: Final = peer.drain() + negotiations: Final = tuple(item["body"] for item in observed if item["body"].get("method") == "initialize") + assert negotiations, "The operation must reach the upstream negotiation" + assert all(request["params"]["protocolVersion"] == upstream for request in negotiations), negotiations + assert len(tool_calls(observed)) == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py new file mode 100644 index 00000000000..f104e655704 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py @@ -0,0 +1,106 @@ +from typing import Final + +import pytest +from mcp import Client +from mcp.server import Server +from mcp.shared.exceptions import MCPError +from mcp.types import PromptsCapability, ResourcesCapability, ServerCapabilities, ToolsCapability +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS + +from litellm.proxy._experimental.mcp_server.capabilities import ( + GATEWAY_OPERATIONS, + REVISION_SUPPORT, + TRANSLATION_PAIRS, + GatewayVersionPolicy, + build_discovery, +) +from litellm.types.mcp import MCPTransport + + +@pytest.mark.parametrize("revision", HANDSHAKE_PROTOCOL_VERSIONS) +@pytest.mark.parametrize("transport", tuple(MCPTransport)) +def test_discovery_only_exposes_authorized_completed_support(revision, transport): + result = build_discovery( + configured=(revision, "2026-07-28", "unknown"), + revision=revision, + transport=transport, + authorized_operations=frozenset({"tools/list", "tools/call"}), + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=ServerCapabilities( + tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability(), + extensions={"io.modelcontextprotocol/ui": {}}, + ), + client_extensions=frozenset({"io.modelcontextprotocol/ui"}), + upstream_extensions=frozenset({"io.modelcontextprotocol/ui"}), + ) + assert result.supported_versions == [revision] + assert result.capabilities.tools is not None + assert result.capabilities.prompts is None + assert result.capabilities.resources is None + assert result.capabilities.extensions is None + assert result.capabilities.tasks is None + assert result.cache_scope == "private" + assert result.ttl_ms == 0 + + +@pytest.mark.parametrize("upstream", [frozenset(), frozenset({"unknown"}), frozenset({"2026-07-28"})]) +def test_unproven_translation_never_advertises_operations(upstream): + result = build_discovery( + configured=HANDSHAKE_PROTOCOL_VERSIONS, + revision="2025-11-25", + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=upstream, + capabilities=ServerCapabilities(tools=ToolsCapability()), + ) + assert result.capabilities.tools is None + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "2024-11-05"]) +def test_unadvertised_revision_never_gains_capabilities(revision): + result = build_discovery( + configured=("2025-11-25",), revision=revision, transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=ServerCapabilities(tools=ToolsCapability()), + ) + assert result.capabilities.tools is None + + +def test_discovery_results_do_not_share_mutable_capabilities(): + capabilities = ServerCapabilities(tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability()) + args = dict( + configured=HANDSHAKE_PROTOCOL_VERSIONS, revision="2025-11-25", transport=MCPTransport.http, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=capabilities, + ) + allowed = build_discovery(**args, authorized_operations=GATEWAY_OPERATIONS) + denied = build_discovery(**args, authorized_operations=frozenset()) + assert allowed.capabilities.prompts is not None + assert allowed.capabilities.resources is not None + assert denied.capabilities.model_dump(exclude_none=True) == {} + assert allowed.capabilities.tools is not None + allowed.capabilities.tools.list_changed = True + assert capabilities.tools.list_changed is not True + + +def test_modern_candidates_do_not_enable_public_serving(): + modern = REVISION_SUPPORT["2026-07-28"] + assert modern.completed is False + assert "input_required" in modern.results + assert MCPTransport.sse not in modern.transports + assert not any("2026-07-28" in pair for pair in TRANSLATION_PAIRS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("versions,accepted", [(("2025-11-25",), True), (("2025-06-18",), False)]) +async def test_version_policy_gates_the_actual_sdk_handshake(versions, accepted): + server: Final = Server("test-gateway", version="1") + server.middleware.append(GatewayVersionPolicy(lambda: versions)) + if accepted: + async with Client(server, mode="legacy") as client: + assert client.protocol_version == "2025-11-25" + result = await client.session.send_ping() + assert result is not None + else: + with pytest.RaisesGroup(pytest.RaisesExc(MCPError, match="Unsupported MCP protocol version"), flatten_subgroups=True): + async with Client(server, mode="legacy"): + pytest.fail("The excluded revision must not initialize") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 97b242831a2..ab00ec4da1e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -10444,11 +10444,12 @@ async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_reque ) @pytest.mark.parametrize("handler", ("handle_streamable_http_mcp", "handle_sse_mcp")) async def test_streamable_http_rejects_modern_protocol_version( - header_value: str, expected_rejected: bool, handler: str + header_value: str, expected_rejected: bool, handler: str, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy._experimental.mcp_server import server as mcp_module from litellm.proxy._experimental.mcp_server.server import unsupported_protocol_version + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) scope: Scope = { "type": "http", "method": "POST", @@ -10659,3 +10660,32 @@ async def test_legacy_sse_mount_emits_message_endpoint( await incoming.put({"type": "http.disconnect"}) await asyncio.wait_for(task, 2) assert await post(initialization) == 404 + + +@pytest.mark.parametrize("revision,rejected", [("2024-11-05", False), ("2025-11-25", True), ("2026-07-28", True)]) +def test_protocol_header_respects_configured_advertisement(revision, rejected): + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.server import unsupported_protocol_version + + with patch.dict(proxy_server.general_settings, {"mcp_advertised_versions": ["2024-11-05"]}): + result = unsupported_protocol_version({"headers": [(b"mcp-protocol-version", revision.encode())]}) + assert result == (revision if rejected else None) + + +@pytest.mark.asyncio +async def test_discovery_adapter_preserves_authenticated_context(_mcp_request_ctx): + from mcp.types import DiscoverResult, RequestParams, ServerCapabilities + from litellm.proxy._experimental.mcp_server import server + + expected = DiscoverResult(supported_versions=["2025-11-25"], capabilities=ServerCapabilities()) + dispatched = AsyncMock(return_value=expected) + auth = UserAPIKeyAuth(user_id="discover-caller") + with ( + patch.object(server, "get_or_extract_auth_context", AsyncMock(return_value=(auth, None, ["allowed"], None, None, None, None))), + patch.object(server.operations.GatewayOperations, "execute", dispatched), + ): + result = await server.discover(_mcp_request_ctx(), RequestParams()) + assert result is expected + context = dispatched.await_args.args[1] + assert context.user_api_key_auth.user_id == "discover-caller" + assert context.mcp_servers == ("allowed",) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 64d94065674..16bffa1a356 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -63,7 +63,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import MCPAuth, MCPAuthType +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPUpstreamProtocol from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer from litellm.caching.caching import DualCache from litellm.caching.llm_caching_handler import LLMClientCache @@ -14637,3 +14637,27 @@ class TestSharedIdentifierPrefixWarning: assert "srv-b" in shared_warnings[0] assert "srv-c" not in shared_warnings[0] assert "'shared'" in shared_warnings[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ["auto", "2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25"]) +async def test_configured_protocol_reaches_the_upstream_client(config_only_mcp_manager_factory, revision): + manager = config_only_mcp_manager_factory() + await manager.load_servers_from_config({"versions": {"url": "http://127.0.0.1:9/mcp", "transport": "http", "protocol_version": revision}}) + server = next(iter(manager.config_mcp_servers.values())) + client = await manager._create_mcp_client(server) + assert server.protocol_version == revision + assert client.protocol_version == revision + + +@pytest.mark.parametrize("revision", ("auto", "2024-11-05", "2025-06-18")) +@pytest.mark.parametrize("explicit", (None, "auto", "2025-11-25")) +def test_runtime_protocol_metadata_preserves_explicit_precedence( + revision: MCPUpstreamProtocol, explicit: MCPUpstreamProtocol | None +) -> None: + server: Final = MCPServer.model_validate({ + "server_id": "preview", "name": "preview", "transport": "http", + "mcp_info": {"protocol_version": revision}, + **({"protocol_version": explicit} if explicit is not None else {}), + }) + assert server.protocol_version == (explicit if explicit is not None else revision) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py index 81f81045740..bb900de4f98 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -542,3 +542,126 @@ async def test_local_tool_json_array_is_converted_once_for_the_caller_revision(c assert [block.text for block in result.content] == [body] assert result.structured_content == (["a", "b"] if compat == "modern" else None) + + +@pytest.mark.asyncio +async def test_discovery_preserves_caller_scope_and_proxy_restrictions(): + from mcp.types import DiscoverRequest, ListToolsResult, Tool + + listed = AsyncMock(return_value=ListToolsResult(tools=[Tool(name="allowed", input_schema={"type": "object"})])) + context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["only-this"], mcp_proxy_mode=True, protocol_version="2025-06-18") + with patch("litellm.proxy._experimental.mcp_server.operations._execute_handle_list_tools", listed): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is not None + assert result.capabilities.resources is None + assert result.capabilities.prompts is None + assert listed.await_args.args[0] is context + assert listed.await_args.args[0].user_api_key_auth.user_id == "scoped" + assert listed.await_args.args[0].mcp_servers == ("only-this",) + + +@pytest.mark.asyncio +async def test_discovery_denial_cannot_advertise_tools(): + from mcp.types import DiscoverRequest + from fastapi import HTTPException + + denied = AsyncMock(side_effect=HTTPException(status_code=403, detail="Forbidden")) + with patch("litellm.proxy._experimental.mcp_server.operations._execute_handle_list_tools", denied): + with pytest.raises(HTTPException) as error: + await GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="denied"))) + assert error.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("available", ["none", "resources", "templates", "prompts"]) +async def test_discovery_lists_each_capability_with_the_same_caller(available): + from mcp.types import ( + DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, + ListResourceTemplatesResult, Prompt, Resource, ResourceTemplate, + ) + from litellm.proxy._experimental.mcp_server import operations + + context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]) + tools = AsyncMock(return_value=ListToolsResult(tools=[])) + prompts = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="allowed")] if available == "prompts" else [])) + resources = AsyncMock(return_value=ListResourcesResult(resources=[Resource(name="allowed", uri="test://allowed")] if available == "resources" else [])) + templates = AsyncMock(return_value=ListResourceTemplatesResult(resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] if available == "templates" else [])) + with ( + patch.object(operations, "_execute_handle_list_tools", tools), + patch.object(operations, "_execute_list_prompts", prompts), + patch.object(operations, "_execute_list_resources", resources), + patch.object(operations, "_execute_list_resource_templates", templates), + ): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is None + assert (result.capabilities.prompts is not None) == (available == "prompts") + assert (result.capabilities.resources is not None) == (available in {"resources", "templates"}) + for listing in (tools, prompts, resources, templates): + assert listing.await_args.args[0] is context + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancel"]) +async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(outcome): + from mcp.types import DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult + from litellm.proxy._experimental.mcp_server import operations + + ready = [asyncio.Event() for _ in range(4)] + closed = [asyncio.Event() for _ in range(4)] + release = asyncio.Event() + responses = (ListToolsResult(tools=[]), ListPromptsResult(prompts=[]), ListResourcesResult(resources=[]), ListResourceTemplatesResult(resource_templates=[])) + + def listing(index): + async def run(*args, **kwargs): + ready[index].set() + try: + await release.wait() + if index == 0 and outcome == "failure": + raise ValueError("discovery failed") + if outcome != "success": + await asyncio.Event().wait() + return responses[index] + finally: + closed[index].set() + return run + + with ( + patch.object(operations, "_execute_handle_list_tools", side_effect=listing(0)) as tools, + patch.object(operations, "_execute_list_prompts", side_effect=listing(1)), + patch.object(operations, "_execute_list_resources", side_effect=listing(2)), + patch.object(operations, "_execute_list_resource_templates", side_effect=listing(3)), + ): + task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped")))) + try: + await asyncio.wait_for(asyncio.gather(*(event.wait() for event in ready)), 1) + if outcome == "cancel": + task.cancel() + else: + release.set() + if outcome == "success": + result = await asyncio.wait_for(task, 1) + assert result.capabilities.model_dump(exclude_none=True) == {} + else: + with pytest.raises(asyncio.CancelledError if outcome == "cancel" else ValueError): + await asyncio.wait_for(task, 1) + assert all(event.is_set() for event in closed) + assert tools.call_args.kwargs["log_list_tools_to_spendlogs"] is False + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("log_enabled", [False, True]) +async def test_tools_listing_preserves_explicit_spend_log_policy(log_enabled): + from mcp.types import PaginatedRequestParams + from litellm.proxy._experimental.mcp_server import operations + + listing = AsyncMock(return_value=operations.AggregateToolListing(tools=[], outcomes={})) + with patch.object(operations, "_list_mcp_tools", listing): + result = await operations._execute_handle_list_tools( + prepare_context(UserAPIKeyAuth(user_id="caller")), PaginatedRequestParams(), + log_list_tools_to_spendlogs=log_enabled, + ) + assert result.tools == [] + assert listing.await_args.kwargs["log_list_tools_to_spendlogs"] is log_enabled diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 4da120cb26f..e82ab28bb4c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -28,7 +28,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp import MCPAuth, MCPTransport, MCPUpstreamProtocol from litellm.types.mcp_server.mcp_server_manager import MCPServer _OK_TOOL_RESULT: Final = CallToolResult(content=[TextContent(type="text", text='{"result": "ok"}')], is_error=False) @@ -1476,7 +1476,7 @@ class TestListToolsRestAPI: monkeypatch, ): """The REST tools/list path should include tools beyond the upstream first page.""" - from mcp.types import ListToolsResult, PaginatedRequestParams + from mcp.types import Implementation, InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities from mcp.types import Tool as MCPTool import litellm.experimental_mcp_client.client as mcp_client_module @@ -1512,7 +1512,11 @@ class TestListToolsRestAPI: mock_session_ctx = AsyncMock() mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock(return_value=None) + mock_session_instance.initialize = AsyncMock(return_value=InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="stub", version="1"), + )) mock_session_instance.list_tools.side_effect = [ ListToolsResult( tools=[ @@ -4628,3 +4632,74 @@ class TestClientAllowlistOnRestRoutes: assert denied.value.detail["error"] == "Forbidden" assert "'claude-code'" in denied.value.detail["details"] acting.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ("auto", "2024-11-05", "2025-06-18")) +async def test_preview_client_honors_protocol_metadata(revision: MCPUpstreamProtocol) -> None: + from litellm.experimental_mcp_client.client import MCPClient + + payload: Final = NewMCPServerRequest( + server_name="preview", url="http://127.0.0.1:9/mcp", transport="http", + auth_type=MCPAuth.none, mcp_info={"protocol_version": revision}, + ) + + async def inspect_client(client: MCPClient) -> dict[str, str]: + return {"protocol_version": client.protocol_version} + + result: Final = await rest_endpoints._execute_with_mcp_client(payload, inspect_client) + assert result == {"protocol_version": revision} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", (MCPAuth.none, MCPAuth.bearer_token, MCPAuth.oauth2)) +@pytest.mark.parametrize( + ("metadata", "expected"), + ( + (None, "2025-11-25"), + ({}, "2025-11-25"), + ({"description": "edited"}, "2025-11-25"), + ({"protocol_version": "auto"}, "auto"), + ({"protocol_version": "2024-11-05"}, "2024-11-05"), + ({"protocol_version": "2025-06-18"}, "2025-06-18"), + ), +) +async def test_saved_preview_protocol_omission_and_explicit_edits( + monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuth, + metadata: dict[str, str] | None, expected: MCPUpstreamProtocol, +) -> None: + from starlette.datastructures import Headers + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview", name="preview", url="https://example.com/mcp", + transport="http", auth_type=auth_type, protocol_version="2025-11-25", + authentication_token="stored-token", + authorization_url="https://example.com/authorize", token_url="https://example.com/token", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, server_name=saved.name, url=saved.url, transport="http", + auth_type=auth_type, mcp_info=metadata, + authorization_url=saved.authorization_url, token_url=saved.token_url, + ) + staged: Final = rest_endpoints._stage_server_test( + payload, Headers({"x-litellm-api-key": "sk-admin", "authorization": "Bearer preview-token"}) + ) + + async def inspect_client(client: MCPClient) -> dict[str, str]: + return {"protocol_version": client.protocol_version} + + result: Final = await rest_endpoints._execute_with_mcp_client( + staged.request, inspect_client, + mcp_auth_header=staged.mcp_auth_header, oauth2_headers=staged.oauth2_headers, + ) + assert result == {"protocol_version": expected} + assert saved.protocol_version == "2025-11-25" + assert payload.mcp_info == metadata diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 7e198bc9131..b2ef327f50e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -4903,3 +4903,19 @@ async def test_model_refresh_updates_availability_catalog_and_retains_it_on_db_f assert await pc._get_models_from_db(client) == [] assert pc.auto_router_db_catalog == () assert find_many.await_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("versions", [None, ["2024-11-05"], [], ["2026-07-28"], ["unknown"]]) +async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, monkeypatch, versions): + config = tmp_path / "mcp-versions.yaml" + config.write_text(json.dumps({"model_list": [], "general_settings": {"mcp_advertised_versions": versions}})) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + if versions is None or versions == ["2024-11-05"]: + _, _, settings = await ProxyConfig().load_config(router=None, config_file_path=str(config)) + assert settings["mcp_advertised_versions"] == versions + return + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=None, config_file_path=str(config)) diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index f5abe0561db..b43a75d3323 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -377,3 +377,21 @@ def test_change_password_request_passwords_hidden_from_repr(): for rendered in (repr(request), str(request)): assert "hunter2hunter2" not in rendered assert "NewP@ssw0rd-2026" not in rendered +@pytest.mark.parametrize("versions", [[], ["2099-01-01"], ["2026-07-28"]]) +def test_mcp_advertised_versions_reject_unavailable_revisions(versions): + from pydantic import ValidationError + + from litellm.proxy._types import ConfigGeneralSettings + + with pytest.raises(ValidationError): + ConfigGeneralSettings(mcp_advertised_versions=versions) + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", None]) +def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision): + from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest + + payload = {"server_id": "test", "transport": "http", "url": "https://example.com/mcp", "mcp_info": {"protocol_version": revision}} + for model in (NewMCPServerRequest, UpdateMCPServerRequest): + with pytest.raises(ValidationError): + model.model_validate(payload) diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 368e34c455d..1a56227b008 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -58,6 +58,15 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer _JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) +def _initialized(instructions: str | None = None) -> InitializeResult: + return InitializeResult( + protocol_version=LATEST_HANDSHAKE_VERSION, + capabilities=ServerCapabilities(), + server_info=Implementation(name="test", version="1"), + instructions=instructions, + ) + + class _MockTransportClient(MCPClient): """An MCPClient whose streamable-HTTP transport runs on an httpx2 MockTransport.""" @@ -125,7 +134,7 @@ class TestMCPClient: mock_stdio_client.return_value = mock_stdio_ctx mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -168,7 +177,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -214,7 +223,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -266,7 +275,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -413,8 +422,7 @@ class TestMCPClientInstructionsCapture: ) mock_session = AsyncMock() - init_result = MagicMock() - init_result.instructions = " upstream says hello " + init_result = _initialized(" upstream says hello ") mock_session.initialize = AsyncMock(return_value=init_result) session_ctx = MagicMock() @@ -442,8 +450,7 @@ class TestMCPClientInstructionsCapture: ) mock_session = AsyncMock() - init_result = MagicMock() - init_result.instructions = None + init_result = _initialized() mock_session.initialize = AsyncMock(return_value=init_result) session_ctx = MagicMock() @@ -600,8 +607,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_cleanup_error_after_success_is_swallowed(self, mock_session_cls): client = MCPClient(server_url="http://example.com/mcp", transport_type="http") - init_result = MagicMock() - init_result.instructions = None + init_result = _initialized() self._make_session(mock_session_cls, AsyncMock(return_value=init_result)) transport_ctx = self._make_transport(_FakeExceptionGroup("late", [httpx2.ConnectError("late cleanup error")])) @@ -634,7 +640,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.parametrize("original_error", (False, True)) @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_session_exit_cancellation_preserves_original_failure(self, session_class, original_error): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) cancelled: Final = asyncio.CancelledError("cancelled while closing session") session_class.return_value.__aexit__ = AsyncMock(side_effect=cancelled) original: Final = RuntimeError("operation failed") @@ -656,7 +662,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.parametrize("phase", ("session", "transport")) @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_cleanup_preserves_process_exit(self, session_class, phase, signal_type): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) signal: Final = signal_type("process stopping") if phase == "session": session_class.return_value.__aexit__ = AsyncMock(side_effect=signal) @@ -670,7 +676,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.asyncio @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_session_and_termination_share_one_cleanup_deadline(self, session_class): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) deleting: Final = asyncio.Event() async def close_session(*args): @@ -1883,16 +1889,17 @@ async def test_sse_read_failure_is_preserved() -> None: @pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ["auto", "2025-06-18"]) @pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) @pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) -async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str, protocol_version: str) -> None: from mcp import ClientSession from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message logging_callback: Final = AsyncMock() read_timeout: Final = 0.2 if mode == "silent" else 30 client: Final = MCPClient( - server_url="https://example.com/sse", transport_type=transport, timeout=read_timeout, logging_callback=logging_callback + server_url="https://example.com/sse", transport_type=transport, timeout=read_timeout, logging_callback=logging_callback, protocol_version=protocol_version ) async def operation(session: ClientSession) -> CallToolResult: @@ -2754,12 +2761,13 @@ async def test_http_close_cancellation_cannot_turn_into_success(original_error: @pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ("auto", "2025-06-18")) @pytest.mark.parametrize("cancel_mode", ("scope", "task", "wait_for", "read_timeout")) @pytest.mark.parametrize("concurrency", (1, 5)) @pytest.mark.parametrize("termination", ("ok", "hang", "hang_body")) @pytest.mark.parametrize("raise_on_error", (False, True)) async def test_cancellation_delivers_termination_over_tcp( - cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool + cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool, protocol_version: str ) -> None: started: Final = asyncio.Event() scope_ready: Final[asyncio.Future[anyio.CancelScope]] = asyncio.get_running_loop().create_future() @@ -2813,6 +2821,8 @@ async def test_cancellation_delivers_termination_over_tcp( await stop.wait() return if payload["method"] == "initialize": + if cancel_mode != "read_timeout": + await asyncio.sleep(0.75) response: Final = json.dumps( { "jsonrpc": "2.0", @@ -2839,7 +2849,7 @@ async def test_cancellation_delivers_termination_over_tcp( listener: Final = await asyncio.start_server(handle_connection, "127.0.0.1", 0) port: Final = listener.sockets[0].getsockname()[1] client: Final = MCPClient( - server_url=f"http://127.0.0.1:{port}/mcp", timeout=2 if cancel_mode == "read_timeout" else 0.5 if termination != "ok" else 30 + server_url=f"http://127.0.0.1:{port}/mcp", protocol_version=protocol_version, timeout=2 if cancel_mode == "read_timeout" else 30 ) async def calls(): @@ -2866,7 +2876,7 @@ async def test_cancellation_delivers_termination_over_tcp( try: task: Final = asyncio.create_task(invoke()) - await asyncio.wait_for(started.wait(), 3) + await asyncio.wait_for(started.wait(), 30) if cancel_mode == "scope": (await scope_ready).deadline = anyio.current_time() + 0.2 if cancel_mode == "task": @@ -2901,3 +2911,52 @@ async def test_cancellation_delivers_termination_over_tcp( closed: Final = await asyncio.wait_for(asyncio.gather(*connections, return_exceptions=True), 2) assert all(result is None or isinstance(result, asyncio.CancelledError) for result in closed), closed await asyncio.wait_for(listener.wait_closed(), 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "auto"]) +@pytest.mark.parametrize("accepted", [True, False]) +@pytest.mark.parametrize("callbacks", [False, True]) +async def test_configured_upstream_revision_is_offered_and_checked(revision, accepted, callbacks): + from mcp.types import JSONRPCRequest + from mcp_types.version import LATEST_HANDSHAKE_VERSION + + offered = LATEST_HANDSHAKE_VERSION if revision == "auto" else revision + + def respond(request): + if request.method == "DELETE": + return httpx2.Response(200) + payload = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + assert payload.params["protocolVersion"] == offered + assert ("sampling" in payload.params["capabilities"]) == callbacks + assert ("elicitation" in payload.params["capabilities"]) == callbacks + return httpx2.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, + "result": {"protocolVersion": offered if accepted else "unsupported", + "capabilities": {"tools": {}}, "serverInfo": {"name": "upstream", "version": "1"}}, + }) + assert accepted, "No operation may execute after failed version negotiation" + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"tools": [{"name": "echo", "inputSchema": {"type": "object"}}]}}) + + client = _MockTransportClient( + respond, server_url="https://example.com/mcp", protocol_version=revision, + sampling_callback=AsyncMock() if callbacks else None, + elicitation_callback=AsyncMock() if callbacks else None, + ) + if accepted: + result = await client.list_tools(raise_on_error=True) + assert [tool.name for tool in result] == ["echo"] + else: + with pytest.raises((MCPError, RuntimeError), match="protocol version"): + await client.list_tools(raise_on_error=True) + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "", None]) +def test_upstream_protocol_configuration_rejects_unavailable_modes(revision): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + MCPClient(protocol_version=revision) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py index 6438525706a..1a592aa1c9a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py @@ -12,7 +12,7 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport from mcp.types import CallToolResult as MCPCallToolResult -from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Implementation, InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities from mcp.types import Tool as MCPTool @@ -128,6 +128,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) client = MCPClient( @@ -163,6 +168,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_tools = [ @@ -204,6 +214,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) first_page_tools = [ @@ -245,6 +260,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_session_instance.list_tools.side_effect = [ @@ -277,6 +297,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_result = MCPCallToolResult(content=[]) @@ -289,7 +314,7 @@ class TestMCPClientUnitTests: assert result == mock_result mock_session_instance.initialize.assert_called_once() mock_session_instance.call_tool.assert_called_once_with( - name="test_tool", arguments={"arg1": "value1"}, progress_callback=ANY + name="test_tool", arguments={"arg1": "value1"}, progress_callback=ANY, allow_input_required=False ) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index dfd250338ff..26674373b4e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -489,7 +489,7 @@ async def test_sse_mcp_handler_mock(): ) with ( - patch("litellm.proxy._experimental.mcp_server.server.server.run", run), + patch("litellm.proxy._experimental.mcp_server.server.serve_loop", run), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -511,7 +511,7 @@ async def test_sse_mcp_handler_mock(): # Call the handler await handle_sse_mcp(mock_scope, mock_receive, mock_send) - assert run.await_args.args[:2] == (read_stream, write_stream) + assert run.await_args.args[1:3] == (read_stream, write_stream) assert mock_sse.connect_sse.call_args.args[0]["path"] == "/mcp/sse" diff --git a/tests/unit/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py index e464402c9d8..b528a75d0df 100644 --- a/tests/unit/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -38,6 +38,7 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl "PATH": f"{shim_dir}{os.pathsep}{os.environ['PATH']}", "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, + "UNIT_FLAG": "", "WORKERS": workers, "UNIT_FLAG": "", }, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bba67bcf6c2..6126733095b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28350,6 +28350,11 @@ export interface components { * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. */ maximum_spend_logs_retention_period?: string | null; + /** + * Mcp Advertised Versions + * @description MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. Modern protocol serving and Apps/Tasks remain disabled. + */ + mcp_advertised_versions?: ("2024-11-05" | "2025-03-26" | "2025-06-18" | "2025-11-25")[] | null; /** * Mcp Allowed Clients * @description MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. From 601c75a475feb5120f6b046bee37ef2e2acd3de0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 13:45:43 -0700 Subject: [PATCH 16/58] fix(proxy): record streamed /v1/responses container ownership before the response.completed frame (#43140) * test(e2e): cover Azure code_interpreter container files by native id with a service-account key * test(e2e): require the code_interpreter tool, skip at collection, and scope the container call timeout * fix(e2e): fail the containers suite when the Azure credentials are missing instead of skipping * fix(proxy): record streamed /v1/responses container ownership before the response.completed frame --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 87 +++++++++---------- .../coverage_registry/llm_conversational.yaml | 1 + .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 3 +- .../llm_translation/test_containers_e2e.py | 43 +++++++-- .../proxy/test_common_request_processing.py | 64 ++++++++++++++ 5 files changed, 144 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7e1989b0aab..4f9b6b3a96f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2781,15 +2781,6 @@ class ProxyBaseLLMRequestProcessing: request=request, ) if route_type == "aresponses": - # Streaming /v1/responses returns here without - # reaching the non-streaming ownership tail below. - # Wrap the SSE generator so container ownership is - # written once the upstream iterator finishes - # assembling ``completed_response`` — otherwise - # code-interpreter containers created during the - # stream stay unregistered and follow-up file API - # calls 403. Covers the background-polling path - # too, which loops ``body_iterator`` end-to-end. selected_data_generator = ( ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( original_stream_response=response, @@ -3011,50 +3002,50 @@ class ProxyBaseLLMRequestProcessing: wrapped_generator: Any, user_api_key_dict: UserAPIKeyAuth, ): - """Forward SSE chunks, then record container ownership at stream end. + """Forward SSE chunks and record container ownership before the terminal chunk goes out. Streaming ``/v1/responses`` short-circuits out of ``base_process_llm_request`` before the non-streaming ownership - tail runs, so without this wrap the - ``LiteLLM_ManagedObjectTable`` row for any container created - during the stream is never written and follow-up file API calls - return 403. + tail runs. The OpenAI SDK closes the connection at ``data: [DONE]`` + and starlette cancels the body task on disconnect, so a write that + waits for the generator to finish never lands. The iterator sets + ``completed_response`` before it hands over its terminal chunk, so + the ``LiteLLM_ManagedObjectTable`` row is written the moment it + appears, ahead of the chunk carrying ``response.completed``. """ - try: - async for chunk in wrapped_generator: + async for chunk in wrapped_generator: + completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response + ) + if completed_obj is None: yield chunk - finally: - try: - completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( - original_stream_response - ) - if completed_obj is not None: - await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( - response=completed_obj, - user_api_key_dict=user_api_key_dict, - ) - else: - # Silent skip caused #30210: the proxy's Router wrapper - # of the responses streaming iterator wasn't propagating - # ``completed_response``, so this hook recorded nothing - # and follow-up /v1/containers//files calls 403'd - # for non-admin keys with no proxy-side hint. Log a - # warning so future regressions of the same shape - # surface in operator logs. - verbose_proxy_logger.warning( - "Container ownership recording skipped on streaming " - "/v1/responses: no completed_response on stream " - "iterator %s. If this stream created any tool " - "container (e.g. code_interpreter), follow-up " - "/v1/containers//files calls will 403 for " - "non-admin keys.", - type(original_stream_response).__name__, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Container ownership recording failed after streaming responses call: %s", - e, - ) + continue + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=completed_obj, + user_api_key_dict=user_api_key_dict, + ) + yield chunk + async for remaining_chunk in wrapped_generator: + yield remaining_chunk + return + late_completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response + ) + if late_completed_obj is not None: + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=late_completed_obj, + user_api_key_dict=user_api_key_dict, + ) + return + verbose_proxy_logger.warning( + "Container ownership recording skipped on streaming " + "/v1/responses: no completed_response on stream " + "iterator %s. If this stream created any tool " + "container (e.g. code_interpreter), follow-up " + "/v1/containers//files calls will 403 for " + "non-admin keys.", + type(original_stream_response).__name__, + ) async def base_passthrough_process_llm_request( self, diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index bc19668e2c9..20c87dbbd74 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -85,6 +85,7 @@ - {id: llm.responses.azure_openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Azure OpenAI (smoke)"} - {id: llm.responses.azure_openai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Responses tool calls w/ Azure OpenAI"} - {id: llm.responses.azure_openai.code_interpreter.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: nonstream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "An implicit code_interpreter container on an Azure deployment that carries its own api_base must serve GET /v1/containers/{id}/files/{fid}/content by its native cntr_ id to a team service-account key, the customer's shape (#27921, #28990)", fail_before_fix: proven} +- {id: llm.responses.azure_openai.code_interpreter.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: stream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "A container created by a streamed /v1/responses code_interpreter call must serve /v1/containers/{id}/files to the same service-account key right after the OpenAI SDK closes at [DONE]; the ownership row used to be written after the stream and the disconnect cancelled it (LIT-8612)", fail_before_fix: proven} - {id: llm.chat_completions.together_ai.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning surfaces as reasoning_content (LIT-5960)"} - {id: llm.chat_completions.together_ai.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning deltas stream as reasoning_content"} - {id: llm.chat_completions.together_ai.thinking.nonstream.template_kwargs_forwarded, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [template_kwargs_forwarded], source: "llm_translation/test_together_ai_e2e.py", rationale: "chat_template_kwargs reaches Together and turns thinking off"} diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index a18c81fa01d..92330db0530 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,7 +48,7 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | -| Responses (Azure code_interpreter container files) | live | gap | live | gap | partial | +| Responses (Azure code_interpreter container files) | live | live | live | gap | partial | | Image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -63,6 +63,7 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | | `test_service_account_key_reads_container_file_by_native_id` | azure responses code_interpreter, non-stream, native container id, service-account key | +| `test_service_account_key_reads_container_file_created_by_a_streamed_response` | azure responses code_interpreter, stream, native container id, service-account key, upload right after `[DONE]` | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/test_containers_e2e.py b/tests/e2e/llm_translation/test_containers_e2e.py index 887aecb8df1..3048a830810 100644 --- a/tests/e2e/llm_translation/test_containers_e2e.py +++ b/tests/e2e/llm_translation/test_containers_e2e.py @@ -31,10 +31,11 @@ A proxy whose env carries ``AZURE_API_BASE`` for the same resource masks the second regression, since the global-credential fallback then reaches the container anyway. -The streaming variant is not here: a streamed ``/v1/responses`` writes the -container ownership row only after the ``[DONE]`` frame, and the OpenAI SDK -closes the connection at ``[DONE]``, so the write is cancelled and every -follow-up container call 403s (LIT-8612). That cell comes with its fix. +The streaming cell repeats the flow with ``stream=True`` and uploads right +after the last event. The OpenAI SDK closes the connection at ``[DONE]``, so an +ownership row written after the stream is cancelled with the body task and every +follow-up container call 403s (LIT-8612); the row has to land before the +``response.completed`` frame goes out. """ from __future__ import annotations @@ -52,7 +53,7 @@ from lifecycle import ResourceManager from management.management_client import ManagementClient, build_client from models import KeyGenerateBody, KeyGenerateResponse, LiteLLMParamsBody, TeamNewBody, UserNewBody from openai import OpenAI -from openai.types.responses import Response, ResponseCodeInterpreterToolCall +from openai.types.responses import Response, ResponseCodeInterpreterToolCall, ResponseCompletedEvent from openai.types.responses.tool_param import CodeInterpreter from proxy_client import ProxyClient from sdk_clients import NO_PROXY_CACHE, SdkClients @@ -120,6 +121,24 @@ def _response_with_code_interpreter(client: OpenAI, model: str) -> Response: ) +def _streamed_response_with_code_interpreter(client: OpenAI, model: str) -> Response: + events: Final = tuple( + client.with_options(timeout=CODE_INTERPRETER_TIMEOUT).responses.create( + model=model, + input=PROMPT, + tools=[CODE_INTERPRETER], + tool_choice="required", + stream=True, + extra_body=NO_PROXY_CACHE, + ) + ) + assert events, "responses stream returned no events" + assert isinstance(events[-1], ResponseCompletedEvent), ( + f"responses stream did not terminate with response.completed: {events[-1].type}" + ) + return events[-1].response + + def _container_id(response: Response) -> str: calls: Final = tuple(item for item in response.output if isinstance(item, ResponseCodeInterpreterToolCall)) assert calls, f"no code_interpreter_call in the responses output: {response.output!r}" @@ -165,3 +184,17 @@ class TestAzureContainerFiles: f"container id is not the provider's own id: {native_id}" ) _assert_file_round_trip(client, native_id, marker) + + @pytest.mark.covers("llm.responses.azure_openai.code_interpreter.stream.works") + def test_service_account_key_reads_container_file_created_by_a_streamed_response( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + marker: Final = unique_marker() + model: Final = _register_two_azure_deployments(proxy, resources, marker) + key: Final = _service_account_key(proxy, resources, build_client(proxy), marker, model) + client: Final = sdk.openai(key) + native_id: Final = _native_container_id( + _container_id(_streamed_response_with_code_interpreter(client, model)) + ) + resources.defer(lambda: client.containers.delete(native_id, extra_query=AZURE_PROVIDER_QUERY)) + _assert_file_round_trip(client, native_id, marker) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7b74e69685c..bdf003085ef 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -9899,3 +9899,67 @@ class TestErrorLogCarriesCallId: record: Final = caplog.records[-1] assert record.litellm_call_id == call_id assert call_id in record.getMessage() + + +class TestStreamingContainerOwnershipRecordedBeforeDone: + """Regression for LIT-8612: the OpenAI SDK closes the connection at + ``data: [DONE]`` and starlette cancels the body task, so an ownership row + written after the SSE generator is exhausted never lands. The row must be + written before the chunk carrying ``response.completed`` is handed to the + client.""" + + CHUNKS: Final = ( + 'data: {"type":"response.created"}\n\n', + 'data: {"type":"response.output_text.delta"}\n\n', + 'data: {"type":"response.completed"}\n\n', + "data: [DONE]\n\n", + ) + TERMINAL_INDEX: Final = 2 + + @staticmethod + def _completed_event() -> SimpleNamespace: + return SimpleNamespace( + type="response.completed", + response=SimpleNamespace( + id="resp_lit8612", + output=[SimpleNamespace(type="code_interpreter_call", container_id="cntr_lit8612")], + ), + ) + + async def _sse(self, stream: SimpleNamespace, populate_at: int) -> AsyncGenerator[str, None]: + for index, chunk in enumerate(self.CHUNKS): + if index == populate_at: + stream.completed_response = self._completed_event() + yield chunk + if populate_at == len(self.CHUNKS): + stream.completed_response = self._completed_event() + + async def _await_counts_per_chunk(self, populate_at: int) -> tuple[tuple[tuple[str, int], ...], AsyncMock]: + stream: Final = SimpleNamespace(completed_response=None, _hidden_params={"custom_llm_provider": "azure"}) + recorder: Final = AsyncMock(return_value=None) + with patch( + "litellm.proxy.container_endpoints.ownership.record_container_owners_from_responses_response", recorder + ): + wrapped: Final = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=stream, + wrapped_generator=self._sse(stream, populate_at), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", team_id="team-1"), + ) + observed: Final = tuple([(chunk, recorder.await_count) async for chunk in wrapped]) + return observed, recorder + + async def test_row_is_written_before_the_terminal_chunk_reaches_the_client(self) -> None: + observed, recorder = await self._await_counts_per_chunk(populate_at=self.TERMINAL_INDEX) + + assert tuple(chunk for chunk, _ in observed) == self.CHUNKS + assert tuple(count for _, count in observed) == (0, 0, 1, 1) + recorder.assert_awaited_once() + assert recorder.await_args.kwargs["response"].output[0].container_id == "cntr_lit8612" + assert recorder.await_args.kwargs["user_api_key_dict"].team_id == "team-1" + + async def test_row_is_still_written_when_the_iterator_completes_only_at_exhaustion(self) -> None: + observed, recorder = await self._await_counts_per_chunk(populate_at=len(self.CHUNKS)) + + assert tuple(chunk for chunk, _ in observed) == self.CHUNKS + assert tuple(count for _, count in observed) == (0, 0, 0, 0) + recorder.assert_awaited_once() From 25de1ab2b02ce724439a1da538dad32ebf2a1bc5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:01:27 -0700 Subject: [PATCH 17/58] fix(tests): drop the repeated UNIT_FLAG key in test_unit_shard_missing_paths (#43212) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- tests/unit/test_unit_shard_missing_paths.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py index b528a75d0df..0360a227142 100644 --- a/tests/unit/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -40,7 +40,6 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl "TEST_PATH": test_path, "UNIT_FLAG": "", "WORKERS": workers, - "UNIT_FLAG": "", }, capture_output=True, text=True, From 8327cd6d4724b6218f85562ba164f0cb2cb71ef6 Mon Sep 17 00:00:00 2001 From: daqiangganjun <93830914+daqiangganjun@users.noreply.github.com> Date: Sat, 26 Sep 2026 05:06:49 +0800 Subject: [PATCH 18/58] fix(router): count provider budget spend on every API surface (#38172) * fix(router): count provider budget spend on every API surface RouterBudgetLimiting read custom_llm_provider from litellm_params, which only chat completions populates. Responses, anthropic_messages, embedding and rerank calls raised inside the success callback before any spend was recorded, so those budgets never moved and a ceiling made up mostly of that traffic was never hit. Read the provider from the standard logging payload, which every surface fills in. Dropping the raise also stops one missing field from taking the deployment and tag budgets down with it. * chore(router): drop the inline comment and type the budget limiter test helper --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router_strategy/budget_limiter.py | 10 +- .../router_strategy/test_budget_limiter.py | 137 ++++++++++++++++++ 2 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 4e84bded9de..64252cbbfb3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -502,12 +502,12 @@ class RouterBudgetLimiting(CustomLogger): response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) - custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None) - if custom_llm_provider is None: - raise ValueError("custom_llm_provider is required") + custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - budget_config: Final = self._get_budget_config_for_provider(custom_llm_provider) - if budget_config: + budget_config: Final = ( + self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None + ) + if custom_llm_provider is not None and budget_config is not None: # increment spend for provider spend_key: Final = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}" start_time_key: Final = f"provider_budget_start_time:{custom_llm_provider}" diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..62de1586fdd --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,137 @@ +""" +Spend tracking in RouterBudgetLimiting.async_log_success_event. + +Only chat completions puts custom_llm_provider into litellm_params. The responses, +anthropic_messages, embedding and rerank surfaces leave it unset, which used to make +the callback raise before any spend was recorded, so those budgets never moved. +""" + +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + + +@pytest.fixture +def disable_budget_sync(monkeypatch): + async def noop(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + +def _success_kwargs( + *, + provider_in_litellm_params: str | None, + provider_in_payload: str | None, + call_type: str = "aresponses", + response_cost: float = 0.25, + model_id: str = "deployment-1", +) -> dict[str, object]: + provider_params: Final[dict[str, str]] = ( + {} if provider_in_litellm_params is None else {"custom_llm_provider": provider_in_litellm_params} + ) + litellm_params: Final[dict[str, str]] = {"model": "openai/gpt-4o", **provider_params} + + return { + "call_type": call_type, + "litellm_params": litellm_params, + "standard_logging_object": { + "response_cost": response_cost, + "model_id": model_id, + "custom_llm_provider": provider_in_payload, + }, + } + + +async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object]) -> None: + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"]) +async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type): + """Non-chat surfaces carry the provider only on the standard logging payload.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type=call_type, + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_chat_completions_spend_still_tracked(disable_budget_sync): + """Chat completions fills in both sources and must keep accumulating.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params="openai", + provider_in_payload="openai", + call_type="acompletion", + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_budget_of_other_provider_is_untouched(disable_budget_sync): + """A provider without its own budget must not bleed into a configured one.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload="anthropic"), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") in (None, 0.0) + + +@pytest.mark.asyncio +async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_budget_sync): + """An unresolvable provider must not abort the deployment and tag budgets that follow it.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config=None, + model_list=[ + { + "model_name": "some-model", + "litellm_params": { + "model": "openai/gpt-4o", + "max_budget": 10.0, + "budget_duration": "1d", + }, + "model_info": {"id": "deployment-1"}, + } + ], + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload=None), + ) + + assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 From 1fd04abb9278ef3b33379009874f75fe32e27315 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:32:41 -0700 Subject: [PATCH 19/58] fix(responses): fall back on pre-output stream drops, fail truncated streams, honor request_timeout (#43133) * fix(responses): fall back on pre-output stream drops, fail truncated streams, honor request_timeout A native /v1/responses stream that drops before any output item now raises the router's fallback-eligible MidStreamFallbackError, so configured fallbacks retry the original input. A stream that ends with a clean EOF or a [DONE] marker but no response.completed, response.incomplete or response.failed event now raises litellm.APIConnectionError instead of ending as if it had completed: fallback-eligible before any output, an explicit error after partial output. The sync iterator mirrors every branch. resolve_llm_passthrough_timeout now consults an explicitly set litellm_settings.request_timeout right after the router timeout and before general_settings.pass_through_request_timeout, so the router's native responses path honors it. * test(responses): give the normal-completion stream tests a terminal event --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/passthrough/timeout_utils.py | 7 +- litellm/responses/streaming_iterator.py | 68 +++++-- ...t_base_responses_api_streaming_iterator.py | 54 ++++-- .../test_pass_through_endpoints.py | 24 +++ .../responses/test_streaming_iterator.py | 168 +++++++++++++++++- tests/unit/test_router/test_router.py | 127 +++++++++++++ 6 files changed, 420 insertions(+), 28 deletions(-) diff --git a/litellm/passthrough/timeout_utils.py b/litellm/passthrough/timeout_utils.py index fc67aa8c553..f600c7817f2 100644 --- a/litellm/passthrough/timeout_utils.py +++ b/litellm/passthrough/timeout_utils.py @@ -5,6 +5,8 @@ from typing import Final from pydantic import TypeAdapter +from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout + DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0 _SECONDS: Final = TypeAdapter(float) @@ -48,8 +50,8 @@ def resolve_llm_passthrough_timeout( Anthropic /v1/messages). Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params - timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout - -> 600s default. + timeout/request_timeout -> router_timeout -> litellm.request_timeout (litellm_settings.request_timeout, + when explicitly set) -> general_settings.pass_through_request_timeout -> 600s default. Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before any generic timeout, matching ``Router._get_stream_timeout`` on the completion route: @@ -73,6 +75,7 @@ def resolve_llm_passthrough_timeout( deployment.get("timeout"), deployment.get("request_timeout"), router_timeout, + get_configured_request_timeout(), ) winner: Final = next((val for val in candidates if val is not None), None) return resolve_pass_through_request_timeout() if winner is None else _SECONDS.validate_python(winner) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index fdc702af005..1ef39775bd3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -10,7 +10,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, overload, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -265,6 +265,9 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 +_PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -292,6 +295,7 @@ class BaseResponsesAPIStreamingIterator: self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False + self._output_started = False self._generated_content = "" self._generated_tool_arguments = "" self._completed_response_cached = False @@ -879,6 +883,46 @@ class BaseResponsesAPIStreamingIterator: except Exception: pass + def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: + self._yielded_first_chunk = True + if event.type not in _PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: + self._output_started = True + + def _fallback_error(self, original: Exception) -> MidStreamFallbackError: + return MidStreamFallbackError( + message=str(original), + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + original_exception=original, + generated_content="", + is_pre_first_chunk=not self._yielded_first_chunk, + ) + + def _stream_ended_early_error(self) -> litellm.APIConnectionError: + return litellm.APIConnectionError( + message=( + f"{self.custom_llm_provider or 'provider'} closed the responses stream before any terminal event " + "(response.completed, response.incomplete or response.failed)" + ), + llm_provider=self.custom_llm_provider or "", + model=self.model or "", + ) + + def _raise_if_ended_without_terminal_event(self) -> None: + if self.completed_response is not None: + return + error: Final = self._stream_ended_early_error() + self._handle_failure(error) + if self._output_started: + raise error + raise self._fallback_error(error) from error + + def _raise_for_transport_error(self, error: httpx.ReadError | httpx.RemoteProtocolError) -> NoReturn: + self._handle_failure(error) + if self._output_started: + raise error + raise self._fallback_error(error) from error + async def call_post_streaming_hooks_for_testing( iterator: object, chunk: ResponsesAPIStreamingResponse @@ -934,12 +978,14 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): sse = await self.stream_iterator.__anext__() except StopAsyncIteration: self.finished = True + self._raise_if_ended_without_terminal_event() raise StopAsyncIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: + self._raise_if_ended_without_terminal_event() raise StopAsyncIteration elif result is not None: self._maybe_raise_for_error_event(result) @@ -948,7 +994,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): result = await self._call_post_streaming_deployment_hook( chunk=result, ) - self._yielded_first_chunk = True + self._note_yielded_event(result) return result # If result is None, continue the loop to get the next chunk @@ -957,10 +1003,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True - if self.completed_response is None: - self._handle_failure(e) - raise - raise StopAsyncIteration from e + if self.completed_response is not None: + raise StopAsyncIteration from e + self._raise_for_transport_error(e) except httpx.HTTPError as e: # Handle HTTP errors self.finished = True @@ -1016,12 +1061,14 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): sse = next(self.stream_iterator) except StopIteration: self.finished = True + self._raise_if_ended_without_terminal_event() raise StopIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: + self._raise_if_ended_without_terminal_event() raise StopIteration elif result is not None: self._maybe_raise_for_error_event(result) @@ -1030,7 +1077,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): async_function=self._call_post_streaming_deployment_hook, chunk=result, ) - self._yielded_first_chunk = True + self._note_yielded_event(result) return result # If result is None, continue the loop to get the next chunk @@ -1039,10 +1086,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True - if self.completed_response is None: - self._handle_failure(e) - raise - raise StopIteration from e + if self.completed_response is not None: + raise StopIteration from e + self._raise_for_transport_error(e) except httpx.HTTPError as e: # Handle HTTP errors self.finished = True diff --git a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py index 47b377dc9a4..da37803b64a 100644 --- a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py +++ b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py @@ -381,6 +381,36 @@ class TestBaseResponsesAPIStreamingIterator: ) raise + @staticmethod + def _config_completing_after_one_delta() -> Mock: + mock_config = Mock(spec=BaseResponsesAPIConfig) + completed_response = ResponsesAPIResponse( + id="resp_123", + created_at=0, + status="completed", + model="gpt-5.5", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2), + ) + + def _transform(model, parsed_chunk, logging_obj): + if parsed_chunk.get("type") == "response.completed": + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed_response, + ) + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_123", + output_index=0, + content_index=0, + delta=parsed_chunk["delta"], + ) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + @pytest.mark.asyncio async def test_stop_async_iteration_not_logged_as_failure(self): """ @@ -399,6 +429,7 @@ class TestBaseResponsesAPIStreamingIterator: async def mock_aiter_bytes(): yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' + yield b'data: {"type": "response.completed", "response": {"id": "resp_123"}}\n\n' mock_response.aiter_bytes = mock_aiter_bytes @@ -408,11 +439,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler = Mock() mock_logging_obj.failure_handler = Mock() - mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_delta_event = Mock() - mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - mock_delta_event.delta = "test" - mock_config.transform_streaming_response.return_value = mock_delta_event + mock_config = self._config_completing_after_one_delta() # Create the iterator instance iterator = ResponsesAPIStreamingIterator( @@ -432,8 +459,9 @@ class TestBaseResponsesAPIStreamingIterator: except StopAsyncIteration: pass # This is expected - # Verify we got the chunk - assert len(chunks_received) == 1 + # Verify we got the delta and the terminal event + assert len(chunks_received) == 2 + assert iterator.completed_response is not None # CRITICAL: Verify that failure handlers were NOT called # StopAsyncIteration is a normal end of stream, not a failure @@ -460,6 +488,7 @@ class TestBaseResponsesAPIStreamingIterator: def mock_iter_bytes(): yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' + yield b'data: {"type": "response.completed", "response": {"id": "resp_123"}}\n\n' mock_response.iter_bytes = mock_iter_bytes @@ -469,11 +498,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler = Mock() mock_logging_obj.failure_handler = Mock() - mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_delta_event = Mock() - mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - mock_delta_event.delta = "test" - mock_config.transform_streaming_response.return_value = mock_delta_event + mock_config = self._config_completing_after_one_delta() # Create the iterator instance iterator = SyncResponsesAPIStreamingIterator( @@ -493,8 +518,9 @@ class TestBaseResponsesAPIStreamingIterator: except StopIteration: pass # This is expected - # Verify we got the chunk - assert len(chunks_received) == 1 + # Verify we got the delta and the terminal event + assert len(chunks_received) == 2 + assert iterator.completed_response is not None # CRITICAL: Verify that failure handlers were NOT called # StopIteration is a normal end of stream, not a failure 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 a40741c8fdb..3469df082e0 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 @@ -23,6 +23,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -1165,6 +1166,29 @@ def test_resolve_llm_passthrough_timeout_precedence(): assert resolve_llm_passthrough_timeout() == 6.0 +def test_resolve_llm_passthrough_timeout_honors_explicit_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert resolve_llm_passthrough_timeout() == 44.0 + assert resolve_llm_passthrough_timeout(kwargs={"stream": True}) == 44.0 + assert resolve_llm_passthrough_timeout(router_timeout=120) == 120.0 + assert resolve_llm_passthrough_timeout(kwargs={"stream": True}, router_stream_timeout=900) == 900.0 + assert resolve_llm_passthrough_timeout(litellm_params={"timeout": 90}) == 90.0 + assert resolve_llm_passthrough_timeout(kwargs={"timeout": 45}) == 45.0 + + +def test_resolve_llm_passthrough_timeout_skips_unset_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("litellm.request_timeout", float(DEFAULT_REQUEST_TIMEOUT_SECONDS), raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", False, raising=False) + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert resolve_llm_passthrough_timeout() == 6.0 + with patch("litellm.proxy.proxy_server.general_settings", {}): + assert resolve_llm_passthrough_timeout() == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS + + def test_resolve_llm_passthrough_timeout_stream_timeout_precedence(): assert ( resolve_llm_passthrough_timeout( diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index dbf54ec3b9b..9dbbc20591e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -6,13 +6,14 @@ completion_start_time = end_time.""" import json from datetime import datetime from typing import Final, Optional -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch import httpx import pytest from pydantic_core import PydanticSerializationError import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( @@ -251,6 +252,171 @@ def test_sync_transport_error_before_completed_event_raises(): pass +_DONE_MARKER: Final = b"data: [DONE]\n\n" +_CREATED_EVENT: Final = _sse_event({"type": "response.created"}) +_IN_PROGRESS_EVENT: Final = _sse_event({"type": "response.in_progress"}) +_PARTIAL_OUTPUT_EVENTS: Final = _COMPLETE_STREAM_EVENTS[:-1] +_PRE_OUTPUT_PREFIXES: Final = [ + pytest.param([], True, id="nothing-yielded"), + pytest.param([_CREATED_EVENT], False, id="created"), + pytest.param([_CREATED_EVENT, _IN_PROGRESS_EVENT], False, id="created-and-in-progress"), +] + + +def _failure_tracking_logging_obj() -> Mock: + logging_obj: Final = _logging_obj_stub() + logging_obj.async_failure_handler = AsyncMock() + return logging_obj + + +def _assert_failure_logged_once(logging_obj: Mock, exception: Exception) -> None: + assert logging_obj.async_failure_handler.await_count == 1 + assert logging_obj.async_failure_handler.await_args.kwargs["exception"] is exception + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix, pre_first_chunk", _PRE_OUTPUT_PREFIXES) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +async def test_transport_error_before_any_output_raises_fallback_error(prefix, pre_first_chunk, trailing_error): + """A connection lost while only lifecycle events (response.created / response.in_progress) + have streamed is fallback-eligible, so it must surface as the MidStreamFallbackError the + router re-routes, carrying the raw transport error and no generated content.""" + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=prefix, logging_obj=logging_obj, trailing_error=trailing_error) + + with pytest.raises(MidStreamFallbackError) as exc_info: + async for _ in iterator: + pass + + assert exc_info.value.original_exception is trailing_error + assert exc_info.value.is_pre_first_chunk is pre_first_chunk + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.asyncio +async def test_transport_error_after_output_started_is_not_fallback_eligible(): + logging_obj: Final = _failure_tracking_logging_obj() + trailing_error: Final = httpx.ReadError("Response payload is not completed") + iterator: Final = _make_iterator( + sse_events=_PARTIAL_OUTPUT_EVENTS, logging_obj=logging_obj, trailing_error=trailing_error + ) + + with pytest.raises(httpx.ReadError) as exc_info: + async for _ in iterator: + pass + + assert exc_info.value is trailing_error + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_stream_ending_after_partial_output_without_terminal_event_raises(trailer): + """A clean EOF or `[DONE]` after output text but with no response.completed / + response.incomplete / response.failed is a truncated answer: the partial events still + reach the caller, then an explicit error follows instead of a normal end of stream.""" + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*_PARTIAL_OUTPUT_EVENTS, *trailer], logging_obj=logging_obj) + + created: Final = await iterator.__anext__() + delta: Final = await iterator.__anext__() + with pytest.raises(litellm.APIConnectionError) as exc_info: + await iterator.__anext__() + + assert (created.type, delta.type) == ("response.created", "response.output_text.delta") + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.llm_provider == "openai" + _assert_failure_logged_once(logging_obj, exc_info.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix, pre_first_chunk", _PRE_OUTPUT_PREFIXES) +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_stream_ending_before_any_output_raises_fallback_error(prefix, pre_first_chunk, trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*prefix, *trailer], logging_obj=logging_obj) + + with pytest.raises(MidStreamFallbackError) as exc_info: + async for _ in iterator: + pass + + assert isinstance(exc_info.value.original_exception, litellm.APIConnectionError) + assert exc_info.value.is_pre_first_chunk is pre_first_chunk + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, exc_info.value.original_exception) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_complete_stream_still_ends_normally(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*_COMPLETE_STREAM_EVENTS, *trailer], logging_obj=logging_obj) + + seen: Final = [event.type async for event in iterator] + + assert seen[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.async_failure_handler.await_count == 0 + + +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +def test_sync_transport_error_before_any_output_raises_fallback_error(trailing_error): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator( + sse_events=[_CREATED_EVENT, _IN_PROGRESS_EVENT], + logging_obj=logging_obj, + trailing_error=trailing_error, + ) + + with pytest.raises(MidStreamFallbackError) as exc_info: + for _ in iterator: + pass + + assert exc_info.value.original_exception is trailing_error + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +def test_sync_stream_ending_after_partial_output_without_terminal_event_raises(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[*_PARTIAL_OUTPUT_EVENTS, *trailer], logging_obj=logging_obj) + + created: Final = next(iterator) + delta: Final = next(iterator) + with pytest.raises(litellm.APIConnectionError) as exc_info: + next(iterator) + + assert (created.type, delta.type) == ("response.created", "response.output_text.delta") + assert not isinstance(exc_info.value, MidStreamFallbackError) + _assert_failure_logged_once(logging_obj, exc_info.value) + + +def test_sync_stream_ending_before_any_output_raises_fallback_error(): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[_CREATED_EVENT], logging_obj=logging_obj) + + with pytest.raises(MidStreamFallbackError) as exc_info: + for _ in iterator: + pass + + assert isinstance(exc_info.value.original_exception, litellm.APIConnectionError) + assert exc_info.value.is_pre_first_chunk is False + _assert_failure_logged_once(logging_obj, exc_info.value.original_exception) + + +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +def test_sync_complete_stream_still_ends_normally(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[*_COMPLETE_STREAM_EVENTS, *trailer], logging_obj=logging_obj) + + seen: Final = [event.type for event in iterator] + + assert seen[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.async_failure_handler.await_count == 0 + + def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch): """ Regression test for LIT-6184 on the /v1/responses streaming surface: the diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 80131534183..10669de9cc8 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -4486,6 +4486,107 @@ async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation( assert fbk["input"] == "Hello" # original input, no continuation messages +def _make_native_responses_iterator(*, sse_payloads: tuple[dict[str, str], ...], trailing_error: Exception | None): + """A real ResponsesAPIStreamingIterator over canned SSE bytes, so the router test covers the + iterator's own transport-error classification instead of a hand-built MidStreamFallbackError.""" + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + async def aiter_bytes(): + for payload in sse_payloads: + yield f"data: {json.dumps(payload)}\n\n".encode() + if trailing_error is not None: + raise trailing_error + + def transform(model, parsed_chunk, logging_obj): + return MagicMock(type=parsed_chunk["type"]) + + response: Final = MagicMock() + response.headers = {} + response.aiter_bytes = aiter_bytes + config: Final = MagicMock(spec=BaseResponsesAPIConfig) + config.transform_streaming_response.side_effect = transform + logging_obj: Final = MagicMock(spec=LiteLLMLogging) + logging_obj.completion_start_time = None + logging_obj.model_call_details = {"litellm_params": {}} + return ResponsesAPIStreamingIterator( + response=response, + model="gpt-4", + responses_api_provider_config=config, + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="openai", + ) + + +_RESPONSES_LIFECYCLE_PAYLOADS: Final = ({"type": "response.created"}, {"type": "response.in_progress"}) + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before_output(): + """A connection lost after response.created but before any output item is re-routed to the + fallback with the original input, the same as a provider error event would be.""" + router: Final = _make_router_with_fallback() + src: Final = _make_native_responses_iterator( + sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, + trailing_error=httpx.ReadError("Response payload is not completed"), + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList([MagicMock(type="response.completed")]), + ) as mock_fallback_utils: + wrapped: Final = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + seen: Final = [chunk.type async for chunk in wrapped] + + assert seen == ["response.created", "response.in_progress", "response.completed"] + assert isinstance(mock_fallback_utils.call_args.kwargs["e"], MidStreamFallbackError) + assert mock_fallback_utils.call_args.kwargs["kwargs"]["input"] == "Hello" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_surfaces_transport_drop_when_no_fallback_lands(): + transport_error: Final = httpx.ReadError("Response payload is not completed") + router: Final = _make_router_with_fallback() + src: Final = _make_native_responses_iterator( + sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, trailing_error=transport_error + ) + + async def reraise_trigger(**kwargs): + raise kwargs["e"] + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=reraise_trigger) + ) as mock_fallback_utils: + wrapped: Final = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + with pytest.raises(httpx.ReadError) as exc_info: + async for _ in wrapped: + pass + + assert exc_info.value is transport_error + assert mock_fallback_utils.await_count == 1 + trigger: Final = mock_fallback_utils.await_args.kwargs["e"] + assert isinstance(trigger, MidStreamFallbackError) + assert trigger.original_exception is transport_error + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: input is rewritten to include user prompt + @@ -6090,6 +6191,32 @@ def test_update_kwargs_with_deployment_passthrough_router_stream_timeout_sources assert _passthrough_timeout(default_router, default_router.model_list[0], stream=False) == 120.0 +def test_update_kwargs_with_deployment_passthrough_honors_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + """litellm_settings.request_timeout must bound the native responses route when neither the + deployment nor the router carries a timeout, while a deployment timeout keeps winning.""" + monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "responses-global-timeout", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key"}, + }, + { + "model_name": "responses-deployment-timeout", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "timeout": 3}, + }, + ], + ) + global_only, per_deployment = router.model_list + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert _passthrough_timeout(router, global_only, stream=True) == 44.0 + assert _passthrough_timeout(router, global_only, stream=False) == 44.0 + assert _passthrough_timeout(router, per_deployment, stream=True) == 3.0 + assert _passthrough_timeout(router, per_deployment, stream=False) == 3.0 + + @pytest.mark.asyncio async def test_router_acompletion_with_unknown_model_and_default_fallback(): """ From a09f8b84a4fc97c57caaf8fb0a446f170c706f10 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:52:35 -0700 Subject: [PATCH 20/58] fix(sentry): scrub PII and secrets inside object reprs and nested locals, add SENTRY_SEND_DEFAULT_PII opt-in (#43123) * fix(sentry): scrub PII and secrets inside object reprs and nested locals, add SENTRY_SEND_DEFAULT_PII opt-in * fix(sentry): keep the SDK denylist and filter the request headers a virtual key arrives in * fix(sentry): leave source context lines unscrubbed * fix(sentry): filter bracketed secret values and cap the JSON walk depth * ci(deps): install sentry-sdk in the proxy-dev group so the unit shards import it * fix(sentry): scrub source-context names outside real stack frames * fix(sentry): tie the key pattern floor to the custom key minimum * fix(sentry): keep the key pattern floor at or below a generated key's length --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 15 + litellm/litellm_core_utils/litellm_logging.py | 17 +- .../litellm_core_utils/sentry_scrubbing.py | 152 ++++++++++ pyproject.toml | 1 + .../code_coverage_tests/recursive_detector.py | 1 + .../test_litellm_logging.py | 148 ++-------- .../test_sentry_scrubbing.py | 278 ++++++++++++++++++ uv.lock | 2 + 8 files changed, 481 insertions(+), 133 deletions(-) create mode 100644 litellm/litellm_core_utils/sentry_scrubbing.py create mode 100644 tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py diff --git a/litellm/constants.py b/litellm/constants.py index e7ba1f6b07f..8316761c95b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1952,6 +1952,15 @@ SENTRY_DENYLIST: Final = [ "auth_token", "jwt_token", "private_key", + "authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "ocp-apim-subscription-key", + "x-litellm-api-key", + "x-mcp-auth", + "cookie", + "set-cookie", "SLACK_WEBHOOK_URL", "ALERTING_WEBHOOK_URL", "webhook_url", @@ -1974,6 +1983,12 @@ SENTRY_DENYLIST: Final = [ ] SENTRY_PII_DENYLIST: Final = [ "user_id", + "user_email", + "end_user_id", + "user_api_key_hash", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_end_user_id", "email", "phone", "address", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 83ab2bc11a2..d8182140a17 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -43,8 +43,6 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, PROVIDER_REQUEST_ID_HEADERS, - SENTRY_DENYLIST, - SENTRY_PII_DENYLIST, ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, @@ -4423,21 +4421,10 @@ def set_callbacks(callback_list, function_id=None): print_verbose("Package 'sentry_sdk' is missing. Installing it...") subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"]) import sentry_sdk - from sentry_sdk.scrubber import EventScrubber + from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options sentry_sdk_instance = sentry_sdk - sentry_trace_rate = os.environ.get("SENTRY_API_TRACE_RATE", "1.0") - sentry_sample_rate = ( - os.environ.get("SENTRY_API_SAMPLE_RATE") if "SENTRY_API_SAMPLE_RATE" in os.environ else "1.0" - ) - sentry_sdk_instance.init( - dsn=os.environ.get("SENTRY_DSN"), - traces_sample_rate=float(sentry_trace_rate), - sample_rate=float(sentry_sample_rate if sentry_sample_rate else 1.0), - send_default_pii=False, # Prevent sending Personal Identifiable Information - event_scrubber=EventScrubber(denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST), - environment=os.environ.get("SENTRY_ENVIRONMENT", "production"), - ) + sentry_sdk_instance.init(**build_sentry_init_options(os.environ)) capture_exception = sentry_sdk_instance.capture_exception add_breadcrumb = sentry_sdk_instance.add_breadcrumb elif callback == "slack": diff --git a/litellm/litellm_core_utils/sentry_scrubbing.py b/litellm/litellm_core_utils/sentry_scrubbing.py new file mode 100644 index 00000000000..4c14cabc2ab --- /dev/null +++ b/litellm/litellm_core_utils/sentry_scrubbing.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import re +from collections.abc import Callable, Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Final, TypeAlias, cast + +from pydantic import JsonValue +from sentry_sdk.scrubber import DEFAULT_DENYLIST, DEFAULT_PII_DENYLIST, EventScrubber +from typing_extensions import ReadOnly, TypedDict + +from litellm.constants import ( + LENGTH_OF_LITELLM_GENERATED_KEY, + MINIMUM_CUSTOM_KEY_LENGTH, + SENTRY_DENYLIST, + SENTRY_PII_DENYLIST, +) +from litellm.secret_managers.main import str_to_bool + +if TYPE_CHECKING: + from sentry_sdk.types import Event, Hint + +EventScrubFn: TypeAlias = "Callable[[Event, Hint], Event]" +JsonPath: TypeAlias = tuple[str, ...] + +FILTERED: Final = "[Filtered]" +SEND_DEFAULT_PII_ENV: Final = "SENTRY_SEND_DEFAULT_PII" +SECRET_FIELD_NAMES: Final = tuple(DEFAULT_DENYLIST) + tuple(SENTRY_DENYLIST) +PII_FIELD_NAMES: Final = tuple(DEFAULT_PII_DENYLIST) + tuple(SENTRY_PII_DENYLIST) + +KEY_PREFIX: Final = "sk-" + + +def build_key_pattern(custom_key_minimum: int, generated_key_bytes: int) -> re.Pattern[str]: + generated_suffix_length: Final = (generated_key_bytes * 4 + 2) // 3 + floor: Final = min(custom_key_minimum - len(KEY_PREFIX), generated_suffix_length) + return re.compile(rf"{KEY_PREFIX}[A-Za-z0-9_-]{{{floor},}}") + + +LITELLM_KEY_PATTERN: Final = build_key_pattern(MINIMUM_CUSTOM_KEY_LENGTH, LENGTH_OF_LITELLM_GENERATED_KEY) +SOURCE_CONTEXT_KEYS: Final = frozenset({"pre_context", "context_line", "post_context"}) +STACK_FRAME_PATHS: Final = frozenset( + { + ("exception", "values", "*", "stacktrace", "frames", "*"), + ("threads", "values", "*", "stacktrace", "frames", "*"), + ("stacktrace", "frames", "*"), + } +) +MAX_SCRUB_DEPTH: Final = 64 +EMAIL_PATTERN: Final = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}") +SHA256_HEX_PATTERN: Final = re.compile(r"(? re.Pattern[str]: + names: Final = "|".join(re.escape(name) for name in field_names) + return re.compile( + rf"(?P(?{QUOTED_VALUE}|{BRACKETED_VALUE}|{BARE_VALUE})", + re.IGNORECASE, + ) + + +def build_string_scrubber(send_default_pii: bool) -> Callable[[str], str]: + field_names: Final = SECRET_FIELD_NAMES if send_default_pii else SECRET_FIELD_NAMES + PII_FIELD_NAMES + field_pattern: Final = build_repr_field_pattern(field_names) + value_patterns: Final = ( + (LITELLM_KEY_PATTERN,) if send_default_pii else (LITELLM_KEY_PATTERN, EMAIL_PATTERN, SHA256_HEX_PATTERN) + ) + + def scrub(text: str) -> str: + fields_scrubbed: Final = field_pattern.sub(_filtered_field, text) + return _substitute_all(value_patterns, fields_scrubbed) + + return scrub + + +def _filtered_field(match: re.Match[str]) -> str: + quote: Final = '"' if match.group("value").startswith('"') else "'" + return f"{match.group('field')}{quote}{FILTERED}{quote}" + + +def _substitute_all(patterns: Sequence[re.Pattern[str]], text: str) -> str: + return reduce(lambda scrubbed, pattern: pattern.sub(FILTERED, scrubbed), patterns, text) + + +def scrub_json_strings(value: JsonValue, scrub: Callable[[str], str], path: JsonPath = ()) -> JsonValue: + if len(path) > MAX_SCRUB_DEPTH: + return FILTERED + if isinstance(value, str): + return scrub(value) + if isinstance(value, dict): + unscrubbed_keys: Final = SOURCE_CONTEXT_KEYS if path in STACK_FRAME_PATHS else frozenset[str]() + return { # mutable-ok: JSON object + key: item if key in unscrubbed_keys else scrub_json_strings(item, scrub, (*path, key)) + for key, item in value.items() + } + if isinstance(value, list): + return [scrub_json_strings(item, scrub, (*path, "*")) for item in value] # mutable-ok: JSON array + return value + + +def build_event_scrubber(send_default_pii: bool) -> EventScrubFn: + scrub: Final = build_string_scrubber(send_default_pii) + + def scrub_event(event: Event, _hint: Hint) -> Event: + json_event: Final = cast("JsonValue", event) # cast-ok: [LIT006] the SDK serialized the event to JSON already + return cast("Event", scrub_json_strings(json_event, scrub)) # cast-ok: [LIT006] same JSON shape going back + + return scrub_event + + +def send_default_pii_from_env(env: Mapping[str, str]) -> bool: + return str_to_bool(env.get(SEND_DEFAULT_PII_ENV)) is True + + +def build_sentry_init_options(env: Mapping[str, str]) -> SentryInitOptions: + send_default_pii: Final = send_default_pii_from_env(env) + scrub_event: Final = build_event_scrubber(send_default_pii) + return SentryInitOptions( + dsn=env.get("SENTRY_DSN"), + traces_sample_rate=float(env.get("SENTRY_API_TRACE_RATE") or "1.0"), + sample_rate=float(env.get("SENTRY_API_SAMPLE_RATE") or "1.0"), + send_default_pii=send_default_pii, + event_scrubber=EventScrubber( + denylist=list(SECRET_FIELD_NAMES), # mutable-ok: EventScrubber appends pii_denylist onto denylist in place + pii_denylist=list(PII_FIELD_NAMES), # mutable-ok: EventScrubber takes List[str] + recursive=True, + send_default_pii=send_default_pii, + ), + before_send=scrub_event, + before_send_transaction=scrub_event, + environment=env.get("SENTRY_ENVIRONMENT", "production"), + ) diff --git a/pyproject.toml b/pyproject.toml index f2364b5e77b..28b00379cc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -249,6 +249,7 @@ proxy-dev = [ "prisma==0.11.0", "hypercorn==0.17.3", "prometheus-client==0.20.0", + "sentry-sdk==2.21.0", "opentelemetry-api==1.33.1", "opentelemetry-sdk==1.33.1", "opentelemetry-exporter-otlp==1.33.1", diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index dc1f8592612..659dc438f2d 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -72,6 +72,7 @@ IGNORE_FUNCTIONS = [ "_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_replace_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible). + "scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap. ] diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3afa31cc801..c4829ced9f3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -20,7 +20,7 @@ from openai._legacy_response import HttpxBinaryResponseContent import litellm from litellm._logging import session_id_var, trace_id_var -from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.constants import SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -357,108 +357,23 @@ def test_post_call_serializes_dict_with_datetime(logging_obj): assert "2026-05-11" in serialized -def test_sentry_sample_rate(monkeypatch): - existing_sample_rate = os.getenv("SENTRY_API_SAMPLE_RATE") - try: - # test with default value by removing the environment variable - if existing_sample_rate: - del os.environ["SENTRY_API_SAMPLE_RATE"] - - set_callbacks(["sentry"]) - # Check if the default sample rate is set to 1.0 - assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "1.0" - - # test with custom value - monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", "0.5") - - set_callbacks(["sentry"]) - # Check if the custom sample rate is set correctly - assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "0.5" - except Exception as e: - print(f"Error: {e}") - finally: - # Restore the original environment variable - if existing_sample_rate: - monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", existing_sample_rate) - else: - if "SENTRY_API_SAMPLE_RATE" in os.environ: - del os.environ["SENTRY_API_SAMPLE_RATE"] - - def test_sentry_environment(monkeypatch): - """Test that SENTRY_ENVIRONMENT is properly handled during Sentry initialization""" - existing_environment = os.getenv("SENTRY_ENVIRONMENT") - existing_dsn = os.getenv("SENTRY_DSN") + import sentry_sdk - # Create mock sentry_sdk module - mock_event_scrubber_instance = MagicMock() - mock_event_scrubber_cls = MagicMock(return_value=mock_event_scrubber_instance) - - mock_scrubber_module = MagicMock() - mock_scrubber_module.EventScrubber = mock_event_scrubber_cls - - mock_sentry_sdk = MagicMock() - mock_sentry_sdk.scrubber = mock_scrubber_module mock_init = MagicMock() - mock_sentry_sdk.init = mock_init + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") + monkeypatch.delenv("SENTRY_ENVIRONMENT", raising=False) - # Inject mocks into sys.modules - sys.modules["sentry_sdk"] = mock_sentry_sdk - sys.modules["sentry_sdk.scrubber"] = mock_scrubber_module - - try: - # Set a mock DSN to allow Sentry initialization - monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") - - # Test with default value (no environment set) - if existing_environment: - del os.environ["SENTRY_ENVIRONMENT"] + set_callbacks(["sentry"]) + assert mock_init.call_args[1]["environment"] == "production" + for environment in ("development", "staging"): + monkeypatch.setenv("SENTRY_ENVIRONMENT", environment) mock_init.reset_mock() set_callbacks(["sentry"]) - # Check that init was called with default environment "production" mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "production" - - # Test with custom environment value - monkeypatch.setenv("SENTRY_ENVIRONMENT", "development") - - mock_init.reset_mock() - set_callbacks(["sentry"]) - # Check that init was called with custom environment "development" - mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "development" - - # Test with staging environment - monkeypatch.setenv("SENTRY_ENVIRONMENT", "staging") - - mock_init.reset_mock() - set_callbacks(["sentry"]) - # Check that init was called with custom environment "staging" - mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "staging" - - except Exception as e: - print(f"Error: {e}") - raise - finally: - # Restore the original environment variables - if existing_environment: - monkeypatch.setenv("SENTRY_ENVIRONMENT", existing_environment) - else: - if "SENTRY_ENVIRONMENT" in os.environ: - del os.environ["SENTRY_ENVIRONMENT"] - - if existing_dsn: - monkeypatch.setenv("SENTRY_DSN", existing_dsn) - else: - if "SENTRY_DSN" in os.environ: - del os.environ["SENTRY_DSN"] - - + assert mock_init.call_args[1]["environment"] == environment def test_use_custom_pricing_for_model(): from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model @@ -3100,37 +3015,34 @@ def test_speech_call_is_still_priced_from_input_characters(call_type): def test_sentry_event_scrubber_initialization(monkeypatch): - # Step 1: Create a fake sentry_sdk.scrubber module - mock_event_scrubber_instance = MagicMock() - mock_event_scrubber_cls = MagicMock(return_value=mock_event_scrubber_instance) + import sentry_sdk - mock_scrubber_module = MagicMock() - mock_scrubber_module.EventScrubber = mock_event_scrubber_cls - - # Step 2: Create a fake sentry_sdk module and insert into sys.modules - mock_sentry_sdk = MagicMock() - mock_sentry_sdk.scrubber = mock_scrubber_module mock_init = MagicMock() - mock_sentry_sdk.init = mock_init + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.delenv("SENTRY_SEND_DEFAULT_PII", raising=False) - # Step 3: Inject both into sys.modules BEFORE import occurs - sys.modules["sentry_sdk"] = mock_sentry_sdk - sys.modules["sentry_sdk.scrubber"] = mock_scrubber_module - - # Step 4: Run the actual sentry setup code set_callbacks(["sentry"]) - # Step 5: Assert the EventScrubber was constructed correctly - mock_event_scrubber_cls.assert_called_once_with( - denylist=SENTRY_DENYLIST, - pii_denylist=SENTRY_PII_DENYLIST, - ) - - # Step 6: Assert the event_scrubber and PII args were passed mock_init.assert_called_once() call_args = mock_init.call_args[1] - assert call_args["event_scrubber"] == mock_event_scrubber_instance assert call_args["send_default_pii"] is False + assert call_args["event_scrubber"].recursive is True + assert {name.lower() for name in SENTRY_PII_DENYLIST} <= {name.lower() for name in call_args["event_scrubber"].denylist} + assert call_args["before_send"] is call_args["before_send_transaction"] + + +def test_sentry_send_default_pii_opt_in(monkeypatch): + import sentry_sdk + + mock_init = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_SEND_DEFAULT_PII", "true") + + set_callbacks(["sentry"]) + + call_args = mock_init.call_args[1] + assert call_args["send_default_pii"] is True + assert not {name.lower() for name in SENTRY_PII_DENYLIST} & {name.lower() for name in call_args["event_scrubber"].denylist} def test_get_masked_values(): diff --git a/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py b/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py new file mode 100644 index 00000000000..9aae3999129 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py @@ -0,0 +1,278 @@ +import hashlib +import json +import secrets +from collections.abc import Callable, Mapping +from functools import reduce +from typing import Final, cast + +import pytest +import sentry_sdk +from pydantic import JsonValue +from sentry_sdk.envelope import Envelope +from sentry_sdk.transport import Transport +from sentry_sdk.utils import event_from_exception + +from litellm.constants import LENGTH_OF_LITELLM_GENERATED_KEY, MINIMUM_CUSTOM_KEY_LENGTH +from litellm.litellm_core_utils.sentry_scrubbing import ( + FILTERED, + MAX_SCRUB_DEPTH, + build_key_pattern, + build_sentry_init_options, + build_string_scrubber, + scrub_json_strings, +) +from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + +EMAIL: Final = "qa.user@example.com" +VIRTUAL_KEY: Final = "sk-virtual-key-under-test" +KEY_HASH: Final = hashlib.sha256(VIRTUAL_KEY.encode()).hexdigest() +MASTER_KEY: Final = "sk-master-key-under-test" +DATABASE_URL: Final = "postgresql://litellm:db-password-under-test@db.internal:5432/litellm" +PII_ON: Final = {"SENTRY_DSN": "https://key@sentry.example/1", "SENTRY_SEND_DEFAULT_PII": "true"} +PII_OFF: Final = {"SENTRY_DSN": "https://key@sentry.example/1"} + + +class RecordingTransport(Transport): + def __init__(self) -> None: + super().__init__() + self.last_envelope: Envelope | None = None + + def capture_envelope(self, envelope: Envelope) -> None: + self.last_envelope = envelope + + +def reject_request( + valid_token: UserAPIKeyAuth, + user_obj: LiteLLM_UserTable, + general_settings: Mapping[str, str], + data: Mapping[str, Mapping[str, str]], + raw_headers: Mapping[str, str], +) -> None: + raise RuntimeError(f"key {valid_token.token} owned by {user_obj.user_email} was rejected") + + +def raise_with_identity_locals() -> None: + reject_request( + valid_token=UserAPIKeyAuth(token=KEY_HASH, key_name="sk-...test", user_id=EMAIL, user_email=EMAIL), + user_obj=LiteLLM_UserTable(user_id=EMAIL, user_email=EMAIL, user_role="internal_user"), + general_settings={"master_key": MASTER_KEY, "database_url": DATABASE_URL}, + data={"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_email": EMAIL}}, + raw_headers={"authorization": f"Bearer {VIRTUAL_KEY}", "x-api-key": VIRTUAL_KEY, "content-type": "application/json"}, + ) + + +def raise_with_source_context_named_locals() -> None: + metadata: Final = {"context_line": f"Bearer {VIRTUAL_KEY}", "pre_context": [EMAIL], "post_context": [KEY_HASH]} + stacktrace: Final = {"frames": [{"context_line": MASTER_KEY, "pre_context": [EMAIL]}]} + raise RuntimeError(f"rejected with {len(metadata)} metadata fields and {len(stacktrace)} stack fields") + + +def capture_serialized_event(env: Mapping[str, str], raiser: Callable[[], None] = raise_with_identity_locals) -> str: + transport: Final = RecordingTransport() + client: Final = sentry_sdk.Client(transport=transport, **build_sentry_init_options(env)) + try: + raiser() + except RuntimeError as error: + event, hint = event_from_exception(error, client_options=client.options) + client.capture_event(event, hint=hint) + assert transport.last_envelope is not None + return json.dumps(transport.last_envelope.items[0].payload.json) + + +def innermost_frame_vars(serialized: str) -> dict[str, JsonValue]: + event: Final = json.loads(serialized) + frames: Final = event["exception"]["values"][0]["stacktrace"]["frames"] + return frames[-1]["vars"] + + +def test_default_event_carries_no_email_hash_or_secret_anywhere() -> None: + serialized: Final = capture_serialized_event(PII_OFF) + assert EMAIL not in serialized + assert KEY_HASH not in serialized + assert MASTER_KEY not in serialized + assert VIRTUAL_KEY not in serialized + assert "db-password-under-test" not in serialized + frame_vars: Final = innermost_frame_vars(serialized) + assert frame_vars["raw_headers"] == {"authorization": FILTERED, "x-api-key": FILTERED, "content-type": "'application/json'"} + assert f"token='{FILTERED}'" in frame_vars["valid_token"] + assert f"user_id='{FILTERED}'" in frame_vars["valid_token"] + assert f"user_email='{FILTERED}'" in frame_vars["user_obj"] + assert frame_vars["general_settings"] == {"master_key": FILTERED, "database_url": FILTERED} + assert frame_vars["data"] == {"metadata": {"user_api_key_hash": FILTERED, "user_api_key_user_email": FILTERED}} + assert "key_name='sk-...test'" in frame_vars["valid_token"] + assert "user_role='internal_user'" in frame_vars["user_obj"] + + +def test_source_context_lines_are_left_readable() -> None: + frames: Final = json.loads(capture_serialized_event(PII_OFF))["exception"]["values"][0]["stacktrace"]["frames"] + source_lines: Final = tuple( + line + for frame in frames + for line in (*frame.get("pre_context", []), frame.get("context_line", ""), *frame.get("post_context", [])) + ) + assert any("token=KEY_HASH" in line for line in source_lines) + assert not any(FILTERED in line for line in source_lines) + + +def test_source_context_names_outside_stack_frames_are_scrubbed() -> None: + serialized: Final = capture_serialized_event(PII_OFF, raise_with_source_context_named_locals) + assert VIRTUAL_KEY not in serialized + assert MASTER_KEY not in serialized + assert EMAIL not in serialized + assert KEY_HASH not in serialized + frame_vars: Final = innermost_frame_vars(serialized) + assert frame_vars["metadata"] == { + "context_line": f"'Bearer {FILTERED}'", + "pre_context": [f"'{FILTERED}'"], + "post_context": [f"'{FILTERED}'"], + } + assert frame_vars["stacktrace"] == {"frames": [{"context_line": f"'{FILTERED}'", "pre_context": [f"'{FILTERED}'"]}]} + innermost_frame: Final = json.loads(serialized)["exception"]["values"][0]["stacktrace"]["frames"][-1] + assert "raise RuntimeError" in innermost_frame["context_line"] + assert FILTERED not in json.dumps(innermost_frame["pre_context"]) + + +def test_default_event_keeps_the_exception_message_shape() -> None: + serialized: Final = capture_serialized_event(PII_OFF) + message: Final = json.loads(serialized)["exception"]["values"][0]["value"] + assert message == f"key {FILTERED} owned by {FILTERED} was rejected" + + +def test_pii_opt_in_keeps_identifiers_and_still_scrubs_secrets() -> None: + serialized: Final = capture_serialized_event(PII_ON) + frame_vars: Final = innermost_frame_vars(serialized) + assert f"user_id='{EMAIL}'" in frame_vars["valid_token"] + assert f"user_email='{EMAIL}'" in frame_vars["user_obj"] + assert frame_vars["data"] == { + "metadata": {"user_api_key_hash": f"'{KEY_HASH}'", "user_api_key_user_email": f"'{EMAIL}'"} + } + assert f"token='{FILTERED}'" in frame_vars["valid_token"] + assert frame_vars["general_settings"] == {"master_key": FILTERED, "database_url": FILTERED} + assert frame_vars["raw_headers"] == {"authorization": FILTERED, "x-api-key": FILTERED, "content-type": "'application/json'"} + assert MASTER_KEY not in serialized + assert VIRTUAL_KEY not in serialized + assert "db-password-under-test" not in serialized + + +def test_transaction_events_are_scrubbed_too() -> None: + transport: Final = RecordingTransport() + client: Final = sentry_sdk.Client(transport=transport, **build_sentry_init_options(PII_OFF)) + client.capture_event( + { + "type": "transaction", + "transaction": "/user/info", + "contexts": {"trace": {"trace_id": "a" * 32, "span_id": "b" * 16}}, + "spans": [{"description": f"lookup {EMAIL} by {KEY_HASH}", "span_id": "c" * 16, "trace_id": "a" * 32}], + } + ) + assert transport.last_envelope is not None + serialized: Final = json.dumps(transport.last_envelope.items[0].payload.json) + assert EMAIL not in serialized + assert KEY_HASH not in serialized + assert f"lookup {FILTERED} by {FILTERED}" in serialized + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ( + "UserAPIKeyAuth(token='abc', key_alias='team-a', user_id=None)", + f"UserAPIKeyAuth(token='{FILTERED}', key_alias='team-a', user_id=None)", + ), + ('{"api_key": "sk-1", "model": "gpt-5"}', f'{{"api_key": "{FILTERED}", "model": "gpt-5"}}'), + ("{'user_id': 'u-1', 'max_budget': 5}", f"{{'user_id': '{FILTERED}', 'max_budget': 5}}"), + ("Config(OPENAI_API_KEY=sk-live, timeout=10)", f"Config(OPENAI_API_KEY='{FILTERED}', timeout=10)"), + ("lookup for somebody@example.com failed", f"lookup for {FILTERED} failed"), + (f"hashed key {KEY_HASH} not found", f"hashed key {FILTERED} not found"), + ("request id 0123456789abcdef0123456789abcdef stays", "request id 0123456789abcdef0123456789abcdef stays"), + ("monkey=banana", "monkey=banana"), + ( + "{'x-api-key': 'k-1', 'cookie': 'session=abc', 'content-type': 'application/json'}", + f"{{'x-api-key': '{FILTERED}', 'cookie': '{FILTERED}', 'content-type': 'application/json'}}", + ), + ( + "headers={'x-tenant-key': 'sk-custom-header-key-0123456789'} key_name='sk-...6789'", + f"headers={{'x-tenant-key': '{FILTERED}'}} key_name='sk-...6789'", + ), + ( + "master_key={'value': 'not-a-litellm-key'} timeout=10", + f"master_key='{FILTERED}' timeout=10", + ), + ( + "credentials=[{'value': ('deep', 'secret')}], model='gpt-5'", + f"credentials='{FILTERED}', model='gpt-5'", + ), + ], +) +def test_string_scrubber_rewrites_field_and_value_forms(text: str, expected: str) -> None: + assert build_string_scrubber(send_default_pii=False)(text) == expected + + +def test_bare_key_floor_follows_the_custom_key_minimum() -> None: + scrub: Final = build_string_scrubber(send_default_pii=False) + shortest_key: Final = "sk-" + "a" * (MINIMUM_CUSTOM_KEY_LENGTH - len("sk-")) + assert scrub(f"label={shortest_key} model=gpt-5") == f"label={FILTERED} model=gpt-5" + assert scrub(f"label={shortest_key[:-1]} model=gpt-5") == f"label={shortest_key[:-1]} model=gpt-5" + + +def test_key_pattern_floor_never_exceeds_a_generated_key() -> None: + generated_key: Final = "sk-" + secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY) + stricter_custom_minimum: Final = len(generated_key) + 10 + assert build_key_pattern(stricter_custom_minimum, LENGTH_OF_LITELLM_GENERATED_KEY).fullmatch(generated_key) + assert build_key_pattern(stricter_custom_minimum, LENGTH_OF_LITELLM_GENERATED_KEY).fullmatch(generated_key[:-1]) is None + + +def test_json_walk_fails_closed_past_the_depth_cap() -> None: + scrub: Final = build_string_scrubber(send_default_pii=False) + nested: Final = reduce(lambda inner, _: [inner], range(MAX_SCRUB_DEPTH + 1), cast("JsonValue", "api_key=sk-1")) + assert FILTERED in json.dumps(scrub_json_strings(nested, scrub)) + assert "sk-1" not in json.dumps(scrub_json_strings(nested, scrub)) + assert scrub_json_strings([["api_key=sk-1"]], scrub) == [[f"api_key='{FILTERED}'"]] + + +def test_string_scrubber_with_pii_on_only_scrubs_secrets() -> None: + scrub: Final = build_string_scrubber(send_default_pii=True) + assert scrub(f"user_id='{EMAIL}', token='{KEY_HASH}', email {EMAIL} hash {KEY_HASH}") == ( + f"user_id='{EMAIL}', token='{FILTERED}', email {EMAIL} hash {KEY_HASH}" + ) + assert scrub(f"headers={{'authorization': 'Bearer {VIRTUAL_KEY}'}} sent {VIRTUAL_KEY}") == ( + f"headers={{'authorization': '{FILTERED}'}} sent {FILTERED}" + ) + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({}, False), + ({"SENTRY_SEND_DEFAULT_PII": "true"}, True), + ({"SENTRY_SEND_DEFAULT_PII": "True"}, True), + ({"SENTRY_SEND_DEFAULT_PII": "false"}, False), + ({"SENTRY_SEND_DEFAULT_PII": "yes please"}, False), + ], +) +def test_send_default_pii_comes_from_the_environment(env: Mapping[str, str], expected: bool) -> None: + assert build_sentry_init_options(env)["send_default_pii"] is expected + + +def test_init_options_read_dsn_rates_and_environment() -> None: + options: Final = build_sentry_init_options( + { + "SENTRY_DSN": "https://key@sentry.example/7", + "SENTRY_API_TRACE_RATE": "0.25", + "SENTRY_API_SAMPLE_RATE": "0.5", + "SENTRY_ENVIRONMENT": "staging", + } + ) + assert options["dsn"] == "https://key@sentry.example/7" + assert options["traces_sample_rate"] == 0.25 + assert options["sample_rate"] == 0.5 + assert options["environment"] == "staging" + assert options["event_scrubber"].recursive is True + + +def test_init_options_defaults() -> None: + options: Final = build_sentry_init_options({}) + assert options["dsn"] is None + assert options["traces_sample_rate"] == 1.0 + assert options["sample_rate"] == 1.0 + assert options["environment"] == "production" diff --git a/uv.lock b/uv.lock index c235171ecb2..527f53bd372 100644 --- a/uv.lock +++ b/uv.lock @@ -4743,6 +4743,7 @@ proxy-dev = [ { name = "opentelemetry-sdk" }, { name = "prisma" }, { name = "prometheus-client" }, + { name = "sentry-sdk" }, ] [package.metadata] @@ -4956,6 +4957,7 @@ proxy-dev = [ { name = "opentelemetry-sdk", specifier = "==1.33.1" }, { name = "prisma", specifier = "==0.11.0" }, { name = "prometheus-client", specifier = "==0.20.0" }, + { name = "sentry-sdk", specifier = "==2.21.0" }, ] [[package]] From b6fcd03848d9245b415a7fb5a92ce82d7770da7b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:18:03 -0700 Subject: [PATCH 21/58] fix(bedrock): surface a converse-stream 200 that decodes to no events as a 502 instead of an empty turn (#43213) * fix(bedrock): surface a converse-stream 200 that decodes to no events as a 502 instead of an empty turn * fix(bedrock): quote the body head only when a stream decoded no events The leftover-bytes error keeps the byte and event counts, the content type and the request id but no longer quotes the first bytes of a stream that already decoded events, since that head is the start of a healthy stream and can hold model output. The anthropic_messages empty-stream warning no longer prints the request's model name. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 3 + litellm/llms/bedrock/chat/converse_handler.py | 4 +- litellm/llms/bedrock/chat/invoke_handler.py | 141 ++++++++++++------ .../anthropic_claude3_transformation.py | 4 +- .../test_litellm_logging.py | 13 ++ .../llms/bedrock/chat/test_invoke_handler.py | 126 ++++++++++++++++ 6 files changed, 243 insertions(+), 48 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d8182140a17..54ed9171dc8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4213,6 +4213,9 @@ class Logging(LiteLLMLoggingBaseClass): json_mode=False, litellm_params={}, ) + elif result is None: + verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") + return litellm.ModelResponse(model=self.model) else: from litellm.types.llms.anthropic import AnthropicResponse diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index bd358805743..65a34f72167 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -69,7 +69,9 @@ def make_sync_call( completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: decoder: Final = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index c7b4018b80b..93804e20041 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1,6 +1,6 @@ import types -from collections.abc import AsyncIterator, Iterator -from typing import Final, cast +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import TYPE_CHECKING, Final, cast import httpx from pydantic import TypeAdapter @@ -51,7 +51,11 @@ from ..common_utils import ( bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_memory=50, default_ttl=600) from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +if TYPE_CHECKING: + from botocore.eventstream import EventStreamMessage + converse_config: Final = AmazonConverseConfig() +_STREAM_HEAD_BYTES: Final = 200 NOVA_INVOKE_STREAM_EVENT_TYPES: Final = ( "messageStart", "contentBlockStart", @@ -162,6 +166,22 @@ class AmazonCohereChatConfig: return optional_params +def _stream_decoder( + bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None, + *, + model: str, + json_mode: bool | None, + sync_stream: bool, +) -> "AWSEventStreamDecoder": + if bedrock_invoke_provider == "anthropic": + return AmazonAnthropicClaudeStreamDecoder(model=model, sync_stream=sync_stream, json_mode=json_mode) + if bedrock_invoke_provider == "deepseek_r1": + return AmazonDeepSeekR1StreamDecoder(model=model, sync_stream=sync_stream) + if bedrock_invoke_provider == "moonshot": + return AmazonOpenAICompatibleStreamDecoder(model=model, sync_stream=sync_stream) + return AWSEventStreamDecoder(model=model, json_mode=json_mode) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -218,28 +238,13 @@ async def make_call( completion_stream: MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] = ( MockResponseIterator(model_response=model_response, json_mode=json_mode) ) - elif bedrock_invoke_provider == "anthropic": - decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( - model=model, - sync_stream=False, - json_mode=json_mode, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "deepseek_r1": - decoder = AmazonDeepSeekR1StreamDecoder( - model=model, - sync_stream=False, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "moonshot": - decoder = AmazonOpenAICompatibleStreamDecoder( - model=model, - sync_stream=False, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) else: - decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) + decoder: Final = _stream_decoder( + bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=False + ) + completion_stream = decoder.aiter_bytes( + response.aiter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( @@ -322,28 +327,13 @@ def make_sync_call( completion_stream: MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] = ( MockResponseIterator(model_response=model_response, json_mode=json_mode) ) - elif bedrock_invoke_provider == "anthropic": - decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( - model=model, - sync_stream=True, - json_mode=json_mode, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "deepseek_r1": - decoder = AmazonDeepSeekR1StreamDecoder( - model=model, - sync_stream=True, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "moonshot": - decoder = AmazonOpenAICompatibleStreamDecoder( - model=model, - sync_stream=True, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) else: - decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + decoder: Final = _stream_decoder( + bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=True + ) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( @@ -370,6 +360,49 @@ def make_sync_call( raise BedrockError(status_code=500, message=str(e)) +def _response_header(response_headers: Mapping[str, str] | None, name: str) -> str | None: + return None if response_headers is None else response_headers.get(name) + + +class _EventStreamTally: + def __init__(self) -> None: + self.bytes_received = 0 + self.bytes_decoded = 0 + self.events = 0 + self.head = b"" + + def add_chunk(self, chunk: bytes) -> None: + self.bytes_received += len(chunk) + if len(self.head) < _STREAM_HEAD_BYTES: + self.head = (self.head + chunk)[:_STREAM_HEAD_BYTES] + + def add_event(self, event: "EventStreamMessage") -> None: + self.events += 1 + self.bytes_decoded += event.prelude.total_length + + def undecoded_stream_error(self, response_headers: Mapping[str, str] | None) -> BedrockError | None: + undecoded: Final = self.bytes_received - self.bytes_decoded + if self.events and not undecoded: + return None + detail: Final = ( + f"content-type={_response_header(response_headers, 'content-type')!r}, " + f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " + f"{self.bytes_received} bytes received" + ) + if not self.events: + return BedrockError( + status_code=502, + message=( + "Bedrock answered the stream with HTTP 200 but its body decoded to no events " + f"({detail}, first bytes={self.head!r})" + ), + ) + return BedrockError( + status_code=502, + message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + ) + + class AWSEventStreamDecoder: def __init__(self, model: str, json_mode: bool | None = False) -> None: from botocore.parsers import EventStreamJSONParser @@ -709,32 +742,48 @@ class AWSEventStreamDecoder: tool_use=None, ) - def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[GChunk | ModelResponseStream | dict]: + def iter_bytes( + self, iterator: Iterator[bytes], *, response_headers: Mapping[str, str] | None = None + ) -> Iterator[GChunk | ModelResponseStream | dict]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() + tally: Final = _EventStreamTally() for chunk in iterator: event_stream_buffer.add_data(chunk) + tally.add_chunk(chunk) for event in event_stream_buffer: + tally.add_event(event) message = self._parse_message_from_event(event) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) + if undecoded_stream_error is not None: + raise undecoded_stream_error - async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[GChunk | ModelResponseStream | dict]: + async def aiter_bytes( + self, iterator: AsyncIterator[bytes], *, response_headers: Mapping[str, str] | None = None + ) -> AsyncIterator[GChunk | ModelResponseStream | dict]: """Given an async iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() + tally: Final = _EventStreamTally() async for chunk in iterator: event_stream_buffer.add_data(chunk) + tally.add_chunk(chunk) for event in event_stream_buffer: + tally.add_event(event) message = self._parse_message_from_event(event) if message: _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) + if undecoded_stream_error is not None: + raise undecoded_stream_error def _parse_message_from_event(self, event) -> str | None: response_stream_shape: Final = get_bedrock_response_stream_shape() diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 14bd2bee6cf..cefc8afed25 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -770,7 +770,9 @@ class AmazonAnthropicClaudeMessagesConfig( aws_decoder: Final = AmazonAnthropicClaudeMessagesStreamDecoder( model=model, ) - completion_stream: Final = aws_decoder.aiter_bytes(httpx_response.aiter_bytes()) + completion_stream: Final = aws_decoder.aiter_bytes( + httpx_response.aiter_bytes(), response_headers=httpx_response.headers + ) # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. return self.bedrock_sse_wrapper( completion_stream=completion_stream, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c4829ced9f3..166eeb53f5f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5308,6 +5308,19 @@ def test_handle_anthropic_messages_response_logging_passes_model_response_throug assert logging_obj._handle_anthropic_messages_response_logging(result=model_response) is model_response +def test_anthropic_messages_logged_response_tolerates_a_stream_that_assembled_nothing(): + """A /v1/messages stream whose upstream yielded no chunks assembles to None; the spend + row must still land under the message id the caller was served instead of crashing.""" + logging_obj = _anthropic_messages_logging_obj() + logging_obj.record_streamed_anthropic_message_id("msg_served") + + result = logging_obj._anthropic_messages_logged_response(result=None) + + assert isinstance(result, ModelResponse) + assert result.id == "msg_served" + assert result.model == "openai/my-local" + + def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_responses_payload(): """If the Responses translation raises (eg. empty output on an incomplete response), the row must still land: a minimal ModelResponse with model + usage is returned.""" diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 466e9b4fda8..ed8b7023977 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -1,5 +1,6 @@ import base64 import binascii +import itertools import datetime import json import struct @@ -14,10 +15,13 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( + AmazonOpenAICompatibleStreamDecoder, AWSEventStreamDecoder, make_call, make_sync_call, ) +from litellm.exceptions import MidStreamFallbackError +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.utils import ModelResponseStream @@ -799,3 +803,125 @@ async def test_moonshot_invoke_async_stream_yields_openai_shaped_chunks(_aws_tes ) _assert_moonshot_stream_content([chunk async for chunk in stream]) + + +def _truncated_frame() -> bytes: + return _bedrock_event_stream_frame(_openai_stream_chunk({"role": "assistant"}))[:-8] + + +def _event_stream_headers() -> httpx.Headers: + return httpx.Headers({"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": "req-empty-1"}) + + +_UNDECODABLE_STREAM_BODIES: Final = ( + pytest.param(b"", id="empty"), + pytest.param(b"\x00\x00\x00\x05", id="shorter-than-a-prelude"), + pytest.param(_truncated_frame(), id="truncated-first-message"), +) + + +def _assert_no_events_error(error: BedrockError, body: bytes) -> None: + assert error.status_code == 502 + assert "HTTP 200" in error.message + assert "decoded to no events" in error.message + assert f"{len(body)} bytes received" in error.message + assert "application/vnd.amazon.eventstream" in error.message + assert "req-empty-1" in error.message + assert f"first bytes={body[:200]!r}" in error.message + + +@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES) +def test_iter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_event_stream_headers())) + + _assert_no_events_error(exc_info.value, body) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES) +async def test_aiter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_event_stream_headers())] + + _assert_no_events_error(exc_info.value, body) + + +def test_iter_bytes_raises_when_the_stream_ends_mid_message() -> None: + decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) + stream: Final = decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM, _truncated_frame()])) + + chunks: Final = list(itertools.islice(stream, 4)) + with pytest.raises(BedrockError) as exc_info: + next(stream) + + _assert_moonshot_stream_content(chunks) + assert exc_info.value.status_code == 502 + assert f"{len(_truncated_frame())} undecoded bytes after 4 events" in exc_info.value.message + assert "first bytes=" not in exc_info.value.message + + +def test_iter_bytes_yields_a_complete_stream_without_raising() -> None: + decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) + + chunks: Final = list(decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM[:100], _MOONSHOT_RAW_STREAM[100:]]))) + + _assert_moonshot_stream_content(chunks) + + +def _assert_empty_stream_surfaced_as_bad_gateway(error: MidStreamFallbackError) -> None: + assert error.status_code == 502 + assert error.is_pre_first_chunk is True + assert isinstance(error.original_exception, litellm.BadGatewayError) + assert "decoded to no events" in str(error) + assert "req-empty-1" in str(error) + + +def test_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn(_aws_test_credentials: None) -> None: + response: Final = MagicMock(status_code=200, headers=_event_stream_headers()) + response.iter_bytes = lambda chunk_size=None: iter([b""]) + client: Final = HTTPHandler() + client.post = MagicMock(return_value=response) + + with pytest.raises(MidStreamFallbackError) as exc_info: + list( + litellm.completion( + model="bedrock/us.moonshotai.kimi-k3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) + ) + + _assert_empty_stream_surfaced_as_bad_gateway(exc_info.value) + + +@pytest.mark.asyncio +async def test_async_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn( + _aws_test_credentials: None, +) -> None: + async def _aiter_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]: + yield b"" + + response: Final = MagicMock(status_code=200, headers=_event_stream_headers()) + response.aiter_bytes = _aiter_bytes + client: Final = AsyncHTTPHandler() + client.post = AsyncMock(return_value=response) + + stream: Final = await litellm.acompletion( + model="bedrock/us.moonshotai.kimi-k3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + _ = [chunk async for chunk in stream] + + _assert_empty_stream_surfaced_as_bad_gateway(exc_info.value) From 7b4fd47c6e665e36ee0669ae7d15887a085dd472 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:28:14 -0700 Subject: [PATCH 22/58] fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim (#43206) * fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(jwt): describe header team selection under fallback_to_db_teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 11 +- litellm/proxy/auth/handle_jwt.py | 58 +++---- .../proxy/auth/test_handle_jwt.py | 141 +++++++++++++++++- 3 files changed, 175 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 89fa0058644..fd13697dab3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5355,11 +5355,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=False, description=( "When True, users whose JWT contains no team claims are authenticated " - "using their database team memberships instead of receiving HTTP 403. " - "Usage is attributed to the user's first resolvable DB team, or to the " - "team specified via the x-litellm-team-id request header (validated " - "against DB membership). Requires user_id_upsert=True so that user " - "records exist before the fallback runs." + "using their database team memberships instead of receiving HTTP 403, " + "with usage attributed to the user's first resolvable DB team. Whether or " + "not the JWT carries team claims, the x-litellm-team-id request header may " + "select any team the user is a member of in the database (validated against " + "DB membership); without the header the JWT team stays the default. Requires " + "user_id_upsert=True so that user records exist before the fallback runs." ), ) issuers: list[JWTIssuerConfig] | None = Field( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index e07b20fd5d5..4f41b283a33 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1930,12 +1930,12 @@ class JWTAuthManager: ) -> HeaderTeam | None: """ The team named by x-litellm-team-id, which may carry a team id or a team - alias. A value that is already an allowed team id (or, under the DB - fallback, an existing team id) never costs an alias lookup; an alias is - accepted only when the team it names would have been accepted by id. - Under the DB fallback only a team row that is provably absent falls - through to the alias lookup; a read that failed for any other reason - keeps the membership denial the id path already gives. + alias. A value that is already an allowed team id never costs a lookup; + under the DB fallback any other value is accepted provisionally, by id + or alias, for the membership check auth_builder runs later. Under the + DB fallback only a team row that is provably absent falls through to + the alias lookup; a read that failed for any other reason keeps the + membership denial the id path already gives. Raises: HTTPException: 403 when neither the value nor the team it aliases is @@ -1948,7 +1948,11 @@ class JWTAuthManager: if not header_value: return None - if fallback_to_db_teams and not allowed_team_ids: + if header_value in allowed_team_ids: + verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value) + return HeaderTeam(header_value=header_value, team_id=header_value) + + if fallback_to_db_teams: try: await get_team_object( team_id=header_value, @@ -1969,10 +1973,6 @@ class JWTAuthManager: JWTAuthManager._raise_header_team_membership_denial(header_value) return HeaderTeam(header_value=header_value, team_id=header_value) - if header_value in allowed_team_ids: - verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value) - return HeaderTeam(header_value=header_value, team_id=header_value) - team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias( header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) @@ -2353,9 +2353,9 @@ class JWTAuthManager: header_value: str, ) -> None: """ - A provisional team_id from the x-litellm-team-id header (accepted without - JWT-team validation when the JWT carries no team claims) must exist in the - user's DB team memberships before it becomes request context. The denial + A provisional team_id from the x-litellm-team-id header (accepted under + fallback_to_db_teams because it is outside the JWT's teams) must exist in + the user's DB team memberships before it becomes request context. The denial names `header_value`, the id or alias the caller sent, not `team_id`. """ user_team_ids: Final = user_object.teams if user_object else [] @@ -2587,22 +2587,30 @@ class JWTAuthManager: if specific_team_id and not db_team_fallback: all_team_ids.add(specific_team_id) + header_db_fallback: Final = handler.litellm_jwtauth.fallback_to_db_teams and team_id is None + header_team: Final = await JWTAuthManager.resolve_team_from_header( request_headers=request_headers, allowed_team_ids=all_team_ids, - fallback_to_db_teams=db_team_fallback, + fallback_to_db_teams=header_db_fallback, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + provisional_header_team: Final = ( + header_team + if header_team is not None and header_db_fallback and header_team.team_id not in all_team_ids + else None + ) if header_team: team_id = header_team.team_id - # A provisional header team (accepted only because the JWT carries no - # team claims) is validated against DB membership further down; never - # upsert it here or an attacker-supplied x-litellm-team-id would create - # an orphaned team row before that check runs. A genuine membership team - # already exists, so suppressing the upsert in that case costs nothing. + # A provisional header team (accepted because it is outside the + # JWT's teams under fallback_to_db_teams) is validated against DB + # membership further down; never upsert it here or an + # attacker-supplied x-litellm-team-id would create an orphaned team + # row before that check runs. A genuine membership team already + # exists, so suppressing the upsert in that case costs nothing. try: team_object = await get_team_object( team_id=team_id, @@ -2610,10 +2618,10 @@ class JWTAuthManager: user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, - team_id_upsert=(team_id_upsert and not db_team_fallback), + team_id_upsert=(team_id_upsert and provisional_header_team is None), ) except HTTPException: - if not db_team_fallback: + if provisional_header_team is None: raise JWTAuthManager._raise_header_team_membership_denial(header_team.header_value) elif not team_id and not db_team_fallback: @@ -2756,11 +2764,11 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - elif db_team_fallback and header_team is not None and team_id == header_team.team_id: + elif provisional_header_team is not None and team_id == provisional_header_team.team_id: JWTAuthManager._validate_header_team_in_db_membership( team_id=team_id, user_object=user_object, - header_value=header_team.header_value, + header_value=provisional_header_team.header_value, ) if not JWTAuthManager._is_team_route_allowed( route=route, @@ -2770,7 +2778,7 @@ class JWTAuthManager: raise HTTPException( status_code=403, detail=( - f"Team '{header_team.header_value}' (from x-litellm-team-id header) " + f"Team '{provisional_header_team.header_value}' (from x-litellm-team-id header) " f"is not allowed to access route '{route}'." ), ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index f8b9043a23f..b1622e0dff0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -5324,16 +5324,22 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym @pytest.mark.asyncio -async def test_resolve_team_from_header_defers_to_db_membership_only_without_jwt_claims(): +async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fallback_even_with_jwt_claims(): """With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing - team is accepted provisionally only when the JWT carries no team claims (allowed - set empty). When the JWT does carry team claims, the header must still be validated - against them, and the flag-off behavior must keep rejecting unknown teams.""" + team is accepted provisionally whether or not the JWT carries team claims; the + union of JWT teams and DB memberships is enforced by auth_builder's later + membership check. Unknown values still 403, and the flag-off behavior keeps + rejecting teams outside the JWT's allowed set.""" known_ids = frozenset({"team-from-db"}) deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404) assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + deferred_with_claims, _, _ = await _resolve_header( + "team-from-db", {"team-1"}, True, _teams_by_id(known_ids), _team_alias_lookup_404 + ) + assert deferred_with_claims == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + with pytest.raises(HTTPException) as exc_info: await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404) assert exc_info.value.status_code == 403 @@ -5849,6 +5855,7 @@ async def _run_auth_builder_with_header_team( allowed_team_ids: set, fake_get_team_by_alias=_team_alias_lookup_404, route: str = "/chat/completions", + send_header: bool = True, ): jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = jwt_auth_config @@ -5909,7 +5916,7 @@ async def _run_auth_builder_with_header_team( user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, - request_headers={"x-litellm-team-id": header_team_id}, + request_headers={"x-litellm-team-id": header_team_id} if send_header else {}, ) @@ -7283,6 +7290,130 @@ async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowe assert allowed["team_id"] == "team_member" +@pytest.mark.asyncio +async def test_auth_builder_header_selects_db_membership_team_when_jwt_also_carries_a_team_claim() -> None: + """Under fallback_to_db_teams, x-litellm-team-id may name a DB-membership + team the JWT does not claim (LIT-8656): the allowed set is the JWT teams + union the user's DB memberships, not the JWT teams alone. The flag-off + path keeps rejecting the same header against the JWT's allowed teams.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member"})) + + by_membership = await _run_auth_builder_with_header_team( + config, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert by_membership["team_id"] == "team_member" + assert by_membership["team_object"].team_id == "team_member" + + by_claim = await _run_auth_builder_with_header_team( + config, token, "team_claimed", user_object, fake_get_team, {"team_claimed"} + ) + assert by_claim["team_id"] == "team_claimed" + + flag_off = LiteLLM_JWTAuth(fallback_to_db_teams=False, team_id_jwt_field="appid") + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + flag_off, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert exc_info.value.status_code == 403 + assert "JWT's allowed teams" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_header_non_member_team_is_denied_when_jwt_also_carries_a_team_claim() -> None: + """A header naming a team the user does not belong to stays a membership + denial even when the JWT carries a team claim, and an existing but + non-member team produces the exact same 403 shape as a nonexistent one so + the response is no oracle for which team ids exist.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member", "team_other"})) + + with pytest.raises(HTTPException) as outsider_exc: + await _run_auth_builder_with_header_team( + config, token, "team_other", user_object, fake_get_team, {"team_claimed"} + ) + with pytest.raises(HTTPException) as missing_exc: + await _run_auth_builder_with_header_team( + config, token, "team_ghost", user_object, fake_get_team, {"team_claimed"} + ) + + assert outsider_exc.value.status_code == 403 + assert missing_exc.value.status_code == 403 + assert outsider_exc.value.detail == ( + "x-litellm-team-id 'team_other' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + assert missing_exc.value.detail.replace("team_ghost", "") == outsider_exc.value.detail.replace( + "team_other", "" + ) + assert "exist" not in missing_exc.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_no_header_keeps_the_jwt_team_when_fallback_to_db_teams_is_on() -> None: + """With no x-litellm-team-id header, fallback_to_db_teams must not disturb + the claim path: the JWT's own team claim still binds the request.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + + result = await _run_auth_builder_with_header_team( + config, + token, + "team_member", + user_object, + _teams_by_id(frozenset({"team_claimed", "team_member"})), + {"team_claimed"}, + send_header=False, + ) + assert result["team_id"] == "team_claimed" + + +@pytest.mark.asyncio +async def test_auth_builder_team_id_default_does_not_widen_the_header_allowed_set() -> None: + """team_id_default fills in a team for claimless tokens but must not widen + the header's allowed set: a header naming the default team is still held + to DB membership under fallback_to_db_teams.""" + user_object = LiteLLM_UserTable( + user_id="u_default", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_default="team_default") + token = {"sub": "u_default", "scope": ""} + + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + config, + token, + "team_default", + user_object, + _teams_by_id(frozenset({"team_default", "team_member"})), + set(), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == ( + "x-litellm-team-id 'team_default' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag(): """Reading the singular team claim during sync is scoped to fallback_to_db_teams. From 191305e6d41ec1cf9496e10505fea9c4afe553b0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:29:23 -0700 Subject: [PATCH 23/58] feat(integrations): add Databricks Zerobus trace logging callback (#42013) * feat(integrations): add Databricks Zerobus trace logging callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(zerobus): escape regex in pytest.raises match Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(zerobus): use unique test module basenames Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(zerobus): hold the queue cap while an insert is in flight Rows arriving during a slow insert are dropped once the queue is at max_queue_size, since trimming the head would corrupt the in-flight batch. Test fakes are typed and record calls as frozen dataclasses; the litellm_logging init and reuse branches are covered. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(zerobus): type the row payload and dashboard config helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(zerobus): assert the trace row survives a JSON round trip Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(zerobus): keep the client secret and access token out of dataclass reprs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 3 + litellm/integrations/callback_configs.json | 39 ++ litellm/integrations/zerobus/__init__.py | 5 + litellm/integrations/zerobus/client.py | 161 +++++++ litellm/integrations/zerobus/logger.py | 230 ++++++++++ litellm/integrations/zerobus/row.py | 156 +++++++ .../custom_logger_registry.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 13 + litellm/proxy/_types.py | 12 + litellm/types/integrations/zerobus.py | 53 +++ .../zerobus/test_zerobus_client.py | 258 ++++++++++++ .../zerobus/test_zerobus_logger.py | 392 ++++++++++++++++++ .../integrations/zerobus/test_zerobus_row.py | 139 +++++++ .../proxy/test_zerobus_dashboard_config.py | 52 +++ .../src/components/callback_info_helpers.tsx | 15 + 15 files changed, 1530 insertions(+) create mode 100644 litellm/integrations/zerobus/__init__.py create mode 100644 litellm/integrations/zerobus/client.py create mode 100644 litellm/integrations/zerobus/logger.py create mode 100644 litellm/integrations/zerobus/row.py create mode 100644 litellm/types/integrations/zerobus.py create mode 100644 tests/test_litellm/integrations/zerobus/test_zerobus_client.py create mode 100644 tests/test_litellm/integrations/zerobus/test_zerobus_logger.py create mode 100644 tests/test_litellm/integrations/zerobus/test_zerobus_row.py create mode 100644 tests/test_litellm/proxy/test_zerobus_dashboard_config.py diff --git a/litellm/__init__.py b/litellm/__init__.py index e334fbe8ca8..676c735b9e8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -50,6 +50,7 @@ from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams from litellm.litellm_core_utils.core_helpers import drop_params_env_flag from litellm.types.integrations.pointfive import PointFiveInitParams +from litellm.types.integrations.zerobus import ZerobusInitParams from litellm._logging import ( set_verbose, _turn_on_debug, @@ -157,6 +158,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "deepeval", "s3_v2", "pointfive", + "zerobus", "aws_sqs", "vector_store_pre_call_hook", "dotprompt", @@ -442,6 +444,7 @@ datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] datadog_params: Optional[Union[DatadogInitParams, Dict]] = None newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None pointfive_params: Optional[Union[PointFiveInitParams, Mapping[str, object]]] = None +zerobus_params: Optional[Union[ZerobusInitParams, Mapping[str, object]]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 5bd8aca55fa..4e72075dc5c 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -406,6 +406,45 @@ }, "description": "PointFive Logging Integration" }, + { + "id": "zerobus", + "displayName": "Databricks Zerobus", + "logo": "databricks.svg", + "supports_key_team_logging": false, + "dynamic_params": { + "ZEROBUS_WORKSPACE_URL": { + "type": "text", + "ui_name": "Workspace URL", + "description": "Databricks workspace URL, e.g. https://dbc-a1b2c3d4-e5f6.cloud.databricks.com", + "required": true + }, + "ZEROBUS_SERVER_ENDPOINT": { + "type": "text", + "ui_name": "Zerobus Endpoint", + "description": "Zerobus ingest endpoint, e.g. https://.zerobus..cloud.databricks.com", + "required": true + }, + "ZEROBUS_CLIENT_ID": { + "type": "text", + "ui_name": "Service Principal Client ID", + "description": "OAuth client id of a service principal with USE CATALOG, USE SCHEMA, SELECT and MODIFY on the table", + "required": true + }, + "ZEROBUS_CLIENT_SECRET": { + "type": "password", + "ui_name": "Service Principal Client Secret", + "description": "OAuth client secret of the service principal", + "required": true + }, + "ZEROBUS_TABLE_NAME": { + "type": "text", + "ui_name": "Table", + "description": "Fully qualified Unity Catalog table, catalog.schema.table, created with the LiteLLM trace schema", + "required": true + } + }, + "description": "Databricks Zerobus Ingest Logging Integration" + }, { "id": "s3", "displayName": "S3", diff --git a/litellm/integrations/zerobus/__init__.py b/litellm/integrations/zerobus/__init__.py new file mode 100644 index 00000000000..b1f5bc2ca40 --- /dev/null +++ b/litellm/integrations/zerobus/__init__.py @@ -0,0 +1,5 @@ +"""Databricks Zerobus logging integration for LiteLLM.""" + +from litellm.integrations.zerobus.logger import ZerobusLogger + +__all__ = ("ZerobusLogger",) diff --git a/litellm/integrations/zerobus/client.py b/litellm/integrations/zerobus/client.py new file mode 100644 index 00000000000..bf3a9e3e269 --- /dev/null +++ b/litellm/integrations/zerobus/client.py @@ -0,0 +1,161 @@ +""" +Writes rows to a Unity Catalog table through the Zerobus Ingest REST API. + +Zerobus only accepts a Databricks OAuth token minted for its own resource and scoped to +the target table's privileges, so the client mints that token itself with the service +principal's client credentials and reuses it until shortly before it expires. +""" + +import asyncio +import base64 +import json +import time +from collections.abc import Callable, Mapping, Sequence +from typing import Final + +import httpx +from pydantic import BaseModel, ValidationError + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.integrations.zerobus import ( + RETRYABLE_INGEST_STATUS_CODES, + TOKEN_REFRESH_LEEWAY_SECONDS, + ZerobusAccessToken, + ZerobusConnection, + ZerobusIngestFailure, +) + +TOKEN_PATH: Final = "/oidc/v1/token" +OAUTH_SCOPE: Final = "all-apis" + + +class _TokenResponse(BaseModel): + access_token: str + expires_in: float = 3600 + + +class ZerobusIngestError(Exception): + """A batch could not be written and the failure is worth retrying.""" + + +def zerobus_resource(workspace_id: str) -> str: + return f"api://databricks/workspaces/{workspace_id}/zerobusDirectWriteApi" + + +def authorization_details(table_name: str) -> str: + """The Unity Catalog privileges Zerobus requires the token to carry, as the token endpoint expects them.""" + catalog, schema, _table = table_name.split(".", 2) + return json.dumps( + ( + { + "type": "unity_catalog_privileges", + "privileges": ("USE CATALOG",), + "object_type": "CATALOG", + "object_full_path": catalog, + }, + { + "type": "unity_catalog_privileges", + "privileges": ("USE SCHEMA",), + "object_type": "SCHEMA", + "object_full_path": f"{catalog}.{schema}", + }, + { + "type": "unity_catalog_privileges", + "privileges": ("SELECT", "MODIFY"), + "object_type": "TABLE", + "object_full_path": table_name, + }, + ) + ) + + +def insert_url(connection: ZerobusConnection) -> str: + return f"{connection.server_endpoint.rstrip('/')}/zerobus/v1/tables/{connection.table_name}/insert" + + +def token_url(connection: ZerobusConnection) -> str: + return f"{connection.workspace_url.rstrip('/')}{TOKEN_PATH}" + + +def _basic_auth(client_id: str, client_secret: str) -> str: + return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + + +def _status_failure(what: str, error: httpx.HTTPStatusError) -> ZerobusIngestFailure: + status: Final = error.response.status_code + return ZerobusIngestFailure( + detail=f"{what} returned {status}: {error.response.text}"[:500], + retryable=status in RETRYABLE_INGEST_STATUS_CODES, + ) + + +class ZerobusIngestClient: + def __init__( + self, + connection: ZerobusConnection, + http_client: AsyncHTTPHandler, + clock: Callable[[], float] = time.time, + ) -> None: + self.connection: Final = connection + self.http_client: Final = http_client + self.clock: Final = clock + self._token: ZerobusAccessToken | None = None + self._token_lock: Final = asyncio.Lock() + + async def insert(self, rows: Sequence[Mapping[str, object]]) -> ZerobusIngestFailure | None: + """Write ``rows`` as one request. ``None`` means Zerobus accepted every row.""" + token: Final = await self.access_token() + if isinstance(token, ZerobusIngestFailure): + return token + try: + await self.http_client.post( + insert_url(self.connection), + content=json.dumps([dict(row) for row in rows]).encode(), + headers={"Content-Type": "application/json", "Authorization": f"Bearer {token.value}"}, + ) + except httpx.HTTPStatusError as error: + if error.response.status_code == 401: + self._token = None + return ZerobusIngestFailure(detail="insert returned 401, token discarded", retryable=True) + return _status_failure("insert", error) + except (httpx.HTTPError, litellm.Timeout) as error: + return ZerobusIngestFailure(detail=f"insert failed: {error}", retryable=True) + return None + + async def access_token(self) -> ZerobusAccessToken | ZerobusIngestFailure: + """The cached token while it has more than the leeway left, otherwise a fresh one.""" + async with self._token_lock: + cached: Final = self._token + if cached is not None and cached.expires_at - self.clock() > TOKEN_REFRESH_LEEWAY_SECONDS: + return cached + minted: Final = await self._mint_token() + if isinstance(minted, ZerobusAccessToken): + self._token = minted + return minted + + async def _mint_token(self) -> ZerobusAccessToken | ZerobusIngestFailure: + connection: Final = self.connection + try: + response: Final = await self.http_client.post( + token_url(connection), + data={ + "grant_type": "client_credentials", + "scope": OAUTH_SCOPE, + "resource": zerobus_resource(connection.workspace_id), + "authorization_details": authorization_details(connection.table_name), + }, + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "Authorization": _basic_auth(connection.client_id, connection.client_secret), + }, + ) + except httpx.HTTPStatusError as error: + return _status_failure("token request", error) + except (httpx.HTTPError, litellm.Timeout) as error: + return ZerobusIngestFailure(detail=f"token request failed: {error}", retryable=True) + try: + parsed: Final = _TokenResponse.model_validate_json(response.text) + except ValidationError as error: + return ZerobusIngestFailure(detail=f"token response was not understood: {error}", retryable=False) + return ZerobusAccessToken(value=parsed.access_token, expires_at=self.clock() + parsed.expires_in) diff --git a/litellm/integrations/zerobus/logger.py b/litellm/integrations/zerobus/logger.py new file mode 100644 index 00000000000..e2007218c8e --- /dev/null +++ b/litellm/integrations/zerobus/logger.py @@ -0,0 +1,230 @@ +"""Databricks Zerobus logging integration.""" + +import asyncio +from collections.abc import Mapping +from datetime import datetime +from typing import Final +from urllib.parse import urlsplit + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.zerobus.client import ZerobusIngestClient, ZerobusIngestError +from litellm.integrations.zerobus.row import trace_row +from litellm.litellm_core_utils.redact_messages import ( + redacted_standard_logging_payload, + should_redact_message_logging, +) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, httpxSpecialProvider +from litellm.secret_managers.main import get_secret_str +from litellm.types.integrations.zerobus import ZerobusConnection, ZerobusInitParams + +_ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def _resolved_secret(value: str | None) -> str | None: + """Resolve a config value that may name a secret; an unset ``os.environ/NAME`` stays unresolved.""" + if value is None: + return None + resolved: Final = get_secret_str(value) + if resolved: + return resolved + return None if value.startswith(_ENV_REFERENCE_PREFIX) else value + + +def _configured_params() -> ZerobusInitParams: + configured: Final = litellm.zerobus_params + if isinstance(configured, ZerobusInitParams): + return configured + if isinstance(configured, Mapping): + return ZerobusInitParams.model_validate(configured) + return ZerobusInitParams() + + +def _setting(configured: str | None, env_var: str) -> str: + """Prefer the configured value, falling back to the environment the proxy UI writes.""" + value: Final = _resolved_secret(configured) or get_secret_str(env_var) + if not value: + raise ValueError( + f"zerobus logging requires {env_var}. Set it in the environment, or " + f"litellm_settings.zerobus_params.{env_var.removeprefix('ZEROBUS_').lower()} in config.yaml" + ) + return value + + +def _workspace_id(server_endpoint: str) -> str: + """The Zerobus endpoint is ``https://.zerobus..``, so the id is its first label.""" + host: Final = urlsplit(server_endpoint).hostname or "" + workspace_id: Final = host.split(".", 1)[0] + if not workspace_id.isdigit(): + raise ValueError( + f"ZEROBUS_SERVER_ENDPOINT {server_endpoint!r} does not look like " + "https://.zerobus..cloud.databricks.com" + ) + return workspace_id + + +def _table_name(configured: str | None) -> str: + table_name: Final = _setting(configured, "ZEROBUS_TABLE_NAME") + if table_name.count(".") != 2: + raise ValueError(f"ZEROBUS_TABLE_NAME {table_name!r} must be fully qualified as catalog.schema.table") + return table_name + + +def connection_for(params: ZerobusInitParams) -> ZerobusConnection: + """The connection configured right now, so a UI edit takes effect without a restart.""" + server_endpoint: Final = _setting(params.server_endpoint, "ZEROBUS_SERVER_ENDPOINT") + return ZerobusConnection( + workspace_url=_setting(params.workspace_url, "ZEROBUS_WORKSPACE_URL"), + workspace_id=_workspace_id(server_endpoint), + server_endpoint=server_endpoint, + client_id=_setting(params.client_id, "ZEROBUS_CLIENT_ID"), + client_secret=_setting(params.client_secret, "ZEROBUS_CLIENT_SECRET"), + table_name=_table_name(params.table_name), + ) + + +class ZerobusLogger(CustomBatchLogger): + preserve_events_added_during_flush = True + + def __init__( + self, + params: ZerobusInitParams | None = None, + client: ZerobusIngestClient | None = None, + start_periodic_flush: bool = True, + ) -> None: + resolved: Final = params if params is not None else _configured_params() + self.params: Final = resolved + self.given_client: Final = client + self._cached_client: ZerobusIngestClient | None = None + if client is None: + connection_for(resolved) + super().__init__( + flush_lock=asyncio.Lock(), + batch_size=resolved.batch_size, + flush_interval=resolved.flush_interval, + turn_off_message_logging=bool(resolved.turn_off_message_logging), + ) + self._flushing: bool = False + self._batch_flush_task: asyncio.Task[None] | None = None + self._periodic_flush_task: asyncio.Task[None] | None = ( + self._start_periodic_flush_task() if start_periodic_flush else None + ) + + @property + def client(self) -> ZerobusIngestClient: + """A client for the current connection, kept while the connection is unchanged so its token is reused.""" + if self.given_client is not None: + return self.given_client + connection: Final = connection_for(self.params) + cached: Final = self._cached_client + if cached is not None and cached.connection == connection: + return cached + fresh: Final = ZerobusIngestClient( + connection=connection, + http_client=get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback), + ) + self._cached_client = fresh + return fresh + + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return None + return loop.create_task(self.periodic_flush()) + + def _start_batch_flush_task(self) -> None: + if self._batch_flush_task is not None and not self._batch_flush_task.done(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._batch_flush_task = loop.create_task(self.flush_queue(skip_if_flushing=True)) + + def _flush_task_is_alive(self) -> bool: + task: Final = self._periodic_flush_task + return task is not None and not task.done() and not task.get_loop().is_closed() + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def _enqueue(self, kwargs: Mapping[str, object]) -> None: + try: + if not self._flush_task_is_alive(): + self._periodic_flush_task = self._start_periodic_flush_task() + + payload: Final = self._payload_for(kwargs) + if payload is None: + verbose_logger.debug("zerobus: event carried no standard_logging_object, skipping") + return + + if self._flushing and len(self.log_queue) >= self.max_queue_size: + verbose_logger.warning("zerobus: queue at %s rows during a flush, dropped a row", self.max_queue_size) + return + + self.log_queue.append(trace_row(payload)) + self._drop_overflow() + if len(self.log_queue) >= self.batch_size: + self._start_batch_flush_task() + except Exception: # noqa: BLE001 # logging must never break the request path + verbose_logger.exception("zerobus: failed to queue an event") + + def _payload_for(self, kwargs: Mapping[str, object]) -> Mapping[str, object] | None: + """The payload to buffer, redacted the way the framework redacts the success path.""" + details: Final = self.redact_standard_logging_payload_from_model_call_details( + dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict + ) + payload: Final = details.get("standard_logging_object") + if not isinstance(payload, dict): + return None + if should_redact_message_logging(details): + return redacted_standard_logging_payload(payload) + return payload + + def _drop_overflow(self) -> None: + """Trim the oldest rows, except mid flush when the in-flight batch is the head of the queue.""" + if self._flushing: + return + overflow: Final = len(self.log_queue) - self.max_queue_size + if overflow <= 0: + return + del self.log_queue[:overflow] + verbose_logger.warning("zerobus: queue over %s rows, dropped %s oldest", self.max_queue_size, overflow) + + async def flush_queue(self, skip_if_flushing: bool = False) -> None: + if skip_if_flushing and self._flushing: + return + self._flushing = True + try: + await super().flush_queue() + finally: + self._flushing = False + + async def async_send_batch(self) -> None: + """A retryable failure propagates so the rows are kept; a permanent one drops them so the queue moves on.""" + rows: Final = tuple(self.log_queue) + if not rows: + return + failure: Final = await self.client.insert(rows) + if failure is None: + return + if failure.retryable: + raise ZerobusIngestError(failure.detail) + verbose_logger.error("zerobus: dropping %s rows, %s", len(rows), failure.detail) diff --git a/litellm/integrations/zerobus/row.py b/litellm/integrations/zerobus/row.py new file mode 100644 index 00000000000..c4da7975c44 --- /dev/null +++ b/litellm/integrations/zerobus/row.py @@ -0,0 +1,156 @@ +""" +Shape of one Delta table row per LiteLLM request. + +Zerobus validates every record against the target table and rejects unknown columns, so +the row is a fixed set of scalar columns for filtering plus JSON-encoded ``VARIANT`` +columns for anything nested. ``create_table_sql`` renders the matching DDL. +""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +TRACE_TABLE_COLUMNS: Final[Mapping[str, str]] = MappingProxyType( + { + "id": "STRING", + "trace_id": "STRING", + "session_id": "STRING", + "litellm_call_id": "STRING", + "call_type": "STRING", + "status": "STRING", + "model": "STRING", + "model_group": "STRING", + "model_id": "STRING", + "custom_llm_provider": "STRING", + "api_base": "STRING", + "stream": "BOOLEAN", + "cache_hit": "BOOLEAN", + "start_time": "TIMESTAMP", + "end_time": "TIMESTAMP", + "completion_start_time": "TIMESTAMP", + "response_time": "DOUBLE", + "prompt_tokens": "LONG", + "completion_tokens": "LONG", + "total_tokens": "LONG", + "response_cost": "DOUBLE", + "saved_cache_cost": "DOUBLE", + "api_key_hash": "STRING", + "api_key_alias": "STRING", + "team_id": "STRING", + "team_alias": "STRING", + "user_id": "STRING", + "org_id": "STRING", + "end_user": "STRING", + "requester_ip_address": "STRING", + "user_agent": "STRING", + "request_tags": "VARIANT", + "messages": "VARIANT", + "response": "VARIANT", + "error_str": "STRING", + "error_information": "VARIANT", + "metadata": "VARIANT", + "model_parameters": "VARIANT", + "hidden_params": "VARIANT", + "guardrail_information": "VARIANT", + "cost_breakdown": "VARIANT", + } +) + +_MICROSECONDS: Final = 1_000_000 + + +def create_table_sql(table_name: str) -> str: + columns: Final = ",\n".join(f" {name} {delta_type}" for name, delta_type in TRACE_TABLE_COLUMNS.items()) + return f"CREATE TABLE {table_name} (\n{columns}\n);" + + +def _text(payload: Mapping[str, object], key: str) -> str | None: + value: Final = payload.get(key) + return value if isinstance(value, str) else None + + +def _flag(payload: Mapping[str, object], key: str) -> bool | None: + value: Final = payload.get(key) + return value if isinstance(value, bool) else None + + +def _number(payload: Mapping[str, object], key: str) -> float | None: + value: Final = payload.get(key) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _count(payload: Mapping[str, object], key: str) -> int | None: + value: Final = _number(payload, key) + return None if value is None else int(value) + + +def _timestamp_micros(payload: Mapping[str, object], key: str) -> int | None: + """Delta ``TIMESTAMP`` over Zerobus is epoch microseconds; LiteLLM keeps epoch seconds.""" + seconds: Final = _number(payload, key) + if seconds is None or seconds <= 0: + return None + return int(seconds * _MICROSECONDS) + + +def _json(payload: Mapping[str, object], key: str) -> str | None: + value: Final = payload.get(key) + return None if value is None else safe_dumps(value) + + +def _metadata(payload: Mapping[str, object]) -> Mapping[str, object]: + value: Final = payload.get("metadata") + return value if isinstance(value, Mapping) else MappingProxyType({}) + + +def trace_row(payload: Mapping[str, object]) -> Mapping[str, object]: + """One ``TRACE_TABLE_COLUMNS`` row for a ``StandardLoggingPayload``.""" + metadata: Final = _metadata(payload) + return MappingProxyType( + { + "id": _text(payload, "id"), + "trace_id": _text(payload, "trace_id"), + "session_id": _text(payload, "session_id"), + "litellm_call_id": _text(payload, "litellm_call_id"), + "call_type": _text(payload, "call_type"), + "status": _text(payload, "status"), + "model": _text(payload, "model"), + "model_group": _text(payload, "model_group"), + "model_id": _text(payload, "model_id"), + "custom_llm_provider": _text(payload, "custom_llm_provider"), + "api_base": _text(payload, "api_base"), + "stream": _flag(payload, "stream"), + "cache_hit": _flag(payload, "cache_hit"), + "start_time": _timestamp_micros(payload, "startTime"), + "end_time": _timestamp_micros(payload, "endTime"), + "completion_start_time": _timestamp_micros(payload, "completionStartTime"), + "response_time": _number(payload, "response_time"), + "prompt_tokens": _count(payload, "prompt_tokens"), + "completion_tokens": _count(payload, "completion_tokens"), + "total_tokens": _count(payload, "total_tokens"), + "response_cost": _number(payload, "response_cost"), + "saved_cache_cost": _number(payload, "saved_cache_cost"), + "api_key_hash": _text(metadata, "user_api_key_hash"), + "api_key_alias": _text(metadata, "user_api_key_alias"), + "team_id": _text(metadata, "user_api_key_team_id"), + "team_alias": _text(metadata, "user_api_key_team_alias"), + "user_id": _text(metadata, "user_api_key_user_id"), + "org_id": _text(metadata, "user_api_key_org_id"), + "end_user": _text(payload, "end_user"), + "requester_ip_address": _text(payload, "requester_ip_address"), + "user_agent": _text(payload, "user_agent"), + "request_tags": _json(payload, "request_tags"), + "messages": _json(payload, "messages"), + "response": _json(payload, "response"), + "error_str": _text(payload, "error_str"), + "error_information": _json(payload, "error_information"), + "metadata": _json(payload, "metadata"), + "model_parameters": _json(payload, "model_parameters"), + "hidden_params": _json(payload, "hidden_params"), + "guardrail_information": _json(payload, "guardrail_information"), + "cost_breakdown": _json(payload, "cost_breakdown"), + } + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 6294f3bc577..7049fdd1f39 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -52,6 +52,7 @@ from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( VectorStorePreCallHook, ) +from litellm.integrations.zerobus import ZerobusLogger from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler from litellm.proxy.hooks.dynamic_rate_limiter_v3 import _PROXY_DynamicRateLimitHandlerV3 @@ -97,6 +98,7 @@ class CustomLoggerRegistry: "deepeval": DeepEvalLogger, "s3_v2": S3Logger, "pointfive": PointFiveLogger, + "zerobus": ZerobusLogger, "aws_sqs": SQSLogger, "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, "dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 54ed9171dc8..f8145cbdc7e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -211,6 +211,7 @@ from ..integrations.s3 import S3Logger from ..integrations.s3_v2 import S3Logger as S3V2Logger from ..integrations.supabase import Supabase from ..integrations.traceloop import TraceloopLogger +from ..integrations.zerobus import ZerobusLogger from .exception_mapping_utils import _get_response_headers from .initialize_dynamic_callback_params import ( get_trusted_callback_params, @@ -4650,6 +4651,14 @@ def _init_custom_logger_compatible_class( _pointfive_logger: Final = PointFiveLogger() _in_memory_loggers.append(_pointfive_logger) return _pointfive_logger + elif logging_integration == "zerobus": + for callback in _in_memory_loggers: + if isinstance(callback, ZerobusLogger): + return callback + + _zerobus_logger: Final = ZerobusLogger() + _in_memory_loggers.append(_zerobus_logger) + return _zerobus_logger elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): @@ -5342,6 +5351,10 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, PointFiveLogger): return callback + elif logging_integration == "zerobus": + for callback in _in_memory_loggers: + if isinstance(callback, ZerobusLogger): + return callback elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fd13697dab3..4597872d84e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4019,6 +4019,18 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ], ) + zerobus: CallbackOnUI = CallbackOnUI( + litellm_callback_name="zerobus", + ui_callback_name="Databricks Zerobus", + litellm_callback_params=[ # mutable-ok: the registry field is typed list + "ZEROBUS_WORKSPACE_URL", + "ZEROBUS_SERVER_ENDPOINT", + "ZEROBUS_CLIENT_ID", + "ZEROBUS_CLIENT_SECRET", + "ZEROBUS_TABLE_NAME", + ], + ) + class HTTPExceptionErrorDetail(TypedDict): """The `{"error": }` shape most proxy endpoints raise as `HTTPException.detail`.""" diff --git a/litellm/types/integrations/zerobus.py b/litellm/types/integrations/zerobus.py new file mode 100644 index 00000000000..217002dbc65 --- /dev/null +++ b/litellm/types/integrations/zerobus.py @@ -0,0 +1,53 @@ +from dataclasses import dataclass, field +from typing import Final + +from pydantic import Field + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + +RETRYABLE_INGEST_STATUS_CODES: Final = frozenset({408, 429, 500, 502, 503, 504}) + +TOKEN_REFRESH_LEEWAY_SECONDS: Final = 60 + + +class ZerobusInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a Databricks Zerobus logger on litellm. + + Every connection field falls back to its ``ZEROBUS_*`` environment variable, which is + what the proxy UI writes. ``table_name`` is the fully qualified ``catalog.schema.table``. + """ + + workspace_url: str | None = None + server_endpoint: str | None = None + client_id: str | None = None + client_secret: str | None = None + table_name: str | None = None + batch_size: int = Field(default=100, gt=0) + flush_interval: int = Field(default=10, gt=0) + + +@dataclass(frozen=True, slots=True) +class ZerobusConnection: + """Everything needed to mint a token for one table and post rows to it.""" + + workspace_url: str + workspace_id: str + server_endpoint: str + client_id: str + client_secret: str = field(repr=False) + table_name: str + + +@dataclass(frozen=True, slots=True) +class ZerobusAccessToken: + value: str = field(repr=False) + expires_at: float + + +@dataclass(frozen=True, slots=True) +class ZerobusIngestFailure: + """Why a batch could not be written, and whether a later attempt could still succeed.""" + + detail: str + retryable: bool diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py new file mode 100644 index 00000000000..ae7610536f3 --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py @@ -0,0 +1,258 @@ +import base64 +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, repeat + +import httpx +import pytest + +from litellm.integrations.zerobus.client import ZerobusIngestClient +from litellm.types.integrations.zerobus import ZerobusAccessToken, ZerobusConnection, ZerobusIngestFailure + +CONNECTION = ZerobusConnection( + workspace_url="https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/", + workspace_id="1234567890123456", + server_endpoint="https://1234567890123456.zerobus.us-west-2.cloud.databricks.com", + client_id="sp-client-id", + client_secret="sp-client-secret", + table_name="main.litellm.traces", +) +ROWS = ({"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}) + + +def _token(value: str = "tok-1", expires_in: float = 3600) -> httpx.Response: + return httpx.Response(200, text=json.dumps({"access_token": value, "expires_in": expires_in})) + + +def _accepted() -> httpx.Response: + return httpx.Response(200, text="{}") + + +@dataclass(frozen=True, slots=True) +class TokenCall: + url: str + data: Mapping[str, str] + headers: Mapping[str, str] + + +@dataclass(frozen=True, slots=True) +class InsertCall: + url: str + content: bytes + headers: Mapping[str, str] + + +def _results(results: Sequence[httpx.Response | Exception]) -> Iterator[httpx.Response | Exception]: + """Results are served in order, and the last one repeats.""" + return chain(results[:-1], repeat(results[-1])) + + +class FakeHTTPClient: + """Stands in for AsyncHTTPHandler, including its habit of raising on error statuses.""" + + def __init__( + self, + token: Sequence[httpx.Response | Exception] = (), + insert: Sequence[httpx.Response | Exception] = (), + ) -> None: + self.token_results = _results(token or (_token(),)) + self.insert_results = _results(insert or (_accepted(),)) + self.token_calls: tuple[TokenCall, ...] = () + self.insert_calls: tuple[InsertCall, ...] = () + + async def post( + self, + url: str, + data: Mapping[str, str] | None = None, + content: bytes | None = None, + headers: Mapping[str, str] | None = None, + ) -> httpx.Response: + if url.endswith("/oidc/v1/token"): + self.token_calls = (*self.token_calls, TokenCall(url, data or {}, headers or {})) + return _raise_like_the_handler(next(self.token_results), url) + self.insert_calls = (*self.insert_calls, InsertCall(url, content or b"", headers or {})) + return _raise_like_the_handler(next(self.insert_results), url) + + +def _raise_like_the_handler(result: httpx.Response | Exception, url: str) -> httpx.Response: + if isinstance(result, Exception): + raise result + if result.status_code >= 300: + raise httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", url), + response=httpx.Response(result.status_code, text=result.text), + ) + return result + + +class FakeClock: + def __init__(self, now: float = 1_000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +def _client(http_client: FakeHTTPClient, clock: FakeClock | None = None) -> ZerobusIngestClient: + return ZerobusIngestClient(connection=CONNECTION, http_client=http_client, clock=clock or FakeClock()) + + +@pytest.mark.asyncio +async def test_rows_are_posted_as_one_json_list_to_the_table_insert_endpoint(): + http_client = FakeHTTPClient() + + outcome = await _client(http_client).insert(ROWS) + + assert outcome is None + (call,) = http_client.insert_calls + # Insert endpoint per the Zerobus Ingest docs, read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == ( + "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com/zerobus/v1/tables/main.litellm.traces/insert" + ) + assert json.loads(call.content) == [{"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}] + assert call.headers["Content-Type"] == "application/json" + assert call.headers["Authorization"] == "Bearer tok-1" + + +@pytest.mark.asyncio +async def test_the_token_is_minted_for_the_zerobus_resource_with_the_table_privileges(): + """Zerobus refuses a plain workspace token: it must name its own resource and the table's UC privileges.""" + http_client = FakeHTTPClient() + + await _client(http_client).insert(ROWS) + + (call,) = http_client.token_calls + # Token form per the Zerobus Ingest docs (REST API authentication), read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/oidc/v1/token" + assert call.data["grant_type"] == "client_credentials" + assert call.data["scope"] == "all-apis" + assert call.data["resource"] == "api://databricks/workspaces/1234567890123456/zerobusDirectWriteApi" + details = json.loads(call.data["authorization_details"]) + assert [(d["object_type"], d["object_full_path"], d["privileges"]) for d in details] == [ + ("CATALOG", "main", ["USE CATALOG"]), + ("SCHEMA", "main.litellm", ["USE SCHEMA"]), + ("TABLE", "main.litellm.traces", ["SELECT", "MODIFY"]), + ] + assert all(d["type"] == "unity_catalog_privileges" for d in details) + + +@pytest.mark.asyncio +async def test_the_service_principal_authenticates_with_http_basic(): + http_client = FakeHTTPClient() + + await _client(http_client).insert(ROWS) + + scheme, credentials = http_client.token_calls[0].headers["Authorization"].split(" ") + assert scheme == "Basic" + assert base64.b64decode(credentials).decode() == "sp-client-id:sp-client-secret" + + +def test_the_client_secret_and_minted_token_stay_out_of_reprs_and_tracebacks(): + token = ZerobusAccessToken(value="tok-secret", expires_at=1.0) + + assert "sp-client-secret" not in repr(CONNECTION) + assert "sp-client-id" in repr(CONNECTION) + assert "tok-secret" not in repr(token) + assert "expires_at=1.0" in repr(token) + + +@pytest.mark.asyncio +async def test_the_token_is_reused_across_inserts_until_it_nears_expiry(): + clock = FakeClock(now=1_000.0) + http_client = FakeHTTPClient(token=[_token("tok-1", expires_in=600), _token("tok-2")]) + client = _client(http_client, clock) + + await client.insert(ROWS) + clock.now = 1_000.0 + 600 - 61 + await client.insert(ROWS) + clock.now = 1_000.0 + 600 - 59 + await client.insert(ROWS) + + assert len(http_client.token_calls) == 2 + assert [call.headers["Authorization"] for call in http_client.insert_calls] == [ + "Bearer tok-1", + "Bearer tok-1", + "Bearer tok-2", + ] + + +@pytest.mark.asyncio +async def test_a_401_discards_the_token_so_the_next_insert_mints_a_fresh_one(): + http_client = FakeHTTPClient( + token=[_token("tok-1"), _token("tok-2")], + insert=[httpx.Response(401, text="expired"), _accepted()], + ) + client = _client(http_client) + + first = await client.insert(ROWS) + second = await client.insert(ROWS) + + assert first == ZerobusIngestFailure(detail="insert returned 401, token discarded", retryable=True) + assert second is None + assert http_client.insert_calls[1].headers["Authorization"] == "Bearer tok-2" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [429, 500, 503]) +async def test_a_transient_insert_status_is_retryable(status: int): + http_client = FakeHTTPClient(insert=[httpx.Response(status, text="later")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + assert str(status) in outcome.detail + + +@pytest.mark.asyncio +async def test_a_schema_rejection_is_not_retryable_and_says_why(): + http_client = FakeHTTPClient(insert=[httpx.Response(400, text="unknown column foo")]) + + outcome = await _client(http_client).insert(ROWS) + + assert outcome == ZerobusIngestFailure(detail="insert returned 400: unknown column foo", retryable=False) + + +@pytest.mark.asyncio +async def test_a_network_failure_on_insert_is_retryable(): + http_client = FakeHTTPClient(insert=[httpx.ConnectError("connection refused")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + + +@pytest.mark.asyncio +async def test_bad_credentials_fail_the_insert_without_posting_rows(): + http_client = FakeHTTPClient(token=[httpx.Response(401, text="invalid_client")]) + + outcome = await _client(http_client).insert(ROWS) + + assert outcome == ZerobusIngestFailure(detail="token request returned 401: invalid_client", retryable=False) + assert http_client.insert_calls == () + + +@pytest.mark.asyncio +async def test_a_token_endpoint_outage_is_retryable(): + http_client = FakeHTTPClient(token=[httpx.Response(503, text="try later")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + + +@pytest.mark.asyncio +async def test_a_token_response_without_a_token_is_reported_not_raised(): + http_client = FakeHTTPClient(token=[httpx.Response(200, text='{"token_type": "Bearer"}')]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is False + assert "token response" in outcome.detail diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py new file mode 100644 index 00000000000..a85a9e2e6a0 --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py @@ -0,0 +1,392 @@ +import asyncio +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import chain, repeat + +import pytest + +import litellm +from litellm.integrations.zerobus.client import ZerobusIngestError +from litellm.integrations.zerobus.logger import ZerobusLogger, connection_for +from litellm.types.integrations.zerobus import ZerobusIngestFailure, ZerobusInitParams + +WORKSPACE_URL = "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com" +SERVER_ENDPOINT = "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com" + + +Row = Mapping[str, object] + + +class FakeIngestClient: + """Records the rows each flush would have written; outcomes are served in order and the last one repeats.""" + + def __init__( + self, + outcomes: Sequence[ZerobusIngestFailure | None] = (None,), + on_insert: Callable[[], None] | None = None, + ) -> None: + self.outcomes: Iterator[ZerobusIngestFailure | None] = chain(outcomes[:-1], repeat(outcomes[-1])) + self.on_insert = on_insert + self.batches: tuple[tuple[Row, ...], ...] = () + + async def insert(self, rows: Sequence[Row]) -> ZerobusIngestFailure | None: + if self.on_insert is not None: + self.on_insert() + self.batches = (*self.batches, tuple(rows)) + return next(self.outcomes) + + def ids(self) -> tuple[object, ...]: + return tuple(row["id"] for batch in self.batches for row in batch) + + +def _logger(client: FakeIngestClient, **params: object) -> ZerobusLogger: + return ZerobusLogger(params=ZerobusInitParams.model_validate(params), client=client) + + +def _event(request_id: str, **payload: object) -> dict[str, object]: + return { + "standard_logging_object": { + "id": request_id, + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "response": {"choices": []}, + **payload, + } + } + + +async def _settle(logger: ZerobusLogger) -> None: + for _ in range(200): + await asyncio.sleep(0.001) + task = logger._batch_flush_task + if (task is None or task.done()) and not logger._flushing: + return + + +@pytest.mark.asyncio +async def test_a_full_batch_is_written_as_one_insert_of_table_rows(): + client = FakeIngestClient() + logger = _logger(client, batch_size=3) + + for request_id in ("a", "b", "c"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert len(client.batches) == 1 + assert client.ids() == ("a", "b", "c") + assert client.batches[0][0]["model"] == "gpt-4o" + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_rows_are_held_until_the_batch_is_full(): + client = FakeIngestClient() + logger = _logger(client, batch_size=3) + + await logger.async_log_success_event(_event("a"), None, None, None) + + assert client.batches == () + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_failed_requests_are_written_too(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + + await logger.async_log_failure_event(_event("failed", status="failure", error_str="boom"), None, None, None) + + await _settle(logger) + assert client.ids() == ("failed",) + assert client.batches[0][0]["status"] == "failure" + assert client.batches[0][0]["error_str"] == "boom" + + +@pytest.mark.asyncio +async def test_an_event_without_a_standard_payload_is_skipped(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + + await logger.async_log_success_event({"kwargs": "but no payload"}, None, None, None) + + assert client.batches == () + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_retryable_failure_keeps_the_rows_for_the_next_flush(): + client = FakeIngestClient([ZerobusIngestFailure("zerobus is down", retryable=True)]) + logger = _logger(client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert [row["id"] for row in logger.log_queue] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_retryable_failure_surfaces_so_the_base_logger_can_preserve_it(): + client = FakeIngestClient([ZerobusIngestFailure("zerobus is down", retryable=True)]) + logger = _logger(client, batch_size=99) + logger.log_queue.append({"id": "a"}) + + with pytest.raises(ZerobusIngestError, match="zerobus is down"): + await logger.async_send_batch() + + +@pytest.mark.asyncio +async def test_a_rejected_batch_is_dropped_rather_than_blocking_the_queue(): + client = FakeIngestClient([ZerobusIngestFailure("unknown column", retryable=False)]) + logger = _logger(client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_row_that_arrives_mid_flush_is_kept_for_the_next_one(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + client.on_insert = lambda: logger.log_queue.append({"id": "late"}) + + await logger.async_log_success_event(_event("first"), None, None, None) + + await _settle(logger) + assert client.ids() == ("first",) + assert [row["id"] for row in logger.log_queue] == ["late"] + + +@pytest.mark.asyncio +async def test_the_queue_cap_holds_while_an_insert_is_in_flight(): + """A slow insert must not let the queue grow past max_queue_size, nor disturb the in-flight head.""" + insert_started = asyncio.Event() + finish_insert = asyncio.Event() + + class SlowClient: + batches: tuple[tuple[Row, ...], ...] = () + + async def insert(self, rows: Sequence[Row]) -> None: + insert_started.set() + await finish_insert.wait() + self.batches = (*self.batches, tuple(rows)) + + client = SlowClient() + logger = ZerobusLogger(params=ZerobusInitParams(batch_size=2), client=client) + logger.max_queue_size = 3 + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + await insert_started.wait() + for request_id in ("c", "d", "e"): + await logger.async_log_success_event(_event(request_id), None, None, None) + finish_insert.set() + await _settle(logger) + + assert [[row["id"] for row in batch] for batch in client.batches] == [["a", "b"]] + assert [row["id"] for row in logger.log_queue] == ["c"] + + +@pytest.mark.asyncio +async def test_a_client_error_does_not_break_the_request_path(): + class ExplodingClient: + async def insert(self, rows: Sequence[Row]) -> None: + raise RuntimeError("bug") + + logger = ZerobusLogger(params=ZerobusInitParams(batch_size=1), client=ExplodingClient()) + + await logger.async_log_success_event(_event("a"), None, None, None) + await _settle(logger) + + assert [row["id"] for row in logger.log_queue] == ["a"] + + +@pytest.mark.asyncio +async def test_turn_off_message_logging_redacts_prompts_and_responses_but_keeps_the_rest(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1, turn_off_message_logging=True) + + await logger.async_log_success_event( + _event("a", prompt_tokens=10, response={"choices": [{"message": {"content": "the secret answer"}}]}), + None, + None, + None, + ) + + await _settle(logger) + (row,) = client.batches[0] + assert row["id"] == "a" + assert row["prompt_tokens"] == 10 + assert '"hi"' not in str(row["messages"]) + assert "the secret answer" not in str(row["response"]) + + +def test_connection_comes_from_the_environment_the_proxy_ui_writes(monkeypatch): + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + + connection = connection_for(ZerobusInitParams()) + + assert connection.workspace_url == WORKSPACE_URL + assert connection.server_endpoint == SERVER_ENDPOINT + assert connection.workspace_id == "1234567890123456" + assert connection.client_id == "sp-id" + assert connection.client_secret == "sp-secret" + assert connection.table_name == "main.litellm.traces" + + +def test_config_yaml_params_win_over_the_environment(monkeypatch): + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "env.schema.table") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "from-env") + + connection = connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="from-config", + table_name="cfg.schema.table", + ) + ) + + assert connection.table_name == "cfg.schema.table" + assert connection.client_secret == "from-config" + + +def test_a_secret_reference_in_config_yaml_is_resolved(monkeypatch): + monkeypatch.setenv("MY_SP_SECRET", "resolved-secret") + + connection = connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="os.environ/MY_SP_SECRET", + table_name="main.litellm.traces", + ) + ) + + assert connection.client_secret == "resolved-secret" + + +def test_a_missing_setting_names_the_env_var_to_set(monkeypatch): + monkeypatch.delenv("ZEROBUS_CLIENT_SECRET", raising=False) + + with pytest.raises(ValueError, match="ZEROBUS_CLIENT_SECRET"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + table_name="main.litellm.traces", + ) + ) + + +def test_a_table_that_is_not_fully_qualified_is_refused(): + with pytest.raises(ValueError, match=r"catalog\.schema\.table"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="sp-secret", + table_name="traces", + ) + ) + + +def test_an_endpoint_without_a_workspace_id_is_refused(): + """The token's resource needs the numeric workspace id, which only the Zerobus hostname carries.""" + with pytest.raises(ValueError, match="ZEROBUS_SERVER_ENDPOINT"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=WORKSPACE_URL, + client_id="sp-id", + client_secret="sp-secret", + table_name="main.litellm.traces", + ) + ) + + +def test_a_misconfigured_logger_fails_at_startup_not_at_first_flush(monkeypatch): + for name in ("WORKSPACE_URL", "SERVER_ENDPOINT", "CLIENT_ID", "CLIENT_SECRET", "TABLE_NAME"): + monkeypatch.delenv(f"ZEROBUS_{name}", raising=False) + monkeypatch.setattr(litellm, "zerobus_params", None) + + with pytest.raises(ValueError, match="ZEROBUS_"): + ZerobusLogger() + + +def test_litellm_zerobus_params_configure_the_logger(monkeypatch): + monkeypatch.setattr( + litellm, + "zerobus_params", + { + "workspace_url": WORKSPACE_URL, + "server_endpoint": SERVER_ENDPOINT, + "client_id": "sp-id", + "client_secret": "sp-secret", + "table_name": "main.litellm.traces", + "batch_size": 7, + "flush_interval": 3, + }, + ) + + logger = ZerobusLogger() + + assert logger.batch_size == 7 + assert logger.flush_interval == 3 + assert logger.client.connection.table_name == "main.litellm.traces" + + +def test_the_client_is_kept_while_the_connection_is_unchanged_and_rebuilt_when_it_changes(monkeypatch): + """The client caches its token, so it must survive across flushes, yet a UI edit must take effect.""" + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + monkeypatch.setattr(litellm, "zerobus_params", None) + logger = ZerobusLogger() + + first = logger.client + unchanged = logger.client + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces_v2") + rebuilt = logger.client + + assert unchanged is first + assert rebuilt is not first + assert rebuilt.connection.table_name == "main.litellm.traces_v2" + + +def test_callbacks_zerobus_builds_one_logger_and_reuses_it(monkeypatch): + """`litellm_settings.callbacks: ["zerobus"]` goes through litellm_logging, which must hand back one instance.""" + from litellm.litellm_core_utils import litellm_logging as logging_module + + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + monkeypatch.setattr(litellm, "zerobus_params", None) + monkeypatch.setattr(logging_module, "_in_memory_loggers", []) + + assert logging_module.get_custom_logger_compatible_class("zerobus") is None + + first = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + second = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + + assert isinstance(first, ZerobusLogger) + assert second is first + assert logging_module.get_custom_logger_compatible_class("zerobus") is first diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_row.py b/tests/test_litellm/integrations/zerobus/test_zerobus_row.py new file mode 100644 index 00000000000..b73c3bae48f --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_row.py @@ -0,0 +1,139 @@ +import json + +from litellm.integrations.zerobus.row import TRACE_TABLE_COLUMNS, create_table_sql, trace_row + + +def _payload() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "trace_id": "trace-1", + "session_id": "session-1", + "litellm_call_id": "call-1", + "call_type": "acompletion", + "status": "success", + "model": "gpt-4o", + "model_group": "gpt-4o-group", + "custom_llm_provider": "openai", + "api_base": "https://api.openai.com", + "stream": False, + "cache_hit": None, + "startTime": 1_700_000_000.25, + "endTime": 1_700_000_001.5, + "completionStartTime": 1_700_000_000.75, + "response_time": 1.25, + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "response_cost": 0.0015, + "saved_cache_cost": 0.0, + "end_user": "end-user-1", + "requester_ip_address": "10.0.0.1", + "user_agent": "curl/8", + "request_tags": ["prod"], + "messages": [{"role": "user", "content": "hi"}], + "response": {"choices": [{"message": {"role": "assistant", "content": "hello"}}]}, + "error_str": None, + "error_information": None, + "metadata": { + "user_api_key_hash": "hash-1", + "user_api_key_alias": "alias-1", + "user_api_key_team_id": "team-1", + "user_api_key_team_alias": "team-alias-1", + "user_api_key_user_id": "user-1", + "user_api_key_org_id": "org-1", + }, + "model_parameters": {"temperature": 0.2}, + "hidden_params": {"response_cost": 0.0015}, + "guardrail_information": None, + "cost_breakdown": {"input_cost": 0.001, "output_cost": 0.0005}, + } + + +def test_every_row_has_exactly_the_documented_columns(): + """Zerobus rejects a record naming a column the table lacks, so the row and the DDL must agree.""" + assert tuple(trace_row(_payload())) == tuple(TRACE_TABLE_COLUMNS) + assert tuple(trace_row({})) == tuple(TRACE_TABLE_COLUMNS) + + +def test_scalars_land_in_their_columns(): + row = trace_row(_payload()) + + assert row["id"] == "chatcmpl-1" + assert row["trace_id"] == "trace-1" + assert row["status"] == "success" + assert row["model"] == "gpt-4o" + assert row["stream"] is False + assert row["prompt_tokens"] == 10 + assert row["total_tokens"] == 15 + assert row["response_cost"] == 0.0015 + assert row["end_user"] == "end-user-1" + + +def test_key_and_team_identity_is_lifted_out_of_metadata(): + """Filtering spend by team or key is the main query, so those live in their own columns.""" + row = trace_row(_payload()) + + assert row["api_key_hash"] == "hash-1" + assert row["api_key_alias"] == "alias-1" + assert row["team_id"] == "team-1" + assert row["team_alias"] == "team-alias-1" + assert row["user_id"] == "user-1" + assert row["org_id"] == "org-1" + + +def test_timestamps_become_epoch_microseconds(): + row = trace_row(_payload()) + + assert row["start_time"] == 1_700_000_000_250_000 + assert row["end_time"] == 1_700_000_001_500_000 + assert row["completion_start_time"] == 1_700_000_000_750_000 + + +def test_a_zero_timestamp_is_null_rather_than_1970(): + """LiteLLM leaves completionStartTime at 0 when there is no first token, which is not a real time.""" + row = trace_row({**_payload(), "completionStartTime": 0}) + + assert row["completion_start_time"] is None + + +def test_nested_fields_are_json_text_for_the_variant_columns(): + row = trace_row(_payload()) + + assert json.loads(str(row["messages"])) == [{"role": "user", "content": "hi"}] + assert json.loads(str(row["metadata"]))["user_api_key_team_id"] == "team-1" + assert json.loads(str(row["request_tags"])) == ["prod"] + assert json.loads(str(row["cost_breakdown"])) == {"input_cost": 0.001, "output_cost": 0.0005} + + +def test_missing_and_null_fields_are_null(): + row = trace_row({**_payload(), "messages": None, "guardrail_information": None}) + + assert row["messages"] is None + assert row["guardrail_information"] is None + assert row["error_str"] is None + assert row["cache_hit"] is None + + +def test_a_wrongly_typed_field_is_null_instead_of_a_rejected_record(): + """One odd payload must not poison the whole batch: the table type wins.""" + row = trace_row({**_payload(), "prompt_tokens": "ten", "stream": "yes", "startTime": "now"}) + + assert row["prompt_tokens"] is None + assert row["stream"] is None + assert row["start_time"] is None + + +def test_the_row_survives_a_json_round_trip_unchanged(): + row = trace_row(_payload()) + + assert json.loads(json.dumps(dict(row))) == dict(row) + + +def test_create_table_sql_declares_every_column_with_its_type(): + sql = create_table_sql("main.litellm.traces") + + assert sql.startswith("CREATE TABLE main.litellm.traces (") + assert " start_time TIMESTAMP," in sql + assert " messages VARIANT," in sql + assert " cost_breakdown VARIANT\n);" in sql + assert sql.count(",") == len(TRACE_TABLE_COLUMNS) - 1 diff --git a/tests/test_litellm/proxy/test_zerobus_dashboard_config.py b/tests/test_litellm/proxy/test_zerobus_dashboard_config.py new file mode 100644 index 00000000000..d2143767480 --- /dev/null +++ b/tests/test_litellm/proxy/test_zerobus_dashboard_config.py @@ -0,0 +1,52 @@ +from pathlib import Path +from typing import Final + +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger + + +class DashboardField(BaseModel): + type: str + required: bool + + +class DashboardCallbackConfig(BaseModel): + id: str + displayName: str + logo: str + supports_key_team_logging: bool + dynamic_params: dict[str, DashboardField] + + +def _zerobus_config() -> DashboardCallbackConfig: + path: Final = Path(litellm.__file__).parent / "integrations" / "callback_configs.json" + configs: Final = TypeAdapter(tuple[DashboardCallbackConfig, ...]).validate_json(path.read_text()) + return next(config for config in configs if config.id == "zerobus") + + +def test_zerobus_appears_in_the_dashboard_callback_dropdown(): + """The dropdown is served from callback_configs.json, so an entry only in the dashboard source is invisible.""" + entry = _zerobus_config() + + assert entry.displayName == "Databricks Zerobus" + assert entry.supports_key_team_logging is False + assert entry.dynamic_params["ZEROBUS_CLIENT_SECRET"].type == "password" + assert all(field.required is True for field in entry.dynamic_params.values()) + + +def test_the_dropdown_logo_asset_exists(): + """A logo the dashboard cannot resolve degrades silently to a letter tile.""" + logo = _zerobus_config().logo + repo_root = Path(litellm.__file__).parent.parent + asset = repo_root / "ui" / "litellm-dashboard" / "public" / "assets" / "logos" / logo + + assert asset.is_file() + + +def test_the_dropdown_fields_are_the_env_vars_the_logger_reads(): + """Naming the fields as stored means the edit form prefills saved values instead of showing blanks.""" + fields = tuple(_zerobus_config().dynamic_params) + + assert fields == tuple(CustomLogger.get_callback_env_vars("zerobus")) diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index f5138d55b5d..bc9889da724 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -10,6 +10,7 @@ import newrelicLogo from "../../public/assets/logos/newrelic.png"; import openmeterLogo from "../../public/assets/logos/openmeter.png"; import otelLogo from "../../public/assets/logos/otel.png"; import pointfiveLogo from "../../public/assets/logos/pointfive.png"; +import databricksLogo from "../../public/assets/logos/databricks.svg"; interface CallbackConfig { id: string; @@ -181,6 +182,20 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [ }, description: "PointFive Logging Integration", }, + { + id: "zerobus", + displayName: "Databricks Zerobus", + logo: databricksLogo.src, + supports_key_team_logging: false, + dynamic_params: { + ZEROBUS_WORKSPACE_URL: "text", + ZEROBUS_SERVER_ENDPOINT: "text", + ZEROBUS_CLIENT_ID: "text", + ZEROBUS_CLIENT_SECRET: "password", + ZEROBUS_TABLE_NAME: "text", + }, + description: "Databricks Zerobus Ingest Logging Integration", + }, { id: "s3", displayName: "S3", From 8b68c3cd0925b95f8e7b8cbe605245313620733b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 25 Sep 2026 18:32:21 -0400 Subject: [PATCH 24/58] fix(vertex_ai): keep legacy bucket_name in credential resolution and add GCS_BATCH_BUCKET_NAME env var (#42803) * fix(vertex_ai): map legacy bucket_name to gcs_bucket_name and add GCS_BATCH_BUCKET_NAME env var * refactor(router): keep legacy bucket_name as a credential field instead of a validator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): pass the RAG corpus bucket to the file upload instead of hopping through GCS_BUCKET_NAME * fix(vertex_ai): accept existing_file_id in the RAG Engine store step so ingest() runs end to end --------- Co-authored-by: Mubashir Osmani Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/vertex_ai/files/handler.py | 14 ++-- .../llms/vertex_ai/files/transformation.py | 5 +- .../llms/vertex_ai/rag_engine/ingestion.py | 49 +++++------- litellm/types/router.py | 1 + .../files/test_vertex_ai_files_handler.py | 35 +++++++++ .../test_vertex_ai_files_transformation.py | 11 +++ .../llms/vertex_ai/rag_engine/__init__.py | 0 .../vertex_ai/rag_engine/test_ingestion.py | 76 +++++++++++++++++++ tests/unit/test_router/test_router.py | 45 +++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 10 files changed, 203 insertions(+), 37 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/rag_engine/__init__.py create mode 100644 tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ac95d1348f9..f2da04a7db7 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -53,14 +53,18 @@ class VertexAIFilesHandler(GCSBucketBase): Sources them from the deployment's ``litellm_params`` (``gcs_bucket_name`` / ``bucket_name`` and ``vertex_credentials``), mirroring the write path in - ``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the global - ``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` env vars. This lets Vertex batch - run entirely at the model-group level, so output written to a per-model bucket is - readable without setting the global env vars. + ``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the + ``GCS_BATCH_BUCKET_NAME`` then ``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` + env vars. This lets Vertex batch run entirely at the model-group level, so output + written to a per-model bucket is readable without setting the global env vars. """ params: Final[Mapping[str, object]] = litellm_params or {} bucket_candidate: Final = params.get("gcs_bucket_name") or params.get("bucket_name") - configured_bucket_name = bucket_candidate if isinstance(bucket_candidate, str) else os.getenv("GCS_BUCKET_NAME") + configured_bucket_name = ( + bucket_candidate + if isinstance(bucket_candidate, str) + else os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME") + ) credentials: Final = params.get("vertex_credentials") or vertex_credentials if isinstance(credentials, dict): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dbb41b57348..2b0694697a4 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -961,7 +961,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_configured_bucket_name(self, litellm_params: dict) -> str: bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + litellm_params.get("gcs_bucket_name") + or litellm_params.get("bucket_name") + or os.getenv("GCS_BATCH_BUCKET_NAME") + or os.getenv("GCS_BUCKET_NAME") ) if not bucket_name: raise ValueError("GCS bucket_name is required") diff --git a/litellm/llms/vertex_ai/rag_engine/ingestion.py b/litellm/llms/vertex_ai/rag_engine/ingestion.py index d9916209a14..c10bac595b6 100644 --- a/litellm/llms/vertex_ai/rag_engine/ingestion.py +++ b/litellm/llms/vertex_ai/rag_engine/ingestion.py @@ -122,41 +122,26 @@ class VertexAIRAGIngestion(BaseRAGIngestion): """ import litellm - # Set GCS_BUCKET_NAME env var for litellm.files.create_file - # The handler uses this to determine where to upload - original_bucket: Final = os.environ.get("GCS_BUCKET_NAME") - if self.gcs_bucket: - os.environ["GCS_BUCKET_NAME"] = self.gcs_bucket + file_tuple: Final = (filename, file_content, content_type) - try: - # Create file tuple for litellm.files.acreate_file - file_tuple: Final = (filename, file_content, content_type) + verbose_logger.debug( + "Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket + ) - verbose_logger.debug( - "Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket - ) + response: Final = await litellm.acreate_file( + file=file_tuple, + purpose="assistants", + custom_llm_provider="vertex_ai", + gcs_bucket_name=self.gcs_bucket, + vertex_project=self.vertex_project, + vertex_location=self.vertex_location, + vertex_credentials=self.vertex_credentials, + ) - # Upload to GCS using LiteLLM's file upload - response: Final = await litellm.acreate_file( - file=file_tuple, - purpose="assistants", # Purpose for file storage - custom_llm_provider="vertex_ai", - vertex_project=self.vertex_project, - vertex_location=self.vertex_location, - vertex_credentials=self.vertex_credentials, - ) + gcs_uri: Final = response.id + verbose_logger.info("Uploaded file to GCS: %s", gcs_uri) - # The response.id should be the GCS URI - gcs_uri: Final = response.id - verbose_logger.info("Uploaded file to GCS: %s", gcs_uri) - - return gcs_uri - finally: - # Restore original env var - if original_bucket is not None: - os.environ["GCS_BUCKET_NAME"] = original_bucket - elif "GCS_BUCKET_NAME" in os.environ: - del os.environ["GCS_BUCKET_NAME"] + return gcs_uri async def _import_file_to_corpus_via_sdk( self, @@ -259,6 +244,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): content_type: str | None, chunks: list[str], embeddings: list[list[float]] | None, + existing_file_id: str | None = None, ) -> tuple[str | None, str | None]: """ Store content in Vertex AI RAG corpus. @@ -274,6 +260,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): content_type: MIME type chunks: Ignored - Vertex AI handles chunking embeddings: Ignored - Vertex AI handles embedding + existing_file_id: Existing provider file ID, unsupported for Vertex AI RAG Engine Returns: Tuple of (corpus_id, gcs_uri) diff --git a/litellm/types/router.py b/litellm/types/router.py index b72809f625f..d545f7ae639 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -345,6 +345,7 @@ class CredentialLiteLLMParams(BaseModel): ## OBJECT STORAGE (files / batches) ## gcs_bucket_name: str | None = None + bucket_name: str | None = None ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: str | None = None diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py index e0f0b7e5c0b..9b7cd127b83 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -208,6 +208,7 @@ class TestVertexAIFilesHandler: assert service_account == "/model/sa.json" def test_resolve_read_gcs_config_falls_back_to_env(self, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket") monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/env/sa.json") @@ -216,6 +217,40 @@ class TestVertexAIFilesHandler: assert bucket == "env-default-bucket" assert service_account == "/env/sa.json" + def test_resolve_read_gcs_config_prefers_batch_env_over_logging_env(self, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + + bucket, _ = self.handler._resolve_read_gcs_config(litellm_params={}, vertex_credentials=None) + + assert bucket == "batch-bucket" + + def test_resolve_read_gcs_config_prefers_per_model_bucket_over_batch_env(self, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"gcs_bucket_name": "my-model-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "my-model-bucket" + + def test_resolve_read_gcs_config_prefers_gcs_bucket_name_over_legacy(self): + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"gcs_bucket_name": "my-model-bucket", "bucket_name": "legacy-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "my-model-bucket" + + def test_resolve_read_gcs_config_accepts_legacy_bucket_name_alone(self): + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"bucket_name": "legacy-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "legacy-bucket" + def test_resolve_read_gcs_config_serializes_dict_credentials(self, monkeypatch): monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) 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 7434eae72a4..6f18a391f7b 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 @@ -1186,10 +1186,21 @@ class TestConfiguredBucketNameResolution: assert config._get_configured_bucket_name({"gcs_bucket_name": "new", "bucket_name": "legacy"}) == "new" def test_should_fall_back_to_env(self, config, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") assert config._get_configured_bucket_name({}) == "env-bucket" + def test_should_prefer_batch_env_over_logging_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + assert config._get_configured_bucket_name({}) == "batch-bucket" + + def test_should_prefer_litellm_params_over_batch_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + assert config._get_configured_bucket_name({"gcs_bucket_name": "per-model-bucket"}) == "per-model-bucket" + def test_should_raise_when_no_bucket_anywhere(self, config, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) with pytest.raises(ValueError, match="GCS bucket_name is required"): config._get_configured_bucket_name({}) diff --git a/tests/unit/llms/vertex_ai/rag_engine/__init__.py b/tests/unit/llms/vertex_ai/rag_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py b/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py new file mode 100644 index 00000000000..3acabc4d14e --- /dev/null +++ b/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py @@ -0,0 +1,76 @@ +import asyncio +import sys +from types import ModuleType, SimpleNamespace + +import litellm +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig +from litellm.llms.vertex_ai.rag_engine.ingestion import VertexAIRAGIngestion + + +def _ingestion_for_bucket(bucket: str) -> VertexAIRAGIngestion: + return VertexAIRAGIngestion( + { + "vector_store": { + "custom_llm_provider": "vertex_ai", + "vector_store_id": "corpus-123", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "gcs_bucket": bucket, + } + } + ) + + +def test_upload_lands_in_the_corpus_bucket_when_batch_bucket_env_is_set(monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + resolver = VertexAIFilesConfig() + + async def acreate_file_through_real_bucket_resolver(**kwargs): + bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs)) + return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}") + + monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver) + + uri = asyncio.run(_ingestion_for_bucket("rag-bucket")._upload_file_to_gcs(b"doc", "doc.txt", "text/plain")) + + assert uri == "gs://rag-bucket/doc.txt" + + +def _vertexai_sdk_stub(import_calls: list[dict[str, object]]) -> ModuleType: + rag = ModuleType("vertexai.rag") + rag.TransformationConfig = lambda chunking_config: chunking_config + rag.ChunkingConfig = lambda chunk_size, chunk_overlap: (chunk_size, chunk_overlap) + + def import_files(**kwargs): + import_calls.append(kwargs) + return SimpleNamespace(imported_rag_files_count=1) + + rag.import_files = import_files + vertexai = ModuleType("vertexai") + vertexai.init = lambda project, location: None + vertexai.rag = rag + return vertexai + + +def test_ingest_runs_end_to_end_through_the_base_pipeline(monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + resolver = VertexAIFilesConfig() + import_calls: list[dict[str, object]] = [] + stub = _vertexai_sdk_stub(import_calls) + monkeypatch.setitem(sys.modules, "vertexai", stub) + monkeypatch.setitem(sys.modules, "vertexai.rag", stub.rag) + + async def acreate_file_through_real_bucket_resolver(**kwargs): + bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs)) + return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}") + + monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver) + + result = asyncio.run(_ingestion_for_bucket("rag-bucket").ingest(file_data=("doc.txt", b"doc", "text/plain"))) + + assert (result["status"], result["vector_store_id"], result["file_id"]) == ("completed", "corpus-123", "gs://rag-bucket/doc.txt") + assert [(c["corpus_name"], c["paths"]) for c in import_calls] == [ + ("projects/test-project/locations/us-central1/ragCorpora/corpus-123", ["gs://rag-bucket/doc.txt"]) + ] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 10669de9cc8..3393c2f0d3c 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -6470,6 +6470,51 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_keeps_legacy_bucket_name(): + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "bucket_name": "my-legacy-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + + assert credentials is not None + assert credentials["bucket_name"] == "my-legacy-bucket" + assert "gcs_bucket_name" not in credentials + + +def test_get_deployment_credentials_with_provider_keeps_both_bucket_keys(): + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "gcs_bucket_name": "new-bucket", + "bucket_name": "legacy-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + + assert credentials is not None + assert credentials["gcs_bucket_name"] == "new-bucket" + assert credentials["bucket_name"] == "legacy-bucket" + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6126733095b..be7094f7f6c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32945,6 +32945,8 @@ export interface components { azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; + /** Bucket Name */ + bucket_name?: string | null; /** Budget Duration */ budget_duration?: string | null; /** Cache Creation Input Audio Token Cost */ @@ -46730,6 +46732,8 @@ export interface components { azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; + /** Bucket Name */ + bucket_name?: string | null; /** Budget Duration */ budget_duration?: string | null; /** Cache Creation Input Audio Token Cost */ From 8ef85a45ce1362fbb6a66ecd52b4e8c25bb26df2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:35:20 -0700 Subject: [PATCH 25/58] feat(xai): add native xAI batches and files support (#42812) * feat(xai): add native xAI batches and files support Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(xai): tighten batch handler typing and avoid Final redeclaration on star import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(xai): walk batch result pages iteratively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(xai): stop paging on empty pagination token and honor litellm.xai_key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(xai): import NotRequired and TypedDict from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(batches): accept image and video endpoints on batch create * test(xai): lock batch endpoint, auth, and result contracts The batches test package collided with litellm/batches under pytest prepend, so the All Other Providers shard could not collect the new tests. * fix(xai): price grok batch usage at xAI's 20 percent batch discount * fix(xai): map not-found file reads to 404, bill batch reasoning tokens, and add 200k batch tier rates * refactor(xai): drop routine prose and move tests under tests/unit * fix(health): hand the resolved provider to list_batches in batch-mode health checks * test(xai): make tests/unit/llms/xai/batches a package * fix(xai): walk every page of the files list by pagination_token --------- Co-authored-by: Mubashir Osmani Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mubashir1osmani Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/batches/batch_utils.py | 4 + litellm/batches/main.py | 77 +++- litellm/files/main.py | 21 +- .../health_check_helpers.py | 12 +- litellm/litellm_core_utils/litellm_logging.py | 3 + litellm/llms/xai/batches/__init__.py | 0 litellm/llms/xai/batches/handler.py | 195 ++++++++++ litellm/llms/xai/batches/transformation.py | 278 ++++++++++++++ litellm/llms/xai/chat/transformation.py | 6 +- litellm/llms/xai/files/__init__.py | 0 litellm/llms/xai/files/transformation.py | 247 +++++++++++++ ...odel_prices_and_context_window_backup.json | 78 ++++ litellm/types/llms/openai.py | 14 +- litellm/types/utils.py | 8 +- litellm/utils.py | 13 + model_prices_and_context_window.json | 78 ++++ model_prices_and_context_window.schema.json | 15 + .../test_health_check_helpers.py | 20 + .../test_litellm_logging.py | 22 +- tests/unit/batches/test_batch_utils.py | 34 ++ tests/unit/llms/xai/batches/__init__.py | 0 .../xai/batches/test_xai_batches_handler.py | 344 ++++++++++++++++++ .../test_xai_batches_transformation.py | 224 ++++++++++++ tests/unit/llms/xai/files/__init__.py | 0 .../files/test_xai_files_transformation.py | 144 ++++++++ .../llms/xai/test_xai_chat_transformation.py | 10 +- tests/unit/test_cost_calculator.py | 49 +++ tests/unit/test_utils.py | 3 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 29 files changed, 1874 insertions(+), 37 deletions(-) create mode 100644 litellm/llms/xai/batches/__init__.py create mode 100644 litellm/llms/xai/batches/handler.py create mode 100644 litellm/llms/xai/batches/transformation.py create mode 100644 litellm/llms/xai/files/__init__.py create mode 100644 litellm/llms/xai/files/transformation.py create mode 100644 tests/unit/llms/xai/batches/__init__.py create mode 100644 tests/unit/llms/xai/batches/test_xai_batches_handler.py create mode 100644 tests/unit/llms/xai/batches/test_xai_batches_transformation.py create mode 100644 tests/unit/llms/xai/files/__init__.py create mode 100644 tests/unit/llms/xai/files/test_xai_files_transformation.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 819a279a43c..246ac4fd369 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -706,6 +706,10 @@ def _get_batch_job_usage_from_response_body( if ResponseAPILoggingUtils._is_response_api_usage(_usage_dict): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_usage_dict) usage: Final[Usage] = Usage(**_usage_dict) + if custom_llm_provider == "xai": + from litellm.llms.xai.chat.transformation import XAIChatConfig + + XAIChatConfig.fold_reasoning_tokens_into_completion(usage) return usage diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 76b6c73b375..f977fc03891 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -31,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.openai import OpenAIBatchesAPI from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction +from litellm.llms.xai.batches.handler import XAIBatchesHandler from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( CancelBatchRequest, @@ -59,6 +60,7 @@ openai_batches_instance: Final = OpenAIBatchesAPI() azure_batches_instance: Final = AzureBatchesAPI() vertex_ai_batches_instance: Final = VertexAIBatchPrediction(gcs_bucket_name="") anthropic_batches_instance: Final = AnthropicBatchesHandler() +xai_batches_instance: Final = XAIBatchesHandler() base_llm_http_handler = BaseLLMHTTPHandler() ################################################# @@ -105,10 +107,22 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], input_file_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -157,10 +171,22 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], input_file_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -243,6 +269,14 @@ def create_batch( model=model, ) return response + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.create_batch( + _is_async=_is_async, + create_batch_data=_create_batch_request, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -345,7 +379,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -393,10 +427,18 @@ def _handle_retrieve_batch_providers_without_provider_config( _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", logging_obj: LiteLLMLoggingObj | None = None, ): + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.retrieve_batch( + _is_async=_is_async, + batch_id=batch_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -518,7 +560,7 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -741,6 +783,15 @@ def list_batches( timeout = 600.0 _is_async: Final = kwargs.pop("alist_batches", False) is True + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.list_batches( + _is_async=_is_async, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + after=after, + limit=limit, + ) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( @@ -837,7 +888,7 @@ def list_batches( async def acancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy", "xai"] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -883,7 +934,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] | str = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy", "xai"] | str = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -933,6 +984,14 @@ def cancel_batch( ) _is_async: Final = kwargs.pop("acancel_batch", False) is True + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.cancel_batch( + _is_async=_is_async, + batch_id=batch_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: api_base = ( diff --git a/litellm/files/main.py b/litellm/files/main.py index 72832aeccc9..723784795b0 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -28,12 +28,15 @@ FileCreateProvider = Literal[ "manus", "anthropic", "mistral", + "xai", ] FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral", "xai" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"] -FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] +FileDeleteProvider = Literal[ + "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral", "xai" +] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral", "xai"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse @@ -49,6 +52,8 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler +from litellm.llms.xai.batches.handler import XAIBatchesHandler +from litellm.llms.xai.batches.transformation import is_xai_batch_results_id from litellm.types.llms.openai import ( CreateFileRequest, FileContentRequest, @@ -103,6 +108,7 @@ openai_files_instance: Final = OpenAIFilesAPI() azure_files_instance: Final = AzureOpenAIFilesAPI() vertex_ai_files_instance: Final = VertexAIFilesHandler() bedrock_files_instance: Final = BedrockFilesHandler() +xai_batch_results_instance: Final = XAIBatchesHandler() ################################################# @@ -920,6 +926,15 @@ def file_content( client=client, ) + if custom_llm_provider == LlmProviders.XAI.value and is_xai_batch_results_id(file_id): + return xai_batch_results_instance.batch_results_content( + _is_async=_is_async, + batch_id=file_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) + # Check if provider has a custom files config (e.g., Anthropic, Manus) provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 22b8d850c83..a0f027cd58f 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -114,10 +114,9 @@ class HealthCheckHelpers: """ Health check for batch mode. - Calls list_batches for providers that support it (openai, hosted_vllm, azure, - vertex_ai). For all other providers (e.g. bedrock) the batch API surface doesn't - include list_batches, so we fall back to acompletion to verify connectivity and - credential validity instead. + Calls list_batches for providers that support it. For all other providers (e.g. bedrock) + the batch API surface doesn't include list_batches, so we fall back to acompletion to + verify connectivity and credential validity instead. """ import litellm @@ -132,10 +131,9 @@ class HealthCheckHelpers: litellm_params={"api_base": api_base} if api_base else None, ) - if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS: - return await litellm.alist_batches(**filtered_model_params) - else: + if custom_llm_provider not in LIST_BATCHES_SUPPORTED_PROVIDERS: return await litellm.acompletion(**model_params) + return await litellm.alist_batches(**{**filtered_model_params, "custom_llm_provider": custom_llm_provider}) @staticmethod async def _image_edit_health_check(edit_request: Callable[[], Awaitable["ImageResponse"]]) -> "ImageResponse": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f8145cbdc7e..0ef5bbf807a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -379,9 +379,12 @@ _DEPLOYMENT_PRICING_KEYS: Final = ( "output_cost_per_token", "input_cost_per_token_batches", "output_cost_per_token_batches", + "input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches", + "output_cost_per_token_above_200k_tokens_batches", "output_cost_per_token_above_272k_tokens_batches", "cache_read_input_token_cost_batches", + "cache_read_input_token_cost_above_200k_tokens_batches", "cache_read_input_token_cost_above_272k_tokens_batches", "cache_creation_input_token_cost_batches", "cache_creation_input_token_cost_above_272k_tokens_batches", diff --git a/litellm/llms/xai/batches/__init__.py b/litellm/llms/xai/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/xai/batches/handler.py b/litellm/llms/xai/batches/handler.py new file mode 100644 index 00000000000..62db1c4833a --- /dev/null +++ b/litellm/llms/xai/batches/handler.py @@ -0,0 +1,195 @@ +from collections.abc import Coroutine +from itertools import chain +from typing import Final + +import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, +) +from litellm.types.llms.openai import CreateBatchRequest, HttpxBinaryResponseContent +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from .transformation import ( + XAI_RESULTS_PAGE_SIZE, + OpenAIBatchListResponse, + XAIBatch, + XAIBatchList, + XAIBatchResult, + XAIBatchResultsPage, + get_xai_auth_headers, + raise_for_xai_status, + results_to_openai_jsonl, + to_create_batch_body, + to_litellm_batch, + to_openai_batch_list, + xai_batches_url, +) + +_JSONL_CONTENT_TYPE: Final = ("content-type", "application/jsonl") + + +class _PageParams(TypedDict): + limit: ReadOnly[int] + pagination_token: NotRequired[ReadOnly[str]] + + +def _results_params(after: str | None, limit: int | None) -> dict[str, object]: # mutable-ok: httpx params + if after is None: + return dict(_PageParams(limit=limit or XAI_RESULTS_PAGE_SIZE)) # mutable-ok: httpx params + return dict(_PageParams(limit=limit or XAI_RESULTS_PAGE_SIZE, pagination_token=after)) # mutable-ok: httpx params + + +def _flatten(pages: list[XAIBatchResultsPage]) -> tuple[XAIBatchResult, ...]: + return tuple(chain.from_iterable(page.results for page in pages)) + + +def _jsonl_response(url: str, results: tuple[XAIBatchResult, ...]) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=results_to_openai_jsonl(results), + headers=(_JSONL_CONTENT_TYPE,), + request=httpx.Request(method="GET", url=url), + ) + ) + + +class XAIBatchesHandler: + def __init__(self, sync_client: HTTPHandler | None = None, async_client: AsyncHTTPHandler | None = None) -> None: + self._sync_client = sync_client + self._async_client = async_client + + def _sync(self, timeout: float | httpx.Timeout) -> HTTPHandler: + return self._sync_client or HTTPHandler(timeout=timeout) + + def _async(self, timeout: float | httpx.Timeout) -> AsyncHTTPHandler: + return self._async_client or get_async_httpx_client( + llm_provider=LlmProviders.XAI, + params={"timeout": timeout}, # mutable-ok: get_async_httpx_client takes a dict + ) + + def create_batch( + self, + _is_async: bool, + create_batch_data: CreateBatchRequest, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base) + headers: Final = get_xai_auth_headers(api_key=api_key) + body: Final = dict(to_create_batch_body(create_batch_data)) # mutable-ok: httpx json body + endpoint: Final = create_batch_data.get("endpoint") or "/v1/chat/completions" + if _is_async: + + async def _acreate() -> LiteLLMBatch: + response: Final = await self._async(timeout).post(url, json=body, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json()), endpoint) + + return _acreate() + response: Final = self._sync(timeout).post(url, json=body, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json()), endpoint) + + def retrieve_batch( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base, batch_id) + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _aretrieve() -> LiteLLMBatch: + response: Final = await self._async(timeout).get(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + return _aretrieve() + response: Final = self._sync(timeout).get(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + def cancel_batch( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base, batch_id, suffix=":cancel") + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _acancel() -> LiteLLMBatch: + response: Final = await self._async(timeout).post(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + return _acancel() + response: Final = self._sync(timeout).post(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + def list_batches( + self, + _is_async: bool, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + after: str | None = None, + limit: int | None = None, + ) -> OpenAIBatchListResponse | Coroutine[None, None, OpenAIBatchListResponse]: + url: Final = xai_batches_url(api_base) + headers: Final = get_xai_auth_headers(api_key=api_key) + params: Final = _results_params(after, limit) + if _is_async: + + async def _alist() -> OpenAIBatchListResponse: + response: Final = await self._async(timeout).get(url, params=params, headers=headers, timeout=timeout) + return to_openai_batch_list(XAIBatchList.model_validate(raise_for_xai_status(response).json())) + + return _alist() + response: Final = self._sync(timeout).get(url, params=params, headers=headers, timeout=timeout) + return to_openai_batch_list(XAIBatchList.model_validate(raise_for_xai_status(response).json())) + + def batch_results_content( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: + url: Final = xai_batches_url(api_base, batch_id, suffix="/results") + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _aresults() -> HttpxBinaryResponseContent: + client: Final = self._async(timeout) + + async def _page(after: str | None) -> XAIBatchResultsPage: + response: Final = await client.get( + url, params=_results_params(after, None), headers=headers, timeout=timeout + ) + return XAIBatchResultsPage.model_validate(raise_for_xai_status(response).json()) + + pages = [await _page(None)] # mutable-ok: page walk terminates on the cursor, not on a fixed count + while pages[-1].pagination_token and pages[-1].results: + pages.append(await _page(pages[-1].pagination_token)) + return _jsonl_response(url, _flatten(pages)) + + return _aresults() + client: Final = self._sync(timeout) + + def _page(after: str | None) -> XAIBatchResultsPage: + response: Final = client.get(url, params=_results_params(after, None), headers=headers, timeout=timeout) + return XAIBatchResultsPage.model_validate(raise_for_xai_status(response).json()) + + pages = [_page(None)] # mutable-ok: page walk terminates on the cursor, not on a fixed count + while pages[-1].pagination_token and pages[-1].results: + pages.append(_page(pages[-1].pagination_token)) + return _jsonl_response(url, _flatten(pages)) diff --git a/litellm/llms/xai/batches/transformation.py b/litellm/llms/xai/batches/transformation.py new file mode 100644 index 00000000000..8f305b8c203 --- /dev/null +++ b/litellm/llms/xai/batches/transformation.py @@ -0,0 +1,278 @@ +""" +xAI Batch API reference: https://docs.x.ai/developers/advanced-api-usage/batch-api + +xAI batches carry request counters, not a status, and no output file: results are paged from +``GET /v1/batches/{id}/results``, so LiteLLM hands back the batch id as ``output_file_id``. +""" + +import json +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import BaseModel, ConfigDict +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.constants import XAI_API_BASE +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import CreateBatchRequest +from litellm.types.utils import LiteLLMBatch + +OpenAIBatchStatus: TypeAlias = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +XAI_BATCH_ID_PREFIX: Final = "batch_" +XAI_RESULTS_PAGE_SIZE: Final = 1000 +DEFAULT_BATCH_NAME: Final = "litellm-batch" +DEFAULT_BATCH_ENDPOINT: Final = "/v1/chat/completions" +_EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + + +class XAIBatchesError(BaseLLMException): + pass + + +def xai_batches_error( + error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers +) -> XAIBatchesError: + return XAIBatchesError( + status_code=status_code, + message=error_message, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(tuple(headers.items())), + ) + + +def raise_for_xai_status(response: httpx.Response) -> httpx.Response: + if response.status_code >= 400: + raise xai_batches_error(response.text, response.status_code, response.headers) + return response + + +def get_xai_api_base(api_base: str | None) -> str: + resolved: Final = (api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE).rstrip("/") + return resolved.removesuffix("/v1") + + +def get_xai_auth_headers( + headers: Mapping[str, str] = _EMPTY_HEADERS, api_key: str | None = None +) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict + resolved_key: Final = XAIModelInfo.get_api_key(api_key) + if resolved_key is None: + raise xai_batches_error( + "Missing xAI API Key. Pass api_key, set litellm.xai_key or XAI_API_KEY", 401, _EMPTY_HEADERS + ) + return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict + + +def xai_batches_url(api_base: str | None, batch_id: str | None = None, suffix: str = "") -> str: + base: Final = f"{get_xai_api_base(api_base)}/v1/batches" + if batch_id is None: + return base + return f"{base}/{encode_url_path_segment(batch_id, field_name='batch_id')}{suffix}" + + +def is_xai_batch_results_id(file_id: str) -> bool: + return file_id.startswith(XAI_BATCH_ID_PREFIX) + + +class XAICreateBatchRequest(TypedDict): + name: ReadOnly[str] + input_file_id: NotRequired[ReadOnly[str]] + + +class XAIBatchState(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + num_requests: int = 0 + num_pending: int = 0 + num_success: int = 0 + num_error: int = 0 + num_cancelled: int = 0 + + +class XAIBatch(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batch_id: str + name: str = "" + create_time: str | None = None + expire_time: str | None = None + cancel_time: str | None = None + cancel_by_xai_message: str | None = None + state: XAIBatchState = XAIBatchState() + input_file_id: str | None = None + + +class XAIBatchList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batches: tuple[XAIBatch, ...] = () + pagination_token: str | None = None + + +class XAIBatchResultError(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + code: int | str | None = None + message: str = "" + + +class XAIBatchResultData(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + response: Mapping[str, Mapping[str, object]] | None = None + error: XAIBatchResultError | None = None + + +class XAIBatchResult(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batch_request_id: str + batch_result: XAIBatchResultData = XAIBatchResultData() + + +class XAIBatchResultsPage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + results: tuple[XAIBatchResult, ...] = () + pagination_token: str | None = None + + +def _to_unix_timestamp(value: str | None) -> int | None: + """xAI returns RFC 3339 timestamps over gRPC but a bare ``YYYY-MM-DD`` over REST.""" + if value is None: + return None + try: + parsed: Final = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return int((parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc)).timestamp()) + + +def xai_batch_status(batch: XAIBatch) -> OpenAIBatchStatus: + """xAI exposes counters, not a status. A batch xAI itself cancelled (input validation failed) is a failure, + a caller-cancelled batch is cancelled, an empty batch is still validating its input file, and a batch + with nothing pending has completed.""" + if batch.cancel_time is not None: + return "failed" if batch.cancel_by_xai_message else "cancelled" + if batch.state.num_requests == 0: + return "validating" + if batch.state.num_pending > 0: + return "in_progress" + return "completed" + + +def to_litellm_batch(batch: XAIBatch, endpoint: str = DEFAULT_BATCH_ENDPOINT) -> LiteLLMBatch: + status: Final = xai_batch_status(batch) + created_at: Final = _to_unix_timestamp(batch.create_time) + cancelled_at: Final = _to_unix_timestamp(batch.cancel_time) + errors: Final = ( + BatchErrors(object="list", data=[BatchError(message=batch.cancel_by_xai_message)]) # mutable-ok: openai type + if batch.cancel_by_xai_message + else None + ) + return LiteLLMBatch( + id=batch.batch_id, + object="batch", + endpoint=endpoint, + input_file_id=batch.input_file_id or "", + completion_window="24h", + status=status, + created_at=created_at if created_at is not None else 0, + expires_at=_to_unix_timestamp(batch.expire_time), + failed_at=cancelled_at if status == "failed" else None, + cancelled_at=cancelled_at if status == "cancelled" else None, + output_file_id=batch.batch_id if status == "completed" else None, + errors=errors, + request_counts=BatchRequestCounts( + total=batch.state.num_requests, + completed=batch.state.num_success, + failed=batch.state.num_error + batch.state.num_cancelled, + ), + metadata={"name": batch.name} if batch.name else None, # mutable-ok: LiteLLMBatch.metadata is a dict + ) + + +class OpenAIBatchListResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + object: Literal["list"] = "list" + data: tuple[LiteLLMBatch, ...] + first_id: str | None + last_id: str | None + has_more: bool + next_page_token: str | None = None + + +def to_openai_batch_list(page: XAIBatchList) -> OpenAIBatchListResponse: + data: Final = tuple(to_litellm_batch(b) for b in page.batches) + return OpenAIBatchListResponse( + data=data, + first_id=data[0].id if data else None, + last_id=data[-1].id if data else None, + has_more=bool(page.pagination_token), + next_page_token=page.pagination_token or None, + ) + + +def to_create_batch_body(create_batch_data: CreateBatchRequest) -> XAICreateBatchRequest: + input_file_id: Final = create_batch_data.get("input_file_id") + if not input_file_id: + raise xai_batches_error("input_file_id is required to create an xAI batch", 400, _EMPTY_HEADERS) + metadata: Final = create_batch_data.get("metadata") + name: Final = metadata.get("name") if metadata else None + return XAICreateBatchRequest(name=name or DEFAULT_BATCH_NAME, input_file_id=input_file_id) + + +class OpenAIBatchOutputError(TypedDict): + code: ReadOnly[str] + message: ReadOnly[str] + + +class OpenAIBatchOutputResponse(TypedDict): + status_code: ReadOnly[int] + request_id: ReadOnly[object] + body: ReadOnly[Mapping[str, object]] + + +class OpenAIBatchOutputLine(TypedDict): + id: ReadOnly[str] + custom_id: ReadOnly[str] + response: ReadOnly[OpenAIBatchOutputResponse | None] + error: ReadOnly[OpenAIBatchOutputError | None] + + +def _result_to_openai_line(result: XAIBatchResult) -> OpenAIBatchOutputLine: + """One output JSONL line. xAI wraps the body in a one-key map named after the endpoint + (``chat_get_completion``, ``responses``, ``image_generation``, ...); the value is the OpenAI body.""" + error: Final = result.batch_result.error + response: Final = result.batch_result.response + body: Final = next(iter(response.values()), None) if response else None + if body is None: + message: Final = error.message if error is not None else "xAI returned no response for this request" + code: Final = str(error.code) if error is not None and error.code is not None else "request_failed" + return OpenAIBatchOutputLine( + id=f"batch_req_{result.batch_request_id}", + custom_id=result.batch_request_id, + response=None, + error=OpenAIBatchOutputError(code=code, message=message), + ) + return OpenAIBatchOutputLine( + id=f"batch_req_{result.batch_request_id}", + custom_id=result.batch_request_id, + response=OpenAIBatchOutputResponse(status_code=200, request_id=body.get("id"), body=body), + error=None, + ) + + +def results_to_openai_jsonl(results: Sequence[XAIBatchResult]) -> bytes: + return "".join(f"{json.dumps(_result_to_openai_line(r), ensure_ascii=False)}\n" for r in results).encode() diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 33ee727dfab..e686d49e689 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -296,7 +296,7 @@ class XAIChatConfig(OpenAIGPTConfig): except Exception as e: verbose_logger.debug("Error extracting X.AI web search usage: %s", e) - self._fold_reasoning_tokens_into_completion(response) + self.fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None)) if restated_usage is not None: @@ -304,7 +304,7 @@ class XAIChatConfig(OpenAIGPTConfig): return response @staticmethod - def _fold_reasoning_tokens_into_completion( + def fold_reasoning_tokens_into_completion( target: ModelResponse | Usage | dict[str, Any] | None, ) -> None: """Reconcile xAI Usage to the OpenAI invariant. @@ -426,7 +426,7 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] if "usage" in chunk and chunk["usage"] is not None: - XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) + XAIChatConfig.fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) parsed_chunk: Final = super().chunk_parser(chunk) diff --git a/litellm/llms/xai/files/__init__.py b/litellm/llms/xai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/xai/files/transformation.py b/litellm/llms/xai/files/transformation.py new file mode 100644 index 00000000000..dbccca47b25 --- /dev/null +++ b/litellm/llms/xai/files/transformation.py @@ -0,0 +1,247 @@ +""" +xAI Files API reference: https://docs.x.ai/developers/rest-api-reference/inference/files + +xAI stores ``purpose`` as an empty string; LiteLLM reports uploads as ``batch``, the only purpose xAI files serve. +""" + +import time +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + OpenAIFilesPurpose, +) +from litellm.types.utils import LlmProviders + +from ..batches.transformation import ( + get_xai_api_base, + get_xai_auth_headers, + raise_for_xai_status, + xai_batches_error, +) + +_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict] +_DEFAULT_PURPOSE: Final[OpenAIFilesPurpose] = "batch" + + +class XAIMultipartUpload(TypedDict): + file: ReadOnly[tuple[str, object, str]] + purpose: ReadOnly[tuple[None, str]] + + +class XAIFile(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + bytes: int = 0 + created_at: int | None = None + filename: str = "" + purpose: str = "" + expires_at: int | None = None + + +class XAIFileList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[XAIFile, ...] = () + pagination_token: str | None = None + + +class XAIFileDeleted(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + deleted: bool = True + + +def _to_openai_file_object(file: XAIFile) -> OpenAIFileObject: + return OpenAIFileObject( + id=file.id, + bytes=file.bytes, + created_at=file.created_at if file.created_at is not None else int(time.time()), + filename=file.filename, + object="file", + purpose=_DEFAULT_PURPOSE, + status="uploaded", + expires_at=file.expires_at, + ) + + +def _api_base_from(litellm_params: Mapping[str, object]) -> str: + api_base: Final = litellm_params.get("api_base") + return get_xai_api_base(api_base if isinstance(api_base, str) else None) + + +class XAIFilesConfig(BaseFilesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.XAI + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return f"{get_xai_api_base(api_base)}/v1/files" + + def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: + encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") + return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: + return xai_batches_error(error_message, status_code, headers) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature + return get_xai_auth_headers(headers, api_key) + + def get_supported_openai_params( + self, model: str + ) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature + return ["purpose"] # mutable-ok: BaseFilesConfig signature + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + return optional_params + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature + if "file" not in create_file_data: + raise ValueError("File data is required") + extracted: Final = extract_file_data(create_file_data["file"]) + filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" + content_type: Final = extracted.get("content_type") or "application/octet-stream" + upload: Final = XAIMultipartUpload( + file=(filename, extracted["content"], content_type), + purpose=(None, create_file_data.get("purpose") or _DEFAULT_PURPOSE), + ) + return dict(upload) # mutable-ok: BaseFilesConfig signature + + def transform_create_file_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> OpenAIFileObject: + return _to_openai_file_object(XAIFile.model_validate(raise_for_xai_status(raw_response).json())) + + def transform_retrieve_file_request( + self, + file_id: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS + + def transform_retrieve_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> OpenAIFileObject: + return _to_openai_file_object(XAIFile.model_validate(raise_for_xai_status(raw_response).json())) + + def transform_delete_file_request( + self, + file_id: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS + + def transform_delete_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> FileDeleted: + deleted: Final = XAIFileDeleted.model_validate(raise_for_xai_status(raw_response).json()) + return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") + + def transform_list_files_request( + self, + purpose: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return f"{_api_base_from(litellm_params)}/v1/files", _NO_QUERY_PARAMS + + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]] | None: # mutable-ok: BaseFilesConfig signature + page: Final = XAIFileList.model_validate(raw_response.json()) + if not page.pagination_token or not page.data: + return None + return f"{_api_base_from(litellm_params)}/v1/files", {"pagination_token": page.pagination_token} + + def transform_list_files_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature + return [ # mutable-ok: BaseFilesConfig signature + _to_openai_file_object(f) + for f in XAIFileList.model_validate(raise_for_xai_status(raw_response).json()).data + ] + + def transform_file_content_request( + self, + file_content_request: FileContentRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS + + def transform_file_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a207dc4c02..a670200c132 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -51436,13 +51436,16 @@ }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -51450,8 +51453,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -51480,9 +51486,13 @@ "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51490,6 +51500,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -51502,9 +51514,13 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51512,6 +51528,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59483,13 +59501,16 @@ }, "xai/grok-4.20-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59497,20 +59518,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, @@ -59519,8 +59546,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ @@ -62787,13 +62817,16 @@ }, "xai/grok-4.20": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62801,21 +62834,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62823,21 +62862,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62845,8 +62890,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -63067,13 +63115,16 @@ }, "xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63081,20 +63132,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63102,20 +63159,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63127,20 +63190,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63152,8 +63221,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, @@ -75459,13 +75531,16 @@ }, "xai/grok-4.20-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -75473,8 +75548,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 6e7e9da3498..99ab5920c4f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -522,7 +522,19 @@ class CreateBatchRequest(TypedDict, total=False): """ completion_window: Literal["24h"] - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"] + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ] input_file_id: str metadata: dict[str, str] | None output_expires_after: FileExpiresAfter diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3e306b48887..cd336c9b989 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -299,6 +299,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens_flex: float | None cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] + cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] @@ -327,8 +328,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_second: float | None # for OpenAI Speech models input_cost_per_token_batches: float | None input_cost_per_video_token_batches: ReadOnly[float | None] + input_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None] input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None] output_cost_per_token_batches: float | None + output_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None] output_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None] output_cost_per_token: Required[float | None] output_cost_per_token_flex: float | None # OpenAI flex service tier pricing @@ -3729,6 +3732,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_token_cost_batches: float | None = None + cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None cache_creation_input_token_cost_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None @@ -3742,6 +3746,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_above_200k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_flex: float | None = None + input_cost_per_token_above_200k_tokens_batches: float | None = None input_cost_per_token_above_272k_tokens_batches: float | None = None input_cost_per_query: float | None = None input_cost_per_image: float | None = None @@ -3766,6 +3771,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_200k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_flex: float | None = None + output_cost_per_token_above_200k_tokens_batches: float | None = None output_cost_per_token_above_272k_tokens_batches: float | None = None output_cost_per_character_above_128k_tokens: float | None = None output_cost_per_image: float | None = None @@ -4140,7 +4146,7 @@ FILE_CONTENT_STREAMING_PROVIDERS: Final[frozenset[str]] = frozenset( LITELLM_EXECUTED_BATCH_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.HOSTED_VLLM.value}) -ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai"] +ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai", "xai"] LIST_BATCHES_SUPPORTED_PROVIDERS: Final[frozenset[str]] = frozenset(get_args(ListBatchesSupportedProvider)) diff --git a/litellm/utils.py b/litellm/utils.py index 4ea0769ea11..be4388802f9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6160,6 +6160,9 @@ def _get_model_info_helper( cache_read_input_token_cost_priority=_model_info.get("cache_read_input_token_cost_priority", None), cache_read_input_token_cost_ultrafast=_model_info.get("cache_read_input_token_cost_ultrafast", None), cache_read_input_token_cost_batches=_model_info.get("cache_read_input_token_cost_batches"), + cache_read_input_token_cost_above_200k_tokens_batches=_model_info.get( + "cache_read_input_token_cost_above_200k_tokens_batches" + ), cache_read_input_token_cost_above_272k_tokens_batches=_model_info.get( "cache_read_input_token_cost_above_272k_tokens_batches" ), @@ -6197,10 +6200,16 @@ def _get_model_info_helper( input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"), input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None), + input_cost_per_token_above_200k_tokens_batches=_model_info.get( + "input_cost_per_token_above_200k_tokens_batches" + ), input_cost_per_token_above_272k_tokens_batches=_model_info.get( "input_cost_per_token_above_272k_tokens_batches" ), output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"), + output_cost_per_token_above_200k_tokens_batches=_model_info.get( + "output_cost_per_token_above_200k_tokens_batches" + ), output_cost_per_token_above_272k_tokens_batches=_model_info.get( "output_cost_per_token_above_272k_tokens_batches" ), @@ -9357,6 +9366,10 @@ class ProviderConfigManager: from litellm.llms.mistral.files.transformation import MistralFilesConfig return MistralFilesConfig() + elif LlmProviders.XAI == provider: + from litellm.llms.xai.files.transformation import XAIFilesConfig + + return XAIFilesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a207dc4c02..a670200c132 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -51436,13 +51436,16 @@ }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -51450,8 +51453,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -51480,9 +51486,13 @@ "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51490,6 +51500,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -51502,9 +51514,13 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51512,6 +51528,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59483,13 +59501,16 @@ }, "xai/grok-4.20-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59497,20 +59518,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, @@ -59519,8 +59546,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ @@ -62787,13 +62817,16 @@ }, "xai/grok-4.20": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62801,21 +62834,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62823,21 +62862,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62845,8 +62890,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -63067,13 +63115,16 @@ }, "xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63081,20 +63132,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63102,20 +63159,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63127,20 +63190,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63152,8 +63221,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, @@ -75459,13 +75531,16 @@ }, "xai/grok-4.20-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -75473,8 +75548,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 35624045fdf..fa1828c780a 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -170,6 +170,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_read_input_token_cost_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "cache_read_input_token_cost_above_200k_tokens_priority": { "type": "number", "minimum": 0, @@ -351,6 +356,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_token_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "input_cost_per_token_above_200k_tokens_priority": { "type": "number", "minimum": 0, @@ -708,6 +718,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "output_cost_per_token_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "output_cost_per_token_above_200k_tokens_priority": { "type": "number", "minimum": 0, diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index 1cc96cb1256..c3478c0d5eb 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py @@ -364,6 +364,26 @@ async def test_batch_health_check_uses_alist_batches_for_supported_providers(): mock_alist.assert_called_once() +@pytest.mark.asyncio +async def test_batch_health_check_hands_the_resolved_provider_to_alist_batches(): + filtered_model_params: Final = { + "model": "xai/grok-4.3", + "api_key": "sk-test", + "litellm_metadata": {"tags": [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME]}, + } + + with patch("litellm.alist_batches", new_callable=AsyncMock, return_value={}) as mock_alist: + await HealthCheckHelpers._batch_health_check( + custom_llm_provider="xai", + model_params={**filtered_model_params, "messages": []}, + filtered_model_params=filtered_model_params, + ) + + assert mock_alist.call_args.kwargs["custom_llm_provider"] == "xai" + assert mock_alist.call_args.kwargs["model"] == "xai/grok-4.3" + assert mock_alist.call_args.kwargs["api_key"] == "sk-test" + + @pytest.mark.asyncio async def test_batch_health_check_falls_back_to_acompletion_for_unsupported(): """Providers not in LIST_BATCHES_SUPPORTED_PROVIDERS fall back to acompletion.""" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 166eeb53f5f..265fdb50836 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -7781,6 +7781,9 @@ _PUBLISHED_BATCH_RATES: Final = MappingProxyType( "output_cost_per_token_batches": 4.1e-6, "cache_read_input_token_cost_batches": 1.2e-7, "cache_creation_input_token_cost_batches": 1.3e-6, + "input_cost_per_token_above_200k_tokens_batches": 2.1e-6, + "output_cost_per_token_above_200k_tokens_batches": 5.1e-6, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.2e-7, "input_cost_per_token_above_272k_tokens_batches": 3.1e-6, "output_cost_per_token_above_272k_tokens_batches": 7.1e-6, "cache_read_input_token_cost_above_272k_tokens_batches": 3.2e-7, @@ -7789,14 +7792,17 @@ _PUBLISHED_BATCH_RATES: Final = MappingProxyType( ) _PUBLISHED_INPUT_BATCH_KEYS: Final = ( "input_cost_per_token_batches", + "input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches", "cache_read_input_token_cost_batches", + "cache_read_input_token_cost_above_200k_tokens_batches", "cache_read_input_token_cost_above_272k_tokens_batches", "cache_creation_input_token_cost_batches", "cache_creation_input_token_cost_above_272k_tokens_batches", ) _PUBLISHED_OUTPUT_BATCH_KEYS: Final = ( "output_cost_per_token_batches", + "output_cost_per_token_above_200k_tokens_batches", "output_cost_per_token_above_272k_tokens_batches", ) @@ -7885,22 +7891,22 @@ def test_batch_cost_calculator_bills_the_carried_output_tier_when_the_deployment ) +@pytest.mark.parametrize( + "tier_key", + ["input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches"], +) def test_deployment_pricing_model_info_honors_a_tier_only_batch_override_over_the_published_flat_rates( - _published_batch_model: None, + _published_batch_model: None, tier_key: str ) -> None: from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info - info: Final = deployment_pricing_model_info( - _batch_deployment_id({"input_cost_per_token_above_272k_tokens_batches": 1e-3}), _PUBLISHED_BATCH_DEPLOYMENT - ) + info: Final = deployment_pricing_model_info(_batch_deployment_id({tier_key: 1e-3}), _PUBLISHED_BATCH_DEPLOYMENT) carried_keys: Final = tuple( - key - for key in (*_PUBLISHED_INPUT_BATCH_KEYS, *_PUBLISHED_OUTPUT_BATCH_KEYS) - if key != "input_cost_per_token_above_272k_tokens_batches" + key for key in (*_PUBLISHED_INPUT_BATCH_KEYS, *_PUBLISHED_OUTPUT_BATCH_KEYS) if key != tier_key ) assert info is not None - assert info["input_cost_per_token_above_272k_tokens_batches"] == 1e-3 + assert info[tier_key] == 1e-3 assert {key: info[key] for key in carried_keys} == {key: _PUBLISHED_BATCH_RATES[key] for key in carried_keys} diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index dd95addac40..b8b922f72a7 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -464,6 +464,40 @@ def test_total_cost_applies_the_long_context_batch_tier_per_line(): assert result.cost == pytest.approx((300_000 * 2e-6) + (10 * 6e-6) + (100 * 1e-6) + (10 * 4e-6)) +def test_xai_output_lines_bill_reasoning_tokens_as_completion_tokens(): + row = _success_row( + model="grok-4.3", + usage={ + "prompt_tokens": 615, + "completion_tokens": 3, + "total_tokens": 993, + "completion_tokens_details": {"reasoning_tokens": 375}, + }, + ) + + result = bu._aggregate_batch_cost_usage_models( + entries=[row], + custom_llm_provider="xai", + model_info=ModelInfo( + key="xai/grok-4.3", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=1.25e-6, + output_cost_per_token=2.5e-6, + litellm_provider="xai", + mode="chat", + supported_openai_params=None, + input_cost_per_token_batches=1e-6, + output_cost_per_token_batches=2e-6, + ), + ) + + assert result.usage.completion_tokens == 378 + assert result.usage.total_tokens == 993 + assert result.cost == pytest.approx((615 * 1e-6) + (378 * 2e-6)) + + def test_total_usage_empty_is_zero(): result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") assert result.cost == 0.0 diff --git a/tests/unit/llms/xai/batches/__init__.py b/tests/unit/llms/xai/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/batches/test_xai_batches_handler.py b/tests/unit/llms/xai/batches/test_xai_batches_handler.py new file mode 100644 index 00000000000..6dcdf06e7ab --- /dev/null +++ b/tests/unit/llms/xai/batches/test_xai_batches_handler.py @@ -0,0 +1,344 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.xai.batches.transformation import XAIBatchesError +from litellm.types.utils import LiteLLMBatch + +API_BASE: Final = "https://api.x.ai" +KEY: Final = "xai-test-key" + + +@pytest.fixture(autouse=True) +def _httpx_transport_so_respx_can_intercept(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +_XAI_BATCH: Final = { + "batch_id": "batch_1", + "name": "litellm-batch", + "create_time": "2026-09-23", + "expire_time": "2026-10-23", + "cancel_time": None, + "cancel_by_xai_message": None, + "state": {"num_requests": 2, "num_pending": 0, "num_success": 2, "num_error": 0, "num_cancelled": 0}, + "input_file_id": "file_1", +} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_create_batch_posts_input_file_id_with_bearer_auth(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + kwargs: Final = { + "completion_window": "24h", + "endpoint": "/v1/embeddings", + "input_file_id": "file_1", + "custom_llm_provider": "xai", + "api_key": KEY, + "api_base": API_BASE, + } + batch: Final = litellm.create_batch(**kwargs) if sync_mode else await litellm.acreate_batch(**kwargs) + + assert isinstance(batch, LiteLLMBatch) + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {KEY}" + assert json.loads(request.content) == {"name": "litellm-batch", "input_file_id": "file_1"} + assert (batch.id, batch.endpoint, batch.status, batch.output_file_id) == ( + "batch_1", + "/v1/embeddings", + "completed", + "batch_1", + ) + + +@pytest.mark.parametrize( + "endpoint", + [ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], +) +@respx.mock +async def test_create_batch_keeps_image_and_video_endpoints_on_the_batch(endpoint: str) -> None: + respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + batch: Final = await litellm.acreate_batch( + completion_window="24h", + endpoint=endpoint, + input_file_id="file_1", + custom_llm_provider="xai", + api_key=KEY, + api_base=API_BASE, + ) + + assert isinstance(batch, LiteLLMBatch) + assert batch.endpoint == endpoint + assert json.loads(respx.calls.last.request.content) == {"name": "litellm-batch", "input_file_id": "file_1"} + + +@respx.mock +async def test_retrieve_after_a_non_chat_create_reports_chat() -> None: + respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + respx.get(f"{API_BASE}/v1/batches/batch_1").respond(200, json=_XAI_BATCH) + + created: Final = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/embeddings", + input_file_id="file_1", + custom_llm_provider="xai", + api_key=KEY, + api_base=API_BASE, + ) + retrieved: Final = await litellm.aretrieve_batch( + batch_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert isinstance(created, LiteLLMBatch) and isinstance(retrieved, LiteLLMBatch) + assert (created.endpoint, retrieved.endpoint) == ("/v1/embeddings", "/v1/chat/completions") + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_retrieve_batch_reads_native_batch_route(sync_mode: bool) -> None: + respx.get(f"{API_BASE}/v1/batches/batch_1").respond( + 200, json={**_XAI_BATCH, "state": {"num_requests": 2, "num_pending": 2}} + ) + + kwargs: Final = {"batch_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + batch: Final = litellm.retrieve_batch(**kwargs) if sync_mode else await litellm.aretrieve_batch(**kwargs) + + assert isinstance(batch, LiteLLMBatch) + assert (batch.status, batch.output_file_id, batch.input_file_id, batch.endpoint) == ( + "in_progress", + None, + "file_1", + "/v1/chat/completions", + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_cancel_batch_uses_colon_cancel_route(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/batches/batch_1:cancel").respond( + 200, json={**_XAI_BATCH, "cancel_time": "2026-09-23", "state": {}} + ) + + kwargs: Final = {"batch_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + batch: Final = litellm.cancel_batch(**kwargs) if sync_mode else await litellm.acancel_batch(**kwargs) + + assert route.called + assert isinstance(batch, LiteLLMBatch) + assert (batch.status, batch.endpoint) == ("cancelled", "/v1/chat/completions") + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_list_batches_forwards_cursor_and_returns_openai_list(sync_mode: bool) -> None: + route: Final = respx.get(f"{API_BASE}/v1/batches").respond( + 200, json={"batches": [_XAI_BATCH], "pagination_token": "next"} + ) + + kwargs: Final = {"custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE, "after": "cur", "limit": 5} + listed: Final = litellm.list_batches(**kwargs) if sync_mode else await litellm.alist_batches(**kwargs) + + assert dict(route.calls.last.request.url.params) == {"limit": "5", "pagination_token": "cur"} + assert listed.object == "list" + assert [(b.id, b.endpoint) for b in listed.data] == [("batch_1", "/v1/chat/completions")] + assert (listed.has_more, listed.next_page_token) == (True, "next") + + +@respx.mock +async def test_list_batches_treats_empty_pagination_token_as_last_page() -> None: + respx.get(f"{API_BASE}/v1/batches").respond(200, json={"batches": [_XAI_BATCH], "pagination_token": ""}) + + listed: Final = await litellm.alist_batches(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + assert (listed.has_more, listed.next_page_token) == (False, None) + assert [batch.endpoint for batch in listed.data] == ["/v1/chat/completions"] + + +@respx.mock +async def test_file_content_stops_paging_on_empty_pagination_token() -> None: + route: Final = respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, + json={ + "results": [{"batch_request_id": "r1", "batch_result": {"error": {"code": 3, "message": "boom"}}}], + "pagination_token": "", + }, + ) + + content: Final = await litellm.afile_content( + file_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert route.call_count == 1 + assert len(content.content.decode().splitlines()) == 1 + + +@pytest.mark.parametrize("operation", ["create", "retrieve", "cancel", "list", "file_content"]) +@respx.mock +async def test_batch_calls_fall_back_to_litellm_xai_key(operation: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", "configured-xai-key") + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + routes: Final = { + "create": respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH), + "retrieve": respx.get(f"{API_BASE}/v1/batches/batch_1").respond(200, json=_XAI_BATCH), + "cancel": respx.post(f"{API_BASE}/v1/batches/batch_1:cancel").respond(200, json=_XAI_BATCH), + "list": respx.get(f"{API_BASE}/v1/batches").respond( + 200, json={"batches": [_XAI_BATCH], "pagination_token": None} + ), + "file_content": respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, json={"results": [], "pagination_token": None} + ), + } + + if operation == "create": + await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file_1", + custom_llm_provider="xai", + api_base=API_BASE, + ) + elif operation == "retrieve": + await litellm.aretrieve_batch(batch_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + elif operation == "cancel": + await litellm.acancel_batch(batch_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + elif operation == "list": + await litellm.alist_batches(custom_llm_provider="xai", api_base=API_BASE) + else: + await litellm.afile_content(file_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + + assert routes[operation].calls.last.request.headers["authorization"] == "Bearer configured-xai-key" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_file_content_of_a_batch_id_walks_every_results_page(sync_mode: bool) -> None: + def _page(request: httpx.Request) -> httpx.Response: + token: Final = request.url.params.get("pagination_token") + if token is None: + return httpx.Response( + 200, + json={ + "results": [ + { + "batch_request_id": "r1", + "batch_result": {"response": {"chat_get_completion": {"id": "c1", "choices": []}}}, + } + ], + "pagination_token": "r1", + }, + ) + assert token == "r1" + return httpx.Response( + 200, + json={ + "results": [ + {"batch_request_id": "r2", "batch_result": {"error": {"code": 3, "message": "boom"}}}, + ], + "pagination_token": None, + }, + ) + + route: Final = respx.get(f"{API_BASE}/v1/batches/batch_1/results").mock(side_effect=_page) + + kwargs: Final = {"file_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + content: Final = litellm.file_content(**kwargs) if sync_mode else await litellm.afile_content(**kwargs) + + assert route.call_count == 2 + assert [dict(c.request.url.params) for c in route.calls] == [ + {"limit": "1000"}, + {"limit": "1000", "pagination_token": "r1"}, + ] + assert [json.loads(line) for line in content.content.decode().splitlines()] == [ + { + "id": "batch_req_r1", + "custom_id": "r1", + "response": {"status_code": 200, "request_id": "c1", "body": {"id": "c1", "choices": []}}, + "error": None, + }, + {"id": "batch_req_r2", "custom_id": "r2", "response": None, "error": {"code": "3", "message": "boom"}}, + ] + + +@respx.mock +async def test_file_content_unwraps_image_and_video_result_bodies() -> None: + respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, + json={ + "results": [ + { + "batch_request_id": "img", + "batch_result": { + "response": {"image_generation": {"data": [{"url": "https://cdn.example/img.png"}]}} + }, + }, + { + "batch_request_id": "vid", + "batch_result": { + "response": {"video_generation": {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}} + }, + }, + ], + "pagination_token": None, + }, + ) + + content: Final = await litellm.afile_content( + file_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert [json.loads(line)["response"]["body"] for line in content.content.decode().splitlines()] == [ + {"data": [{"url": "https://cdn.example/img.png"}]}, + {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}, + ] + + +@respx.mock +async def test_missing_xai_key_is_a_401_before_any_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + route: Final = respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + with pytest.raises(XAIBatchesError) as exc: + await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file_1", + custom_llm_provider="xai", + api_base=API_BASE, + ) + + assert exc.value.status_code == 401 + assert route.called is False + + +@respx.mock +async def test_upstream_error_surfaces_status_code_and_body() -> None: + respx.get(f"{API_BASE}/v1/batches/batch_missing").respond(404, json={"code": "404", "error": "not found"}) + + with pytest.raises(XAIBatchesError) as exc: + await litellm.aretrieve_batch( + batch_id="batch_missing", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert exc.value.status_code == 404 + assert "not found" in exc.value.message diff --git a/tests/unit/llms/xai/batches/test_xai_batches_transformation.py b/tests/unit/llms/xai/batches/test_xai_batches_transformation.py new file mode 100644 index 00000000000..5f2bb6a33ce --- /dev/null +++ b/tests/unit/llms/xai/batches/test_xai_batches_transformation.py @@ -0,0 +1,224 @@ +import json +from typing import Final + +import pytest + +from litellm.llms.xai.batches.transformation import ( + XAIBatch, + XAIBatchesError, + XAIBatchList, + XAIBatchResult, + XAIBatchResultsPage, + get_xai_api_base, + results_to_openai_jsonl, + to_create_batch_body, + to_litellm_batch, + to_openai_batch_list, + xai_batches_url, +) +from litellm.types.llms.openai import CreateBatchRequest + +SEPT_23_2026_UTC: Final = 1790121600 + + +def _xai_batch(**overrides: object) -> XAIBatch: + return XAIBatch.model_validate( + { + "batch_id": "batch_9bdf", + "name": "nightly", + "create_time": "2026-09-23", + "expire_time": "2026-10-23", + "cancel_time": None, + "cancel_by_xai_message": None, + "state": {"num_requests": 2, "num_pending": 0, "num_success": 2, "num_error": 0, "num_cancelled": 0}, + "input_file_id": "file_07", + **overrides, + } + ) + + +def test_completed_batch_exposes_batch_id_as_output_file_and_maps_counts() -> None: + batch: Final = to_litellm_batch(_xai_batch()) + + assert batch.model_dump(exclude_none=True) == { + "id": "batch_9bdf", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file_07", + "completion_window": "24h", + "status": "completed", + "created_at": SEPT_23_2026_UTC, + "expires_at": SEPT_23_2026_UTC + 30 * 86400, + "output_file_id": "batch_9bdf", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + "metadata": {"name": "nightly"}, + } + + +def test_pending_requests_mean_in_progress_and_no_output_file() -> None: + batch: Final = to_litellm_batch( + _xai_batch(state={"num_requests": 3, "num_pending": 1, "num_success": 1, "num_error": 1, "num_cancelled": 0}) + ) + + assert (batch.status, batch.output_file_id) == ("in_progress", None) + assert batch.request_counts is not None + assert batch.request_counts.model_dump() == {"total": 3, "completed": 1, "failed": 1} + + +def test_empty_batch_is_still_validating() -> None: + assert to_litellm_batch(_xai_batch(state={})).status == "validating" + + +def test_batch_cancelled_by_xai_validation_is_failed_with_the_message() -> None: + batch: Final = to_litellm_batch( + _xai_batch( + state={}, + cancel_time="2026-09-23T10:00:00Z", + cancel_by_xai_message="JSONL file validation failed: Model grok-nope is not supported", + ) + ) + + assert batch.status == "failed" + assert batch.failed_at == SEPT_23_2026_UTC + 10 * 3600 + assert batch.cancelled_at is None + assert batch.errors is not None and batch.errors.data is not None + assert [e.message for e in batch.errors.data] == ["JSONL file validation failed: Model grok-nope is not supported"] + + +def test_batch_cancelled_by_caller_is_cancelled() -> None: + batch: Final = to_litellm_batch(_xai_batch(cancel_time="2026-09-23")) + + assert (batch.status, batch.cancelled_at, batch.errors) == ("cancelled", SEPT_23_2026_UTC, None) + + +@pytest.mark.parametrize( + "endpoint", + [ + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos/edits", + "/v1/videos/extensions", + ], +) +def test_create_body_accepts_image_and_video_endpoints(endpoint: str) -> None: + body: Final = to_create_batch_body( + CreateBatchRequest(completion_window="24h", endpoint=endpoint, input_file_id="file_07") + ) + + assert dict(body) == {"name": "litellm-batch", "input_file_id": "file_07"} + + +def test_create_body_uses_input_file_id_and_metadata_name() -> None: + body: Final = to_create_batch_body( + CreateBatchRequest( + completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file_07", metadata={"name": "n1"} + ) + ) + + assert dict(body) == {"name": "n1", "input_file_id": "file_07"} + + +def test_create_body_without_input_file_id_is_a_400() -> None: + with pytest.raises(XAIBatchesError) as exc: + to_create_batch_body(CreateBatchRequest(completion_window="24h", endpoint="/v1/chat/completions")) + + assert exc.value.status_code == 400 + + +def test_results_render_as_openai_output_jsonl_with_errors_per_line() -> None: + page: Final = XAIBatchResultsPage.model_validate( + { + "results": [ + { + "batch_request_id": "r1", + "batch_result": { + "response": { + "chat_get_completion": {"id": "c1", "object": "chat.completion", "choices": [], "usage": {}} + } + }, + }, + {"batch_request_id": "r2", "batch_result": {"error": {"code": 3, "message": "bad model"}}}, + {"batch_request_id": "r3", "batch_result": {}}, + ], + "pagination_token": None, + } + ) + + lines: Final = [json.loads(line) for line in results_to_openai_jsonl(page.results).decode().splitlines()] + + assert lines == [ + { + "id": "batch_req_r1", + "custom_id": "r1", + "response": { + "status_code": 200, + "request_id": "c1", + "body": {"id": "c1", "object": "chat.completion", "choices": [], "usage": {}}, + }, + "error": None, + }, + {"id": "batch_req_r2", "custom_id": "r2", "response": None, "error": {"code": "3", "message": "bad model"}}, + { + "id": "batch_req_r3", + "custom_id": "r3", + "response": None, + "error": {"code": "request_failed", "message": "xAI returned no response for this request"}, + }, + ] + + +@pytest.mark.parametrize( + ("response_key", "body"), + [ + ("responses", {"id": "resp_1", "output": []}), + ("image_generation", {"created": 1, "data": [{"url": "https://cdn.example/img.png"}]}), + ("video_generation", {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}), + ], +) +def test_result_unwraps_the_single_response_key_into_the_openai_body( + response_key: str, body: dict[str, object] +) -> None: + result: Final = XAIBatchResult.model_validate( + {"batch_request_id": "r", "batch_result": {"response": {response_key: body}}} + ) + + line: Final = json.loads(results_to_openai_jsonl((result,)).decode()) + assert line["response"]["body"] == body + assert line["response"]["request_id"] == body.get("id") + assert response_key not in line["response"]["body"] + + +def test_retrieve_and_list_report_chat_because_xai_has_no_batch_endpoint() -> None: + retrieved: Final = to_litellm_batch(_xai_batch()) + listed: Final = to_openai_batch_list(XAIBatchList.model_validate({"batches": [_xai_batch().model_dump()]})) + + assert retrieved.endpoint == "/v1/chat/completions" + assert [batch.endpoint for batch in listed.data] == ["/v1/chat/completions"] + assert retrieved.metadata == {"name": "nightly"} + + +def test_list_page_maps_to_openai_list_with_cursor_flags() -> None: + page: Final = XAIBatchList.model_validate( + {"batches": [_xai_batch().model_dump(), _xai_batch(batch_id="batch_2").model_dump()], "pagination_token": "t"} + ) + + listed: Final = to_openai_batch_list(page) + + assert (listed.object, listed.first_id, listed.last_id, listed.has_more, listed.next_page_token) == ( + "list", + "batch_9bdf", + "batch_2", + True, + "t", + ) + assert [b.id for b in listed.data] == ["batch_9bdf", "batch_2"] + + +@pytest.mark.parametrize( + "api_base", ["https://api.x.ai", "https://api.x.ai/", "https://api.x.ai/v1", "https://api.x.ai/v1/"] +) +def test_api_base_never_doubles_the_v1_segment(api_base: str) -> None: + assert get_xai_api_base(api_base) == "https://api.x.ai" + assert xai_batches_url(api_base, "batch_1", ":cancel") == "https://api.x.ai/v1/batches/batch_1:cancel" + assert xai_batches_url(api_base) == "https://api.x.ai/v1/batches" diff --git a/tests/unit/llms/xai/files/__init__.py b/tests/unit/llms/xai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/files/test_xai_files_transformation.py b/tests/unit/llms/xai/files/test_xai_files_transformation.py new file mode 100644 index 00000000000..5a7d86bdfb7 --- /dev/null +++ b/tests/unit/llms/xai/files/test_xai_files_transformation.py @@ -0,0 +1,144 @@ +from typing import Final + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import OpenAIFileObject + +API_BASE: Final = "https://api.x.ai" +KEY: Final = "xai-test-key" + + +@pytest.fixture(autouse=True) +def _httpx_transport_so_respx_can_intercept(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +_XAI_FILE: Final = { + "bytes": 337, + "created_at": 1790197740, + "expires_at": None, + "filename": "batch.jsonl", + "id": "file_07", + "object": "file", + "purpose": "", +} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_create_file_uploads_multipart_to_xai_and_reports_batch_purpose(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/files").respond(200, json=_XAI_FILE) + + kwargs: Final = { + "file": ("batch.jsonl", b'{"custom_id":"r1"}\n', "application/jsonl"), + "purpose": "batch", + "custom_llm_provider": "xai", + "api_key": KEY, + "api_base": API_BASE, + } + created: Final = litellm.create_file(**kwargs) if sync_mode else await litellm.acreate_file(**kwargs) + + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {KEY}" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'filename="batch.jsonl"' in request.content + assert b'{"custom_id":"r1"}' in request.content + assert created.model_dump(exclude_none=True) == { + "id": "file_07", + "bytes": 337, + "created_at": 1790197740, + "filename": "batch.jsonl", + "object": "file", + "purpose": "batch", + "status": "uploaded", + } + + +@respx.mock +async def test_file_content_of_an_uploaded_file_downloads_original_bytes() -> None: + respx.get(f"{API_BASE}/v1/files/file_07/content").respond(200, content=b'{"custom_id":"r1"}\n') + + content: Final = await litellm.afile_content( + file_id="file_07", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert content.content == b'{"custom_id":"r1"}\n' + + +@respx.mock +async def test_delete_file_maps_xai_deleted_object() -> None: + respx.delete(f"{API_BASE}/v1/files/file_07").respond(200, json={"id": "file_07", "deleted": True, "object": "file"}) + + deleted: Final = await litellm.afile_delete( + file_id="file_07", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert deleted.model_dump() == {"id": "file_07", "deleted": True, "object": "file"} + + +@respx.mock +async def test_create_file_falls_back_to_litellm_xai_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", "configured-xai-key") + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + route: Final = respx.post(f"{API_BASE}/v1/files").respond(200, json=_XAI_FILE) + + await litellm.acreate_file( + file=("batch.jsonl", b'{"custom_id":"r1"}\n', "application/jsonl"), + purpose="batch", + custom_llm_provider="xai", + api_base=API_BASE, + ) + + assert route.calls.last.request.headers["authorization"] == "Bearer configured-xai-key" + + +@respx.mock +async def test_list_files_reads_data_array() -> None: + respx.get(f"{API_BASE}/v1/files").respond(200, json={"data": [_XAI_FILE], "pagination_token": None}) + + listed: Final = await litellm.afile_list(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + files: Final = TypeAdapter(tuple[OpenAIFileObject, ...]).validate_python(listed) + assert [f.id for f in files] == ["file_07"] + + +@respx.mock +async def test_list_files_walks_every_page_by_pagination_token() -> None: + route: Final = respx.get(f"{API_BASE}/v1/files").mock( + side_effect=[ + httpx.Response(200, json={"data": [_XAI_FILE], "pagination_token": "file_07"}), + httpx.Response(200, json={"data": [{**_XAI_FILE, "id": "file_08"}], "pagination_token": "file_08"}), + httpx.Response(200, json={"data": [], "pagination_token": "file_08"}), + ] + ) + + listed: Final = await litellm.afile_list(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + files: Final = TypeAdapter(tuple[OpenAIFileObject, ...]).validate_python(listed) + assert [f.id for f in files] == ["file_07", "file_08"] + assert [call.request.url.params.get("pagination_token") for call in route.calls] == [None, "file_07", "file_08"] + + +async def _retrieve_file(sync_mode: bool, file_id: str) -> None: + if sync_mode: + litellm.file_retrieve(file_id=file_id, custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + return + await litellm.afile_retrieve(file_id=file_id, custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_retrieve_file_maps_xai_not_found_to_a_404_error(sync_mode: bool) -> None: + respx.get(f"{API_BASE}/v1/files/file_gone").respond(404, json={"code": "not-found", "error": "File not found"}) + + with pytest.raises(BaseLLMException) as raised: + await _retrieve_file(sync_mode, "file_gone") + + assert raised.value.status_code == 404 + assert "File not found" in str(raised.value) diff --git a/tests/unit/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py index 3fd666e4f50..704d0061103 100644 --- a/tests/unit/llms/xai/test_xai_chat_transformation.py +++ b/tests/unit/llms/xai/test_xai_chat_transformation.py @@ -16,7 +16,7 @@ from litellm.types.utils import ( class TestXAIReasoningTokenFolding: - """``_fold_reasoning_tokens_into_completion`` re-aligns xAI Usage to the OpenAI invariant.""" + """``fold_reasoning_tokens_into_completion`` re-aligns xAI Usage to the OpenAI invariant.""" @staticmethod def _make_response( @@ -45,7 +45,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) usage = response.usage assert usage.completion_tokens == 322 @@ -59,7 +59,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 322 @@ -71,7 +71,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=0, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 10 @@ -84,7 +84,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 10 assert response.usage.total_tokens == 999 diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 99dea6366f9..62ef9f11c2e 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -4581,6 +4581,55 @@ def test_every_openai_entry_with_a_long_context_rate_and_a_batch_rate_declares_t assert undeclared == [] +@pytest.mark.parametrize("prefix", _BATCH_RATE_PREFIXES) +def test_every_xai_entry_with_a_long_context_rate_and_a_batch_rate_declares_the_batch_tier( + _local_model_cost_map: None, prefix: str +) -> None: + undeclared: Final = [ + name + for name, entry in litellm.model_cost.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") == "xai" + and entry.get(f"{prefix}_above_200k_tokens") is not None + and entry.get(f"{prefix}_batches") is not None + and entry.get(f"{prefix}_above_200k_tokens_batches") is None + ] + + assert undeclared == [] + + +_XAI_TIERED_BATCH_MODEL: Final = "xai/grok-4.3" + + +def test_xai_batch_tier_discounts_the_long_context_rate_like_the_flat_batch_rate(_local_model_cost_map: None) -> None: + info: Final = litellm.get_model_info(_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai") + flat_discount: Final = info["input_cost_per_token_batches"] / info["input_cost_per_token"] + + for prefix in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + tier_discount = info[f"{prefix}_above_200k_tokens_batches"] / info[f"{prefix}_above_200k_tokens"] + assert tier_discount == pytest.approx(flat_discount) + assert info[f"{prefix}_above_200k_tokens_batches"] < info[f"{prefix}_above_200k_tokens"] + + +@pytest.mark.parametrize( + ("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")] +) +def test_xai_batch_cost_calculator_bills_the_200k_batch_tier_inclusively( + _local_model_cost_map: None, prompt_tokens: int, tier: str +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + info: Final = litellm.get_model_info(_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai") + usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=64, total_tokens=prompt_tokens + 64) + + prompt_cost, completion_cost_value = batch_cost_calculator( + usage=usage, model=_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai" + ) + + assert prompt_cost == pytest.approx(prompt_tokens * info[f"input_cost_per_token{tier}"]) + assert completion_cost_value == pytest.approx(64 * info[f"output_cost_per_token{tier}"]) + + def test_batch_cost_calculator_ignores_malformed_batch_tier_keys(): from litellm.cost_calculator import batch_cost_calculator diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 768d8955b8e..0cdc52c9a93 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -774,6 +774,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, @@ -797,6 +798,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_32k_tokens": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, "input_cost_per_token_above_512k_tokens": {"type": "number"}, @@ -897,6 +899,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_above_32k_tokens": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, "output_cost_per_token_above_512k_tokens": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index be7094f7f6c..1e0aed46923 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32981,6 +32981,8 @@ export interface components { cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ cache_read_input_token_cost_above_200k_tokens?: number | null; + /** Cache Read Input Token Cost Above 200K Tokens Batches */ + cache_read_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Read Input Token Cost Above 200K Tokens Priority */ cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ @@ -33059,6 +33061,8 @@ export interface components { input_cost_per_token_above_128k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens */ input_cost_per_token_above_200k_tokens?: number | null; + /** Input Cost Per Token Above 200K Tokens Batches */ + input_cost_per_token_above_200k_tokens_batches?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -33182,6 +33186,8 @@ export interface components { output_cost_per_token_above_128k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens */ output_cost_per_token_above_200k_tokens?: number | null; + /** Output Cost Per Token Above 200K Tokens Batches */ + output_cost_per_token_above_200k_tokens_batches?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ @@ -46768,6 +46774,8 @@ export interface components { cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ cache_read_input_token_cost_above_200k_tokens?: number | null; + /** Cache Read Input Token Cost Above 200K Tokens Batches */ + cache_read_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Read Input Token Cost Above 200K Tokens Priority */ cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ @@ -46846,6 +46854,8 @@ export interface components { input_cost_per_token_above_128k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens */ input_cost_per_token_above_200k_tokens?: number | null; + /** Input Cost Per Token Above 200K Tokens Batches */ + input_cost_per_token_above_200k_tokens_batches?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -46969,6 +46979,8 @@ export interface components { output_cost_per_token_above_128k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens */ output_cost_per_token_above_200k_tokens?: number | null; + /** Output Cost Per Token Above 200K Tokens Batches */ + output_cost_per_token_above_200k_tokens_batches?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ From e065a2575bf31f8b11aa642d7bd6b3adb8edd0ed Mon Sep 17 00:00:00 2001 From: Refael Iliaguyev Date: Sat, 26 Sep 2026 01:36:35 +0300 Subject: [PATCH 26/58] fix(proxy): send a real error event when a /v1/messages stream fails (#41826) * fix(proxy): send a real error event when a /v1/messages stream fails When a stream failed halfway through, the proxy wrote the error as a plain `data: {"error": ...}` line with no `event:` in front of it. Anthropic clients pick stream events by that name, so they skip the line and the request looks like it simply stopped with nothing in it. Write the failure as an `event: error` frame with Anthropic's own payload, and take the error type from the status code * fix(proxy): use the shared Anthropic error mapping for the stream error frame The first pass added a third copy of the status to error-type table, and it disagreed with the documented one: 529 came out as `api_error` rather than `overloaded_error`, and 413 as `invalid_request_error` rather than `request_too_large`, which hides the two failures a client can actually act on. Drop that copy and put the frame builder next to the table litellm already keeps in anthropic_interface/exceptions. The bridged adapter path was building the same frame inline, so it uses the shared one now too * fix(proxy): seal a torn SSE frame before the /v1/messages error event An upstream that drops mid-frame leaves the client inside an open event, so the error frame that follows is glued onto the torn data line and the Anthropic SDK raises a JSON decode error instead of an APIStatusError. Close the open frame with a ping event the SDK skips before writing the error event, and add the e2e stream-cut edge with Bedrock, Anthropic boundary, and Anthropic mid-frame legs. * fix(proxy): keep the SSE tail unchanged on a chunk that is not text A serializer that hands a dict or model object through as-is has no bytes the frame tail can learn from, so advance_sse_tail leaves it alone instead of slicing it. * fix(proxy): answer a /v1/messages stream that fails before its first byte as a JSON error carrying its status * test: move the Anthropic error frame tests into tests/unit * test(e2e): cut the upstream stream only after content has been relayed * test(e2e): carry split SSE lines across chunks and always tear a data line mid-frame * test(e2e): find the next data line across a chunk boundary before tearing it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../exceptions/__init__.py | 4 + .../exceptions/exception_mapping_utils.py | 36 ++- .../adapters/streaming_iterator.py | 8 +- litellm/proxy/common_request_processing.py | 39 ++- litellm/proxy/common_utils/sse_keepalive.py | 26 +- .../coverage_registry/llm_conversational.yaml | 2 + tests/e2e/coverage_registry/schema.py | 1 + .../e2e/llm_translation/test_messages_e2e.py | 260 ++++++++++++++++- tests/e2e/models.py | 10 + tests/e2e/provider_edge.py | 152 +++++++++- tests/e2e/test_provider_edge.py | 32 +++ .../proxy/common_utils/test_sse_keepalive.py | 9 + .../proxy/test_common_request_processing.py | 263 ++++++++++++++++++ .../test_exception_mapping_utils.py | 48 +++- 14 files changed, 869 insertions(+), 21 deletions(-) diff --git a/litellm/anthropic_interface/exceptions/__init__.py b/litellm/anthropic_interface/exceptions/__init__.py index 7f2de0e60dc..7c3cea0a28a 100644 --- a/litellm/anthropic_interface/exceptions/__init__.py +++ b/litellm/anthropic_interface/exceptions/__init__.py @@ -2,7 +2,9 @@ from .exception_mapping_utils import ( ANTHROPIC_ERROR_TYPE_MAP, + AnthropicErrorSseFrame, AnthropicExceptionMapping, + anthropic_error_sse_frame, ) from .exceptions import ( AnthropicErrorDetail, @@ -14,6 +16,8 @@ __all__ = [ "ANTHROPIC_ERROR_TYPE_MAP", "AnthropicErrorDetail", "AnthropicErrorResponse", + "AnthropicErrorSseFrame", "AnthropicErrorType", "AnthropicExceptionMapping", + "anthropic_error_sse_frame", ] diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index d9c9925275b..eb3ec8aaee2 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -4,11 +4,12 @@ Utilities for mapping exceptions to Anthropic error format. Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthropic response format. """ +import json from typing import Final from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from .exceptions import AnthropicErrorResponse, AnthropicErrorType +from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicErrorType # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors @@ -166,3 +167,36 @@ class AnthropicExceptionMapping: message=message, request_id=request_id, ) + + +class AnthropicErrorSseFrame(str): + """One `event: error` frame, for a stream that fails once the response headers are out. + + Anthropic clients pick stream events by the `event:` name, so a frame carrying only a `data:` + line is skipped and the failure never reaches the caller. The frame remembers the status and + body it was built from, so a stream that fails before its first byte can still answer as a + JSON error with that exact status instead of a 200 that only says `api_error` + """ + + status_code: int + error_response: AnthropicErrorResponse + + def __new__(cls, status_code: int, error_response: AnthropicErrorResponse) -> "AnthropicErrorSseFrame": + frame: Final = super().__new__(cls, f"event: error\ndata: {json.dumps(error_response)}\n\n") + frame.status_code = status_code + frame.error_response = error_response + return frame + + def json_body(self, call_id: str | None) -> AnthropicErrorResponse: + if call_id is None: + return self.error_response + detail: Final[AnthropicErrorDetail] = {**self.error_response["error"], "litellm_call_id": call_id} + body: Final[AnthropicErrorResponse] = {**self.error_response, "error": detail} + return body + + +def anthropic_error_sse_frame(status_code: int, raw_message: str) -> AnthropicErrorSseFrame: + return AnthropicErrorSseFrame( + status_code, + AnthropicExceptionMapping.transform_to_anthropic_error(status_code=status_code, raw_message=raw_message), + ) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 20753afee5c..24d5b7f366e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -68,15 +68,11 @@ def _error_status_and_message(exc: Exception) -> tuple[int, str]: def _mid_stream_error_sse_event(exc: Exception) -> bytes: from litellm.anthropic_interface.exceptions.exception_mapping_utils import ( - AnthropicExceptionMapping, + anthropic_error_sse_frame, ) status_code, message = _error_status_and_message(exc) - error_response = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=message, - ) - return f"event: error\ndata: {json.dumps(error_response)}\n\n".encode() + return anthropic_error_sse_frame(status_code=status_code, raw_message=message).encode() def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4f9b6b3a96f..64b0c6c1967 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -32,6 +32,7 @@ from starlette.types import Receive, Scope, Send import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid +from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, DEFAULT_MAX_RECURSE_DEPTH, @@ -102,8 +103,11 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, + SSE_STREAM_START_TAIL, + advance_sse_tail, coerce_keepalive_interval, resolve_ttft_keepalive_interval, + seal_open_sse_frame, wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger @@ -999,6 +1003,17 @@ async def create_response( first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers) + if isinstance(first_chunk_value, AnthropicErrorSseFrame): + with contextlib.suppress(Exception): + await generator.aclose() + return JSONResponse( + status_code=first_chunk_value.status_code, + content=first_chunk_value.json_body( + error_body_call_id(general_settings, resolved_headers.get(LITELLM_CALL_ID_HEADER)) + ), + headers=resolved_headers, + ) + if first_chunk_value is not None: try: error_code_from_chunk: Final = await _parse_event_data_for_error(first_chunk_value) @@ -3852,6 +3867,7 @@ class ProxyBaseLLMRequestProcessing: serialize_error: StreamErrorSerializer, request: Request | None = None, flush_tail: Callable[[], bytes] | None = None, + seal_open_frame: Callable[[bytes], str] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, @@ -3861,6 +3877,12 @@ class ProxyBaseLLMRequestProcessing: ``flush_tail`` runs once after the upstream iterator completes cleanly and its non-empty result is yielded, so a serializer that buffers bytes across chunks can emit anything still held at end of stream. + + ``seal_open_frame`` is given the tail of what has been yielded when the + error frame goes out, and what it returns is written first. A passthrough + relays raw upstream bytes, so an upstream that hangs up mid-frame leaves the + client inside an open frame, where an error frame would be swallowed or + misparsed instead of raised. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3877,6 +3899,7 @@ class ProxyBaseLLMRequestProcessing: stream_completed = False client_disconnected = False delivered_chunk = False + recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes try: str_so_far = "" async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( @@ -3922,7 +3945,9 @@ class ProxyBaseLLMRequestProcessing: # False and refunds. A keepalive ping carries no provider output, # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES - yield serialize_chunk(chunk) + serialized = serialize_chunk(chunk) + recent_tail = advance_sse_tail(recent_tail, serialized) + yield serialized held_tail: Final = flush_tail() if flush_tail is not None else b"" if held_tail: yield serialize_chunk(held_tail) @@ -3970,7 +3995,9 @@ class ProxyBaseLLMRequestProcessing: code=stream_error_status, ) stream_completed = True - yield serialize_error(proxy_exception) + error_frame: Final = serialize_error(proxy_exception) + seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail) + yield seal + error_frame if seal else error_frame finally: await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( request=request, @@ -3992,7 +4019,7 @@ class ProxyBaseLLMRequestProcessing: restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ - Anthropic /messages and Google /generateContent streaming data generator require SSE events. + Anthropic /messages streaming data generator, which requires SSE events. Returns the underlying ``async_streaming_data_generator`` configured with SSE serializers directly (rather than re-wrapping it in another @@ -4010,11 +4037,13 @@ class ProxyBaseLLMRequestProcessing: request_data=request_data, proxy_logging_obj=proxy_logging_obj, serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), - serialize_error=lambda proxy_exc: ( - f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" + serialize_error=lambda proxy_exc: anthropic_error_sse_frame( + status_code=error_status_code(proxy_exc, status.HTTP_500_INTERNAL_SERVER_ERROR), + raw_message=proxy_exc.message, ), request=request, flush_tail=None if restamper is None else restamper.flush, + seal_open_frame=seal_open_sse_frame, ) @overload diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index cf98a7e9224..d9685971f52 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -15,7 +15,7 @@ SSE_COMMENT_PING_BYTES: Final = SSE_COMMENT_PING.encode() # terminates a line with CRLF, LF or CR, so a blank line is any of these three. _SSE_FRAME_DELIMITERS: Final = (b"\r\n\r\n", b"\n\n", b"\r\r") _SSE_DELIMITER_LOOKBACK: Final = max(len(delimiter) for delimiter in _SSE_FRAME_DELIMITERS) -_STREAM_START_TAIL: Final = b"\n\n" +SSE_STREAM_START_TAIL: Final = b"\n\n" _SSE_MEDIA_TYPE: Final = "text/event-stream" @@ -128,7 +128,7 @@ async def _keepalive_ping_byte_stream( # Seeded as a delimiter because a stream starts at a frame boundary, and kept # across chunks because a delimiter can be split between two transport reads, # which testing only the latest chunk would miss for the rest of the stream. - recent_tail = _STREAM_START_TAIL # rebind-ok: rolling window over the relayed bytes + recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the relayed bytes try: while True: await asyncio.wait((pending,), timeout=ping_interval_seconds) @@ -155,6 +155,28 @@ async def _keepalive_ping_byte_stream( await stream.aclose() +def advance_sse_tail(recent_tail: bytes, chunk: object) -> bytes: + written: Final = _sse_tail_bytes(chunk) + if not written: + return recent_tail + return (recent_tail + written)[-_SSE_DELIMITER_LOOKBACK:] + + +def _sse_tail_bytes(chunk: object) -> bytes: + if isinstance(chunk, bytes): + return chunk[-_SSE_DELIMITER_LOOKBACK:] + if isinstance(chunk, str): + return chunk[-_SSE_DELIMITER_LOOKBACK:].encode() + return b"" + + +def seal_open_sse_frame(recent_tail: bytes) -> str: + if recent_tail.endswith(_SSE_FRAME_DELIMITERS): + return "" + line_break: Final = "" if recent_tail.endswith((b"\n", b"\r")) else "\n" + return f"{line_break}{ANTHROPIC_PING_SSE_CHUNK}" + + def resolve_ttft_keepalive_interval( deployments: Iterable[Mapping[str, object]], global_interval: float | str | None, diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 20c87dbbd74..9cfe6e33ed6 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -103,6 +103,8 @@ - {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"} - {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"} - {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"} +- {id: llm.messages.anthropic.upstream_stream_failure.stream.error_event, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: upstream_stream_failure, streaming: stream, assertions: [error_event], source: "customer report", rationale: "An upstream that hangs up mid-stream must reach Anthropic clients as an event: error frame, not an OpenAI-shaped data-only error they silently drop"} +- {id: llm.messages.anthropic.upstream_stream_failure.stream.error_status, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: upstream_stream_failure, streaming: stream, assertions: [error_status], source: "customer report", rationale: "An upstream that hangs up before its first byte must answer as a JSON error carrying its status, so Anthropic clients raise the status-specific error and retry on it instead of reading a 200 stream that only carries an error event"} - {id: llm.messages.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI models served on the Anthropic Messages contract"} - {id: llm.messages.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages streams the Anthropic event grammar"} - {id: llm.messages.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages: cost header and spend row agree"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 5fd19212ab7..fec1934059c 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -86,6 +86,7 @@ LlmCapability = Literal[ "tool_search", "tool_search_history", "tool_use", + "upstream_stream_failure", "vision", "web_search", "web_search_server_tool", diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index d048d1343eb..871fd2f9aef 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -11,8 +11,11 @@ litellm-regression-tests/tests/test_inference_endpoints.py. from __future__ import annotations import time +from collections.abc import Callable +from types import MappingProxyType from typing import Final +import anthropic import pytest from anthropic import Anthropic from anthropic.types import ( @@ -30,12 +33,21 @@ from anthropic.types import ( ToolParam, ToolUseBlock, ) -from e2e_config import STREAM_MIN_LEAD_SECONDS, provider_edge_base, provider_paces_stream, unique_marker +from e2e_config import ( + PROVIDER_EDGE_ADVERTISE_HOST, + PROVIDER_EDGE_BIND_HOST, + STREAM_MIN_LEAD_SECONDS, + provider_edge_base, + provider_paces_stream, + unique_marker, +) from e2e_http import assert_client_error from lifecycle import ResourceManager -from models import ChatMessage, LiteLLMParamsBody, SpendLogRow +from models import AnthropicErrorEvent, AnthropicMessagesBody, ChatMessage, LiteLLMParamsBody, SpendLogRow +from provider_edge import EDGE_MOUNTS, LiveEdge, RunningEdge, StreamCut, start_provider_edge +from provider_edge_bedrock import bedrock_signer from proxy_client import ProxyClient -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @@ -385,3 +397,245 @@ class TestOpenAIMessagesToolContinuation: ) assert _text(continuation).strip() == receipt, "continuation did not consume the correlated tool result" assert all(not isinstance(block, ToolUseBlock) for block in continuation.content) + + +BEDROCK_BACKEND: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +BEDROCK_EDGE_REGION: Final = "us-east-1" +_STREAM_FAILURE_PROMPT: Final = "Count from 1 to 100, one number per line." +_FRAME_PAYLOAD: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_AT_FRAME_BOUNDARY: Final = StreamCut(after_content=True) +_MID_FRAME: Final = StreamCut(after_content=True, mid_chunk=True) +_BEFORE_FIRST_BYTE: Final = StreamCut(after_content=False) + +type _CutRegistration = Callable[[ProxyClient, ResourceManager, StreamCut], tuple[str, str]] + + +def _cut_edge(backend: LiveEdge, mount: str) -> RunningEdge: + return start_provider_edge( + backend, + mounts=MappingProxyType({mount: EDGE_MOUNTS[mount]}), + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + ) + + +def _register_cut_bedrock(proxy: ProxyClient, resources: ResourceManager, cut: StreamCut) -> tuple[str, str]: + mount: Final = f"bedrock/{BEDROCK_EDGE_REGION}" + edge: Final = _cut_edge(LiveEdge(cut=cut, sign=bedrock_signer(BEDROCK_EDGE_REGION)), mount) + resources.defer(edge.shutdown) + return _register( + proxy, + resources, + LiteLLMParamsBody( + model=BEDROCK_BACKEND, + api_base=edge.edge.api_base(mount), + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name=BEDROCK_EDGE_REGION, + ), + prefix="e2e-messages-cut", + ) + + +def _register_cut_anthropic(proxy: ProxyClient, resources: ResourceManager, cut: StreamCut) -> tuple[str, str]: + edge: Final = _cut_edge(LiveEdge(cut=cut), "anthropic") + resources.defer(edge.shutdown) + return _register( + proxy, + resources, + LiteLLMParamsBody( + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=edge.edge.api_base("anthropic") + ), + prefix="e2e-messages-cut", + ) + + +_DROPPED_UPSTREAMS: Final[tuple[tuple[str, _CutRegistration, StreamCut], ...]] = ( + ("bedrock_at_a_frame_boundary", _register_cut_bedrock, _AT_FRAME_BOUNDARY), + ("anthropic_at_a_frame_boundary", _register_cut_anthropic, _AT_FRAME_BOUNDARY), + ("anthropic_mid_frame", _register_cut_anthropic, _MID_FRAME), +) +_DROPPED_BEFORE_FIRST_BYTE: Final[tuple[tuple[str, _CutRegistration, StreamCut], ...]] = ( + ("bedrock_before_the_first_byte", _register_cut_bedrock, _BEFORE_FIRST_BYTE), + ("anthropic_before_the_first_byte", _register_cut_anthropic, _BEFORE_FIRST_BYTE), +) + + +def _payload(frame: str) -> JsonValue | None: + try: + return _FRAME_PAYLOAD.validate_json(frame) + except ValidationError: + return None + + +def _bare_error_frame(frame: str) -> bool: + payload: Final = _payload(frame) + return isinstance(payload, dict) and "error" in payload and payload.get("type") != "error" + + +@pytest.mark.provider_edge_host +@pytest.mark.provider_live +class TestMessagesUpstreamStreamFailure: + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_event") + @pytest.mark.parametrize( + ("register", "cut"), [case[1:] for case in _DROPPED_UPSTREAMS], ids=[case[0] for case in _DROPPED_UPSTREAMS] + ) + def test_interrupted_upstream_stream_raises_in_the_anthropic_sdk( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + register: _CutRegistration, + cut: StreamCut, + ) -> None: + model, key = register(proxy, resources, cut) + client: Final = sdk.anthropic(key) + + stream: Final = client.messages.create( + model=model, + max_tokens=300, + stream=True, + messages=[_user_turn(_STREAM_FAILURE_PROMPT)], + extra_body=NO_PROXY_CACHE, + ) + first: Final = next(stream) + assert first.type == "message_start", ( + f"the stream produced a first event that is not message_start, so this run proves a " + f"startup failure, not an interrupted stream: {first!r}" + ) + with pytest.raises(anthropic.APIStatusError) as raised: + for _ in stream: + pass + try: + AnthropicErrorEvent.model_validate(raised.value.body) + except ValidationError: + pytest.fail( + f"the SDK raised on the interrupted stream but without the Anthropic error envelope a " + f"client reads the failure from: body={raised.value.body!r} message={raised.value}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_event") + @pytest.mark.parametrize( + ("register", "cut"), [case[1:] for case in _DROPPED_UPSTREAMS], ids=[case[0] for case in _DROPPED_UPSTREAMS] + ) + def test_interrupted_upstream_stream_is_an_anthropic_error_event( + self, proxy: ProxyClient, resources: ResourceManager, register: _CutRegistration, cut: StreamCut + ) -> None: + model, key = register(proxy, resources, cut) + + outcome: Final = proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + max_tokens=300, + stream=True, + messages=[ChatMessage(role="user", content=_STREAM_FAILURE_PROMPT)], + ), + ) + frames: Final = outcome.stream_events + assert outcome.is_streaming, ( + f"/v1/messages did not answer with an SSE stream: status={outcome.status_code} body={outcome.body}" + ) + assert frames, ( + f"the proxy sent no SSE data frames although the upstream hung up; stream_error={outcome.stream_error!r}" + ) + assert outcome.stream_error == "event: error", ( + f"the interrupted stream was not announced by an 'event: error' line Anthropic clients read; " + f"stream_error={outcome.stream_error!r} frames={frames}" + ) + try: + AnthropicErrorEvent.model_validate_json(frames[-1]) + except ValidationError: + pytest.fail( + f'the last SSE frame was not an Anthropic {{"type": "error", "error": ...}} envelope; frames={frames}' + ) + torn: Final = tuple(index for index, frame in enumerate(frames) if _payload(frame) is None) + expected_torn: Final = 1 if cut.mid_chunk else 0 + assert len(torn) == expected_torn, ( + f"expected {expected_torn} data line(s) that are not JSON, since the edge tears one only when it " + f"cuts mid-frame, but the proxy relayed {[frames[index] for index in torn]}; all frames={frames}" + ) + for index in torn: + assert _payload(frames[index + 1]) == {"type": "ping"}, ( + f"the frame the upstream tore was not closed as a ping event before the error, so an " + f"Anthropic client parses the error inside it: after {frames[index]!r} came " + f"{frames[index + 1]!r}; all frames={frames}" + ) + bare: Final = tuple(frame for frame in frames if _bare_error_frame(frame)) + assert not bare, ( + f"the proxy emitted error frames without the Anthropic envelope, which Anthropic clients drop: " + f"{bare}; all frames={frames}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_status") + @pytest.mark.parametrize( + ("register", "cut"), + [case[1:] for case in _DROPPED_BEFORE_FIRST_BYTE], + ids=[case[0] for case in _DROPPED_BEFORE_FIRST_BYTE], + ) + def test_upstream_that_hangs_up_before_the_first_byte_raises_with_its_status_in_the_anthropic_sdk( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + register: _CutRegistration, + cut: StreamCut, + ) -> None: + model, key = register(proxy, resources, cut) + client: Final = sdk.anthropic(key) + + with pytest.raises(anthropic.APIStatusError) as raised: + client.messages.create( + model=model, + max_tokens=300, + stream=True, + messages=[_user_turn(_STREAM_FAILURE_PROMPT)], + extra_body=NO_PROXY_CACHE, + ) + assert 500 <= raised.value.status_code < 600, ( + f"an upstream that hung up before sending anything must answer with a server error status the SDK " + f"retries on, not {raised.value.status_code}: {raised.value}" + ) + try: + AnthropicErrorEvent.model_validate(raised.value.body) + except ValidationError: + pytest.fail( + f"the SDK raised with the right status but without the Anthropic error envelope a client reads " + f"the failure from: body={raised.value.body!r} message={raised.value}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_status") + @pytest.mark.parametrize( + ("register", "cut"), + [case[1:] for case in _DROPPED_BEFORE_FIRST_BYTE], + ids=[case[0] for case in _DROPPED_BEFORE_FIRST_BYTE], + ) + def test_upstream_that_hangs_up_before_the_first_byte_is_a_json_error_with_its_status( + self, proxy: ProxyClient, resources: ResourceManager, register: _CutRegistration, cut: StreamCut + ) -> None: + model, key = register(proxy, resources, cut) + + outcome: Final = proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + max_tokens=300, + stream=True, + messages=[ChatMessage(role="user", content=_STREAM_FAILURE_PROMPT)], + ), + ) + assert not outcome.is_streaming, ( + f"nothing had been streamed when the upstream hung up, yet /v1/messages opened a 200 SSE stream " + f"instead of answering with the failure's status: stream_error={outcome.stream_error!r} " + f"frames={outcome.stream_events}" + ) + assert 500 <= outcome.status_code < 600, ( + f"/v1/messages answered {outcome.status_code} for an upstream that hung up before its first byte; " + f"body={outcome.body}" + ) + try: + AnthropicErrorEvent.model_validate_json(outcome.body) + except ValidationError: + pytest.fail( + f'the error body is not an Anthropic {{"type": "error", "error": ...}} envelope; body={outcome.body}' + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index c96f4b0bef1..84399fd6155 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -598,6 +598,16 @@ class CountTokensResponse(BaseModel): input_tokens: int +class AnthropicErrorBody(BaseModel): + type: str + message: str + + +class AnthropicErrorEvent(BaseModel): + type: Literal["error"] + error: AnthropicErrorBody + + # ---------- mcp servers ---------- diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index fc10dde2a77..3680375b6af 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -45,6 +45,7 @@ import hashlib import os import re import threading +import time from collections import deque from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import closing, contextmanager @@ -56,6 +57,7 @@ from types import MappingProxyType from typing import Final, Literal, assert_never from urllib.parse import parse_qsl, urlsplit +from botocore.eventstream import EventStreamBuffer from e2e_http import ( NetworkError, StreamChunk, @@ -96,16 +98,18 @@ from fixture_mode import ( ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity from provider_cache import ( + JSON_VALUE, SIGNATURE_HEADERS, CacheEdge, MountPolicy, RequestSigner, + invoke_chunk_value, is_bedrock, scoped_edge_base, split_test_segment, ) from provider_cache_routing import LIVE_PROVIDER_REQUIRED -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",) @@ -537,10 +541,33 @@ class ReplayEdge: source: ReplaySource +@dataclass(frozen=True, slots=True) +class StreamCut: + """Where a live edge hangs up on a streamed upstream body: before its first byte, or with + ``after_content`` set, right after the first transfer chunk carrying assistant output (a + ``content_block_delta``). That frame is what commits the proxy's mid-stream fallback + wrapper to the client: it holds the lifecycle frames before it back and drops them when + the transport fails first, so a cut after a fixed number of chunks landed on either side + of that commit depending on how the provider batched its frames. With ``mid_chunk`` set + the hang-up comes part way through the next ``data:`` line the provider sends after that, + so the client is left inside an SSE frame the way a dropped transport leaves it. + + Whatever was relayed sits on the wire for ``_CUT_SETTLE_SECONDS`` before the hang-up, so + the client has read it by then instead of receiving the data and the close in one burst, + where its reader can surface the close before what it buffered.""" + + after_content: bool + mid_chunk: bool = False + + +_CUT_SETTLE_SECONDS: Final = 1.0 + + @dataclass(frozen=True, slots=True) class LiveEdge: observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None sign: RequestSigner | None = None + cut: StreamCut | None = None type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge @@ -786,11 +813,128 @@ def _handle_record( assert_never(head) +def _data_line_start(data: bytes) -> int: + if data.startswith(b"data:"): + return 0 + at_line_start: Final = data.find(b"\ndata:") + return -1 if at_line_start < 0 else at_line_start + 1 + + +def _torn_prefix(data: bytes) -> bytes: + start: Final = _data_line_start(data) + line_end: Final = data.find(b"\n", start) + end: Final = len(data) if line_end < 0 else line_end + return data[: start + (end - start) // 2] + + +class _DataLineTearer: + __slots__ = ("_unfinished_line",) + + _unfinished_line: bytes + + def __init__(self) -> None: + self._unfinished_line = b"" + + def observe(self, data: bytes) -> None: + self._unfinished_line = (self._unfinished_line + data).rsplit(b"\n", 1)[-1] + + def tear(self, data: bytes) -> bytes | None: + buffered: Final = self._unfinished_line + data + if _data_line_start(buffered) < 0: + self.observe(data) + return None + return _torn_prefix(buffered)[len(self._unfinished_line):] + + +def _is_content_delta(value: JsonValue | None) -> bool: + return isinstance(value, dict) and value.get("type") == "content_block_delta" + + +def _sse_data_carries_content(line: bytes) -> bool: + if not line.startswith(b"data:"): + return False + try: + return _is_content_delta(JSON_VALUE.validate_json(line[len(b"data:"):].strip())) + except ValidationError: + return False + + +class _AnthropicContentDetector: + __slots__ = ("_unfinished_line",) + + _unfinished_line: bytes + + def __init__(self) -> None: + self._unfinished_line = b"" + + def __call__(self, data: bytes) -> bool: + lines: Final = (self._unfinished_line + data).split(b"\n") + self._unfinished_line = lines[-1] + return any(_sse_data_carries_content(line.rstrip(b"\r")) for line in lines[:-1]) + + +def _invoke_frame_carries_content(payload: bytes) -> bool: + try: + return _is_content_delta(invoke_chunk_value(JSON_VALUE.validate_json(payload))) + except ValidationError: + return False + + +def _bedrock_content_detector() -> Callable[[bytes], bool]: + """Bedrock's invoke stream wraps each Anthropic event in an eventstream frame that a + transfer chunk can split, so the frames are reassembled across chunks before being read.""" + frames: Final = EventStreamBuffer() + + def carries_content(data: bytes) -> bool: + frames.add_data(data) + return any(_invoke_frame_carries_content(frame.payload) for frame in frames) + + return carries_content + + +def _content_detector(mount: str) -> Callable[[bytes], bool]: + return _bedrock_content_detector() if is_bedrock(mount) else _AnthropicContentDetector() + + +def _cut_steps( + steps: Generator[StreamStep, None, None], cut: StreamCut, carries_content: Callable[[bytes], bool] +) -> Generator[StreamStep, None, None]: + with closing(steps) as source: + tearer: Final = _DataLineTearer() + if cut.after_content: + for step in source: + yield step + if isinstance(step, StreamTruncation): + return + tearer.observe(step.data) + if carries_content(step.data): + break + else: + return + if cut.mid_chunk: + for step in source: + if isinstance(step, StreamTruncation): + yield step + return + if (torn := tearer.tear(step.data)) is None: + yield step + continue + if torn: + yield StreamChunk(data=torn) + break + else: + return + if cut.after_content or cut.mid_chunk: + time.sleep(_CUT_SETTLE_SECONDS) + yield StreamTruncation(reason=f"edge cut the upstream stream: {cut!r}") + + def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None, observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None, sign: RequestSigner | None = None, + cut: StreamCut | None = None, ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS @@ -805,6 +949,8 @@ def _handle_live( match head: case NetworkError(message=message): return _recorded_outcome(_network_error_response(message)) + case StreamHead() if cut is not None: + return EdgeStream(head.status_code, _filtered_response_headers(head.headers), _cut_steps(head.steps, cut, _content_detector(mount))) case StreamHead() if _is_streamed(head.headers): return EdgeStream(head.status_code, _filtered_response_headers(head.headers), head.steps) case StreamHead(): @@ -875,10 +1021,10 @@ def handle_edge_request( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, mount, test_key, ) - case LiveEdge(observe_request=observe_request, sign=sign): + case LiveEdge(observe_request=observe_request, sign=sign, cut=cut): return _handle_live( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, - observe_request=observe_request, sign=sign, + mount=mount, observe_request=observe_request, sign=sign, cut=cut, ) case RecordEdge(): return _handle_record( diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index d776c338ef7..978c3671a77 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -54,11 +54,13 @@ from provider_edge import ( EdgeBackend, EdgeReply, EdgeStream, + LiveEdge, ProviderEdge, ProviderRequestObservation, RecordEdge, ReplayEdge, ReplaySource, + StreamCut, edge_request, handle_edge_request, observed_provider_edge, @@ -1000,6 +1002,36 @@ def stream_chunks(response: RecordedStreamedResponse) -> list[bytes]: return [base64.b64decode(chunk) for chunk in response.chunks_b64] +SECOND_DATA_LINE: Final = b'data: {"type":"content_block_delta","delta":{"text":" two"}}' +SPLIT_MARKER_CHUNKS: tuple[bytes, ...] = ( + b'data: {"type":"content_block_delta","delta":{"text":"one"}}\n\nda', + b"ta" + SECOND_DATA_LINE[4:] + b"\n\nda", + b'ta: {"type":"message_delta","usage":{"output_tokens":7}}\n\nda', + b"ta: [DONE]\n\n", +) + + +class TestStreamCut: + def test_a_mid_frame_cut_tears_a_data_line_whose_marker_is_split_across_chunks(self) -> None: + """Every ``data:`` marker after the first content delta straddles a transfer + chunk boundary, so a tearer that inspects each chunk on its own never finds + one and lets the stream finish cleanly instead of cutting it.""" + backend: Final = LiveEdge(cut=StreamCut(after_content=True, mid_chunk=True)) + with chunked_provider(chunks=SPLIT_MARKER_CHUNKS) as provider: + with running_edge(backend, {"openai": provider_url(provider)}) as edge: + head, chunks, ending = raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + + assert head.startswith("HTTP/1.1 200 OK") + assert ending == "truncated" + relayed: Final = b"".join(chunks) + whole: Final = b"".join(SPLIT_MARKER_CHUNKS) + assert whole.startswith(relayed) and relayed != whole + assert relayed.startswith(SPLIT_MARKER_CHUNKS[0]) + torn_line: Final = relayed.rsplit(b"\n", 1)[-1] + assert torn_line and SECOND_DATA_LINE.startswith(torn_line) and torn_line != SECOND_DATA_LINE + assert b"[DONE]" not in relayed + + class TestStreamingFidelity: """LIT-5742: a streamed response records and replays as the chunk sequence the provider actually sent, not as one coalesced body. The unit of fidelity is the 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 69b92f5e4d7..228fd5bcae6 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -6,10 +6,13 @@ import pytest from fastapi.responses import StreamingResponse from litellm.proxy.common_request_processing import create_response +from litellm.types.utils import ModelResponse from litellm.proxy.common_utils.sse_keepalive import ( ANTHROPIC_PING_SSE_CHUNK, SSE_COMMENT_PING_BYTES, + advance_sse_tail, resolve_ttft_keepalive_interval, + seal_open_sse_frame, split_complete_sse_frames, wrap_passthrough_sse_bytes_with_keepalive_pings, wrap_sse_stream_with_keepalive_pings, @@ -32,6 +35,12 @@ 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.parametrize("chunk", [{"content": "hi"}, ModelResponse()]) +def test_advance_sse_tail_ignores_a_chunk_that_is_not_sse_text(chunk: object): + assert advance_sse_tail(b"\n\n", chunk) == b"\n\n" + assert seal_open_sse_frame(advance_sse_tail(b"data: {", chunk)) == "\n" + ANTHROPIC_PING_SSE_CHUNK + + @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/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index bdf003085ef..c17f41a8b8f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7,6 +7,7 @@ from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch +import anthropic import httpx import pytest from fastapi import HTTPException, Request, Response, status @@ -14,6 +15,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid +from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.litellm_core_utils.bug_report import ( DISABLE_ENV_VAR, ISSUE_URL_BASE, @@ -55,6 +57,7 @@ from litellm.proxy.common_request_processing import ( sse_error_payload, ) from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -2543,6 +2546,63 @@ class TestCommonRequestProcessingHelpers: assert response.headers["x-litellm-call-id"] == "call-8302" assert json.loads(response.body) == {"error": {"code": 403, "message": "forbidden"}} + async def test_a_stream_that_fails_before_its_first_byte_answers_as_an_anthropic_json_error(self): + """A /v1/messages stream whose first chunk is already the error frame has nothing + streamed yet, so the failure answers as JSON with the status the upstream gave, + the shape Anthropic clients raise their status-specific errors on""" + + async def stream(): + yield anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + yield ANTHROPIC_PING_SSE_CHUNK + + generator: Final = stream() + response = await create_response(generator, "text/event-stream", {"x-litellm-call-id": "call-8609"}) + + assert isinstance(response, JSONResponse) + assert response.status_code == 503 + assert response.headers["content-type"] == "application/json" + assert response.headers["x-litellm-call-id"] == "call-8609" + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable"}, + } + assert generator.ag_frame is None + + async def test_a_stream_that_fails_before_its_first_byte_names_the_call_when_opted_in(self): + async def stream(): + yield anthropic_error_sse_frame(status_code=429, raw_message="slow down") + + response = await create_response( + stream(), + "text/event-stream", + {"x-litellm-call-id": "call-8609"}, + general_settings={"include_call_id_in_error_body": True}, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "rate_limit_error", "message": "slow down", "litellm_call_id": "call-8609"}, + } + + async def test_an_error_event_after_a_keepalive_ping_still_streams(self): + """Once a keepalive ping went out the headers are committed, so the error frame + streams as an event instead of turning into a JSON answer""" + + async def stream(): + yield ANTHROPIC_PING_SSE_CHUNK + yield anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + response = await create_response(stream(), "text/event-stream", {}) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + assert "".join(await self.consume_stream(response)) == ( + ANTHROPIC_PING_SSE_CHUNK + + 'event: error\ndata: {"type": "error", "error": {"type": "api_error", "message": "upstream unavailable"}}\n\n' + ) + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -9901,6 +9961,209 @@ class TestErrorLogCarriesCallId: assert call_id in record.getMessage() +class TestAnthropicMessagesStreamErrorFrame: + """A ``/v1/messages`` stream that fails after the headers are out has to say so with an + ``event: error`` frame. Anthropic clients pick events by name, so a bare ``data:`` line is + skipped and the request looks like it ended with nothing in it""" + + @staticmethod + def _sse_generator_failing_with(failure: Exception) -> AsyncGenerator[str, None]: + class FailingUpstream: + def __aiter__(self) -> "FailingUpstream": + return self + + async def __anext__(self) -> object: + raise failure + + ProxyLogging._callback_capabilities_cache.clear() + return ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=FailingUpstream(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "claude-sonnet-4-5"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + ) + + @pytest.mark.parametrize( + "status_code, expected_error_type", + [ + (429, "rate_limit_error"), + (529, "overloaded_error"), + (413, "request_too_large"), + (500, "api_error"), + (502, "api_error"), + (400, "invalid_request_error"), + ], + ) + async def test_mid_stream_failure_arrives_as_an_anthropic_error_event( + self, status_code: int, expected_error_type: str + ) -> None: + class UpstreamFailure(Exception): + def __init__(self) -> None: + super().__init__("upstream stopped sending") + self.status_code: Final = status_code + + frames: Final = [frame async for frame in self._sse_generator_failing_with(UpstreamFailure())] + + assert len(frames) == 1 + event_line, data_line, first_blank, second_blank = frames[0].split("\n") + assert isinstance(frames[0], AnthropicErrorSseFrame) + assert frames[0].status_code == status_code + assert event_line == "event: error" + assert (first_blank, second_blank) == ("", "") + payload: Final = json.loads(data_line.removeprefix("data: ")) + assert payload["type"] == "error" + assert payload["error"]["type"] == expected_error_type + assert "upstream stopped sending" in payload["error"]["message"] + + _CONTENT_DELTA_FRAME: Final = ( + b"event: content_block_delta\n" + b'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"1\\n2\\n3"}}\n\n' + ) + _TORN_DATA_LINE: Final = ( + b"event: content_block_delta\n" + b'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"4' + ) + _PING: Final = ANTHROPIC_PING_SSE_CHUNK.encode() + + @staticmethod + def _upstream_failure(status_code: int) -> Exception: + class UpstreamFailure(Exception): + def __init__(self) -> None: + super().__init__("upstream stopped sending") + self.status_code: Final = status_code + + return UpstreamFailure() + + @staticmethod + def _sse_generator_cut_after(relayed: Sequence[bytes], failure: Exception) -> AsyncGenerator[str, None]: + class CutUpstream: + def __init__(self) -> None: + self._remaining: Final = iter(relayed) + + def __aiter__(self) -> "CutUpstream": + return self + + async def __anext__(self) -> object: + chunk: Final = next(self._remaining, None) + if chunk is None: + raise failure + return chunk + + ProxyLogging._callback_capabilities_cache.clear() + return ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=CutUpstream(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "claude-sonnet-4-5"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + ) + + @staticmethod + def _as_bytes(chunk: object) -> bytes: + if isinstance(chunk, bytes): + return chunk + assert isinstance(chunk, str) + return chunk.encode() + + async def _wire_bytes(self, relayed: Sequence[bytes]) -> bytes: + stream: Final = self._sse_generator_cut_after(relayed, self._upstream_failure(500)) + return b"".join([self._as_bytes(chunk) async for chunk in stream]) + + @staticmethod + def _error_frame_after(wire: bytes, relayed: bytes) -> bytes: + assert wire.startswith(relayed), f"the wire did not open with {relayed!r}: {wire!r}" + return wire.removeprefix(relayed) + + @staticmethod + def _assert_error_frame(frame: bytes) -> None: + event_line, data_line, first_blank, second_blank = frame.split(b"\n") + assert event_line == b"event: error" + assert (first_blank, second_blank) == (b"", b"") + payload: Final = json.loads(data_line.removeprefix(b"data: ")) + assert payload["type"] == "error" + assert "upstream stopped sending" in payload["error"]["message"] + + @pytest.mark.parametrize( + "torn, seal", + [ + (_TORN_DATA_LINE, b"\n" + _PING), + (b"event: content_bl", b"\n" + _PING), + (b"event: content_block_delta\n", _PING), + (b'event: content_block_delta\r\ndata: {"type":"content_block_delta"}\r\n', _PING), + ], + ids=["mid_data_line", "mid_event_line", "after_a_complete_line", "after_a_crlf_line"], + ) + async def test_a_frame_the_upstream_tore_is_closed_as_a_ping_before_the_error_event( + self, torn: bytes, seal: bytes + ) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME, torn)) + + self._assert_error_frame(self._error_frame_after(wire, self._CONTENT_DELTA_FRAME + torn + seal)) + + async def test_a_cut_at_a_frame_boundary_gets_the_error_event_alone(self) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME,)) + + self._assert_error_frame(self._error_frame_after(wire, self._CONTENT_DELTA_FRAME)) + + async def test_a_torn_frame_still_raises_the_error_in_the_anthropic_sdk(self) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME, self._TORN_DATA_LINE)) + + def serve(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=wire) + + client: Final = anthropic.Anthropic( + api_key="sk-test", + base_url="http://proxy.test", + http_client=httpx.Client(transport=httpx.MockTransport(serve)), + max_retries=0, + ) + with pytest.raises(anthropic.APIStatusError) as raised: + for _ in client.messages.create( + model="claude-sonnet-4-5", max_tokens=16, messages=[{"role": "user", "content": "count"}], stream=True + ): + pass + body: Final = raised.value.body + assert isinstance(body, dict) + assert body["type"] == "error" + assert "upstream stopped sending" in body["error"]["message"] + + async def test_a_failure_before_the_first_byte_answers_with_its_status_as_json(self) -> None: + response: Final = await create_response( + self._sse_generator_failing_with(self._upstream_failure(502)), "text/event-stream", {} + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 502 + body: Final = json.loads(response.body) + assert body["type"] == "error" + assert body["error"]["type"] == "api_error" + assert "upstream stopped sending" in body["error"]["message"] + + async def test_a_failure_before_the_first_byte_raises_with_its_status_in_the_anthropic_sdk(self) -> None: + response: Final = await create_response( + self._sse_generator_failing_with(self._upstream_failure(502)), "text/event-stream", {} + ) + assert isinstance(response, JSONResponse) + + def serve(request: httpx.Request) -> httpx.Response: + return httpx.Response(response.status_code, headers=dict(response.headers), content=response.body) + + client: Final = anthropic.Anthropic( + api_key="sk-test", + base_url="http://proxy.test", + http_client=httpx.Client(transport=httpx.MockTransport(serve)), + max_retries=0, + ) + with pytest.raises(anthropic.APIStatusError) as raised: + client.messages.create( + model="claude-sonnet-4-5", max_tokens=16, messages=[{"role": "user", "content": "count"}], stream=True + ) + assert raised.value.status_code == 502 + body: Final = raised.value.body + assert isinstance(body, dict) + assert body["type"] == "error" + assert "upstream stopped sending" in body["error"]["message"] + + class TestStreamingContainerOwnershipRecordedBeforeDone: """Regression for LIT-8612: the OpenAI SDK closes the connection at ``data: [DONE]`` and starlette cancels the body task, so an ownership row diff --git a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py index ef092b65f28..0d8e7674e7d 100644 --- a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py +++ b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py @@ -3,8 +3,15 @@ Tests for AnthropicExceptionMapping class in litellm/anthropic_interface/excepti """ import json +from typing import Final -from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping +import pytest + +from litellm.anthropic_interface.exceptions import ( + AnthropicErrorSseFrame, + AnthropicExceptionMapping, + anthropic_error_sse_frame, +) class TestCreateErrorResponse: @@ -206,3 +213,42 @@ class TestTransformToAnthropicError: ) assert result["type"] == "error" assert result["error"]["message"] == '["error1", "error2"]' + + +class TestAnthropicErrorSseFrame: + @pytest.mark.parametrize( + ("status_code", "expected_error_type"), + [(429, "rate_limit_error"), (503, "api_error"), (400, "invalid_request_error")], + ) + def test_the_frame_is_one_error_event_carrying_the_anthropic_envelope( + self, status_code: int, expected_error_type: str + ) -> None: + frame: Final = anthropic_error_sse_frame(status_code=status_code, raw_message="upstream unavailable") + + event_line, data_line, first_blank, second_blank = frame.split("\n") + assert event_line == "event: error" + assert (first_blank, second_blank) == ("", "") + assert json.loads(data_line.removeprefix("data: ")) == { + "type": "error", + "error": {"type": expected_error_type, "message": "upstream unavailable"}, + } + + def test_the_frame_remembers_the_status_and_body_it_was_built_from(self) -> None: + frame: Final = anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + assert isinstance(frame, AnthropicErrorSseFrame) + assert frame.status_code == 503 + data_line: Final = frame.split("\n")[1] + assert data_line == f"data: {json.dumps(frame.json_body(call_id=None))}" + + def test_the_json_body_names_the_call_only_when_asked(self) -> None: + frame: Final = anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + assert frame.json_body(call_id="call-1") == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable", "litellm_call_id": "call-1"}, + } + assert frame.json_body(call_id=None) == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable"}, + } From d86c2e1f4238e0689e9732b1f531d29ea8eb434e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:06:44 -0700 Subject: [PATCH 27/58] fix(logging): redact raw_request when turn_off_message_logging is set in the proxy config (#43219) * fix(logging): redact raw_request when turn_off_message_logging is set in the proxy config The raw request branch bound turn_off_message_logging by name at import, before the proxy config set it, so loggers kept receiving the prompt in metadata.raw_request and raw_request_typed_dict. It now runs the same per-request redaction check messages use, and json_logs is read at call time for the same reason * fix(logging): keep raw_request_typed_dict for the explicit readers and tolerate missing headers in the json debug log * refactor(logging): drop the stale comment above the raw request typed dict --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 61 +++++++------------ .../test_litellm_logging.py | 61 ++++++++++++++++++- tests/unit/test_main.py | 23 +++++++ 3 files changed, 106 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0ef5bbf807a..28d72702f3e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -20,11 +20,7 @@ from httpx import Response from pydantic import BaseModel, JsonValue import litellm -from litellm import ( - _custom_logger_compatible_callbacks_literal, - json_logs, - turn_off_message_logging, -) +from litellm import _custom_logger_compatible_callbacks_literal from litellm._logging import ( _is_debugging_on, _redact_string, @@ -43,6 +39,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, PROVIDER_REQUEST_ID_HEADERS, + REDACTED_BY_LITELLM, ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, @@ -1358,10 +1355,19 @@ class Logging(LiteLLMLoggingBaseClass): _litellm_params: Final = self.model_call_details.get("litellm_params", {}) _metadata: Final = _litellm_params.get("metadata", {}) or {} try: - # [Non-blocking Extra Debug Information in metadata] - if turn_off_message_logging is True: - _metadata["raw_request"] = "redacted by litellm. \ - 'litellm.turn_off_message_logging=True'" + self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( + raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")), + raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, + ) + if should_redact_message_logging(self.model_call_details): + _metadata["raw_request"] = REDACTED_BY_LITELLM else: curl_command: Final = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -1369,20 +1375,7 @@ class Logging(LiteLLMLoggingBaseClass): additional_args=additional_args, data=additional_args.get("complete_input_dict", {}), ) - _metadata["raw_request"] = _redact_string(str(curl_command)) - # split up, so it's easier to parse in the UI - self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( - raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")), - raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) except Exception as e: self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( error=str(e), @@ -1476,7 +1469,7 @@ class Logging(LiteLLMLoggingBaseClass): def _print_llm_call_debugging_log( self, api_base: str, - headers: dict, + headers: dict | None, additional_args: dict, ): """ @@ -1485,8 +1478,8 @@ class Logging(LiteLLMLoggingBaseClass): Prints the RAW curl command sent from LiteLLM """ if _is_debugging_on() or self.litellm_request_debug: - if json_logs: - masked_headers: Final = self._get_masked_headers(headers) + if litellm.json_logs: + masked_headers: Final = self._get_masked_headers(headers or {}) masked_api_base: Final = self._get_masked_api_base(str(api_base or "")) if self.litellm_request_debug: verbose_logger.warning( # .warning ensures this shows up in all environments @@ -1563,20 +1556,12 @@ class Logging(LiteLLMLoggingBaseClass): else: attr = "debug" - if json_logs: - callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug - callattr( - "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get("original_response", self.model_call_details) - ), - ) - else: - callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug - callattr( - "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get("original_response", self.model_call_details) - ) + callattr: Final = verbose_logger.warning if attr == "warning" else verbose_logger.debug + callattr( + "RAW RESPONSE:\n{}\n\n".format( + self.model_call_details.get("original_response", self.model_call_details) ) + ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: self.logger_fn( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 265fdb50836..d717718cba2 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -20,7 +20,7 @@ from openai._legacy_response import HttpxBinaryResponseContent import litellm from litellm._logging import session_id_var, trace_id_var -from litellm.constants import SENTRY_PII_DENYLIST +from litellm.constants import REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -6615,6 +6615,65 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj): assert "key=*****" in raw_api_base +_PRIVATE_RAW_REQUEST_ARGS: Final = { + "api_base": "https://api.openai.com/v1/chat/completions", + "headers": {}, + "complete_input_dict": {"messages": [{"role": "user", "content": "PRIVATE-PHRASE"}]}, +} + + +def _pre_call_with_raw_request_logging(logging_obj) -> dict: + metadata: Final = {"user_api_key_alias": "qa-key"} + logging_obj.model_call_details["litellm_params"] = {"metadata": metadata} + logging_obj.log_raw_request_response = True + logging_obj.pre_call(input="hi", api_key="", additional_args=_PRIVATE_RAW_REQUEST_ARGS) + return metadata + + +def _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata: dict) -> None: + assert metadata["raw_request"] == REDACTED_BY_LITELLM + typed_dict: Final = logging_obj.model_call_details["raw_request_typed_dict"] + assert typed_dict["raw_request_body"] == _PRIVATE_RAW_REQUEST_ARGS["complete_input_dict"] + assert typed_dict["error"] is None + + +def test_pre_call_raw_request_honors_turn_off_message_logging_set_after_import(logging_obj, monkeypatch): + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + + metadata = _pre_call_with_raw_request_logging(logging_obj) + + _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata) + + +def test_pre_call_raw_request_honors_per_request_turn_off_message_logging(logging_obj, monkeypatch): + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + logging_obj.model_call_details["standard_callback_dynamic_params"] = {"turn_off_message_logging": True} + + metadata = _pre_call_with_raw_request_logging(logging_obj) + + _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata) + + +def test_debugging_log_honors_json_logs_set_after_import(logging_obj, monkeypatch): + monkeypatch.setattr(litellm, "json_logs", True) + logging_obj.litellm_request_debug = True + + with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning: + logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers={}, additional_args={}) + + assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"] + + +def test_debugging_log_with_json_logs_tolerates_missing_headers(logging_obj, monkeypatch): + monkeypatch.setattr(litellm, "json_logs", True) + logging_obj.litellm_request_debug = True + + with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning: + logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers=None, additional_args={}) + + assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"] + + def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]): import datetime diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index effc038f85b..c06216e4f4e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -626,6 +626,29 @@ def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter) ] +def test_return_raw_request_ignores_turn_off_message_logging( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "PRIVATE-PHRASE"}] + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + + request: Final = return_raw_request( + endpoint=CallTypes.completion, + kwargs={"model": model, "messages": messages}, + ) + + assert route.call_count == 0 + assert request.get("error") is None + assert request["raw_request_body"]["messages"] == messages + + def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter): """Regression test: completion() must forward the verbosity param to the provider request body.""" from litellm.types.utils import CallTypes From cd1107aac481278561d6d0f9da054ba6bf4095a5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:18:51 -0700 Subject: [PATCH 28/58] fix(router): parse the classifier verdict out of surrounding prose instead of falling to the default tier (#43215) * fix(router): parse the classifier verdict out of prose and fences on every parse path The complexity router's classifier parsers only tolerated a bare JSON object (labeled tier) or a leading Markdown fence (capability, LLM V2), so a json_object classifier that writes its verdict fenced and then explains it in markdown, which Bedrock Haiku 4.5 does on nearly every Claude Code request, failed validation and every request fell to the fallback tier. All three parse paths now extract the first complete JSON object from the reply with json.JSONDecoder.raw_decode, whatever prose or fence surrounds it, and a reply that still fails validation is logged with the pydantic field problems and the raw reply text, withheld when the request turns off message logging. The capability failure reason names the exception type like the labeled path does instead of interpolating str(e), which for a ValidationError carried the whole reply as input_value and for TimeoutError was empty. * fix(router): stop rejecting an LLM V2 forecast over a long explanation field LLMV2Verdict capped crux and each forecast's likely_failure at 512 characters through the ShortText alias, so a verdict whose explanation ran long failed validation and the request fell to the capable tier, even though nothing downstream reads either field. Five of nineteen real Claude Code replies from Bedrock Haiku 4.5 tripped the cap. Both fields keep the strip and non-empty constraints and lose the length cap; the operator-set calibration version keeps ShortText. * fix(complexity_router): withhold the rejected classifier reply under every message-logging opt-out and survive undecodable replies * fix(complexity_router): withhold the rejected classifier reply when the redaction decision cannot be made --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../capability_classifier.py | 29 +++-- .../complexity_router/complexity_router.py | 57 ++++++++- .../complexity_router/llm_v2.py | 5 +- .../router_strategy/test_complexity_router.py | 81 ++++++++++++- .../router_strategy/test_llm_v2.py | 113 +++++++++++++++++- 5 files changed, 261 insertions(+), 24 deletions(-) diff --git a/litellm/router_strategy/complexity_router/capability_classifier.py b/litellm/router_strategy/complexity_router/capability_classifier.py index 21046ff3421..93077af9e47 100644 --- a/litellm/router_strategy/complexity_router/capability_classifier.py +++ b/litellm/router_strategy/complexity_router/capability_classifier.py @@ -202,15 +202,26 @@ def capability_classifier_system_prompt(mode: Literal["json_schema", "json_objec ) -def unwrap_classifier_json(content: str) -> str: - """Remove the optional Markdown fence without repairing or weakening verdict JSON.""" - text: Final = content.strip() - if not text.startswith("```"): - return text - unfenced: Final = text.removeprefix("```").removeprefix("json").lstrip("\n\r") - return unfenced.removesuffix("```").strip() +_JSON_DECODER: Final = json.JSONDecoder() + + +def _complete_json_object_at(content: str, start: int) -> str | None: + try: + _, end = _JSON_DECODER.raw_decode(content, start) + except (ValueError, RecursionError): + return None + return content[start:end] + + +def extract_classifier_json(content: str) -> str: + """Return the first complete JSON object in the reply, whatever prose or fence surrounds it. + + A reply with no complete object comes back stripped so the caller's validation names the defect.""" + object_starts: Final = (index for index, char in enumerate(content) if char == "{") + candidates: Final = (_complete_json_object_at(content, start) for start in object_starts) + return next((candidate for candidate in candidates if candidate is not None), content.strip()) def parse_capability_classifier_verdict(content: str) -> CapabilityClassifierVerdict: - """Parse raw JSON or the fenced JSON shape tolerated by Switchyard.""" - return CapabilityClassifierVerdict.model_validate_json(unwrap_classifier_json(content)) + """Parse the verdict object out of a bare, fenced, or prose-wrapped reply.""" + return CapabilityClassifierVerdict.model_validate_json(extract_classifier_json(content)) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 0f252952a9d..9df6306436b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -29,6 +29,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, TypeAdapter, ValidationError, create_model +from pydantic_core import ErrorDetails from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import claim_affinity_pin @@ -85,8 +86,8 @@ from .capability_classifier import ( CapabilityClassifierForecast, capability_classifier_response_format, capability_classifier_system_prompt, + extract_classifier_json, parse_capability_classifier_verdict, - unwrap_classifier_json, ) from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section from .config import ( @@ -427,6 +428,41 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | N ) +def _classifier_reply_is_private(request_kwargs: Mapping[str, object] | None) -> bool: + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, + ) + from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + kwargs: Final = dict(request_kwargs) if request_kwargs else {} + try: + return should_redact_message_logging( + { + "litellm_params": kwargs, + "standard_callback_dynamic_params": initialize_standard_callback_dynamic_params(kwargs), + } + ) + except AttributeError: + return True + + +def _validation_problem(detail: ErrorDetails) -> str: + location: Final = ".".join(str(part) for part in detail["loc"]) + return f"{location}: {detail['msg']}" if location else detail["msg"] + + +def _log_rejected_classifier_verdict( + error: ValidationError, content: str, request_kwargs: Mapping[str, object] | None +) -> None: + problems: Final = "; ".join(_validation_problem(detail) for detail in error.errors()) + reply: Final = ( + "raw reply withheld (message logging is off)" + if _classifier_reply_is_private(request_kwargs) + else f"raw reply: {content!r}" + ) + verbose_router_logger.warning("ComplexityRouter: classifier verdict rejected (%s); %s", problems, reply) + + _REMINDER_OPEN: Final = "" _REMINDER_CLOSE: Final = "" _DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),) @@ -2040,7 +2076,7 @@ class ComplexityRouter(CustomLogger): except Exception as e: # noqa: BLE001 -- every unavailable or invalid judge verdict must fail closed if breaker is not None and permit is not None: breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e)) - return self._capability_classifier_failure_outcome(f"capability classifier failed ({e})") + return self._capability_classifier_failure_outcome(f"capability classifier failed ({type(e).__name__})") def _capability_classifier_failure_outcome(self, reason: str, signal: str | None = None) -> ClassificationOutcome: """Fail closed to the configured capable tier without consulting another taxonomy.""" @@ -2449,7 +2485,11 @@ class ComplexityRouter(CustomLogger): content, classifier_cost = await self._call_classifier_model( messages_for_call, request_kwargs, encrypted_task=encrypted_task ) - raw_tier: Final = _LabeledTierClassification.model_validate_json(content).tier + try: + raw_tier: Final = _LabeledTierClassification.model_validate_json(extract_classifier_json(content)).tier + except ValidationError as error: + _log_rejected_classifier_verdict(error, content, request_kwargs) + raise tier: Final = self.config.resolve_classified_tier(raw_tier) if tier is None: raise ValueError(f"LLM classifier returned an unrecognized tier: {raw_tier!r}") @@ -2508,7 +2548,11 @@ class ComplexityRouter(CustomLogger): max_output_tokens=capability.max_output_tokens, encrypted_task=encrypted_task, ) - verdict: Final = parse_capability_classifier_verdict(content) + try: + verdict: Final = parse_capability_classifier_verdict(content) + except ValidationError as error: + _log_rejected_classifier_verdict(error, content, request_kwargs) + raise threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step) calibration: Final = capability.calibration forecast: Final = CapabilityClassifierForecast( @@ -2563,8 +2607,9 @@ class ComplexityRouter(CustomLogger): messages_for_call, request_kwargs, encrypted_task=encrypted, max_output_tokens=v2.max_output_tokens ) try: - verdict: Final = LLMV2Verdict.model_validate_json(unwrap_classifier_json(content)) - except ValidationError: + verdict: Final = LLMV2Verdict.model_validate_json(extract_classifier_json(content)) + except ValidationError as error: + _log_rejected_classifier_verdict(error, content, request_kwargs) return self._classifier_failure_outcome("Invalid LLM V2 forecast", prompt, system_prompt)._replace( classifier_cost=classifier_cost ) diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py index 18351237e65..8ef2f554ab2 100644 --- a/litellm/router_strategy/complexity_router/llm_v2.py +++ b/litellm/router_strategy/complexity_router/llm_v2.py @@ -16,6 +16,7 @@ from litellm.llms.base_llm.base_utils import ( from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)] +VerdictText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] class _SolverProfile(TypedDict): @@ -90,7 +91,7 @@ class LLMV2Demands(BaseModel): class LLMV2SolverForecast(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - likely_failure: ShortText + likely_failure: VerdictText p_solve: StrictFloat = Field(ge=0.0, le=1.0) @@ -104,7 +105,7 @@ class LLMV2SolverForecasts(BaseModel): class LLMV2Verdict(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - crux: ShortText + crux: VerdictText demands: LLMV2Demands verification: Literal["relevant", "partial", "unavailable", "unknown"] forecasts: LLMV2SolverForecasts diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 4e146b59b61..3401a335b2f 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2679,6 +2679,25 @@ def _llm_response(content: str, response_cost: float | None = None): return response +_REPLY_SHAPES: Final = ("fenced", "fenced-with-language", "prose-before", "prose-after", "fenced-then-prose") + + +def _wrapped_reply(shape: str, verdict: str) -> str: + match shape: + case "fenced": + return f" ```\n{verdict}\n``` " + case "fenced-with-language": + return f"```json\n{verdict}\n```" + case "prose-before": + return f"Sure {{here}} is the verdict you asked for:\n\n{verdict}" + case "prose-after": + return f"{verdict}\n\nThe efficient solver should handle this {{well}}." + case "fenced-then-prose": + return f"```json\n{verdict}\n```\n\n## Reasoning\n\nThe task is coupled, so the forecasts differ." + case _: + raise AssertionError(shape) + + @pytest.fixture def llm_classifier_config() -> Dict: """Config with an LLM-based classifier wired to a 'haiku-classifier' model.""" @@ -3132,12 +3151,40 @@ class TestCapabilityClassifier: assert outcome.capability_forecast.threshold == pytest.approx(expected_threshold) @pytest.mark.asyncio - async def test_fenced_json_verdict_is_accepted(self, mock_router_instance): - reply = _capability_reply(p_solve=0.8) - mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(f"```json\n{reply}\n```")) + @pytest.mark.parametrize("shape", _REPLY_SHAPES) + async def test_verdict_wrapped_in_fence_or_prose_is_accepted(self, mock_router_instance, shape: str): + reply = _wrapped_reply(shape, _capability_reply(p_solve=0.8)) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply)) outcome = await self._router(mock_router_instance).aclassify("do the task") assert outcome.tier == ComplexityTier.SIMPLE assert outcome.cause == "capability_classifier" + assert outcome.capability_forecast is not None + assert outcome.capability_forecast.p_solve == 0.8 + + @pytest.mark.asyncio + @pytest.mark.parametrize("message_logging_off", (False, True)) + async def test_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off( + self, mock_router_instance, caplog: pytest.LogCaptureFixture, message_logging_off: bool + ): + reply = "The task text is too {vague} for a forecast, sorry." + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply)) + outcome = await self._router(mock_router_instance).aclassify( + "do the task", request_kwargs={"turn_off_message_logging": message_logging_off} + ) + assert outcome.cause == "capability_classifier_fallback" + assert "capability classifier failed (ValidationError)" in caplog.text + assert "classifier verdict rejected (" in caplog.text + assert ("raw reply withheld" in caplog.text) is message_logging_off + assert (reply in caplog.text) is not message_logging_off + + @pytest.mark.asyncio + async def test_call_failure_reason_names_the_exception_type( + self, mock_router_instance, caplog: pytest.LogCaptureFixture + ): + mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError()) + outcome = await self._router(mock_router_instance).aclassify("do the task") + assert outcome.cause == "capability_classifier_fallback" + assert "capability classifier failed (TimeoutError)" in caplog.text @pytest.mark.asyncio async def test_decimal_rounding_does_not_break_inclusive_threshold(self, mock_router_instance): @@ -3954,6 +4001,34 @@ class TestLLMClassifier: assert call_kwargs["model"] == "haiku-classifier" assert call_kwargs["timeout"] == 0.4 + @pytest.mark.asyncio + @pytest.mark.parametrize("shape", _REPLY_SHAPES) + async def test_aclassify_llm_verdict_wrapped_in_fence_or_prose_still_decides_the_tier( + self, llm_complexity_router, mock_router_instance, shape: str + ): + reply = _wrapped_reply(shape, '{"tier": "COMPLEX"}') + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply)) + outcome = await llm_complexity_router.aclassify("hi") + assert outcome.tier == ComplexityTier.COMPLEX + assert outcome.cause == "llm_classifier" + assert "llm-classifier:COMPLEX" in outcome.signals + + @pytest.mark.asyncio + @pytest.mark.parametrize("message_logging_off", (False, True)) + async def test_aclassify_llm_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off( + self, llm_complexity_router, mock_router_instance, caplog: pytest.LogCaptureFixture, message_logging_off: bool + ): + reply = "I would call this COMPLEX, the {tier} field is implied." + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply)) + outcome = await llm_complexity_router.aclassify( + "hi", request_kwargs={"turn_off_message_logging": message_logging_off} + ) + assert outcome.cause != "llm_classifier" + assert "LLM classifier failed (ValidationError)" in caplog.text + assert "classifier verdict rejected (" in caplog.text + assert ("raw reply withheld" in caplog.text) is message_logging_off + assert (reply in caplog.text) is not message_logging_off + @pytest.mark.asyncio async def test_aclassify_llm_success_captures_classifier_cost(self, llm_complexity_router, mock_router_instance): """The classifier call is billed, so its cost must ride the outcome. diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 6fb6df3265d..3fd2e8808e7 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -66,6 +66,25 @@ def _response(content: str) -> ModelResponse: return response +_REPLY_SHAPES: Final = ("fenced", "fenced-with-language", "prose-before", "prose-after", "fenced-then-prose") + + +def _wrapped_reply(shape: str, verdict: str) -> str: + match shape: + case "fenced": + return f" ```\n{verdict}\n``` " + case "fenced-with-language": + return f"```json\n{verdict}\n```" + case "prose-before": + return f"Sure {{here}} is the verdict you asked for:\n\n{verdict}" + case "prose-after": + return f"{verdict}\n\nThe efficient solver should handle this {{well}}." + case "fenced-then-prose": + return f"```json\n{verdict}\n```\n\n## Reasoning\n\nThe task is coupled, so the forecasts differ." + case _: + raise AssertionError(shape) + + def _router(content: str, config: ComplexityRouterConfig | None = None) -> tuple[ComplexityRouter, MagicMock]: client: Final = MagicMock(spec=Router) client.acompletion = AsyncMock(return_value=_response(content)) @@ -334,13 +353,12 @@ async def test_json_object_mode_supplies_schema_in_prompt() -> None: @pytest.mark.asyncio @pytest.mark.parametrize("mode", ("json_schema", "json_object")) -@pytest.mark.parametrize("fence", ("```json", "```")) -async def test_fenced_forecast_routes_by_validated_probabilities(mode: str, fence: str) -> None: +@pytest.mark.parametrize("shape", _REPLY_SHAPES) +async def test_wrapped_forecast_routes_by_validated_probabilities(mode: str, shape: str) -> None: base: Final = _config().llm_v2_config assert base is not None config: Final = _config(llm_v2_config={**base.model_dump(), "response_format": mode}) - content: Final = f" {fence}\n{_verdict().model_dump_json()}\n``` " - router, client = _router(content, config) + router, client = _router(_wrapped_reply(shape, _verdict().model_dump_json()), config) result: Final = await router.async_pre_routing_hook( model="v2-router", messages=[{"role": "user", "content": "Fix nested behavior"}], request_kwargs={} ) @@ -454,6 +472,93 @@ async def test_provider_failure_redacts_prompt_text_from_warning(caplog: pytest. assert "private task text" not in caplog.text +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ("crux", "likely_failure")) +async def test_long_verdict_explanations_still_route_by_validated_probabilities(field: str) -> None: + explanation: Final = "The solver must keep the nested retry behavior intact while it edits. " * 12 + assert len(explanation) > 512 + verdict: Final = _verdict().model_dump() + if field == "crux": + content: Final = json.dumps({**verdict, "crux": explanation}) + else: + forecasts: Final = {**verdict["forecasts"], "efficient": {**verdict["forecasts"]["efficient"], field: explanation}} + content = json.dumps({**verdict, "forecasts": forecasts}) + router, _ = _router(content) + outcome: Final = await router.aclassify("Fix nested behavior") + assert outcome.cause == "llm_v2_classifier" + assert outcome.llm_v2_forecast is not None + assert outcome.llm_v2_forecast.use_efficient + + +@pytest.mark.parametrize("field", ("crux", "likely_failure")) +def test_blank_verdict_explanations_are_still_rejected(field: str) -> None: + verdict: Final = _verdict().model_dump() + blank: Final = ( + {**verdict, "crux": " "} + if field == "crux" + else {**verdict, "forecasts": {**verdict["forecasts"], "capable": {"likely_failure": " ", "p_solve": 0.5}}} + ) + with pytest.raises(ValidationError): + LLMV2Verdict.model_validate(blank) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("message_logging_off", (False, True)) +async def test_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off( + caplog: pytest.LogCaptureFixture, message_logging_off: bool +) -> None: + reply: Final = "I cannot forecast this one, the task text is too {vague} to score." + router, _ = _router(reply) + outcome: Final = await router.aclassify("hi", request_kwargs={"turn_off_message_logging": message_logging_off}) + assert outcome.cause == "llm_v2_fallback" + assert "classifier verdict rejected (" in caplog.text + assert "Invalid LLM V2 forecast" in caplog.text + assert ("raw reply withheld" in caplog.text) is message_logging_off + assert (reply in caplog.text) is not message_logging_off + + +_MESSAGE_LOGGING_OPT_OUTS: Final = ( + pytest.param({"turn_off_message_logging": "True"}, False, id="key-logging-settings-string"), + pytest.param({"metadata": {"headers": {"x-litellm-enable-message-redaction": "true"}}}, False, id="redaction-header"), + pytest.param({}, True, id="global-setting"), + pytest.param({"metadata": {"headers": None}}, False, id="undecidable-headers-fail-closed"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("request_kwargs", "global_off"), _MESSAGE_LOGGING_OPT_OUTS) +async def test_unparseable_reply_text_is_withheld_under_every_message_logging_opt_out( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + request_kwargs: dict[str, object], + global_off: bool, +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", global_off) + reply: Final = "I cannot forecast this one, the task text is too {vague} to score." + router, _ = _router(reply) + outcome: Final = await router.aclassify("hi", request_kwargs=request_kwargs) + assert outcome.cause == "llm_v2_fallback" + assert "raw reply withheld" in caplog.text + assert reply not in caplog.text + + +_REPLIES_THE_JSON_SCANNER_CANNOT_DECODE: Final = ( + pytest.param('{"a":' * 3000, id="deeply-nested"), + pytest.param('{"capability_p": ' + "9" * 5000 + "}", id="integer-over-the-digit-limit"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply", _REPLIES_THE_JSON_SCANNER_CANNOT_DECODE) +async def test_undecodable_reply_is_rejected_as_an_invalid_forecast( + caplog: pytest.LogCaptureFixture, reply: str +) -> None: + router, _ = _router(reply) + outcome: Final = await router.aclassify("hi", request_kwargs={}) + assert outcome.cause == "llm_v2_fallback" + assert "Invalid LLM V2 forecast" in caplog.text + + def test_response_schema_requires_both_model_forecasts() -> None: with pytest.raises(ValidationError): LLMV2Verdict.model_validate( From 474ab91c0972f5074eac297dc63606f607ce5dda Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 25 Sep 2026 16:35:15 -0700 Subject: [PATCH 29/58] test(zerobus): move tests into active CI selection (#43235) * test(zerobus): move tests into the active CI selection * test(zerobus): add package marker for unit test discovery --- tests/unit/integrations/zerobus/__init__.py | 0 .../integrations/zerobus/test_zerobus_client.py | 0 .../integrations/zerobus/test_zerobus_logger.py | 0 .../integrations/zerobus/test_zerobus_row.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/integrations/zerobus/__init__.py rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_client.py (100%) rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_logger.py (100%) rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_row.py (100%) diff --git a/tests/unit/integrations/zerobus/__init__.py b/tests/unit/integrations/zerobus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py b/tests/unit/integrations/zerobus/test_zerobus_client.py similarity index 100% rename from tests/test_litellm/integrations/zerobus/test_zerobus_client.py rename to tests/unit/integrations/zerobus/test_zerobus_client.py diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py b/tests/unit/integrations/zerobus/test_zerobus_logger.py similarity index 100% rename from tests/test_litellm/integrations/zerobus/test_zerobus_logger.py rename to tests/unit/integrations/zerobus/test_zerobus_logger.py diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_row.py b/tests/unit/integrations/zerobus/test_zerobus_row.py similarity index 100% rename from tests/test_litellm/integrations/zerobus/test_zerobus_row.py rename to tests/unit/integrations/zerobus/test_zerobus_row.py From e0fb89bc82195a0d7ad4d7799286ffea61d2b377 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:05:53 -0700 Subject: [PATCH 30/58] fix(proxy): keep the submitted body out of 422 validation errors (#43231) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../common_utils/validation_error_body.py | 13 ++++ litellm/proxy/list_api/common.py | 10 +-- litellm/proxy/proxy_server.py | 16 ++--- .../proxy_setting_endpoints.py | 3 +- .../proxy_server/test_exception_handlers.py | 66 ++++++++++++++++++- .../proxy_server/test_routes_onboarding.py | 19 ++++++ .../test_proxy_setting_endpoints.py | 16 +++++ .../test_validation_error_body.py | 46 +++++++++++++ 8 files changed, 168 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/common_utils/validation_error_body.py create mode 100644 tests/unit/proxy/common_utils/test_validation_error_body.py diff --git a/litellm/proxy/common_utils/validation_error_body.py b/litellm/proxy/common_utils/validation_error_body.py new file mode 100644 index 00000000000..b21f33a2434 --- /dev/null +++ b/litellm/proxy/common_utils/validation_error_body.py @@ -0,0 +1,13 @@ +from collections.abc import Sequence + +from typing_extensions import ReadOnly, TypedDict + + +class ValidationErrorDetail(TypedDict): + type: ReadOnly[str] + loc: ReadOnly[tuple[int | str, ...]] + msg: ReadOnly[str] + + +def public_validation_errors(errors: Sequence[ValidationErrorDetail]) -> tuple[ValidationErrorDetail, ...]: + return tuple(ValidationErrorDetail(type=error["type"], loc=error["loc"], msg=error["msg"]) for error in errors) diff --git a/litellm/proxy/list_api/common.py b/litellm/proxy/list_api/common.py index daa6414fd94..efa8a271459 100644 --- a/litellm/proxy/list_api/common.py +++ b/litellm/proxy/list_api/common.py @@ -8,8 +8,8 @@ from fastapi import Request from fastapi.dependencies.utils import get_flat_params from fastapi.params import ParamTypes from fastapi.responses import JSONResponse -from typing_extensions import ReadOnly, TypedDict +from litellm.proxy.common_utils.validation_error_body import ValidationErrorDetail from litellm.types.proxy.management_endpoints.management_v1 import ( ListLinks, PageLinks, @@ -58,14 +58,6 @@ def escape_like(value: str) -> str: return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") -class ValidationErrorDetail(TypedDict): - """The keys of a pydantic/FastAPI validation error a problem document needs.""" - - type: ReadOnly[str] - loc: ReadOnly[tuple[int | str, ...]] - msg: ReadOnly[str] - - def _is_length_error_of_rejected_items(error: ValidationErrorDetail, errors: Sequence[ValidationErrorDetail]) -> bool: """pydantic counts only items that validated, so a bad item also trips the parent's min_length.""" return error["type"] == "too_short" and any( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61304f0d919..646eca071d1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -476,6 +476,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( project_spend_counter_key, tag_cache_key, ) +from litellm.proxy.common_utils.validation_error_body import public_validation_errors from litellm.proxy.config_resolvers import ( FieldSource, SettingsStore, @@ -551,7 +552,6 @@ from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_sp from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.list_api.common import ( ManagementProblem, - ValidationErrorDetail, problem_response, request_validation_problem, ) @@ -1983,16 +1983,14 @@ class _ExceptionRow(TypedDict, total=False): @app.exception_handler(RequestValidationError) async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError): + public_errors: Final = public_validation_errors(exc.errors()) + public_exc: Final = RequestValidationError(public_errors).with_traceback(exc.__traceback__) if request.url.path.startswith(MANAGEMENT_V1_PREFIX): - validation_errors: Final[Sequence[ValidationErrorDetail]] = exc.errors() - problem: Final = request_validation_problem(validation_errors) - _close_dangling_otel_server_span(request, problem.status, exc=exc) + problem: Final = request_validation_problem(public_errors) + _close_dangling_otel_server_span(request, problem.status, exc=public_exc) return problem_response(problem) - _close_dangling_otel_server_span(request, 422, exc=exc) - return JSONResponse( - status_code=422, - content={"detail": jsonable_encoder(exc.errors())}, - ) + _close_dangling_otel_server_span(request, 422, exc=public_exc) + return JSONResponse(status_code=422, content={"detail": public_errors}) @app.exception_handler(Exception) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index c91b1afd64a..227f0e7f795 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.validation_error_body import public_validation_errors from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError from litellm.proxy.config_resolvers.sso import ( @@ -1875,7 +1876,7 @@ async def update_ui_settings( try: settings: Final = effective_cls.model_validate(settings_body) except ValidationError as e: - raise HTTPException(status_code=422, detail=e.errors()) + raise HTTPException(status_code=422, detail=public_validation_errors(e.errors())) unsupported_team_fields: Final = sorted( frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_PERMISSIONS diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index 089c2d57594..16cb1146ff5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -265,7 +265,7 @@ def test_close_dangling_otel_server_span_logger_raises_state_cleared_error(monke @pytest.mark.asyncio async def test_otel_request_validation_exception_handler_returns_422_detail(): - errors = [{"loc": ["body", "model"], "msg": "field required", "type": "missing"}] + errors = [{"loc": ["body", "model"], "msg": "field required", "type": "missing", "input": {"messages": []}}] exc = RequestValidationError(errors) request = _make_request() @@ -273,7 +273,69 @@ async def test_otel_request_validation_exception_handler_returns_422_detail(): body = json.loads(response.body) assert response.status_code == 422 - assert normalize(body) == {"detail": exc.errors()} + assert body == {"detail": [{"type": "missing", "loc": ["body", "model"], "msg": "field required"}]} + + +_SUBMITTED_PASSWORD: Final = "hunter2-Sup3rSecret!" +_PASSWORD_LEAKING_ERRORS: Final = ( + { + "type": "missing", + "loc": ["body", "new_password"], + "msg": "Field required", + "input": {"current_password": _SUBMITTED_PASSWORD}, + }, + { + "type": "value_error", + "loc": ["body", "password"], + "msg": "Value error, password cannot be set via /user/new", + "input": _SUBMITTED_PASSWORD, + "ctx": {"error": ValueError(_SUBMITTED_PASSWORD)}, + }, +) +_PUBLIC_ERRORS: Final = ( + {"type": "missing", "loc": ["body", "new_password"], "msg": "Field required"}, + {"type": "value_error", "loc": ["body", "password"], "msg": "Value error, password cannot be set via /user/new"}, +) + + +@pytest.mark.asyncio +async def test_otel_request_validation_exception_handler_never_echoes_the_submitted_body(): + """A pydantic error carries the offending value as ``input`` (the whole body for a + ``missing`` error) and input-derived values in ``ctx``; a caller who mistyped a + request holding a password must not get that password back.""" + exc = RequestValidationError(list(_PASSWORD_LEAKING_ERRORS)) + + response = await otel_request_validation_exception_handler(request=_make_request(), exc=exc) + + assert response.status_code == 422 + assert json.loads(response.body) == {"detail": list(_PUBLIC_ERRORS)} + assert _SUBMITTED_PASSWORD.encode() not in response.body + + +@pytest.mark.asyncio +async def test_otel_request_validation_exception_handler_hands_the_span_only_the_public_errors(monkeypatch): + """The OTEL SERVER span's error message is ``str(exc)``, which FastAPI builds from + every error dict ``input`` included, so the span gets the same public-only errors + the caller does, and keeps the traceback the original carried.""" + import litellm.proxy.proxy_server as ps + + fake_logger = MagicMock() + monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False) + exc = RequestValidationError(list(_PASSWORD_LEAKING_ERRORS)) + try: + raise exc + except RequestValidationError as raised: + original_traceback = raised.__traceback__ + request = _make_request(parent_otel_span=MagicMock()) + + await otel_request_validation_exception_handler(request=request, exc=exc) + + (_span, span_exc, status_code) = fake_logger.record_error_attributes_on_span.call_args.args + assert status_code == 422 + assert isinstance(span_exc, RequestValidationError) + assert list(span_exc.errors()) == list(_PUBLIC_ERRORS) + assert _SUBMITTED_PASSWORD not in str(span_exc) + assert span_exc.__traceback__ is original_traceback @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py index 778acc1baab..6c1d869d113 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py @@ -317,6 +317,25 @@ def test_claim_onboarding_link_missing_field_422(client, monkeypatch, mock_prism assert any("password" in str(item) for item in body["detail"]) +def test_claim_onboarding_link_422_never_echoes_the_submitted_password(client): + """A body that fails validation is answered with the field path and message only; + pydantic's ``input`` (the whole submitted body for a missing field, password + included) must never come back to the caller or land in whatever logs the response.""" + password = "hunter2-Sup3rSecret!" + + response = client.post( + "/onboarding/claim_token", + json={"invitation_link": "abc", "password": password}, + ) + + assert response.status_code == 422 + assert password.encode() not in response.content + detail = response.json()["detail"] + assert detail[0]["loc"] == ["body", "user_id"] + assert detail[0]["msg"] + assert set(detail[0]) == {"type", "loc", "msg"} + + def test_claim_onboarding_link_bad_onboarding_jwt_401( client, monkeypatch, mock_prisma ): diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 08d542df16c..0d7a713a380 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3817,6 +3817,22 @@ class TestTeamAdminEditableTeamFieldsSetting: assert response.status_code == 422 + def test_patch_422_never_echoes_the_submitted_value(self, monkeypatch): + self._as_proxy_admin(monkeypatch) + submitted = "hunter2-Sup3rSecret!" + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": submitted}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 422 + assert submitted.encode() not in response.content + detail = response.json()["detail"] + assert detail[0]["loc"] == ["team_admin_editable_team_fields"] + assert detail[0]["msg"] + assert set(detail[0]) == {"type", "loc", "msg"} + def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) general_settings: dict = {"team_admin_editable_team_fields": []} diff --git a/tests/unit/proxy/common_utils/test_validation_error_body.py b/tests/unit/proxy/common_utils/test_validation_error_body.py new file mode 100644 index 00000000000..a86f17b7461 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_validation_error_body.py @@ -0,0 +1,46 @@ +from typing import Final + +from litellm.proxy.common_utils.validation_error_body import public_validation_errors + +_PASSWORD: Final = "hunter2-Sup3rSecret!" + + +def test_public_validation_errors_drops_input_ctx_and_url(): + errors: Final = ( + { + "type": "missing", + "loc": ("body", "user_id"), + "msg": "Field required", + "input": {"invitation_link": "abc", "password": _PASSWORD}, + "url": "https://errors.pydantic.dev/2/v/missing", + }, + { + "type": "value_error", + "loc": ("body", "password"), + "msg": "Value error, password cannot be set here", + "input": _PASSWORD, + "ctx": {"error": ValueError(_PASSWORD)}, + }, + ) + + public: Final = public_validation_errors(errors) + + assert public == ( + {"type": "missing", "loc": ("body", "user_id"), "msg": "Field required"}, + {"type": "value_error", "loc": ("body", "password"), "msg": "Value error, password cannot be set here"}, + ) + assert _PASSWORD not in repr(public) + + +def test_public_validation_errors_keeps_type_loc_and_msg_verbatim_in_order(): + errors: Final = ( + {"type": "int_parsing", "loc": ("body", "litellm_params", "rpm"), "msg": "Input should be a valid integer"}, + {"type": "extra_forbidden", "loc": ("body", "users", 0, "user_emial"), "msg": "Extra inputs are not permitted"}, + {"type": "too_short", "loc": ("body", "users"), "msg": "List should have at least 1 item"}, + ) + + assert public_validation_errors(errors) == errors + + +def test_public_validation_errors_empty_in_empty_out(): + assert public_validation_errors(()) == () From a11a93f44a557dd91100c2e394006b5df0daed65 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 17:10:13 -0700 Subject: [PATCH 31/58] test: move tests/test_litellm core utils, routing, responses, caching and rust_bridge into tests/unit (#43199) * ci: run the unit_selection.sh shard files on every event instead of only fork pull requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rename fork-flag to unit-flag now that it applies on every event * test: move tests/test_litellm root and small trees into tests/unit Pure renames, no content changes. Follow-up commits in this PR fix references, merge the three files that already existed in tests/unit, keep live-provider tests in tests/test_litellm and wire CI. * test: carry tests/test_litellm conftest isolation into tests/unit Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS, proxy-URL and keychain env, and session-end client cleanup now reset for unit tests too. The environment isolation owns its MonkeyPatch so a test's own monkeypatch is undone before the model-cost teardown runs. * test: merge, split and prune the moved root and small-tree tests Merge batches/test_batch_utils.py and the chat_completions and messages dispatch tests into the files that already existed in tests/unit. Keep the live Gemini interactions tests, the async image-fetch format test and the OpenAI embedding scorer test in tests/test_litellm since they need real network or keys. Put test_router.py under tests/unit/test_router so the existing package no longer shadows it. Delete eight tests the audit found superseded by stronger ones kept in this move. * ci: run the moved root and small-tree tests under their legacy flags Add the misc and responses-caching-types flags to unit_selection.sh and CircleCI, extend enterprise-routing and mcp-integration, and point the legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest and change classifier at the new paths. * test: make the new tests/unit directories packages tests/unit/test_package_layout.py requires every directory to carry an __init__.py, and without one the moved and retained test_litellm_responses_bridge.py modules collide on import. * test: scope the unit socket block to tests/unit in shared sessions The GHA shards collect the legacy test-path and the unit selection in one pytest session. The unit conftest's loopback-only block leaked into legacy modules that reach the network at import. The legacy conftest now lifts the restriction at collect and setup time, and the unit conftest re-applies it when collecting its own modules. * test: move tests/test_litellm/llms into tests/unit/llms Rename-only. Moves the provider tests and the fine-tuning fixtures they load, mirroring the old paths. Follow-up commits merge, split and wire them. * test: merge, split and prune the moved llms tests Merges the Databricks chat transformation tests into the existing unit file, keeps the tests that need real keys or the network in tests/test_litellm, deletes the audited tests a stronger unit test already covers, and points imports at tests.unit.llms. * ci: run the moved llms tests under their legacy flags The Vertex AI and All Other Providers shards keep their legacy test-path for the retained files and add the llm-vertex-ai and llm-other-providers unit selections. CircleCI gets matching unit jobs. * test: make the tests/unit/llms directories packages Adds __init__.py to the moved dirs and drops the legacy ones whose directories no longer hold tests. * test: drop script runners and path hacks the llms split left dangling The __main__ runners in the split openai_like files and the Databricks e2e runner called tests that now live in the other half of the split or were deleted. The retained legacy halves also no longer need sys.path edits. * test: give the shard-script tests their own GITHUB_OUTPUT They only passed where the runner set it. The CircleCI unit job's env allowlist drops it, so the script's redirect failed there. * test: point the router and module-deletion checks at tests/unit router_code_coverage and code_qa_check_tests only searched tests/test_litellm, so the moved router tests no longer counted. The two silent-experiment tests the audit deleted were the only direct callers of those methods; they are replaced with tests that assert the forwarded shadow request and the recursion guard. * test: move tests/test_litellm integrations and secret_managers into tests/unit Rename-only. Mirrors the old paths, including the directory conftests and the prompt and JSON fixtures. Follow-up commits prune and wire them. * test: prune and repoint the moved integrations tests Deletes the 7 audited tests a stronger test in the same tree already covers, imports the TLS sink helpers from their new conftest path, and restores os.environ after each integrations test. Some presets write OTEL_EXPORTER_OTLP_HEADERS straight into os.environ, and without the legacy tree's test ordering that header leaked into the AgentOps tests. * ci: run the moved integrations tests under their legacy flag The integrations GHA shard and a new CircleCI job run the integrations unit selection. secret_managers joins the misc selection. * docs: point integrations and secret_managers references at tests/unit * test: make the moved integrations directories packages * test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path The Databricks e2e file is a manual script whose main() calls the tests that were pruned, so pruning them broke the documented run. It is back to its main version. The SageMaker Nova docstring now points at the file's real location in tests/local_testing. * test: move tests/test_litellm core utils, routing, responses, caching and rust_bridge into tests/unit Rename-only. Mirrors the old paths, including fixtures, the stubtest config and the native-route wheel script. Two files that collide with existing unit files are merged in a follow-up commit. * test: merge, prune and repoint the moved core, routing, responses, caching and rust_bridge tests Merges the two files that collided with existing unit files, folding the legacy extra case into test_is_chat_completion_cached_dict, and deletes the 9 audited tests a stronger test in the same file already covers. Keeps what needs the network in tests/test_litellm: test_tokenizers pulls a tokenizer from the Hugging Face hub, and the gpt2 and r50k_base tokenizer cases download their BPE files. The unit core_utils conftest points TIKTOKEN_CACHE_DIR at litellm's bundled encodings so the rest never depend on import order to stay offline, and FakeSecretVault moves to a shared module so both trees can build it. * ci: run the moved core, routing, responses, caching and rust_bridge tests under their flags core_utils gets a core-utils flag and CircleCI job, and its GHA shard keeps the legacy path for the retained network tests. router_utils and router_strategy join enterprise-routing, responses joins responses-caching-types (minus responses/mcp, which mcp-integration owns), caching joins caching-local and rust_bridge joins misc. The redis-compat, test-rust, stubtest and merge-smoke paths follow the move. * docs: point the Rust crate references at tests/unit * test: make the moved core, routing and rust_bridge directories packages * test: keep the no-loop DualCache batch_get_cache regression test It runs the sync path outside any event loop, which the inside-loop test cannot, so a change that picks the Redis client by loop state would only show up there. * test: keep the job's UNIT_FLAG out of the shard-script tests * fix(url_utils): block 192.0.0.0/24 on every Python patch release * test: move the new budget limiter tests into tests/unit/router_strategy * test: move the new sentry scrubbing tests into tests/unit/litellm_core_utils * test: move the new zerobus tests into tests/unit/integrations * test: make tests/unit/integrations/zerobus a package * test: load litellm's own tiktoken cache setup once instead of resetting it per test --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/unit_selection.sh | 9 +- .circleci/tests.yml | 7 + .github/merge-smoke-tests.json | 8 +- .github/workflows/test-redis-compat.yml | 12 +- .github/workflows/test-rust.yml | 6 +- .github/workflows/test-unit.yml | 10 +- Makefile | 6 +- .../crates/callbacks-legacy-python/src/lib.rs | 2 +- litellm-rust/crates/secrets/README.md | 2 +- litellm/litellm_core_utils/url_utils.py | 10 +- .../base_responses_api.py | 2 +- tests/llm_translation/test_gemini.py | 2 +- .../caching/test_caching_handler.py | 867 --------- tests/test_litellm/conftest.py | 68 +- .../litellm_core_utils/__init__.py | 1 - .../litellm_core_utils/test_token_counter.py | 1572 +---------------- .../litellm_core_utils/test_tokenizer.py | 409 +---- tests/test_litellm/proxy/client/test_chat.py | 2 +- .../proxy/hooks/test_tpm_concurrent.py | 2 +- .../test_streaming_handler.py | 4 +- tests/test_litellm/proxy/test_proxy_server.py | 4 +- tests/test_litellm/proxy/test_proxy_utils.py | 2 +- .../rust_bridge/messages/test_route_host.py | 124 -- .../rust_bridge/responses/__init__.py | 0 .../tokenizer/test_fast_count.py | 2 +- .../test_a2a_streaming_iterator.py | 2 +- tests/unit/a2a_protocol/test_main.py | 2 +- .../caching/test_azure_blob_cache.py | 0 .../caching/test_caching.py | 0 tests/unit/caching/test_caching_handler.py | 804 +++++++++ ...test_check_and_fix_namespace_none_guard.py | 0 .../caching/test_disk_cache.py | 0 .../caching/test_dual_cache.py | 0 .../caching/test_embedding_router.py | 0 .../caching/test_evicted_client_closer.py | 0 .../caching/test_gcs_cache.py | 0 .../caching/test_in_memory_cache.py | 0 .../caching/test_llm_caching_handler.py | 0 .../caching/test_llm_client_cache_e2e.py | 0 .../caching/test_qdrant_semantic_cache.py | 2 +- .../caching/test_redis_cache.py | 0 .../caching/test_redis_cluster_cache.py | 0 .../test_redis_cluster_node_isolation.py | 0 .../caching/test_redis_connection_pool.py | 0 .../caching/test_redis_semantic_cache.py | 2 +- .../caching/test_s3_cache.py | 0 .../caching/test_valkey_semantic_cache.py | 0 .../__init__.py | 0 .../azure_shell_tool.json | 0 .../context_management_and_shell.json | 0 .../test_compression_interception_handler.py | 2 +- tests/unit/litellm_core_utils/conftest.py | 15 + .../litellm_core_utils/event_loop_lag.py | 0 .../litellm_core_utils/fake_secret_vault.py | 67 + .../llm_cost_calc}/__init__.py | 0 .../test_azure_assistant_cost_tracking.py | 0 .../llm_cost_calc/test_guardrail_cost.py | 0 .../llm_cost_calc/test_llm_cost_calc_utils.py | 0 .../test_openai_cache_write_cost.py | 0 .../test_responses_cache_cost_breakdown.py | 0 .../test_tool_call_cost_tracking.py | 0 ...est_tool_call_cost_tracking_dict_safety.py | 0 .../test_usage_object_transformation.py | 0 .../test_zero_cost_diagnostic.py | 0 .../llm_response_utils/test_get_api_base.py | 0 .../messages_with_counts.py | 0 .../prompt_templates}/__init__.py | 0 ...edrock_converse_strict_tools_opus_47_48.py | 0 ...ore_utils_prompt_templates_common_utils.py | 0 ...llm_core_utils_prompt_templates_factory.py | 0 ...rompt_templates_mid_conversation_system.py | 0 .../specialty_caches}/__init__.py | 0 .../test_dynamic_logging_cache.py | 0 .../test_agentic_followup_kwargs.py | 0 .../test_anthropic_dedup_factory.py | 0 .../test_api_route_to_call_types.py | 0 .../litellm_core_utils/test_audio_utils.py | 0 .../litellm_core_utils/test_aws_partition.py | 0 .../test_bedrock_converse_dedup_factory.py | 0 .../litellm_core_utils/test_bug_report.py | 0 .../test_chat_completion_agentic_loop.py | 0 .../test_classifier_logging.py | 0 .../test_cli_token_utils.py | 0 .../test_cloud_storage_security.py | 0 .../test_codestral_provider_routing.py | 0 .../litellm_core_utils/test_core_helpers.py | 0 .../test_coroutine_checker.py | 0 .../litellm_core_utils/test_dd_tracing.py | 12 - .../test_decode_special_tokens.py | 0 .../test_dot_notation_indexing.py | 0 .../test_duration_parser.py | 0 .../test_error_normalization.py | 0 .../test_exception_mapping_utils.py | 0 .../test_extract_base64_image.py | 0 .../test_fallback_generalizations.py | 0 .../litellm_core_utils/test_fallback_utils.py | 0 .../test_get_litellm_params.py | 0 .../test_get_llm_provider_endpoint_match.py | 0 .../test_get_llm_provider_logic.py | 0 .../test_get_model_cost_map.py | 0 .../test_get_supported_openai_params.py | 0 .../test_health_check_helpers.py | 0 .../litellm_core_utils/test_image_handling.py | 0 ...test_initialize_dynamic_callback_params.py | 0 .../test_internal_call_metadata.py | 0 .../test_json_fragment_accumulator.py | 0 .../test_json_schema_validation.py | 0 .../test_litellm_logging.py | 0 .../litellm_core_utils/test_llm_judge.py | 0 .../test_llm_request_utils.py | 0 .../litellm_core_utils/test_logging_utils.py | 0 .../litellm_core_utils/test_logging_worker.py | 0 .../test_max_streaming_duration.py | 0 .../test_model_param_helper.py | 0 .../test_model_response_utils.py | 0 .../litellm_core_utils/test_private_json.py | 0 .../test_provider_affinity.py | 0 .../test_provider_specific_headers.py | 0 .../litellm_core_utils/test_ptu_pricing.py | 0 .../test_realtime_errors.py | 0 .../test_realtime_streaming.py | 0 .../test_redact_messages.py | 0 .../test_request_timeout_resolver.py | 0 .../test_retry_after_headers.py | 0 .../test_safe_divide_seconds.py | 0 .../test_safe_json_dumps.py | 0 .../test_sensitive_data_masker.py | 0 .../test_sentry_scrubbing.py | 0 .../test_served_output_texts.py | 0 .../test_streaming_chunk_builder_cursor.py | 0 ...streaming_chunk_builder_server_tool_use.py | 0 .../test_streaming_chunk_builder_utils.py | 0 .../test_streaming_handler.py | 2 +- .../test_streaming_overhead.py | 0 .../test_thread_pool_executor.py | 0 .../litellm_core_utils/test_token_counter.py | 1441 +++++++++++++++ .../test_token_counter_tool.py | 4 +- .../test_token_counter_tool_data.py | 0 .../unit/litellm_core_utils/test_tokenizer.py | 411 +++++ .../test_tool_search_spend_logging.py | 0 .../litellm_core_utils/test_url_utils.py | 0 .../test_xai_oauth_routing.py | 0 .../context_management/test_compact.py | 2 +- .../context_management/test_dispatcher.py | 2 +- .../llms/test_polling_url_origin_match.py | 2 +- .../__init__.py | 0 ...test_function_call_output_normalization.py | 0 .../test_handler.py | 0 .../test_image_generation_output.py | 0 .../test_litellm_completion_responses.py | 0 .../test_reasoning_input_item_preservation.py | 0 .../test_session_handler.py | 0 .../test_session_handler_with_cold_storage.py | 0 .../test_streaming_iterator_transformation.py | 0 ..._tool_output_order_preserved_for_gemini.py | 0 .../mcp/test_chat_completions_handler.py | 0 .../mcp/test_litellm_proxy_mcp_handler.py | 0 .../mcp/test_mcp_streaming_iterator.py | 0 .../responses/test_additional_tools.py | 0 .../responses/test_custom_tool_call.py | 0 .../responses/test_dispatch.py | 0 .../responses/test_metadata_codex_callback.py | 0 .../responses/test_no_duplicate_spend_logs.py | 29 - .../responses/test_null_test_fix.py | 0 .../test_responses_api_bridge_flag.py | 0 .../test_responses_api_request_body.py | 2 +- .../test_responses_prompt_management.py | 0 .../test_responses_router_cooldown.py | 0 .../test_responses_streaming_iterator.py | 0 ...sponses_supported_endpoints_passthrough.py | 0 .../responses/test_responses_utils.py | 0 .../test_responses_websocket_all_providers.py | 91 - .../responses/test_rust_bridge_websocket.py | 0 .../responses/test_sse_output_recovery.py | 0 .../responses/test_streaming_iterator.py | 0 .../test_streaming_iterator_error_events.py | 0 .../responses/test_text_format_conversion.py | 0 .../adaptive_router}/__init__.py | 0 .../adaptive_router/fixtures}/__init__.py | 0 .../fixtures/clean_no_signals.json | 0 .../fixtures/clean_satisfaction.json | 0 .../fixtures/disengagement_giveup.json | 0 .../fixtures/exhaustion_429.json | 0 .../fixtures/exhaustion_context_overflow.json | 0 .../fixtures/failure_tool_error.json | 0 .../fixtures/loop_same_tool.json | 0 .../fixtures/misalignment_rephrase.json | 0 .../mixed_failure_then_satisfaction.json | 0 .../fixtures/stagnation_repeat.json | 0 .../adaptive_router/test_adaptive_router.py | 0 .../adaptive_router/test_async_pre_routing.py | 0 .../adaptive_router/test_bandit.py | 0 .../adaptive_router/test_classifier.py | 0 .../adaptive_router/test_config.py | 0 .../test_e2e_adaptive_router.py | 0 .../adaptive_router/test_hooks.py | 0 .../adaptive_router/test_router_dispatch.py | 0 .../adaptive_router/test_signals.py | 0 .../adaptive_router/test_state_endpoint.py | 0 .../adaptive_router/test_update_queue.py | 0 .../test_context_compaction.py | 0 .../router_strategy/test_auto_router.py | 0 .../test_base_routing_strategy.py | 0 .../router_strategy/test_budget_limiter.py | 0 .../test_budget_limiter_hotpath.py | 0 .../router_strategy/test_complexity_router.py | 0 .../test_complexity_tier_predictor.py | 0 .../router_strategy/test_fuse_presets.py | 0 .../router_strategy/test_lar1_routing.py | 0 .../router_strategy/test_least_busy.py | 0 .../router_strategy/test_litellm_encoder.py | 0 .../router_strategy/test_llm_v2.py | 0 .../router_strategy/test_lowest_cost.py | 0 .../router_strategy/test_lowest_latency.py | 0 .../router_strategy/test_lowest_tpm_rpm.py | 0 .../router_strategy/test_quality_router.py | 0 .../test_router_routing_groups.py | 0 .../test_router_routing_plugins.py | 0 .../test_router_tag_regex_routing.py | 0 .../test_router_tag_routing.py | 0 .../router_strategy/test_savings_baseline.py | 0 .../router_strategy/test_simple_shuffle.py | 0 .../router_strategy/test_stall_detector.py | 0 .../test_prompt_caching_deployment_check.py | 4 +- .../router_utils/test_access_windows.py | 0 .../test_add_retry_fallback_headers.py | 0 .../test_auto_router_model_naming.py | 0 .../test_auto_router_tuning_baseline.py | 0 .../test_client_initalization_utils.py | 0 .../router_utils/test_cooldown_cache.py | 0 .../router_utils/test_cooldown_handlers.py | 0 .../test_fallback_event_handlers.py | 0 .../test_get_retry_from_policy.py | 0 ..._health_check_allowed_fails_integration.py | 0 .../router_utils/test_health_state_cache.py | 0 .../test_pattern_match_deployments.py | 0 .../test_reasoning_effort_capability.py | 0 .../test_router_health_check_routing.py | 0 .../test_router_interactions_endpoints.py | 0 .../test_router_utils_common_utils.py | 0 .../rust_bridge/AGENTS.md | 0 .../rust_bridge/messages/test_route_host.py | 122 ++ .../rust_bridge/messages/test_secrets.py | 0 .../rust_bridge/native_route_wheel_test.py | 0 .../rust_bridge/ocr/test_secrets.py | 0 .../rust_bridge/stubtest.ini | 0 .../rust_bridge/test_bindings.py | 0 .../test_callbacks_legacy_python.py | 0 .../rust_bridge/test_catalog.py | 0 .../rust_bridge/test_configuration.py | 0 .../rust_bridge/test_dispatch.py | 0 .../rust_bridge/test_failures.py | 0 .../rust_bridge/test_fork_guard.py | 0 .../rust_bridge/test_lifecycle.py | 0 .../rust_bridge/test_logger.py | 0 .../rust_bridge/test_runtime.py | 0 .../rust_bridge/test_secret_manager.py | 0 .../rust_bridge/test_settings.py | 0 .../rust_bridge/test_token_counter.py | 0 .../rust_bridge/test_tokenizer.py | 2 +- .../test_verify_linux_native_wheel.py | 0 261 files changed, 2951 insertions(+), 3204 deletions(-) delete mode 100644 tests/test_litellm/caching/test_caching_handler.py delete mode 100644 tests/test_litellm/rust_bridge/messages/test_route_host.py delete mode 100644 tests/test_litellm/rust_bridge/responses/__init__.py rename tests/{test_litellm => unit}/caching/test_azure_blob_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_caching.py (100%) rename tests/{test_litellm => unit}/caching/test_check_and_fix_namespace_none_guard.py (100%) rename tests/{test_litellm => unit}/caching/test_disk_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_dual_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_embedding_router.py (100%) rename tests/{test_litellm => unit}/caching/test_evicted_client_closer.py (100%) rename tests/{test_litellm => unit}/caching/test_gcs_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_in_memory_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_llm_caching_handler.py (100%) rename tests/{test_litellm => unit}/caching/test_llm_client_cache_e2e.py (100%) rename tests/{test_litellm => unit}/caching/test_qdrant_semantic_cache.py (99%) rename tests/{test_litellm => unit}/caching/test_redis_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_redis_cluster_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_redis_cluster_node_isolation.py (100%) rename tests/{test_litellm => unit}/caching/test_redis_connection_pool.py (100%) rename tests/{test_litellm => unit}/caching/test_redis_semantic_cache.py (99%) rename tests/{test_litellm => unit}/caching/test_s3_cache.py (100%) rename tests/{test_litellm => unit}/caching/test_valkey_semantic_cache.py (100%) rename tests/{test_litellm/litellm_core_utils/audio_utils => unit/expected_responses_api_request}/__init__.py (100%) rename tests/{test_litellm => unit}/expected_responses_api_request/azure_shell_tool.json (100%) rename tests/{test_litellm => unit}/expected_responses_api_request/context_management_and_shell.json (100%) create mode 100644 tests/unit/litellm_core_utils/conftest.py rename tests/{test_litellm => unit}/litellm_core_utils/event_loop_lag.py (100%) create mode 100644 tests/unit/litellm_core_utils/fake_secret_vault.py rename tests/{test_litellm/litellm_core_utils/llm_response_utils => unit/litellm_core_utils/llm_cost_calc}/__init__.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_get_api_base.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/messages_with_counts.py (100%) rename tests/{test_litellm/router_strategy/adaptive_router => unit/litellm_core_utils/prompt_templates}/__init__.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py (100%) rename tests/{test_litellm/rust_bridge => unit/litellm_core_utils/specialty_caches}/__init__.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_agentic_followup_kwargs.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_anthropic_dedup_factory.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_api_route_to_call_types.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_audio_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_aws_partition.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_bedrock_converse_dedup_factory.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_bug_report.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_chat_completion_agentic_loop.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_classifier_logging.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_cli_token_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_cloud_storage_security.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_codestral_provider_routing.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_core_helpers.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_coroutine_checker.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_dd_tracing.py (85%) rename tests/{test_litellm => unit}/litellm_core_utils/test_decode_special_tokens.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_dot_notation_indexing.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_duration_parser.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_error_normalization.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_exception_mapping_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_extract_base64_image.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_fallback_generalizations.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_fallback_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_get_litellm_params.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_get_llm_provider_endpoint_match.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_get_llm_provider_logic.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_get_model_cost_map.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_get_supported_openai_params.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_health_check_helpers.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_image_handling.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_initialize_dynamic_callback_params.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_internal_call_metadata.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_json_fragment_accumulator.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_json_schema_validation.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_litellm_logging.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_llm_judge.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_llm_request_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_logging_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_logging_worker.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_max_streaming_duration.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_model_param_helper.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_model_response_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_private_json.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_provider_affinity.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_provider_specific_headers.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_ptu_pricing.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_realtime_errors.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_realtime_streaming.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_redact_messages.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_request_timeout_resolver.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_retry_after_headers.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_safe_divide_seconds.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_safe_json_dumps.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_sensitive_data_masker.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_sentry_scrubbing.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_served_output_texts.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_cursor.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_handler.py (99%) rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_overhead.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_thread_pool_executor.py (100%) create mode 100644 tests/unit/litellm_core_utils/test_token_counter.py rename tests/{test_litellm => unit}/litellm_core_utils/test_token_counter_tool.py (93%) rename tests/{test_litellm => unit}/litellm_core_utils/test_token_counter_tool_data.py (100%) create mode 100644 tests/unit/litellm_core_utils/test_tokenizer.py rename tests/{test_litellm => unit}/litellm_core_utils/test_tool_search_spend_logging.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_url_utils.py (100%) rename tests/{test_litellm => unit}/litellm_core_utils/test_xai_oauth_routing.py (100%) rename tests/{test_litellm/rust_bridge/chat_completions => unit/responses/litellm_completion_transformation}/__init__.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_function_call_output_normalization.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_handler.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_image_generation_output.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_litellm_completion_responses.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_session_handler.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py (100%) rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py (100%) rename tests/{test_litellm => unit}/responses/mcp/test_chat_completions_handler.py (100%) rename tests/{test_litellm => unit}/responses/mcp/test_litellm_proxy_mcp_handler.py (100%) rename tests/{test_litellm => unit}/responses/mcp/test_mcp_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/responses/test_additional_tools.py (100%) rename tests/{test_litellm => unit}/responses/test_custom_tool_call.py (100%) rename tests/{test_litellm => unit}/responses/test_dispatch.py (100%) rename tests/{test_litellm => unit}/responses/test_metadata_codex_callback.py (100%) rename tests/{test_litellm => unit}/responses/test_no_duplicate_spend_logs.py (76%) rename tests/{test_litellm => unit}/responses/test_null_test_fix.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_api_bridge_flag.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_api_request_body.py (99%) rename tests/{test_litellm => unit}/responses/test_responses_prompt_management.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_router_cooldown.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_supported_endpoints_passthrough.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_utils.py (100%) rename tests/{test_litellm => unit}/responses/test_responses_websocket_all_providers.py (97%) rename tests/{test_litellm => unit}/responses/test_rust_bridge_websocket.py (100%) rename tests/{test_litellm => unit}/responses/test_sse_output_recovery.py (100%) rename tests/{test_litellm => unit}/responses/test_streaming_iterator.py (100%) rename tests/{test_litellm => unit}/responses/test_streaming_iterator_error_events.py (100%) rename tests/{test_litellm => unit}/responses/test_text_format_conversion.py (100%) rename tests/{test_litellm/rust_bridge/messages => unit/router_strategy/adaptive_router}/__init__.py (100%) rename tests/{test_litellm/rust_bridge/ocr => unit/router_strategy/adaptive_router/fixtures}/__init__.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/clean_no_signals.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/clean_satisfaction.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/disengagement_giveup.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/exhaustion_429.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/failure_tool_error.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/loop_same_tool.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/stagnation_repeat.json (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_adaptive_router.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_async_pre_routing.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_bandit.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_classifier.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_config.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_e2e_adaptive_router.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_hooks.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_router_dispatch.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_signals.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_state_endpoint.py (100%) rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_update_queue.py (100%) rename tests/{test_litellm => unit}/router_strategy/complexity_router/test_context_compaction.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_auto_router.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_base_routing_strategy.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_budget_limiter.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_budget_limiter_hotpath.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_complexity_router.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_complexity_tier_predictor.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_fuse_presets.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_lar1_routing.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_least_busy.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_litellm_encoder.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_llm_v2.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_lowest_cost.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_lowest_latency.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_lowest_tpm_rpm.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_quality_router.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_router_routing_groups.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_router_routing_plugins.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_router_tag_regex_routing.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_router_tag_routing.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_savings_baseline.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_simple_shuffle.py (100%) rename tests/{test_litellm => unit}/router_strategy/test_stall_detector.py (100%) rename tests/{test_litellm => unit}/router_utils/test_access_windows.py (100%) rename tests/{test_litellm => unit}/router_utils/test_add_retry_fallback_headers.py (100%) rename tests/{test_litellm => unit}/router_utils/test_auto_router_model_naming.py (100%) rename tests/{test_litellm => unit}/router_utils/test_auto_router_tuning_baseline.py (100%) rename tests/{test_litellm => unit}/router_utils/test_client_initalization_utils.py (100%) rename tests/{test_litellm => unit}/router_utils/test_cooldown_cache.py (100%) rename tests/{test_litellm => unit}/router_utils/test_cooldown_handlers.py (100%) rename tests/{test_litellm => unit}/router_utils/test_fallback_event_handlers.py (100%) rename tests/{test_litellm => unit}/router_utils/test_get_retry_from_policy.py (100%) rename tests/{test_litellm => unit}/router_utils/test_health_check_allowed_fails_integration.py (100%) rename tests/{test_litellm => unit}/router_utils/test_health_state_cache.py (100%) rename tests/{test_litellm => unit}/router_utils/test_pattern_match_deployments.py (100%) rename tests/{test_litellm => unit}/router_utils/test_reasoning_effort_capability.py (100%) rename tests/{test_litellm => unit}/router_utils/test_router_health_check_routing.py (100%) rename tests/{test_litellm => unit}/router_utils/test_router_interactions_endpoints.py (100%) rename tests/{test_litellm => unit}/router_utils/test_router_utils_common_utils.py (100%) rename tests/{test_litellm => unit}/rust_bridge/AGENTS.md (100%) rename tests/{test_litellm => unit}/rust_bridge/messages/test_secrets.py (100%) rename tests/{test_litellm => unit}/rust_bridge/native_route_wheel_test.py (100%) rename tests/{test_litellm => unit}/rust_bridge/ocr/test_secrets.py (100%) rename tests/{test_litellm => unit}/rust_bridge/stubtest.ini (100%) rename tests/{test_litellm => unit}/rust_bridge/test_bindings.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_callbacks_legacy_python.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_catalog.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_configuration.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_dispatch.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_failures.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_fork_guard.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_lifecycle.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_logger.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_runtime.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_secret_manager.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_settings.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_token_counter.py (100%) rename tests/{test_litellm => unit}/rust_bridge/test_tokenizer.py (95%) rename tests/{test_litellm => unit}/rust_bridge/test_verify_linux_native_wheel.py (100%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index d56e29fb627..3f4f5620176 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -5,6 +5,7 @@ flag="${1:?usage: unit_selection.sh }" legacy_flags=( caching-local + core-utils enterprise-package enterprise-routing integrations @@ -32,6 +33,7 @@ legacy_flags=( legacy_paths() { case "$1" in caching-local) echo tests/unit/caching ;; + core-utils) echo tests/unit/litellm_core_utils ;; enterprise-package) echo tests/unit/enterprise/integrations echo tests/unit/enterprise/proxy/auth @@ -42,6 +44,8 @@ legacy_paths() { echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; enterprise-routing) echo tests/unit/google_genai + echo tests/unit/router_strategy + echo tests/unit/router_utils echo tests/unit/enterprise/enterprise_callbacks/send_emails echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py @@ -77,6 +81,7 @@ legacy_paths() { echo tests/unit/messages echo tests/unit/rag echo tests/unit/rerank_api + echo tests/unit/rust_bridge echo tests/unit/secret_managers echo tests/unit/vector_stores echo tests/unit/videos ;; @@ -142,7 +147,9 @@ legacy_paths() { proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; proxy-extras) echo tests/unit/litellm_proxy_extras ;; proxy-infra) echo tests/unit/gateway ;; - responses-caching-types) echo tests/unit/types ;; + responses-caching-types) + find tests/unit/responses -name 'test_*.py' -not -path 'tests/unit/responses/mcp/*' + echo tests/unit/types ;; *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; esac } diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 41e9f11cefa..a9cd21bad5e 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -369,6 +369,13 @@ workflows: reruns: 2 base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-core-utils + flag: core-utils + shards: 2 + reruns: 1 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-integrations flag: integrations diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json index a563424c230..727733fa954 100644 --- a/.github/merge-smoke-tests.json +++ b/.github/merge-smoke-tests.json @@ -7,9 +7,9 @@ "MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]", "COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", "COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero", - "LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on", - "LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off", - "CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger", - "CALLBACK-FAILURE": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger" + "LOG-CONTENT-ON": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on", + "LOG-CONTENT-OFF": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off", + "CALLBACK-SUCCESS": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger", + "CALLBACK-FAILURE": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger" } } diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml index 2f5ce4d441a..0423b014ec5 100644 --- a/.github/workflows/test-redis-compat.yml +++ b/.github/workflows/test-redis-compat.yml @@ -12,9 +12,9 @@ on: - "litellm/caching/evicted_client_closer.py" - "tests/unit/test_redis.py" - "tests/local_testing/test_caching.py" - - "tests/test_litellm/caching/test_redis_connection_pool.py" - - "tests/test_litellm/caching/test_redis_cluster_cache.py" - - "tests/test_litellm/caching/test_evicted_client_closer.py" + - "tests/unit/caching/test_redis_connection_pool.py" + - "tests/unit/caching/test_redis_cluster_cache.py" + - "tests/unit/caching/test_evicted_client_closer.py" - ".github/workflows/test-redis-compat.yml" - "pyproject.toml" - "uv.lock" @@ -85,9 +85,9 @@ jobs: redis-server --version uv run --no-sync pytest \ tests/unit/test_redis.py \ - tests/test_litellm/caching/test_redis_connection_pool.py \ - tests/test_litellm/caching/test_redis_cluster_cache.py \ - tests/test_litellm/caching/test_evicted_client_closer.py \ + tests/unit/caching/test_redis_connection_pool.py \ + tests/unit/caching/test_redis_cluster_cache.py \ + tests/unit/caching/test_evicted_client_closer.py \ tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \ tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \ --tb=short -vv \ diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 1f3b5c4d97c..808bb2afd08 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -24,7 +24,7 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" + - "tests/unit/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -52,7 +52,7 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" + - "tests/unit/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" permissions: @@ -171,7 +171,7 @@ jobs: env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - run: python tests/unit/rust_bridge/native_route_wheel_test.py dist/*.whl - name: Run pytest tests/test_litellm_rust with the compiled extension run: make test-rust-extension diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4dca8075440..d75213d37ea 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -62,6 +62,7 @@ jobs: - shard: core-utils artifact-name: core-utils test-path: "tests/test_litellm/litellm_core_utils" + unit-flag: core-utils workers: 2 reruns: 1 timeout-minutes: 20 @@ -69,9 +70,7 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing - test-path: >- - tests/test_litellm/router_utils - tests/test_litellm/router_strategy + test-path: "" unit-flag: enterprise-routing workers: 2 reruns: 2 @@ -111,7 +110,6 @@ jobs: tests/test_litellm/interactions tests/test_litellm/ocr tests/test_litellm/passthrough - tests/test_litellm/rust_bridge tests/test_litellm/test_*.py unit-flag: misc workers: 2 @@ -228,9 +226,7 @@ jobs: - shard: responses-caching-types artifact-name: responses-caching-types - test-path: >- - tests/test_litellm/responses - tests/test_litellm/caching + test-path: "" unit-flag: responses-caching-types workers: 2 reruns: 2 diff --git a/Makefile b/Makefile index f27525b58ff..311a7daef92 100644 --- a/Makefile +++ b/Makefile @@ -301,7 +301,7 @@ test-rust-extension: UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \ $(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \ "$$temporary/venv/bin/python" -I -m mypy.stubtest \ - --mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \ + --mypy-config-file tests/unit/rust_bridge/stubtest.ini \ litellm.rust_bridge._native && \ LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust @@ -329,10 +329,10 @@ test-unit-integrations: install-test-deps $(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20 test-unit-core-utils: install-test-deps - $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 + $(UV_RUN) pytest tests/unit/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 030bf03d4ba..69f72fbc177 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -33,7 +33,7 @@ mod test_support { use crate::{LegacyLogging, LegacySurface, PublicCall}; /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. - /// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python + /// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python /// signatures, and [`namespace`] binds every fake call against it. pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md index c8b01fe9b3a..10619613516 100644 --- a/litellm-rust/crates/secrets/README.md +++ b/litellm-rust/crates/secrets/README.md @@ -30,7 +30,7 @@ The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV Native backends consistently distinguish absence from failure instead of swallowing provider errors. Python-compatible resolution maps these results back to the Python handler contract before applying fallback -`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/test_litellm/rust_bridge/ocr/test_secrets.py` pins this behavior +`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/unit/rust_bridge/ocr/test_secrets.py` pins this behavior Google rejects malformed base64 and mismatched CRC32C values instead of accepting corrupted payloads. Python currently ignores the checksum and uses permissive base64 decoding. Rust follows [RFC 4648](https://www.rfc-editor.org/rfc/rfc4648#section-3.3) and [Google's integrity guidance](https://docs.cloud.google.com/secret-manager/docs/data-integrity); `failed_or_missing_reads_are_not_cached` covers rejection and recovery diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 6c87ef4a3de..43e16599bf6 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -77,13 +77,13 @@ class _CallerHeadersView(TypedDict): headers: ReadOnly[dict[str, str]] -# Globally-routable IPs that are cloud-internal. Everything else -# non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented by -# Python's ``ipaddress`` module). This list only holds IPs that are -# publicly routable *and* point to cloud-fabric services reachable from -# inside a VM via special in-fabric routing. +# Cloud-internal IPs that ``ip.is_global`` can report as public. Everything +# else non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented +# by Python's ``ipaddress`` module). Older Python patch releases (3.12.2, for +# one) treat most of 192.0.0.0/24 as global, so it is listed to block it everywhere. _CLOUD_METADATA_EXCEPTIONS: Final = [ ip_network("168.63.129.16/32"), # Azure Wire Server + ip_network("192.0.0.0/24"), ] _ALLOWED_SCHEMES: Final = ("http", "https") diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 74c0478b08b..fbcf97839b9 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -742,7 +742,7 @@ class BaseResponsesAPITest(ABC): Passes tools=[{"type": "shell", "environment": {"type": "container_auto"}}]; validates that the request is accepted and returns a valid response. Only runs for OpenAI; offline coverage for the Azure route lives in - tests/test_litellm/responses/test_responses_api_request_body.py. + tests/unit/responses/test_responses_api_request_body.py. """ base_completion_call_args = self.get_base_completion_call_args() model = ( diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 0c3eca52dde..1a34e404d7f 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -1800,7 +1800,7 @@ def test_gemini_image_size_limit_exceeded(monkeypatch): that could cause memory issues and pod crashes. The image fetch is mocked (mirroring the LargeImageClient pattern in - tests/test_litellm/litellm_core_utils/test_image_handling.py) so the test + tests/unit/litellm_core_utils/test_image_handling.py) so the test deterministically exercises the size-limit rejection path without any external network dependency. """ diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py deleted file mode 100644 index e5a7f1540ca..00000000000 --- a/tests/test_litellm/caching/test_caching_handler.py +++ /dev/null @@ -1,867 +0,0 @@ -import asyncio -import json -import time -from unittest.mock import MagicMock, patch - -import httpx -import pytest -import respx -from fastapi.testclient import TestClient - -from datetime import datetime -from unittest.mock import AsyncMock - -from litellm.caching.caching_handler import _PENDING_CACHE_WRITES, LLMCachingHandler - - -@pytest.mark.asyncio -async def test_process_async_embedding_cached_response(): - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - args = { - "cached_result": [ - { - "embedding": [-0.025122925639152527, -0.019487135112285614], - "index": 0, - "object": "embedding", - } - ] - } - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=args["cached_result"], - kwargs={"model": "text-embedding-ada-002", "input": "test"}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="text-embedding-ada-002", - ) - - assert cache_hit - - print(f"response: {response}") - assert len(response.data) == 1 - - -@pytest.mark.asyncio -async def test_embedding_cache_preserves_prompt_tokens_details(): - """Test that prompt_tokens_details (including image_count) survives a full cache hit.""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "amazon.titan-embed-image-v1", - "prompt_tokens_details": {"image_count": 1}, - } - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="amazon.titan-embed-image-v1", - ) - - assert cache_hit - assert response.usage is not None - assert response.usage.prompt_tokens_details is not None - assert response.usage.prompt_tokens_details.image_count == 1 - - -@pytest.mark.asyncio -async def test_embedding_cache_backward_compat_no_prompt_tokens_details(): - """Test that old cached items without prompt_tokens_details still work.""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - # Old-format cached item — no prompt_tokens_details field - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "text-embedding-ada-002", - } - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "text-embedding-ada-002", "input": "test"}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="text-embedding-ada-002", - ) - - assert cache_hit - assert response.usage is not None - assert response.usage.prompt_tokens_details is None - - -@pytest.mark.asyncio -async def test_embedding_cache_aggregates_multiple_image_counts(): - """Test that image_count is summed correctly across multiple cached items.""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "amazon.titan-embed-image-v1", - "prompt_tokens_details": {"image_count": 1}, - }, - { - "embedding": [0.031, 0.042], - "index": 1, - "object": "embedding", - "model": "amazon.titan-embed-image-v1", - "prompt_tokens_details": {"image_count": 1}, - }, - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={ - "model": "amazon.titan-embed-image-v1", - "input": ["img1", "img2"], - }, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="amazon.titan-embed-image-v1", - ) - - assert cache_hit - assert response.usage.prompt_tokens_details is not None - assert response.usage.prompt_tokens_details.image_count == 2 - - -def test_combine_usage_merges_prompt_tokens_details(): - """Test that combine_usage merges prompt_tokens_details from both Usage objects.""" - from litellm.types.utils import PromptTokensDetailsWrapper, Usage - - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - usage1 = Usage( - prompt_tokens=10, - completion_tokens=0, - total_tokens=10, - prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), - ) - usage2 = Usage( - prompt_tokens=20, - completion_tokens=0, - total_tokens=20, - prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2), - ) - - combined = llm_caching_handler.combine_usage(usage1, usage2) - - assert combined.prompt_tokens == 30 - assert combined.total_tokens == 30 - assert combined.prompt_tokens_details is not None - assert combined.prompt_tokens_details.image_count == 3 - - -def test_combine_usage_handles_none_details(): - """Test that combine_usage works when one or both sides have null prompt_tokens_details.""" - from litellm.types.utils import PromptTokensDetailsWrapper, Usage - - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - # Both null - usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) - usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20) - combined = llm_caching_handler.combine_usage(usage_a, usage_b) - assert combined.prompt_tokens_details is None - - # Only first has details - usage_c = Usage( - prompt_tokens=10, - completion_tokens=0, - total_tokens=10, - prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), - ) - combined = llm_caching_handler.combine_usage(usage_c, usage_b) - assert combined.prompt_tokens_details is not None - assert combined.prompt_tokens_details.image_count == 1 - - # Only second has details - combined = llm_caching_handler.combine_usage(usage_a, usage_c) - assert combined.prompt_tokens_details is not None - assert combined.prompt_tokens_details.image_count == 1 - - -def test_is_chat_completion_cached_dict(): - from litellm.caching.caching_handler import _is_chat_completion_cached_dict - - assert _is_chat_completion_cached_dict( - {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []} - ) - assert _is_chat_completion_cached_dict( - {"id": "other", "object": "chat.completion.chunk", "choices": []} - ) - assert _is_chat_completion_cached_dict( - {"id": "no-object", "choices": [{"index": 0}]} - ) - assert not _is_chat_completion_cached_dict( - {"id": "resp_abc", "object": "response", "output": []} - ) - - -def _build_logging_obj(call_type: str, stream: bool): - import uuid as _uuid - - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - return LiteLLMLogging( - litellm_call_id=str(datetime.now()), - call_type=call_type, - model="gpt-5.4", - messages=[], - function_id=str(_uuid.uuid4()), - stream=stream, - start_time=datetime.now(), - ) - - -def test_convert_cached_aresponses_bridge_chat_completion_stream(): - """openai/responses chat-completions bridge: streaming cache hit replays as chat stream.""" - from litellm import aresponses - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.types.utils import CallTypes - - caching_handler = LLMCachingHandler( - original_function=aresponses, request_kwargs={}, start_time=datetime.now() - ) - cached_result = { - "id": "chatcmpl-bridge-cache-test", - "object": "chat.completion", - "created": int(time.time()), - "model": "gpt-5.4", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hi!"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, - } - - result = caching_handler._convert_cached_result_to_model_response( - cached_result=cached_result, - call_type=CallTypes.aresponses.value, - kwargs={ - "model": "gpt-5.4", - "stream": True, - "messages": [{"role": "user", "content": "hi"}], - }, - logging_obj=_build_logging_obj(CallTypes.aresponses.value, stream=True), - model="gpt-5.4", - args=(), - ) - - assert isinstance(result, CustomStreamWrapper) - - -def test_convert_cached_responses_bridge_chat_completion_nonstream(): - """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse.""" - from litellm import responses - from litellm.types.utils import CallTypes, ModelResponse - - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) - cached_result = { - "id": "chatcmpl-bridge-nonstream", - "object": "chat.completion", - "created": int(time.time()), - "model": "gpt-5.4", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hi!"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, - } - - result = caching_handler._convert_cached_result_to_model_response( - cached_result=cached_result, - call_type=CallTypes.responses.value, - kwargs={ - "model": "gpt-5.4", - "stream": False, - "messages": [{"role": "user", "content": "hi"}], - }, - logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), - model="gpt-5.4", - args=(), - ) - - assert isinstance(result, ModelResponse) - assert result.choices[0].message.content == "Hi!" - - -def test_convert_cached_responses_legacy_nonstream_path(): - """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path.""" - from litellm import responses - from litellm.types.llms.openai import ResponsesAPIResponse - from litellm.types.utils import CallTypes - - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) - cached_result = { - "id": "resp_legacy_nonstream", - "created_at": int(time.time()), - "status": "completed", - "model": "gpt-4o", - "object": "response", - "output": [ - { - "type": "message", - "id": "msg_legacy", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "legacy response", - "annotations": [], - } - ], - } - ], - } - - result = caching_handler._convert_cached_result_to_model_response( - cached_result=cached_result, - call_type=CallTypes.responses.value, - kwargs={"model": "gpt-4o", "input": "hi", "stream": False}, - logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), - model="gpt-4o", - args=(), - ) - - assert isinstance(result, ResponsesAPIResponse) - assert result.id == "resp_legacy_nonstream" - - -def test_convert_cached_responses_legacy_stream_path(): - """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path.""" - from litellm import responses - from litellm.responses.streaming_iterator import ( - CachedResponsesAPIStreamingIterator, - ) - from litellm.types.utils import CallTypes - - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) - cached_result = { - "id": "resp_legacy_stream", - "created_at": int(time.time()), - "status": "completed", - "model": "gpt-4o", - "object": "response", - "output": [ - { - "type": "message", - "id": "msg_legacy_stream", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "legacy stream", - "annotations": [], - } - ], - } - ], - } - - result = caching_handler._convert_cached_result_to_model_response( - cached_result=cached_result, - call_type=CallTypes.responses.value, - kwargs={"model": "gpt-4o", "input": "hi", "stream": True}, - logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True), - model="gpt-4o", - args=(), - ) - - assert isinstance(result, CachedResponsesAPIStreamingIterator) - - -@pytest.mark.asyncio -async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input(): - """Image-embedding cache hit restores prompt_tokens=0 from the stored value - instead of recomputing a bogus count by tokenizing the base64 input.""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - # base64-like blob — token_counter over this would return a large nonzero count - image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50 - - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "amazon.titan-embed-image-v1", - "prompt_tokens": 0, - "prompt_tokens_details": {"image_count": 1}, - } - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="amazon.titan-embed-image-v1", - ) - - assert cache_hit - assert response.usage is not None - assert response.usage.prompt_tokens == 0 - assert response.usage.total_tokens == 0 - assert response.usage.prompt_tokens_details.image_count == 1 - - -@pytest.mark.asyncio -async def test_embedding_cache_sums_stored_prompt_tokens_across_items(): - """A multi-item cache hit sums the stored per-item prompt_tokens back to the total.""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - cached_result = [ - { - "embedding": [-0.01], - "index": 0, - "object": "embedding", - "model": "text-embedding-3-small", - "prompt_tokens": 5, - }, - { - "embedding": [-0.02], - "index": 1, - "object": "embedding", - "model": "text-embedding-3-small", - "prompt_tokens": 4, - }, - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="text-embedding-3-small", - ) - - assert cache_hit - assert response.usage.prompt_tokens == 9 - assert response.usage.total_tokens == 9 - - -@pytest.mark.asyncio -async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): - """Legacy cache entries with no stored prompt_tokens still recompute via token_counter - for str inputs (backward compatibility).""" - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - # No prompt_tokens key — pre-fix entry - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "text-embedding-ada-002", - }, - ] - - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "text-embedding-ada-002", "input": "hello world"}, - logging_obj=mock_logging_obj, - start_time=datetime.now(), - model="text-embedding-ada-002", - ) - - assert cache_hit - # token_counter over "hello world" yields a nonzero count — fallback path still runs - assert response.usage.prompt_tokens > 0 - - -@pytest.mark.asyncio -async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj(): - """A full embedding cache hit must stamp the resolved provider onto the logging - obj so spend logs record the provider instead of None/unknown.""" - from litellm.types.utils import CallTypes - - llm_caching_handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs={}, - start_time=datetime.now(), - ) - - cached_result = [ - { - "embedding": [-0.025, -0.019], - "index": 0, - "object": "embedding", - "model": "text-embedding-3-small", - "prompt_tokens": 5, - } - ] - - logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False) - logging_obj.async_success_handler = AsyncMock() - - response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( - final_embedding_cached_response=None, - cached_result=cached_result, - kwargs={"model": "text-embedding-3-small", "input": "hello world"}, - logging_obj=logging_obj, - start_time=datetime.now(), - model="text-embedding-3-small", - ) - - assert cache_hit - assert logging_obj.model_call_details["custom_llm_provider"] == "openai" - - -def test_sync_stream_responses_cache_hit_sets_custom_llm_provider_on_logging_obj(monkeypatch): - import litellm - from litellm.caching.caching import Cache - from litellm.types.utils import CallTypes - - monkeypatch.setattr(litellm, "cache", Cache(type="local")) - kwargs = {"model": "azure/gpt-5.4-mini", "input": "hello", "stream": True} - cached_response = { - "id": "resp_sync_stream", - "created_at": int(time.time()), - "status": "completed", - "model": "gpt-5.4-mini", - "object": "response", - "output": [ - { - "type": "message", - "id": "msg_sync_stream", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - } - litellm.cache.add_cache(json.dumps(cached_response), **kwargs) - handler = LLMCachingHandler(original_function=litellm.responses, request_kwargs=kwargs, start_time=datetime.now()) - logging_obj = _build_logging_obj(CallTypes.responses.value, stream=True) - - hit = handler._sync_get_cache( - model="azure/gpt-5.4-mini", - original_function=litellm.responses, - logging_obj=logging_obj, - start_time=datetime.now(), - call_type=CallTypes.responses.value, - kwargs=kwargs, - args=(), - ) - - assert hit.cached_result is not None - assert logging_obj.model_call_details["custom_llm_provider"] == "azure" - assert logging_obj.model_call_details["litellm_params"]["custom_llm_provider"] == "azure" - - -def test_request_kwargs_does_not_retain_logging_obj(): - """ - The caching handler lives on logging_obj._llm_caching_handler, so keeping - litellm_logging_obj inside request_kwargs closes a reference cycle - (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the - full request payload alive until a generational GC pass instead of being - freed by refcount when the request finishes; under bursts of large-token - requests this presents as stepwise RSS growth that never returns to - baseline. Other kwargs (messages included) must be preserved. - """ - logging_obj = MagicMock() - kwargs = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hello"}], - "litellm_logging_obj": logging_obj, - } - - handler = LLMCachingHandler( - original_function=MagicMock(), - request_kwargs=kwargs, - start_time=datetime.now(), - ) - - assert "litellm_logging_obj" not in handler.request_kwargs - assert handler.request_kwargs["messages"] == kwargs["messages"] - assert handler.request_kwargs["model"] == "gpt-4o" - - -def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch): - """ - Regression test for the SDK losing async cache writes in short-lived scripts: - async_set_cache dispatched the write as a bare fire-and-forget task, so - asyncio.run cancelled it at loop close before the write landed (LIT-6184, - deterministic with hiredis installed). The write must survive loop shutdown. - """ - import litellm - - writes = [] - - class _SlowWriteCache: - supported_call_types = ["acompletion"] - cache = None - - async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs): - await asyncio.sleep(0.2) - writes.append(result) - - async def acompletion(**kwargs): - return None - - handler = LLMCachingHandler( - original_function=acompletion, - request_kwargs={}, - start_time=datetime.now(), - ) - monkeypatch.setattr(litellm, "cache", _SlowWriteCache()) - - async def _short_lived_script(): - await handler.async_set_cache( - result=litellm.ModelResponse(), - original_function=acompletion, - kwargs={}, - ) - - asyncio.run(_short_lived_script()) - - assert len(writes) == 1 - - -@pytest.mark.asyncio -async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch): - """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again.""" - import litellm - from litellm.caching.caching import Cache - from litellm.types.utils import CallTypes - - async def acompletion(**kwargs): - return None - - monkeypatch.setattr(litellm, "cache", Cache(type="local")) - kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True} - await litellm.cache.async_add_cache( - litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs - ) - handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now()) - logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False) - logging_obj.async_success_handler = AsyncMock() - - hit = await handler._async_get_cache( - model="gpt-5.4", - original_function=acompletion, - logging_obj=logging_obj, - start_time=datetime.now(), - call_type=CallTypes.acompletion.value, - kwargs=kwargs, - args=(), - ) - - assert hit is not None and hit.cached_result is not None - assert handler.preset_cache_key is not None - assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key - assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key - - -@pytest.mark.asyncio -async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch): - import litellm - from litellm.caching.caching import Cache - from litellm.types.utils import CallTypes - - async def aanthropic_messages(**kwargs): - return None - - monkeypatch.setattr(litellm, "cache", Cache(type="local")) - kwargs = { - "model": "claude-sonnet-5", - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 16, - "caching": True, - "stream": False, - "_websearch_interception_converted_stream": True, - } - cached_message = { - "id": "msg_1", - "type": "message", - "role": "assistant", - "content": [{"type": "text", "text": "hi"}], - } - await litellm.cache.async_add_cache(cached_message, **kwargs) - handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now()) - logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False) - logging_obj.async_success_handler = AsyncMock() - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() - - hit = await handler._async_get_cache( - model="claude-sonnet-5", - original_function=aanthropic_messages, - logging_obj=logging_obj, - start_time=datetime.now(), - call_type=CallTypes.aanthropic_messages.value, - kwargs=kwargs, - args=(), - ) - - assert hit is not None and hit.cached_result == cached_message - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() - assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True - - -@pytest.mark.asyncio -async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch): - import litellm - from litellm.caching.caching import Cache - from litellm.types.utils import CallTypes - - async def acompletion(**kwargs): - return None - - monkeypatch.setattr(litellm, "cache", Cache(type="local")) - kwargs = { - "model": "gpt-5.6", - "messages": [{"role": "user", "content": "run the code"}], - "caching": True, - "stream": False, - "_code_interpreter_interception_converted_stream": True, - "_agentic_loop_depth": 1, - } - await litellm.cache.async_add_cache( - litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs - ) - handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now()) - logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False) - logging_obj.async_success_handler = AsyncMock() - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() - - hit = await handler._async_get_cache( - model="gpt-5.6", - original_function=acompletion, - logging_obj=logging_obj, - start_time=datetime.now(), - call_type=CallTypes.acompletion.value, - kwargs=kwargs, - args=(), - ) - - assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse) - assert hit.cached_result.choices[0].message.content == "done" - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() - assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True - - -@pytest.mark.asyncio -async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch): - import litellm - from litellm import CustomLLM - from litellm.caching.caching import Cache - from litellm.types.utils import Embedding, EmbeddingResponse - - class RecordingEmbedder(CustomLLM): - provider_inputs: tuple[tuple[str, ...], ...] = () - - async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse: - self.provider_inputs = (*self.provider_inputs, tuple(input)) - return EmbeddingResponse( - model=model, - data=[ - Embedding(embedding=[float(len(text))], index=idx, object="embedding") - for idx, text in enumerate(input) - ], - ) - - embedder = RecordingEmbedder() - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}]) - monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"]) - monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"]) - monkeypatch.setattr(litellm, "cache", Cache(type="local")) - - await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"]) - await asyncio.gather(*_PENDING_CACHE_WRITES) - mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"] - response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input) - await asyncio.gather(*_PENDING_CACHE_WRITES) - - assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs - assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4] - assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input] - assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit" - - repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input) - - assert len(embedder.provider_inputs) == 2, embedder.provider_inputs - assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input] diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index f8c7d5273d1..f83c1e76b3a 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -22,19 +22,6 @@ import litellm from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS -from litellm.litellm_core_utils.cli_keyring import ( - KeyringDiscardsWrites, - KeyringUnreachable, - KeyringUnusable, - SecretErase, - SecretErased, - SecretFound, - SecretMissing, - SecretRead, - SecretStored, - SecretStranded, - SecretWrite, -) from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, ) @@ -42,6 +29,7 @@ from litellm.llms.custom_httpx.async_client_cleanup import ( close_litellm_async_clients, ) from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module +from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault def _reset_module_level_aws_auth_caches(): @@ -128,60 +116,6 @@ def isolate_host_os_keychain(monkeypatch): monkeypatch.setenv("LITELLM_CLI_DISABLE_KEYRING", "1") -class FakeSecretVault: - """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised. - - `available=False` models a keychain that is locked or has no backend, `writable=False` one that - refuses to store, `erasable=False` one that will not release what it already holds, and `failure` - picks which unusable state those report. `discards=True` is keyring's null backend, which answers - reads and erases like any other yet keeps nothing it is given, so only writes report it. - """ - - def __init__( - self, - blob: str | None = None, - *, - available: bool = True, - writable: bool = True, - erasable: bool = True, - discards: bool = False, - failure: KeyringUnusable = KeyringUnreachable(), - ) -> None: - self.blob: str | None = blob - self.available: bool = available - self.writable: bool = writable - self.erasable: bool = erasable - self.discards: bool = discards - self.failure: KeyringUnusable = failure - self.reads: int = 0 - self.writes: list[str] = [] - self.erases: int = 0 - - def read(self) -> SecretRead: - self.reads += 1 - if not self.available: - return self.failure - return SecretMissing() if self.blob is None else SecretFound(self.blob) - - def write(self, blob: str) -> SecretWrite: - self.writes.append(blob) - if not (self.available and self.writable): - return self.failure - if self.discards: - return KeyringDiscardsWrites() - self.blob = blob - return SecretStored() - - def erase(self) -> SecretErase: - self.erases += 1 - if not self.available: - return self.failure - if not self.erasable: - return SecretStranded() if self.blob is not None else SecretErased() - self.blob = None - return SecretErased() - - @pytest.fixture def secret_vault_factory(): """Build FakeSecretVault instances; see its docstring for the failure modes it can model.""" diff --git a/tests/test_litellm/litellm_core_utils/__init__.py b/tests/test_litellm/litellm_core_utils/__init__.py index 8c64613a5da..e69de29bb2d 100644 --- a/tests/test_litellm/litellm_core_utils/__init__.py +++ b/tests/test_litellm/litellm_core_utils/__init__.py @@ -1 +0,0 @@ -# This file makes the tests/litellm/litellm_core_utils directory a Python package diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index eccf44a1bda..1e10b7e82b1 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,442 +1,6 @@ -#### What this tests #### -# This tests litellm.token_counter.token_counter() function -import asyncio -import base64 -import importlib -import threading -import time -import traceback -from concurrent.futures import Future, wait -from typing import Final -from unittest.mock import MagicMock - -import anyio.to_thread import pytest -import tiktoken - -from unittest.mock import AsyncMock, patch - -import litellm -from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens -from litellm import token_counter as token_counter_old -import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS -from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.token_counter import ( - _get_exact_count_function, - _get_extrapolating_count_function, - _get_tiktoken_count_function, - calculate_img_tokens, - high_detail_image_token_upper_bound, - offload_token_count, -) -from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new -from tests.large_text import text -from tests.test_litellm.litellm_core_utils.event_loop_lag import ( - assert_loop_stayed_free, - timed_with_loop_lags, - warm_tokenizer, -) -from tests.test_litellm.litellm_core_utils.messages_with_counts import ( - MESSAGES_TEXT, - MESSAGES_WITH_IMAGES, - MESSAGES_WITH_TOOLS, -) - - -def token_counter_both_assert_same(**args): - new = token_counter_new(**args) - old = token_counter_old(**args) - assert new == old, f"New token counter {new} does not match old token counter {old}" - return new - - -## Choose which token_counter the test will use. - -# token_counter = token_counter_new -# token_counter = token_counter_old -token_counter = token_counter_both_assert_same - - -def test_token_counter_basic(): - assert ( - token_counter( - model="claude-2", - messages=[ - { - "role": "user", - "content": "This is a long message that definitely exceeds the token limit.", - } - ], - ) - == 19 - ) - - -def test_token_counter_large_repeated_text_is_fast(): - messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}] - - start_time = time.perf_counter() - tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages) - elapsed = time.perf_counter() - start_time - - assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s" - assert tokens > 0 - - -@pytest.mark.parametrize( - "text", - [ - "Short text", - "This is a normal message with punctuation, numbers, and a few words.", - ], -) -def test_token_counter_short_text_matches_tiktoken(text): - encoding = tiktoken.get_encoding("cl100k_base") - expected = len(encoding.encode(text, disallowed_special=())) - - assert token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) == expected - - -def test_token_counter_default_encoding_matches_cl100k(): - encoding: Final = tiktoken.get_encoding("cl100k_base") - expected: Final = len(encoding.encode("hello world", disallowed_special=())) - - assert token_counter_new(model=None, text="hello world") == expected - - -def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken(): - text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025] - encoding = tiktoken.get_encoding("cl100k_base") - expected = len(encoding.encode(text, disallowed_special=())) - - actual = token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) - - assert abs(actual - expected) <= 4 - - -@pytest.mark.parametrize( - "configured", - ["0", "-1", "-1024", "not-an-int", "", " ", "999999999", "inf", "1e9"], -) -def test_invalid_chunk_size_config_stays_usable(monkeypatch, configured): - """A misconfigured chunk size must not raise, count zero, or restore the quadratic encode cost.""" - monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", configured) - try: - reloaded = importlib.reload(litellm.constants) - chunk_size = reloaded.TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS - assert 1 <= chunk_size <= reloaded.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS - - encoding = tiktoken.get_encoding("cl100k_base") - count_tokens = _get_tiktoken_count_function( - lambda text: len(encoding.encode(text, disallowed_special=())), - chunk_size=chunk_size, - ) - assert count_tokens("The quick brown fox jumps over the lazy dog. " * 40) > 0 - finally: - monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS") - importlib.reload(litellm.constants) - - -def test_valid_chunk_size_config_is_honoured(monkeypatch): - monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", "2048") - try: - assert importlib.reload(litellm.constants).TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS == 2048 - finally: - monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS") - importlib.reload(litellm.constants) - - -async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): - warm_tokenizer("claude-fable-5") - - tokens, took, lags = await timed_with_loop_lags( - lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) - ) - - assert tokens > 0 - assert_loop_stayed_free(took, lags) - - -@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) -def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): - count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) - front_heavy: Final = "a" * 1_000 + "b" * 4_000 - exact: Final = 1_000 + len(front_heavy) - - estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) - - assert abs(estimate - exact) <= exact // 100 - assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars - - -def test_count_at_or_below_the_cap_is_exact(): - count_exactly: Final = MagicMock(side_effect=len) - - assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 - assert count_exactly.call_args_list == [(("a" * 5_000,),)] - - -class _SlowEncoder: - def __init__(self) -> None: - self._lock: Final = threading.Lock() - self.in_flight = 0 - self.peak_in_flight = 0 - - def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: - with self._lock: - self.in_flight += 1 - self.peak_in_flight = max(self.peak_in_flight, self.in_flight) - time.sleep(0.1) - with self._lock: - self.in_flight -= 1 - return [[0] * len(text) for text in texts] - - -@pytest.mark.asyncio -async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): - encoder: Final = _SlowEncoder() - count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) - shared_pool: Final = anyio.to_thread.current_default_thread_limiter() - burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS - - async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: - if counting.done(): - return () - await asyncio.sleep(0.01) - return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) - - counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) - borrowed: Final = await shared_pool_borrowed_until_done(counting) - - assert await counting == [3] * burst - assert len(borrowed) > 1 and max(borrowed) == 0 - assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS - - -def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: - def slow_count(counted: str) -> int: - time.sleep(0.1) - return len(counted) - - result.set_result(asyncio.run(offload_token_count(slow_count)(text))) - - -def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): - loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS - results: Final = tuple(Future[int]() for _ in range(loops)) - threads: Final = tuple( - threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) - for size, result in enumerate(results, start=1) - ) - for thread in threads: - thread.start() - - _, pending = wait(results, timeout=5) - - assert not pending - assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) - - -@pytest.mark.parametrize( - ("configured", "expected"), - [("8", 8), ("0", 4), ("not-an-int", 4)], -) -def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): - monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) - try: - assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected - finally: - monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") - importlib.reload(litellm.constants) - - -def test_token_counter_applies_the_default_cap(): - max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS - prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] - over_the_cap: Final = prose + "a" * 200_000 - exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) - - estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) - - assert estimate != exact - assert abs(estimate - exact) <= exact // 100 - - -@pytest.mark.parametrize( - ("configured", "expected"), - [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], -) -def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): - monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) - try: - assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected - finally: - monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") - importlib.reload(litellm.constants) - - -def test_token_counter_with_prefix(): - messages = [ - {"role": "user", "content": "Who won the world cup in 2022?"}, - {"role": "assistant", "content": "Argentina", "prefix": True}, - ] - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens == 22, f"Expected 22 tokens, got {tokens}" - - -def test_token_counter_normal_plus_function_calling(): - messages = [ - {"role": "system", "content": "System prompt"}, - {"role": "user", "content": "content1"}, - {"role": "assistant", "content": "content2"}, - {"role": "user", "content": "conten3"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_E0lOb1h6qtmflUyok4L06TgY", - "function": { - "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}', - "name": "SearchInternet", - }, - "type": "function", - } - ], - }, - { - "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY", - "role": "tool", - "name": "SearchInternet", - "content": "tool content", - }, - ] - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens == 80 - - -# test_token_counter_normal_plus_function_calling() - - -def test_token_counter_legacy_function_call_counts_arguments(): - """ - Regression for VERIA-492 (Token-counter function_call bypass). - - The legacy OpenAI assistant `function_call` field carries arbitrary text in - `arguments`. Before the fix, `_count_messages` had no branch for - `function_call` and fell through to the unsupported-key `continue`, so an - assistant turn could smuggle unlimited text past `token_counter` and the - proxy `/utils/token_counter` endpoint (and downstream pre-call budget / - `get_modified_max_tokens` math). After the fix it must be counted the - same as the equivalent `tool_calls` payload. - """ - long_arg = "A" * 4000 - fc_messages = [ - {"role": "user", "content": "hi"}, - { - "role": "assistant", - "content": None, - "function_call": {"name": "search", "arguments": long_arg}, - }, - ] - tc_messages = [ - {"role": "user", "content": "hi"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "search", "arguments": long_arg}, - } - ], - }, - ] - fc_tokens = token_counter(model="gpt-3.5-turbo", messages=fc_messages) - tc_tokens = token_counter(model="gpt-3.5-turbo", messages=tc_messages) - assert fc_tokens == tc_tokens, ( - f"function_call arguments must count like tool_calls arguments; " - f"got function_call={fc_tokens}, tool_calls={tc_tokens}" - ) - assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}" - - -@pytest.mark.parametrize( - "message_count_pair", - MESSAGES_TEXT, -) -def test_token_counter_textonly(message_count_pair): - counted_tokens = token_counter( - model="gpt-35-turbo", messages=[message_count_pair["message"]] - ) - assert counted_tokens == message_count_pair["count"] - - -@pytest.mark.parametrize( - "message_count_pair", - MESSAGES_TEXT, -) -def test_token_counter_count_response_tokens(message_count_pair): - counted_tokens = token_counter( - model="gpt-35-turbo", - messages=[message_count_pair["message"]], - count_response_tokens=True, - ) - # 3 tokens are not added because of count_response_tokens=True - expected = message_count_pair["count"] - 3 - assert counted_tokens == expected - - -@pytest.mark.parametrize( - "message_count_pair", - MESSAGES_WITH_IMAGES, -) -def test_token_counter_with_images(message_count_pair): - counted_tokens = token_counter( - model="gpt-4o", messages=[message_count_pair["message"]] - ) - assert counted_tokens == message_count_pair["count"] - - -@pytest.mark.parametrize( - "message_count_pair", - MESSAGES_WITH_TOOLS, -) -def test_token_counter_with_tools(message_count_pair): - counted_tokens = token_counter( - model="gpt-35-turbo", - messages=[message_count_pair["system_message"]], - tools=message_count_pair["tools"], - tool_choice=message_count_pair["tool_choice"], - ) - expected_tokens = message_count_pair["count"] - actual_diff = counted_tokens - expected_tokens - - if "count-tolerate" in message_count_pair: - if message_count_pair["count-tolerate"] == counted_tokens: - pass # expected - else: - tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens - assert ( - actual_diff <= tolerated_diff - ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction." - if actual_diff != tolerated_diff: - raise NeedsToleranceUpdateError( - f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}" - ) - - else: - assert ( - expected_tokens == counted_tokens - ), f"Expected {expected_tokens} tokens, got {counted_tokens}." - - -class NeedsToleranceUpdateError(Exception): - """Custom exception to mark tests that have improved""" - - pass +from litellm import create_pretrained_tokenizer +from tests.unit.litellm_core_utils.test_token_counter import token_counter def test_tokenizers(): @@ -449,32 +13,22 @@ def test_tokenizers(): openai_tokens = token_counter(model="gpt-3.5-turbo", text=sample_text) # claude tokenizer - claude_tokens = token_counter( - model="claude-3-5-haiku-20241022", text=sample_text - ) + claude_tokens = token_counter(model="claude-3-5-haiku-20241022", text=sample_text) # cohere tokenizer cohere_tokens = token_counter(model="command-nightly", text=sample_text) # llama2 tokenizer - llama2_tokens = token_counter( - model="meta-llama/Llama-2-7b-chat", text=sample_text - ) + llama2_tokens = token_counter(model="meta-llama/Llama-2-7b-chat", text=sample_text) # llama3 tokenizer (also testing custom tokenizer) - llama3_tokens_1 = token_counter( - model="meta-llama/llama-3-70b-instruct", text=sample_text - ) + llama3_tokens_1 = token_counter(model="meta-llama/llama-3-70b-instruct", text=sample_text) try: llama3_tokenizer = create_pretrained_tokenizer("Xenova/llama-3-tokenizer") except Exception as e: - pytest.skip( - f"custom tokenizer download failed (HF hub unreachable): {e}" - ) - llama3_tokens_2 = token_counter( - custom_tokenizer=llama3_tokenizer, text=sample_text - ) + pytest.skip(f"custom tokenizer download failed (HF hub unreachable): {e}") + llama3_tokens_2 = token_counter(custom_tokenizer=llama3_tokenizer, text=sample_text) print( f"openai tokens: {openai_tokens}; claude tokens: {claude_tokens}; cohere tokens: {cohere_tokens}; llama2 tokens: {llama2_tokens}; llama3 tokens: {llama3_tokens_1}" @@ -485,1117 +39,13 @@ def test_tokenizers(): # model hub is unreachable (e.g. in CI). In that case the count will # equal the openai count and the differentiation assertion is skipped. if openai_tokens == llama2_tokens: - pytest.skip( - "llama2 fell back to tiktoken (HF hub unreachable); skipping differentiation assertion" - ) + pytest.skip("llama2 fell back to tiktoken (HF hub unreachable); skipping differentiation assertion") assert llama2_tokens != llama3_tokens_1, "Token values are not different." - assert ( - llama3_tokens_1 == llama3_tokens_2 - ), "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same." + assert llama3_tokens_1 == llama3_tokens_2, ( + "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same." + ) print("test tokenizer: It worked!") except Exception as e: pytest.fail(f"An exception occured: {e}") - - -# test_tokenizers() - - -def test_encoding_and_decoding(): - try: - sample_text = "Hellö World, this is my input string!" - # openai encoding + decoding - openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) - openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) - - assert openai_text == sample_text - - # claude encoding + decoding - claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) - - claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) - - assert claude_text == sample_text - - # cohere encoding + decoding - cohere_tokens = encode(model="command-nightly", text=sample_text) - cohere_text = decode(model="command-nightly", tokens=cohere_tokens) - - assert cohere_text == sample_text - - # llama2 encoding + decoding - llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text) - llama2_text = decode(model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens) - - assert llama2_text == sample_text - except Exception as e: - pytest.fail(f"An exception occured: {e}\n{traceback.format_exc()}") - - -# test_encoding_and_decoding() - - -def test_gpt_vision_token_counting(): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What’s in this image?"}, - { - "type": "image_url", - "image_url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", - }, - ], - } - ] - tokens = token_counter(model="gpt-4-vision-preview", messages=messages) - print(f"tokens: {tokens}") - - -# test_gpt_vision_token_counting() - - -@pytest.mark.parametrize( - "model", - [ - "gpt-4-vision-preview", - "gpt-4o", - "claude-3-opus-20240229", - "command-nightly", - "mistral/mistral-tiny", - ], -) -def test_load_test_token_counter(model): - """ - Token count large prompt 100 times. - - Assert time taken is < 1.5s. - """ - import tiktoken - - messages = [{"role": "user", "content": text}] * 10 - - start_time = time.time() - for _ in range(10): - _ = token_counter(model=model, messages=messages) - # enc.encode("".join(m["content"] for m in messages)) - - end_time = time.time() - - total_time = end_time - start_time - print("model={}, total test time={}".format(model, total_time)) - assert total_time < 10, f"Total encoding time > 10s, {total_time}" - - -def test_openai_token_with_image_and_text(): - model = "gpt-4o" - full_request = { - "model": "gpt-4o", - "tools": [ - { - "type": "function", - "function": { - "name": "json", - "parameters": { - "type": "object", - "required": ["clause"], - "properties": {"clause": {"type": "string"}}, - }, - "description": "Respond with a JSON object.", - }, - } - ], - "logprobs": False, - "messages": [ - { - "role": "user", - "content": [ - { - "text": "\n Just some long text, long long text, and you know it will be longer than 7 tokens definetly.", - "type": "text", - } - ], - } - ], - "tool_choice": {"type": "function", "function": {"name": "json"}}, - "exclude_models": [], - "disable_fallback": False, - "exclude_providers": [], - } - messages = full_request.get("messages", []) - - token_count = token_counter(model=model, messages=messages) - print(token_count) - - -@pytest.mark.parametrize( - "model, base_model, input_tokens, user_max_tokens, expected_value", - [ - ("random-model", "random-model", 1024, 1024, 1024), - ("gpt-3.5-turbo", "gpt-3.5-turbo", 4000, 5000, 4096), # model max output = 4096 - ], -) -def test_get_modified_max_tokens( - model, base_model, input_tokens, user_max_tokens, expected_value -): - """ - - Test when max_output is not known => expect user_max_tokens - - Test when max_output == max_input, - - input > max_output, no max_tokens => expect None - - input + max_tokens > max_output => expect remainder - - input + max_tokens < max_output => expect max_tokens - - Test when max_tokens > max_output => expect max_output - """ - args = locals() - import litellm - - litellm.token_counter = MagicMock() - - def _mock_token_counter(*args, **kwargs): - return input_tokens - - litellm.token_counter.side_effect = _mock_token_counter - print(f"_mock_token_counter: {_mock_token_counter()}") - messages = [{"role": "user", "content": "Hello world!"}] - - calculated_value = get_modified_max_tokens( - model=model, - base_model=base_model, - messages=messages, - user_max_tokens=user_max_tokens, - buffer_perc=0, - buffer_num=0, - ) - - if expected_value is None: - assert calculated_value is None - else: - assert ( - calculated_value == expected_value - ), "Got={}, Expected={}, Params={}".format( - calculated_value, expected_value, args - ) - - -def test_empty_tools(): - messages = [{"role": "user", "content": "hey, how's it going?", "tool_calls": None}] - - result = token_counter( - messages=messages, - ) - - print(result) - - -@pytest.mark.skip( - reason="Skipping this test temporarily because it relies on a function being called that I am removing." -) -def test_gpt_4o_token_counter(): - with patch.object( - litellm.utils, "openai_token_counter", new=MagicMock() - ) as mock_client: - token_counter( - model="gpt-4o-2024-05-13", messages=[{"role": "user", "content": "Hey!"}] - ) - - mock_client.assert_called() - - -@pytest.mark.parametrize( - "img_url", - [ - "https://example.com/test-image.png", - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAL0AAAC9CAMAAADRCYwCAAAAh1BMVEX///8AAAD8/Pz5+fkEBAT39/cJCQn09PRNTU3y8vIMDAwzMzPe3t7v7+8QEBCOjo7FxcXR0dHn5+elpaWGhoYYGBivr686OjocHBy0tLQtLS1TU1PY2Ni6urpaWlpERER3d3ecnJxoaGiUlJRiYmIlJSU4ODhBQUFycnKAgIDBwcFnZ2chISE7EjuwAAAI/UlEQVR4nO1caXfiOgz1bhJIyAJhX1JoSzv8/9/3LNlpYd4rhX6o4/N8Z2lKM2cURZau5JsQEhERERERERERERERERERERHx/wBjhDPC3OGN8+Cc5JeMuheaETSdO8vZFyCScHtmz2CsktoeMn7rLM1u3h0PMAEhyYX7v/Q9wQvoGdB0hlbzm45lEq/wd6y6G9aezvBk9AXwp1r3LHJIRsh6s2maxaJpmvqgvkC7WFS3loUnaFJtKRVUCEoV/RpCnHRvAsesVQ1hw+vd7Mpo+424tLs72NplkvQgcdrsvXkW/zJWqH/fA0FT84M/xnQJt4to3+ZLuanbM6X5lfXKHosO9COgREqpCR5i86pf2zPS7j9tTj+9nO7bQz3+xGEyGW9zqgQ1tyQ/VsxEDvce/4dcUPNb5OD9yXvR4Z2QisuP0xiGWPnemgugU5q/troHhGEjIF5sTOyW648aC0TssuaaCEsYEIkGzjWXOp3A0vVsf6kgRyqaDk+T7DIVWrb58b2tT5xpUucKwodOD/5LbrZC1ws6YSaBZJ/8xlh+XZSYXaMJ2ezNqjB3IPXuehPcx2U6b4t1dS/xNdFzguUt8ie7arnPeyCZroxLHzGgGdqVcspwafizPWEXBee+9G1OaufGdvNng/9C+gwgZ3PH3r87G6zXTZ5D5De2G2DeFoANXfbACkT+fxBQ22YFsTTJF9hjFVO6VbqxZXko4WJ8s52P4PnuxO5KRzu0/hlix1ySt8iXjgaQ+4IHPA9nVzNkdduM9LFT/Aacj4FtKrHA7iAw602Vnht6R8Vq1IOS+wNMKLYqayAYfRuufQPGeGb7sZogQQoLZrGPgZ6KoYn70Iw30O92BNEDpvwouCFn6wH2uS+EhRb3WF/HObZk3HuxfRQM3Y/Of/VH0n4MKNHZDiZvO9+m/ABALfkOcuar/7nOo7B95ACGVAFaz4jMiJwJhdaHBkySmzlGTu82gr6FSTik2kJvLnY9nOd/D90qcH268m3I/cgI1xg1maE5CuZYaWLH+UHANCIck0yt7Mx5zBm5vVHXHwChsZ35kKqUpmo5Svq5/fzfAI5g2vDtFPYo1HiEA85QrDeGm9g//LG7K0scO3sdpj2CBDgCa+0OFs0bkvVgnnM/QBDwllOMm+cN7vMSHlB7Uu4haHKaTwgGkv8tlK+hP8fzmFuK/RQTpaLPWvbd58yWIo66HHM0OsPoPhVqmtaEVL7N+wYcTLTbb0DLdgp23Eyy2VYJ2N7bkLFAAibtoLPe5sLt6Oa2bvU+zyeMa8wrixO0gRTn9tO9NCSThTLGqcqtsDvphlfmx/cPBZVvw24jg1LE2lPuEo35Mhi58U0I/Ga8n5w+NS8i34MAQLos5B1u0xL1ZvCVYVRw/Fs2q53KLaXJMWwOZZ/4MPYV19bAHmgGDKB6f01xoeJKFbl63q9J34KdaVNPJWztQyRkzA3KNs1AdAEDowMxh10emXTCx75CkurtbY/ZpdNDGdsn2UcHKHsQ8Ai3WZi48IfkvtjOhsLpuIRSKZTX9FA4o+0d6o/zOWqQzVJMynL9NsxhSJOaourq6nBVQBueMSyubsX2xHrmuABZN2Ns9jr5nwLFlLF/2R6atjW/67Yd11YQ1Z+kA9Zk9dPTM/o6dVo6HHVgC0JR8oUfmI93T9u3gvTG94bAH02Y5xeqRcjuwnKCK6Q2+ajl8KXJ3GSh22P3Zfx6S+n008ROhJn+JRIUVu6o7OXl8w1SeyhuqNDwNI7SjbK08QrqPxS95jy4G7nCXVq6G3HNu0LtK5J0e226CfC005WKK9sVvfxI0eUbcnzutfhWe3rpZHM0nZ/ny/N8tanKYlQ6VEW5Xuym8yV1zZX58vwGhZp/5tFfhybZabdbrQYOs8F+xEhmPsb0/nki6kIyVvzZzUASiOrTfF+Sj9bXC7DoJxeiV8tjQL6loSd0yCx7YyB6rPdLx31U2qCG3F/oXIuDuqd6LFO+4DNIJuxFZqSsU0ea88avovFnWKRYFYRQDfCfcGaBCLn4M4A1ntJ5E57vicwqq2enaZEF5nokCYu9TbKqCC5yCDfL+GhLxT4w4xEJs+anqgou8DOY2q8FMryjb2MehC1dRJ9s4g9NXeTwPkWON4RH+FhIe0AWR/S9ekvQ+t70XHeimGF78LzuU7d7PwrswdIG2VpgF8C53qVQsTDtBJc4CdnkQPbnZY9mbPdDFra3PCXBBQ5QBn2aQqtyhvlyYM4Hb2/mdhsxCUen04GZVvIJZw5PAamMOmjzq8Q+dzAKLXDQ3RUZItWsg4t7W2DP+JDrJDymoMH7E5zQtuEpG03GTIjGCW3LQqOYEsXgFc78x76NeRwY6SNM+IfQoh6myJKRBIcLYxZcwscJ/gI2isTBty2Po9IkYzP0/SS4hGlxRjFAG5z1Jt1LckiB57yWvo35EaolbvA+6fBa24xodL2YjsPpTnj3JgJOqhcgOeLVsYYwoK0wjY+m1D3rGc40CukkaHnkEjarlXrF1B9M6ECQ6Ow0V7R7N4G3LfOHAXtymoyXOb4QhaYHJ/gNBJUkxclpSs7DNcgWWDDmM7Ke5MJpGuioe7w5EOvfTunUKRzOh7G2ylL+6ynHrD54oQO3//cN3yVO+5qMVsPZq0CZIOx4TlcJ8+Vz7V5waL+7WekzUpRFMTnnTlSCq3X5usi8qmIleW/rit1+oQZn1WGSU/sKBYEqMNh1mBOc6PhK8yCfKHdUNQk8o/G19ZPTs5MYfai+DLs5vmee37zEyyH48WW3XA6Xw6+Az8lMhci7N/KleToo7PtTKm+RA887Kqc6E9dyqL/QPTugzMHLbLZtJKqKLFfzVWRNJ63c+95uWT/F7R0U5dDVvuS409AJXhJvD0EwWaWdW8UN11u/7+umaYjT8mJtzZwP/MD4r57fihiHlC5fylHfaqnJdro+Dr7DajvO+vi2EwyD70s8nCH71nzIO1l5Zl+v1DMCb5ebvCMkGHvobXy/hPumGLyX0218/3RyD1GRLOuf9u/OGQyDmto32yMiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIv7GP8YjWPR/czH2AAAAAElFTkSuQmCC", - ], -) -def test_img_url_token_counter(img_url, monkeypatch): - """ - Verify get_image_dimensions returns valid (width, height) for both an - HTTPS URL and a base64 data URI. The HTTPS branch is exercised with a - mocked HTTP fetch so the test is hermetic - it can't break when a - third-party image URL goes away. - """ - import base64 - from litellm.litellm_core_utils.token_counter import get_image_dimensions - - # Minimal valid 1x1 PNG, served by the mocked safe_get for the URL case. - _tiny_png = base64.b64decode( - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" - ) - - if img_url.startswith(("http://", "https://")): - - class _FakeResponse: - headers = {"Content-Length": str(len(_tiny_png))} - - def read(self): - return _tiny_png - - monkeypatch.setattr( - "litellm.litellm_core_utils.token_counter.safe_get", - lambda client, url, **kw: _FakeResponse(), - ) - - width, height = get_image_dimensions(data=img_url) - - print(width, height) - - assert width is not None - assert height is not None - - -def test_token_encode_disallowed_special(): - encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") - token_counter(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") - - -def test_token_counter(): - try: - messages = [{"role": "user", "content": "hi how are you what time is it"}] - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - print("gpt-35-turbo") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="claude-2", messages=messages) - print("claude-2") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="gemini/chat-bison", messages=messages) - print("gemini/chat-bison") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="ollama/llama2", messages=messages) - print("ollama/llama2") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages) - print("anthropic.claude-instant-v1") - print(tokens) - assert tokens > 0 - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -import unittest - -from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding - -# Clear the cache at module load to ensure clean state -_load_huggingface_tokenizer.cache_clear() - - -class TestTokenizerSelection(unittest.TestCase): - def setUp(self): - """Clear the LRU cache before each test method. - - The HuggingFace tokenizers behind _select_tokenizer_helper are cached with - @lru_cache, which can cause cache hits from previous tests when running with - --dist=loadscope (tests from same file run on same worker). - """ - _load_huggingface_tokenizer.cache_clear() - - @patch("litellm.utils.tokenizer_dispatch.from_pretrained") - def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): - # Setup mock to raise an error - mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") - - # Test with llama-3 model - result = _select_tokenizer_helper("llama-3-7b") - - # Verify the attempt to load Llama-3 tokenizer - mock_from_pretrained.assert_called_once_with("Xenova/llama-3-tokenizer") - - # Verify fallback to OpenAI tokenizer - self.assertEqual(result["type"], "openai_tokenizer") - self.assertEqual(result["tokenizer"], encoding) - - @patch("litellm.utils.tokenizer_dispatch.from_pretrained") - def test_cohere_tokenizer_api_failure(self, mock_from_pretrained): - # Setup mock to raise an error - mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") - - # Add Cohere model to the list for testing - litellm.cohere_models = ["command-r-v1"] - - # Test with Cohere model - result = _select_tokenizer_helper("command-r-v1") - - # Verify the attempt to load Cohere tokenizer - mock_from_pretrained.assert_called_once_with( - "Xenova/c4ai-command-r-v01-tokenizer" - ) - - # Verify fallback to OpenAI tokenizer - self.assertEqual(result["type"], "openai_tokenizer") - self.assertEqual(result["tokenizer"], encoding) - - @patch("litellm.utils.tokenizer_dispatch.anthropic") - def test_claude_tokenizer_api_failure(self, mock_anthropic): - # Setup mock to raise an error - mock_anthropic.side_effect = Exception("Failed to load tokenizer") - - # Add Claude model to the list for testing - litellm.anthropic_models = ["claude-2"] - - # Test with Claude model - result = _select_tokenizer_helper("claude-2") - - # Verify the attempt to load Claude tokenizer - mock_anthropic.assert_called_once_with() - - # Verify fallback to OpenAI tokenizer - self.assertEqual(result["type"], "openai_tokenizer") - self.assertEqual(result["tokenizer"], encoding) - - @patch("litellm.utils.tokenizer_dispatch.from_pretrained") - def test_llama2_tokenizer_api_failure(self, mock_from_pretrained): - # Setup mock to raise an error - mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") - - # Test with Llama-2 model - result = _select_tokenizer_helper("llama-2-7b") - - # Verify the attempt to load Llama-2 tokenizer - mock_from_pretrained.assert_called_once_with( - "hf-internal-testing/llama-tokenizer" - ) - - # Verify fallback to OpenAI tokenizer - self.assertEqual(result["type"], "openai_tokenizer") - self.assertEqual(result["tokenizer"], encoding) - - @patch("litellm.utils._return_huggingface_tokenizer") - def test_disable_hf_tokenizer_download(self, mock_return_huggingface_tokenizer): - monkeypatch = pytest.MonkeyPatch() - monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) - try: - result = _select_tokenizer_helper("grok-32r22r") - mock_return_huggingface_tokenizer.assert_not_called() - assert result["type"] == "openai_tokenizer" - assert result["tokenizer"] == encoding - finally: - monkeypatch.undo() - - -@pytest.mark.parametrize( - "model", - [ - "gpt-4o", - "claude-3-opus-20240229", - ], -) -@pytest.mark.parametrize( - "messages", - [ - [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "These are some sample images from a movie. Based on these images, what do you think the tone of the movie is?", - }, - { - "type": "text", - "image_url": { - "url": "https://gratisography.com/wp-content/uploads/2024/11/gratisography-augmented-reality-800x525.jpg", - "detail": "high", - }, - }, - ], - } - ], - ], -) -def test_bad_input_token_counter(model, messages): - """ - Safely handle bad input for token counter. - """ - token_counter( - model=model, - messages=messages, - default_token_count=1000, - ) - - -def test_token_counter_with_anthropic_tool_use(): - """ - Test that _count_anthropic_content() correctly handles tool_use blocks. - - Validates that: - - 'name' field is counted (string) - - 'input' field is counted (dict serialized to string) - - Metadata fields ('type', 'id') are skipped - """ - messages = [ - {"role": "user", "content": "What's the weather in San Francisco?"}, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "I'll check the weather for you."}, - { - "type": "tool_use", - "id": "toolu_01234567890", # Should be skipped - "name": "get_weather", # Should be counted - "input": { # Should be counted (serialized) - "location": "San Francisco, CA", - "unit": "fahrenheit", - }, - }, - ], - }, - ] - - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens > 0, f"Expected positive token count, got {tokens}" - # Should count: user message + "I'll check" text + "get_weather" name + input dict - assert ( - tokens > 15 - ), f"Expected reasonable token count for message with tool_use, got {tokens}" - - -def test_token_counter_with_anthropic_tool_result(): - """ - Test that _count_anthropic_content() correctly handles tool_result blocks. - - Validates that: - - 'content' field (when string) is counted - - Metadata fields ('type', 'tool_use_id') are skipped - - Full conversation with tool_use → tool_result flow works - """ - messages = [ - {"role": "user", "content": "What's the weather in San Francisco?"}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_01234567890", - "name": "get_weather", - "input": {"location": "San Francisco, CA"}, - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_01234567890", # Should be skipped - "content": "The weather in San Francisco is 65°F and sunny.", # Should be counted - } - ], - }, - ] - - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens > 0, f"Expected positive token count, got {tokens}" - assert ( - tokens > 25 - ), f"Expected reasonable token count for conversation with tool_result, got {tokens}" - - -def test_token_counter_with_nested_tool_result(): - """ - Test that _count_anthropic_content() recursively handles nested content lists. - - Validates that: - - tool_result with 'content' as a list (not string) is handled - - Nested content blocks are recursively counted via _count_content_list() - - TypedDict inference correctly identifies list fields - """ - messages = [ - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_01234567890", - "content": [ # Nested list - should recursively count - { - "type": "text", - "text": "The weather in San Francisco is 65°F and sunny.", - }, - {"type": "text", "text": "UV index is moderate."}, - ], - } - ], - } - ] - - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens > 0, f"Expected positive token count, got {tokens}" - # Should count both nested text blocks - assert ( - tokens > 15 - ), f"Expected reasonable token count for nested tool_result, got {tokens}" - - -def test_token_counter_tool_use_and_result_combined(): - """ - Test dynamic field inference with multiple tool_use and tool_result blocks. - - Validates that: - - Multiple tool_use blocks in same message are handled - - Multiple tool_result blocks in same message are handled - - skip_fields correctly filters metadata across all blocks - - Full realistic conversation flow works end-to-end - """ - messages = [ - { - "role": "user", - "content": "What's the weather in San Francisco and New York?", - }, - { - "role": "assistant", - "content": [ - { - "type": "text", - "text": "I'll check the weather in both cities for you.", - }, - { - "type": "tool_use", - "id": "toolu_01A", - "name": "get_weather", - "input": {"location": "San Francisco, CA"}, - }, - { - "type": "tool_use", - "id": "toolu_01B", - "name": "get_weather", - "input": {"location": "New York, NY"}, - }, - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_01A", - "content": "San Francisco: 65°F, sunny", - }, - { - "type": "tool_result", - "tool_use_id": "toolu_01B", - "content": "New York: 45°F, cloudy", - }, - ], - }, - { - "role": "assistant", - "content": "The weather in San Francisco is 65°F and sunny, while New York is cooler at 45°F and cloudy.", - }, - ] - - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - assert tokens > 0, f"Expected positive token count, got {tokens}" - # Should count all text, tool names, inputs, and results - assert ( - tokens > 60 - ), f"Expected substantial token count for full tool conversation, got {tokens}" - - -def test_token_counter_with_image_url(): - """ - Test that _count_image_tokens() correctly handles image_url content blocks. - - Validates that: - - image_url as dict with 'url' and 'detail' is handled - - image_url as string is handled - - 'detail' field validation works ('low', 'high', 'auto') - - calculate_img_tokens is called with correct parameters - """ - # Test with dict format (detail: low) - messages_dict = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What's in this image?"}, - { - "type": "image_url", - "image_url": { - "url": "https://example.com/image.jpg", - "detail": "low", # Should use low token count (85 base tokens) - }, - }, - ], - } - ] - - tokens_dict = token_counter( - model="gpt-3.5-turbo", - messages=messages_dict, - use_default_image_token_count=True, # Avoid actual HTTP request - ) - assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}" - assert tokens_dict > 85, f"Expected at least base image tokens, got {tokens_dict}" - - # Test with string format (defaults to auto/low) - messages_str = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": "https://example.com/image.jpg", # String format - } - ], - } - ] - - tokens_str = token_counter( - model="gpt-3.5-turbo", messages=messages_str, use_default_image_token_count=True - ) - assert ( - tokens_str > 0 - ), f"Expected positive token count for string image_url, got {tokens_str}" - - # Test invalid detail value raises error - messages_invalid = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "https://example.com/image.jpg", - "detail": "invalid", # Should raise ValueError - }, - } - ], - } - ] - - with pytest.raises(ValueError, match="Invalid detail value") as exc_info: - token_counter(model="gpt-3.5-turbo", messages=messages_invalid) - e = exc_info.value - assert "Invalid detail value" in str( - e - ), f"Expected detail validation error, got: {e}" - - -def test_token_counter_with_thinking_content(): - """ - Test that _count_content_list() correctly handles Claude's extended thinking content blocks. - - Validates that: - - 'thinking' content type is recognized and counted - - 'thinking' text field is counted - - 'signature' field is skipped (opaque signature blob) - - Full conversation with thinking blocks works - """ - messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Analyze this complex problem: who came first, chicken or egg", - } - ], - }, - { - "role": "assistant", - "content": [ - { - "type": "thinking", - "thinking": "This is actually a fascinating question that touches on philosophy, biology, and semantics. Let me break this down: The egg came first from an evolutionary biology perspective.", - "signature": "EqcLCkYICxgCKkCrqu6lP...", # Should be skipped - }, - { - "type": "text", - "text": "# The Chicken-or-Egg Question: A Multi-Layered Answer\n\n## **The Short Answer: The Egg Came First**", - }, - ], - }, - {"role": "user", "content": [{"type": "text", "text": "Thanks"}]}, - ] - - tokens = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", messages=messages - ) - assert tokens > 0, f"Expected positive token count, got {tokens}" - # Should count: user message + thinking text + response text + "Thanks" - # The thinking text alone is ~30 tokens, plus other content should be > 50 total - assert ( - tokens > 50 - ), f"Expected substantial token count for message with thinking, got {tokens}" - - # Test that thinking block without 'thinking' field doesn't crash (edge case) - messages_no_thinking = [ - { - "role": "assistant", - "content": [ - { - "type": "thinking", - # No 'thinking' field - should count as 0 tokens - "signature": "EqcLCkYICxgCKkCrqu6lP...", - }, - {"type": "text", "text": "Response"}, - ], - } - ] - - tokens_no_thinking = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", messages=messages_no_thinking - ) - assert ( - tokens_no_thinking > 0 - ), f"Expected positive token count even with empty thinking, got {tokens_no_thinking}" - # Should only count "Response" and message overhead - assert ( - tokens_no_thinking < 15 - ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}" - - - -def test_token_counter_with_redacted_thinking_content(): - """ - A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in - for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking - block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the - prompt_caching pre-call check stop pinning the deployment that held the cached prefix. - """ - model = "anthropic/claude-sonnet-4-5-20250929" - reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."} - redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30} - user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]} - follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]} - - without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up] - with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up] - - assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) - -def test_token_counter_with_tool_reference_block(): - """ - Regression test: a message containing an Anthropic tool-search - `tool_reference` content block must NOT raise. - - Before the fix, token_counter raised - `Invalid content item type: tool_reference`. On the streaming - anthropic_messages proxy path this nulled response_cost and caused the - SpendLogs row to be dropped, silently undercounting cost. token_counter - must instead count the referenced tool name and return a positive count. - """ - messages = [ - { - "role": "assistant", - "content": [ - {"type": "text", "text": "Let me look up the right tool."}, - {"type": "tool_reference", "tool_name": "search_knowledge_base"}, - ], - } - ] - - # Must not raise, and must produce a positive token count. - tokens = token_counter_new( - model="anthropic/claude-sonnet-4-5-20250929", messages=messages - ) - assert tokens > 0, f"Expected positive token count, got {tokens}" - - # A tool_reference with no/empty tool_name must also be handled gracefully. - messages_empty = [ - { - "role": "assistant", - "content": [{"type": "tool_reference", "tool_name": ""}], - } - ] - tokens_empty = token_counter_new( - model="anthropic/claude-sonnet-4-5-20250929", messages=messages_empty - ) - assert tokens_empty >= 0 - - -def test_count_content_list_rejects_unknown_type(): - """ - An unrecognized content block type must raise, and the error message must - enumerate the supported types (including `tool_reference`). This pins the - catch-all contract so a future block type isn't silently dropped. - """ - from litellm.litellm_core_utils.token_counter import _count_content_list - - with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info: - _count_content_list( - count_function=len, - content_list=[{"type": "totally_unknown_block"}], - use_default_image_token_count=False, - default_token_count=None, - ) - - message = str(exc_info.value) - assert "Invalid content item type: totally_unknown_block" in message - assert "tool_reference" in message - - -@pytest.mark.parametrize( - "source", - [ - {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}, - {"type": "url", "url": "https://example.com/image.png"}, - {"type": "file", "file_id": "file-abc123"}, - ], - ids=["base64", "url", "file"], -) -def test_token_counter_with_anthropic_image_block(source: dict[str, str]): - """Anthropic `image` blocks must count for every source variant, not raise `Invalid content item type` (which the router's context-window pre-call check swallows into an unfiltered dispatch).""" - from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT - - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is in this image?"}, - {"type": "image", "source": source}, - ], - } - ] - - tokens = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", - messages=messages, - use_default_image_token_count=True, - ) - assert tokens > DEFAULT_IMAGE_TOKEN_COUNT, ( - f"Expected the image block to contribute tokens, got {tokens}" - ) - - -def test_anthropic_image_block_matches_equivalent_image_url(): - """An Anthropic `image` block prices identically to the OpenAI `image_url` carrying the same bytes.""" - anthropic_messages = [ - { - "role": "user", - "content": [ - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": "iVBORw0KGgo=", - }, - } - ], - } - ] - openai_messages = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, - } - ], - } - ] - - anthropic_tokens = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", messages=anthropic_messages - ) - openai_tokens = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", messages=openai_messages - ) - assert anthropic_tokens == openai_tokens - - -def test_anthropic_image_block_nested_in_tool_result(): - """An `image` block nested in a `tool_result.content` list is counted through the same recursion.""" - messages = [ - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_01", - "content": [ - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": "iVBORw0KGgo=", - }, - } - ], - } - ], - } - ] - - tokens = token_counter( - model="anthropic/claude-sonnet-4-5-20250929", - messages=messages, - use_default_image_token_count=True, - ) - assert tokens > 0 - - -@pytest.mark.parametrize( - ("source", "expected"), - [ - ({"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}, "data:image/jpeg;base64,/9j/4AAQ"), - ({"type": "url", "url": "https://example.com/image.png"}, "https://example.com/image.png"), - ({"type": "file", "file_id": "file-abc123"}, ""), - ], - ids=["base64", "url", "file"], -) -def test_anthropic_image_source_resolves_to_what_the_image_pricer_reads(source: dict[str, str], expected: str): - """base64 sources become a data URI, url sources pass through, file sources resolve to an empty string.""" - from litellm.litellm_core_utils.token_counter import _anthropic_image_source_data - - assert _anthropic_image_source_data(source) == expected - - -def test_anthropic_image_block_with_empty_base64_data(): - """A base64 source with empty `data` prices as an image rather than raising.""" - from litellm.litellm_core_utils.token_counter import _count_content_list - - tokens = _count_content_list( - count_function=len, - content_list=[ - {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}} - ], - use_default_image_token_count=False, - default_token_count=None, - ) - assert tokens > 0 - - -def test_anthropic_image_block_without_source_raises(): - """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior.""" - from litellm.litellm_core_utils.token_counter import _count_content_list - - with pytest.raises(ValueError, match="Error getting number of tokens from content list"): - _count_content_list( - count_function=len, - content_list=[{"type": "image"}], - use_default_image_token_count=False, - default_token_count=None, - ) - - # ... and `default_token_count`, the caller's opt-out from raising, still wins. - assert ( - _count_content_list( - count_function=len, - content_list=[{"type": "image"}], - use_default_image_token_count=False, - default_token_count=7, - ) - == 7 - ) - - -def _count_user_content(content: list[dict]) -> int: - from litellm.litellm_core_utils.token_counter import token_counter - - return token_counter( - model="anthropic/claude-fable-5", - messages=[{"role": "user", "content": content}], - use_default_image_token_count=True, - ) - - -@pytest.mark.parametrize( - "source", - [ - {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, - {"type": "url", "url": "https://example.com/report.pdf"}, - {"type": "file", "file_id": "file-abc123"}, - ], - ids=["base64", "url", "file"], -) -def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]): - """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on.""" - prompt = {"type": "text", "text": "Summarize this file."} - - assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content( - [prompt, {"type": "image", "source": source}] - ) - - -def test_anthropic_document_block_text_sources_count_their_text(): - """`text` and `content` document sources count the text they carry, as inline text blocks would.""" - prompt = {"type": "text", "text": "Summarize this file."} - body = {"type": "text", "text": "Revenue grew eleven percent while churn fell to two percent."} - picture = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}} - - text_source = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": body["text"]}} - assert _count_user_content([prompt, text_source]) == _count_user_content([prompt, body]) - - string_content = {"type": "document", "source": {"type": "content", "content": body["text"]}} - assert _count_user_content([prompt, string_content]) == _count_user_content([prompt, body]) - - block_content = {"type": "document", "source": {"type": "content", "content": [body, picture]}} - assert _count_user_content([prompt, block_content]) == _count_user_content([prompt, body, picture]) - - -def test_anthropic_document_title_and_context_add_their_tokens(): - prompt = {"type": "text", "text": "Summarize this file."} - source = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} - described = {"type": "document", "source": source, "title": "Q3 board packet", "context": "Shared by finance"} - - assert _count_user_content([prompt, described]) == _count_user_content( - [ - prompt, - {"type": "text", "text": "Q3 board packet"}, - {"type": "text", "text": "Shared by finance"}, - {"type": "document", "source": source}, - ] - ) - - -def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): - """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise. - - Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject` - is in the union this counter accepts, so every local count of a Responses `input_file` raised - `Invalid content item type: file` and surfaced as a 500 on /v1/responses/input_tokens. - """ - prompt = {"type": "text", "text": "Summarize this file."} - inline_file = { - "type": "file", - "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, - } - document = { - "type": "document", - "title": "report.pdf", - "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, - } - - assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) - assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) - - -def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): - """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens.""" - prompt = {"type": "text", "text": "Summarize this file."} - - by_id = {"type": "file", "file": {"file_id": "file-abc123"}} - assert _count_user_content([prompt, by_id]) == _count_user_content([prompt]) - - named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}} - assert _count_user_content([prompt, named]) == _count_user_content( - [prompt, {"type": "text", "text": "report.pdf"}] - ) - - -def _png_data_url(width: int, height: int) -> str: - ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") - return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() - - -@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) -def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: - assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() - - -def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: - assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() - assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/litellm_core_utils/test_tokenizer.py b/tests/test_litellm/litellm_core_utils/test_tokenizer.py index aa4a0fc6a1c..2171044970c 100644 --- a/tests/test_litellm/litellm_core_utils/test_tokenizer.py +++ b/tests/test_litellm/litellm_core_utils/test_tokenizer.py @@ -1,403 +1,20 @@ -import copy -import os -import pickle -import subprocess -import sys -from pathlib import Path -from typing import Final, Literal - import pytest -import tiktoken -from tokenizers import Tokenizer as ReferenceTokenizer -import litellm -from litellm.caching._embedding_router import truncate_embedding_input -from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding -from litellm.utils import claude_json_str -from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON - - -@pytest.mark.parametrize( - "name", ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "r50k_base", "gpt2", "o200k_harmony") -) -@pytest.mark.parametrize( - "text", ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64) +from tests.unit.litellm_core_utils.test_tokenizer import ( + UNICODE_TEXTS, + assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface, + assert_openai_encoding_matches_python, ) + +NETWORK_ENCODINGS = ("r50k_base", "gpt2") + + +@pytest.mark.parametrize("name", NETWORK_ENCODINGS) +@pytest.mark.parametrize("text", UNICODE_TEXTS) def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: - reference: Final = tiktoken.get_encoding(name) - encoding: Final = OpenAIEncoding.from_tiktoken(name) - expected: Final = reference.encode(text) - - assert encoding.encode(text) == expected - assert encoding.count(text) == len(expected) - assert encoding.encode_batch([text], num_threads=2) == reference.encode_batch([text], num_threads=2) - assert encoding.encode_ordinary_batch([text]) == reference.encode_ordinary_batch([text]) - assert encoding.decode_batch([expected]) == reference.decode_batch([expected]) - assert encoding.decode_bytes_batch([expected]) == reference.decode_bytes_batch([expected]) + assert_openai_encoding_matches_python(name, text) -@pytest.mark.parametrize("allowed", (frozenset(), frozenset({"<|endoftext|>"}), "all")) -@pytest.mark.parametrize("disallowed", (frozenset(), frozenset({"<|fim_prefix|>"}), "all")) -def test_openai_special_token_options_match_python( - allowed: frozenset[str] | Literal["all"], disallowed: frozenset[str] | Literal["all"] -) -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) - text: Final = "hello<|endoftext|><|fim_prefix|>world" - allowed_set: Final = reference.special_tokens_set if allowed == "all" else allowed - disallowed_set: Final = reference.special_tokens_set - allowed_set if disallowed == "all" else disallowed - if any(token in text for token in disallowed_set): - with pytest.raises(ValueError, match="disallowed special token"): - encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) - return - assert encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) == reference.encode( - text, allowed_special=allowed, disallowed_special=disallowed - ) - assert encoding.special_tokens_set == reference.special_tokens_set - assert encoding.eot_token == reference.eot_token - - -@pytest.mark.parametrize("errors", ("replace", "ignore", "backslashreplace", "strict")) -def test_openai_partial_token_decoding_preserves_error_policy(errors: str) -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) - tokens: Final = reference.encode("🙂")[:1] - assert encoding.decode_bytes(tokens) == reference.decode_bytes(tokens) - if errors == "strict": - with pytest.raises(UnicodeDecodeError): - encoding.decode(tokens, errors=errors) - return - assert encoding.decode(tokens, errors=errors) == reference.decode(tokens, errors=errors) - assert encoding.decode_tokens_bytes(tokens) == reference.decode_tokens_bytes(tokens) - - -def test_public_encoding_and_semantic_cache_preserve_truncated_unicode() -> None: - reference: Final = tiktoken.get_encoding(litellm.encoding.name) - text: Final = "🙂" - tokens: Final = reference.encode(text) - - assert litellm.encoding.encode(text, disallowed_special=()) == tokens - assert litellm.encoding.encode_batch([text]) == [tokens] - assert litellm.decode(tokens=tokens[:1]) == reference.decode(tokens[:1]) - assert truncate_embedding_input(text, "", 1) == reference.decode(tokens[:1]) - - -@pytest.mark.parametrize("add_special_tokens", (True, False)) -def test_huggingface_encoding_preserves_result_fields_and_serialization(add_special_tokens: bool) -> None: - reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) - tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) - expected: Final = reference.encode("Hello World", add_special_tokens=add_special_tokens) - actual: Final = tokenizer.encode("Hello World", add_special_tokens=add_special_tokens) - - assert (actual.ids, actual.tokens, actual.type_ids, actual.offsets, actual.word_ids, actual.sequence_ids) == ( - expected.ids, - expected.tokens, - expected.type_ids, - expected.offsets, - expected.word_ids, - expected.sequence_ids, - ) - assert (actual.attention_mask, actual.special_tokens_mask, actual.n_sequences, len(actual)) == ( - expected.attention_mask, - expected.special_tokens_mask, - expected.n_sequences, - len(expected), - ) - assert copy.deepcopy(actual).ids == expected.ids - assert pickle.loads(pickle.dumps(actual)).offsets == expected.offsets - assert tokenizer.decode(actual.ids, skip_special_tokens=False) == reference.decode( - expected.ids, skip_special_tokens=False - ) - - -def test_huggingface_character_offsets_and_pretokenized_pairs_match_python() -> None: - reference: Final = ReferenceTokenizer.from_str(claude_json_str) - tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) - text: Final = "café 漢字 🙂" - actual: Final = tokenizer.encode(text) - expected: Final = reference.encode(text) - - assert actual.offsets == expected.offsets - assert actual.ids == expected.ids - assert ( - tokenizer.encode(["hello", "world"], ["again"], is_pretokenized=True).ids - == reference.encode(["hello", "world"], ["again"], is_pretokenized=True).ids - ) - - -def test_huggingface_batches_apply_padding_across_inputs() -> None: - reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) - reference.enable_padding(pad_id=0, pad_token="[UNK]") - tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) - inputs: Final = ["Hello", ("Hello World", "World")] - expected: Final = reference.encode_batch(inputs) - actual: Final = tokenizer.encode_batch(inputs) - fast: Final = tokenizer.encode_batch_fast(inputs) - - assert [(item.ids, item.attention_mask, item.offsets) for item in actual] == [ - (item.ids, item.attention_mask, item.offsets) for item in expected - ] - assert [item.ids for item in fast] == [item.ids for item in expected] - assert tokenizer.decode_batch([item.ids for item in actual]) == reference.decode_batch( - [item.ids for item in expected] - ) - - -def test_caller_supplied_huggingface_tokenizer_preserves_public_encode_and_count() -> None: - tokenizer: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) - custom: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer} - expected: Final = tokenizer.encode("Hello World").ids - - assert litellm.encode(text="Hello World", custom_tokenizer=custom) == expected - assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(expected) - assert litellm.decode(tokens=expected, custom_tokenizer=custom) == "Hello World" - - -def test_caller_supplied_tiktoken_treats_special_spellings_as_text() -> None: - tokenizer: Final = tiktoken.get_encoding("cl100k_base") - custom: Final = {"type": "openai_tokenizer", "tokenizer": tokenizer} - text: Final = "<|endoftext|>" - - assert litellm.encode(text=text, custom_tokenizer=custom) == tokenizer.encode(text, disallowed_special=()) - - -def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> None: - custom: Final = litellm.create_tokenizer(TOKENIZER_JSON) - tokenizer: Final = custom["tokenizer"] - path: Final = tmp_path / "tokenizer.json" - tokenizer.save(str(path)) - - assert copy.deepcopy(custom)["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids - assert ( - pickle.loads(pickle.dumps(custom))["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids - ) - assert HuggingFaceTokenizer.from_file(str(path)).encode("Hello World").ids == tokenizer.encode("Hello World").ids - assert copy.deepcopy(litellm.encoding).encode("hello") == litellm.encoding.encode("hello") - assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello") - - -@pytest.mark.parametrize("offline", ("0", "1")) -def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None: - script: Final = """ -import json -import sys -from pathlib import Path -sys.path.insert(0, sys.argv[1]) -import httpx -import huggingface_hub -from huggingface_hub.errors import LocalEntryNotFoundError -import litellm -payload = sys.argv[2].encode() -offline = sys.argv[3] == "1" -observed = [] -def handle(request): - assert not offline, "offline loading issued a request" - if request.url.path.endswith("/tokenizer.json"): - observed.append(request.headers.get("authorization")) - if request.headers.get("authorization") != "Bearer audit-fixture-token": - return httpx.Response(401) - return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"") -if not offline: - huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) -try: - tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"] -except LocalEntryNotFoundError: - assert offline - assert observed == [] -else: - assert not offline - assert "Bearer audit-fixture-token" in observed - assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World" - assert tuple(Path(sys.argv[4]).rglob("tokenizer.json")) -print("compatible") -""" - result: Final = subprocess.run( - [ - sys.executable, - "-I", - "-c", - script, - str(Path(litellm.__file__).parent.parent), - TOKENIZER_JSON, - offline, - str(tmp_path / "cache"), - ], - capture_output=True, - text=True, - timeout=30, - env={ - **os.environ, - "HF_HOME": str(tmp_path / "home"), - "HF_HUB_CACHE": str(tmp_path / "cache"), - "HF_ENDPOINT": "http://127.0.0.1:9", - "HF_TOKEN": "audit-fixture-token", - "HF_HUB_OFFLINE": offline, - "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0", - "LITELLM_LOCAL_MODEL_COST_MAP": "True", - }, - ) - assert result.returncode == 0, result.stdout + result.stderr - assert result.stdout.strip() == "compatible" - - -@pytest.mark.parametrize("rust", (None, "0", "1")) -def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None: - script: Final = """ -import importlib.abc -import sys -sys.path.insert(0, sys.argv[1]) -def reject_network(event, args): - if event == "socket.connect": - raise AssertionError("tokenizer attempted a network connection") -sys.addaudithook(reject_network) -class Block(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "litellm.rust_bridge._native": - raise ImportError("native extension is unavailable") -sys.meta_path.insert(0, Block()) -import litellm -from litellm.rust_bridge.tokenizer import get_encoding -import tiktoken -from tokenizers import Tokenizer -assert isinstance(litellm.encoding, tiktoken.Encoding) -for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"): - encoding = get_encoding(name) - text = "offline café 漢字 🙂" + " " * 64 - assert encoding.decode(encoding.encode(text)) == text -ids = litellm.encode(text="hello world") -assert litellm.decode(tokens=ids) == "hello world" -assert litellm.token_counter(model=None, text="hello world") == len(ids) -custom = litellm.create_tokenizer(sys.argv[2]) -assert isinstance(custom["tokenizer"], Tokenizer) -custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]") -assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World" -print("compatible") -""" - result: Final = subprocess.run( - [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON], - capture_output=True, - text=True, - timeout=30, - cwd=tmp_path, - env={ - **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"}, - **({"LITELLM_RUST": rust} if rust is not None else {}), - "LITELLM_LOCAL_MODEL_COST_MAP": "True", - "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), - }, - ) - assert result.returncode == 0, result.stdout + result.stderr - assert result.stdout.strip() == "compatible" - assert not (tmp_path / "unused-tokenizer-cache").exists() - - -@pytest.mark.parametrize("is_pretokenized", (False, True)) -def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None: - reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) - tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) - inputs: Final = [["Hello", "World"], ("Hello", "World")] - actual: Final = tokenizer.encode_batch(inputs, is_pretokenized=is_pretokenized) - expected: Final = reference.encode_batch(inputs, is_pretokenized=is_pretokenized) - assert [(item.ids, item.type_ids, item.sequence_ids) for item in actual] == [ - (item.ids, item.type_ids, item.sequence_ids) for item in expected - ] - - -@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit", "gpt2")) +@pytest.mark.parametrize("name", ("gpt2",)) def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: - reference: Final = tiktoken.get_encoding(name) - encoding: Final = OpenAIEncoding.from_tiktoken(name) - text: Final = "hello fanta" - - assert repr(encoding) == repr(reference) == f"" - assert (encoding.name, encoding.n_vocab, encoding.max_token_value) == ( - reference.name, - reference.n_vocab, - reference.max_token_value, - ) - assert encoding.token_byte_values() == reference.token_byte_values() - assert encoding.encode_single_token("hello") == reference.encode_single_token("hello") - assert encoding.encode_single_token(b"<|endoftext|>") == reference.eot_token - assert [encoding.is_special_token(token) for token in (0, reference.eot_token)] == [False, True] - assert encoding.decode_with_offsets(reference.encode(text)) == reference.decode_with_offsets(reference.encode(text)) - assert encoding.encode_to_numpy(text).tolist() == reference.encode_to_numpy(text).tolist() - stable, completions = encoding.encode_with_unstable(text) - expected_stable, expected_completions = reference.encode_with_unstable(text) - assert (stable, sorted(completions)) == (expected_stable, sorted(expected_completions)) - with pytest.raises(KeyError): - encoding.encode_single_token("<|not-a-token|>") - - -def test_huggingface_tokenizer_exposes_the_tokenizers_vocabulary_surface() -> None: - reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) - reference.enable_padding(pad_id=0, pad_token="[UNK]", length=4) - reference.enable_truncation(max_length=3, stride=1, strategy="only_first", direction="left") - tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) - - assert tokenizer.token_to_id("Hello") == reference.token_to_id("Hello") == 1 - assert tokenizer.id_to_token(3) == reference.id_to_token(3) == "[BOS]" - assert tokenizer.id_to_token(99) is None - assert tokenizer.get_vocab() == reference.get_vocab() - assert tokenizer.get_vocab(with_added_tokens=False) == reference.get_vocab(with_added_tokens=False) - assert tokenizer.get_vocab_size() == reference.get_vocab_size() == 4 - assert tokenizer.get_vocab_size(with_added_tokens=False) == reference.get_vocab_size(with_added_tokens=False) - added: Final = tokenizer.get_added_tokens_decoder() - expected_added: Final = reference.get_added_tokens_decoder() - assert {token_id: str(token) for token_id, token in added.items()} == { - token_id: str(token) for token_id, token in expected_added.items() - } - assert added[3].special == expected_added[3].special - assert tokenizer.num_special_tokens_to_add(False) == reference.num_special_tokens_to_add(False) == 1 - assert tokenizer.num_special_tokens_to_add(True) == reference.num_special_tokens_to_add(True) == 0 - assert tokenizer.padding == reference.padding - assert tokenizer.truncation == reference.truncation - assert tokenizer.encode_special_tokens == reference.encode_special_tokens is False - assert HuggingFaceTokenizer.from_buffer(TOKENIZER_JSON.encode()).encode("Hello").ids == [3, 1] - assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).padding is None - assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).truncation is None - - -def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface() -> None: - reference: Final = ReferenceTokenizer.from_str(claude_json_str) - tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) - text: Final = "hello wide world" - actual: Final = tokenizer.encode(text, "again") - expected: Final = reference.encode(text, "again") - - lookups: Final = ( - lambda encoding: [encoding.token_to_chars(index) for index in range(len(encoding))], - lambda encoding: [encoding.token_to_word(index) for index in range(len(encoding))], - lambda encoding: [encoding.token_to_sequence(index) for index in range(len(encoding))], - lambda encoding: [encoding.char_to_token(position) for position in range(len(text))], - lambda encoding: [encoding.char_to_word(position) for position in range(len(text))], - lambda encoding: [encoding.char_to_token(position, 1) for position in range(5)], - lambda encoding: [encoding.word_to_tokens(word) for word in range(3)], - lambda encoding: [encoding.word_to_chars(word) for word in range(3)], - lambda encoding: [encoding.word_to_tokens(0, 1), encoding.word_to_chars(0, 1)], - ) - for lookup in lookups: - assert lookup(actual) == lookup(expected) - assert repr(actual) == repr(expected) - - actual.truncate(4, stride=1, direction="left") - expected.truncate(4, stride=1, direction="left") - assert (actual.ids, [item.ids for item in actual.overflowing]) == ( - expected.ids, - [item.ids for item in expected.overflowing], - ) - actual.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") - expected.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") - assert (actual.ids, actual.attention_mask, actual.type_ids, actual.tokens) == ( - expected.ids, - expected.attention_mask, - expected.type_ids, - expected.tokens, - ) - actual.set_sequence_id(3) - expected.set_sequence_id(3) - assert actual.sequence_ids == expected.sequence_ids - merged: Final = type(actual).merge([actual, tokenizer.encode("more")]) - assert merged.ids == type(expected).merge([expected, reference.encode("more")]).ids - assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets - with pytest.raises(ValueError, match="direction"): - actual.pad(8, direction="sideways") + assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name) diff --git a/tests/test_litellm/proxy/client/test_chat.py b/tests/test_litellm/proxy/client/test_chat.py index 67b6ee833f2..8fe1bfcbb2f 100644 --- a/tests/test_litellm/proxy/client/test_chat.py +++ b/tests/test_litellm/proxy/client/test_chat.py @@ -13,7 +13,7 @@ from litellm.proxy.client.exceptions import UnauthorizedError def _load_http_mocking_responses(): """Load the third-party `responses` package even if test collection creates - a top-level `responses` namespace package from `tests/test_litellm/responses`. + a top-level `responses` namespace package from `tests/unit/responses`. """ module = importlib.import_module("responses") if hasattr(module, "activate"): diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index e6795bb22f3..42c1f489bdd 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3674,7 +3674,7 @@ async def test_post_call_success_hook_contains_header_merge_failures( @pytest.mark.asyncio async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py index e91b7ef970c..9d4532df49a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -162,7 +162,7 @@ async def test_interrupted_anthropic_stream_recovers_output_tokens_off_the_event from unittest.mock import AsyncMock from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, @@ -201,7 +201,7 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop( from unittest.mock import AsyncMock from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 884a9c81500..89156cd19a0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -14974,7 +14974,7 @@ def test_settings_store_exposes_dashboard_saved_mcp_client_allowlist_to_the_mcp_ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, @@ -14995,7 +14995,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp from litellm.rust_bridge._native import Tokenizer from litellm import Router - from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 0fc7295a717..ea1870d3b73 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -2021,7 +2021,7 @@ async def test_a_dispatched_failure_is_counted_off_the_event_loop(): from unittest.mock import AsyncMock, patch from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/test_litellm/rust_bridge/messages/test_route_host.py deleted file mode 100644 index c5a442e0709..00000000000 --- a/tests/test_litellm/rust_bridge/messages/test_route_host.py +++ /dev/null @@ -1,124 +0,0 @@ -from dataclasses import astuple -from typing import Final - -import pytest - -import litellm -from litellm.rust_bridge.messages import route_host - -pytestmark = pytest.mark.usefixtures("local_model_cost_map") - - -def _flag_model(monkeypatch: pytest.MonkeyPatch, name: str, **flags: bool) -> None: - monkeypatch.setitem( - litellm.model_cost, - name, - { - "litellm_provider": "anthropic", - "mode": "chat", - "input_cost_per_token": 0, - "output_cost_per_token": 0, - **flags, - }, - ) - - -def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeypatch: pytest.MonkeyPatch) -> None: - _flag_model( - monkeypatch, - "claude-test-adaptive", - supports_reasoning=True, - supports_adaptive_thinking=True, - supports_output_config=True, - supports_xhigh_reasoning_effort=True, - supports_sampling_params=False, - ) - - capabilities: Final = route_host.model_capabilities("anthropic/claude-test-adaptive", None) - - assert capabilities.supports_adaptive_thinking - assert capabilities.supports_output_config - assert not capabilities.supports_legacy_thinking - assert not capabilities.supports_sampling_params - assert capabilities.effort_tiers.xhigh - assert not capabilities.effort_tiers.max - - -def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None: - capabilities: Final = route_host.model_capabilities("anthropic/not-a-real-model", None) - - assert capabilities.supports_sampling_params - assert not capabilities.supports_reasoning - assert not capabilities.supports_adaptive_thinking - assert not any(astuple(capabilities.effort_tiers)) - - -@pytest.mark.parametrize( - ("global_flag", "kwargs", "expected"), - [ - (False, {}, False), - (True, {}, True), - (False, {"drop_params": "true"}, True), - (False, {"drop_params": "nonsense"}, False), - (False, {"drop_params": False}, False), - ], -) -def test_drop_params_merges_the_global_flag_with_the_request( - monkeypatch: pytest.MonkeyPatch, global_flag: bool, kwargs: dict[str, object], expected: bool -) -> None: - monkeypatch.setattr(litellm, "drop_params", global_flag) - - assert route_host.shaping("anthropic/not-a-real-model", None, kwargs)["drop_params"] is expected - - -@pytest.mark.parametrize( - ("configured", "expected"), - [ - (["tools[*].input_examples", 3, "metadata.user_id"], ("tools[*].input_examples", "metadata.user_id")), - ("tools", ()), - (None, ()), - ], -) -def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None: - shaping: Final = route_host.shaping("anthropic/not-a-real-model", None, {"additional_drop_params": configured}) - - assert shaping["additional_drop_params"] == expected - - -def test_native_request_rejections_map_to_the_public_400() -> None: - from types import MappingProxyType - - from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest - - request: Final = LiteLLMMessagesRequest( - model="anthropic/claude-sonnet-5", - messages=(), - max_tokens=8, - stream=None, - api_key=None, - api_base=None, - custom_llm_provider=None, - kwargs=MappingProxyType({}), - ) - rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5") - rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets - - mapped: Final = route_host.map_failure(rejected, request, "anthropic") - - assert isinstance(mapped, litellm.BadRequestError) - assert mapped.status_code == 400 - assert "does not support top_k=5" in mapped.message - assert mapped.model == "claude-sonnet-5" - assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError) - - -def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None: - hidden: Final = route_host.stream_hidden_params( - (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41")) - ) - - additional: Final = hidden["additional_headers"] - assert isinstance(additional, dict) - assert additional["llm_provider-request-id"] == "req_upstream_123" - assert additional["x-ratelimit-remaining-requests"] == "41" - assert "request-id" not in additional diff --git a/tests/test_litellm/rust_bridge/responses/__init__.py b/tests/test_litellm/rust_bridge/responses/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm_rust/tokenizer/test_fast_count.py b/tests/test_litellm_rust/tokenizer/test_fast_count.py index 2902b79dca8..f91f47e4b86 100644 --- a/tests/test_litellm_rust/tokenizer/test_fast_count.py +++ b/tests/test_litellm_rust/tokenizer/test_fast_count.py @@ -7,7 +7,7 @@ from tokenizers import Tokenizer as ReferenceTokenizer from litellm.rust_bridge import _native from litellm.utils import claude_json_str -from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON +from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON pytestmark = pytest.mark.requires_rust_extension diff --git a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py index abf6a6dda31..2e883e91fda 100644 --- a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py @@ -94,7 +94,7 @@ class _AgentChunk: @pytest.mark.asyncio async def test_stream_completion_counts_tokens_off_the_event_loop(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py index c65d171246d..4ba0ef8fa04 100644 --- a/tests/unit/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -469,7 +469,7 @@ class _UsageRecorder(CustomLogger): @pytest.mark.asyncio async def test_asend_message_counts_usage_off_the_event_loop(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/caching/test_azure_blob_cache.py b/tests/unit/caching/test_azure_blob_cache.py similarity index 100% rename from tests/test_litellm/caching/test_azure_blob_cache.py rename to tests/unit/caching/test_azure_blob_cache.py diff --git a/tests/test_litellm/caching/test_caching.py b/tests/unit/caching/test_caching.py similarity index 100% rename from tests/test_litellm/caching/test_caching.py rename to tests/unit/caching/test_caching.py diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index a181ef89fe0..425d657312a 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -39,6 +39,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm._logging import verbose_logger import logging +import json +import httpx +import respx +from fastapi.testclient import TestClient +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES def setup_cache(): @@ -1062,6 +1067,9 @@ def test_is_chat_completion_cached_dict(): assert _is_chat_completion_cached_dict( {"id": "other", "object": "chat.completion.chunk", "choices": []} ) + assert _is_chat_completion_cached_dict( + {"id": "no-object", "choices": [{"index": 0}]} + ) assert not _is_chat_completion_cached_dict( {"id": "resp_abc", "object": "response", "output": []} ) @@ -1432,3 +1440,799 @@ def test_convert_cached_responses_result_parameterized( assert result is not None assert result.id == cached_result["id"] assert result.status == cached_result["status"] + + +@pytest.mark.asyncio +async def test_process_async_embedding_cached_response(): + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + args = { + "cached_result": [ + { + "embedding": [-0.025122925639152527, -0.019487135112285614], + "index": 0, + "object": "embedding", + } + ] + } + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=args["cached_result"], + kwargs={"model": "text-embedding-ada-002", "input": "test"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + + print(f"response: {response}") + assert len(response.data) == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_preserves_prompt_tokens_details(): + """Test that prompt_tokens_details (including image_count) survives a full cache hit.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_backward_compat_no_prompt_tokens_details(): + """Test that old cached items without prompt_tokens_details still work.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Old-format cached item — no prompt_tokens_details field + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "test"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is None + + +@pytest.mark.asyncio +async def test_embedding_cache_aggregates_multiple_image_counts(): + """Test that image_count is summed correctly across multiple cached items.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + { + "embedding": [0.031, 0.042], + "index": 1, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={ + "model": "amazon.titan-embed-image-v1", + "input": ["img1", "img2"], + }, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 2 + + +def test_combine_usage_merges_prompt_tokens_details(): + """Test that combine_usage merges prompt_tokens_details from both Usage objects.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + usage1 = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + usage2 = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2), + ) + + combined = llm_caching_handler.combine_usage(usage1, usage2) + + assert combined.prompt_tokens == 30 + assert combined.total_tokens == 30 + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 3 + + +def test_combine_usage_handles_none_details(): + """Test that combine_usage works when one or both sides have null prompt_tokens_details.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Both null + usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20) + combined = llm_caching_handler.combine_usage(usage_a, usage_b) + assert combined.prompt_tokens_details is None + + # Only first has details + usage_c = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + combined = llm_caching_handler.combine_usage(usage_c, usage_b) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 + + # Only second has details + combined = llm_caching_handler.combine_usage(usage_a, usage_c) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 + + +def _build_logging_obj(call_type: str, stream: bool): + import uuid as _uuid + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + return LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=call_type, + model="gpt-5.4", + messages=[], + function_id=str(_uuid.uuid4()), + stream=stream, + start_time=datetime.now(), + ) + + +def test_convert_cached_responses_bridge_chat_completion_nonstream(): + """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse.""" + from litellm import responses + from litellm.types.utils import CallTypes, ModelResponse + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "chatcmpl-bridge-nonstream", + "object": "chat.completion", + "created": int(time.time()), + "model": "gpt-5.4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}, + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={ + "model": "gpt-5.4", + "stream": False, + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), + model="gpt-5.4", + args=(), + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hi!" + + +def test_convert_cached_responses_legacy_nonstream_path(): + """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path.""" + from litellm import responses + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import CallTypes + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "resp_legacy_nonstream", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_legacy", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "legacy response", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "hi", "stream": False}, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False), + model="gpt-4o", + args=(), + ) + + assert isinstance(result, ResponsesAPIResponse) + assert result.id == "resp_legacy_nonstream" + + +def test_convert_cached_responses_legacy_stream_path(): + """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path.""" + from litellm import responses + from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + ) + from litellm.types.utils import CallTypes + + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + cached_result = { + "id": "resp_legacy_stream", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_legacy_stream", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "legacy stream", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "hi", "stream": True}, + logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True), + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) + + +@pytest.mark.asyncio +async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input(): + """Image-embedding cache hit restores prompt_tokens=0 from the stored value + instead of recomputing a bogus count by tokenizing the base64 input.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # base64-like blob — token_counter over this would return a large nonzero count + image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50 + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens": 0, + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_sums_stored_prompt_tokens_across_items(): + """A multi-item cache hit sums the stored per-item prompt_tokens back to the total.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.01], + "index": 0, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 5, + }, + { + "embedding": [-0.02], + "index": 1, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 4, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-3-small", + ) + + assert cache_hit + assert response.usage.prompt_tokens == 9 + assert response.usage.total_tokens == 9 + + +@pytest.mark.asyncio +async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): + """Legacy cache entries with no stored prompt_tokens still recompute via token_counter + for str inputs (backward compatibility).""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # No prompt_tokens key — pre-fix entry + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "hello world"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + # token_counter over "hello world" yields a nonzero count — fallback path still runs + assert response.usage.prompt_tokens > 0 + + +@pytest.mark.asyncio +async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj(): + """A full embedding cache hit must stamp the resolved provider onto the logging + obj so spend logs record the provider instead of None/unknown.""" + from litellm.types.utils import CallTypes + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 5, + } + ] + + logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False) + logging_obj.async_success_handler = AsyncMock() + + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-3-small", "input": "hello world"}, + logging_obj=logging_obj, + start_time=datetime.now(), + model="text-embedding-3-small", + ) + + assert cache_hit + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + + +def test_sync_stream_responses_cache_hit_sets_custom_llm_provider_on_logging_obj(monkeypatch): + import litellm + from litellm.caching.caching import Cache + from litellm.types.utils import CallTypes + + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + kwargs = {"model": "azure/gpt-5.4-mini", "input": "hello", "stream": True} + cached_response = { + "id": "resp_sync_stream", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-5.4-mini", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_sync_stream", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + } + litellm.cache.add_cache(json.dumps(cached_response), **kwargs) + handler = LLMCachingHandler(original_function=litellm.responses, request_kwargs=kwargs, start_time=datetime.now()) + logging_obj = _build_logging_obj(CallTypes.responses.value, stream=True) + + hit = handler._sync_get_cache( + model="azure/gpt-5.4-mini", + original_function=litellm.responses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.responses.value, + kwargs=kwargs, + args=(), + ) + + assert hit.cached_result is not None + assert logging_obj.model_call_details["custom_llm_provider"] == "azure" + assert logging_obj.model_call_details["litellm_params"]["custom_llm_provider"] == "azure" + + +def test_request_kwargs_does_not_retain_logging_obj(): + """ + The caching handler lives on logging_obj._llm_caching_handler, so keeping + litellm_logging_obj inside request_kwargs closes a reference cycle + (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the + full request payload alive until a generational GC pass instead of being + freed by refcount when the request finishes; under bursts of large-token + requests this presents as stepwise RSS growth that never returns to + baseline. Other kwargs (messages included) must be preserved. + """ + logging_obj = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "litellm_logging_obj": logging_obj, + } + + handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs=kwargs, + start_time=datetime.now(), + ) + + assert "litellm_logging_obj" not in handler.request_kwargs + assert handler.request_kwargs["messages"] == kwargs["messages"] + assert handler.request_kwargs["model"] == "gpt-4o" + + +def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch): + """ + Regression test for the SDK losing async cache writes in short-lived scripts: + async_set_cache dispatched the write as a bare fire-and-forget task, so + asyncio.run cancelled it at loop close before the write landed (LIT-6184, + deterministic with hiredis installed). The write must survive loop shutdown. + """ + import litellm + + writes = [] + + class _SlowWriteCache: + supported_call_types = ["acompletion"] + cache = None + + async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs): + await asyncio.sleep(0.2) + writes.append(result) + + async def acompletion(**kwargs): + return None + + handler = LLMCachingHandler( + original_function=acompletion, + request_kwargs={}, + start_time=datetime.now(), + ) + monkeypatch.setattr(litellm, "cache", _SlowWriteCache()) + + async def _short_lived_script(): + await handler.async_set_cache( + result=litellm.ModelResponse(), + original_function=acompletion, + kwargs={}, + ) + + asyncio.run(_short_lived_script()) + + assert len(writes) == 1 + + +@pytest.mark.asyncio +async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch): + """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again.""" + import litellm + from litellm.caching.caching import Cache + from litellm.types.utils import CallTypes + + async def acompletion(**kwargs): + return None + + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True} + await litellm.cache.async_add_cache( + litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs + ) + handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now()) + logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False) + logging_obj.async_success_handler = AsyncMock() + + hit = await handler._async_get_cache( + model="gpt-5.4", + original_function=acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + args=(), + ) + + assert hit is not None and hit.cached_result is not None + assert handler.preset_cache_key is not None + assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key + assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key + + +@pytest.mark.asyncio +async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch): + import litellm + from litellm.caching.caching import Cache + from litellm.types.utils import CallTypes + + async def aanthropic_messages(**kwargs): + return None + + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + kwargs = { + "model": "claude-sonnet-5", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + "caching": True, + "stream": False, + "_websearch_interception_converted_stream": True, + } + cached_message = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + } + await litellm.cache.async_add_cache(cached_message, **kwargs) + handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now()) + logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False) + logging_obj.async_success_handler = AsyncMock() + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + hit = await handler._async_get_cache( + model="claude-sonnet-5", + original_function=aanthropic_messages, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aanthropic_messages.value, + kwargs=kwargs, + args=(), + ) + + assert hit is not None and hit.cached_result == cached_message + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True + + +@pytest.mark.asyncio +async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch): + import litellm + from litellm.caching.caching import Cache + from litellm.types.utils import CallTypes + + async def acompletion(**kwargs): + return None + + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + kwargs = { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "run the code"}], + "caching": True, + "stream": False, + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_depth": 1, + } + await litellm.cache.async_add_cache( + litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs + ) + handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now()) + logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False) + logging_obj.async_success_handler = AsyncMock() + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + hit = await handler._async_get_cache( + model="gpt-5.6", + original_function=acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + args=(), + ) + + assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse) + assert hit.cached_result.choices[0].message.content == "done" + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True + + +@pytest.mark.asyncio +async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch): + import litellm + from litellm import CustomLLM + from litellm.caching.caching import Cache + from litellm.types.utils import Embedding, EmbeddingResponse + + class RecordingEmbedder(CustomLLM): + provider_inputs: tuple[tuple[str, ...], ...] = () + + async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse: + self.provider_inputs = (*self.provider_inputs, tuple(input)) + return EmbeddingResponse( + model=model, + data=[ + Embedding(embedding=[float(len(text))], index=idx, object="embedding") + for idx, text in enumerate(input) + ], + ) + + embedder = RecordingEmbedder() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}]) + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"]) + monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"]) + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + + await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"]) + await asyncio.gather(*_PENDING_CACHE_WRITES) + mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"] + response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input) + await asyncio.gather(*_PENDING_CACHE_WRITES) + + assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs + assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4] + assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input] + assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit" + + repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input) + + assert len(embedder.provider_inputs) == 2, embedder.provider_inputs + assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input] diff --git a/tests/test_litellm/caching/test_check_and_fix_namespace_none_guard.py b/tests/unit/caching/test_check_and_fix_namespace_none_guard.py similarity index 100% rename from tests/test_litellm/caching/test_check_and_fix_namespace_none_guard.py rename to tests/unit/caching/test_check_and_fix_namespace_none_guard.py diff --git a/tests/test_litellm/caching/test_disk_cache.py b/tests/unit/caching/test_disk_cache.py similarity index 100% rename from tests/test_litellm/caching/test_disk_cache.py rename to tests/unit/caching/test_disk_cache.py diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py similarity index 100% rename from tests/test_litellm/caching/test_dual_cache.py rename to tests/unit/caching/test_dual_cache.py diff --git a/tests/test_litellm/caching/test_embedding_router.py b/tests/unit/caching/test_embedding_router.py similarity index 100% rename from tests/test_litellm/caching/test_embedding_router.py rename to tests/unit/caching/test_embedding_router.py diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/unit/caching/test_evicted_client_closer.py similarity index 100% rename from tests/test_litellm/caching/test_evicted_client_closer.py rename to tests/unit/caching/test_evicted_client_closer.py diff --git a/tests/test_litellm/caching/test_gcs_cache.py b/tests/unit/caching/test_gcs_cache.py similarity index 100% rename from tests/test_litellm/caching/test_gcs_cache.py rename to tests/unit/caching/test_gcs_cache.py diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py similarity index 100% rename from tests/test_litellm/caching/test_in_memory_cache.py rename to tests/unit/caching/test_in_memory_cache.py diff --git a/tests/test_litellm/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py similarity index 100% rename from tests/test_litellm/caching/test_llm_caching_handler.py rename to tests/unit/caching/test_llm_caching_handler.py diff --git a/tests/test_litellm/caching/test_llm_client_cache_e2e.py b/tests/unit/caching/test_llm_client_cache_e2e.py similarity index 100% rename from tests/test_litellm/caching/test_llm_client_cache_e2e.py rename to tests/unit/caching/test_llm_client_cache_e2e.py diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/unit/caching/test_qdrant_semantic_cache.py similarity index 99% rename from tests/test_litellm/caching/test_qdrant_semantic_cache.py rename to tests/unit/caching/test_qdrant_semantic_cache.py index ca7303e4c6d..4f18fb1bca6 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/unit/caching/test_qdrant_semantic_cache.py @@ -1033,7 +1033,7 @@ def test_qdrant_semantic_cache_defaults_embedding_timeout(): @pytest.mark.asyncio async def test_qdrant_async_embedding_truncates_off_the_event_loop(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py similarity index 100% rename from tests/test_litellm/caching/test_redis_cache.py rename to tests/unit/caching/test_redis_cache.py diff --git a/tests/test_litellm/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py similarity index 100% rename from tests/test_litellm/caching/test_redis_cluster_cache.py rename to tests/unit/caching/test_redis_cluster_cache.py diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/unit/caching/test_redis_cluster_node_isolation.py similarity index 100% rename from tests/test_litellm/caching/test_redis_cluster_node_isolation.py rename to tests/unit/caching/test_redis_cluster_node_isolation.py diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/unit/caching/test_redis_connection_pool.py similarity index 100% rename from tests/test_litellm/caching/test_redis_connection_pool.py rename to tests/unit/caching/test_redis_connection_pool.py diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/unit/caching/test_redis_semantic_cache.py similarity index 99% rename from tests/test_litellm/caching/test_redis_semantic_cache.py rename to tests/unit/caching/test_redis_semantic_cache.py index de253b4f10b..461689165bb 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/unit/caching/test_redis_semantic_cache.py @@ -1392,7 +1392,7 @@ def test_redis_semantic_cache_defaults_embedding_timeout(): @pytest.mark.asyncio async def test_redis_async_embedding_truncates_off_the_event_loop(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/unit/caching/test_s3_cache.py similarity index 100% rename from tests/test_litellm/caching/test_s3_cache.py rename to tests/unit/caching/test_s3_cache.py diff --git a/tests/test_litellm/caching/test_valkey_semantic_cache.py b/tests/unit/caching/test_valkey_semantic_cache.py similarity index 100% rename from tests/test_litellm/caching/test_valkey_semantic_cache.py rename to tests/unit/caching/test_valkey_semantic_cache.py diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/__init__.py b/tests/unit/expected_responses_api_request/__init__.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/audio_utils/__init__.py rename to tests/unit/expected_responses_api_request/__init__.py diff --git a/tests/test_litellm/expected_responses_api_request/azure_shell_tool.json b/tests/unit/expected_responses_api_request/azure_shell_tool.json similarity index 100% rename from tests/test_litellm/expected_responses_api_request/azure_shell_tool.json rename to tests/unit/expected_responses_api_request/azure_shell_tool.json diff --git a/tests/test_litellm/expected_responses_api_request/context_management_and_shell.json b/tests/unit/expected_responses_api_request/context_management_and_shell.json similarity index 100% rename from tests/test_litellm/expected_responses_api_request/context_management_and_shell.json rename to tests/unit/expected_responses_api_request/context_management_and_shell.json diff --git a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py index e66cd654f93..d7a1d6f14e1 100644 --- a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py @@ -528,7 +528,7 @@ async def test_pre_call_hook_no_compression_records_no_savings(monkeypatch): @pytest.mark.asyncio async def test_pre_call_hook_counts_tokens_off_the_event_loop(): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/unit/litellm_core_utils/conftest.py b/tests/unit/litellm_core_utils/conftest.py new file mode 100644 index 00000000000..2a1e1f6382c --- /dev/null +++ b/tests/unit/litellm_core_utils/conftest.py @@ -0,0 +1,15 @@ +import importlib + +import pytest + +from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault + + +@pytest.fixture(autouse=True, scope="session") +def bundled_tiktoken_cache() -> None: + importlib.import_module("litellm.litellm_core_utils.default_encoding") + + +@pytest.fixture +def secret_vault_factory() -> type[FakeSecretVault]: + return FakeSecretVault diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/unit/litellm_core_utils/event_loop_lag.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/event_loop_lag.py rename to tests/unit/litellm_core_utils/event_loop_lag.py diff --git a/tests/unit/litellm_core_utils/fake_secret_vault.py b/tests/unit/litellm_core_utils/fake_secret_vault.py new file mode 100644 index 00000000000..75e9d16e9ed --- /dev/null +++ b/tests/unit/litellm_core_utils/fake_secret_vault.py @@ -0,0 +1,67 @@ +from litellm.litellm_core_utils.cli_keyring import ( + KeyringDiscardsWrites, + KeyringUnreachable, + KeyringUnusable, + SecretErase, + SecretErased, + SecretFound, + SecretMissing, + SecretRead, + SecretStored, + SecretStranded, + SecretWrite, +) + + +class FakeSecretVault: + """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised. + + `available=False` models a keychain that is locked or has no backend, `writable=False` one that + refuses to store, `erasable=False` one that will not release what it already holds, and `failure` + picks which unusable state those report. `discards=True` is keyring's null backend, which answers + reads and erases like any other yet keeps nothing it is given, so only writes report it. + """ + + def __init__( + self, + blob: str | None = None, + *, + available: bool = True, + writable: bool = True, + erasable: bool = True, + discards: bool = False, + failure: KeyringUnusable = KeyringUnreachable(), + ) -> None: + self.blob: str | None = blob + self.available: bool = available + self.writable: bool = writable + self.erasable: bool = erasable + self.discards: bool = discards + self.failure: KeyringUnusable = failure + self.reads: int = 0 + self.writes: list[str] = [] + self.erases: int = 0 + + def read(self) -> SecretRead: + self.reads += 1 + if not self.available: + return self.failure + return SecretMissing() if self.blob is None else SecretFound(self.blob) + + def write(self, blob: str) -> SecretWrite: + self.writes.append(blob) + if not (self.available and self.writable): + return self.failure + if self.discards: + return KeyringDiscardsWrites() + self.blob = blob + return SecretStored() + + def erase(self) -> SecretErase: + self.erases += 1 + if not self.available: + return self.failure + if not self.erasable: + return SecretStranded() if self.blob is not None else SecretErased() + self.blob = None + return SecretErased() diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/__init__.py b/tests/unit/litellm_core_utils/llm_cost_calc/__init__.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/__init__.py rename to tests/unit/litellm_core_utils/llm_cost_calc/__init__.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py rename to tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py rename to tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py diff --git a/tests/test_litellm/litellm_core_utils/messages_with_counts.py b/tests/unit/litellm_core_utils/messages_with_counts.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/messages_with_counts.py rename to tests/unit/litellm_core_utils/messages_with_counts.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/__init__.py b/tests/unit/litellm_core_utils/prompt_templates/__init__.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/__init__.py rename to tests/unit/litellm_core_utils/prompt_templates/__init__.py diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/unit/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py rename to tests/unit/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py diff --git a/tests/test_litellm/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 similarity index 100% rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py diff --git a/tests/test_litellm/rust_bridge/__init__.py b/tests/unit/litellm_core_utils/specialty_caches/__init__.py similarity index 100% rename from tests/test_litellm/rust_bridge/__init__.py rename to tests/unit/litellm_core_utils/specialty_caches/__init__.py diff --git a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py b/tests/unit/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py rename to tests/unit/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py diff --git a/tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py b/tests/unit/litellm_core_utils/test_agentic_followup_kwargs.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py rename to tests/unit/litellm_core_utils/test_agentic_followup_kwargs.py diff --git a/tests/test_litellm/litellm_core_utils/test_anthropic_dedup_factory.py b/tests/unit/litellm_core_utils/test_anthropic_dedup_factory.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_anthropic_dedup_factory.py rename to tests/unit/litellm_core_utils/test_anthropic_dedup_factory.py diff --git a/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py b/tests/unit/litellm_core_utils/test_api_route_to_call_types.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py rename to tests/unit/litellm_core_utils/test_api_route_to_call_types.py diff --git a/tests/test_litellm/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_audio_utils.py rename to tests/unit/litellm_core_utils/test_audio_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/unit/litellm_core_utils/test_aws_partition.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_aws_partition.py rename to tests/unit/litellm_core_utils/test_aws_partition.py diff --git a/tests/test_litellm/litellm_core_utils/test_bedrock_converse_dedup_factory.py b/tests/unit/litellm_core_utils/test_bedrock_converse_dedup_factory.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_bedrock_converse_dedup_factory.py rename to tests/unit/litellm_core_utils/test_bedrock_converse_dedup_factory.py diff --git a/tests/test_litellm/litellm_core_utils/test_bug_report.py b/tests/unit/litellm_core_utils/test_bug_report.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_bug_report.py rename to tests/unit/litellm_core_utils/test_bug_report.py diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py rename to tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py diff --git a/tests/test_litellm/litellm_core_utils/test_classifier_logging.py b/tests/unit/litellm_core_utils/test_classifier_logging.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_classifier_logging.py rename to tests/unit/litellm_core_utils/test_classifier_logging.py diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/unit/litellm_core_utils/test_cli_token_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_cli_token_utils.py rename to tests/unit/litellm_core_utils/test_cli_token_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py rename to tests/unit/litellm_core_utils/test_cloud_storage_security.py diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/unit/litellm_core_utils/test_codestral_provider_routing.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py rename to tests/unit/litellm_core_utils/test_codestral_provider_routing.py diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_core_helpers.py rename to tests/unit/litellm_core_utils/test_core_helpers.py diff --git a/tests/test_litellm/litellm_core_utils/test_coroutine_checker.py b/tests/unit/litellm_core_utils/test_coroutine_checker.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_coroutine_checker.py rename to tests/unit/litellm_core_utils/test_coroutine_checker.py diff --git a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py b/tests/unit/litellm_core_utils/test_dd_tracing.py similarity index 85% rename from tests/test_litellm/litellm_core_utils/test_dd_tracing.py rename to tests/unit/litellm_core_utils/test_dd_tracing.py index b55ade5225d..30cae45e250 100644 --- a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py +++ b/tests/unit/litellm_core_utils/test_dd_tracing.py @@ -55,18 +55,6 @@ def test_dd_tracer_when_package_not_exists(): assert result == "test" -def test_null_tracer_context_manager(): - """ - Test that the context manager works without raising exceptions when should_use_dd_tracer is False - """ - with patch("litellm.litellm_core_utils.dd_tracing.should_use_dd_tracer", False): - # Test that the context manager works without raising exceptions - with dd_tracer.trace("test_operation") as span: - # Test that we can call methods on the null span - span.finish() - assert True # If we get here without exceptions, the test passes - - def test_should_use_dd_tracer(): """ Test that the should_use_dd_tracer function works as expected diff --git a/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py b/tests/unit/litellm_core_utils/test_decode_special_tokens.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py rename to tests/unit/litellm_core_utils/test_decode_special_tokens.py diff --git a/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py rename to tests/unit/litellm_core_utils/test_dot_notation_indexing.py diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_duration_parser.py rename to tests/unit/litellm_core_utils/test_duration_parser.py diff --git a/tests/test_litellm/litellm_core_utils/test_error_normalization.py b/tests/unit/litellm_core_utils/test_error_normalization.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_error_normalization.py rename to tests/unit/litellm_core_utils/test_error_normalization.py diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py rename to tests/unit/litellm_core_utils/test_exception_mapping_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/unit/litellm_core_utils/test_extract_base64_image.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_extract_base64_image.py rename to tests/unit/litellm_core_utils/test_extract_base64_image.py diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py rename to tests/unit/litellm_core_utils/test_fallback_generalizations.py diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_utils.py b/tests/unit/litellm_core_utils/test_fallback_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_fallback_utils.py rename to tests/unit/litellm_core_utils/test_fallback_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_get_litellm_params.py rename to tests/unit/litellm_core_utils/test_get_litellm_params.py diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/unit/litellm_core_utils/test_get_llm_provider_endpoint_match.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py rename to tests/unit/litellm_core_utils/test_get_llm_provider_endpoint_match.py diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py rename to tests/unit/litellm_core_utils/test_get_llm_provider_logic.py diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py rename to tests/unit/litellm_core_utils/test_get_model_cost_map.py diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/unit/litellm_core_utils/test_get_supported_openai_params.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py rename to tests/unit/litellm_core_utils/test_get_supported_openai_params.py diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_health_check_helpers.py rename to tests/unit/litellm_core_utils/test_health_check_helpers.py diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_image_handling.py rename to tests/unit/litellm_core_utils/test_image_handling.py diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py rename to tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py diff --git a/tests/test_litellm/litellm_core_utils/test_internal_call_metadata.py b/tests/unit/litellm_core_utils/test_internal_call_metadata.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_internal_call_metadata.py rename to tests/unit/litellm_core_utils/test_internal_call_metadata.py diff --git a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py b/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py rename to tests/unit/litellm_core_utils/test_json_fragment_accumulator.py diff --git a/tests/test_litellm/litellm_core_utils/test_json_schema_validation.py b/tests/unit/litellm_core_utils/test_json_schema_validation.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_json_schema_validation.py rename to tests/unit/litellm_core_utils/test_json_schema_validation.py diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_litellm_logging.py rename to tests/unit/litellm_core_utils/test_litellm_logging.py diff --git a/tests/test_litellm/litellm_core_utils/test_llm_judge.py b/tests/unit/litellm_core_utils/test_llm_judge.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_llm_judge.py rename to tests/unit/litellm_core_utils/test_llm_judge.py diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_llm_request_utils.py rename to tests/unit/litellm_core_utils/test_llm_request_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/unit/litellm_core_utils/test_logging_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_logging_utils.py rename to tests/unit/litellm_core_utils/test_logging_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_logging_worker.py rename to tests/unit/litellm_core_utils/test_logging_worker.py diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py rename to tests/unit/litellm_core_utils/test_max_streaming_duration.py diff --git a/tests/test_litellm/litellm_core_utils/test_model_param_helper.py b/tests/unit/litellm_core_utils/test_model_param_helper.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_model_param_helper.py rename to tests/unit/litellm_core_utils/test_model_param_helper.py diff --git a/tests/test_litellm/litellm_core_utils/test_model_response_utils.py b/tests/unit/litellm_core_utils/test_model_response_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_model_response_utils.py rename to tests/unit/litellm_core_utils/test_model_response_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_private_json.py b/tests/unit/litellm_core_utils/test_private_json.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_private_json.py rename to tests/unit/litellm_core_utils/test_private_json.py diff --git a/tests/test_litellm/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_provider_affinity.py rename to tests/unit/litellm_core_utils/test_provider_affinity.py diff --git a/tests/test_litellm/litellm_core_utils/test_provider_specific_headers.py b/tests/unit/litellm_core_utils/test_provider_specific_headers.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_provider_specific_headers.py rename to tests/unit/litellm_core_utils/test_provider_specific_headers.py diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/unit/litellm_core_utils/test_ptu_pricing.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_ptu_pricing.py rename to tests/unit/litellm_core_utils/test_ptu_pricing.py diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/unit/litellm_core_utils/test_realtime_errors.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_realtime_errors.py rename to tests/unit/litellm_core_utils/test_realtime_errors.py diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_realtime_streaming.py rename to tests/unit/litellm_core_utils/test_realtime_streaming.py diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_redact_messages.py rename to tests/unit/litellm_core_utils/test_redact_messages.py diff --git a/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py b/tests/unit/litellm_core_utils/test_request_timeout_resolver.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py rename to tests/unit/litellm_core_utils/test_request_timeout_resolver.py diff --git a/tests/test_litellm/litellm_core_utils/test_retry_after_headers.py b/tests/unit/litellm_core_utils/test_retry_after_headers.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_retry_after_headers.py rename to tests/unit/litellm_core_utils/test_retry_after_headers.py diff --git a/tests/test_litellm/litellm_core_utils/test_safe_divide_seconds.py b/tests/unit/litellm_core_utils/test_safe_divide_seconds.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_safe_divide_seconds.py rename to tests/unit/litellm_core_utils/test_safe_divide_seconds.py diff --git a/tests/test_litellm/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_safe_json_dumps.py rename to tests/unit/litellm_core_utils/test_safe_json_dumps.py diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/unit/litellm_core_utils/test_sensitive_data_masker.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py rename to tests/unit/litellm_core_utils/test_sensitive_data_masker.py diff --git a/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py b/tests/unit/litellm_core_utils/test_sentry_scrubbing.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py rename to tests/unit/litellm_core_utils/test_sentry_scrubbing.py diff --git a/tests/test_litellm/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_served_output_texts.py rename to tests/unit/litellm_core_utils/test_served_output_texts.py diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_cursor.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_cursor.py diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py similarity index 99% rename from tests/test_litellm/litellm_core_utils/test_streaming_handler.py rename to tests/unit/litellm_core_utils/test_streaming_handler.py index 3af79c709cc..6557811b530 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -4900,7 +4900,7 @@ class TestStableStreamingResponseId: @pytest.mark.asyncio async def test_async_stream_without_usage_counts_tokens_off_the_event_loop(): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_overhead.py b/tests/unit/litellm_core_utils/test_streaming_overhead.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_streaming_overhead.py rename to tests/unit/litellm_core_utils/test_streaming_overhead.py diff --git a/tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py b/tests/unit/litellm_core_utils/test_thread_pool_executor.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py rename to tests/unit/litellm_core_utils/test_thread_pool_executor.py diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py new file mode 100644 index 00000000000..b1a14e61b96 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -0,0 +1,1441 @@ +#### What this tests #### +# This tests litellm.token_counter.token_counter() function +import asyncio +import base64 +import importlib +import threading +import time +import traceback +from concurrent.futures import Future, wait +from typing import Final +from unittest.mock import MagicMock + +import anyio.to_thread +import pytest +import tiktoken + +from unittest.mock import AsyncMock, patch + +import litellm +from litellm import decode, encode, get_modified_max_tokens +from litellm import token_counter as token_counter_old +import litellm.constants +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import ( + _get_exact_count_function, + _get_extrapolating_count_function, + _get_tiktoken_count_function, + calculate_img_tokens, + high_detail_image_token_upper_bound, + offload_token_count, +) +from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new +from tests.large_text import text +from tests.unit.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, +) +from tests.unit.litellm_core_utils.messages_with_counts import ( + MESSAGES_TEXT, + MESSAGES_WITH_IMAGES, + MESSAGES_WITH_TOOLS, +) + + +def token_counter_both_assert_same(**args): + new = token_counter_new(**args) + old = token_counter_old(**args) + assert new == old, f"New token counter {new} does not match old token counter {old}" + return new + + +## Choose which token_counter the test will use. + +# token_counter = token_counter_new +# token_counter = token_counter_old +token_counter = token_counter_both_assert_same + + +def test_token_counter_basic(): + assert ( + token_counter( + model="claude-2", + messages=[ + { + "role": "user", + "content": "This is a long message that definitely exceeds the token limit.", + } + ], + ) + == 19 + ) + + +def test_token_counter_large_repeated_text_is_fast(): + messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}] + + start_time = time.perf_counter() + tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages) + elapsed = time.perf_counter() - start_time + + assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s" + assert tokens > 0 + + +@pytest.mark.parametrize( + "text", + [ + "Short text", + "This is a normal message with punctuation, numbers, and a few words.", + ], +) +def test_token_counter_short_text_matches_tiktoken(text): + encoding = tiktoken.get_encoding("cl100k_base") + expected = len(encoding.encode(text, disallowed_special=())) + + assert token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) == expected + + +def test_token_counter_default_encoding_matches_cl100k(): + encoding: Final = tiktoken.get_encoding("cl100k_base") + expected: Final = len(encoding.encode("hello world", disallowed_special=())) + + assert token_counter_new(model=None, text="hello world") == expected + + +def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken(): + text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025] + encoding = tiktoken.get_encoding("cl100k_base") + expected = len(encoding.encode(text, disallowed_special=())) + + actual = token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) + + assert abs(actual - expected) <= 4 + + +@pytest.mark.parametrize( + "configured", + ["0", "-1", "-1024", "not-an-int", "", " ", "999999999", "inf", "1e9"], +) +def test_invalid_chunk_size_config_stays_usable(monkeypatch, configured): + """A misconfigured chunk size must not raise, count zero, or restore the quadratic encode cost.""" + monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", configured) + try: + reloaded = importlib.reload(litellm.constants) + chunk_size = reloaded.TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS + assert 1 <= chunk_size <= reloaded.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS + + encoding = tiktoken.get_encoding("cl100k_base") + count_tokens = _get_tiktoken_count_function( + lambda text: len(encoding.encode(text, disallowed_special=())), + chunk_size=chunk_size, + ) + assert count_tokens("The quick brown fox jumps over the lazy dog. " * 40) > 0 + finally: + monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS") + importlib.reload(litellm.constants) + + +def test_valid_chunk_size_config_is_honoured(monkeypatch): + monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", "2048") + try: + assert importlib.reload(litellm.constants).TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS == 2048 + finally: + monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS") + importlib.reload(litellm.constants) + + +async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): + warm_tokenizer("claude-fable-5") + + tokens, took, lags = await timed_with_loop_lags( + lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + ) + + assert tokens > 0 + assert_loop_stayed_free(took, lags) + + +@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): + count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) + front_heavy: Final = "a" * 1_000 + "b" * 4_000 + exact: Final = 1_000 + len(front_heavy) + + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) + + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars + + +def test_count_at_or_below_the_cap_is_exact(): + count_exactly: Final = MagicMock(side_effect=len) + + assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 + assert count_exactly.call_args_list == [(("a" * 5_000,),)] + + +class _SlowEncoder: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + time.sleep(0.1) + with self._lock: + self.in_flight -= 1 + return [[0] * len(text) for text in texts] + + +@pytest.mark.asyncio +async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): + encoder: Final = _SlowEncoder() + count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) + shared_pool: Final = anyio.to_thread.current_default_thread_limiter() + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: + if counting.done(): + return () + await asyncio.sleep(0.01) + return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) + + counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) + borrowed: Final = await shared_pool_borrowed_until_done(counting) + + assert await counting == [3] * burst + assert len(borrowed) > 1 and max(borrowed) == 0 + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + +def test_token_counter_applies_the_default_cap(): + max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS + prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prose + "a" * 200_000 + exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) + + estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) + + assert estimate != exact + assert abs(estimate - exact) <= exact // 100 + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], +) +def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") + importlib.reload(litellm.constants) + + +def test_token_counter_with_prefix(): + messages = [ + {"role": "user", "content": "Who won the world cup in 2022?"}, + {"role": "assistant", "content": "Argentina", "prefix": True}, + ] + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens == 22, f"Expected 22 tokens, got {tokens}" + + +def test_token_counter_normal_plus_function_calling(): + messages = [ + {"role": "system", "content": "System prompt"}, + {"role": "user", "content": "content1"}, + {"role": "assistant", "content": "content2"}, + {"role": "user", "content": "conten3"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_E0lOb1h6qtmflUyok4L06TgY", + "function": { + "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}', + "name": "SearchInternet", + }, + "type": "function", + } + ], + }, + { + "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY", + "role": "tool", + "name": "SearchInternet", + "content": "tool content", + }, + ] + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens == 80 + + +# test_token_counter_normal_plus_function_calling() + + +def test_token_counter_legacy_function_call_counts_arguments(): + """ + Regression for VERIA-492 (Token-counter function_call bypass). + + The legacy OpenAI assistant `function_call` field carries arbitrary text in + `arguments`. Before the fix, `_count_messages` had no branch for + `function_call` and fell through to the unsupported-key `continue`, so an + assistant turn could smuggle unlimited text past `token_counter` and the + proxy `/utils/token_counter` endpoint (and downstream pre-call budget / + `get_modified_max_tokens` math). After the fix it must be counted the + same as the equivalent `tool_calls` payload. + """ + long_arg = "A" * 4000 + fc_messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "function_call": {"name": "search", "arguments": long_arg}, + }, + ] + tc_messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": long_arg}, + } + ], + }, + ] + fc_tokens = token_counter(model="gpt-3.5-turbo", messages=fc_messages) + tc_tokens = token_counter(model="gpt-3.5-turbo", messages=tc_messages) + assert fc_tokens == tc_tokens, ( + f"function_call arguments must count like tool_calls arguments; " + f"got function_call={fc_tokens}, tool_calls={tc_tokens}" + ) + assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}" + + +@pytest.mark.parametrize( + "message_count_pair", + MESSAGES_TEXT, +) +def test_token_counter_textonly(message_count_pair): + counted_tokens = token_counter( + model="gpt-35-turbo", messages=[message_count_pair["message"]] + ) + assert counted_tokens == message_count_pair["count"] + + +@pytest.mark.parametrize( + "message_count_pair", + MESSAGES_TEXT, +) +def test_token_counter_count_response_tokens(message_count_pair): + counted_tokens = token_counter( + model="gpt-35-turbo", + messages=[message_count_pair["message"]], + count_response_tokens=True, + ) + # 3 tokens are not added because of count_response_tokens=True + expected = message_count_pair["count"] - 3 + assert counted_tokens == expected + + +@pytest.mark.parametrize( + "message_count_pair", + MESSAGES_WITH_IMAGES, +) +def test_token_counter_with_images(message_count_pair): + counted_tokens = token_counter( + model="gpt-4o", messages=[message_count_pair["message"]] + ) + assert counted_tokens == message_count_pair["count"] + + +@pytest.mark.parametrize( + "message_count_pair", + MESSAGES_WITH_TOOLS, +) +def test_token_counter_with_tools(message_count_pair): + counted_tokens = token_counter( + model="gpt-35-turbo", + messages=[message_count_pair["system_message"]], + tools=message_count_pair["tools"], + tool_choice=message_count_pair["tool_choice"], + ) + expected_tokens = message_count_pair["count"] + actual_diff = counted_tokens - expected_tokens + + if "count-tolerate" in message_count_pair: + if message_count_pair["count-tolerate"] == counted_tokens: + pass # expected + else: + tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens + assert ( + actual_diff <= tolerated_diff + ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction." + if actual_diff != tolerated_diff: + raise NeedsToleranceUpdateError( + f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}" + ) + + else: + assert ( + expected_tokens == counted_tokens + ), f"Expected {expected_tokens} tokens, got {counted_tokens}." + + +class NeedsToleranceUpdateError(Exception): + """Custom exception to mark tests that have improved""" + + pass + + +# test_tokenizers() + + +def test_encoding_and_decoding(): + try: + sample_text = "Hellö World, this is my input string!" + # openai encoding + decoding + openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) + openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) + + assert openai_text == sample_text + + # claude encoding + decoding + claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) + + claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) + + assert claude_text == sample_text + + # cohere encoding + decoding + cohere_tokens = encode(model="command-nightly", text=sample_text) + cohere_text = decode(model="command-nightly", tokens=cohere_tokens) + + assert cohere_text == sample_text + + # llama2 encoding + decoding + llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text) + llama2_text = decode(model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens) + + assert llama2_text == sample_text + except Exception as e: + pytest.fail(f"An exception occured: {e}\n{traceback.format_exc()}") + + +# test_encoding_and_decoding() + + +# test_gpt_vision_token_counting() + + +@pytest.mark.parametrize( + "model", + [ + "gpt-4-vision-preview", + "gpt-4o", + "claude-3-opus-20240229", + "command-nightly", + "mistral/mistral-tiny", + ], +) +def test_load_test_token_counter(model): + """ + Token count large prompt 100 times. + + Assert time taken is < 1.5s. + """ + import tiktoken + + messages = [{"role": "user", "content": text}] * 10 + + start_time = time.time() + for _ in range(10): + _ = token_counter(model=model, messages=messages) + # enc.encode("".join(m["content"] for m in messages)) + + end_time = time.time() + + total_time = end_time - start_time + print("model={}, total test time={}".format(model, total_time)) + assert total_time < 10, f"Total encoding time > 10s, {total_time}" + + +@pytest.mark.parametrize( + "model, base_model, input_tokens, user_max_tokens, expected_value", + [ + ("random-model", "random-model", 1024, 1024, 1024), + ("gpt-3.5-turbo", "gpt-3.5-turbo", 4000, 5000, 4096), # model max output = 4096 + ], +) +def test_get_modified_max_tokens( + model, base_model, input_tokens, user_max_tokens, expected_value +): + """ + - Test when max_output is not known => expect user_max_tokens + - Test when max_output == max_input, + - input > max_output, no max_tokens => expect None + - input + max_tokens > max_output => expect remainder + - input + max_tokens < max_output => expect max_tokens + - Test when max_tokens > max_output => expect max_output + """ + args = locals() + import litellm + + litellm.token_counter = MagicMock() + + def _mock_token_counter(*args, **kwargs): + return input_tokens + + litellm.token_counter.side_effect = _mock_token_counter + print(f"_mock_token_counter: {_mock_token_counter()}") + messages = [{"role": "user", "content": "Hello world!"}] + + calculated_value = get_modified_max_tokens( + model=model, + base_model=base_model, + messages=messages, + user_max_tokens=user_max_tokens, + buffer_perc=0, + buffer_num=0, + ) + + if expected_value is None: + assert calculated_value is None + else: + assert ( + calculated_value == expected_value + ), "Got={}, Expected={}, Params={}".format( + calculated_value, expected_value, args + ) + + +def test_empty_tools(): + messages = [{"role": "user", "content": "hey, how's it going?", "tool_calls": None}] + + result = token_counter( + messages=messages, + ) + + print(result) + + +@pytest.mark.skip( + reason="Skipping this test temporarily because it relies on a function being called that I am removing." +) +def test_gpt_4o_token_counter(): + with patch.object( + litellm.utils, "openai_token_counter", new=MagicMock() + ) as mock_client: + token_counter( + model="gpt-4o-2024-05-13", messages=[{"role": "user", "content": "Hey!"}] + ) + + mock_client.assert_called() + + +@pytest.mark.parametrize( + "img_url", + [ + "https://example.com/test-image.png", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAL0AAAC9CAMAAADRCYwCAAAAh1BMVEX///8AAAD8/Pz5+fkEBAT39/cJCQn09PRNTU3y8vIMDAwzMzPe3t7v7+8QEBCOjo7FxcXR0dHn5+elpaWGhoYYGBivr686OjocHBy0tLQtLS1TU1PY2Ni6urpaWlpERER3d3ecnJxoaGiUlJRiYmIlJSU4ODhBQUFycnKAgIDBwcFnZ2chISE7EjuwAAAI/UlEQVR4nO1caXfiOgz1bhJIyAJhX1JoSzv8/9/3LNlpYd4rhX6o4/N8Z2lKM2cURZau5JsQEhERERERERERERERERERERHx/wBjhDPC3OGN8+Cc5JeMuheaETSdO8vZFyCScHtmz2CsktoeMn7rLM1u3h0PMAEhyYX7v/Q9wQvoGdB0hlbzm45lEq/wd6y6G9aezvBk9AXwp1r3LHJIRsh6s2maxaJpmvqgvkC7WFS3loUnaFJtKRVUCEoV/RpCnHRvAsesVQ1hw+vd7Mpo+424tLs72NplkvQgcdrsvXkW/zJWqH/fA0FT84M/xnQJt4to3+ZLuanbM6X5lfXKHosO9COgREqpCR5i86pf2zPS7j9tTj+9nO7bQz3+xGEyGW9zqgQ1tyQ/VsxEDvce/4dcUPNb5OD9yXvR4Z2QisuP0xiGWPnemgugU5q/troHhGEjIF5sTOyW648aC0TssuaaCEsYEIkGzjWXOp3A0vVsf6kgRyqaDk+T7DIVWrb58b2tT5xpUucKwodOD/5LbrZC1ws6YSaBZJ/8xlh+XZSYXaMJ2ezNqjB3IPXuehPcx2U6b4t1dS/xNdFzguUt8ie7arnPeyCZroxLHzGgGdqVcspwafizPWEXBee+9G1OaufGdvNng/9C+gwgZ3PH3r87G6zXTZ5D5De2G2DeFoANXfbACkT+fxBQ22YFsTTJF9hjFVO6VbqxZXko4WJ8s52P4PnuxO5KRzu0/hlix1ySt8iXjgaQ+4IHPA9nVzNkdduM9LFT/Aacj4FtKrHA7iAw602Vnht6R8Vq1IOS+wNMKLYqayAYfRuufQPGeGb7sZogQQoLZrGPgZ6KoYn70Iw30O92BNEDpvwouCFn6wH2uS+EhRb3WF/HObZk3HuxfRQM3Y/Of/VH0n4MKNHZDiZvO9+m/ABALfkOcuar/7nOo7B95ACGVAFaz4jMiJwJhdaHBkySmzlGTu82gr6FSTik2kJvLnY9nOd/D90qcH268m3I/cgI1xg1maE5CuZYaWLH+UHANCIck0yt7Mx5zBm5vVHXHwChsZ35kKqUpmo5Svq5/fzfAI5g2vDtFPYo1HiEA85QrDeGm9g//LG7K0scO3sdpj2CBDgCa+0OFs0bkvVgnnM/QBDwllOMm+cN7vMSHlB7Uu4haHKaTwgGkv8tlK+hP8fzmFuK/RQTpaLPWvbd58yWIo66HHM0OsPoPhVqmtaEVL7N+wYcTLTbb0DLdgp23Eyy2VYJ2N7bkLFAAibtoLPe5sLt6Oa2bvU+zyeMa8wrixO0gRTn9tO9NCSThTLGqcqtsDvphlfmx/cPBZVvw24jg1LE2lPuEo35Mhi58U0I/Ga8n5w+NS8i34MAQLos5B1u0xL1ZvCVYVRw/Fs2q53KLaXJMWwOZZ/4MPYV19bAHmgGDKB6f01xoeJKFbl63q9J34KdaVNPJWztQyRkzA3KNs1AdAEDowMxh10emXTCx75CkurtbY/ZpdNDGdsn2UcHKHsQ8Ai3WZi48IfkvtjOhsLpuIRSKZTX9FA4o+0d6o/zOWqQzVJMynL9NsxhSJOaourq6nBVQBueMSyubsX2xHrmuABZN2Ns9jr5nwLFlLF/2R6atjW/67Yd11YQ1Z+kA9Zk9dPTM/o6dVo6HHVgC0JR8oUfmI93T9u3gvTG94bAH02Y5xeqRcjuwnKCK6Q2+ajl8KXJ3GSh22P3Zfx6S+n008ROhJn+JRIUVu6o7OXl8w1SeyhuqNDwNI7SjbK08QrqPxS95jy4G7nCXVq6G3HNu0LtK5J0e226CfC005WKK9sVvfxI0eUbcnzutfhWe3rpZHM0nZ/ny/N8tanKYlQ6VEW5Xuym8yV1zZX58vwGhZp/5tFfhybZabdbrQYOs8F+xEhmPsb0/nki6kIyVvzZzUASiOrTfF+Sj9bXC7DoJxeiV8tjQL6loSd0yCx7YyB6rPdLx31U2qCG3F/oXIuDuqd6LFO+4DNIJuxFZqSsU0ea88avovFnWKRYFYRQDfCfcGaBCLn4M4A1ntJ5E57vicwqq2enaZEF5nokCYu9TbKqCC5yCDfL+GhLxT4w4xEJs+anqgou8DOY2q8FMryjb2MehC1dRJ9s4g9NXeTwPkWON4RH+FhIe0AWR/S9ekvQ+t70XHeimGF78LzuU7d7PwrswdIG2VpgF8C53qVQsTDtBJc4CdnkQPbnZY9mbPdDFra3PCXBBQ5QBn2aQqtyhvlyYM4Hb2/mdhsxCUen04GZVvIJZw5PAamMOmjzq8Q+dzAKLXDQ3RUZItWsg4t7W2DP+JDrJDymoMH7E5zQtuEpG03GTIjGCW3LQqOYEsXgFc78x76NeRwY6SNM+IfQoh6myJKRBIcLYxZcwscJ/gI2isTBty2Po9IkYzP0/SS4hGlxRjFAG5z1Jt1LckiB57yWvo35EaolbvA+6fBa24xodL2YjsPpTnj3JgJOqhcgOeLVsYYwoK0wjY+m1D3rGc40CukkaHnkEjarlXrF1B9M6ECQ6Ow0V7R7N4G3LfOHAXtymoyXOb4QhaYHJ/gNBJUkxclpSs7DNcgWWDDmM7Ke5MJpGuioe7w5EOvfTunUKRzOh7G2ylL+6ynHrD54oQO3//cN3yVO+5qMVsPZq0CZIOx4TlcJ8+Vz7V5waL+7WekzUpRFMTnnTlSCq3X5usi8qmIleW/rit1+oQZn1WGSU/sKBYEqMNh1mBOc6PhK8yCfKHdUNQk8o/G19ZPTs5MYfai+DLs5vmee37zEyyH48WW3XA6Xw6+Az8lMhci7N/KleToo7PtTKm+RA887Kqc6E9dyqL/QPTugzMHLbLZtJKqKLFfzVWRNJ63c+95uWT/F7R0U5dDVvuS409AJXhJvD0EwWaWdW8UN11u/7+umaYjT8mJtzZwP/MD4r57fihiHlC5fylHfaqnJdro+Dr7DajvO+vi2EwyD70s8nCH71nzIO1l5Zl+v1DMCb5ebvCMkGHvobXy/hPumGLyX0218/3RyD1GRLOuf9u/OGQyDmto32yMiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIv7GP8YjWPR/czH2AAAAAElFTkSuQmCC", + ], +) +def test_img_url_token_counter(img_url, monkeypatch): + """ + Verify get_image_dimensions returns valid (width, height) for both an + HTTPS URL and a base64 data URI. The HTTPS branch is exercised with a + mocked HTTP fetch so the test is hermetic - it can't break when a + third-party image URL goes away. + """ + import base64 + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + # Minimal valid 1x1 PNG, served by the mocked safe_get for the URL case. + _tiny_png = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + ) + + if img_url.startswith(("http://", "https://")): + + class _FakeResponse: + headers = {"Content-Length": str(len(_tiny_png))} + + def read(self): + return _tiny_png + + monkeypatch.setattr( + "litellm.litellm_core_utils.token_counter.safe_get", + lambda client, url, **kw: _FakeResponse(), + ) + + width, height = get_image_dimensions(data=img_url) + + print(width, height) + + assert width is not None + assert height is not None + + +def test_token_encode_disallowed_special(): + encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") + token_counter(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") + + +def test_token_counter(): + try: + messages = [{"role": "user", "content": "hi how are you what time is it"}] + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + print("gpt-35-turbo") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="claude-2", messages=messages) + print("claude-2") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="gemini/chat-bison", messages=messages) + print("gemini/chat-bison") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="ollama/llama2", messages=messages) + print("ollama/llama2") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages) + print("anthropic.claude-instant-v1") + print(tokens) + assert tokens > 0 + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +import unittest + +from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding + +# Clear the cache at module load to ensure clean state +_load_huggingface_tokenizer.cache_clear() + + +class TestTokenizerSelection(unittest.TestCase): + def setUp(self): + """Clear the LRU cache before each test method. + + The HuggingFace tokenizers behind _select_tokenizer_helper are cached with + @lru_cache, which can cause cache hits from previous tests when running with + --dist=loadscope (tests from same file run on same worker). + """ + _load_huggingface_tokenizer.cache_clear() + + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): + # Setup mock to raise an error + mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") + + # Test with llama-3 model + result = _select_tokenizer_helper("llama-3-7b") + + # Verify the attempt to load Llama-3 tokenizer + mock_from_pretrained.assert_called_once_with("Xenova/llama-3-tokenizer") + + # Verify fallback to OpenAI tokenizer + self.assertEqual(result["type"], "openai_tokenizer") + self.assertEqual(result["tokenizer"], encoding) + + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_cohere_tokenizer_api_failure(self, mock_from_pretrained): + # Setup mock to raise an error + mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") + + # Add Cohere model to the list for testing + litellm.cohere_models = ["command-r-v1"] + + # Test with Cohere model + result = _select_tokenizer_helper("command-r-v1") + + # Verify the attempt to load Cohere tokenizer + mock_from_pretrained.assert_called_once_with( + "Xenova/c4ai-command-r-v01-tokenizer" + ) + + # Verify fallback to OpenAI tokenizer + self.assertEqual(result["type"], "openai_tokenizer") + self.assertEqual(result["tokenizer"], encoding) + + @patch("litellm.utils.tokenizer_dispatch.anthropic") + def test_claude_tokenizer_api_failure(self, mock_anthropic): + # Setup mock to raise an error + mock_anthropic.side_effect = Exception("Failed to load tokenizer") + + # Add Claude model to the list for testing + litellm.anthropic_models = ["claude-2"] + + # Test with Claude model + result = _select_tokenizer_helper("claude-2") + + # Verify the attempt to load Claude tokenizer + mock_anthropic.assert_called_once_with() + + # Verify fallback to OpenAI tokenizer + self.assertEqual(result["type"], "openai_tokenizer") + self.assertEqual(result["tokenizer"], encoding) + + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_llama2_tokenizer_api_failure(self, mock_from_pretrained): + # Setup mock to raise an error + mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") + + # Test with Llama-2 model + result = _select_tokenizer_helper("llama-2-7b") + + # Verify the attempt to load Llama-2 tokenizer + mock_from_pretrained.assert_called_once_with( + "hf-internal-testing/llama-tokenizer" + ) + + # Verify fallback to OpenAI tokenizer + self.assertEqual(result["type"], "openai_tokenizer") + self.assertEqual(result["tokenizer"], encoding) + + @patch("litellm.utils._return_huggingface_tokenizer") + def test_disable_hf_tokenizer_download(self, mock_return_huggingface_tokenizer): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) + try: + result = _select_tokenizer_helper("grok-32r22r") + mock_return_huggingface_tokenizer.assert_not_called() + assert result["type"] == "openai_tokenizer" + assert result["tokenizer"] == encoding + finally: + monkeypatch.undo() + + +def test_token_counter_with_anthropic_tool_use(): + """ + Test that _count_anthropic_content() correctly handles tool_use blocks. + + Validates that: + - 'name' field is counted (string) + - 'input' field is counted (dict serialized to string) + - Metadata fields ('type', 'id') are skipped + """ + messages = [ + {"role": "user", "content": "What's the weather in San Francisco?"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "I'll check the weather for you."}, + { + "type": "tool_use", + "id": "toolu_01234567890", # Should be skipped + "name": "get_weather", # Should be counted + "input": { # Should be counted (serialized) + "location": "San Francisco, CA", + "unit": "fahrenheit", + }, + }, + ], + }, + ] + + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens > 0, f"Expected positive token count, got {tokens}" + # Should count: user message + "I'll check" text + "get_weather" name + input dict + assert ( + tokens > 15 + ), f"Expected reasonable token count for message with tool_use, got {tokens}" + + +def test_token_counter_with_anthropic_tool_result(): + """ + Test that _count_anthropic_content() correctly handles tool_result blocks. + + Validates that: + - 'content' field (when string) is counted + - Metadata fields ('type', 'tool_use_id') are skipped + - Full conversation with tool_use → tool_result flow works + """ + messages = [ + {"role": "user", "content": "What's the weather in San Francisco?"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01234567890", + "name": "get_weather", + "input": {"location": "San Francisco, CA"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234567890", # Should be skipped + "content": "The weather in San Francisco is 65°F and sunny.", # Should be counted + } + ], + }, + ] + + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens > 0, f"Expected positive token count, got {tokens}" + assert ( + tokens > 25 + ), f"Expected reasonable token count for conversation with tool_result, got {tokens}" + + +def test_token_counter_with_nested_tool_result(): + """ + Test that _count_anthropic_content() recursively handles nested content lists. + + Validates that: + - tool_result with 'content' as a list (not string) is handled + - Nested content blocks are recursively counted via _count_content_list() + - TypedDict inference correctly identifies list fields + """ + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234567890", + "content": [ # Nested list - should recursively count + { + "type": "text", + "text": "The weather in San Francisco is 65°F and sunny.", + }, + {"type": "text", "text": "UV index is moderate."}, + ], + } + ], + } + ] + + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens > 0, f"Expected positive token count, got {tokens}" + # Should count both nested text blocks + assert ( + tokens > 15 + ), f"Expected reasonable token count for nested tool_result, got {tokens}" + + +def test_token_counter_tool_use_and_result_combined(): + """ + Test dynamic field inference with multiple tool_use and tool_result blocks. + + Validates that: + - Multiple tool_use blocks in same message are handled + - Multiple tool_result blocks in same message are handled + - skip_fields correctly filters metadata across all blocks + - Full realistic conversation flow works end-to-end + """ + messages = [ + { + "role": "user", + "content": "What's the weather in San Francisco and New York?", + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "I'll check the weather in both cities for you.", + }, + { + "type": "tool_use", + "id": "toolu_01A", + "name": "get_weather", + "input": {"location": "San Francisco, CA"}, + }, + { + "type": "tool_use", + "id": "toolu_01B", + "name": "get_weather", + "input": {"location": "New York, NY"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01A", + "content": "San Francisco: 65°F, sunny", + }, + { + "type": "tool_result", + "tool_use_id": "toolu_01B", + "content": "New York: 45°F, cloudy", + }, + ], + }, + { + "role": "assistant", + "content": "The weather in San Francisco is 65°F and sunny, while New York is cooler at 45°F and cloudy.", + }, + ] + + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens > 0, f"Expected positive token count, got {tokens}" + # Should count all text, tool names, inputs, and results + assert ( + tokens > 60 + ), f"Expected substantial token count for full tool conversation, got {tokens}" + + +def test_token_counter_with_image_url(): + """ + Test that _count_image_tokens() correctly handles image_url content blocks. + + Validates that: + - image_url as dict with 'url' and 'detail' is handled + - image_url as string is handled + - 'detail' field validation works ('low', 'high', 'auto') + - calculate_img_tokens is called with correct parameters + """ + # Test with dict format (detail: low) + messages_dict = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image.jpg", + "detail": "low", # Should use low token count (85 base tokens) + }, + }, + ], + } + ] + + tokens_dict = token_counter( + model="gpt-3.5-turbo", + messages=messages_dict, + use_default_image_token_count=True, # Avoid actual HTTP request + ) + assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}" + assert tokens_dict > 85, f"Expected at least base image tokens, got {tokens_dict}" + + # Test with string format (defaults to auto/low) + messages_str = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "https://example.com/image.jpg", # String format + } + ], + } + ] + + tokens_str = token_counter( + model="gpt-3.5-turbo", messages=messages_str, use_default_image_token_count=True + ) + assert ( + tokens_str > 0 + ), f"Expected positive token count for string image_url, got {tokens_str}" + + # Test invalid detail value raises error + messages_invalid = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image.jpg", + "detail": "invalid", # Should raise ValueError + }, + } + ], + } + ] + + with pytest.raises(ValueError, match="Invalid detail value") as exc_info: + token_counter(model="gpt-3.5-turbo", messages=messages_invalid) + e = exc_info.value + assert "Invalid detail value" in str( + e + ), f"Expected detail validation error, got: {e}" + + +def test_token_counter_with_thinking_content(): + """ + Test that _count_content_list() correctly handles Claude's extended thinking content blocks. + + Validates that: + - 'thinking' content type is recognized and counted + - 'thinking' text field is counted + - 'signature' field is skipped (opaque signature blob) + - Full conversation with thinking blocks works + """ + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this complex problem: who came first, chicken or egg", + } + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "This is actually a fascinating question that touches on philosophy, biology, and semantics. Let me break this down: The egg came first from an evolutionary biology perspective.", + "signature": "EqcLCkYICxgCKkCrqu6lP...", # Should be skipped + }, + { + "type": "text", + "text": "# The Chicken-or-Egg Question: A Multi-Layered Answer\n\n## **The Short Answer: The Egg Came First**", + }, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "Thanks"}]}, + ] + + tokens = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", messages=messages + ) + assert tokens > 0, f"Expected positive token count, got {tokens}" + # Should count: user message + thinking text + response text + "Thanks" + # The thinking text alone is ~30 tokens, plus other content should be > 50 total + assert ( + tokens > 50 + ), f"Expected substantial token count for message with thinking, got {tokens}" + + # Test that thinking block without 'thinking' field doesn't crash (edge case) + messages_no_thinking = [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + # No 'thinking' field - should count as 0 tokens + "signature": "EqcLCkYICxgCKkCrqu6lP...", + }, + {"type": "text", "text": "Response"}, + ], + } + ] + + tokens_no_thinking = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", messages=messages_no_thinking + ) + assert ( + tokens_no_thinking > 0 + ), f"Expected positive token count even with empty thinking, got {tokens_no_thinking}" + # Should only count "Response" and message overhead + assert ( + tokens_no_thinking < 15 + ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}" + + +def test_token_counter_with_redacted_thinking_content(): + """ + A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in + for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking + block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the + prompt_caching pre-call check stop pinning the deployment that held the cached prefix. + """ + model = "anthropic/claude-sonnet-4-5-20250929" + reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."} + redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30} + user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]} + follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]} + + without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up] + with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up] + + assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) + +def test_token_counter_with_tool_reference_block(): + """ + Regression test: a message containing an Anthropic tool-search + `tool_reference` content block must NOT raise. + + Before the fix, token_counter raised + `Invalid content item type: tool_reference`. On the streaming + anthropic_messages proxy path this nulled response_cost and caused the + SpendLogs row to be dropped, silently undercounting cost. token_counter + must instead count the referenced tool name and return a positive count. + """ + messages = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me look up the right tool."}, + {"type": "tool_reference", "tool_name": "search_knowledge_base"}, + ], + } + ] + + # Must not raise, and must produce a positive token count. + tokens = token_counter_new( + model="anthropic/claude-sonnet-4-5-20250929", messages=messages + ) + assert tokens > 0, f"Expected positive token count, got {tokens}" + + # A tool_reference with no/empty tool_name must also be handled gracefully. + messages_empty = [ + { + "role": "assistant", + "content": [{"type": "tool_reference", "tool_name": ""}], + } + ] + tokens_empty = token_counter_new( + model="anthropic/claude-sonnet-4-5-20250929", messages=messages_empty + ) + assert tokens_empty >= 0 + + +def test_count_content_list_rejects_unknown_type(): + """ + An unrecognized content block type must raise, and the error message must + enumerate the supported types (including `tool_reference`). This pins the + catch-all contract so a future block type isn't silently dropped. + """ + from litellm.litellm_core_utils.token_counter import _count_content_list + + with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info: + _count_content_list( + count_function=len, + content_list=[{"type": "totally_unknown_block"}], + use_default_image_token_count=False, + default_token_count=None, + ) + + message = str(exc_info.value) + assert "Invalid content item type: totally_unknown_block" in message + assert "tool_reference" in message + + +@pytest.mark.parametrize( + "source", + [ + {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}, + {"type": "url", "url": "https://example.com/image.png"}, + {"type": "file", "file_id": "file-abc123"}, + ], + ids=["base64", "url", "file"], +) +def test_token_counter_with_anthropic_image_block(source: dict[str, str]): + """Anthropic `image` blocks must count for every source variant, not raise `Invalid content item type` (which the router's context-window pre-call check swallows into an unfiltered dispatch).""" + from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image", "source": source}, + ], + } + ] + + tokens = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", + messages=messages, + use_default_image_token_count=True, + ) + assert tokens > DEFAULT_IMAGE_TOKEN_COUNT, ( + f"Expected the image block to contribute tokens, got {tokens}" + ) + + +def test_anthropic_image_block_matches_equivalent_image_url(): + """An Anthropic `image` block prices identically to the OpenAI `image_url` carrying the same bytes.""" + anthropic_messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgo=", + }, + } + ], + } + ] + openai_messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, + } + ], + } + ] + + anthropic_tokens = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", messages=anthropic_messages + ) + openai_tokens = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", messages=openai_messages + ) + assert anthropic_tokens == openai_tokens + + +def test_anthropic_image_block_nested_in_tool_result(): + """An `image` block nested in a `tool_result.content` list is counted through the same recursion.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgo=", + }, + } + ], + } + ], + } + ] + + tokens = token_counter( + model="anthropic/claude-sonnet-4-5-20250929", + messages=messages, + use_default_image_token_count=True, + ) + assert tokens > 0 + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ({"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}, "data:image/jpeg;base64,/9j/4AAQ"), + ({"type": "url", "url": "https://example.com/image.png"}, "https://example.com/image.png"), + ({"type": "file", "file_id": "file-abc123"}, ""), + ], + ids=["base64", "url", "file"], +) +def test_anthropic_image_source_resolves_to_what_the_image_pricer_reads(source: dict[str, str], expected: str): + """base64 sources become a data URI, url sources pass through, file sources resolve to an empty string.""" + from litellm.litellm_core_utils.token_counter import _anthropic_image_source_data + + assert _anthropic_image_source_data(source) == expected + + +def test_anthropic_image_block_with_empty_base64_data(): + """A base64 source with empty `data` prices as an image rather than raising.""" + from litellm.litellm_core_utils.token_counter import _count_content_list + + tokens = _count_content_list( + count_function=len, + content_list=[ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}} + ], + use_default_image_token_count=False, + default_token_count=None, + ) + assert tokens > 0 + + +def test_anthropic_image_block_without_source_raises(): + """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior.""" + from litellm.litellm_core_utils.token_counter import _count_content_list + + with pytest.raises(ValueError, match="Error getting number of tokens from content list"): + _count_content_list( + count_function=len, + content_list=[{"type": "image"}], + use_default_image_token_count=False, + default_token_count=None, + ) + + # ... and `default_token_count`, the caller's opt-out from raising, still wins. + assert ( + _count_content_list( + count_function=len, + content_list=[{"type": "image"}], + use_default_image_token_count=False, + default_token_count=7, + ) + == 7 + ) + + +def _count_user_content(content: list[dict]) -> int: + from litellm.litellm_core_utils.token_counter import token_counter + + return token_counter( + model="anthropic/claude-fable-5", + messages=[{"role": "user", "content": content}], + use_default_image_token_count=True, + ) + + +@pytest.mark.parametrize( + "source", + [ + {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + {"type": "url", "url": "https://example.com/report.pdf"}, + {"type": "file", "file_id": "file-abc123"}, + ], + ids=["base64", "url", "file"], +) +def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]): + """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on.""" + prompt = {"type": "text", "text": "Summarize this file."} + + assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content( + [prompt, {"type": "image", "source": source}] + ) + + +def test_anthropic_document_block_text_sources_count_their_text(): + """`text` and `content` document sources count the text they carry, as inline text blocks would.""" + prompt = {"type": "text", "text": "Summarize this file."} + body = {"type": "text", "text": "Revenue grew eleven percent while churn fell to two percent."} + picture = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}} + + text_source = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": body["text"]}} + assert _count_user_content([prompt, text_source]) == _count_user_content([prompt, body]) + + string_content = {"type": "document", "source": {"type": "content", "content": body["text"]}} + assert _count_user_content([prompt, string_content]) == _count_user_content([prompt, body]) + + block_content = {"type": "document", "source": {"type": "content", "content": [body, picture]}} + assert _count_user_content([prompt, block_content]) == _count_user_content([prompt, body, picture]) + + +def test_anthropic_document_title_and_context_add_their_tokens(): + prompt = {"type": "text", "text": "Summarize this file."} + source = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} + described = {"type": "document", "source": source, "title": "Q3 board packet", "context": "Shared by finance"} + + assert _count_user_content([prompt, described]) == _count_user_content( + [ + prompt, + {"type": "text", "text": "Q3 board packet"}, + {"type": "text", "text": "Shared by finance"}, + {"type": "document", "source": source}, + ] + ) + + +def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): + """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise. + + Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject` + is in the union this counter accepts, so every local count of a Responses `input_file` raised + `Invalid content item type: file` and surfaced as a 500 on /v1/responses/input_tokens. + """ + prompt = {"type": "text", "text": "Summarize this file."} + inline_file = { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + } + document = { + "type": "document", + "title": "report.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + } + + assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) + assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) + + +def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): + """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens.""" + prompt = {"type": "text", "text": "Summarize this file."} + + by_id = {"type": "file", "file": {"file_id": "file-abc123"}} + assert _count_user_content([prompt, by_id]) == _count_user_content([prompt]) + + named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}} + assert _count_user_content([prompt, named]) == _count_user_content( + [prompt, {"type": "text", "text": "report.pdf"}] + ) + + +def _png_data_url(width: int, height: int) -> str: + ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") + return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() + + +@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) +def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: + assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() + + +def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: + assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() + assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter_tool.py b/tests/unit/litellm_core_utils/test_token_counter_tool.py similarity index 93% rename from tests/test_litellm/litellm_core_utils/test_token_counter_tool.py rename to tests/unit/litellm_core_utils/test_token_counter_tool.py index 9f8c1070a47..f61b7d335c1 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter_tool.py +++ b/tests/unit/litellm_core_utils/test_token_counter_tool.py @@ -5,8 +5,8 @@ import pytest # Use the same token_counter as the main test. -from tests.test_litellm.litellm_core_utils.test_token_counter import token_counter -from tests.test_litellm.litellm_core_utils.test_token_counter_tool_data import * +from tests.unit.litellm_core_utils.test_token_counter import token_counter +from tests.unit.litellm_core_utils.test_token_counter_tool_data import * @pytest.mark.parametrize( diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter_tool_data.py b/tests/unit/litellm_core_utils/test_token_counter_tool_data.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_token_counter_tool_data.py rename to tests/unit/litellm_core_utils/test_token_counter_tool_data.py diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py new file mode 100644 index 00000000000..a9005ff6a86 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -0,0 +1,411 @@ +import copy +import os +import pickle +import subprocess +import sys +from pathlib import Path +from typing import Final, Literal + +import pytest +import tiktoken +from tokenizers import Tokenizer as ReferenceTokenizer + +import litellm +from litellm.caching._embedding_router import truncate_embedding_input +from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding +from litellm.utils import claude_json_str +from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON + + +OFFLINE_ENCODINGS: Final = ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "o200k_harmony") +UNICODE_TEXTS: Final = ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64) + + +@pytest.mark.parametrize("name", OFFLINE_ENCODINGS) +@pytest.mark.parametrize("text", UNICODE_TEXTS) +def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: + assert_openai_encoding_matches_python(name, text) + + +def assert_openai_encoding_matches_python(name: str, text: str) -> None: + reference: Final = tiktoken.get_encoding(name) + encoding: Final = OpenAIEncoding.from_tiktoken(name) + expected: Final = reference.encode(text) + + assert encoding.encode(text) == expected + assert encoding.count(text) == len(expected) + assert encoding.encode_batch([text], num_threads=2) == reference.encode_batch([text], num_threads=2) + assert encoding.encode_ordinary_batch([text]) == reference.encode_ordinary_batch([text]) + assert encoding.decode_batch([expected]) == reference.decode_batch([expected]) + assert encoding.decode_bytes_batch([expected]) == reference.decode_bytes_batch([expected]) + + +@pytest.mark.parametrize("allowed", (frozenset(), frozenset({"<|endoftext|>"}), "all")) +@pytest.mark.parametrize("disallowed", (frozenset(), frozenset({"<|fim_prefix|>"}), "all")) +def test_openai_special_token_options_match_python( + allowed: frozenset[str] | Literal["all"], disallowed: frozenset[str] | Literal["all"] +) -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) + text: Final = "hello<|endoftext|><|fim_prefix|>world" + allowed_set: Final = reference.special_tokens_set if allowed == "all" else allowed + disallowed_set: Final = reference.special_tokens_set - allowed_set if disallowed == "all" else disallowed + if any(token in text for token in disallowed_set): + with pytest.raises(ValueError, match="disallowed special token"): + encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) + return + assert encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) == reference.encode( + text, allowed_special=allowed, disallowed_special=disallowed + ) + assert encoding.special_tokens_set == reference.special_tokens_set + assert encoding.eot_token == reference.eot_token + + +@pytest.mark.parametrize("errors", ("replace", "ignore", "backslashreplace", "strict")) +def test_openai_partial_token_decoding_preserves_error_policy(errors: str) -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) + tokens: Final = reference.encode("🙂")[:1] + assert encoding.decode_bytes(tokens) == reference.decode_bytes(tokens) + if errors == "strict": + with pytest.raises(UnicodeDecodeError): + encoding.decode(tokens, errors=errors) + return + assert encoding.decode(tokens, errors=errors) == reference.decode(tokens, errors=errors) + assert encoding.decode_tokens_bytes(tokens) == reference.decode_tokens_bytes(tokens) + + +def test_public_encoding_and_semantic_cache_preserve_truncated_unicode() -> None: + reference: Final = tiktoken.get_encoding(litellm.encoding.name) + text: Final = "🙂" + tokens: Final = reference.encode(text) + + assert litellm.encoding.encode(text, disallowed_special=()) == tokens + assert litellm.encoding.encode_batch([text]) == [tokens] + assert litellm.decode(tokens=tokens[:1]) == reference.decode(tokens[:1]) + assert truncate_embedding_input(text, "", 1) == reference.decode(tokens[:1]) + + +@pytest.mark.parametrize("add_special_tokens", (True, False)) +def test_huggingface_encoding_preserves_result_fields_and_serialization(add_special_tokens: bool) -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) + expected: Final = reference.encode("Hello World", add_special_tokens=add_special_tokens) + actual: Final = tokenizer.encode("Hello World", add_special_tokens=add_special_tokens) + + assert (actual.ids, actual.tokens, actual.type_ids, actual.offsets, actual.word_ids, actual.sequence_ids) == ( + expected.ids, + expected.tokens, + expected.type_ids, + expected.offsets, + expected.word_ids, + expected.sequence_ids, + ) + assert (actual.attention_mask, actual.special_tokens_mask, actual.n_sequences, len(actual)) == ( + expected.attention_mask, + expected.special_tokens_mask, + expected.n_sequences, + len(expected), + ) + assert copy.deepcopy(actual).ids == expected.ids + assert pickle.loads(pickle.dumps(actual)).offsets == expected.offsets + assert tokenizer.decode(actual.ids, skip_special_tokens=False) == reference.decode( + expected.ids, skip_special_tokens=False + ) + + +def test_huggingface_character_offsets_and_pretokenized_pairs_match_python() -> None: + reference: Final = ReferenceTokenizer.from_str(claude_json_str) + tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) + text: Final = "café 漢字 🙂" + actual: Final = tokenizer.encode(text) + expected: Final = reference.encode(text) + + assert actual.offsets == expected.offsets + assert actual.ids == expected.ids + assert ( + tokenizer.encode(["hello", "world"], ["again"], is_pretokenized=True).ids + == reference.encode(["hello", "world"], ["again"], is_pretokenized=True).ids + ) + + +def test_huggingface_batches_apply_padding_across_inputs() -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + reference.enable_padding(pad_id=0, pad_token="[UNK]") + tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) + inputs: Final = ["Hello", ("Hello World", "World")] + expected: Final = reference.encode_batch(inputs) + actual: Final = tokenizer.encode_batch(inputs) + fast: Final = tokenizer.encode_batch_fast(inputs) + + assert [(item.ids, item.attention_mask, item.offsets) for item in actual] == [ + (item.ids, item.attention_mask, item.offsets) for item in expected + ] + assert [item.ids for item in fast] == [item.ids for item in expected] + assert tokenizer.decode_batch([item.ids for item in actual]) == reference.decode_batch( + [item.ids for item in expected] + ) + + +def test_caller_supplied_huggingface_tokenizer_preserves_public_encode_and_count() -> None: + tokenizer: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + custom: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + expected: Final = tokenizer.encode("Hello World").ids + + assert litellm.encode(text="Hello World", custom_tokenizer=custom) == expected + assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(expected) + assert litellm.decode(tokens=expected, custom_tokenizer=custom) == "Hello World" + + +def test_caller_supplied_tiktoken_treats_special_spellings_as_text() -> None: + tokenizer: Final = tiktoken.get_encoding("cl100k_base") + custom: Final = {"type": "openai_tokenizer", "tokenizer": tokenizer} + text: Final = "<|endoftext|>" + + assert litellm.encode(text=text, custom_tokenizer=custom) == tokenizer.encode(text, disallowed_special=()) + + +def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> None: + custom: Final = litellm.create_tokenizer(TOKENIZER_JSON) + tokenizer: Final = custom["tokenizer"] + path: Final = tmp_path / "tokenizer.json" + tokenizer.save(str(path)) + + assert copy.deepcopy(custom)["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids + assert ( + pickle.loads(pickle.dumps(custom))["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids + ) + assert HuggingFaceTokenizer.from_file(str(path)).encode("Hello World").ids == tokenizer.encode("Hello World").ids + assert copy.deepcopy(litellm.encoding).encode("hello") == litellm.encoding.encode("hello") + assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello") + + +@pytest.mark.parametrize("offline", ("0", "1")) +def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None: + script: Final = """ +import json +import sys +from pathlib import Path +sys.path.insert(0, sys.argv[1]) +import httpx +import huggingface_hub +from huggingface_hub.errors import LocalEntryNotFoundError +import litellm +payload = sys.argv[2].encode() +offline = sys.argv[3] == "1" +observed = [] +def handle(request): + assert not offline, "offline loading issued a request" + if request.url.path.endswith("/tokenizer.json"): + observed.append(request.headers.get("authorization")) + if request.headers.get("authorization") != "Bearer audit-fixture-token": + return httpx.Response(401) + return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"") +if not offline: + huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +try: + tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"] +except LocalEntryNotFoundError: + assert offline + assert observed == [] +else: + assert not offline + assert "Bearer audit-fixture-token" in observed + assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World" + assert tuple(Path(sys.argv[4]).rglob("tokenizer.json")) +print("compatible") +""" + result: Final = subprocess.run( + [ + sys.executable, + "-I", + "-c", + script, + str(Path(litellm.__file__).parent.parent), + TOKENIZER_JSON, + offline, + str(tmp_path / "cache"), + ], + capture_output=True, + text=True, + timeout=30, + env={ + **os.environ, + "HF_HOME": str(tmp_path / "home"), + "HF_HUB_CACHE": str(tmp_path / "cache"), + "HF_ENDPOINT": "http://127.0.0.1:9", + "HF_TOKEN": "audit-fixture-token", + "HF_HUB_OFFLINE": offline, + "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + + +@pytest.mark.parametrize("rust", (None, "0", "1")) +def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None: + script: Final = """ +import importlib.abc +import sys +sys.path.insert(0, sys.argv[1]) +def reject_network(event, args): + if event == "socket.connect": + raise AssertionError("tokenizer attempted a network connection") +sys.addaudithook(reject_network) +class Block(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "litellm.rust_bridge._native": + raise ImportError("native extension is unavailable") +sys.meta_path.insert(0, Block()) +import litellm +from litellm.rust_bridge.tokenizer import get_encoding +import tiktoken +from tokenizers import Tokenizer +assert isinstance(litellm.encoding, tiktoken.Encoding) +for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"): + encoding = get_encoding(name) + text = "offline café 漢字 🙂" + " " * 64 + assert encoding.decode(encoding.encode(text)) == text +ids = litellm.encode(text="hello world") +assert litellm.decode(tokens=ids) == "hello world" +assert litellm.token_counter(model=None, text="hello world") == len(ids) +custom = litellm.create_tokenizer(sys.argv[2]) +assert isinstance(custom["tokenizer"], Tokenizer) +custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]") +assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World" +print("compatible") +""" + result: Final = subprocess.run( + [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON], + capture_output=True, + text=True, + timeout=30, + cwd=tmp_path, + env={ + **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"}, + **({"LITELLM_RUST": rust} if rust is not None else {}), + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + assert not (tmp_path / "unused-tokenizer-cache").exists() + + +@pytest.mark.parametrize("is_pretokenized", (False, True)) +def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) + inputs: Final = [["Hello", "World"], ("Hello", "World")] + actual: Final = tokenizer.encode_batch(inputs, is_pretokenized=is_pretokenized) + expected: Final = reference.encode_batch(inputs, is_pretokenized=is_pretokenized) + assert [(item.ids, item.type_ids, item.sequence_ids) for item in actual] == [ + (item.ids, item.type_ids, item.sequence_ids) for item in expected + ] + + +@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit")) +def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: + assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name) + + +def assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: + reference: Final = tiktoken.get_encoding(name) + encoding: Final = OpenAIEncoding.from_tiktoken(name) + text: Final = "hello fanta" + + assert repr(encoding) == repr(reference) == f"" + assert (encoding.name, encoding.n_vocab, encoding.max_token_value) == ( + reference.name, + reference.n_vocab, + reference.max_token_value, + ) + assert encoding.token_byte_values() == reference.token_byte_values() + assert encoding.encode_single_token("hello") == reference.encode_single_token("hello") + assert encoding.encode_single_token(b"<|endoftext|>") == reference.eot_token + assert [encoding.is_special_token(token) for token in (0, reference.eot_token)] == [False, True] + assert encoding.decode_with_offsets(reference.encode(text)) == reference.decode_with_offsets(reference.encode(text)) + assert encoding.encode_to_numpy(text).tolist() == reference.encode_to_numpy(text).tolist() + stable, completions = encoding.encode_with_unstable(text) + expected_stable, expected_completions = reference.encode_with_unstable(text) + assert (stable, sorted(completions)) == (expected_stable, sorted(expected_completions)) + with pytest.raises(KeyError): + encoding.encode_single_token("<|not-a-token|>") + + +def test_huggingface_tokenizer_exposes_the_tokenizers_vocabulary_surface() -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + reference.enable_padding(pad_id=0, pad_token="[UNK]", length=4) + reference.enable_truncation(max_length=3, stride=1, strategy="only_first", direction="left") + tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) + + assert tokenizer.token_to_id("Hello") == reference.token_to_id("Hello") == 1 + assert tokenizer.id_to_token(3) == reference.id_to_token(3) == "[BOS]" + assert tokenizer.id_to_token(99) is None + assert tokenizer.get_vocab() == reference.get_vocab() + assert tokenizer.get_vocab(with_added_tokens=False) == reference.get_vocab(with_added_tokens=False) + assert tokenizer.get_vocab_size() == reference.get_vocab_size() == 4 + assert tokenizer.get_vocab_size(with_added_tokens=False) == reference.get_vocab_size(with_added_tokens=False) + added: Final = tokenizer.get_added_tokens_decoder() + expected_added: Final = reference.get_added_tokens_decoder() + assert {token_id: str(token) for token_id, token in added.items()} == { + token_id: str(token) for token_id, token in expected_added.items() + } + assert added[3].special == expected_added[3].special + assert tokenizer.num_special_tokens_to_add(False) == reference.num_special_tokens_to_add(False) == 1 + assert tokenizer.num_special_tokens_to_add(True) == reference.num_special_tokens_to_add(True) == 0 + assert tokenizer.padding == reference.padding + assert tokenizer.truncation == reference.truncation + assert tokenizer.encode_special_tokens == reference.encode_special_tokens is False + assert HuggingFaceTokenizer.from_buffer(TOKENIZER_JSON.encode()).encode("Hello").ids == [3, 1] + assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).padding is None + assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).truncation is None + + +def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface() -> None: + reference: Final = ReferenceTokenizer.from_str(claude_json_str) + tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) + text: Final = "hello wide world" + actual: Final = tokenizer.encode(text, "again") + expected: Final = reference.encode(text, "again") + + lookups: Final = ( + lambda encoding: [encoding.token_to_chars(index) for index in range(len(encoding))], + lambda encoding: [encoding.token_to_word(index) for index in range(len(encoding))], + lambda encoding: [encoding.token_to_sequence(index) for index in range(len(encoding))], + lambda encoding: [encoding.char_to_token(position) for position in range(len(text))], + lambda encoding: [encoding.char_to_word(position) for position in range(len(text))], + lambda encoding: [encoding.char_to_token(position, 1) for position in range(5)], + lambda encoding: [encoding.word_to_tokens(word) for word in range(3)], + lambda encoding: [encoding.word_to_chars(word) for word in range(3)], + lambda encoding: [encoding.word_to_tokens(0, 1), encoding.word_to_chars(0, 1)], + ) + for lookup in lookups: + assert lookup(actual) == lookup(expected) + assert repr(actual) == repr(expected) + + actual.truncate(4, stride=1, direction="left") + expected.truncate(4, stride=1, direction="left") + assert (actual.ids, [item.ids for item in actual.overflowing]) == ( + expected.ids, + [item.ids for item in expected.overflowing], + ) + actual.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") + expected.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") + assert (actual.ids, actual.attention_mask, actual.type_ids, actual.tokens) == ( + expected.ids, + expected.attention_mask, + expected.type_ids, + expected.tokens, + ) + actual.set_sequence_id(3) + expected.set_sequence_id(3) + assert actual.sequence_ids == expected.sequence_ids + merged: Final = type(actual).merge([actual, tokenizer.encode("more")]) + assert merged.ids == type(expected).merge([expected, reference.encode("more")]).ids + assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets + with pytest.raises(ValueError, match="direction"): + actual.pad(8, direction="sideways") diff --git a/tests/test_litellm/litellm_core_utils/test_tool_search_spend_logging.py b/tests/unit/litellm_core_utils/test_tool_search_spend_logging.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_tool_search_spend_logging.py rename to tests/unit/litellm_core_utils/test_tool_search_spend_logging.py diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_url_utils.py rename to tests/unit/litellm_core_utils/test_url_utils.py diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/unit/litellm_core_utils/test_xai_oauth_routing.py similarity index 100% rename from tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py rename to tests/unit/litellm_core_utils/test_xai_oauth_routing.py diff --git a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py index d835db63d83..5e2956b532a 100644 --- a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -2733,7 +2733,7 @@ def test_build_summary_messages_keeps_midturn_system_correction_in_place(): async def test_threshold_check_counts_tokens_off_the_event_loop(monkeypatch): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py index a21c22cf5fa..9fad6ca5e66 100644 --- a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py +++ b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py @@ -133,7 +133,7 @@ async def test_malformed_edit_entries_are_skipped(): async def test_sync_editor_counts_tokens_off_the_event_loop(): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/unit/llms/test_polling_url_origin_match.py b/tests/unit/llms/test_polling_url_origin_match.py index ab5f41c757f..2df35131e3d 100644 --- a/tests/unit/llms/test_polling_url_origin_match.py +++ b/tests/unit/llms/test_polling_url_origin_match.py @@ -18,7 +18,7 @@ import pytest # Azure DALL-E sync + async paths route through ``assert_same_origin`` # the same way as the case below. The helper itself is unit-tested in -# ``tests/test_litellm/litellm_core_utils/test_url_utils.py``. +# ``tests/unit/litellm_core_utils/test_url_utils.py``. # ── Black Forest Labs polling ───────────────────────────────────────────────── diff --git a/tests/test_litellm/rust_bridge/chat_completions/__init__.py b/tests/unit/responses/litellm_completion_transformation/__init__.py similarity index 100% rename from tests/test_litellm/rust_bridge/chat_completions/__init__.py rename to tests/unit/responses/litellm_completion_transformation/__init__.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_function_call_output_normalization.py b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_function_call_output_normalization.py rename to tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_handler.py rename to tests/unit/responses/litellm_completion_transformation/test_handler.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py b/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py rename to tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py rename to tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py rename to tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py rename to tests/unit/responses/litellm_completion_transformation/test_session_handler.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py rename to tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py rename to tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py b/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py rename to tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py similarity index 100% rename from tests/test_litellm/responses/mcp/test_chat_completions_handler.py rename to tests/unit/responses/mcp/test_chat_completions_handler.py diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py similarity index 100% rename from tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py rename to tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py similarity index 100% rename from tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py rename to tests/unit/responses/mcp/test_mcp_streaming_iterator.py diff --git a/tests/test_litellm/responses/test_additional_tools.py b/tests/unit/responses/test_additional_tools.py similarity index 100% rename from tests/test_litellm/responses/test_additional_tools.py rename to tests/unit/responses/test_additional_tools.py diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/unit/responses/test_custom_tool_call.py similarity index 100% rename from tests/test_litellm/responses/test_custom_tool_call.py rename to tests/unit/responses/test_custom_tool_call.py diff --git a/tests/test_litellm/responses/test_dispatch.py b/tests/unit/responses/test_dispatch.py similarity index 100% rename from tests/test_litellm/responses/test_dispatch.py rename to tests/unit/responses/test_dispatch.py diff --git a/tests/test_litellm/responses/test_metadata_codex_callback.py b/tests/unit/responses/test_metadata_codex_callback.py similarity index 100% rename from tests/test_litellm/responses/test_metadata_codex_callback.py rename to tests/unit/responses/test_metadata_codex_callback.py diff --git a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py b/tests/unit/responses/test_no_duplicate_spend_logs.py similarity index 76% rename from tests/test_litellm/responses/test_no_duplicate_spend_logs.py rename to tests/unit/responses/test_no_duplicate_spend_logs.py index c98b519ae67..7e4bef5812c 100644 --- a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py +++ b/tests/unit/responses/test_no_duplicate_spend_logs.py @@ -15,35 +15,6 @@ import litellm from litellm.integrations.custom_logger import CustomLogger -def test_logging_object_not_popped(): - """ - Test that litellm_logging_obj is not popped from kwargs. - - This is a regression test for issue #15740. The bug was using - kwargs.pop() which removed the logging object, causing duplicate - spend logs for non-OpenAI providers. - """ - import inspect - - from litellm.responses import main as responses_module - - # Get the source code of the responses function - source = inspect.getsource(responses_module.responses) - - # Check that .pop("litellm_logging_obj") is NOT used - # The bug was using kwargs.pop("litellm_logging_obj") which removes it - assert 'kwargs.pop("litellm_logging_obj")' not in source, ( - "FAIL: Found kwargs.pop('litellm_logging_obj') in responses() function. " - "This causes duplicate spend logs. Use kwargs.get('litellm_logging_obj') instead." - ) - - # Check that .get("litellm_logging_obj") IS used - assert 'kwargs.get("litellm_logging_obj")' in source, ( - "FAIL: Expected kwargs.get('litellm_logging_obj') but not found. " - "The logging object must be accessed with .get() not .pop() to prevent duplication." - ) - - @pytest.mark.asyncio async def test_async_no_duplicate_spend_logs(): """ diff --git a/tests/test_litellm/responses/test_null_test_fix.py b/tests/unit/responses/test_null_test_fix.py similarity index 100% rename from tests/test_litellm/responses/test_null_test_fix.py rename to tests/unit/responses/test_null_test_fix.py diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py similarity index 100% rename from tests/test_litellm/responses/test_responses_api_bridge_flag.py rename to tests/unit/responses/test_responses_api_bridge_flag.py diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py similarity index 99% rename from tests/test_litellm/responses/test_responses_api_request_body.py rename to tests/unit/responses/test_responses_api_request_body.py index 98e74955c6f..b27401d693a 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -20,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler def _expected_dir() -> Path: - """Path to expected_responses_api_request folder (sibling of test_litellm/responses).""" + """Path to expected_responses_api_request folder (sibling of tests/unit/responses).""" return Path(__file__).resolve().parent.parent / "expected_responses_api_request" diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py similarity index 100% rename from tests/test_litellm/responses/test_responses_prompt_management.py rename to tests/unit/responses/test_responses_prompt_management.py diff --git a/tests/test_litellm/responses/test_responses_router_cooldown.py b/tests/unit/responses/test_responses_router_cooldown.py similarity index 100% rename from tests/test_litellm/responses/test_responses_router_cooldown.py rename to tests/unit/responses/test_responses_router_cooldown.py diff --git a/tests/test_litellm/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py similarity index 100% rename from tests/test_litellm/responses/test_responses_streaming_iterator.py rename to tests/unit/responses/test_responses_streaming_iterator.py diff --git a/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py b/tests/unit/responses/test_responses_supported_endpoints_passthrough.py similarity index 100% rename from tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py rename to tests/unit/responses/test_responses_supported_endpoints_passthrough.py diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py similarity index 100% rename from tests/test_litellm/responses/test_responses_utils.py rename to tests/unit/responses/test_responses_utils.py diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py similarity index 97% rename from tests/test_litellm/responses/test_responses_websocket_all_providers.py rename to tests/unit/responses/test_responses_websocket_all_providers.py index 3888a84fb5d..6f346a25d9c 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -2718,97 +2718,6 @@ class TestWebSocketChunkTypes: assert "response.reasoning_content.done" in serialized assert "Complete reasoning" in serialized - def test_extract_output_messages_preserves_multiple_messages(self): - """Test that multiple output messages are all preserved""" - from litellm.responses.streaming_iterator import ( - ManagedResponsesWebSocketHandler, - ) - - completed_event = { - "type": "response.completed", - "response": { - "id": "resp_123", - "output": [ - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": "First message"}], - }, - { - "type": "function_call", - "id": "call_123", - "name": "get_weather", - "arguments": "{}", - }, - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": "Second message"}], - }, - ], - }, - } - - messages = ManagedResponsesWebSocketHandler._extract_output_messages( - completed_event - ) - assert len(messages) == 3 - assert messages[0]["content"][0]["text"] == "First message" - assert messages[1]["type"] == "function_call" - assert messages[2]["content"][0]["text"] == "Second message" - - def test_input_to_messages_with_mixed_content_types(self): - """Test input conversion with mixed content types""" - from litellm.responses.streaming_iterator import ( - ManagedResponsesWebSocketHandler, - ) - - input_list = [ - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_text", "text": "Question"}, - {"type": "input_image", "image_url": "https://example.com/img.png"}, - ], - } - ] - - messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list) - assert len(messages) == 1 - assert len(messages[0]["content"]) == 2 - assert messages[0]["content"][0]["type"] == "input_text" - assert messages[0]["content"][1]["type"] == "input_image" - - def test_extract_output_messages_with_mixed_text_types(self): - """Test that both 'output_text' and 'text' types are extracted""" - from litellm.responses.streaming_iterator import ( - ManagedResponsesWebSocketHandler, - ) - - completed_event = { - "type": "response.completed", - "response": { - "id": "resp_123", - "output": [ - { - "type": "message", - "role": "assistant", - "content": [ - {"type": "output_text", "text": "Part 1"}, - {"type": "text", "text": "Part 2"}, - ], - } - ], - }, - } - - messages = ManagedResponsesWebSocketHandler._extract_output_messages( - completed_event - ) - assert len(messages) == 1 - assert messages[0]["content"][0]["text"] == "Part 1Part 2" - class TestNativeWebSocketUrlConstruction: """Test that native WebSocket URLs include the model query parameter. diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/unit/responses/test_rust_bridge_websocket.py similarity index 100% rename from tests/test_litellm/responses/test_rust_bridge_websocket.py rename to tests/unit/responses/test_rust_bridge_websocket.py diff --git a/tests/test_litellm/responses/test_sse_output_recovery.py b/tests/unit/responses/test_sse_output_recovery.py similarity index 100% rename from tests/test_litellm/responses/test_sse_output_recovery.py rename to tests/unit/responses/test_sse_output_recovery.py diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py similarity index 100% rename from tests/test_litellm/responses/test_streaming_iterator.py rename to tests/unit/responses/test_streaming_iterator.py diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py similarity index 100% rename from tests/test_litellm/responses/test_streaming_iterator_error_events.py rename to tests/unit/responses/test_streaming_iterator_error_events.py diff --git a/tests/test_litellm/responses/test_text_format_conversion.py b/tests/unit/responses/test_text_format_conversion.py similarity index 100% rename from tests/test_litellm/responses/test_text_format_conversion.py rename to tests/unit/responses/test_text_format_conversion.py diff --git a/tests/test_litellm/rust_bridge/messages/__init__.py b/tests/unit/router_strategy/adaptive_router/__init__.py similarity index 100% rename from tests/test_litellm/rust_bridge/messages/__init__.py rename to tests/unit/router_strategy/adaptive_router/__init__.py diff --git a/tests/test_litellm/rust_bridge/ocr/__init__.py b/tests/unit/router_strategy/adaptive_router/fixtures/__init__.py similarity index 100% rename from tests/test_litellm/rust_bridge/ocr/__init__.py rename to tests/unit/router_strategy/adaptive_router/fixtures/__init__.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json b/tests/unit/router_strategy/adaptive_router/fixtures/clean_no_signals.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json rename to tests/unit/router_strategy/adaptive_router/fixtures/clean_no_signals.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json b/tests/unit/router_strategy/adaptive_router/fixtures/clean_satisfaction.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json rename to tests/unit/router_strategy/adaptive_router/fixtures/clean_satisfaction.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json b/tests/unit/router_strategy/adaptive_router/fixtures/disengagement_giveup.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json rename to tests/unit/router_strategy/adaptive_router/fixtures/disengagement_giveup.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json b/tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_429.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json rename to tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_429.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json b/tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json rename to tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json b/tests/unit/router_strategy/adaptive_router/fixtures/failure_tool_error.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json rename to tests/unit/router_strategy/adaptive_router/fixtures/failure_tool_error.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json b/tests/unit/router_strategy/adaptive_router/fixtures/loop_same_tool.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json rename to tests/unit/router_strategy/adaptive_router/fixtures/loop_same_tool.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json b/tests/unit/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json rename to tests/unit/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json b/tests/unit/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json rename to tests/unit/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json b/tests/unit/router_strategy/adaptive_router/fixtures/stagnation_repeat.json similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json rename to tests/unit/router_strategy/adaptive_router/fixtures/stagnation_repeat.json diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py rename to tests/unit/router_strategy/adaptive_router/test_adaptive_router.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py b/tests/unit/router_strategy/adaptive_router/test_async_pre_routing.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py rename to tests/unit/router_strategy/adaptive_router/test_async_pre_routing.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/unit/router_strategy/adaptive_router/test_bandit.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_bandit.py rename to tests/unit/router_strategy/adaptive_router/test_bandit.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_classifier.py b/tests/unit/router_strategy/adaptive_router/test_classifier.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_classifier.py rename to tests/unit/router_strategy/adaptive_router/test_classifier.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_config.py b/tests/unit/router_strategy/adaptive_router/test_config.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_config.py rename to tests/unit/router_strategy/adaptive_router/test_config.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/unit/router_strategy/adaptive_router/test_e2e_adaptive_router.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py rename to tests/unit/router_strategy/adaptive_router/test_e2e_adaptive_router.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_hooks.py rename to tests/unit/router_strategy/adaptive_router/test_hooks.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py rename to tests/unit/router_strategy/adaptive_router/test_router_dispatch.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_signals.py b/tests/unit/router_strategy/adaptive_router/test_signals.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_signals.py rename to tests/unit/router_strategy/adaptive_router/test_signals.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/unit/router_strategy/adaptive_router/test_state_endpoint.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py rename to tests/unit/router_strategy/adaptive_router/test_state_endpoint.py diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py b/tests/unit/router_strategy/adaptive_router/test_update_queue.py similarity index 100% rename from tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py rename to tests/unit/router_strategy/adaptive_router/test_update_queue.py diff --git a/tests/test_litellm/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py similarity index 100% rename from tests/test_litellm/router_strategy/complexity_router/test_context_compaction.py rename to tests/unit/router_strategy/complexity_router/test_context_compaction.py diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/unit/router_strategy/test_auto_router.py similarity index 100% rename from tests/test_litellm/router_strategy/test_auto_router.py rename to tests/unit/router_strategy/test_auto_router.py diff --git a/tests/test_litellm/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py similarity index 100% rename from tests/test_litellm/router_strategy/test_base_routing_strategy.py rename to tests/unit/router_strategy/test_base_routing_strategy.py diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py similarity index 100% rename from tests/test_litellm/router_strategy/test_budget_limiter.py rename to tests/unit/router_strategy/test_budget_limiter.py diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py similarity index 100% rename from tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py rename to tests/unit/router_strategy/test_budget_limiter_hotpath.py diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py similarity index 100% rename from tests/test_litellm/router_strategy/test_complexity_router.py rename to tests/unit/router_strategy/test_complexity_router.py diff --git a/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py b/tests/unit/router_strategy/test_complexity_tier_predictor.py similarity index 100% rename from tests/test_litellm/router_strategy/test_complexity_tier_predictor.py rename to tests/unit/router_strategy/test_complexity_tier_predictor.py diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/unit/router_strategy/test_fuse_presets.py similarity index 100% rename from tests/test_litellm/router_strategy/test_fuse_presets.py rename to tests/unit/router_strategy/test_fuse_presets.py diff --git a/tests/test_litellm/router_strategy/test_lar1_routing.py b/tests/unit/router_strategy/test_lar1_routing.py similarity index 100% rename from tests/test_litellm/router_strategy/test_lar1_routing.py rename to tests/unit/router_strategy/test_lar1_routing.py diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py similarity index 100% rename from tests/test_litellm/router_strategy/test_least_busy.py rename to tests/unit/router_strategy/test_least_busy.py diff --git a/tests/test_litellm/router_strategy/test_litellm_encoder.py b/tests/unit/router_strategy/test_litellm_encoder.py similarity index 100% rename from tests/test_litellm/router_strategy/test_litellm_encoder.py rename to tests/unit/router_strategy/test_litellm_encoder.py diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py similarity index 100% rename from tests/test_litellm/router_strategy/test_llm_v2.py rename to tests/unit/router_strategy/test_llm_v2.py diff --git a/tests/test_litellm/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py similarity index 100% rename from tests/test_litellm/router_strategy/test_lowest_cost.py rename to tests/unit/router_strategy/test_lowest_cost.py diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py similarity index 100% rename from tests/test_litellm/router_strategy/test_lowest_latency.py rename to tests/unit/router_strategy/test_lowest_latency.py diff --git a/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py similarity index 100% rename from tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py rename to tests/unit/router_strategy/test_lowest_tpm_rpm.py diff --git a/tests/test_litellm/router_strategy/test_quality_router.py b/tests/unit/router_strategy/test_quality_router.py similarity index 100% rename from tests/test_litellm/router_strategy/test_quality_router.py rename to tests/unit/router_strategy/test_quality_router.py diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py similarity index 100% rename from tests/test_litellm/router_strategy/test_router_routing_groups.py rename to tests/unit/router_strategy/test_router_routing_groups.py diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/unit/router_strategy/test_router_routing_plugins.py similarity index 100% rename from tests/test_litellm/router_strategy/test_router_routing_plugins.py rename to tests/unit/router_strategy/test_router_routing_plugins.py diff --git a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py b/tests/unit/router_strategy/test_router_tag_regex_routing.py similarity index 100% rename from tests/test_litellm/router_strategy/test_router_tag_regex_routing.py rename to tests/unit/router_strategy/test_router_tag_regex_routing.py diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py similarity index 100% rename from tests/test_litellm/router_strategy/test_router_tag_routing.py rename to tests/unit/router_strategy/test_router_tag_routing.py diff --git a/tests/test_litellm/router_strategy/test_savings_baseline.py b/tests/unit/router_strategy/test_savings_baseline.py similarity index 100% rename from tests/test_litellm/router_strategy/test_savings_baseline.py rename to tests/unit/router_strategy/test_savings_baseline.py diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/unit/router_strategy/test_simple_shuffle.py similarity index 100% rename from tests/test_litellm/router_strategy/test_simple_shuffle.py rename to tests/unit/router_strategy/test_simple_shuffle.py diff --git a/tests/test_litellm/router_strategy/test_stall_detector.py b/tests/unit/router_strategy/test_stall_detector.py similarity index 100% rename from tests/test_litellm/router_strategy/test_stall_detector.py rename to tests/unit/router_strategy/test_stall_detector.py diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index a7006c62438..00462b65bc2 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -565,7 +565,7 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost @pytest.mark.asyncio async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, @@ -589,7 +589,7 @@ async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): @pytest.mark.asyncio async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, timed_with_loop_lags, warm_tokenizer, diff --git a/tests/test_litellm/router_utils/test_access_windows.py b/tests/unit/router_utils/test_access_windows.py similarity index 100% rename from tests/test_litellm/router_utils/test_access_windows.py rename to tests/unit/router_utils/test_access_windows.py diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/unit/router_utils/test_add_retry_fallback_headers.py similarity index 100% rename from tests/test_litellm/router_utils/test_add_retry_fallback_headers.py rename to tests/unit/router_utils/test_add_retry_fallback_headers.py diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/unit/router_utils/test_auto_router_model_naming.py similarity index 100% rename from tests/test_litellm/router_utils/test_auto_router_model_naming.py rename to tests/unit/router_utils/test_auto_router_model_naming.py diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_utils/test_auto_router_tuning_baseline.py similarity index 100% rename from tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py rename to tests/unit/router_utils/test_auto_router_tuning_baseline.py diff --git a/tests/test_litellm/router_utils/test_client_initalization_utils.py b/tests/unit/router_utils/test_client_initalization_utils.py similarity index 100% rename from tests/test_litellm/router_utils/test_client_initalization_utils.py rename to tests/unit/router_utils/test_client_initalization_utils.py diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py similarity index 100% rename from tests/test_litellm/router_utils/test_cooldown_cache.py rename to tests/unit/router_utils/test_cooldown_cache.py diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py similarity index 100% rename from tests/test_litellm/router_utils/test_cooldown_handlers.py rename to tests/unit/router_utils/test_cooldown_handlers.py diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py similarity index 100% rename from tests/test_litellm/router_utils/test_fallback_event_handlers.py rename to tests/unit/router_utils/test_fallback_event_handlers.py diff --git a/tests/test_litellm/router_utils/test_get_retry_from_policy.py b/tests/unit/router_utils/test_get_retry_from_policy.py similarity index 100% rename from tests/test_litellm/router_utils/test_get_retry_from_policy.py rename to tests/unit/router_utils/test_get_retry_from_policy.py diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py similarity index 100% rename from tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py rename to tests/unit/router_utils/test_health_check_allowed_fails_integration.py diff --git a/tests/test_litellm/router_utils/test_health_state_cache.py b/tests/unit/router_utils/test_health_state_cache.py similarity index 100% rename from tests/test_litellm/router_utils/test_health_state_cache.py rename to tests/unit/router_utils/test_health_state_cache.py diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/unit/router_utils/test_pattern_match_deployments.py similarity index 100% rename from tests/test_litellm/router_utils/test_pattern_match_deployments.py rename to tests/unit/router_utils/test_pattern_match_deployments.py diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/unit/router_utils/test_reasoning_effort_capability.py similarity index 100% rename from tests/test_litellm/router_utils/test_reasoning_effort_capability.py rename to tests/unit/router_utils/test_reasoning_effort_capability.py diff --git a/tests/test_litellm/router_utils/test_router_health_check_routing.py b/tests/unit/router_utils/test_router_health_check_routing.py similarity index 100% rename from tests/test_litellm/router_utils/test_router_health_check_routing.py rename to tests/unit/router_utils/test_router_health_check_routing.py diff --git a/tests/test_litellm/router_utils/test_router_interactions_endpoints.py b/tests/unit/router_utils/test_router_interactions_endpoints.py similarity index 100% rename from tests/test_litellm/router_utils/test_router_interactions_endpoints.py rename to tests/unit/router_utils/test_router_interactions_endpoints.py diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py similarity index 100% rename from tests/test_litellm/router_utils/test_router_utils_common_utils.py rename to tests/unit/router_utils/test_router_utils_common_utils.py diff --git a/tests/test_litellm/rust_bridge/AGENTS.md b/tests/unit/rust_bridge/AGENTS.md similarity index 100% rename from tests/test_litellm/rust_bridge/AGENTS.md rename to tests/unit/rust_bridge/AGENTS.md diff --git a/tests/unit/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py index a880cfe3588..1be42e2249d 100644 --- a/tests/unit/rust_bridge/messages/test_route_host.py +++ b/tests/unit/rust_bridge/messages/test_route_host.py @@ -3,6 +3,10 @@ from typing import Final from litellm.rust_bridge.messages.route_host import arguments, response from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest +from dataclasses import astuple +import pytest +import litellm +from litellm.rust_bridge.messages import route_host def test_response_is_a_detached_public_messages_dict() -> None: @@ -40,3 +44,121 @@ def test_arguments_are_the_public_kwargs_view() -> None: ) assert arguments(request) is kwargs + + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +def _flag_model(monkeypatch: pytest.MonkeyPatch, name: str, **flags: bool) -> None: + monkeypatch.setitem( + litellm.model_cost, + name, + { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + **flags, + }, + ) + + +def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeypatch: pytest.MonkeyPatch) -> None: + _flag_model( + monkeypatch, + "claude-test-adaptive", + supports_reasoning=True, + supports_adaptive_thinking=True, + supports_output_config=True, + supports_xhigh_reasoning_effort=True, + supports_sampling_params=False, + ) + + capabilities: Final = route_host.model_capabilities("anthropic/claude-test-adaptive", None) + + assert capabilities.supports_adaptive_thinking + assert capabilities.supports_output_config + assert not capabilities.supports_legacy_thinking + assert not capabilities.supports_sampling_params + assert capabilities.effort_tiers.xhigh + assert not capabilities.effort_tiers.max + + +def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None: + capabilities: Final = route_host.model_capabilities("anthropic/not-a-real-model", None) + + assert capabilities.supports_sampling_params + assert not capabilities.supports_reasoning + assert not capabilities.supports_adaptive_thinking + assert not any(astuple(capabilities.effort_tiers)) + + +@pytest.mark.parametrize( + ("global_flag", "kwargs", "expected"), + [ + (False, {}, False), + (True, {}, True), + (False, {"drop_params": "true"}, True), + (False, {"drop_params": "nonsense"}, False), + (False, {"drop_params": False}, False), + ], +) +def test_drop_params_merges_the_global_flag_with_the_request( + monkeypatch: pytest.MonkeyPatch, global_flag: bool, kwargs: dict[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_flag) + + assert route_host.shaping("anthropic/not-a-real-model", None, kwargs)["drop_params"] is expected + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (["tools[*].input_examples", 3, "metadata.user_id"], ("tools[*].input_examples", "metadata.user_id")), + ("tools", ()), + (None, ()), + ], +) +def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None: + shaping: Final = route_host.shaping("anthropic/not-a-real-model", None, {"additional_drop_params": configured}) + + assert shaping["additional_drop_params"] == expected + + +def test_native_request_rejections_map_to_the_public_400() -> None: + from types import MappingProxyType + + from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest + + request: Final = LiteLLMMessagesRequest( + model="anthropic/claude-sonnet-5", + messages=(), + max_tokens=8, + stream=None, + api_key=None, + api_base=None, + custom_llm_provider=None, + kwargs=MappingProxyType({}), + ) + rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5") + rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = route_host.map_failure(rejected, request, "anthropic") + + assert isinstance(mapped, litellm.BadRequestError) + assert mapped.status_code == 400 + assert "does not support top_k=5" in mapped.message + assert mapped.model == "claude-sonnet-5" + assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError) + + +def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None: + hidden: Final = route_host.stream_hidden_params( + (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41")) + ) + + additional: Final = hidden["additional_headers"] + assert isinstance(additional, dict) + assert additional["llm_provider-request-id"] == "req_upstream_123" + assert additional["x-ratelimit-remaining-requests"] == "41" + assert "request-id" not in additional diff --git a/tests/test_litellm/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py similarity index 100% rename from tests/test_litellm/rust_bridge/messages/test_secrets.py rename to tests/unit/rust_bridge/messages/test_secrets.py diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/unit/rust_bridge/native_route_wheel_test.py similarity index 100% rename from tests/test_litellm/rust_bridge/native_route_wheel_test.py rename to tests/unit/rust_bridge/native_route_wheel_test.py diff --git a/tests/test_litellm/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py similarity index 100% rename from tests/test_litellm/rust_bridge/ocr/test_secrets.py rename to tests/unit/rust_bridge/ocr/test_secrets.py diff --git a/tests/test_litellm/rust_bridge/stubtest.ini b/tests/unit/rust_bridge/stubtest.ini similarity index 100% rename from tests/test_litellm/rust_bridge/stubtest.ini rename to tests/unit/rust_bridge/stubtest.ini diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_bindings.py rename to tests/unit/rust_bridge/test_bindings.py diff --git a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/test_callbacks_legacy_python.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py rename to tests/unit/rust_bridge/test_callbacks_legacy_python.py diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_catalog.py rename to tests/unit/rust_bridge/test_catalog.py diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_configuration.py rename to tests/unit/rust_bridge/test_configuration.py diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/unit/rust_bridge/test_dispatch.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_dispatch.py rename to tests/unit/rust_bridge/test_dispatch.py diff --git a/tests/test_litellm/rust_bridge/test_failures.py b/tests/unit/rust_bridge/test_failures.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_failures.py rename to tests/unit/rust_bridge/test_failures.py diff --git a/tests/test_litellm/rust_bridge/test_fork_guard.py b/tests/unit/rust_bridge/test_fork_guard.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_fork_guard.py rename to tests/unit/rust_bridge/test_fork_guard.py diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/unit/rust_bridge/test_lifecycle.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_lifecycle.py rename to tests/unit/rust_bridge/test_lifecycle.py diff --git a/tests/test_litellm/rust_bridge/test_logger.py b/tests/unit/rust_bridge/test_logger.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_logger.py rename to tests/unit/rust_bridge/test_logger.py diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_runtime.py rename to tests/unit/rust_bridge/test_runtime.py diff --git a/tests/test_litellm/rust_bridge/test_secret_manager.py b/tests/unit/rust_bridge/test_secret_manager.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_secret_manager.py rename to tests/unit/rust_bridge/test_secret_manager.py diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/unit/rust_bridge/test_settings.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_settings.py rename to tests/unit/rust_bridge/test_settings.py diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_token_counter.py rename to tests/unit/rust_bridge/test_token_counter.py diff --git a/tests/test_litellm/rust_bridge/test_tokenizer.py b/tests/unit/rust_bridge/test_tokenizer.py similarity index 95% rename from tests/test_litellm/rust_bridge/test_tokenizer.py rename to tests/unit/rust_bridge/test_tokenizer.py index 188aa81093f..c5093cdb0ce 100644 --- a/tests/test_litellm/rust_bridge/test_tokenizer.py +++ b/tests/unit/rust_bridge/test_tokenizer.py @@ -7,7 +7,7 @@ from tokenizers import Tokenizer from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding from litellm.rust_bridge import tokenizer from litellm.utils import claude_json_str -from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON +from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON TEXTS: Final = ("hello <|endoftext|> world", "café 漢字 🙂", " def f():\n return 1\n", "hello again") diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/unit/rust_bridge/test_verify_linux_native_wheel.py similarity index 100% rename from tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py rename to tests/unit/rust_bridge/test_verify_linux_native_wheel.py From a8fe84bb3e298549edaf6a7abd845d54356b703f Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:20:52 -0700 Subject: [PATCH 32/58] chore(cost-map): add fireworks us-only deepseek v4.1 flash priority prices (#43247) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 ++++++ model_prices_and_context_window.json | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a670200c132..792e7316c67 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60264,13 +60264,16 @@ }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -60360,13 +60363,16 @@ }, "fireworks_ai/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a670200c132..792e7316c67 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60264,13 +60264,16 @@ }, "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -60360,13 +60363,16 @@ }, "fireworks_ai/deepseek-v4p1-flash-us": { "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, From 05f6d97af037c1c9331c1fe1c09370504013c95a Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:34:08 -0700 Subject: [PATCH 33/58] chore(cost-map): sync openrouter prices and add perceptron-mk1.5 (#43246) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 241 ++++++++++-------- model_prices_and_context_window.json | 241 ++++++++++-------- 2 files changed, 272 insertions(+), 210 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 792e7316c67..673e6f15cd3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41477,65 +41477,63 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.44944e-07, + "cache_read_input_token_cost": 3.828e-08, + "input_cost_per_token": 4.5936e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.689888e-06, + "output_cost_per_token": 9.1872e-07, "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.0412e-08, - "supports_audio_input": false, - "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 4.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_vision": true, "supports_pdf_input": false, "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 4.62e-07, + "cache_read_input_token_cost": 8.8e-09, + "input_cost_per_token": 2.64e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.386e-06, + "output_cost_per_token": 7.92e-07, "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.54e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":0.00000132,"output_cost_per_token":0.00000396,"cache_read_input_token_cost":4.4e-8}, - "supports_audio_input": false, - "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, @@ -42755,14 +42753,14 @@ "openrouter/qwen/qwen3-coder-plus": { "cache_creation_input_token_cost": 8.125e-07, "cache_creation_input_token_cost_above_128k_tokens": 2.4375e-06, - "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, - "input_cost_per_token_above_32k_tokens": 1.17e-06, "cache_creation_input_token_cost_above_32k_tokens": 1.4625e-06, - "cache_read_input_token_cost_above_32k_tokens": 2.34e-07, - "output_cost_per_token_above_32k_tokens": 5.85e-06, "cache_read_input_token_cost": 1.3e-07, + "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, + "cache_read_input_token_cost_above_32k_tokens": 2.34e-07, + "deprecation_date": "2026-10-09", "input_cost_per_token": 6.5e-07, "input_cost_per_token_above_128k_tokens": 1.95e-06, + "input_cost_per_token_above_32k_tokens": 1.17e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 65536, @@ -42770,6 +42768,7 @@ "mode": "chat", "output_cost_per_token": 3.25e-06, "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 5.85e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -42802,6 +42801,7 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-235b-a22b-thinking-2507": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, @@ -43099,25 +43099,25 @@ "supports_web_search": false }, "openrouter/z-ai/glm-4.7": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 1.75e-06, "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 8e-08, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, "litellm_provider": "openrouter", "max_input_tokens": 204800, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", + "output_cost_per_token": 2.2e-06, "source": "https://openrouter.ai/api/v1/models", - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_vision": false, - "supports_prompt_caching": true, "supports_assistant_prefill": true, "supports_audio_input": false, + "supports_function_calling": true, "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/z-ai/glm-4.7-flash": { @@ -43162,15 +43162,15 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.1": { - "input_cost_per_token": 9.66e-07, - "output_cost_per_token": 3.036e-06, - "cache_read_input_token_cost": 1.794e-07, "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 1.7914e-07, + "input_cost_per_token": 9.646e-07, "litellm_provider": "openrouter", "max_input_tokens": 204800, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 3.0316e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -66172,23 +66172,23 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { + "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 4.5e-08, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 2.85e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", + "output_cost_per_token": 1.4e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { @@ -66349,24 +66349,24 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 3e-08, - "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 2.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", + "output_cost_per_token": 3.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_response_schema": true, "supports_parallel_function_calling": true, "supports_pdf_input": false, - "supports_vision": false, "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/qwen/qwen3.7-flash": { @@ -66839,46 +66839,47 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-max-preview": { - "input_cost_per_token": 1.027e-06, - "output_cost_per_token": 6.162e-06, "cache_creation_input_token_cost": 1.28375e-06, - "input_cost_per_token_above_128k_tokens": 1.58e-06, - "output_cost_per_token_above_128k_tokens": 9.48e-06, "cache_creation_input_token_cost_above_128k_tokens": 1.975e-06, + "deprecation_date": "2026-10-09", + "input_cost_per_token": 1.027e-06, + "input_cost_per_token_above_128k_tokens": 1.58e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 6.162e-06, + "output_cost_per_token_above_128k_tokens": 9.48e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, "openrouter/qwen/qwen3.6-27b": { - "input_cost_per_token": 3.2e-07, - "output_cost_per_token": 2.7e-06, "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 3.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262140, "max_tokens": 262140, "mode": "chat", + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/openai/gpt-5.5-pro": { @@ -66923,23 +66924,23 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 4.9e-08, - "output_cost_per_token": 9.8e-08, - "cache_read_input_token_cost": 9.8e-09, + "cache_read_input_token_cost": 9.408e-09, + "input_cost_per_token": 4.704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", + "output_cost_per_token": 9.408e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.6": { @@ -66964,22 +66965,22 @@ "supports_web_search": false }, "openrouter/google/gemma-4-26b-a4b-it": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 3.75e-08, + "input_cost_per_token": 6.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, "max_tokens": 235929, "mode": "chat", + "output_cost_per_token": 2.25e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": true, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67264,25 +67265,26 @@ "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 7.8e-07, - "input_cost_per_token_above_32k_tokens": 1.56e-06, - "output_cost_per_token_above_32k_tokens": 7.8e-06, - "output_cost_per_token": 3.9e-06, "input_cost_per_token_above_128k_tokens": 1.95e-06, - "output_cost_per_token_above_128k_tokens": 9.75e-06, + "input_cost_per_token_above_32k_tokens": 1.56e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.9e-06, + "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 7.8e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -67534,59 +67536,62 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-vl-32b-instruct": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.04e-07, - "output_cost_per_token": 4.16e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4.16e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-8b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.8e-07, - "output_cost_per_token": 2.1e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.1e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-8b-instruct": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.17e-07, - "output_cost_per_token": 4.55e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4.55e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67616,40 +67621,41 @@ "supports_web_search": true }, "openrouter/qwen/qwen3-vl-30b-a3b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.4e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 5.2e-07, + "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67673,21 +67679,22 @@ "supports_web_search": true }, "openrouter/qwen/qwen3-vl-235b-a22b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67712,32 +67719,33 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-max": { - "input_cost_per_token": 7.8e-07, - "input_cost_per_token_above_32k_tokens": 1.56e-06, - "cache_creation_input_token_cost_above_32k_tokens": 1.95e-06, - "cache_read_input_token_cost_above_32k_tokens": 3.12e-07, - "output_cost_per_token_above_32k_tokens": 7.8e-06, - "output_cost_per_token": 3.9e-06, - "cache_read_input_token_cost": 1.56e-07, "cache_creation_input_token_cost": 9.75e-07, - "input_cost_per_token_above_128k_tokens": 1.95e-06, - "output_cost_per_token_above_128k_tokens": 9.75e-06, - "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, "cache_creation_input_token_cost_above_128k_tokens": 2.4375e-06, + "cache_creation_input_token_cost_above_32k_tokens": 1.95e-06, + "cache_read_input_token_cost": 1.56e-07, + "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, + "cache_read_input_token_cost_above_32k_tokens": 3.12e-07, + "deprecation_date": "2026-10-09", + "input_cost_per_token": 7.8e-07, + "input_cost_per_token_above_128k_tokens": 1.95e-06, + "input_cost_per_token_above_32k_tokens": 1.56e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.9e-06, + "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 7.8e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_vision": false, "supports_prompt_caching": true, "supports_reasoning": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v3.1-terminus": { @@ -67832,23 +67840,24 @@ "openrouter/qwen/qwen-plus-2025-07-28": { "cache_creation_input_token_cost": 3.25e-07, "cache_read_input_token_cost": 5.2e-08, + "deprecation_date": "2026-10-09", "input_cost_per_token": 2.6e-07, - "output_cost_per_token": 7.8e-07, "input_cost_per_token_above_256k_tokens": 7.8e-07, - "output_cost_per_token_above_256k_tokens": 2.34e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_above_256k_tokens": 2.34e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -67872,21 +67881,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-30b-a3b-thinking-2507": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.4e-06, "litellm_provider": "openrouter", "max_input_tokens": 81920, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -68195,21 +68205,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-8b": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.17e-07, - "output_cost_per_token": 4.55e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", + "output_cost_per_token": 4.55e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -68252,21 +68263,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-235b-a22b": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 4.55e-07, - "output_cost_per_token": 1.82e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", + "output_cost_per_token": 1.82e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -73019,14 +73031,14 @@ "supports_web_search": false }, "openrouter/inclusionai/ling-3.0-flash-vl": { - "cache_read_input_token_cost": 1.2e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-08, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.8e-07, + "output_cost_per_token": 6.16e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -77162,5 +77174,24 @@ "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, "supports_tool_choice": false + }, + "openrouter/perceptron/perceptron-mk1.5": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 36864, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 792e7316c67..673e6f15cd3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41477,65 +41477,63 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.44944e-07, + "cache_read_input_token_cost": 3.828e-08, + "input_cost_per_token": 4.5936e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.689888e-06, + "output_cost_per_token": 9.1872e-07, "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.0412e-08, - "supports_audio_input": false, - "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, + "output_cost_per_token": 4.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_vision": true, "supports_pdf_input": false, "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 4.62e-07, + "cache_read_input_token_cost": 8.8e-09, + "input_cost_per_token": 2.64e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.386e-06, + "output_cost_per_token": 7.92e-07, "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.54e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":0.00000132,"output_cost_per_token":0.00000396,"cache_read_input_token_cost":4.4e-8}, - "supports_audio_input": false, - "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, @@ -42755,14 +42753,14 @@ "openrouter/qwen/qwen3-coder-plus": { "cache_creation_input_token_cost": 8.125e-07, "cache_creation_input_token_cost_above_128k_tokens": 2.4375e-06, - "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, - "input_cost_per_token_above_32k_tokens": 1.17e-06, "cache_creation_input_token_cost_above_32k_tokens": 1.4625e-06, - "cache_read_input_token_cost_above_32k_tokens": 2.34e-07, - "output_cost_per_token_above_32k_tokens": 5.85e-06, "cache_read_input_token_cost": 1.3e-07, + "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, + "cache_read_input_token_cost_above_32k_tokens": 2.34e-07, + "deprecation_date": "2026-10-09", "input_cost_per_token": 6.5e-07, "input_cost_per_token_above_128k_tokens": 1.95e-06, + "input_cost_per_token_above_32k_tokens": 1.17e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 65536, @@ -42770,6 +42768,7 @@ "mode": "chat", "output_cost_per_token": 3.25e-06, "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 5.85e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -42802,6 +42801,7 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-235b-a22b-thinking-2507": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, @@ -43099,25 +43099,25 @@ "supports_web_search": false }, "openrouter/z-ai/glm-4.7": { - "input_cost_per_token": 4e-07, - "output_cost_per_token": 1.75e-06, "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 8e-08, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, "litellm_provider": "openrouter", "max_input_tokens": 204800, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", + "output_cost_per_token": 2.2e-06, "source": "https://openrouter.ai/api/v1/models", - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_vision": false, - "supports_prompt_caching": true, "supports_assistant_prefill": true, "supports_audio_input": false, + "supports_function_calling": true, "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/z-ai/glm-4.7-flash": { @@ -43162,15 +43162,15 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.1": { - "input_cost_per_token": 9.66e-07, - "output_cost_per_token": 3.036e-06, - "cache_read_input_token_cost": 1.794e-07, "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 1.7914e-07, + "input_cost_per_token": 9.646e-07, "litellm_provider": "openrouter", "max_input_tokens": 204800, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token": 3.0316e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -66172,23 +66172,23 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { + "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 4.5e-08, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 2.85e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", + "output_cost_per_token": 1.4e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-vision-exp": { @@ -66349,24 +66349,24 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 3e-08, - "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 2.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", + "output_cost_per_token": 3.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, - "supports_tool_choice": true, - "supports_reasoning": true, - "supports_response_schema": true, "supports_parallel_function_calling": true, "supports_pdf_input": false, - "supports_vision": false, "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/qwen/qwen3.7-flash": { @@ -66839,46 +66839,47 @@ "supports_web_search": false }, "openrouter/qwen/qwen3.6-max-preview": { - "input_cost_per_token": 1.027e-06, - "output_cost_per_token": 6.162e-06, "cache_creation_input_token_cost": 1.28375e-06, - "input_cost_per_token_above_128k_tokens": 1.58e-06, - "output_cost_per_token_above_128k_tokens": 9.48e-06, "cache_creation_input_token_cost_above_128k_tokens": 1.975e-06, + "deprecation_date": "2026-10-09", + "input_cost_per_token": 1.027e-06, + "input_cost_per_token_above_128k_tokens": 1.58e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 6.162e-06, + "output_cost_per_token_above_128k_tokens": 9.48e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, "openrouter/qwen/qwen3.6-27b": { - "input_cost_per_token": 3.2e-07, - "output_cost_per_token": 2.7e-06, "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 3.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262140, "max_tokens": 262140, "mode": "chat", + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/openai/gpt-5.5-pro": { @@ -66923,23 +66924,23 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 4.9e-08, - "output_cost_per_token": 9.8e-08, - "cache_read_input_token_cost": 9.8e-09, + "cache_read_input_token_cost": 9.408e-09, + "input_cost_per_token": 4.704e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", + "output_cost_per_token": 9.408e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, - "supports_prompt_caching": true, "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.6": { @@ -66964,22 +66965,22 @@ "supports_web_search": false }, "openrouter/google/gemma-4-26b-a4b-it": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 3.75e-08, + "input_cost_per_token": 6.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, "max_tokens": 235929, "mode": "chat", + "output_cost_per_token": 2.25e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": true, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67264,25 +67265,26 @@ "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 7.8e-07, - "input_cost_per_token_above_32k_tokens": 1.56e-06, - "output_cost_per_token_above_32k_tokens": 7.8e-06, - "output_cost_per_token": 3.9e-06, "input_cost_per_token_above_128k_tokens": 1.95e-06, - "output_cost_per_token_above_128k_tokens": 9.75e-06, + "input_cost_per_token_above_32k_tokens": 1.56e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.9e-06, + "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 7.8e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -67534,59 +67536,62 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-vl-32b-instruct": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.04e-07, - "output_cost_per_token": 4.16e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4.16e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-8b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.8e-07, - "output_cost_per_token": 2.1e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.1e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-8b-instruct": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.17e-07, - "output_cost_per_token": 4.55e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4.55e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67616,40 +67621,41 @@ "supports_web_search": true }, "openrouter/qwen/qwen3-vl-30b-a3b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.4e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 5.2e-07, + "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 6e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67673,21 +67679,22 @@ "supports_web_search": true }, "openrouter/qwen/qwen3-vl-235b-a22b-thinking": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 4e-07, - "output_cost_per_token": 4e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false }, @@ -67712,32 +67719,33 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-max": { - "input_cost_per_token": 7.8e-07, - "input_cost_per_token_above_32k_tokens": 1.56e-06, - "cache_creation_input_token_cost_above_32k_tokens": 1.95e-06, - "cache_read_input_token_cost_above_32k_tokens": 3.12e-07, - "output_cost_per_token_above_32k_tokens": 7.8e-06, - "output_cost_per_token": 3.9e-06, - "cache_read_input_token_cost": 1.56e-07, "cache_creation_input_token_cost": 9.75e-07, - "input_cost_per_token_above_128k_tokens": 1.95e-06, - "output_cost_per_token_above_128k_tokens": 9.75e-06, - "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, "cache_creation_input_token_cost_above_128k_tokens": 2.4375e-06, + "cache_creation_input_token_cost_above_32k_tokens": 1.95e-06, + "cache_read_input_token_cost": 1.56e-07, + "cache_read_input_token_cost_above_128k_tokens": 3.9e-07, + "cache_read_input_token_cost_above_32k_tokens": 3.12e-07, + "deprecation_date": "2026-10-09", + "input_cost_per_token": 7.8e-07, + "input_cost_per_token_above_128k_tokens": 1.95e-06, + "input_cost_per_token_above_32k_tokens": 1.56e-06, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_token": 3.9e-06, + "output_cost_per_token_above_128k_tokens": 9.75e-06, + "output_cost_per_token_above_32k_tokens": 7.8e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_vision": false, "supports_prompt_caching": true, "supports_reasoning": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v3.1-terminus": { @@ -67832,23 +67840,24 @@ "openrouter/qwen/qwen-plus-2025-07-28": { "cache_creation_input_token_cost": 3.25e-07, "cache_read_input_token_cost": 5.2e-08, + "deprecation_date": "2026-10-09", "input_cost_per_token": 2.6e-07, - "output_cost_per_token": 7.8e-07, "input_cost_per_token_above_256k_tokens": 7.8e-07, - "output_cost_per_token_above_256k_tokens": 2.34e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_above_256k_tokens": 2.34e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": false, - "supports_tool_choice": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -67872,21 +67881,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-30b-a3b-thinking-2507": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 2e-07, - "output_cost_per_token": 2.4e-06, "litellm_provider": "openrouter", "max_input_tokens": 81920, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", + "output_cost_per_token": 2.4e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -68195,21 +68205,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-8b": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 1.17e-07, - "output_cost_per_token": 4.55e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", + "output_cost_per_token": 4.55e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -68252,21 +68263,22 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-235b-a22b": { + "deprecation_date": "2026-10-09", "input_cost_per_token": 4.55e-07, - "output_cost_per_token": 1.82e-06, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", + "output_cost_per_token": 1.82e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, "supports_prompt_caching": false, - "supports_tool_choice": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false }, @@ -73019,14 +73031,14 @@ "supports_web_search": false }, "openrouter/inclusionai/ling-3.0-flash-vl": { - "cache_read_input_token_cost": 1.2e-08, - "input_cost_per_token": 6e-08, + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-08, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.8e-07, + "output_cost_per_token": 6.16e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -77162,5 +77174,24 @@ "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, "supports_tool_choice": false + }, + "openrouter/perceptron/perceptron-mk1.5": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 36864, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false } } From a7f731da22b4642a7e59b13079aee830ce4a35f5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 00:37:55 +0000 Subject: [PATCH 34/58] fix(cost-map): correct fireworks_ai deepseek-v4p1-flash pricing (#43253) The V4.1 Flash rows carried the DeepSeek V4 Flash (0731) prices (0.22/0.66, cache 0.007). Fireworks lists V4.1 Flash at 0.30/1.20 with 0.006 cache read (priority 0.375/1.50/0.0075). Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 28 +++++++++---------- model_prices_and_context_window.json | 28 +++++++++---------- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 673e6f15cd3..59c0ab56a27 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60243,18 +60243,18 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { - "cache_read_input_token_cost": 7e-09, - "cache_read_input_token_cost_priority": 8.75e-09, - "input_cost_per_token": 2.2e-07, - "input_cost_per_token_priority": 2.75e-07, + "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost_priority": 7.5e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 6.6e-07, - "output_cost_per_token_priority": 8.25e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -60342,18 +60342,18 @@ "supports_vision": false }, "fireworks_ai/deepseek-v4p1-flash": { - "cache_read_input_token_cost": 7e-09, - "cache_read_input_token_cost_priority": 8.75e-09, - "input_cost_per_token": 2.2e-07, - "input_cost_per_token_priority": 2.75e-07, + "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost_priority": 7.5e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 6.6e-07, - "output_cost_per_token_priority": 8.25e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 673e6f15cd3..59c0ab56a27 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60243,18 +60243,18 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { - "cache_read_input_token_cost": 7e-09, - "cache_read_input_token_cost_priority": 8.75e-09, - "input_cost_per_token": 2.2e-07, - "input_cost_per_token_priority": 2.75e-07, + "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost_priority": 7.5e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 6.6e-07, - "output_cost_per_token_priority": 8.25e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -60342,18 +60342,18 @@ "supports_vision": false }, "fireworks_ai/deepseek-v4p1-flash": { - "cache_read_input_token_cost": 7e-09, - "cache_read_input_token_cost_priority": 8.75e-09, - "input_cost_per_token": 2.2e-07, - "input_cost_per_token_priority": 2.75e-07, + "cache_read_input_token_cost": 6e-09, + "cache_read_input_token_cost_priority": 7.5e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 6.6e-07, - "output_cost_per_token_priority": 8.25e-07, - "source": "https://api.fireworks.ai/v1/serverless/models", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_priority": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, From 797fddf59c87f9a8cfd626b965ac11062253f3ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:39:21 -0700 Subject: [PATCH 35/58] feat(openrouter): add typesafe/jev-router to the cost map (#43248) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 17 +++++++++++++++++ model_prices_and_context_window.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 59c0ab56a27..8ebef7ae533 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71347,6 +71347,23 @@ "output_cost_per_token": 0.0, "source": "https://openrouter.ai/typesafe/jev-1.13" }, + "openrouter/typesafe/jev-router": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "source": "https://openrouter.ai/typesafe/jev-router", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 59c0ab56a27..8ebef7ae533 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71347,6 +71347,23 @@ "output_cost_per_token": 0.0, "source": "https://openrouter.ai/typesafe/jev-1.13" }, + "openrouter/typesafe/jev-router": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "source": "https://openrouter.ai/typesafe/jev-router", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", From 8694c3cb4b11358bd899fa34e4c5ca180039fe7e Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:43:37 -0700 Subject: [PATCH 36/58] chore(cost-map): add fireworks priority prices for muse glimmer 30b and deepseek v4 flash vision exp (#43252) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++++++++ model_prices_and_context_window.json | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8ebef7ae533..56129fef135 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60284,13 +60284,16 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, + "output_cost_per_token_priority": 8.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -60383,13 +60386,16 @@ }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, + "output_cost_per_token_priority": 8.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -60518,14 +60524,17 @@ }, "fireworks_ai/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "cache_read_input_token_cost_priority": 6e-08, "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -60567,14 +60576,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "cache_read_input_token_cost_priority": 6e-08, "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8ebef7ae533..56129fef135 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60284,13 +60284,16 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, + "output_cost_per_token_priority": 8.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -60383,13 +60386,16 @@ }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "cache_read_input_token_cost_priority": 8.75e-09, "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 6.6e-07, + "output_cost_per_token_priority": 8.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -60518,14 +60524,17 @@ }, "fireworks_ai/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "cache_read_input_token_cost_priority": 6e-08, "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -60567,14 +60576,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "cache_read_input_token_cost_priority": 6e-08, "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_priority": 2.25e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, From e9491d31b5881ae991ffa537fe74f3474379f8ce Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 01:09:36 +0000 Subject: [PATCH 37/58] refactor(rust): move credential inheritance and the SDK limits out of the legacy callback crate into a driver preflight (#43259) The legacy callback crate carried two rewrites that have nothing to do with the Logging contract: litellm_credential_name inheritance and the max_budget and num_retries_per_request checks. Any later callback host would need them unchanged, which is the smell the crate's AGENTS.md now names. They are now a Preflight the driver in litellm-host-python runs on the keyword view begin returned, before the host projects from it, supplied by python-bridge and passed through run_legacy_call. The call order is unchanged (setup, deployment hook, credentials, limits) and a rejection still fails the call as a host failure, so the failure callbacks run as before. The preflight rewrites the adapter's own copy in place, so no extra dict copy and no new lifecycle method Co-authored-by: Yujong Lee Co-authored-by: Claude Fable 5.1 --- litellm-rust/Cargo.lock | 1 + .../crates/callbacks-legacy-python/AGENTS.md | 24 +-- .../python_contract.json | 8 - .../callbacks-legacy-python/src/adapter.rs | 53 +---- .../callbacks-legacy-python/src/call.rs | 10 +- .../crates/callbacks-legacy-python/src/lib.rs | 17 +- .../callbacks-legacy-python/src/python.rs | 10 +- litellm-rust/crates/host-python/AGENTS.md | 2 +- .../crates/host-python/src/adapter.rs | 6 + litellm-rust/crates/host-python/src/driver.rs | 111 +++++++++- litellm-rust/crates/host-python/src/lib.rs | 3 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../python-bridge/preflight_contract.json | 10 + litellm-rust/crates/python-bridge/src/lib.rs | 1 + .../src/preflight.rs} | 203 ++++++++++++++++-- .../python-bridge/src/routes/messages/mod.rs | 1 + .../python-bridge/src/routes/ocr/mod.rs | 1 + .../rust_bridge/callbacks_legacy_python.py | 32 --- litellm/rust_bridge/preflight.py | 45 ++++ .../test_callbacks_legacy_python.py | 30 +-- tests/unit/rust_bridge/test_preflight.py | 63 ++++++ 21 files changed, 460 insertions(+), 172 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/preflight_contract.json rename litellm-rust/crates/{callbacks-legacy-python/src/preparation.rs => python-bridge/src/preflight.rs} (57%) create mode 100644 litellm/rust_bridge/preflight.py create mode 100644 tests/unit/rust_bridge/test_preflight.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 3677d1d654f..ffcc6a5496b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3311,6 +3311,7 @@ dependencies = [ "serde_json", "serde_with", "sha2 0.10.9", + "strum", "thiserror 2.0.19", "tokio", "tokio-tungstenite", diff --git a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md index 8b2e1c15f6e..de6e0c1b225 100644 --- a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md @@ -1,19 +1,19 @@ - Target invariants, not completion claims -- Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits) +- This crate is the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) + - Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here + - SDK request policy (credential inheritance, the budget and retry-count limits) is the driver's preflight, supplied by `python-bridge`; this crate only adopts the keyword view it produces - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call -- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`) +- Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json` - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` - - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy -- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does - - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's +- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; routes hand it over through `run_legacy_call` and keep no copy +- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`; which callbacks run is `Logging`'s decision, never this crate's - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view -- Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once + - Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies; preserve the deployment-hook kwargs view + - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object, resolved through `litellm_host_python::lookup` + - Retain body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only +- Success and failure handlers receive the exact selected public response or exception + - A failure-handler error cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch + - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy releases deferred success at most once - Delivery follows the registry, not the callable's type: direct, awaited, executor-submitted, logging-worker and deferred paths stay distinct - Traverse every retained Python edge; `close` is idempotent and restores the correlation context once diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index 8a7f3b98f47..9ed13ae5ed5 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -6,9 +6,6 @@ "start_time", "asynchronous" ], - "check_limits": [ - "kwargs" - ], "finalize": [ "response", "logger", @@ -76,11 +73,6 @@ ], "custom_pricing_fields": [], "is_internal_call": [], - "credential_list": [], - "warn_unknown_credential": [ - "name", - "loaded" - ], "before_deployment_call": [ "kwargs", "call_type" diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 75a635e9c63..718cc615f30 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -19,7 +19,7 @@ use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, prepare, + finalize, is_internal_call, python::Streaming, setup, }; @@ -117,9 +117,13 @@ impl LegacyLogging { }) } + /// The keyword view the rest of the call reads: a copy, so the deployment hook's own + /// dict is left as the hook returned it, carrying the logger as `@client` injects it. + /// The driver's preflight rewrites this same dict before the host projects from it. fn prepare(&mut self, py: Python<'_>) -> PyResult { - let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind(); - self.call.set_kwargs(prepared); + let prepared = self.call.kwargs().bind(py).copy()?; + prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?; + self.call.set_kwargs(prepared.unbind()); Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py))) } @@ -580,8 +584,6 @@ assert prepared['document'] is replacement assert prepared['pages'] is replaced_kwargs['pages'] assert prepared['litellm_logging_obj'] is logger assert 'litellm_logging_obj' not in replaced_kwargs -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked is prepared ", ); }); @@ -616,8 +618,6 @@ kwargs = {'logger': logger, 'vendor_extension': opaque} &locals, c" assert prepared['vendor_extension'] is opaque -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked['vendor_extension'] is opaque assert hooked == ([opaque] if asynchronous else []), hooked ", ); @@ -733,45 +733,6 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h ); }); } - - #[rstest] - #[case::synchronous(false)] - #[case::asynchronous(true)] - fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -class BudgetExceeded(Exception): - pass - -rejection = BudgetExceeded('over budget') - -class LimitedLogger(StubLogger): - def check_limits(self, arguments): - raise rejection - -logger = LimitedLogger() -logger.hooks = {'pre': lambda kwargs: kwargs} -kwargs = {'logger': logger} -", - ); - let mut logging = legacy_call(py, &locals, asynchronous); - let kwargs = local(&locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - LifecycleStep::Await(_) => { - logging.resume(py, Ok(local(&locals, "kwargs").unbind())) - } - step => Ok(step), - }); - let error = result.err().unwrap(); - assert!(error.value(py).is(local(&locals, "rejection"))); - }); - } } #[cfg(test)] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/call.rs b/litellm-rust/crates/callbacks-legacy-python/src/call.rs index 9b921070839..3fa638ac6d3 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/call.rs @@ -4,7 +4,7 @@ //! this crate holds them. use litellm_host::{machine::Machine, protocol::Protocol}; -use litellm_host_python::{ProtocolHost, lookup, run_call}; +use litellm_host_python::{Preflight, ProtocolHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -39,7 +39,8 @@ impl PublicCall { } /// The keyword view the legacy path currently reads: the caller's copy until - /// `function_setup`, then each rewrite (setup, deployment hook, prepare) in turn. + /// `function_setup`, then each rewrite (setup, deployment hook, the driver's preflight) + /// in turn. pub(crate) fn kwargs(&self) -> &Py { &self.kwargs } @@ -64,13 +65,15 @@ impl PublicCall { } /// Runs one native call under the legacy `Logging` contract: the protocol host projects from -/// the keyword view the contract prepares, and the contract observes the call. +/// the keyword view the contract prepares and `preflight` rewrites, and the contract +/// observes the call. pub fn run_legacy_call( py: Python<'_>, surface: LegacySurface, call: PublicCall, machine: M, host: H, + preflight: Preflight, asynchronous: bool, ) -> PyResult> where @@ -83,6 +86,7 @@ where machine, host, Box::new(LegacyLogging::new(py, surface, call, asynchronous)), + preflight, arguments, asynchronous, ) diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 69f72fbc177..869c534acf4 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -1,9 +1,10 @@ //! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the -//! sync and async callback registries it fans out to, the deployment hooks, the deferred -//! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name -//! inheritance, budget and retry-count limits). All of it sits behind one +//! sync and async callback registries it fans out to, the deployment hooks and the deferred +//! proxy release. All of it sits behind one //! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and -//! core never learn which Python object is on the other end. +//! core never learn which Python object is on the other end. The SDK's own request policy +//! (credential inheritance, the budget and retry limits) is the driver's preflight, not this +//! crate's. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] //! is where those objects live, and [`run_legacy_call`] is how a route hands them over @@ -14,14 +15,12 @@ mod call; mod callbacks; mod deferred; mod logger; -mod preparation; mod python; pub(crate) use adapter::LegacyLogging; pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; -pub(crate) use preparation::prepare; #[cfg(test)] mod test_support { @@ -77,7 +76,6 @@ FAKES = { logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], kwargs=kwargs, ), - 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( kwargs=kwargs, @@ -104,8 +102,6 @@ FAKES = { 'restore_context': lambda logger: logger.record('restore', None), 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), 'is_internal_call': lambda: legacy.is_internal.get(), - 'credential_list': lambda: [], - 'warn_unknown_credential': lambda name, loaded: None, 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( 'success', response, call_type @@ -162,9 +158,6 @@ class StubLogger: self.record(phase + '_hook', call_type) return self.hooks.get(phase, lambda value: 'awaitable')(value) - def check_limits(self, arguments): - self.record('check_limits', arguments) - def failure_handler(self, error, trace, start, end): self.record('failure_handler', error) diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index cb609d52878..47331f369f5 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -19,18 +19,12 @@ pub(crate) enum LegacyPython { Streaming(Streaming), } -/// The `@client` wrapper around the call: `function_setup`, limits, credentials, -/// response metadata and the correlation context. +/// The `@client` wrapper around the call: `function_setup`, response metadata and the +/// correlation context. #[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] pub(crate) enum Wrapper { #[strum(serialize = "setup")] Setup, - #[strum(serialize = "check_limits")] - CheckLimits, - #[strum(serialize = "credential_list")] - CredentialList, - #[strum(serialize = "warn_unknown_credential")] - WarnUnknownCredential, #[strum(serialize = "is_internal_call")] IsInternalCall, #[strum(serialize = "finalize")] diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 7c1919f9f39..cadc55a35a7 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -2,7 +2,7 @@ - Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`ProtocolHost` traits - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - - `ProtocolHost::project` receives the keyword view the adapter's `begin` returned, not the caller's dict; a protocol host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) + - `ProtocolHost::project` receives the keyword view the adapter's `begin` returned, rewritten in place by the route's `Preflight`, not the caller's dict; a protocol host that projects from it inherits the adapter's rewrites (for the legacy adapter: setup, deployment hooks) and the preflight's (credential inheritance) - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 7f07475bc4c..83ed6416d6e 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -9,6 +9,12 @@ pub fn missing_state() -> PyErr { PyRuntimeError::new_err("missing native call state") } +/// The SDK's request policy, run by the driver on the keyword view `begin` returned and +/// before the protocol host projects from it. It rewrites that view in place, so the +/// lifecycle that returned it sees the rewrite too; a rejection fails the call as a host +/// failure, so the lifecycle still observes it. +pub type Preflight = fn(Python<'_>, &Bound<'_, PyDict>) -> PyResult<()>; + /// What an adapter step produced: either the value the driver asked for, or a Python /// awaitable the driver hands back to the caller's task before asking again. pub enum LifecycleStep { diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 372af2843bd..50eae1e0225 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -14,7 +14,8 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, Preflight, ProtocolHost, PythonLifecycle, + missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -83,6 +84,7 @@ where { host: H, adapter: Box, + preflight: Preflight, machine: Option>>>, arguments: Option>, started_at: f64, @@ -95,12 +97,14 @@ where } /// Runs one native call for Python: synchronously, or as a coroutine that awaits every -/// host suspension inline in the caller's task. +/// host suspension inline in the caller's task. `preflight` runs once, on the keyword view +/// the adapter's `begin` returned, before the host projects from it. pub fn run_call( py: Python<'_>, machine: M, host: H, adapter: Box, + preflight: Preflight, arguments: Py, asynchronous: bool, ) -> PyResult> @@ -111,6 +115,7 @@ where let mut driver = PythonDriver { host, adapter, + preflight, machine: Some(Arc::new(Mutex::new(MachineState { machine, result: None, @@ -213,6 +218,9 @@ where match (expect, step) { (Expect::Started, LifecycleStep::Done) => self.begin(py), (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { + if let Err(error) = (self.preflight)(py, arguments.bind(py)) { + return self.adapter_failed(py, error); + } self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) @@ -869,6 +877,21 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri host: SyntheticHost, script: AdapterScript, asynchronous: bool, + ) -> (PyResult>, Vec) { + run_preflighted(py, machine, host, script, no_preflight, asynchronous) + } + + fn no_preflight(_: Python<'_>, _: &Bound<'_, PyDict>) -> PyResult<()> { + Ok(()) + } + + fn run_preflighted( + py: Python<'_>, + machine: CallMachine, + host: SyntheticHost, + script: AdapterScript, + preflight: Preflight, + asynchronous: bool, ) -> (PyResult>, Vec) { let log = Log(host.log.0.clone()); let adapter = SyntheticAdapter { @@ -882,6 +905,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri machine, host, Box::new(adapter), + preflight, arguments.unbind(), asynchronous, ); @@ -1088,6 +1112,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri streaming_machine(), StreamingHost, Box::new(adapter), + no_preflight, PyDict::new(py).unbind(), asynchronous, ) @@ -1291,6 +1316,87 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + /// The rejection a preflight raised, kept so a test can check the caller receives that + /// exact object. A `Preflight` is a plain `fn`, so it cannot capture one itself. + static REJECTION: Mutex>> = Mutex::new(None); + + fn rejecting_preflight(py: Python<'_>, _: &Bound<'_, PyDict>) -> PyResult<()> { + let error = PyValueError::new_err("over budget"); + *REJECTION.lock().unwrap() = Some(error.value(py).clone().unbind()); + Err(error) + } + + fn inheriting_preflight(_: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> { + arguments.set_item("api_key", "inherited") + } + + #[test] + fn a_preflight_rejection_is_the_callers_error_and_the_machine_never_starts() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, log) = run_preflighted( + py, + success_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: false, + }, + AdapterScript::Plain, + rejecting_preflight, + asynchronous, + ); + let error = result.unwrap_err(); + let raised = REJECTION.lock().unwrap().take().unwrap(); + assert!(error.value(py).is(&raised)); + assert_eq!( + log, + [ + "started", + "begin", + "failed:Host:over budget", + "adapter.close", + "host.close" + ] + ); + } + }); + } + + #[test] + fn the_host_projects_from_the_keyword_view_the_preflight_rewrote() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, _) = run_preflighted( + py, + success_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: false, + }, + AdapterScript::Plain, + inheriting_preflight, + asynchronous, + ); + assert_eq!( + result.unwrap().extract::(py).unwrap(), + "project:2|sign|rewritten" + ); + } + }); + } + #[test] fn the_adapters_finalized_response_is_what_the_call_returns_and_reports() { let _guard = PYTHON_GLOBALS @@ -1412,6 +1518,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri success_machine(), host, Box::new(adapter), + no_preflight, PyDict::new(py).unbind(), false, ) diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 7e17c4da51e..2f9e37fe968 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -15,7 +15,8 @@ mod handle; mod marshal; pub use adapter::{ - InvokeError, LifecycleEvent, LifecycleStep, ProtocolHost, PythonLifecycle, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, Preflight, ProtocolHost, PythonLifecycle, + missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a02adfaa064..057cad2f42e 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -55,6 +55,7 @@ pyo3-async-runtimes.workspace = true reqwest.workspace = true redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true +strum.workspace = true veil.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } diff --git a/litellm-rust/crates/python-bridge/preflight_contract.json b/litellm-rust/crates/python-bridge/preflight_contract.json new file mode 100644 index 00000000000..343dea268cd --- /dev/null +++ b/litellm-rust/crates/python-bridge/preflight_contract.json @@ -0,0 +1,10 @@ +{ + "credential_list": [], + "warn_unknown_credential": [ + "name", + "loaded" + ], + "check_limits": [ + "kwargs" + ] +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7c814f540a8..51e112fa1be 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -6,6 +6,7 @@ mod errors; mod http; mod logger; mod marshal; +mod preflight; mod python_settings; mod routes; mod secrets; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/preparation.rs b/litellm-rust/crates/python-bridge/src/preflight.rs similarity index 57% rename from litellm-rust/crates/callbacks-legacy-python/src/preparation.rs rename to litellm-rust/crates/python-bridge/src/preflight.rs index aab654c9893..34813672c09 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/preparation.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -1,9 +1,51 @@ +//! The SDK's request policy the driver runs on every route's keyword view before the host +//! projects from it: credential-name inheritance from `litellm.credential_list`, then the +//! budget and retry-count limits. It is the `@client` prologue after `function_setup` and the +//! deployment hook, and belongs to no callback contract. + use pyo3::{ prelude::*, types::{PyDict, PyList}, }; +use strum::{IntoStaticStr, VariantArray}; -use crate::python::Wrapper; +const MODULE: &str = "litellm.rust_bridge.preflight"; + +/// The litellm globals the preflight still reads through Python. `preflight_contract.json` +/// pins each function's parameters on both sides. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum PythonPreflight { + #[strum(serialize = "credential_list")] + CredentialList, + #[strum(serialize = "warn_unknown_credential")] + WarnUnknownCredential, + #[strum(serialize = "check_limits")] + CheckLimits, +} + +impl PythonPreflight { + fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + py.import(MODULE)?.getattr(<&str>::from(self))?.call1(args) + } +} + +#[cfg(test)] +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.json"); + +/// Rewrites `arguments` in place, in the order the Python wrapper runs: credentials first, +/// so the limits see the same view the provider request is built from. +pub(crate) fn sdk_preflight(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> { + inherit_credentials(py, arguments, || { + Ok(PythonPreflight::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + PythonPreflight::CheckLimits.call(py, (arguments,))?; + Ok(()) +} struct CredentialEntry<'py>(Bound<'py, PyAny>); @@ -17,22 +59,6 @@ impl<'py> CredentialEntry<'py> { } } -pub fn prepare<'py>( - py: Python<'py>, - kwargs: &Bound<'py, PyDict>, - logger: &crate::PythonLogger, -) -> PyResult> { - let arguments = kwargs.copy()?; - arguments.set_item("litellm_logging_obj", logger.object(py))?; - inherit_credentials(py, &arguments, || { - Ok(Wrapper::CredentialList - .call(py, ())? - .cast_into::()?) - })?; - Wrapper::CheckLimits.call(py, (&arguments,))?; - Ok(arguments) -} - fn inherit_credentials<'py>( py: Python<'py>, arguments: &Bound<'py, PyDict>, @@ -54,7 +80,7 @@ fn inherit_credentials<'py>( .map(|credential| CredentialEntry(credential).name()) .collect::>>()?; let Some(index) = names.iter().position(|name| *name == requested) else { - Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?; + PythonPreflight::WarnUnknownCredential.call(py, (requested, names.len()))?; return Ok(()); }; let selected = CredentialEntry(credentials.get_item(index)?); @@ -71,7 +97,42 @@ fn inherit_credentials<'py>( #[cfg(test)] mod tests { + use std::collections::BTreeSet; + use std::sync::Mutex; + use super::*; + use strum::VariantArray; + + /// Tests share one interpreter, and the stub module below is global state, so the + /// tests that install it run one at a time. + static PREFLIGHT_MODULE: Mutex<()> = Mutex::new(()); + + /// A fresh stand-in for `litellm.rust_bridge.preflight` that records every call, then + /// `script` run against it with the module bound as `preflight`. + fn preflight_stubs<'py>(py: Python<'py>, script: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run( + c" +import sys +import types + +for name in ('litellm', 'litellm.rust_bridge'): + sys.modules.setdefault(name, types.ModuleType(name)) +preflight = types.ModuleType('litellm.rust_bridge.preflight') +preflight.warnings = [] +preflight.checked = [] +preflight.credential_list = lambda: [] +preflight.warn_unknown_credential = lambda name, loaded: preflight.warnings.append((name, loaded)) +preflight.check_limits = lambda kwargs: preflight.checked.append(kwargs) +sys.modules['litellm.rust_bridge.preflight'] = preflight +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals + } fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); @@ -312,4 +373,110 @@ arguments = {'litellm_credential_name': 'ocr-test'} } }); } + + #[test] + fn every_borrowed_function_is_in_the_python_contract() { + Python::initialize(); + Python::attach(|py| { + let contract = litellm_host_python::json_loads(py, PYTHON_CONTRACT.as_bytes()).unwrap(); + let declared: BTreeSet = contract + .bind(py) + .cast::() + .unwrap() + .keys() + .extract() + .map(|names: Vec| names.into_iter().collect()) + .unwrap(); + let called: BTreeSet = PythonPreflight::VARIANTS + .iter() + .map(|&function| <&str>::from(function).to_owned()) + .collect(); + assert_eq!( + called.len(), + PythonPreflight::VARIANTS.len(), + "a function is borrowed twice" + ); + assert_eq!(called, declared); + }); + } + + #[test] + fn an_unknown_name_is_reported_with_the_loaded_count_and_leaves_the_arguments_alone() { + let _guard = PREFLIGHT_MODULE + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = preflight_stubs( + py, + c" +class Credential: + credential_name = 'listed' + credential_values = {'api_key': 'listed-key'} +preflight.credential_list = lambda: [Credential(), Credential()] +arguments = {'litellm_credential_name': 'missing'} +", + ); + sdk_preflight(py, &argument_dict(&locals)).unwrap(); + py.run( + c" +assert arguments == {'litellm_credential_name': 'missing'}, arguments +assert preflight.warnings == [('missing', 2)], preflight.warnings +assert preflight.checked == [arguments] +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn limits_are_checked_on_the_arguments_after_credentials_are_inherited() { + let _guard = PREFLIGHT_MODULE + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = preflight_stubs( + py, + c" +class Credential: + credential_name = 'ocr-test' + credential_values = {'api_key': 'inherited'} +preflight.credential_list = lambda: [Credential()] +rejection = RuntimeError('Max retries per request hit!') +def check_limits(arguments): + preflight.checked.append(dict(arguments)) + raise rejection +preflight.check_limits = check_limits +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + let error = sdk_preflight(py, &argument_dict(&locals)).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("rejection").unwrap().unwrap()) + ); + py.run( + c" +assert preflight.checked == [{'litellm_credential_name': 'ocr-test', 'api_key': 'inherited'}], preflight.checked +assert arguments['api_key'] == 'inherited' +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + fn argument_dict<'py>(locals: &Bound<'py, PyDict>) -> Bound<'py, PyDict> { + locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap() + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index dae8623979a..52cebb7c903 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -33,6 +33,7 @@ fn run_messages( PublicCall::capture(&request, &args, &kwargs)?, crate::logger::LoggedMachine::new(messages_machine(secrets)), MessagesPythonHost::new(request.unbind()), + crate::preflight::sdk_preflight, asynchronous, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index a4f2bf851d7..e00c57fad64 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -70,6 +70,7 @@ fn run_ocr( PublicCall::capture(&request, &args, &kwargs)?, crate::logger::LoggedMachine::new(ocr_machine(client)), OcrPythonHost::new(request.unbind()), + crate::preflight::sdk_preflight, asynchronous, ) } diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index 6bbf2ffed6b..e39324d3348 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -22,7 +22,6 @@ from typing import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CredentialItem class MetadataUpdater(Protocol): @@ -72,21 +71,6 @@ def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> Call return call_setup -def check_limits(kwargs: Mapping[str, object]) -> None: - from litellm import ( - BudgetExceededError, - _current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor - max_budget, - num_retries_per_request, - ) - from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit - - if max_budget and _current_cost > max_budget: - raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget) - if max_retries_per_request_hit(kwargs, num_retries_per_request): - raise RuntimeError("Max retries per request hit!") - - def finalize( response: object, logger: Logging, @@ -299,22 +283,6 @@ def is_internal_call() -> bool: return internal.get() -def credential_list() -> list[CredentialItem]: - from litellm import credential_list as credentials - - return credentials - - -def warn_unknown_credential(name: str, loaded: int) -> None: - from litellm._logging import verbose_logger - - verbose_logger.warning( - "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", - name, - loaded, - ) - - def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]: from litellm import utils diff --git a/litellm/rust_bridge/preflight.py b/litellm/rust_bridge/preflight.py new file mode 100644 index 00000000000..e030382bfdc --- /dev/null +++ b/litellm/rust_bridge/preflight.py @@ -0,0 +1,45 @@ +"""The SDK request policy the native driver runs before a route's host projects. + +These are the `@client` prologue steps after `function_setup` and the deployment hook: +credential-name inheritance and the budget and retry-count limits. Rust owns the +inheritance itself; it borrows only the globals below. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from litellm.types.utils import CredentialItem + + +def credential_list() -> list[CredentialItem]: + from litellm import credential_list as credentials + + return credentials + + +def warn_unknown_credential(name: str, loaded: int) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + name, + loaded, + ) + + +def check_limits(kwargs: Mapping[str, object]) -> None: + from litellm import ( + BudgetExceededError, + _current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + max_budget, + num_retries_per_request, + ) + from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit + + if max_budget and _current_cost > max_budget: + raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget) + if max_retries_per_request_hit(kwargs, num_retries_per_request): + raise RuntimeError("Max retries per request hit!") diff --git a/tests/unit/rust_bridge/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/test_callbacks_legacy_python.py index 05f2d13a079..7365679a28c 100644 --- a/tests/unit/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/test_callbacks_legacy_python.py @@ -8,11 +8,10 @@ from typing import Final import pytest from pydantic import TypeAdapter -import litellm from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import callbacks_legacy_python as legacy -from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup +from litellm.rust_bridge.callbacks_legacy_python import failure_handler, setup _OCR_KWARGS: Final = MappingProxyType( { @@ -22,33 +21,6 @@ _OCR_KWARGS: Final = MappingProxyType( ) -@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -@pytest.mark.parametrize( - "cap, request_retry_count, refused", - [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], - ids=[ - "cap-above-four-reached", - "cap-above-four-not-reached", - "first-attempt-passes-cap-of-zero", - "cap-of-zero-refuses-first-retry", - ], -) -def test_check_limits_reads_request_retry_count( - monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool -) -> None: - monkeypatch.setattr(litellm, "num_retries_per_request", cap) - monkeypatch.setattr(litellm, "max_budget", None) - kwargs: Final = { - "model": "mistral/mistral-ocr-latest", - metadata_key: {"request_retry_count": request_retry_count}, - } - if refused: - with pytest.raises(RuntimeError, match="Max retries per request hit!"): - check_limits(kwargs) - else: - check_limits(kwargs) - - def _supplied_logger() -> Logging: return Logging( model="mistral/mistral-ocr-latest", diff --git a/tests/unit/rust_bridge/test_preflight.py b/tests/unit/rust_bridge/test_preflight.py new file mode 100644 index 00000000000..a8b1a40a00b --- /dev/null +++ b/tests/unit/rust_bridge/test_preflight.py @@ -0,0 +1,63 @@ +import inspect +from collections.abc import Callable +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.rust_bridge import preflight +from litellm.rust_bridge.preflight import check_limits + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "cap, request_retry_count, refused", + [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], + ids=[ + "cap-above-four-reached", + "cap-above-four-not-reached", + "first-attempt-passes-cap-of-zero", + "cap-of-zero-refuses-first-retry", + ], +) +def test_check_limits_reads_request_retry_count( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool +) -> None: + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + monkeypatch.setattr(litellm, "max_budget", None) + kwargs: Final = { + "model": "mistral/mistral-ocr-latest", + metadata_key: {"request_retry_count": request_retry_count}, + } + if refused: + with pytest.raises(RuntimeError, match="Max retries per request hit!"): + check_limits(kwargs) + else: + check_limits(kwargs) + + +def test_check_limits_refuses_a_call_over_the_budget(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "num_retries_per_request", None) + monkeypatch.setattr(litellm, "max_budget", 1.0) + monkeypatch.setattr(litellm, "_current_cost", 1.5) + with pytest.raises(litellm.BudgetExceededError): + check_limits({"model": "mistral/mistral-ocr-latest"}) + + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/preflight_contract.json" +_SHIMS: Final[MappingProxyType[str, Callable[..., object]]] = MappingProxyType( + { + "credential_list": preflight.credential_list, + "warn_unknown_credential": preflight.warn_unknown_credential, + "check_limits": preflight.check_limits, + } +) + + +def test_the_rust_contract_matches_the_shim_signatures() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {name: list(inspect.signature(_SHIMS[name]).parameters) for name in contract} From 4adbc13d7996fe93fa8df8b5de3aa1f82a29e1dd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 18:16:05 -0700 Subject: [PATCH 38/58] fix(router): hold Responses lifecycle events until output so a pre-output fallback announces one response (#43238) * fix(router): hold Responses lifecycle events until output so a pre-output fallback announces one response * fix(router): narrow the responses wrapper close guards to Exception and test the hold helpers directly * test(router): type the responses fallback test helpers * fix(router): replay the held lifecycle events when the fallback stream fails before its first event --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 4 +- litellm/router.py | 203 ++++++++++++++-------- tests/unit/test_router/test_router.py | 219 ++++++++++++++++++++++-- 3 files changed, 338 insertions(+), 88 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 1ef39775bd3..9f537d24eaa 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -265,7 +265,7 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 -_PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) +PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) class BaseResponsesAPIStreamingIterator: @@ -885,7 +885,7 @@ class BaseResponsesAPIStreamingIterator: def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True - if event.type not in _PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: + if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: self._output_started = True def _fallback_error(self, original: Exception) -> MidStreamFallbackError: diff --git a/litellm/router.py b/litellm/router.py index 023b99cd64e..6f416c416c0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -614,6 +614,17 @@ def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, bu return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS +MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 + + +def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: + from litellm.responses.streaming_iterator import PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + if held_event_count >= MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: + return False + return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + class FallbackAwareAnthropicMessagesStream: """ Bare async generators can't carry the `_hidden_params` attribute the @@ -3332,100 +3343,140 @@ class Router: await self._async_generator.aclose() async def stream_with_fallbacks(): - fallback_response = None + held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback try: async for item in source_iterator: + if _responses_stream_holds_event(item, len(held_lifecycle_events)): + held_lifecycle_events = (*held_lifecycle_events, item) + continue + for held_event in held_lifecycle_events: + yield held_event + held_lifecycle_events = () yield item + for held_event in held_lifecycle_events: + yield held_event except MidStreamFallbackError as e: - partial_usage: Final = Router._extract_partial_responses_usage(source_iterator) - try: - model_group: Final = cast(str, initial_kwargs.get("model")) - fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", self.fallbacks) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( - "context_window_fallbacks", self.context_window_fallbacks + async with contextlib.aclosing( + self._aresponses_fallback_attempt( + e, source_iterator, initial_kwargs, wrapper.adopt_fallback_headers, held_lifecycle_events ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( - "content_policy_fallbacks", self.content_policy_fallbacks - ) - initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_responses_attempt - if e.is_pre_first_chunk or not e.generated_content: - # No content generated before the error — retry with the - # original input. Adding a continuation prompt would - # waste tokens and confuse the model. - pass - else: - initial_kwargs["input"] = Router._build_responses_continuation_input( - initial_kwargs.get("input"), - e.generated_content, - ) - # The Responses-API path stores observability metadata - # under "litellm_metadata" (not the default "metadata") — - # see _ageneric_api_call_with_fallbacks. Mirroring that - # here ensures model_group, model_group_alias, and trace - # ids land in the same key litellm.aresponses reads from. - self._update_kwargs_before_fallbacks( - model=model_group, - kwargs=initial_kwargs, - metadata_variable_name="litellm_metadata", - ) - # The content-policy dispatch branch matches on the trigger's own type, so a refusal's - # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. - fallback_trigger: Final[Exception] = ( - e.original_exception - if isinstance(e.original_exception, litellm.ContentPolicyViolationError) - else e - ) - fallback_response = await self.async_function_with_fallbacks_common_utils( - e=fallback_trigger, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - - prepared_fallback_hidden_params = wrapper.adopt_fallback_headers(fallback_response) - if hasattr(fallback_response, "__aiter__"): - async for fallback_item in fallback_response: - Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) - if partial_usage is not None: - Router._combine_responses_fallback_usage(fallback_item, partial_usage) - yield fallback_item - else: - yield fallback_response - except Exception as fallback_error: - verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) - if ( - isinstance(fallback_error, MidStreamFallbackError) - and fallback_error.original_exception is not None - ): - raise fallback_error.original_exception from fallback_error - raise fallback_error + ) as fallback_stream: + async for fallback_item in fallback_stream: + yield fallback_item + except Exception: + for held_event in held_lifecycle_events: + yield held_event + raise finally: with anyio.CancelScope(shield=True): if hasattr(source_iterator, "aclose"): try: await source_iterator.aclose() - except BaseException as exc: + except Exception as exc: verbose_router_logger.debug( "stream_with_fallbacks(aresponses): error closing source: %s", exc, ) - if fallback_response is not None and hasattr(fallback_response, "aclose"): - try: - await fallback_response.aclose() - except BaseException as exc: - verbose_router_logger.debug( - "stream_with_fallbacks(aresponses): error closing fallback: %s", - exc, - ) wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) return wrapper + async def _aresponses_fallback_attempt( + self, + e: "MidStreamFallbackError", + source_iterator: "BaseResponsesAPIStreamingIterator", + initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + adopt_headers: Callable[[object], tuple[dict[str, object], dict[str, object]]], # mutable-ok: hidden params + held_lifecycle_events: tuple[object, ...], + ) -> AsyncGenerator[object, None]: + """ + Re-enters the Router's fallback chain for a mid-stream Responses API error and yields + whatever the fallback attempt produces. The lifecycle events the primary stream held + back reach the client only when no fallback lands, so the client sees exactly one + response announced, the one whose id completes. Split out of + _aresponses_streaming_iterator to keep each function's cyclomatic complexity within + the repo's C901 budget. + """ + from litellm.exceptions import MidStreamFallbackError + + partial_usage: Final = Router._extract_partial_responses_usage(source_iterator) + fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted + fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement + try: + model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group + fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param + "fallbacks", self.fallbacks + ) + context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "context_window_fallbacks", self.context_window_fallbacks + ) + content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "content_policy_fallbacks", self.content_policy_fallbacks + ) + initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs + self._ageneric_api_call_with_fallbacks_responses_attempt + ) + if e.generated_content and not e.is_pre_first_chunk: + initial_kwargs["input"] = Router._build_responses_continuation_input( # rebind-ok: fallback hop input + initial_kwargs.get("input"), + e.generated_content, + ) + # The Responses-API path stores observability metadata + # under "litellm_metadata" (not the default "metadata") — + # see _ageneric_api_call_with_fallbacks. Mirroring that + # here ensures model_group, model_group_alias, and trace + # ids land in the same key litellm.aresponses reads from. + self._update_kwargs_before_fallbacks( + model=model_group, + kwargs=initial_kwargs, + metadata_variable_name="litellm_metadata", + ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. + fallback_trigger: Final[Exception] = ( + e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + ) + fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success + e=fallback_trigger, + disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, + ) + prepared_fallback_hidden_params: Final = adopt_headers(fallback_response) + if hasattr(fallback_response, "__aiter__"): + async for fallback_item in fallback_response: + Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) + if partial_usage is not None: + Router._combine_responses_fallback_usage(fallback_item, partial_usage) + fallback_yielded = True + yield fallback_item + else: + fallback_yielded = True # rebind-ok: see the pre-init above + yield fallback_response + except Exception as fallback_error: + verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) + if not fallback_yielded: + for held_event in held_lifecycle_events: + yield held_event + if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: + raise fallback_error.original_exception from fallback_error + raise + finally: + if fallback_response is not None and hasattr(fallback_response, "aclose"): + with anyio.CancelScope(shield=True): + try: + await fallback_response.aclose() + except Exception as exc: + verbose_router_logger.debug( + "stream_with_fallbacks(aresponses): error closing fallback: %s", + exc, + ) + def _completion_streaming_iterator( self, model_response: CustomStreamWrapper, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3393c2f0d3c..e99cefb35dd 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -8,7 +8,7 @@ import os import sys import threading import warnings -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final, Literal @@ -37,6 +37,7 @@ from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth from litellm.router import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, + MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, _anthropic_stream_commits_now, _anthropic_stream_error_is_gateway_verdict, @@ -46,6 +47,7 @@ from litellm.router import ( _anthropic_stream_should_decline_fallback, _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, + _responses_stream_holds_event, ) from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit @@ -4213,6 +4215,7 @@ async def test_aresponses_streaming_iterator_fallback(): hidden_params={"model_id": "src-deployment-1"}, ) fallback_chunks = [ + MagicMock(type="response.created"), MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed"), ] @@ -4235,7 +4238,7 @@ async def test_aresponses_streaming_iterator_fallback(): assert wrapped._hidden_params.get("model_id") == "src-deployment-1" collected = [c async for c in wrapped] - assert len(collected) == 3 # 1 primary chunk + 2 fallback chunks + assert collected == fallback_chunks call_kwargs = mock_fallback_utils.call_args.kwargs fbk = call_kwargs["kwargs"] # Bound methods compare equal when they share the same instance + __func__. @@ -4522,20 +4525,35 @@ def _make_native_responses_iterator(*, sse_payloads: tuple[dict[str, str], ...], _RESPONSES_LIFECYCLE_PAYLOADS: Final = ({"type": "response.created"}, {"type": "response.in_progress"}) +async def _events_until_error(stream: AsyncIterable[object]) -> AsyncIterator[object]: + try: + async for chunk in stream: + yield chunk + except Exception as error: + yield error + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before_output(): """A connection lost after response.created but before any output item is re-routed to the - fallback with the original input, the same as a provider error event would be.""" + fallback with the original input, the same as a provider error event would be, and the client + sees one response lifecycle: the fallback's, whose id the completed event carries.""" router: Final = _make_router_with_fallback() src: Final = _make_native_responses_iterator( sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, trailing_error=httpx.ReadError("Response payload is not completed"), ) + fallback_chunks: Final = [ + MagicMock(type="response.created", response=MagicMock(id="resp_fallback")), + MagicMock(type="response.in_progress", response=MagicMock(id="resp_fallback")), + MagicMock(type="response.output_text.delta"), + MagicMock(type="response.completed", response=MagicMock(id="resp_fallback")), + ] with patch.object( router, "async_function_with_fallbacks_common_utils", - return_value=_AsyncList([MagicMock(type="response.completed")]), + return_value=_AsyncList(fallback_chunks), ) as mock_fallback_utils: wrapped: Final = await router._aresponses_streaming_iterator( response=src, @@ -4546,9 +4564,10 @@ async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before "original_generic_function": litellm.aresponses, }, ) - seen: Final = [chunk.type async for chunk in wrapped] + collected: Final = [chunk async for chunk in wrapped] - assert seen == ["response.created", "response.in_progress", "response.completed"] + assert collected == fallback_chunks + assert [chunk.response.id for chunk in collected if chunk.type == "response.created"] == ["resp_fallback"] assert isinstance(mock_fallback_utils.call_args.kwargs["e"], MidStreamFallbackError) assert mock_fallback_utils.call_args.kwargs["kwargs"]["input"] == "Hello" @@ -4576,17 +4595,197 @@ async def test_aresponses_streaming_iterator_surfaces_transport_drop_when_no_fal "original_generic_function": litellm.aresponses, }, ) - with pytest.raises(httpx.ReadError) as exc_info: - async for _ in wrapped: - pass + outcome: Final = [item async for item in _events_until_error(wrapped)] - assert exc_info.value is transport_error + assert [item.type for item in outcome[:-1]] == ["response.created", "response.in_progress"] + assert outcome[-1] is transport_error assert mock_fallback_utils.await_count == 1 trigger: Final = mock_fallback_utils.await_args.kwargs["e"] assert isinstance(trigger, MidStreamFallbackError) assert trigger.original_exception is transport_error +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_forwards_lifecycle_events_in_order_once_output_starts(): + router: Final = _make_router_with_fallback() + chunks: Final = [ + MagicMock(type="response.created"), + MagicMock(type="response.in_progress"), + MagicMock(type="response.output_text.delta"), + MagicMock(type="response.completed"), + ] + wrapped: Final = await router._aresponses_streaming_iterator( + response=_make_responses_iterator(chunks=chunks), + initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, + ) + + assert [chunk async for chunk in wrapped] == chunks + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_flushes_held_lifecycle_events_when_the_stream_ends_without_output(): + router: Final = _make_router_with_fallback() + chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] + wrapped: Final = await router._aresponses_streaming_iterator( + response=_make_responses_iterator(chunks=chunks), + initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, + ) + + assert [chunk async for chunk in wrapped] == chunks + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_forwards_held_lifecycle_events_before_a_non_fallback_error(): + router: Final = _make_router_with_fallback() + chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] + client_error: Final = litellm.BadRequestError(message="bad input", model="gpt-4", llm_provider="openai") + wrapped: Final = await router._aresponses_streaming_iterator( + response=_make_responses_iterator(chunks=chunks, error=client_error), + initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, + ) + outcome: Final = [item async for item in _events_until_error(wrapped)] + + assert outcome[:-1] == chunks + assert outcome[-1] is client_error + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_commits_held_lifecycle_events_at_the_hold_cap(): + router: Final = _make_router_with_fallback() + chunks: Final = [MagicMock(type="response.in_progress") for _ in range(MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS + 1)] + src: Final = _make_responses_iterator( + chunks=chunks, + error=MidStreamFallbackError( + message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True + ), + ) + fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] + + with patch.object(router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks)): + wrapped: Final = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, + ) + collected: Final = [chunk async for chunk in wrapped] + + assert collected == [*chunks, *fallback_chunks] + + +@pytest.mark.parametrize( + ("event_type", "held_event_count", "expected"), + [ + ("response.created", 0, True), + ("response.in_progress", 1, True), + ("response.queued", MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS - 1, True), + ("response.in_progress", MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, False), + ("response.output_item.added", 0, False), + ("response.output_text.delta", 0, False), + ("response.completed", 0, False), + ], +) +def test_responses_stream_holds_event_holds_only_pre_output_lifecycle_events_under_the_cap( + event_type: str, held_event_count: int, expected: bool +): + assert _responses_stream_holds_event(MagicMock(type=event_type), held_event_count) is expected + + +@pytest.mark.asyncio +async def test_aresponses_fallback_attempt_drops_held_lifecycle_events_when_a_fallback_lands(): + router: Final = _make_router_with_fallback() + trigger: Final = MidStreamFallbackError( + message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True + ) + held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) + fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] + adopt_headers: Final = MagicMock(return_value=({}, {})) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks) + ) as mock_fallback_utils: + collected: Final = [ + chunk + async for chunk in router._aresponses_fallback_attempt( + trigger, + _make_responses_iterator(), + {"model": "gpt-4", "stream": True, "input": "Hello"}, + adopt_headers, + held, + ) + ] + + assert collected == fallback_chunks + adopt_headers.assert_called_once() + assert mock_fallback_utils.await_args.kwargs["e"] is trigger + + +@pytest.mark.asyncio +async def test_aresponses_fallback_attempt_replays_held_lifecycle_events_when_the_fallback_dies_before_its_first_event(): + """A fallback stream that raises before yielding anything announced no response of its own, so the + primary's held created/in_progress pair is replayed ahead of the error and the client sees the + announcement the failure belongs to, the same as when no fallback was attempted at all.""" + router: Final = _make_router_with_fallback() + trigger: Final = MidStreamFallbackError( + message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True + ) + held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) + fallback_error: Final = RuntimeError("fallback closed before its first event") + adopt_headers: Final = MagicMock(return_value=({}, {})) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_make_responses_iterator(error=fallback_error), + ): + outcome: Final = [ + item + async for item in _events_until_error( + router._aresponses_fallback_attempt( + trigger, + _make_responses_iterator(), + {"model": "gpt-4", "stream": True, "input": "Hello"}, + adopt_headers, + held, + ) + ) + ] + + assert outcome == [*held, fallback_error] + + +@pytest.mark.asyncio +async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events_once_the_fallback_announced_itself(): + """Once the fallback has yielded its own created event, a later failure must not replay the + primary's held pair on top of it, or the client would again see two announced response ids.""" + router: Final = _make_router_with_fallback() + trigger: Final = MidStreamFallbackError( + message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True + ) + held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) + fallback_created: Final = MagicMock(type="response.created") + fallback_error: Final = RuntimeError("fallback dropped after announcing itself") + adopt_headers: Final = MagicMock(return_value=({}, {})) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_make_responses_iterator(chunks=(fallback_created,), error=fallback_error), + ): + outcome: Final = [ + item + async for item in _events_until_error( + router._aresponses_fallback_attempt( + trigger, + _make_responses_iterator(), + {"model": "gpt-4", "stream": True, "input": "Hello"}, + adopt_headers, + held, + ) + ) + ] + + assert outcome == [fallback_created, fallback_error] + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: input is rewritten to include user prompt + From 1a581626301c8a09a2cd29578c5175802b1ebb5a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 18:31:15 -0700 Subject: [PATCH 39/58] refactor(http): hand out an owned Client and route all providers through the pool (#43245) * refactor(messages): take the provider client from the injected HTTP pool The messages route kept its own process-wide reqwest client, so it ignored ssl_verify, CA bundles, client certs, proxies and every other setting that litellm-http resolves. The machine now takes the HttpClientPool and the call's HttpClientConfig, as OCR does, and the bridge passes its shared pool. Co-Authored-By: Claude Opus 5.5 * refactor(http): hand out an owned Client and move chat, audio and OIDC onto the pool HttpClientPool now returns litellm_http::Client, a newtype only crates/http can build, so every provider client carries the resolved TLS, proxy and timeout settings. Chat completions and audio transcription drop their process-wide reqwest clients and take the pool and call config like messages; their 600s ceiling moves to the request. OidcResolver takes its client instead of building one, and the bridge hands it the pooled one. Co-Authored-By: Claude Opus 5.5 * refactor(secrets): build Google, Azure and CyberArk manager clients from the pool The native secret managers built bare reqwest clients, so they ignored the host's TLS and proxy settings. load_native_manager now takes the pool and the host config and hands each manager a pooled client. CyberArk's CYBERARK_SSL_VERIFY and CYBERARK_CLIENT_CERT/KEY become an override on the host config instead of a hand-built client. To express a certificate and key in separate files, HttpClientConfig::client_certificate is now a ClientIdentity that is either one PEM or a split pair. Co-Authored-By: Claude Opus 5.5 * chore(clippy): only crates/http may build a reqwest client Fence reqwest::Client, ClientBuilder and the TLS builder methods with disallowed-types and disallowed-methods so new code takes a litellm_http::Client from the pool. crates/http is exempt as the one place clients are built, and testkit as a dev-only installer. Tests move to litellm_http::Client::plain_for_test or a pooled client. Co-Authored-By: Claude Opus 5.5 * fix(secrets-cyberark): keep verifying certificates when the host disables it Python hands CyberArk its own ssl_verify, which wins over the global setting, so CYBERARK_SSL_VERIFY unset or true still verifies even when the host sets ssl_verify false. The pooled client copied the host's Disabled and would send the API key unverified; fall back to the built-in roots instead. Co-Authored-By: Claude Opus 5.5 * fix(python-bridge): treat a missing litellm package as no host HTTP settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 10 +++ litellm-rust/clippy.toml | 12 ++++ litellm-rust/crates/auth-aws/Cargo.toml | 1 + litellm-rust/crates/auth-aws/src/aws.rs | 2 +- .../crates/cache-azure-blob/Cargo.toml | 2 + .../crates/cache-azure-blob/src/cache.rs | 2 +- .../crates/cache-azure-blob/src/transport.rs | 2 +- .../cache-azure-blob/tests/transport.rs | 2 +- litellm-rust/crates/cache-gcs/Cargo.toml | 2 + litellm-rust/crates/cache-gcs/src/cache.rs | 2 +- litellm-rust/crates/cache-gcs/tests/cache.rs | 2 +- .../crates/cache-gcs/tests/support/mod.rs | 2 +- .../crates/cache-qdrant-semantic/Cargo.toml | 2 + .../cache-qdrant-semantic/src/embedder.rs | 2 +- .../cache-qdrant-semantic/tests/embedder.rs | 22 +++++-- litellm-rust/crates/cache-s3/Cargo.toml | 2 + litellm-rust/crates/cache-s3/src/cache.rs | 7 ++- litellm-rust/crates/cache-s3/src/transport.rs | 2 +- .../crates/cache-s3/tests/support/mod.rs | 7 ++- litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/audio_transcription/client.rs | 13 ---- .../core/src/audio_transcription/handler.rs | 20 ++++-- .../core/src/audio_transcription/mod.rs | 13 ++-- .../core/src/chat_completions/client.rs | 14 ----- .../core/src/chat_completions/handler.rs | 20 ++++-- .../crates/core/src/chat_completions/mod.rs | 8 ++- litellm-rust/crates/core/src/constants.rs | 6 -- .../crates/core/src/messages/client.rs | 14 ----- .../crates/core/src/messages/error.rs | 2 + .../crates/core/src/messages/handler.rs | 12 ++-- litellm-rust/crates/core/src/messages/mod.rs | 19 ++++-- .../crates/core/src/messages/route.rs | 15 ++++- litellm-rust/crates/core/src/ocr/prepare.rs | 5 +- .../crates/core/tests/audio_transcription.rs | 20 +++--- .../crates/core/tests/chat_completions.rs | 21 ++++--- .../crates/core/tests/messages/host.rs | 2 +- .../crates/core/tests/messages/main.rs | 11 +++- .../crates/core/tests/messages/response.rs | 35 +++++++---- .../crates/core/tests/messages/stream.rs | 2 +- litellm-rust/crates/core/tests/ocr/main.rs | 7 +-- litellm-rust/crates/core/tests/ocr/mistral.rs | 13 ++-- litellm-rust/crates/core/tests/support/mod.rs | 13 +++- litellm-rust/crates/http/Cargo.toml | 2 + litellm-rust/crates/http/src/client.rs | 38 +++++++++++ litellm-rust/crates/http/src/config.rs | 12 +++- litellm-rust/crates/http/src/lib.rs | 10 ++- litellm-rust/crates/http/src/media.rs | 25 +++----- litellm-rust/crates/http/src/outbound.rs | 2 +- litellm-rust/crates/http/src/pool.rs | 14 +++-- litellm-rust/crates/http/src/tls.rs | 58 ++++++++++++----- litellm-rust/crates/llms/Cargo.toml | 1 + .../document_intelligence/transformation.rs | 4 +- .../crates/llms/src/base_llm/ocr/document.rs | 25 +++++--- .../crates/llms/src/base_llm/ocr/handler.rs | 27 ++++---- .../llms/src/reducto/ocr/transformation.rs | 5 +- litellm-rust/crates/llms/tests/ocr_handler.rs | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../python-bridge/src/cache/activation.rs | 2 +- .../crates/python-bridge/src/cache/config.rs | 2 +- .../crates/python-bridge/src/cache/handle.rs | 2 +- .../crates/python-bridge/src/cache/mod.rs | 10 --- .../crates/python-bridge/src/cache/native.rs | 8 +-- litellm-rust/crates/python-bridge/src/http.rs | 19 ++++-- .../python-bridge/src/python_settings.rs | 12 +++- .../src/routes/audio_transcription.rs | 34 ++++++---- .../src/routes/chat_completions.rs | 42 +++++++++---- .../python-bridge/src/routes/messages/mod.rs | 5 +- .../python-bridge/src/secrets/callback.rs | 2 +- .../crates/python-bridge/src/secrets/mod.rs | 7 ++- .../python-bridge/src/secrets/resolved.rs | 29 ++++++--- .../python-bridge/src/secrets/runtime.rs | 15 ++++- litellm-rust/crates/secrets-azure/Cargo.toml | 2 + .../crates/secrets-azure/src/key_vault.rs | 11 ++-- .../crates/secrets-azure/tests/key_vault.rs | 19 +++--- .../crates/secrets-azure/tests/live.rs | 2 +- .../crates/secrets-cyberark/Cargo.toml | 2 + .../crates/secrets-cyberark/src/error.rs | 2 + .../secrets-cyberark/src/secret_manager.rs | 7 ++- .../src/secret_manager/client.rs | 63 +++++++++++++++---- .../secrets-cyberark/tests/secret_manager.rs | 2 + .../tests/secret_manager/configuration.rs | 26 ++++---- .../tests/secret_manager/support.rs | 14 ++++- .../tests/secret_manager/writes.rs | 6 +- litellm-rust/crates/secrets-google/Cargo.toml | 2 + .../secrets-google/src/secret_manager.rs | 7 ++- .../secrets-google/tests/secret_manager.rs | 16 +++-- litellm-rust/crates/secrets/Cargo.toml | 2 + litellm-rust/crates/secrets/src/error.rs | 2 + litellm-rust/crates/secrets/src/native.rs | 33 ++++++---- litellm-rust/crates/secrets/src/oidc.rs | 34 +++++----- litellm-rust/crates/secrets/src/resolver.rs | 15 +---- litellm-rust/crates/secrets/src/source.rs | 5 +- litellm-rust/crates/secrets/tests/aws.rs | 6 +- litellm-rust/crates/secrets/tests/azure.rs | 6 +- .../secrets/tests/common_read_contract.rs | 8 +-- litellm-rust/crates/secrets/tests/cyberark.rs | 2 +- litellm-rust/crates/secrets/tests/google.rs | 4 +- .../crates/secrets/tests/hashicorp.rs | 4 +- litellm-rust/crates/secrets/tests/oidc.rs | 36 ++++++----- .../crates/secrets/tests/resolution.rs | 14 ++--- litellm-rust/crates/secrets/tests/source.rs | 4 +- litellm-rust/crates/testkit/src/lib.rs | 6 ++ 102 files changed, 745 insertions(+), 403 deletions(-) delete mode 100644 litellm-rust/crates/core/src/audio_transcription/client.rs delete mode 100644 litellm-rust/crates/core/src/chat_completions/client.rs delete mode 100644 litellm-rust/crates/core/src/messages/client.rs create mode 100644 litellm-rust/crates/http/src/client.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ffcc6a5496b..5bb65b6b05f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2911,6 +2911,7 @@ dependencies = [ "litellm-cache", "litellm-cache-response", "litellm-cache-testing", + "litellm-http", "reqwest 0.12.28", "rstest", "serde_json", @@ -2944,6 +2945,7 @@ dependencies = [ "litellm-auth-types", "litellm-cache", "litellm-cache-testing", + "litellm-http", "percent-encoding", "reqwest 0.12.28", "rstest", @@ -2970,6 +2972,7 @@ dependencies = [ "futures-util", "litellm-cache", "litellm-cache-testing", + "litellm-http", "qdrant-client", "reqwest 0.12.28", "rstest", @@ -3043,6 +3046,7 @@ dependencies = [ "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", + "litellm-http", "reqwest 0.12.28", "rstest", "serde_json", @@ -3207,11 +3211,13 @@ dependencies = [ "http 1.4.2", "hyper-util", "litellm-core-utils", + "rcgen", "reqwest 0.12.28", "rstest", "rustls 0.23.42", "serde", "serde_json", + "tempfile", "thiserror 2.0.19", "tokio", "veil", @@ -3345,6 +3351,7 @@ dependencies = [ "google-cloud-auth", "google-cloud-kms-v1", "litellm-core-utils", + "litellm-http", "litellm-python-compat", "litellm-secrets-aws", "litellm-secrets-azure", @@ -3392,6 +3399,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-types", "litellm-core-utils", + "litellm-http", "litellm-secrets-types", "percent-encoding", "reqwest 0.12.28", @@ -3411,6 +3419,7 @@ version = "0.1.0" dependencies = [ "base64 0.22.1", "litellm-core-utils", + "litellm-http", "litellm-secrets-types", "litellm-tracing", "moka", @@ -3439,6 +3448,7 @@ dependencies = [ "litellm-auth-gcp", "litellm-auth-types", "litellm-core-utils", + "litellm-http", "litellm-secrets-types", "moka", "percent-encoding", diff --git a/litellm-rust/clippy.toml b/litellm-rust/clippy.toml index f7e3293069b..0e2ff770d27 100644 --- a/litellm-rust/clippy.toml +++ b/litellm-rust/clippy.toml @@ -7,4 +7,16 @@ disallowed-methods = [ { path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" }, { path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" }, { path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" }, + { path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" }, + { path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" }, + { path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" }, + { path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" }, + { path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" }, +] + +# Every outbound client comes from litellm_http::HttpClientPool so it honors the host's TLS, +# proxy and timeout settings. Only crates/http builds one. +disallowed-types = [ + { path = "reqwest::Client", reason = "take litellm_http::Client from HttpClientPool; only crates/http builds one" }, + { path = "reqwest::ClientBuilder", reason = "HttpClientConfig owns client construction" }, ] diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 1a35af48574..9592f278d94 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -22,5 +22,6 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } reqwest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index cb9195ffeb6..409ff78867f 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -962,7 +962,7 @@ mod tests { &no_env, ) .await?; - let client = reqwest::Client::new(); + let client = litellm_http::Client::plain_for_test(); let mut failures = Vec::new(); for region in ["us-west-2", "us-east-1"] { diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml index baa1b0f5482..5bdfa16ef53 100644 --- a/litellm-rust/crates/cache-azure-blob/Cargo.toml +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true litellm-auth-azure.workspace = true litellm-auth-types.workspace = true litellm-cache.workspace = true @@ -19,6 +20,7 @@ tokio.workspace = true url.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-response.workspace = true litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs index 489b08d485e..c5c1fdd8ab9 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -31,7 +31,7 @@ impl AzureBlobCache { pub async fn connect( account_url: &str, container: &str, - http: reqwest::Client, + http: litellm_http::Client, codec: C, runtime: Handle, ) -> Result { diff --git a/litellm-rust/crates/cache-azure-blob/src/transport.rs b/litellm-rust/crates/cache-azure-blob/src/transport.rs index ed038b8d69d..3914b92365c 100644 --- a/litellm-rust/crates/cache-azure-blob/src/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/src/transport.rs @@ -8,7 +8,7 @@ use azure_core::{ use futures_util::TryStreamExt; #[derive(Debug)] -pub struct ReqwestTransport(pub reqwest::Client); +pub struct ReqwestTransport(pub litellm_http::Client); #[async_trait::async_trait] impl HttpClient for ReqwestTransport { diff --git a/litellm-rust/crates/cache-azure-blob/tests/transport.rs b/litellm-rust/crates/cache-azure-blob/tests/transport.rs index cd1e10aa3d8..c8b14cd8543 100644 --- a/litellm-rust/crates/cache-azure-blob/tests/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/tests/transport.rs @@ -31,7 +31,7 @@ async fn connect(server: &MockServer) -> AzureBlobCache> { None, ClientOptions { transport: Some(Transport::new(Arc::new(ReqwestTransport( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), )))), ..ClientOptions::default() }, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index da0acf554f9..1a06683e615 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true futures-util.workspace = true litellm-auth-gcp.workspace = true litellm-auth-types.workspace = true @@ -15,6 +16,7 @@ reqwest.workspace = true tokio.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index a8a7fbc9a7b..bad81573cb4 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -5,8 +5,8 @@ use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache, }; +use litellm_http::Client; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; -use reqwest::Client; use crate::{GcpTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index cdce6a00bdd..12bb5344570 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -96,7 +96,7 @@ async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) { path_service_account: Some("/secrets/sa.json".into()), ..support::config(&server, Some("folder")) }, - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), litellm_cache::JsonCodec::::new(), ); assert_eq!(cache.bucket_name(), "bucket"); diff --git a/litellm-rust/crates/cache-gcs/tests/support/mod.rs b/litellm-rust/crates/cache-gcs/tests/support/mod.rs index 6097f0ee1bd..beb9aa39d9c 100644 --- a/litellm-rust/crates/cache-gcs/tests/support/mod.rs +++ b/litellm-rust/crates/cache-gcs/tests/support/mod.rs @@ -29,7 +29,7 @@ pub fn cache_with_token( ) -> JsonGcsCache { GcsCache::with_token_source( config(server, gcs_path), - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), JsonCodec::new(), token, ) diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index 950c2db7491..a44bef0a731 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true futures-util.workspace = true litellm-cache.workspace = true qdrant-client = { workspace = true, features = ["serde"] } @@ -17,6 +18,7 @@ tokio.workspace = true uuid.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } futures-executor = "0.3" litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index 340393600f2..b7fbcd9b02d 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_cache::{Error, semantic::Embedder}; -use reqwest::Client; +use litellm_http::Client; use serde_json::Value; pub struct OpenAiEmbedder { diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index de0fab0a66f..24e6e5eba3e 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -5,6 +5,10 @@ use std::{ use litellm_cache::{Error, semantic::Embedder}; use litellm_cache_qdrant_semantic::{OpenAiEmbedder, OpenAiEmbedderConfig}; +use litellm_http::{ + ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, Resolution, + media::PublicDnsResolver, +}; use rstest::rstest; use serde_json::{Value, json}; use tokio::{ @@ -104,7 +108,7 @@ fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { async fn posts_embeddings_request_and_parses_vector() { let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; let embedder = OpenAiEmbedder::new( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), config( format!("{}/", server.base_url()), Some(Duration::from_secs(1)), @@ -156,14 +160,17 @@ async fn status_timeout_and_body_errors_are_unavailable( ) { let server = TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await; - let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), timeout)); + let embedder = OpenAiEmbedder::new( + litellm_http::Client::plain_for_test(), + config(server.base_url(), timeout), + ); assert_eq!(embedder.async_embed("hello", None).await, expected); } #[rstest] fn sync_embedding_is_unsupported() { let embedder = OpenAiEmbedder::new( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), config("http://127.0.0.1:9".to_owned(), None), ); assert_eq!( @@ -176,9 +183,12 @@ fn sync_embedding_is_unsupported() { #[tokio::test] async fn uses_the_injected_client() { let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; - let client = reqwest::Client::builder() - .user_agent("litellm-embedder-test") - .build() + let config_with_agent = HttpClientConfig { + user_agent: Some("litellm-embedder-test".into()), + ..Resolution::from(&HttpSettings::default()).config + }; + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(&config_with_agent, ClientVariant::Provider) .unwrap(); let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None)); assert_eq!( diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index c8150180e7c..680f2da8215 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true litellm-cache.workspace = true litellm-auth-aws.workspace = true aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } @@ -19,6 +20,7 @@ reqwest.workspace = true tokio.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true wiremock = "0.6.5" diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 91cd5e8ef54..ced948e80c6 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -42,7 +42,12 @@ pub struct S3Cache { } impl S3Cache { - pub fn new(config: S3CacheConfig, http: reqwest::Client, codec: C, runtime: Handle) -> Self { + pub fn new( + config: S3CacheConfig, + http: litellm_http::Client, + codec: C, + runtime: Handle, + ) -> Self { let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); let base = aws_sdk_s3::Config::builder() .behavior_version(BehaviorVersion::latest()) diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/cache-s3/src/transport.rs index 3e5ce578c31..eabc54ac9e2 100644 --- a/litellm-rust/crates/cache-s3/src/transport.rs +++ b/litellm-rust/crates/cache-s3/src/transport.rs @@ -9,7 +9,7 @@ use aws_smithy_runtime_api::client::{ use aws_smithy_types::body::SdkBody; #[derive(Clone, Debug)] -pub(crate) struct ReqwestHttpClient(pub(crate) reqwest::Client); +pub(crate) struct ReqwestHttpClient(pub(crate) litellm_http::Client); impl HttpClient for ReqwestHttpClient { fn http_connector( diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 046b042c66a..b628c1df431 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -32,7 +32,12 @@ pub fn config(endpoint: &str) -> S3CacheConfig { } pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache { - S3Cache::new(config, reqwest::Client::new(), JsonCodec::new(), runtime) + S3Cache::new( + config, + litellm_http::Client::plain_for_test(), + JsonCodec::new(), + runtime, + ) } pub fn cache(endpoint: &str) -> JsonS3Cache { diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 12410c187e2..8700c8df308 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -36,6 +36,7 @@ url.workspace = true veil.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-auth-gcp.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/client.rs b/litellm-rust/crates/core/src/audio_transcription/client.rs deleted file mode 100644 index 3cf131839b8..00000000000 --- a/litellm-rust/crates/core/src/audio_transcription/client.rs +++ /dev/null @@ -1,13 +0,0 @@ -use std::{sync::OnceLock, time::Duration}; - -use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS; - -pub(super) fn http_client() -> &'static reqwest::Client { - static CLIENT: OnceLock = OnceLock::new(); - CLIENT.get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)) - .build() - .unwrap_or_else(|_| reqwest::Client::new()) - }) -} diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index a1862f341a5..30900bc14c6 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,10 +1,16 @@ -use litellm_http::request::truncate_error_body; +use std::time::Duration; + +use litellm_http::{Client, request::truncate_error_body}; use serde_json::Value; -use super::{Error, client::http_client}; -use crate::audio_transcription::types::ProviderAudioTranscriptionRequest; +use super::Error; +use crate::{ + audio_transcription::types::ProviderAudioTranscriptionRequest, + constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, +}; pub async fn execute_audio_transcription_provider_call( + http: &Client, request: ProviderAudioTranscriptionRequest, ) -> Result { let response = crate::outbound::outbound_request::( @@ -12,11 +18,15 @@ pub async fn execute_audio_transcription_provider_call( request.url.clone(), request.upstream_headers.clone(), &request.body, - request.timeout, + Some( + request + .timeout + .unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)), + ), &request.optional_params, ) .await? - .send(http_client()) + .send(http) .await .map_err(|error| { Error::Transport(litellm_http::transport::Error::Network(error.to_string())) diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 801fd5e9673..dc75326d5c3 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -1,16 +1,21 @@ mod error; pub mod types; pub use error::Error; -mod client; mod handler; mod prepare; pub use handler::execute_audio_transcription_provider_call; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; pub use prepare::prepare_audio_transcription_provider_call; use serde_json::Value; use crate::audio_transcription::types::AudioTranscriptionRequest; -pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { - execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) - .await +pub async fn audio_transcription( + pool: &HttpClientPool, + config: &HttpClientConfig, + request: AudioTranscriptionRequest<'_>, +) -> Result { + let request = prepare_audio_transcription_provider_call(request)?; + let http = pool.client(config, ClientVariant::Provider)?; + execute_audio_transcription_provider_call(&http, request).await } diff --git a/litellm-rust/crates/core/src/chat_completions/client.rs b/litellm-rust/crates/core/src/chat_completions/client.rs deleted file mode 100644 index d8ad6c49b7b..00000000000 --- a/litellm-rust/crates/core/src/chat_completions/client.rs +++ /dev/null @@ -1,14 +0,0 @@ -use std::{sync::OnceLock, time::Duration}; - -use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS}; - -pub(super) fn http_client() -> &'static reqwest::Client { - static CLIENT: OnceLock = OnceLock::new(); - CLIENT.get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)) - .connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS)) - .build() - .unwrap_or_else(|_| reqwest::Client::new()) - }) -} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 2391ab83a60..f3404fcaa8a 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,20 +1,24 @@ -use litellm_http::{outbound::OutboundRequest, request::truncate_error_body}; +use std::time::Duration; + +use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body}; use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; -use super::{Error, client::http_client, prepare::prepare_provider_request}; -use crate::chat_completions::types::{ - ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest, +use super::{Error, prepare::prepare_provider_request}; +use crate::{ + chat_completions::types::{ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest}, + constants::CHAT_COMPLETIONS_TIMEOUT_SECS, }; pub(super) async fn execute_chat_completions_provider_call( + http: &Client, request: ResolvedChatCompletionsRequest<'_>, ) -> Result { let request = prepare_provider_request(request)?; let outbound = outbound_request(&request).await?; - let response = outbound.send(http_client()).await.map_err(|err| { + let response = outbound.send(http).await.map_err(|err| { // Failing to establish the connection means the request never went out, // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. @@ -72,7 +76,11 @@ pub(super) async fn outbound_request( request.url.clone(), request.upstream_headers.clone(), &request.body, - request.timeout, + Some( + request + .timeout + .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), + ), &request.optional_params, ) .await diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 224c9d8cfed..be22aea5669 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -9,11 +9,11 @@ mod error; pub mod types; pub use error::Error; -mod client; mod common_utils; pub(crate) mod handler; mod prepare; use handler::execute_chat_completions_provider_call; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use litellm_types::utils::ChatCompletionsResponse; use prepare::{parse_messages, resolve_provider_config, resolve_request}; use serde_json::{Map, Value}; @@ -21,9 +21,13 @@ use serde_json::{Map, Value}; use crate::chat_completions::types::ChatCompletionsRequest; pub async fn chat_completions( + pool: &HttpClientPool, + config: &HttpClientConfig, request: ChatCompletionsRequest<'_>, ) -> Result { - execute_chat_completions_provider_call(resolve_request(request)?).await + let request = resolve_request(request)?; + let http = pool.client(config, ClientVariant::Provider)?; + execute_chat_completions_provider_call(&http, request).await } /// Whether the core would accept this request, without resolving credentials or diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 3d740e39677..455c3258799 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -5,9 +5,6 @@ pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com"; /// timeout from the caller still overrides this on the request builder. pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; -/// Connect timeout for Anthropic Messages provider calls, in seconds. -pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10; - /// Provider name used for Anthropic Messages when a deployment's provider model /// does not carry an explicit provider prefix. pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; @@ -16,9 +13,6 @@ pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; /// seconds. Mirrors the Python chat completions default. pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600; -/// Connect timeout for chat completions provider calls, in seconds. -pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; - pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; /// `object` field every non-streaming chat completion response carries. diff --git a/litellm-rust/crates/core/src/messages/client.rs b/litellm-rust/crates/core/src/messages/client.rs deleted file mode 100644 index ca70b1b03eb..00000000000 --- a/litellm-rust/crates/core/src/messages/client.rs +++ /dev/null @@ -1,14 +0,0 @@ -use std::{sync::OnceLock, time::Duration}; - -use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS}; - -pub(super) fn http_client() -> &'static reqwest::Client { - static CLIENT: OnceLock = OnceLock::new(); - CLIENT.get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS)) - .connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS)) - .build() - .unwrap_or_else(|_| reqwest::Client::new()) - }) -} diff --git a/litellm-rust/crates/core/src/messages/error.rs b/litellm-rust/crates/core/src/messages/error.rs index 2a9723beb38..76f8813e330 100644 --- a/litellm-rust/crates/core/src/messages/error.rs +++ b/litellm-rust/crates/core/src/messages/error.rs @@ -17,6 +17,8 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] + Client(#[from] litellm_http::Error), + #[error(transparent)] Transport(#[from] litellm_http::transport::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index de1a5f476ed..f90cb8cb454 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -5,13 +5,15 @@ use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMes use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::Value; -use super::{Error, client::http_client, common_utils::truncate_error_body}; +use super::{Error, common_utils::truncate_error_body}; +use crate::constants::MESSAGES_TIMEOUT_SECS; pub(super) fn network(error: reqwest::Error) -> Error { Error::Transport(TransportError::Network(error.to_string())) } pub(super) async fn send( + http: &litellm_http::Client, url: &str, headers: &[(String, String)], body: &Value, @@ -20,13 +22,11 @@ pub(super) async fn send( let encoded = serde_json::to_vec(body) .map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?; let builder = headers.iter().fold( - http_client().post(url).body(encoded), + http.post(url) + .body(encoded) + .timeout(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), |builder, (key, value)| builder.header(key, value), ); - let builder = match timeout { - Some(duration) => builder.timeout(duration), - None => builder, - }; http_request(builder).await.map_err(network) } diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 180eb08810e..5cb83b4e34d 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -7,13 +7,13 @@ mod error; pub mod types; pub use error::Error; -mod client; mod common_utils; mod handler; mod prepare; pub mod route; use std::sync::Arc; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use litellm_secrets::source::EnvironmentSecrets; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; @@ -21,7 +21,11 @@ use serde_json::Value; use crate::messages::types::MessagesRequest; -pub async fn messages(request: MessagesRequest<'_>) -> Result { +pub async fn messages( + pool: &HttpClientPool, + config: &HttpClientConfig, + request: MessagesRequest<'_>, +) -> Result { let Value::Object(body) = request.body else { return Err(Error::InvalidRequest( "messages body must be an object".into(), @@ -38,8 +42,15 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(*message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 40aff185e81..7f6589cdf3e 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -12,6 +12,7 @@ use litellm_host::{ machine::{CallMachine, HostChannel, MachineFault}, protocol::Protocol, }; +use litellm_http::{Client, ClientVariant, HttpClientConfig, HttpClientPool}; use litellm_secrets::source::SecretSource; use litellm_types::{ llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse, @@ -108,12 +109,20 @@ impl Host for LocalMessagesHost { } } -pub fn messages_machine(secrets: Arc) -> MessagesMachine { - CallMachine::new(move |host| Box::pin(execute(host, secrets.clone()))) +pub fn messages_machine( + pool: &HttpClientPool, + config: &HttpClientConfig, + secrets: Arc, +) -> Result { + let http = pool.client(config, ClientVariant::Provider)?; + Ok(CallMachine::new(move |host| { + Box::pin(execute(host, http.clone(), secrets.clone())) + })) } async fn execute( host: MessagesHost, + http: Client, secrets: Arc, ) -> Result { let call = host.project().await?; @@ -164,7 +173,7 @@ async fn execute( context, ) .await?; - let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?; + let response = send(&http, &wire.url, &wire.headers, &wire.body, request.timeout).await?; if !response.status().is_success() { return Err(provider_error(response).await); } diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 18961ec96fa..f13d6984763 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -110,7 +110,10 @@ mod tests { } fn client() -> OcrClient { - OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()) + OcrClient::for_test( + litellm_http::Client::plain_for_test(), + litellm_http::Client::no_redirect_for_test(), + ) } fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest { diff --git a/litellm-rust/crates/core/tests/audio_transcription.rs b/litellm-rust/crates/core/tests/audio_transcription.rs index 196f085a6c3..612395fe63a 100644 --- a/litellm-rust/crates/core/tests/audio_transcription.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -10,6 +10,10 @@ use support::*; const MODEL: &str = "mistral.voxtral-mini-3b-2507"; +async fn transcribe(request: AudioTranscriptionRequest<'_>) -> Result { + audio_transcription(&http_pool(), &http_config(), request).await +} + fn transcript_response(text: &str) -> ResponseTemplate { json_response(json!({"output": {"message": {"content": [{"text": text}]}}})) } @@ -47,7 +51,7 @@ async fn bedrock_converse_request_is_signed_for_the_requested_region( let upstream = upstream([transcript_response("hello")]).await; let base = upstream.uri(); - let response = audio_transcription(AudioTranscriptionRequest { + let response = transcribe(AudioTranscriptionRequest { api_base: Some(&base), optional_params: aws_params(region), ..request @@ -79,7 +83,7 @@ async fn the_provider_can_come_from_the_model_prefix(request: AudioTranscription let base = upstream.uri(); let model = format!("bedrock/{MODEL}"); - audio_transcription(AudioTranscriptionRequest { + transcribe(AudioTranscriptionRequest { model: &model, custom_llm_provider: None, api_base: Some(&base), @@ -110,7 +114,7 @@ async fn audio_and_transcription_params_reach_the_converse_body( ]) .collect(); - audio_transcription(AudioTranscriptionRequest { + transcribe(AudioTranscriptionRequest { audio: json!({"data": "AQI=", "format": format}), api_base: Some(&base), optional_params, @@ -142,7 +146,7 @@ async fn invalid_audio_is_rejected_before_sending( let upstream = upstream([transcript_response("hello")]).await; let base = upstream.uri(); - let error = audio_transcription(AudioTranscriptionRequest { + let error = transcribe(AudioTranscriptionRequest { audio, api_base: Some(&base), ..request @@ -174,7 +178,7 @@ async fn unsupported_providers_are_rejected_before_sending( #[case] provider: Option<&'static str>, #[case] reported: &str, ) { - let error = audio_transcription(AudioTranscriptionRequest { + let error = transcribe(AudioTranscriptionRequest { model, custom_llm_provider: provider, api_base: Some(UNREACHABLE_BASE), @@ -189,7 +193,7 @@ async fn unsupported_providers_are_rejected_before_sending( #[rstest] #[tokio::test] async fn a_non_string_extra_header_is_rejected(request: AudioTranscriptionRequest<'static>) { - let error = audio_transcription(AudioTranscriptionRequest { + let error = transcribe(AudioTranscriptionRequest { extra_headers: Some(Map::from_iter([("x-count".to_string(), json!(3))])), api_base: Some(UNREACHABLE_BASE), ..request @@ -212,7 +216,7 @@ async fn an_upstream_error_keeps_its_status_and_body( upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await; let base = upstream.uri(); - let error = audio_transcription(AudioTranscriptionRequest { + let error = transcribe(AudioTranscriptionRequest { api_base: Some(&base), ..request }) @@ -239,7 +243,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( let upstream = upstream([response]).await; let base = upstream.uri(); - let error = audio_transcription(AudioTranscriptionRequest { + let error = transcribe(AudioTranscriptionRequest { api_base: Some(&base), ..request }) diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs index ae96509fe2e..d1f6cde19e8 100644 --- a/litellm-rust/crates/core/tests/chat_completions.rs +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -4,6 +4,7 @@ use litellm_core::chat_completions::{ Error, chat_completions, chat_completions_decline_reason, types::ChatCompletionsRequest, }; use litellm_http::transport::Error as TransportError; +use litellm_types::utils::ChatCompletionsResponse; use rstest::{fixture, rstest}; use serde_json::{Map, Value, json}; use wiremock::ResponseTemplate; @@ -13,6 +14,10 @@ use support::*; const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; +async fn complete(request: ChatCompletionsRequest<'_>) -> Result { + chat_completions(&http_pool(), &http_config(), request).await +} + fn object(value: Value) -> Map { let Value::Object(map) = value else { panic!("expected a json object, got {value}"); @@ -50,7 +55,7 @@ async fn anthropic_round_trip_translates_the_conversation_and_normalizes_the_res let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; let base = upstream.uri(); - let response = chat_completions(ChatCompletionsRequest { + let response = complete(ChatCompletionsRequest { messages: json!([ {"role": "system", "content": "be terse"}, {"role": "user", "content": "hi"} @@ -90,7 +95,7 @@ async fn the_deployment_key_replaces_a_caller_supplied_x_api_key( let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; let base = upstream.uri(); - chat_completions(ChatCompletionsRequest { + complete(ChatCompletionsRequest { api_base: Some(&base), extra_headers: Some(object( json!({"x-api-key": "caller-key", "x-trace": "kept"}), @@ -116,7 +121,7 @@ async fn bedrock_round_trip_is_signed_and_normalized(request: ChatCompletionsReq .await; let base = upstream.uri(); - let response = chat_completions(ChatCompletionsRequest { + let response = complete(ChatCompletionsRequest { model: "bedrock/anthropic.claude-sonnet-4-5", optional_params: object(json!({ "aws_access_key_id": "access-key", @@ -167,7 +172,7 @@ async fn a_response_it_cannot_normalize_is_reported_as_already_sent( let upstream = upstream([anthropic_response(body)]).await; let base = upstream.uri(); - let error = chat_completions(ChatCompletionsRequest { + let error = complete(ChatCompletionsRequest { api_base: Some(&base), ..request }) @@ -188,7 +193,7 @@ async fn an_upstream_error_status_keeps_its_code_and_body( let upstream = upstream([ResponseTemplate::new(status).set_body_string("slow down")]).await; let base = upstream.uri(); - let error = chat_completions(ChatCompletionsRequest { + let error = complete(ChatCompletionsRequest { api_base: Some(&base), ..request }) @@ -210,7 +215,7 @@ async fn an_upstream_error_status_keeps_its_code_and_body( async fn a_connection_that_is_never_established_declines_instead_of_failing( request: ChatCompletionsRequest<'static>, ) { - let error = chat_completions(ChatCompletionsRequest { + let error = complete(ChatCompletionsRequest { api_base: Some(UNREACHABLE_BASE), ..request }) @@ -232,7 +237,7 @@ async fn a_timeout_after_sending_is_not_a_pre_send_decline( upstream([anthropic_response(ANTHROPIC_MESSAGE).set_delay(Duration::from_secs(5))]).await; let base = upstream.uri(); - let error = chat_completions(ChatCompletionsRequest { + let error = complete(ChatCompletionsRequest { api_base: Some(&base), timeout: Some(Duration::from_millis(100)), ..request @@ -307,7 +312,7 @@ async fn a_declined_request_fails_the_call_before_sending( let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; let base = upstream.uri(); - let error = chat_completions(ChatCompletionsRequest { + let error = complete(ChatCompletionsRequest { optional_params: object(json!({"stream": true})), api_base: Some(&base), ..request diff --git a/litellm-rust/crates/core/tests/messages/host.rs b/litellm-rust/crates/core/tests/messages/host.rs index ca2aece5ebd..844ada3e1ad 100644 --- a/litellm-rust/crates/core/tests/messages/host.rs +++ b/litellm-rust/crates/core/tests/messages/host.rs @@ -78,7 +78,7 @@ impl Host for RecordingHost { } async fn run_through(host: &RecordingHost) -> Result { - litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await + litellm_host::run::run(machine(Arc::new(RecordingSecrets::empty())), host).await } fn authenticated(call: MessagesCall, api_base: String) -> MessagesCall { diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 21ee678ced3..1ae822e5437 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -2,9 +2,11 @@ use std::{sync::Arc, time::Duration}; use litellm_core::messages::{ Error, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + route::{LocalMessagesHost, MessagesCall, MessagesMachine, MessagesOutput, messages_machine}, types::MessagesShaping, }; +use litellm_http::{HttpSettings, Resolution}; +use litellm_secrets::source::SecretSource; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use rstest::fixture; use serde_json::{Map, Value, json}; @@ -75,11 +77,16 @@ fn headers<'a>(pairs: impl IntoIterator) -> Option) -> MessagesMachine { + messages_machine(&http_pool(), &http_config(), secrets) + .expect("default HTTP settings build a client") +} + async fn run_with( secrets: Arc, call: MessagesCall, ) -> Result { - litellm_host::run::run(messages_machine(secrets), &LocalMessagesHost::new(call)).await + litellm_host::run::run(machine(secrets), &LocalMessagesHost::new(call)).await } /// Runs the route with a secret source that knows nothing, so no environment leaks in. diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index 133b7d2b162..431dd4f4b93 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -190,29 +190,40 @@ fn facade_request(body: Value, api_base: &str) -> MessagesRequest<'_> { } #[tokio::test] -async fn the_facade_runs_the_route_in_process() { +async fn the_facade_sends_through_the_injected_http_pool_configuration() { let upstream = upstream([message_response()]).await; let base = upstream.uri(); + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; - let message = messages(facade_request( - json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), - &base, - )) + let message = messages( + &http_pool(), + &Resolution::from(&settings).config, + facade_request( + json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), + &base, + ), + ) .await .expect("messages request succeeds"); assert_eq!(message.id, "msg_1"); - assert_eq!( - only_request(&upstream).await.header("x-api-key"), - Some("sk-ant") - ); + let sent = only_request(&upstream).await; + assert_eq!(sent.header("x-api-key"), Some("sk-ant")); + assert_eq!(sent.header("user-agent"), Some("host-owned/1")); } #[tokio::test] async fn the_facade_rejects_a_body_that_is_not_an_object() { - let error = messages(facade_request(json!([]), UNREACHABLE_BASE)) - .await - .expect_err("a non-object body is rejected"); + let error = messages( + &http_pool(), + &http_config(), + facade_request(json!([]), UNREACHABLE_BASE), + ) + .await + .expect_err("a non-object body is rejected"); assert_eq!( error, diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index c4be3127d66..4ca6e609052 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -87,7 +87,7 @@ fn sse_response() -> ResponseTemplate { } async fn stream_through(host: &RecordingStreamHost) -> Result { - litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await + litellm_host::run::run(machine(Arc::new(RecordingSecrets::empty())), host).await } #[rstest] diff --git a/litellm-rust/crates/core/tests/ocr/main.rs b/litellm-rust/crates/core/tests/ocr/main.rs index 1a915389b20..e1f6b8cb5c1 100644 --- a/litellm-rust/crates/core/tests/ocr/main.rs +++ b/litellm-rust/crates/core/tests/ocr/main.rs @@ -4,6 +4,7 @@ use litellm_core::ocr::{ types::LiteLLMOcrRequest, wire::{OcrWireRequest, decode_request}, }; +use litellm_http::Client; use litellm_llms::base_llm::ocr::{ error::Error, handler::OcrClient, @@ -37,11 +38,7 @@ fn object(value: Value) -> Map { } fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) + OcrClient::for_test(Client::plain_for_test(), Client::no_redirect_for_test()) } async fn perform(request: LiteLLMOcrRequest) -> Result { diff --git a/litellm-rust/crates/core/tests/ocr/mistral.rs b/litellm-rust/crates/core/tests/ocr/mistral.rs index f80e564b03f..4c3f1c5cc39 100644 --- a/litellm-rust/crates/core/tests/ocr/mistral.rs +++ b/litellm-rust/crates/core/tests/ocr/mistral.rs @@ -1,10 +1,7 @@ use std::sync::Arc; use litellm_auth_gcp::VertexAuth; -use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, -}; +use litellm_http::{HttpSettings, Resolution, media::UrlPolicy}; use litellm_llms::{ base_llm::ocr::{ settings::OcrSettings, @@ -192,12 +189,16 @@ async fn the_client_uses_the_injected_http_pool_configuration() { ..HttpSettings::default() }; let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &http_pool(), &Resolution::from(&settings).config, UrlPolicy::default(), VertexAuth::default(), OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + Arc::new( + litellm_secrets::source::EnvironmentSecrets::python_compatible( + litellm_http::Client::plain_for_test(), + ), + ), ) .unwrap(); diff --git a/litellm-rust/crates/core/tests/support/mod.rs b/litellm-rust/crates/core/tests/support/mod.rs index 4d2fe0232d0..1d9af236811 100644 --- a/litellm-rust/crates/core/tests/support/mod.rs +++ b/litellm-rust/crates/core/tests/support/mod.rs @@ -3,9 +3,12 @@ #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; +use litellm_http::{ + HttpClientConfig, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, +}; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::Value; use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; @@ -13,6 +16,14 @@ use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; /// A port nothing listens on, for calls that must fail before any request is sent. pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1"; +pub fn http_pool() -> HttpClientPool { + HttpClientPool::new(Arc::new(PublicDnsResolver)) +} + +pub fn http_config() -> HttpClientConfig { + Resolution::from(&HttpSettings::default()).config +} + /// Starts an upstream that answers its n-th request with the n-th response and 404s after. pub async fn upstream(responses: impl IntoIterator) -> MockServer { let server = MockServer::start().await; diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index cad5aa87e49..0cb2b15b768 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -22,5 +22,7 @@ veil.workspace = true webpki-roots.workspace = true [dev-dependencies] +rcgen = "0.14.10" +tempfile.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/client.rs b/litellm-rust/crates/http/src/client.rs new file mode 100644 index 00000000000..1f7017d083b --- /dev/null +++ b/litellm-rust/crates/http/src/client.rs @@ -0,0 +1,38 @@ +use std::ops::Deref; + +#[derive(Clone, Debug)] +pub struct Client(reqwest::Client); + +impl Client { + pub(crate) fn new(client: reqwest::Client) -> Self { + Self(client) + } + + #[cfg(any(test, feature = "test-support"))] + pub fn for_test(client: reqwest::Client) -> Self { + Self(client) + } + + #[cfg(any(test, feature = "test-support"))] + pub fn plain_for_test() -> Self { + Self(reqwest::Client::new()) + } + + #[cfg(any(test, feature = "test-support"))] + pub fn no_redirect_for_test() -> Self { + Self( + reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("a client without TLS or proxy settings builds"), + ) + } +} + +impl Deref for Client { + type Target = reqwest::Client; + + fn deref(&self) -> &reqwest::Client { + &self.0 + } +} diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index cb0173369d5..2f36784bc70 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -18,10 +18,16 @@ pub enum Verify { BuiltInRoots, } +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum ClientIdentity { + Pem(PathBuf), + Split { certificate: PathBuf, key: PathBuf }, +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct HttpClientConfig { pub verify: Verify, - pub client_certificate: Option, + pub client_certificate: Option, pub key_exchange_group: Option, pub tls12_cipher_suites: Option>, pub force_ipv4: bool, @@ -67,7 +73,7 @@ impl From<&HttpSettings> for Resolution { Self { config: HttpClientConfig { verify: Verify::from(settings), - client_certificate: settings.ssl_certificate.clone(), + client_certificate: settings.ssl_certificate.clone().map(ClientIdentity::Pem), key_exchange_group: curve.clone().ok().flatten(), tls12_cipher_suites: ciphers.tls12_cipher_suites, force_ipv4: settings.force_ipv4, @@ -276,7 +282,7 @@ mod tests { config, HttpClientConfig { verify: Verify::BuiltInRoots, - client_certificate: Some("/client.pem".into()), + client_certificate: Some(ClientIdentity::Pem("/client.pem".into())), key_exchange_group: None, tls12_cipher_suites: None, force_ipv4: true, diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index a1456208bb3..3e55a1843c8 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,3 +1,10 @@ +#![allow( + clippy::disallowed_types, + clippy::disallowed_methods, + reason = "this crate is the one place reqwest clients are built" +)] + +mod client; mod config; mod error; pub mod media; @@ -9,7 +16,8 @@ mod settings; mod tls; pub mod transport; -pub use config::{HttpClientConfig, Resolution, Verify}; +pub use client::Client; +pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify}; pub use error::{Error, TlsSource}; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 1b9159973ef..1dac68305b0 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -12,7 +12,7 @@ use reqwest::{ dns::{Addrs, Name, Resolve, Resolving}, }; -use crate::{ClientVariant, HttpClientConfig, HttpClientPool}; +use crate::{Client, ClientVariant, HttpClientConfig, HttpClientPool}; #[derive(Debug, thiserror::Error)] pub enum Error { @@ -93,8 +93,8 @@ type ProxyMatch = Arc bool + Send + Sync>; #[derive(Clone)] pub struct MediaFetcher { - pinned: reqwest::Client, - unpinned: reqwest::Client, + pinned: Client, + unpinned: Client, uses_proxy: ProxyMatch, address_resolver: Arc, url_policy: UrlPolicy, @@ -154,7 +154,7 @@ impl MediaFetcher { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(client: reqwest::Client) -> Self { + pub fn for_test(client: Client) -> Self { Self { pinned: client.clone(), unpinned: client, @@ -230,7 +230,7 @@ impl MediaFetcher { } } - async fn client_for(&self, url: &Url) -> Result<&reqwest::Client, Error> { + async fn client_for(&self, url: &Url) -> Result<&Client, Error> { if !self.url_policy.validate { return Ok(&self.unpinned); } @@ -520,10 +520,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test client builds"); + let client = Client::no_redirect_for_test(); let media = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await @@ -539,10 +536,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test client builds"); + let client = Client::no_redirect_for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(2, 0)) .await @@ -557,10 +551,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n", ) .await; - let client = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test client builds"); + let client = Client::no_redirect_for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index d100bdf624b..c2cfb00d79b 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -107,7 +107,7 @@ impl OutboundRequest { self.timeout } - pub async fn send(self, client: &reqwest::Client) -> Result { + pub async fn send(self, client: &crate::Client) -> Result { let builder = with_headers( client.post(&self.url).body(self.body), &self.headers, diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index ee47e5dc52a..1187c34f2d7 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -6,7 +6,7 @@ use std::{ use reqwest::dns::Resolve; -use crate::{config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; +use crate::{client::Client, config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { @@ -48,7 +48,7 @@ impl HttpClientPool { &self, config: &HttpClientConfig, variant: ClientVariant, - ) -> Result { + ) -> Result { let effective = match variant { ClientVariant::Media => HttpClientConfig { client_certificate: None, @@ -65,7 +65,7 @@ impl HttpClientPool { if let Some(pooled) = self.lock().get(&key) && pooled.built_at.elapsed() < self.ttl { - return Ok(pooled.client.clone()); + return Ok(Client::new(pooled.client.clone())); } let client = self .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) @@ -77,7 +77,7 @@ impl HttpClientPool { built_at: Instant::now(), }, ); - Ok(client) + Ok(Client::new(client)) } fn lock(&self) -> MutexGuard<'_, Clients> { @@ -116,7 +116,7 @@ mod tests { }; use super::*; - use crate::{HttpSettings, Resolution, Verify}; + use crate::{ClientIdentity, HttpSettings, Resolution, Verify}; struct FixedResolver(SocketAddr); @@ -288,7 +288,9 @@ mod tests { fn media_variant_never_loads_the_client_certificate() { let pool = pool(); let with_identity = HttpClientConfig { - client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")), + client_certificate: Some(ClientIdentity::Pem( + std::env::temp_dir().join("litellm-http-absent-client.pem"), + )), ..config("a") }; assert!( diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index e2e6d27cd54..c58076607e4 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -8,7 +8,7 @@ use rustls::{ }; use crate::{ - config::{HttpClientConfig, Verify}, + config::{ClientIdentity, HttpClientConfig, Verify}, error::{Error, TlsSource}, }; @@ -203,11 +203,15 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { }; let mut tls = match &config.client_certificate { None => verified.with_no_client_auth(), - Some(path) => { - let (chain, key) = identity(path, TlsSource::ClientIdentity)?; + Some(identity) => { + let (certificate, key) = match identity { + ClientIdentity::Pem(path) => (path, path), + ClientIdentity::Split { certificate, key } => (certificate, key), + }; + let (chain, private_key) = client_identity(certificate, key)?; verified - .with_client_auth_cert(chain, key) - .map_err(|error| invalid_pem(path, TlsSource::ClientIdentity, error))? + .with_client_auth_cert(chain, private_key) + .map_err(|error| invalid_pem(key, TlsSource::ClientIdentity, error))? } }; tls.alpn_protocols = if config.http2 { @@ -233,17 +237,18 @@ fn bundle_roots(path: &Path, source: TlsSource) -> Result Ok(store) } -fn identity( - path: &Path, - source: TlsSource, +fn client_identity( + certificate: &Path, + key: &Path, ) -> Result<(Vec>, PrivateKeyDer<'static>), Error> { - let chain = certificates(path, source)?; + let source = TlsSource::ClientIdentity; + let chain = certificates(certificate, source)?; if chain.is_empty() { - return Err(invalid_pem(path, source, "no certificates found")); + return Err(invalid_pem(certificate, source, "no certificates found")); } - let key = PrivateKeyDer::from_pem_slice(&read(path, source)?) - .map_err(|error| invalid_pem(path, source, error))?; - Ok((chain, key)) + let private_key = PrivateKeyDer::from_pem_slice(&read(key, source)?) + .map_err(|error| invalid_pem(key, source, error))?; + Ok((chain, private_key)) } fn certificates(path: &Path, source: TlsSource) -> Result>, Error> { @@ -405,7 +410,7 @@ mod tests { ) .unwrap(); let result = ClientConfig::try_from(&HttpClientConfig { - client_certificate: Some(path.clone()), + client_certificate: Some(ClientIdentity::Pem(path.clone())), ..config(HttpSettings::default()) }) .map(drop); @@ -419,4 +424,29 @@ mod tests { }) if reported == path )); } + + #[test] + fn split_client_identity_reads_the_key_from_its_own_file() { + let identity = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let directory = tempfile::tempdir().unwrap(); + let certificate = directory.path().join("client.crt"); + let key = directory.path().join("client.key"); + std::fs::write(&certificate, identity.cert.pem()).unwrap(); + std::fs::write(&key, identity.signing_key.serialize_pem()).unwrap(); + + let split = ClientConfig::try_from(&HttpClientConfig { + client_certificate: Some(ClientIdentity::Split { + certificate: certificate.clone(), + key, + }), + ..config(HttpSettings::default()) + }); + let combined = ClientConfig::try_from(&HttpClientConfig { + client_certificate: Some(ClientIdentity::Pem(certificate)), + ..config(HttpSettings::default()) + }); + + assert!(split.unwrap().client_auth_cert_resolver.has_certs()); + assert!(combined.is_err()); + } } diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index ed15d9f7cdb..36ccd18f220 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -35,6 +35,7 @@ tokio = { workspace = true, features = ["sync"] } url.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } aws-smithy-eventstream = "=0.61.4" aws-smithy-types = "1.6.1" rstest.workspace = true diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 51a2668310e..b4e9d01f867 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -450,7 +450,7 @@ fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result: Send + Sync { #[derive(Clone)] pub struct OcrClient { - provider_http: reqwest::Client, - polling_http: reqwest::Client, + provider_http: Client, + polling_http: Client, document_fetcher: MediaFetcher, vertex_auth: VertexAuth, settings: OcrSettings, @@ -60,11 +60,11 @@ impl OcrClient { }) } - pub fn provider_http(&self) -> &reqwest::Client { + pub fn provider_http(&self) -> &Client { &self.provider_http } - pub fn polling_http(&self) -> &reqwest::Client { + pub fn polling_http(&self) -> &Client { &self.polling_http } @@ -85,17 +85,18 @@ impl OcrClient { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { + pub fn for_test(provider_http: Client, no_redirect_http: Client) -> Self { Self { + secrets: Arc::new( + litellm_secrets::source::EnvironmentSecrets::python_compatible( + provider_http.clone(), + ), + ), provider_http, - polling_http: reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test polling client builds"), - document_fetcher: MediaFetcher::for_test(document_http), + polling_http: no_redirect_http.clone(), + document_fetcher: MediaFetcher::for_test(no_redirect_http), vertex_auth: VertexAuth::default(), settings: OcrSettings::default(), - secrets: Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), } } @@ -311,7 +312,7 @@ mod tests { let _connection = listener.accept().await.unwrap(); tokio::time::sleep(Duration::from_secs(1)).await; }); - let error = reqwest::Client::new() + let error = litellm_http::Client::plain_for_test() .get(format!("http://{address}")) .timeout(Duration::from_millis(10)) .send() diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index c3377536545..147056dab8d 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -564,7 +564,10 @@ mod tests { let params = ReductoParseV3Config .map_ocr_params(&overrides, "parse-v3") .unwrap(); - let client = OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()); + let client = OcrClient::for_test( + litellm_http::Client::plain_for_test(), + litellm_http::Client::no_redirect_for_test(), + ); let connection = OcrConnection::default(); let document = serde_json::from_value( json!({"type":"document_url","document_url":"reducto://ready.pdf"}), diff --git a/litellm-rust/crates/llms/tests/ocr_handler.rs b/litellm-rust/crates/llms/tests/ocr_handler.rs index 6e46e6f76d4..5ed3244087d 100644 --- a/litellm-rust/crates/llms/tests/ocr_handler.rs +++ b/litellm-rust/crates/llms/tests/ocr_handler.rs @@ -19,7 +19,7 @@ async fn read_bounded(response: String, limit: usize) -> Result().await; }); - let response = reqwest::Client::new() + let response = litellm_http::Client::plain_for_test() .get(format!("http://{address}")) .send() .await diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 057cad2f42e..d965223bd59 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -62,6 +62,7 @@ tokio = { workspace = true, features = ["rt", "sync"] } url.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-secrets-aws.workspace = true serde.workspace = true serde_with.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/activation.rs b/litellm-rust/crates/python-bridge/src/cache/activation.rs index f77032c579d..58735679554 100644 --- a/litellm-rust/crates/python-bridge/src/cache/activation.rs +++ b/litellm-rust/crates/python-bridge/src/cache/activation.rs @@ -9,10 +9,10 @@ use super::{ cache_error, config::{CacheBackendConfig, NativeCacheConfig, UnsupportedCacheConfig}, embedder::PythonEmbedder, - host_client, native::NativeResponseCache, }; use crate::errors::RustBridgeDeclined; +use crate::http::host_client; fn declined(reason: UnsupportedCacheConfig) -> PyErr { RustBridgeDeclined::new_err(reason.message()) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 6e25f07efa1..e58902b07ee 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1511,7 +1511,7 @@ mod tests { path_service_account: Some("credentials.json".into()), endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(), }, - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), Some("token".into()), ); let matching_config = NativeCacheConfig { diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index e14916b25c6..fcc8aa6218a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,3 +1,4 @@ +use crate::http::host_client; use crate::logger::run_sync_value; use litellm_auth_aws::AwsAuthConfig; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; @@ -19,7 +20,6 @@ use super::{ config::{QdrantSemanticCacheConfig, project_redis_semantic}, embedder::PythonEmbedder, facade::FacadeGuard, - host_client, native::NativeResponseCache, request::duration, }; diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index ac1e00d5273..00b0c71684a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -13,11 +13,9 @@ mod resolver; mod semantic; use litellm_cache::Error; -use litellm_http::ClientVariant; use pyo3::{ exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, - types::PyDict, }; pub(crate) use self::{binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheResolver}; @@ -29,11 +27,3 @@ fn cache_error(error: Error) -> PyErr { _ => PyRuntimeError::new_err(error.to_string()), } } - -/// The host's pooled HTTP client, configured from the proxy's HTTP settings. -fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult { - let http_config = crate::http::call_config(py, &PyDict::new(py), true)?; - crate::http::pool() - .client(&http_config, variant) - .map_err(crate::http::client_error) -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 0e279046812..460136baa1f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -92,7 +92,7 @@ impl NativeResponseCache { )) } - pub async fn s3(config: S3CacheConfig, http: reqwest::Client) -> Self { + pub async fn s3(config: S3CacheConfig, http: litellm_http::Client) -> Self { let runtime = tokio::runtime::Handle::current(); let backend = S3Cache::new(config, http, ResponseCacheCodec, runtime); let identity = BackendIdentity::S3 { @@ -112,7 +112,7 @@ impl NativeResponseCache { Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) } - pub fn gcs(config: GcsConfig, client: reqwest::Client, token: Option) -> Self { + pub fn gcs(config: GcsConfig, client: litellm_http::Client, token: Option) -> Self { let backend = match token { Some(token) => GcsCache::with_token_source( config, @@ -133,7 +133,7 @@ impl NativeResponseCache { pub async fn azure_blob( account_url: &str, container: &str, - http: reqwest::Client, + http: litellm_http::Client, ) -> Result { let backend = AzureBlobCache::connect( account_url, @@ -242,7 +242,7 @@ impl NativeResponseCache { pub async fn qdrant_semantic( config: QdrantSemanticCacheConfig, - client: reqwest::Client, + client: litellm_http::Client, runtime: tokio::runtime::Handle, ) -> Result { let qdrant = qdrant_client::Qdrant::from_url(&config.grpc_url) diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 3dad3447f45..4d8f0fd7147 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -6,8 +6,8 @@ use std::{ use litellm_core_utils::settings::ProcessEnvironment; use litellm_http::{ - HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, - TlsSource, Unsupported, + Client, ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, + Resolution, SslVerify, TlsSource, Unsupported, media::{PublicDnsResolver, UrlPolicy}, }; use pyo3::{ @@ -97,7 +97,10 @@ pub(crate) fn call_config( let settings = HttpSettings::from_layers([ for_call(call_ssl_verify(kwargs)?, asynchronous), HttpSettingsLayer::from_environment(&ProcessEnvironment), - configured(&PythonSettings::Http.read(py)?)?, + match PythonSettings::Http.read_or_unset(py)? { + Some(snapshot) => configured(&snapshot)?, + None => HttpSettingsLayer::default(), + }, ]) .without_missing_files(&|path: &Path| path.exists()); let resolution = Resolution::from(&settings); @@ -107,6 +110,11 @@ pub(crate) fn call_config( Ok(resolution.config) } +pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult { + let config = call_config(py, &PyDict::new(py), true)?; + pool().client(&config, variant).map_err(client_error) +} + pub(crate) fn client_error(error: litellm_http::Error) -> PyErr { match error { litellm_http::Error::Read { @@ -143,7 +151,10 @@ fn unreported( } pub(crate) fn url_policy(py: Python<'_>) -> PyResult { - project_url_policy(&PythonSettings::UrlPolicy.read(py)?) + match PythonSettings::UrlPolicy.read_or_unset(py)? { + Some(snapshot) => project_url_policy(&snapshot), + None => Ok(UrlPolicy::default()), + } } fn project_url_policy(snapshot: &Snapshot<'_>) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index abf664b795d..f03e5fdce7f 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,4 +1,4 @@ -use pyo3::prelude::*; +use pyo3::{exceptions::PyModuleNotFoundError, prelude::*}; use crate::coercion::{FieldSpec, ProjectionError}; @@ -40,6 +40,16 @@ impl PythonSettings { Ok(Snapshot { group: self, value }) } + /// Reads the accessor, or `None` when the litellm package is not installed + /// (a bare extension module), meaning there are no configured values. + pub(crate) fn read_or_unset(self, py: Python<'_>) -> PyResult>> { + match self.read(py) { + Ok(snapshot) => Ok(Some(snapshot)), + Err(error) if error.is_instance_of::(py) => Ok(None), + Err(error) => Err(error), + } + } + #[cfg(test)] pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { Snapshot { group: self, value } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index dec4dcea21c..93d0e11d323 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -3,7 +3,8 @@ use litellm_core::audio_transcription::{ Error, audio_transcription as run_audio_transcription, types::AudioTranscriptionRequest, }; use litellm_host_python::from_py_argument; -use pyo3::prelude::*; +use litellm_http::HttpClientConfig; +use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; use crate::{ @@ -12,6 +13,7 @@ use crate::{ }; async fn execute( + config: HttpClientConfig, audio: Value, optional_params: Map, options: RouteOptions, @@ -24,16 +26,20 @@ async fn execute( extra_headers, timeout, } = options; - run_audio_transcription(AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }) + run_audio_transcription( + crate::http::pool(), + &config, + AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + }, + ) .await } @@ -62,9 +68,10 @@ pub(crate) fn transcription( extra_headers, timeout: optional_timeout(timeout_seconds), }; + let config = crate::http::call_config(py, &PyDict::new(py), false)?; run_sync( py, - execute(audio, optional_params.unwrap_or_default(), options), + execute(config, audio, optional_params.unwrap_or_default(), options), audio_transcription_error_to_pyerr, ) } @@ -94,9 +101,10 @@ pub(crate) fn atranscription<'py>( extra_headers, timeout: optional_timeout(timeout_seconds), }; + let config = crate::http::call_config(py, &PyDict::new(py), true)?; run_async( py, - execute(audio, optional_params.unwrap_or_default(), options), + execute(config, audio, optional_params.unwrap_or_default(), options), audio_transcription_error_to_pyerr, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index b96b12bfc43..6d7fad0d69c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -7,6 +7,7 @@ use litellm_core::chat_completions::{ types::ChatCompletionsRequest, }; use litellm_host_python::from_py_argument; +use litellm_http::HttpClientConfig; use litellm_types::utils::ChatCompletionsResponse; use pyo3::prelude::*; use serde_json::{Map, Value}; @@ -20,6 +21,7 @@ use crate::{ }; async fn execute( + config: HttpClientConfig, messages: Vec, optional_params: Map, options: RouteOptions, @@ -32,16 +34,20 @@ async fn execute( extra_headers, timeout, } = options; - run_chat_completions(ChatCompletionsRequest { - model: &model, - messages: Value::Array(messages), - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) + run_chat_completions( + crate::http::pool(), + &config, + ChatCompletionsRequest { + model: &model, + messages: Value::Array(messages), + optional_params, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }, + ) .await } @@ -87,9 +93,15 @@ pub(crate) fn chat_completions( extra_headers, timeout: optional_timeout(timeout_seconds), }; + let config = crate::http::call_config(py, &PyDict::new(py), false)?; run_sync( py, - execute(messages, optional_params.unwrap_or_default(), options), + execute( + config, + messages, + optional_params.unwrap_or_default(), + options, + ), chat_completions_error_to_pyerr, ) } @@ -119,9 +131,15 @@ pub(crate) fn achat_completions<'py>( extra_headers, timeout: optional_timeout(timeout_seconds), }; + let config = crate::http::call_config(py, &PyDict::new(py), true)?; run_async( py, - execute(messages, optional_params.unwrap_or_default(), options), + execute( + config, + messages, + optional_params.unwrap_or_default(), + options, + ), chat_completions_error_to_pyerr, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 52cebb7c903..a59c9360c36 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -27,11 +27,14 @@ fn run_messages( asynchronous: bool, ) -> PyResult> { let secrets = crate::secrets::source(py)?; + let config = crate::http::call_config(py, &kwargs, asynchronous)?; + let machine = messages_machine(crate::http::pool(), &config, secrets) + .map_err(crate::http::client_error)?; run_legacy_call( py, SURFACE, PublicCall::capture(&request, &args, &kwargs)?, - crate::logger::LoggedMachine::new(messages_machine(secrets)), + crate::logger::LoggedMachine::new(machine), MessagesPythonHost::new(request.unbind()), crate::preflight::sdk_preflight, asynchronous, diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 82bb4443f98..6ba60630b3a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -210,7 +210,7 @@ handler.get_secret_from_manager = get_secret_from_manager KeyManagementSettings::default(), )), Arc::new(move |_: &str| fallback.map(str::to_owned)), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback); (resolver, locals, handler) diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index 439f9ddddd1..ed54c306397 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -27,7 +27,12 @@ const NATIVE: FieldSpec = FieldSpec::new("native", |field| field.schema_bo pub(crate) fn source(py: Python<'_>) -> PyResult> { if PythonSettings::SecretManager.read(py)?.read(&NATIVE)? { let context = litellm_host_python::PythonContext::capture(py)?; - return Ok(Arc::new(ResolvedSecrets::new(config::read(py)?, context))); + let client = crate::http::host_client(py, litellm_http::ClientVariant::Provider)?; + return Ok(Arc::new(ResolvedSecrets::new( + config::read(py)?, + context, + client, + ))); } Ok(Arc::new(PythonSecrets::new(py)?)) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index 40b187c99de..5a606ab1039 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; +use litellm_http::Client; use litellm_secrets::source::SecretSource; use litellm_secrets::{ Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, @@ -15,16 +16,20 @@ pub(crate) struct ResolvedSecrets { } impl ResolvedSecrets { - pub(crate) fn new(snapshot: SecretManagerSnapshot, context: PythonContext) -> Self { - Self::from_state(snapshot.into_state(context)) + pub(crate) fn new( + snapshot: SecretManagerSnapshot, + context: PythonContext, + client: Client, + ) -> Self { + Self::from_state(snapshot.into_state(context), client) } - fn from_state(state: Arc) -> Self { + fn from_state(state: Arc, client: Client) -> Self { Self { resolver: SecretResolver::new_python_compatible( state, Arc::new(ProcessEnvironment), - OidcResolver::default(), + OidcResolver::new(client), ) .with_failure_policy(FailurePolicy::EnvironmentFallback), } @@ -79,7 +84,7 @@ mod tests { } async fn resolve(state: Arc, name: &'static str) -> Option { - ResolvedSecrets::from_state(state) + ResolvedSecrets::from_state(state, litellm_http::Client::plain_for_test()) .resolve(&[name]) .await .unwrap() @@ -175,7 +180,10 @@ mod tests { .expect(1) .mount(&server) .await; - let source = ResolvedSecrets::from_state(state(&server, KeyManagementSettings::default())); + let source = ResolvedSecrets::from_state( + state(&server, KeyManagementSettings::default()), + litellm_http::Client::plain_for_test(), + ); let snapshot = source.resolve(&[declared]).await.unwrap(); assert_eq!(snapshot.get(undeclared), None); let result = source @@ -238,9 +246,12 @@ mod tests { #[tokio::test] async fn oidc_failures_are_not_converted_to_missing_secrets() { - let result = ResolvedSecrets::from_state(Arc::new(SecretManagerState::default())) - .resolve(&["oidc/"]) - .await; + let result = ResolvedSecrets::from_state( + Arc::new(SecretManagerState::default()), + litellm_http::Client::plain_for_test(), + ) + .resolve(&["oidc/"]) + .await; assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc))); } diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index 1a89130ee82..4d2e88115c8 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -10,6 +10,7 @@ use litellm_secrets_types::PythonSecretRead; use pyo3::{ exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, prelude::*, + types::PyDict, }; #[derive(Clone, PartialEq)] @@ -44,10 +45,18 @@ impl NativeSecretManager { let system = configuration.system; let settings = configuration.settings.clone(); let enterprise_enabled = configuration.enterprise_enabled; + let http_config = crate::http::call_config(py, &PyDict::new(py), false)?; let backend = run_sync_value(py, async move { - load_native_manager(system, settings, environment, enterprise_enabled) - .await - .map_err(|error| PyValueError::new_err(error.to_string())) + load_native_manager( + crate::http::pool(), + &http_config, + system, + settings, + environment, + enterprise_enabled, + ) + .await + .map_err(|error| PyValueError::new_err(error.to_string())) })?; Ok(Self { backend, diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 7e8a79f89ef..efdf681e2bc 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true tokio.workspace = true litellm-auth-azure.workspace = true litellm-auth-types.workspace = true @@ -18,6 +19,7 @@ veil.workspace = true percent-encoding = "2.3" [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } wiremock = "0.6.5" rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index 13e59f8e4ac..095c451927c 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -19,7 +19,7 @@ const PATH_SEGMENT: &AsciiSet = &NON_ALPHANUMERIC #[derive(Clone)] pub struct AzureKeyVault { - client: reqwest::Client, + client: litellm_http::Client, vault: reqwest::Url, auth: Arc, inputs: Arc, @@ -33,7 +33,7 @@ struct SecretResponse { impl AzureKeyVault { pub fn with_client( - client: reqwest::Client, + client: litellm_http::Client, vault: reqwest::Url, environment: Arc, ) -> Result { @@ -57,7 +57,10 @@ impl AzureKeyVault { }) } - pub fn new(environment: Arc) -> Result { + pub fn new( + client: litellm_http::Client, + environment: Arc, + ) -> Result { let value = environment .get(AZURE_KEY_VAULT_URI) .ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?; @@ -65,7 +68,7 @@ impl AzureKeyVault { if vault.scheme() != "https" || vault.host_str().is_none() { return Err(Error::VaultUri); } - Self::with_client(reqwest::Client::new(), vault, environment) + Self::with_client(client, vault, environment) } pub fn scope(&self) -> &str { diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index a21149db345..fcbc46092e1 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -130,11 +130,14 @@ fn new_validates_vault_environment( #[case] uri: Option<&'static str>, #[case] missing_environment: bool, ) { - let result = AzureKeyVault::new(Arc::new(move |name: &str| { - (name == "AZURE_KEY_VAULT_URI") - .then(|| uri.map(str::to_owned)) - .flatten() - })); + let result = AzureKeyVault::new( + litellm_http::Client::plain_for_test(), + Arc::new(move |name: &str| { + (name == "AZURE_KEY_VAULT_URI") + .then(|| uri.map(str::to_owned)) + .flatten() + }), + ); if missing_environment { assert!(matches!( @@ -155,7 +158,7 @@ fn new_validates_vault_environment( #[case::local("http://localhost:8080", "https://localhost/.default")] fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) { let manager = AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), uri.parse().unwrap(), Arc::new(|_: &str| None), ) @@ -184,7 +187,7 @@ async fn missing_credentials_do_not_request_vault() { fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| { (name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned()) diff --git a/litellm-rust/crates/secrets-azure/tests/live.rs b/litellm-rust/crates/secrets-azure/tests/live.rs index 18306382613..429cd3013f7 100644 --- a/litellm-rust/crates/secrets-azure/tests/live.rs +++ b/litellm-rust/crates/secrets-azure/tests/live.rs @@ -10,7 +10,7 @@ use rstest::rstest; #[ignore] async fn reads_a_real_secret() { let environment = Arc::new(ProcessEnvironment); - let manager = AzureKeyVault::new(environment).unwrap(); + let manager = AzureKeyVault::new(litellm_http::Client::plain_for_test(), environment).unwrap(); let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap(); let secret = manager.get_secret(&name).await.unwrap().unwrap(); assert!(matches!(&secret, Secret::String(_))); diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index 1c280171f4c..0a91c61ade9 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] litellm-secrets-types.workspace = true litellm-core-utils.workspace = true +litellm-http.workspace = true base64.workspace = true moka.workspace = true reqwest.workspace = true @@ -19,6 +20,7 @@ percent-encoding = "2.3" tokio = { workspace = true, features = ["sync"] } [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } rcgen = "0.14.10" rstest.workspace = true tempfile = "3.27.0" diff --git a/litellm-rust/crates/secrets-cyberark/src/error.rs b/litellm-rust/crates/secrets-cyberark/src/error.rs index 3dfeb95fe26..4f21647225b 100644 --- a/litellm-rust/crates/secrets-cyberark/src/error.rs +++ b/litellm-rust/crates/secrets-cyberark/src/error.rs @@ -18,6 +18,8 @@ pub enum Error { MissingCredentials, #[error("CyberArk client certificate could not be loaded")] ClientCertificate, + #[error("CyberArk Conjur HTTP client could not be built")] + Client(#[redact] Box), #[error("invalid refresh interval")] RefreshInterval, #[error("invalid CyberArk Conjur endpoint")] diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index 252a99c917f..37d547ff6c1 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -2,10 +2,13 @@ mod client; mod read; mod write; -use std::{fs, sync::Arc, time::Duration}; +use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; +use litellm_http::{ + Client, ClientIdentity, ClientVariant, HttpClientConfig, HttpClientPool, TlsSource, Verify, +}; use litellm_secrets_types::{ BaseSecretManager, CyberarkOperationContext, RotationError, SecretCache, SecretDeleter, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, @@ -37,7 +40,7 @@ const SECRET_NAME_SAFE: &AsciiSet = &NON_ALPHANUMERIC #[derive(Clone)] pub struct CyberArkSecretManager { - client: reqwest::Client, + client: Client, endpoint: reqwest::Url, account: String, username: String, diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs index 1d99fe474d5..052d5570896 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs @@ -2,7 +2,7 @@ use super::*; impl CyberArkSecretManager { pub fn with_client( - client: reqwest::Client, + client: Client, endpoint: reqwest::Url, account: String, username: String, @@ -30,6 +30,8 @@ impl CyberArkSecretManager { } pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, environment: Arc, enterprise_enabled: bool, ) -> Result { @@ -46,21 +48,34 @@ impl CyberArkSecretManager { .get(CYBERARK_SSL_VERIFY) .map(|value| !value.trim().eq_ignore_ascii_case("false")) .unwrap_or(true); - let mut builder = reqwest::Client::builder(); if !verify { litellm_tracing::warn!( "CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates." ); - builder = builder.danger_accept_invalid_certs(true); } - if !cert.is_empty() && !key.is_empty() { - let certificate = fs::read(cert).map_err(|_| Error::ClientCertificate)?; - let private_key = fs::read(key).map_err(|_| Error::ClientCertificate)?; - let identity = reqwest::Identity::from_pem(&[certificate, private_key].concat()) - .map_err(|_| Error::ClientCertificate)?; - builder = builder.identity(identity); - } - let client = builder.build()?; + let config = HttpClientConfig { + verify: effective_verify(verify, &config.verify), + client_certificate: (!cert.is_empty() && !key.is_empty()).then(|| { + ClientIdentity::Split { + certificate: cert.into(), + key: key.into(), + } + }), + ..config.clone() + }; + let client = + pool.client(&config, ClientVariant::Provider) + .map_err(|error| match error { + litellm_http::Error::Read { + tls_source: TlsSource::ClientIdentity, + .. + } + | litellm_http::Error::InvalidPem { + tls_source: TlsSource::ClientIdentity, + .. + } => Error::ClientCertificate, + other => Error::Client(Box::new(other)), + })?; let endpoint = reqwest::Url::parse( &environment .get(CYBERARK_API_BASE) @@ -139,9 +154,35 @@ impl CyberArkSecretManager { } } +fn effective_verify(cyberark_verify: bool, host: &Verify) -> Verify { + match (cyberark_verify, host) { + (false, _) => Verify::Disabled, + (true, Verify::Disabled) => Verify::BuiltInRoots, + (true, host) => host.clone(), + } +} + fn normalize_endpoint(mut endpoint: reqwest::Url) -> reqwest::Url { if !endpoint.path().ends_with('/') { endpoint.set_path(&format!("{}/", endpoint.path())); } endpoint } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + + #[test] + fn cyberark_verification_does_not_follow_a_host_that_disabled_it() { + let bundle = Verify::CaBundle(PathBuf::from("/ca.pem")); + assert_eq!( + effective_verify(true, &Verify::Disabled), + Verify::BuiltInRoots + ); + assert_eq!(effective_verify(true, &bundle), bundle); + assert_eq!(effective_verify(false, &bundle), Verify::Disabled); + } +} diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index 2048e067b6e..783ba6ff67a 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -7,6 +7,8 @@ use std::{ }; use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core_utils::settings::Lookup; +use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext, SecretValue}; use rstest::{fixture, rstest}; diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs index fbd4317f446..1bea77e49ae 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs @@ -77,7 +77,7 @@ async fn authentication_encodes_login(#[case] username: &str, #[case] expected_p .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "acct".into(), username.into(), @@ -211,25 +211,25 @@ fn new_validates_credentials_before_license_and_configuration() { let empty: Arc = Arc::new(|_: &str| None); assert!(matches!( - CyberArkSecretManager::new(empty, true), + from_environment(empty, true), Err(Error::MissingCredentials) )); assert!(matches!( - CyberArkSecretManager::new( + from_environment( Arc::new(|name: &str| (name == "CYBERARK_API_KEY").then(|| "k3y".into())), false ), Err(Error::EnterpriseRequired) )); assert!(matches!( - CyberArkSecretManager::new( + from_environment( Arc::new(|name: &str| (name == "CYBERARK_CLIENT_CERT").then(|| "cert".into())), true ), Err(Error::MissingCredentials) )); assert!(matches!( - CyberArkSecretManager::new( + from_environment( Arc::new(|name: &str| match name { "CYBERARK_API_KEY" => Some("k3y".into()), "CYBERARK_REFRESH_INTERVAL" => Some("abc".into()), @@ -240,7 +240,7 @@ fn new_validates_credentials_before_license_and_configuration() { Err(Error::RefreshInterval) )); assert!(matches!( - CyberArkSecretManager::new( + from_environment( Arc::new(|name: &str| match name { "CYBERARK_API_KEY" => Some("k3y".into()), "CYBERARK_API_BASE" => Some("not a url".into()), @@ -254,7 +254,7 @@ fn new_validates_credentials_before_license_and_configuration() { #[rstest] fn certificate_only_credentials_are_validated_as_a_client_identity() { - let result = CyberArkSecretManager::new( + let result = from_environment( Arc::new(|name: &str| match name { "CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()), "CYBERARK_CLIENT_KEY" => Some("/missing/key".into()), @@ -295,7 +295,7 @@ async fn configured_client_identity_preserves_auth_request_and_read_result( let endpoint = server.uri(); let certificate = client_identity_directory.path().join("client.crt"); let key = client_identity_directory.path().join("client.key"); - let manager = CyberArkSecretManager::new( + let manager = from_environment( Arc::new(move |name: &str| match name { "CYBERARK_API_BASE" => Some(endpoint.clone()), "CYBERARK_API_KEY" => Some(api_key.into()), @@ -337,7 +337,7 @@ fn invalid_client_identity_is_not_ignored( let certificate = client_identity_directory.path().join("client.crt"); let key = client_identity_directory.path().join("client.key"); - let result = CyberArkSecretManager::new( + let result = from_environment( Arc::new(move |name: &str| match name { "CYBERARK_API_KEY" => Some(api_key.into()), "CYBERARK_CLIENT_CERT" => Some(certificate.to_str().unwrap().into()), @@ -354,7 +354,7 @@ fn invalid_client_identity_is_not_ignored( #[case::certificate_only("")] #[case::certificate_and_api_key("k3y")] fn client_identity_does_not_bypass_the_enterprise_requirement(#[case] api_key: &'static str) { - let result = CyberArkSecretManager::new( + let result = from_environment( Arc::new(move |name: &str| match name { "CYBERARK_API_KEY" => Some(api_key.into()), "CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()), @@ -381,7 +381,7 @@ async fn new_reads_environment_defaults_end_to_end() { .mount(&server) .await; let endpoint = server.uri(); - let manager = CyberArkSecretManager::new( + let manager = from_environment( Arc::new(move |name: &str| match name { "CYBERARK_API_BASE" => Some(endpoint.clone()), "CYBERARK_API_KEY" => Some("k3y".into()), @@ -404,7 +404,7 @@ async fn new_reads_environment_defaults_end_to_end() { #[rstest] fn new_reports_missing_client_certificate_files() { assert!(matches!( - CyberArkSecretManager::new( + from_environment( Arc::new(|name: &str| match name { "CYBERARK_API_KEY" => Some("k3y".into()), "CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()), @@ -433,7 +433,7 @@ async fn trailing_slash_endpoint_preserves_base_path() { .await; let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap(); let manager = CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), endpoint, "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs index f5ba7a63273..6fa15ecba9b 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs @@ -48,9 +48,21 @@ pub(super) fn client_identity_directory() -> tempfile::TempDir { directory } +pub(super) fn from_environment( + environment: Arc, + enterprise_enabled: bool, +) -> Result { + CyberArkSecretManager::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&HttpSettings::default()).config, + environment, + enterprise_enabled, + ) +} + pub(super) fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager { CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs index 331a26c6119..e9a027091fa 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs @@ -166,7 +166,7 @@ async fn writes_match_python_parity_fixture(parity_fixture: ParityFixture) { .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), parity_fixture.account, parity_fixture.username, @@ -230,7 +230,7 @@ async fn live_conjur_round_trip() { .as_nanos() ); let manager = CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), endpoint.clone(), account.clone(), username.clone(), @@ -245,7 +245,7 @@ async fn live_conjur_round_trip() { .await .unwrap(); let verifier = CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), endpoint.clone(), account.clone(), username.clone(), diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 805eb80740d..208b5ddd03f 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-http.workspace = true moka.workspace = true tokio.workspace = true litellm-auth-gcp = { workspace = true, features = ["google-sdk"] } @@ -24,6 +25,7 @@ serde.workspace = true reqwest.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } google-cloud-auth.workspace = true rstest.workspace = true wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index b8787999e12..08ba466b799 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -23,7 +23,7 @@ const CACHE_CAPACITY: u64 = 200; #[derive(Clone)] pub struct GoogleSecretManager { - client: reqwest::Client, + client: litellm_http::Client, credentials: Arc, endpoint: reqwest::Url, project: String, @@ -46,7 +46,7 @@ struct Payload { impl GoogleSecretManager { pub fn with_client( - client: reqwest::Client, + client: litellm_http::Client, endpoint: reqwest::Url, project: String, environment: Arc, @@ -79,6 +79,7 @@ impl GoogleSecretManager { } pub fn new( + client: litellm_http::Client, environment: Arc, enterprise_enabled: bool, ) -> Result { @@ -104,7 +105,7 @@ impl GoogleSecretManager { .get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER) .is_some_and(|v| v.eq_ignore_ascii_case("true")); Self::with_client( - reqwest::Client::new(), + client, reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"), project, environment, diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index 0d7efc4b1b3..e9bee7633f5 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { GoogleSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "project".into(), Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), @@ -214,11 +214,19 @@ async fn always_read_and_expired_cache_fetch_again( #[rstest] fn google_manager_requires_host_license_and_project_configuration() { assert!(matches!( - GoogleSecretManager::new(Arc::new(|_: &str| None), false), + GoogleSecretManager::new( + litellm_http::Client::plain_for_test(), + Arc::new(|_: &str| None), + false + ), Err(Error::EnterpriseRequired) )); assert!(matches!( - GoogleSecretManager::new(Arc::new(|_: &str| None), true), + GoogleSecretManager::new( + litellm_http::Client::plain_for_test(), + Arc::new(|_: &str| None), + true + ), Err(Error::MissingEnvironment( "GOOGLE_SECRET_MANAGER_PROJECT_ID" )) @@ -236,7 +244,7 @@ fn google_manager_rejects_invalid_refresh_intervals(#[case] variable: &'static s }); assert!(matches!( - GoogleSecretManager::new(environment, true), + GoogleSecretManager::new(litellm_http::Client::plain_for_test(), environment, true), Err(Error::RefreshInterval) )); } diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 17acc01682b..fe61d6cb5b4 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -23,6 +23,7 @@ litellm-secrets-hashicorp = { workspace = true, optional = true } litellm-secrets-azure = { workspace = true, optional = true } litellm-secrets-cyberark = { workspace = true, optional = true } litellm-core-utils.workspace = true +litellm-http.workspace = true base64.workspace = true serde.workspace = true strum.workspace = true @@ -33,6 +34,7 @@ moka.workspace = true tokio = { workspace = true, features = ["fs"] } [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true wiremock = "0.6.5" tempfile = "3" diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 07f2f205bec..7fc756c5c7d 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -10,6 +10,8 @@ pub enum Error { InvalidCiphertext, #[error("decrypted value is not UTF-8")] Utf8, + #[error(transparent)] + Client(#[from] litellm_http::Error), #[error("unsupported OIDC provider or missing build feature")] UnsupportedOidc, #[error("OIDC reference requires a provider and audience")] diff --git a/litellm-rust/crates/secrets/src/native.rs b/litellm-rust/crates/secrets/src/native.rs index 80f0e46245c..f1dc7ccb732 100644 --- a/litellm-rust/crates/secrets/src/native.rs +++ b/litellm-rust/crates/secrets/src/native.rs @@ -1,10 +1,13 @@ use std::sync::Arc; use litellm_core_utils::settings::Lookup; +use litellm_http::{HttpClientConfig, HttpClientPool}; use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; pub async fn load_native_manager( + pool: &HttpClientPool, + config: &HttpClientConfig, system: KeyManagementSystem, settings: KeyManagementSettings, environment: Arc, @@ -29,14 +32,19 @@ pub async fn load_native_manager( } #[cfg(feature = "azure")] (KeyManagementSystem::AzureKeyVault, _, environment, _) => Ok( - SecretManager::AzureKeyVault(crate::azure::AzureKeyVault::new(environment)?), + SecretManager::AzureKeyVault(crate::azure::AzureKeyVault::new( + pool.client(config, litellm_http::ClientVariant::Provider)?, + environment, + )?), ), #[cfg(feature = "google")] - (KeyManagementSystem::GoogleSecretManager, _, environment, enterprise_enabled) => { - Ok(SecretManager::GoogleSecretManager( - crate::google::GoogleSecretManager::new(environment, enterprise_enabled)?, - )) - } + (KeyManagementSystem::GoogleSecretManager, _, environment, enterprise_enabled) => Ok( + SecretManager::GoogleSecretManager(crate::google::GoogleSecretManager::new( + pool.client(config, litellm_http::ClientVariant::Provider)?, + environment, + enterprise_enabled, + )?), + ), #[cfg(feature = "google")] (KeyManagementSystem::GoogleKms, _, environment, _) => { crate::google::load_google_kms(Some(true), environment) @@ -51,11 +59,14 @@ pub async fn load_native_manager( )) } #[cfg(feature = "cyberark")] - (KeyManagementSystem::Cyberark, _, environment, enterprise_enabled) => { - Ok(SecretManager::Cyberark( - crate::cyberark::CyberArkSecretManager::new(environment, enterprise_enabled)?, - )) - } + (KeyManagementSystem::Cyberark, _, environment, enterprise_enabled) => Ok( + SecretManager::Cyberark(crate::cyberark::CyberArkSecretManager::new( + pool, + config, + environment, + enterprise_enabled, + )?), + ), _ => Err(Error::NativeBackendUnavailable), } } diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index f3c1e38ce7b..b6e8dbc123b 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -5,6 +5,7 @@ use std::{ use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_core_utils::settings::Lookup; +use litellm_http::Client; use moka::future::Cache; use serde::Deserialize; @@ -82,7 +83,7 @@ impl NumericDate { } pub struct OidcResolver { - client: reqwest::Client, + client: Client, google_identity_endpoint: reqwest::Url, cache: Cache, clock: fn() -> SystemTime, @@ -90,25 +91,17 @@ pub struct OidcResolver { azure_token_provider: std::sync::Arc, } -impl Default for OidcResolver { - fn default() -> Self { - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(600)) - .connect_timeout(Duration::from_secs(5)) - .build() - .expect("HTTP client configuration"); - Self::new( - client, - reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), - ) - } -} +const GOOGLE_IDENTITY_ENDPOINT: &str = + "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity"; + +const REQUEST_TIMEOUT: Duration = Duration::from_secs(600); impl OidcResolver { - pub fn new(client: reqwest::Client, google_identity_endpoint: reqwest::Url) -> Self { + pub fn new(client: Client) -> Self { Self { client, - google_identity_endpoint, + google_identity_endpoint: reqwest::Url::parse(GOOGLE_IDENTITY_ENDPOINT) + .expect("static URL"), cache: Cache::builder() .max_capacity(200) .time_to_live(GOOGLE_TOKEN_MAX_TTL) @@ -121,6 +114,13 @@ impl OidcResolver { } } + pub fn with_google_identity_endpoint(self, google_identity_endpoint: reqwest::Url) -> Self { + Self { + google_identity_endpoint, + ..self + } + } + #[cfg(feature = "azure")] pub fn with_azure_token_provider( self, @@ -180,6 +180,7 @@ impl OidcResolver { let response = self .client .get(url) + .timeout(REQUEST_TIMEOUT) .query(&[("audience", audience)]) .bearer_auth(authorization) .header("Accept", "application/json; api-version=2.0") @@ -214,6 +215,7 @@ impl OidcResolver { let response = self .client .get(self.google_identity_endpoint.clone()) + .timeout(REQUEST_TIMEOUT) .query(&[("audience", audience)]) .header("Metadata-Flavor", "Google") .send() diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 69445e5b410..830e7c22cd5 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -1,10 +1,7 @@ use std::sync::Arc; use crate::compatibility::python_manager_string; -use litellm_core_utils::{ - serde_compat::parse_str_bool, - settings::{Lookup, ProcessEnvironment}, -}; +use litellm_core_utils::{serde_compat::parse_str_bool, settings::Lookup}; use crate::state::{LookupTarget, normalize_secret_name}; use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; @@ -24,16 +21,6 @@ pub struct SecretResolver { python_compatible: bool, } -impl Default for SecretResolver { - fn default() -> Self { - Self::new( - Arc::new(SecretManagerState::default()), - Arc::new(ProcessEnvironment), - OidcResolver::default(), - ) - } -} - impl SecretResolver { pub fn new( state: Arc, diff --git a/litellm-rust/crates/secrets/src/source.rs b/litellm-rust/crates/secrets/src/source.rs index a1615bad055..3c86fd25e9a 100644 --- a/litellm-rust/crates/secrets/src/source.rs +++ b/litellm-rust/crates/secrets/src/source.rs @@ -37,15 +37,14 @@ impl SecretSource for SecretResolver { } } -#[derive(Default)] pub struct EnvironmentSecrets(SecretResolver); impl EnvironmentSecrets { - pub fn python_compatible() -> Self { + pub fn python_compatible(client: litellm_http::Client) -> Self { Self(SecretResolver::new_python_compatible( Arc::new(crate::SecretManagerState::default()), Arc::new(litellm_core_utils::settings::ProcessEnvironment), - crate::OidcResolver::default(), + crate::OidcResolver::new(client), )) } } diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index 174d0881339..910b1855b77 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -53,7 +53,7 @@ async fn read_results_follow_the_selected_failure_policy( }, )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ) .with_failure_policy(policy); let result = resolver @@ -97,7 +97,7 @@ async fn primary_secret_values_other_than_strings_resolve_to_none( let resolver = SecretResolver::new_python_compatible( Arc::new(state(&server, settings)), Arc::new(|_: &str| Some("fallback".into())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); let text = value.as_str(); assert_eq!( @@ -157,7 +157,7 @@ async fn gating_prediction_matches_actual_lookup( let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index b844b198cd4..60f165add54 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -22,7 +22,7 @@ async fn azure_handler_reads_missing_and_failed_secrets() { .await; let manager = SecretManager::AzureKeyVault( AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -81,7 +81,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ .mount(&server) .await; let manager = AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "token".into())), ) @@ -92,7 +92,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ Default::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index 698e1ad8f63..2e8fbf46008 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -59,7 +59,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Azure => SecretManager::AzureKeyVault( AzureKeyVault::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), environment, ) @@ -67,7 +67,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Google => SecretManager::GoogleSecretManager( GoogleSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "project".into(), environment, @@ -84,7 +84,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { .unwrap(), ), Provider::Cyberark => SecretManager::Cyberark(CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), @@ -240,7 +240,7 @@ async fn python_read_failures_preserve_provider_fallback_rules( KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment_value.map(str::to_owned)), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); let expected = if matches!(provider, Provider::Aws) { None diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index 706c35752d7..b94bd9dc0ad 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -24,7 +24,7 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { .mount(&server) .await; let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 67954fedb78..0b45a90c10b 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -25,7 +25,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) _ => None, }); let manager = GoogleSecretManager::with_client( - reqwest::Client::new(), + litellm_http::Client::plain_for_test(), server.uri().parse().unwrap(), "project".into(), environment.clone(), @@ -40,7 +40,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) let resolver = SecretResolver::new_python_compatible( Arc::new(state), environment, - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ) .with_failure_policy(FailurePolicy::Propagate); let result = resolver.get_secret_str("KEY", None).await; diff --git a/litellm-rust/crates/secrets/tests/hashicorp.rs b/litellm-rust/crates/secrets/tests/hashicorp.rs index bc35b88018e..e10e903c25a 100644 --- a/litellm-rust/crates/secrets/tests/hashicorp.rs +++ b/litellm-rust/crates/secrets/tests/hashicorp.rs @@ -56,7 +56,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }, )), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::default(), + litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( found_resolver @@ -131,7 +131,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let failed_resolver = SecretResolver::new_python_compatible( Arc::new(failed_state), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::default(), + litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), ) .with_failure_policy(FailurePolicy::Propagate); assert!(matches!( diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index afc49e8231d..6f72bd9d645 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -30,7 +30,7 @@ async fn environment_sources_resolve_expected_value( ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), ]); assert_eq!( - OidcResolver::default() + OidcResolver::new(litellm_http::Client::plain_for_test()) .resolve(reference, env.as_ref()) .await .unwrap() @@ -43,7 +43,7 @@ async fn environment_sources_resolve_expected_value( #[tokio::test] async fn environment_sources_bypass_boolean_conversion_and_defaults() { let env = environment(&[("TOKEN", "true")]); - let oidc = OidcResolver::default(); + let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); assert_eq!( resolver @@ -94,7 +94,7 @@ async fn github_requests_are_authenticated_cached_and_revalidate_environment() { ), ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), ]); - let oidc = OidcResolver::default(); + let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); for _ in 0..2 { assert_eq!( oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) @@ -131,7 +131,7 @@ async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explici ("PATH_TOKEN", private.to_str().unwrap()), ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), ]); - let oidc = OidcResolver::default(); + let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); assert_eq!( oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) .await @@ -213,8 +213,9 @@ async fn google_expiry_caps_cache_and_preserves_audience( .expect(calls) .mount(&server) .await; - let oidc = - OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()).with_clock(now); + let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_google_identity_endpoint(server.uri().parse().unwrap()) + .with_clock(now); for _ in 0..2 { assert_eq!( oidc.resolve( @@ -234,7 +235,7 @@ async fn google_expiry_caps_cache_and_preserves_audience( #[tokio::test] async fn google_oidc_requires_its_build_feature() { assert!(matches!( - OidcResolver::default() + OidcResolver::new(litellm_http::Client::plain_for_test()) .resolve("oidc/google/audience", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -245,7 +246,7 @@ async fn google_oidc_requires_its_build_feature() { #[tokio::test] async fn azure_oidc_without_a_token_file_requires_its_build_feature() { assert!(matches!( - OidcResolver::default() + OidcResolver::new(litellm_http::Client::plain_for_test()) .resolve("oidc/azure/scope", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -261,7 +262,7 @@ async fn invalid_references_fail_before_environment_lookup( #[case] reference: &str, #[case] unsupported: bool, ) { - let error = OidcResolver::default() + let error = OidcResolver::new(litellm_http::Client::plain_for_test()) .resolve(reference, &|_: &str| { panic!("invalid reference reached environment lookup") }) @@ -283,7 +284,8 @@ async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { .expect(1) .mount(&server) .await; - let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()); + let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_google_identity_endpoint(server.uri().parse().unwrap()); for _ in 0..2 { assert_eq!( resolver @@ -334,7 +336,8 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] }) } } - let oidc = OidcResolver::default().with_azure_token_provider(Arc::new(Provider(failed))); + let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_azure_token_provider(Arc::new(Provider(failed))); let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), environment(&[("AZURE_CLIENT_ID", "client-id")]), @@ -361,7 +364,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] #[tokio::test] async fn missing_oidc_environment_is_an_error(#[case] reference: &str) { assert!(matches!( - OidcResolver::default() + OidcResolver::new(litellm_http::Client::plain_for_test()) .resolve(reference, environment(&[]).as_ref()) .await, Err(Error::MissingEnvironment) @@ -380,7 +383,8 @@ async fn google_oidc_failures_are_not_cached_or_hidden_by_defaults() { let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), environment(&[]), - OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()), + OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_google_identity_endpoint(server.uri().parse().unwrap()), ); for _ in 0..2 { assert!(matches!( @@ -420,7 +424,8 @@ async fn google_tokens_expire_at_the_python_cache_deadline( .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()) + let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); assert_eq!( resolver @@ -472,7 +477,8 @@ async fn google_cache_uses_payload_expiry_without_requiring_a_jwt_header() { .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()) + let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); for _ in 0..2 { assert_eq!( diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index bed762adc59..37557c36fb5 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -12,7 +12,7 @@ fn resolver(value: Option<&str>) -> SecretResolver { SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), Arc::new(move |_: &str| value.clone()), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ) } @@ -35,7 +35,7 @@ async fn native_reads_preserve_strings_and_report_conversion_errors(#[case] mana let resolver = SecretResolver::new( Arc::new(state), Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver @@ -75,7 +75,7 @@ async fn native_defaults_apply_to_absence_but_never_hide_provider_failures() { KeyManagementSettings::default(), )), Arc::new(|_: &str| None), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); let result = resolver .get_secret_str("key", Some(SecretValue::new("default"))) @@ -189,7 +189,7 @@ fn managed(reply: Result, ()>, environment: Option<&'static str>) KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback) } @@ -281,7 +281,7 @@ async fn prefix_is_removed_once_and_resolved_from_environment() { let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver @@ -340,7 +340,7 @@ async fn excluded_hosted_keys_keep_the_python_manager_conversion_path( }, )), Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver.get_secret("KEY", None).await.unwrap(), @@ -373,7 +373,7 @@ async fn azure_callback_absence_preserves_none_but_errors_fall_back( KeyManagementSettings::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/source.rs b/litellm-rust/crates/secrets/tests/source.rs index b4782c6af86..17210940a67 100644 --- a/litellm-rust/crates/secrets/tests/source.rs +++ b/litellm-rust/crates/secrets/tests/source.rs @@ -15,7 +15,7 @@ mod tests { #[case] expected: Option<&str>, ) { unsafe { std::env::set_var(name, value) }; - let secret = EnvironmentSecrets::python_compatible() + let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::plain_for_test()) .resolve(&[name]) .await .unwrap() @@ -42,7 +42,7 @@ async fn dynamic_names_use_the_same_resolver_and_snapshots_never_do_fresh_lookup reads.fetch_add(1, Ordering::SeqCst); (name != "missing").then(|| name.to_owned()) }), - OidcResolver::default(), + OidcResolver::new(litellm_http::Client::plain_for_test()), ); let snapshot = source.resolve(&["declared", "missing"]).await.unwrap(); let name = format!("runtime-{}", "key"); diff --git a/litellm-rust/crates/testkit/src/lib.rs b/litellm-rust/crates/testkit/src/lib.rs index 9ea6123a176..423adeec426 100644 --- a/litellm-rust/crates/testkit/src/lib.rs +++ b/litellm-rust/crates/testkit/src/lib.rs @@ -1,3 +1,9 @@ +#![allow( + clippy::disallowed_types, + clippy::disallowed_methods, + reason = "a dev-only installer tool that never talks to providers" +)] + mod agent; mod error; mod install; From 7fc22061715487220826135298231fc723e8993a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 19:10:20 -0700 Subject: [PATCH 40/58] test: fix stale and state-leaking tests red on scheduled CircleCI (#43266) test_update_config_success_callback_normalization replaced proxy_server.proxy_logging_obj with a MagicMock and never restored it. Since the proxy unit tests joined tests/unit (#42903), 14 JWT mapping, end-user and MCP tests on the same xdist worker awaited that mock and failed. The test now uses monkeypatch. test_prometheus_logging_callbacks set verbose_logger to DEBUG and litellm.set_verbose at import, so every worker in the unit job ran with DEBUG on. That broke caplog equality in the JEV classifier test, the vertex streaming memory ratio, and four event-loop lag checks. The module-level setup is removed; nothing in the file depended on it. #43081 removed the OCR harness modules but left them in the importability parametrize list. test_get_model_info_bedrock_region reassigned litellm.model_cost and set LITELLM_LOCAL_MODEL_COST_MAP without restoring either, and never cleared the get_model_info caches, so it failed whenever an earlier test had looked up the regional model. It now uses monkeypatch and invalidates the caches; the local_testing isolation fixture also invalidates them after restoring model_cost. The Windows job hit CircleCI's 10 minute no-output limit while cargo compiles the Rust crates inside uv sync and uv build. Those two steps now allow 30 minutes of silence. --- .circleci/config.yml | 2 ++ tests/local_testing/conftest.py | 2 ++ tests/local_testing/test_get_model_info.py | 17 +++++++++-------- tests/test_rust_python_harness.py | 2 -- .../test_prometheus_logging_callbacks.py | 6 ------ tests/unit/proxy/test_proxy_server.py | 8 ++++---- 6 files changed, 17 insertions(+), 20 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 370424dca86..3ff8061fb48 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -323,6 +323,7 @@ jobs: CHOCOLATEY_CONFIRM_ALL: "true" - run: name: Install Dependencies + no_output_timeout: 30m environment: UV_HTTP_TIMEOUT: "300" command: | @@ -381,6 +382,7 @@ jobs: uv run --no-sync python -m pytest tests/windows_tests/ -v - run: name: Guard against MAX_PATH-busting packaged wheel paths + no_output_timeout: 30m environment: UV_HTTP_TIMEOUT: "300" command: | diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index d03f074f557..df3dacac3b2 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -19,6 +19,7 @@ import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.utils import _invalidate_model_cost_lowercase_map # ``litellm.model_cost`` is loaded at import time from the URL pinned to ``main`` # (``LITELLM_MODEL_COST_MAP_URL``). The in-tree backup ships with this branch @@ -232,6 +233,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) + _invalidate_model_cost_lowercase_map() @pytest.fixture(scope="module", autouse=True) diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 1e46a1bf853..dbe1fc3b69b 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm import get_model_info +from litellm.utils import _invalidate_model_cost_lowercase_map from unittest.mock import MagicMock, patch @@ -74,15 +75,15 @@ def test_get_model_info_ollama_chat(): assert mock_client.call_args.kwargs["json"]["name"] == "unknown-model" -def test_get_model_info_bedrock_region(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - args = { - "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "custom_llm_provider": "bedrock", +def test_get_model_info_bedrock_region(monkeypatch): + regional_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + model_cost_without_regional_entry = { + key: value for key, value in litellm.get_model_cost_map(url="").items() if key != regional_model } - litellm.model_cost.pop("us.anthropic.claude-haiku-4-5-20251001-v1:0", None) - info = litellm.get_model_info(**args) + monkeypatch.setattr(litellm, "model_cost", model_cost_without_regional_entry) + _invalidate_model_cost_lowercase_map() + info = litellm.get_model_info(model=regional_model, custom_llm_provider="bedrock") print("info", info) assert info["key"] == "anthropic.claude-haiku-4-5-20251001-v1:0" assert info["litellm_provider"] == "bedrock_converse" diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index a1bb370a074..9af38941684 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -36,8 +36,6 @@ def _case(module: str = "tests.example") -> HarnessCase: @pytest.mark.parametrize( "module", [ - "tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", - "tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case", "tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", "tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case", "tests.rust-python-harness.strategies.trace_parity.sdk.transcription.case", diff --git a/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 92ff3d5813c..f1c80bb11ea 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -1,7 +1,6 @@ import asyncio -import logging from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, call, patch @@ -9,7 +8,6 @@ import pytest from prometheus_client import REGISTRY import litellm -from litellm._logging import verbose_logger from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -27,10 +25,6 @@ except Exception: PrometheusLogger = None from litellm.proxy._types import UserAPIKeyAuth -verbose_logger.setLevel(logging.DEBUG) - -litellm.set_verbose = True - @pytest.fixture def prometheus_logger() -> PrometheusLogger: diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index eae80f311d8..65b368ca9e3 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -2979,7 +2979,7 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): @pytest.mark.asyncio -async def test_update_config_success_callback_normalization(): +async def test_update_config_success_callback_normalization(monkeypatch): """ Ensure success_callback values are normalized to lowercase when updating config. This prevents delete_callback (which searches lowercase) from failing on mixed case inputs like 'SQS'. @@ -2987,7 +2987,7 @@ async def test_update_config_success_callback_normalization(): import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import ConfigYAML - setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) existing_litellm_settings = {"success_callback": ["langfuse"]} @@ -3013,7 +3013,7 @@ async def test_update_config_success_callback_normalization(): self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first) self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert) - setattr(proxy_server, "prisma_client", MockPrisma()) + monkeypatch.setattr(proxy_server, "prisma_client", MockPrisma()) class MockProxyConfig: async def add_deployment(self, prisma_client=None, proxy_logging_obj=None): # noqa: F811 # pytest fixture, not a redefinition @@ -3022,7 +3022,7 @@ async def test_update_config_success_callback_normalization(): def reject_config_owned_writes(self, *, section_name, changed_keys): return None - setattr(proxy_server, "proxy_config", MockProxyConfig()) + monkeypatch.setattr(proxy_server, "proxy_config", MockProxyConfig()) config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]}) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth From 4179860a17ec1b054db1a1e6e1d1402c4693089e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 19:12:36 -0700 Subject: [PATCH 41/58] fix(cost-map): retirement dates, chatgpt reasoning flags, bing pricing, bedrock mantle and mythos, azure gpt-5.6 alias, anthropic batch rates, new nebius, openrouter and xai rows (#42951) --- .../crates/model-catalog/src/model_info.rs | 44 ++ litellm/litellm_core_utils/litellm_logging.py | 1 + ...odel_prices_and_context_window_backup.json | 493 +++++++++++++++--- litellm/types/utils.py | 2 + litellm/utils.py | 3 + model_prices_and_context_window.json | 493 +++++++++++++++--- model_prices_and_context_window.schema.json | 5 + tests/local_testing/test_get_model_info.py | 27 + .../test_bing_grounding_search.py | 3 +- .../llm_cost_calc/test_llm_cost_calc_utils.py | 12 + .../test_litellm_logging.py | 23 + tests/unit/test_model_prices_schema.py | 55 ++ tests/unit/test_utils.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 14 files changed, 1029 insertions(+), 137 deletions(-) diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 361cb56e9b1..9e9a4220e50 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -42,6 +42,9 @@ pub struct ModelInfo { pub cache_creation_input_token_cost_above_200k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. #[serde(skip_serializing_if = "Option::is_none")] + pub cache_creation_input_token_cost_above_200k_tokens_batches: Option, + /// Rate applied once the prompt exceeds the token threshold in the field name. + #[serde(skip_serializing_if = "Option::is_none")] pub cache_creation_input_token_cost_above_256k_tokens: Option, /// Rate applied once the prompt exceeds the token threshold in the field name. #[serde(skip_serializing_if = "Option::is_none")] @@ -78,6 +81,9 @@ pub struct ModelInfo { /// Rate applied once the prompt exceeds the token threshold in the field name. #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_200k_tokens: Option, + /// Rate applied once the prompt exceeds the token threshold in the field name. + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_input_token_cost_above_200k_tokens_batches: Option, /// Priority service-tier rate for the same-named base field. #[serde(skip_serializing_if = "Option::is_none")] pub cache_read_input_token_cost_above_200k_tokens_priority: Option, @@ -113,6 +119,10 @@ pub struct ModelInfo { pub code_interpreter_cost_per_session: Option, #[serde(skip_serializing_if = "Option::is_none")] pub comment: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub computer_use_input_cost_per_1k_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub computer_use_output_cost_per_1k_tokens: Option, /// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'. #[serde(skip_serializing_if = "Option::is_none")] pub default_reasoning_effort: Option, @@ -120,6 +130,10 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub deprecation_date: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub file_search_cost_per_1k_calls: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub file_search_cost_per_gb_per_day: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub gemini_audio_only_live: Option, #[serde(skip_serializing_if = "Option::is_none")] pub gemini_native_audio: Option, @@ -174,6 +188,9 @@ pub struct ModelInfo { /// Rate applied once the prompt exceeds the token threshold in the field name. #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_200k_tokens: Option, + /// Rate applied once the prompt exceeds the token threshold in the field name. + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_token_above_200k_tokens_batches: Option, /// Priority service-tier rate for the same-named base field. #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_token_above_200k_tokens_priority: Option, @@ -265,6 +282,26 @@ pub struct ModelInfo { pub output_cost_per_image_1536: Option, #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_512: Option, + #[serde( + rename = "output_cost_per_image_0.5K", + skip_serializing_if = "Option::is_none" + )] + pub output_cost_per_image_0_5k: Option, + #[serde( + rename = "output_cost_per_image_1K", + skip_serializing_if = "Option::is_none" + )] + pub output_cost_per_image_1k: Option, + #[serde( + rename = "output_cost_per_image_2K", + skip_serializing_if = "Option::is_none" + )] + pub output_cost_per_image_2k: Option, + #[serde( + rename = "output_cost_per_image_4K", + skip_serializing_if = "Option::is_none" + )] + pub output_cost_per_image_4k: Option, #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_token: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -297,6 +334,9 @@ pub struct ModelInfo { /// Rate applied once the prompt exceeds the token threshold in the field name. #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_200k_tokens: Option, + /// Rate applied once the prompt exceeds the token threshold in the field name. + #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_token_above_200k_tokens_batches: Option, /// Priority service-tier rate for the same-named base field. #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_token_above_200k_tokens_priority: Option, @@ -357,6 +397,8 @@ pub struct ModelInfo { /// Provider default requests-per-minute limit. #[serde(skip_serializing_if = "Option::is_none")] pub rpm: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub rules: Option>, /// USD cost per web search query, keyed by search context size. #[serde(skip_serializing_if = "Option::is_none")] pub search_context_cost_per_query: Option, @@ -475,6 +517,8 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub uses_embed_content: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub vector_store_cost_per_gb_per_day: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub vertex_ai_audio_api: Option, /// Whether web search is billed per query or per prompt. #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 28d72702f3e..e955c0157c6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -384,6 +384,7 @@ _DEPLOYMENT_PRICING_KEYS: Final = ( "cache_read_input_token_cost_above_200k_tokens_batches", "cache_read_input_token_cost_above_272k_tokens_batches", "cache_creation_input_token_cost_batches", + "cache_creation_input_token_cost_above_200k_tokens_batches", "cache_creation_input_token_cost_above_272k_tokens_batches", "ocr_cost_per_page", "ocr_cost_per_page_batches", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 56129fef135..ec31039fee0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1279,8 +1279,11 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-mythos-preview": { - "input_cost_per_token": 0, - "output_cost_per_token": 0, + "cache_creation_input_token_cost": 3.4375e-05, + "cache_creation_input_token_cost_above_1hr": 5.5e-05, + "cache_read_input_token_cost": 2.75e-06, + "input_cost_per_token": 2.75e-05, + "output_cost_per_token": 0.0001375, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -1289,10 +1292,11 @@ "thinking_always_on": true, "supports_function_calling": true, "supports_vision": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_output_config": true + "supports_output_config": true, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -7773,27 +7777,27 @@ "output_cost_per_token_above_272k_tokens_batches": 0.000135 }, "azure/gpt-5.6": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.5e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_priority": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_priority": 1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_priority": 6e-05, - "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8579,27 +8583,27 @@ "supports_web_search": true }, "azure/us/gpt-5.6": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.75e-05, - "cache_creation_input_token_cost_priority": 1.375e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2.2e-06, - "cache_read_input_token_cost_priority": 1.1e-06, - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2.2e-05, - "input_cost_per_token_priority": 1.1e-05, + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2.2e-05, + "cache_creation_input_token_cost_priority": 1.1e-05, + "cache_read_input_token_cost": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.76e-06, + "cache_read_input_token_cost_priority": 8.8e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_above_272k_tokens": 8.8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.76e-05, + "input_cost_per_token_priority": 8.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, - "output_cost_per_token_above_272k_tokens_priority": 9.9e-05, - "output_cost_per_token_priority": 6.6e-05, + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_above_272k_tokens": 3.3e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, + "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8985,27 +8989,27 @@ "supports_web_search": true }, "azure/eu/gpt-5.6": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.75e-05, - "cache_creation_input_token_cost_priority": 1.375e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2.2e-06, - "cache_read_input_token_cost_priority": 1.1e-06, - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2.2e-05, - "input_cost_per_token_priority": 1.1e-05, + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2.2e-05, + "cache_creation_input_token_cost_priority": 1.1e-05, + "cache_read_input_token_cost": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.76e-06, + "cache_read_input_token_cost_priority": 8.8e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_above_272k_tokens": 8.8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.76e-05, + "input_cost_per_token_priority": 8.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, - "output_cost_per_token_above_272k_tokens_priority": 9.9e-05, - "output_cost_per_token_priority": 6.6e-05, + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_above_272k_tokens": 3.3e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, + "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14640,14 +14644,18 @@ "claude-haiku-4-5-20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14663,14 +14671,18 @@ "claude-haiku-4-5": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14693,13 +14705,21 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_200k_tokens_batches": 3.75e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens_batches": 1.125e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14727,13 +14747,21 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_200k_tokens_batches": 3.75e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens_batches": 1.125e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14757,14 +14785,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14797,14 +14829,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14865,14 +14901,18 @@ "claude-opus-4-5-20251101": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14895,14 +14935,18 @@ "claude-opus-4-5": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14927,14 +14971,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14966,14 +15014,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15004,14 +15056,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15044,14 +15100,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15083,14 +15143,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15123,14 +15187,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15164,14 +15232,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15207,14 +15279,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15250,14 +15326,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15756,6 +15836,16 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "c4ai-aya-expanse-32b": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "source": "https://docs.cohere.com/docs/models" + }, "command-a-plus-05-2026": { "input_cost_per_token": 0.0, "litellm_provider": "cohere_chat", @@ -22145,11 +22235,11 @@ } }, "bing_grounding/search": { - "input_cost_per_query": 0.035, + "input_cost_per_query": 0.014, "litellm_provider": "bing_grounding", "mode": "search", "metadata": { - "notes": "Grounding with Bing Search (G1 SKU): $35 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." + "notes": "Grounding with Bing Search (G1 SKU): $14 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." } }, "tinyfish/search": { @@ -28036,6 +28126,7 @@ "tpm": 10000000 }, "gemini/gemini-3-pro-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28083,6 +28174,7 @@ "supports_reasoning": false }, "gemini/gemini-3.1-flash-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", @@ -28130,6 +28222,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-05-25", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -30551,7 +30644,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.6-luna": { "litellm_provider": "chatgpt", @@ -30567,7 +30664,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.6-sol": { "litellm_provider": "chatgpt", @@ -30583,7 +30684,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.6-terra": { "litellm_provider": "chatgpt", @@ -30599,7 +30704,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", @@ -30614,7 +30723,12 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "default_reasoning_effort": "none", + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.4-pro": { "litellm_provider": "chatgpt", @@ -30628,7 +30742,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.3-codex": { "litellm_provider": "chatgpt", @@ -30642,7 +30760,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": false, + "supports_minimal_reasoning_effort": true }, "chatgpt/gpt-5.3-codex-spark": { "litellm_provider": "chatgpt", @@ -30686,7 +30808,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.2-codex": { "litellm_provider": "chatgpt", @@ -30700,7 +30823,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.2": { "litellm_provider": "chatgpt", @@ -30715,7 +30839,12 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "default_reasoning_effort": "none", + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.1-codex-max": { "litellm_provider": "chatgpt", @@ -30729,7 +30858,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.1-codex-mini": { "litellm_provider": "chatgpt", @@ -30743,7 +30873,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "gigachat/GigaChat-2": { "input_cost_per_token": 0.0, @@ -39214,6 +39345,17 @@ "supports_function_calling": true, "supports_reasoning": true }, + "nebius/deepseek-ai/DeepSeek-V4.1-Flash": { + "input_cost_per_token": 3e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://tokenfactory.nebius.com/endpoints?modals=endpoint-details&model-id=deepseek-ai/DeepSeek-V4.1-Flash", + "supports_vision": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -56230,7 +56372,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56417,7 +56559,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -59653,14 +59795,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "prompt_cache_min_tokens": 512, "search_context_cost_per_query": { "search_context_size_high": 0.01, @@ -59693,14 +59839,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -63000,6 +63150,22 @@ "image" ] }, + "xai/grok-imagine-image-pro": { + "input_cost_per_image": 0.05, + "litellm_provider": "xai", + "mode": "image_generation", + "source": "https://docs.x.ai/docs/models", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" + ] + }, "xai/grok-imagine-image-2.0": { "input_cost_per_image": 0.06, "litellm_provider": "xai", @@ -64006,6 +64172,55 @@ "supports_response_schema": true, "supports_vision": true }, + "azure_ai/deepseek-r1": { + "input_cost_per_token": 1.35e-06, + "output_cost_per_token": 5.4e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/deepseek-v3-0324": { + "input_cost_per_token": 1.14e-06, + "output_cost_per_token": 4.56e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/deepseek-v3.1": { + "input_cost_per_token": 1.23e-06, + "output_cost_per_token": 4.94e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-3": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-3-mini": { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.27e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-4-fast-non-reasoning": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-4-fast-reasoning": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -64834,6 +65049,131 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/deepseek.v3.1": { + "input_cost_per_token": 5.8e-07, + "output_cost_per_token": 1.68e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-1.html" + }, + "bedrock_mantle/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html" + }, + "bedrock_mantle/qwen.qwen3-235b-a22b-2507": { + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 8.8e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-235b-a22b-2507.html" + }, + "bedrock_mantle/qwen.qwen3-32b": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 32000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html" + }, + "bedrock_mantle/qwen.qwen3-coder-30b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html" + }, + "bedrock_mantle/qwen.qwen3-coder-480b-a35b-instruct": { + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-480b-a35b-instruct.html" + }, + "bedrock_mantle/qwen.qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-next-80b-a3b.html" + }, + "bedrock_mantle/qwen.qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 5.3e-07, + "output_cost_per_token": 2.66e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-vl-235b-a22b.html" + }, "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, "default_reasoning_effort": "none", @@ -76578,6 +76918,23 @@ "supports_vision": true, "supports_web_search": false }, + "openrouter/perceptron/perceptron-mk1.5": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 36864, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "vertex_ai/gemini-2.0-flash": { "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 1e-06, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd336c9b989..2e518af4da4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -302,6 +302,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_batches: ReadOnly[float | None] + cache_creation_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. @@ -3735,6 +3736,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None cache_creation_input_token_cost_batches: float | None = None + cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None cache_read_input_audio_token_cost: float | None = None cache_read_input_image_token_cost: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index be4388802f9..42da2e2a7b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6167,6 +6167,9 @@ def _get_model_info_helper( "cache_read_input_token_cost_above_272k_tokens_batches" ), cache_creation_input_token_cost_batches=_model_info.get("cache_creation_input_token_cost_batches"), + cache_creation_input_token_cost_above_200k_tokens_batches=_model_info.get( + "cache_creation_input_token_cost_above_200k_tokens_batches" + ), cache_creation_input_token_cost_above_272k_tokens_batches=_model_info.get( "cache_creation_input_token_cost_above_272k_tokens_batches" ), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 56129fef135..ec31039fee0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1279,8 +1279,11 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-mythos-preview": { - "input_cost_per_token": 0, - "output_cost_per_token": 0, + "cache_creation_input_token_cost": 3.4375e-05, + "cache_creation_input_token_cost_above_1hr": 5.5e-05, + "cache_read_input_token_cost": 2.75e-06, + "input_cost_per_token": 2.75e-05, + "output_cost_per_token": 0.0001375, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -1289,10 +1292,11 @@ "thinking_always_on": true, "supports_function_calling": true, "supports_vision": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_output_config": true + "supports_output_config": true, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, @@ -7773,27 +7777,27 @@ "output_cost_per_token_above_272k_tokens_batches": 0.000135 }, "azure/gpt-5.6": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.5e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_priority": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_priority": 1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_priority": 6e-05, - "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8579,27 +8583,27 @@ "supports_web_search": true }, "azure/us/gpt-5.6": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.75e-05, - "cache_creation_input_token_cost_priority": 1.375e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2.2e-06, - "cache_read_input_token_cost_priority": 1.1e-06, - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2.2e-05, - "input_cost_per_token_priority": 1.1e-05, + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2.2e-05, + "cache_creation_input_token_cost_priority": 1.1e-05, + "cache_read_input_token_cost": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.76e-06, + "cache_read_input_token_cost_priority": 8.8e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_above_272k_tokens": 8.8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.76e-05, + "input_cost_per_token_priority": 8.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, - "output_cost_per_token_above_272k_tokens_priority": 9.9e-05, - "output_cost_per_token_priority": 6.6e-05, + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_above_272k_tokens": 3.3e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, + "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8985,27 +8989,27 @@ "supports_web_search": true }, "azure/eu/gpt-5.6": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, - "cache_creation_input_token_cost_above_272k_tokens_priority": 2.75e-05, - "cache_creation_input_token_cost_priority": 1.375e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_above_272k_tokens_priority": 2.2e-06, - "cache_read_input_token_cost_priority": 1.1e-06, - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "input_cost_per_token_above_272k_tokens_priority": 2.2e-05, - "input_cost_per_token_priority": 1.1e-05, + "cache_creation_input_token_cost": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2.2e-05, + "cache_creation_input_token_cost_priority": 1.1e-05, + "cache_read_input_token_cost": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.76e-06, + "cache_read_input_token_cost_priority": 8.8e-07, + "input_cost_per_token": 4.4e-06, + "input_cost_per_token_above_272k_tokens": 8.8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.76e-05, + "input_cost_per_token_priority": 8.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, - "output_cost_per_token_above_272k_tokens_priority": 9.9e-05, - "output_cost_per_token_priority": 6.6e-05, + "output_cost_per_token": 2.2e-05, + "output_cost_per_token_above_272k_tokens": 3.3e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, + "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14640,14 +14644,18 @@ "claude-haiku-4-5-20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14663,14 +14671,18 @@ "claude-haiku-4-5": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14693,13 +14705,21 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_200k_tokens_batches": 3.75e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens_batches": 1.125e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14727,13 +14747,21 @@ "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_200k_tokens_batches": 3.75e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens_batches": 1.125e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14757,14 +14785,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14797,14 +14829,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14865,14 +14901,18 @@ "claude-opus-4-5-20251101": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14895,14 +14935,18 @@ "claude-opus-4-5": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14927,14 +14971,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -14966,14 +15014,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15004,14 +15056,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15044,14 +15100,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15083,14 +15143,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15123,14 +15187,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15164,14 +15232,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15207,14 +15279,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15250,14 +15326,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -15756,6 +15836,16 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "c4ai-aya-expanse-32b": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 4000, + "max_tokens": 4000, + "mode": "chat", + "source": "https://docs.cohere.com/docs/models" + }, "command-a-plus-05-2026": { "input_cost_per_token": 0.0, "litellm_provider": "cohere_chat", @@ -22145,11 +22235,11 @@ } }, "bing_grounding/search": { - "input_cost_per_query": 0.035, + "input_cost_per_query": 0.014, "litellm_provider": "bing_grounding", "mode": "search", "metadata": { - "notes": "Grounding with Bing Search (G1 SKU): $35 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." + "notes": "Grounding with Bing Search (G1 SKU): $14 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." } }, "tinyfish/search": { @@ -28036,6 +28126,7 @@ "tpm": 10000000 }, "gemini/gemini-3-pro-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28083,6 +28174,7 @@ "supports_reasoning": false }, "gemini/gemini-3.1-flash-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", @@ -28130,6 +28222,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-05-25", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -30551,7 +30644,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.6-luna": { "litellm_provider": "chatgpt", @@ -30567,7 +30664,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.6-sol": { "litellm_provider": "chatgpt", @@ -30583,7 +30684,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.6-terra": { "litellm_provider": "chatgpt", @@ -30599,7 +30704,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", @@ -30614,7 +30723,12 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "default_reasoning_effort": "none", + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.4-pro": { "litellm_provider": "chatgpt", @@ -30628,7 +30742,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.3-codex": { "litellm_provider": "chatgpt", @@ -30642,7 +30760,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": false, + "supports_minimal_reasoning_effort": true }, "chatgpt/gpt-5.3-codex-spark": { "litellm_provider": "chatgpt", @@ -30686,7 +30808,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.2-codex": { "litellm_provider": "chatgpt", @@ -30700,7 +30823,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.2": { "litellm_provider": "chatgpt", @@ -30715,7 +30839,12 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": true, + "default_reasoning_effort": "none", + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false }, "chatgpt/gpt-5.1-codex-max": { "litellm_provider": "chatgpt", @@ -30729,7 +30858,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "chatgpt/gpt-5.1-codex-mini": { "litellm_provider": "chatgpt", @@ -30743,7 +30873,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true }, "gigachat/GigaChat-2": { "input_cost_per_token": 0.0, @@ -39214,6 +39345,17 @@ "supports_function_calling": true, "supports_reasoning": true }, + "nebius/deepseek-ai/DeepSeek-V4.1-Flash": { + "input_cost_per_token": 3e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://tokenfactory.nebius.com/endpoints?modals=endpoint-details&model-id=deepseek-ai/DeepSeek-V4.1-Flash", + "supports_vision": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -56230,7 +56372,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56417,7 +56559,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -59653,14 +59795,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "prompt_cache_min_tokens": 512, "search_context_cost_per_query": { "search_context_size_high": 0.01, @@ -59693,14 +59839,18 @@ "supports_anthropic_compaction": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -63000,6 +63150,22 @@ "image" ] }, + "xai/grok-imagine-image-pro": { + "input_cost_per_image": 0.05, + "litellm_provider": "xai", + "mode": "image_generation", + "source": "https://docs.x.ai/docs/models", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" + ] + }, "xai/grok-imagine-image-2.0": { "input_cost_per_image": 0.06, "litellm_provider": "xai", @@ -64006,6 +64172,55 @@ "supports_response_schema": true, "supports_vision": true }, + "azure_ai/deepseek-r1": { + "input_cost_per_token": 1.35e-06, + "output_cost_per_token": 5.4e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/deepseek-v3-0324": { + "input_cost_per_token": 1.14e-06, + "output_cost_per_token": 4.56e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/deepseek-v3.1": { + "input_cost_per_token": 1.23e-06, + "output_cost_per_token": 4.94e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-3": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-3-mini": { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.27e-06, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-4-fast-non-reasoning": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/grok-4-fast-reasoning": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -64834,6 +65049,131 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/deepseek.v3.1": { + "input_cost_per_token": 5.8e-07, + "output_cost_per_token": 1.68e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-1.html" + }, + "bedrock_mantle/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html" + }, + "bedrock_mantle/qwen.qwen3-235b-a22b-2507": { + "input_cost_per_token": 2.2e-07, + "output_cost_per_token": 8.8e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-235b-a22b-2507.html" + }, + "bedrock_mantle/qwen.qwen3-32b": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 32000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html" + }, + "bedrock_mantle/qwen.qwen3-coder-30b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html" + }, + "bedrock_mantle/qwen.qwen3-coder-480b-a35b-instruct": { + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-480b-a35b-instruct.html" + }, + "bedrock_mantle/qwen.qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-next-80b-a3b.html" + }, + "bedrock_mantle/qwen.qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 5.3e-07, + "output_cost_per_token": 2.66e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-vl-235b-a22b.html" + }, "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, "default_reasoning_effort": "none", @@ -76578,6 +76918,23 @@ "supports_vision": true, "supports_web_search": false }, + "openrouter/perceptron/perceptron-mk1.5": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 36864, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "vertex_ai/gemini-2.0-flash": { "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 1e-06, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index fa1828c780a..c4048cac905 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -103,6 +103,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_creation_input_token_cost_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "cache_creation_input_token_cost_above_256k_tokens": { "type": "number", "minimum": 0, diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index dbe1fc3b69b..79f6739a423 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -320,6 +320,33 @@ def test_get_model_info_bedrock_cross_region_capability_parity(): assert checked > 0, "no cross-region bedrock profiles found - the filter is inert" + +def test_get_model_info_bedrock_priced_cross_region_profile_has_priced_base(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prefixes = ("us.", "eu.", "apac.", "us-gov.", "au.", "global.") + checked = 0 + + for k, v in litellm.model_cost.items(): + if not str(v.get("litellm_provider", "")).startswith("bedrock"): + continue + base_model_key = next( + (k[len(p) :] for p in prefixes if k.startswith(p)), + None, + ) + if base_model_key is None or base_model_key not in litellm.model_cost: + continue + checked += 1 + base = litellm.model_cost[base_model_key] + for cost_key in ("input_cost_per_token", "output_cost_per_token"): + if (v.get(cost_key) or 0) > 0: + assert ( + base.get(cost_key) or 0 + ) > 0, f"{k} charges {cost_key} but its base {base_model_key} is free" + + assert checked > 0, "no cross-region bedrock profiles found - the filter is inert" + def test_get_model_info_huggingface_models(monkeypatch): from litellm import Router from litellm.types.router import ModelGroupInfo diff --git a/tests/search_tests/test_bing_grounding_search.py b/tests/search_tests/test_bing_grounding_search.py index f532158e462..6ea79076370 100644 --- a/tests/search_tests/test_bing_grounding_search.py +++ b/tests/search_tests/test_bing_grounding_search.py @@ -196,4 +196,5 @@ class TestBingGroundingSearchTransformation: ): response = litellm.search(query="pricing check", search_provider="bing_grounding") - assert response._hidden_params["response_cost"] == pytest.approx(0.035) + # Grounding with Bing Search (G1 SKU): $14 per 1,000 transactions, https://www.microsoft.com/en-us/bing/apis, checked 2026-09-24 + assert response._hidden_params["response_cost"] == pytest.approx(0.014) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 781a3a7c4ed..0afd989272e 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -3789,3 +3789,15 @@ def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base): assert azure_ai_info[field] == base assert azure_us_info[field] == pytest.approx(1.1 * base) assert azure_eu_info[field] == pytest.approx(1.2 * base) + + +@pytest.mark.parametrize("region_prefix", ["azure/", "azure/us/", "azure/eu/"]) +def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_prefix): + """The bare gpt-5.6 alias routes to GPT-5.6 Sol, so every Azure region must bill the + alias exactly like the Sol entry (including the Sept 2026 $4/$20 promo).""" + alias = litellm.model_cost[f"{region_prefix}gpt-5.6"] + sol = litellm.model_cost[f"{region_prefix}gpt-5.6-sol"] + shared_cost_fields = [f for f in alias if "cost" in f and f in sol and not isinstance(alias[f], dict)] + assert shared_cost_fields + for field in shared_cost_fields: + assert alias[field] == sol[field], field diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d717718cba2..cb1e281e356 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7969,6 +7969,29 @@ def test_deployment_pricing_model_info_honors_a_tier_only_batch_override_over_th assert {key: info[key] for key in carried_keys} == {key: _PUBLISHED_BATCH_RATES[key] for key in carried_keys} +@pytest.mark.parametrize( + "override_key", + ( + "output_cost_per_token_above_200k_tokens_batches", + "cache_read_input_token_cost_above_200k_tokens_batches", + "cache_creation_input_token_cost_above_200k_tokens_batches", + ), +) +def test_deployment_pricing_model_info_honors_a_200k_tier_batch_override( + _published_batch_model: None, override_key: str +) -> None: + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info + + info: Final = deployment_pricing_model_info(_batch_deployment_id({override_key: 1e-3}), _PUBLISHED_BATCH_DEPLOYMENT) + carried_keys: Final = tuple( + key for key in (*_PUBLISHED_INPUT_BATCH_KEYS, *_PUBLISHED_OUTPUT_BATCH_KEYS) if key != override_key + ) + + assert info is not None + assert info[override_key] == 1e-3 + assert {key: info[key] for key in carried_keys} == {key: _PUBLISHED_BATCH_RATES[key] for key in carried_keys} + + def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervened(): """LIT-6894: a non-blocking flagged verdict must outrank success in the request-level guardrail_status but never mask an intervention.""" diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index 052278631e2..a05c345b5cd 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -266,6 +266,61 @@ def test_openai_reasoning_family_entries_carry_supports_reasoning(prices: dict): ) +_ABSENT: Final = object() + +REASONING_ANNOTATION_KEYS: Final = ( + "supports_reasoning", + "supports_minimal_reasoning_effort", + "supports_none_reasoning_effort", + "supports_xhigh_reasoning_effort", + "default_reasoning_effort", +) + + +def chatgpt_openai_twins(prices: dict) -> list[tuple[str, str]]: + """`chatgpt/` rows paired with the bare `` row served by the openai provider. + + Scoped to openai twins on purpose. `ChatGPTConfig` and `ChatGPTResponsesAPIConfig` subclass + their openai counterparts, so a chatgpt row's reasoning behaviour is whatever the openai row + describes. The azure rows are a separate registry that already diverges from openai here, and + pinning them to each other would assert something this repository does not control. + """ + pairs = [] + for name, entry in prices.items(): + if not isinstance(entry, dict) or not name.startswith("chatgpt/"): + continue + bare = name.split("/", 1)[1] + twin = prices.get(bare) + if isinstance(twin, dict) and twin.get("litellm_provider") == "openai": + pairs.append((name, bare)) + return pairs + + +def test_chatgpt_rows_carry_their_openai_twin_reasoning_annotations(prices: dict): + """A chatgpt row must not silently drop the reasoning annotations of the model it proxies. + + `litellm.utils._get_model_info_from_generalization` refuses to fall back when an exact cost-map + key exists, so an unannotated `chatgpt/` row wins over its annotated twin and + `/model/info` reports the model as non-reasoning. + """ + twins = chatgpt_openai_twins(prices) + assert twins, "no chatgpt/* row has an openai twin any more; this guard has stopped guarding" + + mismatched = [] + for name, bare in twins: + for key in REASONING_ANNOTATION_KEYS: + if prices[name].get(key, _ABSENT) != prices[bare].get(key, _ABSENT): + mismatched.append( + f"{name}.{key} is {prices[name].get(key)!r}, {bare}.{key} is {prices[bare].get(key)!r}" + ) + + assert mismatched == [], ( + "chatgpt/* entries proxy their openai twin through ChatGPTConfig, so they must carry the " + "same reasoning annotations; an exact cost-map key blocks the generalization fallback, so " + "a missing flag here is reported to callers as 'not a reasoning model':\n" + "\n".join(mismatched) + ) + + def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict): """OpenAI rejects every reasoning.effort on chat-latest except medium, and a reasoning entry with no declared levels resolves to None, which lets /model_group/info and the dashboard offer diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0cdc52c9a93..5b327305e31 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -766,6 +766,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_batches": {"type": "number"}, "cache_creation_input_token_cost_flex": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1e0aed46923..6d46af04730 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32957,6 +32957,8 @@ export interface components { cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 200K Tokens */ cache_creation_input_token_cost_above_200k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 200K Tokens Batches */ + cache_creation_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens */ cache_creation_input_token_cost_above_272k_tokens?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens Batches */ @@ -46750,6 +46752,8 @@ export interface components { cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 200K Tokens */ cache_creation_input_token_cost_above_200k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 200K Tokens Batches */ + cache_creation_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens */ cache_creation_input_token_cost_above_272k_tokens?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens Batches */ From 2ef3250ec3bdd275dae3281d1e0f2347d995e5fe Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 02:16:32 +0000 Subject: [PATCH 42/58] refactor(rust): promote anthropic messages out of experimental_pass_through (#43269) Co-authored-by: Yujong Lee --- litellm-rust/crates/core/src/messages/common_utils.rs | 2 +- litellm-rust/crates/core/src/messages/prepare.rs | 2 +- .../crates/llms/src/anthropic/batches/transformation.rs | 2 +- litellm-rust/crates/llms/src/anthropic/chat/handler.rs | 2 +- .../crates/llms/src/anthropic/chat/transformation.rs | 4 +--- .../llms/src/anthropic/experimental_pass_through/mod.rs | 1 - .../{experimental_pass_through => }/messages/handler.rs | 0 .../{experimental_pass_through => }/messages/headers.rs | 0 .../{experimental_pass_through => }/messages/mod.rs | 0 .../messages/streaming_iterator.rs | 0 .../{experimental_pass_through => }/messages/thinking.rs | 0 .../messages/transformation.rs | 0 litellm-rust/crates/llms/src/anthropic/mod.rs | 5 +++-- .../llms/src/azure_ai/anthropic/messages_transformation.rs | 2 +- .../llms/src/base_llm/anthropic_messages/transformation.rs | 3 +-- 15 files changed, 10 insertions(+), 13 deletions(-) delete mode 100644 litellm-rust/crates/llms/src/anthropic/experimental_pass_through/mod.rs rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/handler.rs (100%) rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/headers.rs (100%) rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/mod.rs (100%) rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/streaming_iterator.rs (100%) rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/thinking.rs (100%) rename litellm-rust/crates/llms/src/anthropic/{experimental_pass_through => }/messages/transformation.rs (100%) diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index d27b79bdc04..95142e87519 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,7 +1,7 @@ use litellm_http::request::string_headers as shared_string_headers; pub(super) use litellm_http::request::truncate_error_body; use litellm_llms::{ - anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, + anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, }; diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index dc4b3562e3f..84884ab279e 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -5,7 +5,7 @@ use litellm_core_utils::{ settings::Lookup, }; use litellm_llms::{ - anthropic::experimental_pass_through::messages::handler::shape_anthropic_messages_request, + anthropic::messages::handler::shape_anthropic_messages_request, base_llm::anthropic_messages::transformation::{ BaseAnthropicMessagesConfig, MessagesTransformContext, }, diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index 94e4dc7838a..1c26684901a 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -5,7 +5,7 @@ use time::OffsetDateTime; use url::Url; use crate::{ - anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base, + anthropic::messages::transformation::resolve_anthropic_api_base, base_llm::chat::transformation::Error, }; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs index a80cfbf28bd..9160cdf28ee 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs @@ -7,7 +7,7 @@ use litellm_types::{ use serde_json::Value; use crate::{ - anthropic::experimental_pass_through::messages::streaming_iterator::{ + anthropic::messages::streaming_iterator::{ AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent, AnthropicStreamUsage, }, diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index fd86c5ca25a..07ed6ba6ed1 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -11,9 +11,7 @@ use serde_json::{Map, Value, json}; use crate::{ anthropic::{ ANTHROPIC_OAUTH_TOKEN_PREFIX, - experimental_pass_through::messages::transformation::{ - complete_anthropic_url, resolve_anthropic_api_key, - }, + messages::transformation::{complete_anthropic_url, resolve_anthropic_api_key}, }, base_llm::chat::transformation::{ BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/mod.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/mod.rs deleted file mode 100644 index ba63992f3cb..00000000000 --- a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod messages; diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs rename to litellm-rust/crates/llms/src/anthropic/messages/handler.rs diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs b/litellm-rust/crates/llms/src/anthropic/messages/headers.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs rename to litellm-rust/crates/llms/src/anthropic/messages/headers.rs diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs rename to litellm-rust/crates/llms/src/anthropic/messages/mod.rs diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs b/litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs rename to litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs rename to litellm-rust/crates/llms/src/anthropic/messages/thinking.rs diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs similarity index 100% rename from litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs rename to litellm-rust/crates/llms/src/anthropic/messages/transformation.rs diff --git a/litellm-rust/crates/llms/src/anthropic/mod.rs b/litellm-rust/crates/llms/src/anthropic/mod.rs index 755bc7d1907..a884c146dca 100644 --- a/litellm-rust/crates/llms/src/anthropic/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/mod.rs @@ -1,7 +1,8 @@ +pub mod common_utils; + pub mod batches; pub mod chat; -pub mod common_utils; pub mod count_tokens; -pub mod experimental_pass_through; +pub mod messages; pub const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat"; diff --git a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs index c409f7f687e..137239bbeaf 100644 --- a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs @@ -6,7 +6,7 @@ use litellm_types::llms::anthropic_messages::{ }; use crate::{ - anthropic::experimental_pass_through::messages::transformation::{ + anthropic::messages::transformation::{ ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }, base_llm::{ diff --git a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs index 8db14687214..eff1dd1cf0b 100644 --- a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs @@ -4,8 +4,7 @@ use litellm_types::llms::anthropic_messages::{ }; use crate::{ - anthropic::experimental_pass_through::messages::thinking::ThinkingContext, - base_llm::chat::transformation::Error, + anthropic::messages::thinking::ThinkingContext, base_llm::chat::transformation::Error, }; pub type Headers = Vec<(String, String)>; From 2530255624904e8e9868bc29dbf2f39f54fd9b70 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 25 Sep 2026 19:27:48 -0700 Subject: [PATCH 43/58] test: stop CI tests from downloading tokenizer files and images (#43257) * test: load the embedding base image from a committed 100x100 PNG instead of downloading it * test: move the volcengine embedding test into tests/unit * test: check gpt2 and r50k_base tokenizer parity against committed tiktoken reference files * test: check hub tokenizer selection against an in-memory Hugging Face hub * test: serve image URLs from respx in the gemini tool-result and format-param tests * ci: drop the emptied legacy core-utils test path * test: cover the cohere and anthropic tokenizer paths in the hub tokenizer test * test: fetch every format-param image through respx and check its bytes reach the request * test: drop the gpt2 and r50k_base parity tests, which no litellm path uses * test: drop comments that restate assertions in the format-param test --- .github/workflows/test-unit.yml | 2 +- .../base_embedding_unit_tests.py | 7 +- .../litellm_core_utils/__init__.py | 0 .../litellm_core_utils/test_token_counter.py | 51 ------ .../litellm_core_utils/test_tokenizer.py | 20 --- .../llms/vertex_ai/gemini/__init__.py | 0 .../test_vertex_ai_gemini_transformation.py | 54 ------ .../test_litellm/llms/volcengine/__init__.py | 1 - tests/test_litellm/test_main.py | 164 ------------------ .../litellm_core_utils/test_token_counter.py | 92 ++++++++++ .../unit/litellm_core_utils/test_tokenizer.py | 12 +- .../test_vertex_ai_gemini_transformation.py | 42 +++++ .../volcengine/test_volcengine_embedding.py | 0 tests/unit/test_main.py | 100 +++++++++++ tests/white_100x100.png | Bin 0 -> 214 bytes 15 files changed, 239 insertions(+), 306 deletions(-) delete mode 100644 tests/test_litellm/litellm_core_utils/__init__.py delete mode 100644 tests/test_litellm/litellm_core_utils/test_token_counter.py delete mode 100644 tests/test_litellm/litellm_core_utils/test_tokenizer.py delete mode 100644 tests/test_litellm/llms/vertex_ai/gemini/__init__.py delete mode 100644 tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py delete mode 100644 tests/test_litellm/llms/volcengine/__init__.py delete mode 100644 tests/test_litellm/test_main.py rename tests/{test_litellm => unit}/llms/volcengine/test_volcengine_embedding.py (100%) create mode 100644 tests/white_100x100.png diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index d75213d37ea..2212b276b0d 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -61,7 +61,7 @@ jobs: - shard: core-utils artifact-name: core-utils - test-path: "tests/test_litellm/litellm_core_utils" + test-path: "" unit-flag: core-utils workers: 2 reruns: 1 diff --git a/tests/llm_translation/base_embedding_unit_tests.py b/tests/llm_translation/base_embedding_unit_tests.py index 1a88f0e9d6b..469416fc0cf 100644 --- a/tests/llm_translation/base_embedding_unit_tests.py +++ b/tests/llm_translation/base_embedding_unit_tests.py @@ -16,15 +16,12 @@ from litellm.utils import ( get_optional_params, get_optional_params_embeddings, ) -import requests import base64 +from pathlib import Path -# test_example.py from abc import ABC, abstractmethod -url = "https://dummyimage.com/100/100/fff&text=Test+image" -response = requests.get(url) -file_data = response.content +file_data = (Path(__file__).parent.parent / "white_100x100.png").read_bytes() encoded_file = base64.b64encode(file_data).decode("utf-8") base64_image = f"data:image/png;base64,{encoded_file}" diff --git a/tests/test_litellm/litellm_core_utils/__init__.py b/tests/test_litellm/litellm_core_utils/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py deleted file mode 100644 index 1e10b7e82b1..00000000000 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ /dev/null @@ -1,51 +0,0 @@ -import pytest -from litellm import create_pretrained_tokenizer -from tests.unit.litellm_core_utils.test_token_counter import token_counter - - -def test_tokenizers(): - try: - ### test the openai, claude, cohere and llama2 tokenizers. - ### The tokenizer value should be different for all - sample_text = "Hellö World, this is my input string! My name is ishaan CTO" - - # openai tokenizer - openai_tokens = token_counter(model="gpt-3.5-turbo", text=sample_text) - - # claude tokenizer - claude_tokens = token_counter(model="claude-3-5-haiku-20241022", text=sample_text) - - # cohere tokenizer - cohere_tokens = token_counter(model="command-nightly", text=sample_text) - - # llama2 tokenizer - llama2_tokens = token_counter(model="meta-llama/Llama-2-7b-chat", text=sample_text) - - # llama3 tokenizer (also testing custom tokenizer) - llama3_tokens_1 = token_counter(model="meta-llama/llama-3-70b-instruct", text=sample_text) - - try: - llama3_tokenizer = create_pretrained_tokenizer("Xenova/llama-3-tokenizer") - except Exception as e: - pytest.skip(f"custom tokenizer download failed (HF hub unreachable): {e}") - llama3_tokens_2 = token_counter(custom_tokenizer=llama3_tokenizer, text=sample_text) - - print( - f"openai tokens: {openai_tokens}; claude tokens: {claude_tokens}; cohere tokens: {cohere_tokens}; llama2 tokens: {llama2_tokens}; llama3 tokens: {llama3_tokens_1}" - ) - - # assert that all token values are different - # llama2 may fall back to the tiktoken tokenizer when the HuggingFace - # model hub is unreachable (e.g. in CI). In that case the count will - # equal the openai count and the differentiation assertion is skipped. - if openai_tokens == llama2_tokens: - pytest.skip("llama2 fell back to tiktoken (HF hub unreachable); skipping differentiation assertion") - assert llama2_tokens != llama3_tokens_1, "Token values are not different." - - assert llama3_tokens_1 == llama3_tokens_2, ( - "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same." - ) - - print("test tokenizer: It worked!") - except Exception as e: - pytest.fail(f"An exception occured: {e}") diff --git a/tests/test_litellm/litellm_core_utils/test_tokenizer.py b/tests/test_litellm/litellm_core_utils/test_tokenizer.py deleted file mode 100644 index 2171044970c..00000000000 --- a/tests/test_litellm/litellm_core_utils/test_tokenizer.py +++ /dev/null @@ -1,20 +0,0 @@ -import pytest - -from tests.unit.litellm_core_utils.test_tokenizer import ( - UNICODE_TEXTS, - assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface, - assert_openai_encoding_matches_python, -) - -NETWORK_ENCODINGS = ("r50k_base", "gpt2") - - -@pytest.mark.parametrize("name", NETWORK_ENCODINGS) -@pytest.mark.parametrize("text", UNICODE_TEXTS) -def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: - assert_openai_encoding_matches_python(name, text) - - -@pytest.mark.parametrize("name", ("gpt2",)) -def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: - assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/__init__.py b/tests/test_litellm/llms/vertex_ai/gemini/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py deleted file mode 100644 index d3a7ba7a1bd..00000000000 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ /dev/null @@ -1,54 +0,0 @@ -import pytest - -from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_result, -) -from litellm.types.llms.vertex_ai import BlobType - - -def test_convert_tool_response_with_url_image(): - """Test tool response with HTTP URL image (will download and convert).""" - # Use a publicly accessible test image URL - test_image_url = "https://via.placeholder.com/1x1.png" - - tool_message = { - "role": "tool", - "tool_call_id": "call_test456", - "content": [ - {"type": "text", "text": '{"url": "https://example.com"}'}, - {"type": "input_image", "image_url": test_image_url}, - ], - } - - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_test456", - "function": { - "name": "type_text_at", - "arguments": '{"x": 300, "y": 400, "text": "hello"}', - }, - } - ] - } - - try: - result = convert_to_gemini_tool_call_result(tool_message, last_message_with_tool_calls) - - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - result_part = result[0] - assert "function_response" in result_part - assert "inline_data" not in result_part - function_response = result_part["function_response"] - assert function_response["name"] == "type_text_at" - - # Check inline_data is nested under functionResponse.parts. - assert "parts" in function_response - assert len(function_response["parts"]) == 1 - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - except Exception as e: - # Skip test if URL download fails (no internet connection, etc.) - pytest.skip(f"Failed to download image from URL: {e}") diff --git a/tests/test_litellm/llms/volcengine/__init__.py b/tests/test_litellm/llms/volcengine/__init__.py deleted file mode 100644 index 825e259b1fc..00000000000 --- a/tests/test_litellm/llms/volcengine/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Volcengine tests diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py deleted file mode 100644 index 78728d6fd58..00000000000 --- a/tests/test_litellm/test_main.py +++ /dev/null @@ -1,164 +0,0 @@ -import json -import os - -import pytest - - -from unittest.mock import MagicMock, patch - -import litellm - - -async def _async_fake_bedrock_image_details(image_url): - return "ZmFrZS1pbWFnZQ==", "image/png" - - -@pytest.fixture(autouse=True) -def clear_client_cache(): - """ - Clear the HTTP client cache before each test to ensure mocks are used. - This prevents cached real clients from being reused across tests. - """ - cache = getattr(litellm, "in_memory_llm_clients_cache", None) - if cache is not None: - cache.flush_cache() - yield - if cache is not None: - cache.flush_cache() - - -@pytest.fixture(autouse=True) -def add_api_keys_to_env(monkeypatch): - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-api03-1234567890") - monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-api03-1234567890") - monkeypatch.setenv("AWS_ACCESS_KEY_ID", "my-fake-aws-access-key-id") - monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "my-fake-aws-secret-access-key") - monkeypatch.setenv("AWS_REGION", "us-east-1") - # Keep these transformation tests on the simple access-key path. A leaked - # session token or role/web-identity env var pushes Bedrock auth down a - # different branch and fails before the mocked HTTP client is exercised. - monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) - monkeypatch.delenv("AWS_ROLE_ARN", raising=False) - monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) - - -@pytest.mark.parametrize( - "model", - [ - "gemini/gemini-1.5-flash", - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", - "anthropic/claude-3-5-sonnet", - ], -) -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_url_with_format_param(model, sync_mode, monkeypatch): - from litellm import acompletion, completion - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory - - if sync_mode: - client = HTTPHandler() - else: - client = AsyncHTTPHandler() - - # This test is about request shaping, not live image downloads. Stub the - # URL->image conversion helpers so suite-level network/client state from - # earlier tests cannot prevent the mocked provider client from being hit. - fake_base64_image = "data:image/png;base64,ZmFrZS1pbWFnZQ==" - monkeypatch.setattr( - prompt_factory, "convert_url_to_base64", lambda url: fake_base64_image - ) - monkeypatch.setattr( - prompt_factory.BedrockImageProcessor, - "get_image_details", - staticmethod(lambda image_url: ("ZmFrZS1pbWFnZQ==", "image/png")), - ) - monkeypatch.setattr( - prompt_factory.BedrockImageProcessor, - "get_image_details_async", - staticmethod(_async_fake_bedrock_image_details), - ) - - args = { - "model": model, - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", - "format": "image/png", - }, - }, - {"type": "text", "text": "Describe this image"}, - ], - } - ], - } - if model.startswith("gemini/"): - args["api_key"] = "test-api-key" - with patch.object(client, "post", new=MagicMock()) as mock_client: - try: - if sync_mode: - response = completion(**args, client=client) - else: - response = await acompletion(**args, client=client) - print(response) - except Exception as e: - pass - - mock_client.assert_called() - - print(mock_client.call_args.kwargs) - - if "data" in mock_client.call_args.kwargs: - json_str = mock_client.call_args.kwargs["data"] - else: - json_str = json.dumps(mock_client.call_args.kwargs["json"]) - - if isinstance(json_str, bytes): - json_str = json_str.decode("utf-8") - - print(f"type of json_str: {type(json_str)}") - - # Bedrock models convert URLs to base64, while direct Anthropic models support URLs - # bedrock/invoke models use Anthropic messages API which supports URLs - if model.startswith("bedrock/invoke/"): - # bedrock/invoke should convert URLs to base64 (doesn't support URL references) - # URL should NOT be in the JSON (it should be converted to base64) - assert "https://awsmp-logos.s3.amazonaws.com" not in json_str - # Should have base64 data in the source (type="base64", not type="url") - assert '"type":"base64"' in json_str or '"type": "base64"' in json_str - # Should have "data" field containing base64 content - assert '"data"' in json_str - elif model.startswith("bedrock/"): - # Regular Bedrock models should convert URLs to base64 (uses "bytes" field) - # URL should NOT be in the JSON (it should be converted to base64) - assert "https://awsmp-logos.s3.amazonaws.com" not in json_str - # Should have "bytes" field (Bedrock uses "bytes" not "base64" in the field name) - assert '"bytes"' in json_str or '"bytes":' in json_str - elif model.startswith("anthropic/"): - # Direct Anthropic models should pass HTTPS URLs directly (HTTP URLs are converted to base64) - # Since we're using HTTPS URL, it should be passed as-is - assert "https://awsmp-logos.s3.amazonaws.com" in json_str - # For Anthropic, URL references use "url" type, not base64 - assert '"type":"url"' in json_str or '"type": "url"' in json_str - else: - # For other models, check format parameter is respected - assert "png" in json_str - assert "jpeg" not in json_str - - -@pytest.fixture(autouse=True) -def set_openrouter_api_key(): - original_api_key = os.environ.get("OPENROUTER_API_KEY") - os.environ["OPENROUTER_API_KEY"] = "fake-key-for-testing" - yield - if original_api_key is not None: - os.environ["OPENROUTER_API_KEY"] = original_api_key - else: - del os.environ["OPENROUTER_API_KEY"] diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index b1a14e61b96..ae9b30d862b 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -3,16 +3,22 @@ import asyncio import base64 import importlib +import json +import os +import subprocess +import sys import threading import time import traceback from concurrent.futures import Future, wait +from pathlib import Path from typing import Final from unittest.mock import MagicMock import anyio.to_thread import pytest import tiktoken +from tokenizers import Regex, Tokenizer, models, pre_tokenizers from unittest.mock import AsyncMock, patch @@ -1439,3 +1445,89 @@ def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() + + +HUB_TOKENIZER_SCRIPT: Final = """ +import json +import sys +sys.path.insert(0, sys.argv[1]) +import httpx +import huggingface_hub +import litellm +served = json.loads(sys.argv[2]) +text = sys.argv[3] +requested = [] +def handle(request): + repo = request.url.path.lstrip("/").split("/resolve/")[0] + if repo not in served or not request.url.path.endswith("/tokenizer.json"): + return httpx.Response(404) + requested.append(repo) + payload = served[repo].encode() + headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40} + return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"") +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +litellm.cohere_models = {"command-r-v1"} +litellm.anthropic_models = {"claude-2"} +custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer") +print(json.dumps({ + "llama2": litellm.token_counter(model="meta-llama/Llama-2-7b-chat", text=text), + "llama3": litellm.token_counter(model="meta-llama/llama-3-70b-instruct", text=text), + "cohere": litellm.token_counter(model="command-r-v1", text=text), + "anthropic": litellm.token_counter(model="claude-2", text=text), + "custom": litellm.token_counter(custom_tokenizer=custom, text=text), + "requested": sorted(set(requested)), +})) +""" + + +def _word_level_tokenizer_json(pre_tokenizer: pre_tokenizers.PreTokenizer) -> str: + tokenizer: Final = Tokenizer(models.WordLevel(vocab={"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizer + return tokenizer.to_str() + + +def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None: + sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!" + served: Final = { + "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()), + "Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")), + "Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()), + } + expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()} + anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids) + tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample) + assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2 + + result: Final = subprocess.run( + [ + sys.executable, + "-I", + "-c", + HUB_TOKENIZER_SCRIPT, + str(Path(litellm.__file__).parent.parent), + json.dumps(served), + sample, + ], + capture_output=True, + text=True, + timeout=60, + env={ + **os.environ, + "HF_HOME": str(tmp_path / "home"), + "HF_HUB_CACHE": str(tmp_path / "cache"), + "HF_ENDPOINT": "http://127.0.0.1:9", + "HF_HUB_OFFLINE": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + }, + ) + + assert result.returncode == 0, result.stdout + result.stderr + counts: Final = json.loads(result.stdout.strip().splitlines()[-1]) + assert counts == { + "llama2": expected["hf-internal-testing/llama-tokenizer"], + "llama3": expected["Xenova/llama-3-tokenizer"], + "cohere": expected["Xenova/c4ai-command-r-v01-tokenizer"], + "anthropic": anthropic_count, + "custom": expected["Xenova/llama-3-tokenizer"], + "requested": sorted(served), + } diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index a9005ff6a86..9d08442b164 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -17,17 +17,13 @@ from litellm.utils import claude_json_str from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON -OFFLINE_ENCODINGS: Final = ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "o200k_harmony") +ENCODINGS: Final = ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "o200k_harmony") UNICODE_TEXTS: Final = ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64) -@pytest.mark.parametrize("name", OFFLINE_ENCODINGS) +@pytest.mark.parametrize("name", ENCODINGS) @pytest.mark.parametrize("text", UNICODE_TEXTS) def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: - assert_openai_encoding_matches_python(name, text) - - -def assert_openai_encoding_matches_python(name: str, text: str) -> None: reference: Final = tiktoken.get_encoding(name) encoding: Final = OpenAIEncoding.from_tiktoken(name) expected: Final = reference.encode(text) @@ -309,10 +305,6 @@ def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: boo @pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit")) def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: - assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name) - - -def assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: reference: Final = tiktoken.get_encoding(name) encoding: Final = OpenAIEncoding.from_tiktoken(name) text: Final = "hello fanta" diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 4f23ac1773a..0b37e033023 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,7 +1,12 @@ import base64 +from pathlib import Path +from typing import Final +import httpx import pytest +import respx +import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) @@ -2727,3 +2732,40 @@ def test_gemini_server_side_tool_signature_not_duplicated_on_text(): assert "thoughtSignature" not in text_part tool_call_part = next(p for p in parts if "toolCall" in p) assert tool_call_part["thoughtSignature"] == "server_side_signature" + + +WHITE_PNG: Final = (Path(__file__).parents[4] / "white_100x100.png").read_bytes() + + +@respx.mock +def test_convert_tool_response_with_url_image(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + image_url: Final = "https://tool-result-images.test/gemini-tool-response.png" + respx.get(image_url).mock(return_value=httpx.Response(200, content=WHITE_PNG, headers={"content-type": "image/png"})) + tool_message: Final = { + "role": "tool", + "tool_call_id": "call_test456", + "content": [ + {"type": "text", "text": '{"url": "https://example.com"}'}, + {"type": "input_image", "image_url": image_url}, + ], + } + last_message_with_tool_calls: Final = { + "tool_calls": [ + { + "id": "call_test456", + "function": {"name": "type_text_at", "arguments": '{"x": 300, "y": 400, "text": "hello"}'}, + } + ] + } + + result: Final = convert_to_gemini_tool_call_result(tool_message, last_message_with_tool_calls) + + assert isinstance(result, list) + assert len(result) == 1 + assert "inline_data" not in result[0] + function_response: Final = result[0]["function_response"] + assert function_response["name"] == "type_text_at" + assert len(function_response["parts"]) == 1 + inline_data: Final[BlobType] = function_response["parts"][0]["inline_data"] + assert inline_data == {"data": base64.b64encode(WHITE_PNG).decode(), "mime_type": "image/png"} diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/unit/llms/volcengine/test_volcengine_embedding.py similarity index 100% rename from tests/test_litellm/llms/volcengine/test_volcengine_embedding.py rename to tests/unit/llms/volcengine/test_volcengine_embedding.py diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index c06216e4f4e..57200a79a8c 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -17,6 +17,7 @@ import respx import urllib.parse from importlib import import_module +from pathlib import Path from unittest.mock import MagicMock, patch import litellm @@ -56,6 +57,9 @@ def add_api_keys_to_env(monkeypatch): monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) +WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() + + @pytest.fixture def openai_api_response(): mock_response_data = { @@ -213,6 +217,102 @@ async def test_url_with_format_param_openai(model, sync_mode): assert "format" not in json_str +@pytest.mark.parametrize( + "model", + [ + "gemini/gemini-1.5-flash", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic/claude-3-5-sonnet", + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_url_with_format_param(model, sync_mode, monkeypatch): + from litellm import acompletion, completion + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + if sync_mode: + client = HTTPHandler() + else: + client = AsyncHTTPHandler() + + image_url: Final = ( + "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png" + f"?case={sync_mode}-{model}" + ) + args = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": image_url, + "format": "image/png", + }, + }, + {"type": "text", "text": "Describe this image"}, + ], + } + ], + } + if model.startswith("gemini/"): + args["api_key"] = "test-api-key" + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport())) + with ( + respx.mock(assert_all_called=False) as image_host, + patch.object(client, "post", new=MagicMock()) as mock_client, + ): + image_route = image_host.get(image_url).mock( + return_value=httpx.Response(200, content=WHITE_PNG, headers={"content-type": "image/png"}) + ) + try: + if sync_mode: + response = completion(**args, client=client) + else: + response = await acompletion(**args, client=client) + print(response) + except Exception as e: + pass + + mock_client.assert_called() + + print(mock_client.call_args.kwargs) + + if "data" in mock_client.call_args.kwargs: + json_str = mock_client.call_args.kwargs["data"] + else: + json_str = json.dumps(mock_client.call_args.kwargs["json"]) + + if isinstance(json_str, bytes): + json_str = json_str.decode("utf-8") + + print(f"type of json_str: {type(json_str)}") + + if model.startswith("bedrock/invoke/"): + assert "https://awsmp-logos.s3.amazonaws.com" not in json_str + assert '"type":"base64"' in json_str or '"type": "base64"' in json_str + assert '"data"' in json_str + elif model.startswith("bedrock/"): + assert "https://awsmp-logos.s3.amazonaws.com" not in json_str + assert '"bytes"' in json_str or '"bytes":' in json_str + elif model.startswith("anthropic/"): + assert "https://awsmp-logos.s3.amazonaws.com" in json_str + assert '"type":"url"' in json_str or '"type": "url"' in json_str + else: + assert "png" in json_str + assert "jpeg" not in json_str + + fetches_image: Final = not model.startswith("anthropic/") + assert image_route.called is fetches_image + assert (base64.b64encode(WHITE_PNG).decode() in json_str) is fetches_image + + def test_bedrock_latency_optimized_inference(): from litellm.llms.custom_httpx.http_handler import HTTPHandler diff --git a/tests/white_100x100.png b/tests/white_100x100.png new file mode 100644 index 0000000000000000000000000000000000000000..fdd268ded88d11837418170a350a045d3395e9d7 GIT binary patch literal 214 zcmeAS@N?(olHy`uVBq!ia0vp^DIm~-MNWmM-(43$OW{8}a1ZHrhoCGsifoebuCXiwvqY Date: Fri, 25 Sep 2026 19:41:53 -0700 Subject: [PATCH 44/58] fix(rust): preserve nested optional import failures (#43265) * fix(rust): preserve nested optional import failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust): restore Python modules after settings tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../python-bridge/src/python_settings.rs | 175 +++++++++++++++++- 1 file changed, 172 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index f03e5fdce7f..6f8388471dc 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -45,8 +45,13 @@ impl PythonSettings { pub(crate) fn read_or_unset(self, py: Python<'_>) -> PyResult>> { match self.read(py) { Ok(snapshot) => Ok(Some(snapshot)), - Err(error) if error.is_instance_of::(py) => Ok(None), - Err(error) => Err(error), + Err(error) => { + if missing_module(py, &error, "litellm")? { + Ok(None) + } else { + Err(error) + } + } } } @@ -56,9 +61,24 @@ impl PythonSettings { } } +fn missing_module(py: Python<'_>, error: &PyErr, expected: &str) -> PyResult { + if !error.is_instance_of::(py) { + return Ok(false); + } + Ok(error + .value(py) + .getattr("name")? + .extract::>()? + .is_some_and(|name| name == expected)) +} + #[cfg(test)] mod tests { - use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; + use pyo3::{ + exceptions::{PyImportError, PyModuleNotFoundError, PyRuntimeError}, + prelude::*, + types::PyDict, + }; use super::PythonSettings; use crate::coercion::FieldSpec; @@ -150,4 +170,153 @@ values = (Descriptor(), SimpleNamespace(flag=Truth())) ); }); } + + #[test] + fn read_or_unset_returns_none_when_litellm_is_missing() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +import sys +class MissingLitellm: + def find_spec(self, fullname, path=None, target=None): + if fullname == 'litellm': + raise ModuleNotFoundError('No module named litellm', name='litellm') +finder = MissingLitellm() +previous_litellm = sys.modules.get('litellm') +had_litellm = 'litellm' in sys.modules +sys.meta_path.insert(0, finder) +sys.modules.pop('litellm', None) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let result = PythonSettings::Http.read_or_unset(py); + assert!(result.unwrap().is_none()); + py.run( + c" +sys.meta_path.remove(finder) +if had_litellm: + sys.modules['litellm'] = previous_litellm +else: + sys.modules.pop('litellm', None) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn read_or_unset_propagates_nested_module_not_found_errors() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +import sys +import types +previous_modules = { + name: sys.modules[name] + for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.settings') + if name in sys.modules +} +litellm = types.ModuleType('litellm') +litellm.__path__ = [] +rust_bridge = types.ModuleType('litellm.rust_bridge') +rust_bridge.__path__ = [] +settings = types.ModuleType('litellm.rust_bridge.settings') +def http_settings(): + raise ModuleNotFoundError('No module named certifi', name='certifi') +settings.http_settings = http_settings +litellm.rust_bridge = rust_bridge +rust_bridge.settings = settings +sys.modules['litellm'] = litellm +sys.modules['litellm.rust_bridge'] = rust_bridge +sys.modules['litellm.rust_bridge.settings'] = settings +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let error = match PythonSettings::Http.read_or_unset(py) { + Ok(_) => panic!("nested module errors must propagate"), + Err(error) => error, + }; + assert!(error.is_instance_of::(py)); + assert_eq!( + error + .value(py) + .getattr("name") + .unwrap() + .extract::() + .unwrap(), + "certifi" + ); + py.run( + c" +for name in ('litellm.rust_bridge.settings', 'litellm.rust_bridge', 'litellm'): + sys.modules.pop(name, None) +sys.modules.update(previous_modules) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn read_or_unset_propagates_import_errors() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +import sys +import types +previous_modules = { + name: sys.modules[name] + for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.settings') + if name in sys.modules +} +litellm = types.ModuleType('litellm') +litellm.__path__ = [] +rust_bridge = types.ModuleType('litellm.rust_bridge') +rust_bridge.__path__ = [] +settings = types.ModuleType('litellm.rust_bridge.settings') +def http_settings(): + raise ImportError('cannot import name setting') +settings.http_settings = http_settings +litellm.rust_bridge = rust_bridge +rust_bridge.settings = settings +sys.modules['litellm'] = litellm +sys.modules['litellm.rust_bridge'] = rust_bridge +sys.modules['litellm.rust_bridge.settings'] = settings +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let error = match PythonSettings::Http.read_or_unset(py) { + Ok(_) => panic!("import errors must propagate"), + Err(error) => error, + }; + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "ImportError: cannot import name setting"); + py.run( + c" +for name in ('litellm.rust_bridge.settings', 'litellm.rust_bridge', 'litellm'): + sys.modules.pop(name, None) +sys.modules.update(previous_modules) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } } From c822c7fffa19b1292a789d147201c694b1b059da Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 02:42:08 +0000 Subject: [PATCH 45/58] ci: drop main and litellm_* branch filters from the CircleCI litellm-main workflows (#43272) Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/config.yml | 121 ++++++++++++------------------------------- 1 file changed, 32 insertions(+), 89 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 3ff8061fb48..d9c85cfa042 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3329,22 +3329,12 @@ workflows: matrix: parameters: suite: [management, accounting, database, providers, extensions, mcp, sdk, cost, browser] - filters: - branches: - only: - - main - - /litellm_.*/ - integration_contracts: name: integration-<< matrix.suite >>-replica matrix: parameters: suite: [management, database] mode: [replica] - filters: - branches: - only: - - main - - /litellm_.*/ build_and_test: unless: or: @@ -3352,101 +3342,60 @@ workflows: - not: equal: ["", << pipeline.parameters.routing_parity_base >>] jobs: - - using_litellm_on_windows: - filters: &main_branches - branches: - only: - - main - - /litellm_.*/ - - unit: - filters: *main_branches + - using_litellm_on_windows + - unit - provider_replay_harness - - base_sdk_install: - filters: *main_branches - - local_testing_part1: - filters: *main_branches - - local_testing_part2: - filters: *main_branches - - langfuse_logging_unit_tests: - filters: *main_branches - - litellm_assistants_api_testing: - filters: *main_branches - - litellm_router_testing: - filters: *main_branches - - litellm_router_unit_testing: - filters: *main_branches - - auth_ui_unit_tests: - filters: *main_branches - - build_docker_database_image: - filters: *main_branches - - e2e_ui_testing: - filters: *main_branches - - e2e_ui_testing_server_root_path: - filters: *main_branches + - base_sdk_install + - local_testing_part1 + - local_testing_part2 + - langfuse_logging_unit_tests + - litellm_assistants_api_testing + - litellm_router_testing + - litellm_router_unit_testing + - auth_ui_unit_tests + - build_docker_database_image + - e2e_ui_testing + - e2e_ui_testing_server_root_path - build_and_test: requires: - build_docker_database_image - filters: *main_branches - e2e_openai_endpoints: requires: - build_docker_database_image - filters: *main_branches - proxy_logging_guardrails_model_info_tests: requires: - build_docker_database_image - filters: *main_branches - proxy_spend_accuracy_tests: requires: - build_docker_database_image - filters: *main_branches - proxy_multi_instance_tests: requires: - build_docker_database_image - filters: *main_branches - proxy_store_model_in_db_tests: requires: - build_docker_database_image - filters: *main_branches - - proxy_build_from_pip_tests: - filters: *main_branches + - proxy_build_from_pip_tests - proxy_pass_through_endpoint_tests: requires: - build_docker_database_image - filters: *main_branches - proxy_e2e_anthropic_messages_tests: requires: - build_docker_database_image - filters: *main_branches - - llm_translation_testing: - filters: *main_branches - - realtime_translation_testing: - filters: *main_branches - - agent_testing: - filters: *main_branches - - guardrails_testing: - filters: *main_branches - - google_generate_content_endpoint_testing: - filters: *main_branches - - llm_responses_api_testing: - filters: *main_branches - - ocr_testing: - filters: *main_branches - - search_testing: - filters: *main_branches - - batches_testing: - filters: *main_branches - - litellm_utils_testing: - filters: *main_branches - - pass_through_unit_testing: - filters: *main_branches - - image_gen_testing: - filters: *main_branches - - logging_testing: - filters: *main_branches - - audio_testing: - filters: *main_branches - - redis_caching_unit_tests: - filters: *main_branches + - llm_translation_testing + - realtime_translation_testing + - agent_testing + - guardrails_testing + - google_generate_content_endpoint_testing + - llm_responses_api_testing + - ocr_testing + - search_testing + - batches_testing + - litellm_utils_testing + - pass_through_unit_testing + - image_gen_testing + - logging_testing + - audio_testing + - redis_caching_unit_tests - upload-coverage: requires: - realtime_translation_testing @@ -3471,18 +3420,12 @@ workflows: - db_migration_disable_update_check: requires: - build_docker_database_image - filters: *main_branches - - installing_litellm_on_python: - filters: *main_branches - - installing_litellm_on_python_3_13: - filters: *main_branches - - installing_litellm_on_python_v2_migration_resolver: - filters: *main_branches + - installing_litellm_on_python + - installing_litellm_on_python_3_13 + - installing_litellm_on_python_v2_migration_resolver - helm_chart_testing: requires: - build_docker_database_image - filters: *main_branches - test_bad_database_url: requires: - build_docker_database_image - filters: *main_branches From d08746feb15d88d97f9d978d2dfba7ef4bb6d607 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 02:42:45 +0000 Subject: [PATCH 46/58] feat(proxy): email alerts at configured percentages of a team member budget (#42665) * feat(proxy): email alerts at configured percentages of a team member budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(alerting): label team member budget crossings as team member budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(auth): cover the team member alert dispatch from _check_team_member_budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(email): drop the emoji from the team member budget alert template Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): ignore team member alert thresholds outside 1 to 100 on both the backend and the dashboard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): bound team member alert threshold key length before int parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop the legacy covers marker from the team member alert test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(team): reject malformed team_member_max_budget_alert_emails on team writes Thresholds outside 1-100, non-list recipients, and invalid emails now return 422 on /team/new, /team/update and PATCH /team/{id} instead of being stored and silently ignored. The value is stored canonically. Read-side LiteLLM_TeamTable is unchanged, and the PATCH body stays a raw merge patch so a null threshold still deletes it. * fix(auth): enforce and alert on team member budgets only in common_checks The builder re-checked the team member budget inline before common_checks ran the same check, so one request that crossed a team_member_max_budget_alert_emails threshold dispatched two alerts. Drop the inline check; common_checks is the single authorization point and already covers per-member rows, the team default member budget, zero-cost skips and the cross-pod spend counter. Its 422 message now uses the TeamMember=user:team form the builder and budget reservation already returned. * Revert "fix(team): reject malformed team_member_max_budget_alert_emails on team writes" This reverts commit 703e754b4615a292b0a947b142751ae33524e329. * fix(alerting): keep BaseBudgetAlertType.get_event_message zero-arg Requiring user_info broke existing callers and out-of-tree subclasses. The team member label now comes from SlackAlerting.budget_alerts, so the interface and its Readme are unchanged from main. * fix(mcp): keep team member budget enforcement on the MCP OAuth auth dependency The MCP OAuth dependency stops at _user_api_key_auth_builder and never reaches common_checks, so removing the builder's inline member budget check would have let over-budget members through there. Enforce it explicitly for that caller. * fix(auth): keep main's team member budget enforcement, alert once per request Restore the builder's team member budget check and 422 message exactly as on main and drop the MCP-only gate. The builder sends the member alert only on the request it rejects; common_checks sends it for requests that get past the builder, so no request alerts twice. * test(integration): read team member alert deliveries without a shared accumulator * test(integration): match team member alert deliveries by subject so other alerts cannot race the count * refactor(proxy): build the team member alert threshold config without mutable collections Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): collapse the alert recipient isinstance checks into one call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): read the SMTP sink through lock-guarded snapshots and assert the exact deliveries --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../send_emails/base_email.py | 53 +++++-- .../SlackAlerting/budget_alert_types.py | 2 + .../SlackAlerting/slack_alerting.py | 6 +- .../integrations/email_templates/templates.py | 22 +++ litellm/proxy/auth/auth_checks.py | 77 ++++++++- litellm/proxy/auth/user_api_key_auth.py | 14 ++ tests/integration/_support/mail.py | 127 +++++++++++++++ .../spend/test_team_member_budget_alerts.py | 93 +++++++++++ .../proxy/auth/test_auth_checks.py | 137 ++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 147 ++++++++++++++++++ .../send_emails/test_base_email.py | 41 +++++ .../SlackAlerting/test_budget_alert_types.py | 33 +++- .../SlackAlerting/test_slack_alerting.py | 27 ++++ .../src/components/team/TeamInfo.test.tsx | 97 ++++++++++++ .../src/components/team/TeamInfo.tsx | 104 +++++++++++++ .../team/teamMemberBudgetAlertEmails.test.ts | 94 +++++++++++ .../team/teamMemberBudgetAlertEmails.ts | 57 +++++++ 17 files changed, 1113 insertions(+), 18 deletions(-) create mode 100644 tests/integration/_support/mail.py create mode 100644 tests/integration/spend/test_team_member_budget_alerts.py create mode 100644 ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 4be09670e92..6e33d9f1bf3 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -32,6 +32,7 @@ from litellm.integrations.email_templates.key_rotated_email import ( from litellm.integrations.email_templates.templates import ( MAX_BUDGET_ALERT_EMAIL_TEMPLATE, SOFT_BUDGET_ALERT_EMAIL_TEMPLATE, + TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE, TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE, ) from litellm.integrations.email_templates.user_invitation_email import ( @@ -48,6 +49,12 @@ from litellm.secret_managers.main import get_secret_bool from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL +def _max_budget_alert_id(user_info: CallInfo) -> str: + if user_info.event_group == Litellm_EntityType.TEAM_MEMBER: + return f"team_member:{user_info.user_id}:{user_info.team_id}" + return user_info.token or user_info.user_id or "default_id" + + def _parse_email_list(raw) -> List[str]: """Parse emails from a list or comma-separated string.""" if isinstance(raw, list): @@ -373,17 +380,31 @@ class BaseEmailLogger(CustomLogger): greeting = html.escape( event.user_email or event.key_alias or event.token or "" ) - email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( - email_logo_url=email_params.logo_url, - recipient_email=greeting, - percentage=percentage, - spend=spend_str, - max_budget=max_budget_str, - alert_threshold=alert_threshold_str, - base_url=email_params.base_url, - email_support_contact=email_params.support_contact, - email_footer=email_params.signature, - ) + if event.event_group == Litellm_EntityType.TEAM_MEMBER: + email_html_content = TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( + email_logo_url=email_params.logo_url, + member=html.escape(event.user_email or event.user_id or ""), + team_alias=html.escape(event.team_alias or event.team_id or ""), + percentage=percentage, + spend=spend_str, + max_budget=max_budget_str, + alert_threshold=alert_threshold_str, + base_url=email_params.base_url, + email_support_contact=email_params.support_contact, + email_footer=email_params.signature, + ) + else: + email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( + email_logo_url=email_params.logo_url, + recipient_email=greeting, + percentage=percentage, + spend=spend_str, + max_budget=max_budget_str, + alert_threshold=alert_threshold_str, + base_url=email_params.base_url, + email_support_contact=email_params.support_contact, + email_footer=email_params.signature, + ) await self.send_email( from_email=self.DEFAULT_LITELLM_EMAIL, to_email=recipient_emails, @@ -607,7 +628,7 @@ class BaseEmailLogger(CustomLogger): if user_info.spend < threshold_amount: continue - _id = user_info.token or user_info.user_id or "default_id" + _id = _max_budget_alert_id(user_info) _cache_key = ( f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}" ) @@ -618,7 +639,7 @@ class BaseEmailLogger(CustomLogger): emails.append(user_info.user_email) if not emails: verbose_proxy_logger.warning( - "No recipients for %d%% threshold on key %s, skipping alert", + "No recipients for %d%% threshold on %s, skipping alert", threshold_pct, _id, ) @@ -633,7 +654,11 @@ class BaseEmailLogger(CustomLogger): if send_count is not None and send_count > 1: continue - event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached" + event_message = ( + f"Team Member Budget Alert - {threshold_pct}% of Team Member Budget Reached" + if user_info.event_group == Litellm_EntityType.TEAM_MEMBER + else f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached" + ) webhook_event = WebhookEvent( event="max_budget_alert", event_message=event_message, diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index f35ff7b5f82..4fe833acecc 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -63,6 +63,8 @@ class TokenBudgetAlert(BaseBudgetAlertType): return "Key Budget: " def get_id(self, user_info: CallInfo) -> str: + if user_info.event_group == Litellm_EntityType.TEAM_MEMBER: + return f"team_member:{user_info.user_id}:{user_info.team_id}" return user_info.token or "default_id" diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 17ec3ed787d..7c608aac8d9 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -555,7 +555,11 @@ class SlackAlerting(CustomBatchLogger): budget_alert_class: Final = get_budget_alert_type(type) _id: Final = budget_alert_class.get_id(user_info) user_info_str: Final = self._get_user_info_str(user_info) - event_message = budget_alert_class.get_event_message() + event_message = ( + "Team Member Budget: " + if user_info.event_group == Litellm_EntityType.TEAM_MEMBER + else budget_alert_class.get_event_message() + ) # Set default event unless we're in projected_limit_exceeded event: ( diff --git a/litellm/integrations/email_templates/templates.py b/litellm/integrations/email_templates/templates.py index 935067c97fc..2bd079ef15d 100644 --- a/litellm/integrations/email_templates/templates.py +++ b/litellm/integrations/email_templates/templates.py @@ -131,3 +131,25 @@ MAX_BUDGET_ALERT_EMAIL_TEMPLATE: Final = """ {email_footer} """ + +TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE: Final = """ + LiteLLM Logo + +

Hi,
+ + Team member {member} has reached {percentage}% of their team member budget in team {team_alias}.

+ + Current Spend: {spend}
+ Team Member Budget: {max_budget}
+ Alert Threshold: {alert_threshold} ({percentage}%)
+ +

+ Warning: Once this member reaches their team member budget of {max_budget}, their requests in this team will be rejected. +

+ + You can view usage and manage team member budgets in the LiteLLM Dashboard.

+ + If you have any questions, please send an email to {email_support_contact}

+ + {email_footer} +""" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f19a8055ae6..12d420141f1 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -18,7 +18,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm @@ -5682,6 +5682,64 @@ async def _virtual_key_max_budget_alert_check( ) +TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_emails" +_TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _is_valid_alert_threshold_pct(pct: str) -> bool: + return pct.isdigit() and len(pct) <= 3 and 1 <= int(pct) <= 100 + + +def _alert_recipients(raw: object) -> Sequence[str] | None: + if isinstance(raw, (str, Sequence)): + return _parse_email_list(raw) + return None + + +def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: + try: + config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) + except ValidationError: + return None + return MappingProxyType( + {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} + ) + + +def _team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None + alert_email_config: Final = _merge_budget_alert_email_configs( + global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) + ) + if not alert_email_config or spend <= 0: + return + min_pct: Final = min(int(pct) for pct in alert_email_config) + if spend < max_budget * (min_pct / 100.0): + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + max_budget_alert_emails=alert_email_config, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) + + async def _check_team_member_budget( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, @@ -5747,7 +5805,22 @@ async def _check_team_member_budget( max_budget=team_member_budget, ) - if math.isfinite(team_member_budget) and team_member_spend >= team_member_budget: + if not math.isfinite(team_member_budget): + return + + _team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=team_object.metadata, + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=team_member_budget, + ) + + if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 22c3a248b9d..e3ce9bcd850 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import ( _get_user_role, _is_model_cost_zero, _is_user_proxy_admin, + _team_member_max_budget_alert_check, _virtual_key_max_budget_alert_check, _virtual_key_max_budget_check, _virtual_key_soft_budget_check, @@ -2287,6 +2288,19 @@ async def _user_api_key_auth_builder( max_budget=team_member_budget, ) if team_member_spend >= team_member_budget: + # common_checks sends this alert on requests that get past here, so only the + # request rejected here sends it from the builder. + _team_member_max_budget_alert_check( + team_id=_team_id, + team_alias=valid_token.team_alias, + team_metadata=valid_token.team_metadata, + organization_id=valid_token.org_id, + user_id=_user_id, + user_email=user_obj.user_email if user_obj is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=team_member_budget, + ) _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, diff --git a/tests/integration/_support/mail.py b/tests/integration/_support/mail.py new file mode 100644 index 00000000000..3894baeccc3 --- /dev/null +++ b/tests/integration/_support/mail.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import socketserver +import threading +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from email import message_from_bytes +from email.message import Message +from queue import SimpleQueue +from typing import Final + + +@dataclass(frozen=True, slots=True) +class Delivery: + sender: str + recipients: tuple[str, ...] + message: Message + + @property + def subject(self) -> str: + return str(self.message["Subject"]) + + @property + def html(self) -> str: + for part in self.message.walk(): + if part.get_content_type() == "text/html": + return part.get_payload(decode=True).decode() + return "" + + +class Mailbox: + def __init__(self, host: str, port: int) -> None: + self.host: Final = host + self.port: Final = port + self._lock: Final = threading.Lock() + self._deliveries: tuple[Delivery, ...] = () + + def record(self, delivery: Delivery) -> None: + with self._lock: + self._deliveries = (*self._deliveries, delivery) + + def deliveries(self) -> tuple[Delivery, ...]: + with self._lock: + return self._deliveries + + +def _address(argument: str) -> str: + return argument.split(":", 1)[1].strip().strip("<>") + + +@contextmanager +def smtp_sink() -> Generator[Mailbox, None, None]: + """Owned plaintext SMTP peer; deliveries traverse the proxy's real smtplib client.""" + errors: Final[SimpleQueue[Exception]] = SimpleQueue() + + class Handler(socketserver.StreamRequestHandler): + timeout = 5 + + def handle(self) -> None: + try: + self._session() + except Exception as error: + errors.put(error) + + def _reply(self, line: str) -> None: + self.wfile.write(f"{line}\r\n".encode()) + self.wfile.flush() + + def _session(self) -> None: + self._reply("220 integration-smtp ready") + # rebind-ok: the SMTP envelope is built across MAIL/RCPT lines and reset after DATA or RSET. + sender = "" + recipients: tuple[str, ...] = () + while True: + raw: Final = self.rfile.readline() + if not raw: + return + line: Final = raw.decode().rstrip("\r\n") + verb: Final = line.split(" ", 1)[0].upper() + if verb in {"EHLO", "HELO"}: + self._reply("250 integration-smtp") + elif verb == "MAIL": + sender = _address(line) + self._reply("250 OK") + elif verb == "RCPT": + recipients = (*recipients, _address(line)) + self._reply("250 OK") + elif verb == "DATA": + self._reply("354 End data with .") + body = bytearray() + while True: + chunk: Final = self.rfile.readline() + if not chunk or chunk == b".\r\n": + break + body.extend(chunk[1:] if chunk.startswith(b"..") else chunk) + mailbox.record(Delivery(sender, recipients, message_from_bytes(bytes(body)))) + sender, recipients = "", () + self._reply("250 OK queued") + elif verb == "RSET": + sender, recipients = "", () + self._reply("250 OK") + elif verb == "NOOP": + self._reply("250 OK") + elif verb == "QUIT": + self._reply("221 Bye") + return + else: + self._reply("502 Command not implemented") + + class OwnedServer(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = False + + with OwnedServer(("127.0.0.1", 0), Handler) as server: + mailbox: Final = Mailbox("127.0.0.1", server.server_address[1]) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield mailbox + finally: + server.shutdown() + thread.join(timeout=6) + assert not thread.is_alive(), "Owned SMTP server survived cleanup" + server.server_close() + failure: Final = None if errors.empty() else errors.get_nowait() + assert failure is None, f"Owned SMTP peer failed: {failure!r}" diff --git a/tests/integration/spend/test_team_member_budget_alerts.py b/tests/integration/spend/test_team_member_budget_alerts.py new file mode 100644 index 00000000000..f12bcb9748a --- /dev/null +++ b/tests/integration/spend/test_team_member_budget_alerts.py @@ -0,0 +1,93 @@ +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.mail import smtp_sink +from integration._support.process import owned_proxy + +MEMBER_BUDGET: Final = 0.10 +CALL_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def _membership_spend(user_id: str, team_id: str) -> float: + rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s', (user_id, team_id) + ) + return float(str(rows[0]["spend"])) if rows else 0.0 + + +def test_team_member_budget_thresholds_email_member_and_configured_recipients(gateway: Gateway, tmp_path: Path) -> None: + member_email: Final = f"member-{uuid.uuid4().hex}@integration.test" + finance_email: Final = f"finance-{uuid.uuid4().hex}@integration.test" + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["general_settings"]["alerting"] = ["email"] + path: Final = tmp_path / "email-alerting.yaml" + path.write_text(yaml.safe_dump(configuration)) + with smtp_sink() as mailbox: + overrides: Final = { + "SMTP_HOST": mailbox.host, + "SMTP_PORT": str(mailbox.port), + "SMTP_TLS": "False", + "SMTP_SENDER_EMAIL": "alerts@integration.test", + } + with owned_proxy(gateway, tmp_path, overrides, config=path) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + user_id: Final = scenario.user(user_email=member_email) + team_id: Final = scenario.team( + models=[model], + team_member_budget=MEMBER_BUDGET, + metadata={"team_member_max_budget_alert_emails": {"50": [], "100": [finance_email]}}, + ) + candidate.post("/team/member_add", {"team_id": team_id, "member": {"user_id": user_id, "role": "user"}}) + key: Final = scenario.key(team_id=team_id, user_id=user_id) + + first: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "first call"}]}, + key=key, + ) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(CALL_COST) + eventually( + lambda: _membership_spend(user_id, team_id), lambda spend: spend == pytest.approx(CALL_COST), seconds=70 + ) + assert mailbox.deliveries() == (), "no threshold is reached before the first call is recorded" + + second: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "second call"}]}, + key=key, + ) + assert second.status_code == 200, second.text + halfway: Final = eventually(mailbox.deliveries, lambda found: len(found) >= 1, seconds=30) + assert [delivery.recipients for delivery in halfway] == [(member_email,)], halfway + assert "50%" in halfway[0].subject, halfway[0].subject + assert f"${MEMBER_BUDGET}" in halfway[0].html, halfway[0].html + eventually( + lambda: _membership_spend(user_id, team_id), + lambda spend: spend == pytest.approx(2 * CALL_COST), + seconds=70, + ) + + third: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "third call"}]}, + key=key, + ) + assert third.status_code == 422 and third.json()["error"]["type"] == "budget_exceeded", third.text + capped: Final = eventually(mailbox.deliveries, lambda found: len(found) >= 3, seconds=30) + hundred: Final = capped[1:] + assert all("100%" in delivery.subject for delivery in hundred), capped + assert {recipient for delivery in hundred for recipient in delivery.recipients} == { + member_email, + finance_email, + }, capped + assert all(member_email in delivery.html and f"${MEMBER_BUDGET}" in delivery.html for delivery in hundred) + assert len(capped) == 3, capped diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index e42a47a1091..f014e9c26d1 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,5 +1,6 @@ import asyncio import json +import sys import time from collections.abc import Iterator, Mapping from types import SimpleNamespace @@ -52,6 +53,7 @@ from litellm.proxy.auth.auth_checks import ( _log_budget_lookup_failure, _tag_max_budget_check, _team_max_budget_check, + _team_member_max_budget_alert_check, _virtual_key_max_budget_alert_check, _check_agent_caller_model_access, _virtual_key_max_budget_check, @@ -3774,6 +3776,141 @@ async def test_virtual_key_max_budget_alert_check_without_user_obj(): assert captured_call_info.user_email is None +@pytest.mark.parametrize( + "spend, team_metadata, expect_alert", + [ + (0.05, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, True), + (0.10, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, True), + (0.049, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, False), + (0.0, {"team_member_max_budget_alert_emails": {"50": []}}, False), + (0.10, {"team_member_max_budget_alert_emails": {"abc": []}}, False), + (0.05, {"team_member_max_budget_alert_emails": {"0": ["finance@co.com"], "100": []}}, False), + (0.10, {"team_member_max_budget_alert_emails": {"101": ["finance@co.com"]}}, False), + (0.10, {"team_member_max_budget_alert_emails": "50"}, False), + (0.10, {"soft_budget_alerting_emails": ["finance@co.com"]}, False), + (0.10, None, False), + ], +) +@pytest.mark.asyncio +async def test_team_member_max_budget_alert_check_dispatches_only_at_configured_thresholds( + spend, team_metadata, expect_alert +): + captured: list[tuple[str, CallInfo]] = [] + + class RecordingProxyLogging: + async def budget_alerts(self, type, user_info): + captured.append((type, user_info)) + + _team_member_max_budget_alert_check( + team_id="team-1", + team_alias="platform", + team_metadata=team_metadata, + organization_id="org-1", + user_id="user-1", + user_email="member@co.com", + proxy_logging_obj=RecordingProxyLogging(), + spend=spend, + max_budget=0.10, + ) + await asyncio.sleep(0) + + if not expect_alert: + assert captured == [], captured + return + assert [type for type, _ in captured] == ["max_budget_alert"], captured + call_info = captured[0][1] + assert call_info.event_group == Litellm_EntityType.TEAM_MEMBER + assert (call_info.spend, call_info.max_budget) == (spend, 0.10) + assert (call_info.user_id, call_info.user_email) == ("user-1", "member@co.com") + assert (call_info.team_id, call_info.team_alias, call_info.organization_id) == ("team-1", "platform", "org-1") + assert call_info.max_budget_alert_emails == {"50": [], "100": ["finance@co.com"]} + assert call_info.token is None + + +@pytest.mark.asyncio +async def test_team_member_max_budget_alert_check_drops_thresholds_outside_1_to_100(): + captured: list[CallInfo] = [] + + class RecordingProxyLogging: + async def budget_alerts(self, type, user_info): + captured.append(user_info) + + _team_member_max_budget_alert_check( + team_id="team-1", + team_alias="platform", + team_metadata={ + "team_member_max_budget_alert_emails": { + "0": ["a@co.com"], + "50": [], + "150": ["b@co.com"], + "1" * (sys.int_info.default_max_str_digits + 1): ["c@co.com"], + } + }, + organization_id="org-1", + user_id="user-1", + user_email="member@co.com", + proxy_logging_obj=RecordingProxyLogging(), + spend=0.05, + max_budget=0.10, + ) + await asyncio.sleep(0) + + assert [call_info.max_budget_alert_emails for call_info in captured] == [{"50": []}], captured + + +@pytest.mark.asyncio +async def test_check_team_member_budget_dispatches_the_configured_alert_before_the_hard_cap(): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership + + captured: list[tuple[str, CallInfo]] = [] + + class RecordingProxyLogging: + async def budget_alerts(self, type, user_info): + captured.append((type, user_info)) + + team_object = LiteLLM_TeamTable( + team_id="team-1", + team_alias="platform", + metadata={"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, + ) + user_object = LiteLLM_UserTable(user_id="user-1", user_email="member@co.com") + valid_token = UserAPIKeyAuth(token="tok-1", user_id="user-1", team_id="team-1") + team_membership = LiteLLM_TeamMembership( + user_id="user-1", + team_id="team-1", + spend=0.10, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def spend_from_fallback(counter_key, fallback_spend, max_budget=None, **kwargs): + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", spend_from_fallback), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, return_value=team_membership + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=RecordingProxyLogging(), + ) + await asyncio.sleep(0) + + assert (exc_info.value.entity_type, exc_info.value.entity_id) == ("team_member", "user-1:team-1") + assert [type for type, _ in captured] == ["max_budget_alert"], captured + call_info = captured[0][1] + assert call_info.event_group == Litellm_EntityType.TEAM_MEMBER + assert (call_info.spend, call_info.max_budget) == (0.10, 0.10) + assert (call_info.user_id, call_info.user_email, call_info.team_id) == ("user-1", "member@co.com", "team-1") + assert call_info.max_budget_alert_emails == {"50": [], "100": ["finance@co.com"]} + + @pytest.mark.parametrize( "spend, max_budget, expect_alert", [ diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index de669449f85..470db99108a 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -29,6 +29,7 @@ from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, + Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -8055,6 +8056,152 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset assert "Max budget: 2.0" in exc_info.value.message +async def _authenticate_and_authorize(mock_request, api_key): + """Builder then the single common_checks gate, the same sequence user_api_key_auth runs.""" + from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request + + request_data = {"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]} + auth_obj = await _user_api_key_auth_builder( + request=mock_request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + recovered = await _authorize_authenticated_request( + user_api_key_auth_obj=auth_obj, + request=mock_request, + request_data=request_data, + route="/v1/messages", + api_key=f"Bearer {api_key}", + ) + return recovered or auth_obj + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_member_spend, expect_blocked, expected_alerts", + [ + (1.1, False, 0), + (1.2, False, 1), + (2.4, True, 1), + ], +) +async def test_cached_key_team_member_budget_emails_configured_thresholds( + team_member_spend, expect_blocked, expected_alerts +): + """The team's team_member_max_budget_alert_emails thresholds fire from the cached-key auth path, + including on the request that trips the hard cap, and stay silent below the lowest threshold.""" + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import ( + team_membership_auth_cache_key, + team_membership_reservation_cache_key, + ) + from litellm.proxy.utils import hash_token + + api_key = "sk-team-member-alert-thresholds" + hashed_token = hash_token(api_key) + team_id = "team-alert-thresholds" + user_id = "user-alert-thresholds" + alert_emails = {"50": [], "100": ["finance@example.com"]} + + user_api_key_cache = DualCache() + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=UserAPIKeyAuth( + token=hashed_token, + team_id=team_id, + team_alias="platform", + team_metadata={"team_member_max_budget_alert_emails": alert_emails}, + user_id=user_id, + team_member_spend=team_member_spend, + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + await user_api_key_cache.async_set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj( + team_id=team_id, + team_alias="platform", + metadata={"team_member_max_budget_alert_emails": alert_emails}, + ), + ) + await user_api_key_cache.async_set_cache( + key=user_id, + value=LiteLLM_UserTable( + user_id=user_id, user_email="member@example.com", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + membership = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=team_member_spend, + budget_id="budget-alert-thresholds", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4), + ) + # A live proxy holds the row under both keys, so any second team-member check in the + # auth flow would find it too and send a duplicate alert. + for membership_cache_key in ( + team_membership_reservation_cache_key(team_id=team_id, user_id=user_id), + team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + ): + await user_api_key_cache.async_set_cache(key=membership_cache_key, value=membership) + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {api_key}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock(return_value=None) + + async def _auth(): + return await _authenticate_and_authorize(mock_request, api_key) + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True} + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: seed the cached key, team and membership without a DB + "litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), + patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=team_member_spend), + ), + ): + if expect_blocked: + with pytest.raises(ProxyException) as exc_info: + await _auth() + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + else: + await _auth() + await asyncio.sleep(0) + + assert proxy_logging_obj.budget_alerts.await_count == expected_alerts + if expected_alerts == 0: + return + call_info = proxy_logging_obj.budget_alerts.await_args.kwargs["user_info"] + assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "max_budget_alert" + assert call_info.event_group == Litellm_EntityType.TEAM_MEMBER + assert (call_info.spend, call_info.max_budget) == (team_member_spend, 2.4) + assert (call_info.user_id, call_info.user_email) == (user_id, "member@example.com") + assert (call_info.team_id, call_info.team_alias) == (team_id, "platform") + assert call_info.max_budget_alert_emails == alert_emails + + async def _proxy_exception_for_key( api_key: str, general_settings: dict[str, bool], diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 8b89c592f02..52e44ca5448 100644 --- a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -1090,6 +1090,47 @@ async def test_multi_threshold_empty_emails_only_owner( assert to_emails == ["owner@co.com"] +@pytest.mark.asyncio +async def test_multi_threshold_team_member_alert_renders_member_template_per_team( + base_email_logger, mock_send_email +): + """A team member budget alert is keyed per member and team, names the member and team, + and goes to the member plus the threshold's configured recipients""" + user_info = CallInfo( + user_id="member_1", + user_email="member@co.com", + team_id="team_a", + team_alias="Platform", + spend=0.10, + max_budget=0.10, + event_group=Litellm_EntityType.TEAM_MEMBER, + max_budget_alert_emails={"50": [], "100": ["finance@co.com"]}, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + cache_keys = sorted(c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list) + assert cache_keys == [ + "email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a", + "email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a", + ] + assert mock_send_email.call_count == 2 + hundred = next( + c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"] + ) + assert hundred["subject"] == "LiteLLM: Team Member Budget Alert - 100% of Team Member Budget Reached" + assert sorted(hundred["to_email"]) == ["finance@co.com", "member@co.com"] + assert "member@co.com" in hundred["html_body"] and "Platform" in hundred["html_body"] + assert "team member budget" in hundred["html_body"] and "$0.1" in hundred["html_body"] + fifty = next(c.kwargs for c in mock_send_email.call_args_list if "50%" in c.kwargs["subject"]) + assert fifty["to_email"] == ["member@co.com"] + + @pytest.mark.asyncio async def test_no_map_preserves_old_single_threshold( base_email_logger, mock_send_email diff --git a/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py index 52b7cc983a7..f3199d9ebf9 100644 --- a/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py +++ b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py @@ -1,4 +1,7 @@ -from litellm.integrations.SlackAlerting.budget_alert_types import SoftBudgetAlert +from litellm.integrations.SlackAlerting.budget_alert_types import ( + SoftBudgetAlert, + TokenBudgetAlert, +) from litellm.proxy._types import CallInfo, Litellm_EntityType @@ -64,3 +67,31 @@ class TestSoftBudgetAlert: result = alert.get_id(user_info) assert result == "default_id" + + +class TestTokenBudgetAlert: + def test_get_id_dedupes_team_member_alerts_per_member_and_team(self): + alert = TokenBudgetAlert() + team_a = CallInfo( + spend=8.0, max_budget=10.0, user_id="member_1", team_id="team_a", event_group=Litellm_EntityType.TEAM_MEMBER + ) + team_b = CallInfo( + spend=8.0, max_budget=10.0, user_id="member_1", team_id="team_b", event_group=Litellm_EntityType.TEAM_MEMBER + ) + + assert alert.get_id(team_a) == "team_member:member_1:team_a" + assert alert.get_id(team_b) == "team_member:member_1:team_b" + + def test_get_id_uses_token_for_key_alerts(self): + alert = TokenBudgetAlert() + user_info = CallInfo( + spend=8.0, + max_budget=10.0, + token="hashed_key", + user_id="member_1", + team_id="team_a", + event_group=Litellm_EntityType.KEY, + ) + + assert alert.get_id(user_info) == "hashed_key" + assert alert.get_event_message() == "Key Budget: " diff --git a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py index b9e5ff2eeb7..0c2b95fd448 100644 --- a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py @@ -393,6 +393,33 @@ def _slack_alerting_with_env_resolution() -> SlackAlerting: return slack_alerting +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event_group, expected_prefix", + [ + (Litellm_EntityType.TEAM_MEMBER, "Team Member Budget: Budget Crossed"), + (Litellm_EntityType.KEY, "Key Budget: Budget Crossed"), + ], +) +async def test_max_budget_alert_labels_team_member_budget(event_group, expected_prefix): + slack_alerting: Final = _slack_alerting_with_env_resolution() + slack_alerting.send_alert = AsyncMock() + + await slack_alerting.budget_alerts( + type="max_budget_alert", + user_info=CallInfo( + spend=10.5, + max_budget=10.0, + token="hashed_key", + user_id="member_1", + team_id="team_a", + event_group=event_group, + ), + ) + + assert slack_alerting.send_alert.await_args.kwargs["message"].startswith(expected_prefix) + + @pytest.mark.asyncio async def test_send_alert_falls_back_to_alerting_webhook_url_env(monkeypatch): monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 37f63433d2a..a693ee971d4 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -2338,6 +2338,103 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { expect(wireBody(payload)).toStrictEqual(expected); }); + const openEditorWithMemberBudgetAlerts = async (user: ReturnType) => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + team_member_budget_table: { max_budget: 42 }, + metadata: { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@test.com"] } }, + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + }; + + const memberBudgetAlertEmails = (payload: Record) => + (wireBody(payload).metadata as Record).team_member_max_budget_alert_emails; + + it("resends the stored team member budget alert thresholds when the section stays closed", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithMemberBudgetAlerts(user); + + const payload = await save(user); + + expect(memberBudgetAlertEmails(payload)).toStrictEqual({ "50": [], "100": ["finance@test.com"] }); + }); + + it("sends the edited team member budget alert thresholds as a percent to recipients map", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithMemberBudgetAlerts(user); + + await user.click(screen.getByText("Team Member Settings")); + await screen.findByLabelText("Default Budget (USD)"); + const thresholds = screen.getAllByPlaceholderText("% of budget"); + const recipients = screen.getAllByPlaceholderText(/Additional recipients/); + expect(thresholds.map((input) => (input as HTMLInputElement).value)).toStrictEqual(["50", "100"]); + expect(recipients.map((input) => (input as HTMLInputElement).value)).toStrictEqual(["", "finance@test.com"]); + + fireEvent.change(thresholds[0], { target: { value: "75" } }); + fireEvent.change(recipients[0], { target: { value: " lead@test.com, finance@test.com " } }); + await user.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" })); + fireEvent.change(screen.getAllByPlaceholderText("% of budget")[2], { target: { value: "90" } }); + + const payload = await save(user); + + expect(memberBudgetAlertEmails(payload)).toStrictEqual({ + "75": ["lead@test.com", "finance@test.com"], + "100": ["finance@test.com"], + "90": [], + }); + }); + + it("drops the team member budget alert thresholds key once every row is removed", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithMemberBudgetAlerts(user); + + await user.click(screen.getByText("Team Member Settings")); + await screen.findByLabelText("Default Budget (USD)"); + const removeButtons = screen.getAllByRole("button", { name: "Remove budget alert threshold" }); + await user.click(removeButtons[1]); + await user.click(removeButtons[0]); + + const payload = await save(user); + + expect(memberBudgetAlertEmails(payload)).toBeUndefined(); + }); + + it("blocks the save when a team member budget alert threshold is above 100", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithMemberBudgetAlerts(user); + + await user.click(screen.getByText("Team Member Settings")); + await screen.findByLabelText("Default Budget (USD)"); + const threshold = screen.getAllByPlaceholderText("% of budget")[0] as HTMLInputElement; + fireEvent.change(threshold, { target: { value: "150" } }); + expect(threshold.validity.rangeOverflow).toBe(true); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).not.toHaveBeenCalled()); + }); + + it("refuses to save a team member budget alert row with no threshold", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithMemberBudgetAlerts(user); + + await user.click(screen.getByText("Team Member Settings")); + await screen.findByLabelText("Default Budget (USD)"); + await user.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await screen.findByText("Enter a whole number from 1 to 100"); + expect(networking.teamUpdateCall).not.toHaveBeenCalled(); + }); + it("carries every typed value to the update payload at the type and shape antd sends today", async () => { const user = userEvent.setup({ delay: null }); await openEditor(user); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 78ed507216a..3845f94593d 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -118,6 +118,13 @@ import { TEAM_INFO_TAB_LABELS, } from "./tabVisibilityUtils"; import TeamMembersComponent from "./TeamMemberTab"; +import { + isValidThreshold, + TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY, + teamMemberBudgetAlertEmailsFromRows, + teamMemberBudgetAlertRowsFromMetadata, + teamMemberBudgetAlertSummary, +} from "./teamMemberBudgetAlertEmails"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; import { customBudgetMemberUserIds, shouldPromptMemberBudgetReset } from "./memberBudgetReset"; @@ -128,6 +135,7 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "logging", "secret_manager_settings", "soft_budget_alerting_emails", + TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY, "model_tpm_limit", "model_rpm_limit", "default_estimated_output_tokens", @@ -355,6 +363,18 @@ const teamUpdateFieldsSchema = z.object({ team_member_key_duration: z.string().optional(), team_member_tpm_limit: numericInputSchema, team_member_rpm_limit: numericInputSchema, + team_member_max_budget_alert_emails: z + .array(z.object({ threshold: z.number().nullable(), emails: z.string() })) + .superRefine((rows, ctx) => { + rows.forEach((row, index) => { + if (!isValidThreshold(row.threshold)) { + ctx.addIssue({ code: "custom", message: "Enter a whole number from 1 to 100", path: [index, "threshold"] }); + } else if (rows.filter((other) => other.threshold === row.threshold).length > 1) { + ctx.addIssue({ code: "custom", message: "Duplicate threshold", path: [index, "threshold"] }); + } + }); + }) + .optional(), budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, @@ -422,6 +442,7 @@ const TEAM_MEMBER_SETTINGS_FIELDS = [ "team_member_key_duration", "team_member_tpm_limit", "team_member_rpm_limit", + "team_member_max_budget_alert_emails", ] as const; const SEARCH_TOOL_SETTINGS_FIELDS = ["object_permission_search_tools"] as const; @@ -437,6 +458,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = { team_member_key_duration: undefined, team_member_tpm_limit: undefined, team_member_rpm_limit: undefined, + team_member_max_budget_alert_emails: [], budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, @@ -487,6 +509,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): team_member_key_duration: info.metadata?.team_member_key_duration, team_member_tpm_limit: info.team_member_budget_table?.tpm_limit, team_member_rpm_limit: info.team_member_budget_table?.rpm_limit, + team_member_max_budget_alert_emails: [...teamMemberBudgetAlertRowsFromMetadata(info.metadata)], budget_duration: info.budget_duration, tpm_limit: info.tpm_limit, rpm_limit: info.rpm_limit, @@ -572,6 +595,11 @@ const TeamInfoView: React.FC = ({ append: appendModelLimit, remove: removeModelLimit, } = useFieldArray({ control: form.control, name: "modelLimits" }); + const { + fields: memberBudgetAlertRows, + append: appendMemberBudgetAlertRow, + remove: removeMemberBudgetAlertRow, + } = useFieldArray({ control: form.control, name: "team_member_max_budget_alert_emails" }); const [teamMemberSettingsOpen, setTeamMemberSettingsOpen] = useState(false); const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false); const [isEditMemberModalVisible, setIsEditMemberModalVisible] = useState(false); @@ -994,6 +1022,15 @@ const TeamInfoView: React.FC = ({ ? { allowed_passthrough_routes: info.metadata.allowed_passthrough_routes } : {}; + const memberBudgetAlertEmails = + values.team_member_max_budget_alert_emails !== undefined + ? teamMemberBudgetAlertEmailsFromRows(values.team_member_max_budget_alert_emails) + : info.metadata?.[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]; + const memberBudgetAlertEmailsMetadata = + memberBudgetAlertEmails !== undefined && Object.keys(memberBudgetAlertEmails).length > 0 + ? { [TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]: memberBudgetAlertEmails } + : {}; + const updateData: any = { team_id: teamId, team_alias: values.team_alias, @@ -1025,6 +1062,7 @@ const TeamInfoView: React.FC = ({ .filter((email: string) => email.length > 0) : values.soft_budget_alerting_emails || [], ...(secretManagerSettings !== undefined ? { secret_manager_settings: secretManagerSettings } : {}), + ...memberBudgetAlertEmailsMetadata, }, ...(values.policies?.length > 0 ? { policies: values.policies } : {}), ...(values.organization_id !== info.organization_id ? { organization_id: values.organization_id ?? null } : {}), @@ -1632,6 +1670,71 @@ const TeamInfoView: React.FC = ({ )} + + + {labelWithHint( + "Budget Alert Thresholds", + "Email each member when their spend reaches a percentage of their team member budget. The member is always notified; add comma-separated addresses to notify others as well. Requires email alerting to be configured on the proxy.", + )} + + {memberBudgetAlertRows.map((row, index) => ( +
+ + {({ ref, value, onChange, ...field }) => ( + ) => + onChange(event.target.value === "" ? null : Number(event.target.value)) + } + placeholder="% of budget" + min={1} + max={100} + step={1} + /> + )} + + + {({ ref, value, ...field }) => ( + + )} + + +
+ ))} + + @@ -2202,6 +2305,7 @@ const TeamInfoView: React.FC = ({
Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}
TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}
RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}
+
Budget Alert Thresholds: {teamMemberBudgetAlertSummary(info.metadata).join("; ") || "None"}

Router Settings

diff --git a/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts b/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts new file mode 100644 index 00000000000..3faa051047b --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it } from "vitest"; +import { + isValidThreshold, + teamMemberBudgetAlertEmailsFromRows, + teamMemberBudgetAlertRowsFromMetadata, + teamMemberBudgetAlertSummary, +} from "./teamMemberBudgetAlertEmails"; + +describe("teamMemberBudgetAlertRowsFromMetadata", () => { + it("turns the stored threshold map into rows sorted by threshold", () => { + const metadata = { + team_member_max_budget_alert_emails: { "100": ["finance@example.com", "cto@example.com"], "50": [] }, + }; + expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([ + { threshold: 50, emails: "" }, + { threshold: 100, emails: "finance@example.com, cto@example.com" }, + ]); + }); + + it("drops non-numeric thresholds and non-list recipients instead of crashing", () => { + const metadata = { + team_member_max_budget_alert_emails: { fifty: [], "75": "finance@example.com", "90": [1], "100": ["a@b.c"] }, + }; + expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([{ threshold: 100, emails: "a@b.c" }]); + }); + + it("drops API-stored thresholds outside 1 to 100 so they never block the form", () => { + const metadata = { + team_member_max_budget_alert_emails: { "0": ["a@b.c"], "50": [], "101": ["a@b.c"] }, + }; + expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([{ threshold: 50, emails: "" }]); + }); + + it.each([undefined, null, "50", { team_member_max_budget_alert_emails: "50" }, { soft_budget_alerting_emails: [] }])( + "returns no rows for unrelated or malformed metadata %j", + (metadata) => { + expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([]); + }, + ); +}); + +describe("teamMemberBudgetAlertEmailsFromRows", () => { + it("builds the threshold map, splitting, trimming and deduplicating recipients", () => { + expect( + teamMemberBudgetAlertEmailsFromRows([ + { threshold: 50, emails: "" }, + { threshold: 100, emails: " finance@example.com,cto@example.com , finance@example.com, " }, + ]), + ).toEqual({ "50": [], "100": ["finance@example.com", "cto@example.com"] }); + }); + + it("skips rows without a valid threshold", () => { + expect( + teamMemberBudgetAlertEmailsFromRows([ + { threshold: null, emails: "finance@example.com" }, + { threshold: 0, emails: "" }, + { threshold: 101, emails: "" }, + { threshold: 12.5, emails: "" }, + { threshold: 80, emails: "" }, + ]), + ).toEqual({ "80": [] }); + }); + + it("round-trips the stored config", () => { + const stored = { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@example.com"] } }; + expect(teamMemberBudgetAlertEmailsFromRows(teamMemberBudgetAlertRowsFromMetadata(stored))).toEqual( + stored.team_member_max_budget_alert_emails, + ); + }); +}); + +describe("isValidThreshold", () => { + it.each([ + [1, true], + [50, true], + [100, true], + [0, false], + [101, false], + [33.3, false], + [null, false], + ])("treats %s as valid=%s", (threshold, valid) => { + expect(isValidThreshold(threshold)).toBe(valid); + }); +}); + +describe("teamMemberBudgetAlertSummary", () => { + it("states that the member is always notified and lists extra recipients", () => { + expect( + teamMemberBudgetAlertSummary({ + team_member_max_budget_alert_emails: { "100": ["finance@example.com"], "50": [] }, + }), + ).toEqual(["50%: member", "100%: member, finance@example.com"]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts b/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts new file mode 100644 index 00000000000..36d5ddbac02 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts @@ -0,0 +1,57 @@ +export const TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY = "team_member_max_budget_alert_emails" as const; + +export interface TeamMemberBudgetAlertRow { + readonly threshold: number | null; + readonly emails: string; +} + +export type TeamMemberBudgetAlertEmails = Readonly>; + +const isEmailList = (value: unknown): value is readonly string[] => + Array.isArray(value) && value.every((email) => typeof email === "string"); + +const splitEmails = (emails: string): readonly string[] => + Array.from( + new Set( + emails + .split(",") + .map((email) => email.trim()) + .filter((email) => email.length > 0), + ), + ); + +const THRESHOLD_MIN = 1; +const THRESHOLD_MAX = 100; + +export const isValidThreshold = (threshold: number | null): threshold is number => { + const isWholeNumber = threshold !== null && Number.isInteger(threshold); + return isWholeNumber && threshold >= THRESHOLD_MIN && threshold <= THRESHOLD_MAX; +}; + +export const teamMemberBudgetAlertRowsFromMetadata = (metadata: unknown): readonly TeamMemberBudgetAlertRow[] => { + if (typeof metadata !== "object" || metadata === null) return []; + const config: unknown = (metadata as Record)[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]; + if (typeof config !== "object" || config === null || Array.isArray(config)) return []; + return Object.entries(config as Record) + .flatMap(([key, emails]) => { + const threshold = Number(key); + return /^\d+$/.test(key) && isValidThreshold(threshold) && isEmailList(emails) + ? [{ threshold, emails: emails.join(", ") }] + : []; + }) + .sort((a, b) => (a.threshold ?? 0) - (b.threshold ?? 0)); +}; + +export const teamMemberBudgetAlertEmailsFromRows = ( + rows: readonly TeamMemberBudgetAlertRow[], +): TeamMemberBudgetAlertEmails => + Object.fromEntries( + rows + .filter((row) => isValidThreshold(row.threshold)) + .map((row) => [String(row.threshold), splitEmails(row.emails)]), + ); + +export const teamMemberBudgetAlertSummary = (metadata: unknown): readonly string[] => + teamMemberBudgetAlertRowsFromMetadata(metadata).map((row) => + row.emails.length > 0 ? `${row.threshold}%: member, ${row.emails}` : `${row.threshold}%: member`, + ); From 4eb340abe0d40543a564b89aa3db79c758649c6a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 20:24:58 -0700 Subject: [PATCH 47/58] fix(callbacks-legacy-python): traverse and release the retained headers dict (#43274) * fix(callbacks-legacy-python): traverse and release the retained headers dict LegacyLogging keeps the headers dict it hands to pre_call and post_call, but its traverse never reported that edge to the collector and close never dropped it. A cycle a callback builds through that dict could not be collected, and a closed call kept the dict alive until the driver dropped the whole adapter. Visit and clear headers like body, with regression tests for both * refactor(callbacks-legacy-python): move the test support module into its own file --------- Co-authored-by: Yujong Lee --- .../callbacks-legacy-python/src/adapter.rs | 144 ++++++++++++- .../crates/callbacks-legacy-python/src/lib.rs | 202 +----------------- .../src/test_support.rs | 199 +++++++++++++++++ 3 files changed, 333 insertions(+), 212 deletions(-) create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/test_support.rs diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 718cc615f30..fe0f6c7dd45 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -469,6 +469,7 @@ impl PythonLifecycle for LegacyLogging { error.write_unraisable(py, None); } self.body = None; + self.headers = None; self.context = None; self.stream = None; } @@ -486,7 +487,8 @@ impl PythonLifecycle for LegacyLogging { visit.call(&stream.chunks)?; visit.call(&stream.first_chunk)?; } - visit.call(&self.body) + visit.call(&self.body)?; + visit.call(&self.headers) } } @@ -743,6 +745,7 @@ mod payload_tests { use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; + use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -832,16 +835,7 @@ check = lambda: None headers: vec![("x-route".into(), "route".into())], body, }; - let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = MachineEvent::ResponseReceived { - raw: RawResponse { - body: "raw response".into(), - }, - }; - assert!(matches!( - logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), - LifecycleStep::Done - )); + let (_, step) = send_and_receive(py, &mut logging, wire, &context); run(py, &locals, c"check()"); let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); @@ -850,6 +844,134 @@ check = lambda: None }) } + /// `before_send` over `wire`, then the provider's raw response the way the driver + /// delivers it, so `pre_call` and `post_call` have both seen the retained payload. + fn send_and_receive<'a>( + py: Python<'_>, + logging: &'a mut LegacyLogging, + wire: WireRequest, + context: &RequestContext, + ) -> (&'a mut LegacyLogging, LifecycleStep) { + let step = logging.before_send(py, Box::new(wire), context).unwrap(); + let raw = MachineEvent::ResponseReceived { + raw: RawResponse { + body: "raw response".into(), + }, + }; + assert!(matches!( + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done + )); + (logging, step) + } + + fn route_context() -> RequestContext { + RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params: json!({}), + secret_fields: vec![], + api_key: Some(SecretValue::new("route-key")), + } + } + + fn route_wire() -> WireRequest { + WireRequest { + url: "https://provider.invalid/ocr".into(), + headers: vec![("x-route".into(), "route".into())], + body: json!({}), + } + } + + /// A Python object owning one `LegacyLogging`, so the interpreter's collector sees the + /// edges the adapter reports and clears them the way the driver's `Execution` does. + #[pyclass(weakref)] + struct Retained { + logging: Option, + } + + #[pymethods] + impl Retained { + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + match &self.logging { + Some(logging) => logging.traverse(&visit), + None => Ok(()), + } + } + + fn __clear__(slf: &Bound<'_, Self>) { + drop(slf.borrow_mut().logging.take()); + } + } + + #[test] + fn a_cycle_through_the_retained_headers_is_collected() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + send_and_receive(py, &mut logging, route_wire(), &route_context()); + let retained = Py::new( + py, + Retained { + logging: Some(logging), + }, + ) + .unwrap(); + locals.set_item("retained", retained).unwrap(); + run( + py, + &locals, + c" +import gc +import weakref + +logger.post[2]['headers']['owner'] = retained +logger.pre = logger.post = None +reference = weakref.ref(retained) +del retained +gc.collect() +assert reference() is None +", + ); + }); + } + + #[test] + fn close_releases_the_retained_headers() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + send_and_receive(py, &mut logging, route_wire(), &route_context()); + run( + py, + &locals, + c" +import weakref + +class Sentinel: + pass + +sentinel = Sentinel() +logger.post[2]['headers']['sentinel'] = sentinel +logger.pre = logger.post = None +reference = weakref.ref(sentinel) +del sentinel +assert reference() is not None +", + ); + logging.close(py); + run(py, &locals, c"assert reference() is None"); + }); + } + #[rstest] #[case::caller_keyword(c" document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 869c534acf4..8fca64d1b0e 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -23,204 +23,4 @@ pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; #[cfg(test)] -mod test_support { - use std::ffi::CStr; - - use pyo3::prelude::*; - use pyo3::types::{PyDict, PyTuple}; - - use crate::{LegacyLogging, LegacySurface, PublicCall}; - - /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. - /// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python - /// signatures, and [`namespace`] binds every fake call against it. - pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); - - /// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests - /// share one interpreter and run concurrently, so each fake is installed idempotently and - /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). - /// Every fake is bound against the contract first, so a call the real module would reject - /// fails here too. - const STUBS: &CStr = c" -import contextvars -import inspect -import json -import sys -import traceback -import types - -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): - sys.modules.setdefault(name, types.ModuleType(name)) - -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] -CONTRACT = json.loads(python_contract) - - -def contracted(name, fake): - signature = inspect.Signature( - [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] - ) - - def checked(*args, **kwargs): - signature.bind(*args, **kwargs) - return fake(*args, **kwargs) - - return checked - - -if not hasattr(legacy, 'is_internal'): - legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) - -FAKES = { - 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - ), - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), - 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=provider, - ), - 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( - original_response, api_key, additional_args - ), - 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), - 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( - response, start, end - ), - 'failure_handler': lambda logger, error, start, end, asynchronous: ( - logger.async_failure_handler if asynchronous else logger.failure_handler - )(error, ''.join(traceback.format_exception(error)), start, end), - 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), - 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), - 'enqueue_logging': lambda coroutine: coroutine.enqueue(), - 'restore_context': lambda logger: logger.record('restore', None), - 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), - 'is_internal_call': lambda: legacy.is_internal.get(), - 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), - 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( - 'success', response, call_type - ), - 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger: logger.record('stream_opened', None), - 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( - 'stream_success', list(chunks) - ), - 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), -} -assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) -for name, fake in FAKES.items(): - setattr(legacy, name, contracted(name, fake)) - - -unraisable = sys.modules.setdefault( - 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') -) -if not hasattr(unraisable, 'events'): - unraisable.events = [] - sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) - - -def unraisable_from(owner): - return [error for source, error in unraisable.events if source is owner] - - -class StubCoroutine: - def __init__(self, logger): - self.logger = logger - - def enqueue(self): - self.logger.record('enqueued', None) - self.logger.on_enqueue(self) - - def close(self): - self.logger.record('closed', None) - - -class StubLogger: - def __init__(self): - self.calls = [] - self.hooks = {} - self.on_enqueue = lambda coroutine: None - - def record(self, name, value): - self.calls.append((name, value)) - - def names(self): - return [name for name, _ in self.calls] - - def hook(self, phase, value, call_type): - self.record(phase + '_hook', call_type) - return self.hooks.get(phase, lambda value: 'awaitable')(value) - - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - - def async_failure_handler(self, error, trace, start, end): - self.record('async_failure_handler', error) - return 'awaitable' - - def success_handler(self, response, start, end): - self.record('success_handler', response) - - def async_success_handler(self, response, start, end): - self.record('async_success_handler', response) - return StubCoroutine(self) - - def handle_sync_success_callbacks_for_async_calls(self, response, start, end): - self.record('sync_success_for_async_call', response) - - -logger = StubLogger() -"; - - /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. - pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); - py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); - py.run(script, Some(&locals), Some(&locals)).unwrap(); - locals - } - - pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { - py.run(code, Some(locals), Some(locals)).unwrap(); - } - - pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - locals.get_item(name).unwrap().unwrap() - } - - /// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). - pub(crate) fn legacy_call( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - asynchronous: bool, - ) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); - LegacyLogging::new( - py, - LegacySurface { - call_type: "test", - input_description: "test input", - stream: None, - }, - call, - asynchronous, - ) - } -} +mod test_support; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs new file mode 100644 index 00000000000..e7973a7e1a0 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -0,0 +1,199 @@ +use std::ffi::CStr; + +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use crate::{LegacyLogging, LegacySurface, PublicCall}; + +/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. +/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. +const STUBS: &CStr = c" +import contextvars +import inspect +import json +import sys +import traceback +import types + +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): + sys.modules.setdefault(name, types.ModuleType(name)) + +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +CONTRACT = json.loads(python_contract) + + +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) + + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) + + +unraisable = sys.modules.setdefault( + 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') +) +if not hasattr(unraisable, 'events'): + unraisable.events = [] + sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) + + +def unraisable_from(owner): + return [error for source, error in unraisable.events if source is owner] + + +class StubCoroutine: + def __init__(self, logger): + self.logger = logger + + def enqueue(self): + self.logger.record('enqueued', None) + self.logger.on_enqueue(self) + + def close(self): + self.logger.record('closed', None) + + +class StubLogger: + def __init__(self): + self.calls = [] + self.hooks = {} + self.on_enqueue = lambda coroutine: None + + def record(self, name, value): + self.calls.append((name, value)) + + def names(self): + return [name for name, _ in self.calls] + + def hook(self, phase, value, call_type): + self.record(phase + '_hook', call_type) + return self.hooks.get(phase, lambda value: 'awaitable')(value) + + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + + def async_failure_handler(self, error, trace, start, end): + self.record('async_failure_handler', error) + return 'awaitable' + + def success_handler(self, response, start, end): + self.record('success_handler', response) + + def async_success_handler(self, response, start, end): + self.record('async_success_handler', response) + return StubCoroutine(self) + + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): + self.record('sync_success_for_async_call', response) + + +logger = StubLogger() +"; + +/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. +pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); + py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals +} + +pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { + py.run(code, Some(locals), Some(locals)).unwrap(); +} + +pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + locals.get_item(name).unwrap().unwrap() +} + +/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). +pub(crate) fn legacy_call( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, +) -> LegacyLogging { + let request = locals + .get_item("request") + .unwrap() + .unwrap_or_else(|| py.None().into_bound(py)); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .map(|kwargs| kwargs.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)); + let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + LegacyLogging::new( + py, + LegacySurface { + call_type: "test", + input_description: "test input", + stream: None, + }, + call, + asynchronous, + ) +} From e11c3f5815b54a3bf76ace48bfead4f2bf0f8da0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 21:54:11 -0700 Subject: [PATCH 48/58] test(integration): port langfuse callbacks-in-db coverage to the local harness (#43282) * test(integration): port langfuse callbacks-in-db coverage to the local harness Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the langfuse db rows after the exported span Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../observability/test_langfuse_delivery.py | 127 +++++++++++++++++- .../test_callbacks_in_db.py | 114 ---------------- 2 files changed, 121 insertions(+), 120 deletions(-) delete mode 100644 tests/store_model_in_db_tests/test_callbacks_in_db.py diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 5a3ffdb9965..9c0e3e407cb 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -2,24 +2,28 @@ import base64 import json import time import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Final import yaml -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows, write_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import KeyValue from opentelemetry.proto.trace.v1.trace_pb2 import Span -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, JsonValue, TypeAdapter PUBLIC_KEY: Final = "pk-lf-integration" SECRET_KEY: Final = "sk-lf-integration" PROJECTS_PATH: Final = "/api/public/projects" TRACES_PATH: Final = "/api/public/otel/v1/traces" PROMPTS_PATH: Final = "/api/public/v2/prompts/" +STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml") +CONFIG_SECTIONS: Final = ("litellm_settings", "environment_variables") +LANGFUSE_ENVIRONMENT: Final = ("LANGFUSE_HOST", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY") _PROXY_CONFIG: Final = TypeAdapter(dict[str, object]) _SETTINGS: Final = TypeAdapter(dict[str, object]) @@ -64,9 +68,7 @@ def _text_prompt(name: str) -> Reply: def _langfuse_config(tmp_path: Path) -> Path: - config: Final = _PROXY_CONFIG.validate_python( - yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - ) + config: Final = _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text())) settings: Final = { **_SETTINGS.validate_python(config["litellm_settings"]), "success_callback": ["langfuse"], @@ -86,6 +88,32 @@ def _langfuse_environment(langfuse: Wire) -> dict[str, str]: } +def _config_rows() -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT param_name, param_value FROM "LiteLLM_Config" WHERE param_name IN (%s, %s) ORDER BY param_name', + CONFIG_SECTIONS, + ) + + +def _restore_config_rows(snapshot: Sequence[Mapping[str, JsonValue]]) -> None: + saved: Final = {string_value(row["param_name"]): row["param_value"] for row in snapshot} + for section in CONFIG_SECTIONS: + if section not in saved: + write_rows('DELETE FROM "LiteLLM_Config" WHERE param_name = %s', (section,)) + elif saved[section] is None: + write_rows( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, NULL) ' + "ON CONFLICT (param_name) DO UPDATE SET param_value = NULL", + (section,), + ) + else: + write_rows( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s::jsonb) ' + "ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value", + (section, json.dumps(saved[section])), + ) + + def _attribute(entries: Sequence[KeyValue], key: str) -> str | list[str] | None: for entry in entries: if entry.key != key: @@ -192,6 +220,93 @@ def test_langfuse_callback_delivers_the_generation_over_otlp_v4_with_the_caller_ ) +def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_generation_over_otlp_v4( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "langfusedb" + uuid.uuid4().hex + provider_secret: Final = "synthetic-provider-secret-" + marker + public_key: Final = "pk-lf-db-" + marker + secret_key: Final = "sk-lf-db-" + marker + assert "langfuse" not in STOCK_CONFIG.read_text() + + def upstream(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {provider_secret}" + return _completion(marker + "-answer") + + def langfuse(request: Request) -> Reply: + if request.method == "GET" and request.target.startswith(PROJECTS_PATH): + return _projects() + return Reply(body=b"", content_type="application/x-protobuf") + + snapshot: Final = _config_rows() + try: + with ( + wire_server(upstream) as provider, + wire_server(langfuse) as destination, + owned_proxy( + gateway, + tmp_path, + {"LANGFUSE_FLUSH_INTERVAL": "1"}, + remove_environment=LANGFUSE_ENVIRONMENT, + ) as candidate, + candidate.scenario() as scenario, + ): + candidate.post( + "/config/update", + { + "litellm_settings": {"success_callback": ["langfuse"]}, + "environment_variables": { + "LANGFUSE_HOST": destination.url, + "LANGFUSE_PUBLIC_KEY": public_key, + "LANGFUSE_SECRET_KEY": secret_key, + }, + }, + ) + model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) + body: Final = candidate.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker + "-question"}], + "metadata": {"generation_name": marker}, + "cache": {"no-cache": True}, + }, + ) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones + + def exported() -> tuple[Span, ...]: + received.extend(destination.drain()) + return tuple(span for span in _spans(received) if span.name == marker) + + spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) + posts: Final = tuple(request for request in received if request.method == "POST") + assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] + basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() + for request in posts: + assert request.headers["authorization"] == basic + assert request.headers["content-type"] == "application/x-protobuf" + assert request.headers["x-langfuse-ingestion-version"] == "4" + assert provider_secret.encode() not in request.body + assert candidate.key.encode() not in request.body + + attributes: Final = spans[0].attributes + assert _attribute(attributes, "langfuse.observation.type") == "generation" + assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) + assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) + assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) + + stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows()} + callbacks: Final = TypeAdapter(list[str]).validate_python( + object_value(stored["litellm_settings"]).get("success_callback") or [] + ) + assert "langfuse" in callbacks, stored + assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored + assert secret_key not in json.dumps(stored["environment_variables"]), stored + finally: + _restore_config_rows(snapshot) + assert _config_rows() == snapshot + + def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_headers_off_the_client( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/store_model_in_db_tests/test_callbacks_in_db.py b/tests/store_model_in_db_tests/test_callbacks_in_db.py deleted file mode 100644 index 6497e4064b7..00000000000 --- a/tests/store_model_in_db_tests/test_callbacks_in_db.py +++ /dev/null @@ -1,114 +0,0 @@ -""" -PROD TEST - DO NOT Delete this Test - -e2e test for langfuse callback in DB -- Add langfuse callback to DB - with /config/update -- wait 20 seconds for the callback to be loaded into the instance -- Make a /chat/completions request to the proxy -- Check if the request is logged in Langfuse -""" - -import pytest -import asyncio -import aiohttp -import os -import dotenv -from dotenv import load_dotenv -from openai import AsyncOpenAI, APIConnectionError -from openai.types.chat import ChatCompletion - -load_dotenv() - -# used for testing -LANGFUSE_BASE_URL = "https://exampleopenaiendpoint-production-c715.up.railway.app" -PROXY_BASE_URL = "http://127.0.0.1:4000" - - -async def wait_for_proxy_ready(session, timeout: int = 60): - for _ in range(timeout): - try: - async with session.get(f"{PROXY_BASE_URL}/health/liveliness") as response: - if response.status == 200: - return - except aiohttp.ClientError: - pass - await asyncio.sleep(1) - raise RuntimeError(f"Proxy at {PROXY_BASE_URL} not ready after {timeout}s") - - -async def config_update(session, routing_strategy=None): - url = f"{PROXY_BASE_URL}/config/update" - headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} - print("routing_strategy: ", routing_strategy) - data = { - "litellm_settings": {"success_callback": ["langfuse"]}, - "environment_variables": { - "LANGFUSE_PUBLIC_KEY": "any-public-key", - "LANGFUSE_SECRET_KEY": "any-secret-key", - "LANGFUSE_HOST": LANGFUSE_BASE_URL, - }, - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print("status: ", status) - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - return await response.json() - - -async def check_langfuse_request(response_id: str): - async with aiohttp.ClientSession() as session: - url = f"{LANGFUSE_BASE_URL}/langfuse/trace/{response_id}" - async with session.get(url) as response: - response_json = await response.json() - assert response.status == 200, f"Expected status 200, got {response.status}" - assert ( - response_json["exists"] == True - ), f"Request {response_id} not found in Langfuse traces" - assert response_json["request_id"] == response_id, f"Request ID mismatch" - - -async def make_chat_completions_request() -> ChatCompletion: - client = AsyncOpenAI(api_key="sk-1234", base_url=PROXY_BASE_URL) - last_error = None - for _ in range(10): - try: - response = await client.chat.completions.create( - model="fake-openai-endpoint", - messages=[{"role": "user", "content": "Hello, world!"}], - ) - print(response) - return response - except APIConnectionError as e: - last_error = e - await asyncio.sleep(2) - raise AssertionError( - f"Proxy at {PROXY_BASE_URL} unreachable after retries: {last_error!r}" - ) - - -@pytest.mark.asyncio -async def test_e2e_langfuse_callbacks_in_db(): - - async with aiohttp.ClientSession() as session: - # add langfuse callback to DB - await config_update(session) - - # wait 20 seconds for the callback to be loaded into the instance - await asyncio.sleep(20) - await wait_for_proxy_ready(session) - - # make a /chat/completions request to the proxy - response = await make_chat_completions_request() - print(response) - response_id = response.id - print("response_id: ", response_id) - - await asyncio.sleep(20) - # check if the request is logged in Langfuse - await check_langfuse_request(response_id) From 3ebf6a1fd8837d479ea7bef7728be9bb8bf84cd3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 22:40:45 -0700 Subject: [PATCH 49/58] test(e2e): accept the otel cost write as a linked root trace (#42931) * test(e2e): accept the otel cost write as a linked root trace Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): window and fail-closed the linked otel trace read-back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): poll the otel read-back without recursion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/logging/test_otel_trace_e2e.py | 191 +++++++++++------------ tests/e2e/otel_client.py | 161 +++++++++++++++---- 2 files changed, 222 insertions(+), 130 deletions(-) diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index 8d154ca0837..8d8595cf221 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -2,8 +2,9 @@ Covers logging.otel.success.exports_metric: a successful non-streaming call must land at the OTEL destination as ONE connected trace - a single root SERVER span -with the auth phase, db lookups, and cost write under it, and the gen-AI CLIENT -span parented into the same tree. The regression this pins: the proxy publishing +with the auth phase and db lookups under it, the gen-AI CLIENT span parented +into the same tree, and the cost write either under it or as the root of its +own trace linked back to the request span. The regression this pins: the proxy publishing the global TracerProvider before callbacks init made server spans export through a different provider than the preset's gen-AI spans, so the destination received the gen-AI span alone, dangling (fixed in #30590; verified failing at its parent @@ -28,7 +29,7 @@ from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, OTEL_EXPORTER_ from lifecycle import ResourceManager from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body from models import LiteLLMParamsBody -from otel_client import JaegerSpan, JaegerTrace, OtelReader +from otel_client import CallTraces, JaegerSpan, JaegerTrace, OtelReader, root_span from pydantic import BaseModel, ConfigDict, ValidationError pytestmark = pytest.mark.e2e @@ -78,12 +79,13 @@ def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool: return False -def _assert_complete_trace( - hits: list[JaegerTrace], *, route: str, genai_span: str, require_cost_span: bool = True -) -> None: - """The enforced behavior: the destination holds exactly one trace for the - call, rooted at the SERVER span, with auth/db/cost children and the gen-AI - span all connected into that one tree - no dangling parent references.""" +def _assert_complete_trace(traces: CallTraces, *, route: str, genai_span: str, require_cost_span: bool = True) -> None: + """The enforced behavior: the destination holds exactly one call-id-tagged + trace for the call, rooted at the SERVER span, with auth/db children and + the gen-AI span all connected into that one tree - no dangling parent + references - and the cost write either in that trace or as the root of + its own trace linked FOLLOWS_FROM to the request SERVER span.""" + hits = traces.hits assert hits, ( "no trace for this call arrived at the destination within the deadline " "(nothing tagged with its call id was found)" @@ -119,8 +121,20 @@ def _assert_complete_trace( assert any(name.startswith(DB_SPAN_PREFIX) for name in names), ( f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}" ) - if require_cost_span: - assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}" + if require_cost_span and COST_SPAN not in names: + cost_traces = [t for t in traces.linked if (r := root_span(t)) is not None and r.operation_name == COST_SPAN] + assert len(cost_traces) == 1, ( + f"cost write span {COST_SPAN!r} reached neither the request trace nor its own " + f"trace linked to the request SERVER span; request spans: {names}; " + f"linked traces: {[(t.trace_id, t.span_names()) for t in traces.linked]}" + ) + cost_root = root_span(cost_traces[0]) + assert cost_root is not None, f"cost write trace has no single root; spans: {cost_traces[0].span_names()}" + link = next(ref for ref in cost_root.references if ref.span_id == root.span_id) + assert link.ref_type == "FOLLOWS_FROM" and link.trace_id == trace.trace_id, ( + f"the cost write trace's root must reference the request SERVER span FOLLOWS_FROM, " + f"got refType={link.ref_type!r} traceID={link.trace_id!r} (request trace {trace.trace_id})" + ) genai = next((span for span in trace.spans if span.operation_name == genai_span), None) assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}" @@ -131,9 +145,15 @@ def _assert_complete_trace( ) -def _settled_names(*, route: str, genai_span: str, require_cost_span: bool = True) -> set[str]: - names = {f"POST {route}", f"auth {route}", genai_span} - return (names | {COST_SPAN}) if require_cost_span else names +def _poll( + otel_reader: OtelReader, *, call_id: str, route: str, genai_span: str, require_cost_span: bool = True +) -> CallTraces: + return otel_reader.poll_traces_for_call( + call_id=call_id, + settled_names={f"POST {route}", f"auth {route}", genai_span}, + settled_prefixes={DB_SPAN_PREFIX}, + linked_names=frozenset({COST_SPAN}) if require_cost_span else frozenset(), + ) def _tag(span: JaegerSpan, key: str) -> str | int | float | bool | None: @@ -174,7 +194,7 @@ def one_served_genai_span(trace: JaegerTrace, genai_span: str) -> JaegerSpan: return served[0] -def _assert_real_ttft(hits: list[JaegerTrace], *, genai_span: str) -> None: +def _assert_real_ttft(hits: tuple[JaegerTrace, ...], *, genai_span: str) -> None: """The enforced behavior: the gen-AI span for the attempt that served the stream records a TTFT that is a real measurement - present, numeric, positive, and strictly less than that span's own total duration. A TTFT of @@ -286,9 +306,10 @@ class TestOtelTraceCompleteness: /chat/completions request produces one complete OTEL trace. The trace should have a single server root span for the incoming request, with - the authentication, database, and cost-recording work beneath it. The span for - the actual model call must also belong to that same trace, rather than being - exported separately with a missing parent. + the authentication and database work beneath it. The span for the actual model + call must also belong to that same trace, rather than being exported separately + with a missing parent, and the cost-recording work must land either in that + trace or in its own trace linked to it. This matters because a split trace is easy to miss: all of the spans may still arrive, but the model call appears without the surrounding request context. @@ -308,12 +329,8 @@ class TestOtelTraceCompleteness: outcome = first_ok(client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16)) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}") + _assert_complete_trace(traces, route=route, genai_span=f"chat {MODEL}") @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["chat_completions"]) @pytest.mark.otel_tls @@ -337,11 +354,7 @@ class TestOtelTraceCompleteness: ) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" - hits: Final = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), - settled_prefixes={DB_SPAN_PREFIX}, - ) + hits: Final = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}") _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["messages"]) @@ -352,8 +365,10 @@ class TestOtelTraceCompleteness: produces exactly one complete OTEL trace. The trace must have a single root span named "POST /v1/messages". The - authentication, database, cost-writing, and model-call spans must all belong to + authentication, database, and model-call spans must all belong to the same trace and have valid parent relationships leading back to that root. + The cost-writing span must land in the request trace or in its own trace + linked to it. The model-call span is expected to be named "chat ". The test fails if the request is split across multiple traces, if any span references a missing @@ -370,12 +385,8 @@ class TestOtelTraceCompleteness: ) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}") + _assert_complete_trace(traces, route=route, genai_span=f"chat {MODEL}") @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["responses"]) def test_responses_exports_complete_trace( @@ -385,12 +396,14 @@ class TestOtelTraceCompleteness: produces exactly one complete OTEL trace. The trace must have a single root span named "POST /v1/responses". The - authentication, database, cost-writing, and model-call spans must all belong to - the same trace and have valid parent relationships leading back to that root. + authentication, database, and model-call spans must all belong to the same + trace and have valid parent relationships leading back to that root. The cost + write finishes after the response, so it lands as the root of its own trace + linked FOLLOWS_FROM to the request SERVER span. - The model-call span is expected to be named "chat ". The test fails if - the request is split across multiple traces, if any span references a missing - parent, or if the model-call span cannot be connected back to the root.""" + The model-call span is expected to be named "chat ". The test fails on + a split request trace, a dangling parent, a disconnected model-call span, or + a cost write that is neither in the request trace nor linked to it.""" route = "/v1/responses" _assert_otel_destination_configured(client) @@ -405,12 +418,8 @@ class TestOtelTraceCompleteness: assert outcome.call_id is not None, "success response must carry x-litellm-call-id" genai_span = f"chat {CHEAP_OPENAI_MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_complete_trace(hits, route=route, genai_span=genai_span) + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span) + _assert_complete_trace(traces, route=route, genai_span=genai_span) @pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["chat_completions"]) def test_chat_completions_stream_exports_complete_trace( @@ -418,8 +427,9 @@ class TestOtelTraceCompleteness: ) -> None: """A successful streamed `/chat/completions` request should export one complete OTEL trace. The trace must contain a single root `SERVER` - span, with the auth, database, cost, and gen-AI `CLIENT` spans all - connected back to that root. + span, with the auth, database, and gen-AI `CLIENT` spans all + connected back to that root, and the cost write in that trace or in + its own trace linked to it. Streaming has an additional lifecycle risk because the gen-AI span is closed by the stream-consumption path after the final chunk has @@ -451,14 +461,10 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_complete_trace(hits, route=route, genai_span=genai_span) + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span) + _assert_complete_trace(traces, route=route, genai_span=genai_span) - served = one_served_genai_span(hits[0], genai_span) + served = one_served_genai_span(traces.hits[0], genai_span) assert _tag(served, "litellm.request.streaming") is True, ( "the gen-AI span must record litellm.request.streaming=true; its absence means " "the stream flag was dropped before the model call" @@ -470,8 +476,9 @@ class TestOtelTraceCompleteness: ) -> None: """A successful streamed `/v1/messages` request should export one complete OTEL trace. The trace must contain a single root `SERVER` - span, with the auth, database, cost, and gen-AI `CLIENT` spans all - connected back to that root. + span, with the auth, database, and gen-AI `CLIENT` spans all + connected back to that root, and the cost write in that trace or in + its own trace linked to it. This endpoint has the same streaming lifecycle risk as `/chat/completions`: the gen-AI span is closed by the @@ -503,14 +510,10 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_complete_trace(hits, route=route, genai_span=genai_span) + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span) + _assert_complete_trace(traces, route=route, genai_span=genai_span) - served = one_served_genai_span(hits[0], genai_span) + served = one_served_genai_span(traces.hits[0], genai_span) assert _tag(served, "litellm.request.streaming") is True, ( "the gen-AI span must record litellm.request.streaming=true; its absence means " "the stream flag was dropped before the model call" @@ -555,14 +558,12 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {CHEAP_OPENAI_MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), - settled_prefixes={DB_SPAN_PREFIX}, + traces = _poll( + otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False ) - _assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False) + _assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False) - one_served_genai_span(hits[0], genai_span) + one_served_genai_span(traces.hits[0], genai_span) spend_row = client.poll_proxy_spend_for_key(key) assert spend_row is not None and spend_row.spend is not None and spend_row.spend > 0, ( @@ -609,12 +610,8 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_real_ttft(hits, genai_span=genai_span) + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span) + _assert_real_ttft(traces.hits, genai_span=genai_span) @pytest.mark.covers("logging.otel.stream.records_ttft", exercised_on=["messages"]) def test_messages_stream_records_real_ttft( @@ -651,12 +648,8 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span), - settled_prefixes={DB_SPAN_PREFIX}, - ) - _assert_real_ttft(hits, genai_span=genai_span) + traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span) + _assert_real_ttft(traces.hits, genai_span=genai_span) @pytest.mark.covers("logging.otel.stream.records_ttft", exercised_on=["responses"]) def test_responses_stream_records_real_ttft( @@ -693,12 +686,10 @@ class TestOtelTraceCompleteness: ) genai_span = f"chat {CHEAP_OPENAI_MODEL}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), - settled_prefixes={DB_SPAN_PREFIX}, + traces = _poll( + otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False ) - _assert_real_ttft(hits, genai_span=genai_span) + _assert_real_ttft(traces.hits, genai_span=genai_span) @pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["chat_completions"]) def test_failed_chat_completions_error_span_attributes( @@ -745,18 +736,16 @@ class TestOtelTraceCompleteness: assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id" genai_span = f"chat {model_name}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), - settled_prefixes={DB_SPAN_PREFIX}, + traces = _poll( + otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False ) - _assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False) + _assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False) - root = next(span for span in hits[0].spans if not span.references) + root = next(span for span in traces.hits[0].spans if not span.references) assert str(_tag(root, "http.status_code")) == "401", ( f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}" ) - genai = next(span for span in hits[0].spans if span.operation_name == genai_span) + genai = next(span for span in traces.hits[0].spans if span.operation_name == genai_span) _assert_error_span_contract(genai) @pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["messages"]) @@ -803,16 +792,14 @@ class TestOtelTraceCompleteness: assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id" genai_span = f"chat {model_name}" - hits = otel_reader.poll_traces_for_call( - call_id=outcome.call_id, - settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), - settled_prefixes={DB_SPAN_PREFIX}, + traces = _poll( + otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False ) - _assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False) + _assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False) - root = next(span for span in hits[0].spans if not span.references) + root = next(span for span in traces.hits[0].spans if not span.references) assert str(_tag(root, "http.status_code")) == "401", ( f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}" ) - genai = next(span for span in hits[0].spans if span.operation_name == genai_span) + genai = next(span for span in traces.hits[0].spans if span.operation_name == genai_span) _assert_error_span_contract(genai) diff --git a/tests/e2e/otel_client.py b/tests/e2e/otel_client.py index b11fddebc9c..c5d709048b3 100644 --- a/tests/e2e/otel_client.py +++ b/tests/e2e/otel_client.py @@ -10,6 +10,11 @@ so the completeness assertions see the whole tree. A failed query is a hard failure, never an empty result - an unreachable destination must not read as "the trace never arrived". +Service spans that end after the response (the cost write is one) carry no +call id and land as the root of their own trace with a link back to the +request span, so they are fetched by operation name and matched by that link +to the request root rather than by the tag query. + External reads go through ``e2e_http`` (the only module allowed to call ``requests.*``). """ @@ -18,7 +23,9 @@ from __future__ import annotations import json import time +from collections.abc import Iterator from dataclasses import dataclass +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -84,18 +91,57 @@ class JaegerTracesPage(BaseModel): class _TracesQuery(BaseModel): service: str - tags: str + tags: str | None = None + operation: str | None = None limit: int = 20 lookback: str = "1h" + start: int | None = None + end: int | None = None + + +def _ticks() -> Iterator[None]: + while True: + yield None + time.sleep(POLL_INTERVAL) def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool: present = set(trace.span_names()) - return names.issubset(present) and all( - any(name.startswith(prefix) for name in present) for prefix in prefixes + return names.issubset(present) and all(any(name.startswith(prefix) for name in present) for prefix in prefixes) + + +def root_span(trace: JaegerTrace) -> JaegerSpan | None: + """The single span whose references all point outside the trace (a span + with no references qualifies). None when there is not exactly one.""" + in_trace = {span.span_id for span in trace.spans} + roots = [span for span in trace.spans if all(ref.span_id not in in_trace for ref in span.references)] + return roots[0] if len(roots) == 1 else None + + +def _follows(trace: JaegerTrace, parent_trace_id: str, parent_span_id: str) -> bool: + root = root_span(trace) + return root is not None and any( + ref.trace_id == parent_trace_id and ref.span_id == parent_span_id for ref in root.references ) +@dataclass(frozen=True, slots=True) +class CallTraces: + hits: tuple[JaegerTrace, ...] + linked: tuple[JaegerTrace, ...] + + +@dataclass(frozen=True, slots=True) +class _Observation: + traces: CallTraces + missing: tuple[str, ...] + unreachable: NetworkError | None + + def settled(self, names: set[str], prefixes: set[str]) -> bool: + hits: Final = self.traces.hits + return self.unreachable is None and len(hits) == 1 and not self.missing and _settled(hits[0], names, prefixes) + + @dataclass(frozen=True, slots=True) class OtelReader: query_url: str @@ -119,36 +165,95 @@ class OtelReader: case failure: pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") + def _query_operation(self, operation: str, *, start: int) -> Result[JaegerTracesPage]: + return get( + URL(f"{self.query_url}/api/traces"), + headers=NoBody(), + params=_TracesQuery( + service=JAEGER_SERVICE, + operation=operation, + limit=200, + start=start, + end=int(time.time() * 1_000_000), + ), + response_type=JaegerTracesPage, + timeout=30.0, + ) + + def linked_traces(self, *, operation: str, parent: JaegerTrace) -> tuple[JaegerTrace, ...] | NetworkError: + """Traces whose root span references the parent trace's root span. + Detached post-response work lands as the root of its own trace with a + link back to the request span instead of the call-id tag, so it is + found by operation name, windowed to start at the parent root's start + time (the detached span always starts after it), and matched on that + link. A NetworkError is handed back so the polling caller can tell an + unreachable read-back endpoint from a span that never arrived.""" + parent_root: Final = root_span(parent) + if parent_root is None: + return () + match self._query_operation(operation, start=parent_root.start_time): + case Success(data=page): + return tuple(t for t in page.data if _follows(t, parent.trace_id, parent_root.span_id)) + case NetworkError() as failure: + return failure + case failure: + pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") + def poll_traces_for_call( - self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str] - ) -> list[JaegerTrace]: - """Poll until exactly one trace holds the call and it carries every span + self, + *, + call_id: str, + settled_names: set[str], + settled_prefixes: set[str], + linked_names: frozenset[str] = frozenset(), + ) -> CallTraces: + """Poll until exactly one trace holds the call, it carries every span name in ``settled_names`` plus at least one name per prefix in - ``settled_prefixes`` (spans flush in batches, the cost write lands after - the response), then return the hits. At the deadline the last hits are - returned as-is so the caller's assertions report the real final state - - on a split trace this never settles and the orphan comes back.""" - deadline = time.monotonic() + POLL_TIMEOUT - hits: list[JaegerTrace] = [] - unreachable: NetworkError | None = None - while time.monotonic() < deadline: - match self._query_traces(call_id): - case Success(data=page): - unreachable = None - hits = page.data - if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes): - return hits - case NetworkError() as failure: - unreachable = failure - case failure: - pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") - time.sleep(POLL_INTERVAL) - if unreachable is not None: + ``settled_prefixes``, and every name in ``linked_names`` is either in + that trace or is the root of its own trace referencing the request + root (post-response work detaches per #42826). At the deadline the + last observed state is returned as-is so the caller's assertions + report the real final state - on a split trace this never settles and + the orphan comes back. A read-back endpoint still failing at the + deadline (either query) is a hard failure, not a missing span.""" + deadline: Final = time.monotonic() + POLL_TIMEOUT + last: Final = self._poll(call_id, settled_names, settled_prefixes, linked_names, deadline) + if last.unreachable is not None: pytest.fail( f"Jaeger query API at {self.query_url} stayed unreachable until the " - f"{POLL_TIMEOUT}s poll deadline: {unreachable}" + f"{POLL_TIMEOUT}s poll deadline: {last.unreachable}" ) - return hits + return last.traces + + def _observe(self, call_id: str, linked_names: frozenset[str]) -> _Observation: + match self._query_traces(call_id): + case NetworkError() as failure: + return _Observation(CallTraces((), ()), tuple(linked_names), failure) + case Success(data=page): + if len(page.data) != 1: + return _Observation(CallTraces(tuple(page.data), ()), tuple(linked_names), None) + hit: Final = page.data[0] + present: Final = frozenset(hit.span_names()) + results: Final = { + name: self.linked_traces(operation=name, parent=hit) for name in linked_names if name not in present + } + unreachable: Final = next((r for r in results.values() if isinstance(r, NetworkError)), None) + linked: Final = tuple(t for r in results.values() if not isinstance(r, NetworkError) for t in r) + missing: Final = tuple(name for name, r in results.items() if isinstance(r, NetworkError) or not r) + return _Observation(CallTraces((hit,), linked), missing, unreachable) + case failure: + pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") + + def _poll( + self, + call_id: str, + names: set[str], + prefixes: set[str], + linked_names: frozenset[str], + deadline: float, + ) -> _Observation: + observations: Final = (self._observe(call_id, linked_names) for _ in _ticks()) + return next(o for o in observations if o.settled(names, prefixes) or time.monotonic() >= deadline) def build_otel_reader() -> OtelReader: From 99655b6f86d0a5584cf4f11ccf56b5a91fe93797 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 22:43:41 -0700 Subject: [PATCH 50/58] test: finish the non-proxy half of tests/test_litellm (#43281) * test: move key-gated tests/test_litellm SDK tests into tests/llm_translation and drop empty folders * test: make token counter and health check unit tests run offline * ci: point unit shards, rust path filter, Makefile and docs at tests/unit * docs: fix stale test_litellm run paths in moved llm_translation tests * fix: correct databricks e2e sys.path depth and contributing example path --------- Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/pull_request_template.md | 2 +- .github/workflows/test-rust.yml | 2 - .github/workflows/test-unit.yml | 7 +- AGENTS.md | 2 +- ARCHITECTURE.md | 2 +- CONTRIBUTING.md | 14 +-- Makefile | 6 +- litellm/containers/README.md | 2 +- tests/README.MD | 2 +- tests/integration/sandbox/test_e2b_sandbox.py | 3 +- .../databricks_config.template.txt | 0 .../interactions/base_interactions_test.py | 0 .../interactions/test_gemini_interactions.py | 2 +- .../test_google_interactions_integration.py | 2 +- .../test_litellm_responses_bridge.py | 2 +- .../test_cometapi_chat_transformation.py | 0 .../test_compression.py | 0 .../test_databricks_e2e.py | 6 +- .../test_json_providers.py | 0 ...tral_audio_transcription_transformation.py | 0 ...loud_audio_transcription_transformation.py | 0 .../test_ovhcloud_chat_transformation.py | 0 ...rtex_ai_image_generation_transformation.py | 0 .../test_xiaomi_mimo.py | 0 .../test_handler_gc_does_not_close_client.py | 5 +- tests/test_litellm/llms/mistral/__init__.py | 0 .../test_litellm/llms/openai_like/__init__.py | 0 tests/test_litellm/llms/vertex_ai/__init__.py | 1 - tests/test_litellm/log.txt | 2 - tests/test_litellm/ocr/__init__.py | 0 tests/test_litellm/passthrough/__init__.py | 0 tests/unit/AGENTS.md | 2 +- .../dotprompt/test_prompt_manager.py | 16 +-- .../test_health_check_helpers.py | 48 ++++---- .../litellm_core_utils/test_token_counter.py | 107 +++++++++++------- .../llms/anthropic/batches/test_handler.py | 3 +- .../test_anthropic_output_format_filter.py | 4 +- 37 files changed, 125 insertions(+), 117 deletions(-) rename tests/{test_litellm/llms/databricks => llm_translation}/databricks_config.template.txt (100%) rename tests/{test_litellm => llm_translation}/interactions/base_interactions_test.py (100%) rename tests/{test_litellm => llm_translation}/interactions/test_gemini_interactions.py (88%) rename tests/{test_litellm => llm_translation}/interactions/test_google_interactions_integration.py (99%) rename tests/{test_litellm => llm_translation}/interactions/test_litellm_responses_bridge.py (91%) rename tests/{test_litellm/llms/cometapi/chat => llm_translation}/test_cometapi_chat_transformation.py (100%) rename tests/{test_litellm => llm_translation}/test_compression.py (100%) rename tests/{test_litellm/llms/databricks => llm_translation}/test_databricks_e2e.py (99%) rename tests/{test_litellm/llms/openai_like => llm_translation}/test_json_providers.py (100%) rename tests/{test_litellm/llms/mistral/audio_transcription => llm_translation}/test_mistral_audio_transcription_transformation.py (100%) rename tests/{test_litellm/llms/ovhcloud => llm_translation}/test_ovhcloud_audio_transcription_transformation.py (100%) rename tests/{test_litellm/llms/ovhcloud => llm_translation}/test_ovhcloud_chat_transformation.py (100%) rename tests/{test_litellm/llms/vertex_ai/image_generation => llm_translation}/test_vertex_ai_image_generation_transformation.py (100%) rename tests/{test_litellm/llms/openai_like => llm_translation}/test_xiaomi_mimo.py (100%) delete mode 100644 tests/test_litellm/llms/mistral/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/__init__.py delete mode 100644 tests/test_litellm/llms/vertex_ai/__init__.py delete mode 100644 tests/test_litellm/log.txt delete mode 100644 tests/test_litellm/ocr/__init__.py delete mode 100644 tests/test_litellm/passthrough/__init__.py diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 1fe0c602036..db46114715d 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -65,7 +65,7 @@ After: the same request comes back with real token counts, so the dashboard show **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests -- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more +- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/unit/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more - [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 808bb2afd08..2d399cca3a4 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -14,7 +14,6 @@ on: - "litellm/ocr/**" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - - "tests/test_litellm/ocr/**" - "tests/test_litellm/conftest.py" - "Makefile" - ".cargo/**" @@ -42,7 +41,6 @@ on: - "litellm/ocr/**" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - - "tests/test_litellm/ocr/**" - "tests/test_litellm/conftest.py" - "Makefile" - ".cargo/**" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2212b276b0d..f55e186e3b2 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -88,7 +88,7 @@ jobs: - shard: Vertex AI artifact-name: llm-vertex-ai - test-path: "tests/test_litellm/llms/vertex_ai" + test-path: "" unit-flag: llm-vertex-ai workers: 1 reruns: 2 @@ -97,7 +97,7 @@ jobs: - shard: All Other Providers artifact-name: llm-other-providers - test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai" + test-path: "" unit-flag: llm-other-providers workers: 2 reruns: 2 @@ -107,9 +107,6 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/test_litellm/interactions - tests/test_litellm/ocr - tests/test_litellm/passthrough tests/test_litellm/test_*.py unit-flag: misc workers: 2 diff --git a/AGENTS.md b/AGENTS.md index 69e034fbdea..a2dcd24bdd1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -27,7 +27,7 @@ Never test structure of code only function of it A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken -`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones +`tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md` diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index f418752d990..c9e046748e8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -255,7 +255,7 @@ Conventions to follow when touching this layer: | Column vs. field names | Where a model field differs from its DB column (for example `org_id` maps to the `organization_id` column), the repository translates in both directions rather than relying on Pydantic to guess. | | Array mutations | Adds use Prisma's atomic `push` (`add_member`, `add_admin`, `add_models`) to avoid read-modify-write races. Removals fall back to read-modify-write because Prisma has no atomic array remove. | -To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/test_litellm/repositories/`. +To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/unit/repositories/`. --- diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 082b7a8fb3e..a5ad6e97f3d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -14,7 +14,7 @@ Here are the core requirements for any PR submitted to LiteLLM: - [ ] **Add testing** - Adding at least 1 test is a hard requirement - [see details](#adding-testing) - [ ] **Ensure your PR passes all checks**: - [ ] [Linting / Formatting](#running-linting-and-formatting-checks) - `make lint` - - [ ] [The tests covering your change](#running-unit-tests) pass, e.g. `uv run pytest tests/test_litellm/.py -v`. CI runs the full unit test matrix, so you don't need to run the whole suite locally + - [ ] [The tests covering your change](#running-unit-tests) pass, e.g. `uv run pytest tests/unit/.py -v`. CI runs the full unit test matrix, so you don't need to run the whole suite locally #### UI PRs @@ -72,7 +72,7 @@ make format make lint # Run the tests covering your change (CI runs the full suite) -uv run pytest tests/test_litellm/.py -v +uv run pytest tests/unit/.py -v # Commit your changes (must follow Conventional Commits — see above) git add . @@ -88,7 +88,7 @@ git push origin feature/your-feature ### Where to Add Tests -Add your tests to the [`tests/test_litellm/` directory](https://github.com/BerriAI/litellm/tree/main/tests/test_litellm). +Add your tests to the [`tests/unit/` directory](https://github.com/BerriAI/litellm/tree/main/tests/unit). - This directory mirrors the structure of the `litellm/` directory - **Only add mocked tests** - no real LLM API calls in this directory @@ -96,10 +96,10 @@ Add your tests to the [`tests/test_litellm/` directory](https://github.com/Berri ### File Naming Convention -The `tests/test_litellm/` directory follows the same structure as `litellm/`: +The `tests/unit/` directory follows the same structure as `litellm/`: - `litellm/proxy/caching_routes.py` → `tests/test_litellm/proxy/test_caching_routes.py` -- `litellm/utils.py` → `tests/test_litellm/test_utils.py` +- `litellm/utils.py` → `tests/unit/test_utils.py` ### Example Test @@ -125,10 +125,10 @@ def test_your_feature(): Run the tests covering your change: ```bash -uv run pytest tests/test_litellm/test_your_file.py -v +uv run pytest tests/unit/test_your_file.py -v ``` -`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets) on beefier boxes, so if, for whatever reason, you must run the whole suite, it's better to rely on CI to do that. +`tests/unit` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets) on beefier boxes, so if, for whatever reason, you must run the whole suite, it's better to rely on CI to do that. If you're running broader test suites, proxy tests, or anything that touches PostgreSQL-backed fixtures/plugins, install the full local test environment first: diff --git a/Makefile b/Makefile index 311a7daef92..79c18f6fe82 100644 --- a/Makefile +++ b/Makefile @@ -42,7 +42,7 @@ help: @echo " make check-circular-imports - Check for circular imports" @echo " make check-import-safety - Check import safety" @echo " make test - Run all tests" - @echo " make test-unit - Run unit tests (tests/test_litellm)" + @echo " make test-unit - Run unit tests (tests/unit and tests/test_litellm)" @echo " make test-unit-llms - Run LLM provider tests (~225 files)" @echo " make test-unit-proxy-guardrails - Run proxy guardrails+mgmt tests (~51 files)" @echo " make test-unit-proxy-core - Run proxy auth+client+db+hooks tests (~52 files)" @@ -310,7 +310,7 @@ test: install-test-deps $(UV_RUN) pytest tests/ test-unit: install-test-deps - $(UV_RUN) pytest tests/test_litellm -x -vv -n 4 + $(UV_RUN) pytest tests/unit tests/test_litellm -x -vv -n 4 # Matrix test targets (matching CI workflow groups) test-unit-llms: install-test-deps @@ -332,7 +332,7 @@ test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/unit/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/litellm/containers/README.md b/litellm/containers/README.md index b54f96b1132..571bcaf415c 100644 --- a/litellm/containers/README.md +++ b/litellm/containers/README.md @@ -213,7 +213,7 @@ Run the container API tests: ```bash cd /Users/ishaanjaffer/github/litellm -python -m pytest tests/test_litellm/containers/ -v +python -m pytest tests/unit/containers/ -v ``` Test via proxy: diff --git a/tests/README.MD b/tests/README.MD index 57275a031f7..6a5da137203 100644 --- a/tests/README.MD +++ b/tests/README.MD @@ -4,6 +4,6 @@ To make it easier to contribute and map what behavior is tested, -we've started mapping the litellm directory in `tests/test_litellm` +we've started mapping the litellm directory in `tests/unit` This folder can only run mock tests. diff --git a/tests/integration/sandbox/test_e2b_sandbox.py b/tests/integration/sandbox/test_e2b_sandbox.py index d1cff2ce178..5adfb99db61 100644 --- a/tests/integration/sandbox/test_e2b_sandbox.py +++ b/tests/integration/sandbox/test_e2b_sandbox.py @@ -2,8 +2,7 @@ e2b code execution sandbox - end-to-end integration tests. These tests make REAL HTTP calls to the e2b API and are skipped automatically -unless E2B_API_KEY is set. Mock-only unit tests live in -tests/test_litellm/sandbox/test_e2b_sandbox.py. +unless E2B_API_KEY is set. Run only these tests: pytest tests/integration/sandbox/test_e2b_sandbox.py -v diff --git a/tests/test_litellm/llms/databricks/databricks_config.template.txt b/tests/llm_translation/databricks_config.template.txt similarity index 100% rename from tests/test_litellm/llms/databricks/databricks_config.template.txt rename to tests/llm_translation/databricks_config.template.txt diff --git a/tests/test_litellm/interactions/base_interactions_test.py b/tests/llm_translation/interactions/base_interactions_test.py similarity index 100% rename from tests/test_litellm/interactions/base_interactions_test.py rename to tests/llm_translation/interactions/base_interactions_test.py diff --git a/tests/test_litellm/interactions/test_gemini_interactions.py b/tests/llm_translation/interactions/test_gemini_interactions.py similarity index 88% rename from tests/test_litellm/interactions/test_gemini_interactions.py rename to tests/llm_translation/interactions/test_gemini_interactions.py index afce77e3ce4..0ab4da952e6 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions.py +++ b/tests/llm_translation/interactions/test_gemini_interactions.py @@ -6,7 +6,7 @@ Inherits from BaseInteractionsTest to run the same test suite against Gemini. import os -from tests.test_litellm.interactions.base_interactions_test import ( +from tests.llm_translation.interactions.base_interactions_test import ( BaseInteractionsTest, ) diff --git a/tests/test_litellm/interactions/test_google_interactions_integration.py b/tests/llm_translation/interactions/test_google_interactions_integration.py similarity index 99% rename from tests/test_litellm/interactions/test_google_interactions_integration.py rename to tests/llm_translation/interactions/test_google_interactions_integration.py index 93429d64789..10e6cf86e6d 100644 --- a/tests/test_litellm/interactions/test_google_interactions_integration.py +++ b/tests/llm_translation/interactions/test_google_interactions_integration.py @@ -5,7 +5,7 @@ Tests the litellm.interactions.create() and related methods against the Google A Per OpenAPI spec: https://ai.google.dev/static/api/interactions.openapi.json -Run with: pytest tests/test_litellm/interactions/test_google_interactions_integration.py -v +Run with: pytest tests/llm_translation/interactions/test_google_interactions_integration.py -v """ import asyncio diff --git a/tests/test_litellm/interactions/test_litellm_responses_bridge.py b/tests/llm_translation/interactions/test_litellm_responses_bridge.py similarity index 91% rename from tests/test_litellm/interactions/test_litellm_responses_bridge.py rename to tests/llm_translation/interactions/test_litellm_responses_bridge.py index 17e7f9fc4ff..ae025ab60b0 100644 --- a/tests/test_litellm/interactions/test_litellm_responses_bridge.py +++ b/tests/llm_translation/interactions/test_litellm_responses_bridge.py @@ -7,7 +7,7 @@ the litellm_responses bridge provider, which calls litellm.responses() internall import os -from tests.test_litellm.interactions.base_interactions_test import ( +from tests.llm_translation.interactions.base_interactions_test import ( BaseInteractionsTest, ) diff --git a/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/llm_translation/test_cometapi_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py rename to tests/llm_translation/test_cometapi_chat_transformation.py diff --git a/tests/test_litellm/test_compression.py b/tests/llm_translation/test_compression.py similarity index 100% rename from tests/test_litellm/test_compression.py rename to tests/llm_translation/test_compression.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_e2e.py b/tests/llm_translation/test_databricks_e2e.py similarity index 99% rename from tests/test_litellm/llms/databricks/test_databricks_e2e.py rename to tests/llm_translation/test_databricks_e2e.py index 669f9e94639..a979988102e 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_e2e.py +++ b/tests/llm_translation/test_databricks_e2e.py @@ -51,7 +51,7 @@ Setup: Run with: cd /path/to/litellm - python tests/test_litellm/llms/databricks/test_databricks_e2e.py + python tests/llm_translation/test_databricks_e2e.py Config Options: TEST_AUTH_METHOD=oauth # Test OAuth M2M authentication @@ -69,12 +69,12 @@ import pytest # These are E2E tests that require real Databricks credentials pytestmark = pytest.mark.skip( reason="E2E tests require real Databricks credentials. Run directly with: " - "python tests/test_litellm/llms/databricks/test_databricks_e2e.py" + "python tests/llm_translation/test_databricks_e2e.py" ) # Add the litellm package to path sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) ) # Config file path - can be overridden with DATABRICKS_TEST_CONFIG env var diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/llm_translation/test_json_providers.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_json_providers.py rename to tests/llm_translation/test_json_providers.py diff --git a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/llm_translation/test_mistral_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py rename to tests/llm_translation/test_mistral_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py rename to tests/llm_translation/test_ovhcloud_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py b/tests/llm_translation/test_ovhcloud_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py rename to tests/llm_translation/test_ovhcloud_chat_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/llm_translation/test_vertex_ai_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py rename to tests/llm_translation/test_vertex_ai_image_generation_transformation.py diff --git a/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py b/tests/llm_translation/test_xiaomi_mimo.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py rename to tests/llm_translation/test_xiaomi_mimo.py diff --git a/tests/local_testing/test_handler_gc_does_not_close_client.py b/tests/local_testing/test_handler_gc_does_not_close_client.py index 63c5694dd89..d6987107fa8 100644 --- a/tests/local_testing/test_handler_gc_does_not_close_client.py +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -23,10 +23,7 @@ test here may keep the client in a local: that inflates the very refcount under test, and the test then passes on a broken handler. They hold weak references instead, which the refcount does not count. -These live here rather than under ``tests/test_litellm/`` because they need a -real connection pool: a mocked transport goes on yielding chunks after its -client is closed, so the very teardown under test is what a mock cannot -reproduce. The server is a hermetic, credential-free ``ThreadingHTTPServer`` on +The server is a hermetic, credential-free ``ThreadingHTTPServer`` on an ephemeral loopback port, and needs no network access beyond it. Related: https://github.com/BerriAI/litellm/issues/24929 diff --git a/tests/test_litellm/llms/mistral/__init__.py b/tests/test_litellm/llms/mistral/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai_like/__init__.py b/tests/test_litellm/llms/openai_like/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/vertex_ai/__init__.py b/tests/test_litellm/llms/vertex_ai/__init__.py deleted file mode 100644 index fc7e977484b..00000000000 --- a/tests/test_litellm/llms/vertex_ai/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Vertex AI tests package.""" diff --git a/tests/test_litellm/log.txt b/tests/test_litellm/log.txt deleted file mode 100644 index 6470b12fedb..00000000000 --- a/tests/test_litellm/log.txt +++ /dev/null @@ -1,2 +0,0 @@ -llms/bedrock/chat/invoke_agent/transformation.py:404: error: Incompatible types in assignment (expression has type "object", variable has type "InvokeAgentModelInvocationOutput | None") [assignment] -llms/bedrock/chat/invoke_agent/transformation.py:405: error: Argument 1 to "get" of "Mapping" has incompatible type "str | InvokeAgentModelInvocationOutput"; expected "str" [typeddict-item] diff --git a/tests/test_litellm/ocr/__init__.py b/tests/test_litellm/ocr/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/passthrough/__init__.py b/tests/test_litellm/passthrough/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md index 191777f3e83..fc9798ad30e 100644 --- a/tests/unit/AGENTS.md +++ b/tests/unit/AGENTS.md @@ -30,7 +30,7 @@ Green if `send_batched` drops every row. pydantic doubles in 12 of 203 files, fa ## Where it goes `tests/unit/` mirrors `litellm/`, so a changed file selects its tests by path, not a mapping -file. Empty today; new unit tests go here. The examples above live in `tests/test_litellm` +file. New unit tests go here ## Writing it so a human can read it diff --git a/tests/unit/integrations/dotprompt/test_prompt_manager.py b/tests/unit/integrations/dotprompt/test_prompt_manager.py index 51e14b61929..dbd4e4a4c4c 100644 --- a/tests/unit/integrations/dotprompt/test_prompt_manager.py +++ b/tests/unit/integrations/dotprompt/test_prompt_manager.py @@ -22,7 +22,7 @@ def test_prompt_manager_initialization(): # Test with the existing prompts directory prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) # Should have loaded at least the sample prompts @@ -56,7 +56,7 @@ def test_render_simple_template(): """Test rendering a simple template with variables.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) # Test sample_prompt rendering @@ -72,7 +72,7 @@ def test_render_chat_prompt(): """Test rendering the chat prompt with conditional content.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) # Test with system context @@ -98,7 +98,7 @@ def test_render_coding_assistant(): """Test rendering the coding assistant prompt with complex logic.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) rendered = manager.render( @@ -159,7 +159,7 @@ def test_prompt_not_found(): """Test error handling for non-existent prompts.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) with pytest.raises(KeyError, match="Prompt 'nonexistent' not found"): @@ -170,7 +170,7 @@ def test_list_prompts(): """Test listing available prompts.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) prompts = manager.list_prompts() @@ -184,7 +184,7 @@ def test_get_prompt_metadata(): """Test retrieving prompt metadata.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) metadata = manager.get_prompt_metadata("sample_prompt") @@ -221,7 +221,7 @@ def test_add_prompt_programmatically(): """Test adding prompts programmatically.""" prompt_dir = Path( __file__ - ).parent # Current directory when running from tests/test_litellm/prompts + ).parent manager = PromptManager(prompt_directory=str(prompt_dir)) initial_count = len(manager.prompts) diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index c3478c0d5eb..47c4576f91f 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -1,5 +1,6 @@ """Test health check helper functions""" +import socket import struct import zlib from types import MappingProxyType @@ -214,24 +215,26 @@ async def test_ahealth_check_failure_masks_raw_request_headers(): This tests the fix for the security vulnerability where Authorization headers were being exposed in health check error responses. """ - # Use a model configuration that will fail (invalid endpoint) test_api_key = "dapi-test-key-1234567890abcdef" test_headers = { "Authorization": f"Bearer {test_api_key}", "Content-Type": "application/json", } - response = await ahealth_check( - model_params={ - "model": "databricks/dbrx-instruct", - "api_base": "https://invalid-endpoint-that-will-fail.com/", - "api_key": test_api_key, - "headers": test_headers, - }, - mode="chat", - ) + with socket.socket() as reserved: + reserved.bind(("127.0.0.1", 0)) + api_base = f"http://127.0.0.1:{reserved.getsockname()[1]}/" + + response = await ahealth_check( + model_params={ + "model": "databricks/dbrx-instruct", + "api_base": api_base, + "api_key": test_api_key, + "headers": test_headers, + }, + mode="chat", + ) - # Should have error and raw_request_typed_dict assert "error" in response assert "raw_request_typed_dict" in response @@ -243,22 +246,15 @@ async def test_ahealth_check_failure_masks_raw_request_headers(): headers = raw_request_dict["raw_request_headers"] assert headers is not None - # Security check: Authorization header should be masked, not show full key - if "Authorization" in headers: - auth_header = headers["Authorization"] - # Should be masked (e.g., "Be****90" or similar) - assert auth_header != f"Bearer {test_api_key}", "Authorization header must be masked" - assert auth_header != test_api_key, "API key must not appear in Authorization header" - # Masked headers typically have asterisks or are truncated - assert "*" in auth_header or len(auth_header) < len(f"Bearer {test_api_key}"), ( - f"Authorization header should be masked but got: {auth_header}" - ) + assert "Authorization" in headers + auth_header = headers["Authorization"] + assert auth_header != f"Bearer {test_api_key}", "Authorization header must be masked" + assert auth_header != test_api_key, "API key must not appear in Authorization header" + assert "*" in auth_header or len(auth_header) < len(f"Bearer {test_api_key}"), ( + f"Authorization header should be masked but got: {auth_header}" + ) - # Content-Type should remain unmasked (not sensitive) - if "Content-Type" in headers: - assert headers["Content-Type"] == "application/json" - - print(f"Masked Authorization header: {headers.get('Authorization', 'NOT FOUND')}") + assert headers["Content-Type"] == "application/json" @pytest.mark.asyncio diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index ae9b30d862b..75d3a23e012 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -9,7 +9,7 @@ import subprocess import sys import threading import time -import traceback +from collections.abc import Mapping from concurrent.futures import Future, wait from pathlib import Path from typing import Final @@ -448,35 +448,46 @@ class NeedsToleranceUpdateError(Exception): # test_tokenizers() -def test_encoding_and_decoding(): - try: - sample_text = "Hellö World, this is my input string!" - # openai encoding + decoding - openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) - openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) +def test_encoding_and_decoding(tmp_path: Path): + sample_text = "Hellö World, this is my input string!" - assert openai_text == sample_text + # openai encoding + decoding + openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) + openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) - # claude encoding + decoding - claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) + assert openai_text == sample_text - claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) + # claude encoding + decoding + claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) - assert claude_text == sample_text + claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) - # cohere encoding + decoding - cohere_tokens = encode(model="command-nightly", text=sample_text) - cohere_text = decode(model="command-nightly", tokens=cohere_tokens) + assert claude_text == sample_text - assert cohere_text == sample_text + # cohere encoding + decoding + cohere_tokens = encode(model="command-nightly", text=sample_text) + cohere_text = decode(model="command-nightly", tokens=cohere_tokens) - # llama2 encoding + decoding - llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text) - llama2_text = decode(model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens) + assert cohere_text == sample_text - assert llama2_text == sample_text - except Exception as e: - pytest.fail(f"An exception occured: {e}\n{traceback.format_exc()}") + # llama2 encoding + decoding + words = sample_text.split() + result = _run_in_memory_hub( + HUB_ROUND_TRIP_SCRIPT, + { + "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json( + pre_tokenizers.WhitespaceSplit(), + vocab={"[UNK]": 0, **{word: i + 1 for i, word in enumerate(words)}}, + ) + }, + sample_text, + tmp_path, + ) + + assert result["decoded"] == sample_text + assert result["requested"] == ["hf-internal-testing/llama-tokenizer"] + assert len(result["tokens"]) == len(words) + assert len(result["tokens"]) != len(encode(model="gpt-3.5-turbo", text=sample_text)) # test_encoding_and_decoding() @@ -1447,7 +1458,7 @@ def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_ assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() -HUB_TOKENIZER_SCRIPT: Final = """ +HUB_SETUP_SCRIPT: Final = """ import json import sys sys.path.insert(0, sys.argv[1]) @@ -1466,6 +1477,9 @@ def handle(request): headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40} return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"") huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +""" + +HUB_TOKENIZER_SCRIPT: Final = HUB_SETUP_SCRIPT + """ litellm.cohere_models = {"command-r-v1"} litellm.anthropic_models = {"claude-2"} custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer") @@ -1479,34 +1493,32 @@ print(json.dumps({ })) """ +HUB_ROUND_TRIP_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) +print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) +""" -def _word_level_tokenizer_json(pre_tokenizer: pre_tokenizers.PreTokenizer) -> str: - tokenizer: Final = Tokenizer(models.WordLevel(vocab={"[UNK]": 0}, unk_token="[UNK]")) + +def _word_level_tokenizer_json( + pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None +) -> str: + tokenizer: Final = Tokenizer( + models.WordLevel(vocab=dict(vocab) if vocab is not None else {"[UNK]": 0}, unk_token="[UNK]") + ) tokenizer.pre_tokenizer = pre_tokenizer return tokenizer.to_str() -def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None: - sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!" - served: Final = { - "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()), - "Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")), - "Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()), - } - expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()} - anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids) - tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample) - assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2 - +def _run_in_memory_hub(script: str, served: dict[str, str], text: str, tmp_path: Path) -> dict: result: Final = subprocess.run( [ sys.executable, "-I", "-c", - HUB_TOKENIZER_SCRIPT, + script, str(Path(litellm.__file__).parent.parent), json.dumps(served), - sample, + text, ], capture_output=True, text=True, @@ -1522,7 +1534,22 @@ def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_t ) assert result.returncode == 0, result.stdout + result.stderr - counts: Final = json.loads(result.stdout.strip().splitlines()[-1]) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None: + sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!" + served: Final = { + "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()), + "Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")), + "Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()), + } + expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()} + anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids) + tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample) + assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2 + + counts: Final = _run_in_memory_hub(HUB_TOKENIZER_SCRIPT, served, sample, tmp_path) assert counts == { "llama2": expected["hf-internal-testing/llama-tokenizer"], "llama3": expected["Xenova/llama-3-tokenizer"], diff --git a/tests/unit/llms/anthropic/batches/test_handler.py b/tests/unit/llms/anthropic/batches/test_handler.py index 6fde6350127..28b84123482 100644 --- a/tests/unit/llms/anthropic/batches/test_handler.py +++ b/tests/unit/llms/anthropic/batches/test_handler.py @@ -10,8 +10,7 @@ env) - and assert exactly which seam fired, with what URL/headers, and that the parsed result is the LiteLLMBatch the transform produced. The sync ``retrieve_batch`` dispatch (``_is_async`` true -> coroutine, false -> -asyncio.run) is exercised directly, mirroring the dispatch-contract discipline in -tests/test_litellm/batches/test_main.py. +asyncio.run) is exercised directly. """ from unittest.mock import AsyncMock, MagicMock, patch diff --git a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py index 90cba035760..b2850192f91 100644 --- a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py +++ b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py @@ -1,9 +1,7 @@ """ Coverage for filter_anthropic_output_schema's array/object constraint stripping. -Mirrors tests/litellm/llms/anthropic/test_anthropic_schema_filter.py, but lives -under tests/test_litellm/ so the coverage-uploading CI job exercises the stripped -keyword handling (uniqueItems / contains / minProperties / maxProperties plus +Exercises the stripped keyword handling (uniqueItems / contains / minProperties / maxProperties plus multipleOf / patternProperties / propertyNames / dependentRequired / dependentSchemas / unevaluatedProperties / if / then / else / not / prefixItems), the ``uniqueItems: false`` branch, the oneOf to anyOf rewrite, and the From 7ae721bf79401cdc954e74d19bff5f15e0ff0506 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:12:48 -0700 Subject: [PATCH 51/58] refactor(rust): prepare inference and auth foundations for the gateway (#43287) * refactor(rust): prepare inference and auth foundations * fix(rust): keep textract operations parsing from kebab-case model names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/AGENTS.md | 6 +- litellm-rust/Cargo.lock | 36 +- litellm-rust/Cargo.toml | 5 +- litellm-rust/crates/auth-aws/Cargo.toml | 1 + litellm-rust/crates/auth-aws/src/aws.rs | 494 ++++++----- litellm-rust/crates/auth-aws/src/signer.rs | 28 +- litellm-rust/crates/auth-gcp/src/lib.rs | 10 +- litellm-rust/crates/auth-types/src/http.rs | 15 - litellm-rust/crates/auth-types/src/lib.rs | 2 +- litellm-rust/crates/auth/src/lib.rs | 3 + litellm-rust/crates/auth/src/services.rs | 9 + litellm-rust/crates/cache-s3/src/auth.rs | 12 +- litellm-rust/crates/core-utils/Cargo.toml | 1 + .../src/prompt_templates/factory.rs | 9 +- litellm-rust/crates/core/AGENTS.md | 10 + litellm-rust/crates/core/Cargo.toml | 2 +- .../core/src/audio_transcription/error.rs | 43 - .../core/src/audio_transcription/handler.rs | 13 +- .../core/src/audio_transcription/mod.rs | 11 +- .../core/src/audio_transcription/prepare.rs | 32 +- .../core/src/audio_transcription/types.rs | 8 +- .../crates/core/src/chat_completions/error.rs | 43 - .../core/src/chat_completions/handler.rs | 17 +- .../crates/core/src/chat_completions/mod.rs | 11 +- .../core/src/chat_completions/prepare.rs | 179 ++-- .../crates/core/src/chat_completions/types.rs | 8 +- litellm-rust/crates/core/src/constants.rs | 4 - litellm-rust/crates/core/src/error.rs | 167 +++- litellm-rust/crates/core/src/lib.rs | 3 +- .../crates/core/src/messages/common_utils.rs | 50 +- .../crates/core/src/messages/error.rs | 82 -- .../crates/core/src/messages/handler.rs | 26 +- litellm-rust/crates/core/src/messages/mod.rs | 32 +- .../crates/core/src/messages/prepare.rs | 266 +++--- .../crates/core/src/messages/route.rs | 148 ++-- .../crates/core/src/messages/types.rs | 38 +- litellm-rust/crates/core/src/outbound.rs | 30 +- litellm-rust/crates/core/src/resources.rs | 38 + .../crates/core/src/responses/error.rs | 17 - litellm-rust/crates/core/src/responses/mod.rs | 3 +- .../crates/core/tests/audio_transcription.rs | 2 +- .../crates/core/tests/chat_completions.rs | 2 +- .../crates/core/tests/messages/host.rs | 17 +- .../crates/core/tests/messages/main.rs | 27 +- .../crates/core/tests/messages/request.rs | 23 +- .../crates/core/tests/messages/response.rs | 55 +- .../crates/core/tests/messages/stream.rs | 28 +- litellm-rust/crates/core/tests/ocr/mistral.rs | 25 +- litellm-rust/crates/core/tests/resources.rs | 157 ++++ litellm-rust/crates/core/tests/support/mod.rs | 4 + litellm-rust/crates/host-python/Cargo.toml | 10 +- litellm-rust/crates/litellm/Cargo.toml | 3 + litellm-rust/crates/litellm/src/lib.rs | 2 + litellm-rust/crates/llms/Cargo.toml | 3 +- .../llms/src/anthropic/batches/AGENTS.md | 1 + .../src/anthropic/batches/transformation.rs | 5 +- .../crates/llms/src/anthropic/chat/handler.rs | 15 +- .../llms/src/anthropic/chat/transformation.rs | 65 +- .../crates/llms/src/anthropic/common_utils.rs | 131 +-- .../llms/src/anthropic/count_tokens/AGENTS.md | 1 + .../anthropic/count_tokens/transformation.rs | 2 +- .../llms/src/anthropic/messages/AGENTS.md | 1 + .../llms/src/anthropic/messages/handler.rs | 80 +- .../llms/src/anthropic/messages/headers.rs | 98 ++- .../anthropic/messages/streaming_iterator.rs | 173 +--- .../llms/src/anthropic/messages/thinking.rs | 770 ++++++++++-------- .../src/anthropic/messages/transformation.rs | 107 +-- .../ocr/analyze_transformation.rs | 9 +- .../llms/src/aws_textract/ocr/common_utils.rs | 6 +- .../src/aws_textract/ocr/transformation.rs | 9 +- .../anthropic/messages_transformation.rs | 112 ++- .../llms/src/azure_ai/ocr/common_utils.rs | 7 +- .../document_intelligence/transformation.rs | 30 +- .../llms/src/azure_ai/ocr/transformation.rs | 52 +- .../src/base_llm/anthropic_messages/mod.rs | 1 + .../base_llm/anthropic_messages/streaming.rs | 141 ++++ .../anthropic_messages/transformation.rs | 224 +---- .../audio_transcription/transformation.rs | 9 +- litellm-rust/crates/llms/src/base_llm/auth.rs | 263 ++++++ .../llms/src/base_llm/base_model_iterator.rs | 153 ++++ .../crates/llms/src/base_llm/chat/mod.rs | 1 + .../llms/src/base_llm/chat/streaming.rs | 54 ++ .../llms/src/base_llm/chat/transformation.rs | 48 +- litellm-rust/crates/llms/src/base_llm/mod.rs | 1 + .../crates/llms/src/base_llm/ocr/handler.rs | 14 +- .../src/base_llm/responses/transformation.rs | 2 +- .../src/bedrock/audio_transcription/mod.rs | 37 +- .../bedrock/chat/converse_transformation.rs | 60 +- .../llms/src/bedrock/chat/invoke_handler.rs | 138 ++++ .../crates/llms/src/bedrock/chat/mod.rs | 1 + .../anthropic_claude3_transformation.rs | 554 +++++++++++++ .../messages/invoke_transformations/mod.rs | 1 + .../crates/llms/src/bedrock/messages/mod.rs | 1 + litellm-rust/crates/llms/src/bedrock/mod.rs | 1 + litellm-rust/crates/llms/src/error.rs | 18 + litellm-rust/crates/llms/src/lib.rs | 3 + .../src/openai/responses/transformation.rs | 6 +- .../llms/src/vertex_ai/ocr/transformation.rs | 3 +- .../tests/anthropic_chat_transformation.rs | 31 +- .../tests/bedrock_converse_transformation.rs | 76 +- litellm-rust/crates/model-catalog/Cargo.toml | 6 +- .../crates/model-catalog/src/capabilities.rs | 14 - .../crates/model-catalog/src/model_info.rs | 5 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 - .../crates/python-bridge/src/errors.rs | 113 +-- litellm-rust/crates/python-bridge/src/http.rs | 13 +- .../python-bridge/src/python_settings.rs | 13 +- .../src/routes/audio_transcription.rs | 8 +- .../src/routes/chat_completions.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 32 +- .../python-bridge/src/routes/messages/mod.rs | 2 +- .../python-bridge/src/routes/ocr/errors.rs | 18 +- .../python-bridge/src/routes/ocr/mod.rs | 19 +- .../python-bridge/src/routes/responses.rs | 13 +- .../python-bridge/src/secrets/callback.rs | 12 +- litellm-rust/crates/secrets-aws/src/auth.rs | 16 +- .../crates/secrets-aws/src/secret_manager.rs | 1 + .../secrets-aws/src/secret_manager/client.rs | 2 + litellm-rust/crates/secrets-types/Cargo.toml | 1 + .../crates/secrets-types/src/config.rs | 4 +- litellm-rust/crates/secrets/Cargo.toml | 2 +- litellm-rust/crates/types/Cargo.toml | 5 + litellm-rust/crates/types/src/lib.rs | 1 + .../anthropic_messages/anthropic_request.rs | 224 ++++- litellm-rust/crates/types/src/llms/openai.rs | 74 ++ litellm-rust/crates/types/src/recognized.rs | 42 + 126 files changed, 4120 insertions(+), 2308 deletions(-) create mode 100644 litellm-rust/crates/auth/src/services.rs delete mode 100644 litellm-rust/crates/core/src/audio_transcription/error.rs delete mode 100644 litellm-rust/crates/core/src/chat_completions/error.rs delete mode 100644 litellm-rust/crates/core/src/messages/error.rs create mode 100644 litellm-rust/crates/core/src/resources.rs delete mode 100644 litellm-rust/crates/core/src/responses/error.rs create mode 100644 litellm-rust/crates/core/tests/resources.rs create mode 100644 litellm-rust/crates/litellm/Cargo.toml create mode 100644 litellm-rust/crates/litellm/src/lib.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md create mode 100644 litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md create mode 100644 litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md create mode 100644 litellm-rust/crates/llms/src/base_llm/anthropic_messages/streaming.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/auth.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/chat/streaming.rs create mode 100644 litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs create mode 100644 litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs create mode 100644 litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs create mode 100644 litellm-rust/crates/llms/src/bedrock/messages/mod.rs create mode 100644 litellm-rust/crates/llms/src/error.rs create mode 100644 litellm-rust/crates/types/src/recognized.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 70fcc367905..bc6a2552e4c 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -9,10 +9,14 @@ - A test for another crate's item belongs in that crate, not in a downstream one - Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own +## Test fixtures and cases + +Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new and updated tests and [`#[fixture]`](https://docs.rs/rstest/latest/rstest/attr.fixture.html) for reusable setup, injected through typed test arguments. Express input variations as named `#[case::name(...)]` cases instead of loops or duplicated tests so each failure identifies its case. Keep behavior assertions in the test body and fixtures focused on setup. Use the workspace `rstest` dependency + ## Error definitions - A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs` -- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each +- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly); failures no caller tells apart share one variant and differ only in its message - Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string - Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return - Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5bb65b6b05f..580b427a97e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -199,7 +199,7 @@ checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] @@ -1053,18 +1053,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.6" +version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.6.6" +version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" dependencies = [ "anstyle", "clap_lex", @@ -2816,6 +2816,10 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "litellm" +version = "0.0.1" + [[package]] name = "litellm-auth" version = "0.1.0" @@ -2840,6 +2844,7 @@ dependencies = [ "litellm-http", "moka", "reqwest 0.12.28", + "rstest", "serde_json", "sha2 0.10.9", "thiserror 2.0.19", @@ -3141,6 +3146,7 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_with", + "strum", "thiserror 2.0.19", "url", ] @@ -3242,6 +3248,7 @@ dependencies = [ "litellm-framing", "litellm-host", "litellm-http", + "litellm-python-compat", "litellm-secrets", "litellm-types", "reqwest 0.12.28", @@ -3263,6 +3270,7 @@ version = "0.1.0" dependencies = [ "indexmap 2.14.0", "jsonschema", + "litellm-types", "rstest", "schemars 1.2.2", "serde", @@ -3282,7 +3290,6 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", - "litellm-auth-gcp", "litellm-cache", "litellm-cache-azure-blob", "litellm-cache-disk", @@ -3490,6 +3497,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "strum", "thiserror 2.0.19", "tokio", "veil", @@ -3587,8 +3595,10 @@ name = "litellm-types" version = "0.1.0" dependencies = [ "rstest", + "schemars 1.2.2", "serde", "serde_json", + "strum", ] [[package]] @@ -4663,7 +4673,7 @@ checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] @@ -5142,7 +5152,7 @@ dependencies = [ "proc-macro2", "quote", "serde_derive_internals", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] @@ -5230,7 +5240,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] @@ -5241,7 +5251,7 @@ checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] @@ -5549,9 +5559,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.0" +version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2fac314a64dc9a36e61a9eb4261a5e9bbfbc922b27e518af97bc32b926cf967" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" dependencies = [ "proc-macro2", "quote", @@ -5663,7 +5673,7 @@ checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.6", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 022e8f13311..442bd620e05 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,7 +10,6 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-tracing = { path = "crates/tracing" } -tracing = "0.1" litellm-core = { path = "crates/core" } litellm-coroutine = { path = "crates/coroutine" } litellm-host = { path = "crates/host" } @@ -48,7 +47,9 @@ litellm-token-counter-fast = { path = "crates/token-counter-fast" } litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } +litellm-python-compat = { path = "crates/python-compat" } +tracing = "0.1" bytes = "1" http = "1" google-cloud-auth = { version = "1.16.0", default-features = false } @@ -57,8 +58,8 @@ hyper-util = { version = "0.1.20", default-features = false, features = ["client proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } -pythonize = "0.29.0" rand = "0.8" +schemars = "1" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } qdrant-client = { version = "1.19.0", default-features = false } uuid = { version = "1", features = ["v4"] } diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 9592f278d94..2a9a9e4768c 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -22,6 +22,7 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } reqwest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 409ff78867f..69eb4265159 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::sync::OnceLock; use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; @@ -26,8 +25,26 @@ use super::constants::{ const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); -static STATIC_CREDENTIALS_CACHE: OnceLock> = OnceLock::new(); -static AMBIENT_CREDENTIALS_CACHE: OnceLock> = OnceLock::new(); +#[derive(Clone)] +pub struct AwsAuthService { + static_credentials: Cache, + ambient_credentials: Cache, +} + +impl Default for AwsAuthService { + fn default() -> Self { + Self { + static_credentials: Cache::builder() + .max_capacity(200) + .time_to_live(STATIC_CREDENTIALS_TTL) + .build(), + ambient_credentials: Cache::builder() + .max_capacity(200) + .time_to_live(AMBIENT_CREDENTIALS_TTL) + .build(), + } + } +} fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { match flow { @@ -108,35 +125,19 @@ fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { format!("{:x}", hasher.finalize()) } -fn static_credentials_cache() -> &'static Cache { - STATIC_CREDENTIALS_CACHE.get_or_init(|| { - Cache::builder() - .max_capacity(200) - .time_to_live(STATIC_CREDENTIALS_TTL) - .build() - }) -} +impl AwsAuthService { + fn get_cached_credentials(&self, key: &str) -> Option { + self.static_credentials + .get(key) + .or_else(|| self.ambient_credentials.get(key)) + } -fn ambient_credentials_cache() -> &'static Cache { - AMBIENT_CREDENTIALS_CACHE.get_or_init(|| { - Cache::builder() - .max_capacity(200) - .time_to_live(AMBIENT_CREDENTIALS_TTL) - .build() - }) -} - -fn get_cached_credentials(key: &str) -> Option { - static_credentials_cache() - .get(key) - .or_else(|| ambient_credentials_cache().get(key)) -} - -fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) { - if ttl == STATIC_CREDENTIALS_TTL { - static_credentials_cache().insert(key, credentials); - } else { - ambient_credentials_cache().insert(key, credentials); + fn set_cached_credentials(&self, key: String, credentials: Credentials, ttl: Duration) { + if ttl == STATIC_CREDENTIALS_TTL { + self.static_credentials.insert(key, credentials); + } else { + self.ambient_credentials.insert(key, credentials); + } } } @@ -214,66 +215,157 @@ pub fn classify_auth( AwsAuthFlow::DefaultChain } -pub async fn resolve_credentials( - config: AwsAuthConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result { - let resolved = config.clone().with_environment(env_lookup); - let flow = classify_auth(config, env_lookup); - match flow { - AwsAuthFlow::SessionToken { - access_key_id, - secret_access_key, - session_token, - } => Ok(Credentials::new( - access_key_id, - secret_access_key, - Some(session_token), - None, - "litellm-static-session", - )), - AwsAuthFlow::StaticKeys { - access_key_id, - secret_access_key, - region_name, - } => { - let flow = AwsAuthFlow::StaticKeys { - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - region_name, - }; - let key = cache_key(&resolved, &flow); - if let Some(credentials) = get_cached_credentials(&key) { - return Ok(credentials); - } - let credentials = Credentials::new( +impl AwsAuthService { + pub async fn resolve_credentials( + &self, + config: AwsAuthConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let resolved = config.clone().with_environment(env_lookup); + let flow = classify_auth(config, env_lookup); + match flow { + AwsAuthFlow::SessionToken { access_key_id, secret_access_key, + session_token, + } => Ok(Credentials::new( + access_key_id, + secret_access_key, + Some(session_token), None, - None, - "litellm-static", - ); - set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), - ); - Ok(credentials) - } - AwsAuthFlow::Profile { name } => { - let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() - .profile_name(name) - .build(); - provider - .provide_credentials() - .await - .map_err(|error| Error::AwsProfile(error.to_string())) - } - AwsAuthFlow::AssumeRole { role, session_name } => { - if is_already_running_as_role(&role, &resolved).await? { - let ambient_flow = AwsAuthFlow::DefaultChain; - let key = cache_key(&resolved, &ambient_flow); - if let Some(credentials) = get_cached_credentials(&key) { + "litellm-static-session", + )), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + region_name, + } => { + let flow = AwsAuthFlow::StaticKeys { + access_key_id: access_key_id.clone(), + secret_access_key: secret_access_key.clone(), + region_name, + }; + let key = cache_key(&resolved, &flow); + if let Some(credentials) = self.get_cached_credentials(&key) { + return Ok(credentials); + } + let credentials = Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "litellm-static", + ); + self.set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), + ); + Ok(credentials) + } + AwsAuthFlow::Profile { name } => { + let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_name(name) + .build(); + provider + .provide_credentials() + .await + .map_err(|error| Error::AwsProfile(error.to_string())) + } + AwsAuthFlow::AssumeRole { role, session_name } => { + if is_already_running_as_role(&role, &resolved).await? { + let ambient_flow = AwsAuthFlow::DefaultChain; + let key = cache_key(&resolved, &ambient_flow); + if let Some(credentials) = self.get_cached_credentials(&key) { + return Ok(credentials); + } + let provider = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await; + let credentials = provider + .provide_credentials() + .await + .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; + self.set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), + ); + return Ok(credentials); + } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + if let (Some(access_key_id), Some(secret_access_key)) = + (resolved.access_key_id, resolved.secret_access_key) + { + loader = loader.credentials_provider(Credentials::new( + access_key_id, + secret_access_key, + resolved.session_token, + None, + "litellm-role-source", + )); + } + let sdk_config = loader.load().await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role); + let builder = match session_name { + Some(name) => builder.session_name(name), + None => builder.session_name(default_session_name()), + }; + let builder = match resolved.external_id { + Some(id) => builder.external_id(id), + None => builder, + }; + let provider = builder.configure(&sdk_config).build().await; + provider + .provide_credentials() + .await + .map_err(|error| Error::AwsAssumeRole(error.to_string())) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => { + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let client = aws_sdk_sts::Client::new(&sdk_config); + let response = client + .assume_role_with_web_identity() + .role_arn(role) + .role_session_name(session_name) + .web_identity_token(token) + .send() + .await + .map_err(|error| Error::AwsWebIdentity(error.to_string()))?; + let credentials = response + .credentials() + .ok_or(Error::AwsMissingWebIdentityCredentials)?; + let expiration = SystemTime::try_from(*credentials.expiration()) + .map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_string()), + Some(expiration), + "litellm-web-identity", + )) + } + AwsAuthFlow::DefaultChain => { + let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); + if let Some(credentials) = self.get_cached_credentials(&key) { return Ok(credentials); } let provider = @@ -284,101 +376,14 @@ pub async fn resolve_credentials( .provide_credentials() .await .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; - set_cached_credentials( + self.set_cached_credentials( key, credentials.clone(), - credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), + credential_cache_ttl(&AwsAuthFlow::DefaultChain) + .unwrap_or(AMBIENT_CREDENTIALS_TTL), ); - return Ok(credentials); + Ok(credentials) } - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - if let (Some(access_key_id), Some(secret_access_key)) = - (resolved.access_key_id, resolved.secret_access_key) - { - loader = loader.credentials_provider(Credentials::new( - access_key_id, - secret_access_key, - resolved.session_token, - None, - "litellm-role-source", - )); - } - let sdk_config = loader.load().await; - let builder = aws_config::sts::AssumeRoleProvider::builder(role); - let builder = match session_name { - Some(name) => builder.session_name(name), - None => builder.session_name(default_session_name()), - }; - let builder = match resolved.external_id { - Some(id) => builder.external_id(id), - None => builder, - }; - let provider = builder.configure(&sdk_config).build().await; - provider - .provide_credentials() - .await - .map_err(|error| Error::AwsAssumeRole(error.to_string())) - } - AwsAuthFlow::WebIdentity { - token, - role, - session_name, - } => { - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let client = aws_sdk_sts::Client::new(&sdk_config); - let response = client - .assume_role_with_web_identity() - .role_arn(role) - .role_session_name(session_name) - .web_identity_token(token) - .send() - .await - .map_err(|error| Error::AwsWebIdentity(error.to_string()))?; - let credentials = response - .credentials() - .ok_or(Error::AwsMissingWebIdentityCredentials)?; - let expiration = SystemTime::try_from(*credentials.expiration()) - .map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?; - Ok(Credentials::new( - credentials.access_key_id(), - credentials.secret_access_key(), - Some(credentials.session_token().to_string()), - Some(expiration), - "litellm-web-identity", - )) - } - AwsAuthFlow::DefaultChain => { - let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); - if let Some(credentials) = get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; - set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - Ok(credentials) } } } @@ -585,6 +590,37 @@ pub fn aws_auth_config( } } +/// Where the credentials that sign a request come from, decided when the request is +/// prepared and resolved when it is sent. +#[derive(Clone, Debug, PartialEq)] +pub enum AwsCredentialSource { + HostSupplied(Credentials), + Chain(AwsAuthConfig), +} + +impl AwsCredentialSource { + pub fn from_params( + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Self { + match host_supplied_credentials(optional_params) { + Some(credentials) => Self::HostSupplied(credentials), + None => Self::Chain(aws_auth_config(optional_params, env_lookup)), + } + } + + pub async fn resolve( + self, + auth: &AwsAuthService, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + match self { + Self::HostSupplied(credentials) => Ok(credentials), + Self::Chain(config) => auth.resolve_credentials(config, env_lookup).await, + } + } +} + /// Credentials a host resolved through its own chain and handed down verbatim. /// /// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads @@ -747,17 +783,18 @@ mod tests { #[tokio::test] async fn static_credentials_do_not_use_network() { - let credentials = resolve_credentials( - AwsAuthConfig { - access_key_id: Some("ak".into()), - secret_access_key: Some("sk".into()), - region_name: Some("us-east-1".into()), - ..Default::default() - }, - &no_env, - ) - .await - .expect("static credentials"); + let credentials = AwsAuthService::default() + .resolve_credentials( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + region_name: Some("us-east-1".into()), + ..Default::default() + }, + &no_env, + ) + .await + .expect("static credentials"); assert_eq!(credentials.access_key_id(), "ak"); assert_eq!(credentials.session_token(), None); } @@ -807,17 +844,67 @@ mod tests { ); } - #[test] + #[rstest::rstest] fn cache_round_trip_preserves_credentials() { + let auth = AwsAuthService::default(); let key = format!("cache-test-{}", std::process::id()); let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); - set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + auth.set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); assert_eq!( - get_cached_credentials(&key).map(|value| value.access_key_id().to_string()), + auth.get_cached_credentials(&key) + .map(|value| value.access_key_id().to_string()), Some("cache-ak".to_string()) ); } + #[rstest::rstest] + #[tokio::test] + async fn cloned_services_reuse_credentials_but_independent_services_do_not() { + let auth = AwsAuthService::default(); + let config = AwsAuthConfig { + access_key_id: Some("configured-key".into()), + secret_access_key: Some("configured-secret".into()), + region_name: Some("us-east-1".into()), + ..AwsAuthConfig::default() + }; + let flow = classify_auth(config.clone(), &no_env); + let cached = Credentials::new("cached-key", "cached-secret", None, None, "test"); + auth.set_cached_credentials( + cache_key(&config, &flow), + cached.clone(), + STATIC_CREDENTIALS_TTL, + ); + + let reused = auth + .clone() + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let independent = AwsAuthService::default() + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let different = AwsAuthConfig { + access_key_id: Some("different-key".into()), + ..config.clone() + }; + let other_identity = auth + .resolve_credentials(different.clone(), &no_env) + .await + .unwrap(); + + assert_eq!(reused.access_key_id(), cached.access_key_id()); + assert_eq!(reused.secret_access_key(), cached.secret_access_key()); + assert_eq!( + Some(independent.access_key_id()), + config.access_key_id.as_deref() + ); + assert_eq!( + Some(other_identity.access_key_id()), + different.access_key_id.as_deref() + ); + } + #[test] fn same_role_comparison_matches_partition_account_and_role() { assert!(same_role_arns( @@ -952,16 +1039,17 @@ mod tests { let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); let headers = BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); - let credentials = resolve_credentials( - AwsAuthConfig { - access_key_id: Some(access_key_id), - secret_access_key: Some(secret_access_key), - region_name: Some("us-west-2".to_string()), - ..Default::default() - }, - &no_env, - ) - .await?; + let credentials = AwsAuthService::default() + .resolve_credentials( + AwsAuthConfig { + access_key_id: Some(access_key_id), + secret_access_key: Some(secret_access_key), + region_name: Some("us-west-2".to_string()), + ..Default::default() + }, + &no_env, + ) + .await?; let client = litellm_http::Client::plain_for_test(); let mut failures = Vec::new(); diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 49a3910c1d5..3868fdd939b 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -1,13 +1,11 @@ use std::{collections::BTreeMap, time::SystemTime}; +use crate::{ + AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header, + sign_post, +}; use aws_credential_types::Credentials; use litellm_http::outbound::{RequestSigner, UnsignedRequest}; -use serde_json::{Map, Value}; - -use crate::{ - Error, aws_auth_config, aws_signature_headers, host_supplied_credentials, - is_sigv4_computed_header, resolve_credentials, sign_post, -}; #[derive(Clone, Debug)] pub struct SigV4Signer { @@ -32,19 +30,17 @@ impl SigV4Signer { } pub async fn resolve( + auth: &AwsAuthService, region: String, service: &'static str, - optional_params: &Map, + credentials: AwsCredentialSource, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { - let credentials = match host_supplied_credentials(optional_params) { - Some(credentials) => credentials, - None => { - resolve_credentials(aws_auth_config(optional_params, env_lookup), env_lookup) - .await? - } - }; - Ok(Self::new(region, service, credentials)) + Ok(Self::new( + region, + service, + credentials.resolve(auth, env_lookup).await?, + )) } } @@ -80,7 +76,7 @@ mod tests { use std::time::{Duration, UNIX_EPOCH}; use litellm_http::outbound::OutboundRequest; - use serde_json::json; + use serde_json::{Value, json}; use super::*; diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 682f1af5fe1..4374dff95aa 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -131,7 +131,7 @@ impl Default for VertexAuth { } impl VertexAuth { - fn new(loader: Arc) -> Self { + pub fn new(loader: Arc) -> Self { Self { providers: Cache::builder().max_capacity(64).build(), loader, @@ -220,16 +220,16 @@ impl VertexAuth { } } -trait VertexTokenSource: Send + Sync { +pub trait VertexTokenSource: Send + Sync { fn project_id(&self) -> VertexAuthFuture<'_, String>; fn token(&self) -> VertexAuthFuture<'_, String>; } -trait VertexProviderLoader: Send + Sync { +pub trait VertexProviderLoader: Send + Sync { fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; } -type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; +pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; struct GcpTokenSource(Arc); @@ -305,7 +305,7 @@ fn validate_request_credentials(configured: &str) -> Result<&str, Error> { } #[derive(Clone, Debug)] -enum CredentialSource { +pub enum CredentialSource { Inline(SecretValue), Trusted(SecretValue), ApplicationCredentials(String), diff --git a/litellm-rust/crates/auth-types/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs index 0cb5839f965..f3c5254b60e 100644 --- a/litellm-rust/crates/auth-types/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -40,21 +40,6 @@ pub fn apply_credential( ) } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RequestAuth { - Header { - name: &'static str, - value: String, - }, - Bearer { - token: String, - }, - AwsSigV4 { - region: String, - service: &'static str, - }, -} - #[cfg(test)] mod tests { use super::{CredentialPlacement, apply_credential}; diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 9d399249c05..ebab5b84d88 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -51,7 +51,7 @@ pub use credential::{ CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; pub use error::Error; -pub use http::{CredentialPlacement, RequestAuth}; +pub use http::CredentialPlacement; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth/src/lib.rs b/litellm-rust/crates/auth/src/lib.rs index 622a5b2d58b..d23bccccc9f 100644 --- a/litellm-rust/crates/auth/src/lib.rs +++ b/litellm-rust/crates/auth/src/lib.rs @@ -2,6 +2,9 @@ pub use litellm_auth_types::*; +mod services; +pub use services::AuthServices; + #[cfg(feature = "aws")] pub use litellm_auth_aws as aws; #[cfg(feature = "azure")] diff --git a/litellm-rust/crates/auth/src/services.rs b/litellm-rust/crates/auth/src/services.rs new file mode 100644 index 00000000000..4c88c9a89a2 --- /dev/null +++ b/litellm-rust/crates/auth/src/services.rs @@ -0,0 +1,9 @@ +#[derive(Default)] +pub struct AuthServices { + #[cfg(feature = "aws")] + pub aws: litellm_auth_aws::AwsAuthService, + #[cfg(feature = "azure")] + pub azure: litellm_auth_azure::AzureAuthService, + #[cfg(feature = "gcp")] + pub gcp: litellm_auth_gcp::VertexAuth, +} diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index fdf71fc011b..f4f06e5371d 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -2,10 +2,11 @@ use aws_credential_types::{ Credentials as AwsCredentials, provider::{ProvideCredentials, error::CredentialsError, future}, }; -use litellm_auth_aws::{AwsAuthConfig, resolve_credentials}; +use litellm_auth_aws::{AwsAuthConfig, AwsAuthService}; #[derive(Clone)] pub struct S3Credentials { + auth: AwsAuthService, config: AwsAuthConfig, env: fn(&str) -> Option, } @@ -16,7 +17,11 @@ impl S3Credentials { } pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> Self { - Self { config, env } + Self { + auth: AwsAuthService::default(), + config, + env, + } } } @@ -38,7 +43,8 @@ impl ProvideCredentials for S3Credentials { "litellm-s3-cache", )); } - resolve_credentials(self.config.clone(), &self.env) + self.auth + .resolve_credentials(self.config.clone(), &self.env) .await .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) }) diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index baf5dd16707..22196979781 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -13,6 +13,7 @@ serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" serde_with.workspace = true +strum.workspace = true thiserror.workspace = true url.workspace = true diff --git a/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs b/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs index 2c4921d26be..63ef79c0fa2 100644 --- a/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs +++ b/litellm-rust/crates/core-utils/src/prompt_templates/factory.rs @@ -11,11 +11,13 @@ //! accepts; anything richer is declined upstream by the capability gate. use litellm_types::llms::openai::{ChatMessage, ChatMessageContent}; +use strum::IntoStaticStr; pub const EMPTY_TEXT_PLACEHOLDER: &str = "[System: Empty message content sanitised to satisfy protocol]"; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)] +#[strum(serialize_all = "snake_case")] pub enum TurnRole { User, Assistant, @@ -23,10 +25,7 @@ pub enum TurnRole { impl TurnRole { pub fn as_str(self) -> &'static str { - match self { - Self::User => "user", - Self::Assistant => "assistant", - } + self.into() } } diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index 0c8a747019d..a40265729c7 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -12,4 +12,14 @@ Each crate mirrors one top-level Python package, so a Rust path reads as its Pyt A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate +## Error placement + +The workspace `Error definitions` rules shape each crate's error; this section decides which crate and module a failure belongs to + +A failure is declared once, by the lowest crate that raises it. Every crate above nests that error unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error`. `RouteError` collects route failures and never re-declares a variant a lower crate raises + +Scope follows the concept, not the first caller. An error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it. A failure two providers or two routes can hit, such as wire framing, stream event decoding, or a malformed provider response, belongs to the crate that owns the concept: `litellm-framing` for framing, `litellm_llms::Error` for the transformation layer + +`litellm_llms::Error` (`crates/llms/src/error.rs`) is the one transformation error for every provider and API. `base_llm/ocr/error.rs` is the recorded exception until OCR folds into it + Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 8700c8df308..0051631f40d 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -13,7 +13,7 @@ litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true base64.workspace = true -litellm-auth.workspace = true +litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-http.workspace = true litellm-llms.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/error.rs b/litellm-rust/crates/core/src/audio_transcription/error.rs deleted file mode 100644 index 81b57af2c6c..00000000000 --- a/litellm-rust/crates/core/src/audio_transcription/error.rs +++ /dev/null @@ -1,43 +0,0 @@ -use litellm_llms::base_llm::chat::transformation::Error as LlmError; - -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, - }, - #[error("missing required field: {0}")] - MissingField(&'static str), - #[error("invalid provider: {0}")] - InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), - #[error(transparent)] - Auth(#[from] litellm_auth::Error), - #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] - Headers(#[from] litellm_http::request::HeaderError), - #[error(transparent)] - Http(#[from] litellm_http::Error), - #[error(transparent)] - Aws(#[from] litellm_auth_aws::Error), -} - -impl From for Error { - fn from(error: LlmError) -> Self { - match error { - LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, - LlmError::MissingField(field) => Self::MissingField(field), - LlmError::InvalidRequest(message) => Self::InvalidRequest(message), - LlmError::InvalidResponse(message) => Self::InvalidResponse(message), - LlmError::Unsupported(reason) => Self::Unsupported(reason), - LlmError::Auth(error) => Self::Auth(error), - } - } -} diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 30900bc14c6..866d08b22e8 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,6 +1,7 @@ use std::time::Duration; use litellm_http::{Client, request::truncate_error_body}; +use litellm_llms::base_llm::auth::resolve_auth; use serde_json::Value; use super::Error; @@ -11,21 +12,21 @@ use crate::{ pub async fn execute_audio_transcription_provider_call( http: &Client, + auth: &litellm_auth::AuthServices, request: ProviderAudioTranscriptionRequest, ) -> Result { - let response = crate::outbound::outbound_request::( - &request.auth, + let env_lookup = |key: &str| std::env::var(key).ok(); + let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?; + let response = crate::outbound::outbound_request( + authenticated, request.url.clone(), - request.upstream_headers.clone(), &request.body, Some( request .timeout .unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)), ), - &request.optional_params, - ) - .await? + )? .send(http) .await .map_err(|error| { diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index dc75326d5c3..3d329ebfc96 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -1,21 +1,20 @@ -mod error; pub mod types; -pub use error::Error; +pub use crate::error::RouteError as Error; mod handler; mod prepare; pub use handler::execute_audio_transcription_provider_call; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{ClientVariant, HttpClientConfig}; pub use prepare::prepare_audio_transcription_provider_call; use serde_json::Value; use crate::audio_transcription::types::AudioTranscriptionRequest; pub async fn audio_transcription( - pool: &HttpClientPool, + resources: &crate::resources::CoreResources, config: &HttpClientConfig, request: AudioTranscriptionRequest<'_>, ) -> Result { let request = prepare_audio_transcription_provider_call(request)?; - let http = pool.client(config, ClientVariant::Provider)?; - execute_audio_transcription_provider_call(&http, request).await + let http = resources.pool.client(config, ClientVariant::Provider)?; + execute_audio_transcription_provider_call(&http, &resources.auth, request).await } diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 807993c38b7..fa50c43d62d 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,7 +1,10 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; -use litellm_http::request::{has_header, string_headers}; +use litellm_http::request::string_headers; use litellm_llms::{ - base_llm::audio_transcription::transformation::{BaseAudioTranscriptionConfig, RequestAuth}, + base_llm::{ + audio_transcription::transformation::BaseAudioTranscriptionConfig, + auth::{ValidatedEnvironment, with_default_headers}, + }, bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG, }; @@ -39,20 +42,13 @@ pub fn prepare_audio_transcription_provider_call( let config = provider_config(provider_info.custom_llm_provider) .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; let env_lookup = |key: &str| std::env::var(key).ok(); - let mut headers = string_headers("audio transcription", request.extra_headers)?; - let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?; - match &auth { - RequestAuth::Bearer { token } if !has_header(&headers, "authorization") => { - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - RequestAuth::Header { name, value } if !has_header(&headers, name) => { - headers.push(((*name).to_string(), value.clone())); - } - RequestAuth::Bearer { .. } | RequestAuth::Header { .. } | RequestAuth::AwsSigV4 { .. } => {} - } - if !has_header(&headers, "content-type") { - headers.push(("Content-Type".to_string(), "application/json".to_string())); - } + let forwarded = string_headers("audio transcription", request.extra_headers)?; + let validated = + config.validate_environment(forwarded, &model, &request.optional_params, &env_lookup)?; + let environment = ValidatedEnvironment { + headers: with_default_headers(validated.headers, &[("Content-Type", "application/json")]), + auth: validated.auth, + }; let url = config.get_complete_url( request.api_base, &model, @@ -68,9 +64,7 @@ pub fn prepare_audio_transcription_provider_call( config, url, body: transformed.body, - upstream_headers: headers, - auth, - optional_params: request.optional_params, + environment, timeout: request.timeout, }) } diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index 0d87483c9bf..eff30c1e19a 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -1,7 +1,7 @@ use std::time::Duration; -use litellm_llms::base_llm::audio_transcription::transformation::{ - BaseAudioTranscriptionConfig, RequestAuth, +use litellm_llms::base_llm::{ + audio_transcription::transformation::BaseAudioTranscriptionConfig, auth::ValidatedEnvironment, }; use serde_json::{Map, Value}; @@ -23,9 +23,7 @@ pub struct ProviderAudioTranscriptionRequest { pub config: &'static dyn BaseAudioTranscriptionConfig, pub url: String, pub body: Value, - pub upstream_headers: Vec<(String, String)>, - pub auth: RequestAuth, - pub optional_params: Map, + pub environment: ValidatedEnvironment, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/chat_completions/error.rs b/litellm-rust/crates/core/src/chat_completions/error.rs deleted file mode 100644 index 81b57af2c6c..00000000000 --- a/litellm-rust/crates/core/src/chat_completions/error.rs +++ /dev/null @@ -1,43 +0,0 @@ -use litellm_llms::base_llm::chat::transformation::Error as LlmError; - -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, - }, - #[error("missing required field: {0}")] - MissingField(&'static str), - #[error("invalid provider: {0}")] - InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), - #[error(transparent)] - Auth(#[from] litellm_auth::Error), - #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] - Headers(#[from] litellm_http::request::HeaderError), - #[error(transparent)] - Http(#[from] litellm_http::Error), - #[error(transparent)] - Aws(#[from] litellm_auth_aws::Error), -} - -impl From for Error { - fn from(error: LlmError) -> Self { - match error { - LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, - LlmError::MissingField(field) => Self::MissingField(field), - LlmError::InvalidRequest(message) => Self::InvalidRequest(message), - LlmError::InvalidResponse(message) => Self::InvalidResponse(message), - LlmError::Unsupported(reason) => Self::Unsupported(reason), - LlmError::Auth(error) => Self::Auth(error), - } - } -} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index f3404fcaa8a..4a1cf7e193e 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body}; -use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData; +use litellm_llms::base_llm::{auth::resolve_auth, chat::transformation::ProviderChatResponseData}; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; @@ -13,10 +13,11 @@ use crate::{ pub(super) async fn execute_chat_completions_provider_call( http: &Client, + auth: &litellm_auth::AuthServices, request: ResolvedChatCompletionsRequest<'_>, ) -> Result { let request = prepare_provider_request(request)?; - let outbound = outbound_request(&request).await?; + let outbound = outbound_request(auth, &request).await?; let response = outbound.send(http).await.map_err(|err| { // Failing to establish the connection means the request never went out, @@ -69,28 +70,28 @@ pub(super) fn as_response_error(err: Error) -> Error { } pub(super) async fn outbound_request( + auth: &litellm_auth::AuthServices, request: &ProviderChatCompletionsRequest, ) -> Result { + let env_lookup = |key: &str| std::env::var(key).ok(); + let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?; crate::outbound::outbound_request( - &request.auth, + authenticated, request.url.clone(), - request.upstream_headers.clone(), &request.body, Some( request .timeout .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), ), - &request.optional_params, ) - .await .map_err(|error| match error { // Python drops the caller's copy and prefers a forwarded Authorization // over the signature, so leave the request to it. - Error::Http(litellm_http::Error::ComputedHeader(_)) => { + litellm_http::Error::ComputedHeader(_) => { Error::Unsupported("request forwards a header AWS SigV4 computes") } - other => other, + other => Error::Http(other), }) } diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index be22aea5669..d7003b5d22a 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -6,14 +6,13 @@ //! credentials, and it resolves the provider, translates the conversation, //! calls the provider, and returns a typed OpenAI-shaped response. -mod error; pub mod types; -pub use error::Error; +pub use crate::error::RouteError as Error; mod common_utils; pub(crate) mod handler; mod prepare; use handler::execute_chat_completions_provider_call; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{ClientVariant, HttpClientConfig}; use litellm_types::utils::ChatCompletionsResponse; use prepare::{parse_messages, resolve_provider_config, resolve_request}; use serde_json::{Map, Value}; @@ -21,13 +20,13 @@ use serde_json::{Map, Value}; use crate::chat_completions::types::ChatCompletionsRequest; pub async fn chat_completions( - pool: &HttpClientPool, + resources: &crate::resources::CoreResources, config: &HttpClientConfig, request: ChatCompletionsRequest<'_>, ) -> Result { let request = resolve_request(request)?; - let http = pool.client(config, ClientVariant::Provider)?; - execute_chat_completions_provider_call(&http, request).await + let http = resources.pool.client(config, ClientVariant::Provider)?; + execute_chat_completions_provider_call(&http, &resources.auth, request).await } /// Whether the core would accept this request, without resolving credentials or diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index b6425773964..8b091d4dd6c 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,6 +1,8 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; -use litellm_http::request::has_header; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; +use litellm_llms::base_llm::{ + auth::{ValidatedEnvironment, with_default_headers}, + chat::transformation::BaseConfig, +}; use litellm_types::llms::openai::ChatMessage; use serde_json::Value; @@ -67,59 +69,26 @@ fn validate_environment( request: &ResolvedChatCompletionsRequest<'_>, model: &str, config: &dyn BaseConfig, -) -> Result<(Vec<(String, String)>, RequestAuth), Error> { +) -> Result { let env_lookup = |key: &str| std::env::var(key).ok(); - let mut headers = string_headers(request.extra_headers.clone())?; - let auth = config.auth( + let forwarded = string_headers(request.extra_headers.clone())?; + let validated = config.validate_environment( + forwarded, request.api_key, model, &request.optional_params, &env_lookup, )?; - match &auth { - RequestAuth::Header { name, value } => { - // The deployment's credential replaces whatever the caller forwarded - // under the same name, mirroring Python's - // `{**headers, **anthropic_headers}`: letting a request header win - // would let its sender choose the principal the call bills to. - // - // The exception is a scheme the provider hands off to entirely, such - // as an Anthropic OAuth bearer, where Python drops `x-api-key` - // instead of resolving one. Re-adding it there would put the - // credential into a header the host removed on purpose. - if !config.defers_to_forwarded_auth(&headers) { - headers.retain(|(header, _)| !header.eq_ignore_ascii_case(name)); - headers.push(((*name).to_string(), value.clone())); - } - } - RequestAuth::Bearer { token } => { - // Bedrock's `get_request_headers` assigns `headers["Authorization"]` - // unconditionally once a bearer token resolves, so the deployment's - // identity outranks whatever the caller forwarded. Keeping the - // caller's would bill and authorize the call as a different - // principal than the same deployment uses on Python. - // - // The `Header` arm below keeps the opposite precedence on purpose: - // Anthropic's transform honours a forwarded OAuth bearer. - headers.retain(|(name, _)| !name.eq_ignore_ascii_case("authorization")); - headers.push(("authorization".to_string(), format!("Bearer {token}"))); - } - // SigV4 signs the serialized body, so the handler adds its headers. - RequestAuth::AwsSigV4 { .. } => {} - } - - for (name, value) in config.default_headers() { - if !has_header(&headers, name) { - headers.push(((*name).to_string(), (*value).to_string())); - } - } - Ok((headers, auth)) + Ok(ValidatedEnvironment { + headers: with_default_headers(validated.headers, config.default_headers()), + auth: validated.auth, + }) } pub(super) fn prepare_provider_request( request: ResolvedChatCompletionsRequest<'_>, ) -> Result { - let (headers, auth) = validate_environment(&request, &request.model, request.config)?; + let environment = validate_environment(&request, &request.model, request.config)?; let model = request.model; let config = request.config; let env_lookup = |key: &str| std::env::var(key).ok(); @@ -130,23 +99,22 @@ pub(super) fn prepare_provider_request( &env_lookup, )?; let transformed = - config.transform_request(&model, request.messages, request.optional_params.clone())?; + config.transform_request(&model, request.messages, request.optional_params)?; Ok(ProviderChatCompletionsRequest { model, config, url, body: transformed.body, - upstream_headers: headers, - auth, - optional_params: request.optional_params, + environment, timeout: request.timeout, }) } #[cfg(test)] mod tests { - use litellm_llms::base_llm::chat::transformation::RequestAuth; + use litellm_auth::CredentialPlacement; + use litellm_llms::base_llm::auth::{AuthScheme, resolve_auth}; use serde_json::{Map, Value, json}; use super::{prepare_provider_request, resolve_request}; @@ -161,6 +129,20 @@ mod tests { prepare_provider_request(resolve_request(request)?) } + /// The headers as they go on the wire, credential applied. + fn wire_headers(prepared: &ProviderChatCompletionsRequest) -> Vec<(String, String)> { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() + .block_on(resolve_auth( + &litellm_auth::AuthServices::default(), + prepared.environment.clone(), + &|_| None, + )) + .unwrap() + .headers + } + fn request<'a>( model: &'a str, provider: Option<&'a str>, @@ -227,19 +209,16 @@ mod tests { )) .expect("prepares"); assert!( - prepared - .upstream_headers - .contains(&("x-api-key".to_string(), "sk-test".to_string())) + wire_headers(&prepared).contains(&("x-api-key".to_string(), "sk-test".to_string())) ); assert!( - prepared - .upstream_headers + wire_headers(&prepared) .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) ); assert!(matches!( - prepared.auth, - RequestAuth::Header { - name: "x-api-key", + prepared.environment.auth, + AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), .. } )); @@ -261,12 +240,12 @@ mod tests { json!("sk-caller"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers + let headers = wire_headers(&prepared); + let keys: Vec<_> = headers .iter() .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys.len(), 1, "got {:?}", headers); assert_eq!(keys[0].1, "sk-test"); } @@ -290,16 +269,14 @@ mod tests { ])); let prepared = prepare_chat_completions_call(call).expect("prepares"); assert!( - !prepared - .upstream_headers + !wire_headers(&prepared) .iter() .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), "the resolved key must not be applied over an OAuth bearer, got {:?}", - prepared.upstream_headers + wire_headers(&prepared) ); assert!( - prepared - .upstream_headers + wire_headers(&prepared) .iter() .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == "Bearer sk-ant-oat01-token") @@ -322,21 +299,20 @@ mod tests { ("X-Api-Key".to_string(), json!("sk-caller")), ])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers + let headers = wire_headers(&prepared); + let keys: Vec<_> = headers .iter() .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys.len(), 1, "got {:?}", headers); assert_eq!(keys[0].1, "sk-test"); assert!( - prepared - .upstream_headers + wire_headers(&prepared) .iter() .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == "Bearer unrelated"), "the unrelated authorization must survive, got {:?}", - prepared.upstream_headers + wire_headers(&prepared) ); } @@ -435,18 +411,16 @@ mod tests { prepared.url, "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" ); - assert_eq!( - prepared.auth, - RequestAuth::AwsSigV4 { - region: "us-east-1".to_string(), - service: "bedrock", - } - ); + assert!(matches!( + &prepared.environment.auth, + AuthScheme::AwsSigV4 { region, service: "bedrock", .. } if region == "us-east-1" + )); // SigV4 signs the serialized body, so prepare must not have added an - // Authorization header; the handler does it. + // Authorization header; the signer does it over the bytes sent. assert!( !prepared - .upstream_headers + .environment + .headers .iter() .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) ); @@ -475,9 +449,12 @@ mod tests { json!("abc-123"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = crate::chat_completions::handler::outbound_request(&prepared) - .await - .expect("signs"); + let signed = crate::chat_completions::handler::outbound_request( + &litellm_auth::AuthServices::default(), + &prepared, + ) + .await + .expect("signs"); let authorization = signed .header("authorization") @@ -525,9 +502,12 @@ mod tests { call.api_key = None; call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = crate::chat_completions::handler::outbound_request(&prepared) - .await - .expect_err("{forwarded} should decline instead of being signed"); + let error = crate::chat_completions::handler::outbound_request( + &litellm_auth::AuthServices::default(), + &prepared, + ) + .await + .expect_err("{forwarded} should decline instead of being signed"); assert!( matches!(error, Error::Unsupported(_)), "{forwarded} declined as {error:?}, which the host would not fall back on" @@ -552,8 +532,8 @@ mod tests { json!("Bearer caller-supplied"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let authorizations: Vec<_> = prepared - .upstream_headers + let headers = wire_headers(&prepared); + let authorizations: Vec<_> = headers .iter() .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) .map(|(_, value)| value.as_str()) @@ -585,16 +565,15 @@ mod tests { json!("Bearer sk-ant-oat01-forwarded"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers + let headers = wire_headers(&prepared); + let keys: Vec<_> = headers .iter() .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) .map(|(_, value)| value.as_str()) .collect(); - assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); + assert!(keys.is_empty(), "got {:?}", headers); assert!( - prepared - .upstream_headers + wire_headers(&prepared) .iter() .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == "Bearer sk-ant-oat01-forwarded") @@ -613,15 +592,13 @@ mod tests { json!({"maxTokens": 16}), )) .expect("prepares"); - assert_eq!( - prepared.auth, - RequestAuth::Bearer { - token: "sk-test".to_string() - } - ); + assert!(matches!( + &prepared.environment.auth, + AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret } + if secret.expose() == "sk-test" + )); assert!( - prepared - .upstream_headers + wire_headers(&prepared) .iter() .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == "Bearer sk-test"), diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 3b74cf5dace..66e9498c749 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; +use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig}; use litellm_types::llms::openai::ChatMessage; use serde_json::{Map, Value}; @@ -37,8 +37,8 @@ pub struct ProviderChatCompletionsRequest { pub config: &'static dyn BaseConfig, pub url: String, pub body: Value, - pub upstream_headers: Vec<(String, String)>, - pub auth: RequestAuth, - pub optional_params: Map, + /// The forwarded and default headers plus how the call authenticates; the credential + /// itself is applied when the request is sent. + pub environment: ValidatedEnvironment, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 455c3258799..c14b54679ff 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -5,10 +5,6 @@ pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com"; /// timeout from the caller still overrides this on the request builder. pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; -/// Provider name used for Anthropic Messages when a deployment's provider model -/// does not carry an explicit provider prefix. -pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; - /// Full-request timeout ceiling for chat completions provider calls, in /// seconds. Mirrors the Python chat completions default. pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index eb4cd2367ec..0d3de6e57c1 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -1,15 +1,166 @@ -use litellm_llms::base_llm::ocr::error::Error as OcrError; +//! One error for every route in this crate. OCR still carries its own, richer enum. +//! +//! A variant is declared by the layer that produces it and nested here as is: +//! credentials by `litellm_auth` (AWS folds into it at that crate's boundary), the wire by +//! `litellm_http`, secrets by `litellm_secrets`. The transformation layer's [`LlmError`] +//! maps onto the same-named variants once, here, so no route re-declares them. -#[derive(Debug, thiserror::Error)] -pub enum Error { +use std::sync::Arc; + +use litellm_http::transport::Error as TransportError; +use litellm_llms::Error as LlmError; + +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum RouteError { + #[error("expected {expected}, got {actual}")] + InvalidType { + expected: &'static str, + actual: &'static str, + }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid provider: {0}")] + InvalidProvider(String), + #[error("invalid request: {0}")] + InvalidRequest(String), + #[error("invalid response: {0}")] + InvalidResponse(String), + #[error("unsupported by the rust path: {0}")] + Unsupported(&'static str), #[error(transparent)] - Ocr(#[from] OcrError), + Auth(#[from] litellm_auth::Error), #[error(transparent)] - Messages(#[from] crate::messages::Error), + Transport(#[from] TransportError), #[error(transparent)] - ChatCompletions(#[from] crate::chat_completions::Error), + Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] - AudioTranscription(#[from] crate::audio_transcription::Error), + Http(#[from] litellm_http::Error), #[error(transparent)] - Responses(#[from] crate::responses::Error), + Secret(#[from] SecretError), +} + +/// Whether the provider had already been called when the route failed. Before the send, a +/// host may retry on another path; after it, the provider has done the work and billed for it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Phase { + BeforeSend, + AfterSend, +} + +impl RouteError { + pub fn phase(&self) -> Phase { + match self { + Self::InvalidResponse(_) + | Self::Transport(TransportError::Http { .. } | TransportError::Network(_)) => { + Phase::AfterSend + } + Self::Transport(TransportError::Connect(_)) + | Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidProvider(_) + | Self::InvalidRequest(_) + | Self::Unsupported(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) => Phase::BeforeSend, + } + } + + /// The caller's request is what is wrong, as opposed to the environment, the wire, or + /// the provider's answer. + pub fn is_request(&self) -> bool { + match self { + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidProvider(_) + | Self::InvalidRequest(_) + | Self::Unsupported(_) + | Self::Headers(_) => true, + Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }), + Self::InvalidResponse(_) | Self::Transport(_) | Self::Http(_) | Self::Secret(_) => { + false + } + } + } +} + +impl From for RouteError { + fn from(error: LlmError) -> Self { + match error { + LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, + LlmError::MissingField(field) => Self::MissingField(field), + LlmError::InvalidRequest(message) => Self::InvalidRequest(message), + LlmError::InvalidResponse(message) => Self::InvalidResponse(message), + LlmError::Unsupported(reason) => Self::Unsupported(reason), + LlmError::Auth(error) => Self::Auth(error), + } + } +} + +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct SecretError(Arc); + +impl SecretError { + pub fn source_error(&self) -> &litellm_secrets::Error { + &self.0 + } +} + +impl From for RouteError { + fn from(error: litellm_secrets::Error) -> Self { + Self::Secret(SecretError(Arc::new(error))) + } +} + +impl PartialEq for SecretError { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for SecretError {} + +#[cfg(test)] +mod tests { + use super::{Phase, RouteError}; + use litellm_http::transport::Error as TransportError; + + #[test] + fn only_a_provider_answer_or_a_lost_connection_counts_as_after_send() { + let after = [ + RouteError::InvalidResponse("bad json".into()), + RouteError::Transport(TransportError::Http { + status: 500, + body: "boom".into(), + }), + RouteError::Transport(TransportError::Network("reset".into())), + ]; + for error in after { + assert_eq!(error.phase(), Phase::AfterSend, "{error:?}"); + } + let before = [ + RouteError::Transport(TransportError::Connect("refused".into())), + RouteError::Unsupported("streaming"), + RouteError::Auth(litellm_auth::Error::InvalidHeader), + ]; + for error in before { + assert_eq!(error.phase(), Phase::BeforeSend, "{error:?}"); + } + } + + #[test] + fn a_missing_api_key_is_the_environment_not_the_request() { + assert!( + !RouteError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + .is_request() + ); + assert!(RouteError::Auth(litellm_auth::Error::InvalidHeader).is_request()); + assert!(RouteError::InvalidRequest("top_k".into()).is_request()); + assert!(!RouteError::InvalidResponse("bad json".into()).is_request()); + } } diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index afe5ea595aa..d373262ae7d 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -5,6 +5,7 @@ pub mod error; pub mod messages; pub mod ocr; mod outbound; +pub mod resources; pub mod responses; -pub use error::Error; +pub use error::{Phase, RouteError}; diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 95142e87519..fc3bbb36098 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -4,20 +4,34 @@ use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, + bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, }; use serde_json::{Map, Value}; +use strum::{EnumString, IntoStaticStr}; use super::Error; const HEADER_CONTEXT: &str = "messages"; -pub(super) fn messages_provider_config( - provider: &str, -) -> Option<&'static dyn BaseAnthropicMessagesConfig> { - match provider { - "anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG), - "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), - _ => None, +#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)] +#[strum(serialize_all = "snake_case")] +pub(crate) enum MessagesProvider { + Anthropic, + AzureAi, + Bedrock, +} + +impl MessagesProvider { + pub(crate) fn as_str(self) -> &'static str { + self.into() + } + + pub(crate) fn config(self) -> &'static dyn BaseAnthropicMessagesConfig { + match self { + Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG, + Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG, + Self::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + } } } @@ -31,14 +45,26 @@ pub(super) fn string_headers( mod tests { use serde_json::json; - use super::{messages_provider_config, string_headers, truncate_error_body}; + use rstest::rstest; + + use super::{MessagesProvider, string_headers, truncate_error_body}; use crate::messages::Error; + #[rstest] + #[case::anthropic("anthropic", MessagesProvider::Anthropic)] + #[case::azure_ai("azure_ai", MessagesProvider::AzureAi)] + #[case::bedrock("bedrock", MessagesProvider::Bedrock)] + fn provider_round_trips_through_its_python_name( + #[case] name: &str, + #[case] provider: MessagesProvider, + ) { + assert_eq!(name.parse::(), Ok(provider)); + assert_eq!(provider.as_str(), name); + } + #[test] - fn provider_config_resolves_anthropic_and_azure_ai() { - assert!(messages_provider_config("anthropic").is_some()); - assert!(messages_provider_config("azure_ai").is_some()); - assert!(messages_provider_config("openai").is_none()); + fn provider_without_a_messages_config_is_rejected() { + assert!("openai".parse::().is_err()); } #[test] diff --git a/litellm-rust/crates/core/src/messages/error.rs b/litellm-rust/crates/core/src/messages/error.rs deleted file mode 100644 index 76f8813e330..00000000000 --- a/litellm-rust/crates/core/src/messages/error.rs +++ /dev/null @@ -1,82 +0,0 @@ -use std::sync::Arc; - -use litellm_llms::base_llm::chat::transformation::Error as LlmError; - -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("invalid provider: {0}")] - InvalidProvider(String), - #[error("missing required field: {0}")] - MissingField(&'static str), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("unsupported by the Rust messages route: {0}")] - Unsupported(&'static str), - #[error(transparent)] - Auth(#[from] litellm_auth::Error), - #[error(transparent)] - Client(#[from] litellm_http::Error), - #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] - Headers(#[from] litellm_http::request::HeaderError), - #[error(transparent)] - Secret(#[from] SecretError), -} - -#[derive(Clone, Debug, thiserror::Error)] -#[error(transparent)] -pub struct SecretError(Arc); - -impl SecretError { - pub fn source_error(&self) -> &litellm_secrets::Error { - &self.0 - } -} - -impl From for Error { - fn from(error: litellm_secrets::Error) -> Self { - Self::Secret(SecretError(Arc::new(error))) - } -} - -impl PartialEq for SecretError { - fn eq(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.0, &other.0) - } -} - -impl Eq for SecretError {} - -impl From for Error { - fn from(error: LlmError) -> Self { - match error { - error @ LlmError::InvalidType { .. } => Self::InvalidRequest(error.to_string()), - LlmError::MissingField(field) => Self::MissingField(field), - LlmError::InvalidRequest(message) => Self::InvalidRequest(message), - LlmError::InvalidResponse(message) => Self::InvalidResponse(message), - LlmError::Unsupported(reason) => Self::Unsupported(reason), - LlmError::Auth(error) => Self::Auth(error), - } - } -} - -impl Error { - pub fn is_request(&self) -> bool { - match self { - Self::InvalidProvider(_) - | Self::MissingField(_) - | Self::InvalidRequest(_) - | Self::Unsupported(_) - | Self::Headers(_) => true, - Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }), - _ => false, - } - } - - pub fn is_response(&self) -> bool { - matches!(self, Self::InvalidResponse(_)) - } -} diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index f90cb8cb454..650447d5abd 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,12 +1,14 @@ use std::time::Duration; -use litellm_http::{request::http_request, transport::Error as TransportError}; -use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; +use litellm_http::transport::Error as TransportError; +use litellm_llms::base_llm::{ + anthropic_messages::transformation::BaseAnthropicMessagesConfig, auth::Authenticated, +}; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::Value; use super::{Error, common_utils::truncate_error_body}; -use crate::constants::MESSAGES_TIMEOUT_SECS; +use crate::{constants::MESSAGES_TIMEOUT_SECS, outbound::outbound_request}; pub(super) fn network(error: reqwest::Error) -> Error { Error::Transport(TransportError::Network(error.to_string())) @@ -14,20 +16,18 @@ pub(super) fn network(error: reqwest::Error) -> Error { pub(super) async fn send( http: &litellm_http::Client, + authenticated: Authenticated, url: &str, - headers: &[(String, String)], body: &Value, timeout: Option, ) -> Result { - let encoded = serde_json::to_vec(body) - .map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?; - let builder = headers.iter().fold( - http.post(url) - .body(encoded) - .timeout(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), - |builder, (key, value)| builder.header(key, value), - ); - http_request(builder).await.map_err(network) + let request = outbound_request( + authenticated, + url.to_string(), + body, + Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), + )?; + request.send(http).await.map_err(network) } pub(super) async fn provider_error(response: reqwest::Response) -> Error { diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 5cb83b4e34d..3c081ff7bbd 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -4,49 +4,29 @@ //! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs //! it in process for a caller that already holds the request and wants the message. -mod error; pub mod types; -pub use error::Error; +pub use crate::error::RouteError as Error; mod common_utils; mod handler; mod prepare; pub mod route; use std::sync::Arc; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{ClientVariant, HttpClientConfig}; use litellm_secrets::source::EnvironmentSecrets; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; -use serde_json::Value; - -use crate::messages::types::MessagesRequest; pub async fn messages( - pool: &HttpClientPool, + resources: &crate::resources::CoreResources, config: &HttpClientConfig, - request: MessagesRequest<'_>, + call: MessagesCall, ) -> Result { - let Value::Object(body) = request.body else { - return Err(Error::InvalidRequest( - "messages body must be an object".into(), - )); - }; - let call = MessagesCall { - model: request.model.into(), - body, - api_key: request.api_key.map(Into::into), - api_base: request.api_base.map(Into::into), - custom_llm_provider: request.custom_llm_provider.map(Into::into), - extra_headers: request.extra_headers, - provider_specific_header: request.provider_specific_header, - timeout: request.timeout, - shaping: request.shaping, - }; let secrets = Arc::new(EnvironmentSecrets::python_compatible( - pool.client(config, ClientVariant::Provider)?, + resources.pool.client(config, ClientVariant::Provider)?, )); match litellm_host::run::run( - messages_machine(pool, config, secrets)?, + messages_machine(resources, config, secrets)?, &LocalMessagesHost::new(call), ) .await? diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 84884ab279e..7cd01a3a84c 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -6,29 +6,29 @@ use litellm_core_utils::{ }; use litellm_llms::{ anthropic::messages::handler::shape_anthropic_messages_request, - base_llm::anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, MessagesTransformContext, + base_llm::{ + anthropic_messages::transformation::MessagesTransformContext, + auth::{ValidatedEnvironment, with_default_headers}, }, }; use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; -use serde_json::{Map, Value}; use super::{ Error, - common_utils::{messages_provider_config, string_headers}, + common_utils::{MessagesProvider, string_headers}, + route::MessagesCall, + types::ProviderMessagesRequest, }; -use crate::messages::types::{MessagesRequest, ProviderMessagesRequest}; -pub(super) struct ResolvedProvider<'a> { - pub(super) model: &'a str, - pub(super) provider: &'a str, - pub(super) config: &'static dyn BaseAnthropicMessagesConfig, +pub(super) struct ResolvedProvider { + pub(super) model: String, + pub(super) provider: MessagesProvider, } -pub(super) fn resolve_provider<'a>( - model: &'a str, - custom_llm_provider: Option<&'a str>, -) -> Result, Error> { +pub(super) fn resolve_provider( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result { let CustomLlmProvider { model, custom_llm_provider: provider, @@ -44,79 +44,79 @@ pub(super) fn resolve_provider<'a>( "unable to resolve custom_llm_provider for messages request".to_string(), ) })?; - let config = messages_provider_config(provider) - .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; + let provider = provider + .parse() + .map_err(|_| Error::InvalidProvider(provider.to_string()))?; Ok(ResolvedProvider { - model, + model: model.to_string(), provider, - config, }) } pub(super) fn prepare_provider_request( - request: MessagesRequest<'_>, - resolved: ResolvedProvider<'_>, + call: MessagesCall, + resolved: ResolvedProvider, secrets: &dyn Lookup, ) -> Result { - let ResolvedProvider { - model, - provider, - config, - } = resolved; - let model = model.to_string(); + let ResolvedProvider { model, provider } = resolved; + let MessagesCall { + body, + api_key, + api_base, + extra_headers, + provider_specific_header, + timeout, + shaping, + .. + } = call; + let config = provider.config(); let env_lookup = |key: &str| secrets.get(key); - let typed_request: AnthropicMessagesRequest = - serde_json::from_value(request.body).map_err(invalid_request)?; let sanitized = shape_anthropic_messages_request( - AnthropicMessagesRequest { - model: model.clone(), - ..typed_request - }, - request.shaping.reasoning_auto_summary, + AnthropicMessagesRequest { model, ..body }, + shaping.reasoning_auto_summary, )?; - let trimmed = - without_additional_drop_params(sanitized, &request.shaping.additional_drop_params)?; + let trimmed = without_additional_drop_params(sanitized, &shaping.additional_drop_params)?; let transformed = config.transform_anthropic_messages_request( trimmed, - &MessagesTransformContext::new(request.shaping.capabilities, request.shaping.drop_params), + &MessagesTransformContext::new(shaping.capabilities, shaping.drop_params), )?; - let scoped = get_provider_specific_headers(request.provider_specific_header.as_ref(), provider); + let scoped = + get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); let forwarded = string_headers(Some( - request - .extra_headers - .into_iter() - .flatten() - .chain(scoped) - .collect(), + extra_headers.into_iter().flatten().chain(scoped).collect(), ))?; - let authenticated = config.authenticate(forwarded, request.api_key, &env_lookup)?; - let headers = config.request_headers( - with_default_headers(authenticated, config.default_headers()), - &transformed, - ); + let validated = config.validate_environment( + forwarded, + api_key.as_deref(), + &transformed.model, + &env_lookup, + )?; + let environment = ValidatedEnvironment { + headers: config.request_headers( + with_default_headers(validated.headers, config.default_headers()), + &transformed, + ), + auth: validated.auth, + }; - let body = serde_json::to_value(transformed).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) - })?; - - let url = config.get_complete_url(request.api_base, &model, &env_lookup)?; + let url = if transformed.params.stream == Some(true) { + config.complete_stream_url(api_base.as_deref(), &transformed.model, &env_lookup)? + } else { + config.get_complete_url(api_base.as_deref(), &transformed.model, &env_lookup)? + }; Ok(ProviderMessagesRequest { - provider: provider.to_string(), - model, - config, + provider, url, - body, - upstream_headers: headers, - timeout: request.timeout, + body: transformed, + environment, + timeout, }) } -fn invalid_request(err: serde_json::Error) -> Error { +pub(super) fn invalid_request(err: serde_json::Error) -> Error { Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) } @@ -127,45 +127,22 @@ fn without_additional_drop_params( if paths.is_empty() { return Ok(request); } - let Value::Object(fields) = serde_json::to_value(request).map_err(invalid_request)? else { - return Err(Error::InvalidRequest( - "Anthropic messages request did not serialize to an object".to_string(), - )); - }; - let (required, optional): (Map, Map) = fields - .into_iter() - .partition(|(key, _)| matches!(key.as_str(), "model" | "messages")); - let trimmed = paths.iter().fold(Value::Object(optional), |body, path| { - delete_nested_value(body, path) - }); - let merged: Map = required - .into_iter() - .chain(trimmed.as_object().cloned().unwrap_or_default()) - .collect(); - serde_json::from_value(Value::Object(merged)).map_err(invalid_request) -} - -fn with_default_headers( - headers: Vec<(String, String)>, - defaults: &[(&str, &str)], -) -> Vec<(String, String)> { - let missing: Vec<(String, String)> = defaults + let params = serde_json::to_value(request.params).map_err(invalid_request)?; + let trimmed = paths .iter() - .filter(|(name, _)| { - !headers - .iter() - .any(|(header, _)| header.eq_ignore_ascii_case(name)) - }) - .map(|(name, value)| ((*name).to_string(), (*value).to_string())) - .collect(); - headers.into_iter().chain(missing).collect() + .fold(params, |params, path| delete_nested_value(params, path)); + Ok(AnthropicMessagesRequest { + params: serde_json::from_value(trimmed).map_err(invalid_request)?, + ..request + }) } #[cfg(test)] mod tests { + use litellm_llms::base_llm::auth::resolve_auth; use litellm_types::utils::ProviderSpecificHeaders; use rstest::{fixture, rstest}; - use serde_json::json; + use serde_json::{Map, Value, json}; use super::*; use crate::messages::types::MessagesShaping; @@ -175,16 +152,34 @@ mod tests { MessagesShaping::default() } - fn prepare(request: MessagesRequest<'_>) -> Result { - prepare_with_secrets(request, &|_: &str| None) + fn body(value: Value) -> AnthropicMessagesRequest { + serde_json::from_value(value).unwrap() + } + + fn prepare(call: MessagesCall) -> Result { + prepare_with_secrets(call, &|_: &str| None) } fn prepare_with_secrets( - request: MessagesRequest<'_>, + call: MessagesCall, secrets: &dyn Lookup, ) -> Result { - let resolved = resolve_provider(request.model, request.custom_llm_provider)?; - prepare_provider_request(request, resolved, secrets) + let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; + prepare_provider_request(call, resolved, secrets) + } + + /// The headers as they go on the wire, credential applied. + fn wire_headers(prepared: &ProviderMessagesRequest) -> Vec<(String, String)> { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() + .block_on(resolve_auth( + &litellm_auth::AuthServices::default(), + prepared.environment.clone(), + &|_| None, + )) + .unwrap() + .headers } #[rstest] @@ -221,12 +216,13 @@ mod tests { .map(|(_, value)| value.to_string()) }; let prepared = prepare_with_secrets( - MessagesRequest { - model: "claude-test", - body: json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + MessagesCall { + body: body( + json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), api_key: None, api_base: None, - custom_llm_provider: Some("anthropic"), + custom_llm_provider: Some("anthropic".into()), extra_headers: None, provider_specific_header: None, timeout: None, @@ -235,8 +231,8 @@ mod tests { &lookup, ) .unwrap(); - let auth: Vec<(&str, &str)> = prepared - .upstream_headers + let headers = wire_headers(&prepared); + let auth: Vec<(&str, &str)> = headers .iter() .filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization")) .map(|(name, value)| (name.as_str(), value.as_str())) @@ -247,48 +243,18 @@ mod tests { ); } - fn prepared_body(body: Value, shaping: MessagesShaping) -> Result { - prepare(MessagesRequest { - model: "anthropic/claude-test", - body, - api_key: Some("sk-test"), - api_base: Some("https://anthropic.test"), - custom_llm_provider: Some("anthropic"), + fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result { + prepare(MessagesCall { + body: body(fields), + api_key: Some("sk-test".into()), + api_base: Some("https://anthropic.test".into()), + custom_llm_provider: Some("anthropic".into()), extra_headers: None, provider_specific_header: None, timeout: None, shaping, }) - .map(|prepared| prepared.body) - } - - #[rstest] - #[case::nothing_forwarded( - &[], - &[("x-version", "1"), ("content-type", "application/json")], - &[("x-version", "1"), ("content-type", "application/json")], - )] - #[case::forwarded_header_wins_in_any_case( - &[("X-Version", "custom"), ("x-api-key", "k")], - &[("x-version", "1"), ("content-type", "application/json")], - &[("X-Version", "custom"), ("x-api-key", "k"), ("content-type", "application/json")], - )] - #[case::no_defaults(&[("x-api-key", "k")], &[], &[("x-api-key", "k")])] - fn default_headers_fill_only_missing_names( - #[case] forwarded: &[(&str, &str)], - #[case] defaults: &[(&str, &str)], - #[case] expected: &[(&str, &str)], - ) { - let owned = |headers: &[(&str, &str)]| -> Vec<(String, String)> { - headers - .iter() - .map(|(name, value)| ((*name).to_string(), (*value).to_string())) - .collect() - }; - assert_eq!( - with_default_headers(owned(forwarded), defaults), - owned(expected) - ); + .map(|prepared| serde_json::to_value(prepared.body).unwrap()) } #[rstest] @@ -380,20 +346,22 @@ mod tests { {"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}} ])) .unwrap(); - let prepared = prepare(MessagesRequest { - model, - body: json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), - api_key: Some("sk-test"), - api_base: Some("https://resource.services.ai.azure.com"), - custom_llm_provider, - extra_headers: Some(serde_json::from_value(json!({"x-priority": "extra"})).unwrap()), + let prepared = prepare(MessagesCall { + body: body( + json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), + api_key: Some("sk-test".into()), + api_base: Some("https://resource.services.ai.azure.com".into()), + custom_llm_provider: custom_llm_provider.map(Into::into), + extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])), provider_specific_header: Some(configured), timeout: None, shaping, }) .unwrap(); let caller_headers: Vec<(&str, &str)> = prepared - .upstream_headers + .environment + .headers .iter() .filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped")) .map(|(name, value)| (name.as_str(), value.as_str())) diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 7f6589cdf3e..f9267dce755 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -5,6 +5,7 @@ use std::{ }; use bytes::Bytes; +use futures_util::StreamExt; use litellm_auth::SecretValue; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, @@ -12,10 +13,16 @@ use litellm_host::{ machine::{CallMachine, HostChannel, MachineFault}, protocol::Protocol, }; -use litellm_http::{Client, ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{Client, ClientVariant, HttpClientConfig}; +use litellm_llms::base_llm::{ + anthropic_messages::streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, + auth::{Authenticated, resolve_auth}, +}; use litellm_secrets::source::SecretSource; use litellm_types::{ - llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse, + llms::anthropic_messages::{ + anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, + }, utils::ProviderSpecificHeaders, }; use serde_json::{Map, Value}; @@ -23,15 +30,13 @@ use serde_json::{Map, Value}; use super::{ Error, handler::{decode_response, network, provider_error, send}, - prepare::{prepare_provider_request, resolve_provider}, - types::{MessagesRequest, MessagesShaping}, + prepare::{invalid_request, prepare_provider_request, resolve_provider}, + types::MessagesShaping, }; -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; /// The caller's request as the host projects it. pub struct MessagesCall { - pub model: String, - pub body: Map, + pub body: AnthropicMessagesRequest, pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, @@ -41,10 +46,9 @@ pub struct MessagesCall { pub shaping: MessagesShaping, } -impl MessagesCall { - fn streams(&self) -> bool { - self.body.get("stream").and_then(Value::as_bool) == Some(true) - } +/// Parses a caller's raw body, failing the way the route fails for any invalid request. +pub fn messages_body(body: Map) -> Result { + serde_json::from_value(Value::Object(body)).map_err(invalid_request) } pub enum MessagesOutput { @@ -110,90 +114,93 @@ impl Host for LocalMessagesHost { } pub fn messages_machine( - pool: &HttpClientPool, + resources: &crate::resources::CoreResources, config: &HttpClientConfig, secrets: Arc, ) -> Result { - let http = pool.client(config, ClientVariant::Provider)?; + let http = resources.pool.client(config, ClientVariant::Provider)?; + let auth = resources.auth.clone(); Ok(CallMachine::new(move |host| { - Box::pin(execute(host, http.clone(), secrets.clone())) + Box::pin(execute(host, http.clone(), auth.clone(), secrets.clone())) })) } async fn execute( host: MessagesHost, http: Client, + auth: Arc, secrets: Arc, ) -> Result { let call = host.project().await?; - let stream = call.streams(); - let resolved = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?; - let secrets = secrets.resolve(resolved.config.secret_names()).await?; - let request = prepare_provider_request( - MessagesRequest { - model: &call.model, - body: Value::Object(call.body.clone()), - api_key: call.api_key.as_deref(), - api_base: call.api_base.as_deref(), - custom_llm_provider: call.custom_llm_provider.as_deref(), - extra_headers: call.extra_headers.clone(), - provider_specific_header: call.provider_specific_header.clone(), - timeout: call.timeout, - shaping: call.shaping.clone(), - }, - resolved, - secrets.as_ref(), - )?; - if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(Error::Unsupported("streaming messages for this provider")); - } + let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; + let secrets = secrets + .resolve(resolved.provider.config().secret_names()) + .await?; + let api_key = call.api_key.clone().map(SecretValue::new); + let request = prepare_provider_request(call, resolved, secrets.as_ref())?; let context = RequestContext { - model: request.model.clone(), - custom_llm_provider: request.provider.clone(), - optional_params: Value::Object( - request - .body - .as_object() - .into_iter() - .flatten() - .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) - .map(|(name, value)| (name.clone(), value.clone())) - .collect(), - ), + model: request.body.model.clone(), + custom_llm_provider: request.provider.as_str().to_string(), + optional_params: serde_json::to_value(&request.body.params).map_err(serialize_failure)?, secret_fields: Vec::new(), - api_key: call.api_key.clone().map(SecretValue::new), + api_key, }; + let stream = request.body.params.stream == Some(true); + let config = request.provider.config(); + let body = serde_json::to_value(&request.body).map_err(serialize_failure)?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let authenticated = resolve_auth(&auth, request.environment, &env_lookup).await?; let wire = host .before_send( WireRequest { url: request.url, - headers: request.upstream_headers, - body: request.body, + headers: authenticated.headers, + body, }, context, ) .await?; - let response = send(&http, &wire.url, &wire.headers, &wire.body, request.timeout).await?; + let response = send( + &http, + Authenticated { + headers: wire.headers, + signer: authenticated.signer, + }, + &wire.url, + &wire.body, + request.timeout, + ) + .await?; if !response.status().is_success() { return Err(provider_error(response).await); } if stream { - return relay(&host, response).await; + return relay(&host, response, config.stream_decoder()).await; } let text = response.text().await.map_err(network)?; host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: text.clone() }, }) .await?; - decode_response(request.config, &request.model, &text) + decode_response(config, &request.body.model, &text) .map(|message| MessagesOutput::Message(Box::new(message))) } +fn serialize_failure(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) +} + /// Hands each upstream chunk to the caller as it arrives. A caller that stops reading /// ends the upstream read, and the call completes with what it delivered. +/// +/// A host on Anthropic SSE is relayed byte for byte. A host on another wire is decoded into +/// Anthropic stream events and re-encoded as Anthropic SSE. async fn relay( host: &MessagesHost, - mut response: reqwest::Response, + response: reqwest::Response, + decoder: Option, ) -> Result { let head = MessagesStreamHead { headers: response @@ -205,6 +212,16 @@ async fn relay( if host.open(head).await? == Demand::Detached { return Ok(MessagesOutput::Streamed); } + match decoder { + None => relay_bytes(host, response).await, + Some(decode) => relay_events(host, response, decode).await, + } +} + +async fn relay_bytes( + host: &MessagesHost, + mut response: reqwest::Response, +) -> Result { while let Some(chunk) = response.chunk().await.map_err(network)? { if host.deliver(chunk).await? == Demand::Detached { break; @@ -212,3 +229,26 @@ async fn relay( } Ok(MessagesOutput::Streamed) } + +async fn relay_events( + host: &MessagesHost, + response: reqwest::Response, + decode: StreamDecoder, +) -> Result { + let bytes: ByteStream = futures_util::stream::unfold(response, |mut response| async move { + match response.chunk().await { + Ok(Some(chunk)) => Some((Ok(chunk), response)), + Ok(None) => None, + Err(error) => Some((Err(std::io::Error::other(error)), response)), + } + }) + .boxed(); + let mut events = decode(bytes); + while let Some(event) = events.next().await { + let chunk = encode_anthropic_sse(&event?)?; + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) +} diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index 4a5dd2926e0..006b1db4efb 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,12 +1,12 @@ use std::time::Duration; use litellm_llms::{ - anthropic::common_utils::AnthropicModelCapabilities, - base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, + anthropic::common_utils::AnthropicModelCapabilities, base_llm::auth::ValidatedEnvironment, }; -use litellm_types::utils::ProviderSpecificHeaders; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; + +use super::common_utils::MessagesProvider; #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { @@ -20,33 +20,21 @@ pub struct MessagesShaping { pub additional_drop_params: Vec, } -pub struct MessagesRequest<'a> { - pub model: &'a str, - pub body: Value, - pub api_key: Option<&'a str>, - pub api_base: Option<&'a str>, - pub custom_llm_provider: Option<&'a str>, - pub extra_headers: Option>, - pub provider_specific_header: Option, - pub timeout: Option, - pub shaping: MessagesShaping, -} - -pub struct ProviderMessagesRequest { - pub provider: String, - pub model: String, - pub config: &'static dyn BaseAnthropicMessagesConfig, - pub url: String, - pub body: Value, - pub upstream_headers: Vec<(String, String)>, - pub timeout: Option, +pub(crate) struct ProviderMessagesRequest { + pub(crate) provider: MessagesProvider, + pub(crate) url: String, + pub(crate) body: AnthropicMessagesRequest, + /// The forwarded, default and feature headers plus how the call authenticates; the + /// credential itself is applied when the request is sent. + pub(crate) environment: ValidatedEnvironment, + pub(crate) timeout: Option, } #[cfg(test)] mod tests { use litellm_llms::anthropic::common_utils::SupportedEffortTiers; use rstest::rstest; - use serde_json::json; + use serde_json::{Value, json}; use super::*; diff --git a/litellm-rust/crates/core/src/outbound.rs b/litellm-rust/crates/core/src/outbound.rs index 7fc90084e6f..0cdbb465f60 100644 --- a/litellm-rust/crates/core/src/outbound.rs +++ b/litellm-rust/crates/core/src/outbound.rs @@ -1,30 +1,20 @@ use std::time::Duration; -use litellm_auth::RequestAuth; -use litellm_auth_aws::SigV4Signer; use litellm_http::outbound::OutboundRequest; -use serde_json::{Map, Value}; +use litellm_llms::base_llm::auth::Authenticated; +use serde_json::Value; /// Header credentials are already in `headers`; SigV4 is applied here, over the /// bytes that are sent. -pub(crate) async fn outbound_request( - auth: &RequestAuth, +pub(crate) fn outbound_request( + authenticated: Authenticated, url: String, - headers: Vec<(String, String)>, body: &Value, timeout: Option, - optional_params: &Map, -) -> Result -where - E: From + From, -{ - let RequestAuth::AwsSigV4 { region, service } = auth else { - return Ok(OutboundRequest::json(url, headers, body, timeout)?); - }; - let env_lookup = |key: &str| std::env::var(key).ok(); - let signer = - SigV4Signer::resolve(region.clone(), service, optional_params, &env_lookup).await?; - Ok(OutboundRequest::signed_json( - url, headers, body, timeout, &signer, - )?) +) -> Result { + let Authenticated { headers, signer } = authenticated; + match signer { + None => OutboundRequest::json(url, headers, body, timeout), + Some(signer) => OutboundRequest::signed_json(url, headers, body, timeout, &signer), + } } diff --git a/litellm-rust/crates/core/src/resources.rs b/litellm-rust/crates/core/src/resources.rs new file mode 100644 index 00000000000..37a29502649 --- /dev/null +++ b/litellm-rust/crates/core/src/resources.rs @@ -0,0 +1,38 @@ +use std::sync::Arc; + +use litellm_auth::AuthServices; +use litellm_http::{HttpClientConfig, HttpClientPool, media::UrlPolicy}; +use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_secrets::source::SecretSource; + +#[derive(Clone)] +pub struct CoreResources { + pub pool: Arc, + pub auth: Arc, +} + +impl CoreResources { + pub fn new(pool: Arc) -> Self { + Self { + pool, + auth: Arc::new(AuthServices::default()), + } + } + + pub fn ocr_client( + &self, + config: &HttpClientConfig, + url_policy: UrlPolicy, + settings: OcrSettings, + secrets: Arc, + ) -> Result { + OcrClient::new( + &self.pool, + config, + url_policy, + self.auth.clone(), + settings, + secrets, + ) + } +} diff --git a/litellm-rust/crates/core/src/responses/error.rs b/litellm-rust/crates/core/src/responses/error.rs deleted file mode 100644 index 1c940d8ed9b..00000000000 --- a/litellm-rust/crates/core/src/responses/error.rs +++ /dev/null @@ -1,17 +0,0 @@ -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("invalid provider: {0}")] - InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("routing error: {0}")] - Routing(String), - #[error(transparent)] - Auth(#[from] litellm_auth::Error), - #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] - Headers(#[from] litellm_http::request::HeaderError), -} diff --git a/litellm-rust/crates/core/src/responses/mod.rs b/litellm-rust/crates/core/src/responses/mod.rs index bc0f71896e5..464a81fe89c 100644 --- a/litellm-rust/crates/core/src/responses/mod.rs +++ b/litellm-rust/crates/core/src/responses/mod.rs @@ -1,3 +1,2 @@ -mod error; -pub use error::Error; +pub use crate::error::RouteError as Error; pub mod websocket; diff --git a/litellm-rust/crates/core/tests/audio_transcription.rs b/litellm-rust/crates/core/tests/audio_transcription.rs index 612395fe63a..c4dfea87319 100644 --- a/litellm-rust/crates/core/tests/audio_transcription.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -11,7 +11,7 @@ use support::*; const MODEL: &str = "mistral.voxtral-mini-3b-2507"; async fn transcribe(request: AudioTranscriptionRequest<'_>) -> Result { - audio_transcription(&http_pool(), &http_config(), request).await + audio_transcription(&support::resources(), &http_config(), request).await } fn transcript_response(text: &str) -> ResponseTemplate { diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs index d1f6cde19e8..f5802f8e305 100644 --- a/litellm-rust/crates/core/tests/chat_completions.rs +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -15,7 +15,7 @@ use support::*; const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; async fn complete(request: ChatCompletionsRequest<'_>) -> Result { - chat_completions(&http_pool(), &http_config(), request).await + chat_completions(&support::resources(), &http_config(), request).await } fn object(value: Value) -> Map { diff --git a/litellm-rust/crates/core/tests/messages/host.rs b/litellm-rust/crates/core/tests/messages/host.rs index 844ada3e1ad..b19ecf11f09 100644 --- a/litellm-rust/crates/core/tests/messages/host.rs +++ b/litellm-rust/crates/core/tests/messages/host.rs @@ -161,10 +161,10 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( #[case] response: ResponseTemplate, ) { let upstream = upstream([response]).await; - let mut body = call.body.clone(); - body.insert("stream".into(), json!(true)); - let host = - RecordingHost::passthrough(authenticated(MessagesCall { body, ..call }, upstream.uri())); + let host = RecordingHost::passthrough(authenticated( + with_fields(call, json!({"stream": true})), + upstream.uri(), + )); let _ = run_through(&host).await; @@ -180,15 +180,8 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages call: MessagesCall, ) { let upstream = upstream([message_response()]).await; - let body: Map = call - .body - .clone() - .into_iter() - .chain([("temperature".to_string(), json!(0.2))]) - .collect(); let host = RecordingHost::passthrough(authenticated( MessagesCall { - body, shaping: MessagesShaping { capabilities: AnthropicModelCapabilities { supports_sampling_params: false, @@ -197,7 +190,7 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages drop_params: true, ..MessagesShaping::default() }, - ..call + ..with_fields(call, json!({"temperature": 0.2})) }, upstream.uri(), )); diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 1ae822e5437..719c86990b0 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -7,7 +7,9 @@ use litellm_core::messages::{ }; use litellm_http::{HttpSettings, Resolution}; use litellm_secrets::source::SecretSource; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use litellm_types::llms::anthropic_messages::{ + anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, +}; use rstest::fixture; use serde_json::{Map, Value, json}; use wiremock::ResponseTemplate; @@ -31,6 +33,24 @@ fn object(value: Value) -> Map { map } +fn body(value: Value) -> AnthropicMessagesRequest { + serde_json::from_value(value).unwrap() +} + +fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall { + let current = object(serde_json::to_value(&call.body).unwrap()); + MessagesCall { + body: body(Value::Object( + current.into_iter().chain(object(fields)).collect(), + )), + ..call + } +} + +fn with_model(call: MessagesCall, model: &str) -> MessagesCall { + with_fields(call, json!({"model": model})) +} + fn message_body() -> Value { json!({ "id": "msg_1", @@ -52,8 +72,7 @@ fn message_response() -> ResponseTemplate { #[fixture] fn call() -> MessagesCall { MessagesCall { - model: MODEL.into(), - body: object(json!({ + body: body(json!({ "model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}] @@ -78,7 +97,7 @@ fn headers<'a>(pairs: impl IntoIterator) -> Option) -> MessagesMachine { - messages_machine(&http_pool(), &http_config(), secrets) + messages_machine(&support::resources(), &http_config(), secrets) .expect("default HTTP settings build a client") } diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index d37910d4ac4..0d44d26d416 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -124,11 +124,10 @@ async fn each_provider_posts_to_its_messages_endpoint( let upstream = upstream([message_response()]).await; run_message(MessagesCall { - model: model.into(), custom_llm_provider: provider.map(Into::into), api_key: Some("sk".into()), api_base: Some(format!("{}{base_suffix}", upstream.uri())), - ..call + ..with_model(call, model) }) .await; @@ -155,11 +154,10 @@ async fn unsupported_providers_are_rejected_before_sending( #[case] reported: &str, ) { let error = run(MessagesCall { - model: model.into(), custom_llm_provider: provider.map(Into::into), api_key: Some("sk".into()), api_base: Some(UNREACHABLE_BASE.into()), - ..call + ..with_model(call, model) }) .await .err() @@ -206,7 +204,7 @@ async fn azure_strips_the_cache_control_scope_anthropic_rejects(call: MessagesCa custom_llm_provider: Some("azure_ai".into()), api_key: Some("sk-azure".into()), api_base: Some(upstream.uri()), - body: object(json!({ + body: body(json!({ "model": MODEL, "max_tokens": 16, "messages": [{ @@ -232,19 +230,15 @@ async fn azure_strips_the_cache_control_scope_anthropic_rejects(call: MessagesCa #[tokio::test] async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) { let upstream = upstream([message_response()]).await; - let mut body = call.body.clone(); - body.insert("temperature".into(), json!(0.5)); - body.insert("top_k".into(), json!(3)); run_message(MessagesCall { api_key: Some("sk".into()), api_base: Some(upstream.uri()), - body, shaping: MessagesShaping { additional_drop_params: vec!["temperature".into()], ..MessagesShaping::default() }, - ..call + ..with_fields(call, json!({"temperature": 0.5, "top_k": 3})) }) .await; @@ -253,11 +247,6 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) assert_eq!(sent["top_k"], 3); } -fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall { - let body: Map = call.body.into_iter().chain(object(fields)).collect(); - MessagesCall { body, ..call } -} - fn sent_betas(request: &wiremock::Request) -> Vec { let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta")) .unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}")); @@ -406,7 +395,6 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i custom_llm_provider: call.custom_llm_provider.clone(), extra_headers: None, provider_specific_header: None, - model: call.model.clone(), timeout: call.timeout, }, fields.clone(), @@ -664,10 +652,9 @@ async fn the_provider_prefix_is_stripped_exactly_once( let upstream = upstream([message_response()]).await; run_message(MessagesCall { - model: model.into(), api_key: Some("sk".into()), api_base: Some(upstream.uri()), - ..call + ..with_model(call, model) }) .await; diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index 431dd4f4b93..ed715e22898 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -1,4 +1,7 @@ -use litellm_core::messages::{messages, types::MessagesRequest}; +use litellm_core::{ + Phase, + messages::{messages, route::messages_body}, +}; use litellm_http::transport::Error as TransportError; use rstest::rstest; @@ -154,7 +157,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( .err() .expect("an unreadable body fails"); - assert!(error.is_response(), "{error:?}"); + assert_eq!(error.phase(), Phase::AfterSend, "{error:?}"); } #[rstest] @@ -175,22 +178,9 @@ async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { assert!(matches!(error, Error::Transport(_)), "{error:?}"); } -fn facade_request(body: Value, api_base: &str) -> MessagesRequest<'_> { - MessagesRequest { - model: MODEL, - body, - api_key: Some("sk-ant"), - api_base: Some(api_base), - custom_llm_provider: Some("anthropic"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - } -} - +#[rstest] #[tokio::test] -async fn the_facade_sends_through_the_injected_http_pool_configuration() { +async fn the_facade_sends_through_the_injected_http_pool_configuration(call: MessagesCall) { let upstream = upstream([message_response()]).await; let base = upstream.uri(); let settings = HttpSettings { @@ -199,12 +189,13 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration() { }; let message = messages( - &http_pool(), + &support::resources(), &Resolution::from(&settings).config, - facade_request( - json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), - &base, - ), + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(base), + ..call + }, ) .await .expect("messages request succeeds"); @@ -215,18 +206,14 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration() { assert_eq!(sent.header("user-agent"), Some("host-owned/1")); } -#[tokio::test] -async fn the_facade_rejects_a_body_that_is_not_an_object() { - let error = messages( - &http_pool(), - &http_config(), - facade_request(json!([]), UNREACHABLE_BASE), - ) - .await - .expect_err("a non-object body is rejected"); +#[rstest] +#[case::mistyped_param(json!({"model": MODEL, "messages": [], "max_tokens": "16"}))] +#[case::missing_messages(json!({"model": MODEL, "max_tokens": 16}))] +fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) { + let error = messages_body(object(raw)).expect_err("the body is rejected"); - assert_eq!( - error, - Error::InvalidRequest("messages body must be an object".into()) + assert!( + matches!(&error, Error::InvalidRequest(message) if message.starts_with("invalid Anthropic messages request: ")), + "{error:?}" ); } diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index 4ca6e609052..49a6f7e87a0 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -69,13 +69,10 @@ impl Host for RecordingStreamHost { } fn streaming(call: MessagesCall, api_base: String) -> MessagesCall { - let mut body = call.body.clone(); - body.insert("stream".into(), json!(true)); MessagesCall { api_key: Some("sk-ant".into()), api_base: Some(api_base), - body, - ..call + ..with_fields(call, json!({"stream": true})) } } @@ -243,7 +240,7 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { #[rstest] #[tokio::test] -async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) { +async fn a_host_on_anthropic_sse_is_relayed_byte_for_byte(call: MessagesCall) { let upstream = upstream([sse_response()]).await; let host = RecordingStreamHost::new( MessagesCall { @@ -253,14 +250,17 @@ async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCal usize::MAX, ); - let error = stream_through(&host) - .await - .err() - .expect("azure streaming is refused"); + let outcome = stream_through(&host).await.expect("azure streams"); - assert_eq!( - error, - Error::Unsupported("streaming messages for this provider") - ); - assert!(received(&upstream).await.is_empty()); + assert!(matches!(outcome, MessagesOutput::Streamed)); + let seen = host.seen.into_inner().unwrap(); + let delivered: Vec = seen + .iter() + .filter_map(|step| match step { + Seen::Deliver(chunk) => Some(chunk.to_vec()), + Seen::Open(_) => None, + }) + .flatten() + .collect(); + assert_eq!(delivered, SSE_BODY.as_bytes()); } diff --git a/litellm-rust/crates/core/tests/ocr/mistral.rs b/litellm-rust/crates/core/tests/ocr/mistral.rs index 4c3f1c5cc39..d542eeaf03a 100644 --- a/litellm-rust/crates/core/tests/ocr/mistral.rs +++ b/litellm-rust/crates/core/tests/ocr/mistral.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use litellm_auth_gcp::VertexAuth; use litellm_http::{HttpSettings, Resolution, media::UrlPolicy}; use litellm_llms::{ base_llm::ocr::{ @@ -181,6 +180,7 @@ async fn missing_credentials_come_from_the_injected_secret_source( ); } +#[rstest] #[tokio::test] async fn the_client_uses_the_injected_http_pool_configuration() { let upstream = upstream([pages_response()]).await; @@ -188,19 +188,18 @@ async fn the_client_uses_the_injected_http_pool_configuration() { user_agent: Some("host-owned/1".into()), ..HttpSettings::default() }; - let client = OcrClient::new( - &http_pool(), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new( - litellm_secrets::source::EnvironmentSecrets::python_compatible( - litellm_http::Client::plain_for_test(), + let client = resources() + .ocr_client( + &Resolution::from(&settings).config, + UrlPolicy::default(), + OcrSettings::default(), + Arc::new( + litellm_secrets::source::EnvironmentSecrets::python_compatible( + litellm_http::Client::plain_for_test(), + ), ), - ), - ) - .unwrap(); + ) + .unwrap(); litellm_core::ocr::client::perform( &client, diff --git a/litellm-rust/crates/core/tests/resources.rs b/litellm-rust/crates/core/tests/resources.rs new file mode 100644 index 00000000000..9764e50de1b --- /dev/null +++ b/litellm-rust/crates/core/tests/resources.rs @@ -0,0 +1,157 @@ +mod support; + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use litellm_auth::AuthServices; +use litellm_auth_gcp::{ + CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource, +}; +use litellm_core::{ + ocr::{ + client::perform, + wire::{OcrWireRequest, decode_request}, + }, + resources::CoreResources, +}; +use litellm_http::{HttpSettings, Resolution}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use rstest::{fixture, rstest}; +use serde_json::json; +use support::{ReceivedRequest, RecordingSecrets, http_pool, json_response, upstream}; + +struct TokenSource(String); + +impl VertexTokenSource for TokenSource { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async { Ok(self.0.clone()) }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async { Ok(self.0.clone()) }) + } +} + +#[derive(Default)] +struct Loader(AtomicUsize); + +impl VertexProviderLoader for Loader { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { + Box::pin(async move { + self.0.fetch_add(1, Ordering::SeqCst); + let identity = match source { + CredentialSource::Trusted(secret) => secret.expose().to_string(), + other => panic!("unexpected credential source: {other:?}"), + }; + Ok(Arc::new(TokenSource(identity)) as Arc) + }) + } +} + +#[fixture] +fn loader() -> Arc { + Arc::new(Loader::default()) +} + +#[fixture] +fn resources(loader: Arc) -> CoreResources { + CoreResources { + auth: Arc::new(AuthServices { + gcp: VertexAuth::new(loader), + ..AuthServices::default() + }), + pool: Arc::new(http_pool()), + } +} + +#[rstest] +#[case::shared_identity(false, "first-identity", 1)] +#[case::different_identity(false, "second-identity", 2)] +#[case::independent_resources(true, "first-identity", 2)] +#[tokio::test] +async fn auth_survives_per_call_clients_without_freezing_settings_or_secrets( + loader: Arc, + #[with(loader.clone())] resources: CoreResources, + #[case] independent: bool, + #[case] second_identity: &str, + #[case] expected_loads: usize, +) { + let response = json_response(json!({"pages": [{"index": 0, "markdown": "hello"}]})); + let upstream = upstream([response.clone(), response]).await; + let second_resources = if independent { + CoreResources { + auth: Arc::new(AuthServices { + gcp: VertexAuth::new(loader.clone()), + ..AuthServices::default() + }), + ..resources.clone() + } + } else { + resources.clone() + }; + for (owner, identity, agent, location) in [ + (&resources, "first-identity", "first-agent", "us-central1"), + ( + &second_resources, + second_identity, + "second-agent", + "europe-west4", + ), + ] { + let http = Resolution::from(&HttpSettings { + user_agent: Some(agent.into()), + ..HttpSettings::default() + }) + .config; + let client = owner + .ocr_client( + &http, + Default::default(), + OcrSettings { + vertex_location: Some(location.into()), + ..OcrSettings::default() + }, + Arc::new(RecordingSecrets::new([("VERTEXAI_CREDENTIALS", identity)])), + ) + .unwrap(); + let request = decode_request(OcrWireRequest { + model: "vertex_ai/mistral-ocr-maas".into(), + document: json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: Some(upstream.uri()), + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(5.0), + }).unwrap(); + let result = perform(&client, request).await.unwrap(); + assert!(!result.pages.is_empty()); + } + let requests = upstream.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + for (request, identity, agent, location) in [ + (&requests[0], "first-identity", "first-agent", "us-central1"), + ( + &requests[1], + second_identity, + "second-agent", + "europe-west4", + ), + ] { + assert_eq!( + request.header("authorization"), + Some(format!("Bearer {identity}").as_str()) + ); + assert_eq!(request.header("user-agent"), Some(agent)); + assert!( + request + .url + .path() + .contains(&format!("/projects/{identity}/locations/{location}/")) + ); + } + assert_eq!(loader.0.load(Ordering::SeqCst), expected_loads); +} diff --git a/litellm-rust/crates/core/tests/support/mod.rs b/litellm-rust/crates/core/tests/support/mod.rs index 1d9af236811..5443437df09 100644 --- a/litellm-rust/crates/core/tests/support/mod.rs +++ b/litellm-rust/crates/core/tests/support/mod.rs @@ -20,6 +20,10 @@ pub fn http_pool() -> HttpClientPool { HttpClientPool::new(Arc::new(PublicDnsResolver)) } +pub fn resources() -> litellm_core::resources::CoreResources { + litellm_core::resources::CoreResources::new(Arc::new(http_pool())) +} + pub fn http_config() -> HttpClientConfig { Resolution::from(&HttpSettings::default()).config } diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index c1b35c0f69d..bb77a1bf330 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -6,15 +6,17 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-host.workspace = true + bytes.workspace = true futures-util.workspace = true -litellm-host.workspace = true -pyo3.workspace = true -pyo3-async-runtimes.workspace = true -pythonize.workspace = true serde.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } +pyo3.workspace = true +pyo3-async-runtimes.workspace = true +pythonize = "0.29.0" + [dev-dependencies] rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/litellm/Cargo.toml b/litellm-rust/crates/litellm/Cargo.toml new file mode 100644 index 00000000000..f6a63227792 --- /dev/null +++ b/litellm-rust/crates/litellm/Cargo.toml @@ -0,0 +1,3 @@ +[package] +name = "litellm" +version = "0.0.1" diff --git a/litellm-rust/crates/litellm/src/lib.rs b/litellm-rust/crates/litellm/src/lib.rs new file mode 100644 index 00000000000..ae6daac0100 --- /dev/null +++ b/litellm-rust/crates/litellm/src/lib.rs @@ -0,0 +1,2 @@ +//! Before publishing this crate, add a registry `version` beside each internal `path` dependency in the workspace manifest. +//! https://crates.io/crates/litellm diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 36ccd18f220..beff99bc73a 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -11,7 +11,7 @@ test-support = ["litellm-http/test-support"] [dependencies] litellm-types.workspace = true litellm-core-utils.workspace = true -litellm-auth.workspace = true +litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true litellm-auth-gcp.workspace = true @@ -19,6 +19,7 @@ litellm-host.workspace = true litellm-framing.workspace = true litellm-http.workspace = true litellm-secrets.workspace = true +litellm-python-compat.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" diff --git a/litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md b/litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md new file mode 100644 index 00000000000..c3a3c234492 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md @@ -0,0 +1 @@ +- https://platform.claude.com/docs/en/api/http/beta/messages/batches/create diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index 1c26684901a..395c2376059 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -4,10 +4,7 @@ use serde_json::Value; use time::OffsetDateTime; use url::Url; -use crate::{ - anthropic::messages::transformation::resolve_anthropic_api_base, - base_llm::chat::transformation::Error, -}; +use crate::{Error, anthropic::messages::transformation::resolve_anthropic_api_base}; const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches"; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs index 9160cdf28ee..f258656494a 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs @@ -7,11 +7,12 @@ use litellm_types::{ use serde_json::Value; use crate::{ + Error, anthropic::messages::streaming_iterator::{ AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent, AnthropicStreamUsage, }, - base_llm::{base_model_iterator::StreamTransformer, chat::transformation::Error}, + base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape}, }; #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -38,7 +39,7 @@ pub struct AnthropicContentBlockDeltaEvent { pub delta: AnthropicContentBlockDelta, } -pub struct AnthropicChatCompletionsStreamTransformer { +pub struct ModelResponseIterator { pub content_blocks: Vec, pub tool_index: i64, pub json_mode: bool, @@ -61,12 +62,8 @@ pub struct AnthropicChatCompletionsStreamTransformer { pub container_id: Option, } -impl AnthropicChatCompletionsStreamTransformer { - pub fn new( - _json_mode: bool, - _speed: Option, - _tool_name_reverse_map: HashMap, - ) -> Self { +impl ModelResponseIterator { + pub fn new(_shape: StreamShape) -> Self { todo!() } @@ -150,7 +147,7 @@ impl AnthropicChatCompletionsStreamTransformer { } } -impl StreamTransformer for AnthropicChatCompletionsStreamTransformer { +impl StreamTransformer for ModelResponseIterator { type Input = AnthropicMessagesStreamEvent; type Output = ChatCompletionChunk; type Error = Error; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 07ed6ba6ed1..b19443a7ff8 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,3 +1,4 @@ +use litellm_auth::{CredentialPlacement, SecretValue}; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, build_conversation}, @@ -9,13 +10,22 @@ use litellm_types::{ use serde_json::{Map, Value, json}; use crate::{ + Error, anthropic::{ ANTHROPIC_OAUTH_TOKEN_PREFIX, + chat::handler::ModelResponseIterator, messages::transformation::{complete_anthropic_url, resolve_anthropic_api_key}, }, - base_llm::chat::transformation::{ - BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, - Unsupported, unsupported_message, unsupported_param, + base_llm::{ + anthropic_messages::streaming::anthropic_sse_event_stream, + auth::AuthScheme, + chat::{ + streaming::{ChatStream, StreamShape}, + transformation::{ + BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, + Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + }, + }, }, }; @@ -40,6 +50,15 @@ pub struct AnthropicConfig; pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig; +fn forwards_oauth_bearer(headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("authorization") + && value + .strip_prefix("Bearer ") + .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) + }) +} + impl BaseConfig for AnthropicConfig { fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS @@ -63,6 +82,7 @@ impl BaseConfig for AnthropicConfig { ) -> Result { Ok(ProviderChatRequestData { body: anthropic_body(model, &build_conversation(&messages), optional_params), + stream_shape: StreamShape::default(), }) } @@ -129,17 +149,35 @@ impl BaseConfig for AnthropicConfig { }) } - fn auth( + /// A forwarded OAuth bearer is the whole credential: Python pops `x-api-key` for it, + /// so the resolved key is not applied over it. Any other forwarded header loses to + /// the deployment's key, which Python writes last. + fn validate_environment( &self, + headers: Headers, api_key: Option<&str>, _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(RequestAuth::Header { - name: "x-api-key", - value: resolve_anthropic_api_key(api_key, env_lookup)?, - }) + ) -> Result { + if forwards_oauth_bearer(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + let auth = AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), + secret: SecretValue::new(resolve_anthropic_api_key(api_key, env_lookup)?), + }; + Ok(ValidatedEnvironment { headers, auth }) + } + + fn model_response_iterator(&self, shape: StreamShape) -> Option { + Some(ChatStream::new( + anthropic_sse_event_stream, + ModelResponseIterator::new(shape), + )) } fn default_headers(&self) -> &'static [(&'static str, &'static str)] { @@ -154,15 +192,6 @@ impl BaseConfig for AnthropicConfig { /// the resolved key must not be applied over the top. Any other forwarded /// `authorization` is unrelated to this header and does not defer, which is /// also what Python does: it sends the deployment's `x-api-key` alongside. - fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, value)| { - name.eq_ignore_ascii_case("authorization") - && value - .strip_prefix("Bearer ") - .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) - }) - } - fn unsupported_reason( &self, messages: &[ChatMessage], diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index a2234e0df03..34c59c6ec5d 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,5 +1,5 @@ use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessage, ContentBlock, MessageContent, + AnthropicMessage, ContentBlock, EffortLevel, MessageContent, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -26,39 +26,6 @@ pub mod beta { pub const PER_TURN_CONTROL_2026_07_01: &str = "per-turn-control-2026-07-01"; } -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum EffortLevel { - Low, - Medium, - High, - Xhigh, - Max, -} - -impl EffortLevel { - pub fn as_str(self) -> &'static str { - match self { - Self::Low => "low", - Self::Medium => "medium", - Self::High => "high", - Self::Xhigh => "xhigh", - Self::Max => "max", - } - } - - pub fn parse(value: &str) -> Option { - match value { - "low" => Some(Self::Low), - "medium" => Some(Self::Medium), - "high" => Some(Self::High), - "xhigh" => Some(Self::Xhigh), - "max" => Some(Self::Max), - _ => None, - } - } -} - #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct SupportedEffortTiers { #[serde(default)] @@ -135,15 +102,11 @@ impl AnthropicModelCapabilities { self.supports_output_config || self.effort_tiers.any() } - pub fn effort_level_rejection(&self, effort: &str, model: &str) -> Option { - match effort { - "max" if !(self.supports_adaptive_thinking || self.effort_tiers.max) => Some(format!( - "effort='max' is not supported by this model. Got model: {model}" - )), - "xhigh" if !self.effort_tiers.xhigh => Some(format!( - "effort='xhigh' is not supported by this model. Got model: {model}" - )), - _ => None, + pub fn accepts_effort(&self, level: EffortLevel) -> bool { + match level { + EffortLevel::Max => self.supports_adaptive_thinking || self.effort_tiers.max, + EffortLevel::Xhigh => self.effort_tiers.xhigh, + EffortLevel::Low | EffortLevel::Medium | EffortLevel::High => true, } } } @@ -1329,34 +1292,6 @@ mod tests { ); } - #[rstest] - #[case::low(EffortLevel::Low, "low")] - #[case::medium(EffortLevel::Medium, "medium")] - #[case::high(EffortLevel::High, "high")] - #[case::xhigh(EffortLevel::Xhigh, "xhigh")] - #[case::max(EffortLevel::Max, "max")] - fn effort_level_names_agree_across_str_parse_and_serde( - #[case] level: EffortLevel, - #[case] name: &str, - ) { - assert_eq!(level.as_str(), name); - assert_eq!(EffortLevel::parse(name), Some(level)); - assert_eq!(serde_json::to_value(level).unwrap(), json!(name)); - assert_eq!( - serde_json::from_value::(json!(name)).unwrap(), - level - ); - } - - #[rstest] - #[case::unknown("ultra")] - #[case::minimal_is_not_an_output_config_level("minimal")] - #[case::uppercase("HIGH")] - #[case::empty("")] - fn effort_level_parse_rejects(#[case] value: &str) { - assert_eq!(EffortLevel::parse(value), None); - } - #[rstest] #[case::minimal_only(tiers(true, false, false, false, false, false), [false, false, false, false, false])] #[case::low_only(tiers(false, true, false, false, false, false), [true, false, false, false, false])] @@ -1450,56 +1385,55 @@ mod tests { } #[rstest] - #[case::max_on_adaptive_thinking_model(true, SupportedEffortTiers::default(), "max", None)] + #[case::max_on_adaptive_thinking_model( + true, + SupportedEffortTiers::default(), + EffortLevel::Max, + true + )] #[case::max_on_max_tier_model( false, tiers(false, false, false, false, false, true), - "max", - None + EffortLevel::Max, + true )] #[case::max_on_output_config_only_model( false, SupportedEffortTiers::default(), - "max", - Some("effort='max' is not supported by this model. Got model: claude-test") + EffortLevel::Max, + false )] #[case::max_on_xhigh_tier_model( false, tiers(false, false, false, false, true, false), - "max", - Some("effort='max' is not supported by this model. Got model: claude-test") + EffortLevel::Max, + false )] #[case::xhigh_on_xhigh_tier_model( false, tiers(false, false, false, false, true, false), - "xhigh", - None + EffortLevel::Xhigh, + true )] #[case::xhigh_on_adaptive_thinking_model( true, SupportedEffortTiers::default(), - "xhigh", - Some("effort='xhigh' is not supported by this model. Got model: claude-test") + EffortLevel::Xhigh, + false )] #[case::xhigh_on_max_tier_model( false, tiers(false, false, false, false, false, true), - "xhigh", - Some("effort='xhigh' is not supported by this model. Got model: claude-test") + EffortLevel::Xhigh, + false )] - #[case::high_on_unmapped_model(false, SupportedEffortTiers::default(), "high", None)] - #[case::low_on_unmapped_model(false, SupportedEffortTiers::default(), "low", None)] - #[case::unknown_level_is_left_to_other_validation( - false, - SupportedEffortTiers::default(), - "ultra", - None - )] - fn effort_level_rejection_cases( + #[case::high_on_unmapped_model(false, SupportedEffortTiers::default(), EffortLevel::High, true)] + #[case::low_on_unmapped_model(false, SupportedEffortTiers::default(), EffortLevel::Low, true)] + fn accepts_effort_cases( #[case] supports_adaptive_thinking: bool, #[case] effort_tiers: SupportedEffortTiers, - #[case] effort: &str, - #[case] expected: Option<&str>, + #[case] level: EffortLevel, + #[case] expected: bool, unmapped: AnthropicModelCapabilities, ) { let capabilities = AnthropicModelCapabilities { @@ -1508,12 +1442,7 @@ mod tests { effort_tiers, ..unmapped }; - assert_eq!( - capabilities - .effort_level_rejection(effort, "claude-test") - .as_deref(), - expected - ); + assert_eq!(capabilities.accepts_effort(level), expected); } #[rstest] diff --git a/litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md b/litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md new file mode 100644 index 00000000000..f08c0c6d017 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md @@ -0,0 +1 @@ +- https://platform.claude.com/docs/en/api/http/messages/count_tokens diff --git a/litellm-rust/crates/llms/src/anthropic/count_tokens/transformation.rs b/litellm-rust/crates/llms/src/anthropic/count_tokens/transformation.rs index a4d8c57ca4f..9fa831b8b66 100644 --- a/litellm-rust/crates/llms/src/anthropic/count_tokens/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/count_tokens/transformation.rs @@ -2,7 +2,7 @@ use litellm_types::llms::anthropic_messages::anthropic_request::{AnthropicMessag use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::{anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, base_llm::chat::transformation::Error}; +use crate::{Error, anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX}; const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens"; const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01"; diff --git a/litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md b/litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md new file mode 100644 index 00000000000..b7832c4b8f3 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md @@ -0,0 +1 @@ +- https://platform.claude.com/docs/en/api/http/messages/create diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 0e2ab97956a..9e187035944 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -1,14 +1,18 @@ -use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessage, AnthropicMessagesRequest, +use litellm_types::{ + llms::anthropic_messages::anthropic_request::{ + AdaptiveThinking, AnthropicMessage, AnthropicMessagesOptionalParams, + AnthropicMessagesRequest, EnabledThinking, ThinkingConfig, ThinkingDisplay, + }, + recognized::Recognized, }; use serde_json::{Value, json}; use crate::{ + Error, anthropic::common_utils::{ flatten_unencrypted_web_search_results, sanitize_tool_use_ids, strip_empty_content_blocks, strip_provider_specific_fields, }, - base_llm::chat::transformation::Error, }; pub fn shape_anthropic_messages_request( @@ -17,12 +21,16 @@ pub fn shape_anthropic_messages_request( ) -> Result { Ok(AnthropicMessagesRequest { messages: sanitize_anthropic_messages(request.messages), - metadata: request - .metadata - .as_ref() - .map(validate_anthropic_api_metadata) - .transpose()?, - thinking: with_reasoning_auto_summary(request.thinking, reasoning_auto_summary), + params: AnthropicMessagesOptionalParams { + metadata: request + .params + .metadata + .as_ref() + .map(validate_anthropic_api_metadata) + .transpose()?, + thinking: with_reasoning_auto_summary(request.params.thinking, reasoning_auto_summary), + ..request.params + }, ..request }) } @@ -48,20 +56,38 @@ fn validate_anthropic_api_metadata(metadata: &Value) -> Result { } } -fn with_reasoning_auto_summary(thinking: Option, enabled: bool) -> Option { - let Some(Value::Object(thinking)) = thinking else { +fn with_reasoning_auto_summary( + thinking: Option>, + enabled: bool, +) -> Option> { + if !enabled { return thinking; - }; - if !enabled || thinking.get("type").and_then(Value::as_str) == Some("disabled") { - return Some(Value::Object(thinking)); } - Some(Value::Object( - thinking - .into_iter() - .filter(|(key, _)| key != "display") - .chain([("display".to_string(), json!("summarized"))]) - .collect(), - )) + let summarized = Some(Recognized::Known(ThinkingDisplay::Summarized)); + match thinking { + Some(Recognized::Known(ThinkingConfig::Enabled(enabled))) => Some(Recognized::Known( + ThinkingConfig::Enabled(EnabledThinking { + display: summarized, + ..enabled + }), + )), + Some(Recognized::Known(ThinkingConfig::Adaptive(adaptive))) => Some(Recognized::Known( + ThinkingConfig::Adaptive(AdaptiveThinking { + display: summarized, + ..adaptive + }), + )), + Some(Recognized::Unrecognized(Value::Object(fields))) => { + Some(Recognized::Unrecognized(Value::Object( + fields + .into_iter() + .filter(|(key, _)| key != "display") + .chain([("display".to_string(), json!("summarized"))]) + .collect(), + ))) + } + other => other, + } } #[cfg(test)] @@ -230,12 +256,22 @@ mod tests { )] #[case::no_thinking(None, true, None)] #[case::non_object_thinking(Some(json!("enabled")), true, Some(json!("enabled")))] + #[case::unknown_type( + Some(json!({"type": "future"})), + true, + Some(json!({"type": "future", "display": "summarized"})), + )] fn reasoning_auto_summary_marks_active_thinking_as_summarized( #[case] thinking: Option, #[case] enabled: bool, #[case] expected: Option, ) { - assert_eq!(with_reasoning_auto_summary(thinking, enabled), expected); + let thinking = thinking.map(|thinking| serde_json::from_value(thinking).unwrap()); + assert_eq!( + with_reasoning_auto_summary(thinking, enabled) + .map(|thinking| serde_json::to_value(thinking).unwrap()), + expected + ); } #[test] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs b/litellm-rust/crates/llms/src/anthropic/messages/headers.rs index 8d48d7a0f5c..bd1b11be92d 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/headers.rs @@ -1,4 +1,7 @@ -use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_types::{ + llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest, recognized::Recognized, +}; use serde_json::Value; use crate::{ @@ -10,7 +13,10 @@ use crate::{ split_beta_values, }, }, - base_llm::anthropic_messages::transformation::Headers, + base_llm::{ + anthropic_messages::transformation::Headers, + auth::{AuthScheme, ValidatedEnvironment}, + }, }; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; @@ -41,13 +47,13 @@ fn existing_betas(headers: &[(String, String)]) -> impl Iterator .flat_map(|(_, value)| split_beta_values(Some(value))) } -fn with_oauth_bearer(headers: Headers, bearer: String) -> Headers { +/// The OAuth headers Python's `optionally_handle_anthropic_oauth` sets next to the bearer. +fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers { let beta = join_beta_values(existing_betas(&headers).chain([ANTHROPIC_OAUTH_BETA_HEADER.to_string()])); - without(headers, &[API_KEY_HEADER, AUTHORIZATION, BETA_HEADER]) + without(headers, &[dropped, &[BETA_HEADER]].concat()) .into_iter() .chain([ - (AUTHORIZATION.to_string(), bearer), (BETA_HEADER.to_string(), beta), (DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string()), ]) @@ -58,38 +64,54 @@ fn non_empty(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } -pub fn authenticate( +fn bearer(token: &str) -> AuthScheme { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), + } +} + +pub fn validate_environment( headers: Headers, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - if let Some(forwarded) = header_value(&headers, AUTHORIZATION) - && forwarded - .strip_prefix("Bearer ") - .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) +) -> Result { + if let Some(token) = header_value(&headers, AUTHORIZATION) + .and_then(|forwarded| forwarded.strip_prefix("Bearer ")) + .filter(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) { - let bearer = forwarded.to_string(); - return Ok(with_oauth_bearer(headers, bearer)); + let auth = bearer(token); + return Ok(ValidatedEnvironment { + headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), + auth, + }); } if let Some(key) = api_key.filter(|key| key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) { - return Ok(with_oauth_bearer(headers, format!("Bearer {key}"))); + return Ok(ValidatedEnvironment { + headers: with_oauth_companions(headers, &[API_KEY_HEADER]), + auth: bearer(key), + }); } if header_value(&headers, API_KEY_HEADER).is_some() || header_value(&headers, AUTHORIZATION).is_some() { - return Ok(headers); + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); } let resolved_key = non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())); let auth = match resolved_key { - Some(key) if is_anthropic_oauth_key(&key) => { - (AUTHORIZATION.to_string(), format!("Bearer {key}")) - } - Some(key) => (API_KEY_HEADER.to_string(), key), + Some(key) if is_anthropic_oauth_key(&key) => bearer(&key), + Some(key) => AuthScheme::Credential { + placement: CredentialPlacement::Header(API_KEY_HEADER), + secret: SecretValue::new(key), + }, None => match env_lookup(ANTHROPIC_AUTH_TOKEN_ENV).filter(|value| !value.trim().is_empty()) { - Some(token) => (AUTHORIZATION.to_string(), format!("Bearer {token}")), + Some(token) => bearer(&token), None => { return Err(litellm_auth::Error::MissingApiKey { provider: "Anthropic", @@ -98,7 +120,7 @@ pub fn authenticate( } }, }; - Ok(headers.into_iter().chain([auth]).collect()) + Ok(ValidatedEnvironment { headers, auth }) } fn context_management_betas( @@ -122,12 +144,13 @@ fn context_management_betas( } fn uses_structured_output(request: &AnthropicMessagesRequest) -> bool { - request.output_format.is_some() + request.params.output_format.is_some() || request + .params .output_config .as_ref() - .and_then(|config| config.get("format")) - .is_some_and(|format| !format.is_null()) + .and_then(Recognized::known) + .is_some_and(|config| config.format.is_some()) } fn messages_carry_output_config(request: &AnthropicMessagesRequest) -> bool { @@ -138,12 +161,12 @@ fn messages_carry_output_config(request: &AnthropicMessagesRequest) -> bool { } pub fn feature_betas(request: &AnthropicMessagesRequest) -> Vec<&'static str> { - let tools = request.tools.as_deref(); + let tools = request.params.tools.as_deref(); [ - requires_native_compaction_beta(request.compaction.as_ref(), &request.messages) + requires_native_compaction_beta(request.params.compaction.as_ref(), &request.messages) .then_some(beta::COMPACT_2026_09_04), uses_structured_output(request).then_some(beta::STRUCTURED_OUTPUT), - (request.speed.as_deref() == Some("fast")).then_some(beta::FAST_MODE_2026_02_01), + (request.params.speed.as_deref() == Some("fast")).then_some(beta::FAST_MODE_2026_02_01), messages_carry_output_config(request).then_some(beta::PER_TURN_CONTROL_2026_07_01), has_advisor_tool(tools).then_some(beta::ADVISOR_TOOL_2026_03_01), is_tool_search_used(tools).then_some(beta::ADVANCED_TOOL_USE_2025_11_20), @@ -151,7 +174,7 @@ pub fn feature_betas(request: &AnthropicMessagesRequest) -> Vec<&'static str> { .into_iter() .flatten() .chain(context_management_betas( - request.context_management.as_ref(), + request.params.context_management.as_ref(), )) .collect() } @@ -179,6 +202,7 @@ mod tests { use serde_json::json; use super::*; + use crate::base_llm::auth::resolve_auth; const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; @@ -230,7 +254,17 @@ mod tests { .find(|(key, _)| *key == name) .map(|(_, value)| value.to_string()) }; - authenticate(headers(forwarded), api_key, &lookup) + let validated = validate_environment(headers(forwarded), api_key, &lookup)?; + let resolved = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() + .block_on(resolve_auth( + &litellm_auth::AuthServices::default(), + validated, + &lookup, + )) + .unwrap(); + Ok(resolved.headers) } #[rstest] @@ -282,9 +316,9 @@ mod tests { .iter() .copied() .chain([ - ("authorization", expected_bearer), ("anthropic-beta", ANTHROPIC_OAUTH_BETA_HEADER), BROWSER_ACCESS, + ("authorization", expected_bearer), ]) .collect::>(); assert_eq!( @@ -318,12 +352,12 @@ mod tests { assert_eq!( authenticate_with(forwarded, api_key, no_env).unwrap(), headers(&[ - ("authorization", OAUTH_BEARER), ( "anthropic-beta", &betas(&[ANTHROPIC_OAUTH_BETA_HEADER, "web-search-2025-03-05"]) ), BROWSER_ACCESS, + ("authorization", OAUTH_BEARER), ]) ); } @@ -621,8 +655,8 @@ mod tests { assert_eq!( with_feature_betas(oauth_headers, &all_features), headers(&[ - ("authorization", OAUTH_BEARER), BROWSER_ACCESS, + ("authorization", OAUTH_BEARER), ( "anthropic-beta", &betas(&[ diff --git a/litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs b/litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs index 3f1b7ed9bcc..4f00cd0af7e 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/streaming_iterator.rs @@ -1,36 +1,16 @@ -use base64::Engine; -use bytes::Buf; -use futures_util::{Stream, StreamExt}; -use litellm_framing::{ - aws_event_stream::{AwsEventStreamCodec, Message}, - frames, - sse::{SseCodec, SseEvent}, -}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("stream framing failed: {0}")] - StreamFraming(String), - #[error("Anthropic stream event is invalid: {0}")] - InvalidStreamEvent(String), - #[error("Bedrock event payload is invalid: {0}")] - InvalidBedrockPayload(String), - #[error("Bedrock event payload has invalid base64: {0}")] - InvalidBedrockBase64(String), -} - #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct AnthropicStreamUsage { - #[serde(default)] - pub input_tokens: u64, - #[serde(default)] - pub output_tokens: u64, - #[serde(default)] - pub cache_creation_input_tokens: u64, - #[serde(default)] - pub cache_read_input_tokens: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_creation_input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_read_input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub server_tool_use: Option, #[serde(flatten)] @@ -151,141 +131,12 @@ pub enum AnthropicMessagesStreamEvent { #[serde(default, skip_serializing_if = "Option::is_none")] context_management: Option, }, - MessageStop, + MessageStop { + #[serde(default, skip_serializing_if = "Option::is_none")] + usage: Option, + }, Ping, Error { error: AnthropicStreamError, }, } - -#[derive(Deserialize)] -struct BedrockChunkPayload { - bytes: String, -} - -pub fn decode_anthropic_sse_frame(event: SseEvent) -> Result { - serde_json::from_str(&event.data).map_err(|error| Error::InvalidStreamEvent(error.to_string())) -} - -pub fn decode_bedrock_anthropic_frame( - message: Message, -) -> Result { - let payload: BedrockChunkPayload = serde_json::from_slice(message.payload()) - .map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?; - let event = base64::engine::general_purpose::STANDARD - .decode(payload.bytes) - .map_err(|error| Error::InvalidBedrockBase64(error.to_string()))?; - serde_json::from_slice(&event).map_err(|error| Error::InvalidStreamEvent(error.to_string())) -} - -pub fn direct_anthropic_event_stream( - input: S, -) -> impl Stream> + Send -where - S: Stream> + Send, - B: Buf + Send, - E: std::error::Error + Send + Sync + 'static, -{ - frames(input, SseCodec::default()).map(|event| { - decode_anthropic_sse_frame(event.map_err(|error| Error::StreamFraming(error.to_string()))?) - }) -} - -pub fn bedrock_anthropic_event_stream( - input: S, -) -> impl Stream> + Send -where - S: Stream> + Send, - B: Buf + Send, - E: std::error::Error + Send + Sync + 'static, -{ - frames(input, AwsEventStreamCodec).map(|message| { - decode_bedrock_anthropic_frame( - message.map_err(|error| Error::StreamFraming(error.to_string()))?, - ) - }) -} - -#[cfg(test)] -mod tests { - use std::io; - - use aws_smithy_eventstream::frame::write_message_to; - use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; - use base64::engine::general_purpose::STANDARD; - use bytes::Bytes; - use futures_util::TryStreamExt; - - use super::*; - - const TEXT_DELTA: &str = - r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; - - #[tokio::test] - async fn direct_anthropic_sse_frames_into_typed_events() { - let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n"); - let events = direct_anthropic_event_stream(futures_util::stream::iter( - wire.as_bytes().chunks(3).map(Ok::<_, io::Error>), - )) - .try_collect::>() - .await - .unwrap(); - - assert_eq!( - events, - vec![AnthropicMessagesStreamEvent::ContentBlockDelta { - index: 0, - delta: AnthropicContentBlockDelta::TextDelta { - text: "hello".into(), - }, - }] - ); - } - - #[test] - fn decodes_citations_delta_events() { - let event = decode_anthropic_sse_frame(SseEvent { - event: Some("content_block_delta".into()), - data: r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"# - .into(), - id: None, - retry: None, - }) - .unwrap(); - - assert!(matches!( - event, - AnthropicMessagesStreamEvent::ContentBlockDelta { - delta: AnthropicContentBlockDelta::Citations { .. }, - .. - } - )); - } - - #[tokio::test] - async fn bedrock_aws_frames_into_the_same_typed_events() { - let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)}); - let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( - Header::new(":event-type", HeaderValue::String("chunk".into())), - ); - let mut wire = Vec::new(); - write_message_to(&message, &mut wire).unwrap(); - - let events = bedrock_anthropic_event_stream(futures_util::stream::iter( - wire.chunks(3).map(Ok::<_, io::Error>), - )) - .try_collect::>() - .await - .unwrap(); - - assert_eq!( - events, - vec![AnthropicMessagesStreamEvent::ContentBlockDelta { - index: 0, - delta: AnthropicContentBlockDelta::TextDelta { - text: "hello".into(), - }, - }] - ); - } -} diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index ffa4c8ffeb8..101ed438738 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -1,15 +1,21 @@ use litellm_core_utils::settings::Lookup; -use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; -use serde_json::{Map, Value, json}; - -use crate::{ - anthropic::common_utils::AnthropicModelCapabilities, base_llm::chat::transformation::Error, +use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy}; +use litellm_types::{ + llms::{ + anthropic_messages::anthropic_request::{ + AnthropicMessagesOptionalParams, AnthropicMessagesRequest, EffortLevel, OutputConfig, + ThinkingConfig, ThinkingDisplay, + }, + openai::ReasoningEffort, + }, + recognized::Recognized, }; +use serde_json::Value; + +use crate::{Error, anthropic::common_utils::AnthropicModelCapabilities}; pub const ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: u64 = 1024; -const EFFORT_NAMES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"; - #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ThinkingBudgets { pub minimal: u64, @@ -50,15 +56,17 @@ impl ThinkingBudgets { } } - fn for_effort(&self, reasoning_effort: &str) -> Option { - match reasoning_effort { - "low" => Some(self.low), - "medium" => Some(self.medium), - "high" => Some(self.high), - "xhigh" => Some(self.xhigh), - "max" => Some(self.max), - "minimal" => Some(self.minimal.max(ANTHROPIC_MIN_THINKING_BUDGET_TOKENS)), - _ => None, + fn for_effort(&self, effort: ReasoningEffort) -> Option { + match effort { + ReasoningEffort::None => None, + ReasoningEffort::Minimal => { + Some(self.minimal.max(ANTHROPIC_MIN_THINKING_BUDGET_TOKENS)) + } + ReasoningEffort::Low => Some(self.low), + ReasoningEffort::Medium => Some(self.medium), + ReasoningEffort::High => Some(self.high), + ReasoningEffort::Xhigh => Some(self.xhigh), + ReasoningEffort::Max => Some(self.max), } } @@ -66,17 +74,17 @@ impl ThinkingBudgets { &self, budget_tokens: u64, capabilities: &AnthropicModelCapabilities, - ) -> &'static str { + ) -> EffortLevel { if budget_tokens >= self.xhigh && capabilities.effort_tiers.xhigh { - return "xhigh"; + return EffortLevel::Xhigh; } if budget_tokens >= self.high { - return "high"; + return EffortLevel::High; } if budget_tokens >= self.medium { - return "medium"; + return EffortLevel::Medium; } - "low" + EffortLevel::Low } } @@ -86,123 +94,163 @@ pub struct ThinkingContext { pub budgets: ThinkingBudgets, } -fn bad_request(message: String) -> Error { - Error::InvalidRequest(message) +fn unmapped_effort(effort: &Value) -> Error { + let choices = ReasoningEffort::ALL + .map(|effort| format!("'{}'", effort.as_str())) + .join(", "); + Error::InvalidRequest(format!( + "Unmapped reasoning effort: {}. Must be one of: {choices}.", + repr(&from_json(effort.clone())) + )) } -fn thinking_type(thinking: Option<&Value>) -> Option<&str> { - thinking?.get("type")?.as_str() +fn unsupported_effort(level: EffortLevel, model: &str) -> Error { + Error::InvalidRequest(format!( + "effort='{}' is not supported by this model. Got model: {model}", + level.as_str() + )) } -fn output_config_effort(output_config: Option<&Value>) -> Option<&str> { - output_config?.get("effort")?.as_str() -} - -fn enabled_thinking(budget_tokens: u64) -> Value { - json!({"type": "enabled", "budget_tokens": budget_tokens}) -} - -fn map_reasoning_effort( - reasoning_effort: &str, - context: &ThinkingContext, -) -> Result, Error> { - if reasoning_effort == "none" { - return Ok(None); +fn output_effort(effort: ReasoningEffort) -> Option { + match effort { + ReasoningEffort::None => None, + ReasoningEffort::Minimal | ReasoningEffort::Low => Some(EffortLevel::Low), + ReasoningEffort::Medium => Some(EffortLevel::Medium), + ReasoningEffort::High => Some(EffortLevel::High), + ReasoningEffort::Xhigh => Some(EffortLevel::Xhigh), + ReasoningEffort::Max => Some(EffortLevel::Max), } - if context.capabilities.supports_adaptive_thinking { - return Ok(Some(json!({"type": "adaptive", "display": "summarized"}))); - } - context - .budgets - .for_effort(reasoning_effort) - .map(|budget| Some(enabled_thinking(budget))) - .ok_or_else(|| { - bad_request(format!( - "Unmapped reasoning effort: '{reasoning_effort}'. Must be one of: {EFFORT_NAMES}." - )) - }) } -fn cap_thinking_budget_to_max_tokens(thinking: Value, max_tokens: Option) -> Option { - let (Some(max_tokens), Some(budget)) = ( - max_tokens, - thinking.get("budget_tokens").and_then(Value::as_u64), - ) else { - return Some(thinking); +fn fit_budget_to_max_tokens(budget_tokens: u64, max_tokens: Option) -> Option { + let Some(max_tokens) = max_tokens else { + return Some(budget_tokens); }; - if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS { - return None; - } - if budget < max_tokens { - return Some(thinking); - } - Some(enabled_thinking(max_tokens - 1)) + (max_tokens > ANTHROPIC_MIN_THINKING_BUDGET_TOKENS).then(|| budget_tokens.min(max_tokens - 1)) } -fn reasoning_effort_to_output_config_effort(reasoning_effort: &str) -> Option<&'static str> { - match reasoning_effort { - "low" | "minimal" => Some("low"), - "medium" => Some("medium"), - "high" => Some("high"), - "xhigh" => Some("xhigh"), - "max" => Some("max"), - _ => None, - } +fn known_thinking(request: &AnthropicMessagesRequest) -> Option<&ThinkingConfig> { + request.params.thinking.as_ref().and_then(Recognized::known) } -fn with_default_effort(output_config: Option, effort: &str) -> Value { - let mut config = match output_config { - Some(Value::Object(config)) => config, - _ => Map::new(), +fn known_effort(request: &AnthropicMessagesRequest) -> Option<&Recognized> { + request + .params + .output_config + .as_ref() + .and_then(Recognized::known) + .and_then(|config| config.effort.as_ref()) +} + +fn with_default_effort( + output_config: Option>, + level: EffortLevel, +) -> Option> { + let config = match output_config { + Some(Recognized::Known(config)) => config, + _ => OutputConfig::default(), }; - if !config.contains_key("effort") { - config.insert("effort".to_string(), Value::String(effort.to_string())); + Some(Recognized::Known(OutputConfig { + effort: Some(config.effort.unwrap_or(Recognized::Known(level))), + ..config + })) +} + +fn without_effort( + output_config: Option>, +) -> Option> { + let Some(Recognized::Known(config)) = output_config else { + return output_config; + }; + if config.effort.is_none() { + return Some(Recognized::Known(config)); + } + let residual = OutputConfig { + effort: None, + ..config + }; + (!residual.is_empty()).then_some(Recognized::Known(residual)) +} + +fn legacy_reasoning_effort( + effort: Option<&Recognized>, +) -> Result { + match effort { + Some(Recognized::Known(level)) => Ok((*level).into()), + Some(Recognized::Unrecognized(value)) if truthy(&from_json(value.clone())) => value + .as_str() + .and_then(ReasoningEffort::parse) + .ok_or_else(|| unmapped_effort(value)), + None | Some(Recognized::Unrecognized(_)) => Ok(ReasoningEffort::Medium), } - Value::Object(config) } fn translate_reasoning_effort( request: AnthropicMessagesRequest, context: &ThinkingContext, ) -> Result { - let Some(reasoning_effort) = request.reasoning_effort.clone() else { + let Some(reasoning_effort) = request.params.reasoning_effort else { return Ok(request); }; let request = AnthropicMessagesRequest { - reasoning_effort: None, + params: AnthropicMessagesOptionalParams { + reasoning_effort: None, + ..request.params + }, ..request }; - let Some(mapped) = map_reasoning_effort(&reasoning_effort, context)? else { + let effort = match reasoning_effort { + Recognized::Known(effort) => effort, + Recognized::Unrecognized(value @ Value::String(_)) => { + return Err(unmapped_effort(&value)); + } + Recognized::Unrecognized(_) => return Ok(request), + }; + let (Some(level), Some(budget)) = (output_effort(effort), context.budgets.for_effort(effort)) + else { return Ok(AnthropicMessagesRequest { - thinking: None, - output_config: None, + params: AnthropicMessagesOptionalParams { + thinking: None, + output_config: None, + ..request.params + }, ..request }); }; - let Some(fitted) = cap_thinking_budget_to_max_tokens(mapped, request.max_tokens) else { + let capabilities = &context.capabilities; + if capabilities.supports_adaptive_thinking { + if !capabilities.accepts_effort(level) { + return Err(unsupported_effort(level, &request.model)); + } + let adaptive = ThinkingConfig::adaptive(Some(ThinkingDisplay::Summarized)); + return Ok(AnthropicMessagesRequest { + params: AnthropicMessagesOptionalParams { + thinking: Some( + request + .params + .thinking + .unwrap_or(Recognized::Known(adaptive)), + ), + output_config: with_default_effort(request.params.output_config, level), + ..request.params + }, + ..request + }); + } + let Some(budget) = fit_budget_to_max_tokens(budget, request.params.max_tokens) else { return Ok(request); }; - let thinking = Some(request.thinking.clone().unwrap_or(fitted)); - if !context.capabilities.supports_adaptive_thinking { - return Ok(AnthropicMessagesRequest { - thinking, - ..request - }); - } - let effort = reasoning_effort_to_output_config_effort(&reasoning_effort).ok_or_else(|| { - bad_request(format!( - "Invalid reasoning_effort: '{reasoning_effort}'. Must be one of: {EFFORT_NAMES}" - )) - })?; - if let Some(rejection) = context - .capabilities - .effort_level_rejection(effort, &request.model) - { - return Err(bad_request(rejection)); - } + let enabled = ThinkingConfig::enabled(budget); Ok(AnthropicMessagesRequest { - thinking, - output_config: Some(with_default_effort(request.output_config.clone(), effort)), + params: AnthropicMessagesOptionalParams { + thinking: Some( + request + .params + .thinking + .unwrap_or(Recognized::Known(enabled)), + ), + ..request.params + }, ..request }) } @@ -212,12 +260,15 @@ fn drop_disabled_thinking( context: &ThinkingContext, ) -> AnthropicMessagesRequest { if !context.capabilities.thinking_always_on - || thinking_type(request.thinking.as_ref()) != Some("disabled") + || !matches!(known_thinking(&request), Some(ThinkingConfig::Disabled(_))) { return request; } AnthropicMessagesRequest { - thinking: None, + params: AnthropicMessagesOptionalParams { + thinking: None, + ..request.params + }, ..request } } @@ -227,40 +278,29 @@ fn translate_legacy_thinking_for_adaptive_model( context: &ThinkingContext, ) -> AnthropicMessagesRequest { let capabilities = &context.capabilities; - if !capabilities.supports_adaptive_thinking - || capabilities.supports_legacy_thinking - || thinking_type(request.thinking.as_ref()) != Some("enabled") - { + if !capabilities.supports_adaptive_thinking || capabilities.supports_legacy_thinking { return request; } - let budget = request - .thinking + let Some(ThinkingConfig::Enabled(enabled)) = known_thinking(&request) else { + return request; + }; + let budget = enabled + .budget_tokens .as_ref() - .and_then(|thinking| thinking.get("budget_tokens")) - .and_then(Value::as_u64) + .and_then(Recognized::known) + .copied() .unwrap_or(0); - let effort = context.budgets.effort_for_budget(budget, capabilities); + let level = context.budgets.effort_for_budget(budget, capabilities); AnthropicMessagesRequest { - thinking: Some(json!({"type": "adaptive"})), - output_config: Some(with_default_effort(request.output_config.clone(), effort)), + params: AnthropicMessagesOptionalParams { + thinking: Some(Recognized::Known(ThinkingConfig::adaptive(None))), + output_config: with_default_effort(request.params.output_config, level), + ..request.params + }, ..request } } -fn output_config_without_effort(output_config: Option) -> Option { - let Some(Value::Object(config)) = output_config else { - return output_config; - }; - if !config.contains_key("effort") { - return Some(Value::Object(config)); - } - let residual: Map = config - .into_iter() - .filter(|(key, _)| key != "effort") - .collect(); - (!residual.is_empty()).then_some(Value::Object(residual)) -} - fn translate_adaptive_effort_for_non_adaptive_model( request: AnthropicMessagesRequest, context: &ThinkingContext, @@ -269,42 +309,43 @@ fn translate_adaptive_effort_for_non_adaptive_model( if capabilities.supports_adaptive_thinking { return Ok(request); } - let effort = output_config_effort(request.output_config.as_ref()).map(str::to_string); - let adaptive_thinking = thinking_type(request.thinking.as_ref()) == Some("adaptive"); + let effort = known_effort(&request).cloned(); + let adaptive_thinking = matches!(known_thinking(&request), Some(ThinkingConfig::Adaptive(_))); if effort.is_none() && !adaptive_thinking { return Ok(request); } - let level_supported = effort.as_deref().is_none_or(|effort| { - capabilities - .effort_level_rejection(effort, &request.model) - .is_none() - }); - if capabilities.supports_effort_param() && (!adaptive_thinking || level_supported) { + let level_accepted = match &effort { + Some(Recognized::Known(level)) => capabilities.accepts_effort(*level), + _ => true, + }; + if capabilities.supports_effort_param() && (!adaptive_thinking || level_accepted) { return Ok(AnthropicMessagesRequest { - thinking: if adaptive_thinking { - None - } else { - request.thinking.clone() + params: AnthropicMessagesOptionalParams { + thinking: if adaptive_thinking { + None + } else { + request.params.thinking + }, + ..request.params }, ..request }); } - let legacy = if capabilities.supports_reasoning { - map_reasoning_effort( - effort - .as_deref() - .filter(|effort| !effort.is_empty()) - .unwrap_or("medium"), - context, - )? + let budget = if capabilities.supports_reasoning { + context + .budgets + .for_effort(legacy_reasoning_effort(effort.as_ref())?) } else { None }; - let capped = - legacy.and_then(|thinking| cap_thinking_budget_to_max_tokens(thinking, request.max_tokens)); Ok(AnthropicMessagesRequest { - thinking: capped, - output_config: output_config_without_effort(request.output_config.clone()), + params: AnthropicMessagesOptionalParams { + thinking: budget + .and_then(|budget| fit_budget_to_max_tokens(budget, request.params.max_tokens)) + .map(|budget| Recognized::Known(ThinkingConfig::enabled(budget))), + output_config: without_effort(request.params.output_config), + ..request.params + }, ..request }) } @@ -317,15 +358,19 @@ fn drop_incompatible_temperature_for_thinking( return request; } let pinned = request + .params .temperature .is_some_and(|temperature| temperature != 1.0); - let thinking_enabled = thinking_type(request.thinking.as_ref()) == Some("enabled"); - let effort_enabled = output_config_effort(request.output_config.as_ref()).is_some(); + let thinking_enabled = matches!(known_thinking(&request), Some(ThinkingConfig::Enabled(_))); + let effort_enabled = known_effort(&request).is_some(); if !pinned || !(thinking_enabled || effort_enabled) { return request; } AnthropicMessagesRequest { - temperature: None, + params: AnthropicMessagesOptionalParams { + temperature: None, + ..request.params + }, ..request } } @@ -348,11 +393,10 @@ mod tests { use super::*; use crate::anthropic::common_utils::SupportedEffortTiers; - const EFFORT_CHOICES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"; + const EFFORT_CHOICES: &str = "'none', 'minimal', 'low', 'medium', 'high', 'xhigh', 'max'"; fn request(fields: Value) -> AnthropicMessagesRequest { - let mut body = - json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); + let mut body = serde_json::json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); body.as_object_mut() .unwrap() .extend(fields.as_object().unwrap().clone()); @@ -386,7 +430,7 @@ mod tests { } fn claude_code_payload(effort: &str, max_tokens: u64) -> Value { - json!({"max_tokens": max_tokens, "thinking": {"type": "adaptive"}, "output_config": {"effort": effort}}) + serde_json::json!({"max_tokens": max_tokens, "thinking": {"type": "adaptive"}, "output_config": {"effort": effort}}) } fn with_temperature(fields: Value, temperature: f64) -> Value { @@ -394,7 +438,7 @@ mod tests { fields .as_object_mut() .unwrap() - .insert("temperature".to_string(), json!(temperature)); + .insert("temperature".to_string(), serde_json::json!(temperature)); fields } @@ -485,9 +529,9 @@ mod tests { assert_eq!( translate( capabilities, - json!({"max_tokens": 1024, "reasoning_effort": reasoning_effort}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": reasoning_effort}) ), - Ok(request(json!({ + Ok(request(serde_json::json!({ "max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": expected_effort} @@ -498,18 +542,18 @@ mod tests { #[rstest] #[case::adaptive_shape_is_not_dropped_for_small_max_tokens( opus_4_7(), - json!({"max_tokens": 64, "reasoning_effort": "high"}), - json!({"max_tokens": 64, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 64, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 64, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) )] #[case::caller_output_config_effort_wins( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": "low", "output_config": {"effort": "max"}}), - json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "max"}}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "low", "output_config": {"effort": "max"}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "max"}}) )] #[case::effort_merges_into_caller_output_config( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": {"format": {"type": "json_schema"}}}), - json!({ + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": {"format": {"type": "json_schema"}}}), + serde_json::json!({ "max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"format": {"type": "json_schema"}, "effort": "high"} @@ -517,18 +561,18 @@ mod tests { )] #[case::non_object_output_config_is_replaced( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": "bogus"}), - json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": "bogus"}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) )] #[case::caller_thinking_and_output_config_win( sonnet_4_6(), - json!({ + serde_json::json!({ "max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}, "output_config": {"effort": "high"} }), - json!({ + serde_json::json!({ "max_tokens": 16000, "thinking": {"type": "enabled", "budget_tokens": 8000}, "output_config": {"effort": "high"} @@ -536,63 +580,86 @@ mod tests { )] #[case::caller_legacy_thinking_is_then_translated_while_reasoning_effort_level_stays( opus_4_7(), - json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), - json!({"max_tokens": 16000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) + serde_json::json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + serde_json::json!({"max_tokens": 16000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) )] #[case::caller_disabled_thinking_is_kept_then_omitted_on_always_on_model( fable_5_1(), - json!({"max_tokens": 1024, "reasoning_effort": "high", "thinking": {"type": "disabled"}}), - json!({"max_tokens": 1024, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "high", "thinking": {"type": "disabled"}}), + serde_json::json!({"max_tokens": 1024, "output_config": {"effort": "high"}}) )] #[case::non_adaptive_model_gets_no_output_config( opus_4_5(), - json!({"max_tokens": 8192, "reasoning_effort": "high"}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::caller_thinking_wins_on_non_adaptive_model( opus_4_5(), - json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), - json!({"max_tokens": 16000, "thinking": {"type": "enabled", "budget_tokens": 8000}}) + serde_json::json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + serde_json::json!({"max_tokens": 16000, "thinking": {"type": "enabled", "budget_tokens": 8000}}) )] #[case::caller_thinking_survives_when_mapped_budget_cannot_fit( opus_4_5(), - json!({"max_tokens": 1024, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 8000}}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 8000}}) )] #[case::missing_max_tokens_leaves_budget_uncapped( haiku_4_5(), - json!({"reasoning_effort": "high"}), - json!({"thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"reasoning_effort": "high"}), + serde_json::json!({"thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::budget_below_max_tokens_is_kept( haiku_4_5(), - json!({"max_tokens": 4097, "reasoning_effort": "high"}), - json!({"max_tokens": 4097, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 4097, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 4097, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::budget_equal_to_max_tokens_is_capped( haiku_4_5(), - json!({"max_tokens": 4096, "reasoning_effort": "high"}), - json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 4095}}) + serde_json::json!({"max_tokens": 4096, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 4095}}) )] #[case::budget_above_max_tokens_is_capped( haiku_4_5(), - json!({"max_tokens": 4000, "reasoning_effort": "xhigh"}), - json!({"max_tokens": 4000, "thinking": {"type": "enabled", "budget_tokens": 3999}}) + serde_json::json!({"max_tokens": 4000, "reasoning_effort": "xhigh"}), + serde_json::json!({"max_tokens": 4000, "thinking": {"type": "enabled", "budget_tokens": 3999}}) )] #[case::max_tokens_just_above_min_budget_caps_to_min_budget( haiku_4_5(), - json!({"max_tokens": 1025, "reasoning_effort": "xhigh"}), - json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + serde_json::json!({"max_tokens": 1025, "reasoning_effort": "xhigh"}), + serde_json::json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) )] #[case::max_tokens_at_min_budget_drops_thinking( haiku_4_5(), - json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), - json!({"max_tokens": 1024}) + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + serde_json::json!({"max_tokens": 1024}) )] #[case::pinned_temperature_is_dropped_after_thinking_is_synthesized( haiku_4_5(), - json!({"max_tokens": 8192, "reasoning_effort": "low", "temperature": 0}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + serde_json::json!({"max_tokens": 8192, "reasoning_effort": "low", "temperature": 0}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + #[case::non_string_reasoning_effort_is_ignored( + opus_4_7(), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": 3, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "adaptive"}}) + )] + #[case::unrecognized_thinking_is_forwarded( + opus_4_5(), + serde_json::json!({"max_tokens": 8192, "reasoning_effort": "high", "thinking": {"type": "future"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "future"}}) + )] + #[case::caller_display_and_block_binding_survive_on_adaptive_model( + opus_4_7(), + serde_json::json!({ + "max_tokens": 1024, + "reasoning_effort": "high", + "thinking": {"type": "adaptive", "display": "omitted", "block_binding": {"prefix_mismatch_behavior": "drop_block"}} + }), + serde_json::json!({ + "max_tokens": 1024, + "thinking": {"type": "adaptive", "display": "omitted", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + "output_config": {"effort": "high"} + }) )] fn reasoning_effort_is_translated( #[case] capabilities: AnthropicModelCapabilities, @@ -617,9 +684,9 @@ mod tests { assert_eq!( translate( haiku_4_5, - json!({"max_tokens": 32000, "reasoning_effort": reasoning_effort}) + serde_json::json!({"max_tokens": 32000, "reasoning_effort": reasoning_effort}) ), - Ok(request(json!({ + Ok(request(serde_json::json!({ "max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": expected_budget} }))) @@ -636,56 +703,56 @@ mod tests { assert_eq!( translate( capabilities, - json!({ + serde_json::json!({ "max_tokens": 1024, "reasoning_effort": "none", "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"} }) ), - Ok(request(json!({"max_tokens": 1024}))) + Ok(request(serde_json::json!({"max_tokens": 1024}))) ); } #[rstest] #[case::bogus_on_budget_model( opus_4_5(), - json!({"max_tokens": 1024, "reasoning_effort": "bogus"}), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "bogus"}), format!("Unmapped reasoning effort: 'bogus'. Must be one of: {EFFORT_CHOICES}.") )] #[case::disabled_on_budget_model( haiku_4_5(), - json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), format!("Unmapped reasoning effort: 'disabled'. Must be one of: {EFFORT_CHOICES}.") )] #[case::empty_on_budget_model( haiku_4_5(), - json!({"max_tokens": 1024, "reasoning_effort": ""}), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": ""}), format!("Unmapped reasoning effort: ''. Must be one of: {EFFORT_CHOICES}.") )] #[case::invalid_on_adaptive_model( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": "invalid"}), - format!("Invalid reasoning_effort: 'invalid'. Must be one of: {EFFORT_CHOICES}") + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "invalid"}), + format!("Unmapped reasoning effort: 'invalid'. Must be one of: {EFFORT_CHOICES}.") )] #[case::disabled_on_adaptive_model( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), - format!("Invalid reasoning_effort: 'disabled'. Must be one of: {EFFORT_CHOICES}") + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), + format!("Unmapped reasoning effort: 'disabled'. Must be one of: {EFFORT_CHOICES}.") )] #[case::empty_on_adaptive_model( opus_4_7(), - json!({"max_tokens": 1024, "reasoning_effort": ""}), - format!("Invalid reasoning_effort: ''. Must be one of: {EFFORT_CHOICES}") + serde_json::json!({"max_tokens": 1024, "reasoning_effort": ""}), + format!("Unmapped reasoning effort: ''. Must be one of: {EFFORT_CHOICES}.") )] #[case::xhigh_without_xhigh_tier_on_4_6( sonnet_4_6(), - json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), "effort='xhigh' is not supported by this model. Got model: claude".to_string() )] #[case::xhigh_without_xhigh_tier_on_unmapped_adaptive_model( newfamily_6(), - json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), "effort='xhigh' is not supported by this model. Got model: claude".to_string() )] #[case::unrecognized_adaptive_effort_on_budget_model( @@ -693,6 +760,16 @@ mod tests { claude_code_payload("turbo", 8192), format!("Unmapped reasoning effort: 'turbo'. Must be one of: {EFFORT_CHOICES}.") )] + #[case::unrecognized_output_config_effort_on_budget_model( + haiku_4_5(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"effort": 5}}), + format!("Unmapped reasoning effort: 5. Must be one of: {EFFORT_CHOICES}.") + )] + #[case::quote_in_effort_is_reprd_like_python( + haiku_4_5(), + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "it's"}), + format!("Unmapped reasoning effort: \"it's\". Must be one of: {EFFORT_CHOICES}.") + )] fn unsupported_effort_is_a_request_error( #[case] capabilities: AnthropicModelCapabilities, #[case] input: Value, @@ -705,13 +782,13 @@ mod tests { } #[rstest] - #[case::omitted_on_always_on_model(fable_5_1(), json!({"type": "disabled"}), None)] - #[case::kept_on_adaptive_model(opus_4_7(), json!({"type": "disabled"}), Some(json!({"type": "disabled"})))] - #[case::kept_on_budget_model(haiku_4_5(), json!({"type": "disabled"}), Some(json!({"type": "disabled"})))] + #[case::omitted_on_always_on_model(fable_5_1(), serde_json::json!({"type": "disabled"}), None)] + #[case::kept_on_adaptive_model(opus_4_7(), serde_json::json!({"type": "disabled"}), Some(serde_json::json!({"type": "disabled"})))] + #[case::kept_on_budget_model(haiku_4_5(), serde_json::json!({"type": "disabled"}), Some(serde_json::json!({"type": "disabled"})))] #[case::adaptive_kept_on_always_on_model( fable_5_1(), - json!({"type": "adaptive"}), - Some(json!({"type": "adaptive"})) + serde_json::json!({"type": "adaptive"}), + Some(serde_json::json!({"type": "adaptive"})) )] fn disabled_thinking_is_omitted_only_for_always_on_models( #[case] capabilities: AnthropicModelCapabilities, @@ -719,46 +796,46 @@ mod tests { #[case] expected_thinking: Option, ) { let expected = match expected_thinking { - Some(thinking) => json!({"max_tokens": 64, "thinking": thinking}), - None => json!({"max_tokens": 64}), + Some(thinking) => serde_json::json!({"max_tokens": 64, "thinking": thinking}), + None => serde_json::json!({"max_tokens": 64}), }; assert_eq!( translate( capabilities, - json!({"max_tokens": 64, "thinking": thinking}) + serde_json::json!({"max_tokens": 64, "thinking": thinking}) ), Ok(request(expected)) ); } #[rstest] - #[case::far_above_xhigh_budget(opus_4_7(), json!(16384), "xhigh")] - #[case::at_xhigh_budget(opus_4_7(), json!(8192), "xhigh")] - #[case::below_xhigh_budget(opus_4_7(), json!(8191), "high")] - #[case::xhigh_budget_without_xhigh_tier(newfamily_6(), json!(8192), "high")] - #[case::large_budget_without_xhigh_tier(newfamily_6(), json!(31999), "high")] - #[case::at_high_budget(opus_4_7(), json!(4096), "high")] - #[case::below_high_budget(opus_4_7(), json!(4095), "medium")] - #[case::at_medium_budget(opus_4_7(), json!(2048), "medium")] - #[case::below_medium_budget(opus_4_7(), json!(2047), "low")] - #[case::tiny_budget(opus_4_7(), json!(1), "low")] + #[case::far_above_xhigh_budget(opus_4_7(), serde_json::json!(16384), "xhigh")] + #[case::at_xhigh_budget(opus_4_7(), serde_json::json!(8192), "xhigh")] + #[case::below_xhigh_budget(opus_4_7(), serde_json::json!(8191), "high")] + #[case::xhigh_budget_without_xhigh_tier(newfamily_6(), serde_json::json!(8192), "high")] + #[case::large_budget_without_xhigh_tier(newfamily_6(), serde_json::json!(31999), "high")] + #[case::at_high_budget(opus_4_7(), serde_json::json!(4096), "high")] + #[case::below_high_budget(opus_4_7(), serde_json::json!(4095), "medium")] + #[case::at_medium_budget(opus_4_7(), serde_json::json!(2048), "medium")] + #[case::below_medium_budget(opus_4_7(), serde_json::json!(2047), "low")] + #[case::tiny_budget(opus_4_7(), serde_json::json!(1), "low")] #[case::missing_budget(opus_4_7(), Value::Null, "low")] - #[case::always_on_model(fable_5_1(), json!(24000), "xhigh")] + #[case::always_on_model(fable_5_1(), serde_json::json!(24000), "xhigh")] fn legacy_thinking_is_bucketed_into_adaptive_effort_on_adaptive_only_models( #[case] capabilities: AnthropicModelCapabilities, #[case] budget_tokens: Value, #[case] expected_effort: &str, ) { let thinking = match budget_tokens { - Value::Null => json!({"type": "enabled"}), - budget_tokens => json!({"type": "enabled", "budget_tokens": budget_tokens}), + Value::Null => serde_json::json!({"type": "enabled"}), + budget_tokens => serde_json::json!({"type": "enabled", "budget_tokens": budget_tokens}), }; assert_eq!( translate( capabilities, - json!({"max_tokens": 1024, "thinking": thinking}) + serde_json::json!({"max_tokens": 1024, "thinking": thinking}) ), - Ok(request(json!({ + Ok(request(serde_json::json!({ "max_tokens": 1024, "thinking": {"type": "adaptive"}, "output_config": {"effort": expected_effort} @@ -769,27 +846,27 @@ mod tests { #[rstest] #[case::verbatim_on_model_accepting_legacy_thinking( sonnet_4_6(), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) )] #[case::verbatim_with_explicit_output_config_on_model_accepting_legacy_thinking( sonnet_4_6(), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}) + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}) )] #[case::verbatim_on_non_adaptive_model( opus_4_5(), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), - json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), + serde_json::json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) )] #[case::caller_output_config_effort_wins( opus_4_7(), - json!({ + serde_json::json!({ "max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low", "format": {"type": "json_schema"}} }), - json!({ + serde_json::json!({ "max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low", "format": {"type": "json_schema"}} @@ -797,12 +874,12 @@ mod tests { )] #[case::effort_merges_into_caller_output_config( opus_4_7(), - json!({ + serde_json::json!({ "max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"format": {"type": "json_schema"}} }), - json!({ + serde_json::json!({ "max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high", "format": {"type": "json_schema"}} @@ -810,8 +887,8 @@ mod tests { )] #[case::adaptive_thinking_is_left_alone( opus_4_7(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}) )] fn legacy_thinking_on_adaptive_capable_models( #[case] capabilities: AnthropicModelCapabilities, @@ -824,42 +901,42 @@ mod tests { #[rstest] #[case::bare_adaptive_becomes_medium_budget_on_budget_model( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::medium_effort_becomes_medium_budget_on_budget_model( haiku_4_5(), claude_code_payload("medium", 8192), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::empty_effort_becomes_medium_budget_on_budget_model( haiku_4_5(), claude_code_payload("", 8192), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::high_effort_becomes_high_budget_on_budget_model( haiku_4_5(), claude_code_payload("high", 8192), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::effort_only_becomes_budget_on_budget_model( haiku_4_5(), - json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::effort_replaces_caller_legacy_budget_on_budget_model( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 3000}, "output_config": {"effort": "high"}}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 3000}, "output_config": {"effort": "high"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::residual_output_config_survives_effort_translation( haiku_4_5(), - json!({ + serde_json::json!({ "max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium", "format": {"type": "json_schema"}} }), - json!({ + serde_json::json!({ "max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}, "output_config": {"format": {"type": "json_schema"}} @@ -867,8 +944,8 @@ mod tests { )] #[case::effortless_output_config_is_kept( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"format": {"type": "json_schema"}}}), - json!({ + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"format": {"type": "json_schema"}}}), + serde_json::json!({ "max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}, "output_config": {"format": {"type": "json_schema"}} @@ -876,88 +953,88 @@ mod tests { )] #[case::empty_output_config_is_kept( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {}}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}, "output_config": {}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}, "output_config": {}}) )] #[case::missing_max_tokens_leaves_budget_uncapped( haiku_4_5(), - json!({"thinking": {"type": "adaptive"}}), - json!({"thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"thinking": {"type": "adaptive"}}), + serde_json::json!({"thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::budget_is_capped_below_max_tokens( haiku_4_5(), claude_code_payload("high", 3000), - json!({"max_tokens": 3000, "thinking": {"type": "enabled", "budget_tokens": 2999}}) + serde_json::json!({"max_tokens": 3000, "thinking": {"type": "enabled", "budget_tokens": 2999}}) )] #[case::max_tokens_just_above_min_budget_caps_to_min_budget( haiku_4_5(), claude_code_payload("medium", 1025), - json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + serde_json::json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) )] #[case::max_tokens_at_min_budget_drops_thinking_and_effort( haiku_4_5(), claude_code_payload("medium", 1024), - json!({"max_tokens": 1024}) + serde_json::json!({"max_tokens": 1024}) )] #[case::max_tokens_below_min_budget_drops_thinking_and_effort( haiku_4_5(), claude_code_payload("medium", 512), - json!({"max_tokens": 512}) + serde_json::json!({"max_tokens": 512}) )] #[case::bare_adaptive_is_dropped_on_non_reasoning_model( haiku_3_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), - json!({"max_tokens": 8192}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 8192}) )] #[case::adaptive_and_effort_are_dropped_on_non_reasoning_model( haiku_3_5(), claude_code_payload("medium", 8192), - json!({"max_tokens": 8192}) + serde_json::json!({"max_tokens": 8192}) )] #[case::effort_only_is_dropped_on_non_reasoning_model( haiku_3_5(), - json!({"max_tokens": 8192, "output_config": {"effort": "high", "format": {"type": "json_schema"}}}), - json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high", "format": {"type": "json_schema"}}}), + serde_json::json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) )] #[case::supported_effort_is_kept_and_adaptive_thinking_dropped_on_effort_model( opus_4_5(), claude_code_payload("medium", 8192), - json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) )] #[case::bare_adaptive_is_dropped_on_effort_model( opus_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), - json!({"max_tokens": 8192}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 8192}) )] #[case::effort_only_is_left_alone_on_effort_model( opus_4_5(), - json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), - json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) )] #[case::unsupported_effort_only_is_left_for_provider_normalization( opus_4_5(), - json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}), - json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}) + serde_json::json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}), + serde_json::json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}) )] #[case::legacy_thinking_is_kept_beside_native_effort_on_effort_model( opus_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}) )] #[case::unsupported_xhigh_with_adaptive_thinking_falls_back_to_budget( opus_4_5(), claude_code_payload("xhigh", 64000), - json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 8192}}) + serde_json::json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 8192}}) )] #[case::unsupported_max_with_adaptive_thinking_falls_back_to_budget( opus_4_5(), claude_code_payload("max", 64000), - json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 16384}}) + serde_json::json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 16384}}) )] #[case::bare_adaptive_is_native_on_4_6( sonnet_4_6(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) )] #[case::adaptive_payload_is_native_on_4_6( sonnet_4_6(), @@ -966,8 +1043,41 @@ mod tests { )] #[case::request_without_adaptive_interface_is_left_alone( haiku_4_5(), - json!({"max_tokens": 1024}), - json!({"max_tokens": 1024}) + serde_json::json!({"max_tokens": 1024}), + serde_json::json!({"max_tokens": 1024}) + )] + #[case::falsy_non_string_effort_becomes_medium_budget_on_budget_model( + haiku_4_5(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"effort": 0}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::minimal_effort_becomes_floored_minimal_budget_on_budget_model( + haiku_4_5(), + claude_code_payload("minimal", 8192), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + #[case::none_effort_drops_thinking_on_budget_model( + haiku_4_5(), + claude_code_payload("none", 8192), + serde_json::json!({"max_tokens": 8192}) + )] + #[case::unrecognized_effort_is_native_on_effort_model( + opus_4_5(), + claude_code_payload("turbo", 8192), + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "turbo"}}) + )] + #[case::task_budget_survives_effort_translation( + haiku_4_5(), + serde_json::json!({ + "max_tokens": 8192, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "high", "task_budget": {"type": "tokens", "total": 4096}} + }), + serde_json::json!({ + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "output_config": {"task_budget": {"type": "tokens", "total": 4096}} + }) )] fn adaptive_interface_is_reshaped_for_non_adaptive_models( #[case] capabilities: AnthropicModelCapabilities, @@ -982,37 +1092,37 @@ mod tests { haiku_4_5(), claude_code_payload("medium", 8192), 0.0, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::bare_adaptive_downgraded_to_enabled_thinking( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), 0.0, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::reasoning_effort_synthesized_enabled_thinking( haiku_4_5(), - json!({"max_tokens": 8192, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 8192, "reasoning_effort": "high"}), 0.2, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] #[case::above_one_with_enabled_thinking( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}), 1.5, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::native_effort_kept_on_effort_model( opus_4_5(), claude_code_payload("medium", 8192), 0.0, - json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) )] #[case::effort_only_on_effort_model( opus_4_5(), - json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), 0.0, - json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) )] fn pinned_temperature_is_dropped_when_thinking_or_effort_survives_on_non_adaptive_model( #[case] capabilities: AnthropicModelCapabilities, @@ -1031,32 +1141,32 @@ mod tests { haiku_4_5(), claude_code_payload("medium", 8192), 1.0, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) )] #[case::thinking_dropped_for_small_max_tokens( haiku_4_5(), claude_code_payload("medium", 512), 0.0, - json!({"max_tokens": 512}) + serde_json::json!({"max_tokens": 512}) )] #[case::thinking_dropped_on_non_reasoning_model( haiku_3_5(), claude_code_payload("medium", 8192), 0.0, - json!({"max_tokens": 8192}) + serde_json::json!({"max_tokens": 8192}) )] #[case::disabled_thinking( haiku_4_5(), - json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}), 0.0, - json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}) )] - #[case::no_thinking(haiku_4_5(), json!({"max_tokens": 8192}), 0.0, json!({"max_tokens": 8192}))] + #[case::no_thinking(haiku_4_5(), serde_json::json!({"max_tokens": 8192}), 0.0, serde_json::json!({"max_tokens": 8192}))] #[case::output_config_without_effort( haiku_4_5(), - json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}), + serde_json::json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}), 0.0, - json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) + serde_json::json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) )] #[case::adaptive_model( opus_4_7(), @@ -1066,9 +1176,9 @@ mod tests { )] #[case::legacy_thinking_on_adaptive_model( sonnet_4_6(), - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}), 0.0, - json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] fn temperature_is_kept( #[case] capabilities: AnthropicModelCapabilities, @@ -1113,56 +1223,56 @@ mod tests { #[case::reasoning_effort_uses_overridden_budget( &[("HIGH", "6000")], haiku_4_5(), - json!({"max_tokens": 32000, "reasoning_effort": "high"}), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}) + serde_json::json!({"max_tokens": 32000, "reasoning_effort": "high"}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}) )] #[case::minimal_override_below_min_budget_is_floored( &[("MINIMAL", "512")], haiku_4_5(), - json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + serde_json::json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 1024}}) )] #[case::minimal_override_above_min_budget_is_used( &[("MINIMAL", "2000")], haiku_4_5(), - json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2000}}) + serde_json::json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2000}}) )] #[case::adaptive_fallback_uses_overridden_medium_budget( &[("MEDIUM", "3000")], haiku_4_5(), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}}), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}) )] #[case::legacy_bucket_below_overridden_high_budget( &[("HIGH", "6000")], opus_4_7(), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 5999}}), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 5999}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) )] #[case::legacy_bucket_at_overridden_high_budget( &[("HIGH", "6000")], opus_4_7(), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) )] #[case::legacy_bucket_below_overridden_xhigh_budget( &[("XHIGH", "20000")], opus_4_7(), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 19999}}), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 19999}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) )] #[case::legacy_bucket_at_overridden_medium_budget( &[("MEDIUM", "3000")], opus_4_7(), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) )] #[case::legacy_bucket_below_overridden_medium_budget( &[("MEDIUM", "3000")], opus_4_7(), - json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2999}}), - json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2999}}), + serde_json::json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) )] fn translation_honors_budget_overrides( #[case] overrides: &[(&str, &str)], diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 59280c04a70..280ea63eefa 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,21 +1,21 @@ use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; -use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use litellm_types::llms::anthropic_messages::anthropic_request::{ + AnthropicMessagesOptionalParams, AnthropicMessagesRequest, +}; use serde_json::{Map, Value, json}; use super::{ - headers::{authenticate, with_feature_betas}, + headers::{validate_environment, with_feature_betas}, thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}, }; use crate::{ + Error, anthropic::common_utils::{ AnthropicModelCapabilities, has_advisor_tool, strip_advisor_blocks, strip_encrypted_reasoning_blocks, }, - base_llm::{ - anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, - }, - chat::transformation::Error, + base_llm::anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, }, }; @@ -65,7 +65,7 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { request: AnthropicMessagesRequest, context: &MessagesTransformContext, ) -> Result { - if request.max_tokens.is_none() { + if request.params.max_tokens.is_none() { return Err(Error::InvalidRequest( "max_tokens is required for Anthropic /v1/messages API".to_string(), )); @@ -73,30 +73,26 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { let request = drop_unsupported_params(request, context)?; let request = translate_thinking(request, &context.thinking)?; let context_management = request + .params .context_management .as_ref() .and_then(map_openai_context_management_to_anthropic) - .or_else(|| request.context_management.clone()); - let messages = if has_advisor_tool(request.tools.as_deref()) { + .or_else(|| request.params.context_management.clone()); + let messages = if has_advisor_tool(request.params.tools.as_deref()) { request.messages } else { strip_advisor_blocks(request.messages) }; Ok(AnthropicMessagesRequest { messages: strip_encrypted_reasoning_blocks(messages), - context_management, + params: AnthropicMessagesOptionalParams { + context_management, + ..request.params + }, ..request }) } - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from) - } - fn secret_names(&self) -> &'static [&'static str] { &[ ANTHROPIC_API_KEY_ENV, @@ -106,13 +102,14 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { ] } - fn authenticate( + fn validate_environment( &self, headers: Headers, api_key: Option<&str>, + _model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - authenticate(headers, api_key, env_lookup).map_err(Error::from) + ) -> Result { + validate_environment(headers, api_key, env_lookup).map_err(Error::from) } fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers { @@ -138,17 +135,21 @@ fn drop_unsupported_params( } Err(unsupported_param(&model, param, &value, hint)) }; - let speed = match request.speed.as_deref() { + let params = request.params; + let speed = match params.speed.as_deref() { Some(speed) if !capabilities.supports_speed => { reject("speed", format!("'{speed}'"), "")?; None } - _ => request.speed.clone(), + _ => params.speed.clone(), }; if capabilities.supports_sampling_params { - return Ok(AnthropicMessagesRequest { speed, ..request }); + return Ok(AnthropicMessagesRequest { + params: AnthropicMessagesOptionalParams { speed, ..params }, + ..request + }); } - let temperature = match request.temperature { + let temperature = match params.temperature { Some(temperature) if temperature != 1.0 => { reject( "temperature", @@ -159,17 +160,20 @@ fn drop_unsupported_params( } temperature => temperature, }; - if let Some(top_p) = request.top_p { + if let Some(top_p) = params.top_p { reject("top_p", json!(top_p).to_string(), "")?; } - if let Some(top_k) = request.top_k { + if let Some(top_k) = params.top_k { reject("top_k", json!(top_k).to_string(), "")?; } Ok(AnthropicMessagesRequest { - speed, - temperature, - top_p: None, - top_k: None, + params: AnthropicMessagesOptionalParams { + speed, + temperature, + top_p: None, + top_k: None, + ..params + }, ..request }) } @@ -251,10 +255,14 @@ pub fn resolve_anthropic_api_base( mod tests { use std::process::Command; + use litellm_auth::CredentialPlacement; use rstest::{fixture, rstest}; use super::*; - use crate::anthropic::common_utils::{ENCRYPTED_REASONING_SIGNATURE_PREFIX, beta}; + use crate::{ + anthropic::common_utils::{ENCRYPTED_REASONING_SIGNATURE_PREFIX, beta}, + base_llm::auth::AuthScheme, + }; type Env = &'static [(&'static str, &'static str)]; @@ -806,25 +814,32 @@ mod tests { #[test] fn config_reports_a_missing_key_as_an_auth_error() { - assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.resolve_api_key(None, &no_env), + assert!(matches!( + ANTHROPIC_MESSAGES_CONFIG.validate_environment(vec![], None, "claude", &no_env), Err(Error::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: ANTHROPIC_API_KEY_ENV, })) - ); + )); } #[test] fn config_authenticates_with_the_anthropic_auth_token() { - assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.authenticate( + let validated = ANTHROPIC_MESSAGES_CONFIG + .validate_environment( vec![], None, - &env(&[("ANTHROPIC_AUTH_TOKEN", "auth-token")]) - ), - Ok(headers(&[("authorization", "Bearer auth-token")])) - ); + "claude", + &env(&[("ANTHROPIC_AUTH_TOKEN", "auth-token")]), + ) + .unwrap(); + assert!(matches!( + validated.auth, + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + ref secret + } if secret.expose() == "auth-token" + )); } #[test] @@ -853,11 +868,7 @@ mod tests { } #[test] - fn auth_strategy_and_default_headers_match_anthropic() { - assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.auth_strategy().header_name(), - "x-api-key" - ); + fn default_headers_match_anthropic() { assert_eq!( ANTHROPIC_MESSAGES_CONFIG.default_headers(), &[ @@ -874,7 +885,7 @@ mod tests { requested.borrow_mut().push(name.to_string()); None }; - let _ = ANTHROPIC_MESSAGES_CONFIG.authenticate(Vec::new(), None, &record); + let _ = ANTHROPIC_MESSAGES_CONFIG.validate_environment(Vec::new(), None, "claude", &record); let _ = ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); let requested = requested.into_inner(); assert!(!requested.is_empty()); diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index 2ce1b0da51b..1defec654bd 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -63,9 +63,14 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { async fn validate_environment( &self, request: &PreparedOcrRequest, - _client: &OcrClient, + client: &OcrClient, ) -> Result { - environment(request, TextractOperation::AnalyzeDocument).await + environment( + &client.auth().aws, + request, + TextractOperation::AnalyzeDocument, + ) + .await } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index 8268ad066a1..678104be982 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -1,5 +1,5 @@ use base64::{Engine, engine::general_purpose::STANDARD}; -use litellm_auth_aws::{SigV4Signer, resolve_aws_region}; +use litellm_auth_aws::{AwsCredentialSource, SigV4Signer, resolve_aws_region}; use litellm_http::outbound::RequestSigner; use serde::{Deserialize, Serialize}; use strum::{EnumString, IntoStaticStr, VariantNames}; @@ -232,6 +232,7 @@ pub(super) fn health_check_document() -> OcrDocument { } pub(super) async fn environment( + auth: &litellm_auth_aws::AwsAuthService, request: &PreparedOcrRequest, operation: TextractOperation, ) -> Result { @@ -244,9 +245,10 @@ pub(super) async fn environment( ) })?; let signer = SigV4Signer::resolve( + auth, region.clone(), TEXTRACT_SERVICE, - &request.optional_params, + AwsCredentialSource::from_params(&request.optional_params, &env_lookup), &env_lookup, ) .await diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index 6eb195defaa..3b4f8e5a7d9 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -48,9 +48,14 @@ impl BaseOcrConfig for TextractDetectTextConfig { async fn validate_environment( &self, request: &PreparedOcrRequest, - _client: &OcrClient, + client: &OcrClient, ) -> Result { - environment(request, TextractOperation::DetectDocumentText).await + environment( + &client.auth().aws, + request, + TextractOperation::DetectDocumentText, + ) + .await } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs index 137239bbeaf..8c768a6a66d 100644 --- a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs @@ -1,19 +1,23 @@ +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_http::request::{has_bearer_auth, has_header}; use litellm_types::llms::anthropic_messages::{ anthropic_request::{ - AnthropicMessage, AnthropicMessagesRequest, ContentBlock, MessageContent, SystemPrompt, + AnthropicMessage, AnthropicMessagesOptionalParams, AnthropicMessagesRequest, ContentBlock, + MessageContent, SystemPrompt, }, anthropic_response::AnthropicMessagesResponse, }; use crate::{ + Error, anthropic::messages::transformation::{ ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }, base_llm::{ anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, Headers, MessagesAuthStrategy, MessagesTransformContext, + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, }, - chat::transformation::Error, + auth::AuthScheme, }, }; @@ -22,6 +26,7 @@ const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; const SYSTEM_ROLE: &str = "system"; +const API_KEY_HEADER: &str = "x-api-key"; pub struct AzureAnthropicMessagesConfig { anthropic: AnthropicMessagesConfig, @@ -48,7 +53,7 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { context: &MessagesTransformContext, ) -> Result { let mut request = fold_system_role_messages(request); - if let Some(system) = request.system.as_mut() { + if let Some(system) = request.params.system.as_mut() { strip_scope_from_system(system); } request @@ -68,24 +73,30 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { .transform_anthropic_messages_response(model, response) } - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_azure_api_key(api_key, env_lookup) - } - fn secret_names(&self) -> &'static [&'static str] { &[AZURE_API_KEY_ENV, AZURE_API_BASE_ENV] } - fn auth_strategy(&self) -> MessagesAuthStrategy { - self.anthropic.auth_strategy() - } - - fn accepts_bearer_auth(&self) -> bool { - true + /// A forwarded `x-api-key` or a non-blank bearer (an Entra ID token) is the credential; + /// otherwise the Azure key goes in `x-api-key`. + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + if has_header(&headers, API_KEY_HEADER) || has_bearer_auth(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + let auth = AuthScheme::Credential { + placement: CredentialPlacement::Header(API_KEY_HEADER), + secret: SecretValue::new(resolve_azure_api_key(api_key, env_lookup)?), + }; + Ok(ValidatedEnvironment { headers, auth }) } fn default_headers(&self) -> &'static [(&'static str, &'static str)] { @@ -181,7 +192,7 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess .into_iter() .partition(|msg| msg.role == SYSTEM_ROLE); - let folded_system: Vec = system_into_blocks(request.system) + let folded_system: Vec = system_into_blocks(request.params.system) .into_iter() .chain( system_messages @@ -192,13 +203,17 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess AnthropicMessagesRequest { messages: chat_messages, - system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)), + params: AnthropicMessagesOptionalParams { + system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)), + ..request.params + }, ..request } } #[cfg(test)] mod tests { + use rstest::rstest; use serde_json::json; use super::*; @@ -292,19 +307,47 @@ mod tests { )); } - #[test] - fn auth_strategy_is_x_api_key() { - assert_eq!( - AZURE_ANTHROPIC_MESSAGES_CONFIG - .auth_strategy() - .header_name(), - "x-api-key" - ); + fn validated(forwarded: &[(&str, &str)], api_key: Option<&str>) -> ValidatedEnvironment { + AZURE_ANTHROPIC_MESSAGES_CONFIG + .validate_environment( + forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + api_key, + "claude", + &|_| None, + ) + .unwrap() } #[test] - fn accepts_bearer_auth_for_entra_id() { - assert!(AZURE_ANTHROPIC_MESSAGES_CONFIG.accepts_bearer_auth()); + fn the_azure_key_goes_in_x_api_key() { + assert!(matches!( + validated(&[], Some("sk-azure")).auth, + AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), + ref secret + } if secret.expose() == "sk-azure" + )); + } + + #[rstest] + #[case::x_api_key(&[("X-Api-Key", "caller")])] + #[case::entra_id_bearer(&[("Authorization", "Bearer eyJ-token")])] + fn a_forwarded_key_or_bearer_is_the_credential(#[case] forwarded: &[(&str, &str)]) { + assert!(matches!( + validated(forwarded, Some("sk-azure")).auth, + AuthScheme::Forwarded + )); + } + + #[test] + fn a_blank_bearer_does_not_count_as_a_credential() { + assert!(matches!( + validated(&[("Authorization", "Bearer ")], Some("sk-azure")).auth, + AuthScheme::Credential { .. } + )); } #[test] @@ -528,7 +571,7 @@ mod tests { assert!(err.is_data()); } - #[rstest::rstest] + #[rstest] #[case::compact_context_management_edit( json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), &[], @@ -608,7 +651,12 @@ mod tests { requested.borrow_mut().push(name.to_string()); None }; - let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.authenticate(Vec::new(), None, &record); + let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.validate_environment( + Vec::new(), + None, + "claude", + &record, + ); let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); let requested = requested.into_inner(); assert!(!requested.is_empty()); diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs index 9c2f3f70b91..f28e5c135b0 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs @@ -1,5 +1,3 @@ -use std::sync::OnceLock; - use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; @@ -20,12 +18,11 @@ pub(crate) fn azure_auth_inputs(request: &PreparedOcrRequest) -> Result Option + Sync), ) -> Result>, Error> { - static SERVICE: OnceLock = OnceLock::new(); - SERVICE - .get_or_init(AzureAuthService::default) + service .get_azure_ad_token(config, env_lookup) .await .or_else(|error| match error { diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index b4e9d01f867..5ee5ab3be94 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -183,12 +183,15 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { async fn validate_environment( &self, request: &PreparedOcrRequest, - _client: &OcrClient, + client: &OcrClient, ) -> Result { let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?; - self.resolve_headers(&request.connection, &config, &|name: &str| { - request.connection.secret(name) - }) + self.resolve_headers( + &client.auth().azure, + &request.connection, + &config, + &|name: &str| request.connection.secret(name), + ) .await } @@ -600,6 +603,7 @@ impl AzureDocumentIntelligenceOcrConfig { async fn resolve_headers( &self, + auth: &litellm_auth_azure::AzureAuthService, connection: &OcrConnection, config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), @@ -635,7 +639,7 @@ impl AzureDocumentIntelligenceOcrConfig { .collect(), ); } - let token = super::super::common_utils::resolve_entra(config, env_lookup) + let token = super::super::common_utils::resolve_entra(auth, config, env_lookup) .await? .ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?; super::super::common_utils::validate_destination(connection, token.source())?; @@ -809,9 +813,12 @@ mod tests { }; let error = AzureDocumentIntelligenceOcrConfig - .resolve_headers(&connection, &Default::default(), &|name| { - (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()) - }) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|name| (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()), + ) .await .unwrap_err(); @@ -833,7 +840,12 @@ mod tests { }; let headers = AzureDocumentIntelligenceOcrConfig - .resolve_headers(&connection, &Default::default(), &|_| None) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|_| None, + ) .await .unwrap(); diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 1556ae2a414..8efc27b0bea 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -60,12 +60,15 @@ impl BaseOcrConfig for AzureAiOcrConfig { async fn validate_environment( &self, request: &PreparedOcrRequest, - _client: &OcrClient, + client: &OcrClient, ) -> Result { let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?; - self.resolve_headers(&request.connection, &config, &|name: &str| { - request.connection.secret(name) - }) + self.resolve_headers( + &client.auth().azure, + &request.connection, + &config, + &|name: &str| request.connection.secret(name), + ) .await } @@ -139,6 +142,7 @@ impl AzureAiOcrConfig { async fn resolve_headers( &self, + auth: &litellm_auth_azure::AzureAuthService, connection: &OcrConnection, config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), @@ -146,7 +150,7 @@ impl AzureAiOcrConfig { Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?; if litellm_http::request::has_header(&connection.extra_headers, "authorization") { if config.azure_ad_token_provider.is_some() { - super::common_utils::resolve_entra(config, env_lookup).await?; + super::common_utils::resolve_entra(auth, config, env_lookup).await?; } super::common_utils::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); @@ -166,7 +170,7 @@ impl AzureAiOcrConfig { super::common_utils::validate_destination(connection, key.source())?; return Ok(bearer_headers(connection, key.value())); } - let key = super::common_utils::resolve_entra(config, env_lookup) + let key = super::common_utils::resolve_entra(auth, config, env_lookup) .await? .ok_or(Error::MissingAzureAiCredentials)?; super::common_utils::validate_destination(connection, key.source())?; @@ -253,9 +257,12 @@ mod tests { }; assert_eq!( AzureAiOcrConfig - .resolve_headers(&connection, &Default::default(), &|_| { - Some("environment-key".into()) - }) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|_| { Some("environment-key".into()) } + ) .await .unwrap(), connection.extra_headers @@ -267,9 +274,12 @@ mod tests { async fn request_key_precedes_environment_key(connection: OcrConnection) { assert_eq!( AzureAiOcrConfig - .resolve_headers(&connection, &Default::default(), &|_| { - Some("environment-key".into()) - }) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|_| { Some("environment-key".into()) } + ) .await .unwrap()[0], ("Authorization".into(), "Bearer request-key".into()) @@ -285,9 +295,12 @@ mod tests { }; let error = AzureAiOcrConfig - .resolve_headers(&connection, &Default::default(), &|name| { - (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) - }) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|name| (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()), + ) .await .unwrap_err(); @@ -309,7 +322,12 @@ mod tests { }; let headers = AzureAiOcrConfig - .resolve_headers(&connection, &Default::default(), &|_| None) + .resolve_headers( + &Default::default(), + &connection, + &Default::default(), + &|_| None, + ) .await .unwrap(); @@ -329,7 +347,7 @@ mod tests { let connection = OcrConnection::default(); let headers = AzureAiOcrConfig - .resolve_headers(&connection, &Default::default(), &env) + .resolve_headers(&Default::default(), &connection, &Default::default(), &env) .await .unwrap(); let url = AzureAiOcrConfig.build_ocr_url(None, &env).unwrap(); diff --git a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/mod.rs b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/mod.rs index f239b6921fa..fa7df180f50 100644 --- a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/mod.rs @@ -1 +1,2 @@ +pub mod streaming; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/streaming.rs b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/streaming.rs new file mode 100644 index 00000000000..abb61297669 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/streaming.rs @@ -0,0 +1,141 @@ +use bytes::Bytes; +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_framing::{frames, sse::SseCodec}; + +pub use crate::base_llm::base_model_iterator::ByteStream; +use crate::{Error, anthropic::messages::streaming_iterator::AnthropicMessagesStreamEvent}; + +pub type EventStream = BoxStream<'static, Result>; +pub type StreamDecoder = fn(ByteStream) -> EventStream; + +pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream { + Box::pin(frames(bytes, SseCodec::default()).map(|event| { + let event = event + .map_err(|error| Error::InvalidResponse(format!("stream framing failed: {error}")))?; + serde_json::from_str(&event.data).map_err(|error| { + Error::InvalidResponse(format!("Anthropic stream event is invalid: {error}")) + }) + })) +} + +pub fn encode_anthropic_sse(event: &AnthropicMessagesStreamEvent) -> Result { + let data = serde_json::to_value(event).map_err(|error| { + Error::InvalidResponse(format!("Anthropic stream event is invalid: {error}")) + })?; + let name = data + .get("type") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + Error::InvalidResponse( + "Anthropic stream event is invalid: stream event has no type".into(), + ) + })?; + Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n"))) +} + +#[cfg(test)] +mod tests { + use futures_util::{StreamExt, TryStreamExt, stream}; + use serde_json::json; + + use super::*; + use crate::anthropic::messages::streaming_iterator::{ + AnthropicContentBlockDelta, AnthropicStreamUsage, + }; + + const TEXT_DELTA: &str = + r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; + + fn in_pieces(wire: &[u8]) -> ByteStream { + let pieces: Vec = wire.chunks(3).map(Bytes::copy_from_slice).collect(); + stream::iter(pieces.into_iter().map(Ok)).boxed() + } + + #[tokio::test] + async fn sse_frames_split_anywhere_decode_into_typed_events() { + let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n"); + let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes())) + .try_collect::>() + .await + .unwrap(); + + assert_eq!( + events, + vec![AnthropicMessagesStreamEvent::ContentBlockDelta { + index: 0, + delta: AnthropicContentBlockDelta::TextDelta { + text: "hello".into(), + }, + }] + ); + } + + #[tokio::test] + async fn decodes_citations_delta_events() { + let wire = concat!( + "event: content_block_delta\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#, + "\n\n", + ); + let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes())) + .try_collect::>() + .await + .unwrap(); + + assert!(matches!( + events.as_slice(), + [AnthropicMessagesStreamEvent::ContentBlockDelta { + delta: AnthropicContentBlockDelta::Citations { .. }, + .. + }] + )); + } + + fn events() -> Vec { + vec![ + AnthropicMessagesStreamEvent::Ping, + AnthropicMessagesStreamEvent::ContentBlockDelta { + index: 1, + delta: AnthropicContentBlockDelta::TextDelta { text: "hi".into() }, + }, + AnthropicMessagesStreamEvent::ContentBlockStop { index: 1 }, + AnthropicMessagesStreamEvent::MessageStop { + usage: Some(AnthropicStreamUsage { + output_tokens: Some(7), + ..AnthropicStreamUsage::default() + }), + }, + ] + } + + #[tokio::test] + async fn encoded_events_decode_back_to_themselves() { + let wire = events() + .iter() + .map(encode_anthropic_sse) + .collect::, _>>() + .unwrap(); + + let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed()) + .try_collect::>() + .await + .unwrap(); + + assert_eq!(decoded, events()); + } + + #[test] + fn an_event_is_named_by_its_type() { + let encoded = + encode_anthropic_sse(&AnthropicMessagesStreamEvent::MessageStop { usage: None }) + .unwrap(); + + assert_eq!( + encoded, + Bytes::from(format!( + "event: message_stop\ndata: {}\n\n", + json!({"type": "message_stop"}) + )) + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs index eff1dd1cf0b..9e9f585263c 100644 --- a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs @@ -1,29 +1,13 @@ -use litellm_http::request::{has_bearer_auth, has_header}; use litellm_types::llms::anthropic_messages::{ anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, }; +pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{ - anthropic::messages::thinking::ThinkingContext, base_llm::chat::transformation::Error, + Error, anthropic::messages::thinking::ThinkingContext, + base_llm::anthropic_messages::streaming::StreamDecoder, }; -pub type Headers = Vec<(String, String)>; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum MessagesAuthStrategy { - Bearer, - Header(&'static str), -} - -impl MessagesAuthStrategy { - pub fn header_name(self) -> &'static str { - match self { - Self::Bearer => "authorization", - Self::Header(header_name) => header_name, - } - } -} - #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct MessagesTransformContext { pub thinking: ThinkingContext, @@ -38,6 +22,15 @@ pub trait BaseAnthropicMessagesConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn complete_stream_url( + &self, + api_base: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + self.get_complete_url(api_base, model, env_lookup) + } + fn transform_anthropic_messages_request( &self, request: AnthropicMessagesRequest, @@ -54,42 +47,24 @@ pub trait BaseAnthropicMessagesConfig: Sync { Ok(response) } - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; - fn secret_names(&self) -> &'static [&'static str]; - fn auth_strategy(&self) -> MessagesAuthStrategy { - MessagesAuthStrategy::Header("x-api-key") - } - - fn accepts_bearer_auth(&self) -> bool { - false - } - - fn authenticate( + /// Shapes the forwarded headers and names the credential, the way Python's + /// `validate_environment` does, without applying it: `resolve_auth` does that once + /// for every config. + fn validate_environment( &self, headers: Headers, api_key: Option<&str>, + model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - let strategy = self.auth_strategy(); - if has_header(&headers, strategy.header_name()) - || (self.accepts_bearer_auth() && has_bearer_auth(&headers)) - { - return Ok(headers); - } - let api_key = self.resolve_api_key(api_key, env_lookup)?; - let auth_header = match strategy { - MessagesAuthStrategy::Bearer => { - ("authorization".to_string(), format!("Bearer {api_key}")) - } - MessagesAuthStrategy::Header(name) => (name.to_string(), api_key), - }; - Ok(headers.into_iter().chain([auth_header]).collect()) + ) -> Result; + + /// `None` relays the upstream bytes untouched, which is right for every host that already + /// speaks Anthropic SSE. A host on another wire returns the decoder that lifts its frames + /// into Anthropic stream events, and the route re-encodes those as Anthropic SSE. + fn stream_decoder(&self) -> Option { + None } fn default_headers(&self) -> &'static [(&'static str, &'static str)] { @@ -106,49 +81,8 @@ pub trait BaseAnthropicMessagesConfig: Sync { #[cfg(test)] mod tests { - use rstest::rstest; - use super::*; - - const X_API_KEY: MessagesAuthStrategy = MessagesAuthStrategy::Header("x-api-key"); - - struct StubConfig { - strategy: MessagesAuthStrategy, - accepts_bearer: bool, - } - - impl BaseAnthropicMessagesConfig for StubConfig { - fn secret_names(&self) -> &'static [&'static str] { - &[] - } - - fn get_complete_url( - &self, - _api_base: Option<&str>, - _model: &str, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(String::new()) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - api_key - .map(str::to_string) - .ok_or(Error::MissingField("api_key")) - } - - fn auth_strategy(&self) -> MessagesAuthStrategy { - self.strategy - } - - fn accepts_bearer_auth(&self) -> bool { - self.accepts_bearer - } - } + use crate::base_llm::auth::AuthScheme; struct DefaultsConfig; @@ -166,32 +100,20 @@ mod tests { Ok(String::new()) } - fn resolve_api_key( + fn validate_environment( &self, - api_key: Option<&str>, + headers: Headers, + _api_key: Option<&str>, + _model: &str, _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - api_key - .map(str::to_string) - .ok_or(Error::MissingField("api_key")) + ) -> Result { + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }) } } - #[test] - fn default_config_adds_its_key_next_to_a_forwarded_bearer() { - assert_eq!( - DefaultsConfig.authenticate( - headers(&[("authorization", "Bearer forwarded")]), - Some("sk"), - &|_| None - ), - Ok(headers(&[ - ("authorization", "Bearer forwarded"), - ("x-api-key", "sk") - ])) - ); - } - #[test] fn default_request_headers_are_the_given_headers() { let request: AnthropicMessagesRequest = serde_json::from_value(serde_json::json!({ @@ -213,82 +135,4 @@ mod tests { .map(|(name, value)| (name.to_string(), value.to_string())) .collect() } - - #[rstest] - #[case::own_header_is_kept( - X_API_KEY, - false, - headers(&[("x-api-key", "forwarded")]), - None, - Ok(headers(&[("x-api-key", "forwarded")])) - )] - #[case::own_header_in_any_casing_is_kept( - X_API_KEY, - false, - headers(&[("X-Api-Key", "forwarded")]), - None, - Ok(headers(&[("X-Api-Key", "forwarded")])) - )] - #[case::accepted_bearer_is_kept( - X_API_KEY, - true, - headers(&[("authorization", "Bearer forwarded")]), - None, - Ok(headers(&[("authorization", "Bearer forwarded")])) - )] - #[case::bearer_the_provider_does_not_accept_gets_the_key_too( - X_API_KEY, - false, - headers(&[("authorization", "Bearer forwarded")]), - Some("sk"), - Ok(headers(&[("authorization", "Bearer forwarded"), ("x-api-key", "sk")])) - )] - #[case::blank_bearer_gets_the_key( - X_API_KEY, - true, - headers(&[("authorization", "Bearer ")]), - Some("sk"), - Ok(headers(&[("authorization", "Bearer "), ("x-api-key", "sk")])) - )] - #[case::key_goes_in_the_provider_header( - X_API_KEY, - false, - headers(&[("content-type", "application/json")]), - Some("sk"), - Ok(headers(&[("content-type", "application/json"), ("x-api-key", "sk")])) - )] - #[case::key_goes_in_a_bearer( - MessagesAuthStrategy::Bearer, - false, - headers(&[]), - Some("sk"), - Ok(headers(&[("authorization", "Bearer sk")])) - )] - #[case::bearer_strategy_keeps_a_forwarded_authorization( - MessagesAuthStrategy::Bearer, - false, - headers(&[("authorization", "Bearer forwarded")]), - None, - Ok(headers(&[("authorization", "Bearer forwarded")])) - )] - #[case::missing_key_is_an_error( - X_API_KEY, - false, - headers(&[]), - None, - Err(Error::MissingField("api_key")) - )] - fn default_authenticate_applies_the_key_unless_a_credential_is_forwarded( - #[case] strategy: MessagesAuthStrategy, - #[case] accepts_bearer: bool, - #[case] forwarded: Headers, - #[case] api_key: Option<&str>, - #[case] expected: Result, - ) { - let config = StubConfig { - strategy, - accepts_bearer, - }; - assert_eq!(config.authenticate(forwarded, api_key, &|_| None), expected); - } } diff --git a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs index 1257bbf0d6a..562902ac6a8 100644 --- a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs @@ -1,7 +1,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use crate::base_llm::chat::transformation::Error; +use crate::Error; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct AudioTranscriptionRequestData { @@ -21,7 +21,7 @@ impl AudioTranscriptionResponseData { } } -pub use litellm_auth::RequestAuth; +pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; pub trait BaseAudioTranscriptionConfig: Sync { fn get_supported_openai_params(&self) -> &'static [&'static str]; @@ -58,10 +58,11 @@ pub trait BaseAudioTranscriptionConfig: Sync { response_json: Value, ) -> Result; - fn auth_strategy( + fn validate_environment( &self, + headers: Headers, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; } diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs new file mode 100644 index 00000000000..897b0f62e82 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -0,0 +1,263 @@ +//! How a provider call authenticates, decided by the provider config when the request is +//! prepared and applied once here when it is sent. +//! +//! Python folds this into `validate_environment` plus `sign_request`. The Rust configs keep +//! that split: `validate_environment` shapes the forwarded headers and names the credential +//! as an [`AuthScheme`], and [`resolve_auth`] turns the scheme into headers and a signer. + +use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; +use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; + +pub type Headers = Vec<(String, String)>; + +#[derive(Clone, Debug)] +pub enum AuthScheme { + /// The caller's own credential is already in the headers and is sent as is. + Forwarded, + /// A credential in hand, placed in its header. A forwarded header of the same name is + /// replaced: the deployment's identity outranks the caller's. + Credential { + placement: CredentialPlacement, + secret: SecretValue, + }, + /// A bearer acquired when the request is sent, from a token source such as a cloud SDK + /// or a caller-supplied callable. + Token { provider: TokenProviderHandle }, + /// AWS SigV4 over the bytes that go on the wire, so the handler signs after the body is + /// serialized. + AwsSigV4 { + region: String, + service: &'static str, + credentials: Box, + }, +} + +/// The outcome of a config's `validate_environment`: the headers it shaped and how the +/// call authenticates. +#[derive(Clone, Debug)] +pub struct ValidatedEnvironment { + pub headers: Headers, + pub auth: AuthScheme, +} + +#[derive(Debug)] +pub struct Authenticated { + pub headers: Headers, + pub signer: Option, +} + +pub async fn resolve_auth( + services: &AuthServices, + validated: ValidatedEnvironment, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result { + let ValidatedEnvironment { headers, auth } = validated; + match auth { + AuthScheme::Forwarded => Ok(Authenticated { + headers, + signer: None, + }), + AuthScheme::Credential { placement, secret } => Ok(Authenticated { + headers: with_credential(headers, placement, secret.expose()), + signer: None, + }), + AuthScheme::Token { provider } => { + let token = provider.acquire().await?; + Ok(Authenticated { + headers: with_credential( + headers, + CredentialPlacement::Bearer, + token.secret().expose(), + ), + signer: None, + }) + } + AuthScheme::AwsSigV4 { + region, + service, + credentials, + } => Ok(Authenticated { + headers, + signer: Some( + SigV4Signer::resolve(&services.aws, region, service, *credentials, env_lookup) + .await?, + ), + }), + } +} + +/// Fills in the defaults the caller did not forward, matching Python's +/// `if name not in headers` checks. +pub fn with_default_headers(headers: Headers, defaults: &[(&str, &str)]) -> Headers { + let missing: Vec<(String, String)> = defaults + .iter() + .filter(|(name, _)| { + !headers + .iter() + .any(|(header, _)| header.eq_ignore_ascii_case(name)) + }) + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + headers.into_iter().chain(missing).collect() +} + +fn with_credential(headers: Headers, placement: CredentialPlacement, credential: &str) -> Headers { + let name = placement.header_name(); + let value = match placement { + CredentialPlacement::Bearer => format!("Bearer {credential}"), + CredentialPlacement::Header(_) => credential.to_string(), + }; + headers + .into_iter() + .filter(|(header, _)| !header.eq_ignore_ascii_case(name)) + .chain([(name.to_ascii_lowercase(), value)]) + .collect() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use litellm_auth::{AuthServices, ResolvedCredential, TokenFuture, TokenProvider}; + use litellm_auth_aws::Credentials; + use rstest::rstest; + + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + async fn resolve(headers: Headers, auth: AuthScheme) -> Authenticated { + resolve_auth( + &AuthServices::default(), + ValidatedEnvironment { headers, auth }, + &no_env, + ) + .await + .unwrap() + } + + #[rstest] + #[case::header_is_appended( + &[("content-type", "application/json")], + CredentialPlacement::Header("x-api-key"), + &[("content-type", "application/json"), ("x-api-key", "sk")], + )] + #[case::forwarded_header_of_the_same_name_is_replaced_in_any_casing( + &[("X-Api-Key", "caller"), ("x-trace", "1")], + CredentialPlacement::Header("x-api-key"), + &[("x-trace", "1"), ("x-api-key", "sk")], + )] + #[case::bearer_replaces_a_forwarded_authorization( + &[("Authorization", "Bearer caller")], + CredentialPlacement::Bearer, + &[("authorization", "Bearer sk")], + )] + #[tokio::test] + async fn a_credential_lands_in_its_header_and_outranks_the_forwarded_one( + #[case] forwarded: &[(&str, &str)], + #[case] placement: CredentialPlacement, + #[case] expected: &[(&str, &str)], + ) { + let authenticated = resolve( + headers(forwarded), + AuthScheme::Credential { + placement, + secret: SecretValue::new("sk"), + }, + ) + .await; + assert_eq!(authenticated.headers, headers(expected)); + assert!(authenticated.signer.is_none()); + } + + #[rstest] + #[case::nothing_forwarded( + &[], + &[("x-version", "1"), ("content-type", "application/json")], + &[("x-version", "1"), ("content-type", "application/json")], + )] + #[case::forwarded_header_wins_in_any_case( + &[("X-Version", "custom"), ("x-api-key", "k")], + &[("x-version", "1"), ("content-type", "application/json")], + &[("X-Version", "custom"), ("x-api-key", "k"), ("content-type", "application/json")], + )] + #[case::no_defaults(&[("x-api-key", "k")], &[], &[("x-api-key", "k")])] + fn default_headers_fill_only_missing_names( + #[case] forwarded: &[(&str, &str)], + #[case] defaults: &[(&str, &str)], + #[case] expected: &[(&str, &str)], + ) { + assert_eq!( + with_default_headers(headers(forwarded), defaults), + headers(expected) + ); + } + + #[tokio::test] + async fn forwarded_auth_sends_the_headers_untouched() { + let forwarded = headers(&[("x-api-key", "caller"), ("authorization", "Bearer caller")]); + let authenticated = resolve(forwarded.clone(), AuthScheme::Forwarded).await; + assert_eq!(authenticated.headers, forwarded); + assert!(authenticated.signer.is_none()); + } + + #[derive(Debug)] + struct StaticToken(&'static str); + + impl TokenProvider for StaticToken { + fn acquire(&self) -> TokenFuture<'_> { + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(self.0), + expires_on: None, + }) + }) + } + } + + #[tokio::test] + async fn a_token_is_acquired_at_send_time_and_sent_as_a_bearer() { + let authenticated = resolve( + headers(&[("authorization", "Bearer stale")]), + AuthScheme::Token { + provider: TokenProviderHandle::new(Arc::new(StaticToken("fresh"))), + }, + ) + .await; + assert_eq!( + authenticated.headers, + headers(&[("authorization", "Bearer fresh")]) + ); + } + + #[tokio::test] + async fn sigv4_leaves_the_headers_to_the_signer() { + let forwarded = headers(&[("x-request-id", "abc")]); + let authenticated = resolve( + forwarded.clone(), + AuthScheme::AwsSigV4 { + region: "us-east-1".into(), + service: "bedrock", + credentials: Box::new(AwsCredentialSource::HostSupplied(Credentials::new( + "AKIDEXAMPLE", + "secret", + None, + None, + "test", + ))), + }, + ) + .await; + assert_eq!(authenticated.headers, forwarded); + assert!(authenticated.signer.is_some()); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs index 928ef80b29a..a283ca84089 100644 --- a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs +++ b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs @@ -1,3 +1,10 @@ +use std::{collections::VecDeque, convert::Infallible, io, pin::Pin}; + +use bytes::Bytes; +use futures_util::{Stream, StreamExt, stream, stream::BoxStream}; + +pub type ByteStream = BoxStream<'static, Result>; + pub trait StreamTransformer { type Input; type Output; @@ -7,3 +14,149 @@ pub trait StreamTransformer { fn finish(&mut self) -> Result, Self::Error>; } + +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum StreamError { + #[error(transparent)] + Decode(D), + #[error(transparent)] + Transform(T), +} + +impl StreamError { + pub fn into_decode(self) -> D { + match self { + Self::Decode(error) => error, + Self::Transform(never) => match never {}, + } + } +} + +struct Driver { + events: Pin>, + transformer: T, + ready: VecDeque, + finished: bool, +} + +/// Drives `transformer` over `events`, then flushes it with `finish`. The first error ends the +/// stream. +pub fn transform_stream( + events: S, + transformer: T, +) -> impl Stream>> + Send +where + S: Stream> + Send, + T: StreamTransformer + Send, + T::Output: Send, + T::Error: Send, + D: Send, +{ + let driver = Driver { + events: Box::pin(events), + transformer, + ready: VecDeque::new(), + finished: false, + }; + stream::unfold(driver, |mut driver| async move { + loop { + if let Some(output) = driver.ready.pop_front() { + return Some((Ok(output), driver)); + } + if driver.finished { + return None; + } + match driver.events.next().await { + Some(Ok(event)) => match driver.transformer.transform(event) { + Ok(outputs) => driver.ready.extend(outputs), + Err(error) => { + driver.finished = true; + return Some((Err(StreamError::Transform(error)), driver)); + } + }, + Some(Err(error)) => { + driver.finished = true; + return Some((Err(StreamError::Decode(error)), driver)); + } + None => { + driver.finished = true; + match driver.transformer.finish() { + Ok(outputs) => driver.ready.extend(outputs), + Err(error) => return Some((Err(StreamError::Transform(error)), driver)), + } + } + } + } + }) +} + +#[cfg(test)] +mod tests { + use futures_util::TryStreamExt; + + use super::*; + + struct Doubler; + + impl StreamTransformer for Doubler { + type Input = u32; + type Output = u32; + type Error = String; + + fn transform(&mut self, input: u32) -> Result, String> { + match input { + 0 => Err("zero".into()), + n => Ok(vec![n, n * 2]), + } + } + + fn finish(&mut self) -> Result, String> { + Ok(vec![u32::MAX]) + } + } + + #[tokio::test] + async fn flat_maps_each_event_and_flushes_at_the_end() { + let output = transform_stream(stream::iter([Ok::<_, String>(1), Ok(2)]), Doubler) + .try_collect::>() + .await + .unwrap(); + + assert_eq!(output, vec![1, 2, 2, 4, u32::MAX]); + } + + #[tokio::test] + async fn a_transform_error_ends_the_stream_without_flushing() { + let output = transform_stream(stream::iter([Ok::<_, String>(1), Ok(0), Ok(3)]), Doubler) + .collect::>() + .await; + + assert_eq!( + output, + vec![ + Ok(1), + Ok(2), + Err(StreamError::Transform("zero".to_string())) + ] + ); + } + + #[tokio::test] + async fn a_decode_error_ends_the_stream_without_flushing() { + let output = transform_stream( + stream::iter([Ok(1), Err("bad frame".to_string()), Ok(3)]), + Doubler, + ) + .collect::>() + .await; + + assert_eq!( + output, + vec![ + Ok(1), + Ok(2), + Err(StreamError::Decode("bad frame".to_string())) + ] + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/chat/mod.rs b/litellm-rust/crates/llms/src/base_llm/chat/mod.rs index f239b6921fa..fa7df180f50 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/mod.rs @@ -1 +1,2 @@ +pub mod streaming; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs b/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs new file mode 100644 index 00000000000..b9d715bcd68 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs @@ -0,0 +1,54 @@ +use std::collections::HashMap; + +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_types::utils::ChatCompletionChunk; + +use crate::{ + Error, + base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream}, +}; + +pub type ChatChunkStream = BoxStream<'static, Result>; + +/// What Python's `map_openai_params` decides about the stream and `completion` +/// hands to `ModelResponseIterator`: it is settled while the request is built, +/// never re-derived from the body. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct StreamShape { + pub json_mode: bool, + pub speed: Option, + pub tool_name_reverse_map: HashMap, +} + +/// A wire decoder paired with the iterator that turns its events into chat chunks. +/// A config names both; the core runs the pair over the response bytes. +pub struct ChatStream { + run: Box ChatChunkStream + Send>, +} + +impl ChatStream { + pub fn new( + decode: fn(ByteStream) -> BoxStream<'static, Result>, + iterator: T, + ) -> Self + where + E: Send + 'static, + T: StreamTransformer + + Send + + 'static, + { + Self { + run: Box::new(move |bytes| { + Box::pin(transform_stream(decode(bytes), iterator).map(|item| { + item.map_err(|error| match error { + StreamError::Decode(error) | StreamError::Transform(error) => error, + }) + })) + }), + } + } + + pub fn run(self, bytes: ByteStream) -> ChatChunkStream { + (self.run)(bytes) + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index c7d1a27c71e..8a074e59207 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -4,30 +4,17 @@ use litellm_types::{ }; use serde_json::{Map, Value}; -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, - }, - #[error("missing required field: {0}")] - MissingField(&'static str), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("unsupported: {0}")] - Unsupported(&'static str), - #[error(transparent)] - Auth(#[from] litellm_auth::Error), -} +use crate::{ + Error, + base_llm::chat::streaming::{ChatStream, StreamShape}, +}; /// The provider-shaped request body a config produces. Named rather than a bare /// `Value` so the transform contract stays a typed one, mirroring /// [`crate::base_llm::audio_transcription::transformation::AudioTranscriptionRequestData`]. pub struct ProviderChatRequestData { pub body: Value, + pub stream_shape: StreamShape, } /// The raw provider response body handed back to a config for normalization. @@ -41,7 +28,7 @@ pub const STREAM_PARAM: &str = "stream"; /// presence does not make a request untranslatable. const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; -pub use litellm_auth::RequestAuth; +pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; /// Why a request cannot be served by the Rust path. /// @@ -78,28 +65,27 @@ pub trait BaseConfig: Sync { response: ProviderChatResponseData, ) -> Result; - fn auth( + /// `None` means this config has no streaming path yet, so the host keeps the request. + fn model_response_iterator(&self, _shape: StreamShape) -> Option { + None + } + + /// Shapes the forwarded headers and names the credential, the way Python's + /// `validate_environment` does, without applying it: `resolve_auth` does that once + /// for every config. + fn validate_environment( &self, + headers: Headers, api_key: Option<&str>, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[("content-type", "application/json")] } - /// Whether an auth header the caller already supplied is the credential this - /// request should authenticate with, so the resolved one is not applied. - /// - /// Defaults to false: the deployment's credential outranks anything - /// forwarded, which is what every provider wants for its own auth header. - /// A provider overrides this only for a scheme it hands off to entirely. - fn defers_to_forwarded_auth(&self, _headers: &[(String, String)]) -> bool { - false - } - /// Parameters consumed as call configuration (credentials, endpoints) /// rather than placed in the body. Accepted, never serialized. fn config_params(&self) -> &'static [&'static str] { diff --git a/litellm-rust/crates/llms/src/base_llm/mod.rs b/litellm-rust/crates/llms/src/base_llm/mod.rs index 8ed37da4573..399b932e9da 100644 --- a/litellm-rust/crates/llms/src/base_llm/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/mod.rs @@ -1,5 +1,6 @@ pub mod anthropic_messages; pub mod audio_transcription; +pub mod auth; pub mod base_model_iterator; pub mod chat; pub mod ocr; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index fe072228234..0148ca2841b 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; -use litellm_auth_gcp::VertexAuth; +use litellm_auth::AuthServices; use litellm_host::event::WireRequest; use litellm_http::{ Client, ClientVariant, HttpClientConfig, HttpClientPool, @@ -36,7 +36,7 @@ pub struct OcrClient { provider_http: Client, polling_http: Client, document_fetcher: MediaFetcher, - vertex_auth: VertexAuth, + auth: Arc, settings: OcrSettings, secrets: Arc, } @@ -46,7 +46,7 @@ impl OcrClient { pool: &HttpClientPool, config: &HttpClientConfig, url_policy: UrlPolicy, - vertex_auth: VertexAuth, + auth: Arc, settings: OcrSettings, secrets: Arc, ) -> Result { @@ -54,7 +54,7 @@ impl OcrClient { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, document_fetcher: MediaFetcher::new(pool, config, url_policy)?, - vertex_auth, + auth, settings, secrets, }) @@ -72,8 +72,8 @@ impl OcrClient { &self.document_fetcher } - pub fn vertex_auth(&self) -> &VertexAuth { - &self.vertex_auth + pub fn auth(&self) -> &AuthServices { + &self.auth } pub fn settings(&self) -> &OcrSettings { @@ -95,7 +95,7 @@ impl OcrClient { provider_http, polling_http: no_redirect_http.clone(), document_fetcher: MediaFetcher::for_test(no_redirect_http), - vertex_auth: VertexAuth::default(), + auth: Arc::new(AuthServices::default()), settings: OcrSettings::default(), } } diff --git a/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs b/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs index 0d9cfcfd4cd..419430250c4 100644 --- a/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs @@ -1,6 +1,6 @@ use litellm_types::responses::streaming_websocket::{ResponsesWsEvent, ResponsesWsTransformResult}; -use crate::base_llm::chat::transformation::Error; +use crate::Error; pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1"; pub const OPENAI_RESPONSES_PATH: &str = "/responses"; diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index cfabcb12341..3525fd6322b 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -1,17 +1,20 @@ use litellm_auth_aws::{ - bedrock_model_id_and_region, + AwsCredentialSource, bedrock_model_id_and_region, constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; use litellm_core_utils::core_helpers::json_type_name; use serde_json::{Map, Value, json}; -use crate::base_llm::{ - audio_transcription::transformation::{ - AudioTranscriptionRequestData, AudioTranscriptionResponseData, - BaseAudioTranscriptionConfig, RequestAuth, +use crate::{ + Error, + base_llm::{ + audio_transcription::transformation::{ + AudioTranscriptionRequestData, AudioTranscriptionResponseData, + BaseAudioTranscriptionConfig, Headers, ValidatedEnvironment, + }, + auth::AuthScheme, }, - chat::transformation::Error, }; const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; @@ -131,16 +134,28 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { )) } - fn auth_strategy( + fn validate_environment( &self, + headers: Headers, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { + ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); - Ok(RequestAuth::AwsSigV4 { - region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), - service: BEDROCK_SERVICE, + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::AwsSigV4 { + region: resolve_bedrock_region( + model_region.as_deref(), + optional_params, + env_lookup, + ), + service: BEDROCK_SERVICE, + credentials: Box::new(AwsCredentialSource::from_params( + optional_params, + env_lookup, + )), + }, }) } } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index b5db88d7dc4..8254165a739 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,5 +1,6 @@ +use litellm_auth::{CredentialPlacement, SecretValue}; use litellm_auth_aws::{ - bedrock_model_id_and_region, + AwsCredentialSource, bedrock_model_id_and_region, constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; @@ -16,9 +17,18 @@ use litellm_types::{ }; use serde_json::{Map, Value, json}; -use crate::base_llm::chat::transformation::{ - BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported, - unsupported_message, unsupported_param, +use crate::{ + Error, + base_llm::{ + auth::AuthScheme, + chat::{ + streaming::StreamShape, + transformation::{ + BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, + Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + }, + }, + }, }; /// Converse parameter names, post `map_openai_params`, that the Rust path can @@ -99,6 +109,7 @@ impl BaseConfig for AmazonConverseConfig { ) -> Result { Ok(ProviderChatRequestData { body: converse_body(&build_conversation(&messages), &optional_params), + stream_shape: StreamShape::default(), }) } @@ -180,31 +191,48 @@ impl BaseConfig for AmazonConverseConfig { }) } - fn auth( + /// Python reads `api_key` as the Bedrock bearer token and consults the env only when + /// the caller passed none, so a caller-supplied empty key falls through to SigV4 + /// without reaching for the environment. An all-whitespace token stays a bearer token + /// here because Python sends it too: treating it as absent would sign as the host + /// principal instead, which is the identity swap this branch exists to prevent. + fn validate_environment( &self, + headers: Headers, api_key: Option<&str>, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - // Python reads `api_key` as the Bedrock bearer token and consults the - // env only when the caller passed none, so a caller-supplied empty key - // falls through to SigV4 without reaching for the environment. An - // all-whitespace token stays a bearer token here because Python sends - // it too: treating it as absent would sign as the host principal - // instead, which is the identity swap this branch exists to prevent. + ) -> Result { let bearer = match api_key { Some(key) => Some(key.to_string()), None => env_lookup(AWS_BEARER_TOKEN_BEDROCK), } .filter(|token| !token.is_empty()); if let Some(token) = bearer { - return Ok(RequestAuth::Bearer { token }); + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), + }, + }); } let (_, model_region) = bedrock_model_id_and_region(model); - Ok(RequestAuth::AwsSigV4 { - region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), - service: BEDROCK_SERVICE, + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::AwsSigV4 { + region: resolve_bedrock_region( + model_region.as_deref(), + optional_params, + env_lookup, + ), + service: BEDROCK_SERVICE, + credentials: Box::new(AwsCredentialSource::from_params( + optional_params, + env_lookup, + )), + }, }) } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs new file mode 100644 index 00000000000..b424dbd358d --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs @@ -0,0 +1,138 @@ +use base64::Engine; +use bytes::Buf; +use futures_util::{Stream, StreamExt}; +use litellm_framing::{ + aws_event_stream::{AwsEventStreamCodec, Message}, + frames, +}; +use serde::Deserialize; +use serde_json::Value; + +use crate::{ + Error, + anthropic::{ + chat::handler::ModelResponseIterator, + messages::streaming_iterator::AnthropicMessagesStreamEvent, + }, + base_llm::{ + anthropic_messages::streaming::{ByteStream, EventStream}, + chat::streaming::{ChatStream, StreamShape}, + }, +}; + +#[derive(Deserialize)] +struct InvokeChunkPayload { + bytes: String, +} + +pub fn decode_invoke_chunk(message: Message) -> Result { + let payload: InvokeChunkPayload = + serde_json::from_slice(message.payload()).map_err(|error| { + Error::InvalidResponse(format!("Bedrock event payload is invalid: {error}")) + })?; + let chunk = base64::engine::general_purpose::STANDARD + .decode(payload.bytes) + .map_err(|error| { + Error::InvalidResponse(format!("Bedrock event payload has invalid base64: {error}")) + })?; + serde_json::from_slice(&chunk).map_err(|error| { + Error::InvalidResponse(format!("Anthropic stream event is invalid: {error}")) + }) +} + +pub fn invoke_chunk_stream(input: S) -> impl Stream> + Send +where + S: Stream> + Send, + B: Buf + Send, + E: std::error::Error + Send + Sync + 'static, +{ + frames(input, AwsEventStreamCodec).map(|message| { + decode_invoke_chunk( + message.map_err(|error| { + Error::InvalidResponse(format!("stream framing failed: {error}")) + })?, + ) + }) +} + +pub fn decode_invoke_anthropic_chunk(chunk: Value) -> Result { + serde_json::from_value(chunk).map_err(|error| { + Error::InvalidResponse(format!("Anthropic stream event is invalid: {error}")) + }) +} + +pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream { + Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?))) +} + +pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result { + match invoke_provider { + "anthropic" => Ok(ChatStream::new( + invoke_anthropic_event_stream, + ModelResponseIterator::new(shape), + )), + "deepseek_r1" | "moonshot" => Err(Error::Unsupported( + "Bedrock invoke streaming for this model family", + )), + _ => Err(Error::Unsupported("Bedrock invoke streaming")), + } +} + +#[cfg(test)] +mod tests { + use aws_smithy_eventstream::frame::write_message_to; + use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; + use base64::engine::general_purpose::STANDARD; + use bytes::Bytes; + use futures_util::TryStreamExt; + + use super::*; + use crate::{ + anthropic::messages::streaming_iterator::AnthropicContentBlockDelta, + base_llm::anthropic_messages::streaming::anthropic_sse_event_stream, + }; + + fn in_pieces(wire: &[u8]) -> ByteStream { + let pieces: Vec = wire.chunks(3).map(Bytes::copy_from_slice).collect(); + futures_util::stream::iter(pieces.into_iter().map(Ok)).boxed() + } + + const TEXT_DELTA: &str = + r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; + + fn aws_wire(chunk: &str) -> Vec { + let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)}); + let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( + Header::new(":event-type", HeaderValue::String("chunk".into())), + ); + let mut wire = Vec::new(); + write_message_to(&message, &mut wire).unwrap(); + wire + } + + #[tokio::test] + async fn aws_and_sse_framing_decode_to_the_same_anthropic_events() { + let aws = aws_wire(TEXT_DELTA); + let sse = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n"); + + let from_aws = invoke_anthropic_event_stream(in_pieces(&aws)) + .try_collect::>() + .await + .unwrap(); + let from_sse = anthropic_sse_event_stream(in_pieces(sse.as_bytes())) + .try_collect::>() + .await + .unwrap(); + + assert_eq!( + from_aws, + vec![AnthropicMessagesStreamEvent::ContentBlockDelta { + index: 0, + delta: AnthropicContentBlockDelta::TextDelta { + text: "hello".into(), + }, + }] + ); + assert_eq!(from_aws, from_sse); + } +} diff --git a/litellm-rust/crates/llms/src/bedrock/chat/mod.rs b/litellm-rust/crates/llms/src/bedrock/chat/mod.rs index a41ad86ef49..a46514aa697 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/mod.rs @@ -1 +1,2 @@ pub mod converse_transformation; +pub mod invoke_handler; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs new file mode 100644 index 00000000000..f2e365e9ed0 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs @@ -0,0 +1,554 @@ +use std::convert::Infallible; + +use futures_util::StreamExt; +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth_aws::{ + AwsCredentialSource, bedrock_model_id_and_region, + constants::{ + AWS_BEARER_TOKEN_BEDROCK, AWS_BEDROCK_RUNTIME_ENDPOINT, AWS_DEFAULT_REGION, AWS_REGION, + AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE, + }, + resolve_bedrock_region, +}; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use serde_json::{Map, Value}; + +use crate::{ + Error, + anthropic::messages::streaming_iterator::{AnthropicMessagesStreamEvent, AnthropicStreamUsage}, + base_llm::{ + anthropic_messages::{ + streaming::{ByteStream, EventStream, StreamDecoder}, + transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, + ValidatedEnvironment, + }, + }, + auth::AuthScheme, + base_model_iterator::{StreamError, StreamTransformer, transform_stream}, + }, + bedrock::chat::invoke_handler::{decode_invoke_anthropic_chunk, invoke_chunk_stream}, +}; + +const INVOCATION_METRICS_KEY: &str = "amazon-bedrock-invocationMetrics"; + +const METRICS_USAGE_KEYS: [(&str, &str); 4] = [ + ("input_tokens", "inputTokenCount"), + ("output_tokens", "outputTokenCount"), + ("cache_read_input_tokens", "cacheReadInputTokenCount"), + ("cache_creation_input_tokens", "cacheWriteInputTokenCount"), +]; + +const INVOKE_PATH: &str = "invoke"; +const INVOKE_STREAM_PATH: &str = "invoke-with-response-stream"; +const INVOKE_MODEL_PREFIX: &str = "invoke/"; + +const SECRET_NAMES: &[&str] = &[ + AWS_BEARER_TOKEN_BEDROCK, + AWS_BEDROCK_RUNTIME_ENDPOINT, + AWS_REGION_NAME, + AWS_REGION, + AWS_DEFAULT_REGION, +]; + +pub struct AmazonAnthropicClaudeMessagesConfig; + +pub const BEDROCK_ANTHROPIC_MESSAGES_CONFIG: AmazonAnthropicClaudeMessagesConfig = + AmazonAnthropicClaudeMessagesConfig; + +fn bearer_token( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + match api_key { + Some(key) => Some(key.to_string()), + None => env_lookup(AWS_BEARER_TOKEN_BEDROCK), + } + .filter(|token| !token.is_empty()) +} + +fn invoke_url( + api_base: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + path: &str, +) -> String { + let (model_id, model_region) = + bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); + let region = resolve_bedrock_region(model_region.as_deref(), &Map::new(), env_lookup); + let endpoint = api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(AWS_BEDROCK_RUNTIME_ENDPOINT)) + .unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion)); + format!("{}/model/{model_id}/{path}", endpoint.trim_end_matches('/')) +} + +impl BaseAnthropicMessagesConfig for AmazonAnthropicClaudeMessagesConfig { + fn get_complete_url( + &self, + api_base: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(invoke_url(api_base, model, env_lookup, INVOKE_PATH)) + } + + fn complete_stream_url( + &self, + api_base: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(invoke_url(api_base, model, env_lookup, INVOKE_STREAM_PATH)) + } + + fn transform_anthropic_messages_request( + &self, + _request: AnthropicMessagesRequest, + _context: &MessagesTransformContext, + ) -> Result { + Err(Error::Unsupported( + "Bedrock invoke messages request shaping", + )) + } + + fn secret_names(&self) -> &'static [&'static str] { + SECRET_NAMES + } + + /// Python reads `api_key` as the Bedrock bearer token and consults the env only when the + /// caller passed none. Without one the request is signed with SigV4. + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + if let Some(token) = bearer_token(api_key, env_lookup) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), + }, + }); + } + let (_, model_region) = + bedrock_model_id_and_region(model.strip_prefix(INVOKE_MODEL_PREFIX).unwrap_or(model)); + let params = Map::new(); + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::AwsSigV4 { + region: resolve_bedrock_region(model_region.as_deref(), ¶ms, env_lookup), + service: BEDROCK_SERVICE, + credentials: Box::new(AwsCredentialSource::from_params(¶ms, env_lookup)), + }, + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[("content-type", "application/json")] + } + + fn stream_decoder(&self) -> Option { + Some(bedrock_anthropic_messages_event_stream) + } +} + +fn with_invocation_usage(chunk: Value) -> Value { + match chunk { + Value::Object(fields) => Value::Object(with_metrics_usage(fields)), + other => other, + } +} + +fn with_metrics_usage(mut fields: Map) -> Map { + let Some(Value::Object(metrics)) = fields.remove(INVOCATION_METRICS_KEY) else { + return fields; + }; + if metrics.is_empty() { + return fields; + } + let preserved = match fields.remove("usage") { + Some(Value::Object(usage)) => usage, + _ => Map::new(), + }; + let usage: Map = METRICS_USAGE_KEYS + .iter() + .filter_map(|(anthropic, metric)| { + Some((anthropic.to_string(), metrics.get(*metric)?.clone())) + }) + .chain(preserved) + .collect(); + fields.insert("usage".to_string(), Value::Object(usage)); + fields +} + +pub fn bedrock_anthropic_messages_event_stream(bytes: ByteStream) -> EventStream { + let events = invoke_chunk_stream(bytes) + .map(|chunk| decode_invoke_anthropic_chunk(with_invocation_usage(chunk?))); + Box::pin( + transform_stream(events, MessageStopUsagePromoter::default()) + .map(|item| item.map_err(StreamError::into_decode)), + ) +} + +#[derive(Default)] +pub struct MessageStopUsagePromoter { + pending_delta: Option, + start_usage: Option, +} + +fn promoted_usage( + delta: Option, + stop: Option<&AnthropicStreamUsage>, + start: Option<&AnthropicStreamUsage>, +) -> Option { + let delta = delta.unwrap_or_default(); + let merged = AnthropicStreamUsage { + input_tokens: stop + .and_then(|stop| stop.input_tokens) + .or(delta.input_tokens), + cache_creation_input_tokens: stop + .and_then(|stop| stop.cache_creation_input_tokens) + .or(delta.cache_creation_input_tokens) + .or_else(|| start.and_then(|start| start.cache_creation_input_tokens)), + cache_read_input_tokens: stop + .and_then(|stop| stop.cache_read_input_tokens) + .or(delta.cache_read_input_tokens) + .or_else(|| start.and_then(|start| start.cache_read_input_tokens)), + extra: delta + .extra + .into_iter() + .chain( + start + .and_then(|start| start.extra.get_key_value("cache_creation")) + .map(|(key, value)| (key.clone(), value.clone())), + ) + .fold(Map::new(), |mut extra, (key, value)| { + extra.entry(key).or_insert(value); + extra + }), + ..delta + }; + (merged != AnthropicStreamUsage::default()).then_some(merged) +} + +fn promoted( + event: AnthropicMessagesStreamEvent, + stop: Option<&AnthropicStreamUsage>, + start: Option<&AnthropicStreamUsage>, +) -> AnthropicMessagesStreamEvent { + match event { + AnthropicMessagesStreamEvent::MessageDelta { + delta, + usage, + context_management, + } => AnthropicMessagesStreamEvent::MessageDelta { + delta, + usage: promoted_usage(usage, stop, start), + context_management, + }, + other => other, + } +} + +impl StreamTransformer for MessageStopUsagePromoter { + type Input = AnthropicMessagesStreamEvent; + type Output = AnthropicMessagesStreamEvent; + type Error = Infallible; + + fn transform( + &mut self, + input: AnthropicMessagesStreamEvent, + ) -> Result, Infallible> { + let pending = self.pending_delta.take(); + match input { + AnthropicMessagesStreamEvent::MessageDelta { .. } => { + self.pending_delta = Some(input); + Ok(pending.into_iter().collect()) + } + AnthropicMessagesStreamEvent::MessageStop { usage } => Ok(pending + .map(|delta| promoted(delta, usage.as_ref(), self.start_usage.as_ref())) + .into_iter() + .chain([AnthropicMessagesStreamEvent::MessageStop { usage }]) + .collect()), + AnthropicMessagesStreamEvent::MessageStart { message } => { + self.start_usage = Some(message.usage.clone()); + Ok(pending + .into_iter() + .chain([AnthropicMessagesStreamEvent::MessageStart { message }]) + .collect()) + } + other => Ok(pending.into_iter().chain([other]).collect()), + } + } + + fn finish(&mut self) -> Result, Infallible> { + Ok(self + .pending_delta + .take() + .map(|delta| promoted(delta, None, self.start_usage.as_ref())) + .into_iter() + .collect()) + } +} + +#[cfg(test)] +mod tests { + use aws_smithy_eventstream::frame::write_message_to; + use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; + use base64::{Engine, engine::general_purpose::STANDARD}; + use bytes::Bytes; + use futures_util::TryStreamExt; + use rstest::rstest; + use serde_json::json; + + use litellm_auth_aws::constants::DEFAULT_BEDROCK_REGION; + + use super::*; + use crate::base_llm::anthropic_messages::streaming::encode_anthropic_sse; + + fn event(value: Value) -> AnthropicMessagesStreamEvent { + serde_json::from_value(value).unwrap() + } + + fn message_start(usage: Value) -> AnthropicMessagesStreamEvent { + event(json!({ + "type": "message_start", + "message": { + "id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": null, "stop_sequence": null, "usage": usage + } + })) + } + + fn message_delta(usage: Value) -> AnthropicMessagesStreamEvent { + event(json!({ + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": usage + })) + } + + fn message_stop(usage: Option) -> AnthropicMessagesStreamEvent { + match usage { + Some(usage) => event(json!({"type": "message_stop", "usage": usage})), + None => event(json!({"type": "message_stop"})), + } + } + + fn promote(events: Vec) -> Vec { + let mut promoter = MessageStopUsagePromoter::default(); + let mut output: Vec<_> = events + .into_iter() + .flat_map(|event| promoter.transform(event).unwrap()) + .collect(); + output.extend(promoter.finish().unwrap()); + output + } + + #[rstest] + #[case::cache_fields_on_message_stop( + json!({"input_tokens": 10, "output_tokens": 0}), + json!({"output_tokens": 5}), + Some(json!({"input_tokens": 3, "cache_read_input_tokens": 100, "cache_creation_input_tokens": 20})), + json!({"input_tokens": 3, "output_tokens": 5, "cache_read_input_tokens": 100, "cache_creation_input_tokens": 20}), + )] + #[case::cache_only_on_message_start( + json!({"input_tokens": 10, "output_tokens": 0, "cache_read_input_tokens": 80, "cache_creation_input_tokens": 4, "cache_creation": {"ephemeral_5m_input_tokens": 4}}), + json!({"output_tokens": 5}), + Some(json!({"input_tokens": 10})), + json!({"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 80, "cache_creation_input_tokens": 4, "cache_creation": {"ephemeral_5m_input_tokens": 4}}), + )] + #[case::message_stop_wins_over_message_start( + json!({"input_tokens": 10, "cache_read_input_tokens": 80}), + json!({"output_tokens": 5}), + Some(json!({"cache_read_input_tokens": 100})), + json!({"output_tokens": 5, "cache_read_input_tokens": 100}), + )] + #[case::delta_cache_fields_are_kept( + json!({"input_tokens": 10, "cache_read_input_tokens": 80}), + json!({"output_tokens": 5, "cache_read_input_tokens": 7}), + None, + json!({"output_tokens": 5, "cache_read_input_tokens": 7}), + )] + fn message_delta_usage_is_completed_from_stop_then_start( + #[case] start: Value, + #[case] delta: Value, + #[case] stop: Option, + #[case] expected: Value, + ) { + let output = promote(vec![ + message_start(start), + message_delta(delta), + message_stop(stop.clone()), + ]); + + assert_eq!(output.len(), 3); + assert_eq!(output[1], message_delta(expected)); + assert_eq!(output[2], message_stop(stop)); + } + + #[test] + fn a_delta_is_flushed_with_start_usage_when_the_stream_ends_without_a_stop() { + let output = promote(vec![ + message_start(json!({"input_tokens": 10, "cache_read_input_tokens": 80})), + message_delta(json!({"output_tokens": 5})), + ]); + + assert_eq!( + output[1], + message_delta(json!({"output_tokens": 5, "cache_read_input_tokens": 80})) + ); + } + + #[test] + fn events_around_the_delta_keep_their_order() { + let ping = event(json!({"type": "ping"})); + let output = promote(vec![ + message_delta(json!({"output_tokens": 5})), + ping.clone(), + message_stop(None), + ]); + + assert_eq!( + output, + vec![ + message_delta(json!({"output_tokens": 5})), + ping, + message_stop(None) + ] + ); + } + + #[rstest] + #[case::metrics_fill_missing_usage( + json!({"type": "message_stop", "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "outputTokenCount": 9}}), + json!({"type": "message_stop", "usage": {"input_tokens": 3, "output_tokens": 9}}), + )] + #[case::the_chunks_own_usage_wins( + json!({"type": "message_stop", "usage": {"input_tokens": 1}, "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "cacheReadInputTokenCount": 40}}), + json!({"type": "message_stop", "usage": {"cache_read_input_tokens": 40, "input_tokens": 1}}), + )] + #[case::no_metrics_leaves_the_chunk( + json!({"type": "message_stop"}), + json!({"type": "message_stop"}), + )] + #[case::empty_metrics_are_dropped( + json!({"type": "message_stop", "amazon-bedrock-invocationMetrics": {}}), + json!({"type": "message_stop"}), + )] + fn invocation_metrics_become_anthropic_usage(#[case] chunk: Value, #[case] expected: Value) { + assert_eq!(with_invocation_usage(chunk), expected); + } + + fn aws_frame(chunk: &Value) -> Vec { + let payload = json!({"bytes": STANDARD.encode(chunk.to_string())}); + let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( + Header::new(":event-type", HeaderValue::String("chunk".into())), + ); + let mut wire = Vec::new(); + write_message_to(&message, &mut wire).unwrap(); + wire + } + + #[tokio::test] + async fn bedrock_stream_yields_the_sse_an_anthropic_client_reads() { + let chunks = [ + json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}), + json!({"type": "message_stop", "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "cacheReadInputTokenCount": 40}}), + ]; + let wire: Vec = chunks.iter().flat_map(aws_frame).collect(); + let bytes: ByteStream = futures_util::stream::iter( + wire.chunks(7) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect::>(), + ) + .boxed(); + + let sse = bedrock_anthropic_messages_event_stream(bytes) + .map_ok(|event| encode_anthropic_sse(&event).unwrap()) + .try_collect::>() + .await + .unwrap() + .concat(); + + let expected: Vec = [ + message_delta( + json!({"output_tokens": 5, "cache_read_input_tokens": 40, "input_tokens": 3}), + ), + message_stop(Some( + json!({"input_tokens": 3, "cache_read_input_tokens": 40}), + )), + ] + .iter() + .flat_map(|event| encode_anthropic_sse(event).unwrap()) + .collect(); + assert_eq!(sse, expected); + } + + #[test] + fn config_uses_the_streaming_url_only_for_streams() { + let env = |_: &str| -> Option { None }; + let config = AmazonAnthropicClaudeMessagesConfig; + + assert_eq!( + config + .get_complete_url(None, "anthropic.claude-3", &env) + .unwrap(), + config + .complete_stream_url(None, "anthropic.claude-3", &env) + .unwrap() + .replace(INVOKE_STREAM_PATH, INVOKE_PATH) + ); + } + + #[rstest] + #[case::an_explicit_key_is_a_bearer_token(Some("token"), None, Some("token"))] + #[case::the_env_token_is_a_bearer_token(None, Some("env-token"), Some("env-token"))] + #[case::no_token_signs_with_sigv4(None, None, None)] + fn requests_sign_only_without_a_bearer_token( + #[case] api_key: Option<&str>, + #[case] env_token: Option<&str>, + #[case] expected_bearer: Option<&str>, + ) { + let env = |name: &str| { + (name == AWS_BEARER_TOKEN_BEDROCK) + .then(|| env_token.map(str::to_string)) + .flatten() + }; + let validated = AmazonAnthropicClaudeMessagesConfig + .validate_environment( + vec![("authorization".into(), "Bearer forwarded".into())], + api_key, + "anthropic.claude-3", + &env, + ) + .unwrap(); + match (validated.auth, expected_bearer) { + ( + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret, + }, + Some(expected), + ) => assert_eq!(secret.expose(), expected), + ( + AuthScheme::AwsSigV4 { + region, service, .. + }, + None, + ) => { + assert_eq!( + (region.as_str(), service), + (DEFAULT_BEDROCK_REGION, BEDROCK_SERVICE) + ); + } + (other, _) => panic!("unexpected auth {other:?}"), + } + } +} diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs new file mode 100644 index 00000000000..4d67a0c0696 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs @@ -0,0 +1 @@ +pub mod anthropic_claude3_transformation; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/mod.rs b/litellm-rust/crates/llms/src/bedrock/messages/mod.rs new file mode 100644 index 00000000000..476a99539ff --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/messages/mod.rs @@ -0,0 +1 @@ +pub mod invoke_transformations; diff --git a/litellm-rust/crates/llms/src/bedrock/mod.rs b/litellm-rust/crates/llms/src/bedrock/mod.rs index 695aeb8af5e..feed6e70e4d 100644 --- a/litellm-rust/crates/llms/src/bedrock/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/mod.rs @@ -1,2 +1,3 @@ pub mod audio_transcription; pub mod chat; +pub mod messages; diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs new file mode 100644 index 00000000000..e885d6f43a1 --- /dev/null +++ b/litellm-rust/crates/llms/src/error.rs @@ -0,0 +1,18 @@ +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("expected {expected}, got {actual}")] + InvalidType { + expected: &'static str, + actual: &'static str, + }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid request: {0}")] + InvalidRequest(String), + #[error("invalid response: {0}")] + InvalidResponse(String), + #[error("unsupported: {0}")] + Unsupported(&'static str), + #[error(transparent)] + Auth(#[from] litellm_auth::Error), +} diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 701eaff4374..e71a9466c0c 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -4,7 +4,10 @@ pub mod azure_ai; pub mod base_llm; pub mod bedrock; pub mod cohere; +mod error; pub mod mistral; pub mod openai; pub mod reducto; pub mod vertex_ai; + +pub use error::Error; diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index f01ec4ad146..1001265413e 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -1,8 +1,8 @@ use litellm_types::responses::streaming_websocket::{ResponsesWsEvent, ResponsesWsTransformResult}; -use crate::base_llm::{ - chat::transformation::Error, - responses::transformation::{ResponsesWebSocketProviderConfig, enforce_model}, +use crate::{ + Error, + base_llm::responses::transformation::{ResponsesWebSocketProviderConfig, enforce_model}, }; pub struct OpenAiResponsesApiConfig; diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index c9342c87e9a..58e2f6cb0ad 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -130,7 +130,8 @@ impl VertexAiOcrConfig { ) -> Result { validate_destination(connection)?; client - .vertex_auth() + .auth() + .gcp .validate_environment( connection.extra_headers.clone(), connection diff --git a/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs index ed22a1d141d..f6f0b8eed42 100644 --- a/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs +++ b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs @@ -1,7 +1,9 @@ use litellm_llms::{ + Error, anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, - base_llm::chat::transformation::{ - BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + base_llm::{ + auth::AuthScheme, + chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, }, }; use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; @@ -430,15 +432,22 @@ fn resolves_the_messages_url_and_x_api_key_auth() { .expect("url builds"), "https://api.anthropic.com/v1/messages" ); - assert_eq!( - config - .auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None) - .expect("auth resolves"), - RequestAuth::Header { - name: "x-api-key", - value: "sk-x".to_string() - } - ); + let validated = config + .validate_environment( + Vec::new(), + Some("sk-x"), + "claude-sonnet-4-5", + &Map::new(), + &|_| None, + ) + .expect("auth resolves"); + assert!(matches!( + validated.auth, + AuthScheme::Credential { + placement: litellm_auth::CredentialPlacement::Header("x-api-key"), + ref secret + } if secret.expose() == "sk-x" + )); assert_eq!( config.default_headers(), &[ diff --git a/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs index 4127bcfa19d..0bd637f602c 100644 --- a/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs +++ b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs @@ -1,6 +1,9 @@ +use litellm_auth::CredentialPlacement; use litellm_llms::{ - base_llm::chat::transformation::{ - BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + Error, + base_llm::{ + auth::AuthScheme, + chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, }, bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, }; @@ -273,22 +276,36 @@ fn prefers_an_explicit_runtime_endpoint_over_the_api_base() { ); } +/// The bearer token a config named, or `None` for a SigV4 scheme in the given region. +fn bearer_or_region(auth: AuthScheme) -> Result { + match auth { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret, + } => Ok(secret.expose().to_string()), + AuthScheme::AwsSigV4 { + region, + service: "bedrock", + .. + } => Err(region), + other => panic!("unexpected auth {other:?}"), + } +} + #[test] fn signs_with_sigv4_in_the_resolved_region() { - let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + let validated = BEDROCK_CHAT_COMPLETIONS_CONFIG + .validate_environment( + Vec::new(), + None, + "eu-central-1/anthropic.claude-v2", + &Map::new(), + &|_| None, + ) + .expect("auth resolves"); assert_eq!( - config - .auth( - None, - "eu-central-1/anthropic.claude-v2", - &Map::new(), - &|_| None - ) - .expect("auth resolves"), - RequestAuth::AwsSigV4 { - region: "eu-central-1".to_string(), - service: "bedrock", - } + bearer_or_region(validated.auth), + Err("eu-central-1".to_string()) ); } @@ -302,22 +319,21 @@ fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() { |key: &str| (key == "AWS_BEARER_TOKEN_BEDROCK").then(|| "from-env".to_string()); let no_env = |_: &str| None; let resolve = |api_key, env: &dyn Fn(&str) -> Option| { - BEDROCK_CHAT_COMPLETIONS_CONFIG - .auth( - api_key, - "eu-central-1/anthropic.claude-v2", - &Map::new(), - env, - ) - .expect("auth resolves") - }; - let bearer = |token: &str| RequestAuth::Bearer { - token: token.to_string(), - }; - let sigv4 = RequestAuth::AwsSigV4 { - region: "eu-central-1".to_string(), - service: "bedrock", + bearer_or_region( + BEDROCK_CHAT_COMPLETIONS_CONFIG + .validate_environment( + Vec::new(), + api_key, + "eu-central-1/anthropic.claude-v2", + &Map::new(), + env, + ) + .expect("auth resolves") + .auth, + ) }; + let bearer = |token: &str| Ok(token.to_string()); + let sigv4 = Err("eu-central-1".to_string()); // A caller-supplied key is the bearer token, and outranks the env. assert_eq!( diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 0b26e398ac8..94a69c94fdf 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -6,11 +6,13 @@ license.workspace = true repository.workspace = true [features] -schema = ["dep:schemars"] +schema = ["dep:schemars", "litellm-types/schema"] [dependencies] +litellm-types.workspace = true + indexmap = { version = "2.14.0", features = ["serde"] } -schemars = { version = "1.0", optional = true } +schemars = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/capabilities.rs b/litellm-rust/crates/model-catalog/src/capabilities.rs index 66b5f1c5d2e..3df68fabc4d 100644 --- a/litellm-rust/crates/model-catalog/src/capabilities.rs +++ b/litellm-rust/crates/model-catalog/src/capabilities.rs @@ -24,20 +24,6 @@ pub enum Mode { VideoGeneration, } -/// Reasoning effort level accepted or applied by the model. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "snake_case")] -pub enum ReasoningEffort { - None, - Minimal, - Low, - Medium, - High, - Xhigh, - Max, -} - /// Gemini audio generation API the model is served through. #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 9e9a4220e50..7b8ce15fbd6 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -1,7 +1,6 @@ -use crate::capabilities::{ - AudioFormat, InputModality, Mode, OutputModality, ReasoningEffort, VertexAiAudioApi, -}; +use crate::capabilities::{AudioFormat, InputModality, Mode, OutputModality, VertexAiAudioApi}; use crate::pricing::{OffPeakPricing, SearchContextCostPerQuery, TieredRate, WebSearchBillingUnit}; +use litellm_types::llms::openai::ReasoningEffort; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::BTreeMap; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d965223bd59..7cfb3f207d4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -42,7 +42,6 @@ litellm-auth-aws.workspace = true litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true litellm-core-utils.workspace = true -litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "hashicorp", "cyberark"] } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 6c5a65173e3..e91dd15beb0 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,6 +1,5 @@ -use litellm_core::{Error, audio_transcription, chat_completions, messages, responses}; +use litellm_core::{Phase, RouteError}; use litellm_http::transport::Error as TransportError; -use litellm_llms::base_llm::ocr::error::Error as OcrError; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, @@ -20,73 +19,16 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); -fn auth_is_value_error(error: &litellm_auth::Error) -> bool { - !matches!(error, litellm_auth::Error::MissingApiKey { .. }) +pub(crate) fn route_error_to_pyerr(error: RouteError) -> PyErr { + by_fault(error.is_request(), error.to_string()) } -pub(crate) fn messages_error_to_pyerr(error: messages::Error) -> PyErr { - core_error_to_pyerr(error.into()) -} - -pub(crate) fn audio_transcription_error_to_pyerr(error: audio_transcription::Error) -> PyErr { - core_error_to_pyerr(error.into()) -} - -pub(crate) fn responses_error_to_pyerr(error: responses::Error) -> PyErr { - core_error_to_pyerr(error.into()) -} - -pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { - let value_error = match &error { - Error::Ocr(error) => { - error.is_request() - || matches!( - error, - OcrError::Auth(_) - | OcrError::InvalidProvider(_) - | OcrError::InvalidRequest(_) - | OcrError::MissingField(_) - | OcrError::MissingDocumentUrl - ) - } - Error::Messages(error) => match error { - messages::Error::Auth(source) => auth_is_value_error(source), - _ => error.is_request(), - }, - Error::AudioTranscription(error) => match error { - audio_transcription::Error::Auth(source) => auth_is_value_error(source), - audio_transcription::Error::InvalidProvider(_) - | audio_transcription::Error::InvalidRequest(_) - | audio_transcription::Error::Headers(_) - | audio_transcription::Error::Http(_) - | audio_transcription::Error::InvalidType { .. } - | audio_transcription::Error::MissingField(_) - | audio_transcription::Error::Aws(_) => true, - _ => false, - }, - Error::ChatCompletions(error) => match error { - chat_completions::Error::Auth(source) => auth_is_value_error(source), - chat_completions::Error::InvalidProvider(_) - | chat_completions::Error::InvalidRequest(_) - | chat_completions::Error::Headers(_) - | chat_completions::Error::Http(_) - | chat_completions::Error::InvalidType { .. } - | chat_completions::Error::MissingField(_) - | chat_completions::Error::Aws(_) => true, - _ => false, - }, - Error::Responses(error) => match error { - responses::Error::Auth(source) => auth_is_value_error(source), - responses::Error::InvalidProvider(_) - | responses::Error::InvalidRequest(_) - | responses::Error::Headers(_) => true, - _ => false, - }, - }; - if value_error { - PyValueError::new_err(error.to_string()) +/// A request the caller got wrong is a `ValueError`; anything else is a `RuntimeError`. +pub(crate) fn by_fault(is_request: bool, message: String) -> PyErr { + if is_request { + PyValueError::new_err(message) } else { - PyRuntimeError::new_err(error.to_string()) + PyRuntimeError::new_err(message) } } @@ -96,27 +38,15 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(error: chat_completions::Error) -> PyErr { - use chat_completions::Error; - match error { - Error::Unsupported(_) - | Error::Auth(_) - | Error::Aws(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::Headers(_) - | Error::Http(_) - | Error::Transport(TransportError::Connect(_)) => { - RustBridgeDeclined::new_err(error.to_string()) - } - Error::Transport(TransportError::Http { status, body }) => { - RustUpstreamError::new_err((status, body)) - } - Error::Transport(TransportError::Network(message)) | Error::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } +pub(crate) fn chat_completions_error_to_pyerr(error: RouteError) -> PyErr { + match error.phase() { + Phase::BeforeSend => RustBridgeDeclined::new_err(error.to_string()), + Phase::AfterSend => RustUpstreamError::new_err(match error { + RouteError::Transport(TransportError::Http { status, body }) => (status, body), + RouteError::Transport(TransportError::Network(message)) + | RouteError::InvalidResponse(message) => (0u16, message), + other => (0u16, other.to_string()), + }), } } @@ -158,15 +88,14 @@ mod tests { fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { Python::initialize(); Python::attach(|py| { - let missing = messages_error_to_pyerr(messages::Error::Auth( - litellm_auth::Error::MissingApiKey { + let missing = + route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: "ANTHROPIC_API_KEY", - }, - )); + })); assert!(missing.is_instance_of::(py)); let invalid = - messages_error_to_pyerr(messages::Error::Auth(litellm_auth::Error::InvalidHeader)); + route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::InvalidHeader)); assert!(invalid.is_instance_of::(py)); }); } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 4d8f0fd7147..e74b9d198a0 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -80,13 +80,20 @@ fn decode_ssl_verify(field: &Field<'_>) -> Result, ProjectionE Err(field.invalid("a Boolean, Boolean string, CA path, or None")) } -static POOL: LazyLock = - LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); +static RESOURCES: LazyLock = LazyLock::new(|| { + litellm_core::resources::CoreResources::new(Arc::new(HttpClientPool::new(Arc::new( + PublicDnsResolver, + )))) +}); + +pub(crate) fn resources() -> &'static litellm_core::resources::CoreResources { + &RESOURCES +} static REPORTED_UNSUPPORTED: LazyLock>> = LazyLock::new(Mutex::default); pub(crate) fn pool() -> &'static HttpClientPool { - &POOL + &resources().pool } pub(crate) fn call_config( diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 6f8388471dc..8f49444f730 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,11 +1,14 @@ use pyo3::{exceptions::PyModuleNotFoundError, prelude::*}; +use strum::IntoStaticStr; use crate::coercion::{FieldSpec, ProjectionError}; const MODULE: &str = "litellm.rust_bridge.settings"; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)] +#[strum(serialize_all = "snake_case")] pub(crate) enum PythonSettings { + #[strum(serialize = "http_settings")] Http, UrlPolicy, ProviderDefaults, @@ -26,13 +29,7 @@ impl Snapshot<'_> { impl PythonSettings { pub(crate) fn name(self) -> &'static str { - match self { - Self::Http => "http_settings", - Self::UrlPolicy => "url_policy", - Self::ProviderDefaults => "provider_defaults", - Self::SecretManager => "secret_manager", - Self::SecretManagerBinding => "secret_manager_binding", - } + self.into() } pub(crate) fn read(self, py: Python<'_>) -> PyResult> { diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 93d0e11d323..ad80659de92 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -8,7 +8,7 @@ use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; use crate::{ - errors::audio_transcription_error_to_pyerr, + errors::route_error_to_pyerr, marshal::{RouteOptions, extra_headers_argument, optional_params_argument, optional_timeout}, }; @@ -27,7 +27,7 @@ async fn execute( timeout, } = options; run_audio_transcription( - crate::http::pool(), + crate::http::resources(), &config, AudioTranscriptionRequest { model: &model, @@ -72,7 +72,7 @@ pub(crate) fn transcription( run_sync( py, execute(config, audio, optional_params.unwrap_or_default(), options), - audio_transcription_error_to_pyerr, + route_error_to_pyerr, ) } @@ -105,6 +105,6 @@ pub(crate) fn atranscription<'py>( run_async( py, execute(config, audio, optional_params.unwrap_or_default(), options), - audio_transcription_error_to_pyerr, + route_error_to_pyerr, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 6d7fad0d69c..f4fd53c61b0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -35,7 +35,7 @@ async fn execute( timeout, } = options; run_chat_completions( - crate::http::pool(), + crate::http::resources(), &config, ChatCompletionsRequest { model: &model, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index a253f4f5670..bb0ec6671e3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,7 @@ use std::convert::Infallible; use bytes::Bytes; use litellm_core::messages::{ Error, - route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead}, + route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead, messages_body}, types::MessagesShaping, }; use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; @@ -18,7 +18,7 @@ use pyo3::{ use serde_json::{Map, Value}; use crate::{ - errors::{RustUpstreamError, messages_error_to_pyerr}, + errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{optional_timeout, python_timeout_seconds}, }; @@ -76,7 +76,7 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult { error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; Ok(error) } - other => Ok(messages_error_to_pyerr(other)), + other => Ok(route_error_to_pyerr(other)), } } @@ -91,7 +91,11 @@ impl MessagesPythonHost { Self { request } } - fn projection(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + fn projection( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult> { let request = self.request.bind(py); let argument = |name: &str| -> PyResult>> { Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) @@ -123,17 +127,20 @@ impl MessagesPythonHost { .flatten(); let custom_llm_provider = string("custom_llm_provider")?; let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; - Ok(MessagesCall { - model, + let api_key = string("api_key")?; + let api_base = string("api_base")?; + let extra_headers = self.merged_headers(py, arguments)?; + let provider_specific_header = self.provider_specific_header(py, arguments)?; + Ok(messages_body(body).map(|body| MessagesCall { body, - api_key: string("api_key")?, - api_base: string("api_base")?, - extra_headers: self.merged_headers(py, arguments)?, - provider_specific_header: self.provider_specific_header(py, arguments)?, + api_key, + api_base, + extra_headers, + provider_specific_header, custom_llm_provider, timeout: optional_timeout(timeout), shaping, - }) + })) } fn merged_headers( @@ -220,7 +227,8 @@ impl ProtocolHost for MessagesPythonHost { arguments: &Bound<'_, PyDict>, ) -> Result> { self.projection(py, arguments) - .map_err(|error| InvokeError::Python(self.map_failure(py, error))) + .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? + .map_err(InvokeError::Native) } fn invoke(&mut self, _: Python<'_>, op: Infallible) -> Result<(), InvokeError> { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index a59c9360c36..96ca9eebecb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -28,7 +28,7 @@ fn run_messages( ) -> PyResult> { let secrets = crate::secrets::source(py)?; let config = crate::http::call_config(py, &kwargs, asynchronous)?; - let machine = messages_machine(crate::http::pool(), &config, secrets) + let machine = messages_machine(crate::http::resources(), &config, secrets) .map_err(crate::http::client_error)?; run_legacy_call( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index b0a6acdebfd..2068b6a6e4b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -4,7 +4,7 @@ use pyo3::{ prelude::*, }; -use crate::errors::{RustUpstreamError, core_error_to_pyerr}; +use crate::errors::{RustUpstreamError, by_fault}; pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); @@ -19,7 +19,7 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { upstream_error(py, status, body, Vec::new())? } Error::RequestFormat => { - let error = core_error_to_pyerr(Error::RequestFormat.into()); + let error = by_fault(true, Error::RequestFormat.to_string()); error .value(py) .setattr("ocr_request_format_error", true) @@ -30,13 +30,25 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) } Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), - other => core_error_to_pyerr(other.into()), + other => by_fault(is_request(&other), other.to_string()), }) }) .unwrap_or_else(|error| error); attach_status(mapped, status) } +fn is_request(error: &Error) -> bool { + error.is_request() + || matches!( + error, + Error::Auth(_) + | Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::MissingField(_) + | Error::MissingDocumentUrl + ) +} + fn upstream_error( py: Python<'_>, status: u16, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index e00c57fad64..d7c54e996ff 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -3,15 +3,12 @@ mod errors; mod host; mod project; -use std::sync::LazyLock; - use host::OcrPythonHost; -use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::{provider_config, route::ocr_machine}; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::to_py; -use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, @@ -44,8 +41,6 @@ const ASYNC_SURFACE: LegacySurface = LegacySurface { ..SURFACE }; -static VERTEX_AUTH: LazyLock = LazyLock::new(VertexAuth::default); - fn run_ocr( py: Python<'_>, request: Bound<'_, PyAny>, @@ -55,15 +50,9 @@ fn run_ocr( ) -> PyResult> { let secrets = secrets::source(py)?; let config = http::call_config(py, &kwargs, asynchronous)?; - let client = OcrClient::new( - http::pool(), - &config, - http::url_policy(py)?, - VERTEX_AUTH.clone(), - ocr_settings(py)?, - secrets, - ) - .map_err(http::client_error)?; + let client = http::resources() + .ocr_client(&config, http::url_policy(py)?, ocr_settings(py)?, secrets) + .map_err(http::client_error)?; run_legacy_call( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses.rs index 5995d64649b..bf17ef6edde 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses.rs @@ -6,7 +6,7 @@ use pyo3::{ use serde_json::Value; use crate::{ - errors::{RustBridgeDeclined, responses_error_to_pyerr}, + errors::{RustBridgeDeclined, route_error_to_pyerr}, marshal::{marshal_headers, optional_timeout}, }; @@ -57,7 +57,7 @@ impl ResponsesWebSocketConnection { crate::logger::run_async_value(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await - .map_err(responses_error_to_pyerr)?; + .map_err(route_error_to_pyerr)?; Ok(ResponsesWebSocketConnection { inner }) }) } @@ -65,24 +65,21 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); crate::logger::run_async_value(py, async move { - inner - .send_text(text) - .await - .map_err(responses_error_to_pyerr) + inner.send_text(text).await.map_err(route_error_to_pyerr) }) } fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); crate::logger::run_async_value(py, async move { - inner.recv_text().await.map_err(responses_error_to_pyerr) + inner.recv_text().await.map_err(route_error_to_pyerr) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); crate::logger::run_async_value(py, async move { - inner.close().await.map_err(responses_error_to_pyerr) + inner.close().await.map_err(route_error_to_pyerr) }) } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 6ba60630b3a..6b61ca22dcd 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -73,17 +73,7 @@ impl PythonClient { /// The `KeyManagementSystem` value as Python spells it. fn python_name(system: KeyManagementSystem) -> &'static str { - match system { - KeyManagementSystem::GoogleKms => "google_kms", - KeyManagementSystem::AzureKeyVault => "azure_key_vault", - KeyManagementSystem::AwsSecretManager => "aws_secret_manager", - KeyManagementSystem::GoogleSecretManager => "google_secret_manager", - KeyManagementSystem::HashicorpVault => "hashicorp_vault", - KeyManagementSystem::Cyberark => "cyberark", - KeyManagementSystem::Local => "local", - KeyManagementSystem::AwsKms => "aws_kms", - KeyManagementSystem::Custom => "custom", - } + system.into() } impl ExternalSecretManager for PythonSecretManager { diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 0c32eb00989..06019cc3c9d 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -2,9 +2,8 @@ use std::sync::Arc; use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; use litellm_auth_aws::{ - AwsAuthConfig, + AwsAuthConfig, AwsAuthService, constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME}, - resolve_credentials, }; use litellm_core_utils::settings::Lookup; use litellm_secrets_types::{AwsOperationContext, KeyManagementSettings}; @@ -13,6 +12,7 @@ use crate::Error; #[derive(Clone)] pub(crate) struct Credentials { + auth: AwsAuthService, config: AwsAuthConfig, environment: Arc, } @@ -22,15 +22,22 @@ impl Credentials { settings: &KeyManagementSettings, environment: Arc, ) -> Self { - Self::with_context(settings, environment, &AwsOperationContext::default()) + Self::with_context( + AwsAuthService::default(), + settings, + environment, + &AwsOperationContext::default(), + ) } pub(crate) fn with_context( + auth: AwsAuthService, settings: &KeyManagementSettings, environment: Arc, context: &AwsOperationContext, ) -> Self { Self { + auth, config: AwsAuthConfig { access_key_id: context .access_key_id @@ -69,7 +76,8 @@ impl ProvideCredentials for Credentials { Self: 'a, { future::ProvideCredentials::new(async { - resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) + self.auth + .resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) .await .map_err(|_| { CredentialsError::provider_error("secret manager authentication failed") diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 508d7da15c1..2220e387838 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -39,6 +39,7 @@ pub struct AwsSecretsManagerV2 { #[derive(Clone)] struct ContextClientFactory { + auth: litellm_auth_aws::AwsAuthService, settings: KeyManagementSettings, environment: Arc, endpoint_url: Option, diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager/client.rs b/litellm-rust/crates/secrets-aws/src/secret_manager/client.rs index aac998c65ab..0f0032ed64f 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager/client.rs @@ -22,6 +22,7 @@ impl AwsSecretsManagerV2 { return Ok(None); } let context_client_factory = ContextClientFactory { + auth: litellm_auth_aws::AwsAuthService::default(), settings: settings.clone(), environment: environment.clone(), endpoint_url: environment @@ -90,6 +91,7 @@ impl ContextClientFactory { self.environment.as_ref(), )?)) .credentials_provider(auth::Credentials::with_context( + self.auth.clone(), &settings, self.environment.clone(), context, diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml index dcd06d1a741..6dd847ec989 100644 --- a/litellm-rust/crates/secrets-types/Cargo.toml +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -11,6 +11,7 @@ moka.workspace = true tokio = { workspace = true, features = ["sync"] } serde.workspace = true serde_json.workspace = true +strum.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 44acf512224..82d48f7b2e1 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -1,11 +1,13 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; +use strum::IntoStaticStr; use crate::SecretValue; -#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, IntoStaticStr, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] pub enum KeyManagementSystem { GoogleKms, AzureKeyVault, diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index fe61d6cb5b4..f855a8a64a6 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -15,7 +15,7 @@ cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] futures-util.workspace = true -litellm-python-compat = { path = "../python-compat" } +litellm-python-compat.workspace = true litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } diff --git a/litellm-rust/crates/types/Cargo.toml b/litellm-rust/crates/types/Cargo.toml index 0a0927386f0..e356c8e127d 100644 --- a/litellm-rust/crates/types/Cargo.toml +++ b/litellm-rust/crates/types/Cargo.toml @@ -5,9 +5,14 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +schema = ["dep:schemars"] + [dependencies] +schemars = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true +strum.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/types/src/lib.rs b/litellm-rust/crates/types/src/lib.rs index da5c9ea893f..dc00ea7128e 100644 --- a/litellm-rust/crates/types/src/lib.rs +++ b/litellm-rust/crates/types/src/lib.rs @@ -1,3 +1,4 @@ pub mod llms; +pub mod recognized; pub mod responses; pub mod utils; diff --git a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs index 2f7a75ba517..342e891a1e3 100644 --- a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs +++ b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs @@ -1,5 +1,8 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use strum::IntoStaticStr; + +use crate::{llms::openai::ReasoningEffort, recognized::Recognized}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(untagged)] @@ -79,10 +82,117 @@ pub struct AnthropicMessage { pub extra: Map, } +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum EffortLevel { + Low, + Medium, + High, + Xhigh, + Max, +} + +impl EffortLevel { + pub fn as_str(self) -> &'static str { + self.into() + } +} + +impl From for ReasoningEffort { + fn from(level: EffortLevel) -> Self { + match level { + EffortLevel::Low => Self::Low, + EffortLevel::Medium => Self::Medium, + EffortLevel::High => Self::High, + EffortLevel::Xhigh => Self::Xhigh, + EffortLevel::Max => Self::Max, + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OutputConfig { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effort: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub format: Option, + #[serde(flatten)] + pub extra: Map, +} + +impl OutputConfig { + pub fn is_empty(&self) -> bool { + self.effort.is_none() && self.format.is_none() && self.extra.is_empty() + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ThinkingDisplay { + Summarized, + Omitted, + Updates, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct EnabledThinking { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub budget_tokens: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display: Option>, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct AdaptiveThinking { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display: Option>, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct DisabledThinking { + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum ThinkingConfig { + Enabled(EnabledThinking), + Adaptive(AdaptiveThinking), + Disabled(DisabledThinking), +} + +impl ThinkingConfig { + pub fn enabled(budget_tokens: u64) -> Self { + Self::Enabled(EnabledThinking { + budget_tokens: Some(Recognized::Known(budget_tokens)), + ..EnabledThinking::default() + }) + } + + pub fn adaptive(display: Option) -> Self { + Self::Adaptive(AdaptiveThinking { + display: display.map(Recognized::Known), + ..AdaptiveThinking::default() + }) + } +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct AnthropicMessagesRequest { pub model: String, pub messages: Vec, + #[serde(flatten)] + pub params: AnthropicMessagesOptionalParams, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct AnthropicMessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -104,7 +214,7 @@ pub struct AnthropicMessagesRequest { #[serde(skip_serializing_if = "Option::is_none")] pub tool_choice: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub thinking: Option, + pub thinking: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub service_tier: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -116,13 +226,13 @@ pub struct AnthropicMessagesRequest { #[serde(skip_serializing_if = "Option::is_none")] pub output_format: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub output_config: Option, + pub output_config: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub speed: Option, #[serde(skip_serializing_if = "Option::is_none")] pub inference_geo: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub reasoning_effort: Option, + pub reasoning_effort: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub compaction: Option, #[serde(flatten)] @@ -182,6 +292,33 @@ mod tests { assert_eq!(round_trip::(&block), block); } + #[test] + fn request_splits_required_fields_from_optional_params() { + let body = json!({ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "stream": true, + "safeguards": [{"type": "dangerous_tool_use"}] + }); + let request: AnthropicMessagesRequest = serde_json::from_value(body.clone()).unwrap(); + + assert_eq!( + ( + request.params.max_tokens, + request.params.stream, + request + .params + .extra + .keys() + .map(String::as_str) + .collect::>(), + ), + (Some(16_u64), Some(true), vec!["safeguards"]) + ); + assert_eq!(serde_json::to_value(request).unwrap(), body); + } + #[test] fn text_constructor_serializes_as_a_text_block() { assert_eq!( @@ -240,7 +377,88 @@ mod tests { "safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}], "metadata": {"user_id": "u"} }))] + #[case::typed_thinking_and_output_config(json!({ + "model": "m", + "messages": [], + "thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + "output_config": {"effort": "xhigh", "format": {"type": "json_schema", "schema": {}}, "task_budget": {"type": "tokens", "total": 4096}} + }))] + #[case::unrecognized_values_are_kept_verbatim(json!({ + "model": "m", + "messages": [], + "reasoning_effort": "turbo", + "thinking": {"type": "adaptive", "display": "loud"}, + "output_config": {"effort": 5} + }))] + #[case::unrecognized_shapes_are_kept_verbatim(json!({ + "model": "m", + "messages": [], + "reasoning_effort": 3, + "thinking": {"type": "future", "budget_tokens": 1}, + "output_config": "bogus" + }))] fn request_round_trips_unchanged(#[case] request: Value) { assert_eq!(round_trip::(&request), request); } + + #[rstest] + #[case::enabled( + json!({"type": "enabled", "budget_tokens": 2048, "display": "omitted"}), + ThinkingConfig::Enabled(EnabledThinking { + budget_tokens: Some(Recognized::Known(2048)), + display: Some(Recognized::Known(ThinkingDisplay::Omitted)), + extra: Map::new(), + }) + )] + #[case::enabled_without_budget( + json!({"type": "enabled"}), + ThinkingConfig::Enabled(EnabledThinking::default()) + )] + #[case::enabled_with_unrecognized_budget( + json!({"type": "enabled", "budget_tokens": "lots"}), + ThinkingConfig::Enabled(EnabledThinking { + budget_tokens: Some(Recognized::Unrecognized(json!("lots"))), + ..EnabledThinking::default() + }) + )] + #[case::adaptive_with_unrecognized_display( + json!({"type": "adaptive", "display": "loud"}), + ThinkingConfig::Adaptive(AdaptiveThinking { + display: Some(Recognized::Unrecognized(json!("loud"))), + extra: Map::new(), + }) + )] + #[case::disabled_keeps_extra_fields( + json!({"type": "disabled", "future": true}), + ThinkingConfig::Disabled(DisabledThinking { + extra: Map::from_iter([("future".to_string(), json!(true))]), + }) + )] + fn thinking_config_parses_every_documented_type_leniently( + #[case] thinking: Value, + #[case] expected: ThinkingConfig, + ) { + assert_eq!( + serde_json::from_value::(thinking).unwrap(), + expected + ); + } + + #[rstest] + fn effort_level_names_match_the_wire( + #[values( + EffortLevel::Low, + EffortLevel::Medium, + EffortLevel::High, + EffortLevel::Xhigh, + EffortLevel::Max + )] + level: EffortLevel, + ) { + assert_eq!(serde_json::to_value(level).unwrap(), json!(level.as_str())); + assert_eq!( + serde_json::to_value(ReasoningEffort::from(level)).unwrap(), + json!(level.as_str()) + ); + } } diff --git a/litellm-rust/crates/types/src/llms/openai.rs b/litellm-rust/crates/types/src/llms/openai.rs index 232f5b9cc51..ee8c882c40c 100644 --- a/litellm-rust/crates/types/src/llms/openai.rs +++ b/litellm-rust/crates/types/src/llms/openai.rs @@ -1,5 +1,43 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use strum::IntoStaticStr; + +/// Reasoning effort level accepted or applied by the model. +#[derive(Clone, Copy, Debug, Deserialize, Eq, IntoStaticStr, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum ReasoningEffort { + None, + Minimal, + Low, + Medium, + High, + Xhigh, + Max, +} + +impl ReasoningEffort { + pub const ALL: [Self; 7] = [ + Self::None, + Self::Minimal, + Self::Low, + Self::Medium, + Self::High, + Self::Xhigh, + Self::Max, + ]; + + pub fn as_str(self) -> &'static str { + self.into() + } + + pub fn parse(value: &str) -> Option { + Self::ALL + .into_iter() + .find(|effort| effort.as_str() == value) + } +} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(untagged)] @@ -56,3 +94,39 @@ pub enum ChatCompletionThinkingBlock { cache_control: Option, }, } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + fn reasoning_effort_names_match_the_wire_and_parse_back( + #[values( + ReasoningEffort::None, + ReasoningEffort::Minimal, + ReasoningEffort::Low, + ReasoningEffort::Medium, + ReasoningEffort::High, + ReasoningEffort::Xhigh, + ReasoningEffort::Max + )] + effort: ReasoningEffort, + ) { + assert_eq!( + serde_json::to_value(effort).unwrap(), + Value::String(effort.as_str().to_string()) + ); + assert_eq!(ReasoningEffort::parse(effort.as_str()), Some(effort)); + assert!(ReasoningEffort::ALL.contains(&effort)); + } + + #[rstest] + #[case::unknown("ultra")] + #[case::uppercase("HIGH")] + #[case::empty("")] + fn reasoning_effort_parse_rejects(#[case] value: &str) { + assert_eq!(ReasoningEffort::parse(value), None); + } +} diff --git a/litellm-rust/crates/types/src/recognized.rs b/litellm-rust/crates/types/src/recognized.rs new file mode 100644 index 00000000000..d82b51f9fde --- /dev/null +++ b/litellm-rust/crates/types/src/recognized.rs @@ -0,0 +1,42 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Recognized { + Known(T), + Unrecognized(Value), +} + +impl Recognized { + pub fn known(&self) -> Option<&T> { + match self { + Self::Known(value) => Some(value), + Self::Unrecognized(_) => None, + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::known(json!(7), Recognized::Known(7))] + #[case::wrong_type(json!("7"), Recognized::Unrecognized(json!("7")))] + #[case::out_of_range(json!(-1), Recognized::Unrecognized(json!(-1)))] + #[case::object(json!({"a": 1}), Recognized::Unrecognized(json!({"a": 1})))] + fn value_is_known_only_when_it_parses_as_the_type( + #[case] value: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(value.clone()).unwrap(), + expected + ); + assert_eq!(serde_json::to_value(expected).unwrap(), value); + } +} From affb5475256af1c8035a10a77e7f66bd222f66de Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:12:48 -0700 Subject: [PATCH 52/58] feat(rust): add config router and gateway crates (#43289) Co-authored-by: Yujong Lee --- litellm-rust/Cargo.lock | 161 ++++++++++++- litellm-rust/Cargo.toml | 6 + litellm-rust/crates/config/Cargo.toml | 16 ++ litellm-rust/crates/config/src/error.rs | 7 + litellm-rust/crates/config/src/lib.rs | 48 ++++ litellm-rust/crates/config/tests/config.rs | 119 +++++++++ litellm-rust/crates/gateway-auth/Cargo.toml | 21 ++ litellm-rust/crates/gateway-auth/src/error.rs | 24 ++ litellm-rust/crates/gateway-auth/src/lib.rs | 74 ++++++ .../crates/gateway-auth/tests/auth.rs | 100 ++++++++ .../crates/gateway-inference/AGENTS.md | 5 + .../crates/gateway-inference/Cargo.toml | 29 +++ .../src/audio_transcription.rs | 59 +++++ .../gateway-inference/src/chat_completions.rs | 83 +++++++ .../crates/gateway-inference/src/error.rs | 225 ++++++++++++++++++ .../crates/gateway-inference/src/lib.rs | 59 +++++ .../gateway-inference/src/messages/host.rs | 69 ++++++ .../gateway-inference/src/messages/mod.rs | 137 +++++++++++ .../crates/gateway-inference/src/ocr.rs | 77 ++++++ .../crates/gateway-inference/src/request.rs | 108 +++++++++ .../gateway-inference/tests/messages.rs | 66 +++++ .../crates/gateway-inference/tests/ocr.rs | 118 +++++++++ .../crates/gateway-inference/tests/routes.rs | 92 +++++++ .../gateway-inference/tests/support/mod.rs | 75 ++++++ litellm-rust/crates/gateway/AGENTS.md | 5 + litellm-rust/crates/gateway/Cargo.toml | 24 ++ litellm-rust/crates/gateway/src/lib.rs | 52 ++++ litellm-rust/crates/gateway/src/main.rs | 18 ++ litellm-rust/crates/gateway/tests/server.rs | 90 +++++++ litellm-rust/crates/router/Cargo.toml | 13 + litellm-rust/crates/router/README.md | 5 + litellm-rust/crates/router/src/deployment.rs | 13 + litellm-rust/crates/router/src/lib.rs | 44 ++++ litellm-rust/crates/router/tests/router.rs | 94 ++++++++ 34 files changed, 2133 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/config/Cargo.toml create mode 100644 litellm-rust/crates/config/src/error.rs create mode 100644 litellm-rust/crates/config/src/lib.rs create mode 100644 litellm-rust/crates/config/tests/config.rs create mode 100644 litellm-rust/crates/gateway-auth/Cargo.toml create mode 100644 litellm-rust/crates/gateway-auth/src/error.rs create mode 100644 litellm-rust/crates/gateway-auth/src/lib.rs create mode 100644 litellm-rust/crates/gateway-auth/tests/auth.rs create mode 100644 litellm-rust/crates/gateway-inference/AGENTS.md create mode 100644 litellm-rust/crates/gateway-inference/Cargo.toml create mode 100644 litellm-rust/crates/gateway-inference/src/audio_transcription.rs create mode 100644 litellm-rust/crates/gateway-inference/src/chat_completions.rs create mode 100644 litellm-rust/crates/gateway-inference/src/error.rs create mode 100644 litellm-rust/crates/gateway-inference/src/lib.rs create mode 100644 litellm-rust/crates/gateway-inference/src/messages/host.rs create mode 100644 litellm-rust/crates/gateway-inference/src/messages/mod.rs create mode 100644 litellm-rust/crates/gateway-inference/src/ocr.rs create mode 100644 litellm-rust/crates/gateway-inference/src/request.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/messages.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/ocr.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/routes.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/support/mod.rs create mode 100644 litellm-rust/crates/gateway/AGENTS.md create mode 100644 litellm-rust/crates/gateway/Cargo.toml create mode 100644 litellm-rust/crates/gateway/src/lib.rs create mode 100644 litellm-rust/crates/gateway/src/main.rs create mode 100644 litellm-rust/crates/gateway/tests/server.rs create mode 100644 litellm-rust/crates/router/Cargo.toml create mode 100644 litellm-rust/crates/router/README.md create mode 100644 litellm-rust/crates/router/src/deployment.rs create mode 100644 litellm-rust/crates/router/src/lib.rs create mode 100644 litellm-rust/crates/router/tests/router.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 580b427a97e..b9d363e72ca 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -710,14 +710,20 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", + "hyper 1.10.1", + "hyper-util", "itoa", "matchit", "memchr", "mime", + "multer", "percent-encoding", "pin-project-lite", "serde_core", + "serde_json", + "serde_path_to_error", "sync_wrapper", + "tokio", "tower", "tower-layer", "tower-service", @@ -1180,7 +1186,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5" dependencies = [ "digest 0.10.7", - "spin", + "spin 0.10.1", ] [[package]] @@ -1581,6 +1587,15 @@ dependencies = [ "serde", ] +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -3096,6 +3111,18 @@ dependencies = [ "strum", ] +[[package]] +name = "litellm-config" +version = "0.1.0" +dependencies = [ + "litellm-auth-types", + "rstest", + "serde", + "serde_yaml_ng", + "tempfile", + "thiserror 2.0.19", +] + [[package]] name = "litellm-core" version = "0.1.0" @@ -3183,6 +3210,66 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "litellm-gateway" +version = "0.1.0" +dependencies = [ + "axum", + "litellm-config", + "litellm-core", + "litellm-gateway-auth", + "litellm-gateway-inference", + "litellm-http", + "litellm-llms", + "litellm-secrets", + "rstest", + "serde_json", + "tokio", + "tower-http 0.7.1", + "tracing", +] + +[[package]] +name = "litellm-gateway-auth" +version = "0.1.0" +dependencies = [ + "axum", + "futures-util", + "litellm-auth-types", + "litellm-config", + "litellm-secrets", + "rstest", + "sha2 0.10.9", + "subtle", + "thiserror 2.0.19", + "tokio", + "tower", +] + +[[package]] +name = "litellm-gateway-inference" +version = "0.1.0" +dependencies = [ + "axum", + "base64 0.22.1", + "bytes", + "futures-util", + "litellm-auth", + "litellm-core", + "litellm-host", + "litellm-http", + "litellm-llms", + "litellm-router", + "litellm-secrets", + "litellm-types", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tower", + "wiremock", +] + [[package]] name = "litellm-host" version = "0.1.0" @@ -3348,6 +3435,15 @@ dependencies = [ "thiserror 2.0.19", ] +[[package]] +name = "litellm-router" +version = "0.1.0" +dependencies = [ + "litellm-config", + "litellm-core", + "rstest", +] + [[package]] name = "litellm-secrets" version = "0.1.0" @@ -3766,6 +3862,23 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "multer" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83e87776546dc87511aa5ee218730c92b666d7264ab6ed41f9d215af9cd5224b" +dependencies = [ + "bytes", + "encoding_rs", + "futures-util", + "http 1.4.2", + "httparse", + "memchr", + "mime", + "spin 0.9.9", + "version_check", +] + [[package]] name = "nom" version = "7.1.3" @@ -4767,7 +4880,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http", + "tower-http 0.6.11", "tower-service", "url", "wasm-bindgen", @@ -4809,7 +4922,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http", + "tower-http 0.6.11", "tower-service", "url", "wasm-bindgen", @@ -5331,6 +5444,19 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "serde_yaml_ng" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f" +dependencies = [ + "indexmap 2.14.0", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "sha1" version = "0.10.7" @@ -5460,6 +5586,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" + [[package]] name = "spin" version = "0.10.1" @@ -6031,6 +6163,23 @@ dependencies = [ "url", ] +[[package]] +name = "tower-http" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" +dependencies = [ + "bitflags 2.13.1", + "bytes", + "http 1.4.2", + "http-body 1.1.0", + "percent-encoding", + "pin-project-lite", + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "tower-layer" version = "0.3.3" @@ -6252,6 +6401,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + [[package]] name = "untrusted" version = "0.9.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 442bd620e05..ed703396c22 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -9,8 +9,13 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +litellm-config = { path = "crates/config" } +litellm-router = { path = "crates/router" } litellm-tracing = { path = "crates/tracing" } litellm-core = { path = "crates/core" } +litellm-gateway = { path = "crates/gateway" } +litellm-gateway-inference = { path = "crates/gateway-inference" } +litellm-gateway-auth = { path = "crates/gateway-auth" } litellm-coroutine = { path = "crates/coroutine" } litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } @@ -50,6 +55,7 @@ litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } tracing = "0.1" +axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } bytes = "1" http = "1" google-cloud-auth = { version = "1.16.0", default-features = false } diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml new file mode 100644 index 00000000000..36bd68fe2a0 --- /dev/null +++ b/litellm-rust/crates/config/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-config" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-types.workspace = true +serde.workspace = true +serde_yaml_ng = "0.10.0" +thiserror.workspace = true + +[dev-dependencies] +rstest.workspace = true +tempfile.workspace = true diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs new file mode 100644 index 00000000000..61d19491abc --- /dev/null +++ b/litellm-rust/crates/config/src/error.rs @@ -0,0 +1,7 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("could not read config")] + Read(#[from] std::io::Error), + #[error("invalid YAML config")] + Parse(#[from] serde_yaml_ng::Error), +} diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs new file mode 100644 index 00000000000..8e86e345025 --- /dev/null +++ b/litellm-rust/crates/config/src/lib.rs @@ -0,0 +1,48 @@ +mod error; + +use std::path::Path; + +use litellm_auth_types::SecretValue; +use serde::Deserialize; + +pub use error::Error; + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Config { + pub model_list: Box<[Model]>, + #[serde(default)] + pub general_settings: GeneralSettings, +} + +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct GeneralSettings { + pub master_key: Option, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Model { + pub model_name: String, + pub litellm_params: LiteLlmParams, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LiteLlmParams { + pub model: String, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, +} + +impl Config { + pub fn from_yaml(yaml: &str) -> Result { + Ok(serde_yaml_ng::from_str(yaml)?) + } + + pub fn load(path: impl AsRef) -> Result { + Self::from_yaml(&std::fs::read_to_string(path)?) + } +} diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs new file mode 100644 index 00000000000..ce6d684ec72 --- /dev/null +++ b/litellm-rust/crates/config/tests/config.rs @@ -0,0 +1,119 @@ +use litellm_config::{Config, Error}; +use rstest::{fixture, rstest}; +use tempfile::TempDir; + +#[fixture] +fn directory() -> TempDir { + tempfile::tempdir().unwrap() +} + +#[fixture] +fn model_list_yaml() -> &'static str { + r#" +model_list: + - model_name: assistant + litellm_params: + model: anthropic/test-model + api_key: os.environ/ANTHROPIC_API_KEY + - model_name: local + litellm_params: + model: test-model + api_base: http://localhost:8000/v1 + custom_llm_provider: openai +"# +} + +#[rstest] +fn loads_model_list_from_file(directory: TempDir, model_list_yaml: &str) { + let path = directory.path().join("config.yaml"); + std::fs::write(&path, model_list_yaml).unwrap(); + + let config = Config::load(path).unwrap(); + assert_eq!(config.model_list.len(), 2); + let anthropic = &config.model_list[0]; + assert_eq!(anthropic.model_name, "assistant"); + assert_eq!(anthropic.litellm_params.model, "anthropic/test-model"); + assert_eq!( + anthropic.litellm_params.api_key.as_ref().unwrap().expose(), + "os.environ/ANTHROPIC_API_KEY" + ); + assert!(anthropic.litellm_params.api_base.is_none()); + assert!(anthropic.litellm_params.custom_llm_provider.is_none()); + let local = &config.model_list[1]; + assert_eq!(local.model_name, "local"); + assert_eq!(local.litellm_params.model, "test-model"); + assert!(local.litellm_params.api_key.is_none()); + assert_eq!( + local.litellm_params.api_base.as_deref(), + Some("http://localhost:8000/v1") + ); + assert_eq!( + local.litellm_params.custom_llm_provider.as_deref(), + Some("openai") + ); +} + +#[rstest] +fn config_debug_redacts_api_keys() { + let config = Config::from_yaml( + "model_list: [{model_name: assistant, litellm_params: {model: anthropic/test-model, api_key: secret-value}}]", + ) + .unwrap(); + assert_eq!( + config.model_list[0] + .litellm_params + .api_key + .as_ref() + .unwrap() + .expose(), + "secret-value" + ); + assert!(!format!("{config:?}").contains("secret-value")); +} + +#[rstest] +#[case::malformed_yaml("model_list: [")] +#[case::missing_model_list("{}")] +#[case::missing_params("model_list: [{model_name: assistant}]")] +#[case::missing_model("model_list: [{model_name: assistant, litellm_params: {api_key: key}}]")] +#[case::unsupported_settings("model_list: []\ngeneral_settings: {unknown: true}")] +#[case::misspelled_param( + "model_list: [{model_name: assistant, litellm_params: {model: test, api_bsae: url}}]" +)] +fn rejects_malformed_incomplete_and_unsupported_config(#[case] yaml: &str) { + assert!(matches!(Config::from_yaml(yaml), Err(Error::Parse(_)))); +} + +#[rstest] +fn distinguishes_read_errors_from_parse_errors(directory: TempDir) { + assert!(matches!( + Config::load(directory.path().join("missing.yaml")), + Err(Error::Read(error)) if error.kind() == std::io::ErrorKind::NotFound + )); +} + +#[rstest] +#[case::literal("secret-master-key")] +#[case::reference("os.environ/LITELLM_MASTER_KEY")] +fn loads_and_redacts_the_master_key(#[case] key: &str) { + let config = Config::from_yaml(&format!( + "model_list: []\ngeneral_settings:\n master_key: {key}\n" + )) + .unwrap(); + assert_eq!( + config + .general_settings + .master_key + .as_ref() + .unwrap() + .expose(), + key + ); + assert!(!format!("{config:?}").contains(key)); +} + +#[rstest] +fn missing_general_settings_has_no_master_key() { + let config = Config::from_yaml("model_list: []").unwrap(); + assert!(config.general_settings.master_key.is_none()); +} diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml new file mode 100644 index 00000000000..340f7224618 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-gateway-auth" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +litellm-auth-types.workspace = true +litellm-config.workspace = true +litellm-secrets.workspace = true +sha2.workspace = true +subtle.workspace = true +thiserror.workspace = true + +[dev-dependencies] +futures-util.workspace = true +rstest.workspace = true +tokio.workspace = true +tower = { version = "0.5.3", features = ["util"] } diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs new file mode 100644 index 00000000000..735eb7741ad --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -0,0 +1,24 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("gateway auth not configured")] + Unconfigured, + #[error("missing or invalid bearer token")] + InvalidToken, + #[error("gateway authentication unavailable")] + Secret(#[from] litellm_secrets::Error), +} + +impl IntoResponse for Error { + fn into_response(self) -> Response { + let status = match &self { + Self::InvalidToken => StatusCode::UNAUTHORIZED, + Self::Unconfigured | Self::Secret(_) => StatusCode::INTERNAL_SERVER_ERROR, + }; + (status, self.to_string()).into_response() + } +} diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs new file mode 100644 index 00000000000..7a1fb244654 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -0,0 +1,74 @@ +mod error; + +use std::sync::Arc; + +use axum::{ + extract::FromRequestParts, + http::{header::AUTHORIZATION, request::Parts}, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_secrets::source::SecretSource; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +pub use error::Error; + +#[derive(Clone)] +pub struct Auth { + master_key: Option, + secrets: Arc, +} + +impl Auth { + pub fn from_config(config: &Config, secrets: Arc) -> Self { + Self { + master_key: config.general_settings.master_key.clone(), + secrets, + } + } + + async fn master_key(&self) -> Result { + let configured = self.master_key.as_ref().ok_or(Error::Unconfigured)?; + let resolved = match configured.expose().strip_prefix("os.environ/") { + Some(name) if !name.is_empty() => self + .secrets + .get_secret_str(name) + .await? + .ok_or(Error::Unconfigured)?, + Some(_) => return Err(Error::Unconfigured), + None => configured.clone(), + }; + if resolved.expose().trim().is_empty() { + return Err(Error::Unconfigured); + } + Ok(resolved) + } +} + +pub fn hash_token(token: &str) -> String { + format!("{:x}", Sha256::digest(token.as_bytes())) +} + +pub struct RequireMasterKey; + +impl FromRequestParts for RequireMasterKey { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &Auth) -> Result { + let expected = state.master_key().await?; + let provided = parts + .headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::trim) + .ok_or(Error::InvalidToken)?; + let actual_hash = Sha256::digest(provided.as_bytes()); + let expected_hash = Sha256::digest(expected.expose().as_bytes()); + match bool::from(actual_hash.ct_eq(&expected_hash)) { + true => Ok(Self), + false => Err(Error::InvalidToken), + } + } +} diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs new file mode 100644 index 00000000000..58625296664 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -0,0 +1,100 @@ +use std::sync::Arc; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, StatusCode}, + middleware::from_extractor_with_state, + routing::get, +}; +use futures_util::future::BoxFuture; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_secrets::source::SecretSource; +use rstest::{fixture, rstest}; +use tower::ServiceExt; + +struct Secrets; + +impl SecretSource for Secrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + match name { + "MASTER_KEY" => Ok(Some(SecretValue::new("resolved-key"))), + "EMPTY" => Ok(Some(SecretValue::new(""))), + "ERROR" => Err(litellm_secrets::Error::ExternalRead(Box::new( + std::io::Error::other("private-backend-detail"), + ))), + _ => Ok(None), + } + }) + } +} + +#[fixture] +fn secrets() -> Arc { + Arc::new(Secrets) +} + +#[rstest] +#[case::literal("literal-key", Some("Bearer literal-key"), 204)] +#[case::reference("os.environ/MASTER_KEY", Some("Bearer resolved-key"), 204)] +#[case::reference_is_not_a_token( + "os.environ/MASTER_KEY", + Some("Bearer os.environ/MASTER_KEY"), + 401 +)] +#[case::wrong("literal-key", Some("Bearer other-key"), 401)] +#[case::missing("literal-key", None, 401)] +#[case::wrong_scheme("literal-key", Some("Basic literal-key"), 401)] +#[case::empty_token("literal-key", Some("Bearer "), 401)] +#[case::missing_reference("os.environ/MISSING", Some("Bearer os.environ/MISSING"), 500)] +#[case::empty_reference("os.environ/EMPTY", Some("Bearer "), 500)] +#[case::empty_key("", Some("Bearer "), 500)] +#[case::whitespace_key(" ", Some("Bearer "), 500)] +#[case::empty_reference_name("os.environ/", Some("Bearer os.environ/"), 500)] +#[case::secret_failure("os.environ/ERROR", Some("Bearer private-backend-detail"), 500)] +#[tokio::test] +async fn enforces_configured_keys_without_exposing_secrets( + secrets: Arc, + #[case] key: &str, + #[case] authorization: Option<&str>, + #[case] status: u16, +) { + let config = Config::from_yaml(&format!( + "model_list: []\ngeneral_settings:\n master_key: '{key}'\n" + )) + .unwrap(); + let app = Router::new() + .route("/protected", get(|| async { StatusCode::NO_CONTENT })) + .layer(from_extractor_with_state::( + Auth::from_config(&config, secrets), + )); + let request = Request::get("/protected"); + let request = match authorization { + Some(value) => request.header("authorization", value), + None => request, + }; + let response = app + .oneshot(request.body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status().as_u16(), status); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + assert!(!text.contains("private-backend-detail")); + assert!(!text.contains("literal-key")); + assert!(!text.contains("resolved-key")); +} + +#[rstest] +fn hash_token_matches_python_sha256_hexdigest() { + assert_eq!( + hash_token("sk-1234"), + "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ); +} diff --git a/litellm-rust/crates/gateway-inference/AGENTS.md b/litellm-rust/crates/gateway-inference/AGENTS.md new file mode 100644 index 00000000000..7dc57380083 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/AGENTS.md @@ -0,0 +1,5 @@ +- Expose a mountable Axum router; listener binding, server lifecycle, and shared inbound middleware belong to `gateway` +- Own the public inference HTTP boundary: endpoint paths, request parsing, model alias resolution, response envelopes, and SSE delivery +- Delegate inference execution to `core` and provider transformations and authentication to `llms` and the auth crates; do not duplicate them in handlers +- Use injected deployments, HTTP pools, settings, and secret sources; do not load process configuration or construct independent clients in handlers +- Test HTTP contracts here, including status codes, forwarded headers, error envelopes, and streaming behavior; keep core and provider tests in their owning crates diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml new file mode 100644 index 00000000000..e3f40ec5354 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "litellm-gateway-inference" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum = { workspace = true, features = ["json", "multipart"] } +base64.workspace = true +bytes.workspace = true +futures-util.workspace = true +litellm-auth.workspace = true +litellm-core.workspace = true +litellm-host.workspace = true +litellm-http.workspace = true +litellm-llms.workspace = true +litellm-router.workspace = true +litellm-secrets.workspace = true +litellm-types.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["sync"] } + +[dev-dependencies] +futures-util.workspace = true +rstest.workspace = true +tower = { version = "0.5.3", features = ["util"] } +wiremock = "0.6.5" diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs new file mode 100644 index 00000000000..d5fd6603e20 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -0,0 +1,59 @@ +use std::{path::Path, sync::Arc}; + +use axum::{ + Json, + extract::{Request, State}, + response::{IntoResponse, Response}, +}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; +use serde_json::{Value, json}; + +use crate::{Error, Gateway, request}; + +pub(crate) async fn create(State(gateway): State>, request: Request) -> Response { + match handle(&gateway, request).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.openai_response(), + } +} + +async fn handle(gateway: &Gateway, request: Request) -> Result { + let (body, upload) = request::parse(request).await?; + let deployment = request::deployment(gateway, &body)?; + let audio = match upload { + Some(upload) => { + let format = upload + .file_name + .as_deref() + .and_then(|name| Path::new(name).extension()) + .and_then(|extension| extension.to_str()) + .ok_or_else(|| { + Error::InvalidBody("audio file requires a filename extension".into()) + })?; + json!({"data": STANDARD.encode(upload.bytes), "format": format.to_ascii_lowercase()}) + } + None => body + .get("audio") + .cloned() + .ok_or_else(|| Error::InvalidBody("audio is required".into()))?, + }; + Ok(audio_transcription( + &gateway.resources, + &gateway.http, + AudioTranscriptionRequest { + model: &deployment.model, + audio, + api_key: deployment.api_key.as_deref(), + api_base: deployment.api_base.as_deref(), + custom_llm_provider: deployment.custom_llm_provider.as_deref(), + extra_headers: None, + optional_params: body + .into_iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "audio")) + .collect(), + timeout: deployment.timeout, + }, + ) + .await?) +} diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs new file mode 100644 index 00000000000..c386f38fe06 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -0,0 +1,83 @@ +use std::sync::Arc; + +use axum::{ + Json, + body::Bytes, + extract::{Path, State}, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use litellm_core::chat_completions::{chat_completions, types::ChatCompletionsRequest}; +use serde_json::{Map, Value}; + +use crate::{Error, Gateway, request}; + +pub(crate) async fn create(State(gateway): State>, body: Bytes) -> Response { + respond(&gateway, request::object(&body)).await +} + +pub(crate) async fn deployment( + State(gateway): State>, + Path(path): Path, + body: Bytes, +) -> Response { + if let Some(model) = path + .strip_suffix("/chat/completions") + .filter(|model| !model.is_empty()) + { + let body = request::object(&body).map(|body| { + if body.get("model").is_some_and(|model| !model.is_null()) { + return body; + } + body.into_iter() + .chain([("model".into(), Value::String(model.into()))]) + .collect() + }); + return respond(&gateway, body).await; + } + if path.ends_with("/embeddings") || path.ends_with("/completions") { + return Error::Unsupported(path).openai_response(); + } + StatusCode::NOT_FOUND.into_response() +} + +async fn respond(gateway: &Gateway, body: Result, Error>) -> Response { + let result = match body { + Ok(body) => handle(gateway, body).await, + Err(error) => Err(error), + }; + match result { + Ok(response) => response, + Err(error) => error.openai_response(), + } +} + +async fn handle(gateway: &Gateway, body: Map) -> Result { + let deployment = request::deployment(gateway, &body)?; + if body.get("stream").and_then(Value::as_bool) == Some(true) { + return Err(Error::Unsupported("streaming chat completions".into())); + } + let messages = body + .get("messages") + .cloned() + .ok_or_else(|| Error::InvalidBody("messages is required".into()))?; + let response = chat_completions( + &gateway.resources, + &gateway.http, + ChatCompletionsRequest { + model: &deployment.model, + messages, + optional_params: body + .into_iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages" | "stream")) + .collect(), + api_key: deployment.api_key.as_deref(), + api_base: deployment.api_base.as_deref(), + custom_llm_provider: deployment.custom_llm_provider.as_deref(), + extra_headers: None, + timeout: deployment.timeout, + }, + ) + .await?; + Ok(Json(response).into_response()) +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs new file mode 100644 index 00000000000..1d40857d962 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -0,0 +1,225 @@ +use axum::http::StatusCode; +use axum::{ + Json, + response::{IntoResponse, Response}, +}; +use litellm_core::RouteError; +use litellm_http::transport::Error as TransportError; +use litellm_llms::base_llm::ocr::error::Error as OcrError; +use serde_json::{Map, Value, json}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid request body: {0}")] + InvalidBody(String), + #[error( + "/v1/messages: Invalid model name passed in model={0}. Call `/v1/models` to view available models for your key." + )] + UnknownModel(String), + #[error(transparent)] + Route(#[from] RouteError), + #[error(transparent)] + Ocr(#[from] OcrError), + #[error("{0} is not implemented by the Rust gateway")] + Unsupported(String), + #[error("request body exceeds the size limit")] + BodyTooLarge, + #[error("{0}")] + Internal(String), +} + +impl Error { + pub fn status(&self) -> StatusCode { + match self { + Self::Unsupported(_) + | Self::Route(RouteError::Unsupported(_)) + | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, + Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, + Self::Ocr( + OcrError::Auth(litellm_auth::Error::MissingApiKey { .. }) + | OcrError::MissingAzureAiCredentials + | OcrError::MissingAzureDocumentIntelligenceCredentials + | OcrError::MissingReductoApiKey, + ) => StatusCode::UNAUTHORIZED, + Self::Ocr(error) => error + .http_status_code() + .and_then(|status| StatusCode::from_u16(status).ok()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST, + Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => { + StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY) + } + Self::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { + StatusCode::UNAUTHORIZED + } + Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST, + Self::Route(_) | Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + pub fn openai_response(self) -> Response { + let status = self.status(); + let message = match &self { + Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"), + _ => self.to_string(), + }; + ( + status, + Json(json!({"error": { + "message": message, + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }})), + ) + .into_response() + } + + /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream + /// body already in that shape passes through, any other has its message extracted. + pub fn body(&self, request_id: Option<&str>) -> Value { + let raw = match self { + Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => body.clone(), + other => other.to_string(), + }; + let parsed = serde_json::from_str::(&raw).ok(); + let envelope = match parsed { + Some(Value::Object(object)) if is_anthropic_error(&object) => object, + Some(Value::Object(object)) => { + envelope(self.status(), provider_message(&object).unwrap_or(&raw)) + } + _ => envelope(self.status(), &raw), + }; + Value::Object(with_request_id(envelope, request_id)) + } + + /// An `event: error` frame, for a stream that fails after its headers went out. + pub fn sse_frame(&self) -> String { + format!("event: error\ndata: {}\n\n", self.body(None)) + } +} + +fn error_type(status: StatusCode) -> &'static str { + match status.as_u16() { + 400 => "invalid_request_error", + 401 => "authentication_error", + 403 => "permission_error", + 404 => "not_found_error", + 413 => "request_too_large", + 429 => "rate_limit_error", + 529 => "overloaded_error", + _ => "api_error", + } +} + +fn envelope(status: StatusCode, message: &str) -> Map { + let Value::Object(envelope) = json!({ + "type": "error", + "error": {"type": error_type(status), "message": message}, + }) else { + unreachable!("a json object literal is an object") + }; + envelope +} + +fn is_anthropic_error(object: &Map) -> bool { + object.get("type").and_then(Value::as_str) == Some("error") + && object + .get("error") + .and_then(Value::as_object) + .is_some_and(|error| error.contains_key("type") && error.contains_key("message")) +} + +fn provider_message(object: &Map) -> Option<&str> { + if let Some(detail) = object.get("detail").and_then(Value::as_object) { + return detail.get("message").and_then(Value::as_str); + } + ["Message", "message"] + .into_iter() + .filter_map(|key| object.get(key).and_then(Value::as_str)) + .find(|message| !message.is_empty()) +} + +fn with_request_id(envelope: Map, request_id: Option<&str>) -> Map { + match request_id { + Some(id) if !id.is_empty() && !envelope.contains_key("request_id") => envelope + .into_iter() + .chain([("request_id".to_string(), Value::from(id))]) + .collect(), + _ => envelope, + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn upstream(status: u16, body: &str) -> Error { + Error::Route(RouteError::Transport(TransportError::Http { + status, + body: body.into(), + })) + } + + #[rstest] + #[case::anthropic_body_passes_through( + upstream(529, r#"{"type":"error","error":{"type":"overloaded_error","message":"busy","extra":1}}"#), + Some("req_1"), + json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy", "extra": 1}, "request_id": "req_1"}), + )] + #[case::upstream_request_id_wins( + upstream(400, r#"{"type":"error","error":{"type":"x","message":"m"},"request_id":"upstream"}"#), + Some("caller"), + json!({"type": "error", "error": {"type": "x", "message": "m"}, "request_id": "upstream"}), + )] + #[case::bedrock_detail( + upstream(403, r#"{"detail":{"message":"denied"}}"#), + None, + json!({"type": "error", "error": {"type": "permission_error", "message": "denied"}}), + )] + #[case::aws_message( + upstream(429, r#"{"Message":"slow down"}"#), + None, + json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}}), + )] + #[case::plain_text_with_unmapped_status( + upstream(502, "bad gateway"), + None, + json!({"type": "error", "error": {"type": "api_error", "message": "bad gateway"}}), + )] + #[case::unknown_model( + Error::UnknownModel("nope".into()), + None, + json!({"type": "error", "error": { + "type": "invalid_request_error", + "message": "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key.", + }}), + )] + fn body_follows_the_anthropic_exception_mapping( + #[case] error: Error, + #[case] request_id: Option<&str>, + #[case] expected: Value, + ) { + assert_eq!(error.body(request_id), expected); + } + + #[rstest] + #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] + #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] + #[case::missing_key( + Error::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + })), + StatusCode::UNAUTHORIZED, + )] + #[case::lost_connection( + Error::Route(RouteError::Transport(TransportError::Network("reset".into()))), + StatusCode::INTERNAL_SERVER_ERROR, + )] + fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { + assert_eq!(error.status(), status); + } +} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs new file mode 100644 index 00000000000..eebe3f34a09 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -0,0 +1,59 @@ +//! The proxy's inference endpoints as an axum [`Router`] a server mounts. +//! +//! Authentication, rate limiting and logging are the mounting server's layers; this crate +//! maps a public model name to its deployment and runs the core route. + +mod audio_transcription; +mod chat_completions; +mod error; +pub mod messages; +mod ocr; +mod request; + +use std::sync::Arc; + +use axum::{Router, routing::post}; +use litellm_core::resources::CoreResources; +use litellm_http::HttpClientConfig; +use litellm_llms::base_llm::ocr::handler::OcrClient; +use litellm_secrets::source::SecretSource; + +pub use error::Error; +pub use litellm_router::{Deployment, Router as ModelList}; + +pub struct Gateway { + pub resources: CoreResources, + pub http: HttpClientConfig, + pub secrets: Arc, + pub models: ModelList, + pub ocr: OcrClient, +} + +pub fn router(gateway: Arc) -> Router { + Router::new() + .route("/v1/messages", post(messages::create)) + .route("/ocr", post(ocr::create)) + .route("/v1/ocr", post(ocr::create)) + .route("/chat/completions", post(chat_completions::create)) + .route("/v1/chat/completions", post(chat_completions::create)) + .route("/engines/{*path}", post(chat_completions::deployment)) + .route( + "/openai/deployments/{*path}", + post(chat_completions::deployment), + ) + .route("/audio/transcriptions", post(audio_transcription::create)) + .route( + "/v1/audio/transcriptions", + post(audio_transcription::create), + ) + .route("/responses", post(request::unsupported)) + .route("/v1/responses", post(request::unsupported)) + .route("/embeddings", post(request::unsupported)) + .route("/v1/embeddings", post(request::unsupported)) + .route("/completions", post(request::unsupported)) + .route("/v1/completions", post(request::unsupported)) + .layer(axum::extract::DefaultBodyLimit::max( + request::MAX_BODY_BYTES, + )) + .with_state(gateway) +} diff --git a/litellm-rust/crates/gateway-inference/src/messages/host.rs b/litellm-rust/crates/gateway-inference/src/messages/host.rs new file mode 100644 index 00000000000..57486d9d157 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/messages/host.rs @@ -0,0 +1,69 @@ +use std::{convert::Infallible, sync::Mutex}; + +use bytes::Bytes; +use litellm_core::messages::{ + Error, + route::{LocalMessagesHost, Messages, MessagesCall, MessagesStreamHead}, +}; +use litellm_host::host::{Demand, Host}; +use tokio::sync::{mpsc, oneshot}; + +/// Hands a streamed response to the HTTP body: the head once, then each chunk. A dropped +/// receiver means the client went away, which detaches the call. +pub(super) struct ChannelHost { + local: LocalMessagesHost, + head: Mutex>>, + pub(super) chunks: mpsc::Sender, +} + +impl ChannelHost { + pub(super) fn new( + call: MessagesCall, + head: oneshot::Sender, + chunks: mpsc::Sender, + ) -> Self { + Self { + local: LocalMessagesHost::new(call), + head: Mutex::new(Some(head)), + chunks, + } + } + + fn take_head(&self) -> Option> { + self.head + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + } + + pub(super) fn opened(&self) -> bool { + self.head + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_none() + } +} + +impl Host for ChannelHost { + async fn project(&self) -> Result { + self.local.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn open(&self, head: MessagesStreamHead) -> Result { + Ok(match self.take_head().map(|sender| sender.send(head)) { + Some(Ok(())) => Demand::More, + Some(Err(_)) | None => Demand::Detached, + }) + } + + async fn deliver(&self, chunk: Bytes) -> Result { + Ok(match self.chunks.send(chunk).await { + Ok(()) => Demand::More, + Err(_) => Demand::Detached, + }) + } +} diff --git a/litellm-rust/crates/gateway-inference/src/messages/mod.rs b/litellm-rust/crates/gateway-inference/src/messages/mod.rs new file mode 100644 index 00000000000..e96cae8ba53 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/messages/mod.rs @@ -0,0 +1,137 @@ +//! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it. + +mod host; + +use std::{convert::Infallible, sync::Arc}; + +use axum::{ + Json, + body::{Body, Bytes}, + extract::State, + http::{HeaderMap, StatusCode, header}, + response::{IntoResponse, Response}, +}; +use host::ChannelHost; +use litellm_core::messages::route::{ + MessagesCall, MessagesOutput, messages_body, messages_machine, +}; +use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; +use serde_json::{Map, Value}; +use tokio::sync::{mpsc, oneshot}; + +use crate::{Deployment, Error, Gateway}; + +/// Client headers Python forwards to Anthropic-speaking providers on every call. +const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"]; +const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai"; + +pub async fn create( + State(gateway): State>, + headers: HeaderMap, + body: Bytes, +) -> Response { + let request_id = headers + .get("x-request-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + match handle(&gateway, &headers, &body).await { + Ok(response) => response, + Err(error) => (error.status(), Json(error.body(request_id.as_deref()))).into_response(), + } +} + +async fn handle(gateway: &Gateway, headers: &HeaderMap, body: &[u8]) -> Result { + let body = match serde_json::from_slice(body) { + Ok(Value::Object(body)) => body, + Ok(_) => return Err(Error::InvalidBody("expected a JSON object".into())), + Err(error) => return Err(Error::InvalidBody(error.to_string())), + }; + let model_name = body + .get("model") + .and_then(Value::as_str) + .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + let deployment = gateway + .models + .get(model_name) + .ok_or_else(|| Error::UnknownModel(model_name.to_owned()))?; + let call = project(deployment, body, headers)?; + let machine = messages_machine(&gateway.resources, &gateway.http, gateway.secrets.clone()) + .map_err(|error| Error::Route(error.into()))?; + + let (head_sender, head) = oneshot::channel(); + let (chunk_sender, chunks) = mpsc::channel(1); + let host = ChannelHost::new(call, head_sender, chunk_sender); + let call = tokio::spawn(async move { + let outcome = litellm_host::run::run(machine, &host).await; + if let Err(error) = &outcome + && host.opened() + { + let _ = host + .chunks + .send(Bytes::from(Error::Route(error.clone()).sse_frame())) + .await; + } + outcome + }); + tokio::select! { + biased; + Ok(_) = head => Ok(stream(chunks)), + joined = call => match joined.map_err(|error| Error::Internal(error.to_string()))?? { + MessagesOutput::Message(message) => Ok(Json(message).into_response()), + MessagesOutput::Streamed => Err(Error::Internal("the stream ended before it opened".into())), + }, + } +} + +fn project( + deployment: &Deployment, + body: Map, + headers: &HeaderMap, +) -> Result { + let body = body + .into_iter() + .map(|(name, value)| match name.as_str() { + "model" => (name, Value::from(deployment.model.as_str())), + _ => (name, value), + }) + .collect(); + Ok(MessagesCall { + body: messages_body(body)?, + api_key: deployment.api_key.clone(), + api_base: deployment.api_base.clone(), + custom_llm_provider: deployment.custom_llm_provider.clone(), + extra_headers: None, + provider_specific_header: anthropic_api_headers(headers), + timeout: deployment.timeout, + shaping: deployment.shaping.clone(), + }) +} + +fn anthropic_api_headers(headers: &HeaderMap) -> Option { + let extra_headers: Map = ANTHROPIC_API_HEADERS + .into_iter() + .filter_map(|name| { + let value = headers.get(name)?.to_str().ok()?; + Some((name.to_owned(), Value::from(value))) + }) + .collect(); + (!extra_headers.is_empty()).then(|| { + ProviderSpecificHeaders::One(ProviderSpecificHeader { + custom_llm_provider: ANTHROPIC_API_HEADER_PROVIDERS.into(), + extra_headers, + }) + }) +} + +fn stream(chunks: mpsc::Receiver) -> Response { + let body = futures_util::stream::unfold(chunks, |mut chunks| async move { + let chunk = chunks.recv().await?; + Some((Ok::<_, Infallible>(chunk), chunks)) + }); + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "text/event-stream")], + Body::from_stream(body), + ) + .into_response() +} diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs new file mode 100644 index 00000000000..d666223e037 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -0,0 +1,77 @@ +use std::sync::Arc; + +use axum::{ + Json, + extract::{Request, State}, + response::{IntoResponse, Response}, +}; +use litellm_auth::SecretValue; +use litellm_core::ocr::{ + client::perform, + types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput}, +}; +use litellm_llms::base_llm::ocr::transformation::OcrDocument; +use serde_json::Value; + +use crate::{Error, Gateway, request}; + +pub(crate) async fn create(State(gateway): State>, request: Request) -> Response { + match handle(&gateway, request).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.openai_response(), + } +} + +async fn handle(gateway: &Gateway, request: Request) -> Result { + let header_format = request + .headers() + .get("x-req-format") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let (body, upload) = request::parse(request).await?; + let deployment = request::deployment(gateway, &body)?; + let document = match upload { + Some(upload) => OcrDocumentInput::Bytes { + bytes: upload.bytes, + file_name: upload.file_name, + mime_type: upload.mime_type, + }, + None => OcrDocument::try_from( + body.get("document") + .cloned() + .ok_or_else(|| Error::InvalidBody("document is required".into()))?, + )? + .into(), + }; + let format = body + .get("req_format") + .filter(|value| !value.is_null()) + .cloned() + .or_else(|| header_format.map(Value::String)); + let format = format.map(|value| match value { + Value::String(value) => Value::String(value.trim().to_ascii_lowercase()), + value => value, + }); + let options = body + .into_iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "document" | "req_format")) + .chain(format.map(|value| ("req_format".into(), value))) + .collect(); + let call = LiteLLMOcrRequest::from_inputs( + deployment.model.clone(), + document, + deployment.custom_llm_provider.as_deref(), + options, + OcrConnectionInputs { + api_key: deployment.api_key.clone().map(SecretValue::new), + api_base: deployment.api_base.clone(), + timeout: deployment.timeout, + ..Default::default() + }, + )?; + let response = perform(&gateway.ocr, call).await?; + match response.provider_native_response { + Some(native) => Ok(Value::Object(native)), + None => Ok(response.into_json()), + } +} diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs new file mode 100644 index 00000000000..f58c7b3ed79 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -0,0 +1,108 @@ +use axum::{ + body::{Bytes, to_bytes}, + extract::{FromRequest, Multipart, Request}, + http::Uri, + response::Response, +}; +use serde_json::{Map, Value}; + +use crate::{Deployment, Error, Gateway}; + +pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024; +pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; + +pub(crate) struct Upload { + pub bytes: Bytes, + pub file_name: Option, + pub mime_type: Option, +} + +pub(crate) fn object(body: &[u8]) -> Result, Error> { + match serde_json::from_slice(body) { + Ok(Value::Object(body)) => Ok(body), + Ok(_) => Err(Error::InvalidBody("expected a JSON object".into())), + Err(error) => Err(Error::InvalidBody(error.to_string())), + } +} + +pub(crate) fn deployment<'a>( + gateway: &'a Gateway, + body: &Map, +) -> Result<&'a Deployment, Error> { + let model = body + .get("model") + .and_then(Value::as_str) + .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + gateway + .models + .get(model) + .ok_or_else(|| Error::UnknownModel(model.to_owned())) +} + +pub(crate) async fn parse(request: Request) -> Result<(Map, Option), Error> { + let multipart = request + .headers() + .get("content-type") + .and_then(|header| header.to_str().ok()) + .is_some_and(|value| { + value + .to_ascii_lowercase() + .starts_with("multipart/form-data") + }); + if !multipart { + let body = to_bytes(request.into_body(), MAX_BODY_BYTES) + .await + .map_err(|_| Error::BodyTooLarge)?; + return Ok((object(&body)?, None)); + } + let mut multipart = Multipart::from_request(request, &()) + .await + .map_err(|error| Error::InvalidBody(error.to_string()))?; + let mut fields = Map::new(); + let mut upload = None; + while let Some(field) = multipart.next_field().await.map_err(multipart_error)? { + let name = field.name().unwrap_or_default().to_owned(); + if name == "file" { + let file_name = field.file_name().map(str::to_owned); + let mime_type = field + .content_type() + .and_then(|value| value.split(';').next()) + .map(str::trim) + .filter(|value| *value != "application/octet-stream") + .map(str::to_owned); + let bytes = field.bytes().await.map_err(multipart_error)?; + if bytes.len() > MAX_FILE_BYTES { + return Err(Error::BodyTooLarge); + } + if bytes.is_empty() { + return Err(Error::InvalidBody("uploaded file is empty".into())); + } + upload = Some(Upload { + bytes, + file_name, + mime_type, + }); + } else if name != "document" { + let text = field.text().await.map_err(multipart_error)?; + let value = serde_json::from_str(&text).unwrap_or(Value::String(text)); + fields.insert(name, value); + } + } + if upload.is_none() { + return Err(Error::InvalidBody( + "multipart request requires a file field".into(), + )); + } + Ok((fields, upload)) +} + +fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error { + if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { + return Error::BodyTooLarge; + } + Error::InvalidBody(error.to_string()) +} + +pub(crate) async fn unsupported(uri: Uri) -> Response { + Error::Unsupported(uri.path().to_owned()).openai_response() +} diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs new file mode 100644 index 00000000000..ef7e7c66681 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -0,0 +1,66 @@ +mod support; + +use axum::{ + body::{Body, to_bytes}, + http::Request, +}; +use rstest::rstest; +use serde_json::json; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header, method, path}, +}; + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn messages_reaches_the_provider_and_preserves_json_or_sse(#[case] streaming: bool) { + let upstream = MockServer::start().await; + let message = json!({"id": "msg_test", "type": "message", "role": "assistant", + "model": "test-model", "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}}); + let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let template = if streaming { + ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream") + } else { + ResponseTemplate::new(200).set_body_json(message.clone()) + }; + let messages = json!([{"role": "user", "content": "hi"}]); + Mock::given(method("POST")).and(path("/v1/messages")) + .and(header("x-api-key", "test-key")) + .and(header("anthropic-beta", "test-feature")) + .and(body_json(json!({"model": "test-model", "messages": messages, "max_tokens": 16, "stream": streaming}))) + .respond_with(template).expect(1).mount(&upstream).await; + let request = Request::post("/v1/messages") + .header("content-type", "application/json").header("anthropic-beta", "test-feature") + .body(Body::from(json!({"model": "public/model", "messages": messages, "max_tokens": 16, "stream": streaming}).to_string())).unwrap(); + let response = support::app("anthropic/test-model", &upstream.uri()) + .oneshot(request) + .await + .unwrap(); + assert_eq!(response.status(), 200); + if streaming { + assert_eq!(response.headers()["content-type"], "text/event-stream"); + assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), sse); + } else { + let body = support::json(response).await; + assert_eq!(body["content"], message["content"]); + assert_eq!(body["usage"], message["usage"]); + } +} + +#[tokio::test] +async fn invalid_messages_stays_an_anthropic_error() { + let response = support::post( + support::app("anthropic/test-model", "http://127.0.0.1:1"), + "/v1/messages", + json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}), + ) + .await; + assert_eq!(response.status(), 400); + let body = support::json(response).await; + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "invalid_request_error"); +} diff --git a/litellm-rust/crates/gateway-inference/tests/ocr.rs b/litellm-rust/crates/gateway-inference/tests/ocr.rs new file mode 100644 index 00000000000..3d5a2d22b09 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/ocr.rs @@ -0,0 +1,118 @@ +mod support; + +use axum::{body::Body, http::Request}; +use rstest::rstest; +use serde_json::{Value, json}; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header, method, path}, +}; + +const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; + +#[rstest] +#[case("/ocr", false)] +#[case("/v1/ocr", true)] +#[tokio::test] +async fn json_and_multipart_reach_ocr_with_the_deployment( + #[case] route: &str, + #[case] multipart: bool, +) { + let upstream = MockServer::start().await; + Mock::given(method("POST")).and(path("/v1/ocr")) + .and(header("authorization", "Bearer test-key")) + .and(body_json(json!({"model": "test-ocr", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"pages": [{"index": 0, "markdown": "recognized text"}]}))) + .expect(1).mount(&upstream).await; + let app = support::app("mistral/test-ocr", &upstream.uri()); + let response = if multipart { + let body = "--boundary\r\nContent-Disposition: form-data; name=\"model\"\r\n\r\npublic/model\r\n--boundary\r\nContent-Disposition: form-data; name=\"pages\"\r\n\r\n[0]\r\n--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.pdf\"\r\nContent-Type: application/pdf\r\n\r\nabc\r\n--boundary--\r\n"; + app.oneshot( + Request::post(route) + .header("content-type", "multipart/form-data; boundary=boundary") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap() + } else { + support::post(app, route, json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]})).await + }; + assert_eq!(response.status(), 200); + let body = support::json(response).await; + assert_eq!(body["pages"][0]["markdown"], "recognized text"); + assert_eq!(body["model"], "test-ocr"); +} + +#[rstest] +#[case(None, true)] +#[case(Some("litellm"), false)] +#[tokio::test] +async fn native_format_header_is_used_unless_the_body_overrides_it( + #[case] format: Option<&str>, + #[case] native: bool, +) { + let upstream = MockServer::start().await; + let payload = json!({"pages": [{"index": 0, "markdown": "text"}], "provider_only": true}); + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(payload.clone())) + .expect(1) + .mount(&upstream) + .await; + let body = json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "req_format": format}); + let response = support::app("mistral/test-ocr", &upstream.uri()) + .oneshot( + Request::post("/ocr") + .header("x-req-format", " Native ") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body = support::json(response).await; + if native { + assert_eq!(body, payload); + } else { + assert_eq!(body["object"], "ocr"); + assert_eq!( + body["pages"][0]["markdown"], + payload["pages"][0]["markdown"] + ); + } +} + +#[tokio::test] +async fn ocr_keeps_upstream_status_in_an_openai_error_envelope() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({"message": "busy"}))) + .expect(1) + .mount(&upstream) + .await; + let response = support::post(support::app("mistral/test-ocr", &upstream.uri()), "/ocr", + json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}})).await; + assert_eq!(response.status(), 429); + let body = support::json(response).await; + assert_eq!(body["error"]["code"], 429); + assert!(body["error"]["message"].as_str().unwrap().contains("busy")); +} + +#[rstest] +#[case(json!({"model": "public/model"}))] +#[case(json!({"model": "public/model", "document": "/etc/passwd"}))] +#[case(json!({"model": "missing", "document": {"type": "document_url", "document_url": DOCUMENT}}))] +#[tokio::test] +async fn invalid_ocr_requests_do_not_call_the_provider(#[case] body: Value) { + let upstream = MockServer::start().await; + let response = support::post( + support::app("mistral/test-ocr", &upstream.uri()), + "/ocr", + body, + ) + .await; + assert_eq!(response.status(), 400); + assert!(support::json(response).await["error"]["message"].is_string()); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs new file mode 100644 index 00000000000..5d3b06ccadb --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -0,0 +1,92 @@ +mod support; + +use rstest::rstest; +use serde_json::{Value, json}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, method}, +}; + +#[rstest] +#[case("/chat/completions", Some("public/model"))] +#[case("/v1/chat/completions", Some("public/model"))] +#[case("/engines/public/model/chat/completions", None)] +#[case("/openai/deployments/public/model/chat/completions", None)] +#[case("/openai/deployments/unused/chat/completions", Some("public/model"))] +#[tokio::test] +async fn chat_aliases_call_core_and_use_the_body_model_before_the_path( + #[case] route: &str, + #[case] model: Option<&str>, +) { + let upstream = MockServer::start().await; + let messages = json!([{"role": "user", "content": "hi"}]); + Mock::given(method("POST")) + .and(body_partial_json( + json!({"model": "test-model", "max_tokens": 16}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "model": "test-model", "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .expect(1) + .mount(&upstream) + .await; + let response = support::post( + support::app("anthropic/test-model", &upstream.uri()), + route, + json!({"model": model, "messages": messages, "max_tokens": 16}), + ) + .await; + assert_eq!(response.status(), 200); + assert_eq!( + support::json(response).await["choices"][0]["message"]["content"], + "hello" + ); +} + +#[rstest] +#[case("/responses")] +#[case("/v1/responses")] +#[case("/embeddings")] +#[case("/v1/embeddings")] +#[case("/completions")] +#[case("/v1/completions")] +#[case("/engines/public/model/embeddings")] +#[case("/openai/deployments/public/model/completions")] +#[tokio::test] +async fn unimplemented_routes_return_an_explicit_error(#[case] path: &str) { + let response = support::post( + support::app("anthropic/test-model", "http://127.0.0.1:1"), + path, + json!({}), + ) + .await; + assert_eq!(response.status(), 501); + assert!( + support::json(response).await["error"]["message"] + .as_str() + .unwrap() + .contains("not implemented") + ); +} + +#[rstest] +#[case("/audio/transcriptions")] +#[case("/v1/audio/transcriptions")] +#[tokio::test] +async fn transcription_aliases_reach_core_validation(#[case] path: &str) { + let response = support::post( + support::app("bedrock/test-model", "http://127.0.0.1:1"), + path, + json!({"model": "public/model", "audio": {"data": "YWJj", "format": "invalid"}}), + ) + .await; + assert_eq!(response.status(), 400); + let body: Value = support::json(response).await; + assert!( + body["error"]["message"] + .as_str() + .unwrap() + .contains("audio.format") + ); +} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs new file mode 100644 index 00000000000..d56489d28cd --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -0,0 +1,75 @@ +use std::{sync::Arc, time::Duration}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::Request, + response::Response, +}; +use futures_util::future::BoxFuture; +use litellm_core::resources::CoreResources; +use litellm_gateway_inference::{Deployment, Gateway, router}; +use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use litellm_secrets::{SecretValue, source::SecretSource}; +use serde_json::Value; +use tower::ServiceExt; + +struct NoSecrets; + +impl SecretSource for NoSecrets { + fn get_secret_str<'a>( + &'a self, + _: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async { Ok(None) }) + } +} + +pub fn app(model: &str, api_base: &str) -> Router { + let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); + let http = Resolution::from(&HttpSettings::default()).config; + let secrets = Arc::new(NoSecrets); + let resources = CoreResources::new(pool); + let ocr = resources + .ocr_client( + &http, + Default::default(), + OcrSettings::default(), + secrets.clone(), + ) + .unwrap(); + router(Arc::new(Gateway { + resources, + http, + secrets, + ocr, + models: [( + "public/model".into(), + Deployment { + model: model.into(), + api_base: Some(api_base.into()), + api_key: Some("test-key".into()), + timeout: Some(Duration::from_secs(5)), + ..Default::default() + }, + )] + .into_iter() + .collect(), + })) +} + +pub async fn post(app: Router, path: &str, body: Value) -> Response { + app.oneshot( + Request::post(path) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + +pub async fn json(response: Response) -> Value { + serde_json::from_slice(&to_bytes(response.into_body(), 1024 * 1024).await.unwrap()).unwrap() +} diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md new file mode 100644 index 00000000000..10a42b9ec3b --- /dev/null +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -0,0 +1,5 @@ +- Keep this crate a thin composition layer: mount endpoint routers and serve the supplied listener +- Server lifecycle and shared inbound middleware belong here, including client authentication, rate limiting, and request logging +- Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates; provider execution belongs to `core` and `llms` +- Inject shared state and infrastructure; avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups +- Test mounting and server lifecycle through public HTTP behavior; test endpoint semantics in the owning crate diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml new file mode 100644 index 00000000000..554186955b4 --- /dev/null +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-gateway" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +litellm-core.workspace = true +litellm-gateway-inference.workspace = true +litellm-gateway-auth.workspace = true +litellm-config.workspace = true +litellm-http.workspace = true +litellm-llms.workspace = true +litellm-secrets.workspace = true +tower-http = { version = "0.7.1", default-features = false, features = ["trace"] } +tracing.workspace = true +tokio.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs new file mode 100644 index 00000000000..3f67f923a5c --- /dev/null +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use axum::{Router, extract::Request}; +use tower_http::trace::{DefaultOnResponse, TraceLayer}; + +use litellm_config::Config; +use litellm_core::resources::CoreResources; +use litellm_gateway_auth::{Auth, RequireMasterKey}; +use litellm_gateway_inference::{Gateway, ModelList}; +use litellm_http::{ + ClientVariant, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, +}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use litellm_secrets::source::EnvironmentSecrets; + +pub fn build_inference(config: &Config) -> Result, litellm_http::Error> { + let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); + let http = Resolution::from(&HttpSettings::default()).config; + let client = pool.client(&http, ClientVariant::Provider)?; + let secrets = Arc::new(EnvironmentSecrets::python_compatible(client)); + let resources = CoreResources::new(pool); + let ocr = resources.ocr_client( + &http, + Default::default(), + OcrSettings::default(), + secrets.clone(), + )?; + + Ok(Arc::new(Gateway { + resources, + http, + secrets, + models: ModelList::from_model_list(&config.model_list), + ocr, + })) +} + +pub fn router(inference: Arc, config: &Config) -> Router { + let auth = Auth::from_config(config, inference.secrets.clone()); + litellm_gateway_inference::router(inference) + .route_layer(axum::middleware::from_extractor_with_state::< + RequireMasterKey, + _, + >(auth)) + .layer( + TraceLayer::new_for_http() + .make_span_with(|request: &Request| { + tracing::info_span!("request", method = %request.method(), path = request.uri().path()) + }) + .on_response(DefaultOnResponse::new().level(tracing::Level::INFO)), + ) +} diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs new file mode 100644 index 00000000000..bae711e16c7 --- /dev/null +++ b/litellm-rust/crates/gateway/src/main.rs @@ -0,0 +1,18 @@ +use std::error::Error; + +use litellm_config::Config; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let config_path = std::env::var("LITELLM_CONFIG").unwrap_or_else(|_| "config.yaml".into()); + let config = Config::load(config_path)?; + let inference = litellm_gateway::build_inference(&config)?; + let host = std::env::var("HOST").unwrap_or_else(|_| "0.0.0.0".into()); + let port = std::env::var("PORT") + .unwrap_or_else(|_| "4000".into()) + .parse::()?; + let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + + axum::serve(listener, litellm_gateway::router(inference, &config)).await?; + Ok(()) +} diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs new file mode 100644 index 00000000000..a19d8c9c4fe --- /dev/null +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -0,0 +1,90 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_config::Config; +use litellm_gateway_inference::{Error, Gateway}; +use litellm_http::ClientVariant; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; +use tokio::{net::TcpListener, sync::oneshot, time::timeout}; + +#[fixture] +fn inference() -> Arc { + litellm_gateway::build_inference(&Config::from_yaml("model_list: []").unwrap()).unwrap() +} + +#[rstest] +#[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)] +#[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)] +#[case::wrong_token("/v1/messages", Some("Bearer wrong"), Some("gateway-key"), 401)] +#[case::ocr("/ocr", None, Some("gateway-key"), 401)] +#[case::chat("/v1/chat/completions", None, Some("gateway-key"), 401)] +#[case::deployment( + "/openai/deployments/model/chat/completions", + None, + Some("gateway-key"), + 401 +)] +#[case::transcription("/audio/transcriptions", None, Some("gateway-key"), 401)] +#[case::unsupported_route("/responses", None, Some("gateway-key"), 401)] +#[case::unknown_path("/unknown", None, Some("gateway-key"), 404)] +#[case::unknown_path_unconfigured("/unknown", None, None, 404)] +#[case::unconfigured("/v1/messages", Some("Bearer gateway-key"), None, 500)] +#[tokio::test] +async fn authenticates_before_serving_mounted_inference_routes( + inference: Arc, + #[case] path: &str, + #[case] authorization: Option<&str>, + #[case] master_key: Option<&str>, + #[case] status: u16, +) { + let config = Config::from_yaml(&format!( + "model_list: []\ngeneral_settings:\n master_key: {}\n", + master_key.unwrap_or("null") + )) + .unwrap(); + let client = inference + .resources + .pool + .client(&inference.http, ClientVariant::Provider) + .unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (shutdown, stopped) = oneshot::channel(); + let server = tokio::spawn(async move { + axum::serve(listener, litellm_gateway::router(inference, &config)) + .with_graceful_shutdown(async move { + let _ = stopped.await; + }) + .await + }); + + let request = client + .post(format!("http://{address}{path}")) + .timeout(Duration::from_secs(5)) + .header("x-request-id", "gateway-test") + .json(&json!({"model": "unconfigured-model"})); + let request = match authorization { + Some(value) => request.header("authorization", value), + None => request, + }; + let response = request.send().await.unwrap(); + assert_eq!(response.status().as_u16(), status); + if status == 400 { + let expected = Error::UnknownModel("unconfigured-model".into()); + assert_eq!( + response.json::().await.unwrap(), + expected.body(Some("gateway-test")) + ); + } else { + let text = response.text().await.unwrap(); + assert!(!text.contains("gateway-key")); + assert!(!text.contains("unconfigured-model")); + } + + shutdown.send(()).unwrap(); + timeout(Duration::from_secs(5), server) + .await + .unwrap() + .unwrap() + .unwrap(); +} diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml new file mode 100644 index 00000000000..cc6972a066e --- /dev/null +++ b/litellm-rust/crates/router/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-router" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-config.workspace = true +litellm-core.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/router/README.md b/litellm-rust/crates/router/README.md new file mode 100644 index 00000000000..ed5b2966fe8 --- /dev/null +++ b/litellm-rust/crates/router/README.md @@ -0,0 +1,5 @@ +`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments. Programmatic callers can collect `(String, Deployment)` entries into a `Router` + +Lookup is exact and returns `None` for an unknown name. This extraction preserves the gateway's existing behavior: the last entry wins when public names repeat. Multiple deployments per model group, routing strategies, retries, cooldowns, and fallbacks are not implemented yet + +The router owns deployment configuration and selection. The gateway handles HTTP errors and responses, while `core` executes provider calls and resolves credentials diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs new file mode 100644 index 00000000000..4904bf4ffdd --- /dev/null +++ b/litellm-rust/crates/router/src/deployment.rs @@ -0,0 +1,13 @@ +use std::time::Duration; + +use litellm_core::messages::types::MessagesShaping; + +#[derive(Clone, Debug, Default)] +pub struct Deployment { + pub model: String, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub timeout: Option, + pub shaping: MessagesShaping, +} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs new file mode 100644 index 00000000000..da33bfb04bd --- /dev/null +++ b/litellm-rust/crates/router/src/lib.rs @@ -0,0 +1,44 @@ +mod deployment; + +use std::collections::HashMap; + +use litellm_config::Model; + +pub use deployment::Deployment; + +#[derive(Clone, Debug, Default)] +pub struct Router(HashMap); + +impl Router { + pub fn from_model_list(model_list: &[Model]) -> Self { + model_list + .iter() + .map(|model| { + ( + model.model_name.clone(), + Deployment { + model: model.litellm_params.model.clone(), + api_key: model + .litellm_params + .api_key + .as_ref() + .map(|value| value.expose().to_string()), + api_base: model.litellm_params.api_base.clone(), + custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), + ..Deployment::default() + }, + ) + }) + .collect() + } + + pub fn get(&self, model_name: &str) -> Option<&Deployment> { + self.0.get(model_name) + } +} + +impl FromIterator<(String, Deployment)> for Router { + fn from_iter>(entries: I) -> Self { + Self(entries.into_iter().collect()) + } +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs new file mode 100644 index 00000000000..3cbd316cab0 --- /dev/null +++ b/litellm-rust/crates/router/tests/router.rs @@ -0,0 +1,94 @@ +use std::time::Duration; + +use litellm_config::Config; +use litellm_core::messages::types::MessagesShaping; +use litellm_router::{Deployment, Router}; +use rstest::rstest; + +#[rstest] +#[case::minimal("")] +#[case::configured( + "api_key: test-key\n api_base: https://provider.example/v1\n custom_llm_provider: test-provider" +)] +#[case::secret_reference("api_key: os.environ/ROUTER_TEST_API_KEY")] +fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}" + )) + .unwrap(); + let router = Router::from_model_list(&config.model_list); + let deployment = router.get(&config.model_list[0].model_name).unwrap(); + let params = &config.model_list[0].litellm_params; + + assert_eq!(deployment.model, params.model); + assert_eq!( + deployment.api_key.as_deref(), + params.api_key.as_ref().map(|key| key.expose()) + ); + assert_eq!(deployment.api_base, params.api_base); + assert_eq!(deployment.custom_llm_provider, params.custom_llm_provider); + assert_eq!(deployment.timeout, Deployment::default().timeout); + assert_eq!(deployment.shaping, Deployment::default().shaping); +} + +#[rstest] +#[case::first("public-a", Some("provider/a"))] +#[case::second("public-b", Some("provider/b"))] +#[case::unknown("missing", None)] +#[case::provider_name_is_not_an_alias("provider/a", None)] +#[case::case_sensitive("PUBLIC-A", None)] +fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) { + let config = Config::from_yaml( + "model_list: + - model_name: public-a + litellm_params: + model: provider/a + - model_name: public-b + litellm_params: + model: provider/b", + ) + .unwrap(); + let router = Router::from_model_list(&config.model_list); + + assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected); +} + +#[rstest] +fn empty_configuration_has_no_deployment() { + let config = Config::from_yaml("model_list: []").unwrap(); + + assert!( + Router::from_model_list(&config.model_list) + .get("") + .is_none() + ); + assert!(Router::default().get("unknown").is_none()); +} + +#[rstest] +fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { + let deployment = Deployment { + model: "provider/selected".into(), + api_key: Some("test-key".into()), + api_base: Some("https://provider.example/v1".into()), + custom_llm_provider: Some("test-provider".into()), + timeout: Some(Duration::from_secs(7)), + shaping: MessagesShaping { + drop_params: true, + additional_drop_params: vec!["metadata.test".into()], + ..Default::default() + }, + }; + let router = Router::from_iter([ + ("public-model".into(), Deployment::default()), + ("public-model".into(), deployment.clone()), + ]); + let selected = router.get("public-model").unwrap(); + + assert_eq!(selected.model, deployment.model); + assert_eq!(selected.api_key, deployment.api_key); + assert_eq!(selected.api_base, deployment.api_base); + assert_eq!(selected.custom_llm_provider, deployment.custom_llm_provider); + assert_eq!(selected.timeout, deployment.timeout); + assert_eq!(selected.shaping, deployment.shaping); +} From 4d7aa89fa3b975eb729b5fc6fcfdca36c97cb64d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:15:16 -0700 Subject: [PATCH 53/58] fix(cost-map): remove duplicate openrouter/perceptron/perceptron-mk1.5 entry (#43273) Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 19 ------------------- model_prices_and_context_window.json | 19 ------------------- 2 files changed, 38 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ec31039fee0..d26ca15450b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -77560,24 +77560,5 @@ "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, "supports_tool_choice": false - }, - "openrouter/perceptron/perceptron-mk1.5": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "openrouter", - "max_input_tokens": 36864, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://openrouter.ai/api/v1/models", - "supports_audio_input": false, - "supports_function_calling": false, - "supports_pdf_input": false, - "supports_prompt_caching": false, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": false, - "supports_vision": true, - "supports_web_search": false } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ec31039fee0..d26ca15450b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -77560,24 +77560,5 @@ "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, "supports_tool_choice": false - }, - "openrouter/perceptron/perceptron-mk1.5": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "openrouter", - "max_input_tokens": 36864, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://openrouter.ai/api/v1/models", - "supports_audio_input": false, - "supports_function_calling": false, - "supports_pdf_input": false, - "supports_prompt_caching": false, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": false, - "supports_vision": true, - "supports_web_search": false } } From a942c343aba0f6bed74438218ba3dbd1ff2e1ee7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 23:40:55 -0700 Subject: [PATCH 54/58] fix(e2e): resolve the blank-S3 gateway repo root from the litellm package location (#42911) Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/batches/bedrock_env_gateway.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/e2e/batches/bedrock_env_gateway.py b/tests/e2e/batches/bedrock_env_gateway.py index fb3ec60c87c..0b1d840eb30 100644 --- a/tests/e2e/batches/bedrock_env_gateway.py +++ b/tests/e2e/batches/bedrock_env_gateway.py @@ -8,6 +8,7 @@ a batch create through it proves blank means unset, not an empty string. from __future__ import annotations +import importlib.util import os import shutil import socket @@ -28,7 +29,13 @@ from pydantic import TypeAdapter STARTUP_TIMEOUT_SECONDS: Final = 240 LOG_TAIL_BYTES: Final = 4000 -REPO_ROOT: Final = Path(__file__).resolve().parents[3] + + +def litellm_root() -> Path: + spec: Final = importlib.util.find_spec("litellm") + assert spec is not None and spec.origin is not None, "litellm must be importable to boot the blank-S3-env gateway" + return Path(spec.origin).resolve().parents[1] + _CONFIG_YAML: Final = """model_list: - model_name: bedrock-blank-s3-batch @@ -68,6 +75,7 @@ class BedrockEnvGateway: @classmethod def start(cls) -> BedrockEnvGateway: assert os.environ.get("DATABASE_URL"), "DATABASE_URL is required for the blank-S3-env gateway" + root: Final = litellm_root() port: Final = available_port() base_url: Final = f"http://127.0.0.1:{port}" master_key: Final = f"sk-e2e-blank-s3-{unique_marker()}" @@ -79,7 +87,7 @@ class BedrockEnvGateway: "DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_MASTER_KEY": master_key, "STORE_MODEL_IN_DB": "False", - "PYTHONPATH": str(REPO_ROOT), + "PYTHONPATH": str(root), "AWS_S3_ENCRYPTION_KEY_ID": "", "AWS_S3_BUCKET_OWNER": "", } @@ -113,13 +121,11 @@ class BedrockEnvGateway: stdout=log, stderr=log, start_new_session=True, - cwd=REPO_ROOT, + cwd=root, ) deadline: Final = time.monotonic() + STARTUP_TIMEOUT_SECONDS while time.monotonic() < deadline: - assert gateway._child.poll() is None, ( - f"blank-S3-env gateway exited early; log tail:\n{gateway.log_tail()}" - ) + assert gateway._child.poll() is None, f"blank-S3-env gateway exited early; log tail:\n{gateway.log_tail()}" result = gateway.proxy.transport.probe("/health/liveliness", params=NoBody()) if result.status_code == 200: return gateway From dd63637322da0fe2055684fa80d73490299f0771 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 26 Sep 2026 00:15:56 -0700 Subject: [PATCH 55/58] test(integration): run the Langfuse DB-callback test on its own scratch database (#43288) * test(integration): run the Langfuse DB-callback test on its own scratch database The test from #43282 wrote success_callback=langfuse and the LANGFUSE_* env into the shared integration LiteLLM_Config. The suite's long-running gateway reloads that table and only ever adds callbacks, so it kept exporting to the test's closed Langfuse fake for the rest of the shard even after the rows were restored. The owned proxy now gets a scratch database, which also removes the snapshot/restore code. scratch_database moves into _support/database.py so test_cache_and_quota and this test share one copy, and the stock-config guard now checks the callback settings instead of the raw YAML text. * test(integration): include failure_callback in the stock-config Langfuse guard --- tests/integration/_support/database.py | 16 ++ .../observability/test_langfuse_delivery.py | 154 ++++++++---------- .../integration/spend/test_cache_and_quota.py | 20 +-- 3 files changed, 87 insertions(+), 103 deletions(-) diff --git a/tests/integration/_support/database.py b/tests/integration/_support/database.py index 461cdbda1ee..3ad4f205e12 100644 --- a/tests/integration/_support/database.py +++ b/tests/integration/_support/database.py @@ -1,7 +1,12 @@ import os +import uuid +from collections.abc import Generator +from contextlib import contextmanager from typing import Final, LiteralString +from urllib.parse import urlsplit, urlunsplit import psycopg +from psycopg import sql from psycopg.rows import dict_row from pydantic import JsonValue, TypeAdapter @@ -19,3 +24,14 @@ def read_rows( def write_rows(query: LiteralString, parameters: tuple[str, ...], *, database_url: str | None = None) -> None: with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: connection.execute(query, parameters) + + +@contextmanager +def scratch_database() -> Generator[str]: + name: Final = f"integration_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) + try: + yield urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{name}")) + finally: + admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 9c0e3e407cb..13be5a95887 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -2,13 +2,13 @@ import base64 import json import time import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from pathlib import Path from typing import Final import yaml from integration._support.client import Gateway, eventually, object_value, string_value -from integration._support.database import read_rows, write_rows +from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest @@ -24,6 +24,7 @@ PROMPTS_PATH: Final = "/api/public/v2/prompts/" STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml") CONFIG_SECTIONS: Final = ("litellm_settings", "environment_variables") LANGFUSE_ENVIRONMENT: Final = ("LANGFUSE_HOST", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY") +INHERITED_ENVIRONMENT: Final = (*LANGFUSE_ENVIRONMENT, "DATABASE_URL_READ_REPLICA") _PROXY_CONFIG: Final = TypeAdapter(dict[str, object]) _SETTINGS: Final = TypeAdapter(dict[str, object]) @@ -88,32 +89,14 @@ def _langfuse_environment(langfuse: Wire) -> dict[str, str]: } -def _config_rows() -> list[dict[str, JsonValue]]: +def _config_rows(database_url: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT param_name, param_value FROM "LiteLLM_Config" WHERE param_name IN (%s, %s) ORDER BY param_name', CONFIG_SECTIONS, + database_url=database_url, ) -def _restore_config_rows(snapshot: Sequence[Mapping[str, JsonValue]]) -> None: - saved: Final = {string_value(row["param_name"]): row["param_value"] for row in snapshot} - for section in CONFIG_SECTIONS: - if section not in saved: - write_rows('DELETE FROM "LiteLLM_Config" WHERE param_name = %s', (section,)) - elif saved[section] is None: - write_rows( - 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, NULL) ' - "ON CONFLICT (param_name) DO UPDATE SET param_value = NULL", - (section,), - ) - else: - write_rows( - 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s::jsonb) ' - "ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value", - (section, json.dumps(saved[section])), - ) - - def _attribute(entries: Sequence[KeyValue], key: str) -> str | list[str] | None: for entry in entries: if entry.key != key: @@ -227,7 +210,12 @@ def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_g provider_secret: Final = "synthetic-provider-secret-" + marker public_key: Final = "pk-lf-db-" + marker secret_key: Final = "sk-lf-db-" + marker - assert "langfuse" not in STOCK_CONFIG.read_text() + stock_settings: Final = _SETTINGS.validate_python( + _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text()))["litellm_settings"] + ) + assert "langfuse" not in json.dumps( + [stock_settings.get(key) for key in ("callbacks", "success_callback", "failure_callback")] + ) def upstream(request: Request) -> Reply: assert request.headers["authorization"] == f"Bearer {provider_secret}" @@ -238,73 +226,69 @@ def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_g return _projects() return Reply(body=b"", content_type="application/x-protobuf") - snapshot: Final = _config_rows() - try: - with ( - wire_server(upstream) as provider, - wire_server(langfuse) as destination, - owned_proxy( - gateway, - tmp_path, - {"LANGFUSE_FLUSH_INTERVAL": "1"}, - remove_environment=LANGFUSE_ENVIRONMENT, - ) as candidate, - candidate.scenario() as scenario, - ): - candidate.post( - "/config/update", - { - "litellm_settings": {"success_callback": ["langfuse"]}, - "environment_variables": { - "LANGFUSE_HOST": destination.url, - "LANGFUSE_PUBLIC_KEY": public_key, - "LANGFUSE_SECRET_KEY": secret_key, - }, + with ( + scratch_database() as scratch_url, + wire_server(upstream) as provider, + wire_server(langfuse) as destination, + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": scratch_url, "LANGFUSE_FLUSH_INTERVAL": "1"}, + remove_environment=INHERITED_ENVIRONMENT, + ) as candidate, + candidate.scenario() as scenario, + ): + candidate.post( + "/config/update", + { + "litellm_settings": {"success_callback": ["langfuse"]}, + "environment_variables": { + "LANGFUSE_HOST": destination.url, + "LANGFUSE_PUBLIC_KEY": public_key, + "LANGFUSE_SECRET_KEY": secret_key, }, - ) - model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) - body: Final = candidate.post( - "/v1/chat/completions", - { - "model": model, - "messages": [{"role": "user", "content": marker + "-question"}], - "metadata": {"generation_name": marker}, - "cache": {"no-cache": True}, - }, - ) - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones + }, + ) + model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) + body: Final = candidate.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker + "-question"}], + "metadata": {"generation_name": marker}, + "cache": {"no-cache": True}, + }, + ) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones - def exported() -> tuple[Span, ...]: - received.extend(destination.drain()) - return tuple(span for span in _spans(received) if span.name == marker) + def exported() -> tuple[Span, ...]: + received.extend(destination.drain()) + return tuple(span for span in _spans(received) if span.name == marker) - spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) - posts: Final = tuple(request for request in received if request.method == "POST") - assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] - basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() - for request in posts: - assert request.headers["authorization"] == basic - assert request.headers["content-type"] == "application/x-protobuf" - assert request.headers["x-langfuse-ingestion-version"] == "4" - assert provider_secret.encode() not in request.body - assert candidate.key.encode() not in request.body + spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) + posts: Final = tuple(request for request in received if request.method == "POST") + assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] + basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() + for request in posts: + assert request.headers["authorization"] == basic + assert request.headers["content-type"] == "application/x-protobuf" + assert request.headers["x-langfuse-ingestion-version"] == "4" + assert provider_secret.encode() not in request.body + assert candidate.key.encode() not in request.body - attributes: Final = spans[0].attributes - assert _attribute(attributes, "langfuse.observation.type") == "generation" - assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) - assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) - assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) + attributes: Final = spans[0].attributes + assert _attribute(attributes, "langfuse.observation.type") == "generation" + assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) + assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) + assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) - stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows()} - callbacks: Final = TypeAdapter(list[str]).validate_python( - object_value(stored["litellm_settings"]).get("success_callback") or [] - ) - assert "langfuse" in callbacks, stored - assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored - assert secret_key not in json.dumps(stored["environment_variables"]), stored - finally: - _restore_config_rows(snapshot) - assert _config_rows() == snapshot + stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows(scratch_url)} + callbacks: Final = TypeAdapter(list[str]).validate_python( + object_value(stored["litellm_settings"]).get("success_callback") or [] + ) + assert "langfuse" in callbacks, stored + assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored + assert secret_key not in json.dumps(stored["environment_variables"]), stored def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_headers_off_the_client( diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1c2b5551855..1cc03838f8d 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -1,27 +1,22 @@ import json -import os import threading import uuid -from collections.abc import Generator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack from hashlib import sha256 from pathlib import Path from typing import Final -from urllib.parse import urlsplit, urlunsplit import httpx -import psycopg import pytest from hypothesis import strategies as st from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test from integration._support.client import Gateway, eventually, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.database_relay import database_relay from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server -from psycopg import sql @pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting") @@ -225,17 +220,6 @@ def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gat RESET_SWEEP_QUERY: Final = b'"LiteLLM_VerificationToken"."budget_reset_at" < $' -@contextmanager -def scratch_database() -> Generator[str]: - name: Final = f"integration_{uuid.uuid4().hex}" - with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: - admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) - try: - yield urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{name}")) - finally: - admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) - - @pytest.mark.covers("quota_management.budget.key.scheduled_reset_survives_transient_db_outage") @pytest.mark.timeout(300) def test_scheduled_budget_reset_reconnects_after_db_transport_failure_and_unblocks_key( From 90873c46de441a80e65203447eba6a241dcf7106 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 08:04:19 +0000 Subject: [PATCH 56/58] refactor(rust): expand logging and test coverage across gateway and Anthropic messages (#43295) * refactor(rust): prepare inference and auth foundations * fix(rust): keep textract operations parsing from kebab-case model names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * done * refactor(types): derive Anthropic beta string conversions with Strum * fix(anthropic): report missing max_tokens as a missing field * refactor(rust): type Anthropic messages headers and auth after the Python layout Delete anthropic/messages/headers.rs. Its OAuth handling, credential ladder and beta merging move to anthropic/common_utils.rs where Python keeps them (optionally_handle_anthropic_oauth, get_auth_header, _merge_beta_headers), and the feature beta injection becomes update_headers_with_anthropic_beta on the messages config, as in Python. The BaseAnthropicMessagesConfig impl is unchanged apart from the bodies of validate_environment and request_headers Beta values are now the AnthropicBeta enum and BetaSet, which sort, dedupe and comma-join by construction. Request params gain typed speed, tools and context_management through Recognized, so the beta logic matches on enums instead of string-comparing JSON. OauthToken parses the sk-ant-oat token once and the chat config shares that detection instead of its own copy Case-insensitive header helpers move next to has_header in litellm-http. One deliberate divergence: a Bearer-prefixed OAuth key configured through api_key or ANTHROPIC_API_KEY is sent with a single Bearer scheme, where Python would emit "Bearer Bearer" Co-Authored-By: Claude Fable 5.1 * done * fix(rust): repair test compilation and clippy failures resolve auth before building the outbound request in prepare tests, give the host hook tests their own error type, and drop the disallowed reqwest client and err().expect() from core tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 --- litellm-rust/Cargo.lock | 30 +- litellm-rust/crates/auth-types/src/http.rs | 2 +- litellm-rust/crates/core/AGENTS.md | 6 +- litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/chat_completions/handler.rs | 219 +++++- .../crates/core/src/chat_completions/mod.rs | 10 +- .../core/src/chat_completions/prepare.rs | 51 +- .../crates/core/src/chat_completions/types.rs | 7 + .../crates/core/src/messages/handler.rs | 210 +++++- litellm-rust/crates/core/src/messages/mod.rs | 38 +- .../crates/core/src/messages/prepare.rs | 47 +- .../crates/core/src/messages/route.rs | 180 +---- .../crates/core/src/messages/types.rs | 51 +- .../crates/core/tests/messages/main.rs | 5 +- .../crates/core/tests/messages/request.rs | 44 +- .../crates/core/tests/messages/response.rs | 8 +- .../crates/core/tests/messages/stream.rs | 207 +++++- .../crates/gateway-inference/Cargo.toml | 3 +- .../gateway-inference/src/messages/host.rs | 69 -- .../gateway-inference/src/messages/mod.rs | 55 +- .../gateway-inference/tests/messages.rs | 55 ++ litellm-rust/crates/gateway/Cargo.toml | 8 +- litellm-rust/crates/gateway/src/lib.rs | 142 +++- litellm-rust/crates/gateway/src/main.rs | 41 +- litellm-rust/crates/gateway/tests/server.rs | 55 +- litellm-rust/crates/host/src/hooks.rs | 145 ++++ litellm-rust/crates/host/src/lib.rs | 1 + litellm-rust/crates/http/src/request.rs | 38 + litellm-rust/crates/litellm/Cargo.toml | 1 + .../src/anthropic/batches/transformation.rs | 2 +- .../llms/src/anthropic/chat/transformation.rs | 21 +- .../crates/llms/src/anthropic/common_utils.rs | 627 ++++++++++++++-- .../llms/src/anthropic/messages/headers.rs | 677 ------------------ .../crates/llms/src/anthropic/messages/mod.rs | 1 - .../src/anthropic/messages/transformation.rs | 634 ++++++++++------ .../anthropic/messages_transformation.rs | 15 +- litellm-rust/crates/llms/src/base_llm/auth.rs | 4 +- .../python-bridge/src/routes/messages/host.rs | 11 +- litellm-rust/crates/router/src/deployment.rs | 2 +- litellm-rust/crates/router/tests/router.rs | 2 +- litellm-rust/crates/secrets/src/native.rs | 12 +- litellm-rust/crates/tracing/Cargo.toml | 1 + litellm-rust/crates/tracing/src/lib.rs | 30 + litellm-rust/crates/tracing/tests/logging.rs | 19 +- .../crates/types/src/llms/anthropic.rs | 240 +++++++ .../anthropic_messages/anthropic_request.rs | 205 +++++- litellm-rust/crates/types/src/llms/mod.rs | 1 + litellm/rust_bridge/catalog.py | 2 +- tests/unit/rust_bridge/test_catalog.py | 14 +- 49 files changed, 2803 insertions(+), 1446 deletions(-) delete mode 100644 litellm-rust/crates/gateway-inference/src/messages/host.rs create mode 100644 litellm-rust/crates/host/src/hooks.rs delete mode 100644 litellm-rust/crates/llms/src/anthropic/messages/headers.rs create mode 100644 litellm-rust/crates/types/src/llms/anthropic.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b9d363e72ca..d67623feffd 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3138,6 +3138,7 @@ dependencies = [ "litellm-http", "litellm-llms", "litellm-secrets", + "litellm-tracing", "litellm-types", "mime_guess", "moka", @@ -3215,6 +3216,8 @@ name = "litellm-gateway" version = "0.1.0" dependencies = [ "axum", + "futures-util", + "http-body-util", "litellm-config", "litellm-core", "litellm-gateway-auth", @@ -3222,11 +3225,13 @@ dependencies = [ "litellm-http", "litellm-llms", "litellm-secrets", + "litellm-tracing", "rstest", "serde_json", "tokio", - "tower-http 0.7.1", + "tower", "tracing", + "uuid", ] [[package]] @@ -3256,7 +3261,6 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-core", - "litellm-host", "litellm-http", "litellm-llms", "litellm-router", @@ -3677,6 +3681,7 @@ dependencies = [ name = "litellm-tracing" version = "0.1.0" dependencies = [ + "base64 0.22.1", "fancy-regex 0.19.2", "percent-encoding", "rstest", @@ -4880,7 +4885,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http 0.6.11", + "tower-http", "tower-service", "url", "wasm-bindgen", @@ -4922,7 +4927,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http 0.6.11", + "tower-http", "tower-service", "url", "wasm-bindgen", @@ -6163,23 +6168,6 @@ dependencies = [ "url", ] -[[package]] -name = "tower-http" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" -dependencies = [ - "bitflags 2.13.1", - "bytes", - "http 1.4.2", - "http-body 1.1.0", - "percent-encoding", - "pin-project-lite", - "tower-layer", - "tower-service", - "tracing", -] - [[package]] name = "tower-layer" version = "0.3.3" diff --git a/litellm-rust/crates/auth-types/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs index f3c5254b60e..1519769521f 100644 --- a/litellm-rust/crates/auth-types/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -7,7 +7,7 @@ pub enum CredentialPlacement { } impl CredentialPlacement { - pub fn header_name(self) -> &'static str { + pub const fn header_name(self) -> &'static str { match self { Self::Bearer => "Authorization", Self::Header(name) => name, diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index a40265729c7..20da74789bb 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,4 +1,6 @@ -litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. +litellm-core is the LiteLLM SDK in Rust. Each top-level call is a module under `src//` exposing a public entrypoint named after the route. `messages::messages()` returns `MessagesResponse::Message` for a completed response or `MessagesResponse::Stream { headers, chunks }` when the request sets `stream: true`. The chunks are Anthropic SSE bytes in a `Stream>`. Dropping the stream cancels the call. The Python bridge drives `messages::route::messages_machine()` instead, because Python has to answer the call's operations on its own thread; the gateway and the Rust SDK call the plain entrypoint + +A route module has the same five pieces, in the order Python runs them. `types.rs` holds the call, the provider request, and the response. `prepare.rs` resolves the provider and credentials and shapes the request (Python's `validate_environment`, `get_complete_url`, `transform_request`). `handler.rs` resolves auth, offers the wire request to `litellm_host::hooks::RouteHooks::before_send`, sends it, reports the raw response through `emit`, and normalizes the response or stream (`pre_call`, `post`, `post_call`, `transform_response`). `mod.rs` exposes the entrypoint that runs prepare then handler with no hooks (`()`). `route.rs`, where a host needs it, wraps the same two calls in a `CallMachine` whose `HostChannel` is the hooks, and pumps a stream through `open` and `deliver`. A handler takes `&impl RouteHooks` and never a `HostChannel` directly, so it runs without a coroutine. Keep provider transport and transformation details out of the machine driver ## Crate layering @@ -10,7 +12,7 @@ Each crate mirrors one top-level Python package, so a Rust path reads as its Pyt - `litellm-llms` mirrors `litellm/llms/`: `base_llm//transformation.rs`, `//transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler) - `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks -A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate +A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::hooks::RouteHooks`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate ## Error placement diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 0051631f40d..6904dcc023c 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -17,6 +17,7 @@ litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-tracing.workspace = true moka.workspace = true mime_guess = "2.0.5" rand.workspace = true diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 4a1cf7e193e..740db2edefb 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,23 +1,68 @@ use std::time::Duration; +use litellm_auth::AuthServices; +use litellm_host::{ + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, + hooks::RouteHooks, +}; use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body}; -use litellm_llms::base_llm::{auth::resolve_auth, chat::transformation::ProviderChatResponseData}; +use litellm_llms::base_llm::{ + auth::{Authenticated, resolve_auth}, + chat::transformation::ProviderChatResponseData, +}; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; -use super::{Error, prepare::prepare_provider_request}; +use super::Error; use crate::{ - chat_completions::types::{ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest}, + chat_completions::types::ProviderChatCompletionsRequest, constants::CHAT_COMPLETIONS_TIMEOUT_SECS, }; -pub(super) async fn execute_chat_completions_provider_call( +pub(super) async fn execute( http: &Client, - auth: &litellm_auth::AuthServices, - request: ResolvedChatCompletionsRequest<'_>, + auth: &AuthServices, + request: ProviderChatCompletionsRequest, + hooks: &impl RouteHooks, ) -> Result { - let request = prepare_provider_request(request)?; - let outbound = outbound_request(auth, &request).await?; + let ProviderChatCompletionsRequest { + model, + custom_llm_provider, + config, + url, + body, + optional_params, + environment, + timeout, + api_key, + } = request; + let context = RequestContext { + model: model.clone(), + custom_llm_provider, + optional_params: Value::Object(optional_params), + secret_fields: Vec::new(), + api_key, + }; + let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?; + let wire = hooks + .before_send( + WireRequest { + url, + headers: authenticated.headers, + body, + }, + context, + ) + .await?; + let outbound = outbound_request( + Authenticated { + headers: wire.headers, + signer: authenticated.signer, + }, + wire.url, + &wire.body, + timeout, + )?; let response = outbound.send(http).await.map_err(|err| { // Failing to establish the connection means the request never went out, @@ -41,13 +86,17 @@ pub(super) async fn execute_chat_completions_provider_call( body: truncate_error_body(&text), })); } + hooks + .emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; let body: Value = serde_json::from_str(&text).map_err(|err| { Error::InvalidResponse(format!("invalid chat completions response JSON: {err}")) })?; - request - .config - .transform_response(&request.model, ProviderChatResponseData { body }) + config + .transform_response(&model, ProviderChatResponseData { body }) .map_err(Error::from) .map_err(as_response_error) } @@ -69,21 +118,17 @@ pub(super) fn as_response_error(err: Error) -> Error { } } -pub(super) async fn outbound_request( - auth: &litellm_auth::AuthServices, - request: &ProviderChatCompletionsRequest, +pub(super) fn outbound_request( + authenticated: Authenticated, + url: String, + body: &Value, + timeout: Option, ) -> Result { - let env_lookup = |key: &str| std::env::var(key).ok(); - let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?; crate::outbound::outbound_request( authenticated, - request.url.clone(), - &request.body, - Some( - request - .timeout - .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), - ), + url, + body, + Some(timeout.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))), ) .map_err(|error| match error { // Python drops the caller's copy and prefers a forwarded Authorization @@ -97,7 +142,133 @@ pub(super) async fn outbound_request( #[cfg(test)] mod tests { - use super::{Error, as_response_error}; + use std::sync::Mutex; + + use rstest::rstest; + use serde_json::json; + use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; + + use super::*; + use crate::chat_completions::{ + prepare::{prepare_provider_request, resolve_request}, + types::ChatCompletionsRequest, + }; + + const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + /// Rewrites the outgoing request and records what the call reports back. + #[derive(Default)] + struct RecordingHooks { + contexts: Mutex>, + raw: Mutex>, + } + + impl RouteHooks for RecordingHooks { + async fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + self.contexts.lock().unwrap().push(context); + let mut body = wire.body; + body["system"] = json!("added by the host"); + Ok(WireRequest { + headers: wire + .headers + .into_iter() + .chain([("x-host".to_string(), "seen".to_string())]) + .collect(), + body, + ..wire + }) + } + + async fn emit(&self, event: MachineEvent) -> Result<(), Error> { + let MachineEvent::ResponseReceived { raw } = event; + self.raw.lock().unwrap().push(raw.body); + Ok(()) + } + } + + fn prepared(api_base: &str) -> ProviderChatCompletionsRequest { + prepare_provider_request( + resolve_request(ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages: json!([{"role": "user", "content": "hi"}]), + optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(), + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: None, + }) + .unwrap(), + ) + .unwrap() + } + + #[rstest] + #[tokio::test] + async fn the_hooks_rewrite_the_wire_request_and_see_the_raw_response() { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with( + ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json"), + ) + .mount(&upstream) + .await; + let hooks = RecordingHooks::default(); + + execute( + &Client::plain_for_test(), + &AuthServices::default(), + prepared(&upstream.uri()), + &hooks, + ) + .await + .expect("chat completions call succeeds"); + + let [request] = <[Request; 1]>::try_from(upstream.received_requests().await.unwrap()) + .unwrap_or_else(|requests| panic!("one request, saw {}", requests.len())); + let sent: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(sent["system"], "added by the host"); + assert_eq!(request.headers["x-host"], "seen"); + assert_eq!(request.headers["x-api-key"], "sk-test"); + let [context] = <[RequestContext; 1]>::try_from(hooks.contexts.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len())); + assert_eq!( + (context.model.as_str(), context.custom_llm_provider.as_str()), + ("claude-sonnet-4-5", "anthropic") + ); + assert_eq!(context.optional_params, json!({"max_tokens": 16})); + assert_eq!(hooks.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]); + } + + #[rstest] + #[tokio::test] + async fn an_upstream_failure_is_not_reported_as_a_received_response() { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(500).set_body_string("boom")) + .mount(&upstream) + .await; + let hooks = RecordingHooks::default(); + + let error = execute( + &Client::plain_for_test(), + &AuthServices::default(), + prepared(&upstream.uri()), + &hooks, + ) + .await + .expect_err("the upstream failure fails the call"); + + assert!(matches!( + error, + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) + )); + assert!(hooks.raw.into_inner().unwrap().is_empty()); + } #[test] fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index d7003b5d22a..dc4e80b816a 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -11,10 +11,9 @@ pub use crate::error::RouteError as Error; mod common_utils; pub(crate) mod handler; mod prepare; -use handler::execute_chat_completions_provider_call; use litellm_http::{ClientVariant, HttpClientConfig}; use litellm_types::utils::ChatCompletionsResponse; -use prepare::{parse_messages, resolve_provider_config, resolve_request}; +use prepare::{parse_messages, prepare_provider_request, resolve_provider_config, resolve_request}; use serde_json::{Map, Value}; use crate::chat_completions::types::ChatCompletionsRequest; @@ -24,9 +23,9 @@ pub async fn chat_completions( config: &HttpClientConfig, request: ChatCompletionsRequest<'_>, ) -> Result { - let request = resolve_request(request)?; let http = resources.pool.client(config, ClientVariant::Provider)?; - execute_chat_completions_provider_call(&http, &resources.auth, request).await + let request = prepare_provider_request(resolve_request(request)?)?; + handler::execute(&http, &resources.auth, request, &()).await } /// Whether the core would accept this request, without resolving credentials or @@ -41,9 +40,10 @@ pub fn chat_completions_decline_reason( messages: Value, optional_params: &Map, ) -> Option<&'static str> { - let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else { + let Ok(resolved) = resolve_provider_config(model, custom_llm_provider) else { return Some("provider is not on the rust chat completions path"); }; + let config = resolved.config; let Ok(messages) = parse_messages(messages) else { return Some("unreadable message list"); }; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 8b091d4dd6c..6ead713eec1 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_llms::base_llm::{ auth::{ValidatedEnvironment, with_default_headers}, @@ -14,10 +15,16 @@ use crate::chat_completions::types::{ ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest, }; +pub(super) struct ResolvedProvider { + pub(super) model: String, + pub(super) custom_llm_provider: String, + pub(super) config: &'static dyn BaseConfig, +} + pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> Result<(String, &'static dyn BaseConfig), Error> { +) -> Result { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -32,7 +39,11 @@ pub(super) fn resolve_provider_config<'a>( })?; let config = chat_completions_provider_config(provider_info.custom_llm_provider) .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; - Ok((provider_info.model.to_string(), config)) + Ok(ResolvedProvider { + model: provider_info.model.to_string(), + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + config, + }) } pub(super) fn parse_messages(messages: Value) -> Result, Error> { @@ -43,7 +54,11 @@ pub(super) fn parse_messages(messages: Value) -> Result, Error> pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, ) -> Result, Error> { - let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; + let ResolvedProvider { + model, + custom_llm_provider, + config, + } = resolve_provider_config(request.model, request.custom_llm_provider)?; let messages = parse_messages(request.messages)?; if messages.is_empty() { return Err(Error::InvalidRequest( @@ -55,6 +70,7 @@ pub(super) fn resolve_request( } Ok(ResolvedChatCompletionsRequest { model, + custom_llm_provider, config, messages, optional_params: request.optional_params, @@ -99,15 +115,18 @@ pub(super) fn prepare_provider_request( &env_lookup, )?; let transformed = - config.transform_request(&model, request.messages, request.optional_params)?; + config.transform_request(&model, request.messages, request.optional_params.clone())?; Ok(ProviderChatCompletionsRequest { model, + custom_llm_provider: request.custom_llm_provider, config, url, body: transformed.body, + optional_params: request.optional_params, environment, timeout: request.timeout, + api_key: request.api_key.map(|key| SecretValue::new(key.to_string())), }) } @@ -449,11 +468,19 @@ mod tests { json!("abc-123"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = crate::chat_completions::handler::outbound_request( + let authenticated = resolve_auth( &litellm_auth::AuthServices::default(), - &prepared, + prepared.environment, + &|_| None, ) .await + .expect("resolves"); + let signed = crate::chat_completions::handler::outbound_request( + authenticated, + prepared.url, + &prepared.body, + prepared.timeout, + ) .expect("signs"); let authorization = signed @@ -502,11 +529,19 @@ mod tests { call.api_key = None; call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = crate::chat_completions::handler::outbound_request( + let authenticated = resolve_auth( &litellm_auth::AuthServices::default(), - &prepared, + prepared.environment, + &|_| None, ) .await + .expect("resolves"); + let error = crate::chat_completions::handler::outbound_request( + authenticated, + prepared.url, + &prepared.body, + prepared.timeout, + ) .expect_err("{forwarded} should decline instead of being signed"); assert!( matches!(error, Error::Unsupported(_)), diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 66e9498c749..8c969dee730 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use litellm_auth::SecretValue; use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig}; use litellm_types::llms::openai::ChatMessage; use serde_json::{Map, Value}; @@ -23,6 +24,7 @@ pub struct ChatCompletionsRequest<'a> { pub struct ResolvedChatCompletionsRequest<'a> { pub model: String, + pub custom_llm_provider: String, pub config: &'static dyn BaseConfig, pub messages: Vec, pub optional_params: Map, @@ -34,11 +36,16 @@ pub struct ResolvedChatCompletionsRequest<'a> { pub struct ProviderChatCompletionsRequest { pub model: String, + pub custom_llm_provider: String, pub config: &'static dyn BaseConfig, pub url: String, pub body: Value, + /// The route's parameters before the provider transformation, reported to the host + /// beside the wire request. + pub optional_params: Map, /// The forwarded and default headers plus how the call authenticates; the credential /// itself is applied when the request is sent. pub environment: ValidatedEnvironment, pub timeout: Option, + pub api_key: Option, } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 650447d5abd..6832a59c7bd 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,20 +1,113 @@ use std::time::Duration; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream::BoxStream}; +use litellm_auth::AuthServices; +use litellm_host::{ + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, + hooks::RouteHooks, +}; use litellm_http::transport::Error as TransportError; use litellm_llms::base_llm::{ - anthropic_messages::transformation::BaseAnthropicMessagesConfig, auth::Authenticated, + anthropic_messages::{ + streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, + transformation::BaseAnthropicMessagesConfig, + }, + auth::{Authenticated, resolve_auth}, }; +use litellm_tracing::{ByteChunk, debug}; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::Value; -use super::{Error, common_utils::truncate_error_body}; +use super::{ + Error, MessagesResponse, common_utils::truncate_error_body, prepare::ProviderMessagesRequest, +}; use crate::{constants::MESSAGES_TIMEOUT_SECS, outbound::outbound_request}; -pub(super) fn network(error: reqwest::Error) -> Error { +pub(super) async fn execute( + http: &litellm_http::Client, + auth: &AuthServices, + request: ProviderMessagesRequest, + hooks: &impl RouteHooks, +) -> Result { + let ProviderMessagesRequest { + provider, + url, + body, + environment, + timeout, + api_key, + } = request; + let stream = body.params.stream == Some(true); + let context = RequestContext { + model: body.model.clone(), + custom_llm_provider: provider.as_str().to_string(), + optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, + secret_fields: Vec::new(), + api_key, + }; + let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?; + let wire = hooks + .before_send( + WireRequest { + url, + headers: authenticated.headers, + body: serde_json::to_value(&body).map_err(serialize_failure)?, + }, + context, + ) + .await?; + let provider_name = provider.as_str(); + debug!(provider = provider_name, stream, body = %wire.body, "provider request"); + let response = send( + http, + Authenticated { + headers: wire.headers, + signer: authenticated.signer, + }, + &wire.url, + &wire.body, + timeout, + ) + .await?; + debug!( + provider = provider_name, + status = response.status().as_u16(), + "provider response headers" + ); + if !response.status().is_success() { + return Err(provider_error(response).await); + } + let config = provider.config(); + if stream { + return Ok(streaming_response( + response, + config.stream_decoder(), + provider_name, + )); + } + let text = response.text().await.map_err(network)?; + debug!(body = text.as_str(), "provider response body"); + hooks + .emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; + decode_response(config, &body.model, &text) + .map(|message| MessagesResponse::Message(Box::new(message))) +} + +fn serialize_failure(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) +} + +fn network(error: reqwest::Error) -> Error { Error::Transport(TransportError::Network(error.to_string())) } -pub(super) async fn send( +async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, @@ -30,18 +123,21 @@ pub(super) async fn send( request.send(http).await.map_err(network) } -pub(super) async fn provider_error(response: reqwest::Response) -> Error { +async fn provider_error(response: reqwest::Response) -> Error { let status = response.status().as_u16(); match response.text().await { - Ok(text) => Error::Transport(TransportError::Http { - status, - body: truncate_error_body(&text), - }), + Ok(text) => { + litellm_tracing::debug!(status, body = text.as_str(), "provider error body"); + Error::Transport(TransportError::Http { + status, + body: truncate_error_body(&text), + }) + } Err(error) => network(error), } } -pub(super) fn decode_response( +fn decode_response( config: &dyn BaseAnthropicMessagesConfig, model: &str, text: &str, @@ -52,3 +148,97 @@ pub(super) fn decode_response( .transform_anthropic_messages_response(model, response) .map_err(Error::from) } + +fn streaming_response( + response: reqwest::Response, + decoder: Option, + provider: &'static str, +) -> MessagesResponse { + let headers = response + .headers() + .iter() + .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) + .collect(); + let chunks = match decoder { + None => futures_util::stream::try_unfold(response, move |mut response| async move { + let chunk = response.chunk().await.map_err(network)?; + Ok(chunk.map(|chunk| { + log_chunk(provider, "provider_response", &chunk); + (chunk, response) + })) + }) + .boxed(), + Some(decode) => decoded_chunks(response, decode, provider), + }; + MessagesResponse::Stream { headers, chunks } +} + +fn decoded_chunks( + response: reqwest::Response, + decode: StreamDecoder, + provider: &'static str, +) -> BoxStream<'static, Result> { + let bytes: ByteStream = response + .bytes_stream() + .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) + .map_err(std::io::Error::other) + .boxed(); + futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { + let Some(event) = events.try_next().await? else { + return Ok(None); + }; + let chunk = encode_anthropic_sse(&event)?; + log_chunk(provider, "client_response", &chunk); + Ok(Some((chunk, events))) + }) + .boxed() +} + +fn log_chunk(provider: &str, stage: &str, data: &Bytes) { + let chunk = ByteChunk::new(data); + debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk"); +} + +#[cfg(test)] +mod tests { + use litellm_llms::base_llm::anthropic_messages::streaming::anthropic_sse_event_stream; + use rstest::rstest; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + + use super::*; + + #[rstest] + #[case::event( + "data: {\"type\":\"ping\"}\n\n", + Some("event: ping\ndata: {\"type\":\"ping\"}\n\n") + )] + #[case::invalid_event("data: invalid\n\ndata: {\"type\":\"ping\"}\n\n", None)] + #[tokio::test] + async fn decoded_streams_encode_events_and_stop_at_the_first_error( + #[case] body: &'static str, + #[case] expected: Option<&str>, + ) { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream")) + .mount(&upstream) + .await; + let response = litellm_http::Client::plain_for_test() + .get(upstream.uri()) + .send() + .await + .unwrap(); + let MessagesResponse::Stream { mut chunks, .. } = + streaming_response(response, Some(anthropic_sse_event_stream), "test") + else { + panic!("a streaming response returns chunks"); + }; + + let chunk = chunks.next().await.unwrap(); + match expected { + Some(expected) => assert_eq!(chunk.unwrap().as_ref(), expected.as_bytes()), + None => assert!(matches!(chunk, Err(Error::InvalidResponse(_))), "{chunk:?}"), + } + assert!(chunks.next().await.is_none()); + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 3c081ff7bbd..f3a57da4d32 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,39 +1,27 @@ -//! The Anthropic Messages call, the Rust equivalent of Python's -//! `litellm.messages()`. +//! The Anthropic Messages call, the Rust equivalent of Python's `litellm.messages()`. //! -//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs -//! it in process for a caller that already holds the request and wants the message. +//! [`messages`] prepares the provider request and sends it in process. [`route`] runs the +//! same two steps as a machine for a host that answers the call's operations itself. -pub mod types; -pub use crate::error::RouteError as Error; mod common_utils; mod handler; mod prepare; pub mod route; -use std::sync::Arc; +mod types; use litellm_http::{ClientVariant, HttpClientConfig}; -use litellm_secrets::source::EnvironmentSecrets; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; +use litellm_secrets::source::SecretSource; + +pub use crate::error::RouteError as Error; +pub use types::{MessagesCall, MessagesResponse, MessagesShaping, messages_body}; pub async fn messages( resources: &crate::resources::CoreResources, config: &HttpClientConfig, + secrets: &dyn SecretSource, call: MessagesCall, -) -> Result { - let secrets = Arc::new(EnvironmentSecrets::python_compatible( - resources.pool.client(config, ClientVariant::Provider)?, - )); - match litellm_host::run::run( - messages_machine(resources, config, secrets)?, - &LocalMessagesHost::new(call), - ) - .await? - { - MessagesOutput::Message(message) => Ok(*message), - MessagesOutput::Streamed => Err(Error::Unsupported( - "streamed responses need a streaming host", - )), - } +) -> Result { + let http = resources.pool.client(config, ClientVariant::Provider)?; + let request = prepare::prepare(call, secrets).await?; + handler::execute(&http, &resources.auth, request, &()).await } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 7cd01a3a84c..c8e90cb5f2c 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,3 +1,6 @@ +use std::time::Duration; + +use litellm_auth::SecretValue; use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}, @@ -11,21 +14,42 @@ use litellm_llms::{ auth::{ValidatedEnvironment, with_default_headers}, }, }; +use litellm_secrets::source::SecretSource; use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use super::{ - Error, + Error, MessagesCall, common_utils::{MessagesProvider, string_headers}, - route::MessagesCall, - types::ProviderMessagesRequest, + types::invalid_request, }; -pub(super) struct ResolvedProvider { - pub(super) model: String, - pub(super) provider: MessagesProvider, +struct ResolvedProvider { + model: String, + provider: MessagesProvider, } -pub(super) fn resolve_provider( +pub(super) struct ProviderMessagesRequest { + pub(super) provider: MessagesProvider, + pub(super) url: String, + pub(super) body: AnthropicMessagesRequest, + pub(super) environment: ValidatedEnvironment, + pub(super) timeout: Option, + /// The caller's own credential, reported to the host beside the wire request. + pub(super) api_key: Option, +} + +pub(super) async fn prepare( + call: MessagesCall, + secrets: &dyn SecretSource, +) -> Result { + let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; + let secrets = secrets + .resolve(resolved.provider.config().secret_names()) + .await?; + prepare_provider_request(call, resolved, secrets.as_ref()) +} + +fn resolve_provider( model: &str, custom_llm_provider: Option<&str>, ) -> Result { @@ -53,7 +77,7 @@ pub(super) fn resolve_provider( }) } -pub(super) fn prepare_provider_request( +fn prepare_provider_request( call: MessagesCall, resolved: ResolvedProvider, secrets: &dyn Lookup, @@ -113,13 +137,10 @@ pub(super) fn prepare_provider_request( body: transformed, environment, timeout, + api_key: api_key.map(SecretValue::new), }) } -pub(super) fn invalid_request(err: serde_json::Error) -> Error { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) -} - fn without_additional_drop_params( request: AnthropicMessagesRequest, paths: &[String], @@ -145,7 +166,7 @@ mod tests { use serde_json::{Map, Value, json}; use super::*; - use crate::messages::types::MessagesShaping; + use crate::messages::MessagesShaping; #[fixture] fn shaping() -> MessagesShaping { diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index f9267dce755..5e5bbc927b6 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -1,55 +1,20 @@ use std::{ convert::Infallible, sync::{Arc, Mutex}, - time::Duration, }; use bytes::Bytes; -use futures_util::StreamExt; -use litellm_auth::SecretValue; +use futures_util::TryStreamExt; use litellm_host::{ - event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, machine::{CallMachine, HostChannel, MachineFault}, protocol::Protocol, }; use litellm_http::{Client, ClientVariant, HttpClientConfig}; -use litellm_llms::base_llm::{ - anthropic_messages::streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, - auth::{Authenticated, resolve_auth}, -}; use litellm_secrets::source::SecretSource; -use litellm_types::{ - llms::anthropic_messages::{ - anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, - }, - utils::ProviderSpecificHeaders, -}; -use serde_json::{Map, Value}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use super::{ - Error, - handler::{decode_response, network, provider_error, send}, - prepare::{invalid_request, prepare_provider_request, resolve_provider}, - types::MessagesShaping, -}; - -/// The caller's request as the host projects it. -pub struct MessagesCall { - pub body: AnthropicMessagesRequest, - pub api_key: Option, - pub api_base: Option, - pub custom_llm_provider: Option, - pub extra_headers: Option>, - pub provider_specific_header: Option, - pub timeout: Option, - pub shaping: MessagesShaping, -} - -/// Parses a caller's raw body, failing the way the route fails for any invalid request. -pub fn messages_body(body: Map) -> Result { - serde_json::from_value(Value::Object(body)).map_err(invalid_request) -} +use super::{Error, MessagesCall, MessagesResponse, handler::execute, prepare::prepare}; pub enum MessagesOutput { Message(Box), @@ -121,134 +86,35 @@ pub fn messages_machine( let http = resources.pool.client(config, ClientVariant::Provider)?; let auth = resources.auth.clone(); Ok(CallMachine::new(move |host| { - Box::pin(execute(host, http.clone(), auth.clone(), secrets.clone())) + Box::pin(drive(host, http, auth, secrets)) })) } -async fn execute( +/// The call as its host sees it: projection first, then the same prepare and execute as +/// [`super::messages`], with each chunk of a stream handed over as it arrives. +async fn drive( host: MessagesHost, http: Client, auth: Arc, secrets: Arc, ) -> Result { let call = host.project().await?; - let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; - let secrets = secrets - .resolve(resolved.provider.config().secret_names()) - .await?; - let api_key = call.api_key.clone().map(SecretValue::new); - let request = prepare_provider_request(call, resolved, secrets.as_ref())?; - let context = RequestContext { - model: request.body.model.clone(), - custom_llm_provider: request.provider.as_str().to_string(), - optional_params: serde_json::to_value(&request.body.params).map_err(serialize_failure)?, - secret_fields: Vec::new(), - api_key, - }; - let stream = request.body.params.stream == Some(true); - let config = request.provider.config(); - let body = serde_json::to_value(&request.body).map_err(serialize_failure)?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let authenticated = resolve_auth(&auth, request.environment, &env_lookup).await?; - let wire = host - .before_send( - WireRequest { - url: request.url, - headers: authenticated.headers, - body, - }, - context, - ) - .await?; - let response = send( - &http, - Authenticated { - headers: wire.headers, - signer: authenticated.signer, - }, - &wire.url, - &wire.body, - request.timeout, - ) - .await?; - if !response.status().is_success() { - return Err(provider_error(response).await); - } - if stream { - return relay(&host, response, config.stream_decoder()).await; - } - let text = response.text().await.map_err(network)?; - host.emit(MachineEvent::ResponseReceived { - raw: RawResponse { body: text.clone() }, - }) - .await?; - decode_response(config, &request.body.model, &text) - .map(|message| MessagesOutput::Message(Box::new(message))) -} - -fn serialize_failure(err: serde_json::Error) -> Error { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) -} - -/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading -/// ends the upstream read, and the call completes with what it delivered. -/// -/// A host on Anthropic SSE is relayed byte for byte. A host on another wire is decoded into -/// Anthropic stream events and re-encoded as Anthropic SSE. -async fn relay( - host: &MessagesHost, - response: reqwest::Response, - decoder: Option, -) -> Result { - let head = MessagesStreamHead { - headers: response - .headers() - .iter() - .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) - .collect(), - }; - if host.open(head).await? == Demand::Detached { - return Ok(MessagesOutput::Streamed); - } - match decoder { - None => relay_bytes(host, response).await, - Some(decode) => relay_events(host, response, decode).await, - } -} - -async fn relay_bytes( - host: &MessagesHost, - mut response: reqwest::Response, -) -> Result { - while let Some(chunk) = response.chunk().await.map_err(network)? { - if host.deliver(chunk).await? == Demand::Detached { - break; + let request = prepare(call, secrets.as_ref()).await?; + match execute(&http, &auth, request, &host).await? { + MessagesResponse::Message(message) => Ok(MessagesOutput::Message(message)), + MessagesResponse::Stream { + headers, + mut chunks, + } => { + if host.open(MessagesStreamHead { headers }).await? == Demand::Detached { + return Ok(MessagesOutput::Streamed); + } + while let Some(chunk) = chunks.try_next().await? { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) } } - Ok(MessagesOutput::Streamed) -} - -async fn relay_events( - host: &MessagesHost, - response: reqwest::Response, - decode: StreamDecoder, -) -> Result { - let bytes: ByteStream = futures_util::stream::unfold(response, |mut response| async move { - match response.chunk().await { - Ok(Some(chunk)) => Some((Ok(chunk), response)), - Ok(None) => None, - Err(error) => Some((Err(std::io::Error::other(error)), response)), - } - }) - .boxed(); - let mut events = decode(bytes); - while let Some(event) = events.next().await { - let chunk = encode_anthropic_sse(&event?)?; - if host.deliver(chunk).await? == Demand::Detached { - break; - } - } - Ok(MessagesOutput::Streamed) } diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index 006b1db4efb..b09cb96a919 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,12 +1,45 @@ use std::time::Duration; -use litellm_llms::{ - anthropic::common_utils::AnthropicModelCapabilities, base_llm::auth::ValidatedEnvironment, +use bytes::Bytes; +use futures_util::stream::BoxStream; +use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities; +use litellm_types::{ + llms::anthropic_messages::{ + anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, + }, + utils::ProviderSpecificHeaders, }; -use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; -use super::common_utils::MessagesProvider; +use super::Error; + +pub struct MessagesCall { + pub body: AnthropicMessagesRequest, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub provider_specific_header: Option, + pub timeout: Option, + pub shaping: MessagesShaping, +} + +pub fn messages_body(body: Map) -> Result { + serde_json::from_value(Value::Object(body)).map_err(invalid_request) +} + +pub(super) fn invalid_request(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) +} + +pub enum MessagesResponse { + Message(Box), + Stream { + headers: Vec<(String, String)>, + chunks: BoxStream<'static, Result>, + }, +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { @@ -20,16 +53,6 @@ pub struct MessagesShaping { pub additional_drop_params: Vec, } -pub(crate) struct ProviderMessagesRequest { - pub(crate) provider: MessagesProvider, - pub(crate) url: String, - pub(crate) body: AnthropicMessagesRequest, - /// The forwarded, default and feature headers plus how the call authenticates; the - /// credential itself is applied when the request is sent. - pub(crate) environment: ValidatedEnvironment, - pub(crate) timeout: Option, -} - #[cfg(test)] mod tests { use litellm_llms::anthropic::common_utils::SupportedEffortTiers; diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 719c86990b0..534af6d7d06 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -1,9 +1,8 @@ use std::{sync::Arc, time::Duration}; use litellm_core::messages::{ - Error, - route::{LocalMessagesHost, MessagesCall, MessagesMachine, MessagesOutput, messages_machine}, - types::MessagesShaping, + Error, MessagesCall, MessagesShaping, + route::{LocalMessagesHost, MessagesMachine, MessagesOutput, messages_machine}, }; use litellm_http::{HttpSettings, Resolution}; use litellm_secrets::source::SecretSource; diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 0d44d26d416..f6ee0e6dfbf 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -1,7 +1,5 @@ -use litellm_llms::anthropic::common_utils::{ - ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities, - SupportedEffortTiers, beta, -}; +use litellm_llms::anthropic::common_utils::{AnthropicModelCapabilities, SupportedEffortTiers}; +use litellm_types::llms::anthropic::{AnthropicBeta, BetaSet}; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use rstest::rstest; @@ -247,41 +245,37 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) assert_eq!(sent["top_k"], 3); } -fn sent_betas(request: &wiremock::Request) -> Vec { +fn sent_betas(request: &wiremock::Request) -> BetaSet { let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta")) .unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}")); - header - .split(',') - .map(str::trim) - .map(str::to_string) - .collect() + header.parse().unwrap() } #[rstest] -#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] -#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] -#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] +#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[AnthropicBeta::StructuredOutputs20251113])] +#[case::fast_mode(json!({"speed": "fast"}), &[AnthropicBeta::FastMode20260201])] +#[case::compaction(json!({"compaction": {"enabled": true}}), &[AnthropicBeta::Compact20260904])] #[case::context_management_edits( json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}), - &[beta::CONTEXT_MANAGEMENT_2025_06_27] + &[AnthropicBeta::ContextManagement20250627] )] #[case::per_message_output_config( json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] + &[AnthropicBeta::PerTurnControl20260701] )] #[case::advisor_tool( - json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}), - &[beta::ADVISOR_TOOL_2026_03_01] + json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": MODEL}]}), + &[AnthropicBeta::AdvisorTool20260301] )] #[case::several_features_at_once( json!({"speed": "fast", "output_format": {"type": "json_schema"}}), - &[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01] + &[AnthropicBeta::StructuredOutputs20251113, AnthropicBeta::FastMode20260201] )] #[tokio::test] async fn feature_betas_join_the_callers_betas_in_one_sorted_header( call: MessagesCall, #[case] fields: Value, - #[case] features: &[&str], + #[case] features: &[AnthropicBeta], ) { let upstream = upstream([message_response()]).await; let capabilities = AnthropicModelCapabilities { @@ -305,12 +299,11 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header( .await; let sent = sent_betas(&only_request(&upstream).await); - let mut expected: Vec = features + let expected: BetaSet = features .iter() - .map(|feature| feature.to_string()) - .chain(["caller-beta-2025-01-01".to_string()]) + .cloned() + .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) .collect(); - expected.sort(); assert_eq!(sent, expected); } @@ -331,7 +324,10 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M request.header("anthropic-dangerous-direct-browser-access"), Some("true") ); - assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]); + assert_eq!( + sent_betas(&request), + BetaSet::from_iter([AnthropicBeta::Oauth20250420]) + ); assert_eq!(request.header("x-api-key"), None); } diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index ed715e22898..14c8eb6b7c6 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -1,6 +1,6 @@ use litellm_core::{ Phase, - messages::{messages, route::messages_body}, + messages::{MessagesResponse, messages, messages_body}, }; use litellm_http::transport::Error as TransportError; use rstest::rstest; @@ -188,9 +188,10 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes ..HttpSettings::default() }; - let message = messages( + let response = messages( &support::resources(), &Resolution::from(&settings).config, + &RecordingSecrets::empty(), MessagesCall { api_key: Some("sk-ant".into()), api_base: Some(base), @@ -200,6 +201,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes .await .expect("messages request succeeds"); + let MessagesResponse::Message(message) = response else { + panic!("a non-streaming request returns a message"); + }; assert_eq!(message.id, "msg_1"); let sent = only_request(&upstream).await; assert_eq!(sent.header("x-api-key"), Some("sk-ant")); diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index 49a6f7e87a0..f0e55eca8dd 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -1,12 +1,21 @@ -use std::{convert::Infallible, sync::Mutex}; +use std::{ + convert::Infallible, + sync::{Mutex, mpsc}, +}; use bytes::Bytes; -use litellm_core::messages::route::{Messages, MessagesStreamHead}; +use futures_util::{StreamExt, TryStreamExt}; +use litellm_core::messages::{ + MessagesResponse, messages, + route::{Messages, MessagesStreamHead}, +}; use litellm_host::host::{Demand, Host}; +use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::rstest; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, + task::JoinHandle, }; use super::*; @@ -23,6 +32,20 @@ enum Seen { Deliver(Bytes), } +struct TraceSink(mpsc::Sender<(String, Value)>); + +impl Sink for TraceSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target().starts_with("litellm_core::messages") + } + + fn emit(&self, record: &Record) { + self.0 + .send((record.message.clone(), Value::Object(record.fields.clone()))) + .unwrap(); + } +} + /// Projects like `LocalMessagesHost`, records every stream op in the order the route /// performs it, and detaches after `detach_after` ops. struct RecordingStreamHost { @@ -120,6 +143,37 @@ async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: Me assert_eq!(delivered, SSE_BODY.as_bytes()); } +#[rstest] +#[tokio::test] +async fn debug_trace_keeps_provider_input_and_every_stream_chunk(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + let (sender, receiver) = mpsc::channel(); + + Logger::new(TraceSink(sender)) + .instrument(stream_through(&host)) + .await + .unwrap(); + + let records: Vec<(String, Value)> = receiver.try_iter().collect(); + let request = records + .iter() + .find(|(message, _)| message == "provider request") + .unwrap(); + let body: Value = serde_json::from_str(request.1["body"].as_str().unwrap()).unwrap(); + assert_eq!(body["messages"][0]["content"], "hi"); + assert_eq!(request.1["stream"], true); + let chunks: String = records + .iter() + .filter(|(message, fields)| { + message == "stream chunk" && fields["stage"] == "provider_response" + }) + .map(|(_, fields)| fields["chunk"].as_str().unwrap()) + .collect(); + assert_eq!(chunks, SSE_BODY); + assert!(!format!("{records:?}").contains("sk-ant")); +} + #[rstest] #[case::at_open(1)] #[case::after_the_first_chunk(2)] @@ -192,10 +246,10 @@ async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: Messages } /// Serves one SSE chunk and then holds the connection open without ever finishing. -async fn stalling_upstream() -> String { +async fn stalling_upstream() -> (String, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base = format!("http://{}", listener.local_addr().unwrap()); - tokio::spawn(async move { + let connection = tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); let mut request = vec![0; 4096]; let _ = socket.read(&mut request).await; @@ -206,15 +260,15 @@ async fn stalling_upstream() -> String { ) .await .unwrap(); - std::future::pending::<()>().await; + let _ = socket.read_to_end(&mut Vec::new()).await; }); - base + (base, connection) } #[rstest] #[tokio::test] async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { - let base = stalling_upstream().await; + let (base, connection) = stalling_upstream().await; let host = RecordingStreamHost::new( MessagesCall { timeout: Some(Duration::from_millis(300)), @@ -236,6 +290,145 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { "the chunk before the stall reached the caller, saw {} ops", seen.len() ); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("timing out closes the upstream connection") + .unwrap(); +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn the_sdk_returns_stream_headers_and_every_sse_byte( + call: MessagesCall, + #[case] provider: &str, +) { + let upstream = upstream([sse_response()]).await; + let response = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + custom_llm_provider: Some(provider.into()), + ..streaming(call, upstream.uri()) + }, + ) + .await + .unwrap(); + + let MessagesResponse::Stream { headers, chunks } = response else { + panic!("a streaming request returns a stream"); + }; + for (name, value) in UPSTREAM_HEADERS { + assert!(headers.contains(&(name.into(), value.into()))); + } + let delivered = chunks.try_collect::>().await.unwrap().concat(); + assert_eq!(delivered, SSE_BODY.as_bytes()); + assert_eq!(only_request(&upstream).await.json()["stream"], true); +} + +#[rstest] +#[tokio::test] +async fn the_sdk_returns_http_errors_before_opening_a_stream(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; + let error = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + streaming(call, upstream.uri()), + ) + .await + .err() + .expect("upstream failure is returned by messages()"); + + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status: 429, + body: "slow down".into(), + }) + ); +} + +#[rstest] +#[case::before_reading(false)] +#[case::after_reading(true)] +#[tokio::test] +async fn dropping_the_sdk_stream_closes_the_unfinished_upstream( + call: MessagesCall, + #[case] read_chunk: bool, +) { + let (base, connection) = stalling_upstream().await; + let response = tokio::time::timeout( + Duration::from_secs(5), + messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + timeout: Some(Duration::from_secs(30)), + ..streaming(call, base) + }, + ), + ) + .await + .expect("messages() returns before the upstream finishes") + .unwrap(); + + let MessagesResponse::Stream { mut chunks, .. } = response else { + panic!("a streaming request returns a stream"); + }; + if read_chunk { + let chunk = tokio::time::timeout(Duration::from_secs(5), chunks.next()) + .await + .expect("the first chunk arrives before the upstream finishes") + .unwrap() + .unwrap(); + assert_eq!(chunk.as_ref(), b"event: message_start\ndata: {}\n\n"); + } + assert!(!connection.is_finished()); + drop(chunks); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("dropping the stream closes the upstream connection") + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) { + let (base, connection) = stalling_upstream().await; + let response = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + ) + .await + .unwrap(); + + let MessagesResponse::Stream { mut chunks, .. } = response else { + panic!("a streaming request returns a stream"); + }; + assert_eq!( + chunks.next().await.unwrap().unwrap().as_ref(), + b"event: message_start\ndata: {}\n\n" + ); + let error = tokio::time::timeout(Duration::from_secs(5), chunks.next()) + .await + .expect("the stalled body times out") + .unwrap() + .unwrap_err(); + assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(chunks.next().await.is_none()); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("the failed stream closes its upstream connection") + .unwrap(); } #[rstest] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index e3f40ec5354..f5ee5e81ba6 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -12,7 +12,6 @@ bytes.workspace = true futures-util.workspace = true litellm-auth.workspace = true litellm-core.workspace = true -litellm-host.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-router.workspace = true @@ -20,10 +19,10 @@ litellm-secrets.workspace = true litellm-types.workspace = true serde_json.workspace = true thiserror.workspace = true -tokio = { workspace = true, features = ["sync"] } [dev-dependencies] futures-util.workspace = true +tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true tower = { version = "0.5.3", features = ["util"] } wiremock = "0.6.5" diff --git a/litellm-rust/crates/gateway-inference/src/messages/host.rs b/litellm-rust/crates/gateway-inference/src/messages/host.rs deleted file mode 100644 index 57486d9d157..00000000000 --- a/litellm-rust/crates/gateway-inference/src/messages/host.rs +++ /dev/null @@ -1,69 +0,0 @@ -use std::{convert::Infallible, sync::Mutex}; - -use bytes::Bytes; -use litellm_core::messages::{ - Error, - route::{LocalMessagesHost, Messages, MessagesCall, MessagesStreamHead}, -}; -use litellm_host::host::{Demand, Host}; -use tokio::sync::{mpsc, oneshot}; - -/// Hands a streamed response to the HTTP body: the head once, then each chunk. A dropped -/// receiver means the client went away, which detaches the call. -pub(super) struct ChannelHost { - local: LocalMessagesHost, - head: Mutex>>, - pub(super) chunks: mpsc::Sender, -} - -impl ChannelHost { - pub(super) fn new( - call: MessagesCall, - head: oneshot::Sender, - chunks: mpsc::Sender, - ) -> Self { - Self { - local: LocalMessagesHost::new(call), - head: Mutex::new(Some(head)), - chunks, - } - } - - fn take_head(&self) -> Option> { - self.head - .lock() - .unwrap_or_else(|error| error.into_inner()) - .take() - } - - pub(super) fn opened(&self) -> bool { - self.head - .lock() - .unwrap_or_else(|error| error.into_inner()) - .is_none() - } -} - -impl Host for ChannelHost { - async fn project(&self) -> Result { - self.local.project().await - } - - async fn custom_op(&self, op: Infallible) -> Result<(), Error> { - match op {} - } - - async fn open(&self, head: MessagesStreamHead) -> Result { - Ok(match self.take_head().map(|sender| sender.send(head)) { - Some(Ok(())) => Demand::More, - Some(Err(_)) | None => Demand::Detached, - }) - } - - async fn deliver(&self, chunk: Bytes) -> Result { - Ok(match self.chunks.send(chunk).await { - Ok(()) => Demand::More, - Err(_) => Demand::Detached, - }) - } -} diff --git a/litellm-rust/crates/gateway-inference/src/messages/mod.rs b/litellm-rust/crates/gateway-inference/src/messages/mod.rs index e96cae8ba53..5d6a8faa0e8 100644 --- a/litellm-rust/crates/gateway-inference/src/messages/mod.rs +++ b/litellm-rust/crates/gateway-inference/src/messages/mod.rs @@ -1,7 +1,5 @@ //! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it. -mod host; - use std::{convert::Infallible, sync::Arc}; use axum::{ @@ -11,13 +9,12 @@ use axum::{ http::{HeaderMap, StatusCode, header}, response::{IntoResponse, Response}, }; -use host::ChannelHost; -use litellm_core::messages::route::{ - MessagesCall, MessagesOutput, messages_body, messages_machine, +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_core::messages::{ + Error as RouteError, MessagesCall, MessagesResponse, messages, messages_body, }; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use serde_json::{Map, Value}; -use tokio::sync::{mpsc, oneshot}; use crate::{Deployment, Error, Gateway}; @@ -55,31 +52,16 @@ async fn handle(gateway: &Gateway, headers: &HeaderMap, body: &[u8]) -> Result Ok(stream(chunks)), - joined = call => match joined.map_err(|error| Error::Internal(error.to_string()))?? { - MessagesOutput::Message(message) => Ok(Json(message).into_response()), - MessagesOutput::Streamed => Err(Error::Internal("the stream ended before it opened".into())), - }, + match messages( + &gateway.resources, + &gateway.http, + gateway.secrets.as_ref(), + call, + ) + .await? + { + MessagesResponse::Message(message) => Ok(Json(message).into_response()), + MessagesResponse::Stream { chunks, .. } => Ok(stream(chunks)), } } @@ -123,10 +105,13 @@ fn anthropic_api_headers(headers: &HeaderMap) -> Option }) } -fn stream(chunks: mpsc::Receiver) -> Response { - let body = futures_util::stream::unfold(chunks, |mut chunks| async move { - let chunk = chunks.recv().await?; - Some((Ok::<_, Infallible>(chunk), chunks)) +/// A chunk that fails after the stream opened is delivered as an SSE error frame, since +/// the status line already went out; the stream ends on it. +fn stream(chunks: BoxStream<'static, Result>) -> Response { + let body = chunks.map(|chunk| { + Ok::<_, Infallible>( + chunk.unwrap_or_else(|error| Bytes::from(Error::Route(error).sse_frame())), + ) }); ( StatusCode::OK, diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index ef7e7c66681..30836498d49 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -6,6 +6,7 @@ use axum::{ }; use rstest::rstest; use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, @@ -64,3 +65,57 @@ async fn invalid_messages_stays_an_anthropic_error() { assert_eq!(body["type"], "error"); assert_eq!(body["error"]["type"], "invalid_request_error"); } + +/// Answers with the SSE head and one event, then drops the connection short of the +/// announced body length. +async fn truncating_upstream() -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{FIRST_EVENT}", + FIRST_EVENT.len() * 2 + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + base +} + +const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n"; + +#[tokio::test] +async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { + let base = truncating_upstream().await; + let request = Request::post("/v1/messages") + .header("content-type", "application/json") + .body(Body::from( + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, "stream": true}) + .to_string(), + )) + .unwrap(); + + let response = support::app("anthropic/test-model", &base) + .oneshot(request) + .await + .unwrap(); + + assert_eq!(response.status(), 200); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + let frame = text + .strip_prefix(FIRST_EVENT) + .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) + .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); + let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); + assert_eq!(error["type"], "error"); + assert_eq!(error["error"]["type"], "api_error"); +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 554186955b4..c27a3f5b17e 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] axum.workspace = true +http-body-util = "0.1" litellm-core.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-auth.workspace = true @@ -14,11 +15,14 @@ litellm-config.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true -tower-http = { version = "0.7.1", default-features = false, features = ["trace"] } +litellm-tracing.workspace = true +serde_json.workspace = true tracing.workspace = true tokio.workspace = true +uuid.workspace = true [dev-dependencies] +futures-util.workspace = true rstest.workspace = true -serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } +tower = { version = "0.5", features = ["util"] } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 3f67f923a5c..fc16dc0de67 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,7 +1,13 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Instant}; -use axum::{Router, extract::Request}; -use tower_http::trace::{DefaultOnResponse, TraceLayer}; +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + middleware::Next, + response::Response, +}; +use http_body_util::BodyExt; use litellm_config::Config; use litellm_core::resources::CoreResources; @@ -12,6 +18,8 @@ use litellm_http::{ }; use litellm_llms::base_llm::ocr::settings::OcrSettings; use litellm_secrets::source::EnvironmentSecrets; +use litellm_tracing::ByteChunk; +use uuid::Uuid; pub fn build_inference(config: &Config) -> Result, litellm_http::Error> { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); @@ -42,11 +50,125 @@ pub fn router(inference: Arc, config: &Config) -> Router { RequireMasterKey, _, >(auth)) - .layer( - TraceLayer::new_for_http() - .make_span_with(|request: &Request| { - tracing::info_span!("request", method = %request.method(), path = request.uri().path()) - }) - .on_response(DefaultOnResponse::new().level(tracing::Level::INFO)), - ) + .layer(axum::middleware::from_fn(log_request)) +} + +async fn log_request(request: Request, next: Next) -> Response { + let request_id = Uuid::new_v4().to_string(); + let log_body_chunks = tracing::enabled!(tracing::Level::DEBUG); + let method = request.method().clone(); + let path = request.uri().path().to_owned(); + let started = Instant::now(); + let request = if log_body_chunks { + request.map(|body| logged_body(body, request_id.clone(), "input")) + } else { + request + }; + let response = next.run(request).await; + tracing::info!( + %request_id, + %method, + %path, + status = response.status().as_u16(), + time_to_headers_ms = started.elapsed().as_secs_f64() * 1000.0, + "response headers" + ); + if log_body_chunks { + response.map(|body| logged_body(body, request_id, "output")) + } else { + response + } +} + +fn logged_body(body: Body, request_id: String, direction: &'static str) -> Body { + Body::new(body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + log_chunk(&request_id, direction, data); + } + frame + })) +} + +fn log_chunk(request_id: &str, direction: &str, data: &Bytes) { + let chunk = ByteChunk::new(data); + tracing::debug!(request_id, direction, encoding = chunk.encoding(), chunk = %chunk, "body chunk"); +} + +#[cfg(test)] +mod tests { + use std::{convert::Infallible, sync::mpsc}; + + use axum::{body::to_bytes, http::StatusCode, routing::post}; + use futures_util::stream; + use litellm_tracing::{Logger, Metadata, Record, Sink}; + use rstest::rstest; + use serde_json::{Value, json}; + use tower::ServiceExt; + + use super::*; + + struct LogSink(mpsc::Sender); + + impl Sink for LogSink { + fn enabled(&self, _: &Metadata<'_>) -> bool { + true + } + + fn emit(&self, record: &Record) { + self.0 + .send(json!({"message": record.message, "fields": record.fields})) + .unwrap(); + } + } + + #[rstest] + #[tokio::test] + async fn logs_each_body_chunk_without_changing_streamed_bytes() { + let app = Router::new() + .route( + "/stream", + post(|_: Bytes| async { + ( + StatusCode::OK, + Body::from_stream(stream::iter([ + Ok::<_, Infallible>(Bytes::from_static(b"event: first\n\n")), + Ok(Bytes::from_static(b"event: second\n\n")), + ])), + ) + }), + ) + .layer(axum::middleware::from_fn(log_request)); + let request_chunks = [ + Ok::<_, Infallible>(Bytes::from_static(b"hello")), + Ok(Bytes::from_static(b" world")), + ]; + let request = Request::post("/stream") + .body(Body::from_stream(stream::iter(request_chunks))) + .unwrap(); + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(LogSink(sender)); + + let output = logger + .instrument(async { + let response = app.oneshot(request).await.unwrap(); + to_bytes(response.into_body(), 1024).await.unwrap() + }) + .await; + + assert_eq!(output, "event: first\n\nevent: second\n\n"); + let records: Vec = receiver.try_iter().collect(); + assert_eq!(records.len(), 5); + assert_eq!(records[0]["fields"]["chunk"], "hello"); + assert_eq!(records[1]["fields"]["chunk"], " world"); + assert_eq!(records[2]["fields"]["status"], 200); + assert_eq!(records[3]["fields"]["chunk"], "event: first\n\n"); + assert_eq!(records[4]["fields"]["chunk"], "event: second\n\n"); + let request_id = &records[2]["fields"]["request_id"]; + assert!(request_id.as_str().is_some()); + assert!( + records + .iter() + .all(|record| &record["fields"]["request_id"] == request_id) + ); + } } diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index bae711e16c7..40f9cb442d5 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -1,9 +1,46 @@ -use std::error::Error; +use std::{ + error::Error, + time::{SystemTime, UNIX_EPOCH}, +}; use litellm_config::Config; +use litellm_tracing::{Level, Logger, Metadata, Record, Sink}; +use serde_json::json; + +struct StderrSink { + level: Level, +} + +impl Sink for StderrSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + *metadata.level() <= self.level && metadata.target().starts_with("litellm") + } + + fn emit(&self, record: &Record) { + let timestamp_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + eprintln!( + "{}", + json!({ + "timestamp_ms": timestamp_ms, + "level": record.metadata.level().as_str(), + "target": record.metadata.target(), + "message": record.message, + "fields": record.fields, + }) + ); + } +} #[tokio::main] async fn main() -> Result<(), Box> { + let level = std::env::var("RUST_LOG") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(Level::INFO); + Logger::new(StderrSink { level }).install_global()?; let config_path = std::env::var("LITELLM_CONFIG").unwrap_or_else(|_| "config.yaml".into()); let config = Config::load(config_path)?; let inference = litellm_gateway::build_inference(&config)?; @@ -13,6 +50,8 @@ async fn main() -> Result<(), Box> { .parse::()?; let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, "gateway listening"); + axum::serve(listener, litellm_gateway::router(inference, &config)).await?; Ok(()) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index a19d8c9c4fe..a91051441ea 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -1,11 +1,31 @@ -use std::{sync::Arc, time::Duration}; +use std::{ + sync::{Arc, mpsc}, + time::Duration, +}; +use axum::{body::Body, http::Request}; use litellm_config::Config; use litellm_gateway_inference::{Error, Gateway}; use litellm_http::ClientVariant; +use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::{fixture, rstest}; use serde_json::{Value, json}; use tokio::{net::TcpListener, sync::oneshot, time::timeout}; +use tower::ServiceExt; + +struct LogSink(mpsc::Sender); + +impl Sink for LogSink { + fn enabled(&self, _: &Metadata<'_>) -> bool { + true + } + + fn emit(&self, record: &Record) { + self.0 + .send(json!({"message": record.message, "fields": record.fields})) + .unwrap(); + } +} #[fixture] fn inference() -> Arc { @@ -88,3 +108,36 @@ async fn authenticates_before_serving_mounted_inference_routes( .unwrap() .unwrap(); } + +#[rstest] +#[tokio::test] +async fn logs_request_outcome_without_credentials_or_query(inference: Arc) { + let config = + Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: gateway-key\n") + .unwrap(); + let request = Request::builder() + .method("POST") + .uri("/v1/messages?token=query-secret") + .header("authorization", "Bearer header-secret") + .body(Body::empty()) + .unwrap(); + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(LogSink(sender)); + + let response = logger + .instrument(litellm_gateway::router(inference, &config).oneshot(request)) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 401); + let record = receiver.try_recv().unwrap(); + assert_eq!(record["message"], "response headers"); + assert_eq!(record["fields"]["method"], "POST"); + assert_eq!(record["fields"]["path"], "/v1/messages"); + assert_eq!(record["fields"]["status"], 401); + assert!(record["fields"]["time_to_headers_ms"].as_f64().unwrap() >= 0.0); + assert!(record["fields"]["request_id"].as_str().is_some()); + assert!(receiver.try_recv().is_err()); + assert!(!record.to_string().contains("header-secret")); + assert!(!record.to_string().contains("query-secret")); +} diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs new file mode 100644 index 00000000000..14b0f1ea08a --- /dev/null +++ b/litellm-rust/crates/host/src/hooks.rs @@ -0,0 +1,145 @@ +use std::future::Future; + +use crate::{ + event::{MachineEvent, RequestContext, WireRequest}, + machine::{HostChannel, MachineFault}, + protocol::Protocol, +}; + +/// What a route reaches for mid-call: the send-time rewrite and the events it reports. +/// Python's `logging_obj.pre_call` and `post_call`, in that order. +pub trait RouteHooks: Send + Sync { + fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> impl Future> + Send; + + fn emit(&self, event: MachineEvent) -> impl Future> + Send; +} + +/// No host: the wire request goes out as prepared and nothing observes the call. +impl RouteHooks for () { + async fn before_send(&self, wire: WireRequest, _: RequestContext) -> Result { + Ok(wire) + } + + async fn emit(&self, _: MachineEvent) -> Result<(), E> { + Ok(()) + } +} + +impl RouteHooks for HostChannel +where + R::Error: From, +{ + async fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + HostChannel::before_send(self, wire, context).await + } + + async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { + HostChannel::emit(self, event).await + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use serde_json::json; + + use super::*; + use crate::{ + event::RawResponse, + host::HostOp, + machine::{CallMachine, Machine, MachineStep}, + }; + + struct Unit; + + #[derive(Clone, Debug)] + struct Fault; + + impl Protocol for Unit { + type Response = (WireRequest, ()); + type Error = Fault; + type Projection = (); + type Op = Infallible; + type Chunk = Infallible; + type StreamHead = Infallible; + } + + impl From for Fault { + fn from(_: MachineFault) -> Self { + Fault + } + } + + fn wire(url: &str) -> WireRequest { + WireRequest { + url: url.into(), + headers: Vec::new(), + body: json!({}), + } + } + + fn context() -> RequestContext { + RequestContext { + model: "m".into(), + custom_llm_provider: "p".into(), + optional_params: json!({}), + secret_fields: Vec::new(), + api_key: None, + } + } + + #[tokio::test] + async fn the_channel_yields_each_hook_as_its_op_and_returns_the_answer() { + let mut machine = CallMachine::::new(|channel| { + Box::pin(async move { + let sent = RouteHooks::before_send(&channel, wire("prepared"), context()).await?; + RouteHooks::emit( + &channel, + MachineEvent::ResponseReceived { + raw: RawResponse { body: "raw".into() }, + }, + ) + .await?; + Ok((sent, ())) + }) + }); + + let Ok(MachineStep::Host(HostOp::BeforeSend { wire, reply, .. })) = machine.resume().await + else { + panic!("before_send yields BeforeSend"); + }; + assert_eq!(wire.url, "prepared"); + reply.send(WireRequest { + url: "rewritten".into(), + ..*wire + }); + + let Ok(MachineStep::Host(HostOp::Emit(event, reply))) = machine.resume().await else { + panic!("emit yields Emit"); + }; + assert!(matches!(event, MachineEvent::ResponseReceived { .. })); + reply.send(()); + + let Ok(MachineStep::Complete((sent, ()))) = machine.resume().await else { + panic!("the call completes with the answers"); + }; + assert_eq!(sent.url, "rewritten"); + } + + #[tokio::test] + async fn no_hooks_pass_the_wire_request_through() { + let sent = RouteHooks::::before_send(&(), wire("prepared"), context()) + .await + .unwrap(); + assert_eq!(sent.url, "prepared"); + } +} diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index c6b9e59b65a..1df68941fa3 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -7,6 +7,7 @@ //! may rewrite the wire request before it is sent. pub mod event; +pub mod hooks; pub mod host; pub mod machine; pub mod protocol; diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index fcf296793a5..17f2e652a92 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -87,6 +87,20 @@ pub fn has_header(headers: &[(String, String)], name: &str) -> bool { .any(|(key, _)| key.eq_ignore_ascii_case(name)) } +pub fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +pub fn without_headers(headers: Vec<(String, String)>, names: &[&str]) -> Vec<(String, String)> { + headers + .into_iter() + .filter(|(key, _)| !names.iter().any(|name| key.eq_ignore_ascii_case(name))) + .collect() +} + pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { headers.iter().any(|(name, value)| { if !name.eq_ignore_ascii_case("authorization") { @@ -194,6 +208,30 @@ mod tests { assert!(!has_header(&headers, "authorization")); } + #[test] + fn header_value_reads_the_first_match_in_any_case() { + let headers = vec![ + ("X-Api-Key".to_string(), "first".to_string()), + ("x-api-key".to_string(), "second".to_string()), + ]; + assert_eq!(header_value(&headers, "x-API-key"), Some("first")); + assert_eq!(header_value(&headers, "authorization"), None); + } + + #[test] + fn without_headers_drops_every_casing_of_the_named_headers_and_keeps_order() { + let headers = vec![ + ("X-Api-Key".to_string(), "k".to_string()), + ("anthropic-version".to_string(), "v".to_string()), + ("AUTHORIZATION".to_string(), "Bearer t".to_string()), + ("x-api-key".to_string(), "k2".to_string()), + ]; + assert_eq!( + without_headers(headers, &["x-api-key", "authorization"]), + vec![("anthropic-version".to_string(), "v".to_string())] + ); + } + #[test] fn auth_header_detection_is_case_insensitive() { let headers = vec![ diff --git a/litellm-rust/crates/litellm/Cargo.toml b/litellm-rust/crates/litellm/Cargo.toml index f6a63227792..41009ceaff1 100644 --- a/litellm-rust/crates/litellm/Cargo.toml +++ b/litellm-rust/crates/litellm/Cargo.toml @@ -1,3 +1,4 @@ [package] name = "litellm" version = "0.0.1" +edition.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index 395c2376059..d1a8bdff55a 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -4,7 +4,7 @@ use serde_json::Value; use time::OffsetDateTime; use url::Url; -use crate::{Error, anthropic::messages::transformation::resolve_anthropic_api_base}; +use crate::{Error, anthropic::common_utils::resolve_anthropic_api_base}; const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches"; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index b19443a7ff8..ba77a6ed790 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,4 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::SecretValue; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, build_conversation}, @@ -12,9 +12,11 @@ use serde_json::{Map, Value, json}; use crate::{ Error, anthropic::{ - ANTHROPIC_OAUTH_TOKEN_PREFIX, chat::handler::ModelResponseIterator, - messages::transformation::{complete_anthropic_url, resolve_anthropic_api_key}, + common_utils::{ + API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, + resolve_anthropic_api_key, + }, }, base_llm::{ anthropic_messages::streaming::anthropic_sse_event_stream, @@ -50,15 +52,6 @@ pub struct AnthropicConfig; pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig; -fn forwards_oauth_bearer(headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, value)| { - name.eq_ignore_ascii_case("authorization") - && value - .strip_prefix("Bearer ") - .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) - }) -} - impl BaseConfig for AnthropicConfig { fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS @@ -160,14 +153,14 @@ impl BaseConfig for AnthropicConfig { _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - if forwards_oauth_bearer(&headers) { + if forwarded_oauth_bearer(&headers).is_some() { return Ok(ValidatedEnvironment { headers, auth: AuthScheme::Forwarded, }); } let auth = AuthScheme::Credential { - placement: CredentialPlacement::Header("x-api-key"), + placement: API_KEY_PLACEMENT, secret: SecretValue::new(resolve_anthropic_api_key(api_key, env_lookup)?), }; Ok(ValidatedEnvironment { headers, auth }) diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 34c59c6ec5d..d8d24ec0402 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,30 +1,33 @@ -use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessage, ContentBlock, EffortLevel, MessageContent, +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_http::request::{has_header, header_value, without_headers}; +use litellm_types::llms::{ + anthropic::{AnthropicBeta, BetaSet}, + anthropic_messages::anthropic_request::{ + AnthropicMessage, AnthropicTool, ContentBlock, EffortLevel, MessageContent, + }, }; +use litellm_types::recognized::Recognized; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX; +use crate::{ + anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, + base_llm::auth::{AuthScheme, Headers}, +}; -pub const ANTHROPIC_OAUTH_BETA_HEADER: &str = "oauth-2025-04-20"; -pub const ANTHROPIC_ADVISOR_TOOL_TYPE: &str = "advisor_20260301"; -pub const ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: [&str; 2] = [ - "tool_search_tool_regex_20251119", - "tool_search_tool_bm25_20251119", -]; +pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +pub const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; pub const ENCRYPTED_REASONING_SIGNATURE_PREFIX: &str = "litellm_encrypted_reasoning:"; const THOUGHT_SIGNATURE_SEPARATOR: &str = "__thought__"; - -pub mod beta { - pub const CONTEXT_MANAGEMENT_2025_06_27: &str = "context-management-2025-06-27"; - pub const COMPACT_2026_01_12: &str = "compact-2026-01-12"; - pub const COMPACT_2026_09_04: &str = "compact-2026-09-04"; - pub const STRUCTURED_OUTPUT: &str = "structured-outputs-2025-11-13"; - pub const ADVANCED_TOOL_USE_2025_11_20: &str = "advanced-tool-use-2025-11-20"; - pub const FAST_MODE_2026_02_01: &str = "fast-mode-2026-02-01"; - pub const ADVISOR_TOOL_2026_03_01: &str = "advisor-tool-2026-03-01"; - pub const PER_TURN_CONTROL_2026_07_01: &str = "per-turn-control-2026-07-01"; -} +const BETA_HEADER: &str = "anthropic-beta"; +pub const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; +pub const ANTHROPIC_BASE_URL_ENV: &str = "ANTHROPIC_BASE_URL"; +pub const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; +pub const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; +pub const API_KEY_PLACEMENT: CredentialPlacement = CredentialPlacement::Header("x-api-key"); +const API_KEY_HEADER: &str = API_KEY_PLACEMENT.header_name(); +const AUTHORIZATION: &str = CredentialPlacement::Bearer.header_name(); +const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct SupportedEffortTiers { @@ -111,42 +114,207 @@ impl AnthropicModelCapabilities { } } -pub fn is_anthropic_oauth_key(value: &str) -> bool { - value - .strip_prefix("Bearer ") - .unwrap_or(value) - .starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) +pub fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) } -pub fn split_beta_values(header: Option<&str>) -> impl Iterator + '_ { - header - .into_iter() - .flat_map(|value| value.split(',')) - .map(str::trim) - .filter(|piece| !piece.is_empty()) +pub fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option, name: &str) -> Option { + env_lookup(name).filter(|value| !value.trim().is_empty()) +} + +/// An Anthropic OAuth access token, which authenticates as a bearer instead of an `x-api-key`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OauthToken<'a>(&'a str); + +impl<'a> OauthToken<'a> { + /// The raw token, as a caller passes it in `api_key`. + pub fn parse(value: &'a str) -> Option { + value + .starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) + .then_some(Self(value)) + } + + /// A configured key, which users paste either raw or already prefixed with `Bearer `. + pub fn parse_key(value: &'a str) -> Option { + Self::parse(value.strip_prefix("Bearer ").unwrap_or(value)) + } + + pub fn as_str(self) -> &'a str { + self.0 + } + + pub fn into_auth(self) -> AuthScheme { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(self.0), + } + } +} + +/// Python's `AnthropicModelInfo.get_api_key`: the param, else `ANTHROPIC_API_KEY`. +pub fn get_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + non_empty(api_key) .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_API_KEY_ENV)) } -pub fn join_beta_values(values: impl IntoIterator) -> String { - let mut values: Vec = values.into_iter().collect(); - values.sort(); - values.dedup(); - values.join(",") +pub fn get_auth_token(env_lookup: &dyn Fn(&str) -> Option) -> Option { + non_empty_env(env_lookup, ANTHROPIC_AUTH_TOKEN_ENV) } -pub fn is_tool_search_used(tools: Option<&[Value]>) -> bool { - tools.into_iter().flatten().any(|tool| { - tool.get("type") - .and_then(Value::as_str) - .is_some_and(|tool_type| ANTHROPIC_TOOL_SEARCH_TOOL_TYPES.contains(&tool_type)) +/// Python's `AnthropicModelInfo.get_auth_header`, naming the credential instead of building +/// the header: the key goes in `x-api-key` unless it is an OAuth token, and without a key +/// `ANTHROPIC_AUTH_TOKEN` is sent as a bearer. +pub fn get_auth_header( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + if let Some(key) = get_api_key(api_key, env_lookup) { + return Some(match OauthToken::parse_key(&key) { + Some(token) => token.into_auth(), + None => AuthScheme::Credential { + placement: API_KEY_PLACEMENT, + secret: SecretValue::new(key), + }, + }); + } + get_auth_token(env_lookup).map(|token| AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), }) } -pub fn has_advisor_tool(tools: Option<&[Value]>) -> bool { +pub fn resolve_anthropic_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + get_api_key(api_key, env_lookup).ok_or(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }) +} + +/// Whether the caller already forwarded an Anthropic credential, in either header. +pub fn has_anthropic_credential(headers: &[(String, String)]) -> bool { + has_header(headers, API_KEY_HEADER) || has_header(headers, AUTHORIZATION) +} + +pub fn resolve_anthropic_api_base( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + non_empty(api_base) + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_API_BASE_ENV)) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_BASE_URL_ENV)) + .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()) +} + +pub fn complete_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + let api_base = resolve_anthropic_api_base(api_base, env_lookup); + + let api_base = api_base.trim_end_matches('/'); + if api_base.ends_with(MESSAGES_PATH_SUFFIX) { + return api_base.to_string(); + } + format!("{api_base}{MESSAGES_PATH_SUFFIX}") +} + +pub fn existing_betas(headers: &[(String, String)]) -> BetaSet { + headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case(BETA_HEADER)) + .flat_map(|(_, value)| { + value + .parse::() + .unwrap_or_else(|never| match never {}) + }) + .collect() +} + +/// Python's `_merge_beta_headers`, over every casing of the header at once: the union of what +/// the caller sent and `added` replaces the header, sorted and deduplicated. Headers without +/// any beta value stay as they are. +pub fn merge_beta_headers(headers: Headers, added: BetaSet) -> Headers { + let merged = existing_betas(&headers).union(added); + if merged.is_empty() { + return headers; + } + without_headers(headers, &[BETA_HEADER]) + .into_iter() + .chain([(BETA_HEADER.to_string(), merged.to_string())]) + .collect() +} + +/// The outcome of Python's `optionally_handle_anthropic_oauth`. +#[derive(Clone, Debug, PartialEq)] +pub enum OauthHandling { + /// An OAuth token is the whole credential. The headers carry its companions and no + /// longer any `x-api-key` or `authorization`, so the bearer is applied on top. + Bearer { + headers: Headers, + token: SecretValue, + }, + Untouched(Headers), +} + +/// The OAuth token a caller forwarded as `Authorization: Bearer sk-ant-oat…`. +pub fn forwarded_oauth_bearer(headers: &[(String, String)]) -> Option> { + header_value(headers, AUTHORIZATION) + .and_then(|value| value.strip_prefix("Bearer ")) + .and_then(OauthToken::parse) +} + +fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers { + merge_beta_headers( + without_headers(headers, dropped), + BetaSet::from_iter([AnthropicBeta::Oauth20250420]), + ) + .into_iter() + .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) + .collect() +} + +pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>) -> OauthHandling { + if let Some(token) = + forwarded_oauth_bearer(&headers).map(|token| SecretValue::new(token.as_str())) + { + return OauthHandling::Bearer { + headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), + token, + }; + } + if let Some(token) = api_key.and_then(OauthToken::parse) { + return OauthHandling::Bearer { + headers: with_oauth_companions(headers, &[API_KEY_HEADER]), + token: SecretValue::new(token.as_str()), + }; + } + OauthHandling::Untouched(headers) +} + +pub fn is_tool_search_used(tools: Option<&[Recognized]>) -> bool { + tools.into_iter().flatten().any(|tool| { + matches!( + tool, + Recognized::Known( + AnthropicTool::ToolSearchRegex { .. } | AnthropicTool::ToolSearchBm25 { .. } + ) + ) + }) +} + +pub fn has_advisor_tool(tools: Option<&[Recognized]>) -> bool { tools .into_iter() .flatten() - .any(|tool| tool.get("type").and_then(Value::as_str) == Some(ANTHROPIC_ADVISOR_TOOL_TYPE)) + .any(|tool| matches!(tool, Recognized::Known(AnthropicTool::Advisor { .. }))) } pub fn requires_native_compaction_beta( @@ -521,8 +689,97 @@ mod tests { serde_json::from_value(messages).unwrap() } - fn tools(value: Option) -> Option> { - value.map(|tools| tools.as_array().unwrap().clone()) + fn tools(value: Option) -> Option>> { + value.map(|tools| serde_json::from_value(tools).unwrap()) + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + fn betas(values: &[&str]) -> BetaSet { + values.join(",").parse().unwrap() + } + + fn env(vars: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + vars.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + const BOTH_BASE_ENVS: &[(&str, &str)] = &[ + (ANTHROPIC_API_BASE_ENV, "https://api-base.example.com"), + (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com"), + ]; + + #[rstest] + #[case::public_endpoint_by_default(None, &[], "https://api.anthropic.com")] + #[case::explicit_api_base_beats_env( + Some("https://explicit.example.com"), + BOTH_BASE_ENVS, + "https://explicit.example.com" + )] + #[case::explicit_api_base_is_trimmed( + Some(" https://explicit.example.com "), + &[], + "https://explicit.example.com" + )] + #[case::blank_api_base_falls_back_to_env( + Some(" "), + BOTH_BASE_ENVS, + "https://api-base.example.com" + )] + #[case::api_base_env_beats_base_url_env(None, BOTH_BASE_ENVS, "https://api-base.example.com")] + #[case::base_url_env_without_api_base_env( + None, + &[(ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_api_base_env_falls_back_to_base_url_env( + None, + &[(ANTHROPIC_API_BASE_ENV, " \t "), (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_envs_fall_back_to_public_endpoint( + None, + &[(ANTHROPIC_API_BASE_ENV, ""), (ANTHROPIC_BASE_URL_ENV, " ")], + "https://api.anthropic.com" + )] + fn api_base_resolution( + #[case] api_base: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: &str, + ) { + assert_eq!(resolve_anthropic_api_base(api_base, &env(vars)), expected); + } + + #[rstest] + #[case::forwarded_api_key(&[("X-Api-Key", "k")], true)] + #[case::forwarded_bearer(&[("Authorization", "Bearer t")], true)] + #[case::nothing_forwarded(&[("anthropic-version", "2023-06-01")], false)] + fn forwarded_credential_is_detected_in_either_header( + #[case] forwarded: &[(&str, &str)], + #[case] expected: bool, + ) { + let headers: Headers = forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + assert_eq!(has_anthropic_credential(&headers), expected); + } + + fn credential(auth: Option) -> Option<(&'static str, String)> { + auth.map(|auth| match auth { + AuthScheme::Credential { placement, secret } => { + (placement.header_name(), secret.expose().to_string()) + } + other => panic!("expected a credential, got {other:?}"), + }) } fn tagged(encrypted: &str) -> String { @@ -1195,50 +1452,268 @@ mod tests { assert_eq!(twice, once); } + const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; + const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; + const REGULAR_KEY: &str = "sk-ant-api03-regular"; + const OAUTH_BETA: &str = "oauth-2025-04-20"; + const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); + #[rstest] - #[case::no_existing_header(None, "b", "b")] - #[case::empty_existing_header(Some(""), "b", "b")] - #[case::whitespace_existing_header(Some(" "), "b", "b")] - #[case::sorted_after_merge(Some("c,a"), "b", "a,b,c")] - #[case::already_present(Some("a,b"), "a", "a,b")] - #[case::trimmed_and_deduplicated(Some("b, a ,b"), "c", "a,b,c")] - #[case::blank_pieces_skipped(Some("a,,b"), "c", "a,b,c")] - fn beta_values_merge_sorted_and_deduplicated( - #[case] existing: Option<&str>, - #[case] new_beta: &str, - #[case] expected: &str, + #[case::no_beta_header(&[("x-api-key", "k")], &[], &[("x-api-key", "k")])] + #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], &[], &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] + #[case::added_to_no_header(&[("x-api-key", "k")], &["b"], &[("x-api-key", "k"), ("anthropic-beta", "b")])] + #[case::added_to_blank_header(&[("anthropic-beta", " ")], &["b"], &[("anthropic-beta", "b")])] + #[case::sorted_after_merge(&[("anthropic-beta", "c,a")], &["b"], &[("anthropic-beta", "a,b,c")])] + #[case::already_present(&[("anthropic-beta", "a,b")], &["a"], &[("anthropic-beta", "a,b")])] + #[case::existing_normalized_without_additions( + &[("Anthropic-Beta", "b, a ,b"), ("x-api-key", "k")], + &[], + &[("x-api-key", "k"), ("anthropic-beta", "a,b")] + )] + #[case::every_casing_unioned_into_one_lowercase_header( + &[("anthropic-beta", "a"), ("ANTHROPIC-BETA", "c"), ("x-api-key", "k")], + &["b"], + &[("x-api-key", "k"), ("anthropic-beta", "a,b,c")] + )] + fn merge_beta_headers_replaces_the_header_with_the_sorted_union( + #[case] input: &[(&str, &str)], + #[case] added: &[&str], + #[case] expected: &[(&str, &str)], ) { assert_eq!( - join_beta_values(split_beta_values(existing).chain([new_beta.to_string()])), + merge_beta_headers(headers(input), betas(added)), + headers(expected) + ); + } + + #[rstest] + #[case::raw_token(OAUTH_TOKEN, Some(OAUTH_TOKEN))] + #[case::bare_prefix(ANTHROPIC_OAUTH_TOKEN_PREFIX, Some(ANTHROPIC_OAUTH_TOKEN_PREFIX))] + #[case::bearer_token(OAUTH_BEARER, None)] + #[case::api_key(REGULAR_KEY, None)] + #[case::empty("", None)] + #[case::uppercase_prefix("sk-ant-OAT01-abc123", None)] + #[case::prefix_not_at_start(" sk-ant-oat01-abc123", None)] + fn oauth_token_parses_only_the_raw_token(#[case] value: &str, #[case] expected: Option<&str>) { + assert_eq!(OauthToken::parse(value).map(OauthToken::as_str), expected); + } + + #[rstest] + #[case::raw_token(OAUTH_TOKEN, Some(OAUTH_TOKEN))] + #[case::bearer_token(OAUTH_BEARER, Some(OAUTH_TOKEN))] + #[case::api_key(REGULAR_KEY, None)] + #[case::bearer_api_key("Bearer sk-ant-api01-abc123", None)] + #[case::empty("", None)] + #[case::shouting_prefix("SK-ANT-OAT01-abc123", None)] + #[case::lowercase_bearer("bearer sk-ant-oat01-abc123", None)] + #[case::bearer_stripped_once("Bearer Bearer sk-ant-oat01-abc123", None)] + fn oauth_key_parses_the_token_behind_an_optional_bearer( + #[case] value: &str, + #[case] expected: Option<&str>, + ) { + assert_eq!( + OauthToken::parse_key(value).map(OauthToken::as_str), expected ); } #[rstest] - #[case::raw_token("sk-ant-oat01-abc123", true)] - #[case::bearer_token("Bearer sk-ant-oat02-xyz789", true)] - #[case::bare_prefix(ANTHROPIC_OAUTH_TOKEN_PREFIX, true)] - #[case::api_key("sk-ant-api01-abc123", false)] - #[case::bearer_api_key("Bearer sk-ant-api01-abc123", false)] - #[case::empty("", false)] - #[case::uppercase_prefix("sk-ant-OAT01-abc123", false)] - #[case::shouting_prefix("SK-ANT-OAT01-abc123", false)] - #[case::lowercase_bearer("bearer sk-ant-oat01-abc123", false)] - #[case::bearer_stripped_once("Bearer Bearer sk-ant-oat01-abc123", false)] - #[case::prefix_not_at_start(" sk-ant-oat01-abc123", false)] - fn anthropic_oauth_key_detection(#[case] value: &str, #[case] expected: bool) { - assert_eq!(is_anthropic_oauth_key(value), expected); + #[case::bearer(&[("authorization", OAUTH_BEARER)], Some(OAUTH_TOKEN))] + #[case::uppercase_header(&[("AUTHORIZATION", OAUTH_BEARER)], Some(OAUTH_TOKEN))] + #[case::non_oauth_bearer(&[("authorization", "Bearer some-proxy-token")], None)] + #[case::token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] + #[case::lowercase_bearer_scheme(&[("authorization", "bearer sk-ant-oat01-token")], None)] + #[case::token_in_x_api_key(&[("x-api-key", OAUTH_TOKEN)], None)] + #[case::no_headers(&[], None)] + fn forwarded_oauth_bearer_reads_the_authorization_header( + #[case] forwarded: &[(&str, &str)], + #[case] expected: Option<&str>, + ) { + assert_eq!( + forwarded_oauth_bearer(&headers(forwarded)).map(OauthToken::as_str), + expected + ); } #[rstest] - #[case::regex_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0], "name": "tool_search_tool_regex"}])), true)] - #[case::bm25_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1], "name": "tool_search_tool_bm25"}])), true)] + #[case::forwarded_bearer_drops_forwarded_and_deployment_keys( + &[("X-Api-Key", REGULAR_KEY), ("Authorization", OAUTH_BEARER)], + Some(REGULAR_KEY), + &[], + )] + #[case::forwarded_bearer_keeps_unrelated_headers_in_place( + &[("anthropic-version", "2023-06-01"), ("authorization", OAUTH_BEARER)], + None, + &[("anthropic-version", "2023-06-01")], + )] + #[case::forwarded_bearer_wins_over_an_oauth_api_key( + &[("authorization", OAUTH_BEARER)], + Some("sk-ant-oat01-deployment"), + &[], + )] + #[case::api_key_alone(&[], Some(OAUTH_TOKEN), &[])] + #[case::api_key_removes_a_forwarded_x_api_key(&[("x-api-key", OAUTH_TOKEN)], Some(OAUTH_TOKEN), &[])] + #[case::api_key_keeps_a_forwarded_non_oauth_bearer( + &[("Authorization", "Bearer some-proxy-token")], + Some(OAUTH_TOKEN), + &[("Authorization", "Bearer some-proxy-token")], + )] + fn oauth_token_is_the_whole_credential( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] kept: &[(&str, &str)], + ) { + let expected = kept + .iter() + .copied() + .chain([("anthropic-beta", OAUTH_BETA), BROWSER_ACCESS]) + .collect::>(); + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Bearer { + headers: headers(&expected), + token: SecretValue::new(OAUTH_TOKEN), + } + ); + } + + #[rstest] + #[case::forwarded_bearer_merges_a_differently_cased_beta_header( + &[("Anthropic-Beta", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::forwarded_bearer_dedupes_an_existing_oauth_beta( + &[("anthropic-beta", "web-search-2025-03-05, oauth-2025-04-20"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::api_key_merges_the_existing_beta_header( + &[("anthropic-beta", " web-search-2025-03-05 ,")], + Some(OAUTH_TOKEN), + )] + #[case::forwarded_bearer_unions_every_beta_header_casing( + &[("anthropic-beta", "oauth-2025-04-20"), ("ANTHROPIC-BETA", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + fn oauth_beta_merges_into_existing_betas( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + ) { + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Bearer { + headers: headers(&[ + ("anthropic-beta", "oauth-2025-04-20,web-search-2025-03-05"), + BROWSER_ACCESS, + ]), + token: SecretValue::new(OAUTH_TOKEN), + } + ); + } + + #[rstest] + #[case::x_api_key(&[("x-api-key", "caller-key")], Some("sk-other"))] + #[case::non_oauth_bearer(&[("Authorization", "Bearer some-proxy-token")], Some(REGULAR_KEY))] + #[case::oauth_token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] + #[case::bearer_prefixed_api_key(&[], Some(OAUTH_BEARER))] + #[case::nothing(&[], None)] + fn without_an_oauth_token_the_headers_are_untouched( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + ) { + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Untouched(headers(forwarded)) + ); + } + + #[rstest] + #[case::api_key_param(Some("sk-param"), &[], Some(("x-api-key", "sk-param")))] + #[case::api_key_param_over_env_key_and_auth_token( + Some("sk-param"), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("x-api-key", "sk-param")), + )] + #[case::env_key_without_a_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], Some(("x-api-key", "sk-env")))] + #[case::env_key_when_the_param_is_blank(Some(" "), &[("ANTHROPIC_API_KEY", "sk-env")], Some(("x-api-key", "sk-env")))] + #[case::env_key_over_auth_token( + None, + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("x-api-key", "sk-env")), + )] + #[case::auth_token_as_a_bearer( + None, + &[("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("Authorization", "env-token")), + )] + #[case::auth_token_when_the_env_key_is_blank( + None, + &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("Authorization", "env-token")), + )] + #[case::oauth_param_as_a_bearer(Some(OAUTH_TOKEN), &[], Some(("Authorization", OAUTH_TOKEN)))] + #[case::bearer_prefixed_oauth_env_key_as_a_bearer_once( + None, + &[("ANTHROPIC_API_KEY", OAUTH_BEARER)], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::no_credentials(None, &[], None)] + #[case::blank_everything(Some(""), &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], None)] + fn auth_header_prefers_the_key_then_the_auth_token( + #[case] api_key: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: Option<(&str, &str)>, + ) { + assert_eq!( + credential(get_auth_header(api_key, &env(vars))), + expected.map(|(header, secret)| (header, secret.to_string())) + ); + } + + #[rstest] + #[case::param(Some("sk-param"), &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-param"))] + #[case::blank_param_falls_back_to_env(Some(" "), &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-env"))] + #[case::env_without_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-env"))] + #[case::blank_env_is_missing(None, &[("ANTHROPIC_API_KEY", " ")], Err(()))] + #[case::nothing_is_missing(None, &[], Err(()))] + fn api_key_resolution( + #[case] api_key: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: Result<&str, ()>, + ) { + assert_eq!( + resolve_anthropic_api_key(api_key, &env(vars)).map_err(|error| { + assert!(matches!( + error, + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + )); + }), + expected.map(str::to_string) + ); + } + + #[rstest] + #[case::absent(None, None)] + #[case::blank(Some(" \t "), None)] + #[case::padded(Some(" value "), Some("value"))] + fn non_empty_trims_and_drops_blank_values( + #[case] value: Option<&str>, + #[case] expected: Option<&str>, + ) { + assert_eq!(non_empty(value), expected); + } + + #[rstest] + #[case::regex_tool(Some(json!([{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}])), true)] + #[case::bm25_tool(Some(json!([{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}])), true)] #[case::after_other_tools( - Some(json!([{"name": "get_weather", "input_schema": {}}, {"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1]}])), + Some(json!([{"name": "get_weather", "input_schema": {}}, {"type": "tool_search_tool_bm25_20251119"}])), true )] #[case::function_tool(Some(json!([{"type": "function", "function": {"name": "get_weather"}}])), false)] - #[case::name_without_type(Some(json!([{"name": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0]}])), false)] + #[case::name_without_type(Some(json!([{"name": "tool_search_tool_regex_20251119"}])), false)] #[case::empty_tools(Some(json!([])), false)] #[case::no_tools(None, false)] fn tool_search_detection(#[case] input: Option, #[case] expected: bool) { @@ -1246,8 +1721,8 @@ mod tests { } #[rstest] - #[case::advisor_tool(Some(json!([{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor"}])), true)] - #[case::after_other_tools(Some(json!([{"name": "f", "input_schema": {}}, {"type": ANTHROPIC_ADVISOR_TOOL_TYPE}])), true)] + #[case::advisor_tool(Some(json!([{"type": "advisor_20260301", "name": "advisor"}])), true)] + #[case::after_other_tools(Some(json!([{"name": "f", "input_schema": {}}, {"type": "advisor_20260301"}])), true)] #[case::tool_named_advisor(Some(json!([{"name": "advisor", "input_schema": {}}])), false)] #[case::other_server_tool(Some(json!([{"type": "web_search_20250305", "name": "web_search"}])), false)] #[case::empty_tools(Some(json!([])), false)] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs b/litellm-rust/crates/llms/src/anthropic/messages/headers.rs deleted file mode 100644 index bd1b11be92d..00000000000 --- a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs +++ /dev/null @@ -1,677 +0,0 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; -use litellm_types::{ - llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest, recognized::Recognized, -}; -use serde_json::Value; - -use crate::{ - anthropic::{ - ANTHROPIC_OAUTH_TOKEN_PREFIX, - common_utils::{ - ANTHROPIC_OAUTH_BETA_HEADER, beta, has_advisor_tool, is_anthropic_oauth_key, - is_tool_search_used, join_beta_values, requires_native_compaction_beta, - split_beta_values, - }, - }, - base_llm::{ - anthropic_messages::transformation::Headers, - auth::{AuthScheme, ValidatedEnvironment}, - }, -}; - -const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; -const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; -const BETA_HEADER: &str = "anthropic-beta"; -const AUTHORIZATION: &str = "authorization"; -const API_KEY_HEADER: &str = "x-api-key"; -const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; - -fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { - headers - .iter() - .find(|(header, _)| header.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) -} - -fn without(headers: Headers, names: &[&str]) -> Headers { - headers - .into_iter() - .filter(|(header, _)| !names.iter().any(|name| header.eq_ignore_ascii_case(name))) - .collect() -} - -fn existing_betas(headers: &[(String, String)]) -> impl Iterator + '_ { - headers - .iter() - .filter(|(header, _)| header.eq_ignore_ascii_case(BETA_HEADER)) - .flat_map(|(_, value)| split_beta_values(Some(value))) -} - -/// The OAuth headers Python's `optionally_handle_anthropic_oauth` sets next to the bearer. -fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers { - let beta = - join_beta_values(existing_betas(&headers).chain([ANTHROPIC_OAUTH_BETA_HEADER.to_string()])); - without(headers, &[dropped, &[BETA_HEADER]].concat()) - .into_iter() - .chain([ - (BETA_HEADER.to_string(), beta), - (DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string()), - ]) - .collect() -} - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn bearer(token: &str) -> AuthScheme { - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: SecretValue::new(token), - } -} - -pub fn validate_environment( - headers: Headers, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - if let Some(token) = header_value(&headers, AUTHORIZATION) - .and_then(|forwarded| forwarded.strip_prefix("Bearer ")) - .filter(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) - { - let auth = bearer(token); - return Ok(ValidatedEnvironment { - headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), - auth, - }); - } - if let Some(key) = api_key.filter(|key| key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) { - return Ok(ValidatedEnvironment { - headers: with_oauth_companions(headers, &[API_KEY_HEADER]), - auth: bearer(key), - }); - } - if header_value(&headers, API_KEY_HEADER).is_some() - || header_value(&headers, AUTHORIZATION).is_some() - { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Forwarded, - }); - } - let resolved_key = non_empty(api_key) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())); - let auth = match resolved_key { - Some(key) if is_anthropic_oauth_key(&key) => bearer(&key), - Some(key) => AuthScheme::Credential { - placement: CredentialPlacement::Header(API_KEY_HEADER), - secret: SecretValue::new(key), - }, - None => match env_lookup(ANTHROPIC_AUTH_TOKEN_ENV).filter(|value| !value.trim().is_empty()) - { - Some(token) => bearer(&token), - None => { - return Err(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }); - } - }, - }; - Ok(ValidatedEnvironment { headers, auth }) -} - -fn context_management_betas( - context_management: Option<&Value>, -) -> impl Iterator { - let edits = context_management - .and_then(|value| value.get("edits")) - .and_then(Value::as_array) - .map(Vec::as_slice) - .unwrap_or(&[]); - let (compact, other) = edits.iter().fold((false, false), |(compact, other), edit| { - match edit.get("type").and_then(Value::as_str) { - Some("compact_20260112") => (true, other), - _ => (compact, true), - } - }); - compact - .then_some(beta::COMPACT_2026_01_12) - .into_iter() - .chain(other.then_some(beta::CONTEXT_MANAGEMENT_2025_06_27)) -} - -fn uses_structured_output(request: &AnthropicMessagesRequest) -> bool { - request.params.output_format.is_some() - || request - .params - .output_config - .as_ref() - .and_then(Recognized::known) - .is_some_and(|config| config.format.is_some()) -} - -fn messages_carry_output_config(request: &AnthropicMessagesRequest) -> bool { - request - .messages - .iter() - .any(|message| message.extra.contains_key("output_config")) -} - -pub fn feature_betas(request: &AnthropicMessagesRequest) -> Vec<&'static str> { - let tools = request.params.tools.as_deref(); - [ - requires_native_compaction_beta(request.params.compaction.as_ref(), &request.messages) - .then_some(beta::COMPACT_2026_09_04), - uses_structured_output(request).then_some(beta::STRUCTURED_OUTPUT), - (request.params.speed.as_deref() == Some("fast")).then_some(beta::FAST_MODE_2026_02_01), - messages_carry_output_config(request).then_some(beta::PER_TURN_CONTROL_2026_07_01), - has_advisor_tool(tools).then_some(beta::ADVISOR_TOOL_2026_03_01), - is_tool_search_used(tools).then_some(beta::ADVANCED_TOOL_USE_2025_11_20), - ] - .into_iter() - .flatten() - .chain(context_management_betas( - request.params.context_management.as_ref(), - )) - .collect() -} - -pub fn with_feature_betas(headers: Headers, request: &AnthropicMessagesRequest) -> Headers { - let existing = existing_betas(&headers).collect::>(); - let features = feature_betas(request); - if existing.is_empty() && features.is_empty() { - return headers; - } - let merged = join_beta_values( - existing - .into_iter() - .chain(features.into_iter().map(str::to_string)), - ); - without(headers, &[BETA_HEADER]) - .into_iter() - .chain([(BETA_HEADER.to_string(), merged)]) - .collect() -} - -#[cfg(test)] -mod tests { - use rstest::{fixture, rstest}; - use serde_json::json; - - use super::*; - use crate::base_llm::auth::resolve_auth; - - const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; - const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; - const REGULAR_KEY: &str = "sk-ant-api03-regular"; - const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); - - type Env = &'static [(&'static str, &'static str)]; - - fn request(fields: Value) -> AnthropicMessagesRequest { - let mut body = - json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); - body.as_object_mut() - .unwrap() - .extend(fields.as_object().unwrap().clone()); - serde_json::from_value(body).unwrap() - } - - fn headers(pairs: &[(&str, &str)]) -> Headers { - pairs - .iter() - .map(|(name, value)| (name.to_string(), value.to_string())) - .collect() - } - - fn betas(values: &[&str]) -> String { - values.join(",") - } - - #[fixture] - fn no_env() -> Env { - &[] - } - - #[fixture] - fn full_env() -> Env { - &[ - ("ANTHROPIC_API_KEY", "sk-env"), - ("ANTHROPIC_AUTH_TOKEN", "env-token"), - ] - } - - fn authenticate_with( - forwarded: &[(&str, &str)], - api_key: Option<&str>, - env: Env, - ) -> Result { - let lookup = |name: &str| { - env.iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()) - }; - let validated = validate_environment(headers(forwarded), api_key, &lookup)?; - let resolved = tokio::runtime::Builder::new_current_thread() - .build() - .unwrap() - .block_on(resolve_auth( - &litellm_auth::AuthServices::default(), - validated, - &lookup, - )) - .unwrap(); - Ok(resolved.headers) - } - - #[rstest] - #[case::forwarded_bearer_drops_forwarded_and_deployment_keys( - &[("X-Api-Key", REGULAR_KEY), ("Authorization", OAUTH_BEARER)], - Some(REGULAR_KEY), - OAUTH_BEARER, - &[], - )] - #[case::forwarded_bearer_in_uppercase_authorization_header( - &[("AUTHORIZATION", OAUTH_BEARER)], - None, - OAUTH_BEARER, - &[], - )] - #[case::forwarded_bearer_keeps_unrelated_headers_in_place( - &[("anthropic-version", "2023-06-01"), ("authorization", OAUTH_BEARER)], - None, - OAUTH_BEARER, - &[("anthropic-version", "2023-06-01")], - )] - #[case::forwarded_bearer_wins_over_an_oauth_api_key( - &[("authorization", OAUTH_BEARER)], - Some("sk-ant-oat01-deployment"), - OAUTH_BEARER, - &[], - )] - #[case::api_key_authenticates_as_a_bearer(&[], Some(OAUTH_TOKEN), OAUTH_BEARER, &[])] - #[case::api_key_removes_a_forwarded_x_api_key( - &[("x-api-key", OAUTH_TOKEN)], - Some(OAUTH_TOKEN), - OAUTH_BEARER, - &[], - )] - #[case::api_key_replaces_a_forwarded_non_oauth_bearer( - &[("Authorization", "Bearer some-proxy-token")], - Some(OAUTH_TOKEN), - OAUTH_BEARER, - &[], - )] - fn oauth_token_is_the_whole_credential( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - #[case] expected_bearer: &str, - #[case] kept: &[(&str, &str)], - full_env: Env, - ) { - let expected = kept - .iter() - .copied() - .chain([ - ("anthropic-beta", ANTHROPIC_OAUTH_BETA_HEADER), - BROWSER_ACCESS, - ("authorization", expected_bearer), - ]) - .collect::>(); - assert_eq!( - authenticate_with(forwarded, api_key, full_env).unwrap(), - headers(&expected) - ); - } - - #[rstest] - #[case::forwarded_bearer_merges_a_differently_cased_beta_header( - &[("Anthropic-Beta", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], - None, - )] - #[case::forwarded_bearer_dedupes_an_existing_oauth_beta( - &[("anthropic-beta", "web-search-2025-03-05, oauth-2025-04-20"), ("authorization", OAUTH_BEARER)], - None, - )] - #[case::api_key_merges_the_existing_beta_header( - &[("anthropic-beta", " web-search-2025-03-05 ,")], - Some(OAUTH_TOKEN), - )] - #[case::forwarded_bearer_unions_every_beta_header_casing( - &[("anthropic-beta", "oauth-2025-04-20"), ("ANTHROPIC-BETA", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], - None, - )] - fn oauth_beta_merges_into_existing_betas( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - no_env: Env, - ) { - assert_eq!( - authenticate_with(forwarded, api_key, no_env).unwrap(), - headers(&[ - ( - "anthropic-beta", - &betas(&[ANTHROPIC_OAUTH_BETA_HEADER, "web-search-2025-03-05"]) - ), - BROWSER_ACCESS, - ("authorization", OAUTH_BEARER), - ]) - ); - } - - #[rstest] - #[case::x_api_key_over_the_deployment_key(&[("x-api-key", "caller-key")], Some("sk-other"))] - #[case::uppercase_x_api_key(&[("X-API-KEY", "caller-key")], None)] - #[case::non_oauth_bearer(&[("Authorization", "Bearer some-proxy-token")], None)] - #[case::non_oauth_bearer_over_a_regular_api_key( - &[("authorization", "Bearer sk-ant-api03-forwarded")], - Some(REGULAR_KEY), - )] - #[case::oauth_token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] - #[case::oauth_token_behind_a_lowercase_bearer_scheme( - &[("authorization", "bearer sk-ant-oat01-token")], - None, - )] - fn forwarded_auth_header_is_kept_untouched( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - full_env: Env, - ) { - assert_eq!( - authenticate_with(forwarded, api_key, full_env).unwrap(), - headers(forwarded) - ); - } - - #[rstest] - #[case::api_key_param(Some("sk-param"), &[], ("x-api-key", "sk-param"))] - #[case::api_key_param_over_env_key_and_auth_token( - Some("sk-param"), - &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("x-api-key", "sk-param"), - )] - #[case::env_key_without_a_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] - #[case::env_key_when_the_param_is_empty(Some(""), &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] - #[case::env_key_when_the_param_is_whitespace( - Some(" "), - &[("ANTHROPIC_API_KEY", "sk-env")], - ("x-api-key", "sk-env"), - )] - #[case::env_key_over_auth_token( - None, - &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("x-api-key", "sk-env"), - )] - #[case::auth_token_as_a_bearer( - None, - &[("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("authorization", "Bearer env-token"), - )] - #[case::auth_token_when_the_env_key_is_whitespace( - None, - &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("authorization", "Bearer env-token"), - )] - #[case::oauth_env_key_as_a_plain_bearer( - None, - &[("ANTHROPIC_API_KEY", "sk-ant-oat01-env")], - ("authorization", "Bearer sk-ant-oat01-env"), - )] - fn credential_is_resolved_after_the_existing_headers( - #[case] api_key: Option<&str>, - #[case] env: Env, - #[case] expected: (&str, &str), - ) { - let forwarded = [("anthropic-beta", "web-search-2025-03-05")]; - assert_eq!( - authenticate_with(&forwarded, api_key, env).unwrap(), - headers(&[forwarded[0], expected]) - ); - } - - #[rstest] - #[case::no_credentials(&[], None, &[])] - #[case::empty_api_key(&[], Some(""), &[])] - #[case::whitespace_only_env_values( - &[], - None, - &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], - )] - #[case::unrelated_forwarded_headers(&[("anthropic-beta", "web-search-2025-03-05")], None, &[])] - fn missing_credentials_are_an_auth_error( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - #[case] env: Env, - ) { - assert!(matches!( - authenticate_with(forwarded, api_key, env), - Err(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - }) - )); - } - - #[rstest] - #[case::no_features(json!({}), &[])] - #[case::output_format(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] - #[case::null_output_format(json!({"output_format": null}), &[])] - #[case::output_config_format( - json!({"output_config": {"format": {"type": "json_schema"}, "effort": "xhigh"}}), - &[beta::STRUCTURED_OUTPUT] - )] - #[case::null_output_config_format(json!({"output_config": {"format": null}}), &[])] - #[case::top_level_output_config_without_format(json!({"output_config": {"effort": "high"}}), &[])] - #[case::fast_speed(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] - #[case::standard_speed(json!({"speed": "standard"}), &[])] - #[case::compaction_param(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] - #[case::empty_compaction_param(json!({"compaction": {}}), &[beta::COMPACT_2026_09_04])] - #[case::signed_compaction_block_in_history( - json!({"messages": [ - {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, - {"role": "user", "content": "Continue"}, - ]}), - &[beta::COMPACT_2026_09_04] - )] - #[case::unsigned_compaction_block_in_history( - json!({"messages": [ - {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": ""}]}, - {"role": "user", "content": "Continue"}, - ]}), - &[] - )] - #[case::advisor_tool( - json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}]}), - &[beta::ADVISOR_TOOL_2026_03_01] - )] - #[case::no_tools(json!({"tools": []}), &[])] - #[case::regex_tool_search( - json!({"tools": [{"type": "tool_search_tool_regex_20251119"}]}), - &[beta::ADVANCED_TOOL_USE_2025_11_20] - )] - #[case::bm25_tool_search( - json!({"tools": [{"type": "tool_search_tool_bm25_20251119"}]}), - &[beta::ADVANCED_TOOL_USE_2025_11_20] - )] - #[case::unrelated_server_tool(json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), &[])] - #[case::only_compact_edits( - json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), - &[beta::COMPACT_2026_01_12] - )] - #[case::only_other_edits( - json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}]}}), - &[beta::CONTEXT_MANAGEMENT_2025_06_27] - )] - #[case::compact_and_other_edits( - json!({"context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_tool_uses_20250919"}]}}), - &[beta::COMPACT_2026_01_12, beta::CONTEXT_MANAGEMENT_2025_06_27] - )] - #[case::edit_without_a_type(json!({"context_management": {"edits": [{}]}}), &[beta::CONTEXT_MANAGEMENT_2025_06_27])] - #[case::empty_edits(json!({"context_management": {"edits": []}}), &[])] - #[case::context_management_without_edits(json!({"context_management": {}}), &[])] - #[case::per_message_output_config( - json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] - )] - #[case::per_message_null_output_config( - json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] - )] - fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { - assert_eq!(feature_betas(&request(fields)), expected); - } - - #[rstest] - #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}))] - #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}))] - fn headers_without_any_beta_value_are_untouched( - #[case] input: &[(&str, &str)], - #[case] fields: Value, - ) { - assert_eq!( - with_feature_betas(headers(input), &request(fields)), - headers(input) - ); - } - - #[rstest] - #[case::feature_beta_is_appended( - &[("x-api-key", "k")], - json!({"speed": "fast"}), - &[("x-api-key", "k"), ("anthropic-beta", beta::FAST_MODE_2026_02_01)], - )] - #[case::existing_betas_are_normalized_without_features( - &[("Anthropic-Beta", "web-search-2025-03-05, interleaved-thinking-2025-05-14 ,web-search-2025-03-05"), ("x-api-key", "k")], - json!({}), - &[("x-api-key", "k"), ("anthropic-beta", "interleaved-thinking-2025-05-14,web-search-2025-03-05")], - )] - #[case::existing_advisor_beta_is_kept_without_an_advisor_tool( - &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], - json!({"tools": []}), - &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], - )] - #[case::feature_already_sent_is_not_duplicated( - &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], - json!({"speed": "fast"}), - &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], - )] - fn feature_betas_merge_into_the_headers( - #[case] input: &[(&str, &str)], - #[case] fields: Value, - #[case] expected: &[(&str, &str)], - ) { - assert_eq!( - with_feature_betas(headers(input), &request(fields)), - headers(expected) - ); - } - - #[test] - fn differently_cased_beta_header_is_replaced_by_one_sorted_header() { - let merged = with_feature_betas( - headers(&[("Anthropic-Beta", "interleaved-thinking-2025-05-14")]), - &request( - json!({"messages": [{"role": "system", "content": "env", "output_config": {"effort": "low"}}]}), - ), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - "interleaved-thinking-2025-05-14", - beta::PER_TURN_CONTROL_2026_07_01 - ]) - )]) - ); - } - - #[test] - fn every_beta_header_casing_is_unioned_into_one_header() { - let merged = with_feature_betas( - headers(&[ - ("anthropic-beta", "interleaved-thinking-2025-05-14"), - ("Anthropic-Beta", "web-search-2025-03-05"), - ]), - &request(json!({"speed": "fast"})), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - beta::FAST_MODE_2026_02_01, - "interleaved-thinking-2025-05-14", - "web-search-2025-03-05" - ]) - )]) - ); - } - - #[test] - fn unknown_client_betas_survive_alongside_the_added_one() { - let client_betas = [ - "claude-code-20250219", - "interleaved-thinking-2025-05-14", - beta::CONTEXT_MANAGEMENT_2025_06_27, - beta::PER_TURN_CONTROL_2026_07_01, - "effort-2025-11-24", - ]; - let merged = with_feature_betas( - headers(&[("anthropic-beta", &betas(&client_betas))]), - &request( - json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - ), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - "claude-code-20250219", - beta::CONTEXT_MANAGEMENT_2025_06_27, - "effort-2025-11-24", - "interleaved-thinking-2025-05-14", - beta::PER_TURN_CONTROL_2026_07_01, - ]) - )]) - ); - } - - #[test] - fn every_feature_merges_with_the_oauth_beta_sorted_and_last() { - let oauth_headers = authenticate_with(&[], Some(OAUTH_TOKEN), &[]).unwrap(); - let all_features = request(json!({ - "compaction": {"enabled": true}, - "output_format": {"type": "json_schema"}, - "speed": "fast", - "tools": [{"type": "advisor_20260301"}, {"type": "tool_search_tool_bm25_20251119"}], - "context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, - "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], - })); - assert_eq!( - with_feature_betas(oauth_headers, &all_features), - headers(&[ - BROWSER_ACCESS, - ("authorization", OAUTH_BEARER), - ( - "anthropic-beta", - &betas(&[ - beta::ADVANCED_TOOL_USE_2025_11_20, - beta::ADVISOR_TOOL_2026_03_01, - beta::COMPACT_2026_01_12, - beta::COMPACT_2026_09_04, - beta::CONTEXT_MANAGEMENT_2025_06_27, - beta::FAST_MODE_2026_02_01, - ANTHROPIC_OAUTH_BETA_HEADER, - beta::PER_TURN_CONTROL_2026_07_01, - beta::STRUCTURED_OUTPUT, - ]) - ), - ]) - ); - } -} diff --git a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs index 5adf5fda16f..dff4bf18bd5 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs @@ -1,5 +1,4 @@ pub mod handler; -pub mod headers; pub mod streaming_iterator; pub mod thinking; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 280ea63eefa..07c2eb46eaa 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,31 +1,35 @@ +use litellm_auth::CredentialPlacement; use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; -use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessagesOptionalParams, AnthropicMessagesRequest, +use litellm_types::{ + llms::{ + anthropic::{AnthropicBeta, BetaSet}, + anthropic_messages::anthropic_request::{ + AnthropicMessage, AnthropicMessagesOptionalParams, AnthropicMessagesRequest, + ContextEdit, ContextManagement, Speed, + }, + }, + recognized::Recognized, }; use serde_json::{Map, Value, json}; -use super::{ - headers::{validate_environment, with_feature_betas}, - thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}, -}; +use super::thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}; use crate::{ Error, anthropic::common_utils::{ - AnthropicModelCapabilities, has_advisor_tool, strip_advisor_blocks, - strip_encrypted_reasoning_blocks, + ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV, + ANTHROPIC_BASE_URL_ENV, AnthropicModelCapabilities, OauthHandling, complete_anthropic_url, + get_auth_header, has_advisor_tool, has_anthropic_credential, is_tool_search_used, + merge_beta_headers, optionally_handle_anthropic_oauth, requires_native_compaction_beta, + strip_advisor_blocks, strip_encrypted_reasoning_blocks, }, - base_llm::anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, + base_llm::{ + anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, + }, + auth::AuthScheme, }, }; -const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; -const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; -const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; -const ANTHROPIC_BASE_URL_ENV: &str = "ANTHROPIC_BASE_URL"; -const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; -const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; - pub struct AnthropicMessagesConfig; pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; @@ -66,18 +70,15 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { context: &MessagesTransformContext, ) -> Result { if request.params.max_tokens.is_none() { - return Err(Error::InvalidRequest( - "max_tokens is required for Anthropic /v1/messages API".to_string(), - )); + return Err(Error::MissingField("max_tokens")); } let request = drop_unsupported_params(request, context)?; let request = translate_thinking(request, &context.thinking)?; let context_management = request .params .context_management - .as_ref() - .and_then(map_openai_context_management_to_anthropic) - .or_else(|| request.params.context_management.clone()); + .clone() + .map(map_openai_context_management_to_anthropic); let messages = if has_advisor_tool(request.params.tools.as_deref()) { request.messages } else { @@ -102,6 +103,8 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { ] } + /// Python's `validate_anthropic_messages_environment` up to the beta merge, which + /// `request_headers` does once the request is final. fn validate_environment( &self, headers: Headers, @@ -109,14 +112,99 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { _model: &str, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - validate_environment(headers, api_key, env_lookup).map_err(Error::from) + let headers = match optionally_handle_anthropic_oauth(headers, api_key) { + OauthHandling::Bearer { headers, token } => { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: token, + }, + }); + } + OauthHandling::Untouched(headers) => headers, + }; + if has_anthropic_credential(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + let auth = get_auth_header(api_key, env_lookup).ok_or(Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }, + ))?; + Ok(ValidatedEnvironment { headers, auth }) } fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers { - with_feature_betas(headers, request) + update_headers_with_anthropic_beta(headers, request) } } +fn update_headers_with_anthropic_beta( + headers: Headers, + request: &AnthropicMessagesRequest, +) -> Headers { + merge_beta_headers(headers, feature_betas(request)) +} + +fn feature_betas(request: &AnthropicMessagesRequest) -> BetaSet { + let params = &request.params; + let tools = params.tools.as_deref(); + [ + requires_native_compaction_beta(params.compaction.as_ref(), &request.messages) + .then_some(AnthropicBeta::Compact20260904), + uses_structured_output(params).then_some(AnthropicBeta::StructuredOutputs20251113), + (params.speed == Some(Recognized::Known(Speed::Fast))) + .then_some(AnthropicBeta::FastMode20260201), + messages_carry_output_config(&request.messages) + .then_some(AnthropicBeta::PerTurnControl20260701), + has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), + is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), + ] + .into_iter() + .flatten() + .chain(context_management_betas(params.context_management.as_ref())) + .collect() +} + +fn is_compact_edit(edit: &Recognized) -> bool { + matches!(edit, Recognized::Known(ContextEdit::Compact { .. })) +} + +fn context_management_betas( + context_management: Option<&Recognized>, +) -> impl Iterator { + let edits = context_management + .and_then(Recognized::known) + .and_then(|context_management| context_management.edits.as_deref()) + .unwrap_or_default(); + let compact = edits.iter().any(is_compact_edit); + let other = edits.iter().any(|edit| !is_compact_edit(edit)); + compact + .then_some(AnthropicBeta::Compact20260112) + .into_iter() + .chain(other.then_some(AnthropicBeta::ContextManagement20250627)) +} + +fn uses_structured_output(params: &AnthropicMessagesOptionalParams) -> bool { + params.output_format.is_some() + || params + .output_config + .as_ref() + .and_then(Recognized::known) + .is_some_and(|config| config.format.is_some()) +} + +fn messages_carry_output_config(messages: &[AnthropicMessage]) -> bool { + messages + .iter() + .any(|message| message.extra.contains_key("output_config")) +} + fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { Error::InvalidRequest(format!( "{model} does not support {param}={value}. {hint}To drop unsupported params, set `litellm.drop_params = True`." @@ -136,9 +224,9 @@ fn drop_unsupported_params( Err(unsupported_param(&model, param, &value, hint)) }; let params = request.params; - let speed = match params.speed.as_deref() { + let speed = match ¶ms.speed { Some(speed) if !capabilities.supports_speed => { - reject("speed", format!("'{speed}'"), "")?; + reject("speed", format!("'{}'", speed_text(speed)), "")?; None } _ => params.speed.clone(), @@ -178,101 +266,73 @@ fn drop_unsupported_params( }) } -pub fn map_openai_context_management_to_anthropic(context_management: &Value) -> Option { - match context_management { - Value::Object(edits) if edits.contains_key("edits") => Some(context_management.clone()), - Value::Array(entries) => { - let edits: Vec = entries - .iter() - .filter_map(Value::as_object) - .filter(|entry| entry.get("type").and_then(Value::as_str) == Some("compaction")) - .map(|entry| { - let trigger = entry.get("compact_threshold").and_then(Value::as_f64).map( - |threshold| json!({"type": "input_tokens", "value": threshold as i64}), - ); - let passthrough = entry - .iter() - .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) - .map(|(key, value)| (key.clone(), value.clone())); - Value::Object( - [("type".to_string(), json!("compact_20260112"))] - .into_iter() - .chain(trigger.map(|trigger| ("trigger".to_string(), trigger))) - .chain(passthrough) - .collect::>(), - ) - }) - .collect(); - (!edits.is_empty()).then(|| json!({"edits": edits})) - } - _ => None, +fn speed_text(speed: &Recognized) -> String { + match speed { + Recognized::Known(speed) => speed.as_str().to_string(), + Recognized::Unrecognized(Value::String(text)) => text.clone(), + Recognized::Unrecognized(other) => other.to_string(), } } -pub fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -pub fn resolve_anthropic_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - non_empty(api_key) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }) -} - -pub fn complete_anthropic_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - let api_base = resolve_anthropic_api_base(api_base, env_lookup); - - let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with(MESSAGES_PATH_SUFFIX) { - return api_base.to_string(); +fn compact_edit_from_openai(entry: &Map) -> Option { + if entry.get("type").and_then(Value::as_str) != Some("compaction") { + return None; } - format!("{api_base}{MESSAGES_PATH_SUFFIX}") + let trigger = entry + .get("compact_threshold") + .and_then(Value::as_f64) + .map(|threshold| json!({"type": "input_tokens", "value": threshold as i64})); + let passthrough = entry + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) + .map(|(key, value)| (key.clone(), value.clone())); + Some(ContextEdit::Compact { + extra: trigger + .map(|trigger| ("trigger".to_string(), trigger)) + .into_iter() + .chain(passthrough) + .collect(), + }) } -pub fn resolve_anthropic_api_base( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - let env = |name: &str| env_lookup(name).filter(|value| !value.trim().is_empty()); - non_empty(api_base) - .map(str::to_string) - .or_else(|| env(ANTHROPIC_API_BASE_ENV)) - .or_else(|| env(ANTHROPIC_BASE_URL_ENV)) - .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()) +/// An OpenAI-style `context_management` list becomes Anthropic `edits` when it holds +/// compaction entries. Anything else, native edits included, is sent as it came. +pub fn map_openai_context_management_to_anthropic( + context_management: Recognized, +) -> Recognized { + let Recognized::Unrecognized(Value::Array(entries)) = &context_management else { + return context_management; + }; + let edits: Vec> = entries + .iter() + .filter_map(Value::as_object) + .filter_map(compact_edit_from_openai) + .map(Recognized::Known) + .collect(); + if edits.is_empty() { + return context_management; + } + Recognized::Known(ContextManagement { + edits: Some(edits), + extra: Map::new(), + }) } #[cfg(test)] mod tests { use std::process::Command; - use litellm_auth::CredentialPlacement; use rstest::{fixture, rstest}; use super::*; - use crate::{ - anthropic::common_utils::{ENCRYPTED_REASONING_SIGNATURE_PREFIX, beta}, - base_llm::auth::AuthScheme, - }; + use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX; type Env = &'static [(&'static str, &'static str)]; - const BOTH_BASE_ENVS: Env = &[ - (ANTHROPIC_API_BASE_ENV, "https://api-base.example.com"), - (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com"), - ]; - const API_KEY_ENV: Env = &[(ANTHROPIC_API_KEY_ENV, "sk-env")]; - const MISSING_API_KEY: &str = - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"; + const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; + const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; + const OAUTH_BETA: &str = "oauth-2025-04-20"; + const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); const LOW_BUDGET_ENV: &str = "DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET"; const PROCESS_ENV_PROBE: &str = "LITELLM_MESSAGES_TRANSFORM_CONTEXT_PROBE"; @@ -377,7 +437,7 @@ mod tests { fn missing_max_tokens_is_rejected(#[case] fields: Value, unmapped: AnthropicModelCapabilities) { assert_eq!( transform(fields, unmapped, false), - invalid("max_tokens is required for Anthropic /v1/messages API") + Err(Error::MissingField("max_tokens")) ); } @@ -569,17 +629,19 @@ mod tests { #[case::empty_list(json!([]), None)] #[case::anthropic_edits_pass_through( json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]}), - Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]})) + None )] #[case::object_without_edits(json!({"type": "compaction"}), None)] #[case::scalar(json!("compaction"), None)] fn openai_context_management_maps_to_anthropic_edits( #[case] context_management: Value, - #[case] expected: Option, + #[case] mapped: Option, ) { + let parsed: Recognized = + serde_json::from_value(context_management.clone()).unwrap(); assert_eq!( - map_openai_context_management_to_anthropic(&context_management), - expected + serde_json::to_value(map_openai_context_management_to_anthropic(parsed)).unwrap(), + mapped.unwrap_or(context_management) ); } @@ -723,47 +785,6 @@ mod tests { ); } - #[rstest] - #[case::public_endpoint_by_default(None, &[], "https://api.anthropic.com")] - #[case::explicit_api_base_beats_env( - Some("https://explicit.example.com"), - BOTH_BASE_ENVS, - "https://explicit.example.com" - )] - #[case::explicit_api_base_is_trimmed( - Some(" https://explicit.example.com "), - &[], - "https://explicit.example.com" - )] - #[case::blank_api_base_falls_back_to_env( - Some(" "), - BOTH_BASE_ENVS, - "https://api-base.example.com" - )] - #[case::api_base_env_beats_base_url_env(None, BOTH_BASE_ENVS, "https://api-base.example.com")] - #[case::base_url_env_without_api_base_env( - None, - &[(ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], - "https://base-url.example.com" - )] - #[case::blank_api_base_env_falls_back_to_base_url_env( - None, - &[(ANTHROPIC_API_BASE_ENV, " \t "), (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], - "https://base-url.example.com" - )] - #[case::blank_envs_fall_back_to_public_endpoint( - None, - &[(ANTHROPIC_API_BASE_ENV, ""), (ANTHROPIC_BASE_URL_ENV, " ")], - "https://api.anthropic.com" - )] - fn api_base_resolution( - #[case] api_base: Option<&str>, - #[case] vars: Env, - #[case] expected: &str, - ) { - assert_eq!(resolve_anthropic_api_base(api_base, &env(vars)), expected); - } - #[rstest] #[case::public_endpoint(None, &[], "https://api.anthropic.com/v1/messages")] #[case::base_url_env( @@ -794,77 +815,274 @@ mod tests { ); } - #[rstest] - #[case::param_beats_env(Some("sk-param"), API_KEY_ENV, Ok("sk-param"))] - #[case::param_is_trimmed(Some(" sk-param "), &[], Ok("sk-param"))] - #[case::blank_param_falls_back_to_env(Some(" "), API_KEY_ENV, Ok("sk-env"))] - #[case::env_without_param(None, API_KEY_ENV, Ok("sk-env"))] - #[case::blank_env_is_missing(None, &[(ANTHROPIC_API_KEY_ENV, " ")], Err(MISSING_API_KEY))] - #[case::nothing_is_missing(None, &[], Err(MISSING_API_KEY))] - fn api_key_resolution( - #[case] api_key: Option<&str>, - #[case] vars: Env, - #[case] expected: Result<&str, &str>, - ) { - assert_eq!( - resolve_anthropic_api_key(api_key, &env(vars)).map_err(|error| error.to_string()), - expected.map(str::to_string).map_err(str::to_string) - ); + fn betas(values: &[&str]) -> BetaSet { + values.join(",").parse().unwrap() } - #[test] - fn config_reports_a_missing_key_as_an_auth_error() { + fn validated( + forwarded: &[(&str, &str)], + api_key: Option<&str>, + vars: Env, + ) -> Result { + ANTHROPIC_MESSAGES_CONFIG.validate_environment( + headers(forwarded), + api_key, + "claude", + &env(vars), + ) + } + + fn credential(auth: &AuthScheme) -> Option<(&'static str, &str)> { + match auth { + AuthScheme::Credential { placement, secret } => { + Some((placement.header_name(), secret.expose())) + } + AuthScheme::Forwarded => None, + other => panic!("unexpected auth scheme {other:?}"), + } + } + + #[rstest] + #[case::forwarded_oauth_bearer( + &[("anthropic-version", "2023-06-01"), ("X-Api-Key", "sk-caller"), ("Authorization", OAUTH_BEARER)], + Some("sk-deployment"), + &[("ANTHROPIC_API_KEY", "sk-env")], + &[("anthropic-version", "2023-06-01"), ("anthropic-beta", OAUTH_BETA), BROWSER_ACCESS], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::oauth_api_key( + &[("x-api-key", OAUTH_TOKEN), ("anthropic-beta", "web-search-2025-03-05")], + Some(OAUTH_TOKEN), + &[], + &[("anthropic-beta", "oauth-2025-04-20,web-search-2025-03-05"), BROWSER_ACCESS], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::forwarded_x_api_key_is_kept_over_the_deployment_key( + &[("X-API-KEY", "caller-key")], + Some("sk-other"), + &[("ANTHROPIC_API_KEY", "sk-env")], + &[("X-API-KEY", "caller-key")], + None, + )] + #[case::forwarded_non_oauth_bearer_is_kept( + &[("Authorization", "Bearer some-proxy-token")], + Some("sk-ant-api03-regular"), + &[], + &[("Authorization", "Bearer some-proxy-token")], + None, + )] + #[case::oauth_token_without_the_bearer_scheme_is_kept( + &[("authorization", OAUTH_TOKEN)], + None, + &[], + &[("authorization", OAUTH_TOKEN)], + None, + )] + #[case::api_key_param( + &[("anthropic-beta", "web-search-2025-03-05")], + Some("sk-param"), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[("anthropic-beta", "web-search-2025-03-05")], + Some(("x-api-key", "sk-param")), + )] + #[case::env_key_when_the_param_is_blank( + &[], + Some(" "), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[], + Some(("x-api-key", "sk-env")), + )] + #[case::auth_token_when_no_key_is_set( + &[], + None, + &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[], + Some(("Authorization", "env-token")), + )] + #[case::oauth_env_key_as_a_bearer( + &[], + None, + &[("ANTHROPIC_API_KEY", "sk-ant-oat01-env")], + &[], + Some(("Authorization", "sk-ant-oat01-env")), + )] + fn validate_environment_shapes_the_headers_and_names_the_credential( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] vars: Env, + #[case] expected_headers: &[(&str, &str)], + #[case] expected_credential: Option<(&str, &str)>, + ) { + let environment = validated(forwarded, api_key, vars).unwrap(); + assert_eq!(environment.headers, headers(expected_headers)); + assert_eq!(credential(&environment.auth), expected_credential); + } + + #[rstest] + #[case::no_credentials(&[], None, &[])] + #[case::empty_api_key(&[], Some(""), &[])] + #[case::whitespace_only_env_values(&[], None, &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")])] + #[case::unrelated_forwarded_headers(&[("anthropic-beta", "web-search-2025-03-05")], None, &[])] + fn missing_credentials_are_an_auth_error( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] vars: Env, + ) { assert!(matches!( - ANTHROPIC_MESSAGES_CONFIG.validate_environment(vec![], None, "claude", &no_env), + validated(forwarded, api_key, vars), Err(Error::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, + environment_variable: "ANTHROPIC_API_KEY", })) )); } - #[test] - fn config_authenticates_with_the_anthropic_auth_token() { - let validated = ANTHROPIC_MESSAGES_CONFIG - .validate_environment( - vec![], - None, - "claude", - &env(&[("ANTHROPIC_AUTH_TOKEN", "auth-token")]), - ) - .unwrap(); - assert!(matches!( - validated.auth, - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - ref secret - } if secret.expose() == "auth-token" - )); - } - - #[test] - fn config_requests_the_betas_the_request_features_need() { - assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.request_headers( - headers(&[("x-api-key", "sk")]), - &request(json!({"speed": "fast"})) - ), - headers(&[ - ("x-api-key", "sk"), - ("anthropic-beta", beta::FAST_MODE_2026_02_01) - ]) - ); + #[rstest] + #[case::no_features(json!({}), &[])] + #[case::output_format(json!({"output_format": {"type": "json_schema"}}), &["structured-outputs-2025-11-13"])] + #[case::null_output_format(json!({"output_format": null}), &[])] + #[case::output_config_format( + json!({"output_config": {"format": {"type": "json_schema"}, "effort": "xhigh"}}), + &["structured-outputs-2025-11-13"] + )] + #[case::null_output_config_format(json!({"output_config": {"format": null}}), &[])] + #[case::top_level_output_config_without_format(json!({"output_config": {"effort": "high"}}), &[])] + #[case::fast_speed(json!({"speed": "fast"}), &["fast-mode-2026-02-01"])] + #[case::standard_speed(json!({"speed": "standard"}), &[])] + #[case::unknown_speed(json!({"speed": "turbo"}), &[])] + #[case::compaction_param(json!({"compaction": {"enabled": true}}), &["compact-2026-09-04"])] + #[case::empty_compaction_param(json!({"compaction": {}}), &["compact-2026-09-04"])] + #[case::signed_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, + {"role": "user", "content": "Continue"}, + ]}), + &["compact-2026-09-04"] + )] + #[case::unsigned_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": ""}]}, + {"role": "user", "content": "Continue"}, + ]}), + &[] + )] + #[case::advisor_tool( + json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}]}), + &["advisor-tool-2026-03-01"] + )] + #[case::no_tools(json!({"tools": []}), &[])] + #[case::regex_tool_search( + json!({"tools": [{"type": "tool_search_tool_regex_20251119"}]}), + &["advanced-tool-use-2025-11-20"] + )] + #[case::bm25_tool_search( + json!({"tools": [{"type": "tool_search_tool_bm25_20251119"}]}), + &["advanced-tool-use-2025-11-20"] + )] + #[case::unrelated_server_tool(json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), &[])] + #[case::only_compact_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), + &["compact-2026-01-12"] + )] + #[case::only_other_edits( + json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}]}}), + &["context-management-2025-06-27"] + )] + #[case::compact_and_other_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_tool_uses_20250919"}]}}), + &["compact-2026-01-12", "context-management-2025-06-27"] + )] + #[case::edit_without_a_type(json!({"context_management": {"edits": [{}]}}), &["context-management-2025-06-27"])] + #[case::unknown_edit_type(json!({"context_management": {"edits": [{"type": "future"}]}}), &["context-management-2025-06-27"])] + #[case::empty_edits(json!({"context_management": {"edits": []}}), &[])] + #[case::context_management_without_edits(json!({"context_management": {}}), &[])] + #[case::unmapped_openai_context_management(json!({"context_management": [{"type": "other"}]}), &[])] + #[case::per_message_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &["per-turn-control-2026-07-01"] + )] + #[case::per_message_null_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), + &["per-turn-control-2026-07-01"] + )] + fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { + assert_eq!(feature_betas(&request(fields)), betas(expected)); } #[rstest] - #[case::absent(None, None)] - #[case::blank(Some(" \t "), None)] - #[case::padded(Some(" value "), Some("value"))] - fn non_empty_trims_and_drops_blank_values( - #[case] value: Option<&str>, - #[case] expected: Option<&str>, + #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])] + #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] + #[case::feature_beta_is_appended( + &[("x-api-key", "k")], + json!({"speed": "fast"}), + &[("x-api-key", "k"), ("anthropic-beta", "fast-mode-2026-02-01")], + )] + #[case::existing_betas_are_normalized_without_features( + &[("Anthropic-Beta", "web-search-2025-03-05, interleaved-thinking-2025-05-14 ,web-search-2025-03-05"), ("x-api-key", "k")], + json!({}), + &[("x-api-key", "k"), ("anthropic-beta", "interleaved-thinking-2025-05-14,web-search-2025-03-05")], + )] + #[case::existing_advisor_beta_is_kept_without_an_advisor_tool( + &[("anthropic-beta", "advisor-tool-2026-03-01")], + json!({"tools": []}), + &[("anthropic-beta", "advisor-tool-2026-03-01")], + )] + #[case::feature_already_sent_is_not_duplicated( + &[("anthropic-beta", "fast-mode-2026-02-01")], + json!({"speed": "fast"}), + &[("anthropic-beta", "fast-mode-2026-02-01")], + )] + #[case::differently_cased_beta_header_is_replaced_by_one_sorted_header( + &[("Anthropic-Beta", "interleaved-thinking-2025-05-14")], + json!({"messages": [{"role": "system", "content": "env", "output_config": {"effort": "low"}}]}), + &[("anthropic-beta", "interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], + )] + #[case::every_beta_header_casing_is_unioned_into_one_header( + &[("anthropic-beta", "interleaved-thinking-2025-05-14"), ("Anthropic-Beta", "web-search-2025-03-05")], + json!({"speed": "fast"}), + &[("anthropic-beta", "fast-mode-2026-02-01,interleaved-thinking-2025-05-14,web-search-2025-03-05")], + )] + #[case::unknown_client_betas_survive_alongside_the_added_one( + &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14,context-management-2025-06-27,per-turn-control-2026-07-01,effort-2025-11-24")], + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], + )] + fn request_headers_merge_the_feature_betas( + #[case] input: &[(&str, &str)], + #[case] fields: Value, + #[case] expected: &[(&str, &str)], ) { - assert_eq!(non_empty(value), expected); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.request_headers(headers(input), &request(fields)), + headers(expected) + ); + } + + #[test] + fn every_feature_merges_with_the_oauth_beta_sorted() { + let environment = validated(&[], Some(OAUTH_TOKEN), &[]).unwrap(); + let all_features = request(json!({ + "compaction": {"enabled": true}, + "output_format": {"type": "json_schema"}, + "speed": "fast", + "tools": [{"type": "advisor_20260301"}, {"type": "tool_search_tool_bm25_20251119"}], + "context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, + "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], + })); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.request_headers(environment.headers, &all_features), + headers(&[ + BROWSER_ACCESS, + ( + "anthropic-beta", + "advanced-tool-use-2025-11-20,advisor-tool-2026-03-01,compact-2026-01-12,compact-2026-09-04,context-management-2025-06-27,fast-mode-2026-02-01,oauth-2025-04-20,per-turn-control-2026-07-01,structured-outputs-2025-11-13" + ), + ]) + ); + assert_eq!( + credential(&environment.auth), + Some(("Authorization", OAUTH_TOKEN)) + ); } #[test] diff --git a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs index 8c768a6a66d..0c1434afb5d 100644 --- a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs @@ -1,4 +1,4 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_types::llms::anthropic_messages::{ anthropic_request::{ @@ -10,8 +10,9 @@ use litellm_types::llms::anthropic_messages::{ use crate::{ Error, - anthropic::messages::transformation::{ - ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, + anthropic::{ + common_utils::{API_KEY_PLACEMENT, MESSAGES_PATH_SUFFIX, non_empty}, + messages::transformation::{ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig}, }, base_llm::{ anthropic_messages::transformation::{ @@ -24,9 +25,7 @@ use crate::{ const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; -const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; const SYSTEM_ROLE: &str = "system"; -const API_KEY_HEADER: &str = "x-api-key"; pub struct AzureAnthropicMessagesConfig { anthropic: AnthropicMessagesConfig, @@ -86,14 +85,14 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { _model: &str, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - if has_header(&headers, API_KEY_HEADER) || has_bearer_auth(&headers) { + if has_header(&headers, API_KEY_PLACEMENT.header_name()) || has_bearer_auth(&headers) { return Ok(ValidatedEnvironment { headers, auth: AuthScheme::Forwarded, }); } let auth = AuthScheme::Credential { - placement: CredentialPlacement::Header(API_KEY_HEADER), + placement: API_KEY_PLACEMENT, secret: SecretValue::new(resolve_azure_api_key(api_key, env_lookup)?), }; Ok(ValidatedEnvironment { headers, auth }) @@ -216,6 +215,8 @@ mod tests { use rstest::rstest; use serde_json::json; + use litellm_auth::CredentialPlacement; + use super::*; use crate::anthropic::common_utils::AnthropicModelCapabilities; diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index 897b0f62e82..7d1b0014dcd 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -7,6 +7,7 @@ use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; +use litellm_http::request::without_headers; pub type Headers = Vec<(String, String)>; @@ -107,9 +108,8 @@ fn with_credential(headers: Headers, placement: CredentialPlacement, credential: CredentialPlacement::Bearer => format!("Bearer {credential}"), CredentialPlacement::Header(_) => credential.to_string(), }; - headers + without_headers(headers, &[name]) .into_iter() - .filter(|(header, _)| !header.eq_ignore_ascii_case(name)) .chain([(name.to_ascii_lowercase(), value)]) .collect() } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index bb0ec6671e3..cf634ab3fc3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -2,9 +2,8 @@ use std::convert::Infallible; use bytes::Bytes; use litellm_core::messages::{ - Error, - route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead, messages_body}, - types::MessagesShaping, + Error, MessagesCall, MessagesShaping, messages_body, + route::{Messages, MessagesOutput, MessagesStreamHead}, }; use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; use litellm_http::transport::Error as TransportError; @@ -76,6 +75,11 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult { error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; Ok(error) } + Error::MissingField(field) => { + let error = PyValueError::new_err(format!("missing required field: {field}")); + error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; + Ok(error) + } other => Ok(route_error_to_pyerr(other)), } } @@ -314,6 +318,7 @@ mod tests { #[rstest] #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] + #[case::missing_field(Error::MissingField("max_tokens"), true)] #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)] #[case::upstream_failure( Error::Transport(TransportError::Http { status: 400, body: "bad".into() }), diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index 4904bf4ffdd..9234728db28 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_core::messages::types::MessagesShaping; +use litellm_core::messages::MessagesShaping; #[derive(Clone, Debug, Default)] pub struct Deployment { diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 3cbd316cab0..2d16f102de3 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_config::Config; -use litellm_core::messages::types::MessagesShaping; +use litellm_core::messages::MessagesShaping; use litellm_router::{Deployment, Router}; use rstest::rstest; diff --git a/litellm-rust/crates/secrets/src/native.rs b/litellm-rust/crates/secrets/src/native.rs index f1dc7ccb732..8edcebe7b12 100644 --- a/litellm-rust/crates/secrets/src/native.rs +++ b/litellm-rust/crates/secrets/src/native.rs @@ -6,8 +6,8 @@ use litellm_http::{HttpClientConfig, HttpClientPool}; use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; pub async fn load_native_manager( - pool: &HttpClientPool, - config: &HttpClientConfig, + _pool: &HttpClientPool, + _config: &HttpClientConfig, system: KeyManagementSystem, settings: KeyManagementSettings, environment: Arc, @@ -33,14 +33,14 @@ pub async fn load_native_manager( #[cfg(feature = "azure")] (KeyManagementSystem::AzureKeyVault, _, environment, _) => Ok( SecretManager::AzureKeyVault(crate::azure::AzureKeyVault::new( - pool.client(config, litellm_http::ClientVariant::Provider)?, + _pool.client(_config, litellm_http::ClientVariant::Provider)?, environment, )?), ), #[cfg(feature = "google")] (KeyManagementSystem::GoogleSecretManager, _, environment, enterprise_enabled) => Ok( SecretManager::GoogleSecretManager(crate::google::GoogleSecretManager::new( - pool.client(config, litellm_http::ClientVariant::Provider)?, + _pool.client(_config, litellm_http::ClientVariant::Provider)?, environment, enterprise_enabled, )?), @@ -61,8 +61,8 @@ pub async fn load_native_manager( #[cfg(feature = "cyberark")] (KeyManagementSystem::Cyberark, _, environment, enterprise_enabled) => Ok( SecretManager::Cyberark(crate::cyberark::CyberArkSecretManager::new( - pool, - config, + _pool, + _config, environment, enterprise_enabled, )?), diff --git a/litellm-rust/crates/tracing/Cargo.toml b/litellm-rust/crates/tracing/Cargo.toml index 41ad20afb3e..914d988d301 100644 --- a/litellm-rust/crates/tracing/Cargo.toml +++ b/litellm-rust/crates/tracing/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true fancy-regex.workspace = true percent-encoding.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/tracing/src/lib.rs b/litellm-rust/crates/tracing/src/lib.rs index 47f97d6db27..4c6ec104f3a 100644 --- a/litellm-rust/crates/tracing/src/lib.rs +++ b/litellm-rust/crates/tracing/src/lib.rs @@ -5,6 +5,7 @@ use std::{ pin::pin, }; +use base64::{Engine, engine::general_purpose::STANDARD}; use serde_json::{Map, Value}; use tracing::{ Dispatch, Event, Subscriber, @@ -20,6 +21,31 @@ pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor}; pub use redaction::{REDACTED, SecretRedactor}; pub use tracing::{Level, Metadata, debug, error, info, trace, warn}; +pub struct ByteChunk<'a>(&'a [u8]); + +impl<'a> ByteChunk<'a> { + pub fn new(data: &'a [u8]) -> Self { + Self(data) + } + + pub fn encoding(&self) -> &'static str { + if std::str::from_utf8(self.0).is_ok() { + "utf8" + } else { + "base64" + } + } +} + +impl fmt::Display for ByteChunk<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match std::str::from_utf8(self.0) { + Ok(text) => formatter.write_str(text), + Err(_) => formatter.write_str(&STANDARD.encode(self.0)), + } + } +} + pub trait Sink: Send + Sync + 'static { fn enabled(&self, metadata: &Metadata<'_>) -> bool; fn emit(&self, record: &Record); @@ -44,6 +70,10 @@ impl Logger { } } + pub fn install_global(&self) -> Result<(), tracing::dispatcher::SetGlobalDefaultError> { + tracing::dispatcher::set_global_default(self.dispatch.clone()) + } + pub fn scope(&self, operation: impl FnOnce() -> T) -> T { if EMITTING.get() { return operation(); diff --git a/litellm-rust/crates/tracing/tests/logging.rs b/litellm-rust/crates/tracing/tests/logging.rs index 585e442dad1..3387f259f8b 100644 --- a/litellm-rust/crates/tracing/tests/logging.rs +++ b/litellm-rust/crates/tracing/tests/logging.rs @@ -4,7 +4,9 @@ use std::sync::{ mpsc, }; -use litellm_tracing::{Level, Logger, Metadata, Record, Sink, info, warn}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_tracing::{ByteChunk, Level, Logger, Metadata, Record, Sink, info, warn}; +use rstest::rstest; use serde_json::{Value, json}; struct Output { @@ -120,3 +122,18 @@ fn nested_scopes_restore_the_previous_sink() { ["inside"] ); } + +#[rstest] +#[case::utf8(b"event: message_stop\n\n", "utf8")] +#[case::binary(&[0xff, 0x00, 0x80], "base64")] +fn byte_chunk_logging_preserves_exact_bytes(#[case] bytes: &[u8], #[case] encoding: &str) { + let chunk = ByteChunk::new(bytes); + assert_eq!(chunk.encoding(), encoding); + let text = chunk.to_string(); + let recovered = match encoding { + "utf8" => text.into_bytes(), + "base64" => STANDARD.decode(text).unwrap(), + _ => unreachable!(), + }; + assert_eq!(recovered, bytes); +} diff --git a/litellm-rust/crates/types/src/llms/anthropic.rs b/litellm-rust/crates/types/src/llms/anthropic.rs new file mode 100644 index 00000000000..3e0c4369b11 --- /dev/null +++ b/litellm-rust/crates/types/src/llms/anthropic.rs @@ -0,0 +1,240 @@ +use std::{ + cmp::Ordering, + collections::BTreeSet, + convert::Infallible, + fmt, + hash::{Hash, Hasher}, + str::FromStr, +}; + +/// One value of the `anthropic-beta` header. Equality, ordering and hashing follow the wire +/// string, so a value parsed from a caller's header never disagrees with the matching variant. +#[derive(Clone, Debug, strum::AsRefStr, strum::Display, strum::EnumString)] +pub enum AnthropicBeta { + #[strum(serialize = "oauth-2025-04-20")] + Oauth20250420, + #[strum(serialize = "web-fetch-2025-09-10")] + WebFetch20250910, + #[strum(serialize = "web-search-2025-03-05")] + WebSearch20250305, + #[strum(serialize = "context-management-2025-06-27")] + ContextManagement20250627, + #[strum(serialize = "compact-2026-01-12")] + Compact20260112, + #[strum(serialize = "compact-2026-09-04")] + Compact20260904, + #[strum(serialize = "structured-outputs-2025-11-13")] + StructuredOutputs20251113, + #[strum(serialize = "advanced-tool-use-2025-11-20")] + AdvancedToolUse20251120, + #[strum(serialize = "fast-mode-2026-02-01")] + FastMode20260201, + #[strum(serialize = "advisor-tool-2026-03-01")] + AdvisorTool20260301, + #[strum(serialize = "per-turn-control-2026-07-01")] + PerTurnControl20260701, + #[strum(serialize = "dangerous-tool-use-2026-09-03")] + DangerousToolUse20260903, + #[strum(default, transparent)] + Other(String), +} + +impl AnthropicBeta { + pub const KNOWN: [Self; 12] = [ + Self::Oauth20250420, + Self::WebFetch20250910, + Self::WebSearch20250305, + Self::ContextManagement20250627, + Self::Compact20260112, + Self::Compact20260904, + Self::StructuredOutputs20251113, + Self::AdvancedToolUse20251120, + Self::FastMode20260201, + Self::AdvisorTool20260301, + Self::PerTurnControl20260701, + Self::DangerousToolUse20260903, + ]; + + pub fn as_str(&self) -> &str { + self.as_ref() + } +} + +impl PartialEq for AnthropicBeta { + fn eq(&self, other: &Self) -> bool { + self.as_str() == other.as_str() + } +} + +impl Eq for AnthropicBeta {} + +impl Hash for AnthropicBeta { + fn hash(&self, state: &mut H) { + self.as_str().hash(state); + } +} + +impl PartialOrd for AnthropicBeta { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for AnthropicBeta { + fn cmp(&self, other: &Self) -> Ordering { + self.as_str().cmp(other.as_str()) + } +} + +/// The values of one `anthropic-beta` header: sorted, deduplicated, comma-joined on the wire. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct BetaSet(BTreeSet); + +impl BetaSet { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn contains(&self, beta: &AnthropicBeta) -> bool { + self.0.contains(beta) + } + + pub fn iter(&self) -> impl Iterator { + self.0.iter() + } + + pub fn union(self, other: Self) -> Self { + self.0.into_iter().chain(other.0).collect() + } +} + +impl FromIterator for BetaSet { + fn from_iter>(betas: I) -> Self { + Self(betas.into_iter().collect()) + } +} + +impl IntoIterator for BetaSet { + type Item = AnthropicBeta; + type IntoIter = std::collections::btree_set::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl FromStr for BetaSet { + type Err = Infallible; + + fn from_str(header: &str) -> Result { + Ok(header + .split(',') + .map(str::trim) + .filter(|piece| !piece.is_empty()) + .map(|piece| AnthropicBeta::from_str(piece).unwrap_or_else(|never| match never {})) + .collect()) + } +} + +impl fmt::Display for BetaSet { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut betas = self.0.iter(); + let Some(first) = betas.next() else { + return Ok(()); + }; + f.write_str(first.as_str())?; + betas.try_for_each(|beta| write!(f, ",{beta}")) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn set(header: &str) -> BetaSet { + header.parse().unwrap_or_else(|never| match never {}) + } + + #[rstest] + fn every_known_beta_parses_back_to_itself( + #[values( + AnthropicBeta::Oauth20250420, + AnthropicBeta::WebFetch20250910, + AnthropicBeta::WebSearch20250305, + AnthropicBeta::ContextManagement20250627, + AnthropicBeta::Compact20260112, + AnthropicBeta::Compact20260904, + AnthropicBeta::StructuredOutputs20251113, + AnthropicBeta::AdvancedToolUse20251120, + AnthropicBeta::FastMode20260201, + AnthropicBeta::AdvisorTool20260301, + AnthropicBeta::PerTurnControl20260701, + AnthropicBeta::DangerousToolUse20260903 + )] + beta: AnthropicBeta, + ) { + let parsed: AnthropicBeta = beta.as_str().parse().unwrap(); + assert!(!matches!(parsed, AnthropicBeta::Other(_))); + assert_eq!(parsed, beta); + assert!(AnthropicBeta::KNOWN.contains(&beta)); + } + + #[test] + fn unknown_values_are_kept_verbatim() { + let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); + assert_eq!( + parsed, + AnthropicBeta::Other("claude-code-20250219".to_string()) + ); + assert_eq!(parsed.to_string(), "claude-code-20250219"); + } + + #[test] + fn a_known_value_spelled_as_other_is_the_same_beta() { + let spelled_out = AnthropicBeta::Other("compact-2026-01-12".to_string()); + assert_eq!(spelled_out, AnthropicBeta::Compact20260112); + assert_eq!( + spelled_out.cmp(&AnthropicBeta::Compact20260112), + Ordering::Equal + ); + assert_eq!( + BetaSet::from_iter([spelled_out, AnthropicBeta::Compact20260112]).to_string(), + "compact-2026-01-12" + ); + } + + #[rstest] + #[case::empty("", "")] + #[case::blank_pieces(" , ,", "")] + #[case::single("b", "b")] + #[case::sorted("c,a", "a,c")] + #[case::trimmed_and_deduplicated("b, a ,b", "a,b")] + #[case::blank_pieces_skipped("a,,b", "a,b")] + #[case::known_and_unknown_sort_together( + "web-search-2025-03-05,claude-code-20250219,fast-mode-2026-02-01", + "claude-code-20250219,fast-mode-2026-02-01,web-search-2025-03-05" + )] + fn header_values_round_trip_sorted_and_deduplicated(#[case] header: &str, #[case] wire: &str) { + assert_eq!(set(header).to_string(), wire); + assert_eq!(set(header).is_empty(), wire.is_empty()); + } + + #[rstest] + #[case::disjoint("a,c", "b", "a,b,c")] + #[case::overlapping("a,b", "b,c", "a,b,c")] + #[case::empty_right("a", "", "a")] + #[case::empty_left("", "a", "a")] + fn union_merges_both_sides(#[case] left: &str, #[case] right: &str, #[case] wire: &str) { + assert_eq!(set(left).union(set(right)).to_string(), wire); + } + + #[test] + fn contains_matches_by_wire_value() { + let betas = set("oauth-2025-04-20,claude-code-20250219"); + assert!(betas.contains(&AnthropicBeta::Oauth20250420)); + assert!(betas.contains(&AnthropicBeta::Other("claude-code-20250219".into()))); + assert!(!betas.contains(&AnthropicBeta::FastMode20260201)); + } +} diff --git a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs index 342e891a1e3..335a03c4b5b 100644 --- a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs +++ b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs @@ -111,6 +111,70 @@ impl From for ReasoningEffort { } } +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum Speed { + Fast, + Standard, +} + +impl Speed { + pub fn as_str(self) -> &'static str { + self.into() + } +} + +/// The tools whose presence changes how the request is sent. Every other tool, custom or +/// server, deserializes as `Recognized::Unrecognized` and passes through verbatim. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum AnthropicTool { + #[serde(rename = "advisor_20260301")] + Advisor { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "tool_search_tool_regex_20251119")] + ToolSearchRegex { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "tool_search_tool_bm25_20251119")] + ToolSearchBm25 { + #[serde(flatten)] + extra: Map, + }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum ContextEdit { + #[serde(rename = "compact_20260112")] + Compact { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "clear_tool_uses_20250919")] + ClearToolUses { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "clear_thinking_20251015")] + ClearThinking { + #[serde(flatten)] + extra: Map, + }, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ContextManagement { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub edits: Option>>, + #[serde(flatten)] + pub extra: Map, +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct OutputConfig { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -210,7 +274,7 @@ pub struct AnthropicMessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub top_k: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, + pub tools: Option>>, #[serde(skip_serializing_if = "Option::is_none")] pub tool_choice: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -222,13 +286,13 @@ pub struct AnthropicMessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub mcp_servers: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub context_management: Option, + pub context_management: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub output_format: Option, #[serde(skip_serializing_if = "Option::is_none")] pub output_config: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub speed: Option, + pub speed: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub inference_geo: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -397,6 +461,33 @@ mod tests { "thinking": {"type": "future", "budget_tokens": 1}, "output_config": "bogus" }))] + #[case::tools_speed_and_context_management(json!({ + "model": "m", + "messages": [], + "speed": "fast", + "tools": [ + {"name": "get_weather", "input_schema": {"type": "object"}}, + {"type": "custom", "name": "f", "input_schema": {}}, + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + {"type": "tool_search_tool_bm25_20251119"} + ], + "context_management": {"edits": [ + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}, + {"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}, + {"type": "clear_thinking_20251015"}, + {"type": "future_edit"}, + {} + ], "future": true} + }))] + #[case::unrecognized_tools_speed_and_context_management_are_kept_verbatim(json!({ + "model": "m", + "messages": [], + "speed": "turbo", + "tools": ["none", 5], + "context_management": [{"type": "compaction", "compact_threshold": 5}] + }))] fn request_round_trips_unchanged(#[case] request: Value) { assert_eq!(round_trip::(&request), request); } @@ -444,6 +535,114 @@ mod tests { ); } + #[rstest] + #[case::advisor( + json!({"type": "advisor_20260301", "name": "advisor"}), + Recognized::Known(AnthropicTool::Advisor { extra: Map::from_iter([("name".to_string(), json!("advisor"))]) }) + )] + #[case::regex_tool_search( + json!({"type": "tool_search_tool_regex_20251119"}), + Recognized::Known(AnthropicTool::ToolSearchRegex { extra: Map::new() }) + )] + #[case::bm25_tool_search( + json!({"type": "tool_search_tool_bm25_20251119"}), + Recognized::Known(AnthropicTool::ToolSearchBm25 { extra: Map::new() }) + )] + #[case::custom_tool_without_a_type( + json!({"name": "advisor", "input_schema": {}}), + Recognized::Unrecognized(json!({"name": "advisor", "input_schema": {}})) + )] + #[case::other_server_tool( + json!({"type": "web_search_20250305", "name": "web_search"}), + Recognized::Unrecognized(json!({"type": "web_search_20250305", "name": "web_search"})) + )] + #[case::not_an_object(json!("advisor_20260301"), Recognized::Unrecognized(json!("advisor_20260301")))] + fn tools_are_recognized_by_their_exact_type( + #[case] tool: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(tool).unwrap(), + expected + ); + } + + #[rstest] + #[case::compact( + json!({"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1}}), + Recognized::Known(ContextEdit::Compact { + extra: Map::from_iter([("trigger".to_string(), json!({"type": "input_tokens", "value": 1}))]), + }) + )] + #[case::clear_tool_uses( + json!({"type": "clear_tool_uses_20250919"}), + Recognized::Known(ContextEdit::ClearToolUses { extra: Map::new() }) + )] + #[case::clear_thinking( + json!({"type": "clear_thinking_20251015"}), + Recognized::Known(ContextEdit::ClearThinking { extra: Map::new() }) + )] + #[case::unknown_type(json!({"type": "future"}), Recognized::Unrecognized(json!({"type": "future"})))] + #[case::no_type(json!({}), Recognized::Unrecognized(json!({})))] + fn context_edits_are_recognized_by_their_exact_type( + #[case] edit: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(edit).unwrap(), + expected + ); + } + + #[rstest] + #[case::edits( + json!({"edits": [{"type": "compact_20260112"}]}), + Recognized::Known(ContextManagement { + edits: Some(vec![Recognized::Known(ContextEdit::Compact { extra: Map::new() })]), + extra: Map::new(), + }) + )] + #[case::object_without_edits( + json!({"future": 1}), + Recognized::Known(ContextManagement { + edits: None, + extra: Map::from_iter([("future".to_string(), json!(1))]), + }) + )] + #[case::openai_list(json!([{"type": "compaction"}]), Recognized::Unrecognized(json!([{"type": "compaction"}])))] + #[case::edits_not_a_list(json!({"edits": 5}), Recognized::Unrecognized(json!({"edits": 5})))] + #[case::scalar(json!("compaction"), Recognized::Unrecognized(json!("compaction")))] + fn context_management_is_known_only_as_an_edits_object( + #[case] value: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(value).unwrap(), + expected + ); + } + + #[rstest] + #[case::fast(json!("fast"), Recognized::Known(Speed::Fast))] + #[case::standard(json!("standard"), Recognized::Known(Speed::Standard))] + #[case::unknown(json!("turbo"), Recognized::Unrecognized(json!("turbo")))] + #[case::wrong_case(json!("Fast"), Recognized::Unrecognized(json!("Fast")))] + #[case::not_a_string(json!(1), Recognized::Unrecognized(json!(1)))] + fn speed_is_known_only_as_a_documented_value( + #[case] value: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(value).unwrap(), + expected + ); + } + + #[rstest] + fn speed_names_match_the_wire(#[values(Speed::Fast, Speed::Standard)] speed: Speed) { + assert_eq!(serde_json::to_value(speed).unwrap(), json!(speed.as_str())); + } + #[rstest] fn effort_level_names_match_the_wire( #[values( diff --git a/litellm-rust/crates/types/src/llms/mod.rs b/litellm-rust/crates/types/src/llms/mod.rs index 09d2207a0ca..19ce0bb77ef 100644 --- a/litellm-rust/crates/types/src/llms/mod.rs +++ b/litellm-rust/crates/types/src/llms/mod.rs @@ -1,2 +1,3 @@ +pub mod anthropic; pub mod anthropic_messages; pub mod openai; diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 6e455817194..46acb94c958 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -109,7 +109,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index eeea364b674..2b3edac612e 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -50,12 +50,13 @@ def test_shipped_decisions( monkeypatch.setenv("LITELLM_RUST", environment) context: Final = RouteContext(route, provider=provider, model="test-model", delivery=delivery) - if route is Route.OCR: - assert catalog.rollout(context) is Rollout.RUST_REQUIRED - assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.TRANSCRIPTION and provider == "bedrock": + if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED + elif route is Route.MESSAGES and provider == "anthropic": + assert catalog.rollout(context) is Rollout.RUST_OPT_IN + opted_in: Final = environment == "1" or (environment is None and process is True) + assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) else: assert catalog.rollout(context) is Rollout.PYTHON_ONLY assert catalog.decision(context) is Decision.PYTHON @@ -145,10 +146,7 @@ def test_ocr_has_no_python_path_to_opt_out_to( monkeypatch.setenv("LITELLM_RUST", environment) assert catalog.decision(RouteContext(Route.OCR, model="m")) is Decision.RUST_REQUIRED - assert ( - catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) - is Decision.RUST_REQUIRED - ) + assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED @pytest.mark.parametrize( From 1e6c98334c6e57a2d7e481a7f5d5de1950a791db Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 01:32:55 -0700 Subject: [PATCH 57/58] refactor: daily fresh tech debt cleanup, rolling PR (2026-09-25) (#43151) * refactor: clean up fresh tech debt from 2026-09-24 (stacked comprehensions, getattr, bare dict) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: drop the budget ratchet from the PR branch, the default-branch automation owns it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 2 +- litellm/integrations/langfuse/langfuse_sdk.py | 2 +- litellm/llms/openai/organization_costs.py | 4 +++- .../mcp_server/mcp_server_manager.py | 9 +++++---- .../proxy/common_utils/model_listing_utils.py | 3 ++- litellm/proxy/hooks/proxy_track_cost_callback.py | 4 +++- .../spend_tracking/key_metadata_recovery.py | 16 +++++++++++----- litellm/types/litellm_params.py | 3 ++- 8 files changed, 28 insertions(+), 15 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 246ac4fd369..9974e77d017 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -43,7 +43,7 @@ def _uses_native_vertex_output( ) -> bool: if custom_llm_provider != "vertex_ai": return False - if model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False): + if model_name and litellm.disable_vertex_batch_output_transformation: return True return first_row is not None and is_native_vertex_batch_output_row(first_row) diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 66819c95ebf..986f35297d2 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -684,7 +684,7 @@ class LangfuseSpanExporter(SpanExporter): def _round(self, halving: _Halving) -> _Halving: sent: Final = tuple((batch, self._send_batch(batch)) for batch in halving.pending) return _Halving( - pending=tuple(part for batch, outcome in sent if outcome == "too_large" for part in _smaller(batch)), + pending=tuple(chain.from_iterable(_smaller(batch) for batch, outcome in sent if outcome == "too_large")), settled=halving.settled + tuple( SpanExportResult.SUCCESS if outcome == "delivered" else SpanExportResult.FAILURE diff --git a/litellm/llms/openai/organization_costs.py b/litellm/llms/openai/organization_costs.py index e7fb22f9b19..856072ddb99 100644 --- a/litellm/llms/openai/organization_costs.py +++ b/litellm/llms/openai/organization_costs.py @@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone +from itertools import chain from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -126,7 +127,8 @@ async def fetch_openai_daily_costs( return MappingProxyType( { day: sum( - result.amount.value for bucket in buckets if _bucket_day(bucket) == day for result in bucket.results + result.amount.value + for result in chain.from_iterable(bucket.results for bucket in buckets if _bucket_day(bucket) == day) ) for day in days } diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6baa695433c..8befc99cad4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6855,12 +6855,13 @@ class MCPServerManager: if not tool_permissions: return {} expanded: Final = tuple( - (server_id, tuple(tools or ())) - for key, tools in tool_permissions.items() - for server_id in self.expand_permission_list([key]) + chain.from_iterable( + ((server_id, tuple(tools or ())) for server_id in self.expand_permission_list([key])) + for key, tools in tool_permissions.items() + ) ) return { - server_id: list(dict.fromkeys(tool for _, tools in group for tool in tools)) + server_id: list(dict.fromkeys(chain.from_iterable(tools for _, tools in group))) for server_id, group in groupby(sorted(expanded, key=itemgetter(0)), key=itemgetter(0)) } diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 8958fb20918..0c702ac4139 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -14,6 +14,7 @@ import re from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass from functools import reduce +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast @@ -191,7 +192,7 @@ def alias_map(aliases: object) -> Mapping[str, str]: def _alias_names(alias_maps: Sequence[Mapping[str, str]]) -> tuple[str, ...]: - return tuple(dict.fromkeys(alias for aliases in alias_maps for alias in aliases)) + return tuple(dict.fromkeys(chain.from_iterable(alias_maps))) def _rewrite(model_id: str, alias_maps: Sequence[Mapping[str, str]]) -> str | None: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e097debde77..0178465739b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -474,7 +474,9 @@ class _ProxyDBLogger(CustomLogger): spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) @staticmethod - async def _enrich_failure_metadata_unless_db_stalled(metadata: dict, original_exception: Exception) -> dict: + async def _enrich_failure_metadata_unless_db_stalled( + metadata: dict[str, object], original_exception: Exception + ) -> dict[str, object]: if isinstance(original_exception, DBLookupDeadlineExceeded): return metadata return await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 08e8fff8f1d..965cded59c4 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -215,17 +215,23 @@ def _meta_with_user_details( return updated +def _user_id_needing_details(api_key: str, meta: KeyMetadataDict) -> str | None: + user_id: Final = meta.get("user_id") + if not isinstance(user_id, str) or not user_id: + return None + if meta.get("user_email") and not (_is_cli_session_key(api_key) and not meta.get("team_id")): + return None + return user_id + + async def attach_user_details( prisma_client: PrismaClient, recovered: Mapping[str, KeyMetadataDict], ) -> Mapping[str, KeyMetadataDict]: needing_details: Final = frozenset( user_id - for api_key, meta in recovered.items() - for user_id in (meta.get("user_id"),) - if isinstance(user_id, str) - and user_id - and (not meta.get("user_email") or (_is_cli_session_key(api_key) and not meta.get("team_id"))) + for user_id in (_user_id_needing_details(api_key, meta) for api_key, meta in recovered.items()) + if user_id is not None ) details: Final = await _details_for_user_ids(prisma_client, needing_details) if not details: diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 83a42c235f9..f5ba9ebd3da 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -3,6 +3,7 @@ models and KWARG_ARTIFACTS into all_litellm_params.""" from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence from dataclasses import dataclass, field, fields, is_dataclass +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias @@ -359,6 +360,6 @@ def owned_wire_names(root: type) -> tuple[str, ...]: return tuple(names()) -OWNED_KWARG_NAMES: Final = tuple(name for root in LITELLM_OWNED_ROOTS for name in owned_wire_names(root)) +OWNED_KWARG_NAMES: Final = tuple(chain.from_iterable(owned_wire_names(root) for root in LITELLM_OWNED_ROOTS)) AGENTIC_LOOP_KWARG_NAMES: Final = (*wire_names(AgenticLoopState), *wire_names(AgenticLoopOptions)) BEDROCK_BATCH_KWARG_NAMES: Final = wire_names(BedrockBatchConnection) From 115668f43ea5f6ed5ce47ad9de0db2df0bf96ff1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 03:12:13 -0700 Subject: [PATCH 58/58] test(proxy_behavior): scope the management proxy fixture to its package so its spend monitor cannot race the spend tests (#43302) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/proxy_behavior/management/conftest.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py index 255b937bdd3..74db323e7f2 100644 --- a/tests/proxy_behavior/management/conftest.py +++ b/tests/proxy_behavior/management/conftest.py @@ -31,7 +31,7 @@ def _write_minimal_proxy_config() -> str: return f.name -@pytest_asyncio.fixture(scope="session") +@pytest_asyncio.fixture(scope="package") async def proxy_app(): from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( @@ -67,7 +67,7 @@ async def proxy_app(): yield app -@pytest_asyncio.fixture(scope="session") +@pytest_asyncio.fixture(scope="package") async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]: transport = httpx.ASGITransport(app=proxy_app) async with httpx.AsyncClient( @@ -76,7 +76,7 @@ async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]: yield client -@pytest_asyncio.fixture(scope="session") +@pytest_asyncio.fixture(scope="package") async def prisma(proxy_app): from litellm.proxy import proxy_server @@ -84,7 +84,7 @@ async def prisma(proxy_app): return proxy_server.prisma_client -@pytest_asyncio.fixture(scope="session") +@pytest_asyncio.fixture(scope="package") async def world(prisma): from .actors import seed_world