mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
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>
This commit is contained in:
Kerry Lu
GitHub
parent
ee39de7107
commit
948445c828
4 changed files with 217 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue