From dff8131f1450f6b2c11ca09d5e6e37d863ec85a4 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Sun, 11 Oct 2026 14:10:16 -0700 Subject: [PATCH] refactor(rust_bridge): route Python state reads through host modules (#45689) * refactor(rust_bridge): group the Python modules Rust imports by name under rust_bridge/host Rust reads Python-owned state (legacy callback registries, model capabilities) through modules it imports by path. Put them in one folder with an AGENTS.md that states the contract, and update the module path constants on the Rust side. Co-Authored-By: Claude Fable 5.1 * refactor(rust_bridge): route Python state reads through host accessors * docs(rust_bridge): point test_support at the moved host contract test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust_bridge): move secret manager, transport and timeout accessors under host Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(greptile): register rust_bridge/host AGENTS.md in files.json Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Claude Fable 5.1 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .greptile/files.json | 7 +++++ .../callbacks-legacy-python/src/python.rs | 4 +-- .../src/test_support.rs | 6 ++-- .../crates/python-bridge/src/cache/AGENTS.md | 2 +- .../python-bridge/src/cache/selection.rs | 7 ++++- .../crates/python-bridge/src/marshal.rs | 4 ++- .../python-bridge/src/python_settings.rs | 30 +++++++++++------- .../python-bridge/src/routes/messages/host.rs | 2 +- .../python-bridge/src/routes/ocr/project.rs | 4 +-- .../python-bridge/src/secrets/callback.rs | 31 ++++++++++--------- .../python-bridge/src/secrets/python.rs | 8 +++-- .../python-bridge/src/secrets/runtime.rs | 4 ++- .../crates/python-bridge/src/secrets/vault.rs | 10 +++--- .../bedrock/audio_transcription/__init__.py | 2 +- litellm/proxy/proxy_server.py | 2 +- litellm/rust_bridge/AGENTS.md | 2 ++ litellm/rust_bridge/host/AGENTS.md | 11 +++++++ litellm/rust_bridge/host/__init__.py | 1 + litellm/rust_bridge/host/cache.py | 12 +++++++ .../{ => host}/callbacks_legacy_python.py | 0 litellm/rust_bridge/host/diagnostics.py | 14 +++++++++ litellm/rust_bridge/host/http_transport.py | 7 +++++ .../{ => host}/model_capabilities.py | 0 .../rust_bridge/{ => host}/secret_manager.py | 6 ++++ litellm/rust_bridge/{ => host}/settings.py | 0 litellm/rust_bridge/{ => host}/timeouts.py | 0 litellm/rust_bridge/responses/websocket.py | 2 +- tests/test_litellm_rust/ocr/test_requests.py | 3 +- .../proxy/proxy_server/test_proxy_config.py | 2 +- tests/unit/rust_bridge/host/__init__.py | 0 tests/unit/rust_bridge/host/test_cache.py | 21 +++++++++++++ .../test_callbacks_legacy_python.py | 6 ++-- .../{ => host}/test_model_capabilities.py | 3 +- .../{ => host}/test_secret_manager.py | 2 +- .../rust_bridge/{ => host}/test_settings.py | 2 +- .../unit/rust_bridge/messages/test_secrets.py | 2 +- tests/unit/rust_bridge/ocr/test_secrets.py | 2 +- 37 files changed, 163 insertions(+), 58 deletions(-) create mode 100644 litellm/rust_bridge/host/AGENTS.md create mode 100644 litellm/rust_bridge/host/__init__.py create mode 100644 litellm/rust_bridge/host/cache.py rename litellm/rust_bridge/{ => host}/callbacks_legacy_python.py (100%) create mode 100644 litellm/rust_bridge/host/diagnostics.py create mode 100644 litellm/rust_bridge/host/http_transport.py rename litellm/rust_bridge/{ => host}/model_capabilities.py (100%) rename litellm/rust_bridge/{ => host}/secret_manager.py (97%) rename litellm/rust_bridge/{ => host}/settings.py (100%) rename litellm/rust_bridge/{ => host}/timeouts.py (100%) create mode 100644 tests/unit/rust_bridge/host/__init__.py create mode 100644 tests/unit/rust_bridge/host/test_cache.py rename tests/unit/rust_bridge/{ => host}/test_callbacks_legacy_python.py (96%) rename tests/unit/rust_bridge/{ => host}/test_model_capabilities.py (96%) rename tests/unit/rust_bridge/{ => host}/test_secret_manager.py (99%) rename tests/unit/rust_bridge/{ => host}/test_settings.py (99%) diff --git a/.greptile/files.json b/.greptile/files.json index 99ef105499e..a1af021999d 100644 --- a/.greptile/files.json +++ b/.greptile/files.json @@ -28,6 +28,13 @@ "litellm/rust_bridge/**" ] }, + { + "path": "litellm/rust_bridge/host/AGENTS.md", + "description": "Conventions for code under litellm/rust_bridge/host/", + "scope": [ + "litellm/rust_bridge/host/**" + ] + }, { "path": "litellm/tracing/AGENTS.md", "description": "Conventions for code under litellm/tracing/", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 47331f369f5..6f1ee9dfd78 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; use strum::{IntoStaticStr, VariantArray}; -const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python"; +const MODULE: &str = "litellm.rust_bridge.host.callbacks_legacy_python"; /// Every litellm Python internal the native call still borrows, grouped by the subsystem it /// belongs to. Rust drives the call; these exist only so behaviour that Python owns today @@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python"; /// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a /// user's own callback is not borrowing and does not belong here. /// -/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and +/// `litellm/rust_bridge/host/callbacks_legacy_python.py` is the only Python module behind it, and /// `python_contract.json` pins each function's parameters on both sides. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum LegacyPython { diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index a7b28f1b0ca..093ce60e510 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -6,7 +6,7 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, PublicCall}; /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. -/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python +/// `tests/unit/rust_bridge/host/test_callbacks_legacy_python.py` pins this file to the Python /// signatures, and [`namespace`] binds every fake call against it. pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); @@ -23,10 +23,10 @@ import sys import traceback import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.callbacks_legacy_python'): sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +legacy = sys.modules['litellm.rust_bridge.host.callbacks_legacy_python'] CONTRACT = json.loads(python_contract) diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index 3f1097f4bc2..41c82187373 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -2,6 +2,6 @@ This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` owns global cache selection, route admission and inference protocol composition for both adapters +`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters `python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index e6e9f8d2d4f..b2dc67fe965 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -9,6 +9,8 @@ use litellm_host::{ use pyo3::{prelude::*, types::PyDict}; use std::sync::Arc; +const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; + pub(crate) struct Cached

(std::marker::PhantomData

); impl Protocol for Cached

{ @@ -56,7 +58,10 @@ fn selected_cache<'py>( kwargs: &Bound<'py, PyDict>, call_type: &str, ) -> PyResult>> { - let configured = py.import("litellm")?.getattr("cache")?; + let configured = py + .import(CACHE_HOST_MODULE)? + .getattr("configured_cache")? + .call0()?; if configured.is_none() || kwargs .get_item("caching")? diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 4a7efeffdd2..038fa8c9630 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -9,6 +9,8 @@ use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; +const TIMEOUTS_MODULE: &str = "litellm.rust_bridge.host.timeouts"; + /// The keyword arguments every value route shares, validated at the Python boundary. pub(crate) struct RouteOptions { pub(crate) model: String, @@ -47,7 +49,7 @@ pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option } pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py) -> PyResult> { - py.import("litellm.rust_bridge.timeouts")? + py.import(TIMEOUTS_MODULE)? .getattr("timeout_to_seconds")? .call1((timeout,))? .extract() diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 88900544779..eacec6859b5 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -3,7 +3,7 @@ use strum::IntoStaticStr; use crate::coercion::{FieldSpec, ProjectionError}; -const MODULE: &str = "litellm.rust_bridge.settings"; +const MODULE: &str = "litellm.rust_bridge.host.settings"; #[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)] pub(crate) enum PythonSettings { @@ -220,22 +220,26 @@ import sys import types previous_modules = { name: sys.modules[name] - for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.settings') + for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.settings') if name in sys.modules } litellm = types.ModuleType('litellm') litellm.__path__ = [] rust_bridge = types.ModuleType('litellm.rust_bridge') rust_bridge.__path__ = [] -settings = types.ModuleType('litellm.rust_bridge.settings') +host = types.ModuleType('litellm.rust_bridge.host') +host.__path__ = [] +settings = types.ModuleType('litellm.rust_bridge.host.settings') def http_settings(): raise ModuleNotFoundError('No module named certifi', name='certifi') settings.http_settings = http_settings litellm.rust_bridge = rust_bridge -rust_bridge.settings = settings +rust_bridge.host = host +host.settings = settings sys.modules['litellm'] = litellm sys.modules['litellm.rust_bridge'] = rust_bridge -sys.modules['litellm.rust_bridge.settings'] = settings +sys.modules['litellm.rust_bridge.host'] = host +sys.modules['litellm.rust_bridge.host.settings'] = settings ", Some(&locals), Some(&locals), @@ -257,7 +261,7 @@ sys.modules['litellm.rust_bridge.settings'] = settings ); py.run( c" -for name in ('litellm.rust_bridge.settings', 'litellm.rust_bridge', 'litellm'): +for name in ('litellm.rust_bridge.host.settings', 'litellm.rust_bridge.host', 'litellm.rust_bridge', 'litellm'): sys.modules.pop(name, None) sys.modules.update(previous_modules) ", @@ -279,22 +283,26 @@ import sys import types previous_modules = { name: sys.modules[name] - for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.settings') + for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.settings') if name in sys.modules } litellm = types.ModuleType('litellm') litellm.__path__ = [] rust_bridge = types.ModuleType('litellm.rust_bridge') rust_bridge.__path__ = [] -settings = types.ModuleType('litellm.rust_bridge.settings') +host = types.ModuleType('litellm.rust_bridge.host') +host.__path__ = [] +settings = types.ModuleType('litellm.rust_bridge.host.settings') def http_settings(): raise ImportError('cannot import name setting') settings.http_settings = http_settings litellm.rust_bridge = rust_bridge -rust_bridge.settings = settings +rust_bridge.host = host +host.settings = settings sys.modules['litellm'] = litellm sys.modules['litellm.rust_bridge'] = rust_bridge -sys.modules['litellm.rust_bridge.settings'] = settings +sys.modules['litellm.rust_bridge.host'] = host +sys.modules['litellm.rust_bridge.host.settings'] = settings ", Some(&locals), Some(&locals), @@ -308,7 +316,7 @@ sys.modules['litellm.rust_bridge.settings'] = settings assert_eq!(error.to_string(), "ImportError: cannot import name setting"); py.run( c" -for name in ('litellm.rust_bridge.settings', 'litellm.rust_bridge', 'litellm'): +for name in ('litellm.rust_bridge.host.settings', 'litellm.rust_bridge.host', 'litellm.rust_bridge', 'litellm'): sys.modules.pop(name, None) sys.modules.update(previous_modules) ", diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index adeba397b26..008f1bde763 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -214,7 +214,7 @@ impl MessagesPythonHost { ) -> PyResult { let module = py.import(ROUTE_HOST_MODULE)?; let capabilities: MessagesModelCapabilities = from_py( - &py.import("litellm.rust_bridge.model_capabilities")? + &py.import("litellm.rust_bridge.host.model_capabilities")? .getattr("anthropic_model_capabilities")? .call1((model, custom_llm_provider))?, )?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 93f04966f08..fe2f83bff46 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -368,11 +368,11 @@ mod tests { c" import sys import types -timeouts = types.ModuleType('litellm.rust_bridge.timeouts') +timeouts = types.ModuleType('litellm.rust_bridge.host.timeouts') timeouts.timeout_to_seconds = lambda timeout: None if timeout is None else float(timeout) sys.modules.setdefault('litellm', types.ModuleType('litellm')) sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) -sys.modules['litellm.rust_bridge.timeouts'] = timeouts +sys.modules['litellm.rust_bridge.host.timeouts'] = timeouts ", ); } diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 6b61ca22dcd..71a06bfdc39 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -14,8 +14,7 @@ use pyo3::{ use super::error::{external_error, read_error}; const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler"; -const ENVIRONMENT_FALLBACK_LOG: &str = - "Defaulting to os.environ value for key=%s. An exception occurred - %s.\n\n%s"; +const DIAGNOSTICS_MODULE: &str = "litellm.rust_bridge.host.diagnostics"; /// A secret manager whose reads execute in Python: a custom manager, a legacy compatible /// client, or a manually assigned SDK client. @@ -112,16 +111,9 @@ impl ExternalSecretManager for PythonSecretManager { } fn log_environment_fallback(py: Python<'_>, name: &str, error: &PyErr) -> PyResult<()> { - let traceback = py - .import("traceback")? - .call_method1("format_exception", (error.value(py),))?; - let traceback = "".into_pyobject(py)?.call_method1("join", (traceback,))?; - py.import("litellm._logging")? - .getattr("verbose_logger")? - .call_method1( - "error", - (ENVIRONMENT_FALLBACK_LOG, name, error.value(py), traceback), - )?; + py.import(DIAGNOSTICS_MODULE)? + .getattr("log_environment_fallback")? + .call1((name, error.value(py)))?; Ok(()) } @@ -139,7 +131,7 @@ mod tests { use litellm_host_python::PythonContext; - use super::{HANDLER_MODULE, PythonSecretManager, python_name}; + use super::{DIAGNOSTICS_MODULE, HANDLER_MODULE, PythonSecretManager, python_name}; use crate::secrets::python_error; /// `sys.modules` is interpreter-global, so tests that install or rely on the handler module @@ -250,13 +242,24 @@ class Logger: self.calls.append(args) logging = types.ModuleType('litellm._logging') logging.verbose_logger = Logger() -sys.modules.setdefault('litellm', types.ModuleType('litellm')) +for module in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): + sys.modules.setdefault(module, types.ModuleType(module)) sys.modules.setdefault('litellm._logging', logging) ", None, None, ) .unwrap(); + pyo3::types::PyModule::from_code( + py, + &std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/host/diagnostics.py" + )) + .unwrap(), + c"diagnostics.py", + &std::ffi::CString::new(DIAGNOSTICS_MODULE).unwrap(), + ) + .unwrap(); py.import("litellm._logging") .unwrap() .getattr("verbose_logger") diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index fb26b1e80fc..2f040bff773 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -7,7 +7,9 @@ use pyo3::prelude::*; use super::error::external_error; -/// Reads each secret through Python's `get_secret_str`, so the configured manager, the key +const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; + +/// Reads each secret through the host module's `read_secret_str`, so the configured manager, the key /// management settings and the environment fallback behave exactly as they do in Python. pub(super) struct PythonSecrets { get_secret_str: Arc>, @@ -17,8 +19,8 @@ pub(super) struct PythonSecrets { impl PythonSecrets { pub(super) fn new(py: Python<'_>) -> PyResult { Ok(Self::reading_with( - py.import("litellm.secret_managers.main")? - .getattr("get_secret_str")? + py.import(SECRET_MANAGER_MODULE)? + .getattr("read_secret_str")? .unbind(), PythonContext::capture(py)?, )) diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index ab8fea3697d..72ba91ded16 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -15,6 +15,8 @@ use pyo3::{ use crate::execution::{run_async_value, run_sync_value}; +const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; + #[derive(Clone, PartialEq)] struct Configuration { system: KeyManagementSystem, @@ -101,7 +103,7 @@ impl NativeSecretManager { return Ok(Some(native)); } let config = py - .import("litellm.rust_bridge.secret_manager")? + .import(SECRET_MANAGER_MODULE)? .getattr("native_secret_manager_config")? .call1((client,))?; if config.is_none() { diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index de0fa95bcef..46b29839fe4 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -7,6 +7,8 @@ use pyo3::prelude::*; use super::mutation::{error_value, http_message, json_value}; +const HTTP_TRANSPORT_MODULE: &str = "litellm.rust_bridge.host.http_transport"; + pub(super) fn failure_value( py: Python<'_>, failure: Failure, @@ -82,7 +84,7 @@ fn cause_message(py: Python<'_>, kind: FailureKind, context: &ErrorContext) -> P kwargs.set_item("model", "default-model-name")?; kwargs.set_item("llm_provider", "litellm-httpx-handler")?; kwargs.set_item("headers", pyo3::types::PyDict::new(py))?; - py.import("litellm")? + py.import("litellm.exceptions")? .getattr("Timeout")? .call((), Some(&kwargs))? .str()? @@ -173,9 +175,9 @@ impl ErrorContext { Ok(Self { timeout: timeout.map(|value| value.str()?.extract()).transpose()?, aiohttp: py - .import("litellm.llms.custom_httpx.http_handler")? - .getattr("AsyncHTTPHandler")? - .call_method0("_should_use_aiohttp_transport")? + .import(HTTP_TRANSPORT_MODULE)? + .getattr("uses_aiohttp_transport")? + .call0()? .extract()?, }) } diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index bf3c5477886..7f7a138398b 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -6,8 +6,8 @@ import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge.host.timeouts import timeout_to_seconds from litellm.rust_bridge.public_call import NativeCall -from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.rust_bridge.transcription.native import ( NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 089647c6297..1c7cf54ab4d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7434,7 +7434,7 @@ class ProxyConfig: else: raise ValueError("Invalid Key Management System selected") - from litellm.rust_bridge.secret_manager import capture_secret_manager + from litellm.rust_bridge.host.secret_manager import capture_secret_manager if litellm.secret_manager_client is not previous_client: capture_secret_manager(litellm.secret_manager_client, key_management_system) diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index 0c313b097cf..c3fb11f0bfb 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -2,6 +2,8 @@ This package owns native rollout and fallback selection, Python public API compatibility, settings projection and the Python bindings supplied to the Rust bridge. Rust core owns provider execution. `litellm-host-python` owns CPython runtime mechanics. `callbacks-legacy-python` owns legacy callback sharing and dispatch policy +`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. `streams.py` supplies the product binding and public stream wrappers through the bridge rather than let the generic Rust host import this package by name. Keep one driver implementing `start`, `resume_value`, `resume_error` 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/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md new file mode 100644 index 00000000000..44d2909a834 --- /dev/null +++ b/litellm/rust_bridge/host/AGENTS.md @@ -0,0 +1,11 @@ +# Rust host modules + +Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name + +- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds) +- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here +- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers +- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green +- Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals +- Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays +- A `*/route_host.py` module is not this: it projects one route's kwargs and maps that route's failures. Shared questions belong here diff --git a/litellm/rust_bridge/host/__init__.py b/litellm/rust_bridge/host/__init__.py new file mode 100644 index 00000000000..5649fa8a1a7 --- /dev/null +++ b/litellm/rust_bridge/host/__init__.py @@ -0,0 +1 @@ +"""Python modules the Rust bridge imports by name. See AGENTS.md.""" diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py new file mode 100644 index 00000000000..cac4b75677f --- /dev/null +++ b/litellm/rust_bridge/host/cache.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from litellm.caching.caching import Cache + + +def configured_cache() -> Cache | None: + import litellm + + return litellm.cache diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py similarity index 100% rename from litellm/rust_bridge/callbacks_legacy_python.py rename to litellm/rust_bridge/host/callbacks_legacy_python.py diff --git a/litellm/rust_bridge/host/diagnostics.py b/litellm/rust_bridge/host/diagnostics.py new file mode 100644 index 00000000000..40c59de075f --- /dev/null +++ b/litellm/rust_bridge/host/diagnostics.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import traceback + + +def log_environment_fallback(name: str, error: Exception) -> None: + from litellm._logging import verbose_logger + + verbose_logger.error( + "Defaulting to os.environ value for key=%s. An exception occurred - %s.\n\n%s", + name, + error, + "".join(traceback.format_exception(error)), + ) diff --git a/litellm/rust_bridge/host/http_transport.py b/litellm/rust_bridge/host/http_transport.py new file mode 100644 index 00000000000..1bf0f890924 --- /dev/null +++ b/litellm/rust_bridge/host/http_transport.py @@ -0,0 +1,7 @@ +from __future__ import annotations + + +def uses_aiohttp_transport() -> bool: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + return AsyncHTTPHandler.should_use_aiohttp_transport() diff --git a/litellm/rust_bridge/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py similarity index 100% rename from litellm/rust_bridge/model_capabilities.py rename to litellm/rust_bridge/host/model_capabilities.py diff --git a/litellm/rust_bridge/secret_manager.py b/litellm/rust_bridge/host/secret_manager.py similarity index 97% rename from litellm/rust_bridge/secret_manager.py rename to litellm/rust_bridge/host/secret_manager.py index 548f65cfc26..4fb3350165c 100644 --- a/litellm/rust_bridge/secret_manager.py +++ b/litellm/rust_bridge/host/secret_manager.py @@ -157,3 +157,9 @@ def native_secret_manager_config(client: object) -> NativeSecretManagerConfig | ): return _capture(client, adapter) return None + + +def read_secret_str(name: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(name) diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/host/settings.py similarity index 100% rename from litellm/rust_bridge/settings.py rename to litellm/rust_bridge/host/settings.py diff --git a/litellm/rust_bridge/timeouts.py b/litellm/rust_bridge/host/timeouts.py similarity index 100% rename from litellm/rust_bridge/timeouts.py rename to litellm/rust_bridge/host/timeouts.py diff --git a/litellm/rust_bridge/responses/websocket.py b/litellm/rust_bridge/responses/websocket.py index 0634867af1c..f10631958bb 100644 --- a/litellm/rust_bridge/responses/websocket.py +++ b/litellm/rust_bridge/responses/websocket.py @@ -8,8 +8,8 @@ from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK +from litellm.rust_bridge.host.timeouts import timeout_to_seconds from litellm.rust_bridge.loader import get_native_bridge -from litellm.rust_bridge.timeouts import timeout_to_seconds class RustResponsesWebSocket(Protocol): diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 88cdd9ff5c1..7c62133202d 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -629,7 +629,8 @@ def test_native_projection_errors_never_select_python( import dataclasses import ssl - from litellm.rust_bridge import runtime, settings + from litellm.rust_bridge import runtime + from litellm.rust_bridge.host import settings from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 720695b0f42..a1e22780d57 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -2320,7 +2320,7 @@ async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_c else: client_type: Final = pytest.importorskip("google.cloud.kms_v1").KeyManagementServiceClient - from litellm.rust_bridge.secret_manager import native_secret_manager_config + from litellm.rust_bridge.host.secret_manager import native_secret_manager_config credentials_file: Final = tmp_path / "credentials.json" credentials_file.write_text( diff --git a/tests/unit/rust_bridge/host/__init__.py b/tests/unit/rust_bridge/host/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py new file mode 100644 index 00000000000..9d532631ef2 --- /dev/null +++ b/tests/unit/rust_bridge/host/test_cache.py @@ -0,0 +1,21 @@ +from typing import Final + +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.rust_bridge.host.cache import configured_cache +from litellm.types.caching import LiteLLMCacheType + + +def test_configured_cache_preserves_identity_and_observes_runtime_replacement(monkeypatch: pytest.MonkeyPatch) -> None: + initial: Final = Cache(type=LiteLLMCacheType.LOCAL) + replacement: Final = Cache(type=LiteLLMCacheType.LOCAL) + monkeypatch.setattr(litellm, "cache", initial) + assert configured_cache() is initial + + monkeypatch.setattr(litellm, "cache", replacement) + assert configured_cache() is replacement + + monkeypatch.setattr(litellm, "cache", None) + assert configured_cache() is None diff --git a/tests/unit/rust_bridge/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py similarity index 96% rename from tests/unit/rust_bridge/test_callbacks_legacy_python.py rename to tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index c78fbb3f246..eb1bfd82715 100644 --- a/tests/unit/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -10,8 +10,8 @@ from pydantic import TypeAdapter from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.rust_bridge import callbacks_legacy_python as legacy -from litellm.rust_bridge.callbacks_legacy_python import failure_handler, setup +from litellm.rust_bridge.host import callbacks_legacy_python as legacy +from litellm.rust_bridge.host.callbacks_legacy_python import failure_handler, setup from litellm.types.utils import ModelResponse _OCR_KWARGS: Final = MappingProxyType( @@ -167,7 +167,7 @@ def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation pending.close() -CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" +CONTRACT_PATH: Final = Path(__file__).parents[4] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" def test_the_rust_contract_matches_the_shim_signatures() -> None: diff --git a/tests/unit/rust_bridge/test_model_capabilities.py b/tests/unit/rust_bridge/host/test_model_capabilities.py similarity index 96% rename from tests/unit/rust_bridge/test_model_capabilities.py rename to tests/unit/rust_bridge/host/test_model_capabilities.py index 93f7dfdc04a..4d73a29d078 100644 --- a/tests/unit/rust_bridge/test_model_capabilities.py +++ b/tests/unit/rust_bridge/host/test_model_capabilities.py @@ -3,8 +3,7 @@ from typing import Final import pytest import litellm -from litellm.rust_bridge.model_capabilities import anthropic_model_capabilities - +from litellm.rust_bridge.host.model_capabilities import anthropic_model_capabilities pytestmark = pytest.mark.usefixtures("local_model_cost_map") diff --git a/tests/unit/rust_bridge/test_secret_manager.py b/tests/unit/rust_bridge/host/test_secret_manager.py similarity index 99% rename from tests/unit/rust_bridge/test_secret_manager.py rename to tests/unit/rust_bridge/host/test_secret_manager.py index 6392e1d1e98..9d876124ffd 100644 --- a/tests/unit/rust_bridge/test_secret_manager.py +++ b/tests/unit/rust_bridge/host/test_secret_manager.py @@ -17,7 +17,7 @@ from botocore.credentials import Credentials from pydantic import JsonValue import litellm -from litellm.rust_bridge.secret_manager import capture_secret_manager, native_secret_manager_config +from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager diff --git a/tests/unit/rust_bridge/test_settings.py b/tests/unit/rust_bridge/host/test_settings.py similarity index 99% rename from tests/unit/rust_bridge/test_settings.py rename to tests/unit/rust_bridge/host/test_settings.py index ea943afe970..5d3486d9a64 100644 --- a/tests/unit/rust_bridge/test_settings.py +++ b/tests/unit/rust_bridge/host/test_settings.py @@ -6,7 +6,7 @@ import pytest import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent -from litellm.rust_bridge import settings +from litellm.rust_bridge.host import settings from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index 1dab9604740..89fd1d1b771 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -10,7 +10,7 @@ import pytest import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.anthropic.pass_through.messages.handler import anthropic_messages -from litellm.rust_bridge import settings +from litellm.rust_bridge.host import settings from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES from litellm.rust_bridge.public_call import NativeCall from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py index 262fba1429b..a9c6072b357 100644 --- a/tests/unit/rust_bridge/ocr/test_secrets.py +++ b/tests/unit/rust_bridge/ocr/test_secrets.py @@ -12,7 +12,7 @@ import pytest import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import settings +from litellm.rust_bridge.host import settings from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR from litellm.rust_bridge.public_call import NativeCall from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem