mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(responses): preserve provider affinity
Co-authored-by: Bytechoreographer <Bytechoreographer@users.noreply.github.com>
This commit is contained in:
AaronHowell
parent
1af7a403c6
commit
d5cf5640b0
2 changed files with 162 additions and 4 deletions
|
|
@ -37,6 +37,7 @@ Safe to enable globally:
|
|||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -48,6 +49,7 @@ from litellm.exceptions import (
|
|||
ServiceUnavailableError,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger, Span
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_utils.cooldown_cache import CooldownCacheValue
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -161,11 +163,13 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
@staticmethod
|
||||
def _encryption_boundary_key(
|
||||
litellm_params: object,
|
||||
) -> tuple | None:
|
||||
) -> tuple[object, object] | None:
|
||||
"""
|
||||
``(api_base, api_key)`` pair identifying an Azure resource. Two
|
||||
deployments sharing both are interchangeable for ``encrypted_content``
|
||||
follow-ups; Azure rejects content produced by any other resource.
|
||||
Missing values are resolved from ``litellm_credential_name`` without
|
||||
modifying the deployment, and explicit deployment values take precedence.
|
||||
|
||||
Accepts any object exposing dict-style ``.get(key, default)``: plain
|
||||
dicts (the common case in ``healthy_deployments``) as well as
|
||||
|
|
@ -180,9 +184,29 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
return None
|
||||
api_base: Final = getter("api_base")
|
||||
api_key: Final = getter("api_key")
|
||||
if not api_base or not 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 and (api_base is None or api_key is None)
|
||||
else None
|
||||
)
|
||||
effective_api_base: Final = (
|
||||
api_base
|
||||
if api_base is not None
|
||||
else credential_values.get("api_base")
|
||||
if credential_values is not None
|
||||
else None
|
||||
)
|
||||
effective_api_key: Final = (
|
||||
api_key
|
||||
if api_key is not None
|
||||
else credential_values.get("api_key")
|
||||
if credential_values is not None
|
||||
else None
|
||||
)
|
||||
if not effective_api_base or not effective_api_key:
|
||||
return None
|
||||
return (api_base, api_key)
|
||||
return (effective_api_base, effective_api_key)
|
||||
|
||||
def _find_deployments_on_same_encryption_boundary(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -21,8 +21,8 @@ from unittest.mock import AsyncMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.models.credentials import CredentialItem
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
|
|
@ -1148,6 +1148,140 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance():
|
|||
)
|
||||
|
||||
|
||||
def test_boundary_key_resolves_missing_values_from_named_credential():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: credential registry is the direct dependency under test
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="account-a",
|
||||
credential_values={
|
||||
"api_base": "https://account-a.example.com",
|
||||
"api_key": "credential-key-a",
|
||||
},
|
||||
credential_info={},
|
||||
)
|
||||
],
|
||||
)
|
||||
):
|
||||
boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"})
|
||||
|
||||
assert boundary == ("https://account-a.example.com", "credential-key-a")
|
||||
|
||||
|
||||
def test_boundary_key_prefers_explicit_values_over_named_credential():
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: credential registry is the direct dependency under test
|
||||
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": "https://deployment.example.com",
|
||||
"litellm_credential_name": "account-a",
|
||||
}
|
||||
)
|
||||
|
||||
assert boundary == ("https://deployment.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 (
|
||||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: credential registry is the direct dependency under test
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="account-a",
|
||||
credential_values={
|
||||
"api_base": "https://account-a.example.com",
|
||||
"api_key": "credential-key-a",
|
||||
},
|
||||
credential_info={},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="account-a-peer",
|
||||
credential_values={
|
||||
"api_base": "https://account-a.example.com",
|
||||
"api_key": "credential-key-a",
|
||||
},
|
||||
credential_info={},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="account-b",
|
||||
credential_values={
|
||||
"api_base": "https://account-b.example.com",
|
||||
"api_key": "credential-key-b",
|
||||
},
|
||||
credential_info={},
|
||||
),
|
||||
],
|
||||
)
|
||||
):
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5.3-codex",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.3-codex",
|
||||
"litellm_credential_name": "account-a",
|
||||
},
|
||||
"model_info": {"id": "origin"},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
check = EncryptedContentAffinityCheck(router=router)
|
||||
healthy_deployments = [
|
||||
{
|
||||
"model_info": {"id": "peer-same-boundary"},
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.4",
|
||||
"litellm_credential_name": "account-a-peer",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_info": {"id": "peer-different-boundary"},
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.4",
|
||||
"litellm_credential_name": "account-b",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
matches, originating = check._find_deployments_on_same_encryption_boundary(
|
||||
healthy_deployments=healthy_deployments,
|
||||
model_id="origin",
|
||||
)
|
||||
|
||||
assert originating is not None
|
||||
assert [deployment["model_info"]["id"] for deployment in matches] == ["peer-same-boundary"]
|
||||
|
||||
|
||||
def test_boundary_key_rejects_non_dict_like_inputs():
|
||||
"""
|
||||
Inputs that don't expose ``.get()`` (None, lists, strings, ints) -> None.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue