diff --git a/litellm/llms/nvidia_nim/passthrough/transformation.py b/litellm/llms/nvidia_nim/passthrough/transformation.py index 6df9681cc81..7de1ce4d631 100644 --- a/litellm/llms/nvidia_nim/passthrough/transformation.py +++ b/litellm/llms/nvidia_nim/passthrough/transformation.py @@ -27,14 +27,26 @@ if TYPE_CHECKING: API_VERSION_SEGMENT: Final = re.compile(r"^v\d+$") NVIDIA_NIM_MODEL_PREFIX: Final = f"{LlmProviders.NVIDIA_NIM.value}/" +NVIDIA_NIM_ROUTE_PREFIX: Final = re.compile(rf"^/{LlmProviders.NVIDIA_NIM.value}/", re.IGNORECASE) + + +def is_nvidia_nim_deployment(deployment: DeploymentTypedDict) -> bool: + litellm_params: Final = deployment["litellm_params"] + return litellm_params.get("custom_llm_provider") == LlmProviders.NVIDIA_NIM.value or litellm_params.get( + "model", "" + ).startswith(NVIDIA_NIM_MODEL_PREFIX) def nvidia_nim_model_groups(deployments: Iterable[DeploymentTypedDict] | None) -> frozenset[str]: - return frozenset( - deployment["model_name"] - for deployment in deployments or () - if deployment["litellm_params"].get("custom_llm_provider") == LlmProviders.NVIDIA_NIM.value - or deployment["litellm_params"].get("model", "").startswith(NVIDIA_NIM_MODEL_PREFIX) + listed: Final = tuple(deployments or ()) + nim_groups: Final = frozenset(d["model_name"] for d in listed if is_nvidia_nim_deployment(d)) + other_groups: Final = frozenset(d["model_name"] for d in listed if not is_nvidia_nim_deployment(d)) + return nim_groups - other_groups + + +def nvidia_nim_model_group_in_path(path: str, deployments: Iterable[DeploymentTypedDict] | None) -> str | None: + return nvidia_nim_router_model_in_endpoint( + NVIDIA_NIM_ROUTE_PREFIX.sub("", path), nvidia_nim_model_groups(deployments) ) @@ -79,7 +91,7 @@ class NvidiaNimPassthroughConfig(BasePassthroughConfig): base_target_url: Final = self.get_api_base(api_base) if base_target_url is None: raise ValueError("NVIDIA NIM api base not found: set `api_base` on the deployment or NVIDIA_NIM_API_BASE") - native_endpoint: Final = strip_leading_model_segment(endpoint, (model_group_from(litellm_params), model)) + native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) root: Final = without_repeated_version_prefix(base_target_url, native_endpoint) return (self.format_url(native_endpoint, root, request_query_params), root) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index dfce3ab206b..e69b0c98ee0 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -28,10 +28,7 @@ from litellm.litellm_core_utils.url_utils import ( validate_url, ) from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint -from litellm.llms.nvidia_nim.passthrough.transformation import ( - nvidia_nim_model_groups, - nvidia_nim_router_model_in_endpoint, -) +from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2045,19 +2042,14 @@ def get_model_from_request( return model if azure_model is None else azure_model if route.lower().startswith("/nvidia_nim/"): - nvidia_nim_model: Final = _router_model_from_nvidia_nim_route(route, llm_router) + nvidia_nim_model: Final = ( + nvidia_nim_model_group_in_path(route, llm_router.get_model_list()) if llm_router else None + ) return model if nvidia_nim_model is None else nvidia_nim_model return model -def _router_model_from_nvidia_nim_route(route: str, llm_router: Router | None) -> str | None: - if llm_router is None: - return None - endpoint: Final = re.sub(r"^/nvidia_nim/", "", route, flags=re.IGNORECASE) - return nvidia_nim_router_model_in_endpoint(endpoint, nvidia_nim_model_groups(llm_router.get_model_list())) - - def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: if llm_router is None: return None diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 65a18c1603e..4fc610f068f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -36,10 +36,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.llms.nvidia_nim.passthrough.transformation import ( - nvidia_nim_model_groups, - nvidia_nim_router_model_in_endpoint, -) +from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse from litellm.proxy._types import * @@ -1654,15 +1651,11 @@ async def relay_nvidia_nim_request( request_body: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, ) -> Response: - model_group: Final = ( - nvidia_nim_router_model_in_endpoint(endpoint, nvidia_nim_model_groups(llm_router.get_model_list())) - if llm_router - else None - ) + model_group: Final = nvidia_nim_model_group_in_path(endpoint, llm_router.get_model_list()) if llm_router else None if llm_router is None or model_group is None: rejection: Final[RelayRejection] = { "error": "no NVIDIA NIM model group in the path; call /nvidia_nim/{model_group}/v1/infer with a model " - "from your `model_list` whose `model` starts with `nvidia_nim/`" + "group from your `model_list` whose deployments all use `nvidia_nim/` models" } raise HTTPException(status_code=400, detail=rejection) diff --git a/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py b/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py index 87de4a795e0..906d6c2b614 100644 --- a/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py +++ b/tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py @@ -8,6 +8,8 @@ import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.nvidia_nim.passthrough.transformation import ( NvidiaNimPassthroughConfig, + nvidia_nim_model_group_in_path, + nvidia_nim_model_groups, nvidia_nim_router_model_in_endpoint, ) from litellm.types.utils import LlmProviders @@ -53,6 +55,12 @@ def test_provider_config_manager_resolves_nvidia_nim_passthrough_config(): (f"{NIM_BASE}/v2", "v1/infer", {}, f"{NIM_BASE}/v2/v1/infer"), (f"{NIM_BASE}/infer", "infer", {}, f"{NIM_BASE}/infer/infer"), (NIM_BASE, "nvidia/nemoretriever-page-elements-v2/v1/infer", {}, f"{NIM_BASE}/v1/infer"), + ( + NIM_BASE, + "nvidia/nemoretriever-page-elements-v2/v1/infer", + {"litellm_metadata": {"model_group": "nvidia"}}, + f"{NIM_BASE}/v1/infer", + ), ], ) def test_relay_url_strips_the_model_group_and_never_doubles_the_api_version( @@ -161,6 +169,45 @@ def test_router_model_in_endpoint_takes_the_longest_leading_model_group(endpoint assert nvidia_nim_router_model_in_endpoint(endpoint, frozenset(router_models)) == expected +def _deployment(model_name: str, model: str, custom_llm_provider: str | None = None): + litellm_params = ( + {"model": model} + if custom_llm_provider is None + else {"model": model, "custom_llm_provider": custom_llm_provider} + ) + return {"model_name": model_name, "litellm_params": litellm_params} + + +MIXED_DEPLOYMENTS = ( + _deployment("nim-page", "nvidia_nim/nvidia/nemoretriever-page-elements-v2"), + _deployment("nim-table", "nvidia/nemoretriever-table-structure-v1", custom_llm_provider="nvidia_nim"), + _deployment("mixed", "nvidia_nim/nvidia/nemoretriever-page-elements-v2"), + _deployment("mixed", "openai/gpt-4o"), + _deployment("gpt-4o", "openai/gpt-4o"), +) + + +def test_model_groups_only_admit_groups_whose_every_deployment_is_nim_backed(): + assert nvidia_nim_model_groups(MIXED_DEPLOYMENTS) == frozenset({"nim-page", "nim-table"}) + assert nvidia_nim_model_groups(None) == frozenset() + + +@pytest.mark.parametrize( + "path, expected", + [ + ("/nvidia_nim/nim-page/v1/infer", "nim-page"), + ("/NVIDIA_NIM/nim-table/v1/infer", "nim-table"), + ("nim-page/v1/infer", "nim-page"), + ("/nvidia_nim/mixed/v1/infer", None), + ("mixed/v1/infer", None), + ("/nvidia_nim/gpt-4o/v1/infer", None), + ("/nvidia_nim/v1/infer", None), + ], +) +def test_model_group_in_path_resolves_the_same_nim_only_groups_for_routes_and_endpoints(path, expected): + assert nvidia_nim_model_group_in_path(path, MIXED_DEPLOYMENTS) == expected + + @pytest.mark.parametrize("request_data, expected", [({"stream": True}, True), ({"stream": False}, False), ({}, False)]) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): assert NvidiaNimPassthroughConfig().is_streaming_request("v1/infer", request_data) is expected @@ -214,3 +261,36 @@ async def test_object_detection_relay_sends_the_native_body_unchanged_to_v1_infe assert response.status_code == 200 assert response.headers["x-nim"] == "1" assert response.json() == {"data": [{"index": 0}, {"index": 1}]} + + +@pytest.mark.asyncio +async def test_router_relay_reaches_v1_infer_when_the_group_name_is_a_leading_segment_of_the_model_id(): + upstream_requests: list[httpx.Request] = [] + + def nim(request: httpx.Request) -> httpx.Response: + upstream_requests.append(request) + return httpx.Response(200, json={"data": [{"index": 0}]}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(nim)) + router = litellm.Router( + model_list=[ + { + "model_name": "nvidia", + "litellm_params": { + "model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2", + "api_base": NIM_BASE, + "api_key": "nvapi-secret", + }, + } + ] + ) + + response = await router.allm_passthrough_route( + model="nvidia", endpoint="nvidia/v1/infer", method="POST", json=dict(INFER_BODY), client=client + ) + + (sent,) = upstream_requests + assert str(sent.url) == f"{NIM_BASE}/v1/infer" + assert json.loads(sent.content) == INFER_BODY + assert response.status_code == 200 diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 0841f82227a..8e592b61f1c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -848,6 +848,18 @@ def _nvidia_nim_relay_router(): "model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, }, + { + "model_name": "detect", + "litellm_params": { + "model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2", + "api_base": "http://nim-a.internal:8000", + "api_key": "k", + }, + }, + { + "model_name": "detect", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + }, ] ) @@ -873,6 +885,7 @@ NIM_INFER_BODY = {"input": [{"type": "image_url", "url": "data:image/png;base64, ("/nvidia_nim/unknown-group/v1/infer", NIM_INFER_BODY, None), ("/nvidia_nim/nim-page-elements-v2/v1/infer", NIM_INFER_BODY, None), ("/nvidia_nim/gpt-4o/v1/infer", NIM_INFER_BODY, None), + ("/nvidia_nim/detect/v1/infer", NIM_INFER_BODY, None), ], ) def test_get_model_from_request_nvidia_nim_relay_routes_use_the_model_group_in_the_path(route, request_data, expected): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4e73a33ecd9..62ecbbbcf51 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5498,6 +5498,32 @@ class TestNvidiaNimProxyRoute: assert exc_info.value.status_code == 400 assert captured == [] + @pytest.mark.asyncio + async def test_a_group_mixing_nim_and_other_deployments_is_rejected_before_any_upstream_call(self): + from fastapi import HTTPException + + captured: list[dict] = [] + + class MixedRouter: + def get_model_list(self): + return [ + { + "model_name": "detect", + "litellm_params": {"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"}, + }, + {"model_name": "detect", "litellm_params": {"model": "openai/gpt-4o"}}, + ] + + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"data": []}) + + with pytest.raises(HTTPException) as exc_info: + await self._relay(MixedRouter(), "detect/v1/infer", NIM_INFER_BODY) + + assert exc_info.value.status_code == 400 + assert captured == [] + @pytest.mark.asyncio async def test_no_router_is_rejected_before_any_upstream_call(self): from fastapi import HTTPException