diff --git a/litellm/llms/watsonx/passthrough/transformation.py b/litellm/llms/watsonx/passthrough/transformation.py index c3697314310..08e286b306b 100644 --- a/litellm/llms/watsonx/passthrough/transformation.py +++ b/litellm/llms/watsonx/passthrough/transformation.py @@ -1,11 +1,19 @@ +from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final +from litellm._logging import verbose_router_logger +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.watsonx.common_utils import IBMWatsonXMixin +from litellm.types.passthrough_endpoints.watsonx import WatsonxPassThroughCredentials if TYPE_CHECKING: from httpx import URL + from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict + class WatsonxPassthroughConfig(IBMWatsonXMixin, BasePassthroughConfig): """ @@ -65,3 +73,123 @@ class WatsonxPassthroughConfig(IBMWatsonXMixin, BasePassthroughConfig): def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: return super().get_models(api_key, api_base) + + # ------------------------------------------------------------------ + # Router deployment credential resolution + # ------------------------------------------------------------------ + + @classmethod + def resolve_credentials_from_router_deployments( + cls, + model_id: str | None = None, + project_id: str | None = None, + space_id: str | None = None, + deployments: "Sequence[DeploymentTypedDict] | None" = None, + ) -> "WatsonxPassThroughCredentials | None": + """ + Resolve watsonx pass-through credentials from router deployments flagged + with ``use_in_pass_through: true``. + + Disambiguation rules (applied in order): + 1. Exact match: model_id + project_id (or space_id) — both must match. + 2. Partial match: model_id only — first deployment whose model matches. + 3. Fallback: first eligible deployment (warns if multiple exist). + + Returns ``None`` when no eligible deployment exists, allowing the caller + to continue with env-var resolution. + """ + candidates: Final = tuple( + creds + for deployment in (deployments or ()) + if (creds := cls.resolve_credentials_from_deployment(deployment["litellm_params"])) is not None + ) + + if not candidates: + return None + first_candidate: Final = candidates[0] + + def model_ok(creds: WatsonxPassThroughCredentials) -> bool: + return bool(model_id and cls.deployment_model_matches(creds, model_id)) + + def project_ok(creds: WatsonxPassThroughCredentials) -> bool: + return bool(project_id) and creds.get("project_id") == project_id + + def space_ok(creds: WatsonxPassThroughCredentials) -> bool: + return bool(space_id) and creds.get("space_id") == space_id + + selection_rules: Final[tuple[Callable[[WatsonxPassThroughCredentials], bool], ...]] = ( + lambda creds: model_ok(creds) and project_ok(creds), + lambda creds: model_ok(creds) and space_ok(creds), + model_ok, + project_ok, + space_ok, + ) + matched: Final = next( + filter(None, (next(filter(rule, candidates), None) for rule in selection_rules)), + None, + ) + if matched is not None: + return matched + + # Fallback: single deployment — unambiguous; multiple — return first with a warning + if len(candidates) > 1: + verbose_router_logger.warning( + "Multiple watsonx use_in_pass_through deployments found but no " + "model_id/project_id/space_id supplied to disambiguate — using the first. " + "Set project_id or space_id on each deployment's litellm_params and pass " + "project_id or space_id in the request body to route to the correct credentials." + ) + return first_candidate + + @staticmethod + def deployment_model_matches(creds: "WatsonxPassThroughCredentials", model_id: str) -> bool: + """Check if the credential's stored model ends with the request's model_id.""" + stored: Final = creds.get("model") + if stored is None: + return False + return stored == model_id or str(stored).endswith(f"/{model_id}") + + @classmethod + def resolve_credentials_from_deployment( + cls, litellm_params: "LiteLLMParamsTypedDict" + ) -> "WatsonxPassThroughCredentials | None": + """ + Extract WatsonxPassThroughCredentials from a single deployment's litellm_params + if it is flagged ``use_in_pass_through: true`` and is a watsonx provider. + Returns ``None`` for non-watsonx or unflagged deployments. + """ + if litellm_params.get("use_in_pass_through") is not True: + return None + model: Final = litellm_params.get("model") + if model is None: + return None + try: + _, provider, _, _ = get_llm_provider( + model=model, + custom_llm_provider=litellm_params.get("custom_llm_provider"), + ) + except BadRequestError: + return None + if provider != "watsonx": + return None + credential_name: Final = litellm_params.get("litellm_credential_name") + credential_values: Final[Mapping[str, object]] = ( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ) + credential_api_key: Final = credential_values.get("api_key") + credential_api_base: Final = credential_values.get("api_base") + api_key: Final = (credential_api_key if isinstance(credential_api_key, str) else None) or litellm_params.get( + "api_key" + ) + api_base: Final = (credential_api_base if isinstance(credential_api_base, str) else None) or litellm_params.get( + "api_base" + ) + # api_key and api_base are both optional: either can be None and fall back to + # WATSONX_API_KEY / WATSONX_API_BASE env vars inside validate_environment / get_complete_url. + return WatsonxPassThroughCredentials( + api_key=api_key, + api_base=api_base, + project_id=litellm_params.get("project_id"), + space_id=litellm_params.get("space_id"), + model=model, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 33b18c3d939..3414c03362d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -159,6 +159,8 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router + from litellm.types.passthrough_endpoints.watsonx import WatsonxPassThroughCredentials + from litellm.types.router import DeploymentTypedDict ProxyConfig = _ProxyConfig # rebind-ok: conditional type alias else: @@ -4439,6 +4441,84 @@ async def handle_gigachat_passthrough_router_model( return result +async def _watsonx_request_body(request: Request) -> Mapping[str, object]: + if request.method != "POST": + return {} + if "multipart/form-data" in request.headers.get("content-type", ""): + return await get_form_data(request) + return await _json_request_body(request) + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) else None + + +async def _watsonx_allowed_deployments( + user_api_key_dict: UserAPIKeyAuth, requested_model: str | None +) -> Sequence[DeploymentTypedDict]: + from fastapi import status + + from litellm.llms.watsonx.passthrough.transformation import WatsonxPassthroughConfig + from litellm.proxy._types import ModelAccessDeniedProxyException, ProxyErrorTypes + from litellm.proxy.auth.auth_checks import model_access_denied_client_message + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # untyped proxy global + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + user_model, # pyright: ignore[reportUnknownVariableType] # untyped proxy global + ) + from litellm.proxy.utils import ( + get_available_models_for_user, # pyright: ignore[reportUnknownVariableType] # untyped settings parameter + ) + + if llm_router is None: + return () + available_models: Final = frozenset( + await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + general_settings=cast(dict[str, object], general_settings), # cast-ok: proxy server settings + user_model=cast(str | None, user_model), # cast-ok: proxy server untyped global + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + ) + ) + allowed_deployments: Final = tuple( + d + for d in llm_router.get_model_list() or () + if d["model_name"] in available_models + or (isinstance(m := d["litellm_params"].get("model"), str) and m in available_models) + ) + if requested_model is None: + return allowed_deployments + model_is_accessible: Final = ( + requested_model in available_models + or f"watsonx/{requested_model}" in available_models + or any( + WatsonxPassthroughConfig.deployment_model_matches(creds, requested_model) + for d in allowed_deployments + if (creds := WatsonxPassthroughConfig.resolve_credentials_from_deployment(d["litellm_params"])) is not None + ) + ) + has_model_restrictions: Final = ( + bool(user_api_key_dict.models) + or bool(user_api_key_dict.team_models) + or bool(user_api_key_dict.access_group_ids) + ) + if has_model_restrictions and not model_is_accessible: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=requested_model), + internal_message=f"Key not allowed to access model {requested_model}", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return allowed_deployments + + @router.api_route( "/watsonx/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -4470,10 +4550,45 @@ async def watsonx_proxy_route( if provider_config is None: raise HTTPException(status_code=404, detail="Watsonx passthrough config not found") + request_body: Final = await _watsonx_request_body(request) + is_streaming_request: Final = bool(request_body.get("stream")) + + from litellm.llms.watsonx.passthrough.transformation import WatsonxPassthroughConfig + + requested_model: Final = ( + _str_or_none(request_body.get("model_id")) + or _str_or_none(request_body.get("model")) + or _str_or_none(request.query_params.get("model_id")) + or _str_or_none(request.query_params.get("model")) + ) + requested_project_id: Final = _str_or_none(request_body.get("project_id")) or _str_or_none( + request.query_params.get("project_id") + ) + requested_space_id: Final = _str_or_none(request_body.get("space_id")) or _str_or_none( + request.query_params.get("space_id") + ) + allowed_deployments: Final = await _watsonx_allowed_deployments( + user_api_key_dict=user_api_key_dict, requested_model=requested_model + ) + watsonx_router_creds: Final[WatsonxPassThroughCredentials | None] = ( + WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + model_id=requested_model, + project_id=requested_project_id, + space_id=requested_space_id, + deployments=allowed_deployments, + ) + ) + _resolved_api_key: Final[str | None] = ( + watsonx_router_creds.get("api_key") if watsonx_router_creds is not None else None + ) + _resolved_api_base: Final[str | None] = ( + watsonx_router_creds.get("api_base") if watsonx_router_creds is not None else None + ) + # Get complete URL with version parameter complete_url, _ = provider_config.get_complete_url( - api_base=None, - api_key=None, + api_base=_resolved_api_base, + api_key=_resolved_api_key, model="", endpoint=endpoint, request_query_params=None, @@ -4487,21 +4602,10 @@ async def watsonx_proxy_route( messages=[], optional_params={}, litellm_params={}, - api_key=None, - api_base=None, + api_key=_resolved_api_key, + api_base=_resolved_api_base, ) - # Check for streaming - is_streaming_request = False - if request.method == "POST": - if "multipart/form-data" not in request.headers.get("content-type", ""): - _request_body = await _json_request_body(request) - else: - _request_body = await get_form_data(request) - - if _request_body.get("stream"): - is_streaming_request = True - request_query_params: Final = dict(request.query_params) if request_query_params.get("version") is None: request_query_params["version"] = litellm.WATSONX_DEFAULT_API_VERSION diff --git a/litellm/types/passthrough_endpoints/watsonx.py b/litellm/types/passthrough_endpoints/watsonx.py new file mode 100644 index 00000000000..d25d2536ded --- /dev/null +++ b/litellm/types/passthrough_endpoints/watsonx.py @@ -0,0 +1,13 @@ +""" +Used for /watsonx/ pass-through endpoints. +""" + +from typing_extensions import NotRequired, ReadOnly, TypedDict + + +class WatsonxPassThroughCredentials(TypedDict): + api_key: ReadOnly[NotRequired[str | None]] + api_base: ReadOnly[NotRequired[str | None]] + project_id: ReadOnly[NotRequired[str | None]] + space_id: ReadOnly[NotRequired[str | None]] + model: ReadOnly[NotRequired[str | None]] diff --git a/litellm/types/router.py b/litellm/types/router.py index 726a328e913..023f102f4e4 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -790,7 +790,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): timeout_by_service_tier: ReadOnly[Mapping[str, float] | None] stream_timeout_by_service_tier: ReadOnly[Mapping[str, float] | None] max_retries: int | None - organization: list | str | None # for openai orgs + organization: list[str] | str | None # for openai orgs provider_affinity_header: ReadOnly[str | None] configurable_clientside_auth_params: ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS # for allowing api base switching on finetuned models diff --git a/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py b/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py index d8976d19f5c..97bc96f37d9 100644 --- a/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py +++ b/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py @@ -5,11 +5,9 @@ Tests the Watsonx-specific passthrough configuration including URL construction, streaming detection, and authentication handling. """ -from unittest.mock import MagicMock, patch +from unittest.mock import patch import httpx -import pytest - import litellm from litellm.llms.watsonx.passthrough.transformation import WatsonxPassthroughConfig @@ -277,3 +275,370 @@ class TestWatsonxPassthroughConfig: # Test with empty dict (defaults to False) assert config.is_streaming_request("endpoint", {}) is False + + +# --------------------------------------------------------------------------- +# Helpers shared by credential-resolution tests +# --------------------------------------------------------------------------- + +def _wx_deployment(model: str, **litellm_params) -> dict: + """Build a minimal router deployment dict for watsonx.""" + return { + "model_name": model.split("/", 1)[-1], + "litellm_params": {"model": model, "use_in_pass_through": True, **litellm_params}, + } + + +def _wx_unflagged_deployment(model: str, **litellm_params) -> dict: + """Build a deployment that does NOT set use_in_pass_through.""" + return { + "model_name": model.split("/", 1)[-1], + "litellm_params": {"model": model, **litellm_params}, + } + + +# --------------------------------------------------------------------------- +# deployment_model_matches +# --------------------------------------------------------------------------- + +class TestDeploymentModelMatches: + """Tests for WatsonxPassthroughConfig.deployment_model_matches.""" + + def test_exact_match(self): + creds = {"model": "ibm/granite-13b-chat-v2"} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-13b-chat-v2") is True + + def test_suffix_match(self): + """watsonx/ prefix stripped — model_id from body is the bare id.""" + creds = {"model": "watsonx/ibm/granite-13b-chat-v2"} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-13b-chat-v2") is True + + def test_no_match_different_model(self): + creds = {"model": "watsonx/ibm/granite-8b-code"} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-13b-chat-v2") is False + + def test_no_match_when_model_none(self): + creds = {"model": None} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-13b-chat-v2") is False + + def test_no_match_when_model_key_absent(self): + creds = {} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-13b-chat-v2") is False + + def test_partial_suffix_does_not_match(self): + """'granite-13b' must not match 'ibm/granite-13b-chat-v2'.""" + creds = {"model": "watsonx/ibm/granite-13b-chat-v2"} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "granite-13b-chat-v2") is True + + def test_wrong_suffix_does_not_match(self): + creds = {"model": "watsonx/ibm/granite-13b-chat-v2"} + assert WatsonxPassthroughConfig.deployment_model_matches(creds, "ibm/granite-8b") is False + + +# --------------------------------------------------------------------------- +# resolve_credentials_from_deployment +# --------------------------------------------------------------------------- + +class TestResolveCredentialsFromDeployment: + """Tests for WatsonxPassthroughConfig.resolve_credentials_from_deployment.""" + + def test_returns_none_when_not_flagged(self): + deployment = _wx_unflagged_deployment( + "watsonx/ibm/granite-13b-chat-v2", + api_key="sk-test", + api_base="https://us-south.ml.cloud.ibm.com", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is None + + def test_returns_none_for_non_watsonx_provider(self): + deployment = _wx_deployment("openai/gpt-4o", api_key="sk-openai") + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is None + + def test_returns_none_when_model_missing(self): + params = {"use_in_pass_through": True} + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment(params) + assert result is None + + def test_returns_none_when_get_llm_provider_raises_bad_request(self): + """get_llm_provider raises BadRequestError for unrecognisable model strings — must return None.""" + from litellm.exceptions import BadRequestError + + params = {"use_in_pass_through": True, "model": "completely-unknown/some-model"} + with patch( + "litellm.llms.watsonx.passthrough.transformation.get_llm_provider", + side_effect=BadRequestError( + message="LLM Provider NOT provided.", model="completely-unknown/some-model", llm_provider="" + ), + ): + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment(params) + assert result is None + + def test_inline_api_key_and_api_base(self): + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + api_key="ibmkey-123", + api_base="https://us-south.ml.cloud.ibm.com", + project_id="proj-abc", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("api_key") == "ibmkey-123" + assert result.get("api_base") == "https://us-south.ml.cloud.ibm.com" + assert result.get("project_id") == "proj-abc" + assert result.get("space_id") is None + assert result.get("model") == "watsonx/ibm/granite-13b-chat-v2" + + def test_space_id_is_read(self): + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + api_key="ibmkey-123", + api_base="https://us-south.ml.cloud.ibm.com", + space_id="space-xyz", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("space_id") == "space-xyz" + assert result.get("project_id") is None + + def test_api_key_none_when_omitted(self): + """api_key is optional — deployment with only api_base is still eligible.""" + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + api_base="https://us-south.ml.cloud.ibm.com", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("api_key") is None + assert result.get("api_base") == "https://us-south.ml.cloud.ibm.com" + + def test_api_base_none_when_omitted(self): + """api_base is optional — deployment with only api_key is still eligible.""" + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + api_key="ibmkey-123", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("api_key") == "ibmkey-123" + assert result.get("api_base") is None + + def test_named_credential_overrides_inline(self, monkeypatch): + """litellm_credential_name values take precedence over inline litellm_params.""" + monkeypatch.setattr(litellm, "credential_list", []) + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.types.utils import CredentialItem + + CredentialAccessor.upsert_credentials([ + CredentialItem( + credential_name="wx-cred", + credential_values={ + "api_key": "cred-key-999", + "api_base": "https://eu-de.ml.cloud.ibm.com", + }, + credential_info={}, + ) + ]) + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + litellm_credential_name="wx-cred", + api_key="inline-key-should-not-win", + api_base="https://us-south.ml.cloud.ibm.com", + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("api_key") == "cred-key-999" + assert result.get("api_base") == "https://eu-de.ml.cloud.ibm.com" + + def test_named_credential_loaded_after_deployment_still_resolves(self, monkeypatch): + """Credential added after router starts is found on next call.""" + monkeypatch.setattr(litellm, "credential_list", []) + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.types.utils import CredentialItem + + deployment = _wx_deployment( + "watsonx/ibm/granite-13b-chat-v2", + litellm_credential_name="wx-late", + ) + # First call — credential not registered yet + result = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result is not None + assert result.get("api_key") is None + + # Register credential after the fact + CredentialAccessor.upsert_credentials([ + CredentialItem( + credential_name="wx-late", + credential_values={"api_key": "late-key", "api_base": "https://us-south.ml.cloud.ibm.com"}, + credential_info={}, + ) + ]) + + result2 = WatsonxPassthroughConfig.resolve_credentials_from_deployment( + deployment["litellm_params"] + ) + assert result2 is not None + assert result2.get("api_key") == "late-key" + + +# --------------------------------------------------------------------------- +# resolve_credentials_from_router_deployments +# --------------------------------------------------------------------------- + +class TestResolveCredentialsFromRouterDeployments: + """Tests for WatsonxPassthroughConfig.resolve_credentials_from_router_deployments.""" + + def _deployments(self, *models_and_params): + """Build a list of deployments from (model, **kwargs) tuples.""" + return [_wx_deployment(m, **p) for m, p in models_and_params] + + def test_returns_none_with_no_deployments(self): + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=[] + ) + assert result is None + + def test_returns_none_with_no_eligible_deployments(self): + deployments = [_wx_unflagged_deployment("watsonx/ibm/granite-13b-chat-v2", api_key="k")] + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=deployments + ) + assert result is None + + def test_single_deployment_returned_as_fallback(self): + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k1", "api_base": "https://us-south.ml.cloud.ibm.com"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=deployments + ) + assert result is not None + assert result.get("api_key") == "k1" + + def test_exact_match_model_and_project_id(self): + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-us", "api_base": "https://us-south.ml.cloud.ibm.com", "project_id": "proj-us"}), + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-eu", "api_base": "https://eu-de.ml.cloud.ibm.com", "project_id": "proj-eu"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + model_id="ibm/granite-13b-chat-v2", + project_id="proj-eu", + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-eu" + assert result.get("api_base") == "https://eu-de.ml.cloud.ibm.com" + + def test_exact_match_model_and_space_id(self): + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-studio", "api_base": "https://us-south.ml.cloud.ibm.com", "project_id": "proj-studio"}), + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-deploy", "api_base": "https://us-south.ml.cloud.ibm.com", "space_id": "space-prod"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + model_id="ibm/granite-13b-chat-v2", + space_id="space-prod", + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-deploy" + + def test_partial_match_model_only(self): + deployments = self._deployments( + ("watsonx/ibm/granite-8b-code", {"api_key": "k-code", "api_base": "https://us-south.ml.cloud.ibm.com"}), + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-chat", "api_base": "https://us-south.ml.cloud.ibm.com"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + model_id="ibm/granite-13b-chat-v2", + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-chat" + + def test_fallback_first_when_no_hints(self, recwarn): + """With multiple deployments and no hints, returns first and emits a warning.""" + deployments = self._deployments( + ("watsonx/ibm/granite-8b-code", {"api_key": "k-first", "api_base": "https://us-south.ml.cloud.ibm.com"}), + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-second", "api_base": "https://us-south.ml.cloud.ibm.com"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-first" + + def test_non_watsonx_deployments_ignored(self): + deployments = [ + _wx_deployment("openai/gpt-4o", api_key="openai-key"), + _wx_deployment("watsonx/ibm/granite-13b-chat-v2", api_key="wx-key", api_base="https://us-south.ml.cloud.ibm.com"), + ] + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "wx-key" + + def test_unflagged_deployments_ignored(self): + deployments = [ + _wx_unflagged_deployment("watsonx/ibm/granite-13b-chat-v2", api_key="should-be-ignored"), + _wx_deployment("watsonx/ibm/granite-8b-code", api_key="eligible", api_base="https://us-south.ml.cloud.ibm.com"), + ] + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "eligible" + + def test_project_id_mismatch_falls_through_to_model_match(self): + """No project_id match → falls through to model-only partial match.""" + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k1", "api_base": "https://us-south.ml.cloud.ibm.com", "project_id": "proj-a"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + model_id="ibm/granite-13b-chat-v2", + project_id="proj-unknown", + deployments=deployments, + ) + # Falls through to partial match on model_id + assert result is not None + assert result.get("api_key") == "k1" + + def test_match_project_id_only_without_model_id(self): + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-proj-a", "api_base": "https://us-south.ml.cloud.ibm.com", "project_id": "proj-a"}), + ("watsonx/ibm/granite-8b-code", {"api_key": "k-proj-b", "api_base": "https://us-south.ml.cloud.ibm.com", "project_id": "proj-b"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + project_id="proj-b", + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-proj-b" + + def test_match_space_id_only_without_model_id(self): + deployments = self._deployments( + ("watsonx/ibm/granite-13b-chat-v2", {"api_key": "k-space-a", "api_base": "https://us-south.ml.cloud.ibm.com", "space_id": "space-a"}), + ("watsonx/ibm/granite-8b-code", {"api_key": "k-space-b", "api_base": "https://us-south.ml.cloud.ibm.com", "space_id": "space-b"}), + ) + result = WatsonxPassthroughConfig.resolve_credentials_from_router_deployments( + space_id="space-b", + deployments=deployments, + ) + assert result is not None + assert result.get("api_key") == "k-space-b" diff --git a/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py b/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py index 5500bb0aad9..06b36999a11 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py +++ b/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py @@ -5,13 +5,11 @@ Tests the Watsonx pass-through endpoint that handles automatic IAM token managem and version parameter injection. """ -import json -from unittest.mock import AsyncMock, MagicMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, Request, Response - import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -437,3 +435,180 @@ class TestWatsonxProxyRoute: call_args["target"] == f"https://us-south.ml.cloud.ibm.com/{endpoint_path}" ) + + @pytest.mark.asyncio + async def test_watsonx_proxy_route_model_restriction_denied(self): + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.query_params = {} + mock_request.headers = {"content-type": "application/json"} + mock_request.json = AsyncMock(return_value={"model_id": "ibm/granite-13b-chat-v2", "input": "test"}) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.models = ["openai/gpt-4o"] + mock_user_api_key_dict.team_models = None + mock_user_api_key_dict.access_group_ids = None + + mock_provider_config = MagicMock() + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "watsonx-chat", + "litellm_params": { + "model": "watsonx/ibm/granite-13b-chat-v2", + "use_in_pass_through": True, + "api_key": "wx-key", + "api_base": "https://us-south.ml.cloud.ibm.com", + }, + } + ] + + from litellm.proxy._types import ModelAccessDeniedProxyException + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_passthrough_config", + return_value=mock_provider_config, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch( + "litellm.proxy.utils.get_available_models_for_user", + AsyncMock(return_value=["openai/gpt-4o"]), + ), + ): + with pytest.raises(ModelAccessDeniedProxyException) as exc_info: + await watsonx_proxy_route( + endpoint="ml/v1/text/generation", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + assert exc_info.value.code == "403" + assert exc_info.value.type == "key_model_access_denied" + + @pytest.mark.asyncio + async def test_watsonx_proxy_route_model_restriction_allowed(self): + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.query_params = {} + mock_request.headers = {"content-type": "application/json"} + mock_request.json = AsyncMock(return_value={"model_id": "ibm/granite-13b-chat-v2", "input": "test"}) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.models = ["watsonx-chat"] + mock_user_api_key_dict.team_models = None + mock_user_api_key_dict.access_group_ids = None + + mock_provider_config = MagicMock() + mock_provider_config.get_complete_url.return_value = ( + "https://us-south.ml.cloud.ibm.com/ml/v1/text/generation", + {}, + ) + mock_provider_config.validate_environment.return_value = {"Authorization": "Bearer test-iam-token"} + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "watsonx-chat", + "litellm_params": { + "model": "watsonx/ibm/granite-13b-chat-v2", + "use_in_pass_through": True, + "api_key": "wx-key", + "api_base": "https://us-south.ml.cloud.ibm.com", + }, + } + ] + mock_endpoint_func = AsyncMock(return_value={"results": []}) + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_passthrough_config", + return_value=mock_provider_config, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch( + "litellm.proxy.utils.get_available_models_for_user", + AsyncMock(return_value=["watsonx-chat"]), + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=mock_endpoint_func, + ), + ): + result = await watsonx_proxy_route( + endpoint="ml/v1/text/generation", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + assert result == {"results": []} + + @pytest.mark.asyncio + async def test_watsonx_proxy_route_match_project_id_from_query_params(self): + mock_request = MagicMock(spec=Request) + mock_request.method = "GET" + mock_request.query_params = {"project_id": "proj-b"} + mock_request.headers = {} + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.models = None + mock_user_api_key_dict.team_models = None + mock_user_api_key_dict.access_group_ids = None + + mock_provider_config = MagicMock() + mock_provider_config.get_complete_url.return_value = ( + "https://us-south.ml.cloud.ibm.com/ml/v1/deployments", + {}, + ) + mock_provider_config.validate_environment.return_value = {"Authorization": "Bearer test-iam-token"} + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "watsonx-chat-a", + "litellm_params": { + "model": "watsonx/ibm/granite-13b-chat-v2", + "use_in_pass_through": True, + "api_key": "wx-key-a", + "project_id": "proj-a", + "api_base": "https://us-south.ml.cloud.ibm.com", + }, + }, + { + "model_name": "watsonx-chat-b", + "litellm_params": { + "model": "watsonx/ibm/granite-8b-code", + "use_in_pass_through": True, + "api_key": "wx-key-b", + "project_id": "proj-b", + "api_base": "https://us-south.ml.cloud.ibm.com", + }, + }, + ] + mock_endpoint_func = AsyncMock(return_value={"resources": []}) + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_passthrough_config", + return_value=mock_provider_config, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch( + "litellm.proxy.utils.get_available_models_for_user", + AsyncMock(return_value=["watsonx-chat-a", "watsonx-chat-b"]), + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=mock_endpoint_func, + ), + ): + result = await watsonx_proxy_route( + endpoint="ml/v1/deployments", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + assert result == {"resources": []} + mock_provider_config.validate_environment.assert_called_once() + call_kwargs = mock_provider_config.validate_environment.call_args[1] + assert call_kwargs["api_key"] == "wx-key-b"