scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

fix(rust): document mutable bridge boundaries

Browse source
This commit is contained in:
Yujong Lee 2026-09-08 09:25:11 -07:00
parent 6fd590e247
commit dbb1d30647
5 changed files with 130 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2429,7 +2429,7 @@ class BaseLLMHTTPHandler:
headers: dict,
request_body: dict,
timeout: float | httpx.Timeout | None,
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates the argument bag
) -> AnthropicMessagesResponse | None:
if custom_llm_provider not in ("azure_ai", "anthropic"):
return None

24
litellm/rust_bridge/_lifecycle.py
View file

@ -20,7 +20,9 @@ TerminalAction = Literal[
"sync_failure",
"async_failure",
]
_OPTIONAL_ARGUMENTS_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(dict[str, object] | None)
_OPTIONAL_ARGUMENTS_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(
dict[str, object] | None
) # mutable-ok: native bridge retains and updates Python argument objects
class NativeOutcome(IntEnum):
@ -51,8 +53,8 @@ class LifecycleHost(Protocol):
class MutableLifecycleHost(LifecycleHost, Protocol):
arguments: dict[str, object]
current: dict[str, object]
arguments: dict[str, object] # mutable-ok: native bridge retains and updates Python argument objects
current: dict[str, object] # mutable-ok: native bridge retains and updates Python argument objects
logger: object | None
response: object
error: BaseException | None
@ -81,7 +83,9 @@ def host_result(host: MutableLifecycleHost) -> object:
raise host.error
async def deployment_pre(arguments: dict[str, object], call_type: str) -> dict[str, object]:
async def deployment_pre(
arguments: dict[str, object], call_type: str
) -> dict[str, object]: # mutable-ok: native bridge retains and updates Python argument objects
from litellm import utils
modified: Final = _OPTIONAL_ARGUMENTS_ADAPTER.validate_python(
@ -90,7 +94,9 @@ async def deployment_pre(arguments: dict[str, object], call_type: str) -> dict[s
return arguments if modified is None else modified
async def deployment_success(arguments: dict[str, object], response: object, call_type: CallTypes) -> object:
async def deployment_success(
arguments: dict[str, object], response: object, call_type: CallTypes
) -> object: # mutable-ok: native bridge retains and updates Python argument objects
from litellm import utils
updated: object = await utils.async_post_call_success_deployment_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook annotations expose an unknown return
@ -99,7 +105,9 @@ async def deployment_success(arguments: dict[str, object], response: object, cal
return updated
async def deployment_failure(arguments: dict[str, object], error: BaseException | None, call_type: str) -> None:
async def deployment_failure(
arguments: dict[str, object], error: BaseException | None, call_type: str
) -> None: # mutable-ok: native bridge retains and updates Python argument objects
from litellm import utils
if not isinstance(error, Exception):
@ -152,7 +160,9 @@ async def drive_async(host: LifecycleHost) -> object:
return host.result()
def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: str) -> object:
def initialize_logging(
arguments: dict[str, object], asynchronous: bool, route: str
) -> object: # mutable-ok: native bridge retains and updates Python argument objects
from litellm.rust_bridge.ocr import initialize_logging as initialize_ocr_logging
return initialize_ocr_logging(arguments, asynchronous, route)

46
litellm/rust_bridge/chat_completions.py
View file

@ -69,12 +69,16 @@ RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
class RustChatCompletions(Protocol):
def __call__(self, arguments: dict[str, object]) -> ModelResponse:
def __call__(
self, arguments: dict[str, object]
) -> ModelResponse: # mutable-ok: native bridge retains and updates Python argument objects
raise NotImplementedError
class RustAchatCompletions(Protocol):
def __call__(self, arguments: dict[str, object]) -> Awaitable[ModelResponse]:
def __call__(
self, arguments: dict[str, object]
) -> Awaitable[ModelResponse]: # mutable-ok: native bridge retains and updates Python argument objects
raise NotImplementedError
@ -374,7 +378,7 @@ def chat_completions(
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates Python argument objects
logging_api_key: str | None = None,
on_response: ResponseObserver | None = None,
) -> ModelResponse | None:
@ -452,7 +456,7 @@ async def achat_completions(
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates Python argument objects
logging_api_key: str | None = None,
on_response: ResponseObserver | None = None,
) -> ModelResponse | None:
@ -533,7 +537,7 @@ async def achat_completions_or_fallback(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
python_fallback: Callable[[], Awaitable[object]],
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates Python argument objects
logging_api_key: str | None = None,
on_response: ResponseObserver | None = None,
) -> object:
@ -564,14 +568,16 @@ async def achat_completions_or_fallback(
return await python_fallback()
def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> object:
def initialize_logging(
arguments: dict[str, object], asynchronous: bool
) -> object: # mutable-ok: native bridge retains and updates Python argument objects
from litellm.rust_bridge._lifecycle import initialize_logging as initialize_lifecycle_logging
return initialize_lifecycle_logging(arguments, asynchronous, "completion")
def _arguments(
arguments: dict[str, object] | None,
arguments: dict[str, object] | None, # mutable-ok: native bridge retains and updates Python argument objects
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
@ -582,7 +588,7 @@ def _arguments(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
logging_api_key: str | None,
) -> dict[str, object]:
) -> dict[str, object]: # mutable-ok: native bridge retains and updates Python argument objects
return {
**(arguments or {}),
"model": model,
@ -599,8 +605,12 @@ def _arguments(
class _ChatCompletionsBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[dict[str, object], bool, bool], NativeLifecycle]
prepare: Callable[[dict[str, object], object], object]
Lifecycle: Callable[
[dict[str, object], bool, bool], NativeLifecycle
] # mutable-ok: native bridge retains and updates Python argument objects
prepare: Callable[
[dict[str, object], object], object
] # mutable-ok: native bridge retains and updates Python argument objects
send: Callable[[object], Awaitable[Mapping[str, object]]]
send_sync: Callable[[object], Mapping[str, object]]
terminal_record: Callable[[object], Mapping[str, object]]
@ -609,7 +619,7 @@ class _ChatCompletionsBindings(NativeLifecycleBindings, Protocol):
class _ChatCompletionsHost:
def __init__(
self,
arguments: dict[str, object],
arguments: dict[str, object], # mutable-ok: native bridge retains and updates Python argument objects
asynchronous: bool,
bindings: _ChatCompletionsBindings,
) -> None:
@ -617,8 +627,12 @@ class _ChatCompletionsHost:
self.bindings: _ChatCompletionsBindings = bindings
self.machine: NativeLifecycle = bindings.Lifecycle(arguments, asynchronous, utils.is_internal_call.get())
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.arguments: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.current: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.state: object | None = None
@ -705,7 +719,8 @@ class _ChatCompletionsHost:
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _ChatCompletionsBindings
arguments: dict[str, object],
bindings: _ChatCompletionsBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> ModelResponse:
result: Final = drive_sync(_ChatCompletionsHost(arguments, False, bindings))
if not isinstance(result, ModelResponse):
@ -714,7 +729,8 @@ def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the nativ
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _ChatCompletionsBindings
arguments: dict[str, object],
bindings: _ChatCompletionsBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> ModelResponse:
result: Final = await drive_async(_ChatCompletionsHost(arguments, True, bindings))
if not isinstance(result, ModelResponse):

62
litellm/rust_bridge/messages.py
View file

@ -34,15 +34,21 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
class RustMessages(Protocol):
def __call__(self, arguments: dict[str, object]) -> AnthropicMessagesResponse: ...
def __call__(
self, arguments: dict[str, object]
) -> AnthropicMessagesResponse: ... # mutable-ok: native bridge retains and updates Python argument objects
class RustAmessages(Protocol):
def __call__(self, arguments: dict[str, object]) -> Awaitable[AnthropicMessagesResponse]: ...
def __call__(
self, arguments: dict[str, object]
) -> Awaitable[
AnthropicMessagesResponse
]: ... # mutable-ok: native bridge retains and updates Python argument objects
class _MessagesLogging(Protocol):
model_call_details: dict[str, object]
model_call_details: dict[str, object] # mutable-ok: native bridge retains and updates Python argument objects
def _handle_anthropic_messages_response_logging(self, result: object) -> object: ...
@ -98,11 +104,15 @@ def load_rust_amessages() -> RustAmessages | None:
return _AMESSAGES.load()
def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> object:
def initialize_logging(
arguments: dict[str, object], asynchronous: bool
) -> object: # mutable-ok: native bridge retains and updates Python argument objects
return initialize_lifecycle_logging(arguments, asynchronous, "messages")
class _RetainedMessagesResponse(dict[str, object]):
class _RetainedMessagesResponse(
dict[str, object]
): # mutable-ok: native bridge retains and updates Python argument objects
def __init__(
self, response: AnthropicMessagesResponse, roots: object, logger: _MessagesLogging, start_time: datetime
) -> None:
@ -158,15 +168,15 @@ def retain_stream_response(
def _arguments(
arguments: dict[str, object],
arguments: dict[str, object], # mutable-ok: native bridge retains and updates Python argument objects
model: str,
body: dict[str, object],
body: dict[str, object], # mutable-ok: native bridge retains and updates Python argument objects
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
extra_headers: dict[str, object] | None, # mutable-ok: native bridge retains and updates Python argument objects
timeout: float | httpx.Timeout | None,
) -> dict[str, object]:
) -> dict[str, object]: # mutable-ok: native bridge retains and updates Python argument objects
return { # mutable-ok: the native bridge requires a concrete argument bag
**arguments,
"model": model,
@ -182,13 +192,13 @@ def _arguments(
def messages(
*,
model: str,
body: dict[str, object],
body: dict[str, object], # mutable-ok: native bridge retains and updates Python argument objects
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
extra_headers: dict[str, object] | None, # mutable-ok: native bridge retains and updates Python argument objects
timeout: float | httpx.Timeout | None,
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates Python argument objects
) -> AnthropicMessagesResponse | None:
implementation: Final = load_rust_messages()
if implementation is None:
@ -203,13 +213,13 @@ def messages(
async def amessages(
*,
model: str,
body: dict[str, object],
body: dict[str, object], # mutable-ok: native bridge retains and updates Python argument objects
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
extra_headers: dict[str, object] | None, # mutable-ok: native bridge retains and updates Python argument objects
timeout: float | httpx.Timeout | None,
arguments: dict[str, object] | None = None,
arguments: dict[str, object] | None = None, # mutable-ok: native bridge retains and updates Python argument objects
) -> AnthropicMessagesResponse | None:
implementation: Final = load_rust_amessages()
if implementation is None:
@ -227,20 +237,28 @@ class _MessagesLifecycle(NativeLifecycle, Protocol):
class _MessagesBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[bool, bool], _MessagesLifecycle]
prepare: Callable[[dict[str, object], object], object]
prepare: Callable[
[dict[str, object], object], object
] # mutable-ok: native bridge retains and updates Python argument objects
send: Callable[[object], Awaitable[AnthropicMessagesResponse]]
send_sync: Callable[[object], AnthropicMessagesResponse]
committed_failure: Callable[[], None]
class _MessagesHost:
def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _MessagesBindings) -> None:
def __init__(
self, arguments: dict[str, object], asynchronous: bool, bindings: _MessagesBindings
) -> None: # mutable-ok: native bridge retains and updates Python argument objects
from litellm import utils
self.bindings: _MessagesBindings = bindings
self.machine: _MessagesLifecycle = bindings.Lifecycle(asynchronous, utils.is_internal_call.get())
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.arguments: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.current: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.lifecycle_owned: bool = self.logger is None
@ -345,13 +363,15 @@ class _MessagesHost:
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _MessagesBindings
arguments: dict[str, object],
bindings: _MessagesBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> AnthropicMessagesResponse:
return cast(AnthropicMessagesResponse, drive_sync(_MessagesHost(arguments, False, bindings)))
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _MessagesBindings
arguments: dict[str, object],
bindings: _MessagesBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> AnthropicMessagesResponse:
return cast(AnthropicMessagesResponse, await drive_async(_MessagesHost(arguments, True, bindings)))

54
litellm/rust_bridge/ocr.py
View file

@ -29,11 +29,15 @@ from litellm.rust_bridge.bindings import NativeBinding
class RustOcr(Protocol):
def __call__(self, arguments: dict[str, object]) -> OCRResponse: ...
def __call__(
self, arguments: dict[str, object]
) -> OCRResponse: ... # mutable-ok: native bridge retains and updates Python argument objects
class RustAocr(Protocol):
def __call__(self, arguments: dict[str, object]) -> Awaitable[OCRResponse]: ...
def __call__(
self, arguments: dict[str, object]
) -> Awaitable[OCRResponse]: ... # mutable-ok: native bridge retains and updates Python argument objects
def _as_ocr(value: object) -> RustOcr | None:
@ -56,21 +60,27 @@ def load_rust_aocr() -> RustAocr | None:
return _AOCR.load()
def ocr(arguments: dict[str, object]) -> OCRResponse:
def ocr(
arguments: dict[str, object],
) -> OCRResponse: # mutable-ok: native bridge retains and updates Python argument objects
implementation: Final = load_rust_ocr()
if implementation is None:
raise RuntimeError("Rust OCR is enabled but the native OCR extension is unavailable")
return implementation(arguments)
async def aocr(arguments: dict[str, object]) -> OCRResponse:
async def aocr(
arguments: dict[str, object],
) -> OCRResponse: # mutable-ok: native bridge retains and updates Python argument objects
implementation: Final = load_rust_aocr()
if implementation is None:
raise RuntimeError("Rust OCR is enabled but the native OCR extension is unavailable")
return await implementation(arguments)
def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: str = "ocr") -> object:
def initialize_logging(
arguments: dict[str, object], asynchronous: bool, route: str = "ocr"
) -> object: # mutable-ok: native bridge retains and updates Python argument objects
import litellm
from litellm import utils
from litellm.integrations.custom_logger import CustomLogger
@ -277,22 +287,36 @@ class _OcrLifecycle(NativeLifecycle, Protocol):
class _OcrBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[dict[str, object], object | None, bool, bool], _OcrLifecycle]
prepare: Callable[[dict[str, object], object, bool], object]
Lifecycle: Callable[
[dict[str, object], object | None, bool, bool], _OcrLifecycle
] # mutable-ok: native bridge retains and updates Python argument objects
prepare: Callable[
[dict[str, object], object, bool], object
] # mutable-ok: native bridge retains and updates Python argument objects
pre_call: Callable[[object], None]
send: Callable[[object], Awaitable[dict[str, object]]]
send: Callable[
[object], Awaitable[dict[str, object]]
] # mutable-ok: native bridge retains and updates Python argument objects
send_sync: Callable[[object], OCRResponse]
finish: Callable[[dict[str, object]], OCRResponse]
finish: Callable[
[dict[str, object]], OCRResponse
] # mutable-ok: native bridge retains and updates Python argument objects
terminal_record: Callable[[object], Mapping[str, object]]
class _OcrHost:
def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _OcrBindings) -> None:
def __init__(
self, arguments: dict[str, object], asynchronous: bool, bindings: _OcrBindings
) -> None: # mutable-ok: native bridge retains and updates Python argument objects
from litellm import utils
self.bindings: _OcrBindings = bindings
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.arguments: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.current: dict[str, object] = (
arguments # mutable-ok: native bridge retains and updates Python argument objects
)
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.machine: _OcrLifecycle = bindings.Lifecycle(
@ -385,12 +409,14 @@ class _OcrHost:
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _OcrBindings
arguments: dict[str, object],
bindings: _OcrBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> OCRResponse:
return cast(OCRResponse, drive_sync(_OcrHost(arguments, False, bindings)))
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _OcrBindings
arguments: dict[str, object],
bindings: _OcrBindings, # mutable-ok: native bridge retains and updates Python argument objects
) -> OCRResponse:
return cast(OCRResponse, await drive_async(_OcrHost(arguments, True, bindings)))