diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 151938d73d7..fa09b229564 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/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_key=None, + api_base=api_base, + 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_base=None, + api_key=api_key, + api_base=api_base, ) # Check for streaming diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 3b49ca1450a..e793dac4d06 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/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( diff --git a/tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py b/tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py new file mode 100644 index 00000000000..f20e6d35cf9 --- /dev/null +++ b/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 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..e0a9bda0967 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 @@ -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"