From c6cd8732665d5606ff71863fe5ff2fa3921f6930 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 6 Sep 2026 18:00:33 -0700 Subject: [PATCH] wip --- .github/workflows/test-rust.yml | 18 +-- Makefile | 23 +++- litellm-rust/Cargo.lock | 24 +--- litellm-rust/README.md | 47 +++++++ litellm-rust/crates/python-bridge/Cargo.toml | 2 +- .../python-bridge/src/routes/bindings.rs | 118 ------------------ .../tests/fixtures/ocr_retained.py | 79 +++--------- .../python-bridge/tests/ocr_retained.rs | 50 ++++++-- .../tests/callback_lifecycle.rs | 37 +++++- .../tests/fixtures/callback_components.py | 13 +- .../tests/fixtures/callback_lifecycle.py | 50 ++++---- tests/test_litellm/ocr/test_rust_bridge.py | 21 ++-- .../rust_bridge/native_route_wheel_test.py | 104 ++++++++++++++- 13 files changed, 300 insertions(+), 286 deletions(-) rename tests/test_litellm/ocr/retained_boundary_fixture.py => litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py (89%) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 31b23d5a5c9..188d5b8400d 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,7 @@ on: push: paths: - "litellm-rust/**" + - "tests/test_litellm/**" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -11,7 +12,6 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,6 +21,7 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "tests/test_litellm/**" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -28,7 +29,6 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" permissions: @@ -128,14 +128,8 @@ jobs: - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - name: Install locked SDK dependencies for retained callback tests - run: uv sync --frozen --no-default-groups --no-install-project --python 3.12 + - name: Check Python fixtures for Cargo tests + run: make lint-rust-python-fixtures - - name: Test retained callbacks with real Python logging - env: - PYO3_PYTHON: ${{ github.workspace }}/.venv/bin/python - PYTHONPATH: ${{ github.workspace }}:${{ github.workspace }}/.venv/lib/python3.12/site-packages - LITELLM_LOCAL_MODEL_COST_MAP: "True" - run: >- - cargo test --manifest-path litellm-rust/Cargo.toml - -p litellm-python-interop --tests --locked -- --include-ignored + - name: Test retained callbacks and OCR with repo Python + run: make test-rust-python diff --git a/Makefile b/Makefile index 91835e19e3c..f5f2193ba86 100644 --- a/Makefile +++ b/Makefile @@ -11,7 +11,7 @@ lint-test-quality lint-test-quality-budget-update \ install-dev install-proxy-dev install-test-deps install-hooks \ install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \ - lint-install lint-fetch-base bootstrap + lint-install lint-fetch-base bootstrap install-rust-python-test-deps test-rust-python lint-rust-python-fixtures # Default target help: @@ -55,7 +55,12 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" +<<<<<<< HEAD @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" +======= + @echo " make test-rust-python - Run ignored Python-integrated Cargo tests (optional TEST_FILTER=substring)" + @echo " make lint-rust-python-fixtures - Check Rust test Python fixtures with Ruff" +>>>>>>> 1f18773bf2 (wip) @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." @@ -115,6 +120,9 @@ install-test-deps: install-proxy-dev $(UV) sync --frozen --all-groups --all-extras $(UV_RUN) prisma generate --schema litellm/proxy/schema.prisma +install-rust-python-test-deps: + $(UV) sync --inexact --frozen --no-default-groups --no-install-project + install-helm-unittest: @helm plugin list | grep -qE '^unittest[[:space:]]+0\.8\.2([[:space:]]|$$)' || { \ helm plugin uninstall unittest >/dev/null 2>&1 || true; \ @@ -302,6 +310,19 @@ test-rust-extension: LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust +test-rust-python: install-rust-python-test-deps + @python=$$($(UV_RUN) python -c 'import sys; print(sys.executable)') && \ + site_packages=$$("$$python" -c 'import os, sysconfig; print(os.pathsep.join(dict.fromkeys(sysconfig.get_path(key) for key in ("purelib", "platlib"))))') && \ + PYO3_PYTHON="$$python" \ + PYTHONPATH="$(CURDIR):$$site_packages$${PYTHONPATH:+:$$PYTHONPATH}" \ + LITELLM_LOCAL_MODEL_COST_MAP=True \ + cargo test --manifest-path litellm-rust/Cargo.toml \ + -p litellm-python-interop -p litellm-python-bridge --tests --locked -- --ignored $(if $(TEST_FILTER),"$(TEST_FILTER)") + +lint-rust-python-fixtures: + $(UV) tool run --from ruff==0.15.3 ruff check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests + $(UV) tool run --from ruff==0.15.3 ruff format --check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests + test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5234f2e7f66..18a08bac8bc 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1471,9 +1471,9 @@ dependencies = [ "litellm-python-interop", "pyo3", "pyo3-async-runtimes", + "rstest", "serde", "serde_json", - "strum", "tokio", "tokio-tungstenite", "tracing", @@ -2368,28 +2368,6 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" -[[package]] -name = "strum" -version = "0.26.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" -dependencies = [ - "strum_macros", -] - -[[package]] -name = "strum_macros" -version = "0.26.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.119", -] - [[package]] name = "subtle" version = "2.6.1" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 650d38753e7..24a5da7aa9c 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -54,3 +54,50 @@ function per top-level route, mirroring the core entrypoints. Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust changes. That list is the single source of truth and matches what GitHub Actions runs for changes under `litellm-rust/`. + +### Python-Integrated Tests + +From the repository root, use the same entrypoint as the Rust CI workflow to run +the ignored Cargo tests that need the repository's Python dependencies: + +```bash +make test-rust-python +make test-rust-python TEST_FILTER=component_contract +make test-rust-python TEST_FILTER=retained +make lint-rust-python-fixtures +``` + +The test target covers `litellm-python-interop` (including `component_contract` +and `prepared_call`) and `litellm-python-bridge` (including `ocr_retained`). +`TEST_FILTER` is an optional Rust test-name substring, not a Python fixture or +Cargo test-binary name. An unmatched filter runs zero tests, so check the test +counts. The underlying command is: + +```bash +cargo test --manifest-path litellm-rust/Cargo.toml \ + -p litellm-python-interop -p litellm-python-bridge --tests --locked -- --ignored +``` + +Install uv and the repository's pinned Rust toolchain first. The +`install-rust-python-test-deps` prerequisite runs +`uv sync --inexact --frozen --no-default-groups --no-install-project` on each +invocation. This installs the locked SDK dependencies without building or +installing LiteLLM, pulling in the full dev groups, or pruning existing venv +packages. No wheel is needed for these embedded-Python tests; the existing wheel +lane checks the installed public interface separately + +The target gets the project interpreter from `uv run --no-sync python`, sets +`PYO3_PYTHON` to that executable, and queries its `sysconfig` for both Python and +platform-specific site-packages. It prepends the repository root and those paths +to `PYTHONPATH`, preserving any existing entries, so embedded Python imports the +checkout and its dependencies. Use uv's `UV_PYTHON` and `UV_PROJECT_ENVIRONMENT` +settings to select a different interpreter or project venv. The target also sets +`LITELLM_LOCAL_MODEL_COST_MAP=True` to use the checked-in model cost map + +Fixture checks are separate from the test target so filtered reruns stay focused. +`make lint-rust-python-fixtures` runs pinned Ruff lint and format checks with +`ruff-tests.toml` over both crates' `tests/` directories, including +`crates/python-bridge/tests/fixtures/ocr_retained.py`, without syncing the project +environment. CI runs both targets; its Python path trigger covers `litellm/**` +so changes to OCR, bridge, logging, streaming, and their shared imports rerun the +integrated tests diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index ab9abe4e3ae..5959ecd73c1 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -30,11 +30,11 @@ pyo3.workspace = true pyo3-async-runtimes.workspace = true serde.workspace = true serde_json.workspace = true -strum.workspace = true tokio.workspace = true [dev-dependencies] criterion = "0.8.2" +rstest.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/bindings.rs b/litellm-rust/crates/python-bridge/src/routes/bindings.rs index 529bda63142..9c0912d5b2a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/routes/bindings.rs @@ -1,23 +1,10 @@ use litellm_python_interop::InvocationMode; -use strum::{Display, EnumIter, IntoStaticStr}; -/// A method name paired inseparably with the mode used to drive it. -/// -/// Keeping the two together is what prevents a `("prepare", Await)`-style -/// mismatch. Callers never assemble these by hand; they ask a catalog entry -/// ([`BoundaryMethod`] or [`LoggingMethod`]) to `resolve` one. pub(crate) struct MethodBinding { pub(crate) name: &'static str, pub(crate) mode: InvocationMode, } -/// OCR route boundary methods: the Python operations invoked with -/// [`litellm_python_interop::PreparedCall`] to carry a request through -/// preparation, encoding, transport and finalization. -/// -/// The `asynchronous` flag drives both the method *name* and its invocation -/// mode together, so a sync/async pair cannot desynchronize. -#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumIter, Display, IntoStaticStr)] pub(crate) enum BoundaryMethod { Prepare, Encode, @@ -50,108 +37,3 @@ impl BoundaryMethod { } } } - -/// Bound methods on the Python `Logging` object that the bridge invokes. -/// -/// Sync hooks (`pre_call`, `success_handler`, `failure_handler`) run inline -/// even on async routes; async hooks return a coroutine for the caller's loop -/// to drive. This mirrors the split LiteLLM keeps in Python between -/// `dynamic_success_callbacks` / `dynamic_async_success_callbacks`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumIter, Display, IntoStaticStr)] -pub(crate) enum LoggingMethod { - PreCall, - PostCall, - SuccessHandler, - FailureHandler, - AsyncSuccessHandler, - AsyncFailureHandler, -} - -// `LoggingMethod` is the forward-looking catalog for the retained-callback -// foundation. Nothing in production consumes it yet (the hooks them are -// exercised by the `component_contract` fixtures via a direct `PreparedCall`), -// so `resolve` is only reached from this module's tests for now. -#[cfg_attr(not(test), allow(dead_code))] -impl LoggingMethod { - pub(crate) fn resolve(self) -> MethodBinding { - match self { - Self::PreCall => MethodBinding { - name: "pre_call", - mode: InvocationMode::Direct, - }, - Self::PostCall => MethodBinding { - name: "post_call", - mode: InvocationMode::Direct, - }, - Self::SuccessHandler => MethodBinding { - name: "success_handler", - mode: InvocationMode::Direct, - }, - Self::FailureHandler => MethodBinding { - name: "failure_handler", - mode: InvocationMode::Direct, - }, - Self::AsyncSuccessHandler => MethodBinding { - name: "async_success_handler", - mode: InvocationMode::Await, - }, - Self::AsyncFailureHandler => MethodBinding { - name: "async_failure_handler", - mode: InvocationMode::Await, - }, - } - } -} - -#[cfg(test)] -mod tests { - use strum::IntoEnumIterator; - - use super::*; - - #[test] - fn boundary_methods_pair_name_and_mode_consistently() { - for method in BoundaryMethod::iter() { - let sync = method.resolve(false); - let asynchronous = method.resolve(true); - - assert!(!sync.name.is_empty()); - assert!(!asynchronous.name.is_empty()); - - if method == BoundaryMethod::Encode { - // `encode` has no async variant; it is always a direct call. - assert_eq!(sync.name, asynchronous.name); - assert_eq!(sync.mode, InvocationMode::Direct); - assert_eq!(asynchronous.mode, InvocationMode::Direct); - } else { - assert_eq!(sync.mode, InvocationMode::Direct); - assert_eq!(asynchronous.mode, InvocationMode::Await); - let async_name = asynchronous.name.strip_prefix('a').unwrap(); - assert_eq!(sync.name, async_name); - } - } - } - - #[test] - fn logging_methods_are_direct_unless_async() { - let mut names = Vec::new(); - for method in LoggingMethod::iter() { - let binding = method.resolve(); - assert!(!binding.name.is_empty()); - assert!(!names.contains(&binding.name), "duplicate method name"); - names.push(binding.name); - - let is_async = matches!( - method, - LoggingMethod::AsyncSuccessHandler | LoggingMethod::AsyncFailureHandler - ); - assert_eq!( - binding.mode == InvocationMode::Await, - is_async, - "{} must be {}", - binding.name, - if is_async { "Await" } else { "Direct" } - ); - } - } -} diff --git a/tests/test_litellm/ocr/retained_boundary_fixture.py b/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py similarity index 89% rename from tests/test_litellm/ocr/retained_boundary_fixture.py rename to litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py index c0561e1b88a..6e002970803 100644 --- a/tests/test_litellm/ocr/retained_boundary_fixture.py +++ b/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py @@ -3,7 +3,6 @@ import asyncio import contextvars import gc -import importlib import inspect import json import threading @@ -13,7 +12,7 @@ from copy import deepcopy from datetime import datetime from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import ModuleType -from unittest.mock import Mock, patch +from unittest.mock import patch import litellm from litellm.integrations.custom_logger import CustomLogger @@ -25,7 +24,6 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.rust_bridge.ocr_retained import OCREncoded, OCRRetainedBoundary, OCRRoots native: ModuleType = globals()["native"] -ocr_main = importlib.import_module("litellm.ocr.main") context = contextvars.ContextVar("retained-real-boundary", default="unset") MODEL = "mistral-ocr-latest" RESPONSE = {"pages": [{"index": 0, "markdown": "local OCR"}], "model": MODEL, "usage_info": {"pages_processed": 1}} @@ -167,30 +165,8 @@ def invoke(mode, kwargs, *, boundary_factory=OCRRetainedBoundary): return handler.async_ocr(**kwargs) if mode == "native-sync": return native.ocr_retained(boundary_factory(handler=handler, **kwargs)) - if mode == "native-async": - return native.aocr_retained(boundary_factory(handler=handler, **kwargs)) - prepared = ocr_main._PreparedOCRRequest( - **{ - k: kwargs[k] - for k in ( - "model", - "document", - "api_key", - "api_base", - "custom_llm_provider", - "provider_config", - "optional_params", - "litellm_params", - ) - }, - extra_headers=kwargs["headers"], - effective_timeout=kwargs["timeout"], - litellm_logging_obj=kwargs["logging_obj"], - ) - if mode == "sdk-sync": - return ocr_main._run_rust_ocr(prepared, lambda _: None, load_retained=lambda: native.ocr_retained) - assert mode == "sdk-async" - return ocr_main._run_rust_aocr(prepared, lambda _: None, load_retained=lambda: native.aocr_retained) + assert mode == "native-async" + return native.aocr_retained(boundary_factory(handler=handler, **kwargs)) class RealBoundaryTests(unittest.TestCase): @@ -299,10 +275,10 @@ class RealBoundaryTests(unittest.TestCase): finally: await async_client.close() - def test_differential_callbacks_wire_and_sdk_dispatch(self): + def test_differential_callbacks_wire(self): async def exercise(): baseline = await self.differential("python-sync") - for mode in ("python-async", "native-sync", "native-async", "sdk-sync", "sdk-async"): + for mode in ("python-async", "native-sync", "native-async"): with self.subTest(mode=mode): self.assertEqual(await self.differential(mode), baseline) @@ -416,7 +392,7 @@ class RealBoundaryTests(unittest.TestCase): def test_public_rust_dispatch_wire_fallback_and_escaping_base_exception(self): async def exercise(): for asynchronous in (False, True): - for outcome in ("success", "missing-symbol", "pre-call-abort"): + for outcome in ("success", "missing-symbol", "pre-call-abort", "disabled"): with self.subTest(asynchronous=asynchronous, outcome=outcome): symbol = "aocr_retained" if asynchronous else "ocr_retained" self.assertTrue(inspect.isbuiltin(getattr(native, symbol))) @@ -424,13 +400,10 @@ class RealBoundaryTests(unittest.TestCase): missing_symbol.__dict__.update( (name, value) for name, value in vars(native).items() if name != symbol ) - handler = Mock(wraps=BaseLLMHTTPHandler()) escaped = PreCallAbort("public pre_call must escape unchanged") before = len(self.server.requests) - def mutate(view): - loader.assert_called_once_with() - self.assertEqual(handler.ocr.call_count, int(outcome == "missing-symbol")) + def mutate(view, outcome=outcome, escaped=escaped): view["headers"]["X-Proof"] = "public-in-place" view["complete_input_dict"]["public_mutation"] = True view["complete_input_dict"]["document"]["document_url"] = ( @@ -441,13 +414,10 @@ class RealBoundaryTests(unittest.TestCase): callback = Callback(mutate) kwargs = inputs(self.server, [callback]) - with ( - patch( - "litellm.rust_bridge.get_native_bridge", - return_value=missing_symbol if outcome == "missing-symbol" else native, - ) as loader, - patch.object(ocr_main, "base_llm_http_handler", handler), - ): + with patch( + "litellm.rust_bridge.get_native_bridge", + return_value=missing_symbol if outcome == "missing-symbol" else native, + ) as loader: public_kwargs = { **{ key: kwargs[key] @@ -462,10 +432,10 @@ class RealBoundaryTests(unittest.TestCase): }, "extra_headers": kwargs["headers"], "litellm_logging_obj": kwargs["logging_obj"], - "rust": True, + "rust": outcome != "disabled", } - async def call(): + async def call(asynchronous=asynchronous, public_kwargs=public_kwargs): if asynchronous: return await litellm.aocr(**public_kwargs) return litellm.ocr(**public_kwargs) @@ -478,22 +448,11 @@ class RealBoundaryTests(unittest.TestCase): response = await call() self.assertEqual(response.pages[0].markdown, "local OCR") - loader.assert_called_once_with() - self.check_callbacks([callback]) - self.assertEqual(handler.ocr.call_count, int(outcome == "missing-symbol")) - self.assertEqual(handler.async_ocr.call_count, 0) - prepare = handler._async_prepare_ocr_request if asynchronous else handler._prepare_ocr_request - if outcome == "missing-symbol": - prepare.assert_not_called() - self.assertIs(handler.ocr.call_args.kwargs["logging_obj"], kwargs["logging_obj"]) - self.assertEqual(handler.ocr.call_args.kwargs["aocr"], asynchronous) + if outcome == "disabled": + loader.assert_not_called() else: - prepare.assert_called_once() - self.assertIs(prepare.call_args.kwargs["logging_obj"], kwargs["logging_obj"]) - unused_prepare = ( - handler._prepare_ocr_request if asynchronous else handler._async_prepare_ocr_request - ) - unused_prepare.assert_not_called() + loader.assert_called_once_with() + self.check_callbacks([callback]) self.assertEqual(len(self.server.requests), before + int(outcome != "pre-call-abort")) if outcome != "pre-call-abort": path, headers, body = self.server.requests[-1] @@ -636,7 +595,3 @@ class RealBoundaryTests(unittest.TestCase): self.assertEqual(len(self.server.requests), 1) asyncio.run(exercise()) - - -result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(RealBoundaryTests)) -assert result.wasSuccessful(), "real retained OCR boundary tests failed" diff --git a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs b/litellm-rust/crates/python-bridge/tests/ocr_retained.rs index c5e0408ea9b..6f9430dddf3 100644 --- a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs +++ b/litellm-rust/crates/python-bridge/tests/ocr_retained.rs @@ -1,19 +1,55 @@ use pyo3::prelude::*; use pyo3::types::PyDict; +use rstest::rstest; -#[test] +#[rstest] +#[case::differential_callbacks_wire("differential_callbacks_wire")] +#[case::negative_control_copied_caller_document("negative_control_copied_caller_document")] +#[case::negative_control_rebound_logging_body("negative_control_rebound_logging_body")] +#[case::negative_control_rebound_logging_headers("negative_control_rebound_logging_headers")] +#[case::differential_callback_retained_mutation_after_post_received( + "differential_callback_retained_mutation_after_post_received" +)] +#[case::public_rust_dispatch_wire_fallback_and_escaping_base_exception( + "public_rust_dispatch_wire_fallback_and_escaping_base_exception" +)] +#[case::collection_after_success_encoding_failure_and_http_error( + "collection_after_success_encoding_failure_and_http_error" +)] +#[case::callback_retained_graph_remains_usable_then_collects( + "callback_retained_graph_remains_usable_then_collects" +)] +#[case::collection_after_cancellation_during_blocked_transport( + "collection_after_cancellation_during_blocked_transport" +)] #[ignore = "requires repo Python"] -fn retained_real_production_boundary_differential_and_lifecycle() -> PyResult<()> { +fn retained_real_production_boundary_differential_and_lifecycle( + #[case] scenario: &str, +) -> PyResult<()> { Python::initialize(); Python::attach(|py| { let module = pyo3::wrap_pymodule!(_native::_native)(py).into_bound(py); let globals = PyDict::new(py); globals.set_item("native", module)?; - let fixture = std::ffi::CString::new(include_str!(concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../tests/test_litellm/ocr/retained_boundary_fixture.py" - )))?; - py.run(&fixture, Some(&globals), Some(&globals)) + let builtins = py.import("builtins")?; + let code = builtins.call_method1( + "compile", + ( + include_str!("fixtures/ocr_retained.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/ocr_retained.py" + ), + "exec", + ), + )?; + builtins.call_method1("exec", (code, &globals))?; + globals + .get_item("RealBoundaryTests")? + .unwrap() + .call1((format!("test_{scenario}"),))? + .call_method0("debug")?; + Ok(()) }) } diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index ed60f8a22a8..61cb1b6b44b 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs @@ -1,5 +1,3 @@ -use std::ffi::CString; - use pyo3::prelude::*; use pyo3::types::PyDict; use rstest::{fixture, rstest}; @@ -9,6 +7,18 @@ mod callback_owner; struct InitializedPython; +fn run_fixture( + py: Python<'_>, + globals: &Bound<'_, PyDict>, + source: &str, + filename: &str, +) -> PyResult<()> { + let builtins = py.import("builtins")?; + let code = builtins.call_method1("compile", (source, filename, "exec"))?; + builtins.call_method1("exec", (code, globals))?; + Ok(()) +} + #[fixture] #[once] fn initialized_python() -> InitializedPython { @@ -27,8 +37,16 @@ fn scenario_scope(initialized_python: &InitializedPython) -> Py { Py::new(py, callback_owner::OwnerFactory::default()).unwrap(), ) .unwrap(); - let source = CString::new(include_str!("fixtures/callback_lifecycle.py")).unwrap(); - py.run(&source, Some(&globals), None).unwrap(); + run_fixture( + py, + &globals, + include_str!("fixtures/callback_lifecycle.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/callback_lifecycle.py" + ), + ) + .unwrap(); globals.unbind() }) } @@ -79,8 +97,15 @@ fn component_contract( ) -> PyResult<()> { Python::attach(|py| { let globals = scenario_scope.bind(py); - let source = CString::new(include_str!("fixtures/callback_components.py")).unwrap(); - py.run(&source, Some(globals), None)?; + run_fixture( + py, + globals, + include_str!("fixtures/callback_components.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/callback_components.py" + ), + )?; globals.get_item("run_scenario")?.unwrap().call1(( scenario, retained, diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py index 38324fb6a66..5c42da2b634 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py @@ -1,5 +1,6 @@ import asyncio from datetime import datetime +from unittest import TestCase from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging @@ -30,8 +31,7 @@ async def real_async_logging(owners): class Retain(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): - observations.append(("retained", kwargs, result)) - assert asyncio.current_task() is task + observations.append(("retained", kwargs, result, asyncio.current_task())) return kwargs, result class MutateThenFail(CustomLogger): @@ -49,8 +49,7 @@ async def real_async_logging(owners): class Observe(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): - observations.append(("observe", kwargs, result)) - assert asyncio.current_task() is task + observations.append(("observe", kwargs, result, asyncio.current_task())) return kwargs, result async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -62,6 +61,7 @@ async def real_async_logging(owners): await owner.invoke() owner.close() assert [entry[0] for entry in observations] == ["retained", "replace", "observe", "success"] + assert observations[0][3] is task and observations[2][3] is task assert observations[0][2] is result and observations[1][2] is result assert observations[2][2] is replacement and observations[3][2] is replacement assert observations[0][1]["retained_shared"] is shared and shared["changed"] @@ -199,11 +199,8 @@ async def real_stream_cancellation(owners): pull.close() await entered.wait() task.cancel() - try: + with TestCase().assertRaises(asyncio.CancelledError): await task - assert False, "pull ignored cancellation" - except asyncio.CancelledError: - pass close = owners.prepare(wrapper.aclose, (), awaited=True) await close.invoke() close.close() diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py index 0b4bbe2f1fc..6a7da52b056 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py @@ -5,6 +5,7 @@ import gc import json import threading import weakref +from unittest import TestCase async def checkpoint(): @@ -76,7 +77,7 @@ async def awaitable_kinds(owners): "async": coroutine, "sync_coroutine": lambda value, *, alias: coroutine(value, alias=alias), "custom": lambda value, *, alias: CustomAwaitable(), - "future": lambda value, *, alias: future, + "future": lambda value, *, alias, future=future: future, }[kind] owner = owners.prepare(callback, (payload,), {"alias": payload}, awaited=True) pending = owner.invoke() @@ -86,11 +87,8 @@ async def awaitable_kinds(owners): assert len(calls) == before + (kind != "future") owner = owners.prepare(lambda: payload, (), awaited=True) - try: + with TestCase().assertRaises(TypeError): await owner.invoke() - assert False, "non-awaitable result accepted" - except TypeError: - pass owner.close() inner = coroutine(payload, alias=payload) @@ -165,24 +163,25 @@ async def exceptions(owners): cause = ValueError("cause") payload = {} - async def callback(): + async def callback(payload=payload, error=error, cause=cause): payload["changed"] = True raise error from cause owner = owners.prepare(callback, (), awaited=True) + caught_error = None try: await owner.invoke() - assert False, "exception lost" except BaseException as caught: - assert caught is error and caught.__cause__ is cause - frames = [] - tb = caught.__traceback__ - while tb: - frames.append(tb.tb_frame.f_code.co_name) - tb = tb.tb_next - assert "callback" in frames + caught_error = caught finally: owner.close() + assert caught_error is error and caught_error.__cause__ is cause + frames = [] + tb = caught_error.__traceback__ + while tb: + frames.append(tb.tb_frame.f_code.co_name) + tb = tb.tb_next + assert "callback" in frames assert payload["changed"] is True @@ -201,14 +200,15 @@ async def exception_ownership(owners): value_ref, callback_ref = weakref.ref(value), weakref.ref(callback) owner = owners.prepare(callback, (value,), awaited=True) del value, callback + caught_error = None try: await owner.invoke() - assert False, "expected callback failure" except RuntimeError as caught: - assert caught is error + caught_error = caught + assert caught_error is error owner.close() assert value_ref().changed and callback_ref() is not None - del error + del error, caught_error gc.collect() assert value_ref() is None and callback_ref() is None @@ -408,7 +408,7 @@ async def stream_lifecycle(owners): item = {"nested": nested} closed = [] - async def source(): + async def source(item=item, nested=nested, terminal=terminal, closed=closed): try: yield item nested["usage"] = 12 @@ -430,11 +430,8 @@ async def stream_lifecycle(owners): await close.invoke() close.close() else: - try: + with TestCase().assertRaises(ValueError if terminal == "failure" else StopAsyncIteration): await pull.invoke() - assert False, "expected stream termination" - except (StopAsyncIteration, ValueError) as error: - assert isinstance(error, ValueError) == (terminal == "failure") pull.close() assert closed == [True] assert retained.invoke() is nested @@ -448,7 +445,7 @@ async def sync_stream_lifecycle(owners): value = {"nested": {"usage": 0}} closed = [] - def source(): + def source(value=value, terminal=terminal, closed=closed): try: yield value value["nested"]["usage"] = 12 @@ -467,11 +464,8 @@ async def sync_stream_lifecycle(owners): close.invoke() close.close() else: - try: + with TestCase().assertRaises(ValueError if terminal == "failure" else StopIteration): pull.invoke() - assert False, "stream did not terminate" - except (StopIteration, ValueError) as error: - assert isinstance(error, ValueError) == (terminal == "failure") pull.close() assert saved.invoke() is value assert value["nested"]["usage"] == (0 if terminal == "close" else 12) @@ -484,7 +478,7 @@ async def repeated_ownership(owners): for batch in range(8): gate = asyncio.Event() - async def work(value): + async def work(value, gate=gate): await gate.wait() return None diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index c34833221cc..bd3a8d186c9 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -816,21 +816,14 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) -def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" - bridge = RecordingBridge() - litellm.rust(False) - rust_bridge._OCR.override(bridge) - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] +@pytest.mark.parametrize("rust_enabled", [False, True], ids=["disabled", "unavailable"]) +def test_ocr_uses_python_when_native_execution_is_unavailable(monkeypatch, rust_enabled): + def load_native(): + assert rust_enabled, "disabled OCR must not load the native module" + return None - -def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): - """Rust enabled but no bridge available (no injected impl, no compiled wheel): - ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI + monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", load_native) + litellm.rust(rust_enabled) captured = {} diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index a7f50a82a99..ee6c5069910 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -10,14 +10,18 @@ import sys import tempfile import threading import zipfile +from collections import Counter from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from queue import SimpleQueue from socket import socket as Socket +from types import FrameType from typing import Final REQUEST_STARTED: Final = threading.Event() REQUEST_CANCELLED: Final = threading.Event() +PUBLIC_OCR_REQUESTS: Final[SimpleQueue[str]] = SimpleQueue() ANTHROPIC_RESPONSE: Final = ( b'{"id":"msg_native","type":"message","role":"assistant",' @@ -35,6 +39,10 @@ class NativeRouteHandler(BaseHTTPRequestHandler): body: Final = json.loads(self.rfile.read(content_length)) route: Final = self.headers.get("x-test-route") outcome: Final = self.headers.get("x-test-outcome") + public_case: Final = self.headers.get("x-test-public-case") + if public_case is not None: + PUBLIC_OCR_REQUESTS.put(public_case) + assert self.headers.get("x-test-callback") == public_case assert_native_request(route, outcome, self.path, self.headers, body) if outcome == "hang": REQUEST_STARTED.set() @@ -227,12 +235,7 @@ async def exercise_async(native: object, api_base: str) -> None: async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather( - *( - native.amessages(**route_kwargs("messages", api_base, "success")) - for _ in range(32) - ) - ), + asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: @@ -249,6 +252,92 @@ def exercise_routes(native_path: Path, api_base: str) -> object: return native +def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.base_llm.ocr.transformation import OCRResponse + from litellm.rust_bridge import get_native_bridge + from litellm.rust_bridge.ocr_retained import OCRRetainedBoundary + + assert case in {"ocr", "aocr"} + assert Path(litellm.__file__).resolve().is_relative_to(install_root.resolve()) + native: Final = get_native_bridge() + assert native is not None + assert Path(native.__file__).resolve().is_relative_to(install_root.resolve()) + retained: Final = getattr(native, f"{case}_retained") + observed: Final[Counter[str]] = Counter() + + def observe(frame: FrameType, event: str, arg: object) -> None: + if event == "c_call" and arg is retained: + observed["retained"] += 1 + if event == "call" and frame.f_code is OCRRetainedBoundary.encode.__code__: + observed["encode"] += 1 + + class MutatingLogger(CustomLogger): + calls = 0 + + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + self.calls += 1 + additional_args: Final = kwargs["additional_args"] + assert isinstance(additional_args, dict) + headers: Final = additional_args["headers"] + body: Final = additional_args["complete_input_dict"] + assert isinstance(headers, dict) and isinstance(body, dict) + assert body["model"] == "mistral-ocr-latest" + assert body["include_image_base64"] is False + assert headers["x-test-callback"] == "before-callback" + headers["x-test-callback"] = case + body["include_image_base64"] = True + + callback: Final = MutatingLogger() + kwargs: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "api_base": api_base, + "api_key": "sk-native", + "extra_headers": { + "x-test-route": "ocr", + "x-test-outcome": "success", + "x-test-public-case": case, + "x-test-callback": "before-callback", + }, + "include_image_base64": False, + "callbacks": [callback], + "rust": True, + "timeout": 3.0, + "num_retries": 0, + } + previous_profile: Final = sys.getprofile() + sys.setprofile(observe) + try: + response: Final = asyncio.run(litellm.aocr(**kwargs)) if case == "aocr" else litellm.ocr(**kwargs) + finally: + sys.setprofile(previous_profile) + assert isinstance(response, OCRResponse) + assert_success("ocr", response.model_dump()) + assert response.model == "mistral-ocr-latest" + assert callback.calls == 1, callback.calls + assert observed == {"retained": 1, "encode": 1}, observed + return 0 + + +def verify_public_ocr(wheel: Path, wheel_root: Path, api_base: str) -> None: + install_root: Final = wheel_root / "sdk-venv" + python: Final = install_root / "bin" / "python" + subprocess.run(("uv", "venv", "--python", sys.executable, str(install_root)), check=True) + subprocess.run(("uv", "pip", "install", "--python", str(python), str(wheel.resolve())), check=True) + for case in ("ocr", "aocr"): + subprocess.run( + (str(python), "-I", str(Path(__file__).resolve()), "public-ocr", str(install_root), api_base, case), + cwd=install_root, + env=os.environ | {"LITELLM_LOCAL_MODEL_COST_MAP": "True", "NO_PROXY": "127.0.0.1", "no_proxy": "127.0.0.1"}, + check=True, + timeout=60, + ) + assert PUBLIC_OCR_REQUESTS.get_nowait() == case + assert PUBLIC_OCR_REQUESTS.empty(), f"{case} sent more than one upstream request" + + def exercise_signal(native: object, api_base: str) -> int: try: native.messages( @@ -320,6 +409,7 @@ def verify_wheel(wheel: Path) -> int: api_base: Final = f"http://127.0.0.1:{server.server_address[1]}" try: verify_sigint(native_path, api_base) + verify_public_ocr(wheel, wheel_root, api_base) finally: server.shutdown() server.server_close() @@ -333,6 +423,8 @@ def main() -> int: if len(sys.argv) == 4 and sys.argv[1] == "child": native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3]) return exercise_signal(native, sys.argv[3]) + if len(sys.argv) == 5 and sys.argv[1] == "public-ocr": + return exercise_public_ocr(Path(sys.argv[2]), sys.argv[3], sys.argv[4]) sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") return 2