mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
feat(anthropic): workload identity federation and pluggable identity sources
Backend half of #38818 (internal copy of the fork PR #38013), rebuilt as one commit on top of litellm_internal_staging without the dashboard changes. Deployments on anthropic/ without a static api_key can exchange an OIDC workload assertion for a short-lived sk-ant-oat01 token through a shared RFC 7523 JWT-bearer engine. The assertion comes from a mounted token file, an env token, a LiteLLM-signed issuer, or Keycloak, chosen per deployment, per named credential, or through ANTHROPIC_IDENTITY_SOURCE. The federation fields are server-owned: refused inline in request bodies and on POST /model/new, proxy-admin only on credentials, and the token exchange is pinned to api.anthropic.com unless LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS adds a host. GET /credentials/{name}/jwks exports the public key set of a LiteLLM-signed credential for the Claude Console. The OpenAI federation trio from #39613 rides along on the backend side with the same server-owned handling. Fixes #28607 Resolves LIT-6107 Co-authored-by: derhornspieler <15236687+derhornspieler@users.noreply.github.com>
This commit is contained in:
parent
a0058ed157
commit
6e44b5b626
113 changed files with 14209 additions and 2683 deletions
|
|
@ -7,7 +7,10 @@ from typing import Any, Final, Literal
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.get_litellm_params import (
|
||||
ANTHROPIC_WIF_KWARGS_KEYS,
|
||||
AWS_CREDENTIAL_KWARGS_KEYS,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
|
||||
from litellm.types.llms.openai import Batch
|
||||
from litellm.types.utils import CallTypes, ModelInfo, Usage
|
||||
|
|
@ -507,6 +510,9 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
"max_retries",
|
||||
"_litellm_internal_model_credentials",
|
||||
*AWS_CREDENTIAL_KWARGS_KEYS,
|
||||
# A federated deployment holds no api_key, so without these the fetch that reads a
|
||||
# finished batch's output has nothing to authenticate with and its cost is never billed.
|
||||
*sorted(ANTHROPIC_WIF_KWARGS_KEYS),
|
||||
)
|
||||
for key in credential_keys:
|
||||
if key in litellm_params:
|
||||
|
|
|
|||
|
|
@ -483,6 +483,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
)
|
||||
api_key = optional_params.api_key or litellm.api_key or litellm.azure_key or get_secret_str("ANTHROPIC_API_KEY")
|
||||
|
||||
batch_params: Final = dict(litellm_params) # mutable-ok: handler contract, copied not shared
|
||||
response = anthropic_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
|
|
@ -490,6 +491,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
api_key=api_key,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
litellm_params=batch_params,
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
|
|
|
|||
|
|
@ -24,10 +24,54 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
|
|||
# The per-deployment Rust opt-in.
|
||||
RUST_KWARG_KEY: Final = "rust"
|
||||
|
||||
# Anthropic workload identity federation config, read from litellm_params by the
|
||||
# Anthropic auth tier. Registered like `rust`: here so the kwargs funnel carries
|
||||
# them, and in `all_litellm_params` so they never leak into the provider body.
|
||||
ANTHROPIC_WIF_KWARGS_KEYS: Final = frozenset(
|
||||
{
|
||||
"anthropic_federation_rule_id",
|
||||
"anthropic_organization_id",
|
||||
"anthropic_service_account_id",
|
||||
"anthropic_workspace_id",
|
||||
"anthropic_identity_token_file",
|
||||
"anthropic_identity_token",
|
||||
# Identity-source selection (Phase 1): absent means the legacy
|
||||
# token_file/env resolver above, byte-identical to today.
|
||||
"anthropic_identity_source",
|
||||
# internal_issuer: litellm self-signs the workload assertion.
|
||||
"anthropic_issuer_url",
|
||||
"anthropic_issuer_subject",
|
||||
"anthropic_issuer_audience",
|
||||
"anthropic_issuer_ttl_seconds",
|
||||
"anthropic_issuer_signing_key_ref",
|
||||
# keycloak: litellm fetches the assertion via client_credentials.
|
||||
"anthropic_keycloak_token_url",
|
||||
"anthropic_keycloak_client_id",
|
||||
"anthropic_keycloak_auth_method",
|
||||
"anthropic_keycloak_client_secret_ref",
|
||||
"anthropic_keycloak_scope",
|
||||
# Set server-side when a client redirects api_base, to stop a federated deployment minting
|
||||
# for a base the caller chose. It has to ride this funnel or it is dropped on the way and
|
||||
# the deployment federates anyway; being carried here also request-bans it, which is right,
|
||||
# since a caller must not be able to set it in either direction.
|
||||
"anthropic_disable_workload_identity_federation",
|
||||
}
|
||||
)
|
||||
|
||||
OPENAI_WIF_KWARGS_KEYS: Final = frozenset(
|
||||
{
|
||||
"openai_identity_provider_id",
|
||||
"openai_service_account_id",
|
||||
"openai_identity_token_file",
|
||||
}
|
||||
)
|
||||
|
||||
# Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
|
||||
# which are otherwise invisible to it because that call site passes explicit
|
||||
# named arguments rather than `**kwargs`.
|
||||
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
|
||||
FORWARDED_KWARGS_KEYS: Final = (
|
||||
AWS_CREDENTIAL_KWARGS_KEYS | ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
|
||||
)
|
||||
|
||||
# Pre-define optional kwargs keys as frozenset for O(1) lookups
|
||||
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
|
||||
|
|
@ -65,6 +109,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
}
|
||||
)
|
||||
| AWS_CREDENTIAL_KWARGS_KEYS
|
||||
| ANTHROPIC_WIF_KWARGS_KEYS
|
||||
| OPENAI_WIF_KWARGS_KEYS
|
||||
)
|
||||
|
||||
# Backward-compatible alias for existing imports/tests.
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ class AnthropicBatchesHandler:
|
|||
timeout: float | httpx.Timeout,
|
||||
max_retries: int | None,
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
Async: Retrieve a batch from Anthropic.
|
||||
|
|
@ -60,9 +61,7 @@ class AnthropicBatchesHandler:
|
|||
# Resolve API credentials
|
||||
api_base = api_base or self.anthropic_model_info.get_api_base(api_base)
|
||||
api_key = api_key or self.anthropic_model_info.get_api_key()
|
||||
|
||||
if not api_key:
|
||||
raise ValueError("Missing Anthropic API Key")
|
||||
resolved_litellm_params: Final = litellm_params if litellm_params is not None else {}
|
||||
|
||||
# Create a minimal logging object if not provided
|
||||
if logging_obj is None:
|
||||
|
|
@ -85,16 +84,18 @@ class AnthropicBatchesHandler:
|
|||
api_base=api_base,
|
||||
batch_id=batch_id,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
litellm_params=resolved_litellm_params,
|
||||
)
|
||||
|
||||
# Validate environment and get headers
|
||||
headers: Final = self.provider_config.validate_environment(
|
||||
# Validate environment and get headers. Offloaded to a worker thread: a WIF token
|
||||
# exchange here would otherwise block the event loop.
|
||||
headers: Final = await asyncio.to_thread(
|
||||
self.provider_config.validate_environment,
|
||||
headers={},
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
litellm_params=resolved_litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
|
@ -130,6 +131,7 @@ class AnthropicBatchesHandler:
|
|||
timeout: float | httpx.Timeout,
|
||||
max_retries: int | None,
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment
|
||||
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
|
||||
"""
|
||||
Retrieve a batch from Anthropic.
|
||||
|
|
@ -154,6 +156,7 @@ class AnthropicBatchesHandler:
|
|||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
else:
|
||||
return asyncio.run(
|
||||
|
|
@ -164,5 +167,6 @@ class AnthropicBatchesHandler:
|
|||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|||
from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest
|
||||
from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse
|
||||
|
||||
from ..common_utils import merge_anthropic_beta_headers, without_caller_credential_headers
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
|
|
@ -70,24 +72,30 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
"""Validate and prepare environment-specific headers and parameters."""
|
||||
if api_base is None and isinstance(litellm_params, dict):
|
||||
api_base = litellm_params.get("api_base")
|
||||
auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base)
|
||||
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
|
||||
if api_base is None and params_mapping is not None:
|
||||
api_base = params_mapping.get("api_base")
|
||||
auth_header: Final = self.anthropic_model_info.get_auth_header(
|
||||
api_key, api_base, litellm_params=params_mapping, allow_workload_identity=True
|
||||
)
|
||||
if auth_header is None:
|
||||
raise ValueError(
|
||||
"Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params"
|
||||
)
|
||||
_headers: Final = {
|
||||
merged_beta: Final = merge_anthropic_beta_headers(
|
||||
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
|
||||
"message-batches-2024-09-24",
|
||||
)
|
||||
# The deployment's own credential is applied below, so a caller-supplied one must not
|
||||
# ride along: without this a minted federation Bearer travels beside the caller's x-api-key.
|
||||
return {
|
||||
**without_caller_credential_headers(headers),
|
||||
"accept": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"content-type": "application/json",
|
||||
**auth_header,
|
||||
"anthropic-beta": merged_beta,
|
||||
}
|
||||
_headers.update(auth_header)
|
||||
# Add beta header for message batches
|
||||
if "anthropic-beta" not in headers:
|
||||
headers["anthropic-beta"] = "message-batches-2024-09-24"
|
||||
headers.update(_headers)
|
||||
return headers
|
||||
|
||||
def get_complete_batch_url(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import re
|
|||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, cast
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
|
|
@ -284,6 +284,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
to pass metadata to anthropic, it's {"user_id": "any-relevant-information"}
|
||||
"""
|
||||
|
||||
_workload_identity_eligible: ClassVar[bool] = True
|
||||
|
||||
max_tokens: int | None = None
|
||||
stop_sequences: list | None = None
|
||||
temperature: int | None = None
|
||||
|
|
|
|||
|
|
@ -7,10 +7,11 @@ import re
|
|||
from collections.abc import Mapping, MutableMapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Any, ClassVar, Final, Literal
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -25,8 +26,14 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
THOUGHT_SIGNATURE_SEPARATOR,
|
||||
)
|
||||
from litellm.llms.anthropic.wif import (
|
||||
aget_anthropic_wif_token,
|
||||
anthropic_base_without_chat_suffix,
|
||||
get_anthropic_wif_token,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.proxy._types import SpecialHeaders
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_HOSTED_TOOLS,
|
||||
ANTHROPIC_OAUTH_BETA_HEADER,
|
||||
|
|
@ -81,6 +88,29 @@ def _strip_bedrock_id_suffixes(model: str) -> str:
|
|||
)
|
||||
|
||||
|
||||
_SERVER_OWNED_AUTH_HEADERS: Final = SpecialHeaders.litellm_credential_header_names()
|
||||
_WIF_ELIGIBILITY_ATTR: Final = "_workload_identity_eligible"
|
||||
|
||||
|
||||
def without_caller_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
|
||||
"""``headers`` minus every header that authenticates the caller to litellm.
|
||||
|
||||
The deployment's own credential is applied on top of the result, so a caller-supplied
|
||||
credential must not survive into the upstream request: without this a minted federation
|
||||
Bearer travels beside the caller's own ``x-api-key``, and Anthropic sees two credentials.
|
||||
"""
|
||||
return MappingProxyType(
|
||||
{name: value for name, value in headers.items() if name.lower() not in _SERVER_OWNED_AUTH_HEADERS}
|
||||
)
|
||||
|
||||
|
||||
def config_allows_workload_identity(config: object) -> bool:
|
||||
"""A federation token is an Anthropic-org credential and its exchange POSTs the workload's OIDC
|
||||
assertion to the deployment's own host, so eligibility is declared per class and read from that
|
||||
class's own ``__dict__``: a subclass written for another provider inherits nothing."""
|
||||
return type(config).__dict__.get(_WIF_ELIGIBILITY_ATTR, False) is True
|
||||
|
||||
|
||||
def is_anthropic_oauth_key(value: str | None) -> bool:
|
||||
"""Check if a value contains an Anthropic OAuth token (sk-ant-oat*)."""
|
||||
if value is None:
|
||||
|
|
@ -90,12 +120,20 @@ def is_anthropic_oauth_key(value: str | None) -> bool:
|
|||
return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX)
|
||||
|
||||
|
||||
def _merge_beta_headers(existing: str | None, new_beta: str) -> str:
|
||||
"""Merge a new beta value into an existing comma-separated anthropic-beta header."""
|
||||
if not existing:
|
||||
return new_beta
|
||||
betas: Final = {b.strip() for b in existing.split(",") if b.strip()}
|
||||
betas.add(new_beta)
|
||||
def merge_anthropic_beta_headers(existing: str | Sequence[str] | None, new_beta: str | Sequence[str] | None) -> str:
|
||||
"""Merge anthropic-beta header values, deduplicated and sorted.
|
||||
|
||||
Either side may arrive as a list rather than a comma-separated string: the Skills surface
|
||||
accepted a list-valued header before it shared this helper, and callers still send one.
|
||||
"""
|
||||
values: Final = (
|
||||
entry
|
||||
for side in (existing, new_beta)
|
||||
if side
|
||||
for entry in ((side,) if isinstance(side, str) else side)
|
||||
if isinstance(entry, str)
|
||||
)
|
||||
betas: Final = frozenset(b.strip() for value in values for b in value.split(",") if b.strip())
|
||||
return ",".join(sorted(betas))
|
||||
|
||||
|
||||
|
|
@ -122,7 +160,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup
|
|||
):
|
||||
headers.pop(name)
|
||||
headers["authorization"] = auth_header
|
||||
headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER)
|
||||
headers["anthropic-beta"] = merge_anthropic_beta_headers(
|
||||
headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER
|
||||
)
|
||||
headers["anthropic-dangerous-direct-browser-access"] = "true"
|
||||
return headers, api_key
|
||||
# Check api_key directly (standard chat/completion flow)
|
||||
|
|
@ -130,7 +170,9 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup
|
|||
for name in tuple(header_name for header_name in headers if header_name.lower() == "x-api-key"):
|
||||
headers.pop(name)
|
||||
headers["authorization"] = f"Bearer {api_key}"
|
||||
headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER)
|
||||
headers["anthropic-beta"] = merge_anthropic_beta_headers(
|
||||
headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER
|
||||
)
|
||||
headers["anthropic-dangerous-direct-browser-access"] = "true"
|
||||
return headers, api_key
|
||||
|
||||
|
|
@ -145,7 +187,79 @@ class AnthropicError(BaseLLMException):
|
|||
super().__init__(status_code=status_code, message=message, headers=headers)
|
||||
|
||||
|
||||
_MODEL_LIST_PAGE_CAP: Final = 20
|
||||
|
||||
|
||||
def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
|
||||
value: Final = litellm_params.get(key) if litellm_params is not None else None
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
class _AnthropicModelListEntry(BaseModel):
|
||||
id: str
|
||||
|
||||
|
||||
class _AnthropicModelsPage(BaseModel):
|
||||
data: Sequence[_AnthropicModelListEntry] = Field(default_factory=tuple)
|
||||
has_more: bool = False
|
||||
last_id: str | None = None
|
||||
|
||||
|
||||
def _sanitized_anthropic_error(response: httpx.Response, detail: str | None = None) -> str:
|
||||
"""A provider error detail built only from structured fields, never ``response.text``
|
||||
verbatim: the raw body is untrusted content the caller of ``/v1/models`` did not ask for
|
||||
and should not have echoed back to it wholesale."""
|
||||
if detail is not None:
|
||||
return f"HTTP {response.status_code}: {detail}"
|
||||
try:
|
||||
body: Final = response.json()
|
||||
except ValueError:
|
||||
return f"HTTP {response.status_code}"
|
||||
error: Final = body.get("error") if isinstance(body, dict) else None
|
||||
message: Final = error.get("message") if isinstance(error, dict) else None
|
||||
return f"HTTP {response.status_code}: {message}" if isinstance(message, str) else f"HTTP {response.status_code}"
|
||||
|
||||
|
||||
def _fetch_anthropic_models_page(
|
||||
api_base: str, headers: Mapping[str, str], after_id: str | None
|
||||
) -> _AnthropicModelsPage:
|
||||
# after_id rides the URL because the client mutates the params mapping it is handed,
|
||||
# which a read-only one cannot support
|
||||
query: Final = f"?after_id={quote(after_id)}" if after_id else ""
|
||||
response: Final = litellm.module_level_client.get(
|
||||
url=f"{api_base}/v1/models{query}",
|
||||
headers=headers,
|
||||
follow_redirects=False,
|
||||
)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError:
|
||||
raise Exception(f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response)}") from None
|
||||
try:
|
||||
return _AnthropicModelsPage.model_validate(response.json())
|
||||
except ValueError as e:
|
||||
raise Exception(
|
||||
f"Failed to fetch models from Anthropic. {_sanitized_anthropic_error(response, detail=str(e))}"
|
||||
) from None
|
||||
|
||||
|
||||
def _fetch_anthropic_model_ids(
|
||||
api_base: str, headers: Mapping[str, str], after_id: str | None, pages_left: int
|
||||
) -> tuple[str, ...]:
|
||||
collected: tuple[str, ...] = () # rebind-ok: accumulates one page of ids per iteration
|
||||
cursor: str | None = after_id # rebind-ok: advances to each page's last_id
|
||||
for _ in range(max(pages_left, 0)):
|
||||
page = _fetch_anthropic_models_page(api_base, headers, cursor) # rebind-ok: one page per iteration
|
||||
collected += tuple(entry.id for entry in page.data)
|
||||
if not page.has_more or page.last_id is None:
|
||||
return collected
|
||||
cursor = page.last_id
|
||||
raise Exception(f"Anthropic /v1/models did not terminate within {_MODEL_LIST_PAGE_CAP} pages.")
|
||||
|
||||
|
||||
class AnthropicModelInfo(BaseLLMModelInfo):
|
||||
_workload_identity_eligible: ClassVar[bool] = True
|
||||
|
||||
def is_cache_control_set(self, messages: list[AllMessageValues]) -> bool:
|
||||
"""
|
||||
Return if {"cache_control": ..} in message content block
|
||||
|
|
@ -720,7 +834,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
return list(set(betas))
|
||||
|
||||
@staticmethod
|
||||
def _make_api_key_auth_header(api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False) -> dict:
|
||||
def _make_api_key_auth_header(
|
||||
api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False
|
||||
) -> Mapping[str, str]:
|
||||
if use_bearer_for_custom_base and (
|
||||
api_base and "api.anthropic.com" not in api_base and not api_key.startswith("sk-ant-")
|
||||
):
|
||||
|
|
@ -728,6 +844,33 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
return {"authorization": value}
|
||||
return {"x-api-key": api_key}
|
||||
|
||||
def _credential_headers(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None,
|
||||
auth_token: str | None,
|
||||
api_base: str | None,
|
||||
use_bearer_for_custom_base: bool,
|
||||
wif_minted: bool,
|
||||
betas: set[str], # mutable-ok: the caller's beta accumulator, appended to by the oauth tier
|
||||
) -> Mapping[str, str]:
|
||||
"""The credential tier walk: a consumer OAuth token, then ANTHROPIC_AUTH_TOKEN, then an api key.
|
||||
|
||||
A server-minted federation token takes the same Bearer shape as a consumer OAuth token but is
|
||||
not browser-forwarded, so it does not get the direct-browser-access header.
|
||||
"""
|
||||
if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX):
|
||||
betas.add(ANTHROPIC_OAUTH_BETA_HEADER)
|
||||
oauth_headers: Final = {"authorization": f"Bearer {api_key}"}
|
||||
if wif_minted:
|
||||
return oauth_headers
|
||||
return {**oauth_headers, "anthropic-dangerous-direct-browser-access": "true"}
|
||||
if auth_token and not api_key:
|
||||
return {"authorization": f"Bearer {auth_token}"}
|
||||
if api_key:
|
||||
return self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base)
|
||||
return {}
|
||||
|
||||
def get_anthropic_headers(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
|
|
@ -749,6 +892,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
container_with_skills_used: bool = False,
|
||||
api_base: str | None = None,
|
||||
use_bearer_for_custom_base: bool = False,
|
||||
wif_minted: bool = False,
|
||||
) -> dict:
|
||||
betas: Final = set()
|
||||
# Anthropic no longer requires the prompt-caching beta header
|
||||
|
|
@ -784,20 +928,21 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
if container_with_skills_used:
|
||||
betas.add("skills-2025-10-02")
|
||||
|
||||
_is_oauth: Final = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX)
|
||||
headers: Final = {
|
||||
"anthropic-version": anthropic_version or "2023-06-01",
|
||||
"accept": "application/json",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
if _is_oauth:
|
||||
headers["authorization"] = f"Bearer {api_key}"
|
||||
headers["anthropic-dangerous-direct-browser-access"] = "true"
|
||||
betas.add(ANTHROPIC_OAUTH_BETA_HEADER)
|
||||
elif auth_token and not api_key:
|
||||
headers["authorization"] = f"Bearer {auth_token}"
|
||||
elif api_key:
|
||||
headers.update(self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base))
|
||||
headers.update(
|
||||
self._credential_headers(
|
||||
api_key=api_key,
|
||||
auth_token=auth_token,
|
||||
api_base=api_base,
|
||||
use_bearer_for_custom_base=use_bearer_for_custom_base,
|
||||
wif_minted=wif_minted,
|
||||
betas=betas,
|
||||
)
|
||||
)
|
||||
|
||||
if user_anthropic_beta_headers is not None:
|
||||
betas.update(user_anthropic_beta_headers)
|
||||
|
|
@ -824,10 +969,11 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
if api_base is None and isinstance(litellm_params, dict):
|
||||
api_base = litellm_params.get("api_base")
|
||||
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
|
||||
if api_base is None and params_mapping is not None:
|
||||
api_base = params_mapping.get("api_base")
|
||||
use_bearer_for_custom_base: Final[bool] = bool(
|
||||
isinstance(litellm_params, dict) and litellm_params.get("use_bearer_for_custom_base", False)
|
||||
params_mapping is not None and params_mapping.get("use_bearer_for_custom_base", False)
|
||||
)
|
||||
# Check for Anthropic OAuth token in headers
|
||||
headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
|
||||
|
|
@ -836,9 +982,23 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
auth_token: str | None = None
|
||||
if api_key is None:
|
||||
auth_token = AnthropicModelInfo.get_auth_token()
|
||||
if api_key is None and auth_token is None:
|
||||
wif_token: Final = (
|
||||
get_anthropic_wif_token(params_mapping, api_base, model)
|
||||
if api_key is None and auth_token is None and config_allows_workload_identity(self)
|
||||
else None
|
||||
)
|
||||
wif_minted: Final = wif_token is not None
|
||||
resolved_api_key: Final = wif_token if wif_token is not None else api_key
|
||||
if resolved_api_key is None and auth_token is None:
|
||||
raise litellm.AuthenticationError(
|
||||
message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` in your environment vars",
|
||||
message=(
|
||||
"Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the "
|
||||
"environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` "
|
||||
"in your environment vars, or configure workload identity federation via "
|
||||
"`ANTHROPIC_FEDERATION_RULE_ID`, `ANTHROPIC_ORGANIZATION_ID`, "
|
||||
"`ANTHROPIC_SERVICE_ACCOUNT_ID` and "
|
||||
"`ANTHROPIC_IDENTITY_TOKEN_FILE` (or `ANTHROPIC_IDENTITY_TOKEN`)"
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model=model,
|
||||
)
|
||||
|
|
@ -863,7 +1023,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
computer_tool_used=computer_tool_used,
|
||||
prompt_caching_set=prompt_caching_set,
|
||||
pdf_used=pdf_used,
|
||||
api_key=api_key,
|
||||
api_key=resolved_api_key,
|
||||
auth_token=auth_token,
|
||||
file_id_used=file_id_used,
|
||||
web_search_tool_used=web_search_tool_used,
|
||||
|
|
@ -878,11 +1038,12 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
container_with_skills_used=container_with_skills_used,
|
||||
api_base=api_base,
|
||||
use_bearer_for_custom_base=use_bearer_for_custom_base,
|
||||
wif_minted=wif_minted,
|
||||
)
|
||||
|
||||
headers = {**headers, **anthropic_headers}
|
||||
caller_headers: Final = without_caller_credential_headers(headers) if wif_minted else headers
|
||||
|
||||
return headers
|
||||
return {**caller_headers, **anthropic_headers}
|
||||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str | None:
|
||||
|
|
@ -917,52 +1078,121 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
use_bearer_for_custom_base: bool = False,
|
||||
) -> dict | None:
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
allow_workload_identity: bool = False,
|
||||
) -> Mapping[str, str] | None:
|
||||
"""Resolve Anthropic credentials and return the appropriate auth header dict.
|
||||
|
||||
Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on
|
||||
use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer).
|
||||
Returns None if neither is available.
|
||||
use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer),
|
||||
then workload identity federation (-> Authorization: Bearer with a minted
|
||||
sk-ant-oat01 token, honoring anthropic_* litellm_params when provided). Every
|
||||
Bearer built from an sk-ant-oat token carries the mandatory oauth anthropic-beta.
|
||||
Returns None if no credential source is available.
|
||||
"""
|
||||
static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base)
|
||||
if static_header is not None:
|
||||
return static_header
|
||||
if not allow_workload_identity:
|
||||
return None
|
||||
wif_token: Final = get_anthropic_wif_token(litellm_params, api_base, "")
|
||||
if wif_token is not None:
|
||||
return AnthropicModelInfo._oauth_bearer_header(wif_token)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
async def aget_auth_header(
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
use_bearer_for_custom_base: bool = False,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
allow_workload_identity: bool = False,
|
||||
) -> Mapping[str, str] | None:
|
||||
"""Async counterpart of get_auth_header: the WIF tier can block on a token
|
||||
exchange POST, so async callers await it off the event loop."""
|
||||
static_header: Final = AnthropicModelInfo._static_auth_header(api_key, api_base, use_bearer_for_custom_base)
|
||||
if static_header is not None:
|
||||
return static_header
|
||||
if not allow_workload_identity:
|
||||
return None
|
||||
wif_token: Final = await aget_anthropic_wif_token(litellm_params, api_base, "")
|
||||
if wif_token is not None:
|
||||
return AnthropicModelInfo._oauth_bearer_header(wif_token)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _static_auth_header(
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
use_bearer_for_custom_base: bool,
|
||||
) -> Mapping[str, str] | None:
|
||||
resolved_key: Final = AnthropicModelInfo.get_api_key(api_key)
|
||||
if resolved_key is not None:
|
||||
if is_anthropic_oauth_key(resolved_key):
|
||||
return {"authorization": f"Bearer {resolved_key}"}
|
||||
return AnthropicModelInfo._oauth_bearer_header(resolved_key)
|
||||
return AnthropicModelInfo._make_api_key_auth_header(resolved_key, api_base, use_bearer_for_custom_base)
|
||||
auth_token: Final = AnthropicModelInfo.get_auth_token()
|
||||
if auth_token is not None:
|
||||
return {"authorization": f"Bearer {auth_token}"}
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _oauth_bearer_header(token: str) -> Mapping[str, str]:
|
||||
return {"authorization": f"Bearer {token}", "anthropic-beta": ANTHROPIC_OAUTH_BETA_HEADER}
|
||||
|
||||
@staticmethod
|
||||
def get_base_model(model: str | None = None) -> str | None:
|
||||
return model.replace("anthropic/", "") if model else None
|
||||
|
||||
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
|
||||
api_base = AnthropicModelInfo.get_api_base(api_base)
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
|
||||
if api_base is None or auth_header is None:
|
||||
raise ValueError(
|
||||
"ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint."
|
||||
)
|
||||
headers: Final = {"anthropic-version": "2023-06-01"}
|
||||
headers.update(auth_header)
|
||||
response: Final = litellm.module_level_client.get(
|
||||
url=f"{api_base}/v1/models",
|
||||
headers=headers,
|
||||
return self._list_models(api_key=api_key, api_base=api_base, litellm_params=None)
|
||||
|
||||
def discover_models(
|
||||
self, litellm_params: Mapping[str, object] | None = None
|
||||
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
"""Live discovery for a configured deployment: unlike ``get_models``, this threads the
|
||||
full ``litellm_params`` into ``get_auth_header`` so a workload-identity-federation source
|
||||
configured on the deployment (rather than the environment) is honored, gated the same way
|
||||
every other Anthropic auth surface is via ``config_allows_workload_identity``."""
|
||||
return self._list_models(
|
||||
api_key=_litellm_params_str(litellm_params, "api_key"),
|
||||
api_base=_litellm_params_str(litellm_params, "api_base"),
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError:
|
||||
raise Exception(
|
||||
f"Failed to fetch models from Anthropic. Status code: {response.status_code}, Response: {response.text}"
|
||||
def _list_models(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base)
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(
|
||||
api_key,
|
||||
resolved_api_base,
|
||||
litellm_params=litellm_params,
|
||||
allow_workload_identity=config_allows_workload_identity(self),
|
||||
)
|
||||
if resolved_api_base is None or auth_header is None:
|
||||
raise ValueError(
|
||||
"ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN (or workload "
|
||||
"identity federation via ANTHROPIC_FEDERATION_RULE_ID/ANTHROPIC_ORGANIZATION_ID/"
|
||||
"ANTHROPIC_IDENTITY_TOKEN_FILE) is not set. Please set the environment variable, to query "
|
||||
"Anthropic's `/models` endpoint."
|
||||
)
|
||||
|
||||
models: Final[Sequence[Mapping[str, str]]] = response.json()["data"]
|
||||
|
||||
litellm_model_names: Final = ["anthropic/" + model["id"] for model in models]
|
||||
return litellm_model_names
|
||||
headers: Final = MappingProxyType({"anthropic-version": "2023-06-01", **auth_header})
|
||||
# /v1/models is appended below, so a base the operator already wrote as .../v1 or
|
||||
# .../v1/messages would otherwise be asked for /v1/v1/models.
|
||||
model_ids: Final = _fetch_anthropic_model_ids(
|
||||
anthropic_base_without_chat_suffix(resolved_api_base),
|
||||
headers,
|
||||
after_id=None,
|
||||
pages_left=_MODEL_LIST_PAGE_CAP,
|
||||
)
|
||||
return [ # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
"anthropic/" + model_id for model_id in model_ids
|
||||
]
|
||||
|
||||
def get_token_counter(self) -> BaseTokenCounter | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
|||
model: The model identifier (e.g., "claude-3-5-sonnet-20241022")
|
||||
messages: The messages to count tokens for
|
||||
api_key: The Anthropic API key
|
||||
api_base: Optional custom API base URL
|
||||
api_base: Optional deployment api_base the count-tokens path is appended to
|
||||
timeout: Optional timeout for the request (defaults to litellm.request_timeout)
|
||||
|
||||
Returns:
|
||||
|
|
@ -67,7 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
|||
verbose_logger.debug("Transformed request: %s", request_body)
|
||||
|
||||
# Get endpoint URL
|
||||
endpoint_url: Final = api_base or self.get_anthropic_count_tokens_endpoint()
|
||||
endpoint_url: Final = self.get_anthropic_count_tokens_endpoint(api_base)
|
||||
|
||||
verbose_logger.debug("Making request to: %s", endpoint_url)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,10 @@
|
|||
Anthropic Token Counter implementation using the CountTokens API.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
|
|
@ -46,28 +46,33 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
Returns:
|
||||
TokenCountResponse with token count, or None if counting fails
|
||||
"""
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo
|
||||
from litellm.llms.anthropic.wif import aget_anthropic_wif_token
|
||||
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
|
||||
# Get Anthropic API key from deployment config or environment
|
||||
api_key = litellm_params.get("api_key")
|
||||
if not api_key:
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No Anthropic API key found for token counting")
|
||||
return None
|
||||
api_base: Final = litellm_params.get("api_base")
|
||||
static_key: Final = AnthropicModelInfo.get_api_key(litellm_params.get("api_key"))
|
||||
auth_token_configured: Final = AnthropicModelInfo.get_auth_token() is not None
|
||||
|
||||
try:
|
||||
api_key: Final = (
|
||||
static_key
|
||||
if static_key or auth_token_configured
|
||||
else await aget_anthropic_wif_token(litellm_params, api_base, model_to_use)
|
||||
)
|
||||
if not api_key:
|
||||
verbose_logger.warning("No Anthropic credential found for token counting")
|
||||
return None
|
||||
|
||||
result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(
|
||||
model=model_to_use,
|
||||
messages=messages,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
tools=tools,
|
||||
system=system,
|
||||
)
|
||||
|
|
@ -80,8 +85,8 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
tokenizer_type="anthropic_api",
|
||||
original_response=result,
|
||||
)
|
||||
except AnthropicError as e:
|
||||
verbose_logger.warning("Anthropic CountTokens API error: status=%s, message=%s", e.status_code, e.message)
|
||||
except (AnthropicError, AuthenticationError) as e:
|
||||
verbose_logger.warning("Anthropic CountTokens error: status=%s, message=%s", e.status_code, e.message)
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ This module handles the transformation of requests to Anthropic's CountTokens AP
|
|||
from typing import Any, Final
|
||||
|
||||
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
from litellm.llms.anthropic.wif import resolve_anthropic_base
|
||||
|
||||
|
||||
class AnthropicCountTokensConfig:
|
||||
|
|
@ -19,14 +20,21 @@ class AnthropicCountTokensConfig:
|
|||
- Response: {"input_tokens": <number>}
|
||||
"""
|
||||
|
||||
def get_anthropic_count_tokens_endpoint(self) -> str:
|
||||
def get_anthropic_count_tokens_endpoint(self, api_base: str | None = None) -> str:
|
||||
"""
|
||||
Get the Anthropic CountTokens API endpoint.
|
||||
|
||||
Args:
|
||||
api_base: The deployment's api_base, which names the chat surface (a host, or a
|
||||
base already carrying ``/v1`` or ``/v1/messages``); the count-tokens path is
|
||||
appended to it, so it is never the full count-tokens URL. Unset or empty falls
|
||||
back to ``ANTHROPIC_API_BASE`` / ``ANTHROPIC_BASE_URL`` and then Anthropic's
|
||||
host, the same resolution chat and the federated exchange use
|
||||
|
||||
Returns:
|
||||
The endpoint URL for the CountTokens API
|
||||
"""
|
||||
return "https://api.anthropic.com/v1/messages/count_tokens"
|
||||
return resolve_anthropic_base(api_base) + "/v1/messages/count_tokens"
|
||||
|
||||
def transform_request_to_count_tokens(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
from typing import Any, ClassVar, Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -23,12 +23,24 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
from ...common_utils import (
|
||||
AnthropicError,
|
||||
AnthropicModelInfo,
|
||||
merge_anthropic_beta_headers,
|
||||
optionally_handle_anthropic_oauth,
|
||||
strip_advisor_blocks_from_messages,
|
||||
)
|
||||
|
||||
DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
|
||||
|
||||
_CALLER_CREDENTIAL_HEADERS: Final = frozenset({"x-api-key", "authorization"})
|
||||
|
||||
|
||||
def _carries_caller_credential(headers: Mapping[str, str]) -> bool:
|
||||
"""Whether the caller sent their own Anthropic credential, in which case this passthrough
|
||||
honors it and never mints. Matched case-insensitively: an SDK caller passing ``X-Api-Key``
|
||||
through extra_headers would otherwise slip the check and end up sending their key beside a
|
||||
minted federation Bearer."""
|
||||
return any(name.lower() in _CALLER_CREDENTIAL_HEADERS for name in headers)
|
||||
|
||||
|
||||
DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING: Final = (
|
||||
"Dropping adaptive `thinking`/`output_config.effort` for model=%s: the model "
|
||||
"does not support extended thinking, or max_tokens is too small to fit the "
|
||||
|
|
@ -42,6 +54,8 @@ DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = (
|
|||
|
||||
|
||||
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
||||
_workload_identity_eligible: ClassVar[bool] = True
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
return "anthropic"
|
||||
|
|
@ -308,32 +322,103 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
# Check for Anthropic OAuth token in Authorization header
|
||||
headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
|
||||
|
||||
header_names: Final = frozenset(name.lower() for name in headers)
|
||||
if "x-api-key" not in header_names and "authorization" not in header_names:
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key)
|
||||
if auth_header is None:
|
||||
raise AuthenticationError(
|
||||
message=(
|
||||
"Missing Anthropic API Key - A call is being made to anthropic but no key is set "
|
||||
"either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` "
|
||||
"or `ANTHROPIC_AUTH_TOKEN` in your environment vars"
|
||||
if not _carries_caller_credential(headers):
|
||||
self._apply_env_auth_header(
|
||||
headers,
|
||||
self._require_auth_header(
|
||||
AnthropicModelInfo.get_auth_header(
|
||||
api_key,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
allow_workload_identity=self._allows_workload_identity,
|
||||
),
|
||||
llm_provider=self._resolved_provider,
|
||||
model=model,
|
||||
)
|
||||
headers.update(auth_header)
|
||||
),
|
||||
)
|
||||
return self._finalize_messages_headers(headers, optional_params), api_base
|
||||
|
||||
async def avalidate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
model: str,
|
||||
messages: list[Any], # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
optional_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
litellm_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]: # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
if type(self).validate_anthropic_messages_environment is not (
|
||||
AnthropicMessagesConfig.validate_anthropic_messages_environment
|
||||
):
|
||||
# a subclass sync override must keep winning on the async path
|
||||
return self.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
oauth_headers, oauth_api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
|
||||
|
||||
if not _carries_caller_credential(oauth_headers):
|
||||
self._apply_env_auth_header(
|
||||
oauth_headers,
|
||||
self._require_auth_header(
|
||||
await AnthropicModelInfo.aget_auth_header(
|
||||
oauth_api_key,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
allow_workload_identity=self._allows_workload_identity,
|
||||
),
|
||||
model=model,
|
||||
),
|
||||
)
|
||||
return self._finalize_messages_headers(oauth_headers, optional_params), api_base
|
||||
|
||||
def _require_auth_header(self, auth_header: Mapping[str, str] | None, model: str) -> Mapping[str, str]:
|
||||
if auth_header is None:
|
||||
raise AuthenticationError(
|
||||
message=(
|
||||
"Missing Anthropic API Key - A call is being made to anthropic but no key is set "
|
||||
"either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` "
|
||||
"or `ANTHROPIC_AUTH_TOKEN` in your environment vars"
|
||||
),
|
||||
llm_provider=self._resolved_provider,
|
||||
model=model,
|
||||
)
|
||||
return auth_header
|
||||
|
||||
@staticmethod
|
||||
def _apply_env_auth_header(headers: dict, auth_header: Mapping[str, str] | None) -> None: # mutable-ok: out-param
|
||||
if auth_header is None:
|
||||
return
|
||||
merged_beta: Final = merge_anthropic_beta_headers(
|
||||
headers.get("anthropic-beta"), auth_header.get("anthropic-beta")
|
||||
)
|
||||
headers.update(auth_header)
|
||||
if merged_beta:
|
||||
headers["anthropic-beta"] = merged_beta
|
||||
|
||||
@property
|
||||
def _allows_workload_identity(self) -> bool:
|
||||
"""Subclasses reuse this validate step for their own /v1/messages-compatible providers, so
|
||||
eligibility is declared per class and never inherited."""
|
||||
from litellm.llms.anthropic.common_utils import config_allows_workload_identity
|
||||
|
||||
return config_allows_workload_identity(self)
|
||||
|
||||
def _finalize_messages_headers(self, headers: dict, optional_params: dict) -> dict: # mutable-ok: out-param
|
||||
if "anthropic-version" not in headers:
|
||||
headers["anthropic-version"] = DEFAULT_ANTHROPIC_API_VERSION
|
||||
if "content-type" not in headers:
|
||||
headers["content-type"] = "application/json"
|
||||
|
||||
headers = self._update_headers_with_anthropic_beta(
|
||||
return self._update_headers_with_anthropic_beta(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
|
||||
return headers, api_base
|
||||
|
||||
@staticmethod
|
||||
def _translate_reasoning_effort_to_anthropic(
|
||||
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str
|
||||
|
|
|
|||
|
|
@ -73,7 +73,9 @@ class AnthropicFilesHandler:
|
|||
|
||||
# Get Anthropic API credentials
|
||||
api_base = self.anthropic_model_info.get_api_base(api_base)
|
||||
auth_header: Final = self.anthropic_model_info.get_auth_header(api_key, api_base)
|
||||
auth_header: Final = await self.anthropic_model_info.aget_auth_header(
|
||||
api_key, api_base, allow_workload_identity=True
|
||||
)
|
||||
|
||||
if auth_header is None:
|
||||
raise ValueError("Missing Anthropic API Key")
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ Anthropic Files API endpoints:
|
|||
|
||||
import calendar
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -35,7 +36,12 @@ from litellm.types.llms.openai import (
|
|||
)
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
from ..common_utils import AnthropicError, AnthropicModelInfo
|
||||
from ..common_utils import (
|
||||
AnthropicError,
|
||||
AnthropicModelInfo,
|
||||
merge_anthropic_beta_headers,
|
||||
without_caller_credential_headers,
|
||||
)
|
||||
|
||||
ANTHROPIC_FILES_API_BASE: Final = "https://api.anthropic.com"
|
||||
ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14"
|
||||
|
|
@ -94,21 +100,55 @@ class AnthropicFilesConfig(BaseFilesConfig):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
if api_base is None and isinstance(litellm_params, dict):
|
||||
api_base = litellm_params.get("api_base")
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
|
||||
params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base)
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(
|
||||
api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True
|
||||
)
|
||||
return self._finalize_headers(headers, auth_header)
|
||||
|
||||
async def avalidate_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
model: str,
|
||||
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
"""Async counterpart of validate_environment: the WIF tier can block on a token
|
||||
exchange POST, so async callers await it off the event loop."""
|
||||
params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base)
|
||||
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
|
||||
api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True
|
||||
)
|
||||
return self._finalize_headers(headers, auth_header)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_params(
|
||||
litellm_params: dict, api_base: str | None
|
||||
) -> tuple[dict | None, str | None]: # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None
|
||||
if api_base is None and params_mapping is not None:
|
||||
api_base = params_mapping.get("api_base")
|
||||
return params_mapping, api_base
|
||||
|
||||
@staticmethod
|
||||
def _finalize_headers(headers: dict, auth_header: Mapping[str, str] | None) -> dict: # mutable-ok: out-param
|
||||
if auth_header is None:
|
||||
raise ValueError(
|
||||
"Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter."
|
||||
)
|
||||
headers.update(
|
||||
{
|
||||
**auth_header,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": ANTHROPIC_FILES_BETA_HEADER,
|
||||
}
|
||||
merged_beta: Final = merge_anthropic_beta_headers(
|
||||
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
|
||||
ANTHROPIC_FILES_BETA_HEADER,
|
||||
)
|
||||
return headers
|
||||
return {
|
||||
**without_caller_credential_headers(headers),
|
||||
**auth_header,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": merged_beta,
|
||||
}
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]:
|
||||
return ["purpose"]
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Anthropic Skills API configuration and transformations
|
||||
"""
|
||||
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -35,40 +36,35 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
|
||||
def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
|
||||
"""Add Anthropic-specific headers"""
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
AnthropicModelInfo,
|
||||
merge_anthropic_beta_headers,
|
||||
without_caller_credential_headers,
|
||||
)
|
||||
|
||||
# Get API key from litellm_params if available
|
||||
api_key = None
|
||||
api_base = None
|
||||
if litellm_params is not None:
|
||||
api_key = litellm_params.api_key
|
||||
api_base = litellm_params.api_base
|
||||
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(api_key, api_base)
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(
|
||||
api_key=litellm_params.api_key if litellm_params is not None else None,
|
||||
api_base=litellm_params.api_base if litellm_params is not None else None,
|
||||
litellm_params=MappingProxyType(dict(litellm_params)) if litellm_params is not None else None,
|
||||
allow_workload_identity=True,
|
||||
)
|
||||
if auth_header is None:
|
||||
raise ValueError("ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API")
|
||||
|
||||
headers.update(auth_header)
|
||||
headers["anthropic-version"] = "2023-06-01"
|
||||
|
||||
# Add beta header for skills API
|
||||
from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION
|
||||
|
||||
if "anthropic-beta" not in headers:
|
||||
headers["anthropic-beta"] = ANTHROPIC_SKILLS_API_BETA_VERSION
|
||||
elif isinstance(headers["anthropic-beta"], list):
|
||||
if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]:
|
||||
headers["anthropic-beta"].append(ANTHROPIC_SKILLS_API_BETA_VERSION)
|
||||
elif isinstance(headers["anthropic-beta"], str):
|
||||
if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]:
|
||||
headers["anthropic-beta"] = [
|
||||
headers["anthropic-beta"],
|
||||
ANTHROPIC_SKILLS_API_BETA_VERSION,
|
||||
]
|
||||
|
||||
headers["content-type"] = "application/json"
|
||||
|
||||
return headers
|
||||
merged_beta: Final = merge_anthropic_beta_headers(
|
||||
merge_anthropic_beta_headers(headers.get("anthropic-beta"), auth_header.get("anthropic-beta")),
|
||||
ANTHROPIC_SKILLS_API_BETA_VERSION,
|
||||
)
|
||||
# The deployment's own credential is applied here, so a caller-supplied one must not ride
|
||||
# along upstream beside a minted federation Bearer.
|
||||
return { # mutable-ok: validate_environment's contract returns a real dict, which httpx then consumes
|
||||
**without_caller_credential_headers(headers),
|
||||
**auth_header,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": merged_beta,
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
501
litellm/llms/anthropic/wif.py
Normal file
501
litellm/llms/anthropic/wif.py
Normal file
|
|
@ -0,0 +1,501 @@
|
|||
"""Anthropic workload identity federation: exchanges an external OIDC identity
|
||||
token for a short-lived ``sk-ant-oat01`` token via the shared RFC 7523 engine."""
|
||||
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, TypeVar
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source
|
||||
from litellm.llms.base_llm.auth.identity_source import (
|
||||
AnthropicIdentitySourceKind,
|
||||
InternalIssuerSource,
|
||||
KeycloakSource,
|
||||
identity_source_ref,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_assertion_source
|
||||
from litellm.llms.base_llm.auth.token_exchange import (
|
||||
JwtBearerTokenExchangeEngine,
|
||||
default_token_exchange_engine,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.types import (
|
||||
AssertionSourceError,
|
||||
ExchangeError,
|
||||
ExchangeResult,
|
||||
InsecureTokenUrl,
|
||||
MalformedTokenResponse,
|
||||
MintedToken,
|
||||
TokenEndpointError,
|
||||
TokenExchangeSpec,
|
||||
TokenTransportError,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_TOKEN_EXCHANGE_PATH
|
||||
|
||||
_JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
|
||||
_DEFAULT_API_BASE: Final = "https://api.anthropic.com"
|
||||
_INLINE_ENV_VAR: Final = "ANTHROPIC_IDENTITY_TOKEN"
|
||||
_DISABLE_WIF_PARAM: Final = "anthropic_disable_workload_identity_federation"
|
||||
_ACCEPTED_REF_PREFIX: Final = "oidc/"
|
||||
_CHAT_BASE_SUFFIXES: Final = ("/v1/messages", "/v1")
|
||||
# Hosts a federated exchange may talk to. api_base decides where the workload's assertion is sent
|
||||
# AND where the minted org-scoped token is presented, so anyone able to write api_base on a
|
||||
# federated deployment could otherwise redirect both. Gating each write path does not terminate:
|
||||
# a deployment, a referenced credential and a future endpoint all reach the same value. This is the
|
||||
# one place a federated exchange is built, so the trust decision is enforced here instead, and the
|
||||
# allowlist is server-owned -- read from the environment, never from a model or credential API.
|
||||
_TRUSTED_EXCHANGE_HOSTS_ENV: Final = "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS"
|
||||
_DEFAULT_TRUSTED_EXCHANGE_HOST: Final = "api.anthropic.com"
|
||||
_REJECTED_REF_PREFIX: Final = "oidc/env_path/"
|
||||
_IDENTITY_SOURCE_PARAM: Final = "anthropic_identity_source"
|
||||
_IDENTITY_SOURCE_ENV: Final = "ANTHROPIC_IDENTITY_SOURCE"
|
||||
_IDENTITY_TOKEN_FILE_PARAM: Final = "anthropic_identity_token_file"
|
||||
_IDENTITY_TOKEN_PARAM: Final = "anthropic_identity_token"
|
||||
|
||||
# litellm_params key -> InternalIssuerSource/KeycloakSource field name. Every key here must
|
||||
# also be listed in ANTHROPIC_WIF_KWARGS_KEYS (get_litellm_params.py), which is what makes it
|
||||
# request-banned and cleared on a client-redirected api_base -- see types/utils.py's
|
||||
# anthropic_wif_litellm_params, derived from that same set.
|
||||
_INTERNAL_ISSUER_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"anthropic_issuer_url": "issuer_url",
|
||||
"anthropic_issuer_subject": "subject",
|
||||
"anthropic_issuer_audience": "audience",
|
||||
"anthropic_issuer_ttl_seconds": "ttl_seconds",
|
||||
"anthropic_issuer_signing_key_ref": "signing_key_ref",
|
||||
}
|
||||
)
|
||||
_KEYCLOAK_FIELD_MAP: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"anthropic_keycloak_token_url": "token_url",
|
||||
"anthropic_keycloak_client_id": "client_id",
|
||||
"anthropic_keycloak_auth_method": "auth_method",
|
||||
"anthropic_keycloak_client_secret_ref": "client_secret_ref",
|
||||
"anthropic_keycloak_scope": "scope",
|
||||
}
|
||||
)
|
||||
_DENIAL_HINT: Final = (
|
||||
" Anthropic answers every denied exchange with the same 401; the reason (for example"
|
||||
" workspace_id_required or jti_reused) is only shown in the Claude Console under"
|
||||
" Settings > Workload identity, in the rule's authentication history."
|
||||
)
|
||||
_WORKSPACE_HINT: Final = (
|
||||
" If the federation rule is enabled in more than one workspace, set anthropic_workspace_id"
|
||||
" (or ANTHROPIC_WORKSPACE_ID) to the wrkspc_ id of the workspace to mint tokens for, or to 'default'."
|
||||
)
|
||||
_SERVICE_ACCOUNT_HINT: Final = (
|
||||
" Anthropic's reference lists service_account_id as required: set anthropic_service_account_id"
|
||||
" (or ANTHROPIC_SERVICE_ACCOUNT_ID) to the svac_ id the federation rule targets."
|
||||
)
|
||||
_MISSING_IDS_HINT: Final = (
|
||||
" Copy them from the federation rule's detail page under Settings > Workload identity in the"
|
||||
" Claude Console, or set ANTHROPIC_FEDERATION_RULE_ID and ANTHROPIC_ORGANIZATION_ID."
|
||||
)
|
||||
_ALLOWLIST_HINT: Final = (
|
||||
" Identity token files must sit under an allowed credential directory"
|
||||
" (/var/run/secrets or /run/secrets by default);"
|
||||
" set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist."
|
||||
)
|
||||
_EMPTY_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
_IdentitySourceVariant = TypeVar("_IdentitySourceVariant", bound="InternalIssuerSource | KeycloakSource")
|
||||
|
||||
|
||||
class AnthropicWifParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
federation_rule_id: str
|
||||
organization_id: str
|
||||
service_account_id: str | None = None
|
||||
workspace_id: str | None = None
|
||||
assertion_ref: str
|
||||
assertion_source: Callable[[], str | None] | None = None
|
||||
|
||||
|
||||
def resolve_anthropic_wif_params(litellm_params: Mapping[str, object] | None) -> AnthropicWifParams | None:
|
||||
if litellm_params is not None and litellm_params.get(_DISABLE_WIF_PARAM) is True:
|
||||
return None
|
||||
federation_rule_id: Final = _config_value(
|
||||
litellm_params, "anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID"
|
||||
)
|
||||
organization_id: Final = _config_value(litellm_params, "anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID")
|
||||
if federation_rule_id is None or organization_id is None:
|
||||
_raise_if_identity_source_configured(litellm_params, federation_rule_id, organization_id)
|
||||
return None
|
||||
identity_source: Final = _resolve_identity_source(litellm_params)
|
||||
if identity_source is None:
|
||||
return None
|
||||
assertion_ref, assertion_source = identity_source
|
||||
return AnthropicWifParams(
|
||||
federation_rule_id=federation_rule_id,
|
||||
organization_id=organization_id,
|
||||
service_account_id=_config_value(
|
||||
litellm_params, "anthropic_service_account_id", "ANTHROPIC_SERVICE_ACCOUNT_ID"
|
||||
),
|
||||
workspace_id=_config_value(litellm_params, "anthropic_workspace_id", "ANTHROPIC_WORKSPACE_ID"),
|
||||
assertion_ref=assertion_ref,
|
||||
assertion_source=assertion_source,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_identity_source(
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
) -> tuple[str, Callable[[], str] | None] | None:
|
||||
"""Dispatches on ``anthropic_identity_source``. Absent (the default) keeps today's
|
||||
token_file/env resolution byte-identical, with no ``assertion_source`` closure -- the engine
|
||||
falls back to its own reader exactly as it does today. A recognized kind builds the matching
|
||||
frozen config, hashes it into the ``oidc/<kind>/<hash>`` cache-key ref (``identity_source_ref``),
|
||||
and closes the source's fetch/mint function over it. An unset-but-invalid config (unknown
|
||||
kind, a missing required field, or a field from the other variant) fails closed here rather
|
||||
than silently falling back to token_file. A deployment whose params carry a legacy token or
|
||||
token_file ref stays on legacy resolution even when ``ANTHROPIC_IDENTITY_SOURCE`` names a
|
||||
fleet-wide kind: the env kind only governs deployments that set no identity params of their own."""
|
||||
source_kind: Final = _resolve_source_kind(litellm_params)
|
||||
if source_kind is None:
|
||||
legacy_ref: Final = _resolve_assertion_ref(litellm_params)
|
||||
return (legacy_ref, None) if legacy_ref is not None else None
|
||||
params: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{key: value for key, value in (litellm_params or _EMPTY_PARAMS).items() if value is not None}
|
||||
)
|
||||
match source_kind:
|
||||
case AnthropicIdentitySourceKind.internal_issuer.value:
|
||||
_reject_foreign_variant_fields(params, foreign_field_map=_KEYCLOAK_FIELD_MAP, chosen_kind=source_kind)
|
||||
issuer_config: Final = _build_variant(InternalIssuerSource, params, _INTERNAL_ISSUER_FIELD_MAP)
|
||||
return identity_source_ref(issuer_config), internal_issuer_assertion_source(issuer_config)
|
||||
case AnthropicIdentitySourceKind.keycloak.value:
|
||||
_reject_foreign_variant_fields(
|
||||
params, foreign_field_map=_INTERNAL_ISSUER_FIELD_MAP, chosen_kind=source_kind
|
||||
)
|
||||
keycloak_config: Final = _build_variant(KeycloakSource, params, _KEYCLOAK_FIELD_MAP)
|
||||
return identity_source_ref(keycloak_config), keycloak_assertion_source(keycloak_config)
|
||||
case _:
|
||||
_raise_unknown_source_kind(source_kind)
|
||||
|
||||
|
||||
def _raise_unknown_source_kind(source_kind: str) -> NoReturn:
|
||||
raise litellm.AuthenticationError(
|
||||
message=(
|
||||
f"{_IDENTITY_SOURCE_PARAM} must be one of "
|
||||
f"{', '.join(kind.value for kind in AnthropicIdentitySourceKind)}; got {source_kind!r}."
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model="",
|
||||
)
|
||||
|
||||
|
||||
def _raise_if_identity_source_configured(
|
||||
litellm_params: Mapping[str, object] | None, federation_rule_id: str | None, organization_id: str | None
|
||||
) -> None:
|
||||
"""A configured identity source is an explicit request to federate, so a missing rule or
|
||||
organization id fails closed with the ids named, rather than silently skipping federation
|
||||
and surfacing later as a missing API key."""
|
||||
source_kind: Final = _resolve_source_kind(litellm_params)
|
||||
if source_kind is None:
|
||||
return
|
||||
if source_kind not in {kind.value for kind in AnthropicIdentitySourceKind}:
|
||||
_raise_unknown_source_kind(source_kind)
|
||||
missing: Final = tuple(
|
||||
param
|
||||
for param, value in (
|
||||
("anthropic_federation_rule_id", federation_rule_id),
|
||||
("anthropic_organization_id", organization_id),
|
||||
)
|
||||
if value is None
|
||||
)
|
||||
raise litellm.AuthenticationError(
|
||||
message=(
|
||||
f"{_IDENTITY_SOURCE_PARAM} is {source_kind!r}, but {' and '.join(missing)} "
|
||||
f"{'is' if len(missing) == 1 else 'are'} not set.{_MISSING_IDS_HINT}"
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model="",
|
||||
)
|
||||
|
||||
|
||||
def _resolve_source_kind(litellm_params: Mapping[str, object] | None) -> str | None:
|
||||
param_kind: Final = _param_str(litellm_params, _IDENTITY_SOURCE_PARAM)
|
||||
if param_kind is not None:
|
||||
return param_kind
|
||||
has_param_legacy_ref: Final = any(
|
||||
_param_str(litellm_params, key) is not None for key in (_IDENTITY_TOKEN_FILE_PARAM, _IDENTITY_TOKEN_PARAM)
|
||||
)
|
||||
return None if has_param_legacy_ref else _env_str(_IDENTITY_SOURCE_ENV)
|
||||
|
||||
|
||||
def _reject_foreign_variant_fields(
|
||||
litellm_params: Mapping[str, object], foreign_field_map: Mapping[str, str], chosen_kind: str
|
||||
) -> None:
|
||||
foreign_keys_present: Final = tuple(param for param in foreign_field_map if param in litellm_params)
|
||||
if foreign_keys_present:
|
||||
raise litellm.AuthenticationError(
|
||||
message=(
|
||||
f"{_IDENTITY_SOURCE_PARAM} is {chosen_kind!r}, but {', '.join(sorted(foreign_keys_present))} "
|
||||
"belongs to a different identity source and cannot be set alongside it."
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model="",
|
||||
)
|
||||
|
||||
|
||||
def _build_variant(
|
||||
model: type[_IdentitySourceVariant],
|
||||
litellm_params: Mapping[str, object],
|
||||
field_map: Mapping[str, str],
|
||||
) -> _IdentitySourceVariant:
|
||||
fields: Final = MappingProxyType(
|
||||
{field_map[key]: value for key, value in litellm_params.items() if key in field_map}
|
||||
)
|
||||
try:
|
||||
return model.model_validate(fields)
|
||||
except ValidationError as e:
|
||||
# hide_input_in_errors=True on both variant models keeps a secret pasted into the
|
||||
# wrong field (e.g. a client_secret typed as signing_key_ref) out of str(e).
|
||||
raise litellm.AuthenticationError(
|
||||
message=f"Invalid {_IDENTITY_SOURCE_PARAM} configuration: {e}",
|
||||
llm_provider="anthropic",
|
||||
model="",
|
||||
) from e
|
||||
|
||||
|
||||
def build_anthropic_wif_spec(params: AnthropicWifParams, api_base: str) -> TokenExchangeSpec:
|
||||
return TokenExchangeSpec(
|
||||
token_url=api_base.rstrip("/") + ANTHROPIC_TOKEN_EXCHANGE_PATH,
|
||||
assertion_ref=params.assertion_ref,
|
||||
assertion_field="assertion",
|
||||
static_body=MappingProxyType(
|
||||
{
|
||||
name: value
|
||||
for name, value in (
|
||||
("grant_type", _JWT_BEARER_GRANT_TYPE),
|
||||
("federation_rule_id", params.federation_rule_id),
|
||||
("organization_id", params.organization_id),
|
||||
("service_account_id", params.service_account_id),
|
||||
("workspace_id", params.workspace_id),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
),
|
||||
body_encoding="json",
|
||||
request_headers=MappingProxyType({}),
|
||||
assertion_source=params.assertion_source,
|
||||
cache_key_identity=(
|
||||
params.federation_rule_id,
|
||||
params.organization_id,
|
||||
params.service_account_id or "",
|
||||
params.workspace_id or "",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def get_anthropic_wif_token(
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
api_base: str | None,
|
||||
model: str,
|
||||
engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine,
|
||||
) -> str | None:
|
||||
params: Final = resolve_anthropic_wif_params(litellm_params)
|
||||
if params is None:
|
||||
return None
|
||||
exchange_base: Final = resolve_anthropic_base(api_base)
|
||||
_raise_if_exchange_host_untrusted(exchange_base, model)
|
||||
result: Final = engine.get_token(build_anthropic_wif_spec(params, exchange_base))
|
||||
return _token_from_result(result, model, params)
|
||||
|
||||
|
||||
async def aget_anthropic_wif_token(
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
api_base: str | None,
|
||||
model: str,
|
||||
engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine,
|
||||
) -> str | None:
|
||||
params: Final = resolve_anthropic_wif_params(litellm_params)
|
||||
if params is None:
|
||||
return None
|
||||
exchange_base: Final = resolve_anthropic_base(api_base)
|
||||
_raise_if_exchange_host_untrusted(exchange_base, model)
|
||||
result: Final = await engine.aget_token(build_anthropic_wif_spec(params, exchange_base))
|
||||
return _token_from_result(result, model, params)
|
||||
|
||||
|
||||
def _token_from_result(result: ExchangeResult, model: str, params: AnthropicWifParams) -> str:
|
||||
match result:
|
||||
case MintedToken():
|
||||
return result.access_token.get_secret_value()
|
||||
case _:
|
||||
_raise_anthropic_wif_error(
|
||||
result,
|
||||
model=model,
|
||||
workspace_id_set=params.workspace_id is not None,
|
||||
service_account_id_set=params.service_account_id is not None,
|
||||
)
|
||||
|
||||
|
||||
def resolve_anthropic_base(api_base: str | None) -> str:
|
||||
"""The base every Anthropic tier derives its URLs from: the deployment api_base when set,
|
||||
else ``ANTHROPIC_API_BASE`` / ``ANTHROPIC_BASE_URL``, else Anthropic's host, with trailing
|
||||
slashes and chat-appended ``/v1/messages`` suffixes stripped, so the token URL, the cache key
|
||||
and the count-tokens URL all agree for the same deployment."""
|
||||
return anthropic_base_without_chat_suffix(api_base or _resolve_default_api_base())
|
||||
|
||||
|
||||
def _trusted_exchange_hosts() -> frozenset[str]:
|
||||
"""Hostnames a federated exchange may reach: Anthropic's own, plus whatever the operator put in
|
||||
the environment. Comma separated, case folded, entries given as a URL reduced to their host."""
|
||||
configured: Final = os.getenv(_TRUSTED_EXCHANGE_HOSTS_ENV) or ""
|
||||
extra: Final = (entry.strip() for entry in configured.split(",") if entry.strip())
|
||||
return frozenset(
|
||||
chain(
|
||||
(_DEFAULT_TRUSTED_EXCHANGE_HOST,),
|
||||
((urlsplit(entry).hostname or entry.split("/")[0]).lower() for entry in extra),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _raise_if_exchange_host_untrusted(exchange_base: str, model: str) -> None:
|
||||
"""The federated exchange refuses any host the operator has not vouched for, whatever wrote the
|
||||
deployment's api_base. Exact hostname match, never a substring: ``api.anthropic.com.evil.test``
|
||||
contains the real host and must not pass."""
|
||||
host: Final = (urlsplit(exchange_base).hostname or "").lower()
|
||||
if host and host in _trusted_exchange_hosts():
|
||||
return
|
||||
raise litellm.AuthenticationError(
|
||||
message=(
|
||||
f"Anthropic workload identity federation refused to use host {host or exchange_base!r}. "
|
||||
f"A federated exchange sends the workload's identity token to this host and presents the "
|
||||
f"minted token to it, so only {_DEFAULT_TRUSTED_EXCHANGE_HOST} is trusted by default. To "
|
||||
f"use a private Anthropic-compatible gateway, add its hostname to the "
|
||||
f"{_TRUSTED_EXCHANGE_HOSTS_ENV} environment variable (comma separated); that is a "
|
||||
f"decision to trust it with org-scoped credentials, so it is deliberately server-owned "
|
||||
f"and cannot be set through the model or credential APIs."
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_default_api_base() -> str:
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
|
||||
return AnthropicModelInfo.get_api_base(None) or _DEFAULT_API_BASE
|
||||
|
||||
|
||||
def anthropic_base_without_chat_suffix(base: str) -> str:
|
||||
"""A deployment base with its chat-surface suffix removed, so the token URL and model
|
||||
discovery both derive from the same value whatever form the operator configured."""
|
||||
parts: Final = urlsplit(base)
|
||||
if not parts.scheme or not parts.netloc:
|
||||
return base.rstrip("/")
|
||||
return urlunsplit((parts.scheme, parts.netloc, _strip_path_suffixes(parts.path), "", ""))
|
||||
|
||||
|
||||
def _strip_path_suffixes(path: str) -> str:
|
||||
"""Drop the chat-surface suffixes a deployment base may carry, so every tier derives the same
|
||||
token URL. Each pass removes at most one suffix, so the loop is bounded by the segment count."""
|
||||
trimmed = path.rstrip("/") # rebind-ok: fixed-point strip, one suffix per pass
|
||||
while True:
|
||||
shortened = next( # rebind-ok: one suffix removed per iteration
|
||||
(trimmed.removesuffix(suffix) for suffix in _CHAT_BASE_SUFFIXES if trimmed.endswith(suffix)),
|
||||
trimmed,
|
||||
)
|
||||
if shortened == trimmed:
|
||||
return trimmed
|
||||
# Re-strip: a doubled suffix leaves a trailing slash that would stop the next match.
|
||||
trimmed = shortened.rstrip("/")
|
||||
|
||||
|
||||
def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None:
|
||||
return _param_str(litellm_params, param_key) or _env_str(env_name)
|
||||
|
||||
|
||||
def _param_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
|
||||
if litellm_params is None:
|
||||
return None
|
||||
value: Final = litellm_params.get(key)
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _env_str(name: str) -> str | None:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
value: Final = get_secret_str(name)
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _resolve_assertion_ref(litellm_params: Mapping[str, object] | None) -> str | None:
|
||||
file_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_FILE_PARAM)
|
||||
if file_param is not None:
|
||||
return f"oidc/file/{file_param}"
|
||||
inline_param: Final = _param_str(litellm_params, _IDENTITY_TOKEN_PARAM)
|
||||
if inline_param is not None:
|
||||
return _validated_inline_ref(inline_param)
|
||||
file_env: Final = _env_str("ANTHROPIC_IDENTITY_TOKEN_FILE")
|
||||
if file_env is not None:
|
||||
return f"oidc/file/{file_env}"
|
||||
if _env_str(_INLINE_ENV_VAR) is not None:
|
||||
return f"oidc/env/{_INLINE_ENV_VAR}"
|
||||
return None
|
||||
|
||||
|
||||
def _validated_inline_ref(value: str) -> str:
|
||||
if value.startswith(_ACCEPTED_REF_PREFIX) and not value.startswith(_REJECTED_REF_PREFIX):
|
||||
return value
|
||||
raise litellm.AuthenticationError(
|
||||
message=(
|
||||
"anthropic_identity_token must be an oidc/ secret reference such as oidc/env/VAR_NAME,"
|
||||
" oidc/file//absolute/path, oidc/github/<audience>, or oidc/google/<audience>."
|
||||
" Raw identity tokens and oidc/env_path/ references are not accepted;"
|
||||
" to pass a token directly, export it and reference it as oidc/env/VAR_NAME"
|
||||
),
|
||||
llm_provider="anthropic",
|
||||
model="",
|
||||
)
|
||||
|
||||
|
||||
def _raise_anthropic_wif_error(
|
||||
error: ExchangeError, model: str, workspace_id_set: bool, service_account_id_set: bool
|
||||
) -> NoReturn:
|
||||
detail: Final = _error_detail(
|
||||
error, workspace_id_set=workspace_id_set, service_account_id_set=service_account_id_set
|
||||
)
|
||||
raise litellm.AuthenticationError(
|
||||
message=f"Anthropic workload identity federation failed. {detail}",
|
||||
llm_provider="anthropic",
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
||||
def _denial_hints(workspace_id_set: bool, service_account_id_set: bool) -> str:
|
||||
return "".join(
|
||||
(
|
||||
_DENIAL_HINT,
|
||||
"" if workspace_id_set else _WORKSPACE_HINT,
|
||||
"" if service_account_id_set else _SERVICE_ACCOUNT_HINT,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _error_detail(error: ExchangeError, workspace_id_set: bool, service_account_id_set: bool) -> str:
|
||||
match error:
|
||||
case AssertionSourceError() if error.kind == "disallowed_path":
|
||||
return f"Could not read the OIDC identity token from {error.source_ref}.{_ALLOWLIST_HINT}"
|
||||
case AssertionSourceError():
|
||||
base: Final = f"Could not obtain the OIDC identity token ({error.kind}) from {error.source_ref}."
|
||||
return f"{base} {error.detail}" if error.detail else base
|
||||
case InsecureTokenUrl():
|
||||
return f"The token endpoint must use https; refusing to send the identity token to host {error.host!r}."
|
||||
case TokenEndpointError() if error.status_code == 401:
|
||||
hints: Final = _denial_hints(workspace_id_set, service_account_id_set)
|
||||
return f"The token endpoint returned HTTP 401: {error.redacted_body}{hints}"
|
||||
case TokenEndpointError():
|
||||
return f"The token endpoint returned HTTP {error.status_code}: {error.redacted_body}"
|
||||
case TokenTransportError():
|
||||
return f"Could not reach the token endpoint: {error.detail}."
|
||||
case MalformedTokenResponse():
|
||||
return f"The token endpoint returned an unusable response: {error.detail}."
|
||||
case _:
|
||||
assert_never(error)
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
|
|
@ -218,6 +219,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
|
|||
aembedding=None,
|
||||
max_retries: int | None = None,
|
||||
shared_session=None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
"""
|
||||
- Separate image url from text
|
||||
|
|
|
|||
|
|
@ -41,6 +41,29 @@ class BaseAnthropicMessagesConfig(ABC):
|
|||
"""
|
||||
return headers, api_base
|
||||
|
||||
async def avalidate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
model: str,
|
||||
messages: list[Any], # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
optional_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
litellm_params: dict, # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]: # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
|
||||
"""Async counterpart used by the async handler. The default delegates to the
|
||||
sync implementation; providers whose sync path can block the event loop
|
||||
(e.g. a WIF token exchange) override this."""
|
||||
return self.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
99
litellm/llms/base_llm/auth/__init__.py
Normal file
99
litellm/llms/base_llm/auth/__init__.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
from litellm.llms.base_llm.auth.client_credentials import (
|
||||
SecretReader,
|
||||
fetch_keycloak_assertion,
|
||||
keycloak_assertion_source,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.identity_source import (
|
||||
AnthropicIdentitySourceConfig,
|
||||
AnthropicIdentitySourceKind,
|
||||
InternalIssuerSource,
|
||||
KeycloakSource,
|
||||
identity_source_config_adapter,
|
||||
identity_source_ref,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.internal_issuer import (
|
||||
SigningKeyReader,
|
||||
internal_issuer_assertion_source,
|
||||
internal_issuer_jwks_document,
|
||||
mint_internal_issuer_assertion,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.jwt_signing import (
|
||||
ALG,
|
||||
build_jwk,
|
||||
build_jwks,
|
||||
jwks_document_json,
|
||||
load_es256_private_key,
|
||||
rfc7638_thumbprint,
|
||||
sign_es256_jwt,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.token_exchange import (
|
||||
ADVISORY_REFRESH_BACKOFF_SECONDS,
|
||||
ADVISORY_REFRESH_SECONDS,
|
||||
MANDATORY_REFRESH_SECONDS,
|
||||
MAX_ASSERTION_BYTES,
|
||||
MAX_RESPONSE_BYTES,
|
||||
JwtBearerTokenExchangeEngine,
|
||||
default_token_exchange_engine,
|
||||
redact_oauth_error_body,
|
||||
validate_token_endpoint_url,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.types import (
|
||||
AssertionReader,
|
||||
AssertionSource,
|
||||
AssertionSourceError,
|
||||
BodyEncoding,
|
||||
ExchangeError,
|
||||
ExchangeResult,
|
||||
InsecureTokenUrl,
|
||||
MalformedTokenResponse,
|
||||
MintedToken,
|
||||
SyncTokenPoster,
|
||||
TokenEndpointError,
|
||||
TokenExchangeSpec,
|
||||
TokenTransportError,
|
||||
)
|
||||
|
||||
__all__ = (
|
||||
"ADVISORY_REFRESH_BACKOFF_SECONDS",
|
||||
"ADVISORY_REFRESH_SECONDS",
|
||||
"ALG",
|
||||
"MANDATORY_REFRESH_SECONDS",
|
||||
"MAX_ASSERTION_BYTES",
|
||||
"MAX_RESPONSE_BYTES",
|
||||
"AnthropicIdentitySourceConfig",
|
||||
"AnthropicIdentitySourceKind",
|
||||
"AssertionReader",
|
||||
"AssertionSource",
|
||||
"AssertionSourceError",
|
||||
"BodyEncoding",
|
||||
"ExchangeError",
|
||||
"ExchangeResult",
|
||||
"InsecureTokenUrl",
|
||||
"InternalIssuerSource",
|
||||
"JwtBearerTokenExchangeEngine",
|
||||
"KeycloakSource",
|
||||
"MalformedTokenResponse",
|
||||
"MintedToken",
|
||||
"SecretReader",
|
||||
"SigningKeyReader",
|
||||
"SyncTokenPoster",
|
||||
"TokenEndpointError",
|
||||
"TokenExchangeSpec",
|
||||
"TokenTransportError",
|
||||
"build_jwk",
|
||||
"build_jwks",
|
||||
"default_token_exchange_engine",
|
||||
"fetch_keycloak_assertion",
|
||||
"identity_source_config_adapter",
|
||||
"identity_source_ref",
|
||||
"internal_issuer_assertion_source",
|
||||
"internal_issuer_jwks_document",
|
||||
"jwks_document_json",
|
||||
"keycloak_assertion_source",
|
||||
"load_es256_private_key",
|
||||
"mint_internal_issuer_assertion",
|
||||
"redact_oauth_error_body",
|
||||
"rfc7638_thumbprint",
|
||||
"sign_es256_jwt",
|
||||
"validate_token_endpoint_url",
|
||||
)
|
||||
227
litellm/llms/base_llm/auth/client_credentials.py
Normal file
227
litellm/llms/base_llm/auth/client_credentials.py
Normal file
|
|
@ -0,0 +1,227 @@
|
|||
"""Fetches a fresh RFC 6749 client_credentials assertion for Anthropic's ``keycloak`` identity
|
||||
source: LiteLLM authenticates to Keycloak as its own confidential client and presents the
|
||||
resulting ``access_token`` as the workload assertion (Phase 1 decision 2).
|
||||
|
||||
The client secret is the operator-supplied pointer at ``KeycloakSource.client_secret_ref``,
|
||||
resolved the same way every other WIF secret pointer already is (env, a Credential, or whatever
|
||||
secret manager ``litellm.secret_manager_client`` is globally configured to, Vault included).
|
||||
Every fetch is a fresh HTTP POST; nothing here caches a fetched token, since the outer
|
||||
token-exchange engine already caches the Anthropic token it buys with one -- see decision 2's
|
||||
"no Keycloak-side cache" ruling.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import threading
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
from urllib.parse import quote, quote_plus, urlencode
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, SecretStr, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.llms.base_llm.auth.identity_source import KeycloakSource, ref_for_error_message
|
||||
from litellm.llms.base_llm.auth.token_exchange import (
|
||||
MAX_RESPONSE_BYTES,
|
||||
endpoint_url_for_error_message,
|
||||
redact_oauth_error_body,
|
||||
require_posted_response,
|
||||
validate_token_endpoint_url,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.types import InsecureTokenUrl, SyncTokenPoster
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
SecretReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list
|
||||
|
||||
_GRANT_TYPE: Final = "client_credentials"
|
||||
_TIMEOUT_SECONDS: Final = 30.0
|
||||
_FORM_CONTENT_TYPE: Final = "application/x-www-form-urlencoded"
|
||||
|
||||
|
||||
class _ClientCredentialsResponse(BaseModel):
|
||||
access_token: str
|
||||
|
||||
|
||||
def _default_secret_reader(ref: str) -> str | None:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return get_secret_str(ref)
|
||||
|
||||
|
||||
def _new_keycloak_handler() -> "HTTPHandler":
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
handler: Final = HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0))
|
||||
handler.client.follow_redirects = False
|
||||
return handler
|
||||
|
||||
|
||||
class _HttpxSyncKeycloakPoster:
|
||||
"""Dedicated HTTPHandler for the Keycloak token POST: no ``logging_obj`` (so litellm's
|
||||
request/response logging never sees the client secret or the fetched token), redirects
|
||||
disabled. A separate instance from the outer engine's own poster, since this is a genuinely
|
||||
new HTTP call site whose no-logging guarantee must be built here, not assumed inherited."""
|
||||
|
||||
def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_keycloak_handler) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._handler_factory: Final = handler_factory
|
||||
self._handler: HTTPHandler | None = None
|
||||
|
||||
def _handler_instance(self) -> "HTTPHandler":
|
||||
with self._lock:
|
||||
if self._handler is None:
|
||||
self._handler = self._handler_factory()
|
||||
return self._handler
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
|
||||
try:
|
||||
response: Final[httpx.Response | None] = self._handler_instance().post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post is legacy-untyped; the result is validated below
|
||||
url,
|
||||
content=content,
|
||||
headers=dict(headers), # mutable-ok: HTTPHandler.post requires a concrete dict
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
return e.response
|
||||
return require_posted_response(response, "keycloak token endpoint")
|
||||
|
||||
|
||||
_DEFAULT_POSTER: Final[SyncTokenPoster] = _HttpxSyncKeycloakPoster()
|
||||
|
||||
|
||||
def _form_encode(value: str) -> str:
|
||||
"""RFC 6749 Appendix B before RFC 6749 2.3.1's base64: application/x-www-form-urlencoded
|
||||
with spaces as ``%20`` rather than ``+``, else a reserved character (":", "+", "%", " ") in
|
||||
the id or secret corrupts the credential the far side decodes back out of Basic auth."""
|
||||
return quote(value, safe="")
|
||||
|
||||
|
||||
def _basic_auth_header(client_id: str, client_secret: str) -> str:
|
||||
encoded_pair: Final = f"{_form_encode(client_id)}:{_form_encode(client_secret)}"
|
||||
return "Basic " + base64.b64encode(encoded_pair.encode()).decode("ascii")
|
||||
|
||||
|
||||
def _prepared_request(config: KeycloakSource, client_secret: str) -> tuple[bytes, Mapping[str, str]]:
|
||||
scope_field: Final[Mapping[str, str]] = (
|
||||
MappingProxyType({"scope": config.scope}) if config.scope else MappingProxyType({})
|
||||
)
|
||||
match config.auth_method:
|
||||
case "client_secret_basic":
|
||||
return (
|
||||
urlencode(MappingProxyType({"grant_type": _GRANT_TYPE, **scope_field})).encode(),
|
||||
MappingProxyType(
|
||||
{
|
||||
"content-type": _FORM_CONTENT_TYPE,
|
||||
"authorization": _basic_auth_header(config.client_id, client_secret),
|
||||
}
|
||||
),
|
||||
)
|
||||
case "client_secret_post":
|
||||
return (
|
||||
urlencode(
|
||||
MappingProxyType(
|
||||
{
|
||||
"grant_type": _GRANT_TYPE,
|
||||
"client_id": config.client_id,
|
||||
"client_secret": client_secret,
|
||||
**scope_field,
|
||||
}
|
||||
)
|
||||
).encode(),
|
||||
MappingProxyType({"content-type": _FORM_CONTENT_TYPE}),
|
||||
)
|
||||
case _:
|
||||
assert_never(config.auth_method)
|
||||
|
||||
|
||||
def _resolve_client_secret(config: KeycloakSource, secret_reader: SecretReader) -> str:
|
||||
secret: Final = secret_reader(config.client_secret_ref)
|
||||
if not secret:
|
||||
raise ValueError(f"keycloak client secret {ref_for_error_message(config.client_secret_ref)} could not be read")
|
||||
return secret
|
||||
|
||||
|
||||
def _wire_forms_of_secret(config: KeycloakSource, client_secret: str) -> tuple[SecretStr, ...]:
|
||||
"""Every shape the secret leaves this process in, so an echo of any of them is caught.
|
||||
|
||||
Neither grant sends the secret verbatim. client_secret_basic base64s ``id:secret``, which
|
||||
decodes straight back to it, and client_secret_post percent-escapes it. An endpoint echoing
|
||||
either shape hands over reversible material a raw comparison would miss.
|
||||
"""
|
||||
raw: Final = SecretStr(client_secret)
|
||||
match config.auth_method:
|
||||
case "client_secret_basic":
|
||||
encoded_pair: Final = f"{_form_encode(config.client_id)}:{_form_encode(client_secret)}"
|
||||
return (raw, SecretStr(base64.b64encode(encoded_pair.encode()).decode("ascii")))
|
||||
case "client_secret_post":
|
||||
# urlencode escapes reserved characters and writes a space as "+", so a secret
|
||||
# containing either leaves in a shape the raw comparison would not recognise coming
|
||||
# back. quote_plus is what urlencode itself applies.
|
||||
return (raw, SecretStr(quote_plus(client_secret)))
|
||||
case _:
|
||||
assert_never(config.auth_method)
|
||||
|
||||
|
||||
def _endpoint_error_message(config: KeycloakSource, response: httpx.Response, client_secret: str) -> str:
|
||||
endpoint_error: Final = redact_oauth_error_body(
|
||||
response.status_code, response.text, _wire_forms_of_secret(config, client_secret)
|
||||
)
|
||||
return (
|
||||
f"keycloak token endpoint {endpoint_url_for_error_message(config.token_url)} "
|
||||
f"returned HTTP {endpoint_error.status_code}: {endpoint_error.redacted_body}"
|
||||
)
|
||||
|
||||
|
||||
def _parse_success_body(response: httpx.Response) -> str:
|
||||
if len(response.content) > MAX_RESPONSE_BYTES:
|
||||
raise ValueError("keycloak token response exceeded the size cap")
|
||||
try:
|
||||
parsed: Final = _ClientCredentialsResponse.model_validate_json(response.content)
|
||||
except ValidationError as e:
|
||||
raise ValueError("keycloak token response failed schema validation") from e
|
||||
token: Final = parsed.access_token.strip()
|
||||
if not token:
|
||||
raise ValueError("keycloak token response carried an empty access_token")
|
||||
return token
|
||||
|
||||
|
||||
def fetch_keycloak_assertion(
|
||||
config: KeycloakSource,
|
||||
*,
|
||||
poster: SyncTokenPoster = _DEFAULT_POSTER,
|
||||
secret_reader: SecretReader = _default_secret_reader,
|
||||
) -> str:
|
||||
"""POSTs one fresh client_credentials grant and returns the resulting ``access_token`` as the
|
||||
workload assertion; the caller must not cache the result -- see the module docstring."""
|
||||
match validate_token_endpoint_url(config.token_url):
|
||||
case InsecureTokenUrl(host=host):
|
||||
raise ValueError(f"keycloak token_url must use https; refusing to send the client secret to host {host!r}")
|
||||
case _:
|
||||
pass
|
||||
client_secret: Final = _resolve_client_secret(config, secret_reader)
|
||||
content, headers = _prepared_request(config, client_secret)
|
||||
try:
|
||||
response: Final = poster.post(config.token_url, content=content, headers=headers, timeout=_TIMEOUT_SECONDS)
|
||||
except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; every failure becomes a ValueError
|
||||
raise ValueError(
|
||||
f"could not reach the keycloak token endpoint {endpoint_url_for_error_message(config.token_url)}: "
|
||||
f"{type(e).__name__}"
|
||||
) from e
|
||||
if not 200 <= response.status_code < 300:
|
||||
raise ValueError(_endpoint_error_message(config, response, client_secret))
|
||||
return _parse_success_body(response)
|
||||
|
||||
|
||||
def keycloak_assertion_source(
|
||||
config: KeycloakSource,
|
||||
*,
|
||||
poster: SyncTokenPoster = _DEFAULT_POSTER,
|
||||
secret_reader: SecretReader = _default_secret_reader,
|
||||
) -> Callable[[], str]:
|
||||
"""A zero-arg closure that fetches fresh on every call: the shape an ``oidc/keycloak/...``
|
||||
ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7)
|
||||
-- the caller parses the config and closes this function over it, with no registry involved."""
|
||||
return lambda: fetch_keycloak_assertion(config, poster=poster, secret_reader=secret_reader)
|
||||
76
litellm/llms/base_llm/auth/identity_source.py
Normal file
76
litellm/llms/base_llm/auth/identity_source.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
"""Tagged-union identity-source configs for Anthropic workload identity federation, beyond the
|
||||
existing token_file/env resolver in ``litellm/llms/anthropic/wif.py``.
|
||||
|
||||
Each variant only ever carries secret *pointer names* (``signing_key_ref``, ``client_secret_ref``),
|
||||
never a resolved secret value, so ``identity_source_ref`` can safely hash a variant into the short,
|
||||
content-derived ``oidc/<kind>/<hash>`` string used elsewhere as a get_secret ref, a token-exchange
|
||||
cache-key discriminator, and an operator-facing error pointer.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from enum import Enum
|
||||
from typing import Annotated, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
|
||||
|
||||
_REF_HASH_HEX_LENGTH: Final = 16
|
||||
_MAX_TTL_SECONDS: Final = 3600
|
||||
_DEFAULT_TTL_SECONDS: Final = 300
|
||||
|
||||
|
||||
class AnthropicIdentitySourceKind(str, Enum):
|
||||
internal_issuer = "internal_issuer"
|
||||
keycloak = "keycloak"
|
||||
|
||||
|
||||
class InternalIssuerSource(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", hide_input_in_errors=True)
|
||||
|
||||
kind: Literal[AnthropicIdentitySourceKind.internal_issuer] = AnthropicIdentitySourceKind.internal_issuer
|
||||
issuer_url: str
|
||||
subject: str
|
||||
audience: str | None = None
|
||||
ttl_seconds: Annotated[int, Field(gt=0, le=_MAX_TTL_SECONDS)] = _DEFAULT_TTL_SECONDS
|
||||
signing_key_ref: str
|
||||
|
||||
|
||||
class KeycloakSource(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", hide_input_in_errors=True)
|
||||
|
||||
kind: Literal[AnthropicIdentitySourceKind.keycloak] = AnthropicIdentitySourceKind.keycloak
|
||||
token_url: str
|
||||
client_id: str
|
||||
auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic"
|
||||
client_secret_ref: str
|
||||
scope: str | None = None
|
||||
|
||||
|
||||
AnthropicIdentitySourceConfig: TypeAlias = Annotated[InternalIssuerSource | KeycloakSource, Field(discriminator="kind")]
|
||||
identity_source_config_adapter: Final = TypeAdapter[AnthropicIdentitySourceConfig](AnthropicIdentitySourceConfig)
|
||||
|
||||
|
||||
def identity_source_ref(config: AnthropicIdentitySourceConfig) -> str:
|
||||
"""``oidc/<kind>/<hash>``: a short, secret-free pointer, stable for identical config and rolling
|
||||
whenever any field does, including a ``*_ref`` pointer NAME (never the secret it points to)."""
|
||||
digest: Final = hashlib.sha256(config.model_dump_json().encode()).hexdigest()[:_REF_HASH_HEX_LENGTH]
|
||||
return f"oidc/{config.kind.value}/{digest}"
|
||||
|
||||
|
||||
_POINTER_REF_PREFIXES: Final = (
|
||||
"oidc/",
|
||||
"os.environ/",
|
||||
"hashicorp_vault/",
|
||||
"aws_secret_manager/",
|
||||
"google_secret_manager/",
|
||||
)
|
||||
|
||||
|
||||
def ref_for_error_message(ref: str) -> str:
|
||||
"""A ``*_ref`` rendered for an operator-facing error.
|
||||
|
||||
Naming the pointer is deliberate: it is what tells an operator which setting failed to
|
||||
resolve. But these fields only ever fail to resolve when what was written is not a pointer,
|
||||
and an operator who pasted the secret itself has made the field's value the secret. So the
|
||||
value is echoed only when it is recognizably a pointer, and withheld otherwise.
|
||||
"""
|
||||
return ref if ref.startswith(_POINTER_REF_PREFIXES) else "<withheld: not a secret reference>"
|
||||
86
litellm/llms/base_llm/auth/internal_issuer.py
Normal file
86
litellm/llms/base_llm/auth/internal_issuer.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
"""Mints a self-issued workload assertion for Anthropic's ``internal_issuer`` identity source:
|
||||
LiteLLM signs its own short-lived ES256 JWT instead of reading one from a mounted OIDC file.
|
||||
|
||||
Signing custody is the operator-supplied PEM at ``InternalIssuerSource.signing_key_ref``,
|
||||
resolved the same way every other WIF secret pointer already is (env, a Credential, or
|
||||
whatever secret manager ``litellm.secret_manager_client`` is globally configured to, Vault
|
||||
included) -- see Phase 1 decision 1. Every mint is fresh; nothing here caches a minted JWT,
|
||||
since the outer token-exchange engine already caches the Anthropic token it buys with one.
|
||||
"""
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource, ref_for_error_message
|
||||
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json, sign_es256_jwt
|
||||
|
||||
SigningKeyReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list
|
||||
|
||||
|
||||
def _default_signing_key_reader(ref: str) -> str | None:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return get_secret_str(ref)
|
||||
|
||||
|
||||
def _claims(config: InternalIssuerSource, issued_at: int) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("sub", config.subject),
|
||||
("iss", config.issuer_url),
|
||||
("aud", config.audience),
|
||||
("iat", issued_at),
|
||||
("exp", issued_at + config.ttl_seconds),
|
||||
("jti", str(uuid.uuid4())),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _resolve_signing_key(config: InternalIssuerSource, key_reader: SigningKeyReader) -> str:
|
||||
pem: Final = key_reader(config.signing_key_ref)
|
||||
if not pem:
|
||||
raise ValueError(
|
||||
f"internal_issuer signing key {ref_for_error_message(config.signing_key_ref)} could not be read"
|
||||
)
|
||||
return pem
|
||||
|
||||
|
||||
def mint_internal_issuer_assertion(
|
||||
config: InternalIssuerSource,
|
||||
*,
|
||||
key_reader: SigningKeyReader = _default_signing_key_reader,
|
||||
clock: Callable[[], float] = time.time,
|
||||
) -> str:
|
||||
"""Signs one fresh, short-lived assertion; the caller must not cache the result, since a
|
||||
cached copy would defeat the point of re-minting on every exchange."""
|
||||
pem: Final = _resolve_signing_key(config, key_reader)
|
||||
return sign_es256_jwt(pem, _claims(config, issued_at=int(clock())))
|
||||
|
||||
|
||||
def internal_issuer_assertion_source(
|
||||
config: InternalIssuerSource,
|
||||
*,
|
||||
key_reader: SigningKeyReader = _default_signing_key_reader,
|
||||
clock: Callable[[], float] = time.time,
|
||||
) -> Callable[[], str]:
|
||||
"""A zero-arg closure that mints fresh on every call: the shape an ``oidc/internal_issuer/...``
|
||||
ref dispatches to once wired into ``TokenExchangeSpec.assertion_source`` (Phase 1 decision 7)
|
||||
-- the caller parses the config and closes this function over it, with no registry involved."""
|
||||
return lambda: mint_internal_issuer_assertion(config, key_reader=key_reader, clock=clock)
|
||||
|
||||
|
||||
def internal_issuer_jwks_document(
|
||||
config: InternalIssuerSource,
|
||||
*,
|
||||
key_reader: SigningKeyReader = _default_signing_key_reader,
|
||||
) -> str:
|
||||
"""The operator-facing JWKS export, resolved from a configured identity source rather than
|
||||
a raw PEM in hand -- the JSON document to register as Anthropic's inline federation issuer."""
|
||||
return jwks_document_json(_resolve_signing_key(config, key_reader))
|
||||
115
litellm/llms/base_llm/auth/jwt_signing.py
Normal file
115
litellm/llms/base_llm/auth/jwt_signing.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
"""ES256 JWT signing primitives for Anthropic workload identity federation's
|
||||
``internal_issuer`` identity source (see ``identity_source.InternalIssuerSource``).
|
||||
|
||||
Pure functions over an already-resolved PEM string: no I/O, no secret-manager awareness, no
|
||||
caching. Given the signing key at, say, $ISSUER_SIGNING_KEY_PEM, an operator publishes the
|
||||
JWKS document Anthropic's inline federation issuer needs with one line:
|
||||
|
||||
python -c "from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json; \\
|
||||
import os; print(jwks_document_json(os.environ['ISSUER_SIGNING_KEY_PEM']))"
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
|
||||
ALG: Final = "ES256"
|
||||
MISSING_SIGNING_DEPENDENCIES_MESSAGE: Final = (
|
||||
"the internal_issuer identity source needs PyJWT and cryptography, which a base litellm install "
|
||||
"does not include: pip install 'litellm[proxy]'"
|
||||
)
|
||||
_JWK_CURVE_NAME: Final = "P-256"
|
||||
_JWK_KEY_TYPE: Final = "EC"
|
||||
_COORDINATE_BYTE_LENGTH: Final = 32 # P-256 field element width, RFC 7518 6.2.1.2/6.2.1.3
|
||||
|
||||
Jwk: TypeAlias = Mapping[str, str]
|
||||
Jwks: TypeAlias = Mapping[str, tuple[Jwk, ...]]
|
||||
|
||||
|
||||
def load_es256_private_key(pem: str) -> "ec.EllipticCurvePrivateKey":
|
||||
"""Parses an unencrypted PEM EC private key. Never echoes the key material in an error."""
|
||||
try:
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
from cryptography.hazmat.primitives.serialization import load_pem_private_key
|
||||
except ImportError as e:
|
||||
raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e
|
||||
try:
|
||||
key: Final = load_pem_private_key(pem.encode(), password=None)
|
||||
except (ValueError, TypeError) as e:
|
||||
raise ValueError("internal_issuer signing key is not a valid unencrypted PEM private key") from e
|
||||
if not isinstance(key, ec.EllipticCurvePrivateKey) or not isinstance(key.curve, ec.SECP256R1):
|
||||
raise ValueError( # noqa: TRY004 # the reader classifies ValueError into a readable config error; TypeError would not
|
||||
"internal_issuer signing key must be an EC P-256 (secp256r1) private key for ES256"
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
def _b64url_coordinate(value: int) -> str:
|
||||
return base64.urlsafe_b64encode(value.to_bytes(_COORDINATE_BYTE_LENGTH, "big")).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
def _jwk_thumbprint_members(public_key: "ec.EllipticCurvePublicKey") -> Jwk:
|
||||
"""RFC 7638 3.2's exact EC member set (crv, kty, x, y) and nothing else: an extra member
|
||||
here would change the thumbprint and desync it from the ``kid`` published in the JWKS."""
|
||||
numbers: Final = public_key.public_numbers()
|
||||
return MappingProxyType(
|
||||
{
|
||||
"crv": _JWK_CURVE_NAME,
|
||||
"kty": _JWK_KEY_TYPE,
|
||||
"x": _b64url_coordinate(numbers.x),
|
||||
"y": _b64url_coordinate(numbers.y),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def rfc7638_thumbprint(public_key: "ec.EllipticCurvePublicKey") -> str:
|
||||
"""RFC 7638: SHA-256 over the lexicographically member-ordered, whitespace-free JSON
|
||||
rendering of the thumbprint members, base64url-encoded without padding."""
|
||||
canonical: Final = json.dumps(
|
||||
dict(sorted(_jwk_thumbprint_members(public_key).items())), # mutable-ok: json.dumps needs a real dict
|
||||
separators=(",", ":"),
|
||||
)
|
||||
return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
def build_jwk(public_key: "ec.EllipticCurvePublicKey", kid: str) -> Jwk:
|
||||
return MappingProxyType({**_jwk_thumbprint_members(public_key), "use": "sig", "alg": ALG, "kid": kid})
|
||||
|
||||
|
||||
def build_jwks(public_key: "ec.EllipticCurvePublicKey") -> Jwks:
|
||||
kid: Final = rfc7638_thumbprint(public_key)
|
||||
return MappingProxyType({"keys": (build_jwk(public_key, kid),)})
|
||||
|
||||
|
||||
def jwks_document_json(pem: str) -> str:
|
||||
"""The operator-facing export: the JSON document to register as Anthropic's inline JWKS.
|
||||
|
||||
``build_jwks`` returns ``MappingProxyType``/tuple values per this repo's no-mutation
|
||||
convention; the ``json`` module only knows plain ``dict``/``list``, so those are converted
|
||||
at this one serialization boundary rather than giving up immutability throughout the module.
|
||||
"""
|
||||
key: Final = load_es256_private_key(pem)
|
||||
jwks: Final = build_jwks(key.public_key())
|
||||
return json.dumps(
|
||||
{"keys": [dict(jwk) for jwk in jwks["keys"]]}, # mutable-ok: json.dumps needs real dicts/lists
|
||||
indent=2,
|
||||
)
|
||||
|
||||
|
||||
def sign_es256_jwt(pem: str, claims: Mapping[str, object]) -> str:
|
||||
"""Signs ``claims`` with the PEM key, stamping ``kid`` as its RFC 7638 thumbprint so a
|
||||
verifier can look the signing key up in the published JWKS by ``kid`` alone."""
|
||||
try:
|
||||
import jwt
|
||||
except ImportError as e:
|
||||
raise ImportError(MISSING_SIGNING_DEPENDENCIES_MESSAGE) from e
|
||||
key: Final = load_es256_private_key(pem)
|
||||
kid: Final = rfc7638_thumbprint(key.public_key())
|
||||
headers: Final = {"kid": kid} # mutable-ok: PyJWT requires a real dict, not a Mapping
|
||||
return jwt.encode(dict(claims), key, algorithm=ALG, headers=headers) # mutable-ok: PyJWT requires a real dict
|
||||
861
litellm/llms/base_llm/auth/token_exchange.py
Normal file
861
litellm/llms/base_llm/auth/token_exchange.py
Normal file
|
|
@ -0,0 +1,861 @@
|
|||
"""RFC 7523 JWT-bearer token exchange engine, shared across providers.
|
||||
|
||||
One sync state machine per process: bounded engine-owned entry map, two-tier
|
||||
refresh (advisory background refresh + mandatory single-flight), HTTPS pinning,
|
||||
response caps, and RFC 6749 5.2 redaction. Providers describe a grant profile as
|
||||
a ``TokenExchangeSpec`` and map the typed ``ExchangeError`` union to their own
|
||||
public exception contract.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Coroutine, Mapping, Sequence
|
||||
from concurrent.futures import Executor, ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from math import inf
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
|
||||
from urllib.parse import unquote, unquote_plus, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.auth.types import (
|
||||
AssertionReader,
|
||||
AssertionSource,
|
||||
AssertionSourceError,
|
||||
ExchangeCallType,
|
||||
ExchangeError,
|
||||
ExchangeResult,
|
||||
InsecureTokenUrl,
|
||||
MalformedTokenResponse,
|
||||
MintedToken,
|
||||
SyncTokenPoster,
|
||||
TokenEndpointError,
|
||||
TokenExchangeMetricsSink,
|
||||
TokenExchangeSpec,
|
||||
TokenTransportError,
|
||||
)
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
CALL_TYPE_COLD_MINT: Final[ExchangeCallType] = "cold_mint"
|
||||
CALL_TYPE_MANDATORY_REFRESH: Final[ExchangeCallType] = "mandatory_refresh"
|
||||
CALL_TYPE_ADVISORY_REFRESH: Final[ExchangeCallType] = "advisory_refresh"
|
||||
CALL_TYPE_CACHE_HIT: Final = "cache_hit"
|
||||
|
||||
ADVISORY_REFRESH_SECONDS: Final = 120.0
|
||||
MANDATORY_REFRESH_SECONDS: Final = 30.0
|
||||
ADVISORY_REFRESH_LIFETIME_FRACTION: Final = 0.5
|
||||
MANDATORY_REFRESH_LIFETIME_FRACTION: Final = 0.125
|
||||
ADVISORY_REFRESH_BACKOFF_SECONDS: Final = 5.0
|
||||
FALLBACK_TOKEN_TTL_SECONDS: Final = 60.0
|
||||
# Metrics are best-effort, so the backlog is capped and further events are dropped. Request volume
|
||||
# must not be able to grow this queue without bound when a telemetry backend stalls.
|
||||
_METRICS_QUEUE_LIMIT: Final = 1000
|
||||
MAX_ASSERTION_BYTES: Final = 16 * 1024
|
||||
MAX_RESPONSE_BYTES: Final = 1024 * 1024
|
||||
|
||||
_REDACTION_CAP: Final = 256
|
||||
_FOLLOWER_WAIT_GRACE_SECONDS: Final = 5.0
|
||||
_LOCAL_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"})
|
||||
_OAUTH_ERROR_FIELDS: Final = ("error", "error_description", "error_uri")
|
||||
_NESTED_ERROR_FIELDS: Final = ("type", "message")
|
||||
_CONTENT_TYPES: Final = MappingProxyType({"json": "application/json", "form": "application/x-www-form-urlencoded"})
|
||||
_OVERSIZED_BODY_MESSAGE: Final = "oversized error response omitted"
|
||||
_NON_OBJECT_BODY_MESSAGE: Final = "non-object error response omitted"
|
||||
_NO_OAUTH_FIELDS_MESSAGE: Final = "error response carried no RFC 6749 fields"
|
||||
_UNSTRUCTURED_BODY_MESSAGE: Final = "non-JSON error response omitted"
|
||||
_REFLECTED_VALUE_MESSAGE: Final = "<redacted: response echoed the request>"
|
||||
# A credential fragment shorter than this is not worth the false positives; longer, and a run
|
||||
# shared with the assertion is reflection rather than coincidence.
|
||||
_REFLECTION_MIN_RUN: Final = 8
|
||||
# Everything a base64url credential is NOT made of, stripped so a fragment split by delimiters
|
||||
# still lines up against the assertion.
|
||||
_CREDENTIAL_CHARS: Final = re.compile(r"[^A-Za-z0-9._~+/=-]")
|
||||
_SENTINEL_BODY_MESSAGES: Final = frozenset({_OVERSIZED_BODY_MESSAGE, _NON_OBJECT_BODY_MESSAGE})
|
||||
|
||||
|
||||
class _TokenExchangeResponse(BaseModel):
|
||||
access_token: str
|
||||
expires_in: int | None = None
|
||||
token_type: str | None = None
|
||||
|
||||
|
||||
_RedactableBody: TypeAlias = Mapping[str, object] | list[object] | str | int | float | bool | None
|
||||
_REDACTABLE_BODY_ADAPTER: Final = TypeAdapter[_RedactableBody](_RedactableBody)
|
||||
|
||||
|
||||
def endpoint_url_for_error_message(url: str) -> str:
|
||||
"""``url`` reduced to scheme, host and path for operator-facing errors.
|
||||
|
||||
A token endpoint is configuration, not a secret, and naming it is what makes these errors
|
||||
actionable. But nothing stops an operator writing a credential into it, as a query parameter
|
||||
or as userinfo, and these errors reach model callers, so neither part is echoed.
|
||||
"""
|
||||
parsed: Final = urlsplit(url)
|
||||
host: Final = parsed.hostname or ""
|
||||
authority: Final = f"{host}:{parsed.port}" if parsed.port is not None else host
|
||||
return urlunsplit((parsed.scheme, authority, parsed.path, "", ""))
|
||||
|
||||
|
||||
def validate_token_endpoint_url(url: str) -> str | InsecureTokenUrl:
|
||||
parsed: Final = urlsplit(url)
|
||||
if parsed.scheme == "https":
|
||||
return url
|
||||
if parsed.scheme == "http" and (parsed.hostname or "") in _LOCAL_HOSTS:
|
||||
return url
|
||||
return InsecureTokenUrl(host=parsed.hostname or "")
|
||||
|
||||
|
||||
def redact_oauth_error_body(
|
||||
status_code: int,
|
||||
body_text: str,
|
||||
assertion: SecretStr | Sequence[SecretStr] | None = None,
|
||||
) -> TokenEndpointError:
|
||||
"""``assertion`` may be every form of the credential that went out on the wire.
|
||||
|
||||
A grant that encodes its credential before sending it (``client_secret_basic`` base64s
|
||||
``id:secret``) can have that encoded form echoed back, and it decodes straight to the secret,
|
||||
so checking only the raw value lets reversible material through.
|
||||
"""
|
||||
rendered: Final = _redact_body_text(body_text)
|
||||
secrets: Final = () if assertion is None else (assertion,) if isinstance(assertion, SecretStr) else tuple(assertion)
|
||||
redacted: Final = next(
|
||||
(
|
||||
_REFLECTED_VALUE_MESSAGE
|
||||
for secret in secrets
|
||||
if _drop_reflected_assertion(rendered, secret) is _REFLECTED_VALUE_MESSAGE
|
||||
),
|
||||
rendered,
|
||||
)
|
||||
return TokenEndpointError(status_code=status_code, redacted_body=redacted)
|
||||
|
||||
|
||||
def _drop_reflected_assertion(rendered: str, assertion: SecretStr | None) -> str:
|
||||
"""Catches an endpoint that echoes the submitted credential back, verbatim or in fragments,
|
||||
however it split or percent-encoded it.
|
||||
|
||||
Both sides are reduced to the characters a credential is made of before comparison. Stripping
|
||||
only the rendered side would stop matching a secret that carries spaces or punctuation of its
|
||||
own, which is exactly the hand-set passphrase most at risk of being echoed.
|
||||
|
||||
This stops an accidental or naive echo. It cannot stop an endpoint that deliberately re-encodes
|
||||
or interleaves the credential, and it is not what keeps the credential from the endpoint, which
|
||||
already holds it. What it protects is blast radius: keeping the value out of the caller's error
|
||||
and out of third-party log sinks.
|
||||
"""
|
||||
if assertion is None:
|
||||
return rendered
|
||||
secret: Final = assertion.get_secret_value()
|
||||
if not secret:
|
||||
return rendered
|
||||
if secret in rendered:
|
||||
return _REFLECTED_VALUE_MESSAGE
|
||||
compacted_secret: Final = _CREDENTIAL_CHARS.sub("", secret)
|
||||
if not compacted_secret:
|
||||
return rendered
|
||||
return _REFLECTED_VALUE_MESSAGE if _shares_a_credential_run(rendered, compacted_secret) else rendered
|
||||
|
||||
|
||||
def _shares_a_credential_run(rendered: str, compacted_secret: str) -> bool:
|
||||
"""``unquote`` covers a credential sent form-encoded, without every caller enumerating that
|
||||
shape for itself: percent-escaping is reversible and applies to any field, query string
|
||||
included.
|
||||
|
||||
A secret shorter than the probe run is compared whole: a window longer than the secret can
|
||||
never be found inside it, which would leave a short client secret unprotected in every shape
|
||||
but the verbatim one.
|
||||
"""
|
||||
# unquote covers %XX; unquote_plus additionally covers the "+" a form-encoded body uses for a
|
||||
# space. Both are kept rather than only the wider one, because "+" is a base64 character and
|
||||
# decoding it away would lose a run that the undecoded candidate still matches on.
|
||||
run: Final = min(_REFLECTION_MIN_RUN, len(compacted_secret))
|
||||
compacted_candidates: Final = tuple(
|
||||
_CREDENTIAL_CHARS.sub("", candidate) for candidate in (rendered, unquote(rendered), unquote_plus(rendered))
|
||||
)
|
||||
return any(
|
||||
compacted[start : start + run] in compacted_secret
|
||||
for compacted in compacted_candidates
|
||||
for start in range(len(compacted) - run + 1)
|
||||
)
|
||||
|
||||
|
||||
def _redact_body_text(body_text: str) -> str:
|
||||
if body_text in _SENTINEL_BODY_MESSAGES:
|
||||
return body_text
|
||||
if len(body_text) > MAX_RESPONSE_BYTES:
|
||||
return _OVERSIZED_BODY_MESSAGE
|
||||
try:
|
||||
parsed: Final = _REDACTABLE_BODY_ADAPTER.validate_json(body_text)
|
||||
except ValidationError:
|
||||
return _UNSTRUCTURED_BODY_MESSAGE
|
||||
match parsed:
|
||||
case Mapping():
|
||||
return _format_oauth_error_fields(parsed)
|
||||
case _:
|
||||
return _NON_OBJECT_BODY_MESSAGE
|
||||
|
||||
|
||||
def _format_oauth_error_fields(body: Mapping[str, object]) -> str:
|
||||
fields: Final = tuple(
|
||||
f"{name}: {_format_oauth_error_value(value)}"
|
||||
for name in _OAUTH_ERROR_FIELDS
|
||||
for value in (body.get(name),)
|
||||
if value is not None
|
||||
)
|
||||
return "; ".join(fields) if fields else _NO_OAUTH_FIELDS_MESSAGE
|
||||
|
||||
|
||||
def _format_oauth_error_value(value: object) -> str:
|
||||
"""RFC 6749 types ``error`` as a string, but Anthropic (and other providers) nest their
|
||||
own ``{"type": ..., "message": ...}`` envelope there; render that rather than a dict repr."""
|
||||
if isinstance(value, Mapping):
|
||||
nested: Final = tuple(
|
||||
f"{str(part)[:_REDACTION_CAP]}"
|
||||
for key in _NESTED_ERROR_FIELDS
|
||||
for part in (value.get(key),)
|
||||
if part is not None
|
||||
)
|
||||
if nested:
|
||||
return " - ".join(nested)
|
||||
return str(value)[:_REDACTION_CAP]
|
||||
|
||||
|
||||
def _error_summary(error: ExchangeError) -> str:
|
||||
match error:
|
||||
case AssertionSourceError():
|
||||
return f"AssertionSourceError: assertion {error.kind} from {error.source_ref}"
|
||||
case InsecureTokenUrl():
|
||||
return f"InsecureTokenUrl: insecure token endpoint host {error.host}"
|
||||
case TokenEndpointError():
|
||||
return f"TokenEndpointError: HTTP {error.status_code}: {error.redacted_body}"
|
||||
case TokenTransportError():
|
||||
return f"TokenTransportError: {error.detail}"
|
||||
case MalformedTokenResponse():
|
||||
return f"MalformedTokenResponse: {error.detail}"
|
||||
case _:
|
||||
assert_never(error)
|
||||
|
||||
|
||||
class _MetricsFailure(Exception):
|
||||
"""Never raised: typed carriers handed to the service failure hook so the prometheus
|
||||
``error_class`` label names the ``ExchangeError`` variant; the message is the redacted
|
||||
``_error_summary`` and carries no credential material."""
|
||||
|
||||
|
||||
class TokenExchangeAssertionSourceFailure(_MetricsFailure): ...
|
||||
|
||||
|
||||
class TokenExchangeInsecureUrlFailure(_MetricsFailure): ...
|
||||
|
||||
|
||||
class TokenExchangeEndpointFailure(_MetricsFailure): ...
|
||||
|
||||
|
||||
class TokenExchangeTransportFailure(_MetricsFailure): ...
|
||||
|
||||
|
||||
class TokenExchangeMalformedResponseFailure(_MetricsFailure): ...
|
||||
|
||||
|
||||
def _failure_exception(error: ExchangeError) -> _MetricsFailure:
|
||||
summary: Final = _error_summary(error)
|
||||
match error:
|
||||
case AssertionSourceError():
|
||||
return TokenExchangeAssertionSourceFailure(summary)
|
||||
case InsecureTokenUrl():
|
||||
return TokenExchangeInsecureUrlFailure(summary)
|
||||
case TokenEndpointError():
|
||||
return TokenExchangeEndpointFailure(summary)
|
||||
case TokenTransportError():
|
||||
return TokenExchangeTransportFailure(summary)
|
||||
case MalformedTokenResponse():
|
||||
return TokenExchangeMalformedResponseFailure(summary)
|
||||
case _:
|
||||
assert_never(error)
|
||||
|
||||
|
||||
def _cache_key(spec: TokenExchangeSpec) -> str:
|
||||
return hashlib.sha256(
|
||||
"\x1f".join((spec.token_url, spec.assertion_ref, *spec.cache_key_identity)).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _assertion_fetch(reader: AssertionReader, spec: TokenExchangeSpec) -> AssertionSource:
|
||||
"""``spec.assertion_source`` (an identity source's own fetch/mint closure) takes priority over
|
||||
the engine-level reader when set; either way, failures are reported against ``spec.assertion_ref``."""
|
||||
if spec.assertion_source is not None:
|
||||
return spec.assertion_source
|
||||
return lambda: reader(spec.assertion_ref)
|
||||
|
||||
|
||||
def _read_assertion(fetch: AssertionSource, ref: str) -> SecretStr | AssertionSourceError:
|
||||
from litellm.secret_managers.main import OidcPathNotAllowedError
|
||||
|
||||
try:
|
||||
raw: Final = fetch()
|
||||
except OidcPathNotAllowedError:
|
||||
return AssertionSourceError(kind="disallowed_path", source_ref=ref)
|
||||
except (ValueError, ImportError) as e:
|
||||
return AssertionSourceError(kind="unreadable", source_ref=ref, detail=str(e)[:_REDACTION_CAP])
|
||||
except Exception: # noqa: BLE001 # injected readers (secret managers) raise arbitrarily; all failures become values
|
||||
return AssertionSourceError(kind="unreadable", source_ref=ref)
|
||||
if raw is None:
|
||||
return AssertionSourceError(kind="missing", source_ref=ref)
|
||||
stripped: Final = raw.strip()
|
||||
if not stripped:
|
||||
return AssertionSourceError(kind="empty", source_ref=ref)
|
||||
if len(stripped.encode("utf-8")) > MAX_ASSERTION_BYTES:
|
||||
return AssertionSourceError(kind="oversized", source_ref=ref)
|
||||
return SecretStr(stripped)
|
||||
|
||||
|
||||
def _serialize_body(spec: TokenExchangeSpec, assertion: SecretStr) -> bytes:
|
||||
if spec.body_encoding == "json":
|
||||
return json.dumps(
|
||||
{ # mutable-ok: transient body dict consumed inline by the serializer
|
||||
**spec.static_body,
|
||||
spec.assertion_field: assertion.get_secret_value(),
|
||||
}
|
||||
).encode()
|
||||
return urlencode(
|
||||
{ # mutable-ok: transient body dict consumed inline by the serializer
|
||||
**spec.static_body,
|
||||
spec.assertion_field: assertion.get_secret_value(),
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def _sanitize_expires_in(expires_in: int | None) -> float:
|
||||
if expires_in is None or expires_in <= 0:
|
||||
return FALLBACK_TOKEN_TTL_SECONDS
|
||||
return float(expires_in)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RefreshWindows:
|
||||
advisory: float
|
||||
mandatory: float
|
||||
|
||||
|
||||
def _refresh_windows(lifetime_seconds: float | None) -> _RefreshWindows:
|
||||
"""A token whose whole life is shorter than the flat windows sits inside them from the moment it
|
||||
is minted, so every request would arm another background exchange against the token endpoint.
|
||||
Scaling each window by a fraction of the observed lifetime makes a 60s token refresh around its
|
||||
half life instead; at a lifetime of 240s and above both fractions reach the flat windows, so
|
||||
ordinary long-lived tokens keep exactly the 120s/30s behaviour."""
|
||||
if lifetime_seconds is None or lifetime_seconds <= 0.0:
|
||||
return _RefreshWindows(advisory=ADVISORY_REFRESH_SECONDS, mandatory=MANDATORY_REFRESH_SECONDS)
|
||||
return _RefreshWindows(
|
||||
advisory=min(ADVISORY_REFRESH_SECONDS, lifetime_seconds * ADVISORY_REFRESH_LIFETIME_FRACTION),
|
||||
mandatory=min(MANDATORY_REFRESH_SECONDS, lifetime_seconds * MANDATORY_REFRESH_LIFETIME_FRACTION),
|
||||
)
|
||||
|
||||
|
||||
def _capped_body_text(response: httpx.Response) -> str:
|
||||
if len(response.content) > MAX_RESPONSE_BYTES:
|
||||
return _OVERSIZED_BODY_MESSAGE
|
||||
return response.text
|
||||
|
||||
|
||||
def _default_assertion_reader(ref: str) -> str | None:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return get_secret_str(ref)
|
||||
|
||||
|
||||
def _new_exchange_handler() -> "HTTPHandler":
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
handler: Final = HTTPHandler(timeout=httpx.Timeout(timeout=30.0, connect=5.0))
|
||||
handler.client.follow_redirects = False
|
||||
return handler
|
||||
|
||||
|
||||
def require_posted_response(response: httpx.Response | None, endpoint_label: str) -> httpx.Response:
|
||||
"""The legacy ``HTTPHandler`` carries no return annotation, so a patched or stubbed client can
|
||||
hand a poster ``None`` back; a transport error beats dereferencing it."""
|
||||
if response is None:
|
||||
raise httpx.TransportError(f"{endpoint_label} returned no response")
|
||||
return response
|
||||
|
||||
|
||||
class _HttpxSyncTokenPoster:
|
||||
"""Default poster: a dedicated HTTPHandler (no logging_obj, so litellm's
|
||||
pre/post-call body logging never sees the exchange POST); returns the
|
||||
response for any status."""
|
||||
|
||||
def __init__(self, handler_factory: Callable[[], "HTTPHandler"] = _new_exchange_handler) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._handler_factory: Final = handler_factory
|
||||
self._handler: HTTPHandler | None = None
|
||||
|
||||
def _handler_instance(self) -> "HTTPHandler":
|
||||
with self._lock:
|
||||
if self._handler is None:
|
||||
self._handler = self._handler_factory()
|
||||
return self._handler
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
|
||||
try:
|
||||
response: Final[httpx.Response | None] = self._handler_instance().post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post is legacy-untyped; the result is validated below
|
||||
url,
|
||||
content=content,
|
||||
headers=dict(headers), # mutable-ok: HTTPHandler.post requires a concrete dict
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
return e.response
|
||||
return require_posted_response(response, "token endpoint")
|
||||
|
||||
|
||||
class _ServiceLoggingHooks(Protocol):
|
||||
"""The slice of ``litellm._service_logger.ServiceLogging`` the metrics sink calls; a protocol
|
||||
so tests inject a recorder instead of monkeypatching."""
|
||||
|
||||
async def async_service_success_hook(self, service: ServiceTypes, call_type: str, duration: float) -> None: ...
|
||||
|
||||
async def async_service_failure_hook(
|
||||
self, service: ServiceTypes, duration: float, error: str | Exception, call_type: str
|
||||
) -> None: ...
|
||||
|
||||
|
||||
_HooksCoroFactory: TypeAlias = Callable[
|
||||
[_ServiceLoggingHooks], # mutable-ok: Callable param-list syntax, not a list
|
||||
Coroutine[object, object, None],
|
||||
]
|
||||
|
||||
|
||||
def _default_service_logging() -> _ServiceLoggingHooks:
|
||||
from litellm._service_logger import ServiceLogging
|
||||
|
||||
return ServiceLogging()
|
||||
|
||||
|
||||
class ServiceLoggingMetricsSink:
|
||||
"""Default sink: bridges engine metrics onto litellm's ServiceTypes pattern
|
||||
(prometheus ``litellm_anthropic_wif_*`` via ``service_callback``). The engine's entry points
|
||||
are sync threads with no event loop, and the service hooks are async, so every emission is
|
||||
fire-and-forget on a dedicated single worker thread that owns its own short-lived loop --
|
||||
the mint path only ever pays for an executor queue put."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
service_logging_factory: Callable[[], _ServiceLoggingHooks] = _default_service_logging,
|
||||
executor: Executor | None = None,
|
||||
) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._service_logging_factory: Final = service_logging_factory
|
||||
self._service_logging: _ServiceLoggingHooks | None = None
|
||||
self._executor: Executor | None = executor
|
||||
self._queued: int = 0 # rebind-ok: backlog depth, guarded by _lock
|
||||
|
||||
def _service_logging_instance(self) -> _ServiceLoggingHooks:
|
||||
with self._lock:
|
||||
if self._service_logging is None:
|
||||
self._service_logging = self._service_logging_factory()
|
||||
return self._service_logging
|
||||
|
||||
def _executor_instance(self) -> Executor:
|
||||
with self._lock:
|
||||
if self._executor is None:
|
||||
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-token-exchange-metrics")
|
||||
return self._executor
|
||||
|
||||
def _emit(self, coro_factory: _HooksCoroFactory) -> None:
|
||||
try:
|
||||
asyncio.run(coro_factory(self._service_logging_instance()))
|
||||
except Exception as e: # noqa: BLE001 # metrics are best-effort; emission failures must never surface
|
||||
verbose_logger.debug("token exchange metrics emission failed: %s", e)
|
||||
|
||||
def _submit(self, coro_factory: _HooksCoroFactory) -> None:
|
||||
"""Drop the event rather than queue it once the backlog is full. A stalled telemetry
|
||||
backend must not let request volume grow an unbounded queue in the proxy: losing a
|
||||
metric sample is always cheaper than losing the process."""
|
||||
with self._lock:
|
||||
if self._queued >= _METRICS_QUEUE_LIMIT:
|
||||
verbose_logger.debug("token exchange metrics queue full, dropping event")
|
||||
return
|
||||
self._queued += 1
|
||||
try:
|
||||
self._executor_instance().submit(self._emit_and_release, coro_factory)
|
||||
except Exception as e: # noqa: BLE001 # a rejected submit must not surface to the mint
|
||||
with self._lock:
|
||||
self._queued -= 1
|
||||
verbose_logger.debug("token exchange metrics submit failed: %s", e)
|
||||
|
||||
def _emit_and_release(self, coro_factory: _HooksCoroFactory) -> None:
|
||||
try:
|
||||
self._emit(coro_factory)
|
||||
finally:
|
||||
with self._lock:
|
||||
self._queued -= 1
|
||||
|
||||
def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None:
|
||||
def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]:
|
||||
return hooks.async_service_success_hook(
|
||||
service=ServiceTypes.ANTHROPIC_WIF, call_type=call_type, duration=duration_seconds
|
||||
)
|
||||
|
||||
self._submit(start)
|
||||
|
||||
def exchange_failure(self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError) -> None:
|
||||
failure: Final = _failure_exception(error)
|
||||
|
||||
def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]:
|
||||
return hooks.async_service_failure_hook(
|
||||
service=ServiceTypes.ANTHROPIC_WIF, duration=duration_seconds, error=failure, call_type=call_type
|
||||
)
|
||||
|
||||
self._submit(start)
|
||||
|
||||
def cache_hit(self) -> None:
|
||||
def start(hooks: _ServiceLoggingHooks) -> Coroutine[object, object, None]:
|
||||
return hooks.async_service_success_hook(
|
||||
service=ServiceTypes.ANTHROPIC_WIF_CACHE, call_type=CALL_TYPE_CACHE_HIT, duration=0.0
|
||||
)
|
||||
|
||||
self._submit(start)
|
||||
|
||||
|
||||
class _Entry:
|
||||
"""Single-flight state for one cache key; mutable by design, confined to the
|
||||
engine, and only ever mutated under the engine lock."""
|
||||
|
||||
__slots__ = ("backoff_until", "done", "force_refresh", "in_flight", "last_error", "lifetime_seconds", "token")
|
||||
|
||||
def __init__(self, force_refresh: bool = False) -> None:
|
||||
self.token: MintedToken | None = None
|
||||
self.lifetime_seconds: float | None = None
|
||||
self.in_flight: bool = False
|
||||
self.done: Final = threading.Event()
|
||||
self.backoff_until: float = float("-inf")
|
||||
self.force_refresh: bool = force_refresh
|
||||
self.last_error: ExchangeError | None = None
|
||||
|
||||
def arm(self) -> None:
|
||||
self.in_flight = True
|
||||
self.last_error = None
|
||||
self.done.clear()
|
||||
|
||||
def _store(self, token: MintedToken, now: float) -> None:
|
||||
self.token = token
|
||||
self.lifetime_seconds = None if token.expires_at is None else max(token.expires_at - now, 0.0)
|
||||
self.last_error = None
|
||||
|
||||
def publish(self, result: ExchangeResult, now: float) -> None:
|
||||
match result:
|
||||
case MintedToken():
|
||||
self._store(result, now)
|
||||
case _:
|
||||
self.last_error = result
|
||||
self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS
|
||||
self.force_refresh = False
|
||||
self.in_flight = False
|
||||
self.done.set()
|
||||
|
||||
def publish_advisory(self, result: ExchangeResult, now: float) -> None:
|
||||
"""A failed advisory refresh records only the backoff, never ``last_error``: a follower whose
|
||||
cached token expires while this runs must be free to re-lead a fresh mint and recover."""
|
||||
match result:
|
||||
case MintedToken():
|
||||
self._store(result, now)
|
||||
case _:
|
||||
self.backoff_until = now + ADVISORY_REFRESH_BACKOFF_SECONDS
|
||||
self.in_flight = False
|
||||
self.done.set()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Serve:
|
||||
token: MintedToken
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ServeAndRefresh:
|
||||
token: MintedToken
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Lead:
|
||||
call_type: ExchangeCallType
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Follow:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Fail:
|
||||
error: ExchangeError
|
||||
|
||||
|
||||
_Decision: TypeAlias = _Serve | _ServeAndRefresh | _Lead | _Follow | _Fail
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Unauthorized:
|
||||
response: httpx.Response
|
||||
assertion: SecretStr
|
||||
|
||||
|
||||
class JwtBearerTokenExchangeEngine:
|
||||
def __init__(
|
||||
self,
|
||||
poster: SyncTokenPoster | None = None,
|
||||
assertion_reader: AssertionReader | None = None,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
refresh_executor: Executor | None = None,
|
||||
max_entries: int = 64,
|
||||
metrics_sink: TokenExchangeMetricsSink | None = None,
|
||||
) -> None:
|
||||
self._poster: Final[SyncTokenPoster] = poster if poster is not None else _HttpxSyncTokenPoster()
|
||||
self._assertion_reader: Final[AssertionReader] = (
|
||||
assertion_reader if assertion_reader is not None else _default_assertion_reader
|
||||
)
|
||||
self._clock: Final = clock
|
||||
self._refresh_executor: Executor | None = refresh_executor
|
||||
self._max_entries: Final = max_entries
|
||||
self._metrics_sink: Final[TokenExchangeMetricsSink] = (
|
||||
metrics_sink if metrics_sink is not None else ServiceLoggingMetricsSink()
|
||||
)
|
||||
self._lock: Final = threading.Lock()
|
||||
self._entries: Final[dict[str, _Entry]] = {} # mutable-ok: engine-owned map guarded by _lock
|
||||
|
||||
def get_token(self, spec: TokenExchangeSpec) -> ExchangeResult:
|
||||
"""A follower whose leader published nothing re-classifies rather than recursing, so a
|
||||
contended entry cannot grow the stack one frame per failed leader."""
|
||||
while True:
|
||||
with self._lock:
|
||||
entry = self._get_or_create_entry_locked(spec) # rebind-ok: re-read per follower round
|
||||
decision = self._classify_and_arm_locked(entry) # rebind-ok: re-read per follower round
|
||||
match decision:
|
||||
case _Serve(token=token):
|
||||
self._report_cache_hit()
|
||||
return token
|
||||
case _ServeAndRefresh(token=token):
|
||||
self._report_cache_hit()
|
||||
self._executor_instance().submit(self._advisory_refresh, spec, entry)
|
||||
return token
|
||||
case _Fail(error=error):
|
||||
return error
|
||||
case _Lead(call_type=call_type):
|
||||
return self._lead(spec, entry, call_type)
|
||||
case _Follow():
|
||||
followed = self._await_leader(spec, entry) # rebind-ok: one leader wait per round
|
||||
if followed is not None:
|
||||
return followed
|
||||
case _:
|
||||
assert_never(decision)
|
||||
|
||||
async def aget_token(self, spec: TokenExchangeSpec) -> ExchangeResult:
|
||||
return await asyncio.to_thread(self.get_token, spec)
|
||||
|
||||
def invalidate(self, spec: TokenExchangeSpec) -> None:
|
||||
key: Final = _cache_key(spec)
|
||||
with self._lock:
|
||||
if key in self._entries:
|
||||
self._entries[key] = _Entry(force_refresh=True)
|
||||
|
||||
def _get_or_create_entry_locked(self, spec: TokenExchangeSpec) -> _Entry:
|
||||
key: Final = _cache_key(spec)
|
||||
existing: Final = self._entries.get(key)
|
||||
if existing is not None:
|
||||
return existing
|
||||
if len(self._entries) >= self._max_entries:
|
||||
self._evict_locked()
|
||||
created: Final = _Entry()
|
||||
self._entries[key] = created
|
||||
return created
|
||||
|
||||
def _evict_locked(self) -> None:
|
||||
now: Final = self._clock()
|
||||
stale: Final = tuple(
|
||||
key
|
||||
for key, entry in self._entries.items()
|
||||
if not entry.in_flight
|
||||
and (entry.token is None or (entry.token.expires_at is not None and entry.token.expires_at <= now))
|
||||
)
|
||||
for key in stale:
|
||||
del self._entries[key]
|
||||
if len(self._entries) < self._max_entries:
|
||||
return
|
||||
# Evict soonest-to-expire first, and take as many as the overshoot needs rather than one, so a
|
||||
# burst of distinct identities does not leave the map permanently above max_entries. An entry
|
||||
# a leader owns or a follower waits on is never a candidate, so a moment where every entry is
|
||||
# in flight still over-inserts; that residue is bounded by the concurrent mints themselves.
|
||||
evictable: Final = sorted(
|
||||
(
|
||||
entry.token.expires_at if entry.token is not None and entry.token.expires_at is not None else -inf,
|
||||
key,
|
||||
)
|
||||
for key, entry in self._entries.items()
|
||||
if not entry.in_flight
|
||||
)
|
||||
for _, key in evictable[: len(self._entries) - self._max_entries + 1]:
|
||||
del self._entries[key]
|
||||
|
||||
def _classify_and_arm_locked(self, entry: _Entry) -> _Decision:
|
||||
token: Final = entry.token
|
||||
if token is not None and not entry.force_refresh:
|
||||
if token.expires_at is None:
|
||||
return _Serve(token=token)
|
||||
windows: Final = _refresh_windows(entry.lifetime_seconds)
|
||||
remaining: Final = token.expires_at - self._clock()
|
||||
if remaining > windows.advisory:
|
||||
return _Serve(token=token)
|
||||
if remaining > windows.mandatory:
|
||||
if entry.in_flight or self._clock() < entry.backoff_until:
|
||||
return _Serve(token=token)
|
||||
entry.arm()
|
||||
return _ServeAndRefresh(token=token)
|
||||
if entry.in_flight:
|
||||
return _Follow()
|
||||
if entry.last_error is not None and self._clock() < entry.backoff_until:
|
||||
return _Fail(error=entry.last_error)
|
||||
entry.arm()
|
||||
return _Lead(call_type=CALL_TYPE_COLD_MINT if token is None else CALL_TYPE_MANDATORY_REFRESH)
|
||||
|
||||
def _executor_instance(self) -> Executor:
|
||||
with self._lock:
|
||||
if self._refresh_executor is None:
|
||||
self._refresh_executor = ThreadPoolExecutor(
|
||||
max_workers=2, thread_name_prefix="litellm-token-exchange-refresh"
|
||||
)
|
||||
return self._refresh_executor
|
||||
|
||||
def _lead(self, spec: TokenExchangeSpec, entry: _Entry, call_type: ExchangeCallType) -> ExchangeResult:
|
||||
started: Final = self._clock()
|
||||
result: Final = self._exchange_never_raises(spec)
|
||||
duration: Final = self._clock() - started
|
||||
with self._lock:
|
||||
entry.publish(result, now=self._clock())
|
||||
self._report_exchange(call_type, duration, result)
|
||||
return result
|
||||
|
||||
def _await_leader(self, spec: TokenExchangeSpec, entry: _Entry) -> "ExchangeResult | None":
|
||||
"""None means the finished round left neither a valid token nor an error
|
||||
(a failed advisory refresh); the caller re-enters and leads a fresh exchange."""
|
||||
leader_finished: Final = entry.done.wait(2 * spec.timeout_seconds + _FOLLOWER_WAIT_GRACE_SECONDS)
|
||||
with self._lock:
|
||||
token: Final = entry.token
|
||||
if token is not None and (token.expires_at is None or token.expires_at > self._clock()):
|
||||
return token
|
||||
if entry.last_error is not None:
|
||||
return entry.last_error
|
||||
if leader_finished:
|
||||
return None
|
||||
return TokenTransportError(detail="timed out waiting for the token exchange leader")
|
||||
|
||||
def _advisory_refresh(self, spec: TokenExchangeSpec, entry: _Entry) -> None:
|
||||
started: Final = self._clock()
|
||||
result: Final = self._exchange_never_raises(spec)
|
||||
duration: Final = self._clock() - started
|
||||
with self._lock:
|
||||
now: Final = self._clock()
|
||||
entry.publish_advisory(result, now=now)
|
||||
stale_expires_at: Final = entry.token.expires_at if entry.token is not None else None
|
||||
stale_mandatory: Final = _refresh_windows(entry.lifetime_seconds).mandatory
|
||||
self._report_exchange(CALL_TYPE_ADVISORY_REFRESH, duration, result)
|
||||
if isinstance(result, MintedToken):
|
||||
return
|
||||
seconds_to_mandatory_wall: Final = (
|
||||
max(stale_expires_at - now - stale_mandatory, 0.0) if stale_expires_at is not None else 0.0
|
||||
)
|
||||
verbose_logger.warning(
|
||||
"Advisory token refresh against %s failed (%s); serving the cached token for up to "
|
||||
"%.0fs before the mandatory refresh wall; next attempt after %.0fs backoff",
|
||||
urlsplit(spec.token_url).hostname or "",
|
||||
_error_summary(result),
|
||||
seconds_to_mandatory_wall,
|
||||
ADVISORY_REFRESH_BACKOFF_SECONDS,
|
||||
)
|
||||
|
||||
def _report_exchange(self, call_type: ExchangeCallType, duration_seconds: float, result: ExchangeResult) -> None:
|
||||
try:
|
||||
match result:
|
||||
case MintedToken():
|
||||
self._metrics_sink.exchange_success(call_type=call_type, duration_seconds=duration_seconds)
|
||||
case _:
|
||||
self._metrics_sink.exchange_failure(
|
||||
call_type=call_type, duration_seconds=duration_seconds, error=result
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a mint
|
||||
verbose_logger.debug("token exchange metrics emission failed: %s", e)
|
||||
|
||||
def _report_cache_hit(self) -> None:
|
||||
try:
|
||||
self._metrics_sink.cache_hit()
|
||||
except Exception as e: # noqa: BLE001 # metrics are best-effort; a sink failure must never fail a serve
|
||||
verbose_logger.debug("token exchange cache-hit metric emission failed: %s", e)
|
||||
|
||||
def _exchange_never_raises(self, spec: TokenExchangeSpec) -> ExchangeResult:
|
||||
"""The single-flight leader and the advisory refresher must always publish a result: an
|
||||
unhandled exception here would leave the entry armed (in_flight, cleared event) forever, so
|
||||
every subsequent caller for this key would follow a leader that never finishes."""
|
||||
try:
|
||||
return self._exchange(spec)
|
||||
except Exception as e: # noqa: BLE001 # a leader must resolve its entry; any failure becomes a value
|
||||
return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP])
|
||||
|
||||
def _exchange(self, spec: TokenExchangeSpec) -> ExchangeResult:
|
||||
first: Final = self._attempt_exchange(spec)
|
||||
if not isinstance(first, _Unauthorized):
|
||||
return first
|
||||
second: Final = self._attempt_exchange(spec)
|
||||
if isinstance(second, _Unauthorized):
|
||||
return redact_oauth_error_body(
|
||||
second.response.status_code, _capped_body_text(second.response), second.assertion
|
||||
)
|
||||
return second
|
||||
|
||||
def _attempt_exchange(self, spec: TokenExchangeSpec) -> "ExchangeResult | _Unauthorized":
|
||||
assertion: Final = _read_assertion(_assertion_fetch(self._assertion_reader, spec), spec.assertion_ref)
|
||||
if isinstance(assertion, AssertionSourceError):
|
||||
return assertion
|
||||
url_check: Final = validate_token_endpoint_url(spec.token_url)
|
||||
if isinstance(url_check, InsecureTokenUrl):
|
||||
return url_check
|
||||
try:
|
||||
response: Final = self._poster.post(
|
||||
spec.token_url,
|
||||
content=_serialize_body(spec, assertion),
|
||||
headers=MappingProxyType({"content-type": _CONTENT_TYPES[spec.body_encoding], **spec.request_headers}),
|
||||
timeout=spec.timeout_seconds,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; transport failures become values
|
||||
return TokenTransportError(detail=f"{type(e).__name__}: {e}"[:_REDACTION_CAP])
|
||||
if response.status_code == 401:
|
||||
return _Unauthorized(response=response, assertion=assertion)
|
||||
return self._parse_response(response, assertion)
|
||||
|
||||
def _parse_response(self, response: httpx.Response, assertion: SecretStr | None = None) -> ExchangeResult:
|
||||
if not 200 <= response.status_code < 300:
|
||||
return redact_oauth_error_body(response.status_code, _capped_body_text(response), assertion)
|
||||
if len(response.content) > MAX_RESPONSE_BYTES:
|
||||
return MalformedTokenResponse(detail="token response body exceeds the 1 MiB cap")
|
||||
try:
|
||||
parsed: Final = _TokenExchangeResponse.model_validate_json(response.content)
|
||||
except ValidationError:
|
||||
return MalformedTokenResponse(detail="token response failed RFC 6749 5.1 schema validation")
|
||||
if parsed.token_type is not None and parsed.token_type.lower() != "bearer":
|
||||
return MalformedTokenResponse(detail="token response carried a non-bearer token_type")
|
||||
if not parsed.access_token.strip():
|
||||
return MalformedTokenResponse(detail="token response carried an empty access_token")
|
||||
return MintedToken(
|
||||
access_token=SecretStr(parsed.access_token),
|
||||
expires_at=self._clock() + _sanitize_expires_in(parsed.expires_in),
|
||||
)
|
||||
|
||||
|
||||
default_token_exchange_engine: Final = JwtBearerTokenExchangeEngine()
|
||||
99
litellm/llms/base_llm/auth/types.py
Normal file
99
litellm/llms/base_llm/auth/types.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""Provider-agnostic types for the RFC 7523 JWT-bearer token exchange engine."""
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import SecretStr
|
||||
|
||||
BodyEncoding: TypeAlias = Literal["json", "form"]
|
||||
AssertionReader: TypeAlias = Callable[[str], str | None] # mutable-ok: Callable param-list syntax, not a list
|
||||
AssertionSource: TypeAlias = Callable[[], str | None] # mutable-ok: Callable param-list syntax, not a list
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenExchangeSpec:
|
||||
"""One grant profile as pure data: one instance per (provider, deployment, identity).
|
||||
|
||||
``token_url`` must be derived from deployment config/env only, never per-request caller
|
||||
input. ``assertion_ref`` is a ``oidc/...`` get_secret ref resolved fresh on every exchange.
|
||||
|
||||
``assertion_source``, when set, is a zero-arg per-config fetch/mint closure that the engine
|
||||
prefers over its own engine-level ``AssertionReader`` -- the dispatch mechanism identity
|
||||
sources beyond token_file/env (e.g. ``internal_issuer``, ``keycloak``) use to plug into the
|
||||
shared engine without a global registry. ``assertion_ref`` still names the cache-key
|
||||
discriminator and the ref echoed into operator-facing errors either way.
|
||||
"""
|
||||
|
||||
token_url: str
|
||||
assertion_ref: str
|
||||
assertion_field: str
|
||||
static_body: Mapping[str, str]
|
||||
body_encoding: BodyEncoding
|
||||
request_headers: Mapping[str, str]
|
||||
cache_key_identity: tuple[str, ...]
|
||||
timeout_seconds: float = 30.0
|
||||
assertion_source: AssertionSource | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MintedToken:
|
||||
access_token: SecretStr
|
||||
expires_at: float | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AssertionSourceError:
|
||||
kind: Literal["missing", "empty", "oversized", "unreadable", "disallowed_path"]
|
||||
source_ref: str
|
||||
detail: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InsecureTokenUrl:
|
||||
host: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenEndpointError:
|
||||
status_code: int
|
||||
redacted_body: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenTransportError:
|
||||
detail: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MalformedTokenResponse:
|
||||
detail: str
|
||||
|
||||
|
||||
ExchangeError: TypeAlias = (
|
||||
AssertionSourceError | InsecureTokenUrl | TokenEndpointError | TokenTransportError | MalformedTokenResponse
|
||||
)
|
||||
ExchangeResult: TypeAlias = MintedToken | ExchangeError
|
||||
|
||||
ExchangeCallType: TypeAlias = Literal["cold_mint", "mandatory_refresh", "advisory_refresh"]
|
||||
|
||||
|
||||
class TokenExchangeMetricsSink(Protocol):
|
||||
"""Observability seam for the exchange engine. Implementations must be best-effort: never raise
|
||||
into the mint path, never block the calling thread, and never receive credential material --
|
||||
``ExchangeError`` values are redacted by construction."""
|
||||
|
||||
def exchange_success(self, *, call_type: ExchangeCallType, duration_seconds: float) -> None: ...
|
||||
|
||||
def exchange_failure(
|
||||
self, *, call_type: ExchangeCallType, duration_seconds: float, error: ExchangeError
|
||||
) -> None: ...
|
||||
|
||||
def cache_hit(self) -> None: ...
|
||||
|
||||
|
||||
class SyncTokenPoster(Protocol):
|
||||
"""Returns the response for ANY status; never raises for status."""
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: ...
|
||||
|
|
@ -5,6 +5,7 @@ Utility functions for base LLM classes.
|
|||
import copy
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from openai.lib import _parsing, _pydantic
|
||||
|
|
@ -57,6 +58,22 @@ class BaseLLMModelInfo(ABC):
|
|||
"""
|
||||
return []
|
||||
|
||||
def discover_models(
|
||||
self, litellm_params: Mapping[str, object] | None = None
|
||||
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
"""
|
||||
Live model discovery for a configured deployment. Defaults to the api_key/api_base
|
||||
facade every provider already implements via ``get_models``; a provider whose
|
||||
discovery needs more of ``litellm_params`` (e.g. Anthropic's workload identity
|
||||
federation) overrides this instead of widening ``get_models`` for every provider.
|
||||
"""
|
||||
api_key: Final = litellm_params.get("api_key") if litellm_params is not None else None
|
||||
api_base: Final = litellm_params.get("api_base") if litellm_params is not None else None
|
||||
return self.get_models(
|
||||
api_key=api_key if isinstance(api_key, str) else None,
|
||||
api_base=api_base if isinstance(api_base, str) else None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_api_key(api_key: str | None = None) -> str | None:
|
||||
|
|
|
|||
|
|
@ -1289,7 +1289,7 @@ class HTTPHandler:
|
|||
self,
|
||||
url: str,
|
||||
params: dict | None = None,
|
||||
headers: dict | None = None,
|
||||
headers: Mapping[str, Any] | None = None,
|
||||
follow_redirects: bool | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -1,11 +1,25 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import ssl
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
get_type_hints,
|
||||
runtime_checkable,
|
||||
)
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
|
|
@ -219,6 +233,55 @@ class _MediaUploadKwargs(TypedDict, total=False):
|
|||
timeout: float | httpx.Timeout
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _AsyncFilesEnvironmentValidator(Protocol):
|
||||
async def avalidate_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
model: str,
|
||||
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict: ... # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
|
||||
|
||||
async def _avalidate_files_environment(
|
||||
provider_config: BaseFilesConfig | BaseBatchesConfig,
|
||||
*,
|
||||
headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
model: str,
|
||||
messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
api_key: str | None,
|
||||
) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides
|
||||
"""Await the provider's async credential hook when it has one (e.g. Anthropic's workload
|
||||
identity token exchange); otherwise offload the sync hook to a worker thread. Either way
|
||||
the caller, an async file handler, never blocks the event loop on it."""
|
||||
if isinstance(provider_config, _AsyncFilesEnvironmentValidator) and inspect.iscoroutinefunction(
|
||||
provider_config.avalidate_environment
|
||||
):
|
||||
return await provider_config.avalidate_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
)
|
||||
return await asyncio.to_thread(
|
||||
provider_config.validate_environment,
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
|
||||
def _google_genai_streaming_hidden_params(
|
||||
*,
|
||||
api_base: str,
|
||||
|
|
@ -2130,7 +2193,7 @@ class BaseLLMHTTPHandler:
|
|||
(
|
||||
headers,
|
||||
api_base,
|
||||
) = anthropic_messages_provider_config.validate_anthropic_messages_environment(
|
||||
) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment(
|
||||
headers=merged_headers or {},
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
@ -3520,6 +3583,19 @@ class BaseLLMHTTPHandler:
|
|||
"""
|
||||
Creates a file using Gemini's two-step upload process
|
||||
"""
|
||||
if _is_async:
|
||||
return self._avalidate_and_create_file(
|
||||
create_file_data=create_file_data,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# get config from model, custom llm provider
|
||||
headers = provider_config.validate_environment(
|
||||
api_key=api_key,
|
||||
|
|
@ -3549,18 +3625,6 @@ class BaseLLMHTTPHandler:
|
|||
optional_params={},
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
return self.async_create_file(
|
||||
transformed_request=transformed_request,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
sync_httpx_client = _get_httpx_client()
|
||||
else:
|
||||
|
|
@ -3688,6 +3752,52 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params_with_url,
|
||||
)
|
||||
|
||||
async def _avalidate_and_create_file(
|
||||
self,
|
||||
*,
|
||||
create_file_data: CreateFileRequest,
|
||||
litellm_params: dict, # mutable-ok: mirrors the create_file contract this dispatches for
|
||||
provider_config: BaseFilesConfig,
|
||||
headers: dict, # mutable-ok: mirrors the create_file contract this dispatches for
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> OpenAIFileObject:
|
||||
validated_headers: Final = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
headers=headers,
|
||||
model="",
|
||||
messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
)
|
||||
complete_api_base: Final = provider_config.get_complete_file_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model="",
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
litellm_params=litellm_params,
|
||||
data=create_file_data,
|
||||
)
|
||||
return await self.async_create_file(
|
||||
transformed_request=provider_config.transform_create_file_request(
|
||||
model="",
|
||||
create_file_data=create_file_data,
|
||||
litellm_params=litellm_params,
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
),
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=validated_headers,
|
||||
api_base=complete_api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def async_create_file(
|
||||
self,
|
||||
transformed_request: Union[bytes, str, dict, "TwoStepFileUploadConfig"],
|
||||
|
|
@ -3938,6 +4048,20 @@ class BaseLLMHTTPHandler:
|
|||
if model is None:
|
||||
raise ValueError("model is required for create_batch")
|
||||
|
||||
if _is_async:
|
||||
return self._avalidate_and_create_batch(
|
||||
create_batch_data=create_batch_data,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
model=model,
|
||||
)
|
||||
|
||||
headers = provider_config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers=headers,
|
||||
|
|
@ -3966,19 +4090,6 @@ class BaseLLMHTTPHandler:
|
|||
optional_params={},
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
return self.async_create_batch(
|
||||
transformed_request=transformed_request,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
create_batch_data=create_batch_data,
|
||||
)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
sync_httpx_client = _get_httpx_client()
|
||||
else:
|
||||
|
|
@ -4115,6 +4226,54 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
async def _avalidate_and_create_batch(
|
||||
self,
|
||||
*,
|
||||
create_batch_data: "CreateBatchRequest",
|
||||
litellm_params: dict, # mutable-ok: mirrors the create_batch contract this dispatches for
|
||||
provider_config: "BaseBatchesConfig",
|
||||
headers: dict, # mutable-ok: mirrors the create_batch contract this dispatches for
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
client: Union["HTTPHandler", "AsyncHTTPHandler"] | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
model: str,
|
||||
) -> "LiteLLMBatch":
|
||||
validated_headers: Final = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=[], # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
)
|
||||
complete_api_base: Final = provider_config.get_complete_batch_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
litellm_params=litellm_params,
|
||||
data=create_batch_data,
|
||||
)
|
||||
return await self.async_create_batch(
|
||||
transformed_request=provider_config.transform_create_batch_request(
|
||||
model=model,
|
||||
create_batch_data=create_batch_data,
|
||||
litellm_params=litellm_params,
|
||||
optional_params={}, # mutable-ok: fresh per call; the legacy files/batches contract types this mutable and may mutate it
|
||||
),
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=validated_headers,
|
||||
api_base=complete_api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
create_batch_data=create_batch_data,
|
||||
)
|
||||
|
||||
async def async_create_batch(
|
||||
self,
|
||||
transformed_request: bytes | str | dict,
|
||||
|
|
@ -4712,7 +4871,8 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
# Validate environment and get headers
|
||||
headers = provider_config.validate_environment(
|
||||
headers = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
api_key=litellm_params.get("api_key"),
|
||||
headers=headers,
|
||||
model="",
|
||||
|
|
@ -4836,7 +4996,8 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
# Validate environment and get headers
|
||||
headers = provider_config.validate_environment(
|
||||
headers = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
api_key=litellm_params.get("api_key"),
|
||||
headers=headers,
|
||||
model="",
|
||||
|
|
@ -4960,7 +5121,8 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
# Validate environment and get headers
|
||||
headers = provider_config.validate_environment(
|
||||
headers = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
api_key=litellm_params.get("api_key"),
|
||||
headers=headers,
|
||||
model="",
|
||||
|
|
@ -5091,7 +5253,8 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
# Validate environment and get headers
|
||||
headers = provider_config.validate_environment(
|
||||
headers = await _avalidate_files_environment(
|
||||
provider_config,
|
||||
api_key=litellm_params.get("api_key"),
|
||||
headers=headers,
|
||||
model="",
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from litellm.types.utils import (
|
|||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
from ..common_utils import OpenAIError
|
||||
from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
|
@ -70,6 +71,11 @@ else:
|
|||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
|
||||
value: Final = litellm_params.get(key) if litellm_params is not None else None
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
||||
"""
|
||||
Reference: https://platform.openai.com/docs/api-reference/chat/create
|
||||
|
|
@ -747,28 +753,39 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
"""
|
||||
Calls OpenAI's `/v1/models` endpoint and returns the list of models.
|
||||
"""
|
||||
|
||||
if api_base is None:
|
||||
api_base = "https://api.openai.com"
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("OPENAI_API_KEY")
|
||||
|
||||
# Strip api_base to just the base URL (scheme + host + port)
|
||||
parsed_url: Final = httpx.URL(api_base)
|
||||
base_url = f"{parsed_url.scheme}://{parsed_url.host}"
|
||||
if parsed_url.port:
|
||||
base_url += f":{parsed_url.port}"
|
||||
|
||||
response: Final = litellm.module_level_client.get(
|
||||
url=f"{base_url}/v1/models",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
return self._fetch_model_ids(
|
||||
api_base=api_base, bearer_token=get_secret_str("OPENAI_API_KEY") if api_key is None else api_key
|
||||
)
|
||||
|
||||
def discover_models(
|
||||
self, litellm_params: Mapping[str, object] | None = None
|
||||
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
if type(self) is not OpenAIGPTConfig:
|
||||
return super().discover_models(litellm_params)
|
||||
api_key: Final = _litellm_params_str(litellm_params, "api_key")
|
||||
api_base: Final = _litellm_params_str(litellm_params, "api_base")
|
||||
workload_identity_config: Final = resolve_openai_workload_identity_config(
|
||||
api_key=api_key, api_base=api_base, litellm_params=litellm_params
|
||||
)
|
||||
if workload_identity_config is None:
|
||||
return self.get_models(api_key=api_key, api_base=api_base)
|
||||
return self._fetch_model_ids(
|
||||
api_base=api_base, bearer_token=get_workload_identity_bearer_token(workload_identity_config)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fetch_model_ids(
|
||||
api_base: str | None, bearer_token: str | None
|
||||
) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override
|
||||
parsed_url: Final = httpx.URL(api_base or "https://api.openai.com")
|
||||
port_suffix: Final = f":{parsed_url.port}" if parsed_url.port else ""
|
||||
response: Final = litellm.module_level_client.get(
|
||||
url=f"{parsed_url.scheme}://{parsed_url.host}{port_suffix}/v1/models",
|
||||
headers={"Authorization": f"Bearer {bearer_token}"},
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to get models: {response.text}")
|
||||
|
||||
models: Final = response.json()["data"]
|
||||
return [model["id"] for model in models]
|
||||
return [model["id"] for model in response.json()["data"]]
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(api_key: str | None = None) -> str | None:
|
||||
|
|
|
|||
|
|
@ -382,8 +382,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
organization: str | None = None,
|
||||
client: OpenAI | AsyncOpenAI | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> OpenAI | AsyncOpenAI | None:
|
||||
workload_identity_config: Final = resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base)
|
||||
workload_identity_config: Final = resolve_openai_workload_identity_config(
|
||||
api_key=api_key, api_base=api_base, litellm_params=litellm_params
|
||||
)
|
||||
client_initialization_params: Final[dict] = locals()
|
||||
if client is None:
|
||||
if not isinstance(max_retries, int):
|
||||
|
|
@ -773,6 +776,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
max_retries=max_retries,
|
||||
organization=organization,
|
||||
stream_options=stream_options,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
else:
|
||||
if not isinstance(max_retries, int):
|
||||
|
|
@ -786,6 +790,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
max_retries=max_retries,
|
||||
organization=organization,
|
||||
client=client,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -927,6 +932,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
organization=organization,
|
||||
client=client,
|
||||
shared_session=shared_session,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -1022,6 +1028,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
max_retries=None,
|
||||
headers=None,
|
||||
stream_options: dict | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
):
|
||||
data["stream"] = True
|
||||
data.update(self.get_stream_options(stream_options=stream_options, api_base=api_base))
|
||||
|
|
@ -1035,6 +1042,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
max_retries=max_retries,
|
||||
organization=organization,
|
||||
client=client,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
@ -1107,6 +1115,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
organization=organization,
|
||||
client=client,
|
||||
shared_session=shared_session,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
@ -1241,6 +1250,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
client: AsyncOpenAI | None = None,
|
||||
max_retries=None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
):
|
||||
try:
|
||||
openai_aclient: Final[AsyncOpenAI] = self._get_openai_client(
|
||||
|
|
@ -1251,6 +1261,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
max_retries=max_retries,
|
||||
client=client,
|
||||
shared_session=shared_session,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
raw_response: Final = await self.make_openai_embedding_request(
|
||||
openai_aclient=openai_aclient,
|
||||
|
|
@ -1314,6 +1325,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
aembedding=None,
|
||||
max_retries: int | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
super().embedding()
|
||||
try:
|
||||
|
|
@ -1340,6 +1352,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
client=client,
|
||||
max_retries=max_retries,
|
||||
shared_session=shared_session,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
openai_client: Final[OpenAI] = self._get_openai_client(
|
||||
|
|
@ -1349,6 +1362,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
client=client,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
## embedding CALL
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
_safe_convert_created_field,
|
||||
)
|
||||
|
|
@ -494,7 +495,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
|
||||
headers.setdefault("Content-Type", "application/json")
|
||||
workload_identity_config: Final = (
|
||||
resolve_openai_workload_identity_config(api_key=api_key, api_base=litellm_params.api_base)
|
||||
resolve_openai_workload_identity_config(
|
||||
api_key=api_key,
|
||||
api_base=litellm_params.api_base,
|
||||
litellm_params=litellm_params.model_dump(include=set(OPENAI_WIF_KWARGS_KEYS)),
|
||||
)
|
||||
if self.custom_llm_provider is LlmProviders.OPENAI
|
||||
else None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
|
@ -43,6 +44,7 @@ class OpenAIWorkloadIdentityConfig:
|
|||
def resolve_openai_workload_identity_config(
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> OpenAIWorkloadIdentityConfig | None:
|
||||
static_api_key: Final = normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str(
|
||||
get_secret_str("OPENAI_API_KEY")
|
||||
|
|
@ -54,10 +56,12 @@ def resolve_openai_workload_identity_config(
|
|||
)
|
||||
if not _targets_openai_api(effective_api_base):
|
||||
return None
|
||||
identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID")
|
||||
service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID")
|
||||
token_file: Final = get_secret_str("OPENAI_IDENTITY_TOKEN_FILE")
|
||||
if not identity_provider_id or not service_account_id or not token_file:
|
||||
identity_provider_id: Final = _config_value(
|
||||
litellm_params, "openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID"
|
||||
)
|
||||
service_account_id: Final = _config_value(litellm_params, "openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID")
|
||||
token_file: Final = _config_value(litellm_params, "openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE")
|
||||
if identity_provider_id is None or service_account_id is None or token_file is None:
|
||||
return None
|
||||
return OpenAIWorkloadIdentityConfig(
|
||||
identity_provider_id=identity_provider_id,
|
||||
|
|
@ -70,6 +74,13 @@ def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) ->
|
|||
return _workload_identity_auth(config).get_token()
|
||||
|
||||
|
||||
def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None:
|
||||
param_value: Final = litellm_params.get(param_key) if litellm_params is not None else None
|
||||
if isinstance(param_value, str) and param_value:
|
||||
return param_value
|
||||
return normalize_nonempty_secret_str(get_secret_str(env_name))
|
||||
|
||||
|
||||
def _targets_openai_api(api_base: str | None) -> bool:
|
||||
if api_base is None:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -6326,6 +6326,7 @@ def embedding(
|
|||
aembedding=aembedding,
|
||||
max_retries=max_retries,
|
||||
shared_session=shared_session,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
elif custom_llm_provider == "databricks":
|
||||
api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE")
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ These are the canonical credential types for the proxy. They live in the model
|
|||
layer; ``litellm.types.utils`` re-exports them for backwards compatibility.
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, model_validator
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
|
||||
class CredentialBase(BaseModel):
|
||||
|
|
@ -15,6 +15,10 @@ class CredentialBase(BaseModel):
|
|||
|
||||
class CredentialItem(CredentialBase):
|
||||
credential_values: dict
|
||||
# PATCH-only instruction naming keys to drop from the stored credential_values. It describes an
|
||||
# edit rather than the credential, so it stays out of dumps: those feed config loading, the DB
|
||||
# write, and the in-memory list, none of which have a place for it.
|
||||
credential_values_to_delete: tuple[str, ...] | None = Field(default=None, exclude=True)
|
||||
|
||||
|
||||
class CreateCredentialItem(CredentialBase):
|
||||
|
|
|
|||
|
|
@ -33,7 +33,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
|
||||
)
|
||||
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams
|
||||
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams, server_owned_wif_litellm_params
|
||||
|
||||
|
||||
def is_invalid_virtual_key_error(exception: BaseException | None) -> bool:
|
||||
|
|
@ -226,6 +227,16 @@ def _allow_model_level_clientside_configurable_parameters(
|
|||
# ``extra_body.aws_web_identity_token``) without re-validating, so the
|
||||
# banned-key check has to descend into it the same way it descends into
|
||||
# ``litellm_embedding_config``.
|
||||
_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED: Final[tuple[str, ...]] = server_owned_wif_litellm_params
|
||||
# The Bedrock Claude Platform route reads a workspace from workspace_id or aws_workspace_id as
|
||||
# well, and neither is a federation parameter, so say so rather than leaving that caller stuck.
|
||||
|
||||
|
||||
# Re-exported from litellm.types.router, where it lives so the router can call it on a
|
||||
# post-authentication merge without core importing from the proxy package.
|
||||
reject_server_owned_wif_params = _reject_server_owned_wif_params
|
||||
|
||||
|
||||
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
|
||||
|
||||
# Metadata containers that carry per-request configuration consumed by the
|
||||
|
|
@ -379,6 +390,7 @@ def _check_banned_params(
|
|||
Shared between the root-level check and the nested-config check so a
|
||||
new banned param only needs to be added in one place.
|
||||
"""
|
||||
reject_server_owned_wif_params(body)
|
||||
for param in _BANNED_REQUEST_BODY_PARAMS:
|
||||
if param not in body:
|
||||
continue
|
||||
|
|
@ -519,6 +531,7 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router:
|
|||
_reject_url_valued_fallback_target(target)
|
||||
litellm_params: Final = _coerce_metadata_to_dict(request_body.get("litellm_params"))
|
||||
if litellm_params is not None:
|
||||
reject_server_owned_wif_params(litellm_params)
|
||||
litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
|
||||
if litellm_params_metadata is not None:
|
||||
_check_banned_params(
|
||||
|
|
|
|||
158
litellm/proxy/common_utils/credential_hydration.py
Normal file
158
litellm/proxy/common_utils/credential_hydration.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""Shared helper for resolving a named Credential's values server-side.
|
||||
|
||||
Memory first (``litellm.credential_list``, already decrypted -- matching
|
||||
``CredentialAccessor.get_credential_values``), then a DB decrypt fallback for a pod whose
|
||||
in-memory list has not yet picked up a credential another pod just wrote or updated.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.types.router import (
|
||||
GenericLiteLLMParams,
|
||||
server_owned_wif_fields_named,
|
||||
server_owned_wif_fields_present,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem, LlmProviders
|
||||
|
||||
_LITELLM_PROVIDER_IDS: Final = frozenset(provider.value for provider in LlmProviders)
|
||||
|
||||
|
||||
def stored_credential_provider(credential_provider: object) -> str | None:
|
||||
"""The dashboard stores its display casing (``Anthropic``) on credentials it creates, so the
|
||||
provider a credential names is the lowercased value when that is a litellm provider id."""
|
||||
if not isinstance(credential_provider, str):
|
||||
return None
|
||||
lowered: Final = credential_provider.lower()
|
||||
return lowered if lowered in _LITELLM_PROVIDER_IDS else None
|
||||
|
||||
|
||||
def decrypted_or_stored(key: str, value: str) -> str:
|
||||
"""The stored value decrypted, or as stored when it was never encrypted (a config.yaml value)."""
|
||||
decrypted: Final = decrypt_value_helper(value=value, key=key)
|
||||
return value if decrypted is None else decrypted
|
||||
|
||||
|
||||
def _decrypted(db_credential: CredentialItem) -> CredentialItem:
|
||||
"""The stored credential with every value decrypted, leaving already-plaintext values alone."""
|
||||
decrypted_values: Final = MappingProxyType(
|
||||
{key: decrypted_or_stored(key, value) for key, value in db_credential.credential_values.items()}
|
||||
)
|
||||
return CredentialItem(
|
||||
credential_name=db_credential.credential_name,
|
||||
credential_values=decrypted_values, # pyright: ignore[reportArgumentType] # declared dict[str, str], and pydantic copies this mapping into one on validation; LIT002 rules out building that dict here
|
||||
credential_info=db_credential.credential_info,
|
||||
)
|
||||
|
||||
|
||||
async def hydrate_named_credential_authoritative(
|
||||
credential_name: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> CredentialItem | None:
|
||||
"""The stored credential, preferring the row over this pod's in-memory copy.
|
||||
|
||||
``hydrate_named_credential`` reads memory first, which is right when serving a request. A
|
||||
management operation cannot: on a pod whose in-memory copy predates another pod's update, it
|
||||
would export the superseded JWKS, or discover models against superseded values. Same reason
|
||||
``named_credential_wif_fields`` reads both.
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return await hydrate_named_credential(credential_name, prisma_client)
|
||||
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
|
||||
if db_credential is None:
|
||||
return await hydrate_named_credential(credential_name, prisma_client)
|
||||
return _decrypted(db_credential)
|
||||
|
||||
|
||||
async def hydrate_named_credential(
|
||||
credential_name: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> CredentialItem | None:
|
||||
for credential in litellm.credential_list:
|
||||
if credential.credential_name == credential_name:
|
||||
return credential
|
||||
if prisma_client is None:
|
||||
return None
|
||||
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
|
||||
if db_credential is None:
|
||||
return None
|
||||
return _decrypted(db_credential)
|
||||
|
||||
|
||||
async def named_credential_wif_fields(
|
||||
credential_name: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Federation field names a write to ``credential_name`` would touch, from memory AND the row.
|
||||
|
||||
Resolution reads memory first and stops there, which is right when serving a request. An
|
||||
authorization decision cannot: a pod whose in-memory copy predates an admin adding federation
|
||||
fields would see none and allow the write. This reads both and returns the union, so the gate
|
||||
refuses whenever either side says the credential is server-owned.
|
||||
"""
|
||||
in_memory: Final = tuple(
|
||||
name
|
||||
for credential in litellm.credential_list
|
||||
if credential.credential_name == credential_name
|
||||
for name in server_owned_wif_fields_named(credential.credential_values)
|
||||
)
|
||||
if prisma_client is None:
|
||||
return in_memory
|
||||
db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name)
|
||||
stored: Final = () if db_credential is None else server_owned_wif_fields_named(db_credential.credential_values)
|
||||
return tuple(dict.fromkeys(in_memory + stored))
|
||||
|
||||
|
||||
async def effective_server_owned_wif_fields(
|
||||
stored: Mapping[str, object] | None,
|
||||
incoming: GenericLiteLLMParams | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Federation field names the deployment would carry AFTER this write.
|
||||
|
||||
Authorization has to read the resulting deployment, not the submitted payload. A patch that
|
||||
names no federation field still lands on a deployment that has them, and a patch that only
|
||||
attaches ``litellm_credential_name`` inherits whatever that credential holds.
|
||||
|
||||
The two sides are matched differently on purpose. ``stored`` is matched by VALUE, because
|
||||
``GenericLiteLLMParams`` declares every federation field, so matching it by key would
|
||||
report every deployment on the proxy as federated. ``incoming`` is matched by the keys the
|
||||
write actually set, so an explicit null still counts as touching the field.
|
||||
"""
|
||||
from_stored: Final = () if stored is None else server_owned_wif_fields_present(stored)
|
||||
from_incoming: Final = () if incoming is None else server_owned_wif_fields_named(incoming.model_fields_set)
|
||||
from_credential: Final = tuple(
|
||||
chain.from_iterable(
|
||||
await asyncio.gather(
|
||||
*(
|
||||
named_credential_wif_fields(credential_name, prisma_client)
|
||||
for credential_name in _effective_credential_names(stored, incoming)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
return tuple(dict.fromkeys(from_stored + from_incoming + from_credential))
|
||||
|
||||
|
||||
def _effective_credential_names(
|
||||
stored: Mapping[str, object] | None,
|
||||
incoming: GenericLiteLLMParams | None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Both the credential the deployment already carries and the one this write names.
|
||||
|
||||
Taking only the incoming name would let a write clear its way out: detaching a federated
|
||||
credential, by sending ``litellm_credential_name: null`` alongside an api_key or api_base of
|
||||
the caller's choosing, would leave nothing federated to find and the write would be allowed.
|
||||
Detaching an administrator's federated credential is itself an administrator's action, so the
|
||||
stored name counts whatever the write says.
|
||||
"""
|
||||
from_stored: Final = None if stored is None else stored.get("litellm_credential_name")
|
||||
from_incoming: Final = None if incoming is None else incoming.litellm_credential_name
|
||||
return tuple(dict.fromkeys(name for name in (from_stored, from_incoming) if isinstance(name, str)))
|
||||
|
|
@ -13,16 +13,111 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.llms.anthropic.wif import (
|
||||
_IDENTITY_SOURCE_PARAM, # pyright: ignore[reportPrivateUsage] # one canonical param name, shared with the litellm_params identity-source resolver
|
||||
_INTERNAL_ISSUER_FIELD_MAP, # pyright: ignore[reportPrivateUsage] # one canonical field map, shared with the litellm_params identity-source resolver
|
||||
_build_variant, # pyright: ignore[reportPrivateUsage] # one canonical builder, shared with the litellm_params identity-source resolver
|
||||
)
|
||||
from litellm.llms.base_llm.auth.identity_source import (
|
||||
AnthropicIdentitySourceKind,
|
||||
InternalIssuerSource,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_jwks_document
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.credential_hydration import (
|
||||
hydrate_named_credential,
|
||||
hydrate_named_credential_authoritative,
|
||||
named_credential_wif_fields,
|
||||
stored_credential_provider,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.types.router import server_owned_wif_fields_named
|
||||
from litellm.types.utils import CreateCredentialItem, CredentialItem
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def _reject_non_admin_wif_fields(
|
||||
wif_fields: tuple[str, ...],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""A credential referenced by ``litellm_credential_name`` feeds its values into the same
|
||||
workload identity federation resolution as a deployment's own ``litellm_params``. Only proxy
|
||||
admins may touch a server-owned WIF field, whether they write it, drop it, or edit a stored
|
||||
credential that already carries one.
|
||||
"""
|
||||
if not wif_fields or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
|
||||
"error": (
|
||||
f"Only proxy admins can change {wif_fields[0]!r}, a server-owned workload identity federation "
|
||||
"parameter."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _incoming_wif_fields(credential: CredentialItem) -> tuple[str, ...]:
|
||||
"""WIF fields the request payload itself touches: the ones it sets (to any value, ``None``
|
||||
included, since the key alone is what the federation resolver reacts to), plus the ones it
|
||||
names in ``credential_values_to_delete``, since dropping a federation field off the stored
|
||||
credential breaks every deployment referencing it just as installing one would redirect them.
|
||||
"""
|
||||
return server_owned_wif_fields_named(credential.credential_values) + server_owned_wif_fields_named(
|
||||
credential.credential_values_to_delete or ()
|
||||
)
|
||||
|
||||
|
||||
def _stored_wif_fields(stored_credential: CredentialItem) -> tuple[str, ...]:
|
||||
return server_owned_wif_fields_named(stored_credential.credential_values)
|
||||
|
||||
|
||||
def _reject_overlapping_credential_values(credential: CredentialItem) -> None:
|
||||
overlap: Final = frozenset(credential.credential_values) & frozenset(credential.credential_values_to_delete or ())
|
||||
if overlap:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"credential_values_to_delete overlaps credential_values for key(s): {sorted(overlap)}",
|
||||
)
|
||||
|
||||
|
||||
def _sync_in_memory_credential(credential: CredentialItem, credential_name: str, new_name: str) -> None:
|
||||
"""Mirror a DB credential update into the in-memory ``credential_list`` used by request-time
|
||||
resolution; a no-op if the credential isn't loaded in memory (e.g. proxy restarted since boot).
|
||||
"""
|
||||
existing_in_memory: CredentialItem | None = None
|
||||
for cred in litellm.credential_list:
|
||||
if cred.credential_name == credential_name:
|
||||
existing_in_memory = cred
|
||||
break
|
||||
|
||||
if existing_in_memory is None:
|
||||
return
|
||||
|
||||
in_memory_values: Final = dict(existing_in_memory.credential_values or {})
|
||||
if credential.credential_values:
|
||||
in_memory_values.update(credential.credential_values)
|
||||
for key in credential.credential_values_to_delete or ():
|
||||
in_memory_values.pop(key, None)
|
||||
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
|
||||
if credential.credential_info:
|
||||
in_memory_info.update(credential.credential_info)
|
||||
updated_in_memory: Final = CredentialItem(
|
||||
credential_name=new_name,
|
||||
credential_values=in_memory_values,
|
||||
credential_info=in_memory_info,
|
||||
)
|
||||
# Remove old entry if renamed, then use upsert_credentials to handle duplicates
|
||||
if new_name != credential_name:
|
||||
litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name]
|
||||
CredentialAccessor.upsert_credentials([updated_in_memory])
|
||||
|
||||
|
||||
class CredentialHelperUtils:
|
||||
@staticmethod
|
||||
def encrypt_credential_values(credential: CredentialItem, new_encryption_key: str | None = None) -> CredentialItem:
|
||||
|
|
@ -84,13 +179,19 @@ async def create_credential(
|
|||
status_code=400,
|
||||
detail="Credential values are required. Unable to infer credential values from model ID.",
|
||||
)
|
||||
_reject_non_admin_wif_fields(server_owned_wif_fields_named(credential.credential_values), user_api_key_dict)
|
||||
_reject_non_admin_wif_fields(
|
||||
await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict
|
||||
)
|
||||
processed_credential: Final = CredentialItem(
|
||||
credential_name=credential.credential_name,
|
||||
credential_values=credential.credential_values,
|
||||
credential_info=credential.credential_info,
|
||||
)
|
||||
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
|
||||
credentials_dict: Final = encrypted_credential.model_dump()
|
||||
# exclude_none: wif.py rejects foreign-variant fields by presence, so persisting a null
|
||||
# for every unset variant field would fail the next request against this credential
|
||||
credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True)
|
||||
credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
|
||||
"dict[str, object]", jsonify_object(credentials_dict)
|
||||
)
|
||||
|
|
@ -175,6 +276,75 @@ async def get_credential_by_name(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/credentials/{credential_name:path}/jwks",
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
tags=["credential management"], # mutable-ok: FastAPI's include_router does self.tags.copy(), needs a real list
|
||||
)
|
||||
async def get_credential_internal_issuer_jwks(
|
||||
credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default
|
||||
):
|
||||
"""
|
||||
Export the public JWKS for an anthropic ``internal_issuer`` credential, so the operator can
|
||||
register it on the Anthropic federation issuer from the UI. Never touches the private signing
|
||||
key: only its derived public JWKS leaves this process. 404s for any other credential shape.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
|
||||
"error": "Only proxy admins can export a credential's JWKS."
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
credential: Final = await hydrate_named_credential_authoritative(credential_name, prisma_client)
|
||||
credential_provider: Final = (
|
||||
None
|
||||
if credential is None
|
||||
else stored_credential_provider(credential.credential_info.get("custom_llm_provider"))
|
||||
)
|
||||
if credential is None or credential_provider != "anthropic":
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
|
||||
"error": f"No anthropic credential named {credential_name!r}."
|
||||
},
|
||||
)
|
||||
configured_source: Final = credential.credential_values.get(_IDENTITY_SOURCE_PARAM)
|
||||
if configured_source != AnthropicIdentitySourceKind.internal_issuer.value:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
|
||||
"error": (
|
||||
f"Credential {credential_name!r} is not configured with "
|
||||
f"{_IDENTITY_SOURCE_PARAM}={AnthropicIdentitySourceKind.internal_issuer.value!r}."
|
||||
)
|
||||
},
|
||||
)
|
||||
try:
|
||||
issuer_source: Final = _build_variant(
|
||||
InternalIssuerSource, credential.credential_values, _INTERNAL_ISSUER_FIELD_MAP
|
||||
)
|
||||
jwks_document: Final = internal_issuer_jwks_document(issuer_source)
|
||||
except (litellm.AuthenticationError, ValueError) as e:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={ # mutable-ok: starlette json.dumps()s HTTPException.detail raw, needs a real dict
|
||||
"error": str(e)
|
||||
},
|
||||
) from e
|
||||
return Response(content=jwks_document, media_type="application/json")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy's error contract
|
||||
verbose_proxy_logger.exception(e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/credentials/by_model/{model_id}",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -239,6 +409,9 @@ async def delete_credential(
|
|||
status_code=500,
|
||||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||||
)
|
||||
_reject_non_admin_wif_fields(
|
||||
await named_credential_wif_fields(credential_name, prisma_client), user_api_key_dict
|
||||
)
|
||||
deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name)
|
||||
if deleted is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -249,6 +422,8 @@ async def delete_credential(
|
|||
## DELETE FROM LITELLM ##
|
||||
litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name]
|
||||
return {"success": True, "message": "Credential deleted successfully"}
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
|
@ -282,6 +457,9 @@ def update_db_credential(
|
|||
|
||||
merged_credential.credential_values.update(encrypted_params)
|
||||
|
||||
for key in updated_patch.credential_values_to_delete or ():
|
||||
merged_credential.credential_values.pop(key, None)
|
||||
|
||||
# update model info
|
||||
if encrypted_credential.credential_info:
|
||||
"""Update credential info"""
|
||||
|
|
@ -310,6 +488,8 @@ async def update_credential(
|
|||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
try:
|
||||
_reject_overlapping_credential_values(credential)
|
||||
_reject_non_admin_wif_fields(_incoming_wif_fields(credential), user_api_key_dict)
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
@ -319,9 +499,14 @@ async def update_credential(
|
|||
db_credential: Final = await credentials_repository.find_by_name(credential_name)
|
||||
if db_credential is None:
|
||||
raise HTTPException(status_code=404, detail="Credential not found in DB.")
|
||||
_reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict)
|
||||
if credential.credential_name != credential_name:
|
||||
shadowed_credential: Final = await hydrate_named_credential(credential.credential_name, prisma_client)
|
||||
if shadowed_credential is not None:
|
||||
_reject_non_admin_wif_fields(_stored_wif_fields(shadowed_credential), user_api_key_dict)
|
||||
merged_credential: Final = update_db_credential(db_credential, credential)
|
||||
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
|
||||
"dict[str, object]", jsonify_object(merged_credential.model_dump())
|
||||
"dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True))
|
||||
)
|
||||
await credentials_repository.update_by_name(
|
||||
credential_name,
|
||||
|
|
@ -332,29 +517,7 @@ async def update_credential(
|
|||
)
|
||||
|
||||
# Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted)
|
||||
new_name: Final = merged_credential.credential_name
|
||||
existing_in_memory: CredentialItem | None = None
|
||||
for cred in litellm.credential_list:
|
||||
if cred.credential_name == credential_name:
|
||||
existing_in_memory = cred
|
||||
break
|
||||
|
||||
if existing_in_memory is not None:
|
||||
in_memory_values: Final = dict(existing_in_memory.credential_values or {})
|
||||
if credential.credential_values:
|
||||
in_memory_values.update(credential.credential_values)
|
||||
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
|
||||
if credential.credential_info:
|
||||
in_memory_info.update(credential.credential_info)
|
||||
updated_in_memory: Final = CredentialItem(
|
||||
credential_name=new_name,
|
||||
credential_values=in_memory_values,
|
||||
credential_info=in_memory_info,
|
||||
)
|
||||
# Remove old entry if renamed, then use upsert_credentials to handle duplicates
|
||||
if new_name != credential_name:
|
||||
litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name]
|
||||
CredentialAccessor.upsert_credentials([updated_in_memory])
|
||||
_sync_in_memory_credential(credential, credential_name, merged_credential.credential_name)
|
||||
|
||||
return {"success": True, "message": "Credential updated successfully"}
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check
|
||||
reject_server_owned_wif_params,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
|
|
@ -2065,6 +2066,7 @@ async def test_model_connection(
|
|||
"Could not find model %s in router: %s. Proceeding with request params only.", model_name, e
|
||||
)
|
||||
|
||||
reject_server_owned_wif_params(request_litellm_params)
|
||||
# Merge: config params (from proxy config) as base, request params override
|
||||
litellm_params = {
|
||||
**_config_base_for_health_check(
|
||||
|
|
@ -2091,6 +2093,9 @@ async def test_model_connection(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
# The Deployment above already carries the caller's merged params, so the effective
|
||||
# state is model_params itself; there is no separate incoming patch here.
|
||||
incoming_params=None,
|
||||
)
|
||||
mode = mode or litellm_params.pop("mode", None)
|
||||
|
||||
|
|
|
|||
|
|
@ -56,6 +56,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
|
|||
coordination_redis_cache,
|
||||
publish_config_change,
|
||||
)
|
||||
from litellm.proxy.common_utils.credential_hydration import (
|
||||
effective_server_owned_wif_fields,
|
||||
hydrate_named_credential,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
|
|
@ -289,6 +293,23 @@ def _raise_on_strategy_router_write_violation(
|
|||
)
|
||||
|
||||
|
||||
def _reject_non_admin_blocked_flag_on_create(
|
||||
blocked: bool | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""Same proxy-admin-only rule patch_model applies to the blocked flag: a team admin passed
|
||||
the team-scoped auth check above, but must not be able to create a model already paused
|
||||
(or explicitly unpaused) out from under the proxy admin.
|
||||
"""
|
||||
if blocked is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise ProxyException(
|
||||
message="Only proxy admins can set a model's blocked flag.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param="blocked",
|
||||
)
|
||||
|
||||
|
||||
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301
|
||||
_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
|
||||
_STORED_LITELLM_PARAMS_SQL: Final = (
|
||||
|
|
@ -831,6 +852,7 @@ async def patch_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
incoming_params=patch_data.litellm_params,
|
||||
)
|
||||
|
||||
# Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins
|
||||
|
|
@ -1141,6 +1163,8 @@ async def _add_model_to_db(
|
|||
}
|
||||
if model_params.model_info.id is not None:
|
||||
_data["model_id"] = model_params.model_info.id
|
||||
if model_params.blocked is not None:
|
||||
_data["blocked"] = model_params.blocked
|
||||
_create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above
|
||||
if not should_create_model_in_db:
|
||||
return LiteLLM_ProxyModelTable(**_data)
|
||||
|
|
@ -1691,14 +1715,65 @@ class ModelManagementAuthChecks:
|
|||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
async def _reject_non_admin_wif_write(
|
||||
*,
|
||||
model_params: Deployment,
|
||||
incoming_params: GenericLiteLLMParams | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
stored: Final = model_params.litellm_params.model_dump(exclude_none=True)
|
||||
wif_fields: Final = await effective_server_owned_wif_fields(stored, incoming_params, prisma_client)
|
||||
if wif_fields:
|
||||
# ProxyException rather than HTTPException so the offending field stays a structured
|
||||
# `param`, which is the contract the narrower gate this replaced already published.
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Only proxy admins can modify a deployment configured for workload identity "
|
||||
f"federation ({wif_fields[0]!r})."
|
||||
),
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param=wif_fields[0],
|
||||
)
|
||||
# A name the caller expects an admin to create later would resolve to nothing today and
|
||||
# start federating the moment it exists, so a non-admin may only attach one that is already there.
|
||||
if incoming_params is not None and "litellm_credential_name" in incoming_params.model_fields_set:
|
||||
named: Final = incoming_params.litellm_credential_name
|
||||
if isinstance(named, str) and await hydrate_named_credential(named, prisma_client) is None:
|
||||
raise ProxyException(
|
||||
message=f"No credential named {named!r} exists.",
|
||||
type=ProxyErrorTypes.bad_request_error.value,
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
param="litellm_credential_name",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def can_user_make_model_call(
|
||||
model_params: Deployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
premium_user: bool,
|
||||
*,
|
||||
incoming_params: GenericLiteLLMParams | None,
|
||||
allow_missing_team: bool = False,
|
||||
) -> Literal[True]:
|
||||
# Federation fields choose which server-side secret is read and where the org-scoped token
|
||||
# it buys is sent, so only a proxy admin may touch a deployment that has them. Evaluated on
|
||||
# the RESULTING deployment: a patch naming no federation field still lands on one that has
|
||||
# them, and a patch attaching a credential by name inherits whatever that credential holds.
|
||||
# `incoming_params` is keyword-only with no default so a new write path cannot typecheck
|
||||
# without deciding what it writes.
|
||||
await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
## Check team model auth
|
||||
if model_params.model_info is not None and model_params.model_info.team_id is not None:
|
||||
team_obj_row: Final = await _repo_team_table(prisma_client).find_unique(
|
||||
|
|
@ -1793,6 +1868,7 @@ async def delete_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
incoming_params=None,
|
||||
allow_missing_team=True,
|
||||
)
|
||||
|
||||
|
|
@ -1914,7 +1990,6 @@ async def delete_team_model_alias(
|
|||
return removed_model_aliases
|
||||
|
||||
|
||||
#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
|
||||
@router.post(
|
||||
"/model/new",
|
||||
description="Allows adding new models to the model list in the config.yaml",
|
||||
|
|
@ -1983,8 +2058,11 @@ async def add_new_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
incoming_params=model_params.litellm_params,
|
||||
)
|
||||
|
||||
_reject_non_admin_blocked_flag_on_create(model_params.blocked, user_api_key_dict)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2166,6 +2244,7 @@ async def update_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
incoming_params=model_params.litellm_params,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
|
|
|
|||
|
|
@ -32,7 +32,11 @@ from litellm.constants import (
|
|||
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
_SERVER_OWNED_AUTH_HEADERS, # pyright: ignore[reportPrivateUsage] # canonical set, must not be duplicated here
|
||||
AnthropicModelInfo,
|
||||
merge_anthropic_beta_headers,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
||||
|
|
@ -644,6 +648,57 @@ async def is_streaming_request_fn(request: Request) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _anthropic_passthrough_headers(auth_header: Mapping[str, str] | None, client_beta: str | None) -> Mapping[str, str]:
|
||||
"""Custom headers take priority over forwarded client headers, so merge the
|
||||
client's anthropic-beta into the auth header's instead of clobbering it."""
|
||||
if auth_header is None:
|
||||
return MappingProxyType({})
|
||||
auth_beta: Final = auth_header.get("anthropic-beta")
|
||||
if auth_beta is None or client_beta is None:
|
||||
return auth_header
|
||||
return MappingProxyType({**auth_header, "anthropic-beta": merge_anthropic_beta_headers(client_beta, auth_beta)})
|
||||
|
||||
|
||||
def _configured_litellm_key_header_name() -> str | None:
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings, # pyright: ignore[reportUnknownVariableType] # proxy_server declares it as a bare dict
|
||||
)
|
||||
|
||||
configured: Final = general_settings.get( # pyright: ignore[reportUnknownMemberType] # proxy_server general_settings is a bare dict
|
||||
"litellm_key_header_name"
|
||||
)
|
||||
return configured if isinstance(configured, str) else None
|
||||
|
||||
|
||||
def _anthropic_passthrough_header_plan(
|
||||
request: Request, auth_header: Mapping[str, str] | None, litellm_key_header_name: str | None
|
||||
) -> tuple[Mapping[str, str], bool]:
|
||||
"""Returns the headers to send upstream plus whether the relay should still forward the
|
||||
caller's own headers. Once the server owns the Anthropic credential, none of the headers
|
||||
the proxy accepts a LiteLLM key in (``SpecialHeaders`` plus the configured custom name)
|
||||
may ride upstream beside it, so the forward merge runs here with those stripped and the
|
||||
relay is told not to merge again. With no server credential the caller's key is the only
|
||||
one there is, so forwarding stays on (BYOK)."""
|
||||
server_headers: Final = _anthropic_passthrough_headers(auth_header, request.headers.get("anthropic-beta"))
|
||||
if auth_header is None:
|
||||
return server_headers, True
|
||||
caller_owned: Final = _SERVER_OWNED_AUTH_HEADERS | frozenset(
|
||||
(litellm_key_header_name.lower(),) if litellm_key_header_name else ()
|
||||
)
|
||||
caller_headers: Final[dict[str, str]] = { # mutable-ok: forward_headers_from_request takes a concrete dict
|
||||
name: value for name, value in request.headers.items() if name.lower() not in caller_owned
|
||||
}
|
||||
merged: Final = cast( # cast-ok: forward_headers_from_request is untyped upstream, its result is a header dict
|
||||
"dict[str, str]",
|
||||
HttpPassThroughEndpointHelpers.forward_headers_from_request( # pyright: ignore[reportUnknownMemberType] # untyped upstream
|
||||
request_headers=caller_headers,
|
||||
headers=dict(server_headers), # mutable-ok: forward_headers_from_request takes a concrete dict
|
||||
forward_headers=True,
|
||||
),
|
||||
)
|
||||
return MappingProxyType(merged), False
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/anthropic/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
@ -682,12 +737,17 @@ async def anthropic_proxy_route(
|
|||
is_streaming_request: Final = await is_streaming_request_fn(request)
|
||||
|
||||
## CREATE PASS-THROUGH
|
||||
auth_header: Final = AnthropicModelInfo.get_auth_header(anthropic_api_key or None)
|
||||
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
|
||||
anthropic_api_key or None, allow_workload_identity=True
|
||||
)
|
||||
upstream_headers, forward_caller_headers = _anthropic_passthrough_header_plan(
|
||||
request, auth_header, _configured_litellm_key_header_name()
|
||||
)
|
||||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=endpoint,
|
||||
target=str(updated_url),
|
||||
custom_headers=auth_header if auth_header is not None else {},
|
||||
_forward_headers=True,
|
||||
custom_headers=upstream_headers,
|
||||
_forward_headers=forward_caller_headers,
|
||||
is_streaming_request=is_streaming_request,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
received_value: Final = await endpoint_func(
|
||||
|
|
|
|||
|
|
@ -336,6 +336,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
|||
)
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber
|
||||
from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored
|
||||
from litellm.proxy.common_utils.debug_utils import init_verbose_loggers
|
||||
from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
|
|
@ -719,6 +720,7 @@ from litellm.types.router import (
|
|||
RouterGeneralSettings,
|
||||
RoutingPlugin,
|
||||
SearchToolTypedDict,
|
||||
holds_secret_pointer,
|
||||
updateDeployment,
|
||||
)
|
||||
from litellm.types.router import ModelInfo as RouterModelInfo
|
||||
|
|
@ -4679,7 +4681,7 @@ class ProxyConfig:
|
|||
if isinstance(item, dict):
|
||||
item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth)
|
||||
# if the value is a string and starts with "os.environ/" - then it's an environment variable
|
||||
elif isinstance(value, str) and value.startswith("os.environ/"):
|
||||
elif isinstance(value, str) and value.startswith("os.environ/") and not holds_secret_pointer(key):
|
||||
resolved = get_secret(value)
|
||||
if resolved is None and secret_manager_would_be_consulted(value):
|
||||
verbose_proxy_logger.warning("%s is absent from the configured secret manager", value)
|
||||
|
|
@ -5777,7 +5779,7 @@ class ProxyConfig:
|
|||
for model in model_list:
|
||||
### LOAD FROM os.environ/ ###
|
||||
for k, v in model["litellm_params"].items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
|
||||
model["litellm_params"][k] = get_secret(v)
|
||||
validate_deployment_max_agentic_loops(model)
|
||||
validate_deployment_complexity_router_placement(model)
|
||||
|
|
@ -6170,7 +6172,7 @@ class ProxyConfig:
|
|||
for model in model_list:
|
||||
### LOAD FROM os.environ/ ###
|
||||
for k, v in model["litellm_params"].items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
|
||||
model["litellm_params"][k] = get_secret(v)
|
||||
|
||||
## check if they have model-id's ##
|
||||
|
|
@ -6198,7 +6200,11 @@ class ProxyConfig:
|
|||
return value
|
||||
|
||||
decrypted_value: Final = decrypt_value_helper(value=value, key=key, return_original_value=True)
|
||||
if isinstance(decrypted_value, str) and decrypted_value.startswith("os.environ/"):
|
||||
if (
|
||||
isinstance(decrypted_value, str)
|
||||
and decrypted_value.startswith("os.environ/")
|
||||
and not holds_secret_pointer(key)
|
||||
):
|
||||
return get_secret(decrypted_value)
|
||||
return decrypted_value
|
||||
|
||||
|
|
@ -7967,7 +7973,7 @@ class ProxyConfig:
|
|||
|
||||
decrypted_credential_values: Final = {}
|
||||
for k, v in credential_object.credential_values.items():
|
||||
decrypted_credential_values[k] = decrypt_value_helper(value=v, key=k) or v
|
||||
decrypted_credential_values[k] = decrypted_or_stored(k, v)
|
||||
|
||||
credential_object.credential_values = decrypted_credential_values
|
||||
return credential_object
|
||||
|
|
|
|||
|
|
@ -235,6 +235,7 @@ from litellm.types.router import (
|
|||
RoutingStrategy,
|
||||
SearchToolTypedDict,
|
||||
TaggedPreRoutingStrategy,
|
||||
holds_secret_pointer,
|
||||
)
|
||||
from litellm.types.services import ServiceTypes
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -8650,13 +8651,10 @@ class Router:
|
|||
if ptu_error is not None and is_ptu_cost_attribution_enabled():
|
||||
raise ValueError(ptu_error)
|
||||
zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(
|
||||
**(
|
||||
_litellm_params
|
||||
if zeroed_pricing is None
|
||||
else MappingProxyType({**_litellm_params, **zeroed_pricing})
|
||||
)
|
||||
merged_params: Final[Mapping[str, Any]] = (
|
||||
_litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing})
|
||||
)
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**merged_params)
|
||||
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params)
|
||||
deployment = Deployment(
|
||||
**deployment_info,
|
||||
|
|
@ -9226,7 +9224,7 @@ class Router:
|
|||
## check if litellm params in os.environ
|
||||
if isinstance(_litellm_params, dict):
|
||||
for k, v in _litellm_params.items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k):
|
||||
_litellm_params[k] = get_secret(v)
|
||||
|
||||
_model_info: dict = model.pop("model_info", {})
|
||||
|
|
|
|||
|
|
@ -13,8 +13,16 @@ Ensures cooldowns are applied correctly.
|
|||
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.utils import server_owned_wif_litellm_params
|
||||
|
||||
clientside_credential_keys: Final = ["api_key", "api_base", "base_url"]
|
||||
|
||||
# Set on a deployment whose api_base was client-redirected, so the Anthropic auth path refuses to
|
||||
# mint a federation token there even when WIF is configured only through ANTHROPIC_* env vars (which
|
||||
# cannot be cleared from litellm_params).
|
||||
DISABLE_WORKLOAD_IDENTITY_PARAM: Final = "anthropic_disable_workload_identity_federation"
|
||||
_WIF_CLEAR_ON_BASE_OVERRIDE: Final = tuple(sorted(server_owned_wif_litellm_params))
|
||||
|
||||
|
||||
def _admin_config_fields_to_clear_on_base_override() -> list[str]:
|
||||
"""
|
||||
|
|
@ -59,6 +67,14 @@ def _admin_config_fields_to_clear_on_base_override() -> list[str]:
|
|||
# ``api_base`` for the same reason as the OCI entries above.
|
||||
"nvcf_function_id",
|
||||
"use_ssl",
|
||||
# Workload-identity federation minting fields, restated here from
|
||||
# server_owned_wif_litellm_params the same way azure_ad_token above is restated
|
||||
# despite also being declared on CredentialLiteLLMParams (hence covered by
|
||||
# typed_fields too): a federation token minted for a client-redirected api_base
|
||||
# would send the workload's OIDC assertion, and then the minted bearer, to the
|
||||
# caller-chosen host, so this list must stay correct even if a field is ever
|
||||
# dropped from the typed model.
|
||||
*_WIF_CLEAR_ON_BASE_OVERRIDE,
|
||||
]
|
||||
return typed_fields + kwargs_only_fields
|
||||
|
||||
|
|
@ -101,5 +117,6 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
|
|||
litellm_params.pop(field, None)
|
||||
if field in request_kwargs:
|
||||
litellm_params[field] = request_kwargs[field]
|
||||
litellm_params[DISABLE_WORKLOAD_IDENTITY_PARAM] = True
|
||||
|
||||
return litellm_params
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
||||
increment_deployment_failures_for_current_minute,
|
||||
)
|
||||
from litellm.types.router import LiteLLMParamsTypedDict
|
||||
from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router as _Router
|
||||
|
|
@ -507,6 +507,14 @@ async def run_async_fallback(
|
|||
failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group
|
||||
attempted.record(failed_model_group)
|
||||
|
||||
# A dict target is merged straight into kwargs below, and kwargs win over the deployment's own
|
||||
# params, so a stored key/team/global fallback could otherwise set a federation field that the
|
||||
# request itself is forbidden to carry. Checked here rather than at the merge: inside the loop
|
||||
# the refusal would be caught as a per-target failure and quietly skipped to the next one.
|
||||
for target in fallback_model_group:
|
||||
if isinstance(target, dict):
|
||||
reject_server_owned_wif_params(target)
|
||||
|
||||
for mg in fallback_model_group:
|
||||
if mg == failed_model_group:
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -49,6 +49,10 @@ def _oidc_token_cache_ttl(oidc_token: str, max_ttl: int) -> int:
|
|||
_DEFAULT_OIDC_ALLOWED_CREDENTIAL_DIRS: Final = ("/var/run/secrets", "/run/secrets")
|
||||
|
||||
|
||||
class OidcPathNotAllowedError(ValueError):
|
||||
"""An ``oidc/file/`` path was rejected by the credential-directory allowlist."""
|
||||
|
||||
|
||||
def _get_oidc_allowed_credential_dirs() -> list[str]:
|
||||
"""
|
||||
Return the absolute, normalized list of directories from which
|
||||
|
|
@ -73,7 +77,7 @@ def _resolve_oidc_file_path(requested_path: str) -> str:
|
|||
credential directories. Raises ``ValueError`` otherwise.
|
||||
"""
|
||||
if not os.path.isabs(requested_path):
|
||||
raise ValueError(
|
||||
raise OidcPathNotAllowedError(
|
||||
"oidc/file path must be absolute. Use the format "
|
||||
"'oidc/file//var/run/secrets/<name>' (note the leading slash "
|
||||
"after 'oidc/file/')."
|
||||
|
|
@ -87,7 +91,7 @@ def _resolve_oidc_file_path(requested_path: str) -> str:
|
|||
# commonpath raises when paths are on different drives (Windows);
|
||||
# treat as not-matching and continue.
|
||||
continue
|
||||
raise ValueError(
|
||||
raise OidcPathNotAllowedError(
|
||||
"oidc/file path is outside the allowed credential directories. "
|
||||
"Set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -750,5 +750,6 @@ ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24"
|
|||
# OAuth constants
|
||||
ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat"
|
||||
ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20"
|
||||
ANTHROPIC_TOKEN_EXCHANGE_PATH: Final = "/v1/oauth/token"
|
||||
|
||||
ANTHROPIC_PROMPT_CACHING_SCOPE_BETA_HEADER: Final = "prompt-caching-scope-2026-01-05"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
|
|||
|
||||
import datetime
|
||||
import enum
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Container, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints
|
||||
|
||||
|
|
@ -27,6 +27,10 @@ from .utils import (
|
|||
ModelResponse,
|
||||
StandardLoggingRoutingDecision,
|
||||
)
|
||||
from .utils import (
|
||||
# private alias: `from .types.router import *` would rebind a public Final in litellm/__init__.py
|
||||
server_owned_wif_litellm_params as _server_owned_wif_litellm_params,
|
||||
)
|
||||
|
||||
|
||||
class ConfigurableClientsideParamsCustomAuth(TypedDict):
|
||||
|
|
@ -297,6 +301,69 @@ class CredentialLiteLLMParams(BaseModel):
|
|||
## IBM WATSONX ##
|
||||
watsonx_region_name: str | None = None
|
||||
|
||||
## ANTHROPIC WORKLOAD IDENTITY FEDERATION ##
|
||||
# Without these, get_deployment_credentials_with_provider silently drops a
|
||||
# litellm_params-configured WIF setup before files/batches/passthrough callers see
|
||||
# it, the same #30235-shaped gap azure_ad_token above was added to close.
|
||||
anthropic_federation_rule_id: str | None = None
|
||||
anthropic_organization_id: str | None = None
|
||||
anthropic_service_account_id: str | None = None
|
||||
anthropic_workspace_id: str | None = None
|
||||
anthropic_identity_token_file: str | None = None
|
||||
anthropic_identity_token: str | None = None
|
||||
anthropic_identity_source: str | None = None
|
||||
anthropic_issuer_url: str | None = None
|
||||
anthropic_issuer_subject: str | None = None
|
||||
anthropic_issuer_audience: str | None = None
|
||||
anthropic_issuer_ttl_seconds: int | None = None
|
||||
anthropic_issuer_signing_key_ref: str | None = None
|
||||
anthropic_keycloak_token_url: str | None = None
|
||||
anthropic_keycloak_client_id: str | None = None
|
||||
anthropic_keycloak_auth_method: str | None = None
|
||||
anthropic_keycloak_client_secret_ref: str | None = None
|
||||
anthropic_keycloak_scope: str | None = None
|
||||
# Server-set when a client redirects api_base. Declared so it survives the strict dump the
|
||||
# other federation fields above are declared for, rather than being rebuilt away in transit.
|
||||
anthropic_disable_workload_identity_federation: bool | None = None
|
||||
|
||||
## OPENAI WORKLOAD IDENTITY FEDERATION ##
|
||||
openai_identity_provider_id: str | None = None
|
||||
openai_service_account_id: str | None = None
|
||||
openai_identity_token_file: str | None = None
|
||||
|
||||
|
||||
def server_owned_wif_fields_present(fields: Mapping[str, object]) -> tuple[str, ...]:
|
||||
"""Server-owned workload identity federation field names set in ``fields``.
|
||||
|
||||
``fields`` is a ``litellm_params`` dict (or a credential's ``credential_values`` mapping,
|
||||
which feeds the same resolution when referenced by name). Derived from
|
||||
``server_owned_wif_litellm_params`` rather than hand-copied, so a persistence gate built on
|
||||
this stays correct when a new WIF field is added there.
|
||||
"""
|
||||
return tuple(name for name in _server_owned_wif_litellm_params if fields.get(name) is not None)
|
||||
|
||||
|
||||
def server_owned_wif_fields_named(keys: Container[str]) -> tuple[str, ...]:
|
||||
"""Server-owned workload identity federation field names that appear in ``keys``, whatever
|
||||
value they carry.
|
||||
|
||||
The write gates on credentials need this key-based sibling of ``server_owned_wif_fields_present``:
|
||||
``get_litellm_params`` forwards a WIF kwarg on key presence and the federation resolver rejects
|
||||
a foreign variant's field by key, so a persisted ``{"anthropic_issuer_url": None}`` wedges every
|
||||
deployment that references the credential even though no value is set. Pass a mapping (its keys
|
||||
are tested) or a plain collection of key names.
|
||||
"""
|
||||
return tuple(name for name in _server_owned_wif_litellm_params if name in keys)
|
||||
|
||||
|
||||
_WIF_POINTER_FIELDS: Final = frozenset(name for name in _server_owned_wif_litellm_params if name.endswith("_ref"))
|
||||
|
||||
|
||||
def holds_secret_pointer(param_name: str) -> bool:
|
||||
"""A ``*_ref`` federation field is a secret POINTER the identity source dereferences at use
|
||||
time, so a loader expanding ``os.environ/`` values must leave it as written."""
|
||||
return param_name in _WIF_POINTER_FIELDS
|
||||
|
||||
|
||||
_RESERVED_INIT_KEYS: Final = frozenset({"self", "params", "__class__"})
|
||||
|
||||
|
|
@ -542,6 +609,9 @@ class Deployment(BaseModel):
|
|||
model_name: str
|
||||
litellm_params: LiteLLM_Params
|
||||
model_info: ModelInfo
|
||||
# admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked. None means "don't set it
|
||||
# on create" -- the Prisma column defaults to False -- rather than "explicitly unblocked".
|
||||
blocked: bool | None = None
|
||||
|
||||
model_config = ConfigDict(extra="allow", protected_namespaces=())
|
||||
|
||||
|
|
@ -1064,3 +1134,23 @@ class AdaptiveRouterPreferences(BaseModel):
|
|||
|
||||
quality_tier: int = Field(ge=1, le=3)
|
||||
strengths: list[RequestType] = Field(default_factory=list)
|
||||
|
||||
|
||||
_BEDROCK_WORKSPACE_HINT: Final = " On the Bedrock Claude Platform route, pass workspace_id or aws_workspace_id instead."
|
||||
|
||||
|
||||
def reject_server_owned_wif_params(body: Mapping[str, object]) -> None:
|
||||
"""Raise ``ValueError`` if a mapping that did not come from deployment config carries a
|
||||
server-owned workload identity federation field.
|
||||
|
||||
These are never client-settable on any surface, with or without a client-side credential
|
||||
opt-in. This lives here rather than under ``litellm.proxy`` so the router can call it on a
|
||||
post-authentication merge without core importing from the proxy package.
|
||||
"""
|
||||
for param in _server_owned_wif_litellm_params:
|
||||
if param in body:
|
||||
raise ValueError(
|
||||
f"Rejected Request: {param} is a server-owned workload identity federation parameter "
|
||||
"and cannot be set in a request body; configure it on the deployment instead."
|
||||
+ (_BEDROCK_WORKSPACE_HINT if param == "anthropic_workspace_id" else "")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ class ServiceTypes(str, enum.Enum):
|
|||
AUTH = "auth"
|
||||
PROXY_PRE_CALL = "proxy_pre_call"
|
||||
POD_LOCK_MANAGER = "pod_lock_manager"
|
||||
ANTHROPIC_WIF = "anthropic_wif"
|
||||
ANTHROPIC_WIF_CACHE = "anthropic_wif_cache"
|
||||
|
||||
"""
|
||||
Operational metrics for DB Transaction Queues
|
||||
|
|
@ -67,6 +69,13 @@ DEFAULT_SERVICE_CONFIGS: Final = {
|
|||
ServiceTypes.ROUTER.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
|
||||
ServiceTypes.AUTH.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
|
||||
ServiceTypes.PROXY_PRE_CALL.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]},
|
||||
ServiceTypes.ANTHROPIC_WIF.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape
|
||||
"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] # mutable-ok: ServiceConfig mandates a list
|
||||
},
|
||||
# cache hits are counter-only: no HTTP call happens, so observing a latency would be a lie
|
||||
ServiceTypes.ANTHROPIC_WIF_CACHE.value: { # mutable-ok: ServiceConfig mandates the dict-of-list shape
|
||||
"metrics": [ServiceMetrics.COUNTER] # mutable-ok: ServiceConfig mandates a list
|
||||
},
|
||||
# Operational metrics for DB Transaction Queues
|
||||
ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]},
|
||||
ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]},
|
||||
|
|
|
|||
|
|
@ -3643,9 +3643,27 @@ bedrock_batch_litellm_params: Final = (
|
|||
"bedrock_tags",
|
||||
)
|
||||
|
||||
# Anthropic workload identity federation config, read from litellm_params by the
|
||||
# Anthropic auth tier. Listed for the same reason as the fields above: an
|
||||
# unrecognized top-level key is swept into extra_body and sent to /v1/messages.
|
||||
# Derived from get_litellm_params.ANTHROPIC_WIF_KWARGS_KEYS (not hand-typed) so the
|
||||
# request-body ban list and the clear-on-api_base-override list can never drift from
|
||||
# the set the kwargs funnel actually forwards. Imported here rather than at module top:
|
||||
# get_litellm_params.py's own import chain (llms/openai/data_residency -> llms/__init__)
|
||||
# reaches back into this module for CallTypes, which by this point in the file is
|
||||
# already bound on the partially-initialized module.
|
||||
from ..litellm_core_utils.get_litellm_params import ( # noqa: E402 # deferred past CallTypes to break the import cycle
|
||||
ANTHROPIC_WIF_KWARGS_KEYS,
|
||||
OPENAI_WIF_KWARGS_KEYS,
|
||||
)
|
||||
|
||||
anthropic_wif_litellm_params: Final = tuple(sorted(ANTHROPIC_WIF_KWARGS_KEYS))
|
||||
openai_wif_litellm_params: Final = tuple(sorted(OPENAI_WIF_KWARGS_KEYS))
|
||||
server_owned_wif_litellm_params: Final = anthropic_wif_litellm_params + openai_wif_litellm_params
|
||||
|
||||
all_litellm_params = (
|
||||
agentic_loop_internal_litellm_params
|
||||
+ [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params]
|
||||
+ [TRUSTED_CALLBACK_VARS_FIELD, *bedrock_batch_litellm_params, *server_owned_wif_litellm_params]
|
||||
+ [
|
||||
"metadata",
|
||||
"litellm_metadata",
|
||||
|
|
|
|||
|
|
@ -7361,9 +7361,8 @@ def _get_valid_models_from_provider_api(
|
|||
|
||||
if cached_result is not None:
|
||||
return cached_result
|
||||
models: Final = provider_config.get_models(
|
||||
api_key=litellm_params.api_key if litellm_params is not None else None,
|
||||
api_base=litellm_params.api_base if litellm_params is not None else None,
|
||||
models: Final = provider_config.discover_models(
|
||||
litellm_params=litellm_params.model_dump(exclude_none=True) if litellm_params is not None else None
|
||||
)
|
||||
|
||||
_model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models)
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ GET /user/spend/report
|
|||
|
||||
# Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state
|
||||
GET /budget/settings
|
||||
GET /credentials/{credential_name}/jwks
|
||||
GET /router/fields
|
||||
GET /guardrails/ui/add_guardrail_settings
|
||||
GET /guardrails/ui/category_yaml/{category_name}
|
||||
|
|
|
|||
|
|
@ -1235,6 +1235,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
|
|||
result set - zero cost, zero usage, no models - instead of letting the file
|
||||
fetch raise "Output file id is None" on every aretrieve_batch logging poll.
|
||||
"""
|
||||
|
||||
# The output-file fetch must not even be attempted when there is no output file.
|
||||
async def _must_not_fetch(*args, **kwargs):
|
||||
pytest.fail("_fetch_batch_output_file_content should not be called")
|
||||
|
|
@ -1361,7 +1362,10 @@ def test_anthropic_response_body_is_result_message():
|
|||
|
||||
|
||||
def test_anthropic_usage_conversion_includes_cache_tokens():
|
||||
body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)}
|
||||
body = {
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000),
|
||||
}
|
||||
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic")
|
||||
assert usage.prompt_tokens == 11000
|
||||
assert usage.completion_tokens == 200
|
||||
|
|
@ -1376,7 +1380,9 @@ def test_bedrock_model_output_line_success_check():
|
|||
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
|
||||
}
|
||||
assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True
|
||||
assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
|
||||
assert (
|
||||
bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_cost_uses_deployment_model_name():
|
||||
|
|
@ -1430,7 +1436,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
|
|||
rows = [
|
||||
{
|
||||
"custom_id": "req-1",
|
||||
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
|
||||
"response": {
|
||||
"status_code": 200,
|
||||
"body": {
|
||||
"model": "gpt-5.2",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
|
|
@ -1674,7 +1686,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) ->
|
|||
|
||||
|
||||
def test_bedrock_converse_shaped_batch_usage_is_parsed():
|
||||
body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}}
|
||||
body = {
|
||||
"model": "us.amazon.nova-lite-v1:0",
|
||||
"usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742},
|
||||
}
|
||||
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742)
|
||||
|
||||
|
|
@ -1718,3 +1733,37 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
|
|||
assert usage.total_tokens == 0
|
||||
assert "does not understand" in caplog.text
|
||||
assert "inputTextTokenCount" in caplog.text
|
||||
|
||||
|
||||
class TestFileAccessCredentialsCarryFederation:
|
||||
"""A federated deployment holds no api_key, so the fetch that reads a finished batch's output
|
||||
has to inherit the federation fields or it cannot authenticate and the batch is never billed."""
|
||||
|
||||
def test_federation_fields_survive_extraction(self):
|
||||
from litellm.batches.batch_utils import _extract_file_access_credentials
|
||||
|
||||
credentials = _extract_file_access_credentials(
|
||||
{
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "fdrl_x",
|
||||
"anthropic_organization_id": "org-x",
|
||||
"anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token",
|
||||
"something_unrelated": "dropped",
|
||||
}
|
||||
)
|
||||
|
||||
assert credentials["anthropic_federation_rule_id"] == "fdrl_x"
|
||||
assert credentials["anthropic_organization_id"] == "org-x"
|
||||
assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token"
|
||||
assert "something_unrelated" not in credentials
|
||||
|
||||
def test_every_federation_field_is_carried(self):
|
||||
"""Derived from the kwargs set, so a new federation field is carried without an edit here."""
|
||||
from litellm.batches.batch_utils import _extract_file_access_credentials
|
||||
from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS
|
||||
|
||||
params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS}
|
||||
|
||||
credentials = _extract_file_access_credentials(params)
|
||||
|
||||
assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS)
|
||||
|
|
|
|||
|
|
@ -135,3 +135,28 @@ def test_services_logger_custom_latency_buckets():
|
|||
REGISTRY.unregister(collector)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def test_anthropic_wif_services_are_wired_into_the_registry():
|
||||
"""Reverting the ANTHROPIC_WIF/ANTHROPIC_WIF_CACHE ServiceTypes members or their
|
||||
DEFAULT_SERVICE_CONFIGS entries must fail here: the exchange service gets counters plus a
|
||||
latency histogram, while the cache-hit service is counter-only so a hit can never fake a latency."""
|
||||
from litellm.types.services import DEFAULT_SERVICE_CONFIGS
|
||||
|
||||
assert ServiceTypes.ANTHROPIC_WIF.value == "anthropic_wif"
|
||||
assert ServiceTypes.ANTHROPIC_WIF_CACHE.value == "anthropic_wif_cache"
|
||||
assert DEFAULT_SERVICE_CONFIGS["anthropic_wif"]["metrics"] == [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
|
||||
assert DEFAULT_SERVICE_CONFIGS["anthropic_wif_cache"]["metrics"] == [ServiceMetrics.COUNTER]
|
||||
|
||||
pl = PrometheusServicesLogger()
|
||||
wif_names = {obj._name for obj in pl.payload_to_prometheus_map["anthropic_wif"]}
|
||||
assert wif_names == {
|
||||
"litellm_anthropic_wif_latency",
|
||||
"litellm_anthropic_wif_failed_requests",
|
||||
"litellm_anthropic_wif_total_requests",
|
||||
}
|
||||
cache_names = {obj._name for obj in pl.payload_to_prometheus_map["anthropic_wif_cache"]}
|
||||
assert cache_names == {
|
||||
"litellm_anthropic_wif_cache_failed_requests",
|
||||
"litellm_anthropic_wif_cache_total_requests",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -244,3 +244,129 @@ class TestRustOptIn:
|
|||
from litellm.types.utils import all_litellm_params
|
||||
|
||||
assert "rust" in all_litellm_params
|
||||
|
||||
|
||||
class TestAnthropicWifKeys:
|
||||
"""The six anthropic_* WIF keys need the same dual registration as `rust`:
|
||||
carried by the kwargs funnel into litellm_params (where the Anthropic auth
|
||||
tier reads them) AND listed in all_litellm_params (so the extra_body sweep
|
||||
never sends them to /v1/messages)."""
|
||||
|
||||
SIX_KEYS = {
|
||||
"anthropic_federation_rule_id": "fdrl_1",
|
||||
"anthropic_organization_id": "org-1",
|
||||
"anthropic_service_account_id": "svcacct_1",
|
||||
"anthropic_workspace_id": "wrkspc_1",
|
||||
"anthropic_identity_token_file": "/var/run/secrets/tok",
|
||||
"anthropic_identity_token": "oidc/env/TOK",
|
||||
}
|
||||
|
||||
def test_keys_survive_into_litellm_params(self):
|
||||
params = get_litellm_params(**self.SIX_KEYS)
|
||||
for key, value in self.SIX_KEYS.items():
|
||||
assert params[key] == value
|
||||
|
||||
def test_keys_are_forwarded_from_completion_kwargs(self):
|
||||
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
|
||||
|
||||
assert set(self.SIX_KEYS) <= FORWARDED_KWARGS_KEYS
|
||||
|
||||
def test_keys_stay_out_of_the_provider_body(self):
|
||||
from litellm.types.utils import all_litellm_params
|
||||
|
||||
for key in self.SIX_KEYS:
|
||||
assert key in all_litellm_params
|
||||
|
||||
def test_keys_absent_when_not_configured(self):
|
||||
params = get_litellm_params()
|
||||
for key in self.SIX_KEYS:
|
||||
assert key not in params
|
||||
|
||||
|
||||
class TestAnthropicWifIdentitySourceKeys:
|
||||
"""Phase 1 adds 11 more anthropic_* WIF keys (the anthropic_identity_source discriminator
|
||||
plus the internal_issuer/keycloak identity-source fields) that need the same dual
|
||||
registration as the original six tested above."""
|
||||
|
||||
NEW_KEYS = {
|
||||
"anthropic_identity_source": "keycloak",
|
||||
"anthropic_issuer_url": "https://issuer.example",
|
||||
"anthropic_issuer_subject": "svc-account",
|
||||
"anthropic_issuer_audience": "https://api.anthropic.com",
|
||||
"anthropic_issuer_ttl_seconds": "300",
|
||||
"anthropic_issuer_signing_key_ref": "oidc/env/ISSUER_KEY",
|
||||
"anthropic_keycloak_token_url": "https://kc.example/realms/r/protocol/openid-connect/token",
|
||||
"anthropic_keycloak_client_id": "litellm",
|
||||
"anthropic_keycloak_auth_method": "client_secret_basic",
|
||||
"anthropic_keycloak_client_secret_ref": "oidc/env/KC_SECRET",
|
||||
"anthropic_keycloak_scope": "anthropic-wif",
|
||||
# Server-set when a client redirects api_base; carried here so it is not dropped in transit
|
||||
"anthropic_disable_workload_identity_federation": True,
|
||||
}
|
||||
|
||||
def test_new_keys_are_exactly_the_non_legacy_registered_set(self):
|
||||
"""Fails the moment a key is added to ANTHROPIC_WIF_KWARGS_KEYS without a matching entry
|
||||
here (or vice versa), catching drift between what wif.py dispatches on and what this
|
||||
test (and the funnel/provider-body tests below) actually exercises."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS
|
||||
|
||||
assert set(self.NEW_KEYS) == ANTHROPIC_WIF_KWARGS_KEYS - set(TestAnthropicWifKeys.SIX_KEYS)
|
||||
|
||||
def test_keys_survive_into_litellm_params(self):
|
||||
params = get_litellm_params(**self.NEW_KEYS)
|
||||
for key, value in self.NEW_KEYS.items():
|
||||
assert params[key] == value
|
||||
|
||||
def test_keys_are_forwarded_from_completion_kwargs(self):
|
||||
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
|
||||
|
||||
assert set(self.NEW_KEYS) <= FORWARDED_KWARGS_KEYS
|
||||
|
||||
def test_keys_stay_out_of_the_provider_body(self):
|
||||
from litellm.types.utils import all_litellm_params
|
||||
|
||||
for key in self.NEW_KEYS:
|
||||
assert key in all_litellm_params
|
||||
|
||||
def test_keys_absent_when_not_configured(self):
|
||||
params = get_litellm_params()
|
||||
for key in self.NEW_KEYS:
|
||||
assert key not in params
|
||||
|
||||
|
||||
class TestOpenAIWifKeys:
|
||||
"""The three openai_* WIF keys carry a deployment's federation identity through the kwargs
|
||||
funnel into litellm_params (where the OpenAI client factory reads them) and stay out of the
|
||||
provider body, exactly like the anthropic_* keys above."""
|
||||
|
||||
THREE_KEYS = {
|
||||
"openai_identity_provider_id": "idp_1",
|
||||
"openai_service_account_id": "user-1",
|
||||
"openai_identity_token_file": "/var/run/secrets/tokens/openai",
|
||||
}
|
||||
|
||||
def test_keys_are_exactly_the_registered_set(self):
|
||||
from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS
|
||||
|
||||
assert set(self.THREE_KEYS) == OPENAI_WIF_KWARGS_KEYS
|
||||
|
||||
def test_keys_survive_into_litellm_params(self):
|
||||
params = get_litellm_params(**self.THREE_KEYS)
|
||||
for key, value in self.THREE_KEYS.items():
|
||||
assert params[key] == value
|
||||
|
||||
def test_keys_are_forwarded_from_completion_kwargs(self):
|
||||
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
|
||||
|
||||
assert set(self.THREE_KEYS) <= FORWARDED_KWARGS_KEYS
|
||||
|
||||
def test_keys_stay_out_of_the_provider_body(self):
|
||||
from litellm.types.utils import all_litellm_params
|
||||
|
||||
for key in self.THREE_KEYS:
|
||||
assert key in all_litellm_params
|
||||
|
||||
def test_keys_absent_when_not_configured(self):
|
||||
params = get_litellm_params()
|
||||
for key in self.THREE_KEYS:
|
||||
assert key not in params
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ asyncio.run) is exercised directly, mirroring the dispatch-contract discipline i
|
|||
tests/test_litellm/batches/test_main.py.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -35,9 +37,7 @@ def _ok_batch_response():
|
|||
"ended_at": "2024-09-24T11:00:00Z",
|
||||
"request_counts": {"succeeded": 2, "errored": 0},
|
||||
},
|
||||
request=httpx.Request(
|
||||
"GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
|
||||
),
|
||||
request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/msgbatch_abc"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -59,9 +59,7 @@ def patched_client():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
|
||||
handler, patched_client
|
||||
):
|
||||
async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(handler, patched_client):
|
||||
fake_client, factory = patched_client
|
||||
|
||||
batch = await handler.aretrieve_batch(
|
||||
|
|
@ -76,9 +74,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
|
|||
fake_client.get.assert_awaited_once()
|
||||
_, call_kwargs = fake_client.get.call_args
|
||||
# Exact URL built by get_retrieve_batch_url.
|
||||
assert call_kwargs["url"] == (
|
||||
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
|
||||
)
|
||||
assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
|
||||
# Auth + version + beta headers built by validate_environment.
|
||||
headers = call_kwargs["headers"]
|
||||
assert headers["x-api-key"] == "sk-ant-test"
|
||||
|
|
@ -93,9 +89,7 @@ async def test_aretrieve_batch_fires_get_with_correct_url_and_headers(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aretrieve_batch_uses_anthropic_provider_for_client(
|
||||
handler, patched_client
|
||||
):
|
||||
async def test_aretrieve_batch_uses_anthropic_provider_for_client(handler, patched_client):
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_, factory = patched_client
|
||||
|
|
@ -111,14 +105,10 @@ async def test_aretrieve_batch_uses_anthropic_provider_for_client(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aretrieve_batch_resolves_api_key_from_model_info(
|
||||
handler, patched_client
|
||||
):
|
||||
async def test_aretrieve_batch_resolves_api_key_from_model_info(handler, patched_client):
|
||||
fake_client, _ = patched_client
|
||||
# api_key=None -> handler falls back to AnthropicModelInfo.get_api_key().
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="sk-from-env"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="sk-from-env"):
|
||||
await handler.aretrieve_batch(
|
||||
batch_id="msgbatch_abc",
|
||||
api_base="https://api.anthropic.com",
|
||||
|
|
@ -134,9 +124,7 @@ async def test_aretrieve_batch_resolves_api_key_from_model_info(
|
|||
async def test_aretrieve_batch_missing_api_key_raises(handler, patched_client):
|
||||
fake_client, _ = patched_client
|
||||
# No api_key and resolver yields None -> hard error before any network call.
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value=None
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None):
|
||||
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
|
||||
await handler.aretrieve_batch(
|
||||
batch_id="msgbatch_abc",
|
||||
|
|
@ -165,9 +153,7 @@ async def test_aretrieve_batch_resolves_default_api_base(handler, patched_client
|
|||
max_retries=0,
|
||||
)
|
||||
_, call_kwargs = fake_client.get.call_args
|
||||
assert call_kwargs["url"] == (
|
||||
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
|
||||
)
|
||||
assert call_kwargs["url"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -176,9 +162,7 @@ async def test_aretrieve_batch_raises_for_status(handler):
|
|||
error_response = httpx.Response(
|
||||
status_code=404,
|
||||
json={"error": "not found"},
|
||||
request=httpx.Request(
|
||||
"GET", "https://api.anthropic.com/v1/messages/batches/missing"
|
||||
),
|
||||
request=httpx.Request("GET", "https://api.anthropic.com/v1/messages/batches/missing"),
|
||||
)
|
||||
fake_client = MagicMock()
|
||||
fake_client.get = AsyncMock(return_value=error_response)
|
||||
|
|
@ -213,21 +197,15 @@ async def test_aretrieve_batch_invokes_pre_call_logging(handler, patched_client)
|
|||
assert pre_kwargs["input"] == "msgbatch_abc"
|
||||
assert pre_kwargs["api_key"] == "sk-ant-test"
|
||||
# The logged api_base is the full retrieve URL, not the bare base.
|
||||
assert pre_kwargs["additional_args"]["api_base"] == (
|
||||
"https://api.anthropic.com/v1/messages/batches/msgbatch_abc"
|
||||
)
|
||||
assert pre_kwargs["additional_args"]["api_base"] == ("https://api.anthropic.com/v1/messages/batches/msgbatch_abc")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aretrieve_batch_builds_default_logging_obj_when_absent(
|
||||
handler, patched_client
|
||||
):
|
||||
async def test_aretrieve_batch_builds_default_logging_obj_when_absent(handler, patched_client):
|
||||
# logging_obj=None -> handler constructs a real Logging object; the call
|
||||
# must still complete (no AttributeError on a missing logger).
|
||||
_, _ = patched_client
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.Logging"
|
||||
) as logging_cls:
|
||||
with patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls:
|
||||
logging_cls.return_value = MagicMock()
|
||||
batch = await handler.aretrieve_batch(
|
||||
batch_id="msgbatch_abc",
|
||||
|
|
@ -281,3 +259,99 @@ def test_retrieve_batch_sync_runs_to_result(handler, patched_client):
|
|||
assert isinstance(batch, LiteLLMBatch)
|
||||
assert batch.id == "msgbatch_abc"
|
||||
assert batch.status == "completed"
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# aretrieve_batch must not block the event loop on a WIF token exchange
|
||||
# =========================================================================== #
|
||||
|
||||
_WIF_ENV = {
|
||||
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_batches_seam",
|
||||
"ANTHROPIC_ORGANIZATION_ID": "org-batches-seam",
|
||||
"ANTHROPIC_IDENTITY_TOKEN": "batches-seam-inline-jwt",
|
||||
}
|
||||
|
||||
|
||||
class _BlockingPoster:
|
||||
"""A token-endpoint poster that blocks until released, so the test can prove
|
||||
the exchange ran off the event loop's own thread instead of freezing it."""
|
||||
|
||||
def __init__(self):
|
||||
self.release = threading.Event()
|
||||
self.thread_ids = []
|
||||
|
||||
def post(self, url, *, content, headers, timeout):
|
||||
self.thread_ids.append(threading.get_ident())
|
||||
self.release.wait(timeout=5)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "sk-ant-oat01-batches-seam",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, patched_client, monkeypatch):
|
||||
"""Regression: aretrieve_batch called the synchronous validate_environment
|
||||
directly, so a cold WIF mint ran inline on the event loop and froze every
|
||||
other concurrent coroutine until the exchange finished."""
|
||||
from litellm.llms.anthropic import common_utils as anthropic_common_utils
|
||||
from litellm.llms.anthropic.wif import get_anthropic_wif_token
|
||||
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine
|
||||
|
||||
fake_client, _ = patched_client
|
||||
for name in (
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN",
|
||||
"ANTHROPIC_API_BASE",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
for name, value in _WIF_ENV.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
|
||||
poster = _BlockingPoster()
|
||||
engine = JwtBearerTokenExchangeEngine(poster=poster)
|
||||
|
||||
def routed_through_injected_engine(litellm_params, api_base, model):
|
||||
return get_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", routed_through_injected_engine)
|
||||
|
||||
ticks = []
|
||||
|
||||
async def ticker():
|
||||
for i in range(20):
|
||||
await asyncio.sleep(0.005)
|
||||
ticks.append(i)
|
||||
|
||||
ticker_task = asyncio.create_task(ticker())
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
retrieve_task = asyncio.create_task(
|
||||
handler.aretrieve_batch(
|
||||
batch_id="msgbatch_abc",
|
||||
api_base="https://api.anthropic.com",
|
||||
api_key=None,
|
||||
timeout=60.0,
|
||||
max_retries=0,
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
# The ticker kept advancing while the token exchange was still blocked on
|
||||
# poster.release, proving the exchange did not run on the event loop.
|
||||
assert len(ticks) > 0
|
||||
assert not retrieve_task.done()
|
||||
|
||||
poster.release.set()
|
||||
batch = await retrieve_task
|
||||
await ticker_task
|
||||
|
||||
assert batch.id == "msgbatch_abc"
|
||||
assert poster.thread_ids
|
||||
assert poster.thread_ids[0] != threading.get_ident()
|
||||
sent_headers = fake_client.get.call_args.kwargs["headers"]
|
||||
assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam"
|
||||
|
|
|
|||
|
|
@ -80,8 +80,11 @@ def test_validate_environment_preserves_existing_beta_header(config):
|
|||
litellm_params={},
|
||||
api_key="sk-ant-test",
|
||||
)
|
||||
# Existing beta header must NOT be overwritten.
|
||||
assert headers["anthropic-beta"] == "custom-beta-value"
|
||||
# Existing beta values are preserved and the batches beta is merged in.
|
||||
assert set(headers["anthropic-beta"].split(",")) == {
|
||||
"custom-beta-value",
|
||||
"message-batches-2024-09-24",
|
||||
}
|
||||
|
||||
|
||||
def test_validate_environment_oauth_key_uses_bearer(config):
|
||||
|
|
@ -100,9 +103,7 @@ def test_validate_environment_oauth_key_uses_bearer(config):
|
|||
|
||||
def test_validate_environment_missing_key_raises(config):
|
||||
# No api_key passed and no env credentials -> get_auth_header returns None.
|
||||
with patch.object(
|
||||
config.anthropic_model_info, "get_auth_header", return_value=None
|
||||
):
|
||||
with patch.object(config.anthropic_model_info, "get_auth_header", return_value=None):
|
||||
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
|
||||
config.validate_environment(
|
||||
headers={},
|
||||
|
|
@ -241,12 +242,7 @@ def test_get_retrieve_batch_url_uses_default_api_base(config):
|
|||
|
||||
|
||||
def test_transform_retrieve_batch_request_returns_empty_dict(config):
|
||||
assert (
|
||||
config.transform_retrieve_batch_request(
|
||||
batch_id="msgbatch_123", optional_params={}, litellm_params={}
|
||||
)
|
||||
== {}
|
||||
)
|
||||
assert config.transform_retrieve_batch_request(batch_id="msgbatch_123", optional_params={}, litellm_params={}) == {}
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -455,9 +451,7 @@ def test_transform_retrieve_response_unparseable_json_raises(config):
|
|||
|
||||
|
||||
def test_get_error_class_with_dict_headers(config):
|
||||
err = config.get_error_class(
|
||||
error_message="rate limited", status_code=429, headers={"x-ratelimit": "0"}
|
||||
)
|
||||
err = config.get_error_class(error_message="rate limited", status_code=429, headers={"x-ratelimit": "0"})
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
|
||||
assert isinstance(err, AnthropicError)
|
||||
|
|
@ -467,9 +461,7 @@ def test_get_error_class_with_dict_headers(config):
|
|||
|
||||
def test_get_error_class_with_httpx_headers(config):
|
||||
hdrs = httpx.Headers({"retry-after": "5"})
|
||||
err = config.get_error_class(
|
||||
error_message="server error", status_code=500, headers=hdrs
|
||||
)
|
||||
err = config.get_error_class(error_message="server error", status_code=500, headers=hdrs)
|
||||
assert err.status_code == 500
|
||||
assert err.message == "server error"
|
||||
|
||||
|
|
@ -543,9 +535,7 @@ def test_transform_response_skips_malformed_lines(config):
|
|||
|
||||
def fake_transform_parsed(*, completion_response, raw_response, model_response):
|
||||
mr = ModelResponse()
|
||||
setattr(
|
||||
mr, "usage", Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10)
|
||||
)
|
||||
setattr(mr, "usage", Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10))
|
||||
return mr
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -588,13 +578,16 @@ def test_transform_response_reraises_unexpected_error(config):
|
|||
|
||||
# A non-JSONDecodeError raised during usage aggregation must propagate
|
||||
# (the outer `except Exception: raise e`), not be swallowed.
|
||||
with patch.object(
|
||||
config.anthropic_chat_config,
|
||||
"transform_parsed_response",
|
||||
side_effect=fake_transform_parsed,
|
||||
), patch(
|
||||
"litellm.cost_calculator.BaseTokenUsageProcessor.combine_usage_objects",
|
||||
side_effect=RuntimeError("boom"),
|
||||
with (
|
||||
patch.object(
|
||||
config.anthropic_chat_config,
|
||||
"transform_parsed_response",
|
||||
side_effect=fake_transform_parsed,
|
||||
),
|
||||
patch(
|
||||
"litellm.cost_calculator.BaseTokenUsageProcessor.combine_usage_objects",
|
||||
side_effect=RuntimeError("boom"),
|
||||
),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
config.transform_response(
|
||||
|
|
|
|||
|
|
@ -48,9 +48,7 @@ class MockDynamicGuardrail(CustomGuardrail):
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(
|
||||
request_data
|
||||
)
|
||||
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data)
|
||||
return inputs
|
||||
|
||||
|
||||
|
|
@ -197,9 +195,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
|
|||
|
||||
assert guardrail.request_data is not None
|
||||
assert guardrail.request_data["response"] is mock_response
|
||||
assert (
|
||||
guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
|
||||
)
|
||||
assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_chunk_passes_responses_so_far_and_metadata(self):
|
||||
|
|
@ -211,9 +207,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
|
|||
|
||||
with (
|
||||
patch.object(handler, "_check_streaming_has_ended", return_value=False),
|
||||
patch.object(
|
||||
handler, "get_streaming_string_so_far", return_value="partial text"
|
||||
),
|
||||
patch.object(handler, "get_streaming_string_so_far", return_value="partial text"),
|
||||
):
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
|
|
@ -225,9 +219,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData:
|
|||
|
||||
assert guardrail.request_data is not None
|
||||
assert guardrail.request_data["responses"] is responses_so_far
|
||||
assert (
|
||||
guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
|
||||
)
|
||||
assert guardrail.request_data["litellm_metadata"]["user_api_key_user_id"] == "u-1"
|
||||
|
||||
|
||||
class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
||||
|
|
@ -276,17 +268,11 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
data = {
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_metadata": {
|
||||
"guardrails": [
|
||||
{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}
|
||||
]
|
||||
},
|
||||
"litellm_metadata": {"guardrails": [{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}]},
|
||||
}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
await handler.process_input_messages(
|
||||
data=data, guardrail_to_apply=guardrail
|
||||
)
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert data.get("litellm_metadata", {}).get("guardrails")
|
||||
assert guardrail.dynamic_params == {"policy_id": "policy-123"}
|
||||
|
|
@ -1183,9 +1169,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
# Mock _check_streaming_has_ended to return False (stream not ended)
|
||||
with (
|
||||
patch.object(handler, "_check_streaming_has_ended", return_value=False),
|
||||
patch.object(
|
||||
handler, "get_streaming_string_so_far", return_value="partial text"
|
||||
),
|
||||
patch.object(handler, "get_streaming_string_so_far", return_value="partial text"),
|
||||
):
|
||||
responses_so_far = [b"data: some chunk"]
|
||||
|
||||
|
|
@ -1216,9 +1200,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
|
||||
data = {
|
||||
"model": "claude-opus-4-6",
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the weather in San Francisco?"}
|
||||
],
|
||||
"messages": [{"role": "user", "content": "What is the weather in San Francisco?"}],
|
||||
"tools": [
|
||||
{
|
||||
"type": "tool_search_tool_regex_20251119",
|
||||
|
|
@ -1381,17 +1363,11 @@ class TestAnthropicMessagesIncrementalScan:
|
|||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(
|
||||
data=self._data(turn1, sid), guardrail_to_apply=guardrail
|
||||
)
|
||||
await handler.process_input_messages(data=self._data(turn1, sid), guardrail_to_apply=guardrail)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"What is the capital of France?"
|
||||
]
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["What is the capital of France?"]
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(
|
||||
data=self._data(turn2, sid), guardrail_to_apply=guardrail
|
||||
)
|
||||
await handler.process_input_messages(data=self._data(turn2, sid), guardrail_to_apply=guardrail)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"Paris.",
|
||||
|
|
@ -1830,9 +1806,7 @@ class TestAnthropicMessagesScanOnlyToolResults:
|
|||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.seen_texts == ["fetched POISON page"], (
|
||||
"only the tool_result payload may reach the guardrail"
|
||||
)
|
||||
assert guardrail.seen_texts == ["fetched POISON page"], "only the tool_result payload may reach the guardrail"
|
||||
assert guardrail.captured_inputs is not None
|
||||
assert guardrail.captured_inputs.get("tools") is None
|
||||
assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"]
|
||||
|
|
|
|||
|
|
@ -23,9 +23,7 @@ async def test_make_call_passes_logging_obj_to_client_post():
|
|||
mock_client = AsyncMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.aiter_lines = MagicMock(
|
||||
return_value=iter(
|
||||
[b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']
|
||||
)
|
||||
return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'])
|
||||
)
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
|
|
@ -94,9 +92,7 @@ def test_redacted_thinking_content_block_delta():
|
|||
"data": "EuoBCoYBGAIiQJ/SxkPAgqxhKok29YrpJHRUJ0OT8ahCHKAwyhmRuUhtdmDX9+mn4gDzKNv3fVpQdB01zEPMzNY3QuTCd+1bdtEqQK6JuKHqdndbwpr81oVWb4wxd1GqF/7Jkw74IlQa27oobX+KuRkopr9Dllt/RDe7Se0sI1IkU7tJIAQCoP46OAwSDF51P09q67xhHlQ3ihoM2aOVlkghq/X0w8NlIjBMNvXYNbjhyrOcIg6kPFn2ed/KK7Cm5prYAtXCwkb4Wr5tUSoSHu9T5hKdJRbr6WsqEc7Lle7FULqMLZGkhqXyc3BA",
|
||||
},
|
||||
}
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=False, json_mode=False
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
|
||||
model_response = model_response_iterator.chunk_parser(chunk=chunk)
|
||||
print(f"\n\nmodel_response: {model_response}\n\n")
|
||||
assert model_response.choices[0].delta.thinking_blocks is not None
|
||||
|
|
@ -104,19 +100,14 @@ def test_redacted_thinking_content_block_delta():
|
|||
print(
|
||||
f"\n\nmodel_response.choices[0].delta.thinking_blocks[0]: {model_response.choices[0].delta.thinking_blocks[0]}\n\n"
|
||||
)
|
||||
assert (
|
||||
model_response.choices[0].delta.thinking_blocks[0]["type"]
|
||||
== "redacted_thinking"
|
||||
)
|
||||
assert model_response.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking"
|
||||
|
||||
assert model_response.choices[0].delta.provider_specific_fields is not None
|
||||
assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields
|
||||
|
||||
|
||||
def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
chunks = [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
|
|
@ -140,17 +131,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
|
|||
},
|
||||
]
|
||||
|
||||
parsed_chunks = [
|
||||
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
|
||||
]
|
||||
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
|
||||
reasoning_content = "".join(
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
|
||||
for chunk in parsed_chunks
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
|
||||
)
|
||||
thinking_blocks = tuple(
|
||||
block
|
||||
for chunk in parsed_chunks
|
||||
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
)
|
||||
expected_delta_blocks = (
|
||||
{"type": "thinking", "thinking": "Step 1. "},
|
||||
|
|
@ -164,18 +150,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
|
|||
|
||||
assert reasoning_content == "Step 1. Step 2."
|
||||
assert thinking_blocks == (*expected_delta_blocks, expected_thinking_block)
|
||||
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {
|
||||
"thinking_blocks": [expected_delta_blocks[0]]
|
||||
}
|
||||
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {
|
||||
"thinking_blocks": [expected_thinking_block]
|
||||
}
|
||||
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]}
|
||||
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_thinking_block]}
|
||||
|
||||
|
||||
def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
chunks = [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
|
|
@ -195,17 +175,12 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
|
|||
{"type": "content_block_stop", "index": 0},
|
||||
]
|
||||
|
||||
parsed_chunks = [
|
||||
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
|
||||
]
|
||||
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
|
||||
reasoning_content = "".join(
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
|
||||
for chunk in parsed_chunks
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
|
||||
)
|
||||
thinking_blocks = tuple(
|
||||
block
|
||||
for chunk in parsed_chunks
|
||||
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
)
|
||||
|
||||
assert reasoning_content == "Step 1. Step 2."
|
||||
|
|
@ -216,9 +191,7 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
|
|||
|
||||
|
||||
def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
chunks = [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
|
|
@ -237,17 +210,12 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
|
|||
},
|
||||
]
|
||||
|
||||
parsed_chunks = [
|
||||
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
|
||||
]
|
||||
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
|
||||
reasoning_content = "".join(
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
|
||||
for chunk in parsed_chunks
|
||||
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
|
||||
)
|
||||
thinking_blocks = tuple(
|
||||
block
|
||||
for chunk in parsed_chunks
|
||||
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
|
||||
)
|
||||
|
||||
assert reasoning_content == "Step 1. Step 2."
|
||||
|
|
@ -258,9 +226,7 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
|
|||
|
||||
|
||||
def test_handle_json_mode_chunk_response_format_tool():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
response_format_tool = ChatCompletionToolCallChunk(
|
||||
id="tool_123",
|
||||
type="function",
|
||||
|
|
@ -271,9 +237,7 @@ def test_handle_json_mode_chunk_response_format_tool():
|
|||
index=0,
|
||||
)
|
||||
|
||||
text, tool_use = model_response_iterator._handle_json_mode_chunk(
|
||||
"", response_format_tool
|
||||
)
|
||||
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
|
||||
print(f"\n\nresponse_format_tool text: {text}\n\n")
|
||||
print(f"\n\nresponse_format_tool tool_use: {tool_use}\n\n")
|
||||
|
||||
|
|
@ -282,15 +246,11 @@ def test_handle_json_mode_chunk_response_format_tool():
|
|||
|
||||
|
||||
def test_handle_json_mode_chunk_regular_tool():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
regular_tool = ChatCompletionToolCallChunk(
|
||||
id="tool_456",
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name="get_weather", arguments='{"location": "San Francisco, CA"}'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
|
@ -304,17 +264,13 @@ def test_handle_json_mode_chunk_regular_tool():
|
|||
|
||||
|
||||
def test_handle_json_mode_chunk_streaming_response_format_tool():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
|
||||
# First chunk: response_format tool with id and name, but no arguments
|
||||
first_chunk = ChatCompletionToolCallChunk(
|
||||
id="tool_123",
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=RESPONSE_FORMAT_TOOL_NAME, arguments=""
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=""),
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
|
@ -322,9 +278,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
|
|||
second_chunk = ChatCompletionToolCallChunk(
|
||||
id=None,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=None, arguments='{"question": "What is the weather?"'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"question": "What is the weather?"'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
|
@ -332,9 +286,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
|
|||
third_chunk = ChatCompletionToolCallChunk(
|
||||
id=None,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=None, arguments=', "answer": "It is sunny"}'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name=None, arguments=', "answer": "It is sunny"}'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
|
@ -365,9 +317,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
|
|||
|
||||
|
||||
def test_handle_json_mode_chunk_streaming_regular_tool():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
|
||||
# First chunk: regular tool with id and name, but no arguments
|
||||
first_chunk = ChatCompletionToolCallChunk(
|
||||
|
|
@ -381,9 +331,7 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
|
|||
second_chunk = ChatCompletionToolCallChunk(
|
||||
id=None,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=None, arguments='{"location": "San Francisco, CA"}'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"location": "San Francisco, CA"}'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
|
@ -408,27 +356,19 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
|
|||
|
||||
|
||||
def test_response_format_tool_finish_reason():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
|
||||
# First chunk: response_format tool
|
||||
response_format_tool = ChatCompletionToolCallChunk(
|
||||
id="tool_123",
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
# Process the tool call (should set converted_response_format_tool flag)
|
||||
text, tool_use = model_response_iterator._handle_json_mode_chunk(
|
||||
"", response_format_tool
|
||||
)
|
||||
print(
|
||||
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
|
||||
)
|
||||
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
|
||||
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
|
||||
|
||||
# Simulate message_delta chunk with tool_use stop_reason
|
||||
message_delta_chunk = {
|
||||
|
|
@ -447,25 +387,19 @@ def test_response_format_tool_finish_reason():
|
|||
|
||||
|
||||
def test_regular_tool_finish_reason():
|
||||
model_response_iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=True
|
||||
)
|
||||
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
|
||||
|
||||
# First chunk: regular tool (not response_format)
|
||||
regular_tool = ChatCompletionToolCallChunk(
|
||||
id="tool_456",
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name="get_weather", arguments='{"location": "San Francisco, CA"}'
|
||||
),
|
||||
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
|
||||
index=0,
|
||||
)
|
||||
|
||||
# Process the tool call (should NOT set converted_response_format_tool flag)
|
||||
text, tool_use = model_response_iterator._handle_json_mode_chunk("", regular_tool)
|
||||
print(
|
||||
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
|
||||
)
|
||||
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
|
||||
|
||||
# Simulate message_delta chunk with tool_use stop_reason
|
||||
message_delta_chunk = {
|
||||
|
|
@ -525,9 +459,7 @@ def test_text_only_streaming_has_index_zero():
|
|||
for chunk in chunks:
|
||||
parsed = iterator.chunk_parser(chunk)
|
||||
if parsed.choices:
|
||||
assert (
|
||||
parsed.choices[0].index == 0
|
||||
), f"Expected index=0, got {parsed.choices[0].index}"
|
||||
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
|
||||
|
||||
|
||||
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
|
||||
|
|
@ -704,9 +636,7 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
|
|||
]
|
||||
self._write_response(
|
||||
content_type="text/event-stream",
|
||||
body="".join(
|
||||
f"data: {json.dumps(event)}\n\n" for event in events
|
||||
).encode("utf-8"),
|
||||
body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode("utf-8"),
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -787,13 +717,9 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
|
|||
assert content_chunks == [answer_text]
|
||||
assert stream_usage is not None
|
||||
stream_completion_details = stream_usage["completion_tokens_details"]
|
||||
assert (
|
||||
stream_completion_details["reasoning_tokens"]
|
||||
== non_stream_details.reasoning_tokens
|
||||
)
|
||||
assert stream_completion_details["reasoning_tokens"] == non_stream_details.reasoning_tokens
|
||||
assert stream_completion_details["text_tokens"] == (
|
||||
stream_usage["completion_tokens"]
|
||||
- stream_completion_details["reasoning_tokens"]
|
||||
stream_usage["completion_tokens"] - stream_completion_details["reasoning_tokens"]
|
||||
)
|
||||
assert requests_seen == [
|
||||
{
|
||||
|
|
@ -885,9 +811,9 @@ def test_text_and_tool_streaming_has_index_zero():
|
|||
for chunk in chunks:
|
||||
parsed = iterator.chunk_parser(chunk)
|
||||
if parsed.choices:
|
||||
assert (
|
||||
parsed.choices[0].index == 0
|
||||
), f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
|
||||
assert parsed.choices[0].index == 0, (
|
||||
f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
|
||||
)
|
||||
|
||||
|
||||
def test_multiple_tools_streaming_has_index_zero():
|
||||
|
|
@ -940,15 +866,11 @@ def test_multiple_tools_streaming_has_index_zero():
|
|||
for chunk in chunks:
|
||||
parsed = iterator.chunk_parser(chunk)
|
||||
if parsed.choices:
|
||||
assert (
|
||||
parsed.choices[0].index == 0
|
||||
), f"Expected index=0, got {parsed.choices[0].index}"
|
||||
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
|
||||
|
||||
|
||||
def test_streaming_chunks_have_stable_ids():
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=False, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
|
||||
first_chunk = {
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
|
|
@ -973,9 +895,7 @@ def test_partial_json_chunk_accumulation():
|
|||
This tests the fix for https://github.com/BerriAI/litellm/issues/17473
|
||||
where network fragmentation can cause SSE data to arrive in partial chunks.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel'
|
||||
partial_chunk_2 = 'lo"}}'
|
||||
|
|
@ -983,31 +903,21 @@ def test_partial_json_chunk_accumulation():
|
|||
# First partial chunk should return None (still accumulating)
|
||||
result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}")
|
||||
assert result1 is None, "First partial chunk should return None while accumulating"
|
||||
assert (
|
||||
iterator.chunk_type == "accumulated_json"
|
||||
), "Should switch to accumulated_json mode"
|
||||
assert (
|
||||
iterator.accumulated_json == partial_chunk_1
|
||||
), "Should have accumulated first part"
|
||||
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
|
||||
assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part"
|
||||
|
||||
# Second partial chunk should complete the JSON and return a parsed result
|
||||
result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}")
|
||||
assert result2 is not None, "Second chunk should return parsed result"
|
||||
assert (
|
||||
iterator.accumulated_json == ""
|
||||
), "Buffer should be cleared after successful parse"
|
||||
assert (
|
||||
result2.choices[0].delta.content == "Hello"
|
||||
), f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
|
||||
assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse"
|
||||
assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
|
||||
|
||||
|
||||
def test_complete_json_chunk_no_accumulation():
|
||||
"""
|
||||
Test that complete JSON chunks are parsed immediately without accumulation.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}'
|
||||
|
||||
|
|
@ -1015,18 +925,14 @@ def test_complete_json_chunk_no_accumulation():
|
|||
assert result is not None, "Complete chunk should return parsed result immediately"
|
||||
assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode"
|
||||
assert iterator.accumulated_json == "", "Buffer should remain empty"
|
||||
assert (
|
||||
result.choices[0].delta.content == "Hello"
|
||||
), f"Expected 'Hello', got '{result.choices[0].delta.content}'"
|
||||
assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'"
|
||||
|
||||
|
||||
def test_multiple_partial_chunks_accumulation():
|
||||
"""
|
||||
Test that multiple partial chunks can be accumulated across several iterations.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Split a JSON chunk into three parts
|
||||
part1 = '{"type":"content_block_del'
|
||||
|
|
@ -1194,9 +1100,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
|
|||
The issue was that web_search_tool_result blocks have input_json_delta events with {}
|
||||
that were incorrectly being converted to tool calls.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Simulate the streaming sequence:
|
||||
# 1. server_tool_use block starts (web_search)
|
||||
|
|
@ -1271,9 +1175,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
|
|||
# Should have exactly 2 tool calls:
|
||||
# 1. From content_block_start (server_tool_use) with id and name
|
||||
# 2. From content_block_delta with the actual query
|
||||
assert (
|
||||
len(tool_calls_emitted) == 2
|
||||
), f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
|
||||
assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
|
||||
|
||||
# First tool call should have the id and name
|
||||
assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123"
|
||||
|
|
@ -1289,9 +1191,7 @@ def test_current_content_block_type_tracking():
|
|||
"""
|
||||
Test that current_content_block_type is properly tracked and reset.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Initially should be None
|
||||
assert iterator.current_content_block_type is None
|
||||
|
|
@ -1344,9 +1244,7 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
|
|||
The web_search_tool_result content comes ALL AT ONCE in content_block_start,
|
||||
not in deltas, so we need to capture it there.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Simulate the streaming sequence with web_search_tool_result
|
||||
chunks = [
|
||||
|
|
@ -1417,23 +1315,15 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
|
|||
and parsed.choices[0].delta.provider_specific_fields
|
||||
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
|
||||
):
|
||||
web_search_results = parsed.choices[0].delta.provider_specific_fields[
|
||||
"web_search_results"
|
||||
]
|
||||
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
|
||||
|
||||
# Verify web_search_results was captured
|
||||
assert web_search_results is not None, "web_search_results should be captured"
|
||||
assert len(web_search_results) == 1, "Should have 1 web_search_tool_result block"
|
||||
assert (
|
||||
web_search_results[0]["type"] == "web_search_tool_result"
|
||||
), "Block type should be web_search_tool_result"
|
||||
assert (
|
||||
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
|
||||
), "tool_use_id should match"
|
||||
assert web_search_results[0]["type"] == "web_search_tool_result", "Block type should be web_search_tool_result"
|
||||
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
|
||||
assert len(web_search_results[0]["content"]) == 2, "Should have 2 search results"
|
||||
assert (
|
||||
web_search_results[0]["content"][0]["title"] == "Fun Otter Facts"
|
||||
), "First result title should match"
|
||||
assert web_search_results[0]["content"][0]["title"] == "Fun Otter Facts", "First result title should match"
|
||||
|
||||
|
||||
def test_web_fetch_tool_result_captured_in_provider_specific_fields():
|
||||
|
|
@ -1447,9 +1337,7 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
|
|||
The web_fetch_tool_result content comes ALL AT ONCE in content_block_start,
|
||||
not in deltas, so we need to capture it there.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Simulate the streaming sequence with web_fetch_tool_result
|
||||
chunks = [
|
||||
|
|
@ -1520,25 +1408,15 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
|
|||
and parsed.choices[0].delta.provider_specific_fields
|
||||
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
|
||||
):
|
||||
web_search_results = parsed.choices[0].delta.provider_specific_fields[
|
||||
"web_search_results"
|
||||
]
|
||||
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
|
||||
|
||||
# Verify web_fetch_tool_result was captured (stored in web_search_results list)
|
||||
assert web_search_results is not None, "web_search_results should be captured"
|
||||
assert len(web_search_results) == 1, "Should have 1 web_fetch_tool_result block"
|
||||
assert (
|
||||
web_search_results[0]["type"] == "web_fetch_tool_result"
|
||||
), "Block type should be web_fetch_tool_result"
|
||||
assert (
|
||||
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
|
||||
), "tool_use_id should match"
|
||||
assert (
|
||||
web_search_results[0]["content"]["url"] == "https://example.com"
|
||||
), "URL should match"
|
||||
assert (
|
||||
web_search_results[0]["content"]["content"]["title"] == "Example Page"
|
||||
), "Title should match"
|
||||
assert web_search_results[0]["type"] == "web_fetch_tool_result", "Block type should be web_fetch_tool_result"
|
||||
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
|
||||
assert web_search_results[0]["content"]["url"] == "https://example.com", "URL should match"
|
||||
assert web_search_results[0]["content"]["content"]["title"] == "Example Page", "Title should match"
|
||||
|
||||
|
||||
def test_web_fetch_tool_result_no_extra_tool_calls():
|
||||
|
|
@ -1551,9 +1429,7 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
|
|||
The issue was that web_fetch_tool_result blocks have input_json_delta events with {}
|
||||
that were incorrectly being converted to tool calls.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# to verify it doesn't emit tool calls
|
||||
chunks = [
|
||||
|
|
@ -1597,9 +1473,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
|
|||
tool_call_count += 1
|
||||
|
||||
# Should have 0 tool calls - web_fetch_tool_result should not emit tool calls
|
||||
assert (
|
||||
tool_call_count == 0
|
||||
), f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
|
||||
assert tool_call_count == 0, (
|
||||
f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
|
||||
)
|
||||
|
||||
|
||||
def test_container_in_provider_specific_fields_streaming():
|
||||
|
|
@ -1609,9 +1485,7 @@ def test_container_in_provider_specific_fields_streaming():
|
|||
When container with skills is used, the container field should be present in
|
||||
the provider_specific_fields of the message_delta chunk.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
|
||||
|
||||
# Simulate streaming chunks
|
||||
chunks = [
|
||||
|
|
@ -1679,20 +1553,12 @@ def test_container_in_provider_specific_fields_streaming():
|
|||
and parsed.choices[0].delta.provider_specific_fields
|
||||
and "container" in parsed.choices[0].delta.provider_specific_fields
|
||||
):
|
||||
container_field = parsed.choices[0].delta.provider_specific_fields[
|
||||
"container"
|
||||
]
|
||||
container_field = parsed.choices[0].delta.provider_specific_fields["container"]
|
||||
|
||||
# Verify container was captured
|
||||
assert (
|
||||
container_field is not None
|
||||
), "container should be captured in provider_specific_fields"
|
||||
assert (
|
||||
container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p"
|
||||
), "container id should match"
|
||||
assert (
|
||||
container_field["expires_at"] == "2025-12-16T04:57:16.913181Z"
|
||||
), "expires_at should match"
|
||||
assert container_field is not None, "container should be captured in provider_specific_fields"
|
||||
assert container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p", "container id should match"
|
||||
assert container_field["expires_at"] == "2025-12-16T04:57:16.913181Z", "expires_at should match"
|
||||
assert len(container_field["skills"]) == 1, "Should have 1 skill"
|
||||
assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx"
|
||||
assert container_field["skills"][0]["version"] == "20251013", "version should match"
|
||||
|
|
@ -1705,9 +1571,7 @@ def test_container_in_provider_specific_fields_non_streaming():
|
|||
When container with skills is used in non-streaming, the container field should be
|
||||
present in the provider_specific_fields of the response.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=False, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
|
||||
|
||||
# Simulate a message_delta chunk with container (as it would appear in non-streaming)
|
||||
message_delta_chunk = {
|
||||
|
|
@ -1743,21 +1607,13 @@ def test_container_in_provider_specific_fields_non_streaming():
|
|||
# Verify container is in provider_specific_fields
|
||||
assert model_response.choices[0].delta.provider_specific_fields is not None
|
||||
assert "container" in model_response.choices[0].delta.provider_specific_fields
|
||||
container_field = model_response.choices[0].delta.provider_specific_fields[
|
||||
"container"
|
||||
]
|
||||
container_field = model_response.choices[0].delta.provider_specific_fields["container"]
|
||||
|
||||
assert container_field["id"] == "container_abc123xyz", "container id should match"
|
||||
assert (
|
||||
container_field["expires_at"] == "2025-12-20T10:30:00.000000Z"
|
||||
), "expires_at should match"
|
||||
assert container_field["expires_at"] == "2025-12-20T10:30:00.000000Z", "expires_at should match"
|
||||
assert len(container_field["skills"]) == 2, "Should have 2 skills"
|
||||
assert (
|
||||
container_field["skills"][0]["skill_id"] == "code_execution"
|
||||
), "First skill_id should be code_execution"
|
||||
assert (
|
||||
container_field["skills"][1]["skill_id"] == "pptx"
|
||||
), "Second skill_id should be pptx"
|
||||
assert container_field["skills"][0]["skill_id"] == "code_execution", "First skill_id should be code_execution"
|
||||
assert container_field["skills"][1]["skill_id"] == "pptx", "Second skill_id should be pptx"
|
||||
|
||||
|
||||
def test_container_absent_when_not_provided():
|
||||
|
|
@ -1766,9 +1622,7 @@ def test_container_absent_when_not_provided():
|
|||
|
||||
This ensures we don't add empty or None container fields.
|
||||
"""
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(), sync_stream=False, json_mode=False
|
||||
)
|
||||
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
|
||||
|
||||
# message_delta without container
|
||||
message_delta_chunk = {
|
||||
|
|
@ -1787,9 +1641,9 @@ def test_container_absent_when_not_provided():
|
|||
|
||||
# Verify container is NOT in provider_specific_fields when not provided
|
||||
if model_response.choices[0].delta.provider_specific_fields:
|
||||
assert (
|
||||
"container" not in model_response.choices[0].delta.provider_specific_fields
|
||||
), "container should not be present when not provided in delta"
|
||||
assert "container" not in model_response.choices[0].delta.provider_specific_fields, (
|
||||
"container should not be present when not provided in delta"
|
||||
)
|
||||
|
||||
|
||||
def test_streaming_code_execution_produces_code_interpreter_results():
|
||||
|
|
@ -1985,8 +1839,7 @@ def test_streaming_multiple_code_executions_no_duplicates():
|
|||
# Second (final) emission: cumulative list with BOTH results
|
||||
# This is what stream_chunk_builder will pick as "last value wins"
|
||||
assert len(emissions[1]) == 2, (
|
||||
f"Expected final emission to have 2 results, got {len(emissions[1])}. "
|
||||
f"IDs: {[r.id for r in emissions[1]]}"
|
||||
f"Expected final emission to have 2 results, got {len(emissions[1])}. IDs: {[r.id for r in emissions[1]]}"
|
||||
)
|
||||
assert emissions[1][0].id == "srvtoolu_01AAA"
|
||||
assert emissions[1][0].code == "echo first"
|
||||
|
|
@ -2150,9 +2003,7 @@ def test_empty_output_produces_null_outputs():
|
|||
assert code_results is not None, "No code_interpreter_results emitted"
|
||||
assert len(code_results) == 1
|
||||
assert code_results[0].id == "srvtoolu_01AAA"
|
||||
assert (
|
||||
code_results[0].outputs is None
|
||||
), f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
|
||||
assert code_results[0].outputs is None, f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
|
||||
|
||||
|
||||
def test_non_bash_tool_result_skipped():
|
||||
|
|
@ -2215,12 +2066,10 @@ def test_non_bash_tool_result_skipped():
|
|||
code_results = psf["code_interpreter_results"]
|
||||
|
||||
# code_interpreter_results should be emitted but empty (no bash results)
|
||||
assert (
|
||||
code_results is not None
|
||||
), "Expected code_interpreter_results key to be emitted"
|
||||
assert (
|
||||
len(code_results) == 0
|
||||
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
|
||||
assert code_results is not None, "Expected code_interpreter_results key to be emitted"
|
||||
assert len(code_results) == 0, (
|
||||
f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
|
||||
)
|
||||
|
||||
|
||||
class TestRustChatCompletionsHook:
|
||||
|
|
@ -2257,13 +2106,9 @@ class TestRustChatCompletionsHook:
|
|||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
monkeypatch.delenv("LITELLM_RUST", raising=False)
|
||||
bridge.set_rust_chat_completions(
|
||||
chat_completions=None, achat_completions=None, decline=None
|
||||
)
|
||||
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
|
||||
yield
|
||||
bridge.set_rust_chat_completions(
|
||||
chat_completions=None, achat_completions=None, decline=None
|
||||
)
|
||||
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
|
||||
|
||||
@staticmethod
|
||||
def _completion_kwargs(**overrides):
|
||||
|
|
@ -2361,9 +2206,7 @@ class TestRustChatCompletionsHook:
|
|||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(optional_params={"max_tokens": 7})
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={"max_tokens": 7}))
|
||||
assert seen["call"][0]["optional_params"]["max_tokens"] == 7
|
||||
|
||||
def test_without_the_opt_in_the_core_is_never_consulted(self):
|
||||
|
|
@ -2371,15 +2214,14 @@ class TestRustChatCompletionsHook:
|
|||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
seen = self._inject()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
) as transform, patch.object(
|
||||
AnthropicChatCompletion, "acompletion_function"
|
||||
with (
|
||||
patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
) as transform,
|
||||
patch.object(AnthropicChatCompletion, "acompletion_function"),
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(litellm_params={})
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(litellm_params={}))
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; reaching it is
|
||||
# the assertion, so the network failure below is expected.
|
||||
|
|
@ -2393,9 +2235,7 @@ class TestRustChatCompletionsHook:
|
|||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
seen = self._inject(decline_reason="unrecognized request parameter")
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs())
|
||||
except Exception:
|
||||
|
|
@ -2408,9 +2248,7 @@ class TestRustChatCompletionsHook:
|
|||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
seen = self._inject()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True})
|
||||
|
|
@ -2424,9 +2262,7 @@ class TestRustChatCompletionsHook:
|
|||
|
||||
seen = self._inject()
|
||||
logging_obj = MagicMock()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
|
||||
assert logging_obj.pre_call.call_count == 1
|
||||
assert len(seen["call"]) == 1
|
||||
|
||||
|
|
@ -2440,9 +2276,7 @@ class TestRustChatCompletionsHook:
|
|||
|
||||
self._inject()
|
||||
logging_obj = MagicMock()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
|
||||
|
||||
assert logging_obj.post_call.call_count == 1
|
||||
logged = logging_obj.post_call.call_args.kwargs["original_response"]
|
||||
|
|
@ -2467,18 +2301,12 @@ class TestRustChatCompletionsHook:
|
|||
raise _Declined("blank message text")
|
||||
|
||||
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
|
||||
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; the log count is
|
||||
# the assertion, so a failure past this point is expected.
|
||||
|
|
@ -2503,21 +2331,15 @@ class TestRustChatCompletionsHook:
|
|||
async def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=declining_native
|
||||
)
|
||||
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native)
|
||||
|
||||
sentinel = object()
|
||||
|
||||
async def python_path(**_kwargs):
|
||||
return sentinel
|
||||
|
||||
with patch.object(
|
||||
AnthropicChatCompletion, "acompletion_function", side_effect=python_path
|
||||
) as python_call:
|
||||
result = await AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(acompletion=True)
|
||||
)
|
||||
with patch.object(AnthropicChatCompletion, "acompletion_function", side_effect=python_path) as python_call:
|
||||
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
|
||||
|
||||
assert result is sentinel
|
||||
assert python_call.called, "a failing rust call must re-enter the python path"
|
||||
|
|
@ -2530,20 +2352,15 @@ class TestRustChatCompletionsHook:
|
|||
async def native(**_kwargs):
|
||||
return dict(self.RUST_RESPONSE)
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=native
|
||||
)
|
||||
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native)
|
||||
|
||||
with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call:
|
||||
result = await AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(acompletion=True)
|
||||
)
|
||||
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
|
||||
|
||||
assert result.choices[0].message.content == "hello from rust"
|
||||
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert not python_call.called
|
||||
|
||||
|
||||
def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch):
|
||||
"""One request, one pre_call, on the synchronous path too. Without the
|
||||
suppression the Python path logs a second time for the same attempt."""
|
||||
|
|
@ -2563,27 +2380,19 @@ class TestRustChatCompletionsHook:
|
|||
def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
|
||||
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; the log count is
|
||||
# the assertion, so a failure past this point is expected.
|
||||
pass
|
||||
|
||||
assert len(calls["pre_call"]) == 1
|
||||
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == (
|
||||
"claude-sonnet-4-5"
|
||||
)
|
||||
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ("claude-sonnet-4-5")
|
||||
|
||||
def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
|
||||
"""The suppression must not swallow the log on the ordinary path."""
|
||||
|
|
@ -2592,9 +2401,7 @@ class TestRustChatCompletionsHook:
|
|||
|
||||
self._inject()
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(litellm_params={}, logging_obj=logging_obj)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
|
||||
import pytest
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
|
@ -33,13 +32,9 @@ def test_response_format_transformation_unit_test():
|
|||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
result = config._create_json_tool_call_for_response_format(
|
||||
json_schema=response_format_json_schema
|
||||
)
|
||||
result = config._create_json_tool_call_for_response_format(json_schema=response_format_json_schema)
|
||||
|
||||
assert result["input_schema"]["properties"] == {
|
||||
"agent_doing": {"title": "Agent Doing", "type": "string"}
|
||||
}
|
||||
assert result["input_schema"]["properties"] == {"agent_doing": {"title": "Agent Doing", "type": "string"}}
|
||||
print(result)
|
||||
|
||||
|
||||
|
|
@ -550,9 +545,7 @@ def test_extract_response_content_with_citations():
|
|||
},
|
||||
}
|
||||
|
||||
_, citations, _, _, _, _, _, _ = config.extract_response_content(
|
||||
completion_response
|
||||
)
|
||||
_, citations, _, _, _, _, _, _ = config.extract_response_content(completion_response)
|
||||
assert citations == [
|
||||
[
|
||||
{
|
||||
|
|
@ -625,12 +618,8 @@ def test_web_search_tool_transformation():
|
|||
assert anthropic_web_search_tool["user_location"]["city"] == "San Francisco"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)]
|
||||
)
|
||||
def test_web_search_tool_transformation_with_search_context_size(
|
||||
search_context_size, expected_max_uses
|
||||
):
|
||||
@pytest.mark.parametrize("search_context_size, expected_max_uses", [("low", 1), ("medium", 5), ("high", 10)])
|
||||
def test_web_search_tool_transformation_with_search_context_size(search_context_size, expected_max_uses):
|
||||
from litellm.types.llms.openai import OpenAIWebSearchOptions
|
||||
|
||||
config = AnthropicConfig()
|
||||
|
|
@ -805,10 +794,7 @@ def test_web_search_tool_result_in_provider_specific_fields():
|
|||
assert "web_search_results" in provider_fields
|
||||
assert len(provider_fields["web_search_results"]) == 1
|
||||
assert provider_fields["web_search_results"][0]["type"] == "web_search_tool_result"
|
||||
assert (
|
||||
provider_fields["web_search_results"][0]["tool_use_id"]
|
||||
== "srvtoolu_provider_test"
|
||||
)
|
||||
assert provider_fields["web_search_results"][0]["tool_use_id"] == "srvtoolu_provider_test"
|
||||
|
||||
|
||||
def test_multiple_web_search_tool_results():
|
||||
|
|
@ -1032,10 +1018,7 @@ def test_transform_response_with_prefix_prompt():
|
|||
)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result.choices[0].message.content
|
||||
== "You are a helpful assistant. The grass is green."
|
||||
)
|
||||
assert result.choices[0].message.content == "You are a helpful assistant. The grass is green."
|
||||
|
||||
|
||||
def test_get_supported_params_thinking():
|
||||
|
|
@ -1150,18 +1133,12 @@ def test_anthropic_beta_header_merging_with_output_format():
|
|||
}
|
||||
}
|
||||
|
||||
result_headers = config.update_headers_with_optional_anthropic_beta(
|
||||
headers, optional_params
|
||||
)
|
||||
result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
|
||||
|
||||
# Both beta headers should be present
|
||||
beta_value = result_headers["anthropic-beta"]
|
||||
assert (
|
||||
"context-1m-2025-08-07" in beta_value
|
||||
), f"User's context-1m beta header missing from: {beta_value}"
|
||||
assert (
|
||||
"structured-outputs-2025-11-13" in beta_value
|
||||
), f"Structured output beta header missing from: {beta_value}"
|
||||
assert "context-1m-2025-08-07" in beta_value, f"User's context-1m beta header missing from: {beta_value}"
|
||||
assert "structured-outputs-2025-11-13" in beta_value, f"Structured output beta header missing from: {beta_value}"
|
||||
|
||||
|
||||
def test_anthropic_beta_header_merging_with_multiple_features():
|
||||
|
|
@ -1183,9 +1160,7 @@ def test_anthropic_beta_header_merging_with_multiple_features():
|
|||
"tools": [{"type": "web_fetch_20250910", "name": "web_fetch"}],
|
||||
}
|
||||
|
||||
result_headers = config.update_headers_with_optional_anthropic_beta(
|
||||
headers, optional_params
|
||||
)
|
||||
result_headers = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
|
||||
|
||||
beta_value = result_headers["anthropic-beta"]
|
||||
|
||||
|
|
@ -1228,9 +1203,7 @@ def test_anthropic_structured_output_beta_header():
|
|||
"strict": True,
|
||||
"schema": {
|
||||
"description": 'Progress report for the thinking process\n\nThis model represents a snapshot of the agent\'s current progress during\nthe thinking process, providing a brief description of the current activity.\n\nAttributes:\n agent_doing: Brief description of what the agent is currently doing.\n Should be kept under 10 words. Example: "Learning about home automation"',
|
||||
"properties": {
|
||||
"agent_doing": {"title": "Agent Doing", "type": "string"}
|
||||
},
|
||||
"properties": {"agent_doing": {"title": "Agent Doing", "type": "string"}},
|
||||
"required": ["agent_doing"],
|
||||
"title": "ThinkingStep",
|
||||
"type": "object",
|
||||
|
|
@ -1244,10 +1217,7 @@ def test_anthropic_structured_output_beta_header():
|
|||
assert response is not None
|
||||
print(f"response: {response}")
|
||||
print(f"raw_request_headers: {response['raw_request_headers']}")
|
||||
assert (
|
||||
"structured-outputs-2025-11-13"
|
||||
in response["raw_request_headers"]["anthropic-beta"]
|
||||
)
|
||||
assert "structured-outputs-2025-11-13" in response["raw_request_headers"]["anthropic-beta"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -1383,9 +1353,7 @@ def test_tool_search_regex_detection():
|
|||
config = AnthropicModelInfo()
|
||||
|
||||
# Test with tool search regex tool
|
||||
tools = [
|
||||
{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}
|
||||
]
|
||||
tools = [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}]
|
||||
assert config.is_tool_search_used(tools) is True
|
||||
|
||||
# Test without tool search
|
||||
|
|
@ -1400,9 +1368,7 @@ def test_tool_search_bm25_detection():
|
|||
config = AnthropicModelInfo()
|
||||
|
||||
# Test with tool search BM25 tool
|
||||
tools = [
|
||||
{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}
|
||||
]
|
||||
tools = [{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}]
|
||||
assert config.is_tool_search_used(tools) is True
|
||||
|
||||
|
||||
|
|
@ -1594,9 +1560,7 @@ def test_tool_search_complete_response_parsing():
|
|||
"tool_use_id": "srvtoolu_015i6aVA2niwzv4RG4DtnxDJ",
|
||||
"content": {
|
||||
"type": "tool_search_tool_search_result",
|
||||
"tool_references": [
|
||||
{"type": "tool_reference", "tool_name": "get_weather"}
|
||||
],
|
||||
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "Great! I found a weather tool."},
|
||||
|
|
@ -1647,9 +1611,7 @@ def test_tool_search_complete_response_parsing():
|
|||
|
||||
assert usage.server_tool_use is not None
|
||||
assert usage.server_tool_use.web_search_requests == 0
|
||||
assert (
|
||||
usage.server_tool_use.tool_search_requests == 1
|
||||
) # Counted from server_tool_use blocks
|
||||
assert usage.server_tool_use.tool_search_requests == 1 # Counted from server_tool_use blocks
|
||||
|
||||
|
||||
def test_allowed_callers_field_preservation():
|
||||
|
|
@ -1701,9 +1663,7 @@ def test_programmatic_tool_calling_beta_header():
|
|||
assert is_programmatic is True
|
||||
|
||||
# Test header generation
|
||||
headers = model_info.get_anthropic_headers(
|
||||
api_key="test-key", programmatic_tool_calling_used=True
|
||||
)
|
||||
headers = model_info.get_anthropic_headers(api_key="test-key", programmatic_tool_calling_used=True)
|
||||
|
||||
assert "anthropic-beta" in headers
|
||||
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
|
||||
|
|
@ -1847,9 +1807,7 @@ def test_input_examples_beta_header():
|
|||
assert is_examples_used is True
|
||||
|
||||
# Test header generation
|
||||
headers = model_info.get_anthropic_headers(
|
||||
api_key="test-key", input_examples_used=True
|
||||
)
|
||||
headers = model_info.get_anthropic_headers(api_key="test-key", input_examples_used=True)
|
||||
|
||||
assert "anthropic-beta" in headers
|
||||
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
|
||||
|
|
@ -1935,10 +1893,7 @@ def test_input_examples_empty_list_not_added():
|
|||
transformed_tool, _ = config._map_tool_helper(tool)
|
||||
assert transformed_tool is not None
|
||||
# Empty list should not be added
|
||||
assert (
|
||||
"input_examples" not in transformed_tool
|
||||
or len(transformed_tool.get("input_examples", [])) == 0
|
||||
)
|
||||
assert "input_examples" not in transformed_tool or len(transformed_tool.get("input_examples", [])) == 0
|
||||
|
||||
|
||||
# ============ Effort Parameter Tests ============
|
||||
|
|
@ -1998,9 +1953,7 @@ def test_effort_beta_header_injection():
|
|||
effort_used = model_info.is_effort_used(optional_params=optional_params, custom_llm_provider="anthropic")
|
||||
assert effort_used is True
|
||||
|
||||
headers = model_info.get_anthropic_headers(
|
||||
api_key="test-key", effort_used=effort_used
|
||||
)
|
||||
headers = model_info.get_anthropic_headers(api_key="test-key", effort_used=effort_used)
|
||||
|
||||
assert "anthropic-beta" in headers
|
||||
assert "effort-2025-11-24" in headers["anthropic-beta"]
|
||||
|
|
@ -2026,9 +1979,7 @@ def test_effort_validation():
|
|||
|
||||
optional_params = {"output_config": {"effort": "invalid"}}
|
||||
|
||||
with pytest.raises(
|
||||
litellm.exceptions.BadRequestError, match="Invalid effort value"
|
||||
):
|
||||
with pytest.raises(litellm.exceptions.BadRequestError, match="Invalid effort value"):
|
||||
config.transform_request(
|
||||
model="claude-opus-4-5-20251101",
|
||||
messages=messages,
|
||||
|
|
@ -2264,16 +2215,8 @@ def test_anthropic_model_supports_speed_param_rejects_non_anthropic_providers(
|
|||
):
|
||||
"""Fast mode is direct-Anthropic-only. Vertex/Azure/Bedrock strip their prefix
|
||||
before the shared transform runs, so the bare Opus id must still be rejected."""
|
||||
assert (
|
||||
AnthropicConfig._model_supports_speed_param(
|
||||
"claude-opus-4-8", custom_llm_provider
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic")
|
||||
is True
|
||||
)
|
||||
assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", custom_llm_provider) is False
|
||||
assert AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic") is True
|
||||
|
||||
|
||||
def test_vertex_anthropic_drops_speed_for_opus_with_drop_params(monkeypatch):
|
||||
|
|
@ -2572,9 +2515,7 @@ def test_supports_effort_level_handles_provider_prefixes(model, level, expected)
|
|||
("claude-opus-4-5-20251101", None, False),
|
||||
],
|
||||
)
|
||||
def test_validate_effort_for_model_centralises_per_model_gating(
|
||||
model, effort, expect_error
|
||||
):
|
||||
def test_validate_effort_for_model_centralises_per_model_gating(model, effort, expect_error):
|
||||
err = AnthropicConfig._validate_effort_for_model(model, effort, "anthropic")
|
||||
if expect_error:
|
||||
assert err is not None
|
||||
|
|
@ -2623,11 +2564,7 @@ def test_transform_request_injects_dummy_tool_without_tools_param():
|
|||
litellm.modify_params = prev_modify_params
|
||||
|
||||
assert "tools" in result
|
||||
names = [
|
||||
t.get("name")
|
||||
for t in result["tools"]
|
||||
if isinstance(t, dict) and t.get("name") is not None
|
||||
]
|
||||
names = [t.get("name") for t in result["tools"] if isinstance(t, dict) and t.get("name") is not None]
|
||||
assert "dummy_tool" in names
|
||||
|
||||
|
||||
|
|
@ -2716,13 +2653,9 @@ def test_calculate_usage_completion_tokens_details_with_reasoning():
|
|||
"output_tokens": 500,
|
||||
}
|
||||
# Simulating reasoning content that would count as ~50 tokens
|
||||
reasoning_content = (
|
||||
"Let me think about this step by step. " * 10
|
||||
) # Roughly 50 tokens
|
||||
reasoning_content = "Let me think about this step by step. " * 10 # Roughly 50 tokens
|
||||
|
||||
usage = config.calculate_usage(
|
||||
usage_object=usage_object, reasoning_content=reasoning_content
|
||||
)
|
||||
usage = config.calculate_usage(usage_object=usage_object, reasoning_content=reasoning_content)
|
||||
|
||||
# completion_tokens_details should be populated with both reasoning and text tokens
|
||||
assert usage.completion_tokens_details is not None
|
||||
|
|
@ -2773,9 +2706,7 @@ def test_reasoning_effort_maps_to_adaptive_thinking_for_claude_4_6_models():
|
|||
# reasoning_effort should not be in the result (it's transformed to thinking)
|
||||
assert "reasoning_effort" not in result
|
||||
# Should set output_config with the mapped effort value
|
||||
assert (
|
||||
"output_config" in result
|
||||
), f"output_config missing for {model} with effort={effort}"
|
||||
assert "output_config" in result, f"output_config missing for {model} with effort={effort}"
|
||||
assert result["output_config"]["effort"] == effort_map[effort]
|
||||
|
||||
|
||||
|
|
@ -2851,7 +2782,6 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model():
|
|||
assert result["thinking"] == {"type": "adaptive"}
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected",
|
||||
[
|
||||
|
|
@ -2877,9 +2807,7 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model():
|
|||
("gpt-4o", False),
|
||||
],
|
||||
)
|
||||
def test_is_adaptive_thinking_model_is_sourced_from_cost_map(
|
||||
local_model_cost_map, model, expected
|
||||
):
|
||||
def test_is_adaptive_thinking_model_is_sourced_from_cost_map(local_model_cost_map, model, expected):
|
||||
"""Adaptive thinking resolves from the cost map first (an explicit
|
||||
supports_adaptive_thinking entry, or the anthropic-claude fallback rule for unmapped
|
||||
future Claudes), then from a date-safe opus/sonnet/haiku >= 4.6 name version as a
|
||||
|
|
@ -2995,9 +2923,7 @@ def test_reasoning_effort_sets_output_config_for_46_models():
|
|||
drop_params=False,
|
||||
)
|
||||
|
||||
assert (
|
||||
"output_config" in result
|
||||
), f"output_config missing for {model} with effort={effort}"
|
||||
assert "output_config" in result, f"output_config missing for {model} with effort={effort}"
|
||||
assert result["output_config"]["effort"] == effort
|
||||
|
||||
|
||||
|
|
@ -3036,9 +2962,7 @@ def test_reasoning_effort_does_not_set_output_config_for_older_models():
|
|||
drop_params=False,
|
||||
)
|
||||
|
||||
assert (
|
||||
"output_config" not in result
|
||||
), f"output_config should not be set for {model}"
|
||||
assert "output_config" not in result, f"output_config should not be set for {model}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -3078,14 +3002,10 @@ def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort
|
|||
)
|
||||
|
||||
# thinking must be set (adaptive for 4.6+)
|
||||
assert (
|
||||
"thinking" in result
|
||||
), f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert result["thinking"]["type"] == "adaptive"
|
||||
# output_config must carry the mapped effort
|
||||
assert (
|
||||
"output_config" in result
|
||||
), f"output_config missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert "output_config" in result, f"output_config missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert result["output_config"]["effort"] == "low"
|
||||
|
||||
|
||||
|
|
@ -3114,16 +3034,13 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(
|
|||
drop_params=False,
|
||||
)
|
||||
|
||||
assert (
|
||||
"thinking" in result
|
||||
), f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert "thinking" in result, f"thinking missing for reasoning_effort={reasoning_effort_value!r}"
|
||||
assert result["thinking"]["type"] == "enabled"
|
||||
assert "budget_tokens" in result["thinking"]
|
||||
assert result["thinking"]["budget_tokens"] > 0
|
||||
# Older models must not get adaptive-thinking output_config
|
||||
assert "output_config" not in result, (
|
||||
f"output_config should not be set for non-adaptive model "
|
||||
f"(reasoning_effort={reasoning_effort_value!r})"
|
||||
f"output_config should not be set for non-adaptive model (reasoning_effort={reasoning_effort_value!r})"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3174,12 +3091,8 @@ def test_reasoning_effort_unparseable_dict_is_dropped(bad_value):
|
|||
model="claude-sonnet-4-6-20260219",
|
||||
drop_params=False,
|
||||
)
|
||||
assert (
|
||||
"thinking" not in result
|
||||
), f"thinking should not be set for bad value {bad_value!r}"
|
||||
assert (
|
||||
"output_config" not in result
|
||||
), f"output_config should not be set for bad value {bad_value!r}"
|
||||
assert "thinking" not in result, f"thinking should not be set for bad value {bad_value!r}"
|
||||
assert "output_config" not in result, f"output_config should not be set for bad value {bad_value!r}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -3310,9 +3223,7 @@ def test_reasoning_effort_garbage_raises_bad_request(effort):
|
|||
("max", DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET),
|
||||
],
|
||||
)
|
||||
def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model(
|
||||
effort, expected_budget
|
||||
):
|
||||
def test_reasoning_effort_xhigh_max_maps_to_budget_on_budget_model(effort, expected_budget):
|
||||
"""``xhigh`` / ``max`` extend the budget_tokens progression on budget-mode models."""
|
||||
config = AnthropicConfig()
|
||||
|
||||
|
|
@ -3459,17 +3370,11 @@ def test_code_execution_tool_results_extraction():
|
|||
|
||||
# Verify first tool call
|
||||
assert transformed_response.choices[0].message.tool_calls[0].id == "srvtoolu_01ABC"
|
||||
assert (
|
||||
transformed_response.choices[0].message.tool_calls[0].function.name
|
||||
== "bash_code_execution"
|
||||
)
|
||||
assert transformed_response.choices[0].message.tool_calls[0].function.name == "bash_code_execution"
|
||||
|
||||
# Verify second tool call
|
||||
assert transformed_response.choices[0].message.tool_calls[1].id == "srvtoolu_01DEF"
|
||||
assert (
|
||||
transformed_response.choices[0].message.tool_calls[1].function.name
|
||||
== "text_editor_code_execution"
|
||||
)
|
||||
assert transformed_response.choices[0].message.tool_calls[1].function.name == "text_editor_code_execution"
|
||||
|
||||
# Verify tool results are in provider_specific_fields
|
||||
provider_fields = transformed_response.choices[0].message.provider_specific_fields
|
||||
|
|
@ -3492,10 +3397,7 @@ def test_code_execution_tool_results_extraction():
|
|||
assert editor_result["content"]["is_file_update"] is False
|
||||
|
||||
# Verify text content is properly concatenated
|
||||
assert (
|
||||
"I'll calculate that for you."
|
||||
in transformed_response.choices[0].message.content
|
||||
)
|
||||
assert "I'll calculate that for you." in transformed_response.choices[0].message.content
|
||||
assert "Done!" in transformed_response.choices[0].message.content
|
||||
|
||||
|
||||
|
|
@ -3563,10 +3465,7 @@ def test_code_execution_tool_results_in_hidden_params():
|
|||
assert "provider_specific_fields" in hidden
|
||||
assert "tool_results" in hidden["provider_specific_fields"]
|
||||
assert len(hidden["provider_specific_fields"]["tool_results"]) == 1
|
||||
assert (
|
||||
hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"]
|
||||
== "hello\n"
|
||||
)
|
||||
assert hidden["provider_specific_fields"]["tool_results"][0]["content"]["stdout"] == "hello\n"
|
||||
|
||||
|
||||
def test_tool_search_tool_result_not_in_tool_results():
|
||||
|
|
@ -3762,10 +3661,7 @@ def test_compaction_block_in_provider_specific_fields():
|
|||
assert "compaction_blocks" in provider_fields
|
||||
assert len(provider_fields["compaction_blocks"]) == 1
|
||||
assert provider_fields["compaction_blocks"][0]["type"] == "compaction"
|
||||
assert (
|
||||
"Summary of the conversation"
|
||||
in provider_fields["compaction_blocks"][0]["content"]
|
||||
)
|
||||
assert "Summary of the conversation" in provider_fields["compaction_blocks"][0]["content"]
|
||||
|
||||
|
||||
def test_multiple_compaction_blocks():
|
||||
|
|
@ -3813,9 +3709,7 @@ def test_compaction_block_request_transformation():
|
|||
{"role": "user", "content": "What is the weather in San Francisco?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "I don't have access to real-time data."}
|
||||
],
|
||||
"content": [{"type": "text", "text": "I don't have access to real-time data."}],
|
||||
"provider_specific_fields": {
|
||||
"compaction_blocks": [
|
||||
{
|
||||
|
|
@ -3828,9 +3722,7 @@ def test_compaction_block_request_transformation():
|
|||
{"role": "user", "content": "What about New York?"},
|
||||
]
|
||||
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages, model="claude-opus-4-6", llm_provider="anthropic"
|
||||
)
|
||||
result = anthropic_messages_pt(messages=messages, model="claude-opus-4-6", llm_provider="anthropic")
|
||||
|
||||
# Find the assistant message
|
||||
assistant_message = None
|
||||
|
|
@ -3944,9 +3836,7 @@ def test_map_openai_context_management_to_anthropic():
|
|||
"instructions": "Focus on preserving code snippets",
|
||||
}
|
||||
]
|
||||
result = config.map_openai_context_management_to_anthropic(
|
||||
openai_format_with_instructions
|
||||
)
|
||||
result = config.map_openai_context_management_to_anthropic(openai_format_with_instructions)
|
||||
|
||||
assert result is not None
|
||||
assert result["edits"][0]["trigger"]["value"] == 150000
|
||||
|
|
@ -3973,9 +3863,7 @@ def test_map_openai_params_with_context_management():
|
|||
config = AnthropicConfig()
|
||||
|
||||
# Test with OpenAI list format
|
||||
non_default_params = {
|
||||
"context_management": [{"type": "compaction", "compact_threshold": 200000}]
|
||||
}
|
||||
non_default_params = {"context_management": [{"type": "compaction", "compact_threshold": 200000}]}
|
||||
optional_params = {}
|
||||
|
||||
result = config.map_openai_params(
|
||||
|
|
@ -4012,10 +3900,7 @@ def test_map_openai_params_with_context_management():
|
|||
)
|
||||
|
||||
assert "context_management" in result
|
||||
assert (
|
||||
result["context_management"]
|
||||
== non_default_params_anthropic["context_management"]
|
||||
)
|
||||
assert result["context_management"] == non_default_params_anthropic["context_management"]
|
||||
|
||||
|
||||
def test_cache_control_in_supported_params():
|
||||
|
|
@ -4126,10 +4011,7 @@ def test_compaction_block_empty_list_not_added():
|
|||
# Verify compaction_blocks is not in provider_specific_fields when there are none
|
||||
provider_fields = result.choices[0].message.provider_specific_fields
|
||||
if provider_fields:
|
||||
assert (
|
||||
"compaction_blocks" not in provider_fields
|
||||
or provider_fields.get("compaction_blocks") is None
|
||||
)
|
||||
assert "compaction_blocks" not in provider_fields or provider_fields.get("compaction_blocks") is None
|
||||
|
||||
|
||||
def test_fast_mode_beta_header():
|
||||
|
|
@ -4178,9 +4060,7 @@ def test_fast_mode_usage_calculation():
|
|||
"output_tokens": 500,
|
||||
}
|
||||
|
||||
usage = config.calculate_usage(
|
||||
usage_object=usage_object, reasoning_content=None, speed="fast"
|
||||
)
|
||||
usage = config.calculate_usage(usage_object=usage_object, reasoning_content=None, speed="fast")
|
||||
|
||||
assert usage.prompt_tokens == 1000
|
||||
assert usage.completion_tokens == 500
|
||||
|
|
@ -4201,9 +4081,7 @@ def test_fast_mode_cost_calculation():
|
|||
base_completion = 0.025
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.cost_calculation.generic_cost_per_token"
|
||||
) as mock_cost,
|
||||
patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost,
|
||||
patch("litellm.get_model_info") as mock_info,
|
||||
):
|
||||
mock_cost.return_value = (base_prompt, base_completion)
|
||||
|
|
@ -4243,9 +4121,7 @@ def test_fast_mode_with_inference_geo():
|
|||
base_completion = 0.025
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.cost_calculation.generic_cost_per_token"
|
||||
) as mock_cost,
|
||||
patch("litellm.llms.anthropic.cost_calculation.generic_cost_per_token") as mock_cost,
|
||||
patch("litellm.get_model_info") as mock_info,
|
||||
):
|
||||
mock_cost.return_value = (base_prompt, base_completion)
|
||||
|
|
@ -4436,9 +4312,7 @@ def test_map_tool_helper_enforces_object_type_when_missing():
|
|||
"name": "search_code",
|
||||
"description": "Search for code patterns",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search query"}
|
||||
},
|
||||
"properties": {"query": {"type": "string", "description": "Search query"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
|
|
@ -4451,9 +4325,9 @@ def test_map_tool_helper_enforces_object_type_when_missing():
|
|||
assert "properties" in result["input_schema"]
|
||||
assert "query" in result["input_schema"]["properties"]
|
||||
# Original parameters dict must not be modified in place
|
||||
assert (
|
||||
tool["function"]["parameters"] == original_params
|
||||
), "parameters dict was mutated; _map_tool_helper should not modify caller data"
|
||||
assert tool["function"]["parameters"] == original_params, (
|
||||
"parameters dict was mutated; _map_tool_helper should not modify caller data"
|
||||
)
|
||||
|
||||
|
||||
def test_map_tool_helper_enforces_object_type_when_wrong_type():
|
||||
|
|
@ -4479,13 +4353,13 @@ def test_map_tool_helper_enforces_object_type_when_wrong_type():
|
|||
result, _ = config._map_tool_helper(tool)
|
||||
assert result is not None
|
||||
assert result["input_schema"]["type"] == "object"
|
||||
assert (
|
||||
result["input_schema"].get("properties") == {}
|
||||
), "properties should be injected as {} when schema has non-object type and no properties key"
|
||||
assert result["input_schema"].get("properties") == {}, (
|
||||
"properties should be injected as {} when schema has non-object type and no properties key"
|
||||
)
|
||||
# Original parameters dict must not be modified in place
|
||||
assert (
|
||||
tool["function"]["parameters"] == original_params
|
||||
), "parameters dict was mutated; _map_tool_helper should not modify caller data"
|
||||
assert tool["function"]["parameters"] == original_params, (
|
||||
"parameters dict was mutated; _map_tool_helper should not modify caller data"
|
||||
)
|
||||
|
||||
|
||||
def test_map_tool_helper_preserves_valid_object_schema():
|
||||
|
|
@ -4552,12 +4426,8 @@ def test_extract_response_content_thinking_block_null_thinking():
|
|||
{"type": "text", "text": "Hello"},
|
||||
]
|
||||
}
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
|
||||
completion_response_null
|
||||
)
|
||||
assert (
|
||||
thinking_blocks is not None
|
||||
), "thinking blocks should not be None when thinking=null"
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_null)
|
||||
assert thinking_blocks is not None, "thinking blocks should not be None when thinking=null"
|
||||
assert len(thinking_blocks) == 1
|
||||
assert "Hello" in text
|
||||
|
||||
|
|
@ -4568,12 +4438,8 @@ def test_extract_response_content_thinking_block_null_thinking():
|
|||
{"type": "text", "text": "World"},
|
||||
]
|
||||
}
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
|
||||
completion_response_missing
|
||||
)
|
||||
assert (
|
||||
thinking_blocks is not None
|
||||
), "thinking blocks should not be None when thinking key is absent"
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_missing)
|
||||
assert thinking_blocks is not None, "thinking blocks should not be None when thinking key is absent"
|
||||
assert len(thinking_blocks) == 1
|
||||
assert "World" in text
|
||||
|
||||
|
|
@ -4584,9 +4450,7 @@ def test_extract_response_content_thinking_block_null_thinking():
|
|||
{"type": "text", "text": "Done"},
|
||||
]
|
||||
}
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(
|
||||
completion_response_text
|
||||
)
|
||||
text, _, thinking_blocks, _, _, _, _, _ = config.extract_response_content(completion_response_text)
|
||||
assert thinking_blocks is not None
|
||||
assert len(thinking_blocks) == 1
|
||||
assert thinking_blocks[0]["thinking"] == "Let me think..."
|
||||
|
|
@ -4645,12 +4509,8 @@ def test_advisor_beta_header_injected():
|
|||
}
|
||||
]
|
||||
}
|
||||
result = config.update_headers_with_optional_anthropic_beta(
|
||||
headers, optional_params
|
||||
)
|
||||
assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get(
|
||||
"anthropic-beta", ""
|
||||
)
|
||||
result = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
|
||||
assert ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value in result.get("anthropic-beta", "")
|
||||
|
||||
|
||||
def test_advisor_beta_header_not_injected_without_tool():
|
||||
|
|
@ -4658,9 +4518,7 @@ def test_advisor_beta_header_not_injected_without_tool():
|
|||
config = AnthropicConfig()
|
||||
headers: dict = {}
|
||||
optional_params: dict = {"tools": []}
|
||||
result = config.update_headers_with_optional_anthropic_beta(
|
||||
headers, optional_params
|
||||
)
|
||||
result = config.update_headers_with_optional_anthropic_beta(headers, optional_params)
|
||||
assert "advisor-tool-2026-03-01" not in result.get("anthropic-beta", "")
|
||||
|
||||
|
||||
|
|
@ -4687,9 +4545,7 @@ def test_advisor_tool_result_preserved_in_response():
|
|||
{"type": "text", "text": "Here is the implementation."},
|
||||
]
|
||||
}
|
||||
text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content(
|
||||
completion_response
|
||||
)
|
||||
text, _, _, _, tool_calls, _, tool_results, _ = config.extract_response_content(completion_response)
|
||||
assert "Consulting advisor." in text
|
||||
assert "Here is the implementation." in text
|
||||
# server_tool_use (advisor) should be a tool_call
|
||||
|
|
@ -4804,9 +4660,7 @@ def test_basic_sanitize_anthropic_tool_name_replaces_invalid_chars():
|
|||
)
|
||||
|
||||
assert (
|
||||
_basic_sanitize_anthropic_tool_name(
|
||||
"github_openapi_mcp-actions/download-job-logs-for-workflow-run"
|
||||
)
|
||||
_basic_sanitize_anthropic_tool_name("github_openapi_mcp-actions/download-job-logs-for-workflow-run")
|
||||
== "github_openapi_mcp-actions_download-job-logs-for-workflow-run"
|
||||
)
|
||||
# other punctuation
|
||||
|
|
@ -4835,9 +4689,7 @@ def test_build_anthropic_tool_name_maps_no_collisions():
|
|||
]
|
||||
)
|
||||
assert forward == {
|
||||
"actions/download-job-logs-for-workflow-run": (
|
||||
"actions_download-job-logs-for-workflow-run"
|
||||
),
|
||||
"actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run"),
|
||||
"pulls/list-files": "pulls_list-files",
|
||||
}
|
||||
assert reverse == {v: k for k, v in forward.items()}
|
||||
|
|
@ -4888,9 +4740,7 @@ def test_build_anthropic_tool_name_maps_three_way_collision():
|
|||
_build_anthropic_tool_name_maps,
|
||||
)
|
||||
|
||||
forward, reverse = _build_anthropic_tool_name_maps(
|
||||
["foo_bar", "foo/bar", "foo.bar"]
|
||||
)
|
||||
forward, reverse = _build_anthropic_tool_name_maps(["foo_bar", "foo/bar", "foo.bar"])
|
||||
assert "foo_bar" not in forward # untouched
|
||||
assert forward["foo/bar"] == "foo_bar_2"
|
||||
assert forward["foo.bar"] == "foo_bar_3"
|
||||
|
|
@ -4963,16 +4813,13 @@ def test_map_openai_params_does_not_pollute_optional_params_with_internal_keys()
|
|||
)
|
||||
# No internal keys may appear in optional_params for ANY input.
|
||||
for key in optional_params:
|
||||
assert not key.startswith(
|
||||
"_anthropic_tool_name"
|
||||
), f"optional_params leaked internal key {key!r}: {optional_params}"
|
||||
assert not key.startswith("_anthropic_tool_name"), (
|
||||
f"optional_params leaked internal key {key!r}: {optional_params}"
|
||||
)
|
||||
# And no key starting with `_` either; optional_params should only
|
||||
# contain documented Anthropic Messages API parameters.
|
||||
for key in optional_params:
|
||||
assert not key.startswith("_"), (
|
||||
f"optional_params leaked underscore-prefixed key {key!r}: "
|
||||
f"{optional_params}"
|
||||
)
|
||||
assert not key.startswith("_"), f"optional_params leaked underscore-prefixed key {key!r}: {optional_params}"
|
||||
|
||||
|
||||
def test_map_openai_params_no_maps_when_all_names_already_valid():
|
||||
|
|
@ -5001,11 +4848,7 @@ def test_map_openai_params_no_maps_when_all_names_already_valid():
|
|||
|
||||
def test_rewrite_tool_names_in_messages_uses_forward_map():
|
||||
config = AnthropicConfig()
|
||||
forward_map = {
|
||||
"actions/download-job-logs-for-workflow-run": (
|
||||
"actions_download-job-logs-for-workflow-run"
|
||||
)
|
||||
}
|
||||
forward_map = {"actions/download-job-logs-for-workflow-run": ("actions_download-job-logs-for-workflow-run")}
|
||||
messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{
|
||||
|
|
@ -5028,15 +4871,9 @@ def test_rewrite_tool_names_in_messages_uses_forward_map():
|
|||
out = config._rewrite_tool_names_in_messages(messages, forward_map)
|
||||
|
||||
# input list must not be mutated
|
||||
assert (
|
||||
messages[1]["tool_calls"][0]["function"]["name"]
|
||||
== "actions/download-job-logs-for-workflow-run"
|
||||
)
|
||||
assert messages[1]["tool_calls"][0]["function"]["name"] == "actions/download-job-logs-for-workflow-run"
|
||||
# output rewritten according to forward map
|
||||
assert (
|
||||
out[1]["tool_calls"][0]["function"]["name"]
|
||||
== "actions_download-job-logs-for-workflow-run"
|
||||
)
|
||||
assert out[1]["tool_calls"][0]["function"]["name"] == "actions_download-job-logs-for-workflow-run"
|
||||
# non-tool-call messages pass through unchanged (same object)
|
||||
assert out[0] is messages[0]
|
||||
assert out[2] is messages[2]
|
||||
|
|
@ -5112,9 +4949,7 @@ def test_sanitize_tool_names_in_request_does_not_mutate_caller_tool_dicts():
|
|||
caller_tools = [caller_tool]
|
||||
optional_params: dict = {"tools": caller_tools}
|
||||
|
||||
forward, reverse = config._sanitize_tool_names_in_request(
|
||||
optional_params=optional_params
|
||||
)
|
||||
forward, reverse = config._sanitize_tool_names_in_request(optional_params=optional_params)
|
||||
|
||||
assert forward.get(original_name)
|
||||
sanitized = forward[original_name]
|
||||
|
|
@ -5263,10 +5098,7 @@ def test_streaming_iterator_reverse_maps_tool_use_name():
|
|||
parsed = iterator.chunk_parser(chunk=chunk)
|
||||
tool_calls = parsed.choices[0].delta.tool_calls
|
||||
assert tool_calls is not None and len(tool_calls) == 1
|
||||
assert (
|
||||
tool_calls[0]["function"]["name"]
|
||||
== "actions/download-job-logs-for-workflow-run"
|
||||
)
|
||||
assert tool_calls[0]["function"]["name"] == "actions/download-job-logs-for-workflow-run"
|
||||
|
||||
|
||||
def test_streaming_iterator_passthrough_when_name_not_in_map():
|
||||
|
|
@ -5362,9 +5194,9 @@ def test_transform_request_does_not_leak_internal_keys_into_body():
|
|||
for tool in data.get("tools", []):
|
||||
name = tool.get("name")
|
||||
assert isinstance(name, str)
|
||||
assert _re.fullmatch(
|
||||
r"[a-zA-Z0-9_-]{1,128}", name
|
||||
), f"sanitized tool name {name!r} still violates Anthropic regex"
|
||||
assert _re.fullmatch(r"[a-zA-Z0-9_-]{1,128}", name), (
|
||||
f"sanitized tool name {name!r} still violates Anthropic regex"
|
||||
)
|
||||
|
||||
# Sent name for the bad tool is the disambiguated form, valid name passes through.
|
||||
sent_names = {t["name"] for t in data["tools"]}
|
||||
|
|
@ -5500,9 +5332,7 @@ def test_transform_request_rewrites_tool_names_in_history():
|
|||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "tool_use":
|
||||
tool_use_names.append(block.get("name"))
|
||||
assert (
|
||||
tool_use_names
|
||||
), "expected at least one tool_use block in transformed messages"
|
||||
assert tool_use_names, "expected at least one tool_use block in transformed messages"
|
||||
for name in tool_use_names:
|
||||
assert name == "actions_download-job-logs-for-workflow-run", (
|
||||
f"history tool_use.name {name!r} not rewritten -- Anthropic will "
|
||||
|
|
@ -5526,19 +5356,12 @@ def test_sanitize_tool_names_in_request_skips_hosted_tools():
|
|||
}
|
||||
forward, reverse = AnthropicConfig._sanitize_tool_names_in_request(optional_params)
|
||||
# Only the custom tool was rewritten.
|
||||
assert forward == {
|
||||
"actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run"
|
||||
}
|
||||
assert reverse == {
|
||||
"actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run"
|
||||
}
|
||||
assert forward == {"actions/download-job-logs-for-workflow-run": "actions_download-job-logs-for-workflow-run"}
|
||||
assert reverse == {"actions_download-job-logs-for-workflow-run": "actions/download-job-logs-for-workflow-run"}
|
||||
# Hosted tool's name unchanged.
|
||||
assert optional_params["tools"][0]["name"] == "web_search"
|
||||
# Custom tool's name updated in place.
|
||||
assert (
|
||||
optional_params["tools"][1]["name"]
|
||||
== "actions_download-job-logs-for-workflow-run"
|
||||
)
|
||||
assert optional_params["tools"][1]["name"] == "actions_download-job-logs-for-workflow-run"
|
||||
|
||||
|
||||
def test_sanitize_tool_names_in_request_no_tools_is_noop():
|
||||
|
|
@ -5772,9 +5595,7 @@ def test_translate_system_message_keeps_billing_header_for_first_party_anthropic
|
|||
assert config.should_strip_billing_metadata() is False
|
||||
|
||||
result = config.translate_system_message(
|
||||
messages=_system_with_billing_header(
|
||||
"You are Claude Code, Anthropic's official CLI for Claude."
|
||||
)
|
||||
messages=_system_with_billing_header("You are Claude Code, Anthropic's official CLI for Claude.")
|
||||
)
|
||||
|
||||
texts = [block["text"] for block in result]
|
||||
|
|
@ -5790,9 +5611,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock():
|
|||
config = BedrockClaudePlatformConfig()
|
||||
assert config.should_strip_billing_metadata() is True
|
||||
|
||||
result = config.translate_system_message(
|
||||
messages=_system_with_billing_header("real system prompt")
|
||||
)
|
||||
result = config.translate_system_message(messages=_system_with_billing_header("real system prompt"))
|
||||
|
||||
texts = [block["text"] for block in result]
|
||||
assert all(not t.startswith("x-anthropic-billing-header:") for t in texts)
|
||||
|
|
@ -5858,9 +5677,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke():
|
|||
config = AmazonAnthropicClaudeConfig()
|
||||
assert config.should_strip_billing_metadata() is True
|
||||
|
||||
result = config.translate_system_message(
|
||||
messages=_system_with_billing_header("real system prompt")
|
||||
)
|
||||
result = config.translate_system_message(messages=_system_with_billing_header("real system prompt"))
|
||||
|
||||
texts = [block["text"] for block in result]
|
||||
assert all(not t.startswith("x-anthropic-billing-header:") for t in texts)
|
||||
|
|
@ -5914,9 +5731,7 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke():
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_should_strip_billing_metadata_by_provider(
|
||||
module_path, class_name, expected_strip
|
||||
):
|
||||
def test_should_strip_billing_metadata_by_provider(module_path, class_name, expected_strip):
|
||||
import importlib
|
||||
|
||||
config_cls = getattr(importlib.import_module(module_path), class_name)
|
||||
|
|
@ -6088,12 +5903,8 @@ def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch):
|
|||
"""The drop/raise decision must come from ``supports_sampling_params`` in
|
||||
the model map, not just name matching: a flagged entry gates a model whose
|
||||
name says nothing, and an explicit ``true`` overrides the name fallback."""
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True}
|
||||
)
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True})
|
||||
config = AnthropicConfig()
|
||||
|
||||
flagged_off = config.map_openai_params(
|
||||
|
|
@ -6213,9 +6024,7 @@ def test_is_anthropic_usage_object_rejects_responses_api_usage():
|
|||
("claude-sonnet-4-5-20250929", False),
|
||||
],
|
||||
)
|
||||
def test_disabled_thinking_omitted_only_for_always_on_models(
|
||||
local_model_cost_map, model, expected_dropped
|
||||
):
|
||||
def test_disabled_thinking_omitted_only_for_always_on_models(local_model_cost_map, model, expected_dropped):
|
||||
"""``thinking={"type": "disabled"}`` is omitted for always-on-thinking models
|
||||
(Fable/Mythos, which 400 on it: the API remedy is to omit the param) and is
|
||||
forwarded verbatim for every model that accepts it."""
|
||||
|
|
|
|||
|
|
@ -136,22 +136,14 @@ def test_in_place_substitution_preserves_ordering():
|
|||
responses_output = [msg_item, fc_exec1, fc_regular, fc_exec2]
|
||||
|
||||
# Apply the same logic as _transform_chat_completion_choices_to_responses_output
|
||||
tool_result_items = (
|
||||
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
|
||||
)
|
||||
tool_result_items = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
|
||||
if tool_result_items:
|
||||
result_by_id = {
|
||||
(item.get("id") if isinstance(item, dict) else item.id): item
|
||||
for item in tool_result_items
|
||||
}
|
||||
result_by_id = {(item.get("id") if isinstance(item, dict) else item.id): item for item in tool_result_items}
|
||||
replaced_ids = set(result_by_id.keys())
|
||||
responses_output = [
|
||||
(
|
||||
result_by_id[getattr(item, "call_id", None)]
|
||||
if (
|
||||
getattr(item, "type", None) == "function_call"
|
||||
and getattr(item, "call_id", None) in replaced_ids
|
||||
)
|
||||
if (getattr(item, "type", None) == "function_call" and getattr(item, "call_id", None) in replaced_ids)
|
||||
else item
|
||||
)
|
||||
for item in responses_output
|
||||
|
|
@ -255,9 +247,7 @@ def test_end_to_end_streaming_chunks_to_code_interpreter_output():
|
|||
assert code_results[0]["code"] == "echo e2e_test"
|
||||
|
||||
# Step 3: Extract via _extract_tool_result_output_items (Responses API layer)
|
||||
tool_result_items = (
|
||||
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(assembled)
|
||||
)
|
||||
tool_result_items = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(assembled)
|
||||
assert len(tool_result_items) == 1
|
||||
item = tool_result_items[0]
|
||||
# Items are reconstructed as Pydantic OutputCodeInterpreterCall objects
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ import pytest
|
|||
import litellm
|
||||
|
||||
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
TOOL_RESULT_IMAGE_PLACEHOLDER,
|
||||
)
|
||||
|
|
@ -53,9 +52,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
|
|||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call_d581d130-e234-4315-94e8-27e7ff7c4e55",
|
||||
function=Function(
|
||||
arguments='{"location": "Boston"}', name="get_weather"
|
||||
),
|
||||
function=Function(arguments='{"location": "Boston"}', name="get_weather"),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
|
|
@ -69,9 +66,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
print(content_block_start)
|
||||
|
||||
|
|
@ -101,9 +96,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
|
|||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=combined,
|
||||
function=Function(
|
||||
arguments='{"a": 17, "b": 25}', name="add_numbers"
|
||||
),
|
||||
function=Function(arguments='{"a": 17, "b": 25}', name="add_numbers"),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
|
|
@ -117,9 +110,7 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "tool_use"
|
||||
assert content_block_start["id"] == base
|
||||
|
|
@ -164,9 +155,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block():
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "thinking"
|
||||
assert content_block_start == {
|
||||
|
|
@ -202,9 +191,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_only_co
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "thinking"
|
||||
assert content_block_start == {
|
||||
|
|
@ -250,9 +237,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block(
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "thinking"
|
||||
assert content_block_start == {
|
||||
|
|
@ -305,9 +290,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_an
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "thinking"
|
||||
|
||||
|
|
@ -350,10 +333,7 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks():
|
|||
assert "thinking_blocks" in result[1]
|
||||
assert len(result[1]["thinking_blocks"]) == 2
|
||||
assert result[1]["thinking_blocks"][0]["type"] == "thinking"
|
||||
assert (
|
||||
result[1]["thinking_blocks"][0]["thinking"]
|
||||
== "I will call the get_weather tool."
|
||||
)
|
||||
assert result[1]["thinking_blocks"][0]["thinking"] == "I will call the get_weather tool."
|
||||
assert result[1]["thinking_blocks"][0]["signature"] == "sigsig"
|
||||
assert result[1]["thinking_blocks"][1]["type"] == "redacted_thinking"
|
||||
assert result[1]["thinking_blocks"][1]["data"] == "REDACTED"
|
||||
|
|
@ -456,9 +436,7 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement():
|
|||
|
||||
assert tool_message_idx is not None, "Tool message not found"
|
||||
assert user_message_idx is not None, "User message not found"
|
||||
assert (
|
||||
tool_message_idx < user_message_idx
|
||||
), "Tool message should be placed before user message"
|
||||
assert tool_message_idx < user_message_idx, "Tool message should be placed before user message"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -733,9 +711,7 @@ def test_translate_anthropic_to_openai_skips_prompt_cache_key_when_provider_lack
|
|||
|
||||
|
||||
def test_translate_anthropic_to_openai_skips_prompt_cache_key_for_chained_litellm_proxy():
|
||||
assert "prompt_cache_key" in litellm.get_supported_openai_params(
|
||||
model="xai", custom_llm_provider="litellm_proxy"
|
||||
)
|
||||
assert "prompt_cache_key" in litellm.get_supported_openai_params(model="xai", custom_llm_provider="litellm_proxy")
|
||||
openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": "session-abc"}, "litellm_proxy")
|
||||
assert openai_request["user"] == "session-abc"
|
||||
assert "prompt_cache_key" not in openai_request
|
||||
|
|
@ -780,7 +756,8 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
|
|||
id="call_empty_args",
|
||||
type="function",
|
||||
function=Function(
|
||||
name="test_function", arguments="" # empty arguments string
|
||||
name="test_function",
|
||||
arguments="", # empty arguments string
|
||||
),
|
||||
)
|
||||
],
|
||||
|
|
@ -795,9 +772,7 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
|
|||
assert result[0]["type"] == "tool_use"
|
||||
assert result[0]["id"] == "call_empty_args"
|
||||
assert result[0]["name"] == "test_function"
|
||||
assert (
|
||||
result[0]["input"] == {}
|
||||
), "Empty function arguments should result in empty dict"
|
||||
assert result[0]["input"] == {}, "Empty function arguments should result in empty dict"
|
||||
|
||||
|
||||
def test_translate_openai_content_to_anthropic_text_and_tool_calls():
|
||||
|
|
@ -917,9 +892,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
|
|||
ChatCompletionAssistantToolCall(
|
||||
id="call_tool_combo",
|
||||
type="function",
|
||||
function=Function(
|
||||
name="get_weather", arguments='{"location": "Paris"}'
|
||||
),
|
||||
function=Function(name="get_weather", arguments='{"location": "Paris"}'),
|
||||
)
|
||||
],
|
||||
),
|
||||
|
|
@ -929,9 +902,7 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
|
|||
)
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
anthropic_response = adapter.translate_openai_response_to_anthropic(
|
||||
response=openai_response
|
||||
)
|
||||
anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
|
||||
|
||||
anthropic_content = anthropic_response.get("content")
|
||||
assert anthropic_content is not None
|
||||
|
|
@ -972,9 +943,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_partial_json():
|
|||
(
|
||||
type_of_content,
|
||||
content_block_delta,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
|
||||
|
||||
print("Type of content:", type_of_content)
|
||||
print("Content block delta:", content_block_delta)
|
||||
|
|
@ -1083,9 +1052,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta():
|
|||
(
|
||||
type_of_content,
|
||||
content_block_delta,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
|
||||
|
||||
assert type_of_content == "thinking_delta"
|
||||
assert content_block_delta["type"] == "thinking_delta"
|
||||
|
|
@ -1128,9 +1095,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking():
|
|||
(
|
||||
type_of_content,
|
||||
content_block_delta,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
|
||||
|
||||
assert type_of_content == "signature_delta"
|
||||
assert content_block_delta["type"] == "signature_delta"
|
||||
|
|
@ -1194,9 +1159,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thin
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "thinking"
|
||||
|
||||
|
|
@ -1236,9 +1199,7 @@ def test_translate_anthropic_messages_to_openai_user_message_with_base64_image()
|
|||
# Check image content
|
||||
assert result[0]["content"][1]["type"] == "image_url"
|
||||
assert "image_url" in result[0]["content"][1]
|
||||
assert result[0]["content"][1]["image_url"]["url"].startswith(
|
||||
"data:image/png;base64,"
|
||||
)
|
||||
assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
|
||||
assert (
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
in result[0]["content"][1]["image_url"]["url"]
|
||||
|
|
@ -1276,18 +1237,14 @@ def test_translate_anthropic_messages_to_openai_user_message_with_url_image():
|
|||
# Check image content
|
||||
assert result[0]["content"][1]["type"] == "image_url"
|
||||
assert "image_url" in result[0]["content"][1]
|
||||
assert (
|
||||
result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
|
||||
)
|
||||
assert result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
|
||||
|
||||
|
||||
def test_translate_anthropic_messages_to_openai_tool_result_with_base64_image():
|
||||
"""Test that base64 images in tool results are correctly translated to OpenAI format."""
|
||||
|
||||
anthropic_messages = [
|
||||
AnthropicMessagesUserMessageParam(
|
||||
role="user", content=[{"type": "text", "text": "Take a screenshot"}]
|
||||
),
|
||||
AnthropicMessagesUserMessageParam(role="user", content=[{"type": "text", "text": "Take a screenshot"}]),
|
||||
AnthopicMessagesAssistantMessageParam(
|
||||
role="assistant",
|
||||
content=[
|
||||
|
|
@ -1439,9 +1396,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
|
|||
|
||||
# Check first image (base64)
|
||||
assert result[0]["content"][1]["type"] == "image_url"
|
||||
assert result[0]["content"][1]["image_url"]["url"].startswith(
|
||||
"data:image/png;base64,"
|
||||
)
|
||||
assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
|
||||
|
||||
# Check middle text
|
||||
assert result[0]["content"][2]["type"] == "text"
|
||||
|
|
@ -1449,9 +1404,7 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
|
|||
|
||||
# Check second image (URL)
|
||||
assert result[0]["content"][3]["type"] == "image_url"
|
||||
assert (
|
||||
result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
|
||||
)
|
||||
assert result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
|
||||
|
||||
# Check final text
|
||||
assert result[0]["content"][4]["type"] == "text"
|
||||
|
|
@ -1497,10 +1450,7 @@ def test_translate_anthropic_messages_to_openai_tool_use_with_signature():
|
|||
assert tool_call["id"] == "call_386f67af31f9415781bc35071405"
|
||||
assert "function" in tool_call
|
||||
assert "provider_specific_fields" in tool_call["function"]
|
||||
assert (
|
||||
tool_call["function"]["provider_specific_fields"]["thought_signature"]
|
||||
== test_signature
|
||||
)
|
||||
assert tool_call["function"]["provider_specific_fields"]["thought_signature"] == test_signature
|
||||
|
||||
|
||||
def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_content_items():
|
||||
|
|
@ -1558,9 +1508,7 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
|
|||
result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
|
||||
|
||||
# Count how many tool messages have the same tool_call_id
|
||||
tool_messages = [
|
||||
msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
|
||||
]
|
||||
tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
|
||||
tool_call_ids = [msg.get("tool_call_id") for msg in tool_messages]
|
||||
|
||||
# The critical assertion: each tool_call_id should appear only ONCE
|
||||
|
|
@ -1576,12 +1524,8 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
|
|||
# The content should be a list with all items combined
|
||||
tool_message = tool_messages[0]
|
||||
assert tool_message["tool_call_id"] == "toolu_016hYHBkTf4JDF3p22UoYk5C"
|
||||
assert isinstance(
|
||||
tool_message["content"], list
|
||||
), "Multiple content items should be combined into a list"
|
||||
assert (
|
||||
len(tool_message["content"]) == 3
|
||||
), f"Expected 3 content items, got {len(tool_message['content'])}"
|
||||
assert isinstance(tool_message["content"], list), "Multiple content items should be combined into a list"
|
||||
assert len(tool_message["content"]) == 3, f"Expected 3 content items, got {len(tool_message['content'])}"
|
||||
|
||||
# Verify content types
|
||||
assert tool_message["content"][0]["type"] == "text"
|
||||
|
|
@ -1630,17 +1574,14 @@ def test_translate_anthropic_messages_to_openai_tool_result_single_item_backward
|
|||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
|
||||
|
||||
tool_messages = [
|
||||
msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
|
||||
]
|
||||
tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
|
||||
|
||||
assert len(tool_messages) == 1
|
||||
tool_message = tool_messages[0]
|
||||
|
||||
# Single item should be a string for backward compatibility
|
||||
assert isinstance(tool_message["content"], str), (
|
||||
f"Single content item should be a string for backward compatibility, "
|
||||
f"got {type(tool_message['content'])}"
|
||||
f"Single content item should be a string for backward compatibility, got {type(tool_message['content'])}"
|
||||
)
|
||||
assert tool_message["content"] == "72°F and sunny"
|
||||
|
||||
|
|
@ -1689,9 +1630,7 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238():
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "tool_use"
|
||||
assert content_block_start["name"] == "Bash"
|
||||
|
|
@ -1735,9 +1674,7 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
|
|||
(
|
||||
block_type,
|
||||
content_block_start,
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
|
||||
choices=choices
|
||||
)
|
||||
) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
|
||||
|
||||
assert block_type == "text"
|
||||
assert content_block_start == {"type": "text", "text": ""}
|
||||
|
|
@ -1748,15 +1685,12 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
|
|||
# ============================================================================
|
||||
|
||||
# Model constant for cache control tests
|
||||
CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = (
|
||||
"bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
|
||||
)
|
||||
CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
|
||||
CACHE_CONTROL_NON_ANTHROPIC_MODEL = "gpt-4"
|
||||
# Bedrock Application Inference Profile ARN: the string contains neither
|
||||
# "anthropic" nor "claude", so the model can only be recognized via its ARN shape
|
||||
CACHE_CONTROL_BEDROCK_ARN_MODEL = (
|
||||
"bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:"
|
||||
"application-inference-profile/abcdef123456"
|
||||
"bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1772,9 +1706,7 @@ def test_should_add_cache_control_for_anthropic_model():
|
|||
"vertex_ai/claude-3-sonnet@20240229",
|
||||
]:
|
||||
target = {}
|
||||
adapter._add_cache_control_if_applicable(
|
||||
{"cache_control": cache_control}, target, model
|
||||
)
|
||||
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
|
||||
assert "cache_control" in target
|
||||
assert target["cache_control"] == cache_control
|
||||
|
||||
|
|
@ -1790,9 +1722,7 @@ def test_should_not_add_cache_control_for_non_anthropic_model():
|
|||
"gemini-pro",
|
||||
]:
|
||||
target = {}
|
||||
adapter._add_cache_control_if_applicable(
|
||||
{"cache_control": cache_control}, target, model
|
||||
)
|
||||
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
|
||||
assert "cache_control" not in target
|
||||
|
||||
|
||||
|
|
@ -1807,9 +1737,7 @@ def test_should_not_add_cache_control_when_none():
|
|||
{},
|
||||
]:
|
||||
target = {}
|
||||
adapter._add_cache_control_if_applicable(
|
||||
source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL
|
||||
)
|
||||
adapter._add_cache_control_if_applicable(source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL)
|
||||
assert "cache_control" not in target
|
||||
|
||||
|
||||
|
|
@ -1820,9 +1748,7 @@ def test_should_not_add_cache_control_when_model_none():
|
|||
|
||||
for model in [None, ""]:
|
||||
target = {}
|
||||
adapter._add_cache_control_if_applicable(
|
||||
{"cache_control": cache_control}, target, model
|
||||
)
|
||||
adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
|
||||
assert "cache_control" not in target
|
||||
|
||||
|
||||
|
|
@ -1928,12 +1854,7 @@ def test_cache_control_fix_does_not_broaden_claude_detection():
|
|||
make is_anthropic_claude_model treat ARN profiles as Claude, which would route
|
||||
thinking params through unmodified and break non-Claude Bedrock profiles.
|
||||
"""
|
||||
assert (
|
||||
LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(
|
||||
CACHE_CONTROL_BEDROCK_ARN_MODEL
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(CACHE_CONTROL_BEDROCK_ARN_MODEL) is False
|
||||
|
||||
|
||||
def test_thinking_preserved_for_bedrock_arn_inference_profile():
|
||||
|
|
@ -2531,9 +2452,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_without
|
|||
(
|
||||
type_of_content,
|
||||
content_block_delta,
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
|
||||
choices=choices
|
||||
)
|
||||
) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
|
||||
|
||||
assert type_of_content == "thinking_delta"
|
||||
assert content_block_delta["type"] == "thinking_delta"
|
||||
|
|
@ -2565,9 +2484,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
|
|||
)
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
anthropic_response = adapter.translate_openai_response_to_anthropic(
|
||||
response=openai_response
|
||||
)
|
||||
anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
|
||||
|
||||
anthropic_content = anthropic_response.get("content")
|
||||
assert anthropic_content is not None
|
||||
|
|
@ -2580,9 +2497,7 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
|
|||
|
||||
# Second block should be text
|
||||
assert anthropic_content[1]["type"] == "text"
|
||||
assert (
|
||||
anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
|
||||
)
|
||||
assert anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
|
||||
|
||||
assert anthropic_response.get("stop_reason") == "end_turn"
|
||||
|
||||
|
|
@ -2634,9 +2549,7 @@ def test_truncate_tool_name_deterministic():
|
|||
def test_truncate_tool_name_avoids_collisions():
|
||||
"""Similar long names should produce different truncated names."""
|
||||
name1 = "process_user_data_with_validation_and_error_handling_for_production_environment"
|
||||
name2 = (
|
||||
"process_user_data_with_validation_and_error_handling_for_staging_environment"
|
||||
)
|
||||
name2 = "process_user_data_with_validation_and_error_handling_for_staging_environment"
|
||||
|
||||
result1 = truncate_tool_name(name1)
|
||||
result2 = truncate_tool_name(name2)
|
||||
|
|
@ -2656,9 +2569,7 @@ def test_create_tool_name_mapping_no_long_names():
|
|||
|
||||
def test_create_tool_name_mapping_with_long_names():
|
||||
"""Mapping should contain entries for truncated names."""
|
||||
long_name = (
|
||||
"a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
|
||||
)
|
||||
long_name = "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
|
||||
tools = [
|
||||
{"name": "short_name"},
|
||||
{"name": long_name},
|
||||
|
|
@ -2683,9 +2594,7 @@ def test_translate_anthropic_tools_with_long_names():
|
|||
]
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
|
||||
tools=tools, model="gpt-4"
|
||||
)
|
||||
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
|
||||
|
||||
assert len(result) == 1
|
||||
# The tool name should be truncated
|
||||
|
|
@ -2707,9 +2616,7 @@ def test_translate_anthropic_tools_mixed_names():
|
|||
]
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
|
||||
tools=tools, model="gpt-4"
|
||||
)
|
||||
result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
|
||||
|
||||
assert len(result) == 2
|
||||
# Short name unchanged
|
||||
|
|
@ -2723,9 +2630,7 @@ def test_translate_anthropic_tools_mixed_names():
|
|||
|
||||
def test_translate_openai_response_restores_tool_names():
|
||||
"""Tool names in responses should be restored to original."""
|
||||
original_name = (
|
||||
"a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
|
||||
)
|
||||
original_name = "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
|
||||
truncated_name = truncate_tool_name(original_name)
|
||||
tool_name_mapping = {truncated_name: original_name}
|
||||
|
||||
|
|
@ -2757,9 +2662,7 @@ def test_translate_openai_response_restores_tool_names():
|
|||
)
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result = adapter.translate_openai_response_to_anthropic(
|
||||
response=response, tool_name_mapping=tool_name_mapping
|
||||
)
|
||||
result = adapter.translate_openai_response_to_anthropic(response=response, tool_name_mapping=tool_name_mapping)
|
||||
|
||||
# Find the tool_use block in the response
|
||||
tool_use_blocks = [c for c in result["content"] if c.get("type") == "tool_use"]
|
||||
|
|
@ -2925,9 +2828,7 @@ def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with
|
|||
"cache_write_tokens": 20.0,
|
||||
}
|
||||
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
|
||||
usage
|
||||
)
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
|
||||
|
||||
assert anthropic_usage["input_tokens"] == 70
|
||||
assert anthropic_usage["output_tokens"] == 50
|
||||
|
|
@ -2946,9 +2847,7 @@ def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens():
|
|||
"cache_creation_tokens": 20.25,
|
||||
}
|
||||
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
|
||||
usage
|
||||
)
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
|
||||
|
||||
assert anthropic_usage["input_tokens"] == 120
|
||||
assert anthropic_usage["output_tokens"] == 50
|
||||
|
|
@ -2965,9 +2864,7 @@ def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens():
|
|||
usage.cache_read_input_tokens = True
|
||||
usage.cache_creation_input_tokens = True
|
||||
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
|
||||
usage
|
||||
)
|
||||
anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
|
||||
|
||||
assert anthropic_usage["input_tokens"] == 120
|
||||
assert anthropic_usage["output_tokens"] == 50
|
||||
|
|
@ -3186,9 +3083,7 @@ def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_appl
|
|||
assert message_delta["usage"]["output_tokens"] == 50
|
||||
assert message_delta["usage"]["cache_read_input_tokens"] == 30
|
||||
assert message_delta["usage"]["cache_creation_input_tokens"] == 20
|
||||
assert message_delta["context_management"]["applied_edits"][0]["type"] == (
|
||||
"compact_20260112"
|
||||
)
|
||||
assert message_delta["context_management"]["applied_edits"][0]["type"] == ("compact_20260112")
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
@ -3363,15 +3258,8 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
|
|||
assert schema["required"] == ["user"]
|
||||
assert schema["properties"]["user"]["additionalProperties"] is False
|
||||
assert schema["properties"]["user"]["required"] == ["name", "address"]
|
||||
assert (
|
||||
schema["properties"]["user"]["properties"]["address"][
|
||||
"additionalProperties"
|
||||
]
|
||||
is False
|
||||
)
|
||||
assert schema["properties"]["user"]["properties"]["address"]["required"] == [
|
||||
"city"
|
||||
]
|
||||
assert schema["properties"]["user"]["properties"]["address"]["additionalProperties"] is False
|
||||
assert schema["properties"]["user"]["properties"]["address"]["required"] == ["city"]
|
||||
|
||||
def test_array_items_object_adds_additional_properties_false(self):
|
||||
output_format = {
|
||||
|
|
@ -3446,19 +3334,9 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
|
|||
assert sorted(schema["required"]) == ["age", "email", "name"]
|
||||
|
||||
def test_invalid_output_format_returns_none(self):
|
||||
assert (
|
||||
self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
|
||||
)
|
||||
assert (
|
||||
self.adapter.translate_anthropic_output_format_to_openai({"type": "text"})
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
self.adapter.translate_anthropic_output_format_to_openai(
|
||||
{"type": "json_schema"}
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
|
||||
assert self.adapter.translate_anthropic_output_format_to_openai({"type": "text"}) is None
|
||||
assert self.adapter.translate_anthropic_output_format_to_openai({"type": "json_schema"}) is None
|
||||
|
||||
|
||||
class TestAnthropicStreamWrapperToolArgs:
|
||||
|
|
@ -3662,9 +3540,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_compaction_block()
|
|||
)
|
||||
response = _make_simple_openai_response(text="Hello after compaction.")
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result = adapter.translate_openai_response_to_anthropic(
|
||||
response=response, polyfill_result=polyfill
|
||||
)
|
||||
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
|
||||
|
||||
content = result.get("content")
|
||||
assert content is not None
|
||||
|
|
@ -3696,9 +3572,7 @@ def test_translate_openai_response_to_anthropic_with_polyfill_iterations_usage()
|
|||
)
|
||||
response = _make_simple_openai_response(prompt_tokens=100, completion_tokens=30)
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result = adapter.translate_openai_response_to_anthropic(
|
||||
response=response, polyfill_result=polyfill
|
||||
)
|
||||
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
|
||||
|
||||
usage = result.get("usage")
|
||||
assert usage is not None
|
||||
|
|
@ -3753,13 +3627,9 @@ def test_translate_openai_response_to_anthropic_with_polyfill_both_compaction_an
|
|||
{"type": "compaction", "input_tokens": 300, "output_tokens": 75},
|
||||
],
|
||||
)
|
||||
response = _make_simple_openai_response(
|
||||
text="After compaction.", prompt_tokens=120, completion_tokens=40
|
||||
)
|
||||
response = _make_simple_openai_response(text="After compaction.", prompt_tokens=120, completion_tokens=40)
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
result = adapter.translate_openai_response_to_anthropic(
|
||||
response=response, polyfill_result=polyfill
|
||||
)
|
||||
result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
|
||||
|
||||
# compaction block must come first
|
||||
content = result.get("content")
|
||||
|
|
@ -3861,7 +3731,9 @@ def test_translate_anthropic_tools_to_openai_omits_unset_strict():
|
|||
assert function["parameters"]["required"] == ["query"]
|
||||
|
||||
|
||||
TOOL_RESULT_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
TOOL_RESULT_IMAGE_B64 = (
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
)
|
||||
TOOL_RESULT_IMAGE_URL = "https://example.com/screenshot.png"
|
||||
|
||||
|
||||
|
|
@ -3869,8 +3741,7 @@ def _anthropic_tool_use_turn(*tool_use_ids):
|
|||
return AnthopicMessagesAssistantMessageParam(
|
||||
role="assistant",
|
||||
content=[
|
||||
{"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}}
|
||||
for tid in tool_use_ids
|
||||
{"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}} for tid in tool_use_ids
|
||||
],
|
||||
)
|
||||
|
||||
|
|
@ -3990,9 +3861,7 @@ def test_tool_result_parallel_tool_calls_keep_tool_message_adjacency():
|
|||
result = _run_chat_completions_pipeline(
|
||||
[
|
||||
_anthropic_tool_use_turn("toolu_01", "toolu_02"),
|
||||
_anthropic_tool_result_turn(
|
||||
{"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}
|
||||
),
|
||||
_anthropic_tool_result_turn({"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}),
|
||||
]
|
||||
)
|
||||
|
||||
|
|
@ -4156,7 +4025,9 @@ def test_translate_anthropic_to_openai_without_prompt_cache_breakpoint_adds_noth
|
|||
def test_translate_anthropic_messages_to_openai_carries_midturn_system_prompt_cache_breakpoint():
|
||||
explicit = {"mode": "explicit"}
|
||||
result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
|
||||
messages=[{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}],
|
||||
messages=[
|
||||
{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}
|
||||
],
|
||||
model="gpt-5.6",
|
||||
)
|
||||
assert result == [
|
||||
|
|
|
|||
|
|
@ -34,9 +34,7 @@ import pytest
|
|||
# Anchor sys.path to this file's location — not the working-directory-relative
|
||||
# pattern Greptile flagged on PR #23706. Resolves correctly regardless of
|
||||
# where pytest is invoked from.
|
||||
sys.path.insert(
|
||||
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))
|
||||
)
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")))
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.handler import (
|
||||
ANTHROPIC_ONLY_REQUEST_KEYS,
|
||||
|
|
@ -174,9 +172,7 @@ class TestOutputConfigStrippedFromCompletionKwargs:
|
|||
result = _call_prepare(
|
||||
extra_kwargs={
|
||||
"custom_llm_provider": "azure",
|
||||
"output_config": {
|
||||
"format": {"type": "json_schema", "schema": losing_schema}
|
||||
},
|
||||
"output_config": {"format": {"type": "json_schema", "schema": losing_schema}},
|
||||
},
|
||||
output_format={"type": "json_schema", "schema": winning_schema},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -31,9 +31,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _build_fake_stream(
|
||||
content: str, finish_reason: str = "stop"
|
||||
) -> MockResponseIterator:
|
||||
def _build_fake_stream(content: str, finish_reason: str = "stop") -> MockResponseIterator:
|
||||
"""Mimic a Vertex Gemma `:predict` fake stream: one collapsed chunk."""
|
||||
model_response = ModelResponse()
|
||||
model_response.choices = [
|
||||
|
|
@ -133,9 +131,7 @@ def test_delayed_usage_chunk_preserves_cache_tokens():
|
|||
wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o")
|
||||
events = list(wrapper)
|
||||
|
||||
message_delta = next(
|
||||
event for event in events if event.get("type") == "message_delta"
|
||||
)
|
||||
message_delta = next(event for event in events if event.get("type") == "message_delta")
|
||||
|
||||
assert message_delta["usage"]["input_tokens"] == 70
|
||||
assert message_delta["usage"]["output_tokens"] == 5
|
||||
|
|
@ -145,13 +141,7 @@ def test_delayed_usage_chunk_preserves_cache_tokens():
|
|||
|
||||
def test_splitter_passes_through_non_combined_chunks():
|
||||
"""A chunk with content but no finish_reason is not split."""
|
||||
chunk = ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0, delta=Delta(content="partial"), finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)])
|
||||
chunks = list(_CombinedChunkSplitter(iter([chunk])))
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].choices[0].delta.content == "partial"
|
||||
|
|
@ -159,11 +149,7 @@ def test_splitter_passes_through_non_combined_chunks():
|
|||
|
||||
def test_splitter_splits_combined_chunk_into_content_then_finish():
|
||||
"""A chunk with both content and finish_reason becomes two chunks."""
|
||||
chunk = ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop")
|
||||
]
|
||||
)
|
||||
chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop")])
|
||||
content_chunk, finish_chunk = list(_CombinedChunkSplitter(iter([chunk])))
|
||||
|
||||
assert content_chunk.choices[0].delta.content == "done"
|
||||
|
|
@ -193,9 +179,7 @@ def test_split_clears_reasoning_and_thinking_on_finish_chunk():
|
|||
reasoning_content="some reasoning",
|
||||
thinking_blocks=[{"type": "thinking"}],
|
||||
)
|
||||
chunk = SimpleNamespace(
|
||||
choices=[SimpleNamespace(finish_reason="stop", delta=delta)]
|
||||
)
|
||||
chunk = SimpleNamespace(choices=[SimpleNamespace(finish_reason="stop", delta=delta)])
|
||||
|
||||
content_chunk, finish_chunk = _CombinedChunkSplitter._split(chunk)
|
||||
|
||||
|
|
|
|||
|
|
@ -22,9 +22,7 @@ def _make_text_chunk(
|
|||
StreamingChoices(
|
||||
finish_reason=finish_reason,
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=text, role="assistant" if text else None, tool_calls=None
|
||||
),
|
||||
delta=Delta(content=text, role="assistant" if text else None, tool_calls=None),
|
||||
logprobs=None,
|
||||
)
|
||||
]
|
||||
|
|
@ -73,34 +71,23 @@ async def test_stream_emits_compaction_block_before_text():
|
|||
compaction_start = next(
|
||||
e
|
||||
for e in events
|
||||
if e.get("type") == "content_block_start"
|
||||
and e.get("content_block", {}).get("type") == "compaction"
|
||||
if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "compaction"
|
||||
)
|
||||
assert compaction_start["index"] == 0
|
||||
|
||||
compaction_delta = next(
|
||||
e
|
||||
for e in events
|
||||
if e.get("type") == "content_block_delta"
|
||||
and e.get("delta", {}).get("type") == "compaction_delta"
|
||||
if e.get("type") == "content_block_delta" and e.get("delta", {}).get("type") == "compaction_delta"
|
||||
)
|
||||
assert compaction_delta["index"] == 0
|
||||
assert (
|
||||
compaction_delta["delta"]["content"] == "Summary of prior conversation turns."
|
||||
)
|
||||
assert compaction_delta["delta"]["content"] == "Summary of prior conversation turns."
|
||||
|
||||
compaction_stop = next(
|
||||
e
|
||||
for e in events
|
||||
if e.get("type") == "content_block_stop" and e.get("index") == 0
|
||||
)
|
||||
compaction_stop = next(e for e in events if e.get("type") == "content_block_stop" and e.get("index") == 0)
|
||||
assert compaction_stop is not None
|
||||
|
||||
text_start = next(
|
||||
e
|
||||
for e in events
|
||||
if e.get("type") == "content_block_start"
|
||||
and e.get("content_block", {}).get("type") == "text"
|
||||
e for e in events if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "text"
|
||||
)
|
||||
assert text_start["index"] == 1
|
||||
|
||||
|
|
@ -177,14 +164,9 @@ async def test_stream_without_compaction_block_unchanged():
|
|||
events = await _collect_events_async(wrapper)
|
||||
|
||||
assert not any(
|
||||
e.get("content_block", {}).get("type") == "compaction"
|
||||
for e in events
|
||||
if e.get("type") == "content_block_start"
|
||||
e.get("content_block", {}).get("type") == "compaction" for e in events if e.get("type") == "content_block_start"
|
||||
)
|
||||
text_start = next(
|
||||
e
|
||||
for e in events
|
||||
if e.get("type") == "content_block_start"
|
||||
and e.get("content_block", {}).get("type") == "text"
|
||||
e for e in events if e.get("type") == "content_block_start" and e.get("content_block", {}).get("type") == "text"
|
||||
)
|
||||
assert text_start["index"] == 0
|
||||
|
|
|
|||
|
|
@ -19,9 +19,7 @@ from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Us
|
|||
|
||||
|
||||
def _text_chunk(text: str) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)]
|
||||
)
|
||||
return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)])
|
||||
|
||||
|
||||
def _finish_chunk() -> ModelResponseStream:
|
||||
|
|
@ -61,9 +59,7 @@ def test_leading_metadata_chunk_without_choices_does_not_kill_stream():
|
|||
wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model")
|
||||
events = list(wrapper)
|
||||
|
||||
text = "".join(
|
||||
event["delta"]["text"] for event in events if event.get("type") == "content_block_delta"
|
||||
)
|
||||
text = "".join(event["delta"]["text"] for event in events if event.get("type") == "content_block_delta")
|
||||
assert text == "Hello there"
|
||||
assert events[-1]["type"] == "message_stop"
|
||||
|
||||
|
|
|
|||
|
|
@ -507,11 +507,7 @@ def _thinking_first_chunks() -> List[MagicMock]:
|
|||
|
||||
|
||||
def _assert_thinking_first_block_opens_at_index_zero(events: List[dict]) -> None:
|
||||
starts = [
|
||||
(e["index"], e["content_block"]["type"])
|
||||
for e in events
|
||||
if e.get("type") == "content_block_start"
|
||||
]
|
||||
starts = [(e["index"], e["content_block"]["type"]) for e in events if e.get("type") == "content_block_start"]
|
||||
assert starts == [(0, "thinking"), (1, "text")], starts
|
||||
assert "" not in _text_deltas(events)
|
||||
assert _thinking_deltas(events) == ["Let me think", "about it."]
|
||||
|
|
@ -980,9 +976,7 @@ def test_tool_block_start_emitted_without_awaiting_the_next_chunk_sync():
|
|||
"name": "Write",
|
||||
"input": {},
|
||||
}
|
||||
assert stream.pulled == 1, (
|
||||
f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
|
||||
)
|
||||
assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -997,9 +991,7 @@ async def test_tool_block_start_emitted_without_awaiting_the_next_chunk_async():
|
|||
start = await wrapper.__anext__()
|
||||
assert start["type"] == "content_block_start"
|
||||
assert start["content_block"]["name"] == "Write"
|
||||
assert stream.pulled == 1, (
|
||||
f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
|
||||
)
|
||||
assert stream.pulled == 1, f"content_block_start was withheld until {stream.pulled} upstream chunks had arrived"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_async", [False, True])
|
||||
|
|
|
|||
|
|
@ -127,21 +127,17 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args():
|
|||
):
|
||||
input_json_delta_idx = i
|
||||
|
||||
assert (
|
||||
tool_start_idx is not None
|
||||
), f"Expected content_block_start with type=tool_use; events: {event_types}"
|
||||
assert (
|
||||
input_json_delta_idx is not None
|
||||
), f"Expected content_block_delta with input_json_delta; events: {event_types}"
|
||||
assert (
|
||||
input_json_delta_idx == tool_start_idx + 1
|
||||
), "input_json_delta should immediately follow the tool_use content_block_start"
|
||||
assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}"
|
||||
assert input_json_delta_idx is not None, (
|
||||
f"Expected content_block_delta with input_json_delta; events: {event_types}"
|
||||
)
|
||||
assert input_json_delta_idx == tool_start_idx + 1, (
|
||||
"input_json_delta should immediately follow the tool_use content_block_start"
|
||||
)
|
||||
|
||||
# Verify the delta carries the tool arguments
|
||||
delta_event = events[input_json_delta_idx]
|
||||
assert delta_event["delta"][
|
||||
"partial_json"
|
||||
], "input_json_delta should have non-empty partial_json"
|
||||
assert delta_event["delta"]["partial_json"], "input_json_delta should have non-empty partial_json"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -230,8 +226,7 @@ async def test_async_stream_no_extra_delta_when_tool_args_empty():
|
|||
and e["delta"].get("type") == "input_json_delta"
|
||||
]
|
||||
assert len(input_json_deltas) == 1, (
|
||||
f"Expected exactly 1 input_json_delta (from the follow-up chunk), "
|
||||
f"got {len(input_json_deltas)}"
|
||||
f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}"
|
||||
)
|
||||
assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}'
|
||||
|
||||
|
|
@ -291,15 +286,13 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args():
|
|||
):
|
||||
input_json_delta_idx = i
|
||||
|
||||
assert (
|
||||
tool_start_idx is not None
|
||||
), f"Expected content_block_start with type=tool_use; events: {event_types}"
|
||||
assert (
|
||||
input_json_delta_idx is not None
|
||||
), f"Expected content_block_delta with input_json_delta; events: {event_types}"
|
||||
assert (
|
||||
input_json_delta_idx == tool_start_idx + 1
|
||||
), "input_json_delta should immediately follow the tool_use content_block_start"
|
||||
assert tool_start_idx is not None, f"Expected content_block_start with type=tool_use; events: {event_types}"
|
||||
assert input_json_delta_idx is not None, (
|
||||
f"Expected content_block_delta with input_json_delta; events: {event_types}"
|
||||
)
|
||||
assert input_json_delta_idx == tool_start_idx + 1, (
|
||||
"input_json_delta should immediately follow the tool_use content_block_start"
|
||||
)
|
||||
assert events[input_json_delta_idx]["delta"]["partial_json"]
|
||||
|
||||
|
||||
|
|
@ -343,9 +336,7 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty():
|
|||
)
|
||||
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=iter(
|
||||
[text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk]
|
||||
),
|
||||
completion_stream=iter([text_chunk, tool_name_chunk, tool_args_chunk, finish_chunk]),
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
|
|
@ -374,7 +365,6 @@ def test_sync_stream_no_extra_delta_when_tool_args_empty():
|
|||
and e["delta"].get("type") == "input_json_delta"
|
||||
]
|
||||
assert len(input_json_deltas) == 1, (
|
||||
f"Expected exactly 1 input_json_delta (from the follow-up chunk), "
|
||||
f"got {len(input_json_deltas)}"
|
||||
f"Expected exactly 1 input_json_delta (from the follow-up chunk), got {len(input_json_deltas)}"
|
||||
)
|
||||
assert input_json_deltas[0]["delta"]["partial_json"] == '{"location": "NYC"}'
|
||||
|
|
|
|||
|
|
@ -258,9 +258,7 @@ def test_tool_result_list_content_shape_preserved():
|
|||
{"role": "user", "content": "Hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}}
|
||||
],
|
||||
"content": [{"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
|
|
@ -274,9 +272,7 @@ def test_tool_result_list_content_shape_preserved():
|
|||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}}
|
||||
],
|
||||
"content": [{"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
|
|
|
|||
|
|
@ -193,10 +193,7 @@ def test_select_last_user_question_strips_tool_result_from_mixed_turn():
|
|||
content = selected[0]["content"]
|
||||
assert isinstance(content, list)
|
||||
assert all(b.get("type") != "tool_result" for b in content)
|
||||
assert any(
|
||||
b.get("type") == "text" and b.get("text") == "follow-up question"
|
||||
for b in content
|
||||
)
|
||||
assert any(b.get("type") == "text" and b.get("text") == "follow-up question" for b in content)
|
||||
|
||||
|
||||
def test_select_last_user_question_skips_pure_tool_result_turn():
|
||||
|
|
@ -425,9 +422,7 @@ def test_client_compaction_block_history_without_context_management():
|
|||
|
||||
|
||||
def test_client_compaction_block_history_no_compaction_returns_none():
|
||||
result = apply_client_compaction_block_history(
|
||||
messages=_simple_messages(), system="base"
|
||||
)
|
||||
result = apply_client_compaction_block_history(messages=_simple_messages(), system="base")
|
||||
assert result is None
|
||||
|
||||
|
||||
|
|
@ -512,9 +507,7 @@ async def test_slice_only_no_compaction_block_under_threshold():
|
|||
async def test_full_summary_path():
|
||||
"""Over threshold: summary call fires, compaction_block and iterations_usage returned."""
|
||||
messages = _simple_messages()
|
||||
mock_response = _make_mock_response(
|
||||
"<summary>Condensed history</summary>", prompt_tokens=200, completion_tokens=50
|
||||
)
|
||||
mock_response = _make_mock_response("<summary>Condensed history</summary>", prompt_tokens=200, completion_tokens=50)
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -1068,13 +1061,11 @@ async def test_summary_call_does_not_emit_consecutive_user_turns():
|
|||
)
|
||||
|
||||
summary_messages = captured_calls[0]["summary_messages"]
|
||||
user_indices = [
|
||||
idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"
|
||||
]
|
||||
user_indices = [idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"]
|
||||
# No two adjacent indices.
|
||||
assert all(
|
||||
b - a > 1 for a, b in zip(user_indices, user_indices[1:])
|
||||
), f"two consecutive user turns produced: {summary_messages}"
|
||||
assert all(b - a > 1 for a, b in zip(user_indices, user_indices[1:])), (
|
||||
f"two consecutive user turns produced: {summary_messages}"
|
||||
)
|
||||
|
||||
|
||||
async def test_summary_call_sends_default_max_tokens():
|
||||
|
|
@ -1157,9 +1148,9 @@ def test_summary_max_tokens_setting_falls_back_for_invalid_values():
|
|||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"context_management_summary_max_tokens": bad},
|
||||
):
|
||||
assert (
|
||||
_read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS
|
||||
), f"expected default for invalid override {bad!r}"
|
||||
assert _read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS, (
|
||||
f"expected default for invalid override {bad!r}"
|
||||
)
|
||||
|
||||
|
||||
async def test_summary_call_sends_default_timeout():
|
||||
|
|
@ -1282,9 +1273,7 @@ async def test_summary_model_denied_when_team_not_in_allowlist():
|
|||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=_fake_user_api_key_auth(
|
||||
key_models=["all-proxy-models"], team_models=["gpt-4o"]
|
||||
),
|
||||
user_api_key_auth=_fake_user_api_key_auth(key_models=["all-proxy-models"], team_models=["gpt-4o"]),
|
||||
)
|
||||
|
||||
mock_call.assert_not_awaited()
|
||||
|
|
@ -1313,9 +1302,7 @@ async def test_summary_model_allowed_when_in_key_allowlist():
|
|||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=_fake_user_api_key_auth(
|
||||
key_models=["claude-haiku-4-5", "gpt-4o"]
|
||||
),
|
||||
user_api_key_auth=_fake_user_api_key_auth(key_models=["claude-haiku-4-5", "gpt-4o"]),
|
||||
)
|
||||
|
||||
mock_call.assert_awaited_once()
|
||||
|
|
@ -1521,9 +1508,7 @@ async def test_summary_model_denied_when_key_over_model_budget():
|
|||
|
||||
limiter = MagicMock()
|
||||
limiter.is_key_within_model_budget = AsyncMock(
|
||||
side_effect=litellm.BudgetExceededError(
|
||||
message="over budget", current_cost=10, max_budget=5
|
||||
)
|
||||
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
|
||||
)
|
||||
|
||||
with (
|
||||
|
|
@ -1574,9 +1559,7 @@ async def test_summary_model_denied_when_user_over_model_budget():
|
|||
|
||||
limiter = MagicMock()
|
||||
limiter.is_user_within_model_budget = AsyncMock(
|
||||
side_effect=litellm.BudgetExceededError(
|
||||
message="over budget", current_cost=10, max_budget=5
|
||||
)
|
||||
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
|
||||
)
|
||||
|
||||
with (
|
||||
|
|
@ -1617,9 +1600,7 @@ async def test_summary_model_denied_when_user_over_model_budget():
|
|||
_PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
)
|
||||
|
||||
real_params = inspect.signature(
|
||||
_PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget
|
||||
).parameters
|
||||
real_params = inspect.signature(_PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget).parameters
|
||||
for kwarg in ("user_id", "user_model_max_budget", "model"):
|
||||
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts"
|
||||
|
||||
|
|
@ -1641,9 +1622,7 @@ async def test_summary_model_denied_when_end_user_over_model_budget():
|
|||
limiter = MagicMock()
|
||||
limiter.is_key_within_model_budget = AsyncMock(return_value=True)
|
||||
limiter.is_end_user_within_model_budget = AsyncMock(
|
||||
side_effect=litellm.BudgetExceededError(
|
||||
message="over budget", current_cost=10, max_budget=5
|
||||
)
|
||||
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
|
||||
)
|
||||
|
||||
with (
|
||||
|
|
@ -1956,9 +1935,7 @@ async def test_model_budget_metadata_propagated_to_summary_call():
|
|||
parent_litellm_metadata = {
|
||||
"user_api_key": "sk-test",
|
||||
"user_api_key_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 5}},
|
||||
"user_api_key_end_user_model_max_budget": {
|
||||
"claude-haiku-4-5": {"budget_limit": 2}
|
||||
},
|
||||
"user_api_key_end_user_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 2}},
|
||||
}
|
||||
|
||||
with (
|
||||
|
|
@ -1983,12 +1960,8 @@ async def test_model_budget_metadata_propagated_to_summary_call():
|
|||
)
|
||||
|
||||
propagated = mock_call.call_args.kwargs["metadata"]
|
||||
assert propagated["user_api_key_model_max_budget"] == {
|
||||
"claude-haiku-4-5": {"budget_limit": 5}
|
||||
}
|
||||
assert propagated["user_api_key_end_user_model_max_budget"] == {
|
||||
"claude-haiku-4-5": {"budget_limit": 2}
|
||||
}
|
||||
assert propagated["user_api_key_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 5}}
|
||||
assert propagated["user_api_key_end_user_model_max_budget"] == {"claude-haiku-4-5": {"budget_limit": 2}}
|
||||
|
||||
|
||||
async def test_summary_call_propagates_allowed_model_region():
|
||||
|
|
@ -2460,9 +2433,7 @@ def test_endpoint_returns_anthropic_400_on_context_management_error():
|
|||
mock_proxy_server.version = "test"
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
with patch(
|
||||
"litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_cls:
|
||||
with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls:
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.base_process_llm_request = AsyncMock(
|
||||
side_effect=AnthropicContextManagementError(
|
||||
|
|
@ -2521,9 +2492,7 @@ def test_endpoint_runs_failure_hook_on_500_context_management_error():
|
|||
mock_proxy_server.version = "test"
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
with patch(
|
||||
"litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_cls:
|
||||
with patch("litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_cls:
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.base_process_llm_request = AsyncMock(
|
||||
side_effect=AnthropicContextManagementError(
|
||||
|
|
|
|||
|
|
@ -50,9 +50,7 @@ async def test_unknown_edit_type_is_noop():
|
|||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
context_management_spec={
|
||||
"edits": [{"type": "totally_not_a_real_edit_20999999"}]
|
||||
},
|
||||
context_management_spec={"edits": [{"type": "totally_not_a_real_edit_20999999"}]},
|
||||
)
|
||||
assert result.applied_edits == []
|
||||
assert result.messages == messages
|
||||
|
|
|
|||
|
|
@ -39,9 +39,7 @@ def _text_resp(text: str, model: str = "gpt-4o-mini") -> Dict:
|
|||
}
|
||||
|
||||
|
||||
def _advisor_call_resp(
|
||||
question: str = "How do I approach this?", tool_id: str = "tid_01"
|
||||
) -> Dict:
|
||||
def _advisor_call_resp(question: str = "How do I approach this?", tool_id: str = "tid_01") -> Dict:
|
||||
return {
|
||||
"id": "msg_int_test",
|
||||
"type": "message",
|
||||
|
|
@ -106,14 +104,10 @@ async def test_full_dispatch_interceptor_fires_and_loop_completes():
|
|||
assert isinstance(result, dict)
|
||||
content = result.get("content", [])
|
||||
text_blocks = [b for b in content if b.get("type") == "text"]
|
||||
advisor_uses = [
|
||||
b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"
|
||||
]
|
||||
advisor_uses = [b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"]
|
||||
|
||||
assert len(text_blocks) >= 1, "Final response must have text"
|
||||
assert (
|
||||
len(advisor_uses) == 0
|
||||
), "No advisor tool_use blocks must appear in final output"
|
||||
assert len(advisor_uses) == 0, "No advisor tool_use blocks must appear in final output"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -221,9 +215,7 @@ async def test_named_params_forwarded_into_advisor_executor_subcall():
|
|||
|
||||
captured_executor_kwargs: Dict = {}
|
||||
|
||||
async def mock_handler(
|
||||
model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs
|
||||
):
|
||||
async def mock_handler(model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs):
|
||||
# First call is the executor sub-call (returns advisor tool_use).
|
||||
# Capture its kwargs so we can assert the forwarded params.
|
||||
if not captured_executor_kwargs:
|
||||
|
|
@ -267,8 +259,7 @@ async def test_named_params_forwarded_into_advisor_executor_subcall():
|
|||
)
|
||||
|
||||
assert captured_executor_kwargs["thinking"] == {"type": "adaptive"}, (
|
||||
"thinking must be forwarded into executor sub-call — see "
|
||||
"anthropic_messages.handler interceptor invocation."
|
||||
"thinking must be forwarded into executor sub-call — see anthropic_messages.handler interceptor invocation."
|
||||
)
|
||||
# The advisor enriches metadata with `advisor_sub_call` / `parent_request_id`,
|
||||
# but the original caller fields must survive into the executor sub-call.
|
||||
|
|
@ -304,9 +295,7 @@ async def test_pre_request_hook_override_does_not_collide_with_explicit_kwargs()
|
|||
|
||||
captured: Dict = {}
|
||||
|
||||
async def mock_handler(
|
||||
model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs
|
||||
):
|
||||
async def mock_handler(model, messages, tools, stream, max_tokens, custom_llm_provider, **kwargs):
|
||||
if not captured:
|
||||
captured.update(
|
||||
{
|
||||
|
|
@ -320,9 +309,7 @@ async def test_pre_request_hook_override_does_not_collide_with_explicit_kwargs()
|
|||
return _text_resp("Some advice.", model="claude-opus-4-6")
|
||||
return _text_resp("Final answer.")
|
||||
|
||||
async def fake_pre_request_hooks(
|
||||
model, messages, tools, stream, custom_llm_provider, **hook_kwargs
|
||||
):
|
||||
async def fake_pre_request_hooks(model, messages, tools, stream, custom_llm_provider, **hook_kwargs):
|
||||
# Simulate a CustomLogger.async_pre_request_hook that overrides several
|
||||
# named params on its way through. Without the request_kwargs.pop()
|
||||
# extraction in handler.py, these would collide with the explicit
|
||||
|
|
|
|||
|
|
@ -73,9 +73,7 @@ def _build_simple_text_stream() -> List[bytes]:
|
|||
},
|
||||
)
|
||||
)
|
||||
chunks.append(
|
||||
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
|
||||
)
|
||||
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}))
|
||||
chunks.append(
|
||||
_sse_event(
|
||||
"message_delta",
|
||||
|
|
@ -148,9 +146,7 @@ def _build_tool_use_stream() -> List[bytes]:
|
|||
},
|
||||
)
|
||||
)
|
||||
chunks.append(
|
||||
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
|
||||
)
|
||||
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}))
|
||||
# tool_use block
|
||||
chunks.append(
|
||||
_sse_event(
|
||||
|
|
@ -190,9 +186,7 @@ def _build_tool_use_stream() -> List[bytes]:
|
|||
},
|
||||
)
|
||||
)
|
||||
chunks.append(
|
||||
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1})
|
||||
)
|
||||
chunks.append(_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1}))
|
||||
chunks.append(
|
||||
_sse_event(
|
||||
"message_delta",
|
||||
|
|
@ -284,9 +278,7 @@ def _build_hold_back_iterator(
|
|||
|
||||
class TestParseSSEEvents:
|
||||
def test_should_parse_single_event(self):
|
||||
raw = _sse_event(
|
||||
"message_start", {"type": "message_start", "message": {"id": "1"}}
|
||||
)
|
||||
raw = _sse_event("message_start", {"type": "message_start", "message": {"id": "1"}})
|
||||
events = _parse_sse_events(raw)
|
||||
assert len(events) == 1
|
||||
assert events[0][0] == "message_start"
|
||||
|
|
@ -457,9 +449,7 @@ class TestHandleMessageDelta:
|
|||
class TestRebuildAnthropicResponse:
|
||||
def test_should_rebuild_simple_text_response(self):
|
||||
raw_bytes = _build_simple_text_stream()
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
raw_bytes
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
|
||||
assert result is not None
|
||||
assert result["id"] == "msg_123"
|
||||
assert result["model"] == "claude-sonnet-4-20250514"
|
||||
|
|
@ -472,9 +462,7 @@ class TestRebuildAnthropicResponse:
|
|||
|
||||
def test_should_rebuild_tool_use_response(self):
|
||||
raw_bytes = _build_tool_use_stream()
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
raw_bytes
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
|
||||
assert result is not None
|
||||
assert result["id"] == "msg_tool_456"
|
||||
assert result["stop_reason"] == "tool_use"
|
||||
|
|
@ -502,23 +490,17 @@ class TestRebuildAnthropicResponse:
|
|||
},
|
||||
)
|
||||
]
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
raw_bytes
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
|
||||
assert result is None
|
||||
|
||||
def test_should_handle_empty_bytes(self):
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
[]
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse([])
|
||||
assert result is None
|
||||
|
||||
def test_should_handle_multi_event_chunks(self):
|
||||
"""When multiple SSE events arrive in a single bytes chunk."""
|
||||
combined = b"".join(_build_simple_text_stream())
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
[combined]
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse([combined])
|
||||
assert result is not None
|
||||
assert result["content"][0]["text"] == "Hello, world!"
|
||||
|
||||
|
|
@ -550,9 +532,7 @@ class TestRebuildAnthropicResponse:
|
|||
),
|
||||
_sse_event("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
raw_bytes
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
|
||||
assert result is not None
|
||||
assert result["usage"]["cache_creation_input_tokens"] == 50
|
||||
assert result["usage"]["cache_read_input_tokens"] == 30
|
||||
|
|
@ -593,9 +573,7 @@ class TestRebuildAnthropicResponse:
|
|||
),
|
||||
_sse_event("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
|
||||
raw_bytes
|
||||
)
|
||||
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(raw_bytes)
|
||||
assert result is not None
|
||||
assert result["content"][0]["type"] == "redacted_thinking"
|
||||
|
||||
|
|
@ -722,9 +700,7 @@ class TestAgenticStreamingIteratorPhase2:
|
|||
}
|
||||
|
||||
mock_handler = MagicMock()
|
||||
mock_handler._call_agentic_completion_hooks = AsyncMock(
|
||||
return_value=fake_response
|
||||
)
|
||||
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=fake_response)
|
||||
|
||||
iterator = AgenticAnthropicStreamingIterator(
|
||||
completion_stream=mock_stream,
|
||||
|
|
@ -757,9 +733,7 @@ class TestAgenticStreamingIteratorErrorHandling:
|
|||
mock_stream = MockAsyncStream(chunks)
|
||||
|
||||
mock_handler = MagicMock()
|
||||
mock_handler._call_agentic_completion_hooks = AsyncMock(
|
||||
side_effect=RuntimeError("hook exploded")
|
||||
)
|
||||
mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded"))
|
||||
|
||||
mock_logging = MagicMock()
|
||||
mock_logging.litellm_call_id = "test_call_123"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,118 @@
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.llms.anthropic.wif import get_anthropic_wif_token
|
||||
from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig
|
||||
from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig
|
||||
from tests.test_litellm.llms.anthropic.test_anthropic_wif import (
|
||||
ScriptedPoster,
|
||||
make_engine,
|
||||
token_response,
|
||||
write_token_file,
|
||||
)
|
||||
|
||||
_WIF_PARAMS: Final[dict] = {
|
||||
"anthropic_federation_rule_id": "fdrl_abc123",
|
||||
"anthropic_organization_id": "org-uuid-1",
|
||||
"anthropic_identity_token_file": "/var/run/secrets/identity-token",
|
||||
}
|
||||
|
||||
|
||||
def test_workload_identity_allowed_for_anthropic() -> None:
|
||||
assert AnthropicMessagesConfig()._allows_workload_identity is True
|
||||
|
||||
|
||||
def test_workload_identity_blocked_for_minimax() -> None:
|
||||
assert MinimaxMessagesConfig()._allows_workload_identity is False
|
||||
|
||||
|
||||
def test_workload_identity_blocked_for_tencent() -> None:
|
||||
assert TencentAnthropicMessagesConfig()._allows_workload_identity is False
|
||||
|
||||
|
||||
def test_minimax_validate_environment_never_attaches_anthropic_wif_credential(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Regression test: before the fix, an Anthropic-WIF-configured proxy would mint a real
|
||||
Anthropic federation token inside MiniMax's inherited validate_anthropic_messages_environment
|
||||
and send it as the Authorization header on the MiniMax-routed request. With no MiniMax
|
||||
credential of its own the deployment must fail closed on the missing key instead."""
|
||||
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
|
||||
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
monkeypatch.delenv("MINIMAX_API_KEY", raising=False)
|
||||
token_file = write_token_file(tmp_path, "jwt-assertion-value")
|
||||
litellm_params = {"anthropic_identity_token_file": str(token_file)}
|
||||
|
||||
with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"):
|
||||
MinimaxMessagesConfig().validate_anthropic_messages_environment(
|
||||
headers={},
|
||||
model="MiniMax-M2.1",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
api_key=None,
|
||||
api_base="https://api.minimax.io/anthropic",
|
||||
)
|
||||
|
||||
|
||||
def test_tencent_validate_environment_never_attaches_anthropic_wif_credential(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
|
||||
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
monkeypatch.delenv("TENCENT_API_KEY", raising=False)
|
||||
token_file = write_token_file(tmp_path, "jwt-assertion-value")
|
||||
litellm_params = {"anthropic_identity_token_file": str(token_file)}
|
||||
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
with pytest.raises(litellm.AuthenticationError, match="Missing Anthropic API Key"):
|
||||
TencentAnthropicMessagesConfig().validate_anthropic_messages_environment(
|
||||
headers={},
|
||||
model="deepseek-v4-pro",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
api_key=None,
|
||||
api_base="https://tokenhub-intl.tencentcloudmaas.com",
|
||||
)
|
||||
|
||||
|
||||
def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""get_anthropic_wif_token's engine parameter is the only DI seam in the WIF minting chain;
|
||||
validate_anthropic_messages_environment always uses the module's default engine, so this
|
||||
drives that seam directly with the exact litellm_params AnthropicModelInfo.get_auth_header
|
||||
would receive from each config, proving MiniMax/Tencent never reach the token endpoint even
|
||||
when a mint would otherwise succeed."""
|
||||
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod")
|
||||
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid")
|
||||
monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path))
|
||||
token_file = write_token_file(tmp_path, "jwt-assertion-value")
|
||||
litellm_params = {"anthropic_identity_token_file": str(token_file)}
|
||||
poster = ScriptedPoster([token_response("sk-ant-oat01-canary")])
|
||||
engine = make_engine(poster)
|
||||
|
||||
minted: Final = get_anthropic_wif_token(
|
||||
litellm_params,
|
||||
"https://api.anthropic.com",
|
||||
"claude-sonnet-4-5",
|
||||
engine,
|
||||
)
|
||||
assert minted == "sk-ant-oat01-canary"
|
||||
assert len(poster.requests) == 1
|
||||
|
||||
for config in (MinimaxMessagesConfig(), TencentAnthropicMessagesConfig()):
|
||||
assert config._allows_workload_identity is False
|
||||
|
||||
assert len(poster.requests) == 1
|
||||
|
|
@ -93,13 +93,11 @@ def test_messages_drops_speed_for_vertex_opus_with_drop_params(monkeypatch):
|
|||
"""Regression: a vertex_ai Opus passthrough must drop ``speed`` even though the
|
||||
prefix-stripped model id maps to a fast-mode-capable direct-Anthropic entry."""
|
||||
monkeypatch.setattr(litellm, "drop_params", True)
|
||||
optional_params = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"max_tokens": 1024, "speed": "fast"},
|
||||
model="claude-opus-4-8",
|
||||
drop_params=False,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"max_tokens": 1024, "speed": "fast"},
|
||||
model="claude-opus-4-8",
|
||||
drop_params=False,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
assert "speed" not in optional_params
|
||||
|
|
|
|||
|
|
@ -13,9 +13,7 @@ def test_output_format_supported_and_transforms_correctly():
|
|||
config = AnthropicMessagesConfig()
|
||||
|
||||
# 1. Verify it's in supported parameters
|
||||
supported_params = config.get_supported_anthropic_messages_params(
|
||||
"claude-sonnet-4-5"
|
||||
)
|
||||
supported_params = config.get_supported_anthropic_messages_params("claude-sonnet-4-5")
|
||||
assert "output_format" in supported_params
|
||||
|
||||
# 2. Verify transformation preserves output_format and adds beta header
|
||||
|
|
|
|||
|
|
@ -30,26 +30,14 @@ class MockCompletionStreamWithContentAfterStopReason:
|
|||
self.responses = [
|
||||
# Initial text content
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content="Hello"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
|
||||
),
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=" world"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content=" world"), index=0, finish_reason=None)],
|
||||
),
|
||||
# Message delta with stop_reason AND usage (this is how it actually comes from the API)
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=""), index=0, finish_reason="stop"
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")],
|
||||
usage=Usage(prompt_tokens=230, completion_tokens=65, total_tokens=295),
|
||||
),
|
||||
# Additional content after the stop_reason - this simulates the scenario
|
||||
|
|
@ -118,9 +106,9 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
|
|||
print(f"Expected chunk types: {expected_types}")
|
||||
|
||||
# Verify we have the expected number of chunks
|
||||
assert len(chunk_types) >= len(
|
||||
expected_types
|
||||
), f"Expected at least {len(expected_types)} chunks, got {len(chunk_types)}"
|
||||
assert len(chunk_types) >= len(expected_types), (
|
||||
f"Expected at least {len(expected_types)} chunks, got {len(chunk_types)}"
|
||||
)
|
||||
|
||||
# Verify key chunk types are present
|
||||
assert "message_start" in chunk_types
|
||||
|
|
@ -143,15 +131,9 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
|
|||
delta = message_delta_chunk.get("delta", {})
|
||||
usage = message_delta_chunk.get("usage", {})
|
||||
|
||||
assert (
|
||||
delta.get("stop_reason") == "end_turn"
|
||||
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
|
||||
assert (
|
||||
usage.get("input_tokens") == 230
|
||||
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
|
||||
assert (
|
||||
usage.get("output_tokens") == 65
|
||||
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
|
||||
assert delta.get("stop_reason") == "end_turn", f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
|
||||
assert usage.get("input_tokens") == 230, f"Expected input_tokens 230, got {usage.get('input_tokens')}"
|
||||
assert usage.get("output_tokens") == 65, f"Expected output_tokens 65, got {usage.get('output_tokens')}"
|
||||
|
||||
# Verify content_block_stop comes before message_delta
|
||||
content_block_stop_index = None
|
||||
|
|
@ -165,9 +147,7 @@ def test_anthropic_stream_wrapper_content_after_stop_reason():
|
|||
|
||||
assert content_block_stop_index is not None, "content_block_stop not found"
|
||||
assert message_delta_index is not None, "message_delta not found"
|
||||
assert (
|
||||
content_block_stop_index < message_delta_index
|
||||
), "content_block_stop should come before message_delta"
|
||||
assert content_block_stop_index < message_delta_index, "content_block_stop should come before message_delta"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -210,15 +190,9 @@ async def test_async_anthropic_stream_wrapper_content_after_stop_reason():
|
|||
delta = message_delta_chunk.get("delta", {})
|
||||
usage = message_delta_chunk.get("usage", {})
|
||||
|
||||
assert (
|
||||
delta.get("stop_reason") == "end_turn"
|
||||
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
|
||||
assert (
|
||||
usage.get("input_tokens") == 230
|
||||
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
|
||||
assert (
|
||||
usage.get("output_tokens") == 65
|
||||
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
|
||||
assert delta.get("stop_reason") == "end_turn", f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
|
||||
assert usage.get("input_tokens") == 230, f"Expected input_tokens 230, got {usage.get('input_tokens')}"
|
||||
assert usage.get("output_tokens") == 65, f"Expected output_tokens 65, got {usage.get('output_tokens')}"
|
||||
|
||||
|
||||
def test_usage_merging_behavior():
|
||||
|
|
@ -234,18 +208,10 @@ def test_usage_merging_behavior():
|
|||
for chunk in wrapper:
|
||||
chunks.append(chunk)
|
||||
# If this is a message_delta with stop_reason, verify it has usage
|
||||
if (
|
||||
chunk.get("type") == "message_delta"
|
||||
and chunk.get("delta", {}).get("stop_reason") is not None
|
||||
):
|
||||
|
||||
if chunk.get("type") == "message_delta" and chunk.get("delta", {}).get("stop_reason") is not None:
|
||||
usage = chunk.get("usage", {})
|
||||
assert (
|
||||
usage.get("input_tokens") is not None
|
||||
), "Usage should be merged with stop_reason chunk"
|
||||
assert (
|
||||
usage.get("output_tokens") is not None
|
||||
), "Usage should be merged with stop_reason chunk"
|
||||
assert usage.get("input_tokens") is not None, "Usage should be merged with stop_reason chunk"
|
||||
assert usage.get("output_tokens") is not None, "Usage should be merged with stop_reason chunk"
|
||||
break
|
||||
|
||||
|
||||
|
|
@ -273,12 +239,8 @@ def test_sse_wrapper_with_content_after_stop_reason():
|
|||
lines = chunk_str.split("\n")
|
||||
|
||||
# Should have event and data lines
|
||||
assert any(
|
||||
line.startswith("event: ") for line in lines
|
||||
), f"Missing event line in: {chunk_str}"
|
||||
assert any(
|
||||
line.startswith("data: ") for line in lines
|
||||
), f"Missing data line in: {chunk_str}"
|
||||
assert any(line.startswith("event: ") for line in lines), f"Missing event line in: {chunk_str}"
|
||||
assert any(line.startswith("data: ") for line in lines), f"Missing data line in: {chunk_str}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -306,12 +268,8 @@ async def test_async_sse_wrapper_with_content_after_stop_reason():
|
|||
lines = chunk_str.split("\n")
|
||||
|
||||
# Should have event and data lines
|
||||
assert any(
|
||||
line.startswith("event: ") for line in lines
|
||||
), f"Missing event line in: {chunk_str}"
|
||||
assert any(
|
||||
line.startswith("data: ") for line in lines
|
||||
), f"Missing data line in: {chunk_str}"
|
||||
assert any(line.startswith("event: ") for line in lines), f"Missing event line in: {chunk_str}"
|
||||
assert any(line.startswith("data: ") for line in lines), f"Missing data line in: {chunk_str}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ def test_anthropic_messages_handler_skips_the_gateway_on_recursion():
|
|||
"litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp",
|
||||
new=AsyncMock(return_value={"routed": True}),
|
||||
) as routed:
|
||||
with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'):
|
||||
with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"):
|
||||
anthropic_messages_handler(
|
||||
max_tokens=100,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -78,7 +78,7 @@ def test_anthropic_messages_handler_leaves_native_tools_alone():
|
|||
"litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp",
|
||||
new=AsyncMock(return_value={"routed": True}),
|
||||
) as routed:
|
||||
with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'):
|
||||
with pytest.raises(ValueError, match="anthropic_messages_handler is not implemented for sync calls"):
|
||||
anthropic_messages_handler(
|
||||
max_tokens=100,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -115,9 +115,7 @@ def test_build_tool_result_message_uses_anthropic_tool_result_blocks():
|
|||
message = _build_tool_result_message([{"tool_call_id": "toolu_1", "result": "9 sections", "name": "read_wiki"}])
|
||||
|
||||
assert message["role"] == "user"
|
||||
assert list(message["content"]) == [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"}
|
||||
]
|
||||
assert list(message["content"]) == [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from typing import List
|
||||
|
||||
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
|
||||
AnthropicStreamWrapper,
|
||||
)
|
||||
|
|
@ -53,9 +52,7 @@ def construct_text_chunk(text: str) -> ModelResponseStream:
|
|||
)
|
||||
|
||||
|
||||
def construct_split_tool_call(
|
||||
id: str, function_name: str, function_arg_parts: List[str]
|
||||
) -> List[ModelResponseStream]:
|
||||
def construct_split_tool_call(id: str, function_name: str, function_arg_parts: List[str]) -> List[ModelResponseStream]:
|
||||
return [
|
||||
# https://platform.openai.com/docs/guides/function-calling#streaming
|
||||
ModelResponseStream(
|
||||
|
|
@ -144,10 +141,7 @@ def test_anthropic_stream_wrapper_single_tool_call():
|
|||
get_weather_calls = 0
|
||||
|
||||
for chunk in chunks:
|
||||
if (
|
||||
chunk.get("type") == "content_block_start"
|
||||
and chunk["content_block"]["type"] == "tool_use"
|
||||
):
|
||||
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
|
||||
if chunk["content_block"]["name"] == "get_weather":
|
||||
get_weather_calls += 1
|
||||
|
||||
|
|
@ -203,10 +197,7 @@ def test_anthropic_stream_wrapper_back_to_back_tool_calls():
|
|||
get_weather_calls = 0
|
||||
|
||||
for chunk in chunks:
|
||||
if (
|
||||
chunk.get("type") == "content_block_start"
|
||||
and chunk["content_block"]["type"] == "tool_use"
|
||||
):
|
||||
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
|
||||
if chunk["content_block"]["name"] == "get_weather":
|
||||
get_weather_calls += 1
|
||||
|
||||
|
|
@ -218,9 +209,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
*construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']),
|
||||
construct_text_chunk("The weather is nice today."),
|
||||
*construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"SF"}']),
|
||||
*construct_split_tool_call(
|
||||
"tooluse_bar", "get_weather", ['{"city":', '"CHI"}']
|
||||
),
|
||||
*construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"CHI"}']),
|
||||
construct_text_chunk("The weather is not so nice today."),
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
|
|
@ -280,8 +269,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
text_deltas = [
|
||||
chunk["delta"]["text"]
|
||||
for chunk in chunks
|
||||
if chunk.get("type") == "content_block_delta"
|
||||
and chunk["delta"].get("type") == "text_delta"
|
||||
if chunk.get("type") == "content_block_delta" and chunk["delta"].get("type") == "text_delta"
|
||||
]
|
||||
assert text_deltas == [
|
||||
"The weather is nice today.",
|
||||
|
|
@ -291,10 +279,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text():
|
|||
get_weather_calls = 0
|
||||
|
||||
for chunk in chunks:
|
||||
if (
|
||||
chunk.get("type") == "content_block_start"
|
||||
and chunk["content_block"]["type"] == "tool_use"
|
||||
):
|
||||
if chunk.get("type") == "content_block_start" and chunk["content_block"]["type"] == "tool_use":
|
||||
if chunk["content_block"]["name"] == "get_weather":
|
||||
get_weather_calls += 1
|
||||
|
||||
|
|
|
|||
|
|
@ -29,13 +29,12 @@ def _call_handler_and_capture_optional_params(thinking=None, **extra_kwargs):
|
|||
"""
|
||||
captured = {}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.handler."
|
||||
"base_llm_http_handler"
|
||||
) as mock_handler, patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.handler."
|
||||
"ProviderConfigManager"
|
||||
) as mock_pcm:
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.handler.base_llm_http_handler"
|
||||
) as mock_handler,
|
||||
patch("litellm.llms.anthropic.experimental_pass_through.messages.handler.ProviderConfigManager") as mock_pcm,
|
||||
):
|
||||
# Make get_provider_anthropic_messages_config return a non-None config
|
||||
# so the handler takes the native Anthropic path
|
||||
mock_pcm.get_provider_anthropic_messages_config.return_value = MagicMock()
|
||||
|
|
@ -71,9 +70,7 @@ class TestReasoningAutoSummaryMessages:
|
|||
def test_adaptive_thinking_gets_display_summarized(self):
|
||||
"""reasoning_auto_summary=True + thinking.type='adaptive' -> display='summarized'."""
|
||||
with patch.object(litellm, "reasoning_auto_summary", True):
|
||||
params = _call_handler_and_capture_optional_params(
|
||||
thinking={"type": "adaptive", "budget_tokens": 5000}
|
||||
)
|
||||
params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000})
|
||||
thinking = params.get("thinking", {})
|
||||
assert thinking.get("display") == "summarized"
|
||||
assert thinking.get("type") == "adaptive"
|
||||
|
|
@ -82,9 +79,7 @@ class TestReasoningAutoSummaryMessages:
|
|||
def test_enabled_thinking_gets_display_summarized(self):
|
||||
"""reasoning_auto_summary=True + thinking.type='enabled' -> display='summarized'."""
|
||||
with patch.object(litellm, "reasoning_auto_summary", True):
|
||||
params = _call_handler_and_capture_optional_params(
|
||||
thinking={"type": "enabled", "budget_tokens": 10000}
|
||||
)
|
||||
params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000})
|
||||
thinking = params.get("thinking", {})
|
||||
assert thinking.get("display") == "summarized"
|
||||
assert thinking.get("type") == "enabled"
|
||||
|
|
@ -92,18 +87,14 @@ class TestReasoningAutoSummaryMessages:
|
|||
def test_disabled_thinking_no_display(self):
|
||||
"""reasoning_auto_summary=True + thinking.type='disabled' -> display NOT set."""
|
||||
with patch.object(litellm, "reasoning_auto_summary", True):
|
||||
params = _call_handler_and_capture_optional_params(
|
||||
thinking={"type": "disabled"}
|
||||
)
|
||||
params = _call_handler_and_capture_optional_params(thinking={"type": "disabled"})
|
||||
thinking = params.get("thinking", {})
|
||||
assert "display" not in thinking
|
||||
|
||||
def test_no_injection_when_flag_false(self):
|
||||
"""reasoning_auto_summary=False + active thinking -> display NOT set."""
|
||||
with patch.object(litellm, "reasoning_auto_summary", False):
|
||||
params = _call_handler_and_capture_optional_params(
|
||||
thinking={"type": "enabled", "budget_tokens": 10000}
|
||||
)
|
||||
params = _call_handler_and_capture_optional_params(thinking={"type": "enabled", "budget_tokens": 10000})
|
||||
thinking = params.get("thinking", {})
|
||||
assert "display" not in thinking
|
||||
|
||||
|
|
@ -117,12 +108,11 @@ class TestReasoningAutoSummaryMessages:
|
|||
|
||||
def test_env_var_enables_auto_summary(self):
|
||||
"""LITELLM_REASONING_AUTO_SUMMARY=true env var enables the feature."""
|
||||
with patch.object(litellm, "reasoning_auto_summary", False), patch.dict(
|
||||
os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}
|
||||
with (
|
||||
patch.object(litellm, "reasoning_auto_summary", False),
|
||||
patch.dict(os.environ, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}),
|
||||
):
|
||||
params = _call_handler_and_capture_optional_params(
|
||||
thinking={"type": "adaptive", "budget_tokens": 5000}
|
||||
)
|
||||
params = _call_handler_and_capture_optional_params(thinking={"type": "adaptive", "budget_tokens": 5000})
|
||||
thinking = params.get("thinking", {})
|
||||
assert thinking.get("display") == "summarized"
|
||||
|
||||
|
|
|
|||
|
|
@ -23,11 +23,7 @@ def test_optional_param_filtering_unchanged():
|
|||
"not_a_real_param": "drop me", # invalid key dropped
|
||||
"stream": True,
|
||||
}
|
||||
result = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params
|
||||
)
|
||||
)
|
||||
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(params)
|
||||
assert result == {"temperature": 0.5, "tools": [{"name": "x"}], "stream": True}
|
||||
assert "top_p" not in result
|
||||
assert "not_a_real_param" not in result
|
||||
|
|
@ -37,9 +33,7 @@ def test_valid_keys_are_memoized():
|
|||
_anthropic_messages_optional_param_keys.cache_clear()
|
||||
first = _anthropic_messages_optional_param_keys()
|
||||
for _ in range(50):
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
{"temperature": 0.1}
|
||||
)
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param({"temperature": 0.1})
|
||||
info = _anthropic_messages_optional_param_keys.cache_info()
|
||||
# Resolved exactly once despite many calls.
|
||||
assert info.misses == 1
|
||||
|
|
@ -51,23 +45,16 @@ def test_valid_keys_are_memoized():
|
|||
|
||||
|
||||
def test_empty_params():
|
||||
assert (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
{}
|
||||
)
|
||||
== {}
|
||||
)
|
||||
assert AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param({}) == {}
|
||||
|
||||
|
||||
def test_drop_params_strips_speed_for_unsupported_model():
|
||||
original = litellm.drop_params
|
||||
litellm.drop_params = True
|
||||
try:
|
||||
result = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"speed": "fast", "temperature": 0.5},
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"speed": "fast", "temperature": 0.5},
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
finally:
|
||||
litellm.drop_params = original
|
||||
|
|
@ -80,11 +67,9 @@ def test_drop_params_keeps_speed_for_supporting_model():
|
|||
original = litellm.drop_params
|
||||
litellm.drop_params = True
|
||||
try:
|
||||
result = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"speed": "fast"},
|
||||
model="claude-opus-4-6",
|
||||
)
|
||||
result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params={"speed": "fast"},
|
||||
model="claude-opus-4-6",
|
||||
)
|
||||
finally:
|
||||
litellm.drop_params = original
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
|
@ -14,25 +13,13 @@ class MockCompletionStream:
|
|||
def __init__(self):
|
||||
self.responses = [
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content="Hello"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
|
||||
),
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=" World"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content=" World"), index=0, finish_reason=None)],
|
||||
),
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=""), index=0, finish_reason="stop"
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")],
|
||||
),
|
||||
]
|
||||
self.index = 0
|
||||
|
|
@ -50,9 +37,7 @@ class MockCompletionStream:
|
|||
|
||||
def test_anthropic_sse_wrapper_format():
|
||||
"""Test that the SSE wrapper produces proper event and data formatting"""
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=MockCompletionStream(), model="claude-3"
|
||||
)
|
||||
wrapper = AnthropicStreamWrapper(completion_stream=MockCompletionStream(), model="claude-3")
|
||||
|
||||
# Get the first chunk from the SSE wrapper
|
||||
first_chunk = next(wrapper.anthropic_sse_wrapper())
|
||||
|
|
@ -73,9 +58,7 @@ def test_anthropic_sse_wrapper_format():
|
|||
|
||||
def test_anthropic_sse_wrapper_event_types():
|
||||
"""Test that different chunk types produce correct event types"""
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=MockCompletionStream(), model="claude-3"
|
||||
)
|
||||
wrapper = AnthropicStreamWrapper(completion_stream=MockCompletionStream(), model="claude-3")
|
||||
|
||||
chunks = []
|
||||
for chunk in wrapper.anthropic_sse_wrapper():
|
||||
|
|
@ -104,18 +87,10 @@ async def test_async_anthropic_sse_wrapper():
|
|||
def __init__(self):
|
||||
self.responses = [
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content="Hello"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content="Hello"), index=0, finish_reason=None)],
|
||||
),
|
||||
ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=" World"), index=0, finish_reason=None
|
||||
)
|
||||
],
|
||||
choices=[StreamingChoices(delta=Delta(content=" World"), index=0, finish_reason=None)],
|
||||
),
|
||||
]
|
||||
self.index = 0
|
||||
|
|
@ -130,9 +105,7 @@ async def test_async_anthropic_sse_wrapper():
|
|||
self.index += 1
|
||||
return response
|
||||
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=AsyncMockCompletionStream(), model="claude-3"
|
||||
)
|
||||
wrapper = AnthropicStreamWrapper(completion_stream=AsyncMockCompletionStream(), model="claude-3")
|
||||
|
||||
# Get the first chunk from the async SSE wrapper
|
||||
first_chunk = None
|
||||
|
|
|
|||
|
|
@ -165,7 +165,7 @@ async def test_async_sse_wrapper_treats_message_stop_bytes_as_complete():
|
|||
def test_is_message_stop_chunk():
|
||||
assert _is_message_stop_chunk({"type": "message_stop"}) is True
|
||||
assert _is_message_stop_chunk({"type": "message_delta"}) is False
|
||||
assert _is_message_stop_chunk(b'event: message_stop\ndata: {}\n\n') is True
|
||||
assert _is_message_stop_chunk(b"event: message_stop\ndata: {}\n\n") is True
|
||||
assert _is_message_stop_chunk(b"raw-bytes") is False
|
||||
assert _is_message_stop_chunk("message_stop") is False
|
||||
|
||||
|
|
@ -177,7 +177,7 @@ def test_is_message_stop_chunk_ignores_substring_in_payload():
|
|||
not be treated as a terminal stop event.
|
||||
"""
|
||||
delta_frame_with_substring = (
|
||||
b'event: content_block_delta\n'
|
||||
b"event: content_block_delta\n"
|
||||
b'data: {"type": "content_block_delta", "delta": '
|
||||
b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n'
|
||||
)
|
||||
|
|
@ -281,10 +281,11 @@ async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_mes
|
|||
payload text contains `message_stop` (but never emits the actual
|
||||
`event: message_stop` frame) must still be flagged as incomplete.
|
||||
"""
|
||||
|
||||
async def _byte_stream():
|
||||
yield b'event: message_start\ndata: {"type": "message_start"}\n\n'
|
||||
yield (
|
||||
b'event: content_block_delta\n'
|
||||
b"event: content_block_delta\n"
|
||||
b'data: {"type": "content_block_delta", "delta": '
|
||||
b'{"type": "input_json_delta", "partial_json": "\\"message_stop\\""}}\n\n'
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1481,9 +1481,7 @@ class TestToolResultImages:
|
|||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}
|
||||
],
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}],
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -1857,7 +1855,9 @@ class TestPromptCacheBreakpointToResponses:
|
|||
]
|
||||
|
||||
def test_system_without_breakpoint_still_becomes_instructions(self):
|
||||
request = _make_request(system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}])
|
||||
request = _make_request(
|
||||
system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}]
|
||||
)
|
||||
kwargs = _ADAPTER.translate_request(request)
|
||||
assert kwargs["instructions"] == "Be concise.\nBe helpful."
|
||||
assert kwargs["input"] == [
|
||||
|
|
|
|||
|
|
@ -5,7 +5,9 @@ Tests the AnthropicFilesConfig class which transforms between
|
|||
OpenAI-compatible file operations and Anthropic's Files API format.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import threading
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
|
@ -90,6 +92,38 @@ class TestAnthropicFilesConfig:
|
|||
api_key=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_avalidate_environment_sets_headers(self):
|
||||
headers = {}
|
||||
result = await self.config.avalidate_environment(
|
||||
headers=headers,
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="sk-ant-test-key",
|
||||
)
|
||||
assert result["x-api-key"] == "sk-ant-test-key"
|
||||
assert result["anthropic-version"] == "2023-06-01"
|
||||
assert result["anthropic-beta"] == ANTHROPIC_FILES_BETA_HEADER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.dict("os.environ", {}, clear=True)
|
||||
@patch(
|
||||
"litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key",
|
||||
return_value=None,
|
||||
)
|
||||
async def test_avalidate_environment_missing_api_key(self, mock_get_key):
|
||||
with pytest.raises(ValueError, match="Anthropic API key is required"):
|
||||
await self.config.avalidate_environment(
|
||||
headers={},
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
def test_get_supported_openai_params(self):
|
||||
params = self.config.get_supported_openai_params(model="")
|
||||
assert "purpose" in params
|
||||
|
|
@ -187,10 +221,7 @@ class TestAnthropicFilesConfig:
|
|||
litellm_params={},
|
||||
)
|
||||
|
||||
assert (
|
||||
url
|
||||
== f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag"
|
||||
)
|
||||
assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag"
|
||||
assert params == {}
|
||||
|
||||
def test_transform_retrieve_file_response(self):
|
||||
|
|
@ -411,6 +442,108 @@ class TestAnthropicFilesConfig:
|
|||
assert error.message == "Not found"
|
||||
|
||||
|
||||
_WIF_ENV = {
|
||||
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_files_seam",
|
||||
"ANTHROPIC_ORGANIZATION_ID": "org-files-seam",
|
||||
"ANTHROPIC_IDENTITY_TOKEN": "files-seam-inline-jwt",
|
||||
}
|
||||
|
||||
|
||||
class _BlockingPoster:
|
||||
"""A token-endpoint poster that blocks until released, so the test can prove
|
||||
the exchange ran off the event loop's own thread instead of freezing it."""
|
||||
|
||||
def __init__(self):
|
||||
self.release = threading.Event()
|
||||
self.thread_ids = []
|
||||
|
||||
def post(self, url, *, content, headers, timeout):
|
||||
self.thread_ids.append(threading.get_ident())
|
||||
self.release.wait(timeout=5)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "sk-ant-oat01-files-seam",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TestAnthropicFilesConfigWifAsyncSeam:
|
||||
"""Regression (Greptile P1): avalidate_environment must resolve workload identity
|
||||
federation through the async token-exchange facade, never the blocking sync one,
|
||||
so a cold WIF mint on async file retrieval doesn't freeze the event loop."""
|
||||
|
||||
def setup_method(self):
|
||||
self.config = AnthropicFilesConfig()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_avalidate_environment_wif_exchange_does_not_block_event_loop(self, monkeypatch):
|
||||
from litellm.llms.anthropic import common_utils as anthropic_common_utils
|
||||
from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token
|
||||
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine
|
||||
|
||||
for name in (
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN",
|
||||
"ANTHROPIC_API_BASE",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
for name, value in _WIF_ENV.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
|
||||
poster = _BlockingPoster()
|
||||
engine = JwtBearerTokenExchangeEngine(poster=poster)
|
||||
sync_calls = []
|
||||
|
||||
def sync_shim(litellm_params, api_base, model):
|
||||
sync_calls.append(model)
|
||||
return get_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
async def async_shim(litellm_params, api_base, model):
|
||||
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
|
||||
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
|
||||
|
||||
ticks = []
|
||||
|
||||
async def ticker():
|
||||
for i in range(20):
|
||||
await asyncio.sleep(0.005)
|
||||
ticks.append(i)
|
||||
|
||||
ticker_task = asyncio.create_task(ticker())
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
validate_task = asyncio.create_task(
|
||||
self.config.avalidate_environment(
|
||||
headers={},
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key=None,
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
# The ticker kept advancing while the exchange was still blocked on
|
||||
# poster.release, proving avalidate_environment did not run it inline.
|
||||
assert len(ticks) > 0
|
||||
assert not validate_task.done()
|
||||
|
||||
poster.release.set()
|
||||
headers = await validate_task
|
||||
await ticker_task
|
||||
|
||||
assert headers["authorization"] == "Bearer sk-ant-oat01-files-seam"
|
||||
assert sync_calls == []
|
||||
assert poster.thread_ids
|
||||
assert poster.thread_ids[0] != threading.get_ident()
|
||||
|
||||
|
||||
class TestProviderConfigRegistration:
|
||||
"""Test that AnthropicFilesConfig is properly registered."""
|
||||
|
||||
|
|
|
|||
|
|
@ -101,9 +101,7 @@ async def test_anthropic_native_interceptor_skipped():
|
|||
)
|
||||
|
||||
h = AdvisorOrchestrationHandler()
|
||||
assert not h.can_handle(
|
||||
[ADVISOR_TOOL], "anthropic"
|
||||
), "Interceptor must NOT trigger for anthropic provider"
|
||||
assert not h.can_handle([ADVISOR_TOOL], "anthropic"), "Interceptor must NOT trigger for anthropic provider"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -204,9 +202,7 @@ async def test_loop_one_advisor_call():
|
|||
assert "is_prime" in texts[0]["text"]
|
||||
|
||||
# No advisor tool_use blocks in final response
|
||||
advisor_uses = [
|
||||
b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"
|
||||
]
|
||||
advisor_uses = [b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"]
|
||||
assert len(advisor_uses) == 0
|
||||
|
||||
|
||||
|
|
@ -366,9 +362,7 @@ async def test_prior_advisor_blocks_replaced_in_history():
|
|||
|
||||
# Text block with advisor feedback must be present
|
||||
text_blocks = [b for b in content if b.get("type") == "text"]
|
||||
feedback_blocks = [
|
||||
b for b in text_blocks if "advisor_feedback" in b.get("text", "")
|
||||
]
|
||||
feedback_blocks = [b for b in text_blocks if "advisor_feedback" in b.get("text", "")]
|
||||
assert len(feedback_blocks) >= 1
|
||||
assert "trial division" in feedback_blocks[0]["text"]
|
||||
|
||||
|
|
@ -707,11 +701,7 @@ async def test_advisor_ignores_tool_credentials_when_clientside_disabled():
|
|||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"litellm.proxy.proxy_server": _fake_proxy_server(
|
||||
{"allow_client_side_credentials": False}
|
||||
)
|
||||
},
|
||||
{"litellm.proxy.proxy_server": _fake_proxy_server({"allow_client_side_credentials": False})},
|
||||
):
|
||||
captured = await _run_advisor_and_capture_subcall_kwargs()
|
||||
assert captured["api_key"] is None
|
||||
|
|
@ -726,11 +716,7 @@ async def test_advisor_uses_tool_credentials_when_clientside_enabled():
|
|||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"litellm.proxy.proxy_server": _fake_proxy_server(
|
||||
{"allow_client_side_credentials": True}
|
||||
)
|
||||
},
|
||||
{"litellm.proxy.proxy_server": _fake_proxy_server({"allow_client_side_credentials": True})},
|
||||
):
|
||||
captured = await _run_advisor_and_capture_subcall_kwargs()
|
||||
assert captured["api_key"] == "sk-other"
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,4 +1,9 @@
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountTokensConfig,
|
||||
)
|
||||
|
|
@ -88,3 +93,72 @@ def test_transform_no_system_no_tools():
|
|||
|
||||
assert "system" not in result
|
||||
assert "tools" not in result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected"),
|
||||
[
|
||||
(None, "https://api.anthropic.com/v1/messages/count_tokens"),
|
||||
("", "https://api.anthropic.com/v1/messages/count_tokens"),
|
||||
("https://gateway.example", "https://gateway.example/v1/messages/count_tokens"),
|
||||
("https://gateway.example/", "https://gateway.example/v1/messages/count_tokens"),
|
||||
("https://gateway.example/v1", "https://gateway.example/v1/messages/count_tokens"),
|
||||
("https://gateway.example/anthropic/v1/messages", "https://gateway.example/anthropic/v1/messages/count_tokens"),
|
||||
],
|
||||
)
|
||||
def test_endpoint_appends_count_tokens_path_to_deployment_api_base(api_base, expected, monkeypatch):
|
||||
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
|
||||
assert AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("env_name", ["ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"])
|
||||
@pytest.mark.parametrize("api_base", [None, ""])
|
||||
def test_endpoint_without_deployment_api_base_follows_env_base(env_name, api_base, monkeypatch):
|
||||
"""Chat and the federated exchange resolve an unset deployment base through the environment,
|
||||
so an env-only gateway must receive the count too, never Anthropic's public host."""
|
||||
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
|
||||
monkeypatch.setenv(env_name, "https://env-gateway.example/v1/messages/")
|
||||
assert (
|
||||
AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint(api_base)
|
||||
== "https://env-gateway.example/v1/messages/count_tokens"
|
||||
)
|
||||
|
||||
|
||||
def test_endpoint_prefers_deployment_api_base_over_env_base(monkeypatch):
|
||||
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://env-gateway.example")
|
||||
assert (
|
||||
AnthropicCountTokensConfig().get_anthropic_count_tokens_endpoint("https://gateway.example/v1")
|
||||
== "https://gateway.example/v1/messages/count_tokens"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def httpx_transport_clients(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
client_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
|
||||
if client_cache is not None:
|
||||
client_cache.flush_cache()
|
||||
yield
|
||||
if client_cache is not None:
|
||||
client_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_posts_to_count_tokens_path_under_deployment_api_base(httpx_transport_clients):
|
||||
"""A deployment api_base names the chat host, so a handler that posts to it verbatim lands on
|
||||
the host root, gets a 404, and the official count silently degrades to the local tokenizer."""
|
||||
with respx.mock:
|
||||
route = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 7})
|
||||
)
|
||||
result = await AnthropicCountTokensHandler().handle_count_tokens_request(
|
||||
model="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_key="sk-ant-api03-test-key",
|
||||
api_base="https://gateway.example",
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert result == {"input_tokens": 7}
|
||||
|
|
|
|||
|
|
@ -70,9 +70,7 @@ class TestAnthropicFilesHandler:
|
|||
@pytest.fixture
|
||||
def mock_anthropic_batch_results_canceled(self):
|
||||
"""Mock Anthropic batch results with canceled status"""
|
||||
return json.dumps(
|
||||
{"custom_id": "test-request-3", "result": {"type": "canceled"}}
|
||||
).encode("utf-8")
|
||||
return json.dumps({"custom_id": "test-request-3", "result": {"type": "canceled"}}).encode("utf-8")
|
||||
|
||||
@pytest.fixture
|
||||
def mock_anthropic_batch_results_mixed(self):
|
||||
|
|
@ -114,9 +112,7 @@ class TestAnthropicFilesHandler:
|
|||
return "\n".join(lines).encode("utf-8")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_success(
|
||||
self, handler, mock_anthropic_batch_results_succeeded
|
||||
):
|
||||
async def test_afile_content_success(self, handler, mock_anthropic_batch_results_succeeded):
|
||||
"""Test successful file content retrieval and transformation"""
|
||||
file_content_request: FileContentRequest = {
|
||||
"file_id": "batch_123",
|
||||
|
|
@ -135,16 +131,14 @@ class TestAnthropicFilesHandler:
|
|||
),
|
||||
)
|
||||
|
||||
with patch(
|
||||
with patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -161,9 +155,7 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
# Verify transformation to OpenAI format
|
||||
content = result.response.content.decode("utf-8")
|
||||
lines = [
|
||||
line for line in content.strip().split("\n") if line.strip()
|
||||
]
|
||||
lines = [line for line in content.strip().split("\n") if line.strip()]
|
||||
assert len(lines) == 1
|
||||
|
||||
transformed_result = json.loads(lines[0])
|
||||
|
|
@ -172,18 +164,13 @@ class TestAnthropicFilesHandler:
|
|||
assert "body" in transformed_result["response"]
|
||||
# Verify body has required OpenAI format fields
|
||||
assert "id" in transformed_result["response"]["body"]
|
||||
assert (
|
||||
transformed_result["response"]["body"]["object"]
|
||||
== "chat.completion"
|
||||
)
|
||||
assert transformed_result["response"]["body"]["object"] == "chat.completion"
|
||||
assert "choices" in transformed_result["response"]["body"]
|
||||
# Verify request_id matches the original message id
|
||||
assert transformed_result["response"]["request_id"] == "msg_123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_with_prefix(
|
||||
self, handler, mock_anthropic_batch_results_succeeded
|
||||
):
|
||||
async def test_afile_content_with_prefix(self, handler, mock_anthropic_batch_results_succeeded):
|
||||
"""Test file content retrieval with anthropic_batch_results: prefix"""
|
||||
file_content_request: FileContentRequest = {
|
||||
"file_id": "anthropic_batch_results:batch_123",
|
||||
|
|
@ -203,14 +190,12 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -228,9 +213,7 @@ class TestAnthropicFilesHandler:
|
|||
assert "batch_123" in call_url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_errored_result(
|
||||
self, handler, mock_anthropic_batch_results_errored
|
||||
):
|
||||
async def test_afile_content_errored_result(self, handler, mock_anthropic_batch_results_errored):
|
||||
"""Test transformation of errored batch results"""
|
||||
file_content_request: FileContentRequest = {
|
||||
"file_id": "batch_123",
|
||||
|
|
@ -250,14 +233,12 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -269,29 +250,17 @@ class TestAnthropicFilesHandler:
|
|||
)
|
||||
|
||||
content = result.response.content.decode("utf-8")
|
||||
lines = [
|
||||
line for line in content.strip().split("\n") if line.strip()
|
||||
]
|
||||
lines = [line for line in content.strip().split("\n") if line.strip()]
|
||||
assert len(lines) == 1
|
||||
|
||||
transformed_result = json.loads(lines[0])
|
||||
assert transformed_result["custom_id"] == "test-request-2"
|
||||
assert (
|
||||
transformed_result["response"]["status_code"] == 400
|
||||
) # invalid_request_error maps to 400
|
||||
assert (
|
||||
transformed_result["response"]["body"]["error"]["type"]
|
||||
== "invalid_request_error"
|
||||
)
|
||||
assert (
|
||||
transformed_result["response"]["body"]["error"]["message"]
|
||||
== "Invalid request"
|
||||
)
|
||||
assert transformed_result["response"]["status_code"] == 400 # invalid_request_error maps to 400
|
||||
assert transformed_result["response"]["body"]["error"]["type"] == "invalid_request_error"
|
||||
assert transformed_result["response"]["body"]["error"]["message"] == "Invalid request"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_canceled_result(
|
||||
self, handler, mock_anthropic_batch_results_canceled
|
||||
):
|
||||
async def test_afile_content_canceled_result(self, handler, mock_anthropic_batch_results_canceled):
|
||||
"""Test transformation of canceled batch results"""
|
||||
file_content_request: FileContentRequest = {
|
||||
"file_id": "batch_123",
|
||||
|
|
@ -311,14 +280,12 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -330,23 +297,16 @@ class TestAnthropicFilesHandler:
|
|||
)
|
||||
|
||||
content = result.response.content.decode("utf-8")
|
||||
lines = [
|
||||
line for line in content.strip().split("\n") if line.strip()
|
||||
]
|
||||
lines = [line for line in content.strip().split("\n") if line.strip()]
|
||||
assert len(lines) == 1
|
||||
|
||||
transformed_result = json.loads(lines[0])
|
||||
assert transformed_result["custom_id"] == "test-request-3"
|
||||
assert transformed_result["response"]["status_code"] == 400
|
||||
assert (
|
||||
"Batch request was canceled"
|
||||
in transformed_result["response"]["body"]["error"]["message"]
|
||||
)
|
||||
assert "Batch request was canceled" in transformed_result["response"]["body"]["error"]["message"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_mixed_results(
|
||||
self, handler, mock_anthropic_batch_results_mixed
|
||||
):
|
||||
async def test_afile_content_mixed_results(self, handler, mock_anthropic_batch_results_mixed):
|
||||
"""Test transformation of mixed batch results (succeeded, errored, expired)"""
|
||||
file_content_request: FileContentRequest = {
|
||||
"file_id": "batch_123",
|
||||
|
|
@ -366,14 +326,12 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -385,9 +343,7 @@ class TestAnthropicFilesHandler:
|
|||
)
|
||||
|
||||
content = result.response.content.decode("utf-8")
|
||||
lines = [
|
||||
line for line in content.strip().split("\n") if line.strip()
|
||||
]
|
||||
lines = [line for line in content.strip().split("\n") if line.strip()]
|
||||
assert len(lines) == 3
|
||||
|
||||
# Check first result (succeeded)
|
||||
|
|
@ -396,9 +352,7 @@ class TestAnthropicFilesHandler:
|
|||
|
||||
# Check second result (errored)
|
||||
result2 = json.loads(lines[1])
|
||||
assert (
|
||||
result2["response"]["status_code"] == 429
|
||||
) # rate_limit_error maps to 429
|
||||
assert result2["response"]["status_code"] == 429 # rate_limit_error maps to 429
|
||||
|
||||
# Check third result (expired)
|
||||
result3 = json.loads(lines[2])
|
||||
|
|
@ -415,12 +369,12 @@ class TestAnthropicFilesHandler:
|
|||
}
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_auth_header", return_value=None
|
||||
handler.anthropic_model_info,
|
||||
"aget_auth_header",
|
||||
new=AsyncMock(return_value=None),
|
||||
):
|
||||
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
|
||||
await handler.afile_content(
|
||||
file_content_request=file_content_request, api_key=None
|
||||
)
|
||||
await handler.afile_content(file_content_request=file_content_request, api_key=None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_missing_file_id(self, handler):
|
||||
|
|
@ -432,9 +386,7 @@ class TestAnthropicFilesHandler:
|
|||
}
|
||||
|
||||
with pytest.raises(ValueError, match="file_id is required"):
|
||||
await handler.afile_content(
|
||||
file_content_request=file_content_request, api_key="test-api-key"
|
||||
)
|
||||
await handler.afile_content(file_content_request=file_content_request, api_key="test-api-key")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_http_error(self, handler):
|
||||
|
|
@ -454,21 +406,17 @@ class TestAnthropicFilesHandler:
|
|||
),
|
||||
)
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"Not Found", request=mock_response.request, response=mock_response
|
||||
)
|
||||
side_effect=httpx.HTTPStatusError("Not Found", request=mock_response.request, response=mock_response)
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with patch.object(
|
||||
handler.anthropic_model_info, "get_api_key", return_value="test-api-key"
|
||||
):
|
||||
with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"):
|
||||
with patch.object(
|
||||
handler.anthropic_model_info,
|
||||
"get_api_base",
|
||||
|
|
@ -480,6 +428,86 @@ class TestAnthropicFilesHandler:
|
|||
api_key="test-api-key",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_resolves_wif_via_async_facade(
|
||||
self, handler, mock_anthropic_batch_results_succeeded, monkeypatch
|
||||
):
|
||||
"""Regression: afile_content ran the blocking WIF mint on the event loop
|
||||
through the sync get_auth_header; it must go through the async facade."""
|
||||
import threading
|
||||
|
||||
from litellm.llms.anthropic import common_utils as anthropic_common_utils
|
||||
from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token
|
||||
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine
|
||||
|
||||
for name in (
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN",
|
||||
"ANTHROPIC_API_BASE",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_files")
|
||||
monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-files")
|
||||
monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "files-inline-jwt")
|
||||
|
||||
minted = "sk-ant-oat01-files-minted"
|
||||
thread_ids = []
|
||||
|
||||
class ThreadRecordingPoster:
|
||||
def post(self, url, *, content, headers, timeout):
|
||||
thread_ids.append(threading.get_ident())
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"access_token": minted, "token_type": "Bearer", "expires_in": 3600},
|
||||
)
|
||||
|
||||
engine = JwtBearerTokenExchangeEngine(poster=ThreadRecordingPoster())
|
||||
sync_calls = []
|
||||
|
||||
def sync_shim(litellm_params, api_base, model):
|
||||
sync_calls.append(model)
|
||||
return get_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
async def async_shim(litellm_params, api_base, model):
|
||||
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
|
||||
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=mock_anthropic_batch_results_succeeded,
|
||||
headers={"content-type": "application/json"},
|
||||
request=httpx.Request(
|
||||
method="GET",
|
||||
url="https://api.anthropic.com/v1/messages/batches/batch_123/results",
|
||||
),
|
||||
)
|
||||
|
||||
with patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.llms.anthropic.files.handler.get_async_httpx_client"
|
||||
) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
await handler.afile_content(
|
||||
file_content_request={
|
||||
"file_id": "batch_123",
|
||||
"extra_headers": None,
|
||||
"extra_body": None,
|
||||
},
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
sent_headers = mock_client.get.call_args.kwargs["headers"]
|
||||
|
||||
assert sent_headers["authorization"] == f"Bearer {minted}"
|
||||
assert "oauth-2025-04-20" in sent_headers["anthropic-beta"]
|
||||
assert sync_calls == []
|
||||
assert thread_ids and thread_ids[0] != threading.get_ident()
|
||||
|
||||
|
||||
class TestAnthropicBatchesConfig:
|
||||
"""Test Anthropic Batches Config for batch retrieval transformation"""
|
||||
|
|
@ -562,15 +590,11 @@ class TestAnthropicBatchesConfig:
|
|||
)
|
||||
assert url == "https://api.anthropic.com/v1/messages/batches/batch_123"
|
||||
|
||||
def test_transform_retrieve_batch_response_in_progress(
|
||||
self, config, mock_anthropic_batch_response_in_progress
|
||||
):
|
||||
def test_transform_retrieve_batch_response_in_progress(self, config, mock_anthropic_batch_response_in_progress):
|
||||
"""Test transformation of in_progress batch response"""
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=json.dumps(mock_anthropic_batch_response_in_progress).encode(
|
||||
"utf-8"
|
||||
),
|
||||
content=json.dumps(mock_anthropic_batch_response_in_progress).encode("utf-8"),
|
||||
request=httpx.Request(
|
||||
method="GET",
|
||||
url="https://api.anthropic.com/v1/messages/batches/batch_123",
|
||||
|
|
@ -596,9 +620,7 @@ class TestAnthropicBatchesConfig:
|
|||
assert batch.in_progress_at is not None
|
||||
assert batch.completed_at is None
|
||||
|
||||
def test_transform_retrieve_batch_response_completed(
|
||||
self, config, mock_anthropic_batch_response_completed
|
||||
):
|
||||
def test_transform_retrieve_batch_response_completed(self, config, mock_anthropic_batch_response_completed):
|
||||
"""Test transformation of completed batch response"""
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
|
|
@ -624,9 +646,7 @@ class TestAnthropicBatchesConfig:
|
|||
assert batch.request_counts.completed == 10
|
||||
assert batch.request_counts.failed == 0
|
||||
|
||||
def test_transform_retrieve_batch_response_canceling(
|
||||
self, config, mock_anthropic_batch_response_canceling
|
||||
):
|
||||
def test_transform_retrieve_batch_response_canceling(self, config, mock_anthropic_batch_response_canceling):
|
||||
"""Test transformation of canceling batch response"""
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
|
|
@ -663,9 +683,7 @@ class TestAnthropicBatchesConfig:
|
|||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
with pytest.raises(
|
||||
ValueError, match="Failed to parse Anthropic batch response"
|
||||
):
|
||||
with pytest.raises(ValueError, match="Failed to parse Anthropic batch response"):
|
||||
config.transform_retrieve_batch_response(
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
raw_response=mock_response,
|
||||
|
|
|
|||
1222
tests/test_litellm/llms/anthropic/test_anthropic_wif.py
Normal file
1222
tests/test_litellm/llms/anthropic/test_anthropic_wif.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -3,8 +3,6 @@ Test that Azure AI Anthropic models have cache pricing configured.
|
|||
Verifies the fix for issue #19532.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm import get_model_info
|
||||
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
|
||||
|
|
@ -27,16 +25,11 @@ def reload_model_costs():
|
|||
("claude-sonnet-4-5", 3.75e-06, 3e-07),
|
||||
],
|
||||
)
|
||||
def test_azure_ai_claude_cache_pricing(
|
||||
model, expected_cache_creation_cost, expected_cache_read_cost
|
||||
):
|
||||
def test_azure_ai_claude_cache_pricing(model, expected_cache_creation_cost, expected_cache_read_cost):
|
||||
"""Test that Azure AI Claude models have correct cache pricing."""
|
||||
model_info = get_model_info(model=model, custom_llm_provider="azure_ai")
|
||||
|
||||
assert model_info.get("cache_creation_input_token_cost") is not None
|
||||
assert model_info.get("cache_read_input_token_cost") is not None
|
||||
assert (
|
||||
model_info.get("cache_creation_input_token_cost")
|
||||
== expected_cache_creation_cost
|
||||
)
|
||||
assert model_info.get("cache_creation_input_token_cost") == expected_cache_creation_cost
|
||||
assert model_info.get("cache_read_input_token_cost") == expected_cache_read_cost
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ being either a ``dict`` or a ``ServerToolUse`` pydantic instance.
|
|||
See https://github.com/BerriAI/litellm/issues/26153.
|
||||
"""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
|
||||
|
|
@ -54,7 +53,8 @@ def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use():
|
|||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
model_info=info,
|
||||
usage=usage, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.03)
|
||||
|
|
@ -65,7 +65,8 @@ def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use():
|
|||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
model_info=info,
|
||||
usage=usage, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.03)
|
||||
|
|
@ -76,7 +77,8 @@ def test_get_cost_for_anthropic_web_search_with_none_server_tool_use():
|
|||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
model_info=info,
|
||||
usage=usage, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
|
|
|
|||
|
|
@ -10,9 +10,9 @@ Regression test for https://github.com/BerriAI/litellm/issues/22040
|
|||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
|
||||
)
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
|
||||
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountTokensConfig,
|
||||
|
|
@ -78,9 +78,177 @@ class TestCountTokensOAuthHeaders:
|
|||
headers = config.get_required_headers(FAKE_OAUTH_TOKEN)
|
||||
|
||||
beta_value = headers.get("anthropic-beta", "")
|
||||
assert (
|
||||
"token-counting" in beta_value
|
||||
), f"token-counting beta missing from OAuth headers: {beta_value}"
|
||||
assert (
|
||||
"oauth-2025-04-20" in beta_value
|
||||
), f"oauth beta missing from OAuth headers: {beta_value}"
|
||||
assert "token-counting" in beta_value, f"token-counting beta missing from OAuth headers: {beta_value}"
|
||||
assert "oauth-2025-04-20" in beta_value, f"oauth beta missing from OAuth headers: {beta_value}"
|
||||
|
||||
|
||||
class TestCountTokensUsesWorkloadIdentity:
|
||||
"""A federated deployment holds no static key. Without minting one, count_tokens returns None
|
||||
and the caller silently falls back to the local tokenizer, so the number a federated
|
||||
deployment reports would never come from Anthropic."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_federated_deployment_mints_and_counts(self, monkeypatch):
|
||||
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
|
||||
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
minted = "sk-ant-oat01-minted-for-count"
|
||||
|
||||
async def fake_mint(_params, _api_base, _model):
|
||||
return minted
|
||||
|
||||
monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
async def fake_request(**kwargs):
|
||||
seen.update(kwargs)
|
||||
return {"input_tokens": 42}
|
||||
|
||||
monkeypatch.setattr(
|
||||
token_counter_module.anthropic_count_tokens_handler,
|
||||
"handle_count_tokens_request",
|
||||
fake_request,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "fdrl_x",
|
||||
"anthropic_organization_id": "org-x",
|
||||
}
|
||||
},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.total_tokens == 42
|
||||
assert seen["api_key"] == minted
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_auth_token_deployment_never_mints(self, monkeypatch):
|
||||
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
|
||||
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.setenv("ANTHROPIC_AUTH_TOKEN", "bearer-token-for-testing")
|
||||
mint_calls: list[str] = []
|
||||
|
||||
async def fake_mint(_params, _api_base, model):
|
||||
mint_calls.append(model)
|
||||
return "sk-ant-oat01-should-not-be-minted"
|
||||
|
||||
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
|
||||
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "fdrl_x",
|
||||
"anthropic_organization_id": "org-x",
|
||||
}
|
||||
},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert mint_calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch):
|
||||
import litellm
|
||||
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
|
||||
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
|
||||
async def failing_mint(_params, _api_base, model):
|
||||
raise litellm.AuthenticationError(
|
||||
message="federation_rule_id is not a well-formed fdrl_ tagged ID",
|
||||
llm_provider="anthropic",
|
||||
model=model,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", failing_mint, raising=False)
|
||||
|
||||
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "not-a-rule",
|
||||
"anthropic_organization_id": "org-x",
|
||||
}
|
||||
},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 401
|
||||
assert result.total_tokens == 0
|
||||
assert "fdrl_" in (result.error_message or "")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_vault_backed_static_key_never_mints(self, monkeypatch):
|
||||
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
|
||||
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
vault_key = "sk-ant-api03-only-in-the-vault"
|
||||
|
||||
def vault_only(secret_name, default_value=None):
|
||||
return vault_key if secret_name == "ANTHROPIC_API_KEY" else None
|
||||
|
||||
monkeypatch.setattr("litellm.secret_managers.main.get_secret_str", vault_only, raising=False)
|
||||
|
||||
mint_calls: list[str] = []
|
||||
|
||||
async def fake_mint(_params, _api_base, model):
|
||||
mint_calls.append(model)
|
||||
return "sk-ant-oat01-should-not-be-minted"
|
||||
|
||||
monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False)
|
||||
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
async def fake_request(**kwargs):
|
||||
seen.update(kwargs)
|
||||
return {"input_tokens": 7}
|
||||
|
||||
monkeypatch.setattr(
|
||||
token_counter_module.anthropic_count_tokens_handler,
|
||||
"handle_count_tokens_request",
|
||||
fake_request,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
result = await token_counter_module.AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"anthropic_federation_rule_id": "fdrl_x",
|
||||
"anthropic_organization_id": "org-x",
|
||||
}
|
||||
},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.total_tokens == 7
|
||||
assert seen["api_key"] == vault_key
|
||||
assert mint_calls == []
|
||||
|
|
|
|||
|
|
@ -12,9 +12,7 @@ import sys
|
|||
import os
|
||||
|
||||
# Add the parent directory to the path so we can import litellm
|
||||
sys.path.insert(
|
||||
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
|
||||
)
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -68,10 +66,7 @@ class TestMessageSanitization:
|
|||
assert sanitized[1]["role"] == "assistant"
|
||||
assert sanitized[2]["role"] == "tool"
|
||||
assert sanitized[2]["tool_call_id"] == "toolu_01Kus2cC3ydjBW7UK4GJqBP4"
|
||||
assert (
|
||||
"skipped" in sanitized[2]["content"].lower()
|
||||
or "interrupted" in sanitized[2]["content"].lower()
|
||||
)
|
||||
assert "skipped" in sanitized[2]["content"].lower() or "interrupted" in sanitized[2]["content"].lower()
|
||||
assert "get_weather" in sanitized[2]["content"]
|
||||
|
||||
def test_case_a_orphaned_tool_call_multiple(self):
|
||||
|
|
@ -115,12 +110,8 @@ class TestMessageSanitization:
|
|||
assert len(sanitized) == 4
|
||||
assert sanitized[0]["role"] == "user"
|
||||
assert sanitized[1]["role"] == "assistant"
|
||||
assert (
|
||||
sanitized[2]["tool_call_id"] == "call_1"
|
||||
) # Original tool result (first in tool_calls)
|
||||
assert (
|
||||
sanitized[3]["tool_call_id"] == "call_2"
|
||||
) # Dummy added for missing call_2
|
||||
assert sanitized[2]["tool_call_id"] == "call_1" # Original tool result (first in tool_calls)
|
||||
assert sanitized[3]["tool_call_id"] == "call_2" # Dummy added for missing call_2
|
||||
|
||||
def test_case_b_orphaned_tool_result(self):
|
||||
"""
|
||||
|
|
@ -188,10 +179,7 @@ class TestMessageSanitization:
|
|||
|
||||
assert len(sanitized) == 2
|
||||
assert sanitized[0]["role"] == "user"
|
||||
assert (
|
||||
sanitized[0]["content"]
|
||||
== "[System: Empty message content sanitised to satisfy protocol]"
|
||||
)
|
||||
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
|
||||
|
||||
def test_case_c_whitespace_only_content(self):
|
||||
"""
|
||||
|
|
@ -206,14 +194,8 @@ class TestMessageSanitization:
|
|||
sanitized = sanitize_messages_for_tool_calling(messages)
|
||||
|
||||
assert len(sanitized) == 2
|
||||
assert (
|
||||
sanitized[0]["content"]
|
||||
== "[System: Empty message content sanitised to satisfy protocol]"
|
||||
)
|
||||
assert (
|
||||
sanitized[1]["content"]
|
||||
== "[System: Empty message content sanitised to satisfy protocol]"
|
||||
)
|
||||
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
|
||||
assert sanitized[1]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
|
||||
|
||||
def test_case_c_valid_content_preserved(self):
|
||||
"""
|
||||
|
|
@ -270,10 +252,7 @@ class TestMessageSanitization:
|
|||
assert sanitized[2]["role"] == "tool"
|
||||
assert sanitized[2]["tool_call_id"] == "call_1" # Dummy added
|
||||
assert sanitized[3]["role"] == "user"
|
||||
assert (
|
||||
sanitized[3]["content"]
|
||||
== "[System: Empty message content sanitised to satisfy protocol]"
|
||||
)
|
||||
assert sanitized[3]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
|
||||
assert sanitized[4]["role"] == "assistant"
|
||||
|
||||
def test_modify_params_false_no_sanitization(self):
|
||||
|
|
@ -329,9 +308,7 @@ class TestMessageSanitization:
|
|||
]
|
||||
|
||||
# This should not raise an error and should add dummy tool result
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
|
||||
)
|
||||
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
|
||||
|
||||
# Should have at least 2 messages (user and assistant)
|
||||
# The tool result will be merged into user content
|
||||
|
|
@ -355,23 +332,17 @@ class TestMessageSanitization:
|
|||
{"role": "user", "content": ""},
|
||||
]
|
||||
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
|
||||
)
|
||||
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
|
||||
|
||||
# All three user messages get merged into one user turn for Anthropic.
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
text_blocks = [
|
||||
b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"
|
||||
]
|
||||
text_blocks = [b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"]
|
||||
assert len(text_blocks) == 3
|
||||
# No text block may be empty — that's the contract Anthropic enforces.
|
||||
for block in text_blocks:
|
||||
assert block["text"].strip() != ""
|
||||
assert text_blocks[2]["text"] == (
|
||||
"[System: Empty message content sanitised to satisfy protocol]"
|
||||
)
|
||||
assert text_blocks[2]["text"] == ("[System: Empty message content sanitised to satisfy protocol]")
|
||||
|
||||
def test_empty_text_block_in_list_content_sanitized(self):
|
||||
"""
|
||||
|
|
@ -392,14 +363,10 @@ class TestMessageSanitization:
|
|||
},
|
||||
]
|
||||
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
|
||||
)
|
||||
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
|
||||
|
||||
assert len(result) == 1
|
||||
text_blocks = [
|
||||
b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"
|
||||
]
|
||||
text_blocks = [b for b in result[0]["content"] if isinstance(b, dict) and b.get("type") == "text"]
|
||||
assert len(text_blocks) == 3
|
||||
assert text_blocks[0]["text"] == "real content"
|
||||
for block in text_blocks[1:]:
|
||||
|
|
@ -418,9 +385,7 @@ class TestMessageSanitization:
|
|||
{"role": "user", "content": "How are you?"},
|
||||
]
|
||||
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
|
||||
)
|
||||
result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic")
|
||||
|
||||
# Two user turns + one assistant turn (alternation preserved).
|
||||
assert len(result) == 3
|
||||
|
|
|
|||
0
tests/test_litellm/llms/base_llm/auth/__init__.py
Normal file
0
tests/test_litellm/llms/base_llm/auth/__init__.py
Normal file
484
tests/test_litellm/llms/base_llm/auth/test_client_credentials.py
Normal file
484
tests/test_litellm/llms/base_llm/auth/test_client_credentials.py
Normal file
|
|
@ -0,0 +1,484 @@
|
|||
import base64
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, unquote
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.base_llm.auth.client_credentials import (
|
||||
_HttpxSyncKeycloakPoster,
|
||||
_default_secret_reader,
|
||||
_new_keycloak_handler,
|
||||
fetch_keycloak_assertion,
|
||||
keycloak_assertion_source,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.identity_source import KeycloakSource, identity_source_ref
|
||||
from litellm.llms.base_llm.auth.token_exchange import MAX_RESPONSE_BYTES
|
||||
|
||||
TOKEN_URL: Final = "https://keycloak.example/realms/litellm/protocol/openid-connect/token"
|
||||
CLIENT_ID: Final = "litellm"
|
||||
CLIENT_SECRET_REF: Final = "oidc/env/KEYCLOAK_CLIENT_SECRET"
|
||||
CLIENT_SECRET: Final = "s3cr3t-client-value"
|
||||
|
||||
|
||||
class RecordedRequest:
|
||||
def __init__(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> None:
|
||||
self.url = url
|
||||
self.content = content
|
||||
self.headers = dict(headers)
|
||||
self.timeout = timeout
|
||||
|
||||
def form_body(self) -> dict[str, str]:
|
||||
return dict(parse_qsl(self.content.decode()))
|
||||
|
||||
|
||||
class ScriptedPoster:
|
||||
"""Returns one scripted response per call; records every request it receives."""
|
||||
|
||||
def __init__(self, responses: list[httpx.Response]) -> None:
|
||||
self.requests: list[RecordedRequest] = []
|
||||
self._responses = list(responses)
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
|
||||
self.requests.append(RecordedRequest(url, content, headers, timeout))
|
||||
return self._responses.pop(0) if len(self._responses) > 1 else self._responses[0]
|
||||
|
||||
|
||||
class RaisingPoster:
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.calls = 0
|
||||
self._error = error
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
|
||||
self.calls += 1
|
||||
raise self._error
|
||||
|
||||
|
||||
def make_config(
|
||||
auth_method: str = "client_secret_basic",
|
||||
scope: str | None = None,
|
||||
token_url: str = TOKEN_URL,
|
||||
client_secret_ref: str = CLIENT_SECRET_REF,
|
||||
client_id: str = CLIENT_ID,
|
||||
) -> KeycloakSource:
|
||||
return KeycloakSource(
|
||||
token_url=token_url,
|
||||
client_id=client_id,
|
||||
client_secret_ref=client_secret_ref,
|
||||
auth_method=auth_method, # pyright: ignore[reportArgumentType] # test-only string widened for parametrization
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
|
||||
def secret_reader_returning(secret: str | None):
|
||||
def reader(ref: str) -> str | None:
|
||||
assert ref == CLIENT_SECRET_REF
|
||||
return secret
|
||||
|
||||
return reader
|
||||
|
||||
|
||||
DEFAULT_SECRET_READER: Final = secret_reader_returning(CLIENT_SECRET)
|
||||
|
||||
|
||||
def token_response(access_token: str = "keycloak-minted-token") -> httpx.Response:
|
||||
return httpx.Response(200, json={"access_token": access_token, "token_type": "Bearer", "expires_in": 300})
|
||||
|
||||
|
||||
class TestClientSecretBasic:
|
||||
def test_sends_basic_auth_header_and_no_secret_in_body(self):
|
||||
poster = ScriptedPoster([token_response("minted-1")])
|
||||
|
||||
token = fetch_keycloak_assertion(
|
||||
make_config(auth_method="client_secret_basic"), poster=poster, secret_reader=DEFAULT_SECRET_READER
|
||||
)
|
||||
|
||||
assert token == "minted-1"
|
||||
request = poster.requests[0]
|
||||
assert request.url == TOKEN_URL
|
||||
expected_auth = "Basic " + base64.b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode("ascii")
|
||||
assert request.headers["authorization"] == expected_auth
|
||||
assert request.headers["content-type"] == "application/x-www-form-urlencoded"
|
||||
body = request.form_body()
|
||||
assert body["grant_type"] == "client_credentials"
|
||||
assert "client_secret" not in body
|
||||
assert "client_id" not in body
|
||||
|
||||
def test_reserved_characters_are_form_encoded_before_basic(self):
|
||||
"""RFC 6749 2.3.1 requires the client id and secret be application/x-www-form-urlencoded
|
||||
(Appendix B) before being base64'd into the Basic header; a raw join lets a reserved
|
||||
character in either value corrupt the ':'-joined pair Keycloak decodes back out."""
|
||||
client_id = "id:with+reserved% chars"
|
||||
client_secret = "secret:with+reserved% chars"
|
||||
poster = ScriptedPoster([token_response("minted-reserved")])
|
||||
|
||||
fetch_keycloak_assertion(
|
||||
make_config(auth_method="client_secret_basic", client_id=client_id),
|
||||
poster=poster,
|
||||
secret_reader=secret_reader_returning(client_secret),
|
||||
)
|
||||
|
||||
header = poster.requests[0].headers["authorization"]
|
||||
assert header.startswith("Basic ")
|
||||
decoded = base64.b64decode(header.removeprefix("Basic ")).decode("ascii")
|
||||
encoded_id, _, encoded_secret = decoded.partition(":")
|
||||
assert unquote(encoded_id) == client_id
|
||||
assert unquote(encoded_secret) == client_secret
|
||||
|
||||
def test_scope_included_only_when_set(self):
|
||||
poster = ScriptedPoster([token_response()])
|
||||
fetch_keycloak_assertion(
|
||||
make_config(scope="openid profile"), poster=poster, secret_reader=DEFAULT_SECRET_READER
|
||||
)
|
||||
|
||||
assert poster.requests[0].form_body()["scope"] == "openid profile"
|
||||
|
||||
poster_no_scope = ScriptedPoster([token_response()])
|
||||
fetch_keycloak_assertion(make_config(scope=None), poster=poster_no_scope, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert "scope" not in poster_no_scope.requests[0].form_body()
|
||||
|
||||
|
||||
class TestClientSecretPost:
|
||||
def test_sends_client_id_and_secret_in_body_with_no_basic_header(self):
|
||||
poster = ScriptedPoster([token_response("minted-2")])
|
||||
|
||||
token = fetch_keycloak_assertion(
|
||||
make_config(auth_method="client_secret_post"), poster=poster, secret_reader=DEFAULT_SECRET_READER
|
||||
)
|
||||
|
||||
assert token == "minted-2"
|
||||
request = poster.requests[0]
|
||||
assert "authorization" not in request.headers
|
||||
body = request.form_body()
|
||||
assert body["grant_type"] == "client_credentials"
|
||||
assert body["client_id"] == CLIENT_ID
|
||||
assert body["client_secret"] == CLIENT_SECRET
|
||||
|
||||
|
||||
class TestOnePostPerExchange:
|
||||
def test_exactly_one_post_per_call_no_cache(self):
|
||||
poster = ScriptedPoster([token_response("first"), token_response("second")])
|
||||
|
||||
first = fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
second = fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert first == "first"
|
||||
assert second == "second"
|
||||
assert len(poster.requests) == 2
|
||||
|
||||
|
||||
class TestInvalidClient:
|
||||
def test_400_invalid_client_surfaces_redacted_detail(self):
|
||||
poster = ScriptedPoster(
|
||||
[httpx.Response(400, json={"error": "invalid_client", "error_description": "unauthorized client"})]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="invalid_client") as exc_info:
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert "unauthorized client" in str(exc_info.value)
|
||||
assert "400" in str(exc_info.value)
|
||||
assert CLIENT_SECRET not in str(exc_info.value)
|
||||
|
||||
def test_echoed_client_secret_is_never_reflected_into_the_error(self):
|
||||
"""A misbehaving Keycloak that echoes the submitted client_secret back in its error body
|
||||
must never leak it into the exception the caller sees."""
|
||||
long_secret: Final = "reflectable-secret-0123456789"
|
||||
poster = ScriptedPoster(
|
||||
[httpx.Response(400, json={"error": "invalid_client", "error_description": f"got {long_secret} in body"})]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="keycloak") as exc_info:
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(long_secret))
|
||||
|
||||
assert long_secret not in str(exc_info.value)
|
||||
assert "redacted" in str(exc_info.value)
|
||||
|
||||
def test_echoed_short_client_secret_is_never_reflected_into_the_error(self):
|
||||
"""Real Keycloak client secrets are often shorter than a JWT: the reflection probe must
|
||||
not silently stop protecting a secret just because it is under the probe's usual length."""
|
||||
short_secret: Final = "hand-set-14ch"
|
||||
poster = ScriptedPoster(
|
||||
[httpx.Response(400, json={"error": "invalid_client", "error_description": f"got {short_secret} in body"})]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="keycloak") as exc_info:
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(short_secret))
|
||||
|
||||
assert short_secret not in str(exc_info.value)
|
||||
assert "redacted" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestUnreachable:
|
||||
def test_transport_failure_raises_diagnosable_value_error(self):
|
||||
poster = RaisingPoster(httpx.ConnectError("connection refused"))
|
||||
|
||||
with pytest.raises(ValueError, match="ConnectError") as exc_info:
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert poster.calls == 1
|
||||
assert CLIENT_SECRET not in str(exc_info.value)
|
||||
|
||||
|
||||
class TestNon2xx:
|
||||
def test_500_raises_value_error_with_status_code(self):
|
||||
poster = ScriptedPoster([httpx.Response(500, json={"error": "server_error"})])
|
||||
|
||||
with pytest.raises(ValueError, match="500"):
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
|
||||
class TestResponseValidation:
|
||||
def test_missing_access_token_is_a_value_error(self):
|
||||
poster = ScriptedPoster([httpx.Response(200, json={"token_type": "Bearer"})])
|
||||
|
||||
with pytest.raises(ValueError, match="schema validation"):
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
def test_empty_access_token_is_a_value_error(self):
|
||||
poster = ScriptedPoster([httpx.Response(200, json={"access_token": " "})])
|
||||
|
||||
with pytest.raises(ValueError, match="empty access_token"):
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
|
||||
class TestInsecureTokenUrl:
|
||||
def test_http_url_is_rejected_before_any_post(self):
|
||||
poster = ScriptedPoster([token_response()])
|
||||
|
||||
with pytest.raises(ValueError, match="https"):
|
||||
fetch_keycloak_assertion(
|
||||
make_config(token_url="http://keycloak.example/token"),
|
||||
poster=poster,
|
||||
secret_reader=DEFAULT_SECRET_READER,
|
||||
)
|
||||
|
||||
assert poster.requests == []
|
||||
|
||||
|
||||
class TestMissingClientSecret:
|
||||
def test_unresolvable_secret_ref_raises_value_error_naming_the_ref_not_a_secret(self):
|
||||
poster = ScriptedPoster([token_response()])
|
||||
|
||||
with pytest.raises(ValueError, match=CLIENT_SECRET_REF):
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=secret_reader_returning(None))
|
||||
|
||||
assert poster.requests == []
|
||||
|
||||
|
||||
class TestKeycloakAssertionSource:
|
||||
def test_returns_a_callable_that_fetches_fresh_each_call(self):
|
||||
poster = ScriptedPoster([token_response("first"), token_response("second")])
|
||||
source = keycloak_assertion_source(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert source() == "first"
|
||||
assert source() == "second"
|
||||
assert len(poster.requests) == 2
|
||||
|
||||
def test_propagates_the_underlying_fetch_failure(self):
|
||||
poster = ScriptedPoster([httpx.Response(400, json={"error": "invalid_client"})])
|
||||
source = keycloak_assertion_source(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
with pytest.raises(ValueError, match="invalid_client"):
|
||||
source()
|
||||
|
||||
|
||||
class TestClientSecretNeverLeaks:
|
||||
"""Regression coverage for the load-bearing property: a Keycloak client_secret must never
|
||||
surface in the assertion_ref, in any error message, or in a log record, however it fails."""
|
||||
|
||||
def test_never_in_the_assertion_ref(self):
|
||||
config = make_config(client_secret_ref=CLIENT_SECRET_REF)
|
||||
|
||||
ref = identity_source_ref(config)
|
||||
|
||||
assert CLIENT_SECRET not in ref
|
||||
assert CLIENT_SECRET_REF not in ref
|
||||
|
||||
def test_never_in_any_raised_error_message_across_every_failure_mode(self):
|
||||
config = make_config()
|
||||
failures = [
|
||||
lambda: fetch_keycloak_assertion(
|
||||
config,
|
||||
poster=ScriptedPoster([httpx.Response(400, json={"error": "invalid_client"})]),
|
||||
secret_reader=DEFAULT_SECRET_READER,
|
||||
),
|
||||
lambda: fetch_keycloak_assertion(
|
||||
config, poster=RaisingPoster(httpx.ConnectError("boom")), secret_reader=DEFAULT_SECRET_READER
|
||||
),
|
||||
lambda: fetch_keycloak_assertion(
|
||||
config,
|
||||
poster=ScriptedPoster([httpx.Response(500, json={"error": "server_error"})]),
|
||||
secret_reader=DEFAULT_SECRET_READER,
|
||||
),
|
||||
lambda: fetch_keycloak_assertion(
|
||||
config, poster=ScriptedPoster([token_response()]), secret_reader=secret_reader_returning(None)
|
||||
),
|
||||
]
|
||||
for fail in failures:
|
||||
with pytest.raises(ValueError, match="keycloak") as exc_info:
|
||||
fail()
|
||||
assert CLIENT_SECRET not in str(exc_info.value)
|
||||
|
||||
def test_never_in_a_log_record(self, caplog: pytest.LogCaptureFixture):
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
poster = ScriptedPoster(
|
||||
[httpx.Response(400, json={"error": "invalid_client", "error_description": CLIENT_SECRET})]
|
||||
)
|
||||
with pytest.raises(ValueError, match="keycloak"):
|
||||
fetch_keycloak_assertion(make_config(), poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
fetch_keycloak_assertion(
|
||||
make_config(), poster=ScriptedPoster([token_response()]), secret_reader=DEFAULT_SECRET_READER
|
||||
)
|
||||
|
||||
assert CLIENT_SECRET not in caplog.text
|
||||
|
||||
|
||||
class StubHandler:
|
||||
"""Stands in for the HTTPHandler the default poster builds, so the poster's own contract
|
||||
(redirects off, error responses returned rather than raised, no-response guarded) is testable
|
||||
without a socket."""
|
||||
|
||||
def __init__(self, result: httpx.Response | Exception | None) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
self._result = result
|
||||
|
||||
def post(self, url: str, *, content: bytes, headers: dict[str, str], timeout: float) -> httpx.Response | None:
|
||||
self.calls.append({"url": url, "content": content, "headers": headers, "timeout": timeout})
|
||||
if isinstance(self._result, Exception):
|
||||
raise self._result
|
||||
return self._result
|
||||
|
||||
|
||||
class TestDefaultKeycloakPoster:
|
||||
def test_builds_its_handler_once_with_redirects_disabled(self):
|
||||
built: list[StubHandler] = []
|
||||
|
||||
def factory() -> StubHandler:
|
||||
handler = StubHandler(httpx.Response(200, json={"access_token": "kc-token"}))
|
||||
built.append(handler)
|
||||
return handler
|
||||
|
||||
poster: Final = _HttpxSyncKeycloakPoster(handler_factory=factory) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
|
||||
for _ in range(3):
|
||||
poster.post(TOKEN_URL, content=b"grant_type=client_credentials", headers={}, timeout=1.0)
|
||||
|
||||
assert len(built) == 1, "the handler is built once and reused"
|
||||
assert len(built[0].calls) == 3
|
||||
|
||||
def test_the_real_handler_refuses_to_follow_redirects(self):
|
||||
handler: Final = _new_keycloak_handler()
|
||||
assert handler.client.follow_redirects is False, (
|
||||
"a redirected token POST would replay the client secret to whatever host the redirect names"
|
||||
)
|
||||
|
||||
def test_an_http_status_error_becomes_its_response_rather_than_an_exception(self):
|
||||
response: Final = httpx.Response(
|
||||
401, json={"error": "invalid_client"}, request=httpx.Request("POST", TOKEN_URL)
|
||||
)
|
||||
poster: Final = _HttpxSyncKeycloakPoster(
|
||||
handler_factory=lambda: StubHandler(
|
||||
httpx.HTTPStatusError("boom", request=response.request, response=response)
|
||||
) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
|
||||
)
|
||||
|
||||
assert poster.post(TOKEN_URL, content=b"", headers={}, timeout=1.0).status_code == 401
|
||||
|
||||
def test_a_missing_response_is_a_transport_error_not_a_none_deref(self):
|
||||
poster: Final = _HttpxSyncKeycloakPoster(handler_factory=lambda: StubHandler(None)) # pyright: ignore[reportArgumentType] # StubHandler stands in for the legacy-untyped HTTPHandler
|
||||
|
||||
with pytest.raises(httpx.TransportError):
|
||||
poster.post(TOKEN_URL, content=b"", headers={}, timeout=1.0)
|
||||
|
||||
|
||||
class TestDefaultSecretReader:
|
||||
def test_reads_through_litellm_secret_resolution(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("KEYCLOAK_CLIENT_SECRET_FOR_DEFAULT_READER", CLIENT_SECRET)
|
||||
|
||||
assert _default_secret_reader("os.environ/KEYCLOAK_CLIENT_SECRET_FOR_DEFAULT_READER") == CLIENT_SECRET
|
||||
|
||||
def test_an_unset_reference_reads_as_none_so_the_caller_raises(self):
|
||||
assert _default_secret_reader("os.environ/DEFINITELY_NOT_SET_KEYCLOAK_SECRET_REF") is None
|
||||
|
||||
|
||||
class TestOversizedSuccessBody:
|
||||
def test_a_success_body_over_the_cap_is_refused_before_it_is_parsed(self):
|
||||
oversized: Final = httpx.Response(200, content=b'{"access_token": "' + b"x" * MAX_RESPONSE_BYTES + b'"}')
|
||||
|
||||
with pytest.raises(ValueError, match="exceeded the size cap"):
|
||||
fetch_keycloak_assertion(
|
||||
make_config(), poster=ScriptedPoster([oversized]), secret_reader=DEFAULT_SECRET_READER
|
||||
)
|
||||
|
||||
|
||||
class TestUnresolvedSecretRefIsNotEchoed:
|
||||
"""An operator who pastes the secret itself into the *_ref field turns that field INTO the
|
||||
secret, and this error reaches model callers, so it must never echo the value."""
|
||||
|
||||
def test_keycloak_ref_value_is_not_in_the_error(self):
|
||||
from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source
|
||||
from litellm.llms.base_llm.auth.identity_source import KeycloakSource
|
||||
|
||||
pasted_secret = "sUp3r-s3cret-value-not-a-pointer"
|
||||
config = KeycloakSource(
|
||||
token_url="https://keycloak.example.com/realms/p/protocol/openid-connect/token",
|
||||
client_id="litellm",
|
||||
client_secret_ref=pasted_secret,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="could not be read") as excinfo:
|
||||
keycloak_assertion_source(config, secret_reader=lambda _ref: None)()
|
||||
|
||||
assert pasted_secret not in str(excinfo.value)
|
||||
assert "withheld" in str(excinfo.value)
|
||||
|
||||
def test_internal_issuer_ref_value_is_not_in_the_error(self):
|
||||
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource
|
||||
from litellm.llms.base_llm.auth.internal_issuer import internal_issuer_assertion_source
|
||||
|
||||
pasted_pem = "-----BEGIN PRIVATE KEY-----MIGHAgEA-----END PRIVATE KEY-----"
|
||||
config = InternalIssuerSource(
|
||||
issuer_url="https://proxy.example.com",
|
||||
subject="litellm-proxy",
|
||||
signing_key_ref=pasted_pem,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="could not be read") as excinfo:
|
||||
internal_issuer_assertion_source(config, key_reader=lambda _ref: None)()
|
||||
|
||||
assert pasted_pem not in str(excinfo.value)
|
||||
assert "withheld" in str(excinfo.value)
|
||||
|
||||
|
||||
class TestTokenUrlIsNotEchoedWholesale:
|
||||
"""A token endpoint is configuration and naming it makes the error actionable, but nothing
|
||||
stops an operator putting a credential in the URL, and these errors reach model callers."""
|
||||
|
||||
def test_query_string_is_dropped_from_a_status_error(self):
|
||||
from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message
|
||||
|
||||
rendered = endpoint_url_for_error_message("https://idp.example/token?client_secret=supersecret")
|
||||
|
||||
assert "supersecret" not in rendered
|
||||
assert rendered == "https://idp.example/token"
|
||||
|
||||
def test_userinfo_is_dropped_too(self):
|
||||
from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message
|
||||
|
||||
rendered = endpoint_url_for_error_message("https://user:pw@idp.example:8443/token")
|
||||
|
||||
assert "pw" not in rendered
|
||||
assert rendered == "https://idp.example:8443/token"
|
||||
|
||||
def test_transport_failure_message_carries_no_query_secret(self):
|
||||
poster = RaisingPoster(httpx.ConnectTimeout("timed out"))
|
||||
config = make_config(token_url="https://idp.example/token?client_secret=supersecret")
|
||||
|
||||
with pytest.raises(ValueError, match="could not reach the keycloak token endpoint") as excinfo:
|
||||
fetch_keycloak_assertion(config, poster=poster, secret_reader=DEFAULT_SECRET_READER)
|
||||
|
||||
assert "supersecret" not in str(excinfo.value)
|
||||
assert "idp.example/token" in str(excinfo.value)
|
||||
239
tests/test_litellm/llms/base_llm/auth/test_identity_source.py
Normal file
239
tests/test_litellm/llms/base_llm/auth/test_identity_source.py
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.llms.base_llm.auth.identity_source import (
|
||||
AnthropicIdentitySourceKind,
|
||||
InternalIssuerSource,
|
||||
KeycloakSource,
|
||||
identity_source_config_adapter,
|
||||
identity_source_ref,
|
||||
)
|
||||
|
||||
SIGNING_KEY_REF: Final = "oidc/env/ISSUER_SIGNING_KEY_PEM"
|
||||
OTHER_SIGNING_KEY_REF: Final = "oidc/env/OTHER_SIGNING_KEY_PEM"
|
||||
CLIENT_SECRET_REF: Final = "oidc/env/KEYCLOAK_CLIENT_SECRET"
|
||||
ISSUER_URL: Final = "https://issuer.internal.example"
|
||||
SUBJECT: Final = "workload-a"
|
||||
TOKEN_URL: Final = "https://keycloak.example/realms/litellm/protocol/openid-connect/token"
|
||||
CLIENT_ID: Final = "litellm"
|
||||
|
||||
|
||||
def make_issuer(
|
||||
issuer_url: str = ISSUER_URL,
|
||||
subject: str = SUBJECT,
|
||||
signing_key_ref: str = SIGNING_KEY_REF,
|
||||
ttl_seconds: int = 300,
|
||||
) -> InternalIssuerSource:
|
||||
return InternalIssuerSource(
|
||||
issuer_url=issuer_url, subject=subject, signing_key_ref=signing_key_ref, ttl_seconds=ttl_seconds
|
||||
)
|
||||
|
||||
|
||||
def make_keycloak(
|
||||
token_url: str = TOKEN_URL,
|
||||
client_id: str = CLIENT_ID,
|
||||
client_secret_ref: str = CLIENT_SECRET_REF,
|
||||
auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic",
|
||||
scope: str | None = None,
|
||||
) -> KeycloakSource:
|
||||
return KeycloakSource(
|
||||
token_url=token_url,
|
||||
client_id=client_id,
|
||||
client_secret_ref=client_secret_ref,
|
||||
auth_method=auth_method,
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
|
||||
class TestIdentitySourceRefHashing:
|
||||
def test_identical_config_hashes_idempotently(self):
|
||||
assert identity_source_ref(make_issuer()) == identity_source_ref(make_issuer())
|
||||
|
||||
def test_ref_is_prefixed_by_kind(self):
|
||||
assert identity_source_ref(make_issuer()).startswith("oidc/internal_issuer/")
|
||||
assert identity_source_ref(make_keycloak()).startswith("oidc/keycloak/")
|
||||
|
||||
def test_pointer_name_change_changes_ref(self):
|
||||
"""Two configs differing only in which secret a pointer names must never collide, since a
|
||||
stale ref would let the token exchange's outer cache key alias two different credentials."""
|
||||
first: Final = identity_source_ref(make_issuer(signing_key_ref=SIGNING_KEY_REF))
|
||||
second: Final = identity_source_ref(make_issuer(signing_key_ref=OTHER_SIGNING_KEY_REF))
|
||||
|
||||
assert first != second
|
||||
|
||||
def test_non_pointer_field_change_changes_ref(self):
|
||||
first: Final = identity_source_ref(make_keycloak(scope="openid"))
|
||||
second: Final = identity_source_ref(make_keycloak(scope="openid profile"))
|
||||
|
||||
assert first != second
|
||||
|
||||
def test_ref_never_contains_the_pointer_field_values(self):
|
||||
"""The ref is a fixed-width hash, not a serialization of the config, so no field value -
|
||||
pointer name or otherwise - can leak into the secret-free string echoed into errors."""
|
||||
ref: Final = identity_source_ref(make_issuer())
|
||||
|
||||
assert SIGNING_KEY_REF not in ref
|
||||
assert "issuer.internal.example" not in ref
|
||||
|
||||
def test_different_kinds_with_disjoint_fields_never_collide(self):
|
||||
assert identity_source_ref(make_issuer()) != identity_source_ref(make_keycloak())
|
||||
|
||||
|
||||
class TestInternalIssuerSourceValidation:
|
||||
def test_defaults(self):
|
||||
source: Final = make_issuer()
|
||||
|
||||
assert source.kind == AnthropicIdentitySourceKind.internal_issuer
|
||||
assert source.ttl_seconds == 300
|
||||
assert source.audience is None
|
||||
|
||||
def test_ttl_seconds_over_one_hour_is_rejected(self):
|
||||
with pytest.raises(ValidationError):
|
||||
make_issuer(ttl_seconds=3601)
|
||||
|
||||
def test_ttl_seconds_at_one_hour_is_accepted(self):
|
||||
assert make_issuer(ttl_seconds=3600).ttl_seconds == 3600
|
||||
|
||||
def test_non_positive_ttl_seconds_is_rejected(self):
|
||||
with pytest.raises(ValidationError):
|
||||
make_issuer(ttl_seconds=0)
|
||||
|
||||
def test_missing_signing_key_ref_is_rejected(self):
|
||||
missing_field: Final = MappingProxyType({"issuer_url": ISSUER_URL, "subject": SUBJECT})
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
InternalIssuerSource.model_validate(missing_field)
|
||||
|
||||
def test_keycloak_only_field_is_rejected_as_extra(self):
|
||||
mixed_variant: Final = MappingProxyType(
|
||||
{
|
||||
"issuer_url": ISSUER_URL,
|
||||
"subject": SUBJECT,
|
||||
"signing_key_ref": SIGNING_KEY_REF,
|
||||
"client_secret_ref": CLIENT_SECRET_REF,
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
InternalIssuerSource.model_validate(mixed_variant)
|
||||
|
||||
def test_is_frozen(self):
|
||||
source: Final = make_issuer()
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
source.subject = "workload-b"
|
||||
|
||||
def test_secret_pasted_into_wrong_typed_field_is_not_echoed_in_the_error(self):
|
||||
"""hide_input_in_errors keeps a value the operator pasted into a mistyped field out of the
|
||||
validation error, so a client_secret headed for the wrong field isn't logged in the raise."""
|
||||
leaked_secret: Final = "shh-do-not-log-me"
|
||||
wrong_type: Final = MappingProxyType(
|
||||
{
|
||||
"issuer_url": ISSUER_URL,
|
||||
"subject": SUBJECT,
|
||||
"signing_key_ref": SIGNING_KEY_REF,
|
||||
"ttl_seconds": leaked_secret,
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
InternalIssuerSource.model_validate(wrong_type)
|
||||
|
||||
assert leaked_secret not in str(exc_info.value)
|
||||
|
||||
|
||||
class TestKeycloakSourceValidation:
|
||||
def test_defaults(self):
|
||||
source: Final = make_keycloak()
|
||||
|
||||
assert source.kind == AnthropicIdentitySourceKind.keycloak
|
||||
assert source.auth_method == "client_secret_basic"
|
||||
assert source.scope is None
|
||||
|
||||
def test_client_secret_post_is_accepted(self):
|
||||
assert make_keycloak(auth_method="client_secret_post").auth_method == "client_secret_post"
|
||||
|
||||
def test_private_key_jwt_is_not_a_supported_auth_method_yet(self):
|
||||
unshipped_auth_method: Final = MappingProxyType(
|
||||
{
|
||||
"token_url": TOKEN_URL,
|
||||
"client_id": CLIENT_ID,
|
||||
"client_secret_ref": CLIENT_SECRET_REF,
|
||||
"auth_method": "private_key_jwt",
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
KeycloakSource.model_validate(unshipped_auth_method)
|
||||
|
||||
def test_audience_field_was_dropped(self):
|
||||
dropped_field: Final = MappingProxyType(
|
||||
{
|
||||
"token_url": TOKEN_URL,
|
||||
"client_id": CLIENT_ID,
|
||||
"client_secret_ref": CLIENT_SECRET_REF,
|
||||
"audience": "https://anthropic.example",
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
KeycloakSource.model_validate(dropped_field)
|
||||
|
||||
def test_missing_client_secret_ref_is_rejected(self):
|
||||
missing_field: Final = MappingProxyType({"token_url": TOKEN_URL, "client_id": CLIENT_ID})
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
KeycloakSource.model_validate(missing_field)
|
||||
|
||||
|
||||
class TestDiscriminatedUnionParsing:
|
||||
def test_parses_internal_issuer_variant(self):
|
||||
parsed: Final = identity_source_config_adapter.validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
"kind": "internal_issuer",
|
||||
"issuer_url": ISSUER_URL,
|
||||
"subject": SUBJECT,
|
||||
"signing_key_ref": SIGNING_KEY_REF,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(parsed, InternalIssuerSource)
|
||||
|
||||
def test_parses_keycloak_variant(self):
|
||||
parsed: Final = identity_source_config_adapter.validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
"kind": "keycloak",
|
||||
"token_url": TOKEN_URL,
|
||||
"client_id": CLIENT_ID,
|
||||
"client_secret_ref": CLIENT_SECRET_REF,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(parsed, KeycloakSource)
|
||||
|
||||
def test_unknown_kind_is_a_hard_error(self):
|
||||
with pytest.raises(ValidationError):
|
||||
identity_source_config_adapter.validate_python(MappingProxyType({"kind": "token_file"}))
|
||||
|
||||
def test_mixed_variant_fields_are_a_hard_error(self):
|
||||
"""A keycloak field on an internal_issuer-tagged payload must fail closed rather than be
|
||||
silently dropped or silently accepted as if it selected the other variant."""
|
||||
with pytest.raises(ValidationError):
|
||||
identity_source_config_adapter.validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
"kind": "internal_issuer",
|
||||
"issuer_url": ISSUER_URL,
|
||||
"subject": SUBJECT,
|
||||
"signing_key_ref": SIGNING_KEY_REF,
|
||||
"client_secret_ref": CLIENT_SECRET_REF,
|
||||
}
|
||||
)
|
||||
)
|
||||
188
tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py
Normal file
188
tests/test_litellm/llms/base_llm/auth/test_internal_issuer.py
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
|
||||
from litellm.llms.base_llm.auth.identity_source import InternalIssuerSource
|
||||
from litellm.llms.base_llm.auth.internal_issuer import (
|
||||
internal_issuer_assertion_source,
|
||||
internal_issuer_jwks_document,
|
||||
mint_internal_issuer_assertion,
|
||||
)
|
||||
from litellm.llms.base_llm.auth.jwt_signing import build_jwks, rfc7638_thumbprint
|
||||
|
||||
SIGNING_KEY_REF: Final = "oidc/env/ISSUER_SIGNING_KEY_PEM"
|
||||
ISSUER_URL: Final = "https://issuer.internal.example"
|
||||
SUBJECT: Final = "workload-a"
|
||||
|
||||
|
||||
_PRIVATE_VALUE: Final = 90123456789012345678901234567890123456789012345678901234567890
|
||||
|
||||
|
||||
def signing_key() -> ec.EllipticCurvePrivateKey:
|
||||
return ec.derive_private_key(_PRIVATE_VALUE, ec.SECP256R1())
|
||||
|
||||
|
||||
def pem_of(key: ec.EllipticCurvePrivateKey) -> str:
|
||||
return key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
).decode()
|
||||
|
||||
|
||||
def make_config(
|
||||
issuer_url: str = ISSUER_URL,
|
||||
subject: str = SUBJECT,
|
||||
audience: str | None = None,
|
||||
ttl_seconds: int = 300,
|
||||
signing_key_ref: str = SIGNING_KEY_REF,
|
||||
) -> InternalIssuerSource:
|
||||
return InternalIssuerSource(
|
||||
issuer_url=issuer_url,
|
||||
subject=subject,
|
||||
audience=audience,
|
||||
ttl_seconds=ttl_seconds,
|
||||
signing_key_ref=signing_key_ref,
|
||||
)
|
||||
|
||||
|
||||
def key_reader_returning(pem: str | None):
|
||||
def reader(ref: str) -> str | None:
|
||||
assert ref == SIGNING_KEY_REF
|
||||
return pem
|
||||
|
||||
return reader
|
||||
|
||||
|
||||
class FakeClock:
|
||||
def __init__(self, value: float) -> None:
|
||||
self._value: Final = value
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self._value
|
||||
|
||||
|
||||
def decode_ignoring_wall_clock(token: str, public_key: ec.EllipticCurvePublicKey) -> dict:
|
||||
"""Tests mint with a fixed past ``FakeClock`` and no expected audience, so PyJWT's
|
||||
real-wall-clock ``exp``/``aud`` checks (irrelevant to what these tests verify) are disabled."""
|
||||
return jwt.decode(token, public_key, algorithms=["ES256"], options={"verify_exp": False, "verify_aud": False})
|
||||
|
||||
|
||||
class TestMintInternalIssuerAssertion:
|
||||
def test_required_claims_and_asymmetric_alg(self):
|
||||
key: Final = signing_key()
|
||||
config: Final = make_config(ttl_seconds=300)
|
||||
|
||||
token: Final = mint_internal_issuer_assertion(
|
||||
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
|
||||
)
|
||||
header: Final = jwt.get_unverified_header(token)
|
||||
claims: Final = decode_ignoring_wall_clock(token, key.public_key())
|
||||
|
||||
assert header["alg"] == "ES256"
|
||||
assert claims["sub"] == SUBJECT
|
||||
assert claims["iss"] == ISSUER_URL
|
||||
assert claims["iat"] == 1_700_000_000
|
||||
assert claims["exp"] == 1_700_000_300
|
||||
|
||||
def test_kid_matches_the_published_jwks(self):
|
||||
key: Final = signing_key()
|
||||
config: Final = make_config()
|
||||
|
||||
token: Final = mint_internal_issuer_assertion(
|
||||
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
|
||||
)
|
||||
|
||||
header_kid: Final = jwt.get_unverified_header(token)["kid"]
|
||||
published_kid: Final = build_jwks(key.public_key())["keys"][0]["kid"]
|
||||
assert header_kid == published_kid == rfc7638_thumbprint(key.public_key())
|
||||
|
||||
def test_ttl_bounds_exp_minus_iat(self):
|
||||
key: Final = signing_key()
|
||||
config: Final = make_config(ttl_seconds=120)
|
||||
|
||||
token: Final = mint_internal_issuer_assertion(
|
||||
config, key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
|
||||
)
|
||||
claims: Final = decode_ignoring_wall_clock(token, key.public_key())
|
||||
|
||||
assert claims["exp"] - claims["iat"] == 120
|
||||
|
||||
def test_audience_included_only_when_set(self):
|
||||
key: Final = signing_key()
|
||||
without_audience: Final = mint_internal_issuer_assertion(
|
||||
make_config(audience=None), key_reader=key_reader_returning(pem_of(key)), clock=FakeClock(1_700_000_000.0)
|
||||
)
|
||||
with_audience: Final = mint_internal_issuer_assertion(
|
||||
make_config(audience="urn:anthropic:federation"),
|
||||
key_reader=key_reader_returning(pem_of(key)),
|
||||
clock=FakeClock(1_700_000_000.0),
|
||||
)
|
||||
|
||||
claims_without: Final = decode_ignoring_wall_clock(without_audience, key.public_key())
|
||||
claims_with: Final = decode_ignoring_wall_clock(with_audience, key.public_key())
|
||||
assert "aud" not in claims_without
|
||||
assert claims_with["aud"] == "urn:anthropic:federation"
|
||||
|
||||
def test_jti_is_present_and_fresh_on_every_mint(self):
|
||||
key: Final = signing_key()
|
||||
config: Final = make_config()
|
||||
reader: Final = key_reader_returning(pem_of(key))
|
||||
|
||||
first: Final = decode_ignoring_wall_clock(
|
||||
mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0)),
|
||||
key.public_key(),
|
||||
)
|
||||
second: Final = decode_ignoring_wall_clock(
|
||||
mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0)),
|
||||
key.public_key(),
|
||||
)
|
||||
|
||||
assert first["jti"] and second["jti"]
|
||||
assert first["jti"] != second["jti"]
|
||||
|
||||
def test_missing_signing_key_raises_value_error_naming_the_ref_not_a_secret(self):
|
||||
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
|
||||
mint_internal_issuer_assertion(make_config(), key_reader=key_reader_returning(None))
|
||||
|
||||
def test_malformed_signing_key_raises_value_error(self):
|
||||
with pytest.raises(ValueError, match="not a valid unencrypted PEM"):
|
||||
mint_internal_issuer_assertion(make_config(), key_reader=key_reader_returning("not-a-pem"))
|
||||
|
||||
|
||||
class TestInternalIssuerAssertionSource:
|
||||
def test_returns_a_callable_that_mints_fresh_each_call(self):
|
||||
key: Final = signing_key()
|
||||
source: Final = internal_issuer_assertion_source(make_config(), key_reader=key_reader_returning(pem_of(key)))
|
||||
|
||||
first: Final = jwt.decode(source(), key.public_key(), algorithms=["ES256"])
|
||||
second: Final = jwt.decode(source(), key.public_key(), algorithms=["ES256"])
|
||||
|
||||
assert first["jti"] != second["jti"]
|
||||
|
||||
def test_propagates_the_underlying_mint_failure(self):
|
||||
source: Final = internal_issuer_assertion_source(make_config(), key_reader=key_reader_returning(None))
|
||||
|
||||
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
|
||||
source()
|
||||
|
||||
|
||||
class TestInternalIssuerJwksDocument:
|
||||
def test_matches_the_key_used_to_mint(self):
|
||||
key: Final = signing_key()
|
||||
config: Final = make_config()
|
||||
reader: Final = key_reader_returning(pem_of(key))
|
||||
|
||||
document: Final = json.loads(internal_issuer_jwks_document(config, key_reader=reader))
|
||||
token: Final = mint_internal_issuer_assertion(config, key_reader=reader, clock=FakeClock(1_700_000_000.0))
|
||||
|
||||
assert document["keys"][0]["kid"] == jwt.get_unverified_header(token)["kid"]
|
||||
assert decode_ignoring_wall_clock(token, key.public_key())
|
||||
|
||||
def test_missing_signing_key_raises_value_error(self):
|
||||
with pytest.raises(ValueError, match=SIGNING_KEY_REF):
|
||||
internal_issuer_jwks_document(make_config(), key_reader=key_reader_returning(None))
|
||||
213
tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py
Normal file
213
tests/test_litellm/llms/base_llm/auth/test_jwt_signing.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import ec, rsa
|
||||
|
||||
from litellm.llms.base_llm.auth.jwt_signing import (
|
||||
MISSING_SIGNING_DEPENDENCIES_MESSAGE,
|
||||
build_jwk,
|
||||
build_jwks,
|
||||
jwks_document_json,
|
||||
load_es256_private_key,
|
||||
rfc7638_thumbprint,
|
||||
sign_es256_jwt,
|
||||
)
|
||||
|
||||
_FIXED_PRIVATE_VALUE: Final = 55090612345678901234567890123456789012345678901234567890123456
|
||||
_OTHER_PRIVATE_VALUE: Final = 1
|
||||
|
||||
|
||||
def fixed_private_key(value: int = _FIXED_PRIVATE_VALUE) -> ec.EllipticCurvePrivateKey:
|
||||
return ec.derive_private_key(value, ec.SECP256R1())
|
||||
|
||||
|
||||
def pem_of(key: ec.EllipticCurvePrivateKey) -> str:
|
||||
return key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
).decode()
|
||||
|
||||
|
||||
def independent_thumbprint(public_key: ec.EllipticCurvePublicKey) -> str:
|
||||
"""Recomputes RFC 7638 by hand, deliberately not sharing a single line of code with
|
||||
``jwt_signing.rfc7638_thumbprint`` -- a mutation that broke the real implementation must not
|
||||
also break this reference, or the two would trivially agree by sharing the bug."""
|
||||
numbers: Final = public_key.public_numbers()
|
||||
x: Final = base64.urlsafe_b64encode(numbers.x.to_bytes(32, "big")).rstrip(b"=").decode()
|
||||
y: Final = base64.urlsafe_b64encode(numbers.y.to_bytes(32, "big")).rstrip(b"=").decode()
|
||||
canonical: Final = f'{{"crv":"P-256","kty":"EC","x":"{x}","y":"{y}"}}'
|
||||
return base64.urlsafe_b64encode(hashlib.sha256(canonical.encode()).digest()).rstrip(b"=").decode()
|
||||
|
||||
|
||||
class TestLoadEs256PrivateKey:
|
||||
def test_valid_ec_p256_pem_loads(self):
|
||||
key: Final = load_es256_private_key(pem_of(fixed_private_key()))
|
||||
|
||||
assert isinstance(key, ec.EllipticCurvePrivateKey)
|
||||
assert isinstance(key.curve, ec.SECP256R1)
|
||||
|
||||
def test_garbage_pem_is_rejected(self):
|
||||
with pytest.raises(ValueError, match="not a valid unencrypted PEM"):
|
||||
load_es256_private_key("not a pem")
|
||||
|
||||
def test_rsa_key_is_rejected(self):
|
||||
rsa_pem: Final = (
|
||||
rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
.decode()
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="P-256"):
|
||||
load_es256_private_key(rsa_pem)
|
||||
|
||||
def test_non_p256_curve_is_rejected(self):
|
||||
secp384_pem: Final = (
|
||||
ec.generate_private_key(ec.SECP384R1())
|
||||
.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
.decode()
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="P-256"):
|
||||
load_es256_private_key(secp384_pem)
|
||||
|
||||
def test_error_never_echoes_key_material(self):
|
||||
pem: Final = pem_of(fixed_private_key())
|
||||
|
||||
with pytest.raises(ValueError, match="P-256"):
|
||||
load_es256_private_key(pem_of(ec.generate_private_key(ec.SECP384R1())))
|
||||
with pytest.raises(ValueError, match="not a valid unencrypted PEM") as exc_info:
|
||||
load_es256_private_key("garbage-not-a-pem")
|
||||
|
||||
assert pem not in str(exc_info.value)
|
||||
|
||||
|
||||
class TestRfc7638Thumbprint:
|
||||
def test_matches_independent_recomputation(self):
|
||||
public_key: Final = fixed_private_key().public_key()
|
||||
|
||||
assert rfc7638_thumbprint(public_key) == independent_thumbprint(public_key)
|
||||
|
||||
def test_different_keys_have_different_thumbprints(self):
|
||||
first: Final = fixed_private_key(_FIXED_PRIVATE_VALUE).public_key()
|
||||
second: Final = fixed_private_key(_OTHER_PRIVATE_VALUE).public_key()
|
||||
|
||||
assert rfc7638_thumbprint(first) != rfc7638_thumbprint(second)
|
||||
|
||||
def test_thumbprint_is_deterministic(self):
|
||||
public_key: Final = fixed_private_key().public_key()
|
||||
|
||||
assert rfc7638_thumbprint(public_key) == rfc7638_thumbprint(public_key)
|
||||
|
||||
|
||||
class TestBuildJwks:
|
||||
def test_jwks_contains_one_key_matching_the_thumbprint(self):
|
||||
public_key: Final = fixed_private_key().public_key()
|
||||
|
||||
jwks: Final = build_jwks(public_key)
|
||||
|
||||
assert len(jwks["keys"]) == 1
|
||||
assert jwks["keys"][0]["kid"] == rfc7638_thumbprint(public_key)
|
||||
assert jwks["keys"][0]["kty"] == "EC"
|
||||
assert jwks["keys"][0]["crv"] == "P-256"
|
||||
assert jwks["keys"][0]["alg"] == "ES256"
|
||||
|
||||
def test_build_jwk_stamps_the_given_kid_verbatim(self):
|
||||
jwk: Final = build_jwk(fixed_private_key().public_key(), kid="caller-supplied-kid")
|
||||
|
||||
assert jwk["kid"] == "caller-supplied-kid"
|
||||
|
||||
def test_jwks_document_json_round_trips_through_build_jwks(self):
|
||||
key: Final = fixed_private_key()
|
||||
|
||||
document: Final = json.loads(jwks_document_json(pem_of(key)))
|
||||
jwks: Final = build_jwks(key.public_key())
|
||||
|
||||
assert document == {"keys": [dict(jwk) for jwk in jwks["keys"]]}
|
||||
|
||||
|
||||
class TestSignEs256Jwt:
|
||||
def test_minted_token_verifies_against_the_matching_public_key(self):
|
||||
key: Final = fixed_private_key()
|
||||
now: Final = int(time.time())
|
||||
claims: Final = {"sub": "workload-a", "iss": "https://issuer.example", "iat": now, "exp": now + 300}
|
||||
|
||||
token: Final = sign_es256_jwt(pem_of(key), claims)
|
||||
decoded: Final = jwt.decode(token, key.public_key(), algorithms=["ES256"])
|
||||
|
||||
assert decoded == claims
|
||||
|
||||
def test_header_alg_is_es256(self):
|
||||
token: Final = sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"})
|
||||
|
||||
assert jwt.get_unverified_header(token)["alg"] == "ES256"
|
||||
|
||||
def test_header_kid_matches_the_published_jwks(self):
|
||||
key: Final = fixed_private_key()
|
||||
|
||||
token: Final = sign_es256_jwt(pem_of(key), {"sub": "x"})
|
||||
|
||||
header_kid: Final = jwt.get_unverified_header(token)["kid"]
|
||||
published_kid: Final = build_jwks(key.public_key())["keys"][0]["kid"]
|
||||
assert header_kid == published_kid == rfc7638_thumbprint(key.public_key())
|
||||
|
||||
def test_wrong_key_fails_verification(self):
|
||||
signing_key: Final = fixed_private_key(_FIXED_PRIVATE_VALUE)
|
||||
other_key: Final = fixed_private_key(_OTHER_PRIVATE_VALUE)
|
||||
|
||||
token: Final = sign_es256_jwt(pem_of(signing_key), {"sub": "x"})
|
||||
|
||||
with pytest.raises(jwt.exceptions.InvalidSignatureError):
|
||||
jwt.decode(token, other_key.public_key(), algorithms=["ES256"])
|
||||
|
||||
|
||||
class TestBaseSdkImport:
|
||||
"""A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras),
|
||||
and ``litellm/__init__`` reaches this module through the Anthropic provider, so a
|
||||
module-level import of either would break ``import litellm`` for every base SDK user."""
|
||||
|
||||
def test_module_imports_with_pyjwt_and_cryptography_absent(self):
|
||||
script: Final = textwrap.dedent(
|
||||
"""
|
||||
import sys
|
||||
|
||||
class Blocker:
|
||||
def find_spec(self, name, path=None, target=None):
|
||||
if name.split(".")[0] in {"jwt", "cryptography"}:
|
||||
raise ModuleNotFoundError(f"No module named {name!r}")
|
||||
|
||||
sys.meta_path.insert(0, Blocker())
|
||||
import litellm
|
||||
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json
|
||||
try:
|
||||
jwks_document_json("not a key")
|
||||
except ImportError as e:
|
||||
print(e)
|
||||
"""
|
||||
)
|
||||
result: Final = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False)
|
||||
assert result.returncode == 0, result.stderr[-2000:]
|
||||
assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE
|
||||
|
||||
def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setitem(sys.modules, "jwt", None)
|
||||
with pytest.raises(ImportError, match="litellm\\[proxy\\]"):
|
||||
sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"})
|
||||
|
||||
1752
tests/test_litellm/llms/base_llm/auth/test_token_exchange.py
Normal file
1752
tests/test_litellm/llms/base_llm/auth/test_token_exchange.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
|
|
@ -17,7 +18,9 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
|
|||
AudioTranscriptionRequestData,
|
||||
BaseAudioTranscriptionConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
BaseLLMHTTPHandler,
|
||||
|
|
@ -455,7 +458,7 @@ async def test_async_anthropic_messages_handler_extra_headers():
|
|||
|
||||
# Mock the config
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
|
|
@ -502,7 +505,7 @@ async def test_async_anthropic_messages_handler_extra_headers():
|
|||
captured_headers.update(kwargs.get("headers", {}))
|
||||
return ({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
|
||||
mock_config.validate_anthropic_messages_environment = capture_validate
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate)
|
||||
|
||||
try:
|
||||
await handler.async_anthropic_messages_handler(
|
||||
|
|
@ -758,7 +761,7 @@ async def test_async_anthropic_messages_handler_passes_litellm_metadata():
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
|
|
@ -837,7 +840,7 @@ async def test_async_anthropic_messages_handler_forwards_router_model_info():
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
|
|
@ -929,7 +932,7 @@ async def test_async_anthropic_messages_handler_header_priority():
|
|||
captured_headers.update(kwargs.get("headers", {}))
|
||||
return ({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
|
||||
mock_config.validate_anthropic_messages_environment = capture_validate
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(side_effect=capture_validate)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
return_value={"model": "claude-3-opus-20240229", "messages": []}
|
||||
)
|
||||
|
|
@ -968,7 +971,7 @@ async def test_async_anthropic_messages_handler_drops_top_level_and_nested_param
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com")
|
||||
)
|
||||
|
||||
|
|
@ -1181,9 +1184,7 @@ def test_sync_delete_responses_sets_json_content_type():
|
|||
({}, True, None, None),
|
||||
],
|
||||
)
|
||||
def test_resolve_anthropic_messages_timeout(
|
||||
monkeypatch, litellm_params_kwargs, stream, global_timeout, expected
|
||||
):
|
||||
def test_resolve_anthropic_messages_timeout(monkeypatch, litellm_params_kwargs, stream, global_timeout, expected):
|
||||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
|
||||
if global_timeout is None:
|
||||
|
|
@ -1199,9 +1200,7 @@ def test_resolve_anthropic_messages_timeout(
|
|||
)
|
||||
else:
|
||||
monkeypatch.setattr("litellm.request_timeout", global_timeout, raising=False)
|
||||
monkeypatch.setattr(
|
||||
"litellm.request_timeout_explicitly_set", True, raising=False
|
||||
)
|
||||
monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False)
|
||||
|
||||
resolved = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout(
|
||||
litellm_params=GenericLiteLLMParams(**litellm_params_kwargs),
|
||||
|
|
@ -1222,13 +1221,11 @@ async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeyp
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "k"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
return_value={"model": "claude", "messages": []}
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []})
|
||||
mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages")
|
||||
mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None))
|
||||
mock_config.max_retry_on_anthropic_messages_http_error = 1
|
||||
|
|
@ -1270,13 +1267,11 @@ async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypa
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "k"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
return_value={"model": "claude", "messages": []}
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(return_value={"model": "claude", "messages": []})
|
||||
mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages")
|
||||
mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None))
|
||||
mock_config.max_retry_on_anthropic_messages_http_error = 1
|
||||
|
|
@ -1678,7 +1673,7 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
|
|||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"x-api-key": "sk-test"}, "https://api.anthropic.com")
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
|
|
@ -1686,7 +1681,13 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
|
|||
)
|
||||
mock_config.sign_request = Mock(return_value=({}, None))
|
||||
|
||||
fake_raw_response = {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "stop_reason": "end_turn"}
|
||||
fake_raw_response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
mock_config.transform_anthropic_messages_response = Mock(return_value=fake_raw_response)
|
||||
|
||||
mock_logging_obj = Mock()
|
||||
|
|
@ -1706,10 +1707,17 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
|
|||
mock_httpx_response.status_code = 200
|
||||
|
||||
with (
|
||||
patch.object(handler, "_async_post_anthropic_messages_with_http_error_retry", new=AsyncMock(return_value=mock_httpx_response)),
|
||||
patch.object(
|
||||
handler,
|
||||
"_async_post_anthropic_messages_with_http_error_retry",
|
||||
new=AsyncMock(return_value=mock_httpx_response),
|
||||
),
|
||||
patch.object(handler, "_call_agentic_completion_hooks", side_effect=fake_agentic_hooks),
|
||||
patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client"),
|
||||
patch("litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", return_value=None),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
result = await handler.async_anthropic_messages_handler(
|
||||
model="claude-haiku",
|
||||
|
|
@ -1951,7 +1959,9 @@ def test_audio_transcriptions_sends_dict_data_as_json_body():
|
|||
form-encodes it and silently ignores json=; JSON-body providers (e.g.
|
||||
Google Speech-to-Text) need an application/json body."""
|
||||
captured = {}
|
||||
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured))))
|
||||
client = HTTPHandler(
|
||||
client=httpx.Client(transport=httpx.MockTransport(_capture_json_transcription_request(captured)))
|
||||
)
|
||||
|
||||
response = BaseLLMHTTPHandler().audio_transcriptions(
|
||||
client=client,
|
||||
|
|
@ -2110,6 +2120,105 @@ def test_sync_retrieve_file_content_raises_on_http_error():
|
|||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
_FILE_CONTENT_WIF_ENV = {
|
||||
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_llm_http_handler_seam",
|
||||
"ANTHROPIC_ORGANIZATION_ID": "org-llm-http-handler-seam",
|
||||
"ANTHROPIC_IDENTITY_TOKEN": "llm-http-handler-seam-inline-jwt",
|
||||
}
|
||||
|
||||
|
||||
class _BlockingWifPoster:
|
||||
"""A token-endpoint poster that blocks until released, so the test can prove
|
||||
the exchange ran off the event loop's own thread instead of freezing it."""
|
||||
|
||||
def __init__(self):
|
||||
self.release = threading.Event()
|
||||
self.thread_ids = []
|
||||
|
||||
def post(self, url, *, content, headers, timeout):
|
||||
self.thread_ids.append(threading.get_ident())
|
||||
self.release.wait(timeout=5)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "sk-ant-oat01-llm-http-handler-seam",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_retrieve_file_content_wif_exchange_does_not_block_event_loop(monkeypatch):
|
||||
"""Regression (Greptile P1): async_retrieve_file_content called the synchronous
|
||||
validate_environment directly, so a cold WIF mint on this call site froze the
|
||||
event loop until the exchange finished. It must resolve credentials through the
|
||||
async facade instead."""
|
||||
from litellm.llms.anthropic import common_utils as anthropic_common_utils
|
||||
from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig
|
||||
from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token
|
||||
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine
|
||||
|
||||
for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
for name, value in _FILE_CONTENT_WIF_ENV.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
|
||||
poster = _BlockingWifPoster()
|
||||
engine = JwtBearerTokenExchangeEngine(poster=poster)
|
||||
sync_calls = []
|
||||
|
||||
def sync_shim(litellm_params, api_base, model):
|
||||
sync_calls.append(model)
|
||||
return get_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
async def async_shim(litellm_params, api_base, model):
|
||||
return await aget_anthropic_wif_token(litellm_params, api_base, model, engine)
|
||||
|
||||
monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim)
|
||||
monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim)
|
||||
|
||||
handler = BaseLLMHTTPHandler()
|
||||
client = Mock(spec=AsyncHTTPHandler)
|
||||
client.get = AsyncMock(return_value=httpx.Response(status_code=200, content=b"file bytes"))
|
||||
|
||||
ticks = []
|
||||
|
||||
async def ticker():
|
||||
for i in range(20):
|
||||
await asyncio.sleep(0.005)
|
||||
ticks.append(i)
|
||||
|
||||
ticker_task = asyncio.create_task(ticker())
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
retrieve_task = asyncio.create_task(
|
||||
handler.async_retrieve_file_content(
|
||||
file_content_request={"file_id": "file-abc"},
|
||||
provider_config=AnthropicFilesConfig(),
|
||||
litellm_params={},
|
||||
headers={},
|
||||
logging_obj=Mock(),
|
||||
client=client,
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
# The ticker kept advancing while the token exchange was still blocked on
|
||||
# poster.release, proving the exchange did not run inline on the event loop.
|
||||
assert len(ticks) > 0
|
||||
assert not retrieve_task.done()
|
||||
|
||||
poster.release.set()
|
||||
await retrieve_task
|
||||
await ticker_task
|
||||
|
||||
assert sync_calls == []
|
||||
assert poster.thread_ids
|
||||
assert poster.thread_ids[0] != threading.get_ident()
|
||||
sent_headers = client.get.call_args.kwargs["headers"]
|
||||
assert sent_headers["authorization"] == "Bearer sk-ant-oat01-llm-http-handler-seam"
|
||||
|
||||
|
||||
_UPSTREAM_NOT_FOUND_BODY = {
|
||||
"error": {
|
||||
"message": "Response with id 'resp_abc' not found.",
|
||||
|
|
@ -2257,9 +2366,7 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
|
|||
ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url))
|
||||
|
||||
class FakeAsyncClient:
|
||||
async def post(
|
||||
self, url, headers, data, stream=False, logging_obj=None, timeout=None
|
||||
):
|
||||
async def post(self, url, headers, data, stream=False, logging_obj=None, timeout=None):
|
||||
posts.append({"headers": dict(headers), "data": data})
|
||||
return invalid_signature_response if len(posts) == 1 else ok_response
|
||||
|
||||
|
|
@ -2576,7 +2683,7 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati
|
|||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
mock_config = Mock()
|
||||
mock_config.validate_anthropic_messages_environment = Mock(
|
||||
mock_config.avalidate_anthropic_messages_environment = AsyncMock(
|
||||
return_value=({"authorization": "Bearer t"}, "https://us-east5-aiplatform.googleapis.com")
|
||||
)
|
||||
mock_config.transform_anthropic_messages_request = Mock(
|
||||
|
|
@ -3100,6 +3207,70 @@ async def test_a_provider_that_keeps_rejecting_is_not_retried_forever_on_the_asy
|
|||
assert len(recorder.bodies) == 2
|
||||
|
||||
|
||||
def _async_client_returning(response: Mock) -> AsyncMock:
|
||||
client = AsyncMock(spec=AsyncHTTPHandler)
|
||||
client.post.return_value = response
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_file_async_awaits_the_provider_credential_hook_instead_of_blocking():
|
||||
provider_config = Mock(spec=BaseFilesConfig)
|
||||
provider_config.validate_environment.side_effect = AssertionError("sync validate_environment ran on the event loop")
|
||||
provider_config.avalidate_environment = AsyncMock(return_value={"x-api-key": "federated"})
|
||||
provider_config.get_complete_file_url.return_value = "https://files.example/v1/files"
|
||||
provider_config.transform_create_file_request.return_value = {"file": ("batch.jsonl", b"{}", "application/jsonl")}
|
||||
file_object = object()
|
||||
provider_config.transform_create_file_response.return_value = file_object
|
||||
client = _async_client_returning(Mock(spec=httpx.Response))
|
||||
|
||||
result = await BaseLLMHTTPHandler().create_file(
|
||||
create_file_data={"file": b"{}", "purpose": "batch"},
|
||||
litellm_params={},
|
||||
provider_config=provider_config,
|
||||
headers={},
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
logging_obj=Mock(),
|
||||
_is_async=True,
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert result is file_object
|
||||
provider_config.validate_environment.assert_not_called()
|
||||
provider_config.avalidate_environment.assert_awaited_once()
|
||||
assert client.post.call_args.kwargs["headers"] == {"x-api-key": "federated"}
|
||||
assert client.post.call_args.kwargs["url"] == "https://files.example/v1/files"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_async_validates_credentials_off_the_event_loop():
|
||||
provider_config = Mock(spec=BaseBatchesConfig)
|
||||
provider_config.validate_environment.side_effect = lambda **_: {"x-validated-on": str(threading.get_ident())}
|
||||
provider_config.get_complete_batch_url.return_value = "https://batches.example/v1/messages/batches"
|
||||
provider_config.transform_create_batch_request.return_value = {"requests": []}
|
||||
batch = object()
|
||||
provider_config.transform_create_batch_response.return_value = batch
|
||||
client = _async_client_returning(Mock(spec=httpx.Response))
|
||||
|
||||
result = await BaseLLMHTTPHandler().create_batch(
|
||||
create_batch_data={"input_file_id": "file_1", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
litellm_params={},
|
||||
provider_config=provider_config,
|
||||
headers={},
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
logging_obj=Mock(),
|
||||
_is_async=True,
|
||||
client=client,
|
||||
model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is batch
|
||||
validated_on = client.post.call_args.kwargs["headers"]["x-validated-on"]
|
||||
assert validated_on != str(threading.get_ident())
|
||||
assert client.post.call_args.kwargs["url"] == "https://batches.example/v1/messages/batches"
|
||||
|
||||
CONTAINER_NOT_FOUND_BODY = {
|
||||
"error": {
|
||||
"message": "Container with id 'cntr_gone' not found.",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from openai import AsyncOpenAI, OpenAI
|
|||
|
||||
import litellm
|
||||
from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError
|
||||
from litellm.llms.openai.openai import OpenAIChatCompletion
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
|
|
@ -22,6 +23,17 @@ from litellm.llms.openai.workload_identity import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token"
|
||||
CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions"
|
||||
EMBEDDINGS_URL: Final = "https://api.openai.com/v1/embeddings"
|
||||
MODELS_URL: Final = "https://api.openai.com/v1/models"
|
||||
CHAT_COMPLETION_BODY: Final = {
|
||||
"id": "chatcmpl-wif",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -279,3 +291,342 @@ class TestResponsesValidateEnvironment:
|
|||
headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer None"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def deployment_wif(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> dict[str, str]:
|
||||
token_file: Final = tmp_path / "deployment_subject_token.jwt"
|
||||
token_file.write_text("subject-token-from-deployment-file")
|
||||
for name in (
|
||||
"OPENAI_API_KEY",
|
||||
"OPENAI_BASE_URL",
|
||||
"OPENAI_API_BASE",
|
||||
"OPENAI_IDENTITY_PROVIDER_ID",
|
||||
"OPENAI_SERVICE_ACCOUNT_ID",
|
||||
"OPENAI_IDENTITY_TOKEN_FILE",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.setattr(litellm, "api_base", None)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
_workload_identity_auth.cache_clear()
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
return {
|
||||
"openai_identity_provider_id": "idp_deployment",
|
||||
"openai_service_account_id": "user-deployment",
|
||||
"openai_identity_token_file": str(token_file),
|
||||
}
|
||||
|
||||
|
||||
def deployment_config(deployment_wif: dict[str, str]) -> OpenAIWorkloadIdentityConfig:
|
||||
return OpenAIWorkloadIdentityConfig(
|
||||
identity_provider_id="idp_deployment",
|
||||
service_account_id="user-deployment",
|
||||
token_file=deployment_wif["openai_identity_token_file"],
|
||||
)
|
||||
|
||||
|
||||
def mock_chat_completions() -> respx.Route:
|
||||
return respx.post(CHAT_COMPLETIONS_URL).mock(return_value=httpx.Response(200, json=CHAT_COMPLETION_BODY))
|
||||
|
||||
|
||||
def mock_streaming_chat_completions() -> respx.Route:
|
||||
chunk: Final = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
|
||||
events: Final = (
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}]},
|
||||
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
||||
)
|
||||
body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n"
|
||||
return respx.post(CHAT_COMPLETIONS_URL).mock(
|
||||
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body)
|
||||
)
|
||||
|
||||
|
||||
class TestResolveConfigFromDeployment:
|
||||
def test_resolves_from_litellm_params_without_env(self, deployment_wif: dict[str, str]) -> None:
|
||||
assert resolve_openai_workload_identity_config(
|
||||
api_key=None, api_base=None, litellm_params=deployment_wif
|
||||
) == deployment_config(deployment_wif)
|
||||
|
||||
def test_env_alone_disables_nothing_when_params_are_absent(self, deployment_wif: dict[str, str]) -> None:
|
||||
assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=None) is None
|
||||
|
||||
def test_unrelated_litellm_params_do_not_resolve(self, deployment_wif: dict[str, str]) -> None:
|
||||
assert (
|
||||
resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params={"model": "gpt-4o"})
|
||||
is None
|
||||
)
|
||||
|
||||
def test_litellm_params_beat_env(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
|
||||
config: Final = resolve_openai_workload_identity_config(
|
||||
api_key=None,
|
||||
api_base=None,
|
||||
litellm_params={
|
||||
"openai_identity_provider_id": "idp_deployment",
|
||||
"openai_service_account_id": "user-deployment",
|
||||
"openai_identity_token_file": wif_env.token_file,
|
||||
},
|
||||
)
|
||||
assert config == OpenAIWorkloadIdentityConfig(
|
||||
identity_provider_id="idp_deployment",
|
||||
service_account_id="user-deployment",
|
||||
token_file=wif_env.token_file,
|
||||
)
|
||||
|
||||
def test_partial_litellm_params_fill_from_env_per_field(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
|
||||
config: Final = resolve_openai_workload_identity_config(
|
||||
api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": "idp_deployment"}
|
||||
)
|
||||
assert config == OpenAIWorkloadIdentityConfig(
|
||||
identity_provider_id="idp_deployment",
|
||||
service_account_id=wif_env.service_account_id,
|
||||
token_file=wif_env.token_file,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("blank", ["", None, 7])
|
||||
def test_blank_or_non_string_param_falls_back_to_env(
|
||||
self, wif_env: OpenAIWorkloadIdentityConfig, blank: object
|
||||
) -> None:
|
||||
config: Final = resolve_openai_workload_identity_config(
|
||||
api_key=None, api_base=None, litellm_params={"openai_identity_provider_id": blank}
|
||||
)
|
||||
assert config == wif_env
|
||||
|
||||
def test_partial_litellm_params_without_env_disable(self, deployment_wif: dict[str, str]) -> None:
|
||||
partial: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"}
|
||||
assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=partial) is None
|
||||
|
||||
def test_static_api_key_beats_litellm_params(self, deployment_wif: dict[str, str]) -> None:
|
||||
assert (
|
||||
resolve_openai_workload_identity_config(api_key="sk-static", api_base=None, litellm_params=deployment_wif)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_env_openai_api_key_beats_litellm_params(
|
||||
self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env")
|
||||
assert (
|
||||
resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=deployment_wif) is None
|
||||
)
|
||||
|
||||
def test_foreign_api_base_disables_deployment_wif(self, deployment_wif: dict[str, str]) -> None:
|
||||
assert (
|
||||
resolve_openai_workload_identity_config(
|
||||
api_key=None, api_base="https://my-vllm.internal/v1", litellm_params=deployment_wif
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
class TestDeploymentClientConstruction:
|
||||
def test_sync_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None:
|
||||
client: Final = OpenAIChatCompletion()._get_openai_client(
|
||||
is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif
|
||||
)
|
||||
assert isinstance(client, OpenAI)
|
||||
assert client.api_key == "workload-identity-auth"
|
||||
assert client._workload_identity_auth is not None
|
||||
|
||||
def test_async_client_from_deployment_params(self, deployment_wif: dict[str, str]) -> None:
|
||||
client: Final = OpenAIChatCompletion()._get_openai_client(
|
||||
is_async=True, api_key=None, api_base=None, litellm_params=deployment_wif
|
||||
)
|
||||
assert isinstance(client, AsyncOpenAI)
|
||||
assert client._workload_identity_auth is not None
|
||||
|
||||
def test_distinct_deployments_get_distinct_cached_clients(self, deployment_wif: dict[str, str]) -> None:
|
||||
other_deployment: Final = {**deployment_wif, "openai_service_account_id": "user-other"}
|
||||
handler: Final = OpenAIChatCompletion()
|
||||
first: Final = handler._get_openai_client(
|
||||
is_async=False, api_key=None, api_base=None, litellm_params=deployment_wif
|
||||
)
|
||||
second: Final = handler._get_openai_client(
|
||||
is_async=False, api_key=None, api_base=None, litellm_params=other_deployment
|
||||
)
|
||||
again: Final = handler._get_openai_client(
|
||||
is_async=False, api_key=None, api_base=None, litellm_params=dict(deployment_wif)
|
||||
)
|
||||
assert first is not second
|
||||
assert again is first
|
||||
|
||||
@respx.mock
|
||||
def test_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("deployment-bearer")
|
||||
completion_route: Final = mock_chat_completions()
|
||||
|
||||
response: Final = litellm.completion(
|
||||
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], **deployment_wif
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "ok"
|
||||
request: Final = completion_route.calls.last.request
|
||||
assert request.headers["Authorization"] == "Bearer deployment-bearer"
|
||||
assert not any(key.startswith("openai_") for key in json.loads(request.content))
|
||||
|
||||
@respx.mock
|
||||
def test_streaming_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("stream-bearer")
|
||||
stream_route: Final = mock_streaming_chat_completions()
|
||||
|
||||
chunks: Final = tuple(
|
||||
litellm.completion(
|
||||
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif
|
||||
)
|
||||
)
|
||||
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
|
||||
assert stream_route.calls.last.request.headers["Authorization"] == "Bearer stream-bearer"
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_streaming_completion_kwargs_carry_exchanged_bearer(
|
||||
self, deployment_wif: dict[str, str]
|
||||
) -> None:
|
||||
mock_token_exchange("async-stream-bearer")
|
||||
stream_route: Final = mock_streaming_chat_completions()
|
||||
|
||||
stream: Final = await litellm.acompletion(
|
||||
model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **deployment_wif
|
||||
)
|
||||
chunks: Final = tuple([chunk async for chunk in stream])
|
||||
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
|
||||
assert stream_route.calls.last.request.headers["Authorization"] == "Bearer async-stream-bearer"
|
||||
|
||||
@respx.mock
|
||||
def test_router_deployment_without_api_key_authenticates_via_token_exchange(
|
||||
self, deployment_wif: dict[str, str]
|
||||
) -> None:
|
||||
exchange_route: Final = mock_token_exchange("router-bearer")
|
||||
completion_route: Final = mock_chat_completions()
|
||||
router: Final = litellm.Router(
|
||||
model_list=[{"model_name": "wif-gpt", "litellm_params": {"model": "openai/gpt-4o-mini", **deployment_wif}}]
|
||||
)
|
||||
|
||||
response: Final = router.completion(model="wif-gpt", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert response.choices[0].message.content == "ok"
|
||||
assert exchange_route.called
|
||||
assert completion_route.calls.last.request.headers["Authorization"] == "Bearer router-bearer"
|
||||
|
||||
@respx.mock
|
||||
def test_embedding_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("embedding-bearer")
|
||||
embeddings_route: Final = respx.post(EMBEDDINGS_URL).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"object": "list",
|
||||
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 1, "total_tokens": 1},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
litellm.embedding(model="openai/text-embedding-3-small", input=["hi"], **deployment_wif)
|
||||
|
||||
assert embeddings_route.calls.last.request.headers["Authorization"] == "Bearer embedding-bearer"
|
||||
|
||||
|
||||
class TestResponsesValidateEnvironmentFromDeployment:
|
||||
@respx.mock
|
||||
def test_mints_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("responses-bearer")
|
||||
headers: Final = OpenAIResponsesAPIConfig().validate_environment(
|
||||
headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams(**deployment_wif)
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer responses-bearer"
|
||||
|
||||
def test_static_key_in_litellm_params_wins(self, deployment_wif: dict[str, str]) -> None:
|
||||
headers: Final = OpenAIResponsesAPIConfig().validate_environment(
|
||||
headers={},
|
||||
model="gpt-4o-mini",
|
||||
litellm_params=GenericLiteLLMParams(api_key="sk-responses", **deployment_wif),
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer sk-responses"
|
||||
|
||||
|
||||
class TestDiscoverModels:
|
||||
@staticmethod
|
||||
def mock_models() -> respx.Route:
|
||||
return respx.get(MODELS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o-mini"}, {"id": "gpt-4.1"}]})
|
||||
)
|
||||
|
||||
@respx.mock
|
||||
def test_discovers_with_exchanged_bearer_from_litellm_params(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("discovery-bearer")
|
||||
models_route: Final = self.mock_models()
|
||||
|
||||
assert OpenAIGPTConfig().discover_models(deployment_wif) == ["gpt-4o-mini", "gpt-4.1"]
|
||||
assert models_route.calls.last.request.headers["Authorization"] == "Bearer discovery-bearer"
|
||||
|
||||
@respx.mock
|
||||
def test_discovers_with_env_wif_when_params_carry_no_key(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
|
||||
mock_token_exchange("env-discovery-bearer")
|
||||
models_route: Final = self.mock_models()
|
||||
|
||||
OpenAIGPTConfig().discover_models({})
|
||||
|
||||
assert models_route.calls.last.request.headers["Authorization"] == "Bearer env-discovery-bearer"
|
||||
|
||||
@respx.mock
|
||||
def test_static_api_key_in_params_skips_token_exchange(self, deployment_wif: dict[str, str]) -> None:
|
||||
exchange_route: Final = mock_token_exchange()
|
||||
models_route: Final = self.mock_models()
|
||||
|
||||
OpenAIGPTConfig().discover_models({**deployment_wif, "api_key": "sk-discovery"})
|
||||
|
||||
assert models_route.calls.last.request.headers["Authorization"] == "Bearer sk-discovery"
|
||||
assert not exchange_route.called
|
||||
|
||||
@respx.mock
|
||||
def test_blank_api_base_in_params_discovers_from_openai(self, deployment_wif: dict[str, str]) -> None:
|
||||
mock_token_exchange("blank-base-bearer")
|
||||
models_route: Final = self.mock_models()
|
||||
|
||||
assert OpenAIGPTConfig().discover_models({**deployment_wif, "api_base": ""}) == ["gpt-4o-mini", "gpt-4.1"]
|
||||
assert models_route.calls.last.request.headers["Authorization"] == "Bearer blank-base-bearer"
|
||||
|
||||
@respx.mock
|
||||
def test_openai_compatible_subclass_never_mints_wif(self, deployment_wif: dict[str, str]) -> None:
|
||||
exchange_route: Final = mock_token_exchange()
|
||||
models_route: Final = self.mock_models()
|
||||
|
||||
class CompatibleConfig(OpenAIGPTConfig):
|
||||
pass
|
||||
|
||||
CompatibleConfig().discover_models(deployment_wif)
|
||||
|
||||
assert models_route.calls.last.request.headers["Authorization"] == "Bearer None"
|
||||
assert not exchange_route.called
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_empty_static_key_never_borrows_the_env_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-key-that-must-stay-home")
|
||||
foreign_models: Final = respx.get("https://third-party.example/v1/models").mock(
|
||||
return_value=httpx.Response(200, json={"data": [{"id": "other-model"}]})
|
||||
)
|
||||
|
||||
assert OpenAIGPTConfig().get_models(api_key="", api_base="https://third-party.example") == ["other-model"]
|
||||
assert foreign_models.calls.last.request.headers["Authorization"] == "Bearer "
|
||||
|
||||
|
||||
class TestClientsideBaseOverride:
|
||||
def test_client_api_base_override_clears_deployment_wif(self, deployment_wif: dict[str, str]) -> None:
|
||||
from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params
|
||||
|
||||
redirected: Final = get_dynamic_litellm_params(
|
||||
litellm_params={"model": "openai/gpt-4o-mini", **deployment_wif},
|
||||
request_kwargs={"api_base": "https://not-openai.example/v1"},
|
||||
)
|
||||
|
||||
assert not any(key in redirected for key in deployment_wif)
|
||||
assert (
|
||||
resolve_openai_workload_identity_config(
|
||||
api_key=None, api_base=redirected["api_base"], litellm_params=redirected
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue