scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

refactor(rust_bridge): remove rule-gated native secret-manager selection (#44906)

Browse source 
* refactor(rust_bridge): remove rule-gated native secret-manager selection

Mirror the cache treatment: SecretManagerRule/SecretManagerContext and the resolve_native_* binding plumbing are gone. The Rust bridge now selects the native backend from the explicitly configured client (capture_secret_manager / _SecretManagerRuntime.from_client), and get_secret_from_manager is the plain Python handler path.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore(deps): bump sharp to 0.35.5 for GHSA-wq5f-xc86-pv6w

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Yujong Lee <yujong@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-06 19:14:55 +00:00 • committed by GitHub
parent fb554d34f3
commit b6d587d23e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 155 additions and 891 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/crates/python-bridge/src/secrets/mod.rs
View file

@ -21,10 +21,10 @@ use crate::{coercion::FieldSpec, python_settings::PythonSettings};
const NATIVE: FieldSpec<bool> = FieldSpec::new("native", |field| field.schema_bool());
/// Where a Rust route reads provider secrets from. Python's `get_secret_str` until a
/// `SecretManagerRule` in `catalog.py` moves the configured system off `PYTHON_ONLY`, then the
/// native secret manager. 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` 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.
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(

30
litellm/rust_bridge/catalog.py
View file

@ -1,4 +1,4 @@
"""Ordered rollout policy for routes, loggers, and secret managers.
"""Ordered rollout policy for routes and loggers.
The first matching rule wins; unmatched contexts stay on Python. Native
admission separately decides whether the selected implementation can execute.
@ -12,7 +12,6 @@ from typing import Final, TypeAlias
from litellm.rust_bridge.configuration import Decision, Rollout
from litellm.rust_bridge.configuration import decision as _decision
from litellm.types.secret_managers.main import KeyManagementSystem
class Route(str, Enum):
@ -49,20 +48,6 @@ class RouteRule:
)
@dataclass(frozen=True, slots=True)
class SecretManagerContext:
system: str
@dataclass(frozen=True, slots=True)
class SecretManagerRule:
rollout: Rollout
systems: frozenset[str] | None = None
def matches(self, context: Context) -> bool:
return isinstance(context, SecretManagerContext) and (self.systems is None or context.system in self.systems)
@dataclass(frozen=True, slots=True)
class LoggerContext:
pass
@ -76,8 +61,8 @@ class LoggerRule:
return isinstance(context, LoggerContext)
Context: TypeAlias = RouteContext | SecretManagerContext | LoggerContext
Rule: TypeAlias = RouteRule | SecretManagerRule | LoggerRule
Context: TypeAlias = RouteContext | LoggerContext
Rule: TypeAlias = RouteRule | LoggerRule
Rules: TypeAlias = tuple[Rule, ...]
RULES: Final[Rules] = (
@ -91,15 +76,6 @@ RULES: Final[Rules] = (
RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY),
RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY),
RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_KMS.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AZURE_KEY_VAULT.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_SECRET_MANAGER.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_SECRET_MANAGER.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.HASHICORP_VAULT.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CYBERARK.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.LOCAL.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_KMS.value})),
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CUSTOM.value})),
)

130
litellm/rust_bridge/secret_manager.py
View file

