diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 52b7c74412b..ab6d39dd4af 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -7,13 +7,11 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_litellm_params import ( - ANTHROPIC_WIF_KWARGS_KEYS, - AWS_CREDENTIAL_KWARGS_KEYS, -) +from litellm.litellm_core_utils.get_litellm_params import 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 +from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS from litellm.utils import token_counter diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 7fbab6e59ee..9a592665376 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import Final from litellm.llms.openai.data_residency import infer_openai_data_residency +from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( { @@ -24,48 +25,6 @@ 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`. diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py index 97dd5014caa..e9ef47da59c 100644 --- a/litellm/llms/anthropic/wif.py +++ b/litellm/llms/anthropic/wif.py @@ -58,7 +58,7 @@ _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 +# also be listed in ANTHROPIC_WIF_KWARGS_KEYS (types/workload_identity.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( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 110e359f4af..89b5de62f9e 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -10,7 +10,6 @@ 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, ) @@ -24,6 +23,7 @@ from litellm.types.llms.openai import * from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS from ..common_utils import OpenAIError from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1142b8263f4..0d5b794b507 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -49,6 +49,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo +from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from .agents import LiteLLMSendMessageResponse @@ -3643,20 +3644,6 @@ 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 diff --git a/litellm/types/workload_identity.py b/litellm/types/workload_identity.py new file mode 100644 index 00000000000..26362a7e442 --- /dev/null +++ b/litellm/types/workload_identity.py @@ -0,0 +1,43 @@ +"""litellm_params keys that configure workload identity federation. + +The kwargs funnel (``litellm_core_utils.get_litellm_params``) and the request-body ban list +(``types.utils.all_litellm_params``) both derive from these sets, so they live in a module with +no litellm imports that either side can reach without a cycle. Every key here rides the funnel +into ``litellm_params`` and is banned from request bodies, which also covers +``anthropic_disable_workload_identity_federation``: the proxy sets it when a client redirects +``api_base`` so a federated deployment stops minting for a base the caller chose, and a caller +must not be able to set it in either direction. +""" + +from typing import Final + +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", + "anthropic_identity_source", + "anthropic_issuer_url", + "anthropic_issuer_subject", + "anthropic_issuer_audience", + "anthropic_issuer_ttl_seconds", + "anthropic_issuer_signing_key_ref", + "anthropic_keycloak_token_url", + "anthropic_keycloak_client_id", + "anthropic_keycloak_auth_method", + "anthropic_keycloak_client_secret_ref", + "anthropic_keycloak_scope", + "anthropic_disable_workload_identity_federation", + } +) + +OPENAI_WIF_KWARGS_KEYS: Final = frozenset( + { + "openai_identity_provider_id", + "openai_service_account_id", + "openai_identity_token_file", + } +) diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index ad4e417a29f..0ad3e3afa24 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -308,7 +308,7 @@ class TestAnthropicWifIdentitySourceKeys: """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 + from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS assert set(self.NEW_KEYS) == ANTHROPIC_WIF_KWARGS_KEYS - set(TestAnthropicWifKeys.SIX_KEYS) @@ -346,7 +346,7 @@ class TestOpenAIWifKeys: } def test_keys_are_exactly_the_registered_set(self): - from litellm.litellm_core_utils.get_litellm_params import OPENAI_WIF_KWARGS_KEYS + from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS assert set(self.THREE_KEYS) == OPENAI_WIF_KWARGS_KEYS diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f5c7ac713c1..87e4bcc8191 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -31,16 +31,13 @@ from litellm.proxy.auth.auth_utils import ( def test_every_server_owned_wif_kwarg_key_is_request_banned(): """server_owned_wif_litellm_params (types/utils.py) is derived from ANTHROPIC_WIF_KWARGS_KEYS - and OPENAI_WIF_KWARGS_KEYS (get_litellm_params.py) precisely so a new WIF field can never be + and OPENAI_WIF_KWARGS_KEYS (types/workload_identity.py) precisely so a new WIF field can never be added to the kwargs funnel without automatically joining the request-body ban list; this guards that invariant itself, independent of today's field count, so it fails if the derivation is ever reverted to a hand-typed list that drifts.""" - from litellm.litellm_core_utils.get_litellm_params import ( - ANTHROPIC_WIF_KWARGS_KEYS, - OPENAI_WIF_KWARGS_KEYS, - ) from litellm.proxy.auth.auth_utils import _SERVER_OWNED_WIF_UNCONDITIONAL_BANNED + from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS assert ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS == set(_SERVER_OWNED_WIF_UNCONDITIONAL_BANNED) diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 09d5e23e2e6..a1b07155304 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -69,9 +69,7 @@ def _verify_only_requested_name_imported(name: str, all_names: tuple): litellm_globals = sys.modules["litellm"].__dict__ for other_name in all_names: if other_name != name: - assert ( - other_name not in litellm_globals - ), f"{other_name} should not be imported when importing {name}" + assert other_name not in litellm_globals, f"{other_name} should not be imported when importing {name}" def _verify_only_requested_name_imported_in_utils(name: str, all_names: tuple): @@ -80,9 +78,7 @@ def _verify_only_requested_name_imported_in_utils(name: str, all_names: tuple): utils_globals = sys.modules["litellm.utils"].__dict__ for other_name in all_names: if other_name != name: - assert ( - other_name not in utils_globals - ), f"{other_name} should not be imported when importing {name}" + assert other_name not in utils_globals, f"{other_name} should not be imported when importing {name}" def test_cost_calculator_lazy_imports(): @@ -394,6 +390,17 @@ def test_proxy_private_submodule_resolves_in_fresh_process(): assert result.stdout.strip() == "litellm.proxy._types" +@pytest.mark.parametrize( + "module", + ["litellm.litellm_core_utils.get_litellm_params", "litellm.batches.batch_utils", "litellm.types.utils"], +) +def test_kwargs_funnel_and_its_importers_load_first_in_fresh_process(module: str): + """With `import litellm` lazy, these modules are often the first to pull in litellm.types.utils, and + the WIF key sets shared between the funnel and all_litellm_params must not turn that into a cycle.""" + result = subprocess.run([sys.executable, "-c", f"import {module}"], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + + def test_lazy_instances_are_singletons(): """Lazily created instances are cached, so repeated access returns the same object.""" assert litellm._key_management_settings is litellm._key_management_settings