mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(router): keep encrypted reasoning when a fallback hops between same-key deployments with no api_base (#49396)
* fix(router): keep encrypted reasoning on a hop between same-key deployments on the provider's default endpoint * test(router): type the default-endpoint boundary test helpers * fix(guardrails): restore the google_text_moderation embeddings fix an earlier main merge dropped * fix(router): match encryption boundaries on the key, account, and region the request really uses A deployment's named credential now fills only the values the deployment leaves out, through the same helper the request path uses, so the boundary names the key and base the request sends. A deployment with no api_key takes the provider's default key, the boundary carries the OpenAI organization and project a deployment sends, and a Bedrock Mantle region named in the model routes through the provider resolver before the Responses URL builder * fix(router): keep header-signed accounts apart on a shared boundary and skip boundary work no reasoning needs Headers the request logger masks as secrets (Authorization, api-key, x-api-key) now join the boundary's account scope, with extra_headers overriding headers the way the request merges them, so deployments that sign in through a header on one base never share encrypted reasoning. Target boundaries are built only when a reasoning item from another deployment needs the comparison, and the Responses input strip asks only about encrypted reasoning items, so a pinned turn or fallback hop with nothing to compare makes no secret-manager reads. * fix(router): resolve default boundaries only for openai and drop the unsent organization param Aggregators and Bedrock Mantle serve several vendors on one key, so a missing key or base there no longer yields a shared boundary. The Responses request never sends the organization param, so it no longer splits the boundary, while a Cookie header now does. * fix(router): match explicit pairs as on main and key default boundaries by model An explicit api_base plus api_key pair is the whole boundary again, so tag headers no longer split deployments that main kept together. The default OpenAI boundary now carries the model, since one OPENAI_BASE_URL can front several vendors, and a deployment that declares another provider gets no boundary before any secret lookup. * fix(router): match an explicit key and the environment key on one api_base A deployment with an explicit api_base and api_key now shares a boundary with one on the same api_base that takes the same key from OPENAI_API_KEY, when the model and account headers also match. Two explicit pairs still match on base and key alone, as on main. * fix(router): skip a deployment whose params cannot be read in boundary checks A named credential with a value of the wrong type, or an extra header whose name is not a string, made the boundary check raise and fail the whole request. That deployment now has no boundary, so the others still keep the reasoning they can decrypt. Two explicit api_base plus api_key pairs also compare without resolving the provider, so they read no secrets, as on main. * fix(router): ask a secret manager for default boundaries at most once a minute per deployment * fix(router): resolve a changed named credential at once in the cached boundary and take an injected clock --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
572832dea1
commit
ff91b89b8d
5 changed files with 917 additions and 106 deletions
|
|
@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider(
|
|||
return model, custom_llm_provider
|
||||
|
||||
|
||||
def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
|
||||
return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
|
||||
|
||||
|
||||
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
|
||||
"""The authenticating provider this pair already names, or None.
|
||||
|
||||
|
|
@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str
|
|||
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
|
||||
adopt the declaration instead of resolving.
|
||||
"""
|
||||
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
|
||||
declared: Final = declared_provider(model, custom_llm_provider)
|
||||
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -678,20 +678,25 @@ class ResponsesAPIRequestUtils:
|
|||
return
|
||||
items: Final = cast(list[object], request_input) # cast-ok: untyped client json
|
||||
stripped: Final = tuple(
|
||||
ResponsesAPIRequestUtils._without_encrypted_reasoning(item)
|
||||
if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item)))
|
||||
ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning)
|
||||
if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None
|
||||
and (should_strip is None or should_strip(reasoning))
|
||||
else item
|
||||
for item in items
|
||||
)
|
||||
items[:] = (item for item in stripped if item is not None)
|
||||
|
||||
@staticmethod
|
||||
def _without_encrypted_reasoning(item: object) -> object | None:
|
||||
def _encrypted_reasoning_item(item: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(item, dict):
|
||||
return item
|
||||
return None
|
||||
reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json
|
||||
if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"):
|
||||
return reasoning
|
||||
return None
|
||||
return reasoning
|
||||
|
||||
@staticmethod
|
||||
def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None:
|
||||
readable: Final = any(
|
||||
ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -36,31 +36,75 @@ Safe to enable globally:
|
|||
- No cache required.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from functools import cache
|
||||
import hashlib
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from functools import cache, lru_cache, partial
|
||||
from time import monotonic
|
||||
from typing import TYPE_CHECKING, Final, Optional, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.exceptions import AuthenticationError, BadRequestError
|
||||
from litellm.integrations.custom_logger import CustomLogger, Span
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import has_sensitive_keyword
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
anthropic_content_lists,
|
||||
encrypted_content_of_block,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
)
|
||||
from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param
|
||||
from litellm.secret_managers.main import secret_manager_would_be_consulted
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.utils import get_order_filtered_deployments
|
||||
from litellm.types.router import Deployment, GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_api import RoutingContext
|
||||
|
||||
_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[object, object])
|
||||
_ACCOUNT_SCOPE_HEADERS: Final = frozenset({"openai-organization", "openai-project", "cookie"})
|
||||
_SINGLE_VENDOR_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value})
|
||||
_DEFAULT_ROUTE_SECRETS: Final = ("USE_LITELLM_PROXY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE")
|
||||
_SECRET_MANAGER_REFRESH_SECONDS: Final = 60
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderRoute:
|
||||
model: str
|
||||
provider: str
|
||||
api_base: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _EncryptionBoundary:
|
||||
api_base: object
|
||||
api_key: object
|
||||
model: Callable[[], str | None] = field(compare=False)
|
||||
account_scope: frozenset[tuple[str, object]]
|
||||
explicit_pair: bool
|
||||
|
||||
def is_shared_with(self, other: "_EncryptionBoundary") -> bool:
|
||||
if (self.api_base, self.api_key) != (other.api_base, other.api_key):
|
||||
return False
|
||||
if self.explicit_pair and other.explicit_pair:
|
||||
return True
|
||||
return self.account_scope == other.account_scope and self.model() == other.model()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ParamsInRefreshWindow:
|
||||
digest: str
|
||||
window: int
|
||||
params: Mapping[str, object] = field(compare=False)
|
||||
|
||||
|
||||
class EncryptedContentAffinityCheck(CustomLogger):
|
||||
|
|
@ -83,9 +127,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
ctx: Optional["RoutingContext"] = None,
|
||||
enable_global_affinity: bool = True,
|
||||
model_group_affinity_config: dict[str, list[str]] | None = None,
|
||||
clock: Callable[[], float] = monotonic,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self._clock = clock
|
||||
self.enable_global_affinity = enable_global_affinity
|
||||
self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {}
|
||||
|
||||
|
|
@ -209,44 +255,156 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
@staticmethod
|
||||
def _encryption_boundary_key(
|
||||
litellm_params: object,
|
||||
) -> tuple[object, object] | None:
|
||||
clock: Callable[[], float] = monotonic,
|
||||
) -> _EncryptionBoundary | None:
|
||||
"""
|
||||
``(api_base, api_key)`` identifies an upstream encryption boundary.
|
||||
The values are resolved from the deployment and its named credential
|
||||
without modifying the deployment.
|
||||
The upstream encryption boundary a deployment's encrypted reasoning belongs to. The values
|
||||
are resolved the way the request resolves them, the deployment's own values winning over its
|
||||
named credential's, without modifying the deployment. Two explicit ``(api_base, api_key)``
|
||||
pairs match on those values alone. On a provider that serves one vendor's models, a deployment
|
||||
with no ``api_key`` takes the provider's default key, one with no ``api_base`` takes the endpoint
|
||||
its Responses URL builder resolves, and a match involving either also needs the same model and
|
||||
the same headers that pick an account.
|
||||
|
||||
Accepts any object exposing dict-style ``.get(key, default)``: plain
|
||||
dicts (the common case in ``healthy_deployments``) as well as
|
||||
``LiteLLM_Params``-style Pydantic instances, which define a custom
|
||||
``.get()``. A stricter ``isinstance(dict)`` guard would silently drop
|
||||
the latter from boundary matching and fall back to the full pool —
|
||||
i.e. trigger the exact ``invalid_encrypted_content`` failure this
|
||||
check exists to prevent.
|
||||
Accepts plain dicts (the common case in ``healthy_deployments``) as
|
||||
well as ``LiteLLM_Params`` Pydantic instances. A stricter
|
||||
``isinstance(dict)`` guard would silently drop the latter from
|
||||
boundary matching and fall back to the full pool — i.e. trigger the
|
||||
exact ``invalid_encrypted_content`` failure this check exists to prevent.
|
||||
|
||||
When a secret manager would be asked for the provider defaults, the boundary of one set of
|
||||
deployment and named credential values is reused for up to a minute, so a run of pin misses
|
||||
does not ask it again for every candidate, while a changed credential resolves at once.
|
||||
"""
|
||||
getter: Final = getattr(litellm_params, "get", None)
|
||||
if not callable(getter):
|
||||
params: Final = EncryptedContentAffinityCheck._params_mapping(litellm_params)
|
||||
if params is None:
|
||||
return None
|
||||
api_base: Final = getter("api_base")
|
||||
api_key: Final = getter("api_key")
|
||||
credential_name: Final = getter("litellm_credential_name")
|
||||
credential_values: Final[Mapping[str, object] | None] = (
|
||||
CredentialAccessor.get_credential_values(credential_name)
|
||||
if isinstance(credential_name, str) and credential_name
|
||||
else None
|
||||
effective_params: Final = with_named_credential(params)
|
||||
if not any(secret_manager_would_be_consulted(name) for name in _DEFAULT_ROUTE_SECRETS):
|
||||
return EncryptedContentAffinityCheck._resolve_encryption_boundary(effective_params)
|
||||
return EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window(
|
||||
_ParamsInRefreshWindow(
|
||||
hashlib.sha256(repr(effective_params).encode()).hexdigest(),
|
||||
int(clock() // _SECRET_MANAGER_REFRESH_SECONDS),
|
||||
effective_params,
|
||||
)
|
||||
)
|
||||
effective_api_base: Final = (
|
||||
credential_values.get("api_base")
|
||||
if credential_values is not None and "api_base" in credential_values
|
||||
else api_base
|
||||
)
|
||||
effective_api_key: Final = (
|
||||
credential_values.get("api_key")
|
||||
if credential_values is not None and "api_key" in credential_values
|
||||
else api_key
|
||||
)
|
||||
if not effective_api_base or not effective_api_key:
|
||||
|
||||
@staticmethod
|
||||
@lru_cache(maxsize=1024)
|
||||
def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None:
|
||||
return EncryptedContentAffinityCheck._resolve_encryption_boundary(key.params)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_encryption_boundary(effective_params: Mapping[str, object]) -> _EncryptionBoundary | None:
|
||||
explicit_api_key: Final = effective_params.get("api_key")
|
||||
explicit_api_base: Final = effective_params.get("api_base")
|
||||
account_scope: Final = EncryptedContentAffinityCheck._account_scope(effective_params)
|
||||
if explicit_api_key and explicit_api_base:
|
||||
return _EncryptionBoundary(
|
||||
explicit_api_base,
|
||||
explicit_api_key,
|
||||
cache(partial(EncryptedContentAffinityCheck._single_vendor_model, effective_params)),
|
||||
account_scope,
|
||||
explicit_pair=True,
|
||||
)
|
||||
route: Final = EncryptedContentAffinityCheck._single_vendor_route(effective_params)
|
||||
if route is None:
|
||||
return None
|
||||
api_key: Final = explicit_api_key or EncryptedContentAffinityCheck._default_api_key(route)
|
||||
if not api_key:
|
||||
return None
|
||||
api_base: Final = explicit_api_base or EncryptedContentAffinityCheck._default_responses_endpoint(
|
||||
route, effective_params
|
||||
)
|
||||
if not api_base:
|
||||
return None
|
||||
return _EncryptionBoundary(api_base, api_key, lambda: route.model, account_scope, explicit_pair=False)
|
||||
|
||||
@staticmethod
|
||||
def _params_mapping(litellm_params: object) -> Mapping[str, object] | None:
|
||||
if isinstance(litellm_params, BaseModel):
|
||||
return litellm_params.model_dump(exclude_none=True)
|
||||
if isinstance(litellm_params, Mapping):
|
||||
return _REQUEST_KWARGS_ADAPTER.validate_python(litellm_params)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _account_scope(litellm_params: Mapping[str, object]) -> frozenset[tuple[str, object]]:
|
||||
return frozenset(
|
||||
(name.lower(), str(value))
|
||||
for name, value in EncryptedContentAffinityCheck._header_items(litellm_params)
|
||||
if name.lower() in _ACCOUNT_SCOPE_HEADERS or has_sensitive_keyword(name)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _header_items(litellm_params: Mapping[str, object]) -> Iterator[tuple[str, object]]:
|
||||
headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("headers"))
|
||||
extra_headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("extra_headers"))
|
||||
overridden: Final = frozenset(name.lower() for name in extra_headers)
|
||||
yield from ((name, value) for name, value in headers.items() if name.lower() not in overridden)
|
||||
yield from extra_headers.items()
|
||||
|
||||
@staticmethod
|
||||
def _header_mapping(headers: object) -> Mapping[str, object]:
|
||||
if isinstance(headers, Mapping):
|
||||
return {str(name): value for name, value in _HEADERS_ADAPTER.validate_python(headers).items()}
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def _single_vendor_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None:
|
||||
route: Final = EncryptedContentAffinityCheck._provider_route(litellm_params)
|
||||
return route if route is not None and route.provider in _SINGLE_VENDOR_PROVIDERS else None
|
||||
|
||||
@staticmethod
|
||||
def _single_vendor_model(litellm_params: Mapping[str, object]) -> str | None:
|
||||
route: Final = EncryptedContentAffinityCheck._single_vendor_route(litellm_params)
|
||||
return route.model if route is not None else None
|
||||
|
||||
@staticmethod
|
||||
def _provider_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None:
|
||||
model: Final = litellm_params.get("model")
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
||||
declared: Final = declared_provider(
|
||||
model, custom_llm_provider if isinstance(custom_llm_provider, str) else None
|
||||
)
|
||||
if declared in litellm.provider_list and declared not in _SINGLE_VENDOR_PROVIDERS:
|
||||
return None
|
||||
try:
|
||||
provider_model, provider, _, provider_api_base = get_llm_provider(
|
||||
model=model, litellm_params=GenericLiteLLMParams.model_validate(dict(litellm_params))
|
||||
)
|
||||
except (BadRequestError, ValidationError):
|
||||
return None
|
||||
return _ProviderRoute(provider_model, provider, provider_api_base)
|
||||
|
||||
@staticmethod
|
||||
def _default_api_key(route: _ProviderRoute) -> str | None:
|
||||
model_info: Final = ProviderConfigManager.get_provider_model_info(
|
||||
model=route.model, provider=LlmProviders(route.provider)
|
||||
)
|
||||
if model_info is None:
|
||||
return None
|
||||
try:
|
||||
return model_info.get_api_key()
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _default_responses_endpoint(route: _ProviderRoute, litellm_params: Mapping[str, object]) -> str | None:
|
||||
responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(
|
||||
provider=route.provider, model=route.model
|
||||
)
|
||||
if responses_config is None:
|
||||
return None
|
||||
try:
|
||||
return responses_config.get_complete_url( # pyright: ignore[reportUnknownMemberType] # the base types litellm_params as a bare dict
|
||||
api_base=route.api_base, litellm_params=dict(litellm_params)
|
||||
)
|
||||
except (ValueError, BaseLLMException, AuthenticationError):
|
||||
return None
|
||||
return (effective_api_base, effective_api_key)
|
||||
|
||||
def _find_deployments_on_same_encryption_boundary(
|
||||
self,
|
||||
|
|
@ -255,7 +413,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
) -> tuple[list[dict], Deployment | None]:
|
||||
"""
|
||||
Deployments in ``healthy_deployments`` sharing the originating
|
||||
deployment's ``(api_base, api_key)``, alongside the originating
|
||||
deployment's encryption boundary, alongside the originating
|
||||
deployment object (or ``None`` if it was removed / router unavailable).
|
||||
"""
|
||||
if self.ctx is None:
|
||||
|
|
@ -263,11 +421,16 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
originating: Final = self.ctx.registry().deployment(model_id)
|
||||
if originating is None:
|
||||
return [], None
|
||||
boundary: Final = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True))
|
||||
boundary: Final = self._encryption_boundary_key(
|
||||
originating.litellm_params.model_dump(exclude_none=True), self._clock
|
||||
)
|
||||
if boundary is None:
|
||||
return [], originating
|
||||
matches: Final = [
|
||||
d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary
|
||||
d
|
||||
for d in healthy_deployments
|
||||
if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None
|
||||
and boundary.is_shared_with(candidate)
|
||||
]
|
||||
return matches, originating
|
||||
|
||||
|
|
@ -280,6 +443,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
target_deployments: Sequence[Mapping[str, object]],
|
||||
*,
|
||||
unmarked_origin: str | None,
|
||||
clock: Callable[[], float] = monotonic,
|
||||
) -> None:
|
||||
"""
|
||||
Drop the encrypted reasoning that none of ``target_deployments`` minted or shares an
|
||||
|
|
@ -293,12 +457,19 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
for target in target_deployments
|
||||
if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None
|
||||
)
|
||||
target_boundaries: Final = frozenset(
|
||||
boundary
|
||||
for target in target_deployments
|
||||
if (boundary := EncryptedContentAffinityCheck._encryption_boundary_key(target.get("litellm_params")))
|
||||
is not None
|
||||
)
|
||||
|
||||
@cache
|
||||
def target_boundaries() -> tuple[_EncryptionBoundary, ...]:
|
||||
return tuple(
|
||||
boundary
|
||||
for target in target_deployments
|
||||
if (
|
||||
boundary := EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
target.get("litellm_params"), clock
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
@cache
|
||||
def target_can_decrypt(marked_origin: str | None) -> bool:
|
||||
|
|
@ -312,12 +483,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
origin: Final = ctx.registry().deployment(origin_model_id)
|
||||
origin_boundary: Final = (
|
||||
EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
origin.litellm_params.model_dump(exclude_none=True)
|
||||
origin.litellm_params.model_dump(exclude_none=True), clock
|
||||
)
|
||||
if origin is not None
|
||||
else None
|
||||
)
|
||||
return origin_boundary is not None and origin_boundary in target_boundaries
|
||||
return origin_boundary is not None and any(
|
||||
origin_boundary.is_shared_with(target_boundary) for target_boundary in target_boundaries()
|
||||
)
|
||||
|
||||
def should_strip_input_item(item: Mapping[str, object]) -> bool:
|
||||
return not target_can_decrypt(EncryptedContentAffinityCheck._model_id_of_input_item(item))
|
||||
|
|
@ -424,7 +597,12 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
)
|
||||
request_kwargs["_encrypted_content_affinity_pinned"] = True
|
||||
self.strip_reasoning_the_targets_cannot_decrypt(
|
||||
self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id
|
||||
self.ctx,
|
||||
request_input,
|
||||
anthropic_messages,
|
||||
(deployment,),
|
||||
unmarked_origin=model_id,
|
||||
clock=self._clock,
|
||||
)
|
||||
return [deployment]
|
||||
|
||||
|
|
@ -442,7 +620,12 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
)
|
||||
request_kwargs["_encrypted_content_affinity_pinned"] = True
|
||||
self.strip_reasoning_the_targets_cannot_decrypt(
|
||||
self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id
|
||||
self.ctx,
|
||||
request_input,
|
||||
anthropic_messages,
|
||||
boundary_matches,
|
||||
unmarked_origin=model_id,
|
||||
clock=self._clock,
|
||||
)
|
||||
return boundary_matches
|
||||
|
||||
|
|
|
|||
|
|
@ -303,7 +303,7 @@ def scenario_api_base(scenario: Scenario, scenario_id: str, response: StoredResp
|
|||
def deployment(
|
||||
scenario: Scenario,
|
||||
group: str,
|
||||
api_base: str,
|
||||
api_base: str | None,
|
||||
*,
|
||||
order: int | None,
|
||||
model: str = RESPONSES_MODEL,
|
||||
|
|
@ -316,7 +316,7 @@ def deployment(
|
|||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_key": PROVIDER_KEY,
|
||||
"api_base": api_base,
|
||||
**({} if api_base is None else {"api_base": api_base}),
|
||||
**({} if order is None else {"order": order}),
|
||||
**extra_params,
|
||||
},
|
||||
|
|
@ -863,6 +863,38 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway
|
|||
assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies
|
||||
|
||||
|
||||
def test_hop_between_same_key_deployments_on_the_providers_default_endpoint_keeps_the_reasoning(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
shared: Final = f"default-{uuid.uuid4().hex[:12]}"
|
||||
with (
|
||||
scratch_database() as database,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"DATABASE_URL": database, "OPENAI_BASE_URL": f"{gateway.upstream_url}/{shared}"},
|
||||
remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE"),
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}})
|
||||
one: Final = scenario_api_base(scenario, f"{shared}-o1", failing()).rsplit("/", 1)[-1]
|
||||
two: Final = scenario_api_base(scenario, f"{shared}-o2", healthy_json()).rsplit("/", 1)[-1]
|
||||
order_one: Final = deployment(
|
||||
scenario, shared, None, order=1, extra_params={"extra_headers": {"x-scripted-scenario": one}}
|
||||
)
|
||||
order_two: Final = deployment(
|
||||
scenario, shared, None, order=2, extra_params={"extra_headers": {"x-scripted-scenario": two}}
|
||||
)
|
||||
mark: Final = marker()
|
||||
answer: Final = responses_answer(
|
||||
owned.gateway, HTTPX, Target(shared, frozenset({order_one, order_two})), history(mark)
|
||||
)
|
||||
assert_served_by(answer, order_two)
|
||||
bodies: Final = bodies_for(observed(owned.gateway), shared)
|
||||
assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies
|
||||
|
||||
|
||||
def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None:
|
||||
records: Final = observed(gateway)
|
||||
assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records
|
||||
|
|
|
|||
|
|
@ -15,17 +15,21 @@ The mechanism works without any cache and supports two encoding strategies:
|
|||
encrypted_content back to their original forms before sending to the upstream provider.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.models.credentials import CredentialItem
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -1100,14 +1104,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance():
|
|||
plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params)
|
||||
|
||||
assert pydantic_key is not None
|
||||
assert (
|
||||
pydantic_key
|
||||
== plain_key
|
||||
== (
|
||||
"https://mateo-resource.openai.azure.com",
|
||||
"fake-azure-resource-key-a",
|
||||
)
|
||||
)
|
||||
assert pydantic_key == plain_key
|
||||
|
||||
|
||||
def test_boundary_key_resolves_missing_values_from_named_credential():
|
||||
|
|
@ -1131,14 +1128,35 @@ def test_boundary_key_resolves_missing_values_from_named_credential():
|
|||
):
|
||||
boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"})
|
||||
|
||||
assert boundary == ("https://account-a.example.com", "credential-key-a")
|
||||
assert boundary is not None
|
||||
assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
{"api_base": "https://account-a.example.com", "api_key": "credential-key-a"}
|
||||
)
|
||||
assert boundary != EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
{"api_base": "https://account-a.example.com", "api_key": "credential-key-b"}
|
||||
)
|
||||
|
||||
|
||||
def test_boundary_key_matches_named_credential_precedence():
|
||||
@pytest.mark.parametrize(
|
||||
"deployment_values",
|
||||
[
|
||||
pytest.param({"api_base": "https://deployment.example.com", "api_key": "deployment-key"}, id="set values"),
|
||||
pytest.param({"api_base": "", "api_key": ""}, id="empty values"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.usefixtures("no_process_wide_api_base")
|
||||
def test_boundary_key_resolves_named_credential_values_the_way_the_request_does(
|
||||
monkeypatch: pytest.MonkeyPatch, deployment_values: dict[str, str]
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
from litellm.utils import with_named_credential
|
||||
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "env-key")
|
||||
deployment: Final = {"model": "openai/gpt-5.5", **deployment_values, "litellm_credential_name": "account-a"}
|
||||
with patch.object(
|
||||
litellm,
|
||||
"credential_list",
|
||||
|
|
@ -1153,46 +1171,14 @@ def test_boundary_key_matches_named_credential_precedence():
|
|||
)
|
||||
],
|
||||
):
|
||||
boundary = EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
{
|
||||
"api_base": "https://deployment.example.com",
|
||||
"api_key": "deployment-key",
|
||||
"litellm_credential_name": "account-a",
|
||||
}
|
||||
)
|
||||
boundary: Final = EncryptedContentAffinityCheck._encryption_boundary_key(deployment)
|
||||
sent_params: Final = with_named_credential(deployment)
|
||||
|
||||
assert boundary == ("https://credential.example.com", "credential-key-a")
|
||||
|
||||
|
||||
def test_boundary_key_resolves_credential_when_explicit_values_are_empty():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
assert boundary is not None
|
||||
assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
{key: value for key, value in sent_params.items() if key != "litellm_credential_name"}
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="account-a",
|
||||
credential_values={
|
||||
"api_base": "https://credential.example.com",
|
||||
"api_key": "credential-key-a",
|
||||
},
|
||||
credential_info={},
|
||||
)
|
||||
],
|
||||
):
|
||||
boundary = EncryptedContentAffinityCheck._encryption_boundary_key(
|
||||
{
|
||||
"api_base": "",
|
||||
"api_key": "",
|
||||
"litellm_credential_name": "account-a",
|
||||
}
|
||||
)
|
||||
|
||||
assert boundary == ("https://credential.example.com", "credential-key-a")
|
||||
|
||||
|
||||
def test_boundary_fallback_matches_deployments_with_same_named_credential_values():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
|
|
@ -2728,3 +2714,604 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th
|
|||
|
||||
assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"]
|
||||
assert reasoning_input[0] == original_reasoning
|
||||
|
||||
|
||||
def _default_endpoint_deployment(model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]:
|
||||
return {"model_name": "gpt-5.5", "litellm_params": dict(litellm_params), "model_info": {"id": model_id}}
|
||||
|
||||
|
||||
def _deployment_of(ctx: RoutingContext, model_id: str) -> dict[str, object]:
|
||||
deployment: Final = ctx.registry().deployment(model_id)
|
||||
assert deployment is not None
|
||||
return deployment.model_dump(exclude_none=True)
|
||||
|
||||
|
||||
SUMMARY_ONLY_REASONING: Final = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]}
|
||||
UNMARKED_REASONING: Final = {**SUMMARY_ONLY_REASONING, "id": "rs_test", "encrypted_content": "gAAAAA-blob"}
|
||||
OPENAI_SHARED_KEY: Final = {"model": "openai/gpt-5.5", "api_key": "shared-key"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_process_wide_api_base(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "api_base", None)
|
||||
monkeypatch.delenv("AZURE_API_BASE", raising=False)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_process_wide_api_base")
|
||||
@pytest.mark.parametrize("marked", [True, False], ids=["marked origin", "unmarked origin"])
|
||||
@pytest.mark.parametrize(
|
||||
"origin_params, target_params, target_can_decrypt",
|
||||
[
|
||||
pytest.param(OPENAI_SHARED_KEY, OPENAI_SHARED_KEY, True, id="same key on the provider default endpoint"),
|
||||
pytest.param(
|
||||
{"model": "gpt-5.5", "custom_llm_provider": "openai", "api_key": "shared-key"},
|
||||
OPENAI_SHARED_KEY,
|
||||
True,
|
||||
id="provider named apart from the model",
|
||||
),
|
||||
pytest.param(OPENAI_SHARED_KEY, {**OPENAI_SHARED_KEY, "api_key": "other-key"}, False, id="different key"),
|
||||
pytest.param(
|
||||
OPENAI_SHARED_KEY,
|
||||
{**OPENAI_SHARED_KEY, "model": "openai/gpt-5.5-mini"},
|
||||
False,
|
||||
id="same key on the default endpoint for another model",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"},
|
||||
{"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"},
|
||||
False,
|
||||
id="same key on Bedrock Mantle, which serves several vendors",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "openrouter/openai/gpt-6.1-sol", "api_key": "shared-key"},
|
||||
{"model": "openrouter/anthropic/claude-haiku-5.5", "api_key": "shared-key"},
|
||||
False,
|
||||
id="same key on an aggregator that serves several vendors",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"},
|
||||
{"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"},
|
||||
False,
|
||||
id="provider without a Responses URL builder",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "azure/gpt-5.5", "api_key": "shared-key"},
|
||||
{"model": "azure/gpt-5.5", "api_key": "shared-key"},
|
||||
False,
|
||||
id="provider whose builder needs an api_base",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "github_copilot/gpt-5.5", "api_key": "shared-key"},
|
||||
{"model": "github_copilot/gpt-5.5", "api_key": "shared-key"},
|
||||
False,
|
||||
id="same key on a provider that authenticates with its own token",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "groq/llama-4-scout"},
|
||||
{"model": "groq/llama-4-scout"},
|
||||
False,
|
||||
id="provider with no default key",
|
||||
),
|
||||
pytest.param(
|
||||
{"model": "vertex_ai/gemini-3-pro"},
|
||||
{"model": "vertex_ai/gemini-3-pro"},
|
||||
False,
|
||||
id="provider whose default key is a credentials file",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "api_base": "https://openai-proxy.example.com/v1"},
|
||||
OPENAI_SHARED_KEY,
|
||||
False,
|
||||
id="explicit api_base beside the default endpoint",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-a"}},
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"openai-organization": "org-b"}},
|
||||
False,
|
||||
id="same key sent to different OpenAI organizations",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Project": "proj-a"}},
|
||||
OPENAI_SHARED_KEY,
|
||||
False,
|
||||
id="same key with an OpenAI project beside one without",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"Cookie": "session=account-a"}},
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"cookie": "session=account-b"}},
|
||||
False,
|
||||
id="same key signed in to different accounts by cookie",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "headers": {"api-key": "account-a"}},
|
||||
{**OPENAI_SHARED_KEY, "headers": {"api-key": "account-b"}},
|
||||
False,
|
||||
id="explicit key beside different api-key headers",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "organization": "org-a"},
|
||||
{**OPENAI_SHARED_KEY, "organization": "org-b"},
|
||||
True,
|
||||
id="organization param the Responses request never sends",
|
||||
),
|
||||
pytest.param(
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "a", "OpenAI-Organization": "org-a"}},
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "b", "OpenAI-Organization": "org-a"}},
|
||||
True,
|
||||
id="same OpenAI organization beside unrelated headers",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
**OPENAI_SHARED_KEY,
|
||||
"headers": {"OpenAI-Organization": "org-a"},
|
||||
"extra_headers": {"openai-organization": "org-b"},
|
||||
},
|
||||
{**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-b"}},
|
||||
True,
|
||||
id="extra headers override the organization the headers set",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_keeps_reasoning_between_same_key_deployments_on_the_default_endpoint(
|
||||
origin_params: dict[str, object], target_params: dict[str, object], target_can_decrypt: bool, marked: bool
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", origin_params),
|
||||
_default_endpoint_deployment("target", target_params),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
reasoning: Final = _reasoning_input_minted_by("origin")[0] if marked else dict(UNMARKED_REASONING)
|
||||
original_reasoning: Final = dict(reasoning)
|
||||
request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx,
|
||||
request_input,
|
||||
None,
|
||||
(_deployment_of(ctx, "target"),),
|
||||
unmarked_origin=None if marked else "origin",
|
||||
)
|
||||
|
||||
assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_process_wide_api_base")
|
||||
@pytest.mark.parametrize(
|
||||
"target_api_key, target_can_decrypt",
|
||||
[
|
||||
pytest.param(None, True, id="both on the environment key"),
|
||||
pytest.param("env-key", True, id="explicit key equal to the environment key"),
|
||||
pytest.param("other-key", False, id="explicit key other than the environment key"),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_resolves_a_missing_key_from_the_environment(
|
||||
monkeypatch: pytest.MonkeyPatch, target_api_key: str | None, target_can_decrypt: bool
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "env-key")
|
||||
target_params: Final = {"model": "openai/gpt-5.5", **({"api_key": target_api_key} if target_api_key else {})}
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", {"model": "openai/gpt-5.5"}),
|
||||
_default_endpoint_deployment("target", target_params),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
reasoning: Final = _reasoning_input_minted_by("origin")[0]
|
||||
original_reasoning: Final = dict(reasoning)
|
||||
request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None
|
||||
)
|
||||
|
||||
assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING)
|
||||
|
||||
|
||||
def test_hop_strip_leaves_input_items_that_are_not_objects_as_sent():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
request_input: Final = ["and sunsets?", *_reasoning_input_minted_by("origin")]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
None, request_input, None, (), unmarked_origin=None
|
||||
)
|
||||
|
||||
assert request_input == ["and sunsets?", SUMMARY_ONLY_REASONING, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
|
||||
GATEWAY_BASE: Final = {"model": "openai/gpt-5.5", "api_base": "https://gateway.example.com/v1"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"origin_params, target_params, target_can_decrypt",
|
||||
[
|
||||
pytest.param(
|
||||
{**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}},
|
||||
{**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-b"}},
|
||||
False,
|
||||
id="environment key replaced by different Authorization headers",
|
||||
),
|
||||
pytest.param(
|
||||
{**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "a"}},
|
||||
{**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "b"}},
|
||||
True,
|
||||
id="explicit pair beside different tag headers",
|
||||
),
|
||||
pytest.param(
|
||||
{**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}},
|
||||
{**GATEWAY_BASE, "extra_headers": {"authorization": "Bearer sk-account-a"}},
|
||||
True,
|
||||
id="same Authorization header on both",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_keeps_accounts_that_sign_in_through_headers_apart(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
origin_params: dict[str, object],
|
||||
target_params: dict[str, object],
|
||||
target_can_decrypt: bool,
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "env-key")
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", origin_params),
|
||||
_default_endpoint_deployment("target", target_params),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
reasoning: Final = _reasoning_input_minted_by("origin")[0]
|
||||
original_reasoning: Final = dict(reasoning)
|
||||
request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None
|
||||
)
|
||||
|
||||
assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"environment_key, target_model, target_can_decrypt",
|
||||
[
|
||||
pytest.param("shared-key", GATEWAY_BASE["model"], True, id="environment key equal to the explicit key"),
|
||||
pytest.param("other-key", GATEWAY_BASE["model"], False, id="environment key other than the explicit key"),
|
||||
pytest.param("shared-key", "openai/gpt-5.5-mini", False, id="equal keys serving another model"),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_matches_an_explicit_key_and_the_environment_key_on_one_api_base(
|
||||
monkeypatch: pytest.MonkeyPatch, environment_key: str, target_model: str, target_can_decrypt: bool
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", environment_key)
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", {**GATEWAY_BASE, "api_key": "shared-key"}),
|
||||
_default_endpoint_deployment("target", {**GATEWAY_BASE, "model": target_model}),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
reasoning: Final = _reasoning_input_minted_by("origin")[0]
|
||||
original_reasoning: Final = dict(reasoning)
|
||||
request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None
|
||||
)
|
||||
|
||||
assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING)
|
||||
|
||||
|
||||
class _SecretReadCounter(CustomSecretManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(secret_manager_name="read-counter")
|
||||
self.reads: int = 0
|
||||
|
||||
async def async_read_secret(
|
||||
self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None
|
||||
) -> str | None:
|
||||
self.reads += 1
|
||||
return None
|
||||
|
||||
def sync_read_secret(
|
||||
self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None
|
||||
) -> str | None:
|
||||
self.reads += 1
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"deployment_params, reasoning_minted_by",
|
||||
[
|
||||
pytest.param(OPENAI_SHARED_KEY, None, id="no reasoning in the input"),
|
||||
pytest.param(OPENAI_SHARED_KEY, "target", id="reasoning the hop's target minted"),
|
||||
pytest.param(
|
||||
{"model": "groq/llama-4-scout"}, "origin", id="reasoning from a provider with no default boundary"
|
||||
),
|
||||
pytest.param({**GATEWAY_BASE, "api_key": "shared-key"}, "origin", id="reasoning between explicit pairs"),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_reads_no_secrets_when_no_default_boundary_needs_resolving(
|
||||
monkeypatch: pytest.MonkeyPatch, deployment_params: dict[str, object], reasoning_minted_by: str | None
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", deployment_params),
|
||||
_default_endpoint_deployment("target", deployment_params),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
secret_reads: Final = _SecretReadCounter()
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", secret_reads)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
request_input: Final = (
|
||||
_reasoning_input_minted_by(reasoning_minted_by)
|
||||
if reasoning_minted_by
|
||||
else [{"role": "user", "content": "and sunsets?"}]
|
||||
)
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin="origin"
|
||||
)
|
||||
|
||||
assert secret_reads.reads == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"unreadable_params",
|
||||
[
|
||||
pytest.param(
|
||||
{"model": "openai/gpt-5.5", "litellm_credential_name": "mistyped"},
|
||||
id="named credential with a value of the wrong type",
|
||||
),
|
||||
pytest.param(
|
||||
{**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {1: "one"}},
|
||||
id="extra header whose name is not a string",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_hop_strip_keeps_reasoning_a_target_can_decrypt_beside_a_target_with_unreadable_params(
|
||||
monkeypatch: pytest.MonkeyPatch, unreadable_params: dict[str, object]
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="mistyped",
|
||||
credential_values={"api_key": "shared-key", "api_version": datetime.date(2024, 10, 21)},
|
||||
credential_info={},
|
||||
)
|
||||
],
|
||||
)
|
||||
shared_pair: Final = {**GATEWAY_BASE, "api_key": "shared-key"}
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", shared_pair),
|
||||
_default_endpoint_deployment("unreadable", unreadable_params),
|
||||
_default_endpoint_deployment("target", shared_pair),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
reasoning: Final = _reasoning_input_minted_by("origin")[0]
|
||||
original_reasoning: Final = dict(reasoning)
|
||||
request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}]
|
||||
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
ctx,
|
||||
request_input,
|
||||
None,
|
||||
(_deployment_of(ctx, "unreadable"), _deployment_of(ctx, "target")),
|
||||
unmarked_origin=None,
|
||||
)
|
||||
|
||||
assert request_input[0] == original_reasoning
|
||||
|
||||
|
||||
def test_boundary_key_puts_pydantic_and_plain_params_on_one_default_endpoint_boundary():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
pydantic_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(LiteLLM_Params(**OPENAI_SHARED_KEY))
|
||||
plain_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(dict(OPENAI_SHARED_KEY))
|
||||
|
||||
assert pydantic_key is not None
|
||||
assert plain_key is not None
|
||||
assert pydantic_key.is_shared_with(plain_key)
|
||||
assert plain_key.is_shared_with(pydantic_key)
|
||||
|
||||
|
||||
def test_boundary_key_has_no_default_endpoint_for_a_model_no_provider_claims():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
unmapped: Final = {"model": "unmapped-model-name", "api_key": "shared-key"}
|
||||
|
||||
assert EncryptedContentAffinityCheck._encryption_boundary_key(unmapped) is None
|
||||
|
||||
|
||||
def test_boundary_key_has_no_default_key_for_a_custom_provider(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom"])
|
||||
monkeypatch.setattr(litellm, "_custom_providers", ["qa-custom"])
|
||||
|
||||
assert EncryptedContentAffinityCheck._encryption_boundary_key({"model": "qa-custom/some-model"}) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_affinity_pins_to_the_same_key_peer_on_the_default_endpoint_when_the_origin_cannot_serve():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", OPENAI_SHARED_KEY),
|
||||
_default_endpoint_deployment("peer", OPENAI_SHARED_KEY),
|
||||
_default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
peer: Final = _deployment_of(ctx, "peer")
|
||||
reasoning_input: Final = _reasoning_input_minted_by("origin")
|
||||
original_reasoning: Final = dict(reasoning_input[0])
|
||||
|
||||
result: Final[list[dict[str, object]]] = await EncryptedContentAffinityCheck(ctx=ctx).async_filter_deployments(
|
||||
model="gpt-5.5",
|
||||
healthy_deployments=[peer, _deployment_of(ctx, "other-org")],
|
||||
messages=None,
|
||||
request_kwargs={"litellm_metadata": {}, "input": reasoning_input},
|
||||
)
|
||||
|
||||
assert result == [peer]
|
||||
assert reasoning_input[0] == original_reasoning
|
||||
|
||||
|
||||
class _HandAdvancedClock:
|
||||
def __init__(self) -> None:
|
||||
self.now: float = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
def _count_secret_reads(monkeypatch: pytest.MonkeyPatch) -> _SecretReadCounter:
|
||||
secret_reads: Final = _SecretReadCounter()
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", secret_reads)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
return secret_reads
|
||||
|
||||
|
||||
async def _pin_a_turn_minted_by_the_origin(
|
||||
affinity: object, healthy_deployments: list[dict[str, object]]
|
||||
) -> tuple[list[dict[str, object]], Mapping[str, object]]:
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
assert isinstance(affinity, EncryptedContentAffinityCheck)
|
||||
reasoning_input: Final = _reasoning_input_minted_by("origin")
|
||||
result: Final[list[dict[str, object]]] = await affinity.async_filter_deployments(
|
||||
model="gpt-5.5",
|
||||
healthy_deployments=list(healthy_deployments),
|
||||
messages=None,
|
||||
request_kwargs={"litellm_metadata": {}, "input": reasoning_input},
|
||||
)
|
||||
return result, reasoning_input[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_pins_to_a_default_endpoint_peer_ask_the_secret_manager_once_a_minute(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment("origin", OPENAI_SHARED_KEY),
|
||||
_default_endpoint_deployment("peer", OPENAI_SHARED_KEY),
|
||||
_default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
peer: Final = _deployment_of(ctx, "peer")
|
||||
healthy_deployments: Final = [peer, _deployment_of(ctx, "other-org")]
|
||||
secret_reads: Final = _count_secret_reads(monkeypatch)
|
||||
EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear()
|
||||
clock: Final = _HandAdvancedClock()
|
||||
affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=clock)
|
||||
|
||||
first_pin, first_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments)
|
||||
reads_for_the_first_pin: Final = secret_reads.reads
|
||||
clock.now = 59.0
|
||||
second_pin, second_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments)
|
||||
reads_within_the_minute: Final = secret_reads.reads
|
||||
clock.now = 60.0
|
||||
third_pin, third_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments)
|
||||
|
||||
assert first_pin == second_pin == third_pin == [peer]
|
||||
assert first_reasoning == second_reasoning == third_reasoning == _reasoning_input_minted_by("origin")[0]
|
||||
assert reads_for_the_first_pin > 0
|
||||
assert reads_within_the_minute == reads_for_the_first_pin
|
||||
assert secret_reads.reads > reads_within_the_minute
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manager_refresh_minute(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
def origin_credential(api_key: str) -> list[CredentialItem]:
|
||||
return [
|
||||
CredentialItem(credential_name="origin-account", credential_values={"api_key": api_key}, credential_info={})
|
||||
]
|
||||
|
||||
monkeypatch.setattr(litellm, "credential_list", origin_credential("shared-key"))
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_default_endpoint_deployment(
|
||||
"origin", {"model": "openai/gpt-5.5", "litellm_credential_name": "origin-account"}
|
||||
),
|
||||
_default_endpoint_deployment("peer", OPENAI_SHARED_KEY),
|
||||
_default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}),
|
||||
]
|
||||
)
|
||||
ctx: Final = RoutingContext(router)
|
||||
peer: Final = _deployment_of(ctx, "peer")
|
||||
other_org: Final = _deployment_of(ctx, "other-org")
|
||||
_count_secret_reads(monkeypatch)
|
||||
EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear()
|
||||
affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=_HandAdvancedClock())
|
||||
|
||||
pin_before_the_change, _ = await _pin_a_turn_minted_by_the_origin(affinity, [peer, other_org])
|
||||
monkeypatch.setattr(litellm, "credential_list", origin_credential("other-key"))
|
||||
pin_after_the_change, reasoning_after_the_change = await _pin_a_turn_minted_by_the_origin(
|
||||
affinity, [peer, other_org]
|
||||
)
|
||||
|
||||
assert pin_before_the_change == [peer]
|
||||
assert pin_after_the_change == [other_org]
|
||||
assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue