mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(rust): read OCR secrets from the process environment and decline when a secret manager is readable
The OCR route called back into Python's get_secret_str for every env fallback. With no secret manager configured that is os.environ behind a GIL hop, and with one configured it blocked a tokio worker on vault I/O and also sent the Azure and GCP identity variables, which Python reads with os.getenv, to the vault. The other Rust routes already read the process environment. Read the process environment here too. When litellm would read secrets from a secret manager, decline the Rust route so the Python route serves the call with the vault-backed keys
This commit is contained in:
parent
b341d21a76
commit
1669213eb5
5 changed files with 100 additions and 86 deletions
|
|
@ -19,5 +19,8 @@
|
|||
"vertex_project",
|
||||
"vertex_location",
|
||||
"enable_azure_ad_token_refresh"
|
||||
],
|
||||
"secret_manager": [
|
||||
"readable"
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
use litellm_core_utils::settings::Lookup;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.settings";
|
||||
|
|
@ -8,17 +7,24 @@ pub(crate) enum PythonSettings {
|
|||
Http,
|
||||
UrlPolicy,
|
||||
ProviderDefaults,
|
||||
SecretManager,
|
||||
}
|
||||
|
||||
impl PythonSettings {
|
||||
#[cfg(test)]
|
||||
pub(crate) const ALL: [Self; 3] = [Self::Http, Self::UrlPolicy, Self::ProviderDefaults];
|
||||
pub(crate) const ALL: [Self; 4] = [
|
||||
Self::Http,
|
||||
Self::UrlPolicy,
|
||||
Self::ProviderDefaults,
|
||||
Self::SecretManager,
|
||||
];
|
||||
|
||||
pub(crate) fn name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Http => "http_settings",
|
||||
Self::UrlPolicy => "url_policy",
|
||||
Self::ProviderDefaults => "provider_defaults",
|
||||
Self::SecretManager => "secret_manager",
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -32,23 +38,6 @@ impl PythonSettings {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) struct PythonSecrets;
|
||||
|
||||
impl Lookup for PythonSecrets {
|
||||
fn get(&self, name: &str) -> Option<String> {
|
||||
Python::attach(|py| {
|
||||
py.import(MODULE)
|
||||
.and_then(|module| module.getattr("secret")?.call1((name,)))
|
||||
.and_then(|value| value.extract::<Option<String>>())
|
||||
.unwrap_or_else(|error| {
|
||||
let _ =
|
||||
PythonSettings::warn(py, &format!("reading secret {name} failed: {error}"));
|
||||
None
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
|
||||
|
||||
|
|
@ -56,10 +45,9 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
|
|||
mod tests {
|
||||
use std::{collections::BTreeSet, ffi::CString};
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{CONTRACT, PythonSecrets, PythonSettings};
|
||||
use super::{CONTRACT, PythonSettings};
|
||||
|
||||
#[test]
|
||||
fn every_settings_group_is_in_the_python_contract() {
|
||||
|
|
@ -83,48 +71,4 @@ mod tests {
|
|||
assert_eq!(read, declared);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
py.run(
|
||||
c"
|
||||
import sys
|
||||
import types
|
||||
settings = types.ModuleType('litellm.rust_bridge.settings')
|
||||
settings.warnings = []
|
||||
def secret(name):
|
||||
if name == 'BROKEN':
|
||||
raise RuntimeError('vault down')
|
||||
return {'MISTRAL_API_KEY': 'from-vault'}.get(name)
|
||||
settings.secret = secret
|
||||
settings.warn = settings.warnings.append
|
||||
sys.modules.setdefault('litellm', types.ModuleType('litellm'))
|
||||
sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge'))
|
||||
sys.modules['litellm.rust_bridge.settings'] = settings
|
||||
",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
assert_eq!(
|
||||
PythonSecrets.get("MISTRAL_API_KEY").as_deref(),
|
||||
Some("from-vault")
|
||||
);
|
||||
assert_eq!(PythonSecrets.get("ABSENT"), None);
|
||||
assert_eq!(PythonSecrets.get("BROKEN"), None);
|
||||
Python::attach(|py| {
|
||||
let warnings: Vec<String> = py
|
||||
.import("litellm.rust_bridge.settings")
|
||||
.unwrap()
|
||||
.getattr("warnings")
|
||||
.unwrap()
|
||||
.extract()
|
||||
.unwrap();
|
||||
assert_eq!(warnings.len(), 1);
|
||||
assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down"));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,17 +10,16 @@ use litellm_auth_gcp::VertexAuth;
|
|||
use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call};
|
||||
use litellm_core::ocr::route::ocr_machine;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings};
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
handler::OcrClient,
|
||||
settings::{OcrSettings, Secrets},
|
||||
};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
errors::RustBridgeDeclined,
|
||||
http,
|
||||
python_settings::{PythonSecrets, PythonSettings},
|
||||
};
|
||||
use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings};
|
||||
|
||||
const SURFACE: LegacySurface = LegacySurface {
|
||||
call_type: "ocr",
|
||||
|
|
@ -42,6 +41,7 @@ fn run_ocr(
|
|||
kwargs: Bound<'_, PyDict>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let secrets = process_environment_secrets(&PythonSettings::SecretManager.read(py)?)?;
|
||||
let config = http::call_config(py, &kwargs, asynchronous)?;
|
||||
let client = OcrClient::new(
|
||||
http::pool(),
|
||||
|
|
@ -49,7 +49,7 @@ fn run_ocr(
|
|||
http::url_policy(py)?,
|
||||
VERTEX_AUTH.clone(),
|
||||
ocr_settings(py)?,
|
||||
Arc::new(PythonSecrets),
|
||||
secrets,
|
||||
)
|
||||
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
|
||||
run_legacy_call(
|
||||
|
|
@ -62,6 +62,21 @@ fn run_ocr(
|
|||
)
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
struct PythonSecretManager {
|
||||
readable: bool,
|
||||
}
|
||||
|
||||
fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Secrets> {
|
||||
let manager: PythonSecretManager = secret_manager.extract()?;
|
||||
if manager.readable {
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"a readable secret manager is configured and the Rust route only reads the process environment",
|
||||
));
|
||||
}
|
||||
Ok(Arc::new(ProcessEnvironment))
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
struct PythonProviderDefaults {
|
||||
vertex_project: Option<String>,
|
||||
|
|
@ -105,3 +120,47 @@ pub(crate) fn aocr(
|
|||
) -> PyResult<Py<PyAny>> {
|
||||
run_ocr(py, request, args, kwargs, true)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::process_environment_secrets;
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("readable", readable).unwrap();
|
||||
py.run(
|
||||
c"import types\nmanager = types.SimpleNamespace(readable=readable)",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
locals.get_item("manager").unwrap().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readable_secret_manager_sends_the_call_back_to_python() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let declined = process_environment_secrets(&secret_manager(py, true))
|
||||
.err()
|
||||
.expect("the Rust route declines");
|
||||
assert!(declined.is_instance_of::<RustBridgeDeclined>(py));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_a_readable_secret_manager_secrets_are_the_process_environment() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap();
|
||||
assert_eq!(
|
||||
secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"),
|
||||
None
|
||||
);
|
||||
assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,16 +31,21 @@ class ProviderDefaults:
|
|||
enable_azure_ad_token_refresh: bool | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManager:
|
||||
readable: bool
|
||||
|
||||
|
||||
def warn(message: str) -> None:
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
verbose_logger.warning("%s", message)
|
||||
|
||||
|
||||
def secret(name: str) -> str | None:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
def secret_manager() -> SecretManager:
|
||||
from litellm.secret_managers.main import _should_read_secret_from_secret_manager
|
||||
|
||||
return get_secret_str(name)
|
||||
return SecretManager(readable=_should_read_secret_from_secret_manager())
|
||||
|
||||
|
||||
def provider_defaults() -> ProviderDefaults:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.custom_httpx.http_handler import default_user_agent
|
||||
from litellm.rust_bridge import settings
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
|
||||
CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json"
|
||||
|
|
@ -23,6 +24,7 @@ def test_the_rust_contract_matches_the_returned_fields() -> None:
|
|||
"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())],
|
||||
"url_policy": [field.name for field in dataclasses.fields(settings.url_policy())],
|
||||
"provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())],
|
||||
"secret_manager": [field.name for field in dataclasses.fields(settings.secret_manager())],
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -101,25 +103,26 @@ class _VaultSecrets(CustomSecretManager):
|
|||
return self.secrets.get(secret_name)
|
||||
|
||||
|
||||
def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
@pytest.mark.parametrize(
|
||||
("access_mode", "readable"),
|
||||
[("read_only", True), ("read_and_write", True), ("write_only", False)],
|
||||
)
|
||||
def test_secret_manager_is_readable_only_when_litellm_would_read_secrets_from_it(
|
||||
monkeypatch: pytest.MonkeyPatch, access_mode: str, readable: bool
|
||||
) -> None:
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "stale-env-key")
|
||||
monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key")
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "env-key")
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"}))
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
|
||||
|
||||
assert settings.secret("MISTRAL_API_KEY") == "vault-key"
|
||||
assert settings.secret("REDUCTO_API_KEY") == "env-only-key"
|
||||
assert settings.secret("ABSENT_KEY") is None
|
||||
assert settings.secret_manager() == settings.SecretManager(readable=readable)
|
||||
assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable
|
||||
|
||||
|
||||
def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "env-key")
|
||||
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", None)
|
||||
|
||||
assert settings.secret("MISTRAL_API_KEY") == "env-key"
|
||||
assert settings.secret_manager() == settings.SecretManager(readable=False)
|
||||
|
||||
|
||||
def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue