diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index e4861c3954c..05c74a05924 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -4,9 +4,10 @@ mod service; use axum::Router; use axum::body::Body; -use axum::extract::{Json, State}; +use axum::extract::{Json, Request, State}; use axum::http::StatusCode; use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue}; +use axum::middleware::{self, Next}; use axum::response::{IntoResponse, Response}; use axum::routing::post; use litellm_core::Error; @@ -16,9 +17,23 @@ use crate::auth::RequireMasterKey; use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH}; use crate::state::AppState; +const CORE_ENGINE_HEADER: &str = "x-litellm-core"; +const RUST_CORE_ENGINE: &str = "rust"; + /// This route's contribution to the app router. pub fn router() -> Router { - Router::new().route(MESSAGES_ROUTE_PATH, post(handle)) + Router::new() + .route(MESSAGES_ROUTE_PATH, post(handle)) + .route_layer(middleware::from_fn(core_engine_header)) +} + +async fn core_engine_header(request: Request, next: Next) -> Response { + let mut response = next.run(request).await; + response.headers_mut().insert( + CORE_ENGINE_HEADER, + HeaderValue::from_static(RUST_CORE_ENGINE), + ); + response } async fn handle( @@ -142,6 +157,7 @@ mod tests { use tower::ServiceExt; use super::super::app; + use super::{CORE_ENGINE_HEADER, RUST_CORE_ENGINE}; use crate::io::realtime_pool::RealtimePool; use crate::state::AppState; @@ -289,6 +305,7 @@ mod tests { .await .expect("route responds"); assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()[CORE_ENGINE_HEADER], RUST_CORE_ENGINE); let body = axum::body::to_bytes(response.into_body(), usize::MAX) .await .expect("response body reads"); 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 66e36dab2ba..4424ef4a081 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -6,7 +6,7 @@ from typing import Any, Final, Protocol, runtime_checkable import httpx from pydantic import TypeAdapter -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.constants import ( ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, @@ -19,6 +19,7 @@ from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) +from litellm.rust_bridge.runtime import CoreEngine, execution_additional_headers from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -325,7 +326,8 @@ def _anthropic_content_block_events(index: int, block: Mapping[str, object]) -> class AnthropicMessagesStreamHiddenParams(TypedDict): - additional_headers: dict[str, str] + additional_headers: ReadOnly[Mapping[str, str]] + core_engine: ReadOnly[str] @runtime_checkable @@ -344,8 +346,10 @@ _RESPONSE_HEADERS_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict def anthropic_messages_stream_hidden_params( response_headers: httpx.Headers, ) -> AnthropicMessagesStreamHiddenParams: + additional_headers: Final = _RESPONSE_HEADERS_ADAPTER.validate_python(process_response_headers(response_headers)) return AnthropicMessagesStreamHiddenParams( - additional_headers=_RESPONSE_HEADERS_ADAPTER.validate_python(process_response_headers(response_headers)) + additional_headers=execution_additional_headers(additional_headers, CoreEngine.PYTHON), + core_engine=CoreEngine.PYTHON.value, ) diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index b1f8c957ff4..e1504a18f71 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -5,6 +5,7 @@ import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.rust_bridge import transcription as rust_transcription_bridge +from litellm.rust_bridge.runtime import CoreEngine, execution_hidden_params from litellm.types.utils import FileTypes, TranscriptionResponse @@ -55,7 +56,9 @@ class BedrockAudioTranscriptionRustDispatch: ) if rust_response is None: raise RuntimeError("Rust audio transcription bridge is unavailable") - return TranscriptionResponse(**rust_response) + response: Final = TranscriptionResponse(**rust_response) + response["_hidden_params"] = execution_hidden_params(None, CoreEngine.RUST) + return response async def async_audio_transcriptions( self, @@ -81,4 +84,6 @@ class BedrockAudioTranscriptionRustDispatch: ) if rust_response is None: raise RuntimeError("Rust audio transcription bridge is unavailable") - return TranscriptionResponse(**rust_response) + response: Final = TranscriptionResponse(**rust_response) + response["_hidden_params"] = execution_hidden_params(None, CoreEngine.RUST) + return response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3b08550a0ba..64d9241062f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -19,7 +19,6 @@ import litellm.types.utils from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES -from litellm.exceptions import APIError from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -89,6 +88,7 @@ from litellm.responses.streaming_iterator import ( ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, ) +from litellm.rust_bridge.runtime import CoreEngine, execution_additional_headers, execution_hidden_params from litellm.types.containers.main import ( ContainerFileListResponse, ContainerListResponse, @@ -166,6 +166,18 @@ def _rust_responses_websocket_enabled( return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) +def _anthropic_messages_with_core_engine( + response: AnthropicMessagesResponse, + source: CoreEngine, +) -> AnthropicMessagesResponse: + existing_hidden_params: Final = response.get("_hidden_params") + response["_hidden_params"] = execution_hidden_params( + existing_hidden_params if isinstance(existing_hidden_params, dict) else None, + source, + ) + return response + + from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: @@ -2361,8 +2373,21 @@ class BaseLLMHTTPHandler: kwargs=kwargs_for_agentic, ) + selected_response: Final = final_response if final_response is not None else initial_response + if not isinstance(selected_response, dict): + return self._maybe_wrap_in_fake_stream( + selected_response, + logging_obj, + "anthropic_messages", + ) + typed_response: Final = cast(AnthropicMessagesResponse, selected_response) + existing_hidden_params: Final = typed_response.get("_hidden_params") + existing_source: Final = ( + existing_hidden_params.get("core_engine") if isinstance(existing_hidden_params, dict) else None + ) + source: Final = CoreEngine.RUST if existing_source == CoreEngine.RUST.value else CoreEngine.PYTHON return self._maybe_wrap_in_fake_stream( - final_response if final_response is not None else initial_response, + _anthropic_messages_with_core_engine(typed_response, source), logging_obj, "anthropic_messages", ) @@ -2407,7 +2432,7 @@ class BaseLLMHTTPHandler: return None response_obj: Final = cast(AnthropicMessagesResponse, dict(rust_response)) - return response_obj + return _anthropic_messages_with_core_engine(response_obj, CoreEngine.RUST) @staticmethod def _rust_anthropic_messages_fake_stream( @@ -2422,7 +2447,10 @@ class BaseLLMHTTPHandler: ) completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) - hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={}) + hidden_params: Final = AnthropicMessagesStreamHiddenParams( + additional_headers=execution_additional_headers(None, CoreEngine.RUST), + core_engine=CoreEngine.RUST.value, + ) return AnthropicMessagesStreamingResponse( completion_stream=completion_stream, hidden_params=hidden_params, @@ -6483,11 +6511,12 @@ class BaseLLMHTTPHandler: if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params): from litellm.rust_bridge import responses_websocket as rust_responses_websocket - rust_backend: Final = await rust_responses_websocket.connect( + rust_execution: Final = await rust_responses_websocket.connect( url=ws_url, headers={str(key): str(value) for key, value in headers.items()}, timeout=timeout, ) + rust_backend: Final = rust_execution.value if rust_backend is not None: yield rust_backend return diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b260ec6e06f..e361f68aad6 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.runtime import CoreEngine, execution_hidden_params from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -67,6 +68,24 @@ _RUST_OCR_PROVIDERS: Final = { } +def _with_core_engine(response: OCRResponse, source: CoreEngine) -> OCRResponse: + raw_hidden_params: Final = cast(object, response._hidden_params) + hidden_params: Final = ( + cast(Mapping[str, object], raw_hidden_params) if isinstance(raw_hidden_params, Mapping) else None + ) + response._hidden_params = execution_hidden_params( # pyright: ignore[reportPrivateUsage] # OCRResponse has no public metadata setter + hidden_params, source + ) + return response + + +async def _await_with_core_engine( + response: Coroutine[object, object, OCRResponse], + source: CoreEngine, +) -> OCRResponse: + return _with_core_engine(await response, source) + + def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -306,7 +325,7 @@ def _run_rust_ocr( ) if rust_response is None: return None - return OCRResponse.model_validate(rust_response) + return _with_core_engine(OCRResponse.model_validate(rust_response), CoreEngine.RUST) async def _run_rust_aocr( @@ -331,7 +350,7 @@ async def _run_rust_aocr( ) if rust_response is None: return None - return OCRResponse.model_validate(rust_response) + return _with_core_engine(OCRResponse.model_validate(rust_response), CoreEngine.RUST) @client @@ -461,7 +480,7 @@ async def aocr( if response is None: raise ValueError(f"Got an unexpected None response from the OCR API: {response}") - return response + return _with_core_engine(response, CoreEngine.PYTHON) except Exception as e: raise litellm.exception_type( model=model, @@ -727,7 +746,9 @@ def ocr( litellm_params=prepared.litellm_params, ) - return response + if asyncio.iscoroutine(response): + return _await_with_core_engine(response, CoreEngine.PYTHON) + return _with_core_engine(response, CoreEngine.PYTHON) except Exception as e: raise litellm.exception_type( model=model, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 05ddef822f1..c631cc7d74a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1553,6 +1553,16 @@ class ProxyBaseLLMRequestProcessing: logging_obj=litellm_logging_obj, use_logging_obj=read_timing_from_logging_obj, ) + core_engine_value: Final = hidden_params.get("core_engine") + core_engine: Final = ( + core_engine_value + if isinstance(core_engine_value, str) and core_engine_value in ("python", "rust") + else "python" + ) + reserved_core_headers: Final = frozenset({"x-litellm-core", "x-litellm-rust"}) + forwarded_headers: Final = { + key: value for key, value in kwargs.items() if key.lower() not in reserved_core_headers + } cost_breakdown: Final = _get_cost_breakdown_from_logging_obj( litellm_logging_obj=litellm_logging_obj, response_cost=response_cost @@ -1637,7 +1647,9 @@ class ProxyBaseLLMRequestProcessing: str(fastest_response_batch_completion) if fastest_response_batch_completion is not None else None ), "x-litellm-timeout": str(timeout) if timeout is not None else None, - **{k: str(v) for k, v in kwargs.items()}, + **{key: str(value) for key, value in forwarded_headers.items()}, + "x-litellm-core": core_engine, + **({"x-litellm-rust": "true"} if core_engine == "rust" else {}), } if request_data: remaining_tokens_header: Final = get_remaining_tokens_and_requests_from_request_data(request_data) diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index efdb6eaa9c5..5fb89ae5848 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -14,7 +14,7 @@ from __future__ import annotations import json from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol, cast import httpx from pydantic import TypeAdapter, ValidationError @@ -28,12 +28,15 @@ from litellm.rust_bridge.bindings import UNSET, NativeBinding, Unset from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import ( BridgeErrorContext, + CoreEngine, + ExecutionResult, FallbackMode, RustDeclined, RustHandled, aattempt, ainvoke, attempt, + execution_hidden_params, identity, invoke, ) @@ -51,6 +54,7 @@ RUST_CHAT_COMPLETIONS_PROVIDERS: Final = frozenset({"anthropic", "bedrock"}) # rather than narrowing an unparameterized `Mapping` and typing the result Any. _LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + class RustChatCompletions(Protocol): def __call__( self, @@ -265,12 +269,31 @@ def _build_model_response( built: Final = convert_to_model_response_object( response_object=dict(rust_response), # mutable-ok: the converter takes a real dict and rewrites it model_response_object=model_response, + hidden_params=execution_hidden_params(None, CoreEngine.RUST), # mutable-ok: rewritten by the converter ) if not isinstance(built, ModelResponse): raise TypeError(f"expected a ModelResponse from the rust path, got {type(built).__name__}") return built +def _unwrap_execution(result: ExecutionResult[object]) -> object: + value: Final = result.value + if isinstance(value, ModelResponse): + raw_hidden_params: Final = cast( + object, + value._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse has no public metadata getter + ) + hidden_params: Final = ( + cast(Mapping[str, object], raw_hidden_params) # cast-ok: guarded by the mapping check + if isinstance(raw_hidden_params, Mapping) + else None + ) + value._hidden_params = execution_hidden_params( # pyright: ignore[reportPrivateUsage] # ModelResponse has no public metadata setter + hidden_params, result.source + ) + return value + + def chat_completions( *, model: str, @@ -387,12 +410,14 @@ def chat_completions_or_fallback( on_response(rust_response) return _build_model_response(rust_response, model_response) - return invoke( - native_call=native_call, - fallback=python_fallback, - adapt=adapt, - mode=FallbackMode.PYTHON, - context=BridgeErrorContext(route="chat completions", provider=custom_llm_provider or "", model=model), + return _unwrap_execution( + invoke( + native_call=native_call, + fallback=python_fallback, + adapt=adapt, + mode=FallbackMode.PYTHON, + context=BridgeErrorContext(route="chat completions", provider=custom_llm_provider or "", model=model), + ) ) @@ -430,10 +455,12 @@ async def achat_completions_or_fallback( on_response(rust_response) return _build_model_response(rust_response, model_response) - return await ainvoke( - native_call=native_call, - fallback=python_fallback, - adapt=adapt, - mode=FallbackMode.PYTHON, - context=BridgeErrorContext(route="chat completions", provider=custom_llm_provider or "", model=model), + return _unwrap_execution( + await ainvoke( + native_call=native_call, + fallback=python_fallback, + adapt=adapt, + mode=FallbackMode.PYTHON, + context=BridgeErrorContext(route="chat completions", provider=custom_llm_provider or "", model=model), + ) ) diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index e0bdef20dfe..82ff6cabcb7 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -105,7 +105,7 @@ def messages( adapt=identity, mode=FallbackMode.PYTHON, context=_context(model, custom_llm_provider), - ) + ).value async def amessages( @@ -119,22 +119,24 @@ async def amessages( timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: native: Final = load_rust_amessages() - return await ainvoke( - native_call=( - None - if native is None - else lambda: native( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), - ) - ), - fallback=async_none, - adapt=identity, - mode=FallbackMode.PYTHON, - context=_context(model, custom_llm_provider), - ) + return ( + await ainvoke( + native_call=( + None + if native is None + else lambda: native( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + ), + fallback=async_none, + adapt=identity, + mode=FallbackMode.PYTHON, + context=_context(model, custom_llm_provider), + ) + ).value diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 355254fb237..401ca3448f9 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -106,7 +106,7 @@ def ocr( adapt=identity, mode=FallbackMode.PYTHON, context=BridgeErrorContext(route="ocr", provider=custom_llm_provider or "", model=model), - ) + ).value async def aocr( @@ -135,10 +135,12 @@ async def aocr( timeout_seconds=_timeout_to_seconds(timeout), ) ) - return await ainvoke( - native_call=native_call, - fallback=async_none, - adapt=identity, - mode=FallbackMode.PYTHON, - context=BridgeErrorContext(route="ocr", provider=custom_llm_provider or "", model=model), - ) + return ( + await ainvoke( + native_call=native_call, + fallback=async_none, + adapt=identity, + mode=FallbackMode.PYTHON, + context=BridgeErrorContext(route="ocr", provider=custom_llm_provider or "", model=model), + ) + ).value diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 4d505a757c9..4c1c2330f5c 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -10,6 +10,8 @@ from websockets.exceptions import ConnectionClosedOK from litellm.rust_bridge.bindings import UNSET, NativeBinding, Unset from litellm.rust_bridge.runtime import ( BridgeErrorContext, + CoreEngine, + ExecutionResult, FallbackMode, acall, ainvoke, @@ -36,9 +38,7 @@ class RustResponsesWebSocketConnection(Protocol): ) -> RustResponsesWebSocket: ... -_CONNECTION: Final = NativeBinding[type[RustResponsesWebSocketConnection]]( - "ResponsesWebSocketConnection" -) +_CONNECTION: Final = NativeBinding[type[RustResponsesWebSocketConnection]]("ResponsesWebSocketConnection") def set_rust_responses_websocket( @@ -55,6 +55,7 @@ def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | class _ConnectionAdapter: def __init__(self, connection: RustResponsesWebSocket): self._connection: Final = connection + self.core_engine: Final = CoreEngine.RUST async def send(self, text: str) -> None: await acall( @@ -83,7 +84,7 @@ async def connect( url: str, headers: dict[str, str], timeout: float | httpx.Timeout | None, -) -> _ConnectionAdapter | None: +) -> ExecutionResult[_ConnectionAdapter | None]: connection_type: Final = load_rust_responses_websocket() native_call: Final = ( None diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 0bf60b1aa0b..e27099d4d22 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from enum import Enum from typing import Final, Generic, NoReturn, TypeAlias, TypeVar, cast @@ -11,12 +11,27 @@ from litellm.rust_bridge.bindings import native_exception_types NativeT = TypeVar("NativeT") ResultT = TypeVar("ResultT") +CORE_ENGINE_HIDDEN_PARAM: Final = "core_engine" +CORE_ENGINE_HEADER: Final = "x-litellm-core" +LEGACY_RUST_HEADER: Final = "x-litellm-rust" + class FallbackMode(Enum): PYTHON = "python" RUST_REQUIRED = "rust_required" +class CoreEngine(str, Enum): + PYTHON = "python" + RUST = "rust" + + +@dataclass(frozen=True, slots=True) +class ExecutionResult(Generic[ResultT]): + value: ResultT + source: CoreEngine + + @dataclass(frozen=True, slots=True) class RustHandled(Generic[ResultT]): value: ResultT @@ -42,6 +57,48 @@ class BridgeErrorContext: model: str +def execution_headers(source: CoreEngine) -> dict[str, str]: + if source is CoreEngine.RUST: + return { # mutable-ok: response adapters require a mutable header dict + CORE_ENGINE_HEADER: source.value, + LEGACY_RUST_HEADER: "true", + } + return {CORE_ENGINE_HEADER: source.value} # mutable-ok: response adapters require a mutable header dict + + +def execution_additional_headers( + additional_headers: Mapping[str, object] | None, + source: CoreEngine, +) -> dict[str, str]: + existing: Final = additional_headers or {} # mutable-ok: empty default is local and never mutated + reserved: Final = frozenset({CORE_ENGINE_HEADER, LEGACY_RUST_HEADER}) + preserved: Final = { # mutable-ok: response adapters require a mutable header dict + str(name): str(value) for name, value in existing.items() if str(name).lower() not in reserved + } + return { # mutable-ok: response adapters require a mutable header dict + **preserved, + **execution_headers(source), + } + + +def execution_hidden_params( + hidden_params: Mapping[str, object] | None, + source: CoreEngine, +) -> dict[str, object]: + existing: Final = hidden_params or {} # mutable-ok: empty default is local and never mutated + raw_headers: Final = existing.get("additional_headers") + additional_headers: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_headers) # cast-ok: the isinstance check validates the mapping boundary + if isinstance(raw_headers, Mapping) + else {} # mutable-ok: empty default is local and never mutated + ) + return { # mutable-ok: response objects require mutable hidden params + **existing, + CORE_ENGINE_HIDDEN_PARAM: source.value, + "additional_headers": execution_additional_headers(additional_headers, source), + } + + def invoke( *, native_call: Callable[[], NativeT] | None, @@ -49,12 +106,12 @@ def invoke( adapt: Callable[[NativeT], ResultT], mode: FallbackMode, context: BridgeErrorContext, -) -> ResultT: +) -> ExecutionResult[ResultT]: result: Final = attempt(native_call=native_call, adapt=adapt, context=context) if isinstance(result, RustHandled): - return result.value + return ExecutionResult(value=result.value, source=CoreEngine.RUST) if mode is FallbackMode.PYTHON: - return fallback() + return ExecutionResult(value=fallback(), source=CoreEngine.PYTHON) _raise_required(result, context) @@ -65,12 +122,12 @@ async def ainvoke( adapt: Callable[[NativeT], ResultT], mode: FallbackMode, context: BridgeErrorContext, -) -> ResultT: +) -> ExecutionResult[ResultT]: result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context) if isinstance(result, RustHandled): - return result.value + return ExecutionResult(value=result.value, source=CoreEngine.RUST) if mode is FallbackMode.PYTHON: - return await fallback() + return ExecutionResult(value=await fallback(), source=CoreEngine.PYTHON) _raise_required(result, context) @@ -164,12 +221,16 @@ def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoRetu message_value: Final = args[1] if len(args) > 1 else str(error) status: Final = status_value if isinstance(status_value, int) else 0 message: Final = message_value if isinstance(message_value, str) else str(message_value) - raise APIError( + api_error: Final = APIError( status_code=status or 500, message=f"litellm rust {context.route}: {message}", llm_provider=context.provider, model=context.model, - ) from error + ) + api_error.headers = execution_headers( # pyright: ignore[reportAttributeAccessIssue] # proxy reads exception headers + CoreEngine.RUST + ) + raise api_error from error def identity(value: ResultT) -> ResultT: diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index bf1e3fb60d8..295b0d4ad01 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -102,7 +102,7 @@ def transcription( adapt=identity, mode=FallbackMode.RUST_REQUIRED, context=BridgeErrorContext(route="audio transcription", provider=custom_llm_provider or "", model=model), - ) + ).value async def atranscription( @@ -131,10 +131,12 @@ async def atranscription( timeout_seconds=timeout_to_seconds(timeout), ) ) - return await ainvoke( - native_call=native_call, - fallback=async_none, - adapt=identity, - mode=FallbackMode.RUST_REQUIRED, - context=BridgeErrorContext(route="audio transcription", provider=custom_llm_provider or "", model=model), - ) + return ( + await ainvoke( + native_call=native_call, + fallback=async_none, + adapt=identity, + mode=FallbackMode.RUST_REQUIRED, + context=BridgeErrorContext(route="audio transcription", provider=custom_llm_provider or "", model=model), + ) + ).value diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 362d3a5f5a0..f1a4d584ab6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -431,7 +431,11 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)) stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) - assert stream._hidden_params["additional_headers"] == {} + assert stream._hidden_params["core_engine"] == "rust" + assert stream._hidden_params["additional_headers"] == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + } chunks = [chunk async for chunk in stream] joined = b"".join(chunks) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index e8f3ea3a0a9..395b151b224 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -450,6 +450,11 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" + assert response._hidden_params["core_engine"] == "rust" + assert response._hidden_params["additional_headers"] == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + } assert bridge.calls[0] == { "model": "mistral-ocr-latest", "document": DOCUMENT, @@ -643,6 +648,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" + assert response._hidden_params["core_engine"] == "rust" assert len(fake_bridge.calls) == 1 call = fake_bridge.calls[0] assert call["model"] == "mistral-ocr-latest" @@ -792,6 +798,8 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): assert captured.get("called") is True # Python path was used assert isinstance(response, OCRResponse) + assert response._hidden_params["core_engine"] == "python" + assert response._hidden_params["additional_headers"] == {"x-litellm-core": "python"} def test_ocr_unsupported_provider_skips_rust(monkeypatch): @@ -810,6 +818,7 @@ def test_ocr_unsupported_provider_skips_rust(monkeypatch): ) assert isinstance(response, OCRResponse) + assert response._hidden_params["core_engine"] == "python" assert bridge.calls == [] diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 9aa9c4acd3d..12f75bc32f0 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -78,14 +78,13 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) - assert ( - await responses_websocket.connect( - url="wss://example.test/responses", - headers={}, - timeout=None, - ) - is None + result = await responses_websocket.connect( + url="wss://example.test/responses", + headers={}, + timeout=None, ) + assert result.value is None + assert result.source is responses_websocket.CoreEngine.PYTHON @pytest.mark.asyncio @@ -94,13 +93,16 @@ async def test_enabled_bridge_connects_and_adapts_socket( ) -> None: responses_websocket.set_rust_responses_websocket(connection=_FakeNativeBridge) - connection = await responses_websocket.connect( + result = await responses_websocket.connect( url="wss://example.test/responses", headers={"Authorization": "Bearer key"}, timeout=1.0, ) + assert result.source is responses_websocket.CoreEngine.RUST + connection = result.value assert connection is not None + assert connection.core_engine is responses_websocket.CoreEngine.RUST await connection.send("response.create") assert await connection.recv() == "response.completed" await connection.close() diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index a7c8570e78d..fd114c0b9a3 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -264,6 +264,11 @@ class TestSyncCall: assert result.usage.completion_tokens == 4 assert result.usage.total_tokens == 15 assert result.id == original_id, "the rust path must keep the chatcmpl id litellm already minted" + assert result._hidden_params["core_engine"] == "rust" + assert result._hidden_params["additional_headers"] == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + } def test_passes_the_timeout_through_as_seconds(self): native = _RecordingCall() @@ -291,6 +296,8 @@ class TestSyncCall: ) assert result is fallback_response + assert result._hidden_params["core_engine"] == "python" + assert result._hidden_params["additional_headers"] == {"x-litellm-core": "python"} class TestAsyncCall: @@ -300,6 +307,7 @@ class TestAsyncCall: result = await bridge.achat_completions_or_fallback(**_async_call_kwargs(ModelResponse())) assert result is not None assert result.choices[0].message.content == "hello from rust" + assert result._hidden_params["core_engine"] == "rust" @pytest.mark.asyncio async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index b0fa510069b..9d0fd6d2098 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -44,7 +44,8 @@ def test_invoke_tags_native_decline_before_running_fallback() -> None: context=context(), ) - assert value == "fallback" + assert value.value == "fallback" + assert value.source is runtime.CoreEngine.PYTHON assert calls == ["rust", "python"] @@ -62,6 +63,10 @@ def test_invoke_translates_upstream_without_fallback() -> None: ) assert caught.value.status_code == 429 + assert caught.value.headers == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + } @pytest.mark.asyncio @@ -72,16 +77,38 @@ async def test_ainvoke_handles_native_success() -> None: async def fallback() -> str: pytest.fail("fallback must not run") - assert ( - await runtime.ainvoke( - native_call=native, - fallback=fallback, - adapt=str, - mode=runtime.FallbackMode.PYTHON, - context=context(), - ) - == "3" + result = await runtime.ainvoke( + native_call=native, + fallback=fallback, + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), ) + assert result.value == "3" + assert result.source is runtime.CoreEngine.RUST + + +def test_execution_hidden_params_overwrites_reserved_provider_headers() -> None: + hidden_params = runtime.execution_hidden_params( + { + "provider": "anthropic", + "additional_headers": { + "request-id": "req-1", + "X-LiteLLM-Core": "spoofed", + "x-litellm-rust": "spoofed", + }, + }, + runtime.CoreEngine.PYTHON, + ) + + assert hidden_params == { + "provider": "anthropic", + "core_engine": "python", + "additional_headers": { + "request-id": "req-1", + "x-litellm-core": "python", + }, + } def test_required_mode_rejects_unavailable_bridge() -> None: diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 92e86f48c31..f9fdfe84c25 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -133,6 +133,11 @@ def test_bedrock_transcription_uses_rust_only_path() -> None: rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) assert response.text == "rust" + assert response._hidden_params["core_engine"] == "rust" + assert response._hidden_params["additional_headers"] == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + } @pytest.mark.asyncio @@ -150,3 +155,8 @@ async def test_bedrock_atranscription_uses_rust_only_path() -> None: rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) assert response.text == "rust" + assert response._hidden_params["core_engine"] == "rust" + assert response._hidden_params["additional_headers"] == { + "x-litellm-core": "rust", + "x-litellm-rust": "true", + }