scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 69

fix(proxy): reject mixed NIM model groups and strip the deployment model before the group in /nvidia_nim URLs

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-15 23:34:22 +00:00
parent b50a22b370
commit 771b2509d1
6 changed files with 144 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/llms/nvidia_nim/passthrough/transformation.py
View file

@ -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)

16
litellm/proxy/auth/auth_utils.py
View file

@ -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

13
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -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)

80
tests/test_litellm/llms/nvidia_nim/passthrough/test_nvidia_nim_passthrough_transformation.py
View file

@ -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

13
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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):

26
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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