scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

wip

Browse source
This commit is contained in:
Yujong Lee 2026-09-06 18:00:33 -07:00
parent f1553ae9c3
commit c6cd873266
13 changed files with 300 additions and 286 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
.github/workflows/test-rust.yml vendored
View file

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

23
Makefile
View file

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

24
litellm-rust/Cargo.lock generated
View file

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

47
litellm-rust/README.md
View file

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

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

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

118
litellm-rust/crates/python-bridge/src/routes/bindings.rs
View file

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

79
tests/test_litellm/ocr/retained_boundary_fixture.py → litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py vendored
View file

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

50
litellm-rust/crates/python-bridge/tests/ocr_retained.rs
View file

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

37
litellm-rust/crates/python-interop/tests/callback_lifecycle.rs
View file

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

13
litellm-rust/crates/python-interop/tests/fixtures/callback_components.py vendored
View file

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

50
litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py vendored
View file

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

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

@ -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 = {}

104
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

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