mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(watsonx): read pass-through credentials from use_in_pass_through deployments (#49462)
Backport of #49462 to lts/1.108.x.
Cherry-picked from 948445c828 (main). get_proxy_wide_deployment_credentials hand-ported onto the line's _resolve_matching_deployment_api_key, since the line has no _get_deployment_credentials.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
kerry
parent
7dffccccae
commit
de425fb956
4 changed files with 244 additions and 4 deletions
|
|
@ -4259,10 +4259,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,
|
||||
|
|
@ -4276,8 +4282,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
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -57,6 +58,12 @@ def _resolve_oidc_reference(credential: str | None) -> str | None:
|
|||
) from exc
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PassThroughCredentials:
|
||||
api_key: str
|
||||
api_base: str | None
|
||||
|
||||
|
||||
class PassthroughEndpointRouter:
|
||||
"""
|
||||
Use this class to Get credentials for pass-through endpoints
|
||||
|
|
@ -89,6 +96,35 @@ class PassthroughEndpointRouter:
|
|||
)
|
||||
return _resolve_oidc_reference(get_secret_str(_env_variable_name))
|
||||
|
||||
def get_proxy_wide_deployment_credentials(self, custom_llm_provider: str) -> _PassThroughCredentials | None:
|
||||
llm_router: Final = self.llm_router_getter()
|
||||
if llm_router is None:
|
||||
return None
|
||||
for deployment in llm_router.get_model_list() or ():
|
||||
litellm_params: LiteLLMParamsTypedDict = deployment["litellm_params"]
|
||||
api_key: str | None = self._resolve_matching_deployment_api_key(
|
||||
litellm_params=litellm_params,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
region_name=None,
|
||||
)
|
||||
if api_key is None:
|
||||
continue
|
||||
resolved_api_key: str | None = _resolve_oidc_reference(api_key)
|
||||
if not resolved_api_key:
|
||||
return None
|
||||
return _PassThroughCredentials(
|
||||
api_key=resolved_api_key,
|
||||
api_base=self._resolve_matching_deployment_api_base(litellm_params),
|
||||
)
|
||||
return None
|
||||
|
||||
def _resolve_matching_deployment_api_base(self, litellm_params: LiteLLMParamsTypedDict) -> str | None:
|
||||
credential_name: Final = litellm_params.get("litellm_credential_name")
|
||||
credential_values: Final = (
|
||||
CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else None
|
||||
)
|
||||
return _get_str_value(credential_values, "api_base") or litellm_params.get("api_base")
|
||||
|
||||
def _get_deployment_api_key(
|
||||
self,
|
||||
custom_llm_provider: str,
|
||||
|
|
|
|||
|
|
@ -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,120 @@ 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