scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(proxy): report Rust execution provenance

Browse source
This commit is contained in:
Yujong Lee 2026-09-02 06:56:42 -07:00 • committed by GitHub
parent 5178766454
commit 0f579cd28e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 341 additions and 98 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -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<AppState> {
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");

10
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -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,
)

9
litellm/llms/bedrock/audio_transcription/__init__.py
View file

@ -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

39
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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

29
litellm/ocr/main.py
View file

@ -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,

14
litellm/proxy/common_request_processing.py
View file

@ -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)

53
litellm/rust_bridge/chat_completions.py
View file

@ -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),
)
)

42
litellm/rust_bridge/messages.py
View file

@ -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

18
litellm/rust_bridge/ocr.py
View file

@ -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

9
litellm/rust_bridge/responses_websocket.py
View file

@ -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

79
litellm/rust_bridge/runtime.py
View file

@ -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:

18
litellm/rust_bridge/transcription.py
View file

@ -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

6
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -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)

9
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -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 == []

18
tests/test_litellm/responses/test_rust_bridge_websocket.py
View file

@ -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()

8
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -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):

47
tests/test_litellm/rust_bridge/test_runtime.py
View file

@ -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:

10
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -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",
}