@ -1,17 +1,11 @@
from __future__ import annotations
import os
from collections.abc import Awaitable, Mapping
from collections.abc import Mapping
from dataclasses import dataclass, field
from importlib import import_module
from typing import Final, Protocol, runtime_checkable
from typing import Final
import httpx
from pydantic import JsonValue
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Rules, SecretManagerContext, decision
from litellm.rust_bridge.configuration import Decision
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -163,123 +157,3 @@ def native_secret_manager_config(client: object) -> NativeSecretManagerConfig |
):
return _capture(client, adapter)
return None
class NativeSecretManagerRuntime(Protocol):
@property
def system(self) -> str: ...
def read_secret(self, name: str, settings: Mapping[str, object] | None = None) -> JsonValue: ...
@runtime_checkable
class NativeSecretManagerFactory(Protocol):
@staticmethod
def from_client(client: object) -> NativeSecretManagerRuntime | None: ...
def _factory(value: object) -> NativeSecretManagerFactory | None:
return value if isinstance(value, NativeSecretManagerFactory) and callable(value.from_client) else None
NATIVE_SECRET_MANAGER: Final = NativeBinding("_SecretManagerRuntime", validate=_factory)
def resolve_native_secret_manager(
client: object,
system: str,
rules: Rules | None = None,
*,
binding: NativeBinding[NativeSecretManagerFactory] = NATIVE_SECRET_MANAGER,
) -> NativeSecretManagerRuntime | None:
if system in ("custom", "local"):
return None
selected: Final = decision(SecretManagerContext(system=system), rules)
if selected is Decision.PYTHON:
return None
factory: Final = binding.load()
if factory is None:
if selected is Decision.RUST_REQUIRED:
raise RuntimeError("Rust secret manager runtime is unavailable")
return None
runtime: Final = factory.from_client(client)
if runtime is not None and runtime.system != system:
raise ValueError("Native secret manager system does not match configuration")
return runtime
@runtime_checkable
class NativeProviderReader(Protocol):
def sync_read_secret(
self,
secret_name: str,
optional_params: Mapping[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None: ...
def async_read_secret(
self,
secret_name: str,
optional_params: Mapping[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> Awaitable[str | None]: ...
def resolve_native_provider_reader(
client: object,
system: str,
rules: Rules | None = None,
*,
binding: NativeBinding[NativeSecretManagerFactory] = NATIVE_SECRET_MANAGER,
) -> NativeProviderReader | None:
runtime: Final = resolve_native_secret_manager(client, system, rules, binding=binding)
if runtime is None:
return None
if not isinstance(runtime, NativeProviderReader):
raise TypeError("Rust secret manager provider reads are unavailable")
return runtime
@runtime_checkable
class NativeProviderWriter(Protocol):
def async_write_secret(
self,
secret_name: str,
secret_value: str,
description: str | None = None,
optional_params: Mapping[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
tags: object = None,
) -> Awaitable[dict[str, JsonValue]]: ...
def async_delete_secret(
self,
secret_name: str,
recovery_window_in_days: int | None = 7,
optional_params: Mapping[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> Awaitable[dict[str, JsonValue]]: ...
def async_rotate_secret(
self,
current_secret_name: str,
new_secret_name: str,
new_secret_value: str,
optional_params: Mapping[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> Awaitable[dict[str, JsonValue]]: ...
def resolve_native_provider_writer(
client: object,
system: str,
rules: Rules | None = None,
*,
binding: NativeBinding[NativeSecretManagerFactory] = NATIVE_SECRET_MANAGER,
) -> NativeProviderWriter | None:
runtime: Final = resolve_native_secret_manager(client, system, rules, binding=binding)
if runtime is None:
return None
if not isinstance(runtime, NativeProviderWriter):
raise TypeError("Rust secret manager provider writes are unavailable")
return runtime

16
litellm/rust_bridge/settings.py
View file

@ -1,10 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from litellm.rust_bridge.catalog import Rules
from typing import Final
@dataclass(frozen=True, slots=True)
@ -62,10 +59,8 @@ class SecretManagerBinding:
settings_object: object
def secret_manager(rules: Rules | None = None) -> SecretManager:
def secret_manager() -> SecretManager:
import litellm
from litellm.rust_bridge.catalog import SecretManagerContext, decision
from litellm.rust_bridge.configuration import Decision
from litellm.secret_managers.main import (
_should_read_secret_from_secret_manager, # pyright: ignore[reportPrivateUsage] # canonical resolver is private
)
@ -74,12 +69,7 @@ def secret_manager(rules: Rules | None = None) -> SecretManager:
system: Final = (
litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
)
native: Final = (
readable
and system is not None
and decision(SecretManagerContext(system=system.value), rules) is not Decision.PYTHON
)
return SecretManager(readable=readable, native=native)
return SecretManager(readable=readable, native=readable and system is not None)
def secret_manager_binding() -> SecretManagerBinding:

9
litellm/secret_managers/aws_secret_manager_v2.py
View file

@ -30,7 +30,6 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
from litellm.rust_bridge.secret_manager import resolve_native_provider_reader
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -140,10 +139,6 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
secret_name=secret_name, primary_secret_name=primary_secret_name
)
native: Final = resolve_native_provider_reader(self, "aws_secret_manager")
if native is not None:
return await native.async_read_secret(secret_name, optional_params, timeout)
endpoint_url, headers, body = self._prepare_request(
action="GetSecretValue",
secret_name=secret_name,
@ -196,10 +191,6 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
secret_name=secret_name, primary_secret_name=primary_secret_name
)
native: Final = resolve_native_provider_reader(self, "aws_secret_manager")
if native is not None:
return native.sync_read_secret(secret_name, optional_params, timeout)
endpoint_url, headers, body = self._prepare_request(
action="GetSecretValue",
secret_name=secret_name,

28
litellm/secret_managers/cyberark_secret_manager.py
View file

@ -16,7 +16,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.rust_bridge.secret_manager import resolve_native_provider_reader, resolve_native_provider_writer
from litellm.types.secret_managers.main import KeyManagementSystem
from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name
@ -199,10 +198,6 @@ class CyberArkSecretManager(BaseSecretManager):
Returns:
Optional[str]: The secret value if found, None otherwise
"""
native: Final = resolve_native_provider_reader(self, "cyberark")
if native is not None:
return await native.async_read_secret(secret_name, optional_params, timeout)
# Check cache first
if self.cache.get_cache(secret_name) is not None:
return self.cache.get_cache(secret_name)
@ -249,10 +244,6 @@ class CyberArkSecretManager(BaseSecretManager):
Returns:
Optional[str]: The secret value if found, None otherwise
"""
native: Final = resolve_native_provider_reader(self, "cyberark")
if native is not None:
return native.sync_read_secret(secret_name, optional_params, timeout)
# Check cache first
if self.cache.get_cache(secret_name) is not None:
return self.cache.get_cache(secret_name)
@ -302,12 +293,6 @@ class CyberArkSecretManager(BaseSecretManager):
Returns:
dict: Response containing status and details of the operation
"""
native: Final = resolve_native_provider_writer(self, "cyberark")
if native is not None:
return await native.async_write_secret(
secret_name, secret_value, description, optional_params, timeout, tags
)
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"ssl_verify": self.ssl_verify},
@ -353,10 +338,6 @@ class CyberArkSecretManager(BaseSecretManager):
Returns:
dict: Response indicating operation not supported
"""
native: Final = resolve_native_provider_writer(self, "cyberark")
if native is not None:
return await native.async_delete_secret(secret_name, recovery_window_in_days, optional_params, timeout)
verbose_logger.warning(
"CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates."
)
@ -377,15 +358,6 @@ class CyberArkSecretManager(BaseSecretManager):
optional_params: dict | None = None,
timeout: float | httpx.Timeout | None = None,
) -> dict:
native: Final = resolve_native_provider_writer(self, "cyberark")
if native is not None:
return await native.async_rotate_secret(
current_secret_name,
new_secret_name,
new_secret_value,
optional_params,
timeout,
)
return await super().async_rotate_secret(
current_secret_name,
new_secret_name,

30
litellm/secret_managers/dispatch.py
View file

@ -1,30 +0,0 @@
from typing import Final
from pydantic import JsonValue
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Rules
from litellm.rust_bridge.secret_manager import (
NATIVE_SECRET_MANAGER,
NativeSecretManagerFactory,
resolve_native_secret_manager,
)
from litellm.secret_managers.secret_manager_handler import get_secret_from_manager as python_get_secret_from_manager
from litellm.types.secret_managers.main import KeyManagementSettings
def get_secret_from_manager(
client: object,
key_manager: str,
secret_name: str,
key_management_settings: KeyManagementSettings | None = None,
*,
rules: Rules | None = None,
binding: NativeBinding[NativeSecretManagerFactory] = NATIVE_SECRET_MANAGER,
) -> JsonValue:
native: Final = resolve_native_secret_manager(client, key_manager, rules, binding=binding)
if native is None:
return python_get_secret_from_manager(client, key_manager, secret_name, key_management_settings)
return native.read_secret(
secret_name, key_management_settings.model_dump(mode="json") if key_management_settings is not None else None
)

5
litellm/secret_managers/google_secret_manager.py
View file

@ -9,7 +9,6 @@ from litellm.constants import SECRET_MANAGER_REFRESH_INTERVAL
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.proxy._types import CommonProxyErrors
from litellm.rust_bridge.secret_manager import resolve_native_provider_reader
from litellm.types.secret_managers.main import KeyManagementSystem
@ -62,10 +61,6 @@ class GoogleSecretManager(GCSBucketBase):
Returns:
str: The secret value if successful, None otherwise.
"""
native: Final = resolve_native_provider_reader(self, "google_secret_manager")
if native is not None:
return native.sync_read_secret(secret_name)
if self.always_read_secret_manager is not True:
cached_secret: Final = self.cache.get_cache(secret_name)
if cached_secret is not None:

25
litellm/secret_managers/hashicorp_secret_manager.py
View file

@ -15,7 +15,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.rust_bridge.secret_manager import resolve_native_provider_reader, resolve_native_provider_writer
from litellm.types.secret_managers.main import KeyManagementSystem
from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name
@ -406,10 +405,6 @@ class HashicorpSecretManager(BaseSecretManager):
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
Returns the entire data dict from data.data, or None on failure.
"""
native: Final = resolve_native_provider_reader(self, "hashicorp_vault")
if native is not None:
return await native.async_read_secret(secret_name, optional_params, timeout)
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
)
@ -441,10 +436,6 @@ class HashicorpSecretManager(BaseSecretManager):
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
Returns the entire data dict from data.data, or None on failure.
"""
native: Final = resolve_native_provider_reader(self, "hashicorp_vault")
if native is not None:
return native.sync_read_secret(secret_name, optional_params, timeout)
sync_client: Final = _get_httpx_client()
try:
target: Final = self._build_secret_target(secret_name, optional_params)
@ -485,12 +476,6 @@ class HashicorpSecretManager(BaseSecretManager):
Returns:
dict: Response containing status and details of the operation
"""
native: Final = resolve_native_provider_writer(self, "hashicorp_vault")
if native is not None:
return await native.async_write_secret(
secret_name, secret_value, description, optional_params, timeout, tags
)
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},
@ -540,12 +525,6 @@ class HashicorpSecretManager(BaseSecretManager):
On success, returns the response from async_write_secret.
On error, returns {"status": "error", "message": "error message"}
"""
native: Final = resolve_native_provider_writer(self, "hashicorp_vault")
if native is not None:
return await native.async_rotate_secret(
current_secret_name, new_secret_name, new_secret_value, optional_params, timeout
)
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},
@ -692,10 +671,6 @@ class HashicorpSecretManager(BaseSecretManager):
Returns:
dict: Response containing status and details of the operation
"""
native: Final = resolve_native_provider_writer(self, "hashicorp_vault")
if native is not None:
return await native.async_delete_secret(secret_name, recovery_window_in_days, optional_params, timeout)
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},

2
litellm/secret_managers/main.py
View file

@ -12,10 +12,10 @@ import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.secret_managers.dispatch import get_secret_from_manager
from litellm.secret_managers.get_azure_ad_token_provider import (
get_azure_ad_token_provider,
)
from litellm.secret_managers.secret_manager_handler import get_secret_from_manager
from litellm.types.llms.base import LiteLLMBaseModel
oidc_cache: Final = DualCache()

36
tests/unit/rust_bridge/test_catalog.py
View file

@ -9,15 +9,13 @@ from litellm.rust_bridge import catalog, configuration
from litellm.rust_bridge.catalog import (
Context,
LoggerContext,
LoggerRule,
Route,
RouteContext,
RouteRule,
Rules,
SecretManagerContext,
SecretManagerRule,
)
from litellm.rust_bridge.configuration import Decision, Rollout
from litellm.types.secret_managers.main import KeyManagementSystem
@pytest.fixture(autouse=True)
@ -65,23 +63,6 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt
assert catalog.decision(RouteContext(route), rules=()) is Decision.PYTHON
@pytest.mark.parametrize(
"context",
(
*(SecretManagerContext(system.value) for system in KeyManagementSystem),
SecretManagerContext("unknown"),
),
)
def test_backend_rollouts_stay_on_python_when_global_rust_is_enabled(
monkeypatch: pytest.MonkeyPatch, context: Context
) -> None:
configuration.rust(True)
monkeypatch.setenv("LITELLM_RUST", "1")
assert catalog.rollout(context) is Rollout.PYTHON_ONLY
assert catalog.decision(context) is Decision.PYTHON
def test_logger_rollout_obeys_the_global_switch() -> None:
assert catalog.rollout(LoggerContext()) is Rollout.RUST_OPT_IN
assert catalog.decision(LoggerContext()) is Decision.PYTHON
@ -134,14 +115,12 @@ def test_ocr_has_no_python_path_to_opt_out_to(
(RouteContext(Route.OCR, provider="local"), Decision.RUST_REQUIRED),
(RouteContext(Route.OCR, provider="other"), Decision.PYTHON),
(RouteContext(Route.MESSAGES, provider="local"), Decision.PYTHON),
(SecretManagerContext("local"), Decision.PYTHON),
(SecretManagerContext("other"), Decision.RUST_REQUIRED),
(LoggerContext(), Decision.RUST_REQUIRED),
),
)
def test_mixed_rules_select_only_the_matching_domain(context: Context, expected: Decision) -> None:
rules: Final[Rules] = (
SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})),
SecretManagerRule(Rollout.RUST_REQUIRED),
LoggerRule(Rollout.RUST_REQUIRED),
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"local"})),
RouteRule(Route.OCR, Rollout.PYTHON_ONLY),
)
@ -149,7 +128,7 @@ def test_mixed_rules_select_only_the_matching_domain(context: Context, expected:
assert catalog.decision(context, rules) is expected
@pytest.mark.parametrize("context", (RouteContext(Route.OCR), SecretManagerContext("local")))
@pytest.mark.parametrize("context", (RouteContext(Route.OCR), LoggerContext()))
@pytest.mark.parametrize(
("rollout", "process", "environment", "expected"),
(
@ -176,20 +155,19 @@ def test_all_domains_share_rollout_switches_and_first_match(
monkeypatch.setenv("LITELLM_RUST", environment)
rules: Final[Rules] = (
RouteRule(Route.OCR, rollout),
SecretManagerRule(rollout),
LoggerRule(rollout),
RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
SecretManagerRule(Rollout.RUST_REQUIRED),
)
assert catalog.decision(context, rules) is expected
assert catalog.decision(context, ()) is Decision.PYTHON
@pytest.mark.parametrize("context", (RouteContext(Route.OCR), SecretManagerContext("local")))
@pytest.mark.parametrize("context", (RouteContext(Route.OCR),))
def test_empty_constraints_match_nothing(context: Context) -> None:
rules: Final[Rules] = (
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()),
SecretManagerRule(Rollout.RUST_REQUIRED, systems=frozenset()),
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, models=frozenset()),
)
assert catalog.decision(context, rules) is Decision.PYTHON

8
tests/unit/rust_bridge/test_dispatch.py
View file

@ -6,7 +6,7 @@ import pytest
from litellm.rust_bridge import configuration
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules, SecretManagerRule
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.runtime import NO_PYTHON, NoPythonImplementationError
@ -23,7 +23,7 @@ def binding() -> NativeBinding[object]:
return bound
@pytest.mark.parametrize("rules", ((), (SecretManagerRule(Rollout.RUST_REQUIRED),)))
@pytest.mark.parametrize("rules", ((), (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)))
def test_route_without_rules_forwards_before_request_projection(rules: Rules) -> None:
stream: Final[Iterator[int]] = iter((1, 2))
@ -97,7 +97,7 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None:
request: Final = Request(model="streaming-model")
stream: Final[Iterator[int]] = iter((1, 2))
rules: Final[Rules] = (
SecretManagerRule(Rollout.PYTHON_ONLY),
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),
RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),
)
dispatch: Final = PublicDispatch(
@ -125,7 +125,7 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None:
@pytest.mark.asyncio
@pytest.mark.parametrize("rules", ((), (SecretManagerRule(Rollout.RUST_REQUIRED),)))
@pytest.mark.parametrize("rules", ((), (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)))
async def test_async_route_without_rules_preserves_async_iterator_result(rules: Rules) -> None:
async def chunks() -> AsyncGenerator[int, None]:
yield 1

684
tests/unit/rust_bridge/test_secret_manager.py
View file

File diff suppressed because it is too large Load diff

35
tests/unit/rust_bridge/test_settings.py
View file

@ -6,9 +6,7 @@ import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.custom_httpx.http_handler import default_user_agent
from litellm.rust_bridge import catalog, settings
from litellm.rust_bridge.catalog import SecretManagerRule
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge import settings
from litellm.secret_managers.main import get_secret_str
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -81,11 +79,6 @@ class _VaultSecrets(CustomSecretManager):
return self.secrets.get(secret_name)
_RUST_FOR_CUSTOM: Final = (
SecretManagerRule(Rollout.RUST_REQUIRED, systems=frozenset({KeyManagementSystem.CUSTOM.value})),
)
@pytest.mark.parametrize(
("access_mode", "readable"),
[("read_only", True), ("read_and_write", True), ("write_only", False)],
@ -98,38 +91,30 @@ def test_secret_manager_is_readable_only_when_litellm_would_read_secrets_from_it
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
assert settings.secret_manager(rules=()) == settings.SecretManager(readable=readable, native=False)
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", "rules", "native"),
("system", "access_mode", "native"),
[
(KeyManagementSystem.CUSTOM, "read_only", _RUST_FOR_CUSTOM, True),
(KeyManagementSystem.CUSTOM, "read_only", (), False),
(
KeyManagementSystem.CUSTOM,
"read_only",
(SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CUSTOM.value})),),
False,
),
(KeyManagementSystem.CUSTOM, "write_only", _RUST_FOR_CUSTOM, False),
(None, "read_only", _RUST_FOR_CUSTOM, False),
(KeyManagementSystem.AWS_SECRET_MANAGER, "read_only", _RUST_FOR_CUSTOM, False),
(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_the_rules_select_rust_for_its_system(
def test_secret_manager_is_native_only_when_a_readable_system_is_configured(
monkeypatch: pytest.MonkeyPatch,
system: KeyManagementSystem | None,
access_mode: str,
rules: catalog.Rules,
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(rules=rules) == settings.SecretManager(
assert settings.secret_manager() == settings.SecretManager(
readable=access_mode != "write_only", native=native
)
@ -137,7 +122,7 @@ def test_secret_manager_is_native_only_when_the_rules_select_rust_for_its_system
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(rules=_RUST_FOR_CUSTOM) == settings.SecretManager(readable=False, native=False)
assert settings.secret_manager() == settings.SecretManager(readable=False, native=False)
def test_secret_manager_projects_custom_settings(monkeypatch: pytest.MonkeyPatch) -> None: