From 11af2eb22dc799dff033c2d3ba71be991b6a8201 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 11 Nov 2026 10:57:54 +0900 Subject: [PATCH] refactor(rust_bridge): route every Rust-to-Python call through host modules (#50822) * refactor(rust_bridge): route every Rust-to-Python call through host modules Rust now reaches Python only through litellm/rust_bridge/host/ (plus the execution driver), and no Python module outside host/ imports it. Library calls that reproduce Python behavior (httpx and OSError error text, huggingface_hub downloads, the random generator, os.environ) go through host functions too. The rule is written into host/AGENTS.md The live litellm.model_cost table, its version and every Python writer move to litellm_core_utils/model_cost_table.py, and host/model_cost.py keeps only the functions Rust calls. The cost map fetch passes use_http_settings instead of handing Rust an http_settings() snapshot * refactor(rust_bridge): read REDIS_USERNAME from the process environment in Rust The Redis bridge asked Python for all of os.environ to read one key. CPython's os.environ writes through to the process environment, so Rust reads it with ProcessEnvironment like the rest of the module. Also drops doc comments that restated their code. Co-Authored-By: Claude Opus 5.5 * ci(rust): raise the native extension size cap to 55 MB Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- .github/scripts/verify_linux_native_wheel.py | 2 +- litellm-rust/crates/python-bridge/AGENTS.md | 1 + litellm-rust/crates/python-bridge/src/http.rs | 15 +- .../crates/python-bridge/src/logger/mod.rs | 2 +- .../python-bridge/src/model_cost/fetch.rs | 9 +- .../crates/python-bridge/src/preflight.rs | 10 +- .../python-bridge/src/python_settings.rs | 9 +- .../crates/python-bridge/src/redis/breaker.rs | 2 +- .../src/redis/client/credentials.rs | 4 +- .../python-bridge/src/redis/client/mod.rs | 65 +----- .../crates/python-bridge/src/router/host.rs | 10 +- .../python-bridge/src/routes/messages/host.rs | 17 +- .../src/routes/responses/host.rs | 9 +- .../python-bridge/src/secrets/callback.rs | 38 ++-- .../crates/python-bridge/src/secrets/mod.rs | 1 + .../python-bridge/src/secrets/mutation.rs | 28 +-- .../python-bridge/src/secrets/python.rs | 4 +- .../python-bridge/src/secrets/runtime.rs | 4 +- .../crates/python-bridge/src/secrets/vault.rs | 32 +-- .../crates/python-bridge/src/tokenizer.rs | 24 +-- litellm/litellm_core_utils/AGENTS.md | 16 +- .../litellm_core_utils/get_model_cost_map.py | 27 ++- .../litellm_core_utils/model_cost_table.py | 112 ++++++++++ .../bedrock/audio_transcription/__init__.py | 3 +- .../llms/fireworks_ai/chat/transformation.py | 2 +- .../proxy/common_utils/model_deprecation.py | 2 +- litellm/proxy/proxy_server.py | 2 +- litellm/rust_bridge/AGENTS.md | 2 +- litellm/rust_bridge/_native.pyi | 29 +-- litellm/rust_bridge/host/AGENTS.md | 50 +++-- litellm/rust_bridge/{ => host}/logger.py | 0 .../rust_bridge/host/model_capabilities.py | 2 +- litellm/rust_bridge/host/model_cost.py | 117 +---------- litellm/rust_bridge/{ => host}/preflight.py | 0 litellm/rust_bridge/host/randomness.py | 11 + litellm/rust_bridge/host/routes/responses.py | 6 + litellm/rust_bridge/host/secret_manager.py | 193 ++++-------------- litellm/rust_bridge/host/timeouts.py | 10 +- litellm/rust_bridge/host/tokenizers.py | 17 ++ litellm/rust_bridge/public_call.py | 10 + litellm/rust_bridge/responses/detectors.py | 10 +- litellm/rust_bridge/responses/websocket.py | 2 +- litellm/secret_managers/native_config.py | 159 +++++++++++++++ litellm/utils.py | 14 +- ...edrock_native_output_register_model_sdk.py | 2 +- .../sdk/test_model_cost_fetch_sdk.py | 2 +- .../unit/caching/test_llm_caching_handler.py | 2 +- tests/unit/conftest.py | 2 +- tests/unit/integrations/arize/test_arize.py | 2 +- .../integrations/arize/test_arize_phoenix.py | 2 +- tests/unit/integrations/opik/test_opik.py | 2 +- .../test_anthropic_cache_control_hook.py | 2 +- tests/unit/integrations/test_helicone.py | 2 +- ...llm_core_utils_prompt_templates_factory.py | 2 +- .../test_get_llm_provider_logic.py | 2 +- .../test_model_cost_table.py | 57 ++++++ .../native/test_anthropic_messages_handler.py | 2 +- .../test_azure_chat_gpt_transformation.py | 2 +- .../test_azure_ai_flux2_image_generation.py | 2 +- .../test_azure_ai_fw_models_metadata.py | 2 +- .../test_azure_ai_kimi_k26_metadata.py | 2 +- .../test_amazon_nova_canvas_image_edit.py | 2 +- .../llms/custom_httpx/test_http_handler.py | 2 +- tests/unit/llms/openai/test_openai.py | 2 +- .../parallel_ai/test_parallel_ai_search.py | 2 +- tests/unit/llms/test_custom_llm.py | 2 +- ...test_vertex_and_google_ai_studio_gemini.py | 2 +- .../common_utils/test_config_sync_pubsub.py | 2 +- .../common_utils/test_http_parsing_utils.py | 2 +- .../common_utils/test_model_deprecation.py | 2 +- .../proxy/hooks/test_dynamic_rate_limiter.py | 2 +- .../hooks/test_prompt_injection_detection.py | 2 +- .../test_sso_helper_utils.py | 2 +- .../proxy_server/test_routes_model_info.py | 2 +- .../proxy/proxy_server/test_routes_utils.py | 2 +- ...test_proxy_server_endpoints_and_startup.py | 2 +- tests/unit/responses/test_responses_utils.py | 2 +- .../unit/router_strategy/test_lowest_cost.py | 2 +- .../router_strategy/test_lowest_tpm_rpm_v2.py | 2 +- .../unit/rust_bridge/host/test_model_cost.py | 61 +----- .../rust_bridge/{ => host}/test_preflight.py | 6 +- .../rust_bridge/host/test_secret_manager.py | 3 +- tests/unit/rust_bridge/test_logger.py | 3 +- tests/unit/test_cost_calculator.py | 2 +- tests/unit/test_main.py | 2 +- tests/unit/test_model_cost_aliases.py | 2 +- .../test_no_top_level_test_invocations.py | 2 +- tests/unit/test_pydantic_namespaces.py | 2 +- .../test_register_model_custom_pricing.py | 4 +- tests/unit/test_router/test_router.py | 2 +- .../unit/test_router_model_cost_isolation.py | 2 +- tests/unit/test_scheduler.py | 2 +- tests/unit/test_utils.py | 18 +- tests/unit/test_utils_get_model_info.py | 2 +- tests/unit/types/test_files.py | 2 +- 95 files changed, 701 insertions(+), 621 deletions(-) create mode 100644 litellm/litellm_core_utils/model_cost_table.py rename litellm/rust_bridge/{ => host}/logger.py (100%) rename litellm/rust_bridge/{ => host}/preflight.py (100%) create mode 100644 litellm/rust_bridge/host/randomness.py create mode 100644 litellm/rust_bridge/host/tokenizers.py create mode 100644 litellm/secret_managers/native_config.py create mode 100644 tests/unit/litellm_core_utils/test_model_cost_table.py rename tests/unit/rust_bridge/{ => host}/test_preflight.py (92%) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 4c4cb399dce..571092263ed 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -214,7 +214,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 48_000_000 + native_size_limit: Final = 55_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 26858b2afe2..b18a1e2ca19 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -81,6 +81,7 @@ GIL handling to `litellm-host-python`. ## Bridge Shape +- Reach Python only through `litellm/rust_bridge/host/` and the `litellm.rust_bridge.execution` driver. Never import another LiteLLM module or read a `litellm` global. `litellm/rust_bridge/host/AGENTS.md` lists what may be imported directly - Prefer one stable method per top-level LiteLLM route, for example `messages(...)`, calling the matching `litellm-inference-` entrypoint. - Do not add one exported PyO3 function per provider helper unless there is a diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index c6a8b98747c..57c58c3f7cd 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -112,19 +112,22 @@ pub(crate) fn call_config( Ok(resolved(py, settings)) } -/// The client config for a value of `litellm.rust_bridge.host.settings.http_settings()` the caller handed in, under -/// the process environment. `None` reads the environment alone pub(crate) fn settings_config( py: Python<'_>, - http_settings: Option<&Bound<'_, PyAny>>, + use_http_settings: bool, ) -> PyResult { - let configured = match http_settings { - Some(value) => configured(&PythonSettings::Http.snapshot(value.clone()))?, + let snapshot = if use_http_settings { + PythonSettings::Http.read_or_unset(py)? + } else { + None + }; + let layer = match snapshot { + Some(snapshot) => configured(&snapshot)?, None => HttpSettingsLayer::default(), }; let settings = HttpSettings::from_layers([ HttpSettingsLayer::from_environment(&ProcessEnvironment), - configured, + layer, ]); Ok(resolved(py, settings)) } diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index 8f40d94f458..5032b418404 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -7,7 +7,7 @@ use litellm_tracing::{DiagnosticInput, Level, Logger, Metadata, Policy, Processo use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; -const MODULE: &str = "litellm.rust_bridge.logger"; +const MODULE: &str = "litellm.rust_bridge.host.logger"; type NativeDiagnosticOutput = (String, Option, Option, Vec, bool); #[pyclass] diff --git a/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs index 77ec88cc1b9..dff6296cf91 100644 --- a/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs +++ b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs @@ -20,13 +20,12 @@ pub(crate) struct NativeCostMapSource { #[pymethods] impl NativeCostMapSource { - /// `http_settings` is `litellm.rust_bridge.host.settings.http_settings()`, or `None` for the environment alone #[new] - #[pyo3(signature = (url, *, http_settings))] - fn new(py: Python<'_>, url: String, http_settings: Option>) -> PyResult { + #[pyo3(signature = (url, *, use_http_settings))] + fn new(py: Python<'_>, url: String, use_http_settings: bool) -> PyResult { Ok(Self { url, - http: crate::http::settings_config(py, http_settings.as_ref())?, + http: crate::http::settings_config(py, use_http_settings)?, }) } @@ -226,7 +225,7 @@ mod tests { }); Python::initialize(); Python::attach(|py| { - let source = NativeCostMapSource::new(py, url, None).unwrap(); + let source = NativeCostMapSource::new(py, url, false).unwrap(); let policy = Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap(); let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap(); diff --git a/litellm-rust/crates/python-bridge/src/preflight.rs b/litellm-rust/crates/python-bridge/src/preflight.rs index bc692498880..2a07723b94c 100644 --- a/litellm-rust/crates/python-bridge/src/preflight.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -7,7 +7,7 @@ use pyo3::{ }; use strum::{IntoStaticStr, VariantArray}; -const MODULE: &str = "litellm.rust_bridge.preflight"; +const MODULE: &str = "litellm.rust_bridge.host.preflight"; /// The litellm globals the preflight still reads through Python. `preflight_contract.json` /// pins each function's parameters on both sides. @@ -115,7 +115,7 @@ mod tests { /// tests that install it run one at a time. static PREFLIGHT_MODULE: Mutex<()> = Mutex::new(()); - /// A fresh stand-in for `litellm.rust_bridge.preflight` that records every call, then + /// A fresh stand-in for `litellm.rust_bridge.host.preflight` that records every call, then /// `script` run against it with the module bound as `preflight`. fn preflight_stubs<'py>(py: Python<'py>, script: &std::ffi::CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); @@ -124,15 +124,15 @@ mod tests { import sys import types -for name in ('litellm', 'litellm.rust_bridge'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -preflight = types.ModuleType('litellm.rust_bridge.preflight') +preflight = types.ModuleType('litellm.rust_bridge.host.preflight') preflight.warnings = [] preflight.checked = [] preflight.credential_list = lambda: [] preflight.warn_unknown_credential = lambda name, loaded: preflight.warnings.append((name, loaded)) preflight.check_limits = lambda kwargs: preflight.checked.append(kwargs) -sys.modules['litellm.rust_bridge.preflight'] = preflight +sys.modules['litellm.rust_bridge.host.preflight'] = preflight ", Some(&locals), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 93377abb7e9..1cb1813049c 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -26,10 +26,15 @@ pub(crate) struct Snapshot<'py> { value: Bound<'py, PyAny>, } -impl Snapshot<'_> { +impl<'py> Snapshot<'py> { pub(crate) fn read(&self, spec: &FieldSpec) -> Result { spec.read(&self.value, self.group.into()) } + + pub(crate) fn attribute(&self, name: &str) -> PyResult>> { + let value = self.value.getattr(name)?; + Ok((!value.is_none()).then_some(value)) + } } impl PythonSettings { @@ -56,7 +61,7 @@ impl PythonSettings { } } - /// A snapshot of a value Python handed in, read with the same field rules as the accessor's + #[cfg(test)] pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { Snapshot { group: self, value } } diff --git a/litellm-rust/crates/python-bridge/src/redis/breaker.rs b/litellm-rust/crates/python-bridge/src/redis/breaker.rs index 733cf1bd969..5fdac76cb02 100644 --- a/litellm-rust/crates/python-bridge/src/redis/breaker.rs +++ b/litellm-rust/crates/python-bridge/src/redis/breaker.rs @@ -72,7 +72,7 @@ impl NativeRedisBreaker { redis_kwargs: RedisSpec, settings: BreakerSettings, ) -> PyResult { - let endpoint = resolved_endpoint(py, &redis_kwargs)?; + let endpoint = resolved_endpoint(&redis_kwargs); Self::build(py, &settings, endpoint) } diff --git a/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs index 21a37947c30..6995981a03e 100644 --- a/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs @@ -5,7 +5,7 @@ use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, use litellm_secret_value::SecretValue; use pyo3::prelude::*; -const MODULE: &str = "litellm.rust_bridge.host.redis_credentials"; +const CREDENTIALS_MODULE: &str = "litellm.rust_bridge.host.redis_credentials"; /// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host /// module the way redis-py's async connection asks it. @@ -29,7 +29,7 @@ impl CredentialSource for PythonCredentialSource { let context = self.context.clone(); Box::pin(async move { let (username, password) = attach_blocking(context, move |py| { - py.import(MODULE)? + py.import(CREDENTIALS_MODULE)? .getattr("credentials")? .call1((provider.bind(py),))? .extract::<(Option, String)>() diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs index 989f3dc9247..5e5ab7e6f5a 100644 --- a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -3,7 +3,7 @@ mod credentials; use std::{collections::BTreeMap, sync::Arc, time::Duration}; use litellm_circuit_breaker::Config as BreakerConfig; -use litellm_core_utils::settings::ProcessEnvironment; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; use litellm_host_python::PythonContext; use litellm_python_compat::number::{self, SecondsError}; use litellm_redis_client::{RedisClient, RedisEvents}; @@ -92,7 +92,7 @@ fn shared_client( ) -> PyResult> { let environment: SharedEnvironment = Arc::new(ProcessEnvironment); let logger = crate::logger::capture(py); - let resolved = match logger.scope(|| resolve_spec(py, &spec))? { + let resolved = match logger.scope(|| resolve_spec(&spec)) { Ok(resolution) => { logger.scope(|| { for ignored in &resolution.ignored { @@ -174,39 +174,23 @@ impl NativeRedisClient { } /// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it. -pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult> { - Ok(resolve_spec(py, spec)? +pub(crate) fn resolved_endpoint(spec: &RedisSpec) -> Option { + resolve_spec(spec) .ok() - .and_then(|resolution| resolution.target.endpoint())) + .and_then(|resolution| resolution.target.endpoint()) } // A live registry entry's client keeps its provider alive, so no live entry shares its address -fn resolve_spec( - py: Python<'_>, - spec: &RedisSpec, -) -> PyResult> { +fn resolve_spec(spec: &RedisSpec) -> Result { let handle = spec .credential_provider .as_ref() .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); - resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?) -} - -/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth -/// reads `os.environ`. -fn resolve( - kwargs: &Kwargs, - provider: Option, - environ: &Bound<'_, PyAny>, -) -> PyResult> { - let redis_username: Option = environ - .call_method1("get", ("REDIS_USERNAME",))? - .extract()?; - Ok(litellm_redis_resolve::interpret( - kwargs, - provider, - redis_username.as_deref(), - )) + litellm_redis_resolve::interpret( + &spec.kwargs, + handle, + ProcessEnvironment.get("REDIS_USERNAME").as_deref(), + ) } fn unresolved( @@ -257,9 +241,6 @@ fn seconds_error(error: SecondsError) -> PyErr { #[cfg(test)] mod tests { - use std::ffi::CStr; - - use litellm_redis_types::Auth; use rstest::rstest; use serde_json::json; @@ -270,30 +251,6 @@ mod tests { litellm_redis_resolve::kwargs(fields) } - #[rstest] - #[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))] - #[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)] - fn azure_auth_takes_its_username_from_the_environ_like_python( - #[case] source: &CStr, - #[case] expected: Option<&str>, - ) { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run(source, Some(&locals), Some(&locals)).unwrap(); - let environ = locals.get_item("environ").unwrap().unwrap(); - let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); - - let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap(); - - assert!( - matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), - "{:?}", - resolved.auth - ); - }); - } - #[rstest] #[case::no_target(json!({"port": 6379}), Some("ValueError"))] #[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))] diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 27f9fe347e0..04d19e70edc 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -717,6 +717,8 @@ pub(crate) fn deployment_record(deployment: &Deployment) -> Value { deployment.record() } +const RANDOMNESS_MODULE: &str = "litellm.rust_bridge.host.randomness"; + /// The process's `random` module: a Python router draws its shuffle and retry jitter there, so /// a seeded process picks the same deployments on either core. The bare module functions share /// one generator, so these calls interleave with every other draw the process makes. @@ -725,7 +727,9 @@ pub(crate) struct PythonRandomness; impl Randomness for PythonRandomness { fn next_f64(&self) -> f64 { Python::attach(|py| -> PyResult { - py.import("random")?.call_method0("random")?.extract() + py.import(RANDOMNESS_MODULE)? + .call_method0("next_float")? + .extract() }) .unwrap_or_else(|_| ThreadRandomness.next_f64()) } @@ -733,8 +737,8 @@ impl Randomness for PythonRandomness { fn index_below(&self, len: usize) -> usize { // `random.choice` is `seq[self._randbelow(len(seq))]`, and so is `randrange(len)`. Python::attach(|py| -> PyResult { - py.import("random")? - .call_method1("randrange", (len,))? + py.import(RANDOMNESS_MODULE)? + .call_method1("index_below", (len,))? .extract() }) .unwrap_or_else(|_| ThreadRandomness.index_below(len)) 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 2a5f79ca393..c975ac339c2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -24,11 +24,12 @@ use crate::{ ArgumentError, optional_timeout, project_optional_fields, public_response, python_timeout_seconds, }, - python_settings::missing_module, + python_settings::PythonSettings, routes::RequestView, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; +const MODEL_CAPABILITIES_MODULE: &str = "litellm.rust_bridge.host.model_capabilities"; fn merge_headers( forwarded: Option>, @@ -40,16 +41,10 @@ fn merge_headers( (!merged.is_empty()).then_some(merged) } -/// Reads `litellm.`, the module global a param spec names, the way -/// `VertexBase.get_vertex_ai_project` reads `litellm.vertex_project`. fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult>> { - let module = match py.import("litellm") { - Ok(module) => module, - Err(error) if missing_module(py, &error, "litellm")? => return Ok(None), - Err(error) => return Err(error), - }; - let value = module.getattr(name)?; - Ok((!value.is_none()).then_some(value)) + PythonSettings::ProviderDefaults + .read_or_unset(py)? + .map_or(Ok(None), |defaults| defaults.attribute(name)) } /// The caller's litellm params, read from the kwargs by the names the typed params declare, @@ -173,7 +168,7 @@ impl MessagesPythonHost { ) -> PyResult { let module = py.import(ROUTE_HOST_MODULE)?; let capabilities: MessagesModelCapabilities = from_py( - &py.import("litellm.rust_bridge.host.model_capabilities")? + &py.import(MODEL_CAPABILITIES_MODULE)? .getattr("anthropic_model_capabilities")? .call1((model, custom_llm_provider))?, )?; diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index b797e334c6a..019f8452737 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -44,12 +44,9 @@ pub(super) fn project( let value = match (name.as_str(), value) { ("previous_response_id", serde_json::Value::String(id)) => { let decoded: String = py - .import("litellm.responses.utils")? - .getattr("ResponsesAPIRequestUtils")? - .call_method1( - "decode_previous_response_id_to_original_previous_response_id", - (id,), - )? + .import(super::ROUTE_HOST_MODULE)? + .getattr("original_previous_response_id")? + .call1((id,))? .extract()?; serde_json::Value::String(decoded) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index af6485a56e1..a2afb42b540 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -15,7 +15,6 @@ use pyo3::{ use super::error::{external_error, read_error}; -const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler"; const DIAGNOSTICS_MODULE: &str = "litellm.rust_bridge.host.diagnostics"; /// A secret manager whose reads execute in Python: a custom manager, a legacy compatible @@ -61,8 +60,8 @@ impl PythonClient { None => kwargs.set_item("key_management_settings", py.None())?, } let result = py - .import(HANDLER_MODULE)? - .getattr("get_secret_from_manager")? + .import(super::SECRET_MANAGER_MODULE)? + .getattr("read_from_manager")? .call((), Some(&kwargs))?; if result.is_instance_of::() { result.extract().map(Some) @@ -133,7 +132,8 @@ mod tests { use litellm_host_python::PythonContext; - use super::{DIAGNOSTICS_MODULE, HANDLER_MODULE, PythonSecretManager, python_name}; + use super::{DIAGNOSTICS_MODULE, PythonSecretManager, python_name}; + use crate::secrets::SECRET_MANAGER_MODULE; use crate::secrets::python_error; /// `sys.modules` is interpreter-global, so tests that install or rely on the handler module @@ -169,12 +169,12 @@ class Manager: raise failure manager = Manager() import sys, types -for name in ('litellm', 'litellm.secret_managers'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -handler = sys.modules.setdefault('litellm.secret_managers.secret_manager_handler', types.ModuleType('litellm.secret_managers.secret_manager_handler')) -def get_secret_from_manager(**kwargs): +handler = sys.modules.setdefault('litellm.rust_bridge.host.secret_manager', types.ModuleType('litellm.rust_bridge.host.secret_manager')) +def read_from_manager(**kwargs): return kwargs['client'].sync_read_secret(kwargs['secret_name']) -handler.get_secret_from_manager = get_secret_from_manager +handler.read_from_manager = read_from_manager ", Some(&locals), Some(&locals), @@ -317,7 +317,7 @@ sys.modules.setdefault('litellm._logging', logging) }); } - /// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and + /// Installs a fake `read_from_manager` that records its kwargs, runs `body`, and /// removes the fake handler again; parent package stubs persist for concurrent tests. /// Callers hold `handler_guard` before attaching so the GIL is never held while waiting on it. fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) { @@ -325,16 +325,16 @@ sys.modules.setdefault('litellm._logging', logging) py.run( c" import sys, types -previous_handler = sys.modules.get('litellm.secret_managers.secret_manager_handler') +previous_handler = sys.modules.get('litellm.rust_bridge.host.secret_manager') calls = [] -def get_secret_from_manager(**kwargs): +def read_from_manager(**kwargs): calls.append(kwargs) return 'handled-' + kwargs['secret_name'] -handler = types.ModuleType('litellm.secret_managers.secret_manager_handler') -handler.get_secret_from_manager = get_secret_from_manager -for name in ('litellm', 'litellm.secret_managers'): +handler = types.ModuleType('litellm.rust_bridge.host.secret_manager') +handler.read_from_manager = read_from_manager +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -sys.modules['litellm.secret_managers.secret_manager_handler'] = handler +sys.modules['litellm.rust_bridge.host.secret_manager'] = handler ", Some(&locals), Some(&locals), @@ -344,9 +344,9 @@ sys.modules['litellm.secret_managers.secret_manager_handler'] = handler py.run( c" if previous_handler is None: - sys.modules.pop('litellm.secret_managers.secret_manager_handler', None) + sys.modules.pop('litellm.rust_bridge.host.secret_manager', None) else: - sys.modules['litellm.secret_managers.secret_manager_handler'] = previous_handler + sys.modules['litellm.rust_bridge.host.secret_manager'] = previous_handler ", Some(&locals), Some(&locals), @@ -366,7 +366,7 @@ else: with_fake_handler(py, |locals| { locals.set_item("expression", expression).unwrap(); py.run( - c"handler.get_secret_from_manager = lambda **kwargs: eval(expression)", + c"handler.read_from_manager = lambda **kwargs: eval(expression)", Some(locals), Some(locals), ) @@ -417,7 +417,7 @@ else: reader.client.read(py, "API_KEY").unwrap().as_deref(), Some("handled-API_KEY") ); - assert!(py.import(HANDLER_MODULE).is_ok()); + assert!(py.import(SECRET_MANAGER_MODULE).is_ok()); let calls = locals.get_item("calls").unwrap().unwrap(); let call = calls.get_item(0).unwrap().cast_into::().unwrap(); assert!(call.get_item("client").unwrap().unwrap().is(&client)); diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index 57042e6a32e..0840ecdfc6c 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -21,6 +21,7 @@ use resolved::ResolvedSecrets; use crate::{coercion::FieldSpec, python_settings::PythonSettings}; +const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; const NATIVE: FieldSpec = FieldSpec::new("native", |field| field.schema_bool()); /// Where a Rust route reads provider secrets from. Python's `get_secret_str` when the settings diff --git a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs index b04060651ac..47cd36c4015 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs @@ -2,7 +2,7 @@ use super::operations::{PythonMutationError, PythonMutationResponse}; use litellm_error::Report; use litellm_host_python::{json_loads, to_py}; use litellm_secrets::cyberark; -use pyo3::{prelude::*, types::PyDict}; +use pyo3::prelude::*; pub(super) fn mutation_value( result: Result, @@ -62,14 +62,11 @@ fn cyberark_message( http_message(py, "POST", url, status)? } cyberark::Error::Http(source) if authentication => match os_error_code(&source) { - Some(code) => { - let reason = py.import("os")?.getattr("strerror")?.call1((code,))?; - py.import("builtins")? - .getattr("OSError")? - .call1((code, reason))? - .str()? - .extract()? - } + Some(code) => py + .import(super::SECRET_MANAGER_MODULE)? + .getattr("os_error_message")? + .call1((code,))? + .extract()?, None => Report::new(&cyberark::Error::Http(source)).to_string(), }, cyberark::Error::Http(source) if source.is_connect() => { @@ -103,13 +100,8 @@ pub(super) fn http_message( url: &str, status: u16, ) -> PyResult { - let httpx = py.import("httpx")?; - let request = httpx.getattr("Request")?.call1((method, url))?; - let kwargs = PyDict::new(py); - kwargs.set_item("request", request)?; - let response = httpx.getattr("Response")?.call((status,), Some(&kwargs))?; - match response.call_method0("raise_for_status") { - Err(error) => error.value(py).str()?.extract(), - Ok(_) => Ok(format!("HTTP {status}")), - } + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("http_status_message")? + .call1((method, url, status))? + .extract() } diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index c5d9e3565b3..a92cf9a497a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -9,8 +9,6 @@ use pyo3::prelude::*; use super::error::external_error; -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 { @@ -21,7 +19,7 @@ pub(super) struct PythonSecrets { impl PythonSecrets { pub(super) fn new(py: Python<'_>) -> PyResult { Ok(Self::reading_with( - py.import(SECRET_MANAGER_MODULE)? + py.import(super::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 0d52b0aa629..8061d1682ed 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -12,8 +12,6 @@ use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict}; use super::ManagerError; 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, @@ -98,7 +96,7 @@ impl NativeSecretManager { return Ok(Some(native)); } let config = py - .import(SECRET_MANAGER_MODULE)? + .import(super::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 30e3bd2c413..5f14727abaa 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -134,31 +134,21 @@ fn timeout_message( String::new() }); } - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item( - "message", - transport_timeout_message( - TimeoutPhase::Read, - context.timeout, - Some(elapsed.as_secs_f64()), - ), - )?; - kwargs.set_item("model", "default-model-name")?; - kwargs.set_item("llm_provider", "litellm-httpx-handler")?; - py.import("litellm.exceptions")? - .getattr("Timeout")? - .call((), Some(&kwargs))? - .str()? + let message = transport_timeout_message( + TimeoutPhase::Read, + context.timeout, + Some(elapsed.as_secs_f64()), + ); + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("timeout_message")? + .call1((message,))? .extract() } fn response_text(py: Python<'_>, body: &[u8]) -> PyResult { - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item("content", pyo3::types::PyBytes::new(py, body))?; - py.import("httpx")? - .getattr("Response")? - .call((200,), Some(&kwargs))? - .getattr("text")? + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("response_text")? + .call1((pyo3::types::PyBytes::new(py, body),))? .extract() } diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index 2026b2fdc42..11b21bd5673 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -38,16 +38,19 @@ use litellm_token_counter::huggingface::{ #[cfg(feature = "tiktoken")] use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary}; +#[cfg(any(feature = "tiktoken", feature = "huggingface"))] +const HOST_MODULE: &str = "litellm.rust_bridge.host.tokenizers"; + #[cfg(feature = "tiktoken")] pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { enter_native()?; - let resource: std::path::PathBuf = - PyModule::import(py, "litellm.litellm_core_utils.tokenizers")? - .getattr("__file__")? - .extract()?; + let resources: std::path::PathBuf = PyModule::import(py, HOST_MODULE)? + .getattr("tiktoken_resource_dir")? + .call0()? + .extract()?; release_gil(py, || { TiktokenTokenizer::from_cached_ranks(encoding, |file| { - std::fs::read_to_string(resource.with_file_name(file)) + std::fs::read_to_string(resources.join(file)) }) }) .map_err(|error| token_count_error_to_pyerr(error.into())) @@ -197,14 +200,9 @@ impl Tokenizer { #[cfg(feature = "huggingface")] { enter_native()?; - let kwargs = PyDict::new(py); - kwargs.set_item("repo_id", identifier)?; - kwargs.set_item("filename", "tokenizer.json")?; - kwargs.set_item("revision", revision)?; - kwargs.set_item("token", token)?; - let path: String = PyModule::import(py, "huggingface_hub")? - .getattr("hf_hub_download")? - .call((), Some(&kwargs))? + let path: String = PyModule::import(py, HOST_MODULE)? + .getattr("download_tokenizer")? + .call1((identifier, revision, token))? .extract()?; let json = release_gil(py, || std::fs::read_to_string(path)) .map_err(TokenizerError::Download)?; diff --git a/litellm/litellm_core_utils/AGENTS.md b/litellm/litellm_core_utils/AGENTS.md index e817eca0ac4..9cfeccba44c 100644 --- a/litellm/litellm_core_utils/AGENTS.md +++ b/litellm/litellm_core_utils/AGENTS.md @@ -3,7 +3,7 @@ - This module decides when to fetch and install. Fetching, parsing, integrity checks and alias expansion run in Rust - The install order is publish, `add_known_models`, router replay, then `model_cost_inputs.reapply` - Registrations land by lookup in the published table, so they are written after publishing -- Startup and `register_model(url)` fetch with `http_settings=None`, which reads env vars only like main's bare httpx. Reload passes `http_settings()` +- Startup and `register_model(url)` fetch with `use_http_settings=False`, which reads env vars only like main's bare httpx. Reload passes `use_http_settings=True`, so Rust layers the configured HTTP settings on top - Startup falls back to the bundled map on any failure. Reload never falls back and keeps the current map - `register_model(url)` makes one attempt and registers nothing on failure - `utils._register_on_current_map` reruns the whole pass on the new table after `NativeTableSwapped`, and its last attempt passes `rows=None`, so no entry is lost @@ -12,4 +12,16 @@ - `GetModelCostMap.validate_model_cost_map` and `_check_*`, since Rust checks fetched documents - `ModelCostMapReloaded.model_cost_map`, use `GetModelCostMap.loaded_model_cost_map()` - the `sleep`, `rng` and `client` parameters, which are ignored with a `DeprecationWarning` - - `AliasedCostEntry` re-exported here, which now lives in `rust_bridge/host/model_cost.py` + - `AliasedCostEntry` re-exported here, which now lives in `model_cost_table.py` + +## model_cost_table.py + +- This module owns the live `litellm.model_cost` table. Python code reads and writes it here, and `rust_bridge/host/model_cost.py` calls it for Rust +- `publish`, `write` and `changed_in_place` are the only writers. Each moves `ModelCostVersion` +- The version is `(id(table), len(table), writes)`, so replacing the dict or changing its size moves it without a writer + - An in-place value edit does not, so code that edits rows in place calls `changed_in_place()` +- Anything derived from the table is cached under the version (`@VersionCached`, `model_cost_version()`), never cleared by hand +- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own + - `utils._invalidate_model_cost_lowercase_map` forwards to `changed_in_place` + - `utils.get_model_cost_mutation_generation` forwards to `model_cost_version().writes` + - `utils._model_cost_lowercase_map` is always `None` diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 80393a72631..6f4405066ca 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -32,9 +32,13 @@ from litellm.constants import ( MODEL_COST_MAP_MAX_SHRINK_RATIO, MODEL_COST_MAP_MIN_MODEL_COUNT, ) +from litellm.litellm_core_utils import model_cost_table from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +from litellm.litellm_core_utils.model_cost_table import ( + AliasedCostEntry as AliasedCostEntry, # noqa: PLC0414 # main exported the row type from here +) from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.rust_bridge._native import ( NativeCostMapSource, @@ -44,11 +48,6 @@ from litellm.rust_bridge._native import ( NativeModelCostMap, NativeRejected, ) -from litellm.rust_bridge.host import model_cost as model_cost_host -from litellm.rust_bridge.host.model_cost import ( - AliasedCostEntry as AliasedCostEntry, # noqa: PLC0414 # main exported the row type from here -) -from litellm.rust_bridge.host.settings import HttpSettings _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) @@ -259,8 +258,8 @@ def get_model_cost_map_loaded_at() -> "datetime | None": return _LastLoad.current.loaded_at -def cost_map_source(url: str, http_settings: HttpSettings | None) -> NativeCostMapSource: - return NativeCostMapSource(url, http_settings=http_settings) +def cost_map_source(url: str, use_http_settings: bool) -> NativeCostMapSource: + return NativeCostMapSource(url, use_http_settings=use_http_settings) def fetch_policy(timeout: float, max_attempts: int) -> NativeFetchPolicy: @@ -288,7 +287,7 @@ def adopt_model_cost_map(new_model_cost_map: dict) -> int: # mutable-ok: main's import litellm from litellm import utils - model_cost_host.publish(new_model_cost_map) + model_cost_table.publish(new_model_cost_map) litellm.add_known_models(model_cost_map=new_model_cost_map) utils.reapply_runtime_model_cost_registrations() return len(new_model_cost_map) @@ -354,15 +353,13 @@ async def refetch_model_cost_map( on failure it returns ``ModelCostMapReloadUnavailable`` and the map already installed stays. """ - from litellm.rust_bridge.host.settings import http_settings - _warn_ignored_fetch_hooks(sleep, rng, client) if _uses_bundled_map_only(): _mark_rows_with_rules_set("local", None, True, None, datetime.now(timezone.utc)) return ModelCostMapReloaded(model_count=_install(model_cost_inputs.install_bundled())) - fetched: Final = await cost_map_source(url, http_settings()).fetch(fetch_policy(timeout, max_attempts)) + fetched: Final = await cost_map_source(url, use_http_settings=True).fetch(fetch_policy(timeout, max_attempts)) if isinstance(fetched, NativeFetchFailed): verbose_logger.warning( "LiteLLM: model cost map reload failed: %s. Keeping the currently loaded map.", fetched.reason @@ -382,7 +379,9 @@ def _retry_remote_fetch_in_background( url: str, timeout: int, max_attempts: int, first_failure: NativeFetchFailed ) -> None: try: - fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(timeout, max_attempts), first_failure) + fetched: Final = cost_map_source(url, use_http_settings=False).fetch_blocking( + fetch_policy(timeout, max_attempts), first_failure + ) if isinstance(fetched, NativeFetchFailed): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", @@ -425,7 +424,7 @@ def get_model_cost_map( _mark_rows_with_rules_set("local", None, True, None, started) return _rows_with_rules_set(model_cost_inputs.install_bundled()) - fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(timeout, 1)) + fetched: Final = cost_map_source(url, use_http_settings=False).fetch_blocking(fetch_policy(timeout, 1)) if isinstance(fetched, NativeFetchFailed): if fetched.retryable and max_attempts > 1: threading.Thread( @@ -459,7 +458,7 @@ def get_model_cost_map( def fetch_registration_document(url: str) -> dict[str, dict[str, object]] | None: """The document behind ``register_model(url)``, fetched in one attempt, or ``None`` when the fetch failed.""" - fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(5, 1)) + fetched: Final = cost_map_source(url, use_http_settings=False).fetch_blocking(fetch_policy(5, 1)) if isinstance(fetched, NativeFetchFailed): verbose_logger.warning("LiteLLM: register_model could not fetch %s: %s", url, fetched.reason) return None diff --git a/litellm/litellm_core_utils/model_cost_table.py b/litellm/litellm_core_utils/model_cost_table.py new file mode 100644 index 00000000000..c28098d674c --- /dev/null +++ b/litellm/litellm_core_utils/model_cost_table.py @@ -0,0 +1,112 @@ +"""The live ``litellm.model_cost`` table, every write to it, and its version.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import lru_cache +from typing import Final, Generic, ParamSpec, TypeVar + +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE + +P = ParamSpec("P") +R = TypeVar("R") + + +class AliasedCostEntry(dict[str, object]): + """A ``litellm.model_cost`` row, the one object every alias name of a model is served under.""" + + __slots__ = ("names", "registered_fields") + + def __init__( + self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset() + ) -> None: + super().__init__(entry) + self.names: tuple[str, ...] = names + self.registered_fields: frozenset[str] = registered_fields + + +@dataclass(frozen=True, slots=True) +class ModelCostVersion: + table: int + size: int + writes: int + + +class _Writes: + __slots__ = ("count",) + + def __init__(self) -> None: + self.count: int = 0 + + def record(self) -> None: + self.count += 1 + + +_WRITES: Final = _Writes() + + +def model_cost_version() -> ModelCostVersion: + import litellm + + table: Final = litellm.model_cost + return ModelCostVersion(table=id(table), size=len(table), writes=_WRITES.count) + + +def publish(table: dict[str, object]) -> None: + """Makes ``table`` the live ``litellm.model_cost``.""" + import litellm + + litellm.model_cost = table + _WRITES.record() + + +def write(rows: dict[str, object], written: Mapping[str, object]) -> None: + """Writes one registration's rows into ``rows``, the published table.""" + rows.update(written) + _WRITES.record() + + +def changed_in_place() -> None: + """Moves the version and drops the name index after an in-place edit of ``litellm.model_cost``.""" + from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs + + model_cost_inputs.invalidate() + _WRITES.record() + + +class VersionCached(Generic[P, R]): + """Decorates ``function`` with an ``lru_cache`` keyed on the ``litellm.model_cost`` version each call read, and + exposes that cache's ``cache_clear`` and ``cache_info``.""" + + __slots__ = ("_at", "cache_clear", "cache_info") + + def __init__(self, function: Callable[P, R]) -> None: + def at( + _version: ModelCostVersion, + *args: P.args, + **kwargs: P.kwargs, # kwargs-ok: forwards a ParamSpec unchanged + ) -> R: + return function(*args, **kwargs) + + cached: Final = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(at) + self._at: Final = cached + self.cache_clear: Final = cached.cache_clear + self.cache_info: Final = cached.cache_info + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: # kwargs-ok: forwards a ParamSpec unchanged + return self._at(model_cost_version(), *args, **kwargs) + + +@dataclass(frozen=True, slots=True) +class Landing: + key: str + base: Mapping[str, object] + matched_builtin: bool + + +def published() -> dict[str, object]: + """The table published as ``litellm.model_cost`` now.""" + import litellm + + return litellm.model_cost diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index bc74e9c38ab..dcda72c9a60 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -10,8 +10,7 @@ from pydantic import BaseModel, ConfigDict 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.public_call import NativeCall, timeout_to_seconds from litellm.rust_bridge.transcription.native import ( NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION, diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 9db1d50d730..55a75423042 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -13,13 +13,13 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) +from litellm.litellm_core_utils.model_cost_table import ModelCostVersion, model_cost_version from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, unpack_legacy_defs, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index b50890089e4..f33bb82b49b 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap -from litellm.rust_bridge.host.model_cost import AliasedCostEntry +from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry from litellm.types.proxy.model_deprecation import ( DEFAULT_DEPRECATION_WARN_DAYS, DeprecationStatus, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 888276f7c31..f7a27e4896c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7896,7 +7896,7 @@ class ProxyConfig: else: raise ValueError("Invalid Key Management System selected") - from litellm.rust_bridge.host.secret_manager import capture_secret_manager + from litellm.secret_managers.native_config 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 c2a030ca4d9..167b91c9b89 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -9,7 +9,7 @@ - provider execution, which is Rust core - CPython runtime mechanics, which is `litellm-host-python` - legacy callback sharing and dispatch policy, which is `callbacks-legacy-python` -- `host/` holds the modules Rust imports by name to read Python-owned state. `host/AGENTS.md` states that contract +- `host/` holds the modules Rust imports by name to reach Python-owned state and behavior. Only Rust calls them, and Rust reaches Python nowhere else. `host/AGENTS.md` states that contract # Routing diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 221968e07ed..ebbd45c7382 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -8,7 +8,6 @@ from typing_extensions import disjoint_base from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings from litellm.rust_bridge.execution import Execution -from litellm.rust_bridge.host.settings import HttpSettings from litellm.rust_bridge.public_call import NativeCall _T_co = TypeVar("_T_co", covariant=True) @@ -22,10 +21,9 @@ class NativeRedisStoreError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... - @final class NativeCostMapSource: - def __new__(cls, url: str, *, http_settings: HttpSettings | None) -> NativeCostMapSource: ... + def __new__(cls, url: str, *, use_http_settings: bool) -> NativeCostMapSource: ... def fetch(self, policy: NativeFetchPolicy) -> Future[NativeFetched | NativeFetchFailed]: ... def fetch_blocking( self, policy: NativeFetchPolicy, resume_from: NativeFetchFailed | None = None @@ -230,9 +228,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... def increment_deployment( self, deployment_id: str, @@ -241,12 +237,8 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> Future[None]: ... - def get_many( - self, counters: Sequence[_UsageCounter] - ) -> Future[list[float | None]]: ... - def increment_sync( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... def increment_deployment_sync( self, deployment_id: str, @@ -255,9 +247,7 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> None: ... - def get_many_sync( - self, counters: Sequence[_UsageCounter] - ) -> list[float | None]: ... + def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -433,7 +423,9 @@ class NativeAuthCache: def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... - def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ... + def declare_set( + self, key: str, value: object, ttl_seconds: float | None, send: NativeSend + ) -> NativePending[None]: ... def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ... @@ -550,7 +542,6 @@ def messages( def amessages( call: NativeCall, ) -> Execution: ... - @final class ResponsesWebSocketConnection: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @@ -1189,7 +1180,9 @@ class RouterCoreHandle: callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... # Count a success against any id, driven like `report_failure`; `done` carries `key`. - def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + def count_success( + self, id: str, callbacks: Mapping[str, bool] | None = None + ) -> RouterCall[_SuccessReportSteps]: ... # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 4cb87a85c28..c249eb82ee0 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -1,24 +1,48 @@ # 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 +Python behavior 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), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with), `model_cost` (the live `litellm.model_cost` table, every write to it and its version) -- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here +## Boundary + +- This package is the only place Rust reaches into Python. `python-bridge` and `callbacks-legacy-python` import no other LiteLLM module, and never do `py.import("litellm").getattr(...)` on a global + - The one exception is `litellm.rust_bridge.execution`, the driver's step classes. It is runtime protocol, not Python state, and the Python entrypoints import it too +- Only Rust calls these modules. No Python module outside `host/` imports one. Tests may, to set up the Python state Rust reads + - Logic both sides need lives with its Python owner (`public_call`, `secret_managers`, the route's types), and the host module calls it +- A Python call that reproduces Python behavior goes through a host function, including standard library and third-party calls: error text built from `httpx`, `OSError` or `litellm.exceptions`, `huggingface_hub` downloads, the process `random` generator + - Rust reads the process environment itself (`ProcessEnvironment`). CPython's `os.environ` writes through to it, so there is nothing to ask Python for + - Rust imports a module directly only to get a type or a Python builtin it applies itself: `isinstance` checks (`os.PathLike`, `ssl.SSLContext`), exception classes it raises (`redis.exceptions`), `builtins.object` and `inspect.getattr_static` for Python attribute semantics + - `litellm-host-python` is the generic runtime and keeps its own `asyncio`, `contextvars` and `json` imports +- Check these invariants by review whenever a Rust file adds a `py.import`, an `import_exception!` or a `*_MODULE` constant, or a Python file imports from `litellm.rust_bridge.host` + +## Modules + +- 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 + - `logger`: where native diagnostics are emitted + - `preflight`: the SDK globals the native preflight borrows + - `secret_manager`: the configured secret manager, its native config and the error text Python's managers produce + - `http_transport`: which async transport is in use + - `timeouts`: timeout to seconds + - `tokenizers`: bundled tiktoken files and Hugging Face tokenizer downloads + - `randomness`: the process `random` generator the Python router draws from + - `redis_events`: reporting native Redis calls and breaker transitions to `ServiceLogging` + - `redis_credentials`: what a caller's redis-py `CredentialProvider` authenticates with + - `model_cost`: where a `register_model` entry lands in the published table and the row it writes +- `routes/.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead - 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 + +## Contract + +- The module path and function signature are a contract. Each Rust crate declares a host module's path once, in a `const *_MODULE` next to the code that asks (`python-bridge/src/secrets/mod.rs`, `python-bridge/src/routes//`, `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 -- `routes/.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead ## model_cost -- `publish`, `write` and `changed_in_place` are the only writers. Each moves `ModelCostVersion` -- The version is `(id(table), len(table), writes)`, so replacing the dict or changing its size moves it without the host - - An in-place value edit does not, so code that edits rows in place calls `changed_in_place()` -- Anything derived from the table is cached under the version (`@VersionCached`, `model_cost_version()`), never cleared by hand +- `model_cost` answers what Rust needs while registering models: `entry`, `published`, `write`, `resolve`, `merged` and `kept`. The live table, its version and every Python writer belong to `litellm_core_utils/model_cost_table.py` - `resolve`, `merged` and `kept` are shims until `get_model_info` and a merge generic over Python values move to Rust. Then their bodies become native calls and the signatures stay - `resolve` returns `None` when it was handed a table that is no longer published. That is how a pass learns a reload swapped the table -- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own - - `utils._invalidate_model_cost_lowercase_map` forwards to `changed_in_place` - - `utils.get_model_cost_mutation_generation` forwards to `model_cost_version().writes` - - `utils._model_cost_lowercase_map` is always `None` diff --git a/litellm/rust_bridge/logger.py b/litellm/rust_bridge/host/logger.py similarity index 100% rename from litellm/rust_bridge/logger.py rename to litellm/rust_bridge/host/logger.py diff --git a/litellm/rust_bridge/host/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py index 6fd8b42520e..00a00b9e79b 100644 --- a/litellm/rust_bridge/host/model_capabilities.py +++ b/litellm/rust_bridge/host/model_capabilities.py @@ -6,7 +6,7 @@ from typing import Final, TypeAlias import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version +from litellm.litellm_core_utils.model_cost_table import ModelCostVersion, model_cost_version class _Held: diff --git a/litellm/rust_bridge/host/model_cost.py b/litellm/rust_bridge/host/model_cost.py index 328e0c6252e..cfedbb64592 100644 --- a/litellm/rust_bridge/host/model_cost.py +++ b/litellm/rust_bridge/host/model_cost.py @@ -1,118 +1,21 @@ -"""The live ``litellm.model_cost`` table, every write to it, and its version.""" +"""The ``litellm.model_cost`` questions Rust asks while registering models.""" from __future__ import annotations -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from functools import lru_cache -from typing import Final, Generic, ParamSpec, TypeVar +from collections.abc import Mapping +from typing import Final -from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE - -P = ParamSpec("P") -R = TypeVar("R") +from litellm.litellm_core_utils import model_cost_table +from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry, Landing -class AliasedCostEntry(dict[str, object]): - """A ``litellm.model_cost`` row, the one object every alias name of a model is served under.""" - - __slots__ = ("names", "registered_fields") - - def __init__( - self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset() - ) -> None: - super().__init__(entry) - self.names: tuple[str, ...] = names - self.registered_fields: frozenset[str] = registered_fields +def entry(fields: Mapping[str, object], names: tuple[str, ...]) -> AliasedCostEntry: + """The row object every name in ``names`` shares, the owner first.""" + return AliasedCostEntry(fields, names) -entry: Final = AliasedCostEntry - - -@dataclass(frozen=True, slots=True) -class ModelCostVersion: - table: int - size: int - writes: int - - -class _Writes: - __slots__ = ("count",) - - def __init__(self) -> None: - self.count: int = 0 - - def record(self) -> None: - self.count += 1 - - -_WRITES: Final = _Writes() - - -def model_cost_version() -> ModelCostVersion: - import litellm - - table: Final = litellm.model_cost - return ModelCostVersion(table=id(table), size=len(table), writes=_WRITES.count) - - -def publish(table: dict[str, object]) -> None: - """Makes ``table`` the live ``litellm.model_cost``.""" - import litellm - - litellm.model_cost = table - _WRITES.record() - - -def write(rows: dict[str, object], written: Mapping[str, object]) -> None: - """Writes one registration's rows into ``rows``, the published table.""" - rows.update(written) - _WRITES.record() - - -def changed_in_place() -> None: - """Moves the version and drops the name index after an in-place edit of ``litellm.model_cost``.""" - from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs - - model_cost_inputs.invalidate() - _WRITES.record() - - -class VersionCached(Generic[P, R]): - """Decorates ``function`` with an ``lru_cache`` keyed on the ``litellm.model_cost`` version each call read, and - exposes that cache's ``cache_clear`` and ``cache_info``.""" - - __slots__ = ("_at", "cache_clear", "cache_info") - - def __init__(self, function: Callable[P, R]) -> None: - def at( - _version: ModelCostVersion, - *args: P.args, - **kwargs: P.kwargs, # kwargs-ok: forwards a ParamSpec unchanged - ) -> R: - return function(*args, **kwargs) - - cached: Final = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(at) - self._at: Final = cached - self.cache_clear: Final = cached.cache_clear - self.cache_info: Final = cached.cache_info - - def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: # kwargs-ok: forwards a ParamSpec unchanged - return self._at(model_cost_version(), *args, **kwargs) - - -@dataclass(frozen=True, slots=True) -class Landing: - key: str - base: Mapping[str, object] - matched_builtin: bool - - -def published() -> dict[str, object]: - """The table published as ``litellm.model_cost`` now.""" - import litellm - - return litellm.model_cost +published: Final = model_cost_table.published +write: Final = model_cost_table.write def resolve( diff --git a/litellm/rust_bridge/preflight.py b/litellm/rust_bridge/host/preflight.py similarity index 100% rename from litellm/rust_bridge/preflight.py rename to litellm/rust_bridge/host/preflight.py diff --git a/litellm/rust_bridge/host/randomness.py b/litellm/rust_bridge/host/randomness.py new file mode 100644 index 00000000000..71d841833e4 --- /dev/null +++ b/litellm/rust_bridge/host/randomness.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +import random + + +def next_float() -> float: + return random.random() + + +def index_below(length: int) -> int: + return random.randrange(length) diff --git a/litellm/rust_bridge/host/routes/responses.py b/litellm/rust_bridge/host/routes/responses.py index 22e8cbbf3b0..0568326d18b 100644 --- a/litellm/rust_bridge/host/routes/responses.py +++ b/litellm/rust_bridge/host/routes/responses.py @@ -20,6 +20,12 @@ def connection_defaults(provider: str) -> tuple[str | None, str | None]: return litellm.api_key or provider_key, litellm.api_base +def original_previous_response_id(previous_response_id: str) -> str: + from litellm.responses.utils import ResponsesAPIRequestUtils + + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id) + + def response(value: Mapping[str, object]) -> ResponsesAPIResponse: return ResponsesAPIResponse.model_validate(value) diff --git a/litellm/rust_bridge/host/secret_manager.py b/litellm/rust_bridge/host/secret_manager.py index 4fb3350165c..b236b4ee19d 100644 --- a/litellm/rust_bridge/host/secret_manager.py +++ b/litellm/rust_bridge/host/secret_manager.py @@ -1,165 +1,58 @@ from __future__ import annotations import os -from collections.abc import Mapping -from dataclasses import dataclass, field -from importlib import import_module from typing import Final -from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +import httpx - -@dataclass(frozen=True, slots=True) -class NativeSecretManagerConfig: - system: str - environment: tuple[tuple[str, str], ...] = field(repr=False) - settings: Mapping[str, object] = field(repr=False) - enterprise_enabled: bool - owner_type: type[object] - environment_attributes: tuple[tuple[str, str], ...] - settings_attributes: tuple[str, ...] - methods: tuple[tuple[str, object], ...] = field(repr=False) - - -@dataclass(frozen=True, slots=True) -class _ClientAdapter: - system: KeyManagementSystem - module: str - name: str - methods: tuple[str, ...] - environment_attributes: tuple[tuple[str, str], ...] = () - settings_attributes: tuple[str, ...] = () - enterprise_enabled: bool = False - - -_ADAPTERS: Final = ( - _ClientAdapter( - KeyManagementSystem.AWS_SECRET_MANAGER, - "litellm.secret_managers.aws_secret_manager_v2", - "AWSSecretsManagerV2", - ("sync_read_secret", "async_read_secret"), - settings_attributes=( - "aws_region_name", - "aws_role_name", - "aws_session_name", - "aws_external_id", - "aws_profile_name", - "aws_web_identity_token", - "aws_sts_endpoint", - "replica_regions", - "kms_key_id", - ), - ), - _ClientAdapter( - KeyManagementSystem.HASHICORP_VAULT, - "litellm.secret_managers.hashicorp_secret_manager", - "HashicorpSecretManager", - ("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"), - environment_attributes=( - ("HCP_VAULT_ADDR", "vault_addr"), - ("HCP_VAULT_TOKEN", "vault_token"), - ("HCP_VAULT_NAMESPACE", "vault_namespace"), - ("HCP_VAULT_LOGIN_NAMESPACE", "login_namespace_override"), - ("HCP_VAULT_SECRET_NAMESPACE", "secret_namespace_override"), - ("HCP_VAULT_MOUNT_NAME", "vault_mount_name"), - ("HCP_VAULT_PATH_PREFIX", "vault_path_prefix"), - ("HCP_VAULT_CLIENT_CERT", "tls_cert_path"), - ("HCP_VAULT_CLIENT_KEY", "tls_key_path"), - ("HCP_VAULT_CERT_ROLE", "vault_cert_role"), - ("HCP_VAULT_APPROLE_ROLE_ID", "approle_role_id"), - ("HCP_VAULT_APPROLE_SECRET_ID", "approle_secret_id"), - ("HCP_VAULT_APPROLE_MOUNT_PATH", "approle_mount_path"), - ("HCP_VAULT_REFRESH_INTERVAL", "cache.default_ttl"), - ), - enterprise_enabled=True, - ), - _ClientAdapter( - KeyManagementSystem.CYBERARK, - "litellm.secret_managers.cyberark_secret_manager", - "CyberArkSecretManager", - ("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"), - environment_attributes=( - ("CYBERARK_API_BASE", "conjur_addr"), - ("CYBERARK_ACCOUNT", "conjur_account"), - ("CYBERARK_USERNAME", "conjur_username"), - ("CYBERARK_API_KEY", "conjur_api_key"), - ("CYBERARK_CLIENT_CERT", "tls_cert_path"), - ("CYBERARK_CLIENT_KEY", "tls_key_path"), - ("CYBERARK_SSL_VERIFY", "ssl_verify"), - ("CYBERARK_REFRESH_INTERVAL", "cache.default_ttl"), - ), - enterprise_enabled=True, - ), - _ClientAdapter( - KeyManagementSystem.GOOGLE_SECRET_MANAGER, - "litellm.secret_managers.google_secret_manager", - "GoogleSecretManager", - ("get_secret_from_google_secret_manager",), - environment_attributes=( - ("GOOGLE_SECRET_MANAGER_PROJECT_ID", "PROJECT_ID"), - ("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", "cache.default_ttl"), - ("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", "always_read_secret_manager"), - ), - enterprise_enabled=True, - ), -) - -_SDK_ADAPTERS: Final = ( - _ClientAdapter(KeyManagementSystem.AZURE_KEY_VAULT, "azure.keyvault.secrets", "SecretClient", ("get_secret",)), - _ClientAdapter(KeyManagementSystem.GOOGLE_KMS, "google.cloud.kms_v1", "KeyManagementServiceClient", ("decrypt",)), -) - - -def _capture(client: object, adapter: _ClientAdapter) -> NativeSecretManagerConfig: - prefixes: Final = ("AWS_", "AZURE_", "GOOGLE_", "VERTEX_", "GCS_", "HCP_VAULT_", "CYBERARK_") - config: Final = NativeSecretManagerConfig( - system=adapter.system.value, - environment=tuple( - (name, value) - for name, value in os.environ.items() - if name.startswith(prefixes) or name == "SECRET_MANAGER_REFRESH_INTERVAL" - ), - settings=KeyManagementSettings().model_dump(mode="json"), - enterprise_enabled=adapter.enterprise_enabled, - owner_type=type(client), - environment_attributes=adapter.environment_attributes, - settings_attributes=adapter.settings_attributes, - methods=tuple((name, getattr(type(client), name)) for name in adapter.methods), - ) - vars(client)["_litellm_native_secret_config"] = config - return config - - -def capture_secret_manager(client: object, system: str) -> None: - for adapter in (*_ADAPTERS, *_SDK_ADAPTERS): - if ( - adapter.system.value == system - and type(client).__name__ == adapter.name - and type(client) is getattr(import_module(adapter.module), adapter.name) - ): - _capture(client, adapter) - return - if ( - system == KeyManagementSystem.AWS_KMS.value - and type(client).__module__ == "botocore.client" - and type(client).__name__ == "KMS" - ): - _capture(client, _ClientAdapter(KeyManagementSystem.AWS_KMS, "botocore.client", "KMS", ("decrypt",))) +from litellm.secret_managers.native_config import NativeSecretManagerConfig, captured_config +from litellm.types.secret_managers.main import KeyManagementSettings def native_secret_manager_config(client: object) -> NativeSecretManagerConfig | None: - captured: Final = getattr(client, "_litellm_native_secret_config", None) - if isinstance(captured, NativeSecretManagerConfig): - return captured - for adapter in _ADAPTERS: - if type(client).__module__ == adapter.module and type(client) is getattr( - import_module(adapter.module), adapter.name - ): - return _capture(client, adapter) - return None + return captured_config(client) def read_secret_str(name: str) -> str | None: from litellm.secret_managers.main import get_secret_str return get_secret_str(name) + + +def read_from_manager( + client: object, + key_manager: str, + secret_name: str, + key_management_settings: KeyManagementSettings | None, +) -> str | None: + from litellm.secret_managers.secret_manager_handler import get_secret_from_manager + + return get_secret_from_manager( + client=client, + key_manager=key_manager, + secret_name=secret_name, + key_management_settings=key_management_settings, + ) + + +def timeout_message(message: str) -> str: + from litellm.exceptions import Timeout + + return str(Timeout(message=message, model="default-model-name", llm_provider="litellm-httpx-handler")) + + +def response_text(body: bytes) -> str: + return httpx.Response(200, content=body).text + + +def http_status_message(method: str, url: str, status: int) -> str: + response: Final = httpx.Response(status, request=httpx.Request(method, url)) + try: + _ = response.raise_for_status() + except httpx.HTTPStatusError as error: + return str(error) + return f"HTTP {status}" + + +def os_error_message(code: int) -> str: + return str(OSError(code, os.strerror(code))) diff --git a/litellm/rust_bridge/host/timeouts.py b/litellm/rust_bridge/host/timeouts.py index 5fde4397b55..fe9abd71c1e 100644 --- a/litellm/rust_bridge/host/timeouts.py +++ b/litellm/rust_bridge/host/timeouts.py @@ -1,13 +1,9 @@ -"""Shared timeout conversion for the native Rust bridges.""" - from __future__ import annotations import httpx +from litellm.rust_bridge import public_call + def timeout_to_seconds(timeout: float | httpx.Timeout | None) -> float | None: - if timeout is None: - return None - if isinstance(timeout, httpx.Timeout): - return timeout.read - return float(timeout) + return public_call.timeout_to_seconds(timeout) diff --git a/litellm/rust_bridge/host/tokenizers.py b/litellm/rust_bridge/host/tokenizers.py new file mode 100644 index 00000000000..e86e9b54aa7 --- /dev/null +++ b/litellm/rust_bridge/host/tokenizers.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from pathlib import Path + + +def tiktoken_resource_dir() -> str: + from litellm.litellm_core_utils import tokenizers + + return str(Path(tokenizers.__file__).parent) + + +def download_tokenizer(repo_id: str, revision: str, token: str | None) -> str: + import huggingface_hub + + return huggingface_hub.hf_hub_download( # pyright: ignore[reportUnknownMemberType] # untyped user_agent dict in its signature + repo_id=repo_id, filename="tokenizer.json", revision=revision, token=token + ) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index 1ab3d5e10c9..0572d498b7d 100644 --- a/litellm/rust_bridge/public_call.py +++ b/litellm/rust_bridge/public_call.py @@ -9,6 +9,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them +import httpx + from litellm.utils import ( _get_wrapper_timeout, # pyright: ignore[reportPrivateUsage] # the @client wrapper's own failure stamp ) @@ -49,6 +51,14 @@ def optional_sequence(value: object) -> Sequence[object] | None: return cast("Sequence[object]", value) # cast-ok: the same caller-owned object is handed on unchanged +def timeout_to_seconds(timeout: float | httpx.Timeout | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + def known_params(names: Collection[str], request: Mapping[str, object]) -> Mapping[str, object]: return {name: value for name, value in request.items() if name in names and value is not None} diff --git a/litellm/rust_bridge/responses/detectors.py b/litellm/rust_bridge/responses/detectors.py index 2fdb2c44b54..f8d5b4462a1 100644 --- a/litellm/rust_bridge/responses/detectors.py +++ b/litellm/rust_bridge/responses/detectors.py @@ -12,11 +12,12 @@ import litellm from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, call_callbacks, overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams if TYPE_CHECKING: from openai.types.responses import ResponseInputParam - from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponsesAPIOptionalRequestParams + from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS def client_headers(request: NativeCall) -> bool: @@ -160,9 +161,12 @@ _PARAMS: Final = TypeAdapter(dict[str, object]) def model_param_gating(request: NativeCall) -> bool: """Python clamps, drops or remaps Responses params according to the model's capabilities.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - from litellm.rust_bridge.host.routes.responses import PARAMETERS - params: Final = {name: request.resolved[name] for name in PARAMETERS if request.resolved.get(name) is not None} + params: Final = { + name: request.resolved[name] + for name in ResponsesAPIOptionalRequestParams.__annotations__ + if request.resolved.get(name) is not None + } mapped: Final = _PARAMS.validate_python( OpenAIResponsesAPIConfig().map_openai_params( # pyright: ignore[reportUnknownMemberType] # returns a bare dict cast("ResponsesAPIOptionalRequestParams", params), # cast-ok: the handler passes these kwargs diff --git a/litellm/rust_bridge/responses/websocket.py b/litellm/rust_bridge/responses/websocket.py index 83a1fa53376..e27e88fa9aa 100644 --- a/litellm/rust_bridge/responses/websocket.py +++ b/litellm/rust_bridge/responses/websocket.py @@ -10,7 +10,7 @@ import httpx from websockets.exceptions import ConnectionClosedOK from litellm.rust_bridge import _native -from litellm.rust_bridge.host.timeouts import timeout_to_seconds +from litellm.rust_bridge.public_call import timeout_to_seconds class RustResponsesWebSocket(Protocol): diff --git a/litellm/secret_managers/native_config.py b/litellm/secret_managers/native_config.py new file mode 100644 index 00000000000..c75fe94bae0 --- /dev/null +++ b/litellm/secret_managers/native_config.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass, field +from importlib import import_module +from typing import Final + +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem + + +@dataclass(frozen=True, slots=True) +class NativeSecretManagerConfig: + system: str + environment: tuple[tuple[str, str], ...] = field(repr=False) + settings: Mapping[str, object] = field(repr=False) + enterprise_enabled: bool + owner_type: type[object] + environment_attributes: tuple[tuple[str, str], ...] + settings_attributes: tuple[str, ...] + methods: tuple[tuple[str, object], ...] = field(repr=False) + + +@dataclass(frozen=True, slots=True) +class _ClientAdapter: + system: KeyManagementSystem + module: str + name: str + methods: tuple[str, ...] + environment_attributes: tuple[tuple[str, str], ...] = () + settings_attributes: tuple[str, ...] = () + enterprise_enabled: bool = False + + +_ADAPTERS: Final = ( + _ClientAdapter( + KeyManagementSystem.AWS_SECRET_MANAGER, + "litellm.secret_managers.aws_secret_manager_v2", + "AWSSecretsManagerV2", + ("sync_read_secret", "async_read_secret"), + settings_attributes=( + "aws_region_name", + "aws_role_name", + "aws_session_name", + "aws_external_id", + "aws_profile_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "replica_regions", + "kms_key_id", + ), + ), + _ClientAdapter( + KeyManagementSystem.HASHICORP_VAULT, + "litellm.secret_managers.hashicorp_secret_manager", + "HashicorpSecretManager", + ("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"), + environment_attributes=( + ("HCP_VAULT_ADDR", "vault_addr"), + ("HCP_VAULT_TOKEN", "vault_token"), + ("HCP_VAULT_NAMESPACE", "vault_namespace"), + ("HCP_VAULT_LOGIN_NAMESPACE", "login_namespace_override"), + ("HCP_VAULT_SECRET_NAMESPACE", "secret_namespace_override"), + ("HCP_VAULT_MOUNT_NAME", "vault_mount_name"), + ("HCP_VAULT_PATH_PREFIX", "vault_path_prefix"), + ("HCP_VAULT_CLIENT_CERT", "tls_cert_path"), + ("HCP_VAULT_CLIENT_KEY", "tls_key_path"), + ("HCP_VAULT_CERT_ROLE", "vault_cert_role"), + ("HCP_VAULT_APPROLE_ROLE_ID", "approle_role_id"), + ("HCP_VAULT_APPROLE_SECRET_ID", "approle_secret_id"), + ("HCP_VAULT_APPROLE_MOUNT_PATH", "approle_mount_path"), + ("HCP_VAULT_REFRESH_INTERVAL", "cache.default_ttl"), + ), + enterprise_enabled=True, + ), + _ClientAdapter( + KeyManagementSystem.CYBERARK, + "litellm.secret_managers.cyberark_secret_manager", + "CyberArkSecretManager", + ("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"), + environment_attributes=( + ("CYBERARK_API_BASE", "conjur_addr"), + ("CYBERARK_ACCOUNT", "conjur_account"), + ("CYBERARK_USERNAME", "conjur_username"), + ("CYBERARK_API_KEY", "conjur_api_key"), + ("CYBERARK_CLIENT_CERT", "tls_cert_path"), + ("CYBERARK_CLIENT_KEY", "tls_key_path"), + ("CYBERARK_SSL_VERIFY", "ssl_verify"), + ("CYBERARK_REFRESH_INTERVAL", "cache.default_ttl"), + ), + enterprise_enabled=True, + ), + _ClientAdapter( + KeyManagementSystem.GOOGLE_SECRET_MANAGER, + "litellm.secret_managers.google_secret_manager", + "GoogleSecretManager", + ("get_secret_from_google_secret_manager",), + environment_attributes=( + ("GOOGLE_SECRET_MANAGER_PROJECT_ID", "PROJECT_ID"), + ("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", "cache.default_ttl"), + ("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", "always_read_secret_manager"), + ), + enterprise_enabled=True, + ), +) + +_SDK_ADAPTERS: Final = ( + _ClientAdapter(KeyManagementSystem.AZURE_KEY_VAULT, "azure.keyvault.secrets", "SecretClient", ("get_secret",)), + _ClientAdapter(KeyManagementSystem.GOOGLE_KMS, "google.cloud.kms_v1", "KeyManagementServiceClient", ("decrypt",)), +) + + +def _capture(client: object, adapter: _ClientAdapter) -> NativeSecretManagerConfig: + prefixes: Final = ("AWS_", "AZURE_", "GOOGLE_", "VERTEX_", "GCS_", "HCP_VAULT_", "CYBERARK_") + config: Final = NativeSecretManagerConfig( + system=adapter.system.value, + environment=tuple( + (name, value) + for name, value in os.environ.items() + if name.startswith(prefixes) or name == "SECRET_MANAGER_REFRESH_INTERVAL" + ), + settings=KeyManagementSettings().model_dump(mode="json"), + enterprise_enabled=adapter.enterprise_enabled, + owner_type=type(client), + environment_attributes=adapter.environment_attributes, + settings_attributes=adapter.settings_attributes, + methods=tuple((name, getattr(type(client), name)) for name in adapter.methods), + ) + vars(client)["_litellm_native_secret_config"] = config + return config + + +def capture_secret_manager(client: object, system: str) -> None: + for adapter in (*_ADAPTERS, *_SDK_ADAPTERS): + if ( + adapter.system.value == system + and type(client).__name__ == adapter.name + and type(client) is getattr(import_module(adapter.module), adapter.name) + ): + _capture(client, adapter) + return + if ( + system == KeyManagementSystem.AWS_KMS.value + and type(client).__module__ == "botocore.client" + and type(client).__name__ == "KMS" + ): + _capture(client, _ClientAdapter(KeyManagementSystem.AWS_KMS, "botocore.client", "KMS", ("decrypt",))) + + +def captured_config(client: object) -> NativeSecretManagerConfig | None: + captured: Final = getattr(client, "_litellm_native_secret_config", None) + if isinstance(captured, NativeSecretManagerConfig): + return captured + for adapter in _ADAPTERS: + if type(client).__module__ == adapter.module and type(client) is getattr( + import_module(adapter.module), adapter.name + ): + return _capture(client, adapter) + return None diff --git a/litellm/utils.py b/litellm/utils.py index 4a4f2e8edfe..c6ade2603c4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -107,6 +107,13 @@ from litellm.litellm_core_utils.get_model_cost_map import ( fetch_registration_document, model_cost_inputs, ) +from litellm.litellm_core_utils.model_cost_table import ( + AliasedCostEntry, + Landing, + VersionCached, + changed_in_place, + model_cost_version, +) from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode @@ -121,13 +128,6 @@ from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge._native import NativeModelCostMap, NativeRegistered from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision -from litellm.rust_bridge.host.model_cost import ( - AliasedCostEntry, - Landing, - VersionCached, - changed_in_place, - model_cost_version, -) _CachingHandlerResponse = None _LLMCachingHandler = None diff --git a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py index 921552455b7..0aa36635c65 100644 --- a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py +++ b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py @@ -10,7 +10,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm import ModelResponse, Router -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish _OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" _JSON_TOOL: Final = "json_tool_call" diff --git a/tests/integration/sdk/test_model_cost_fetch_sdk.py b/tests/integration/sdk/test_model_cost_fetch_sdk.py index 206ffa99909..fadd8386d17 100644 --- a/tests/integration/sdk/test_model_cost_fetch_sdk.py +++ b/tests/integration/sdk/test_model_cost_fetch_sdk.py @@ -28,7 +28,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ( model_cost_inputs, refetch_model_cost_map, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index be0c5e311da..16ad2235a57 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -28,7 +28,7 @@ from litellm.llms.custom_httpx.http_handler import( HTTPHandler, ) from litellm.types.utils import LlmProviders -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index fe21cb23caf..88784814ba4 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -134,7 +134,7 @@ os.environ["PYTHON_DOTENV_DISABLED"] = "1" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import -from litellm.rust_bridge.host.model_cost import publish # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.model_cost_table import publish # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency from litellm.anthropic_beta_headers_manager import ( reload_beta_headers_config, # noqa: E402 # same import-time dependency diff --git a/tests/unit/integrations/arize/test_arize.py b/tests/unit/integrations/arize/test_arize.py index 42dd5e01ce9..331e0503db0 100644 --- a/tests/unit/integrations/arize/test_arize.py +++ b/tests/unit/integrations/arize/test_arize.py @@ -16,7 +16,7 @@ from litellm.integrations.arize.arize import ArizeConfig, ArizeLogger from litellm.integrations.opentelemetry import OpenTelemetryConfig from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py index 73e6b62ec4c..08aa400c339 100644 --- a/tests/unit/integrations/arize/test_arize_phoenix.py +++ b/tests/unit/integrations/arize/test_arize_phoenix.py @@ -17,7 +17,7 @@ from litellm.integrations.arize.arize_phoenix import ( ) from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/integrations/opik/test_opik.py b/tests/unit/integrations/opik/test_opik.py index 0bf35ffdabf..30d6fffee78 100644 --- a/tests/unit/integrations/opik/test_opik.py +++ b/tests/unit/integrations/opik/test_opik.py @@ -12,7 +12,7 @@ import pytest import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 4b6f59a78ab..c6f99a36ec4 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -19,7 +19,7 @@ from openai.types.chat import ChatCompletionMessage from pydantic import BaseModel, ConfigDict import litellm -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from litellm import Router from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( diff --git a/tests/unit/integrations/test_helicone.py b/tests/unit/integrations/test_helicone.py index 90fc2f97ee8..f979090dc60 100644 --- a/tests/unit/integrations/test_helicone.py +++ b/tests/unit/integrations/test_helicone.py @@ -7,7 +7,7 @@ from collections.abc import Iterator from pathlib import Path from typing import Final from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import MagicMock diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index fafc098ed5d..01b94f01891 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -86,7 +86,7 @@ from litellm.utils import ( function_setup, validate_and_fix_openai_messages, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index ea270fbf7c8..3eb87947f0c 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -21,7 +21,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.types.router import LiteLLM_Params -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import patch diff --git a/tests/unit/litellm_core_utils/test_model_cost_table.py b/tests/unit/litellm_core_utils/test_model_cost_table.py new file mode 100644 index 00000000000..1873dab8129 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_model_cost_table.py @@ -0,0 +1,57 @@ +from typing import Final + +import pytest + +import litellm +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs +from litellm.litellm_core_utils.model_cost_table import changed_in_place, publish, write + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +_MODEL: Final = "lit-table-model" + + +def _row(price: float) -> dict[str, object]: + return {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": price, "output_cost_per_token": price} + + +def _billed() -> float: + prompt, completion = cost_per_token( + model=_MODEL, custom_llm_provider="openai", prompt_tokens=1, completion_tokens=0 + ) + return prompt + completion + + +def test_a_published_table_is_read_in_place_of_what_was_cached_for_the_one_before() -> None: + publish({**litellm.model_cost, _MODEL: _row(1.0)}) + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 1.0 + assert _billed() == 1.0 + + publish({**litellm.model_cost, _MODEL: _row(2.0)}) + + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 2.0 + assert _billed() == 2.0 + + +def test_rows_written_into_the_live_table_replace_a_cached_value_of_the_same_size() -> None: + publish({**litellm.model_cost, _MODEL: _row(1.0)}) + assert _billed() == 1.0 + + write(litellm.model_cost, {_MODEL: _row(3.0)}) + + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 3.0 + assert _billed() == 3.0 + + +def test_a_provider_changed_in_place_regroups_the_key_once_the_change_is_reported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") + + monkeypatch.setitem(litellm.model_cost["gpt-4o"], "litellm_provider", "lit-moved") + changed_in_place() + + assert model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") == ["gpt-4o"] + assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "openai") diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index b4bf1120935..f6300b2e70c 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -32,7 +32,7 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, StreamingChoices, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish import unittest.mock diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index 2cb8b9cf6c5..cec8c07ca1a 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -27,7 +27,7 @@ from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5 from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig from litellm.router import Router from litellm.utils import get_optional_params -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome _MAPPED_PARAMS: Final = TypeAdapter(dict[str, object]) diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 5ead4ad3986..90a8939e2dd 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -16,7 +16,7 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import ( ) from litellm.types.utils import ImageObject, ImageResponse from litellm.utils import get_optional_params_image_gen -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish @pytest.fixture(autouse=True) diff --git a/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py index 927833b530d..9dc29e9860d 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -19,7 +19,7 @@ def use_local_model_cost_map(): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") import litellm - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish original_model_cost = litellm.model_cost litellm.model_cost = json.loads( diff --git a/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py index e23be01df55..be5e3a5ce18 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py @@ -14,7 +14,7 @@ def use_local_model_cost_map(): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") import litellm - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish original_model_cost = litellm.model_cost litellm.model_cost = json.loads( diff --git a/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py index d17c2569b92..690350223ea 100644 --- a/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py +++ b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py @@ -27,7 +27,7 @@ def ensure_nova_canvas_image_edit_model_cost_flags(monkeypatch): Full ``model_prices_and_context_window.json`` includes these flags, but CI or alternate cost maps may omit them—merge minimal entries so tests match production. """ - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish for key in ( "amazon.nova-canvas-v1:0", diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index b0071287746..4e031531180 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -42,7 +42,7 @@ from litellm.types.llms.custom_http import VerifyTypes from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.exceptions import Timeout as LitellmTimeout from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index bdeff226af2..1936fc1d2df 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -42,7 +42,7 @@ from litellm.types.utils import ( PromptTokensDetails, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from openai.types.beta.assistant import Assistant from openai.types.beta.assistant_deleted import AssistantDeleted from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/llms/parallel_ai/test_parallel_ai_search.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py index 1400fcc7d4e..d067aa49352 100644 --- a/tests/unit/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py @@ -57,7 +57,7 @@ def bundled_cost_map(monkeypatch): litellm caches model-info lookups, so swapping ``model_cost`` only takes effect once those caches are invalidated -- on the way in and back out. """ - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) publish(litellm.model_cost) diff --git a/tests/unit/llms/test_custom_llm.py b/tests/unit/llms/test_custom_llm.py index e731d79ed99..a9668ebb485 100644 --- a/tests/unit/llms/test_custom_llm.py +++ b/tests/unit/llms/test_custom_llm.py @@ -36,7 +36,7 @@ from litellm.types.utils import ( StreamingChoices, ) from litellm.utils import ModelResponseIterator -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 8c2cb643943..531c5d0b1e6 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -37,7 +37,7 @@ from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage from litellm.utils import CustomStreamWrapper -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger from litellm.llms.vertex_ai.gemini.transformation import ( diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index feda96c46b7..bd5941b880e 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -1086,7 +1086,7 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None: from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache from litellm.proxy.utils import litellm_config_cache - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish litellm_config_cache.flush_cache() prisma_client = _reload_config_prisma_client() diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 588411cd57a..8423dee9886 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -34,7 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from fastapi import Request as Request_http_parsing from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from starlette.types import Message from starlette.websockets import WebSocket from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index 57839d65cfb..8961c0a1a34 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -15,7 +15,7 @@ import pytest import litellm from litellm import Router from litellm._logging import verbose_logger -from litellm.rust_bridge.host.model_cost import AliasedCostEntry +from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry from litellm.proxy.common_utils.model_deprecation import ( _classify, _parse_deprecation_date, diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py index c24d04b3a8b..87bb193dfbf 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py @@ -14,7 +14,7 @@ from litellm import DualCache as DualCache_dynamic_rate, Router from litellm._uuid import uuid from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import UserAPIKeyAuth -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from typing import Optional diff --git a/tests/unit/proxy/hooks/test_prompt_injection_detection.py b/tests/unit/proxy/hooks/test_prompt_injection_detection.py index 1b50de938b6..1db5e9e6e18 100644 --- a/tests/unit/proxy/hooks/test_prompt_injection_detection.py +++ b/tests/unit/proxy/hooks/test_prompt_injection_detection.py @@ -17,7 +17,7 @@ from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm import Router as Router_prompt_injection from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py b/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py index 0b488618636..1e5bb5a8531 100644 --- a/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py +++ b/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py @@ -9,7 +9,7 @@ import os import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from litellm.proxy._types import LitellmUserRoles diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index a988250159f..f0459e7b927 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -27,7 +27,7 @@ from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server from litellm.proxy._types import LiteLLM_ProxyModelTable, UserAPIKeyAuth -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from .conftest import normalize, route_entry_points_to_stubs # type: ignore[import-not-found] diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 551c4b76766..cdbdb70e3ac 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -17,7 +17,7 @@ import pytest import respx import litellm -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from litellm.litellm_core_utils import get_llm_provider_logic from litellm.proxy import proxy_server diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9bbdc7165dd..5b423421458 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -59,7 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY from litellm.redis.stores.connection import connect diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index e3eb9961228..94f60fa5b30 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -13,7 +13,7 @@ from litellm.responses.litellm_completion_transformation.transformation import L from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse from litellm.types.utils import ModelInfo, Usage -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish class TestResponsesAPIRequestUtils: diff --git a/tests/unit/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py index f71d1ad9144..e23b7628c35 100644 --- a/tests/unit/router_strategy/test_lowest_cost.py +++ b/tests/unit/router_strategy/test_lowest_cost.py @@ -7,7 +7,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome DEPLOYMENT_ID = "9876" diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index 1139de6c0fc..60cad628598 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -11,7 +11,7 @@ import traceback from typing import Dict, Final, Iterator from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import patch diff --git a/tests/unit/rust_bridge/host/test_model_cost.py b/tests/unit/rust_bridge/host/test_model_cost.py index 82a52a941a1..66f411da83e 100644 --- a/tests/unit/rust_bridge/host/test_model_cost.py +++ b/tests/unit/rust_bridge/host/test_model_cost.py @@ -3,16 +3,8 @@ from typing import Final import pytest import litellm -from litellm.cost_calculator import cost_per_token -from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs -from litellm.rust_bridge.host.model_cost import ( - changed_in_place, - entry, - merged, - publish, - resolve, - write, -) +from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry +from litellm.rust_bridge.host.model_cost import merged, resolve pytestmark = pytest.mark.usefixtures("local_model_cost_map") @@ -25,56 +17,9 @@ def test_resolve_lands_on_the_row_a_lookup_in_the_published_table_resolves() -> def test_merged_keeps_the_alias_names_and_adds_only_the_fields_a_registration_set() -> None: - current: Final = entry({"mode": "chat", "x": 1}, ("owner", "alias"), frozenset({"mode"})) + current: Final = AliasedCostEntry({"mode": "chat", "x": 1}, ("owner", "alias"), frozenset({"mode"})) row: Final = merged(current, {"x": 2}, {"litellm_provider": "openai", "output_cost_per_token": None}) assert row == {"mode": "chat", "x": 2, "litellm_provider": "openai"} assert (row.names, row.registered_fields) == (("owner", "alias"), frozenset({"mode", "litellm_provider"})) - - -_MODEL: Final = "lit-table-model" - - -def _row(price: float) -> dict[str, object]: - return {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": price, "output_cost_per_token": price} - - -def _billed() -> float: - prompt, completion = cost_per_token( - model=_MODEL, custom_llm_provider="openai", prompt_tokens=1, completion_tokens=0 - ) - return prompt + completion - - -def test_a_published_table_is_read_in_place_of_what_was_cached_for_the_one_before() -> None: - publish({**litellm.model_cost, _MODEL: _row(1.0)}) - assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 1.0 - assert _billed() == 1.0 - - publish({**litellm.model_cost, _MODEL: _row(2.0)}) - - assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 2.0 - assert _billed() == 2.0 - - -def test_rows_written_into_the_live_table_replace_a_cached_value_of_the_same_size() -> None: - publish({**litellm.model_cost, _MODEL: _row(1.0)}) - assert _billed() == 1.0 - - write(litellm.model_cost, {_MODEL: _row(3.0)}) - - assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 3.0 - assert _billed() == 3.0 - - -def test_a_provider_changed_in_place_regroups_the_key_once_the_change_is_reported( - monkeypatch: pytest.MonkeyPatch, -) -> None: - assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") - - monkeypatch.setitem(litellm.model_cost["gpt-4o"], "litellm_provider", "lit-moved") - changed_in_place() - - assert model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") == ["gpt-4o"] - assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "openai") diff --git a/tests/unit/rust_bridge/test_preflight.py b/tests/unit/rust_bridge/host/test_preflight.py similarity index 92% rename from tests/unit/rust_bridge/test_preflight.py rename to tests/unit/rust_bridge/host/test_preflight.py index a8b1a40a00b..ebf8638f72e 100644 --- a/tests/unit/rust_bridge/test_preflight.py +++ b/tests/unit/rust_bridge/host/test_preflight.py @@ -8,8 +8,8 @@ import pytest from pydantic import TypeAdapter import litellm -from litellm.rust_bridge import preflight -from litellm.rust_bridge.preflight import check_limits +from litellm.rust_bridge.host import preflight +from litellm.rust_bridge.host.preflight import check_limits @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @@ -47,7 +47,7 @@ def test_check_limits_refuses_a_call_over_the_budget(monkeypatch: pytest.MonkeyP check_limits({"model": "mistral/mistral-ocr-latest"}) -CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/preflight_contract.json" +CONTRACT_PATH: Final = Path(__file__).parents[4] / "litellm-rust/crates/python-bridge/preflight_contract.json" _SHIMS: Final[MappingProxyType[str, Callable[..., object]]] = MappingProxyType( { "credential_list": preflight.credential_list, diff --git a/tests/unit/rust_bridge/host/test_secret_manager.py b/tests/unit/rust_bridge/host/test_secret_manager.py index 1f8acc5792e..0705b3c806d 100644 --- a/tests/unit/rust_bridge/host/test_secret_manager.py +++ b/tests/unit/rust_bridge/host/test_secret_manager.py @@ -18,7 +18,8 @@ from pydantic import JsonValue import litellm from litellm.rust_bridge import _native as native -from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config +from litellm.rust_bridge.host.secret_manager import native_secret_manager_config +from litellm.secret_managers.native_config import capture_secret_manager 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_logger.py b/tests/unit/rust_bridge/test_logger.py index df3aedf4cb9..97a32609538 100644 --- a/tests/unit/rust_bridge/test_logger.py +++ b/tests/unit/rust_bridge/test_logger.py @@ -18,7 +18,8 @@ from litellm.litellm_core_utils.secret_redaction import ( python_redact_string, python_redact_structured_value, ) -from litellm.rust_bridge import diagnostics, logger +from litellm.rust_bridge import diagnostics +from litellm.rust_bridge.host import logger def test_native_records_preserve_metadata_and_redact_before_custom_handlers( diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 8af9ae29c67..07b6111aa67 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -39,7 +39,7 @@ from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject from litellm.utils import cached_get_model_info_helper, supports_prompt_caching -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish import os import threading from unittest.mock import MagicMock, patch diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index b300562b029..c4879587038 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5653,7 +5653,7 @@ def _priced_at(prompt_tokens, completion_tokens): def local_cost_map(monkeypatch): """The prices these tests assert are the checked-in ones. Setting the environment variable alone does not reload the map, so pin the map itself.""" - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) diff --git a/tests/unit/test_model_cost_aliases.py b/tests/unit/test_model_cost_aliases.py index c01e4d07e3e..3dfee102004 100644 --- a/tests/unit/test_model_cost_aliases.py +++ b/tests/unit/test_model_cost_aliases.py @@ -12,7 +12,7 @@ from typing import Final import pytest from litellm.rust_bridge._native import NativeInstalled, NativeModelCostMap, NativeRejected -from litellm.rust_bridge.host.model_cost import AliasedCostEntry +from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry def _install(document: dict[str, dict[str, object]]) -> NativeInstalled: diff --git a/tests/unit/test_no_top_level_test_invocations.py b/tests/unit/test_no_top_level_test_invocations.py index 386c71d779e..eae026bbf71 100644 --- a/tests/unit/test_no_top_level_test_invocations.py +++ b/tests/unit/test_no_top_level_test_invocations.py @@ -8,7 +8,7 @@ import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome LOCAL_TESTING_DIR = Path(__file__).parents[1] / "local_testing" diff --git a/tests/unit/test_pydantic_namespaces.py b/tests/unit/test_pydantic_namespaces.py index fc21ea501db..3a4c1459d14 100644 --- a/tests/unit/test_pydantic_namespaces.py +++ b/tests/unit/test_pydantic_namespaces.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index d4f591e7291..969ca0269cf 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -28,7 +28,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ( model_cost_inputs, ) from litellm.rust_bridge._native import NativeTableSwapped -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from litellm.types.utils import ( Choices, CompletionTokensDetailsWrapper, @@ -258,7 +258,7 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa litellm.model_cost.pop(registered_key, None) litellm.model_cost = original_model_cost os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish publish(litellm.model_cost) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 6557db7be32..e84d187f937 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -101,7 +101,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) from litellm.types.utils import CallTypes, CredentialItem, Usage -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from tests.large_text import text diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 8d15ba02a38..8bc201b030a 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -48,7 +48,7 @@ from litellm.types.router import ( from litellm.utils import ( reapply_runtime_model_cost_registrations, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish _COPILOT_TOKEN_EXPIRES_AT: Final = 4_102_444_800 diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index 5856caf7740..832b73a0730 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -21,7 +21,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 8d8c123ef49..6fece4d8c66 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -120,7 +120,7 @@ from litellm.utils import ( validate_environment, validate_chat_completion_tool_choice, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish @pytest.mark.parametrize("root_base_model", [42, ["gemini-3.5-flash"], {"model": "gemini-3.5-flash"}]) @@ -2529,7 +2529,7 @@ def test_register_model_with_scientific_notation(): test_model_name = f"test-scientific-notation-model-{uuid.uuid4().hex[:12]}" # Clear LRU caches that might have stale data - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish publish(litellm.model_cost) model_cost_dict = { @@ -5173,7 +5173,7 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( from litellm.utils import ( reapply_runtime_model_cost_registrations, ) - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish # Only the recorded half is under test here; the live-router rebuild is covered # in test_router_model_cost_isolation.py. Routers built by earlier tests in this @@ -5230,7 +5230,7 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey from litellm.utils import ( reapply_runtime_model_cost_registrations, ) - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish saved_model_cost = litellm.model_cost try: @@ -5257,7 +5257,7 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish class _EntryRegisteredMidWalk(dict): def get(self, field: str, default: object = None) -> object: @@ -5306,7 +5306,7 @@ def test_request_path_cost_map_walks_survive_a_registration_mid_walk( from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.ollama.common_utils import OllamaModelInfo from litellm.llms.openai.chat.gpt_reasoning_transformation import _catalogue_declares_default_effort - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish walks: Final[Mapping[str, Callable[[], object]]] = { "fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"), @@ -8537,7 +8537,7 @@ def test_get_model_info_gemini(monkeypatch): def test_register_model_without_a_provider_never_inherits_another_providers_entry( monkeypatch: pytest.MonkeyPatch, ) -> None: - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish key: Final = "azure/claude-sonnet-4-6" monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) @@ -8570,7 +8570,7 @@ def test_register_model_without_a_provider_never_inherits_another_providers_entr def test_register_model_over_a_builtin_row_bills_explicit_cache_hits( monkeypatch: pytest.MonkeyPatch, registered_prices: dict[str, float], billed_cache_read_rate: float ) -> None: - from litellm.rust_bridge.host.model_cost import publish + from litellm.litellm_core_utils.model_cost_table import publish key: Final = "dashscope/qwen-register-explicit-cache-test" monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) @@ -11659,7 +11659,7 @@ def test_the_catalog_with_register_model_entries_lists_each_model_once_and_leave ) def test_runtime_model_cost_registrations_resolve_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None: from litellm import utils as model_utils - from litellm.rust_bridge.host.model_cost import AliasedCostEntry + from litellm.litellm_core_utils.model_cost_table import AliasedCostEntry rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} alias: Final = "openai.test-converse-alias" diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index cde0a8978c0..0ebe291cf62 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -14,7 +14,7 @@ from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.utils import ModelInfoBase from litellm.utils import supports_native_structured_output -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish import os from unittest.mock import MagicMock, patch diff --git a/tests/unit/types/test_files.py b/tests/unit/types/test_files.py index 5b2c82b69db..b55dde3d44b 100644 --- a/tests/unit/types/test_files.py +++ b/tests/unit/types/test_files.py @@ -16,7 +16,7 @@ from litellm.types.files import ( get_file_mime_type_from_extension, get_file_type_from_extension, ) -from litellm.rust_bridge.host.model_cost import publish +from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome