From 87b7be97a20793e70f09fcb68709bff38f4c73f6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 13 Oct 2026 15:40:25 -0700 Subject: [PATCH] refactor(rust-bridge): fold lifecycle and public_execution into one execution module The public stream wrappers only forwarded to the lifecycle streams and added the header dict, so the streams now carry _hidden_params themselves. lifecycle.py no longer held any call lifecycle, just the Execution driver, so it is renamed to execution.py and absorbs the synchronous and asynchronous adapters Co-Authored-By: Claude Opus 5.5 --- litellm-rust/crates/coroutine/AGENTS.md | 2 +- litellm-rust/crates/host-python/AGENTS.md | 2 +- .../crates/host-python/tests/driver.rs | 2 +- .../crates/host-python/tests/support/mod.rs | 4 +- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- .../crates/python-bridge/src/routes/mod.rs | 2 +- litellm/caching/redis_cache.py | 3 +- litellm/rust_bridge/AGENTS.md | 2 +- litellm/rust_bridge/_native.pyi | 2 +- .../chat_completions/entrypoints.py | 10 ++-- litellm/rust_bridge/embeddings/entrypoints.py | 10 ++-- .../{lifecycle.py => execution.py} | 33 +++++++++-- litellm/rust_bridge/messages/entrypoints.py | 10 ++-- litellm/rust_bridge/ocr/entrypoints.py | 10 ++-- litellm/rust_bridge/public_execution.py | 58 ------------------- litellm/rust_bridge/responses/entrypoints.py | 10 ++-- tests/test_litellm_rust/test_inference.py | 12 ++-- tests/unit/caching/test_redis_cache.py | 12 ++-- tests/unit/rust_bridge/test_bindings.py | 2 +- .../{test_lifecycle.py => test_execution.py} | 6 +- tests/unit/rust_bridge/test_runtime.py | 3 +- 21 files changed, 79 insertions(+), 118 deletions(-) rename litellm/rust_bridge/{lifecycle.py => execution.py} (80%) delete mode 100644 litellm/rust_bridge/public_execution.py rename tests/unit/rust_bridge/{test_lifecycle.py => test_execution.py} (95%) diff --git a/litellm-rust/crates/coroutine/AGENTS.md b/litellm-rust/crates/coroutine/AGENTS.md index 403d53a0a6d..0c3c737b0ef 100644 --- a/litellm-rust/crates/coroutine/AGENTS.md +++ b/litellm-rust/crates/coroutine/AGENTS.md @@ -3,7 +3,7 @@ Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks - R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime -- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context +- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/execution.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context - R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime - R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states - R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index e5324cd4688..04e8dab997c 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -37,7 +37,7 @@ Terminal notifications share the selected response or exception. An ordinary not - Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal - Keep coroutine driving in the shared Python driver and the native handle - - Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding + - Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding - Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html) - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html) diff --git a/litellm-rust/crates/host-python/tests/driver.rs b/litellm-rust/crates/host-python/tests/driver.rs index bc8c5a0afbc..157cfc92148 100644 --- a/litellm-rust/crates/host-python/tests/driver.rs +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -393,7 +393,7 @@ fn a_stream_carries_its_head_before_the_first_chunk( call_options(asynchronous), ); let hidden: std::collections::HashMap> = - stream.getattr("head").unwrap().extract().unwrap(); + stream.getattr("_hidden_params").unwrap().extract().unwrap(); assert_eq!( hidden["additional_headers"], std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())]) diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index e71592124a6..ef09135698e 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -70,14 +70,14 @@ pub fn call_options(asynchronous: bool) -> CallOptions { pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> { let source = std::ffi::CString::new(concat!( - include_str!("../../../../../litellm/rust_bridge/lifecycle.py"), + include_str!("../../../../../litellm/rust_bridge/execution.py"), "\n_drive = drive\n\n\nasync def drive(execution):\n return await _drive(execution, AsyncStream)\n", )) .unwrap(); PyModule::from_code( py, &source, - pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("execution.py"), pyo3::ffi::c_str!("host_test_lifecycle"), ) .unwrap() diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 28ea62468f9..c435ba37e78 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -48,7 +48,7 @@ Each step needs focused regression tests in the owning crate and Python integrat - Preserve provider-specific upload/submission/poll observation and encoding boundaries. Signed/build-captured bytes must not be silently reserialized - Conversion errors and every failure after the call starts are terminal - Disabled/unavailable native execution may select legacy once. Callback exceptions never authorize fallback or replay -- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle and call driver in `litellm-host-python` +- Use one ordinary inline `async def` driver in `litellm/rust_bridge/execution.py`, with the native handle and call driver in `litellm-host-python` - Contract: generator-style `send` (`send(None)` starts), `throw`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values - Validate Created/Running/Suspended/Closed protocol states. The machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy - Defer effectful setup/context reads/timestamps until start. Unstarted-handle destruction releases inputs independently of Python `finally` diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 976608a8d69..08c44efc58e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -130,7 +130,7 @@ fn call_hooks( } fn execution_steps(py: Python<'_>) -> PyResult> { - py.import("litellm.rust_bridge.lifecycle") + py.import("litellm.rust_bridge.execution") } fn run_public_call( diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index af1d750df68..ecfeb02b37d 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -100,8 +100,7 @@ _GENERIC_CALLER_MODULES: Final = frozenset( "litellm.caching.redis_batch", "litellm.caching.dual_cache", "litellm.caching.caching", - "litellm.rust_bridge.lifecycle", - "litellm.rust_bridge.public_execution", + "litellm.rust_bridge.execution", "contextlib", } ) diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index a62366ee84c..607bb257a0e 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -4,7 +4,7 @@ This package owns native rollout and fallback selection, Python public API compa `host/` holds the Python modules Rust imports by name to read Python-owned state (provider resolution, model capabilities, settings, caches, diagnostics, legacy callbacks, the secret manager, transport and timeout conversions). Its AGENTS.md states that contract -`lifecycle.py` owns generic inline execution and stream iteration, plus the step classes Rust constructs. Native entrypoints return an unstarted `Execution`, and `public_execution.py` hands it to the public caller through `lifecycle.run` or `lifecycle.drive` with the product stream wrappers, so Rust never names a Python driver or stream class. Keep one driver implementing generator-style `send` (`send(None)` starts), `throw` and idempotent `close`. Do not create a second implementation +`execution.py` owns generic inline execution and stream iteration, plus the step classes Rust constructs. Native entrypoints return an unstarted `Execution`, and `execution.synchronous` or `execution.asynchronous` hands it to the public caller through `run` or `drive` with the one stream pair, so Rust never names a Python driver or stream class. Keep one driver implementing generator-style `send` (`send(None)` starts), `throw` and idempotent `close`. Do not create a second implementation Generic execution steps and inline suspension handling must not depend on LiteLLM response metadata. Public stream construction, `_hidden_params` and header compatibility remain product responsibilities. Keep generic stream heads opaque and preserve public header behavior in the product wrappers diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f42b41e9477..7364a1d72f2 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -5,7 +5,7 @@ from typing import Never, final import httpx from pydantic import JsonValue -from litellm.rust_bridge.lifecycle import Execution +from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall class RustBridgeDeclined(Exception): ... diff --git a/litellm/rust_bridge/chat_completions/entrypoints.py b/litellm/rust_bridge/chat_completions/entrypoints.py index 3f40592cb74..8f6a3033103 100644 --- a/litellm/rust_bridge/chat_completions/entrypoints.py +++ b/litellm/rust_bridge/chat_completions/entrypoints.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables -from litellm.rust_bridge import public_execution +from litellm.rust_bridge import execution from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.public_call import NativeCall from litellm.types.utils import ModelResponse @@ -26,15 +26,15 @@ class NativeAcompletion(Protocol): def _completion_binding(value: object) -> NativeCompletion | None: if not callable(value): return None - native: Final = public_execution.synchronous(value) - return cast("NativeCompletion", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.synchronous(value) + return cast("NativeCompletion", native) # cast-ok: the execution adapter returns the route's result def _acompletion_binding(value: object) -> NativeAcompletion | None: if not callable(value): return None - native: Final = public_execution.asynchronous(value) - return cast("NativeAcompletion", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.asynchronous(value) + return cast("NativeAcompletion", native) # cast-ok: the execution adapter returns the route's result NATIVE_COMPLETION: Final = NativeBinding("completion", validate=_completion_binding) diff --git a/litellm/rust_bridge/embeddings/entrypoints.py b/litellm/rust_bridge/embeddings/entrypoints.py index a512f7d2a92..8b027a6bd6a 100644 --- a/litellm/rust_bridge/embeddings/entrypoints.py +++ b/litellm/rust_bridge/embeddings/entrypoints.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables -from litellm.rust_bridge import public_execution +from litellm.rust_bridge import execution from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.public_call import NativeCall from litellm.types.utils import EmbeddingResponse @@ -26,15 +26,15 @@ class NativeAembedding(Protocol): def _embedding_binding(value: object) -> NativeEmbedding | None: if not callable(value): return None - native: Final = public_execution.synchronous(value) - return cast("NativeEmbedding", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.synchronous(value) + return cast("NativeEmbedding", native) # cast-ok: the execution adapter returns the route's result def _aembedding_binding(value: object) -> NativeAembedding | None: if not callable(value): return None - native: Final = public_execution.asynchronous(value) - return cast("NativeAembedding", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.asynchronous(value) + return cast("NativeAembedding", native) # cast-ok: the execution adapter returns the route's result NATIVE_EMBEDDING: Final = NativeBinding("embedding", validate=_embedding_binding) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/execution.py similarity index 80% rename from litellm/rust_bridge/lifecycle.py rename to litellm/rust_bridge/execution.py index 1cfe4afd2a5..04e934ab8f8 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/execution.py @@ -1,9 +1,13 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from contextlib import ExitStack from dataclasses import dataclass -from typing import Final, Protocol +from typing import Final, Protocol, TypeVar, cast # noqa: TID251 # trusts the native stub at the binding boundary + +from pydantic import TypeAdapter + +_Call = TypeVar("_Call") @dataclass(frozen=True, slots=True) @@ -40,6 +44,11 @@ class Execution(Protocol): _REOPENED: Final = "execution opened a stream twice" +_HEADERS: Final = TypeAdapter(Mapping[str, object]) + + +def _headers(head: object) -> Mapping[str, object]: + return _HEADERS.validate_python(head) class StreamClosed(Exception): @@ -91,13 +100,25 @@ def run(execution: Execution, open_stream: Callable[[Execution, object], object] return _hand_off(execution, _settled(execution.send(None)), open_stream, cleanup) +def synchronous(native: Callable[[_Call], object]) -> Callable[[_Call], object]: + return lambda call: run(_execution(native(call)), SyncStream) + + +def asynchronous(native: Callable[[_Call], object]) -> Callable[[_Call], Awaitable[object]]: + return lambda call: drive(_execution(native(call)), AsyncStream) + + +def _execution(value: object) -> Execution: + return cast("Execution", value) # cast-ok: native entrypoints return Execution per _native.pyi + + class AsyncStream(AsyncIterator[object]): """A streamed native call: each read resumes the execution until its next chunk.""" - def __init__(self, execution: Execution, head: object = None) -> None: + def __init__(self, execution: Execution, head: object) -> None: self._execution: Final = execution self._done = False - self.head: Final = head + self._hidden_params: dict[str, object] = dict(_headers(head)) # mutable-ok: header writers mutate it def __aiter__(self) -> AsyncStream: return self @@ -136,10 +157,10 @@ class AsyncStream(AsyncIterator[object]): class SyncStream(Iterator[object]): """The sync form of `AsyncStream`; its execution never suspends on an awaitable.""" - def __init__(self, execution: Execution, head: object = None) -> None: + def __init__(self, execution: Execution, head: object) -> None: self._execution: Final = execution self._done = False - self.head: Final = head + self._hidden_params: dict[str, object] = dict(_headers(head)) # mutable-ok: header writers mutate it def __iter__(self) -> SyncStream: return self diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py index 63a5c2489c9..756d62b747a 100644 --- a/litellm/rust_bridge/messages/entrypoints.py +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import AsyncIterator, Awaitable, Iterator from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables -from litellm.rust_bridge import public_execution +from litellm.rust_bridge import execution from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -26,15 +26,15 @@ class NativeAmessages(Protocol): def _messages_binding(value: object) -> NativeMessages | None: if not callable(value): return None - native: Final = public_execution.synchronous(value) - return cast("NativeMessages", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.synchronous(value) + return cast("NativeMessages", native) # cast-ok: the execution adapter returns the route's result def _amessages_binding(value: object) -> NativeAmessages | None: if not callable(value): return None - native: Final = public_execution.asynchronous(value) - return cast("NativeAmessages", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.asynchronous(value) + return cast("NativeAmessages", native) # cast-ok: the execution adapter returns the route's result NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding) diff --git a/litellm/rust_bridge/ocr/entrypoints.py b/litellm/rust_bridge/ocr/entrypoints.py index 7dabd89e0e4..09a98d8dbfc 100644 --- a/litellm/rust_bridge/ocr/entrypoints.py +++ b/litellm/rust_bridge/ocr/entrypoints.py @@ -4,7 +4,7 @@ from collections.abc import Awaitable, Mapping from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRResponse -from litellm.rust_bridge import public_execution +from litellm.rust_bridge import execution from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.public_call import NativeCall @@ -26,15 +26,15 @@ class NativeAocr(Protocol): def _ocr_binding(value: object) -> NativeOcr | None: if not callable(value): return None - native: Final = public_execution.synchronous(value) - return cast("NativeOcr", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.synchronous(value) + return cast("NativeOcr", native) # cast-ok: the execution adapter returns the route's result def _aocr_binding(value: object) -> NativeAocr | None: if not callable(value): return None - native: Final = public_execution.asynchronous(value) - return cast("NativeAocr", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.asynchronous(value) + return cast("NativeAocr", native) # cast-ok: the execution adapter returns the route's result class NativeOcrHealthCheckDocument(Protocol): diff --git a/litellm/rust_bridge/public_execution.py b/litellm/rust_bridge/public_execution.py deleted file mode 100644 index c035e68dd25..00000000000 --- a/litellm/rust_bridge/public_execution.py +++ /dev/null @@ -1,58 +0,0 @@ -from __future__ import annotations - -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping -from typing import Final, cast # noqa: TID251 # trusts the native stub at the binding boundary - -from pydantic import TypeAdapter - -from litellm.rust_bridge import lifecycle -from litellm.rust_bridge.lifecycle import Execution -from litellm.rust_bridge.public_call import NativeCall - -_HEADERS: Final = TypeAdapter(Mapping[str, object]) - - -def synchronous(native: Callable[[NativeCall], object]) -> Callable[[NativeCall], object]: - return lambda call: lifecycle.run(_execution(native(call)), SyncStream) - - -def asynchronous(native: Callable[[NativeCall], object]) -> Callable[[NativeCall], Awaitable[object]]: - return lambda call: lifecycle.drive(_execution(native(call)), AsyncStream) - - -def _execution(value: object) -> Execution: - return cast("Execution", value) # cast-ok: native entrypoints return Execution per _native.pyi - - -class AsyncStream(AsyncIterator[object]): - def __init__(self, execution: Execution, head: object) -> None: - self._stream: Final = lifecycle.AsyncStream(execution, head) - self._hidden_params: dict[str, object] = dict(_headers(head)) # mutable-ok: header writers mutate it - - def __aiter__(self) -> AsyncStream: - return self - - async def __anext__(self) -> object: - return await self._stream.__anext__() - - async def aclose(self) -> None: - await self._stream.aclose() - - -class SyncStream(Iterator[object]): - def __init__(self, execution: Execution, head: object) -> None: - self._stream: Final = lifecycle.SyncStream(execution, head) - self._hidden_params: dict[str, object] = dict(_headers(head)) # mutable-ok: header writers mutate it - - def __iter__(self) -> SyncStream: - return self - - def __next__(self) -> object: - return next(self._stream) - - def close(self) -> None: - self._stream.close() - - -def _headers(head: object) -> Mapping[str, object]: - return _HEADERS.validate_python(head) diff --git a/litellm/rust_bridge/responses/entrypoints.py b/litellm/rust_bridge/responses/entrypoints.py index 43411b4dc2c..909d00d8131 100644 --- a/litellm/rust_bridge/responses/entrypoints.py +++ b/litellm/rust_bridge/responses/entrypoints.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables -from litellm.rust_bridge import public_execution +from litellm.rust_bridge import execution from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.openai import ResponsesAPIResponse @@ -26,15 +26,15 @@ class NativeAresponses(Protocol): def _responses_binding(value: object) -> NativeResponses | None: if not callable(value): return None - native: Final = public_execution.synchronous(value) - return cast("NativeResponses", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.synchronous(value) + return cast("NativeResponses", native) # cast-ok: the execution adapter returns the route's result def _aresponses_binding(value: object) -> NativeAresponses | None: if not callable(value): return None - native: Final = public_execution.asynchronous(value) - return cast("NativeAresponses", native) # cast-ok: the public execution adapter returns the route's result + native: Final = execution.asynchronous(value) + return cast("NativeAresponses", native) # cast-ok: the execution adapter returns the route's result NATIVE_RESPONSES: Final = NativeBinding("responses", validate=_responses_binding) diff --git a/tests/test_litellm_rust/test_inference.py b/tests/test_litellm_rust/test_inference.py index 7ef9acd9ac5..716dd6b2574 100644 --- a/tests/test_litellm_rust/test_inference.py +++ b/tests/test_litellm_rust/test_inference.py @@ -11,7 +11,7 @@ from litellm.chat_completions import dispatch as chat_dispatch from litellm.integrations.custom_logger import CustomLogger from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.rust_bridge import _native, public_execution +from litellm.rust_bridge import _native, execution from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import CallTypes, ModelResponse @@ -48,7 +48,7 @@ def route(request: pytest.FixtureRequest) -> Route: def _public(native: Callable[[NativeCall], object], asynchronous: bool) -> Callable[[NativeCall], object]: - return public_execution.asynchronous(native) if asynchronous else public_execution.synchronous(native) + return execution.asynchronous(native) if asynchronous else execution.synchronous(native) def native_call(route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]) -> object: @@ -309,13 +309,13 @@ async def test_native_projection_reads_positional_parameters(route: Route, recor args: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.25) request: Final = chat_dispatch.request(args, kwargs) assert request is not None - await asyncio.to_thread(public_execution.synchronous(_native.completion), request) + await asyncio.to_thread(execution.synchronous(_native.completion), request) assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.25 else: response_args: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16) response_request: Final = responses_dispatch.request(response_args, kwargs) assert response_request is not None - await asyncio.to_thread(public_execution.synchronous(_native.responses), response_request) + await asyncio.to_thread(execution.synchronous(_native.responses), response_request) body: Final = _OBJECT.validate_python(recording_server.requests[0].body) assert body["instructions"] == "Be brief" assert body["max_output_tokens"] == 16 @@ -389,9 +389,9 @@ async def test_native_chat_uses_bound_positional_parameters( call: Final = chat_dispatch._DISPATCH.request(arguments, supplied) # pyright: ignore[reportPrivateUsage] # exercise the native request produced by public binding assert call is not None result: Final = ( - await public_execution.asynchronous(_native.acompletion)(call) + await execution.asynchronous(_native.acompletion)(call) if asynchronous - else await asyncio.to_thread(public_execution.synchronous(_native.completion), call) + else await asyncio.to_thread(execution.synchronous(_native.completion), call) ) assert isinstance(result, ModelResponse) assert result.choices[0].message.content == "Hello from native Messages" diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 57c32ace509..1d68feaadef 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1599,18 +1599,18 @@ def drive(): return _settle() -def test_call_stack_info_skips_native_lifecycle_frames(): - """The Rust execution awaits the response-cache coroutine from ``lifecycle._settle`` inside +def test_call_stack_info_skips_native_execution_frames(): + """The Rust execution awaits the response-cache coroutine from ``execution._settle`` inside ``drive``; those frames forward every native suspension, so the chain names the code that started the native call instead of ``_settle <- drive``.""" - lifecycle_globals = {"__name__": "litellm.rust_bridge.lifecycle", "_native_probe": _native_probe} - native_settle = types.FunctionType(_settle.__code__, lifecycle_globals, "_settle") - native_drive = types.FunctionType(drive.__code__, {**lifecycle_globals, "_settle": native_settle}, "drive") + execution_globals = {"__name__": "litellm.rust_bridge.execution", "_native_probe": _native_probe} + native_settle = types.FunctionType(_settle.__code__, execution_globals, "_settle") + native_drive = types.FunctionType(drive.__code__, {**execution_globals, "_settle": native_settle}, "drive") def anthropic_messages(): return native_drive() - assert anthropic_messages() == "anthropic_messages <- test_call_stack_info_skips_native_lifecycle_frames" + assert anthropic_messages() == "anthropic_messages <- test_call_stack_info_skips_native_execution_frames" class _FormatCountingStr(str): diff --git a/tests/unit/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py index eed2947a48c..39a157220dc 100644 --- a/tests/unit/rust_bridge/test_bindings.py +++ b/tests/unit/rust_bridge/test_bindings.py @@ -9,7 +9,7 @@ import pytest from litellm.rust_bridge import bindings from litellm.rust_bridge.chat_completions import entrypoints as chat_completions from litellm.rust_bridge.embeddings import entrypoints as embeddings -from litellm.rust_bridge.lifecycle import Complete, Step +from litellm.rust_bridge.execution import Complete, Step from litellm.rust_bridge.messages import entrypoints as messages from litellm.rust_bridge.ocr import entrypoints as ocr from litellm.rust_bridge.public_call import NativeCall diff --git a/tests/unit/rust_bridge/test_lifecycle.py b/tests/unit/rust_bridge/test_execution.py similarity index 95% rename from tests/unit/rust_bridge/test_lifecycle.py rename to tests/unit/rust_bridge/test_execution.py index b4d62786ec4..88c77904d8e 100644 --- a/tests/unit/rust_bridge/test_lifecycle.py +++ b/tests/unit/rust_bridge/test_execution.py @@ -6,7 +6,7 @@ from typing import Final import pytest -from litellm.rust_bridge.lifecycle import ( +from litellm.rust_bridge.execution import ( AsyncStream, Await, Complete, @@ -95,11 +95,11 @@ def test_stream_read_rejects_a_second_open_and_closes_the_execution(asynchronous execution: Final = ScriptedExecution([StreamChunk("first"), StreamOpen("head")]) async def read() -> tuple[object, ...]: - stream: Final = AsyncStream(execution) + stream: Final = AsyncStream(execution, {}) return (await anext(stream), await anext(stream)) def read_sync() -> tuple[object, ...]: - stream: Final = SyncStream(execution) + stream: Final = SyncStream(execution, {}) return (next(stream), next(stream)) with pytest.raises(RuntimeError, match="opened a stream twice"): diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index 2159993d2ac..95fa3e2cc7b 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/tests/unit/rust_bridge/test_runtime.py @@ -12,8 +12,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from litellm.rust_bridge import bindings, configuration, runtime from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.lifecycle import Complete, StreamChunk -from litellm.rust_bridge.public_execution import AsyncStream, SyncStream +from litellm.rust_bridge.execution import AsyncStream, Complete, StreamChunk, SyncStream class RustBridgeDeclined(Exception):