mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 4351906ca7 into df78b16941
This commit is contained in:
commit
5f39586cf2
23 changed files with 160 additions and 1995 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -4279,7 +4279,6 @@ dependencies = [
|
|||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-sdk-secretsmanager",
|
||||
"bytes",
|
||||
"criterion",
|
||||
"fancy-regex 0.19.2",
|
||||
|
|
@ -4306,7 +4305,6 @@ dependencies = [
|
|||
"litellm-router-types",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
"litellm-secrets-types",
|
||||
"litellm-spend-clickhouse",
|
||||
"litellm-storage-clickhouse",
|
||||
|
|
|
|||
|
|
@ -63,7 +63,6 @@ tokio = { workspace = true, features = ["rt", "sync"] }
|
|||
[dev-dependencies]
|
||||
tracing.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-secrets-aws.workspace = true
|
||||
serde.workspace = true
|
||||
serde_with.workspace = true
|
||||
criterion.workspace = true
|
||||
|
|
@ -73,7 +72,6 @@ sqlx = { workspace = true, features = ["migrate"] }
|
|||
tempfile.workspace = true
|
||||
tokio-tungstenite.workspace = true
|
||||
wiremock.workspace = true
|
||||
aws-sdk-secretsmanager = "1.117.0"
|
||||
|
||||
[[bench]]
|
||||
name = "serialization"
|
||||
|
|
|
|||
|
|
@ -23,12 +23,6 @@ pub(crate) enum ProjectionDetail {
|
|||
expected: &'static str,
|
||||
actual: String,
|
||||
},
|
||||
#[error("{field}")]
|
||||
InvalidField {
|
||||
field: &'static str,
|
||||
#[source]
|
||||
source: serde_json::Error,
|
||||
},
|
||||
#[error("{group}.{name}: missing snapshot field")]
|
||||
MissingSnapshotField {
|
||||
group: &'static str,
|
||||
|
|
@ -160,15 +154,6 @@ impl<'py> Field<'py> {
|
|||
self.value.is(PyBool::new(self.value.py(), true))
|
||||
}
|
||||
|
||||
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
|
||||
if !self.value.is_instance_of::<PyBool>() {
|
||||
return Err(ProjectionError::InternalSchemaFailure(
|
||||
self.expected("a Boolean")?,
|
||||
));
|
||||
}
|
||||
Ok(self.exact_true())
|
||||
}
|
||||
|
||||
pub(crate) fn strict_string(&self) -> Result<String, ProjectionError> {
|
||||
let Ok(value) = self.value.cast::<PyString>() else {
|
||||
return Err(self.invalid("a string"));
|
||||
|
|
@ -235,19 +220,6 @@ impl<'py> Field<'py> {
|
|||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn optional_string_collection(
|
||||
&self,
|
||||
) -> Result<Option<Vec<String>>, ProjectionError> {
|
||||
if self.value.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
self.string_collection().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn python_binding(&self) -> Option<Py<PyAny>> {
|
||||
(!self.value.is_none()).then(|| self.value.clone().unbind())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
|
|||
|
|
@ -375,7 +375,7 @@ mod tests {
|
|||
Python::attach(|py| {
|
||||
let raised = pyo3::exceptions::PyKeyError::new_err("vault down");
|
||||
let identity = raised.value(py).clone().unbind();
|
||||
let error = Error::from(crate::secrets::read_error(py, raised));
|
||||
let error = Error::from(crate::secrets::external_error(py, raised));
|
||||
let surfaced = route_failure(py, error, RequestMarker::Unmarked, |_| {
|
||||
panic!("the raised exception must not be remapped")
|
||||
})
|
||||
|
|
|
|||
|
|
@ -13,10 +13,6 @@ pub(crate) enum PythonSettings {
|
|||
UrlPolicy,
|
||||
#[strum(serialize = "provider_defaults")]
|
||||
ProviderDefaults,
|
||||
#[strum(serialize = "secret_manager")]
|
||||
SecretManager,
|
||||
#[strum(serialize = "secret_manager_binding")]
|
||||
SecretManagerBinding,
|
||||
}
|
||||
|
||||
pub(crate) struct Snapshot<'py> {
|
||||
|
|
|
|||
|
|
@ -1,508 +0,0 @@
|
|||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_host_python::{PythonContext, attach_blocking};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use litellm_secrets::{
|
||||
Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret,
|
||||
};
|
||||
|
||||
use pyo3::{
|
||||
exceptions::PyException,
|
||||
prelude::*,
|
||||
types::{PyDict, PyString},
|
||||
};
|
||||
|
||||
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
|
||||
/// client, or a manually assigned SDK client.
|
||||
pub(crate) struct PythonSecretManager {
|
||||
client: Arc<PythonClient>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
struct PythonClient {
|
||||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl PythonSecretManager {
|
||||
pub(crate) fn new(
|
||||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
) -> Self {
|
||||
Self {
|
||||
client: Arc::new(PythonClient {
|
||||
client,
|
||||
system,
|
||||
settings,
|
||||
}),
|
||||
context,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonClient {
|
||||
fn read(&self, py: Python<'_>, name: &str) -> PyResult<Option<String>> {
|
||||
let client = self.client.bind(py);
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("client", client)?;
|
||||
kwargs.set_item("key_manager", self.system.map_or("local", python_name))?;
|
||||
kwargs.set_item("secret_name", name)?;
|
||||
match &self.settings {
|
||||
Some(settings) => kwargs.set_item("key_management_settings", settings.bind(py))?,
|
||||
None => kwargs.set_item("key_management_settings", py.None())?,
|
||||
}
|
||||
let result = py
|
||||
.import(HANDLER_MODULE)?
|
||||
.getattr("get_secret_from_manager")?
|
||||
.call((), Some(&kwargs))?;
|
||||
if result.is_instance_of::<PyString>() {
|
||||
result.extract().map(Some)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The `KeyManagementSystem` value as Python spells it.
|
||||
fn python_name(system: KeyManagementSystem) -> &'static str {
|
||||
system.into()
|
||||
}
|
||||
|
||||
impl ExternalSecretManager for PythonSecretManager {
|
||||
fn system(&self) -> KeyManagementSystem {
|
||||
self.client.system.unwrap_or(KeyManagementSystem::Custom)
|
||||
}
|
||||
|
||||
fn read_secret<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
_settings: &'a KeyManagementSettings,
|
||||
_environment: &'a (dyn Lookup + Send + Sync),
|
||||
) -> Pin<Box<dyn Future<Output = Result<Option<Secret>, Error>> + Send + 'a>> {
|
||||
let client = Arc::clone(&self.client);
|
||||
let context = self.context.clone();
|
||||
let name = name.to_owned();
|
||||
Box::pin(async move {
|
||||
match attach_blocking(context, move |py| match client.read(py, &name) {
|
||||
Ok(value) => Ok(value.map(SecretValue::new).map(Secret::String)),
|
||||
// `get_secret` answers a failed manager read from the process environment, but
|
||||
// only for `Exception`: cancellation and other `BaseException`s propagate.
|
||||
Err(error) if error.is_instance_of::<PyException>(py) => {
|
||||
log_environment_fallback(py, &name, &error)
|
||||
.map_err(|error| external_error(py, error))?;
|
||||
Err(read_error(py, error))
|
||||
}
|
||||
Err(error) => Err(external_error(py, error)),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(error) => Python::attach(|py| Err(external_error(py, error))),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn log_environment_fallback(py: Python<'_>, name: &str, error: &PyErr) -> PyResult<()> {
|
||||
py.import(DIAGNOSTICS_MODULE)?
|
||||
.getattr("log_environment_fallback")?
|
||||
.call1((name, error.value(py)))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
|
||||
use litellm_secrets::{
|
||||
FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager,
|
||||
SecretManagerState, SecretResolver,
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
use rstest::rstest;
|
||||
|
||||
use litellm_host_python::PythonContext;
|
||||
|
||||
use super::{DIAGNOSTICS_MODULE, HANDLER_MODULE, PythonSecretManager, python_name};
|
||||
use crate::secrets::python_error;
|
||||
|
||||
/// `sys.modules` is interpreter-global, so tests that install or rely on the handler module
|
||||
/// cannot overlap with any other test on this list.
|
||||
static HANDLER_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
fn handler_guard() -> MutexGuard<'static, ()> {
|
||||
HANDLER_LOCK.lock().expect("handler lock poisoned")
|
||||
}
|
||||
|
||||
/// A resolver over a Python manager whose reads raise `failure_type`, with the chained
|
||||
/// exceptions Python attaches, and `fallback` as the process environment. The returned guard
|
||||
/// keeps other module-mutating tests out for the lifetime of the returned resolver.
|
||||
fn failing_resolver(
|
||||
failure_type: &str,
|
||||
fallback: Option<&'static str>,
|
||||
) -> (SecretResolver, Py<PyDict>, MutexGuard<'static, ()>) {
|
||||
let handler = handler_guard();
|
||||
Python::initialize();
|
||||
let (reader, locals) = Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("failure_type", failure_type).unwrap();
|
||||
py.run(
|
||||
c"
|
||||
import asyncio
|
||||
failure = eval(failure_type)('secret manager failed')
|
||||
cause = RuntimeError('original cause')
|
||||
context = RuntimeError('original context')
|
||||
failure.__cause__ = cause
|
||||
failure.__context__ = context
|
||||
class Manager:
|
||||
def sync_read_secret(self, secret_name):
|
||||
raise failure
|
||||
manager = Manager()
|
||||
import sys, types
|
||||
for name in ('litellm', 'litellm.secret_managers'):
|
||||
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):
|
||||
return kwargs['client'].sync_read_secret(kwargs['secret_name'])
|
||||
handler.get_secret_from_manager = get_secret_from_manager
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
locals.get_item("manager").unwrap().unwrap().unbind(),
|
||||
None,
|
||||
None,
|
||||
PythonContext::capture(py).unwrap(),
|
||||
);
|
||||
(reader, locals.unbind())
|
||||
});
|
||||
let resolver = SecretResolver::new_python_compatible(
|
||||
Arc::new(SecretManagerState::new(
|
||||
SecretManager::External(Arc::new(reader)),
|
||||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(move |_: &str| fallback.map(str::to_owned)),
|
||||
OidcResolver::new(litellm_http::Client::for_test()),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback);
|
||||
(resolver, locals, handler)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::cancelled("asyncio.CancelledError", None)]
|
||||
#[case::cancelled_with_fallback("asyncio.CancelledError", Some("environment-key"))]
|
||||
#[case::keyboard_interrupt("KeyboardInterrupt", Some("environment-key"))]
|
||||
#[tokio::test]
|
||||
async fn base_exceptions_propagate_unchanged_even_with_environment_fallback(
|
||||
#[case] failure_type: &str,
|
||||
#[case] fallback: Option<&'static str>,
|
||||
) {
|
||||
let (resolver, locals, _handler) = failing_resolver(failure_type, fallback);
|
||||
let error = resolver.get_secret("API_KEY", None).await.unwrap_err();
|
||||
Python::attach(|py| {
|
||||
let original = python_error(py, &error).unwrap();
|
||||
let locals = locals.bind(py);
|
||||
assert!(
|
||||
original
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
for (attribute, name) in [("__cause__", "cause"), ("__context__", "context")] {
|
||||
assert!(
|
||||
original
|
||||
.value(py)
|
||||
.getattr(attribute)
|
||||
.unwrap()
|
||||
.is(locals.get_item(name).unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
assert!(original.traceback(py).is_some());
|
||||
});
|
||||
}
|
||||
|
||||
/// Installs a persistent `litellm._logging` stub whose `verbose_logger.error` records its
|
||||
/// arguments, and returns those recorded for `name`.
|
||||
fn logged_errors<'py>(py: Python<'py>, name: &str) -> Vec<Bound<'py, PyAny>> {
|
||||
py.run(
|
||||
c"
|
||||
import sys, types
|
||||
class Logger:
|
||||
calls = []
|
||||
def error(self, *args):
|
||||
self.calls.append(args)
|
||||
logging = types.ModuleType('litellm._logging')
|
||||
logging.verbose_logger = Logger()
|
||||
for module in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'):
|
||||
sys.modules.setdefault(module, types.ModuleType(module))
|
||||
sys.modules.setdefault('litellm._logging', logging)
|
||||
",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
pyo3::types::PyModule::from_code(
|
||||
py,
|
||||
&std::ffi::CString::new(include_str!(
|
||||
"../../../../../litellm/rust_bridge/host/diagnostics.py"
|
||||
))
|
||||
.unwrap(),
|
||||
c"diagnostics.py",
|
||||
&std::ffi::CString::new(DIAGNOSTICS_MODULE).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
py.import("litellm._logging")
|
||||
.unwrap()
|
||||
.getattr("verbose_logger")
|
||||
.unwrap()
|
||||
.getattr("calls")
|
||||
.unwrap()
|
||||
.try_iter()
|
||||
.unwrap()
|
||||
.map(Result::unwrap)
|
||||
.filter(|call| call.get_item(1).unwrap().extract::<String>().unwrap() == name)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::value_error("ValueError", None, "FALLBACK_VALUE_ERROR")]
|
||||
#[case::value_error_with_fallback(
|
||||
"ValueError",
|
||||
Some("environment-key"),
|
||||
"FALLBACK_VALUE_ERROR_WITH_ENVIRONMENT"
|
||||
)]
|
||||
#[case::runtime_error_with_fallback(
|
||||
"RuntimeError",
|
||||
Some("environment-key"),
|
||||
"FALLBACK_RUNTIME_ERROR_WITH_ENVIRONMENT"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn exceptions_are_logged_and_answered_from_the_environment(
|
||||
#[case] failure_type: &str,
|
||||
#[case] fallback: Option<&'static str>,
|
||||
#[case] name: &str,
|
||||
) {
|
||||
let (resolver, _locals, _handler) = failing_resolver(failure_type, fallback);
|
||||
Python::attach(|py| assert!(logged_errors(py, name).is_empty()));
|
||||
let secret = resolver.get_secret(name, None).await.unwrap();
|
||||
assert_eq!(
|
||||
secret.map(|secret| match secret {
|
||||
litellm_secrets::Secret::String(value) => value.expose().to_owned(),
|
||||
other => panic!("unexpected secret {other:?}"),
|
||||
}),
|
||||
fallback.map(str::to_owned)
|
||||
);
|
||||
Python::attach(|py| {
|
||||
let calls = logged_errors(py, name);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert!(
|
||||
calls[0]
|
||||
.get_item(3)
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap()
|
||||
.contains("sync_read_secret")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
/// Installs a fake `get_secret_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>)) {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
import sys, types
|
||||
previous_handler = sys.modules.get('litellm.secret_managers.secret_manager_handler')
|
||||
calls = []
|
||||
def get_secret_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'):
|
||||
sys.modules.setdefault(name, types.ModuleType(name))
|
||||
sys.modules['litellm.secret_managers.secret_manager_handler'] = handler
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
body(&locals);
|
||||
py.run(
|
||||
c"
|
||||
if previous_handler is None:
|
||||
sys.modules.pop('litellm.secret_managers.secret_manager_handler', None)
|
||||
else:
|
||||
sys.modules['litellm.secret_managers.secret_manager_handler'] = previous_handler
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None")]
|
||||
#[case("True")]
|
||||
#[case("123")]
|
||||
#[case("{'key': 'value'}")]
|
||||
fn nonstring_results_are_absent_without_a_read_failure(#[case] expression: &str) {
|
||||
let _handler = handler_guard();
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
with_fake_handler(py, |locals| {
|
||||
locals.set_item("expression", expression).unwrap();
|
||||
py.run(
|
||||
c"handler.get_secret_from_manager = lambda **kwargs: eval(expression)",
|
||||
Some(locals),
|
||||
Some(locals),
|
||||
)
|
||||
.unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
py.None(),
|
||||
None,
|
||||
None,
|
||||
PythonContext::capture(py).unwrap(),
|
||||
);
|
||||
assert_eq!(reader.client.read(py, "KEY").unwrap(), None);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::google_kms(KeyManagementSystem::GoogleKms)]
|
||||
#[case::azure_key_vault(KeyManagementSystem::AzureKeyVault)]
|
||||
#[case::aws_secret_manager(KeyManagementSystem::AwsSecretManager)]
|
||||
#[case::google_secret_manager(KeyManagementSystem::GoogleSecretManager)]
|
||||
#[case::hashicorp_vault(KeyManagementSystem::HashicorpVault)]
|
||||
#[case::cyberark(KeyManagementSystem::Cyberark)]
|
||||
#[case::local(KeyManagementSystem::Local)]
|
||||
#[case::aws_kms(KeyManagementSystem::AwsKms)]
|
||||
#[case::custom(KeyManagementSystem::Custom)]
|
||||
fn python_names_round_trip_through_serde(#[case] system: KeyManagementSystem) {
|
||||
assert_eq!(
|
||||
serde_json::to_value(system).unwrap(),
|
||||
serde_json::Value::String(python_name(system).to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() {
|
||||
let _handler = handler_guard();
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
with_fake_handler(py, |locals| {
|
||||
let client = py.eval(c"object()", None, None).unwrap();
|
||||
let settings = py.eval(c"object()", None, None).unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
client.clone().unbind(),
|
||||
Some(KeyManagementSystem::AzureKeyVault),
|
||||
Some(settings.clone().unbind()),
|
||||
PythonContext::capture(py).unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
reader.client.read(py, "API_KEY").unwrap().as_deref(),
|
||||
Some("handled-API_KEY")
|
||||
);
|
||||
assert!(py.import(HANDLER_MODULE).is_ok());
|
||||
let calls = locals.get_item("calls").unwrap().unwrap();
|
||||
let call = calls.get_item(0).unwrap().cast_into::<PyDict>().unwrap();
|
||||
assert!(call.get_item("client").unwrap().unwrap().is(&client));
|
||||
assert!(
|
||||
call.get_item("key_management_settings")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.is(&settings)
|
||||
);
|
||||
assert_eq!(
|
||||
call.get_item("key_manager")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"azure_key_vault"
|
||||
);
|
||||
assert_eq!(
|
||||
call.get_item("secret_name")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"API_KEY"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::manually_assigned(None, "local")]
|
||||
#[case::custom(Some(KeyManagementSystem::Custom), "custom")]
|
||||
fn direct_readers_dispatch_through_the_python_handler_like_get_secret(
|
||||
#[case] system: Option<KeyManagementSystem>,
|
||||
#[case] key_manager: &str,
|
||||
) {
|
||||
let _handler = handler_guard();
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
with_fake_handler(py, |locals| {
|
||||
py.run(
|
||||
c"
|
||||
class Manager:
|
||||
def __init__(self):
|
||||
self.names = []
|
||||
def sync_read_secret(self, secret_name, optional_params=None, timeout=None):
|
||||
self.names.append(secret_name)
|
||||
return 'direct-' + secret_name
|
||||
manager = Manager()
|
||||
",
|
||||
Some(locals),
|
||||
Some(locals),
|
||||
)
|
||||
.unwrap();
|
||||
let manager = locals.get_item("manager").unwrap().unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
manager.clone().unbind(),
|
||||
system,
|
||||
None,
|
||||
PythonContext::capture(py).unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
reader.client.read(py, "API_KEY").unwrap().as_deref(),
|
||||
Some("handled-API_KEY")
|
||||
);
|
||||
assert_eq!(
|
||||
manager
|
||||
.getattr("names")
|
||||
.unwrap()
|
||||
.extract::<Vec<String>>()
|
||||
.unwrap(),
|
||||
Vec::<String>::new()
|
||||
);
|
||||
let calls = locals.get_item("calls").unwrap().unwrap();
|
||||
let call = calls.get_item(0).unwrap().cast_into::<PyDict>().unwrap();
|
||||
assert!(call.get_item("client").unwrap().unwrap().is(&manager));
|
||||
assert_eq!(
|
||||
call.get_item("key_manager")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
key_manager
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -1,385 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_secret_value::SecretValue;
|
||||
use litellm_secrets::{SecretManager, SecretManagerState};
|
||||
use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem};
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
|
||||
use litellm_host_python::PythonContext;
|
||||
|
||||
use super::callback::PythonSecretManager;
|
||||
use crate::{
|
||||
coercion::{Field, FieldSpec, ProjectionError},
|
||||
python_settings::{PythonSettings, Snapshot},
|
||||
};
|
||||
|
||||
const SYSTEM: FieldSpec<Option<KeyManagementSystem>> =
|
||||
FieldSpec::new("system", parse_optional_system);
|
||||
const ACCESS_MODE: FieldSpec<AccessMode> = FieldSpec::new("access_mode", parse_access_mode);
|
||||
const HOSTED_KEYS: FieldSpec<Option<Vec<String>>> =
|
||||
FieldSpec::new("hosted_keys", |field| field.optional_string_collection());
|
||||
const STORE_VIRTUAL_KEYS: FieldSpec<bool> =
|
||||
FieldSpec::new("store_virtual_keys", |field| field.truthy());
|
||||
const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec<String> =
|
||||
FieldSpec::new("prefix_for_stored_virtual_keys", |field| {
|
||||
field.strict_string()
|
||||
});
|
||||
const PRIMARY_SECRET_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string());
|
||||
const KMS_KEY_ID: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("kms_key_id", |field| field.falsy_optional_string());
|
||||
const CUSTOM_SECRET_MANAGER: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("custom_secret_manager", |field| {
|
||||
field.falsy_optional_string()
|
||||
});
|
||||
const AWS_REGION_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_region_name", |field| field.falsy_optional_string());
|
||||
const AWS_ROLE_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_role_name", |field| field.falsy_optional_string());
|
||||
const AWS_SESSION_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_session_name", |field| field.falsy_optional_string());
|
||||
const AWS_EXTERNAL_ID: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_external_id", |field| field.falsy_optional_string());
|
||||
const AWS_PROFILE_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string());
|
||||
const AWS_WEB_IDENTITY_TOKEN: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_web_identity_token", |field| {
|
||||
field.falsy_optional_string()
|
||||
});
|
||||
const AWS_STS_ENDPOINT: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string());
|
||||
const REPLICA_REGIONS: FieldSpec<Option<Vec<String>>> =
|
||||
FieldSpec::new("replica_regions", |field| {
|
||||
field.optional_string_collection()
|
||||
});
|
||||
const CLIENT: FieldSpec<Option<Py<PyAny>>> =
|
||||
FieldSpec::new("client", |field| Ok(field.python_binding()));
|
||||
const SETTINGS_OBJECT: FieldSpec<Option<Py<PyAny>>> =
|
||||
FieldSpec::new("settings_object", |field| Ok(field.python_binding()));
|
||||
|
||||
/// `litellm.secret_manager_client` as the bridge classifies it.
|
||||
pub(crate) enum SecretManagerClient {
|
||||
/// `None`: reads come from the process environment.
|
||||
Local,
|
||||
/// A custom manager, legacy compatible client, or manually assigned SDK client that keeps
|
||||
/// executing in Python.
|
||||
PythonCallback(Py<PyAny>),
|
||||
Native(Box<SecretManager>),
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for SecretManagerClient {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(match self {
|
||||
Self::Local => "Local",
|
||||
Self::Native(_) => "Native",
|
||||
Self::PythonCallback(_) => "PythonCallback",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// One operation-local capture of the secret manager globals, taken while attached to Python.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SecretManagerSnapshot {
|
||||
pub(crate) client: SecretManagerClient,
|
||||
pub(crate) system: Option<KeyManagementSystem>,
|
||||
/// Typed settings that drive native routing: access mode and hosted keys.
|
||||
pub(crate) settings: KeyManagementSettings,
|
||||
/// The original `KeyManagementSettings` object, handed back to Python callbacks unchanged.
|
||||
pub(crate) settings_object: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl SecretManagerSnapshot {
|
||||
pub(crate) fn into_state(self, context: PythonContext) -> Arc<SecretManagerState> {
|
||||
match self.client {
|
||||
SecretManagerClient::Native(backend) => {
|
||||
Arc::new(SecretManagerState::new(*backend, self.settings))
|
||||
}
|
||||
SecretManagerClient::Local => Arc::new(SecretManagerState::default()),
|
||||
SecretManagerClient::PythonCallback(client) => Arc::new(SecretManagerState::new(
|
||||
SecretManager::External(Arc::new(PythonSecretManager::new(
|
||||
client,
|
||||
self.system,
|
||||
self.settings_object,
|
||||
context,
|
||||
))),
|
||||
self.settings,
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads and projects the secret manager settings group in one attached operation.
|
||||
pub(crate) fn read(py: Python<'_>) -> PyResult<SecretManagerSnapshot> {
|
||||
let snapshot = project(&PythonSettings::SecretManagerBinding.read(py)?)?;
|
||||
let SecretManagerClient::PythonCallback(client) = &snapshot.client else {
|
||||
return Ok(snapshot);
|
||||
};
|
||||
if matches!(
|
||||
snapshot.system,
|
||||
Some(KeyManagementSystem::Custom | KeyManagementSystem::Local)
|
||||
) {
|
||||
return Ok(snapshot);
|
||||
}
|
||||
let Some(native) = super::runtime::NativeSecretManager::from_client(client.bind(py))? else {
|
||||
return Ok(snapshot);
|
||||
};
|
||||
let backend = native.borrow(py).backend()?;
|
||||
if snapshot
|
||||
.system
|
||||
.is_some_and(|system| system != backend.system())
|
||||
{
|
||||
return Err(super::ManagerError::SystemMismatch.into());
|
||||
}
|
||||
Ok(SecretManagerSnapshot {
|
||||
client: SecretManagerClient::Native(Box::new(backend)),
|
||||
..snapshot
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn project(snapshot: &Snapshot<'_>) -> Result<SecretManagerSnapshot, ProjectionError> {
|
||||
let system = snapshot.read(&SYSTEM)?;
|
||||
let access_mode = snapshot.read(&ACCESS_MODE)?;
|
||||
let settings = KeyManagementSettings {
|
||||
hosted_keys: snapshot.read(&HOSTED_KEYS)?,
|
||||
store_virtual_keys: Some(snapshot.read(&STORE_VIRTUAL_KEYS)?),
|
||||
prefix_for_stored_virtual_keys: snapshot.read(&PREFIX_FOR_STORED_VIRTUAL_KEYS)?,
|
||||
access_mode,
|
||||
primary_secret_name: snapshot.read(&PRIMARY_SECRET_NAME)?,
|
||||
kms_key_id: snapshot.read(&KMS_KEY_ID)?,
|
||||
custom_secret_manager: snapshot.read(&CUSTOM_SECRET_MANAGER)?,
|
||||
aws_region_name: snapshot.read(&AWS_REGION_NAME)?,
|
||||
aws_role_name: snapshot.read(&AWS_ROLE_NAME)?,
|
||||
aws_session_name: snapshot.read(&AWS_SESSION_NAME)?,
|
||||
aws_external_id: snapshot.read(&AWS_EXTERNAL_ID)?.map(SecretValue::new),
|
||||
aws_profile_name: snapshot.read(&AWS_PROFILE_NAME)?,
|
||||
aws_web_identity_token: snapshot
|
||||
.read(&AWS_WEB_IDENTITY_TOKEN)?
|
||||
.map(SecretValue::new),
|
||||
aws_sts_endpoint: snapshot.read(&AWS_STS_ENDPOINT)?,
|
||||
replica_regions: snapshot.read(&REPLICA_REGIONS)?,
|
||||
..KeyManagementSettings::default()
|
||||
};
|
||||
let client = match snapshot.read(&CLIENT)? {
|
||||
None => SecretManagerClient::Local,
|
||||
Some(client) => SecretManagerClient::PythonCallback(client),
|
||||
};
|
||||
Ok(SecretManagerSnapshot {
|
||||
client,
|
||||
system,
|
||||
settings,
|
||||
settings_object: snapshot.read(&SETTINGS_OBJECT)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_system(
|
||||
field: &Field<'_>,
|
||||
) -> Result<Option<KeyManagementSystem>, ProjectionError> {
|
||||
let Some(value) = field.falsy_optional_string()? else {
|
||||
return Ok(None);
|
||||
};
|
||||
serde_json::from_value(Value::String(value))
|
||||
.map(Some)
|
||||
.map_err(|error| {
|
||||
ProjectionError::InvalidConfiguration(crate::coercion::ProjectionDetail::InvalidField {
|
||||
field: "secret manager system",
|
||||
source: error,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_access_mode(field: &Field<'_>) -> Result<AccessMode, ProjectionError> {
|
||||
let value = field.strict_string()?;
|
||||
serde_json::from_value(Value::String(value)).map_err(|error| {
|
||||
ProjectionError::InvalidConfiguration(crate::coercion::ProjectionDetail::InvalidField {
|
||||
field: "secret manager access mode",
|
||||
source: error,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use super::{SecretManagerClient, project};
|
||||
use crate::python_settings::PythonSettings;
|
||||
|
||||
fn snapshot<'py>(
|
||||
py: Python<'py>,
|
||||
system: &str,
|
||||
access_mode: &str,
|
||||
store_virtual_keys: Bound<'py, PyAny>,
|
||||
hosted_keys: Bound<'py, PyAny>,
|
||||
) -> crate::python_settings::Snapshot<'py> {
|
||||
snapshot_with_client(
|
||||
py,
|
||||
system,
|
||||
access_mode,
|
||||
store_virtual_keys,
|
||||
hosted_keys,
|
||||
py.None().into_bound(py),
|
||||
)
|
||||
}
|
||||
|
||||
fn snapshot_with_client<'py>(
|
||||
py: Python<'py>,
|
||||
system: &str,
|
||||
access_mode: &str,
|
||||
store_virtual_keys: Bound<'py, PyAny>,
|
||||
hosted_keys: Bound<'py, PyAny>,
|
||||
client: Bound<'py, PyAny>,
|
||||
) -> crate::python_settings::Snapshot<'py> {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("client", client).unwrap();
|
||||
locals.set_item("system", system).unwrap();
|
||||
locals.set_item("access_mode", access_mode).unwrap();
|
||||
locals
|
||||
.set_item("store_virtual_keys", store_virtual_keys)
|
||||
.unwrap();
|
||||
locals.set_item("hosted_keys", hosted_keys).unwrap();
|
||||
py.run(
|
||||
cr#"
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManager:
|
||||
system: object
|
||||
access_mode: object
|
||||
hosted_keys: object
|
||||
primary_secret_name: object
|
||||
store_virtual_keys: object
|
||||
prefix_for_stored_virtual_keys: object
|
||||
kms_key_id: object
|
||||
custom_secret_manager: object
|
||||
aws_region_name: object
|
||||
aws_role_name: object
|
||||
aws_session_name: object
|
||||
aws_external_id: object
|
||||
aws_profile_name: object
|
||||
aws_web_identity_token: object
|
||||
aws_sts_endpoint: object
|
||||
replica_regions: object
|
||||
client: object
|
||||
settings_object: object
|
||||
|
||||
root = SimpleNamespace(secret_manager=SecretManager(
|
||||
system=system,
|
||||
access_mode=access_mode,
|
||||
hosted_keys=hosted_keys,
|
||||
primary_secret_name=None,
|
||||
store_virtual_keys=store_virtual_keys,
|
||||
prefix_for_stored_virtual_keys="litellm/",
|
||||
kms_key_id=None,
|
||||
custom_secret_manager=None,
|
||||
aws_region_name=None,
|
||||
aws_role_name=None,
|
||||
aws_session_name=None,
|
||||
aws_external_id=None,
|
||||
aws_profile_name=None,
|
||||
aws_web_identity_token=None,
|
||||
aws_sts_endpoint=None,
|
||||
replica_regions=None,
|
||||
client=client,
|
||||
settings_object=None,
|
||||
))
|
||||
"#,
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
PythonSettings::SecretManagerBinding.snapshot(
|
||||
locals
|
||||
.get_item("root")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.getattr("secret_manager")
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::string_true(Some("true"), false, true)]
|
||||
#[case::string_one(Some("1"), false, true)]
|
||||
#[case::true_value(None, true, true)]
|
||||
#[case::false_value(None, false, false)]
|
||||
#[case::string_false(Some("false"), false, true)]
|
||||
fn python_compatible_boolean_coercion(
|
||||
#[case] string_value: Option<&str>,
|
||||
#[case] bool_value: bool,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let store_virtual_keys = match string_value {
|
||||
Some(value) => value.into_pyobject(py).unwrap().into_any(),
|
||||
None => bool_value.into_pyobject(py).unwrap().to_owned().into_any(),
|
||||
};
|
||||
let hosted_keys = PyTuple::new(py, ["ONE"]).unwrap().into_any();
|
||||
let projected = project(&snapshot(
|
||||
py,
|
||||
"local",
|
||||
"read_only",
|
||||
store_virtual_keys,
|
||||
hosted_keys,
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(projected.settings.store_virtual_keys, Some(expected));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_system_is_rejected() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = project(&snapshot(
|
||||
py,
|
||||
"unknown",
|
||||
"read_only",
|
||||
false.into_pyobject(py).unwrap().to_owned().into_any(),
|
||||
PyTuple::empty(py).into_any(),
|
||||
))
|
||||
.unwrap_err();
|
||||
let error: PyErr = error.into();
|
||||
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_identity_selects_local_or_python_callback() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let falsy = false.into_pyobject(py).unwrap().to_owned().into_any();
|
||||
let local = project(&snapshot(
|
||||
py,
|
||||
"local",
|
||||
"read_only",
|
||||
falsy.clone(),
|
||||
PyTuple::empty(py).into_any(),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(matches!(local.client, SecretManagerClient::Local));
|
||||
assert!(local.settings_object.is_none());
|
||||
|
||||
let manager = py.eval(c"object()", None, None).unwrap();
|
||||
let custom = project(&snapshot_with_client(
|
||||
py,
|
||||
"custom",
|
||||
"read_only",
|
||||
falsy,
|
||||
PyTuple::empty(py).into_any(),
|
||||
manager.clone(),
|
||||
))
|
||||
.unwrap();
|
||||
let SecretManagerClient::PythonCallback(client) = custom.client else {
|
||||
panic!("a live client must stay a Python callback");
|
||||
};
|
||||
assert!(client.bind(py).is(&manager));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -16,8 +16,6 @@ pub(crate) enum ManagerError {
|
|||
UnknownSystem(#[source] serde_json::Error),
|
||||
#[error("invalid secret manager settings")]
|
||||
Settings(#[source] litellm_error::Source),
|
||||
#[error("native secret manager system does not match configuration")]
|
||||
SystemMismatch,
|
||||
#[error("secret manager does not support provider reads")]
|
||||
ProviderReads,
|
||||
#[error("invalid timeout")]
|
||||
|
|
@ -49,7 +47,6 @@ impl From<ManagerError> for PyErr {
|
|||
ManagerError::Forked => PyRuntimeError::new_err(message),
|
||||
ManagerError::UnknownSystem(_)
|
||||
| ManagerError::Settings(_)
|
||||
| ManagerError::SystemMismatch
|
||||
| ManagerError::ProviderReads
|
||||
| ManagerError::Timeout(_)
|
||||
| ManagerError::CyberarkSecretMissing(_)
|
||||
|
|
@ -72,18 +69,12 @@ impl fmt::Debug for PythonSecretError {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) fn external_error(py: Python<'_>, error: PyErr) -> Error {
|
||||
pub(crate) fn external_error(py: Python<'_>, error: PyErr) -> Error {
|
||||
Error::ExternalManager(litellm_error::Source::new(PythonSecretError(
|
||||
error.into_value(py),
|
||||
)))
|
||||
}
|
||||
|
||||
pub(crate) fn read_error(py: Python<'_>, error: PyErr) -> Error {
|
||||
Error::ExternalRead(litellm_error::Source::new(PythonSecretError(
|
||||
error.into_value(py),
|
||||
)))
|
||||
}
|
||||
|
||||
/// The exception a Python secret manager raised anywhere in the chain, so a route surfaces it
|
||||
/// unchanged instead of wrapping it.
|
||||
pub(crate) fn python_error(
|
||||
|
|
|
|||
|
|
@ -1,46 +1,33 @@
|
|||
pub(crate) mod callback;
|
||||
pub(crate) mod config;
|
||||
mod error;
|
||||
mod mutation;
|
||||
mod operations;
|
||||
mod provider;
|
||||
mod python;
|
||||
pub(crate) mod resolved;
|
||||
pub(crate) mod runtime;
|
||||
mod vault;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use error::read_error;
|
||||
pub(crate) use error::external_error;
|
||||
pub(crate) use error::{ManagerError, python_error};
|
||||
use litellm_secrets::source::{EnvironmentSecrets, SecretSource};
|
||||
use pyo3::prelude::*;
|
||||
use python::PythonSecrets;
|
||||
use resolved::ResolvedSecrets;
|
||||
|
||||
use crate::{coercion::FieldSpec, python_settings::PythonSettings};
|
||||
use crate::python_settings::missing_module;
|
||||
|
||||
const NATIVE: FieldSpec<bool> = FieldSpec::new("native", |field| field.schema_bool());
|
||||
|
||||
/// Where a Rust route reads provider secrets from. Python's `get_secret_str` when the settings
|
||||
/// projection reports no configured secret manager, otherwise the resolved manager: the native
|
||||
/// backend when the configured client captures one, the Python callback when it does not. A bare
|
||||
/// extension module without the litellm package reads the process environment.
|
||||
/// Where a Rust route reads provider secrets from: Python's `get_secret_str`, the one resolver for
|
||||
/// the configured secret manager. A bare extension module without the litellm package reads the
|
||||
/// process environment.
|
||||
pub(crate) fn source(py: Python<'_>) -> PyResult<Arc<dyn SecretSource>> {
|
||||
let Some(snapshot) = PythonSettings::SecretManager.read_or_unset(py)? else {
|
||||
return Ok(Arc::new(EnvironmentSecrets::python_compatible(
|
||||
crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
|
||||
)));
|
||||
};
|
||||
if snapshot.read(&NATIVE)? {
|
||||
let context = litellm_host_python::PythonContext::capture(py)?;
|
||||
let client = crate::http::host_client(py, litellm_http::ClientVariant::Provider)?;
|
||||
return Ok(Arc::new(ResolvedSecrets::new(
|
||||
config::read(py)?,
|
||||
context,
|
||||
client,
|
||||
)));
|
||||
match PythonSecrets::new(py) {
|
||||
Ok(secrets) => Ok(Arc::new(secrets)),
|
||||
Err(error) if missing_module(py, &error, "litellm")? => {
|
||||
Ok(Arc::new(EnvironmentSecrets::python_compatible(
|
||||
crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
|
||||
)))
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
Ok(Arc::new(PythonSecrets::new(py)?))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,20 +1,25 @@
|
|||
use std::sync::Arc;
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_host_python::{PythonContext, attach_blocking};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use litellm_secrets::{Error, source::SecretSource};
|
||||
use litellm_secrets::{
|
||||
Error,
|
||||
source::{SecretSource, Secrets},
|
||||
};
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::{prelude::*, types::PyTuple};
|
||||
|
||||
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.
|
||||
/// Reads secrets through the host module's `read_secrets`, so the configured manager, the key
|
||||
/// management settings and the environment fallback behave exactly as they do in Python. A batch of
|
||||
/// names is one host call.
|
||||
pub(super) struct PythonSecrets {
|
||||
get_secret_str: Arc<Py<PyAny>>,
|
||||
read_secrets: Arc<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
|
|
@ -22,18 +27,47 @@ impl PythonSecrets {
|
|||
pub(super) fn new(py: Python<'_>) -> PyResult<Self> {
|
||||
Ok(Self::reading_with(
|
||||
py.import(SECRET_MANAGER_MODULE)?
|
||||
.getattr("read_secret_str")?
|
||||
.getattr("read_secrets")?
|
||||
.unbind(),
|
||||
PythonContext::capture(py)?,
|
||||
))
|
||||
}
|
||||
|
||||
fn reading_with(get_secret_str: Py<PyAny>, context: PythonContext) -> Self {
|
||||
fn reading_with(read_secrets: Py<PyAny>, context: PythonContext) -> Self {
|
||||
Self {
|
||||
get_secret_str: Arc::new(get_secret_str),
|
||||
read_secrets: Arc::new(read_secrets),
|
||||
context,
|
||||
}
|
||||
}
|
||||
|
||||
async fn read(&self, names: Vec<String>) -> Result<Vec<Option<SecretValue>>, Error> {
|
||||
let read_secrets = Arc::clone(&self.read_secrets);
|
||||
match attach_blocking(self.context.clone(), move |py| {
|
||||
PyTuple::new(py, names)
|
||||
.and_then(|names| read_secrets.bind(py).call1((names,)))
|
||||
.and_then(|values| values.extract::<Vec<Option<String>>>())
|
||||
.map(|values| {
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| value.map(SecretValue::new))
|
||||
.collect()
|
||||
})
|
||||
.map_err(|error| external_error(py, error))
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(error) => Python::attach(|py| Err(external_error(py, error))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ReadSecrets(HashMap<String, SecretValue>);
|
||||
|
||||
impl Lookup for ReadSecrets {
|
||||
fn get(&self, name: &str) -> Option<String> {
|
||||
self.0.get(name).map(|value| value.expose().to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
impl SecretSource for PythonSecrets {
|
||||
|
|
@ -41,23 +75,27 @@ impl SecretSource for PythonSecrets {
|
|||
&'a self,
|
||||
name: &'a str,
|
||||
) -> BoxFuture<'a, Result<Option<SecretValue>, Error>> {
|
||||
let get_secret_str = Arc::clone(&self.get_secret_str);
|
||||
let context = self.context.clone();
|
||||
let name = name.to_owned();
|
||||
Box::pin(async move {
|
||||
match attach_blocking(context, move |py| {
|
||||
get_secret_str
|
||||
.bind(py)
|
||||
.call1((name,))
|
||||
.and_then(|value| value.extract::<Option<String>>())
|
||||
.map(|value| value.map(SecretValue::new))
|
||||
.map_err(|error| external_error(py, error))
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(error) => Python::attach(|py| Err(external_error(py, error))),
|
||||
}
|
||||
Ok(self
|
||||
.read(vec![name.to_owned()])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
.flatten())
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve<'a>(&'a self, names: &'a [&str]) -> BoxFuture<'a, Result<Secrets, Error>> {
|
||||
Box::pin(async move {
|
||||
let values = self
|
||||
.read(names.iter().map(|name| (*name).to_owned()).collect())
|
||||
.await?;
|
||||
let found = names
|
||||
.iter()
|
||||
.zip(values)
|
||||
.filter_map(|(name, value)| Some(((*name).to_owned(), value?)))
|
||||
.collect();
|
||||
Ok(Arc::new(ReadSecrets(found)) as Secrets)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -84,14 +122,17 @@ import threading
|
|||
read_on = None
|
||||
request_var = contextvars.ContextVar('request_var', default=None)
|
||||
seen_context_values = []
|
||||
calls = []
|
||||
raised = KeyboardInterrupt('secret manager stopped')
|
||||
def get_secret_str(name):
|
||||
values = {'MISTRAL_API_KEY': 'vault-key', 'MISTRAL_API_BASE': 'https://mistral.internal'}
|
||||
def read_secrets(names):
|
||||
global read_on
|
||||
read_on = threading.get_ident()
|
||||
seen_context_values.append(request_var.get())
|
||||
if name == 'RAISING':
|
||||
calls.append(names)
|
||||
if 'RAISING' in names:
|
||||
raise raised
|
||||
return {'MISTRAL_API_KEY': 'vault-key'}.get(name)
|
||||
return tuple(values.get(name) for name in names)
|
||||
",
|
||||
Some(&namespace),
|
||||
None,
|
||||
|
|
@ -113,7 +154,7 @@ def get_secret_str(name):
|
|||
.unwrap();
|
||||
(
|
||||
namespace
|
||||
.get_item("get_secret_str")
|
||||
.get_item("read_secrets")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unbind(),
|
||||
|
|
@ -141,7 +182,7 @@ def get_secret_str(name):
|
|||
#[case::found("MISTRAL_API_KEY", Some("vault-key"))]
|
||||
#[case::missing("OTHER", None)]
|
||||
#[tokio::test]
|
||||
async fn returns_what_get_secret_str_returns(
|
||||
async fn returns_what_read_secrets_returns(
|
||||
secrets: (PythonSecrets, Py<PyDict>),
|
||||
#[case] name: &str,
|
||||
#[case] expected: Option<&str>,
|
||||
|
|
@ -151,12 +192,46 @@ def get_secret_str(name):
|
|||
assert_eq!(value.as_ref().map(|value| value.expose()), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn resolve_reads_every_name_in_one_python_call(secrets: (PythonSecrets, Py<PyDict>)) {
|
||||
let resolved = secrets
|
||||
.0
|
||||
.resolve(&["MISTRAL_API_KEY", "OTHER", "MISTRAL_API_BASE"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved.get("MISTRAL_API_KEY").as_deref(),
|
||||
Some("vault-key")
|
||||
);
|
||||
assert_eq!(resolved.get("OTHER"), None);
|
||||
assert_eq!(
|
||||
resolved.get("MISTRAL_API_BASE").as_deref(),
|
||||
Some("https://mistral.internal")
|
||||
);
|
||||
let calls: Vec<Vec<String>> = Python::attach(|py| global(&secrets.1, py, "calls"));
|
||||
assert_eq!(
|
||||
calls,
|
||||
vec![vec![
|
||||
"MISTRAL_API_KEY".to_owned(),
|
||||
"OTHER".to_owned(),
|
||||
"MISTRAL_API_BASE".to_owned()
|
||||
]]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn exceptions_surface_as_the_original_python_object(
|
||||
secrets: (PythonSecrets, Py<PyDict>),
|
||||
) {
|
||||
let error = secrets.0.get_secret_str("RAISING").await.unwrap_err();
|
||||
let error = secrets
|
||||
.0
|
||||
.resolve(&["MISTRAL_API_KEY", "RAISING"])
|
||||
.await
|
||||
.err()
|
||||
.unwrap();
|
||||
|
||||
Python::attach(|py| {
|
||||
let surfaced = python_error(py, &error).expect("the Python exception is preserved");
|
||||
|
|
@ -186,7 +261,7 @@ def get_secret_str(name):
|
|||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_see_the_callers_contextvars(secrets: (PythonSecrets, Py<PyDict>)) {
|
||||
secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap();
|
||||
secrets.0.resolve(&["MISTRAL_API_KEY"]).await.unwrap();
|
||||
|
||||
let seen: Vec<String> = Python::attach(|py| global(&secrets.1, py, "seen_context_values"));
|
||||
assert_eq!(seen, vec!["request-value".to_owned()]);
|
||||
|
|
|
|||
|
|
@ -1,277 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_host_python::PythonContext;
|
||||
use litellm_http::Client;
|
||||
use litellm_secret_value::SecretValue;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver};
|
||||
|
||||
use super::config::SecretManagerSnapshot;
|
||||
|
||||
pub(crate) struct ResolvedSecrets {
|
||||
resolver: SecretResolver,
|
||||
}
|
||||
|
||||
impl ResolvedSecrets {
|
||||
pub(crate) fn new(
|
||||
snapshot: SecretManagerSnapshot,
|
||||
context: PythonContext,
|
||||
client: Client,
|
||||
) -> Self {
|
||||
Self::from_state(snapshot.into_state(context), client)
|
||||
}
|
||||
|
||||
fn from_state(state: Arc<SecretManagerState>, client: Client) -> Self {
|
||||
Self {
|
||||
resolver: SecretResolver::new_python_compatible(
|
||||
state,
|
||||
Arc::new(ProcessEnvironment),
|
||||
OidcResolver::new(client),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SecretSource for ResolvedSecrets {
|
||||
fn get_secret_str<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
) -> BoxFuture<'a, Result<Option<SecretValue>, Error>> {
|
||||
Box::pin(self.resolver.get_secret_str(name, None))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aws_sdk_secretsmanager::Client;
|
||||
use aws_sdk_secretsmanager::config::{
|
||||
BehaviorVersion, Credentials, Region, retry::RetryConfig,
|
||||
};
|
||||
use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState};
|
||||
use litellm_secrets_aws::AwsSecretsManagerV2;
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_partial_json, header},
|
||||
};
|
||||
|
||||
use super::ResolvedSecrets;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
||||
fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc<SecretManagerState> {
|
||||
let client = Client::from_conf(
|
||||
aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.build(),
|
||||
);
|
||||
Arc::new(SecretManagerState::new(
|
||||
SecretManager::AwsSecretsManagerV2(AwsSecretsManagerV2::new(
|
||||
client,
|
||||
(&settings).into(),
|
||||
)),
|
||||
settings,
|
||||
))
|
||||
}
|
||||
|
||||
async fn resolve(state: Arc<SecretManagerState>, name: &'static str) -> Option<String> {
|
||||
ResolvedSecrets::from_state(state, litellm_http::Client::for_test())
|
||||
.resolve(&[name])
|
||||
.await
|
||||
.unwrap()
|
||||
.get(name)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hosted_key_miss_falls_back_to_environment() {
|
||||
let name = "LITELLM_RUST_BRIDGE_HOSTED_KEY_MISS";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": name})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
hosted_keys: Some(vec!["OTHER".into()]),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn aws_read_failure_preserves_absence_without_environment_fallback() {
|
||||
let name = "LITELLM_RUST_BRIDGE_MANAGER_FAILURE";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(state(&server, KeyManagementSettings::default()), name).await;
|
||||
let missing = resolve(
|
||||
state(&server, KeyManagementSettings::default()),
|
||||
"LITELLM_RUST_BRIDGE_MANAGER_FAILURE_MISSING",
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result, None);
|
||||
assert_eq!(missing, None);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::capitalized_true("True")]
|
||||
#[case::parenthesized_false("(False)")]
|
||||
#[tokio::test]
|
||||
async fn boolean_manager_values_are_absent_like_get_secret_str(#[case] value: &str) {
|
||||
let name = "LITELLM_RUST_BRIDGE_BOOLEAN_VALUE";
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString": value})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert_eq!(
|
||||
resolve(state(&server, KeyManagementSettings::default()), name).await,
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn undeclared_names_are_read_from_the_manager() {
|
||||
let declared = "LITELLM_RUST_BRIDGE_DECLARED";
|
||||
let undeclared = "LITELLM_RUST_BRIDGE_UNDECLARED_MANAGED";
|
||||
unsafe { std::env::set_var(undeclared, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": declared})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "declared-key"})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": undeclared})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let source = ResolvedSecrets::from_state(
|
||||
state(&server, KeyManagementSettings::default()),
|
||||
litellm_http::Client::for_test(),
|
||||
);
|
||||
let snapshot = source.resolve(&[declared]).await.unwrap();
|
||||
assert_eq!(snapshot.get(undeclared), None);
|
||||
let result = source
|
||||
.get_secret_str(undeclared)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|value| value.expose().to_owned());
|
||||
unsafe { std::env::remove_var(undeclared) };
|
||||
assert_eq!(result.as_deref(), Some("manager-key"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_only_mode_never_consults_the_manager() {
|
||||
let name = "LITELLM_RUST_BRIDGE_WRITE_ONLY";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
access_mode: AccessMode::WriteOnly,
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_only_mode_resolves_from_the_manager() {
|
||||
let name = "LITELLM_RUST_BRIDGE_READ_ONLY";
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": name})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert_eq!(
|
||||
resolve(state(&server, KeyManagementSettings::default()), name)
|
||||
.await
|
||||
.as_deref(),
|
||||
Some("manager-key")
|
||||
);
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oidc_failures_are_not_converted_to_missing_secrets() {
|
||||
let result = ResolvedSecrets::from_state(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
litellm_http::Client::for_test(),
|
||||
)
|
||||
.resolve(&["oidc/"])
|
||||
.await;
|
||||
assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn names_excluded_by_hosted_keys_read_the_process_environment() {
|
||||
let name = "LITELLM_RUST_BRIDGE_UNDECLARED";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
hosted_keys: Some(vec!["OTHER".into()]),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
}
|
||||
|
|
@ -7,14 +7,11 @@ use litellm_secrets::{
|
|||
read_secret_from_python_manager,
|
||||
};
|
||||
use litellm_secrets_types::PythonSecretRead;
|
||||
use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict};
|
||||
use pyo3::{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,
|
||||
settings: KeyManagementSettings,
|
||||
|
|
@ -90,100 +87,6 @@ impl NativeSecretManager {
|
|||
)
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
pub(super) fn from_client(client: &Bound<'_, PyAny>) -> PyResult<Option<Py<Self>>> {
|
||||
let py = client.py();
|
||||
if let Ok(native) = client.extract::<Py<Self>>() {
|
||||
native.borrow(py).backend()?;
|
||||
return Ok(Some(native));
|
||||
}
|
||||
let config = py
|
||||
.import(SECRET_MANAGER_MODULE)?
|
||||
.getattr("native_secret_manager_config")?
|
||||
.call1((client,))?;
|
||||
if config.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
if !config.getattr("owner_type")?.is(client.get_type()) {
|
||||
return Ok(None);
|
||||
}
|
||||
let methods = config
|
||||
.getattr("methods")?
|
||||
.extract::<Vec<(String, Py<PyAny>)>>()?;
|
||||
for (name, original) in methods {
|
||||
let current = client.getattr(name.as_str())?;
|
||||
let implementation = optional_attribute(¤t, "__func__")?.unwrap_or(current);
|
||||
if !implementation.is(original.bind(py)) {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
let environment_attributes: BTreeMap<String, String> = config
|
||||
.getattr("environment_attributes")?
|
||||
.extract::<Vec<(String, String)>>()?
|
||||
.into_iter()
|
||||
.collect();
|
||||
let captured = config
|
||||
.getattr("environment")?
|
||||
.extract::<Vec<(String, String)>>()?;
|
||||
let overrides = environment_attributes
|
||||
.iter()
|
||||
.map(|(key, attribute)| {
|
||||
let value = attribute_path(client, attribute)?;
|
||||
Ok((
|
||||
key.clone(),
|
||||
if value.is_none() {
|
||||
None
|
||||
} else {
|
||||
Some(value.str()?.extract::<String>()?)
|
||||
},
|
||||
))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let settings =
|
||||
from_py::<serde_json::Map<String, serde_json::Value>>(&config.getattr("settings")?)
|
||||
.map_err(ManagerError::settings)?;
|
||||
let attributes = config
|
||||
.getattr("settings_attributes")?
|
||||
.extract::<Vec<String>>()?;
|
||||
let setting_overrides = attributes
|
||||
.into_iter()
|
||||
.map(|name| {
|
||||
let value = from_py::<serde_json::Value>(&client.getattr(name.as_str())?)?;
|
||||
Ok((name, value))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let configuration = Configuration {
|
||||
system: serde_json::from_value(serde_json::Value::String(
|
||||
config.getattr("system")?.extract()?,
|
||||
))
|
||||
.map_err(ManagerError::UnknownSystem)?,
|
||||
settings: serde_json::from_value(serde_json::Value::Object(
|
||||
settings.into_iter().chain(setting_overrides).collect(),
|
||||
))
|
||||
.map_err(ManagerError::settings)?,
|
||||
environment: captured
|
||||
.into_iter()
|
||||
.filter(|(key, _)| !environment_attributes.contains_key(key))
|
||||
.chain(
|
||||
overrides
|
||||
.into_iter()
|
||||
.filter_map(|(key, value)| value.map(|value| (key, value))),
|
||||
)
|
||||
.collect(),
|
||||
enterprise_enabled: config.getattr("enterprise_enabled")?.extract()?,
|
||||
};
|
||||
if let Some(native) = cached(client, &configuration)? {
|
||||
return Ok(Some(native));
|
||||
}
|
||||
let runtime = Self::build(py, configuration)?;
|
||||
if let Some(native) = cached(client, &runtime.configuration)? {
|
||||
return Ok(Some(native));
|
||||
}
|
||||
let native = Py::new(py, runtime)?;
|
||||
client.setattr("_litellm_native_secret_manager", native.bind(py))?;
|
||||
Ok(Some(native))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn system(&self) -> &'static str {
|
||||
self.configuration.system.into()
|
||||
|
|
@ -395,37 +298,6 @@ impl NativeSecretManager {
|
|||
}
|
||||
}
|
||||
|
||||
fn optional_attribute<'py>(
|
||||
object: &Bound<'py, PyAny>,
|
||||
name: &str,
|
||||
) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
match object.getattr(name) {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(error) if error.is_instance_of::<PyAttributeError>(object.py()) => Ok(None),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn attribute_path<'py>(object: &Bound<'py, PyAny>, path: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
match path.split_once('.') {
|
||||
Some((head, tail)) => attribute_path(&object.getattr(head)?, tail),
|
||||
None => object.getattr(path),
|
||||
}
|
||||
}
|
||||
|
||||
fn cached(
|
||||
client: &Bound<'_, PyAny>,
|
||||
configuration: &Configuration,
|
||||
) -> PyResult<Option<Py<NativeSecretManager>>> {
|
||||
let Some(value) = optional_attribute(client, "_litellm_native_secret_manager")? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let native = value.extract::<Py<NativeSecretManager>>()?;
|
||||
let same_configuration = native.borrow(client.py()).pid == std::process::id()
|
||||
&& &native.borrow(client.py()).configuration == configuration;
|
||||
Ok(same_configuration.then_some(native))
|
||||
}
|
||||
|
||||
fn parse_settings(value: Option<&Bound<'_, PyAny>>) -> PyResult<KeyManagementSettings> {
|
||||
value
|
||||
.map(|value| {
|
||||
|
|
|
|||
|
|
@ -7764,7 +7764,6 @@ class ProxyConfig:
|
|||
"""
|
||||
Initialize the relevant secret manager if `key_management_system` is provided
|
||||
"""
|
||||
previous_client: Final[object] = litellm.secret_manager_client
|
||||
if key_management_system is not None:
|
||||
if key_management_system == KeyManagementSystem.AZURE_KEY_VAULT.value:
|
||||
### LOAD FROM AZURE KEY VAULT ###
|
||||
|
|
@ -7813,11 +7812,6 @@ class ProxyConfig:
|
|||
else:
|
||||
raise ValueError("Invalid Key Management System selected")
|
||||
|
||||
from litellm.rust_bridge.host.secret_manager import capture_secret_manager
|
||||
|
||||
if litellm.secret_manager_client is not previous_client:
|
||||
capture_secret_manager(litellm.secret_manager_client, key_management_system)
|
||||
|
||||
def get_model_info_with_id(self, model, db_model=False) -> RouterModelInfo:
|
||||
"""
|
||||
Common logic across add + delete router models
|
||||
|
|
|
|||
|
|
@ -252,8 +252,6 @@ class _SecretManagerRuntime:
|
|||
settings: Mapping[str, object] | None = None,
|
||||
enterprise_enabled: bool = False,
|
||||
) -> _SecretManagerRuntime: ...
|
||||
@staticmethod
|
||||
def from_client(client: object) -> _SecretManagerRuntime | None: ...
|
||||
@property
|
||||
def system(self) -> str: ...
|
||||
def read_secret(self, name: str, settings: Mapping[str, object] | None = None) -> JsonValue: ...
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
|
||||
|
||||
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
|
||||
- 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` (a batch of secrets read through `get_secret_str`), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
|
||||
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
|
||||
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
|
||||
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
|
||||
|
|
|
|||
|
|
@ -1,165 +1,7 @@
|
|||
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 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
|
||||
|
||||
|
||||
def read_secret_str(name: str) -> str | None:
|
||||
def read_secrets(names: tuple[str, ...]) -> tuple[str | None, ...]:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return get_secret_str(name)
|
||||
return tuple(get_secret_str(name) for name in names)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -31,83 +30,6 @@ class ProviderDefaults:
|
|||
enable_azure_ad_token_refresh: object
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManager:
|
||||
readable: bool
|
||||
native: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManagerBinding:
|
||||
system: object
|
||||
access_mode: object
|
||||
hosted_keys: object
|
||||
primary_secret_name: object
|
||||
store_virtual_keys: object
|
||||
prefix_for_stored_virtual_keys: object
|
||||
kms_key_id: object
|
||||
custom_secret_manager: object
|
||||
aws_region_name: object
|
||||
aws_role_name: object
|
||||
aws_session_name: object
|
||||
aws_external_id: object
|
||||
aws_profile_name: object
|
||||
aws_web_identity_token: object
|
||||
aws_sts_endpoint: object
|
||||
replica_regions: object
|
||||
client: object
|
||||
settings_object: object
|
||||
|
||||
|
||||
def secret_manager() -> SecretManager:
|
||||
import litellm
|
||||
from litellm.secret_managers.main import (
|
||||
_should_read_secret_from_secret_manager, # pyright: ignore[reportPrivateUsage] # canonical resolver is private
|
||||
)
|
||||
|
||||
readable: Final = _should_read_secret_from_secret_manager()
|
||||
system: Final = (
|
||||
litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
|
||||
)
|
||||
return SecretManager(readable=readable, native=readable and system is not None)
|
||||
|
||||
|
||||
def secret_manager_binding() -> SecretManagerBinding:
|
||||
import litellm
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings
|
||||
|
||||
configured_system: Final = (
|
||||
litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
|
||||
)
|
||||
configured_settings: Final = (
|
||||
litellm._key_management_settings # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
|
||||
)
|
||||
settings: Final = configured_settings or KeyManagementSettings()
|
||||
system: Final = (
|
||||
configured_system.value if litellm.secret_manager_client is not None and configured_system is not None else None
|
||||
)
|
||||
return SecretManagerBinding(
|
||||
system=system,
|
||||
access_mode=settings.access_mode,
|
||||
hosted_keys=settings.hosted_keys,
|
||||
primary_secret_name=settings.primary_secret_name,
|
||||
store_virtual_keys=settings.store_virtual_keys,
|
||||
prefix_for_stored_virtual_keys=settings.prefix_for_stored_virtual_keys,
|
||||
kms_key_id=settings.kms_key_id,
|
||||
custom_secret_manager=settings.custom_secret_manager,
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
aws_external_id=settings.aws_external_id,
|
||||
aws_profile_name=settings.aws_profile_name,
|
||||
aws_web_identity_token=settings.aws_web_identity_token,
|
||||
aws_sts_endpoint=settings.aws_sts_endpoint,
|
||||
replica_regions=settings.replica_regions,
|
||||
client=litellm.secret_manager_client,
|
||||
settings_object=configured_settings,
|
||||
)
|
||||
|
||||
|
||||
def provider_defaults() -> ProviderDefaults:
|
||||
import litellm
|
||||
|
||||
|
|
|
|||
|
|
@ -2553,7 +2553,7 @@ def test_ProxyConfig__load_environment_variables_blocks_dangerous_keys(monkeypat
|
|||
@pytest.mark.parametrize(
|
||||
("flag", "system"), (("use_google_kms", "google_kms"), ("use_azure_key_vault", "azure_key_vault"))
|
||||
)
|
||||
async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_client(
|
||||
async def test_load_config_legacy_secret_manager_flags_initialize_the_client(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, flag: str, system: str
|
||||
) -> None:
|
||||
if system == "azure_key_vault":
|
||||
|
|
@ -2561,8 +2561,6 @@ async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_c
|
|||
else:
|
||||
client_type: Final = pytest.importorskip("google.cloud.kms_v1").KeyManagementServiceClient
|
||||
|
||||
from litellm.rust_bridge.host.secret_manager import native_secret_manager_config
|
||||
|
||||
credentials_file: Final = tmp_path / "credentials.json"
|
||||
credentials_file.write_text(
|
||||
json.dumps(
|
||||
|
|
@ -2594,10 +2592,6 @@ async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_c
|
|||
client: Final = litellm.secret_manager_client
|
||||
assert isinstance(client, client_type)
|
||||
try:
|
||||
captured: Final = native_secret_manager_config(client)
|
||||
assert captured is not None
|
||||
assert captured.system == system
|
||||
assert dict(captured.environment)["GOOGLE_APPLICATION_CREDENTIALS"] == str(credentials_file)
|
||||
assert litellm._key_management_system is not None
|
||||
assert litellm._key_management_system.value == system
|
||||
finally:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
Test what each side of the bridge does, not the rollout policy that picks a side. `LITELLM_RUST` and `catalog.RULES` change every time a route or backend rolls forward, so a test that sets the env var or patches the catalog to reach a path goes red on a policy change even when the code under test is fine
|
||||
|
||||
Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting such as `settings.secret_manager().native`, pin that field in the test instead of deriving it from the catalog. `ocr/test_secrets.py` shows the pattern
|
||||
Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting, pin that field in the test instead of deriving it from the catalog
|
||||
|
||||
Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py` and `test_configuration.py`, tested against rules the test builds rather than the shipped `catalog.RULES`. Which path a public call takes under those rules is a dispatch test and lives in `tests/test_litellm_rust/` for now
|
||||
|
||||
|
|
|
|||
|
|
@ -3,9 +3,10 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from functools import partial
|
||||
from types import SimpleNamespace
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
|
|
@ -18,7 +19,6 @@ 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.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
|
||||
|
|
@ -67,10 +67,24 @@ def _vault_body(value: str) -> dict[str, object]:
|
|||
}
|
||||
|
||||
|
||||
def _native_handle(client: object) -> "_SecretManagerRuntime":
|
||||
handle: Final = native._SecretManagerRuntime.from_client(client)
|
||||
assert handle is not None
|
||||
return handle
|
||||
_NATIVE_SYSTEMS: Final = MappingProxyType(
|
||||
{
|
||||
AWSSecretsManagerV2: KeyManagementSystem.AWS_SECRET_MANAGER,
|
||||
HashicorpSecretManager: KeyManagementSystem.HASHICORP_VAULT,
|
||||
CyberArkSecretManager: KeyManagementSystem.CYBERARK,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _native_handle(
|
||||
client: AWSSecretsManagerV2 | HashicorpSecretManager | CyberArkSecretManager,
|
||||
) -> _SecretManagerRuntime:
|
||||
return native._SecretManagerRuntime.from_config(
|
||||
_NATIVE_SYSTEMS[type(client)].value,
|
||||
{name: value for name, value in os.environ.items() if name.startswith(("AWS_", "HCP_VAULT_", "CYBERARK_"))},
|
||||
KeyManagementSettings(aws_region_name="us-east-1").model_dump(mode="json"),
|
||||
enterprise_enabled=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark"))
|
||||
|
|
@ -107,64 +121,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s
|
|||
assert handle.read_secret("KEY") == expected
|
||||
|
||||
|
||||
def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": "native-value"})
|
||||
server.expected_requests = 2
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "captured-access")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "captured-secret")
|
||||
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", server.base_url)
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(aws_region_name="us-east-1"))
|
||||
ProxyConfig().initialize_secret_manager(KeyManagementSystem.AWS_SECRET_MANAGER.value)
|
||||
manager: Final = litellm.secret_manager_client
|
||||
assert isinstance(manager, AWSSecretsManagerV2)
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "changed-access")
|
||||
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "http://127.0.0.1:1")
|
||||
first: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert first is not None
|
||||
assert native._SecretManagerRuntime.from_client(manager) is first
|
||||
assert first.read_secret("KEY") == "native-value"
|
||||
manager.aws_region_name = "us-west-2"
|
||||
second: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert second is not None
|
||||
assert second is not first
|
||||
assert second.read_secret("KEY") == "native-value"
|
||||
assert "Credential=captured-access/" in server.requests[0].headers["authorization"]
|
||||
assert "/us-east-1/" in server.requests[0].headers["authorization"]
|
||||
assert "/us-west-2/" in server.requests[1].headers["authorization"]
|
||||
|
||||
|
||||
def test_configuration_replacement_rebuilds_without_invalidating_existing_handles(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
with recording_service() as first_server, recording_service() as second_server:
|
||||
first_server.default_response = ResponseSpec(body=_vault_body("first"))
|
||||
second_server.default_response = ResponseSpec(body=_vault_body("second"))
|
||||
manager: Final = _vault(monkeypatch, first_server.base_url)
|
||||
first: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert first is not None
|
||||
assert first.read_secret("KEY") == "first"
|
||||
manager.vault_addr = second_server.base_url
|
||||
second: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert second is not None
|
||||
assert second is not first
|
||||
assert second.read_secret("KEY") == "second"
|
||||
assert first.read_secret("KEY") == "first"
|
||||
|
||||
|
||||
def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
|
||||
|
||||
class CustomManager(AWSSecretsManagerV2):
|
||||
def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str:
|
||||
return f"custom:{secret_name}"
|
||||
|
||||
manager: Final = CustomManager(aws_region_name="us-east-1")
|
||||
assert native._SecretManagerRuntime.from_client(manager) is None
|
||||
assert manager.sync_read_secret("KEY") == "custom:KEY"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark"))
|
||||
def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None:
|
||||
with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"):
|
||||
|
|
@ -176,19 +132,6 @@ def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None:
|
|||
native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"})
|
||||
|
||||
|
||||
def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=_vault_body("native-value"))
|
||||
manager: Final = _vault(monkeypatch, server.base_url)
|
||||
handle: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert handle is not None
|
||||
assert handle.read_secret("KEY") == "native-value"
|
||||
|
||||
|
||||
def test_uncaptured_clients_have_no_native_handle() -> None:
|
||||
assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None
|
||||
|
||||
|
||||
def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]})
|
||||
|
|
@ -200,62 +143,6 @@ def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatc
|
|||
assert len(server.requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit_capture", (False, True))
|
||||
def test_config_capture_preserves_credentials_and_excludes_unrelated_environment(
|
||||
monkeypatch: pytest.MonkeyPatch, explicit_capture: bool
|
||||
) -> None:
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "initial-access")
|
||||
monkeypatch.setenv("SECRET_MANAGER_REFRESH_INTERVAL", "45")
|
||||
monkeypatch.setenv("UNRELATED_PRIVATE_TOKEN", "unrelated-secret")
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
if explicit_capture:
|
||||
capture_secret_manager(manager, "aws_secret_manager")
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "replacement-access")
|
||||
captured: Final = native_secret_manager_config(manager)
|
||||
assert captured is not None
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "replacement-access")
|
||||
|
||||
retained: Final = native_secret_manager_config(manager)
|
||||
|
||||
assert retained is captured
|
||||
assert retained.system == "aws_secret_manager"
|
||||
assert dict(retained.environment)["AWS_ACCESS_KEY_ID"] == "initial-access"
|
||||
assert dict(retained.environment)["SECRET_MANAGER_REFRESH_INTERVAL"] == "45"
|
||||
assert "UNRELATED_PRIVATE_TOKEN" not in dict(retained.environment)
|
||||
assert "initial-access" not in repr(retained)
|
||||
assert retained.settings == KeyManagementSettings().model_dump(mode="json")
|
||||
|
||||
|
||||
def test_kms_sdk_client_capture_preserves_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import boto3
|
||||
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "captured-region")
|
||||
client: Final = boto3.client(
|
||||
"kms", region_name="us-east-1", aws_access_key_id="access", aws_secret_access_key="secret"
|
||||
)
|
||||
try:
|
||||
capture_secret_manager(client, "aws_kms")
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "replacement-region")
|
||||
|
||||
captured: Final = native_secret_manager_config(client)
|
||||
|
||||
assert captured is not None
|
||||
assert captured.system == "aws_kms"
|
||||
assert dict(captured.environment)["AWS_REGION_NAME"] == "captured-region"
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_same_named_custom_client_is_not_captured() -> None:
|
||||
class AWSSecretsManagerV2:
|
||||
pass
|
||||
|
||||
client: Final = AWSSecretsManagerV2()
|
||||
capture_secret_manager(client, "aws_secret_manager")
|
||||
|
||||
assert native_secret_manager_config(client) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None))
|
||||
async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None:
|
||||
with recording_service() as server:
|
||||
|
|
@ -267,8 +154,7 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke
|
|||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
settings: Final = KeyManagementSettings(primary_secret_name="primary")
|
||||
reference: Final = get_secret_from_manager(manager, "aws_secret_manager", "KEY", settings)
|
||||
handle: Final = native._SecretManagerRuntime.from_client(manager)
|
||||
assert handle is not None
|
||||
handle: Final = _native_handle(manager)
|
||||
actual: Final = handle.read_secret("KEY", settings.model_dump(mode="json"))
|
||||
asynchronous: Final = await handle.read_secret_async("KEY", settings.model_dump(mode="json"))
|
||||
assert type(actual) is type(reference) is type(value)
|
||||
|
|
@ -295,7 +181,9 @@ async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyP
|
|||
asynchronous: Final = await handle.read_secret_async("KEY", settings_dump)
|
||||
assert synchronous == reference
|
||||
assert asynchronous == reference
|
||||
assert tuple(json.loads(request.raw_body) for request in server.requests) == ({"SecretId": primary or "KEY"},) * 3
|
||||
assert (
|
||||
tuple(json.loads(request.raw_body) for request in server.requests) == ({"SecretId": primary or "KEY"},) * 3
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("primary", (None, "primary"))
|
||||
|
|
|
|||
|
|
@ -1,14 +1,10 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.custom_httpx.http_handler import default_user_agent
|
||||
from litellm.rust_bridge.host import settings
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
|
||||
|
||||
def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -57,143 +53,6 @@ def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyP
|
|||
assert result.ssl_verify is True
|
||||
|
||||
|
||||
class _VaultSecrets(CustomSecretManager):
|
||||
def __init__(self, secrets: dict[str, str]) -> None:
|
||||
super().__init__(secret_manager_name="rust_bridge_settings_test")
|
||||
self.secrets = secrets
|
||||
|
||||
async def async_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> str | None:
|
||||
return self.secrets.get(secret_name)
|
||||
|
||||
def sync_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> str | None:
|
||||
return self.secrets.get(secret_name)
|
||||
|
||||
|
||||
@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", "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=access_mode))
|
||||
|
||||
assert settings.secret_manager() == settings.SecretManager(readable=readable, native=readable)
|
||||
assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("system", "access_mode", "native"),
|
||||
[
|
||||
(KeyManagementSystem.CUSTOM, "read_only", True),
|
||||
(KeyManagementSystem.AWS_SECRET_MANAGER, "read_only", True),
|
||||
(KeyManagementSystem.CUSTOM, "write_only", False),
|
||||
(None, "read_only", False),
|
||||
],
|
||||
)
|
||||
def test_secret_manager_is_native_only_when_a_readable_system_is_configured(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
system: KeyManagementSystem | None,
|
||||
access_mode: str,
|
||||
native: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({}))
|
||||
monkeypatch.setattr(litellm, "_key_management_system", system)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
|
||||
|
||||
assert settings.secret_manager() == settings.SecretManager(
|
||||
readable=access_mode != "write_only", native=native
|
||||
)
|
||||
|
||||
|
||||
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", None)
|
||||
|
||||
assert settings.secret_manager() == settings.SecretManager(readable=False, native=False)
|
||||
|
||||
|
||||
def test_secret_manager_projects_custom_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
manager_settings: Final = KeyManagementSettings(
|
||||
access_mode="read_and_write",
|
||||
hosted_keys=["MISTRAL_API_KEY"],
|
||||
primary_secret_name="primary",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
client: Final = _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", client)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", manager_settings)
|
||||
|
||||
assert settings.secret_manager_binding() == settings.SecretManagerBinding(
|
||||
system="custom",
|
||||
access_mode="read_and_write",
|
||||
hosted_keys=["MISTRAL_API_KEY"],
|
||||
primary_secret_name="primary",
|
||||
store_virtual_keys=manager_settings.store_virtual_keys,
|
||||
prefix_for_stored_virtual_keys=manager_settings.prefix_for_stored_virtual_keys,
|
||||
kms_key_id=manager_settings.kms_key_id,
|
||||
custom_secret_manager=manager_settings.custom_secret_manager,
|
||||
aws_region_name="us-east-1",
|
||||
aws_role_name=manager_settings.aws_role_name,
|
||||
aws_session_name=manager_settings.aws_session_name,
|
||||
aws_external_id=manager_settings.aws_external_id,
|
||||
aws_profile_name=manager_settings.aws_profile_name,
|
||||
aws_web_identity_token=manager_settings.aws_web_identity_token,
|
||||
aws_sts_endpoint=manager_settings.aws_sts_endpoint,
|
||||
replica_regions=manager_settings.replica_regions,
|
||||
client=client,
|
||||
settings_object=manager_settings,
|
||||
)
|
||||
|
||||
|
||||
def test_secret_manager_without_a_client_has_no_system(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", None)
|
||||
|
||||
assert settings.secret_manager_binding().system is None
|
||||
|
||||
|
||||
def test_secret_manager_uses_key_management_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", None)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", None)
|
||||
|
||||
defaults: Final = KeyManagementSettings()
|
||||
result: Final = settings.secret_manager_binding()
|
||||
|
||||
assert result == settings.SecretManagerBinding(
|
||||
system=None,
|
||||
access_mode=defaults.access_mode,
|
||||
hosted_keys=defaults.hosted_keys,
|
||||
primary_secret_name=defaults.primary_secret_name,
|
||||
store_virtual_keys=defaults.store_virtual_keys,
|
||||
prefix_for_stored_virtual_keys=defaults.prefix_for_stored_virtual_keys,
|
||||
kms_key_id=defaults.kms_key_id,
|
||||
custom_secret_manager=defaults.custom_secret_manager,
|
||||
aws_region_name=defaults.aws_region_name,
|
||||
aws_role_name=defaults.aws_role_name,
|
||||
aws_session_name=defaults.aws_session_name,
|
||||
aws_external_id=defaults.aws_external_id,
|
||||
aws_profile_name=defaults.aws_profile_name,
|
||||
aws_web_identity_token=defaults.aws_web_identity_token,
|
||||
aws_sts_endpoint=defaults.aws_sts_endpoint,
|
||||
replica_regions=defaults.replica_regions,
|
||||
client=None,
|
||||
settings_object=None,
|
||||
)
|
||||
|
||||
|
||||
def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "vertex_project", "configured-project")
|
||||
monkeypatch.setattr(litellm, "vertex_location", "europe-west4")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from dataclasses import replace
|
||||
from typing import Final, Protocol, cast # noqa: TID251 # narrows the parametrized path to its protocol
|
||||
|
||||
import httpx
|
||||
|
|
@ -10,7 +9,6 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.anthropic.messages.native.handler import anthropic_messages
|
||||
from litellm.rust_bridge.host import settings
|
||||
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
|
|
@ -97,8 +95,6 @@ async def test_secret_manager_supplies_the_anthropic_key_and_base(
|
|||
)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
configured: Final = settings.secret_manager
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=True))
|
||||
|
||||
await messages()
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
from collections.abc import Awaitable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import replace
|
||||
from typing import Final, Literal, Protocol, TypeAlias, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -12,8 +11,6 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.rust_bridge import _native as native
|
||||
from litellm.rust_bridge.host import settings
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
|
|
@ -119,7 +116,6 @@ def _configure(
|
|||
*,
|
||||
manager: _VaultSecrets,
|
||||
key_management: KeyManagementSettings,
|
||||
native_secret_manager: bool = True,
|
||||
environment_key: str | None = "environment-key",
|
||||
) -> None:
|
||||
if environment_key is None:
|
||||
|
|
@ -129,8 +125,6 @@ def _configure(
|
|||
monkeypatch.setattr(litellm, "secret_manager_client", manager)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", key_management)
|
||||
configured: Final = settings.secret_manager
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=native_secret_manager))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -220,7 +214,7 @@ async def test_rust_reads_a_python_only_secret_manager_through_python(
|
|||
) -> None:
|
||||
manager: Final = _VaultSecrets()
|
||||
key_management: Final = KeyManagementSettings(access_mode="read_only")
|
||||
_configure(monkeypatch, manager=manager, key_management=key_management, native_secret_manager=False)
|
||||
_configure(monkeypatch, manager=manager, key_management=key_management)
|
||||
|
||||
with _mistral_service() as server:
|
||||
await rust_ocr(server.base_url)
|
||||
|
|
@ -245,7 +239,7 @@ class _PythonManagerReads:
|
|||
return "python-key" if secret_name == "MISTRAL_API_KEY" else None
|
||||
|
||||
|
||||
async def test_python_only_builtin_manager_is_read_through_python_not_its_native_backend(
|
||||
async def test_builtin_manager_is_read_through_python_not_its_native_backend(
|
||||
monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr
|
||||
) -> None:
|
||||
from litellm.secret_managers import main as secret_manager_main
|
||||
|
|
@ -255,19 +249,18 @@ async def test_python_only_builtin_manager_is_read_through_python_not_its_native
|
|||
monkeypatch.setattr(secret_manager_main, "get_secret_from_manager", python_reads)
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "native-access")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "native-secret")
|
||||
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "http://127.0.0.1:1")
|
||||
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", manager)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(hosted_keys=["MISTRAL_API_KEY"]))
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: settings.SecretManager(readable=True, native=False))
|
||||
|
||||
with _mistral_service() as provider:
|
||||
await rust_ocr(provider.base_url)
|
||||
|
||||
assert provider.requests[0].headers["authorization"] == "Bearer python-key"
|
||||
assert (manager, "aws_secret_manager", "MISTRAL_API_KEY") in python_reads.reads
|
||||
assert getattr(manager, "_litellm_native_secret_manager", None) is None
|
||||
|
||||
|
||||
async def test_no_secret_client_leaves_dormant_binding_settings_unread(
|
||||
|
|
@ -341,8 +334,6 @@ def _configure_client(
|
|||
monkeypatch.setattr(litellm, "secret_manager_client", client)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", system)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", key_management)
|
||||
configured: Final = settings.secret_manager
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=True))
|
||||
|
||||
|
||||
async def _assert_missing_key(ocr: Ocr) -> None:
|
||||
|
|
@ -401,7 +392,9 @@ async def test_manager_without_the_key_does_not_fall_back_to_the_environment(
|
|||
await _assert_missing_key(ocr)
|
||||
|
||||
|
||||
async def test_rust_hosted_keys_exclude_azure_sdk_clients_too(monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr) -> None:
|
||||
async def test_azure_sdk_clients_are_read_even_for_names_outside_hosted_keys(
|
||||
monkeypatch: pytest.MonkeyPatch, ocr: Ocr
|
||||
) -> None:
|
||||
_configure_client(
|
||||
monkeypatch,
|
||||
client=_azure_sdk_client("vault-key"),
|
||||
|
|
@ -410,11 +403,9 @@ async def test_rust_hosted_keys_exclude_azure_sdk_clients_too(monkeypatch: pytes
|
|||
)
|
||||
|
||||
with _mistral_service() as server:
|
||||
await rust_ocr(server.base_url)
|
||||
await ocr(server.base_url)
|
||||
|
||||
assert server.requests[0].headers["authorization"] == "Bearer environment-key", (
|
||||
"recorded divergence: Python's get_secret_from_manager recognizes Azure SDK clients by type and ignores hosted_keys"
|
||||
)
|
||||
assert server.requests[0].headers["authorization"] == "Bearer vault-key"
|
||||
|
||||
|
||||
async def test_custom_system_with_a_foreign_client_falls_back_to_the_environment(
|
||||
|
|
@ -431,41 +422,3 @@ async def test_custom_system_with_a_foreign_client_falls_back_to_the_environment
|
|||
await ocr(server.base_url)
|
||||
|
||||
assert server.requests[0].headers["authorization"] == "Bearer environment-key"
|
||||
|
||||
|
||||
async def test_native_backend_supplies_ocr_credentials_without_a_python_reader(
|
||||
monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr
|
||||
) -> None:
|
||||
from litellm.secret_managers import main as secret_manager_main
|
||||
from litellm.secret_managers import secret_manager_handler
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
|
||||
|
||||
def reject_python_read(
|
||||
client: object,
|
||||
key_manager: str,
|
||||
secret_name: str,
|
||||
key_management_settings: KeyManagementSettings | None = None,
|
||||
) -> str | None:
|
||||
raise AssertionError("Rust must read the native backend directly")
|
||||
|
||||
monkeypatch.setattr(secret_manager_handler, "get_secret_from_manager", reject_python_read)
|
||||
monkeypatch.setattr(secret_manager_main, "get_secret_from_manager", reject_python_read)
|
||||
with recording_service() as secrets, _mistral_service(expected_requests=2) as provider:
|
||||
secrets.default_response = ResponseSpec(body={"SecretString": "native-key"})
|
||||
secrets.expected_requests = None
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "native-access")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "native-secret")
|
||||
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", secrets.base_url)
|
||||
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", manager)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(hosted_keys=["MISTRAL_API_KEY"]))
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: settings.SecretManager(readable=True, native=True))
|
||||
await rust_ocr(provider.base_url)
|
||||
await rust_ocr(provider.base_url)
|
||||
|
||||
assert len(secrets.requests) == 2, [(request.path, request.body) for request in secrets.requests]
|
||||
assert all(request.headers["authorization"] == "Bearer native-key" for request in provider.requests)
|
||||
assert all("Credential=native-access/" in request.headers["authorization"] for request in secrets.requests)
|
||||
assert native._SecretManagerRuntime.from_client(manager) is getattr(manager, "_litellm_native_secret_manager")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue