fix(watsonx): read pass-through credentials from use_in_pass_through deployments (#49462)

* fix(watsonx): read passthrough credentials from use_in_pass_through deployments

/watsonx/{endpoint} always called the provider config with api_key=None and api_base=None, so it only ever
worked from WATSONX_API_KEY / WATSONX_API_BASE. It now asks PassthroughEndpointRouter for the proxy-wide
watsonx deployment flagged use_in_pass_through, like the other passthrough routes, and keeps the env
fallback when none is flagged.

Co-authored-by: T K Chandra Hasan <t.k.chandra.hasan@ibm.com>

* fix(watsonx): keep env key precedence when no deployment is flagged

Resolve only use_in_pass_through deployments in the route. get_proxy_wide_credentials also falls back to
WATSONX_API_KEY, which put it ahead of WX_API_KEY in generate_iam_token

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: T K Chandra Hasan <t.k.chandra.hasan@ibm.com>

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -4474,10 +4474,16 @@ async def watsonx_proxy_route(
     if provider_config is None:
         raise HTTPException(status_code=404, detail="Watsonx passthrough config not found")
 
+    watsonx_credentials: Final = passthrough_endpoint_router.get_proxy_wide_deployment_credentials(
+        custom_llm_provider="watsonx"
+    )
+    api_key: Final = watsonx_credentials.api_key if watsonx_credentials is not None else None
+    api_base: Final = watsonx_credentials.api_base if watsonx_credentials is not None else None
+
     # Get complete URL with version parameter
     complete_url, _ = provider_config.get_complete_url(
-        api_base=None,
+        api_base=api_base,
-        api_key=None,
+        api_key=api_key,
         model="",
         endpoint=endpoint,
         request_query_params=None,
@@ -4491,8 +4497,8 @@ async def watsonx_proxy_route(
         messages=[],
         optional_params={},
         litellm_params={},
-        api_key=None,
+        api_key=api_key,
-        api_base=None,
+        api_base=api_base,
     )
 
     # Check for streaming

litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py

@@ -132,6 +132,19 @@ class PassthroughEndpointRouter:
             api_base=deployment_credentials.api_base if deployment_credentials is not None else None,
         )
 
+    def get_proxy_wide_deployment_credentials(self, custom_llm_provider: str) -> _PassThroughCredentials | None:
+        deployment_credentials: Final = self._get_deployment_credentials(
+            custom_llm_provider=custom_llm_provider,
+            region_name=None,
+            include_team_deployments=False,
+        )
+        if deployment_credentials is None:
+            return None
+        api_key: Final = _resolve_oidc_reference(deployment_credentials.api_key)
+        if not api_key:
+            return None
+        return _PassThroughCredentials(api_key=api_key, api_base=deployment_credentials.api_base)
+
     def has_credentials(self, custom_llm_provider: str) -> bool:
         if (
             self._get_deployment_credentials(

tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py

@@ -0,0 +1,81 @@
+import json
+import uuid
+from pathlib import Path
+from typing import Final
+
+import yaml
+from integration._support.client import Gateway, JsonValue, object_value
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, wire_server
+
+_DEPLOYMENT_KEY: Final = "synthetic-watsonx-deployment-key"
+_IAM_TOKEN: Final = "synthetic-watsonx-iam-token"
+_WATSONX_ENV: Final = (
+    "WATSONX_API_KEY",
+    "WATSONX_APIKEY",
+    "WX_API_KEY",
+    "WATSONX_API_BASE",
+    "WATSONX_URL",
+    "WX_URL",
+    "WML_URL",
+    "WATSONX_TOKEN",
+    "WATSONX_ZENAPIKEY",
+)
+
+
+def _config(directory: Path, api_base: str) -> Path:
+    stock: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    config: Final[dict[str, JsonValue]] = {
+        **stock,
+        "model_list": [
+            {
+                "model_name": "granite",
+                "litellm_params": {
+                    "model": "watsonx/ibm/granite-3-8b-instruct",
+                    "api_key": _DEPLOYMENT_KEY,
+                    "api_base": api_base,
+                    "project_id": "synthetic-project",
+                    "use_in_pass_through": True,
+                },
+            }
+        ],
+    }
+    path: Final = directory / "watsonx-passthrough.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+def test_watsonx_passthrough_uses_the_flagged_deployment_key_and_base(gateway: Gateway, tmp_path: Path) -> None:
+    marker: Final = "watsonx-" + uuid.uuid4().hex
+
+    def respond(request: Request) -> Reply:
+        if request.method == "POST" and request.target == "/identity/token":
+            assert f"apikey={_DEPLOYMENT_KEY}".encode() in request.body, request.body
+            return Reply(body=json.dumps({"access_token": _IAM_TOKEN, "expires_in": 3600}).encode())
+        if request.method == "POST" and request.target.startswith("/ml/v1/text/generation?"):
+            assert request.headers["authorization"] == f"Bearer {_IAM_TOKEN}", request.headers
+            assert marker.encode() in request.body, request.body[:400]
+            return Reply(
+                body=json.dumps(
+                    {"model_id": "ibm/granite-3-8b-instruct", "results": [{"generated_text": marker}]}
+                ).encode()
+            )
+        return Reply(status=404)
+
+    with wire_server(respond) as wire:
+        with owned_proxy(
+            gateway,
+            tmp_path,
+            {"WATSONX_IAM_URL": f"{wire.url}/identity/token"},
+            config=_config(tmp_path, wire.url),
+            remove_environment=_WATSONX_ENV,
+        ) as candidate:
+            response: Final = candidate.request(
+                "POST",
+                "/watsonx/ml/v1/text/generation",
+                {"model_id": "ibm/granite-3-8b-instruct", "input": marker, "project_id": "synthetic-project"},
+            )
+            assert response.status_code == 200, response.text
+            assert response.json()["results"][0]["generated_text"] == marker
+        forwarded: Final = tuple((request.method, request.target.split("?")[0]) for request in wire.drain())
+        assert forwarded == (("POST", "/identity/token"), ("POST", "/ml/v1/text/generation")), forwarded

tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py

@@ -437,3 +437,116 @@ class TestWatsonxProxyRoute:
                     call_args["target"]
                     == f"https://us-south.ml.cloud.ibm.com/{endpoint_path}"
                 )
+
+
+_WATSONX_ENV = (
+    "WATSONX_API_KEY",
+    "WATSONX_APIKEY",
+    "WX_API_KEY",
+    "WATSONX_API_BASE",
+    "WATSONX_URL",
+    "WX_URL",
+    "WML_URL",
+    "WATSONX_TOKEN",
+    "WATSONX_ZENAPIKEY",
+)
+
+
+def _watsonx_deployment(api_key: str, api_base: str, use_in_pass_through: bool) -> dict:
+    return {
+        "model_name": "granite",
+        "litellm_params": {
+            "model": "watsonx/ibm/granite-3-8b-instruct",
+            "api_key": api_key,
+            "api_base": api_base,
+            "project_id": "proj",
+            "use_in_pass_through": use_in_pass_through,
+        },
+    }
+
+
+def _post_request(body: dict) -> MagicMock:
+    request = MagicMock(spec=Request)
+    request.method = "POST"
+    request.query_params = {}
+    request.headers = {"content-type": "application/json"}
+    request.json = AsyncMock(return_value=body)
+    return request
+
+
+def _iam_response(token: str) -> MagicMock:
+    response = MagicMock()
+    response.json.return_value = {"access_token": token, "expires_in": 3600}
+    return response
+
+
+async def _route_target_and_iam_apikey(
+    monkeypatch, llm_router: litellm.Router, env: dict[str, str]
+) -> tuple[str, str, str]:
+    from litellm.llms.watsonx.common_utils import iam_token_cache
+    from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import PassthroughEndpointRouter
+
+    for name in _WATSONX_ENV:
+        monkeypatch.delenv(name, raising=False)
+    for name, value in env.items():
+        monkeypatch.setenv(name, value)
+    monkeypatch.setattr(litellm, "credential_list", [])
+    iam_token_cache.flush_cache()
+    monkeypatch.setattr(
+        "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router",
+        PassthroughEndpointRouter(llm_router_getter=lambda: llm_router),
+    )
+    endpoint_func = AsyncMock(return_value={"results": []})
+    with (
+        patch.object(litellm.module_level_client, "post", return_value=_iam_response("iam-token-1")) as iam_post,
+        patch(
+            "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
+            return_value=endpoint_func,
+        ) as create_route,
+    ):
+        await watsonx_proxy_route(
+            endpoint="ml/v1/text/generation",
+            request=_post_request({"input": "hi", "project_id": "proj"}),
+            fastapi_response=MagicMock(spec=Response),
+            user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
+        )
+    route_kwargs = create_route.call_args.kwargs
+    return (
+        route_kwargs["target"],
+        route_kwargs["custom_headers"]["Authorization"],
+        iam_post.call_args.kwargs["data"]["apikey"],
+    )
+
+
+@pytest.mark.asyncio
+async def test_watsonx_route_uses_the_use_in_pass_through_deployment_key_and_base(monkeypatch):
+    llm_router = litellm.Router(
+        model_list=[_watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=True)]
+    )
+
+    target, authorization, iam_apikey = await _route_target_and_iam_apikey(monkeypatch, llm_router, env={})
+
+    assert target == "https://eu-de.ml.cloud.ibm.com/ml/v1/text/generation"
+    assert authorization == "Bearer iam-token-1"
+    assert iam_apikey == "wx-deployment-key"
+
+
+@pytest.mark.asyncio
+async def test_watsonx_route_falls_back_to_env_when_no_deployment_is_flagged(monkeypatch):
+    llm_router = litellm.Router(
+        model_list=[_watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=False)]
+    )
+
+    target, authorization, iam_apikey = await _route_target_and_iam_apikey(
+        monkeypatch,
+        llm_router,
+        env={
+            "WX_API_KEY": "wx-env-key",
+            "WATSONX_API_KEY": "watsonx-env-key",
+            "WATSONX_API_BASE": "https://us-south.ml.cloud.ibm.com",
+        },
+    )
+
+    assert target == "https://us-south.ml.cloud.ibm.com/ml/v1/text/generation"
+    assert authorization == "Bearer iam-token-1"
+    assert iam_apikey == "wx-env-key"