Merge pull request #41316 from BerriAI/litellm_nvidia_nim_infer_passthrough

feat(proxy): add /nvidia_nim passthrough route for NIM object detection and OCR /v1/infer
This commit is contained in:
Yassin Kortam 2026-09-15 17:28:25 -07:00 committed by GitHub
commit 7eeba69016
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1247 additions and 23 deletions

View file

@ -96,6 +96,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/langfuse/",
"/vllm/",
"/mistral/",
"/nvidia_nim/",
"/groq/",
"/voyage/",
"/cursor/",

View file

@ -5,7 +5,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.llms.azure_ai.common_utils import (
@ -18,6 +18,8 @@ from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
RelayShape,
logged_relay_shape,
model_group_from,
relayed_body,
strip_leading_model_segment,
)
from litellm.types.llms.openai import AllMessageValues
@ -35,19 +37,6 @@ if TYPE_CHECKING:
EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
class PassthroughMetadata(BaseModel):
model_config = ConfigDict(extra="ignore")
model_group: str = ""
def model_group_from(litellm_params: Mapping[str, object]) -> str:
try:
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
except ValidationError:
return ""
def api_version_from(litellm_params: Mapping[str, object]) -> str | None:
try:
return TypeAdapter(str | None).validate_python(litellm_params.get("api_version"))
@ -96,14 +85,6 @@ def relay_query_params(
return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version})
def relayed_body(httpx_response: Response) -> str | dict:
try:
body: Final[object] = httpx_response.json()
except ValueError:
return httpx_response.text
return body if isinstance(body, dict) else httpx_response.text
FOUNDRY_RELAY_SHAPES: Final = (
RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate),
RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate),

View file

@ -6,7 +6,7 @@ from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from pydantic import TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.types.utils import CallTypes
@ -29,6 +29,19 @@ if TYPE_CHECKING:
RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
class PassthroughMetadata(BaseModel):
model_config = ConfigDict(extra="ignore")
model_group: str = ""
def model_group_from(litellm_params: Mapping[str, object]) -> str:
try:
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
except ValidationError:
return ""
def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str:
path: Final = endpoint.lstrip("/")
for model_name in model_names:
@ -55,6 +68,14 @@ def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None
return None
def relayed_body(httpx_response: Response) -> str | dict:
try:
body: Final[object] = httpx_response.json()
except ValueError:
return httpx_response.text
return body if isinstance(body, dict) else httpx_response.text
@dataclass(frozen=True, slots=True)
class RelayShape:
path_suffix: str

View file

@ -0,0 +1,139 @@
from __future__ import annotations
import re
from collections.abc import Collection, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
model_group_from,
relayed_body,
strip_leading_model_segment,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import DeploymentTypedDict
from litellm.types.utils import LlmProviders, StandardPassThroughResponseObject
if TYPE_CHECKING:
from httpx import URL, Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
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]:
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)
)
def nvidia_nim_router_model_in_endpoint(endpoint: str, router_models: Collection[str]) -> str | None:
segments: Final = tuple(segment for segment in endpoint.split("/") if segment)
return next(
(
"/".join(segments[:length])
for length in range(len(segments), 0, -1)
if "/".join(segments[:length]) in router_models
),
None,
)
def without_repeated_version_prefix(api_base: str, native_endpoint: str) -> str:
url: Final = httpx.URL(api_base)
base_segments: Final = tuple(segment for segment in url.path.split("/") if segment)
first_native_segment: Final = native_endpoint.lstrip("/").split("/", 1)[0]
repeated: Final = (
bool(base_segments)
and API_VERSION_SEGMENT.match(first_native_segment) is not None
and base_segments[-1] == first_native_segment
)
kept_segments: Final = base_segments[:-1] if repeated else base_segments
return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/")
class NvidiaNimPassthroughConfig(BasePassthroughConfig):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return bool(request_data.get("stream", False))
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
endpoint: str,
request_query_params: dict | None,
litellm_params: dict,
) -> tuple[URL, str]:
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, 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)
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx
if api_key is None:
return dict(headers) # mutable-ok: base class contract returns dict for httpx
return {
**headers,
"Authorization": f"Bearer {api_key}",
} # mutable-ok: base class contract returns dict for httpx
@staticmethod
def get_api_base(api_base: str | None = None) -> str | None:
return api_base or get_secret_str("NVIDIA_NIM_API_BASE")
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
return api_key or get_secret_str("NVIDIA_NIM_API_KEY")
@staticmethod
def get_base_model(model: str) -> str | None:
return model
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
return []
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: Mapping[str, object],
logging_obj: Logging,
endpoint: str,
) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None:
return StandardPassThroughResponseObject(response=relayed_body(httpx_response))

View file

@ -205,6 +205,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/gigachat/",
"/milvus/",
"/mistral/",
"/nvidia_nim/",
"/openai/",
"/openai_passthrough/",
"/vertex-ai/",

View file

@ -18945,6 +18945,228 @@
]
}
},
"/nvidia_nim/{endpoint}": {
"delete": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__delete",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"get": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__get",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"patch": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__patch",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"post": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__post",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"put": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__put",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
}
},
"/openai/deployments/{model}/chat/completions": {
"post": {
"description": "Follows the exact same API spec as `OpenAI's Chat API https://platform.openai.com/docs/api-reference/chat`\n\n```bash\ncurl -X POST http://localhost:4000/v1/chat/completions \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"gpt-4o\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"Hello!\"\n }\n ]\n}'\n```",

View file

@ -486,6 +486,7 @@ class LiteLLMRoutes(enum.Enum):
"/milvus",
"/gigachat",
"/watsonx",
"/nvidia_nim",
]
#########################################################

View file

@ -28,6 +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_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 (
@ -2043,6 +2044,12 @@ def get_model_from_request(
azure_model: Final = _router_model_from_azure_route(route, llm_router)
return model if azure_model is None else azure_model
if route.lower().startswith("/nvidia_nim/"):
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

View file

@ -36,6 +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_group_in_path
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
from litellm.proxy._types import *
@ -1550,6 +1551,26 @@ async def _relay_azure_router_model(
"put the model group name in the deployments segment"
}
raise HTTPException(status_code=400, detail=rejection)
return await _relay_router_model(
llm_router=llm_router,
model=model,
endpoint=endpoint,
request=request,
request_body=request_body,
is_streaming_request=is_streaming_request,
user_api_key_dict=user_api_key_dict,
)
async def _relay_router_model(
llm_router: litellm.Router,
model: str,
endpoint: str,
request: Request,
request_body: Mapping[str, object],
is_streaming_request: bool,
user_api_key_dict: UserAPIKeyAuth,
) -> Response:
try:
result: Final = await llm_router.allm_passthrough_route(
model=model,
@ -1599,6 +1620,65 @@ async def _relay_azure_router_model(
)
@router.api_route(
"/nvidia_nim/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
tags=["NVIDIA NIM Pass-through", "pass-through"],
)
async def nvidia_nim_proxy_route(
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Relay a native NVIDIA NIM request through a LiteLLM model group.
`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
virtual key auth, model access checks, and spend logging.
"""
from litellm.proxy.proxy_server import llm_router
return await relay_nvidia_nim_request(
llm_router=llm_router,
endpoint=endpoint,
request=request,
request_body=await get_request_body(request),
user_api_key_dict=user_api_key_dict,
)
async def relay_nvidia_nim_request(
llm_router: litellm.Router | None,
endpoint: str,
request: Request,
request_body: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
) -> Response:
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 "
"group from your `model_list` whose deployments all use `nvidia_nim/` models"
}
raise HTTPException(status_code=400, detail=rejection)
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
return await open_sse_before_first_byte(
_relay_router_model(
llm_router=llm_router,
model=model_group,
endpoint=endpoint,
request=request,
request_body=request_body,
is_streaming_request=is_streaming_request,
user_api_key_dict=user_api_key_dict,
),
ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None),
)
@router.api_route(
"/azure_ai/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],

View file

@ -8998,6 +8998,12 @@ class ProviderConfigManager:
)
return WatsonxPassthroughConfig()
elif LlmProviders.NVIDIA_NIM == provider:
from litellm.llms.nvidia_nim.passthrough.transformation import (
NvidiaNimPassthroughConfig,
)
return NvidiaNimPassthroughConfig()
return None
@staticmethod

View file

@ -0,0 +1,296 @@
import json
from types import MappingProxyType
import httpx
import pytest
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
from litellm.utils import ProviderConfigManager
NIM_BASE = "http://nim.internal:8000"
INFER_BODY = {
"input": [
{"type": "image_url", "url": "data:image/png;base64,AAAA"},
{"type": "image_url", "url": "data:image/png;base64,BBBB"},
]
}
@pytest.fixture(autouse=True)
def clear_nvidia_nim_env(monkeypatch):
for env_var in ("NVIDIA_NIM_API_BASE", "NVIDIA_NIM_API_KEY"):
monkeypatch.delenv(env_var, raising=False)
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setattr(litellm, "api_key", None)
def test_provider_config_manager_resolves_nvidia_nim_passthrough_config():
config = ProviderConfigManager.get_provider_passthrough_config(
model="nvidia/nemoretriever-page-elements-v2", provider=LlmProviders.NVIDIA_NIM
)
assert isinstance(config, NvidiaNimPassthroughConfig)
@pytest.mark.parametrize(
"api_base, endpoint, litellm_params, expected",
[
(NIM_BASE, "nim-page/v1/infer", {"litellm_metadata": {"model_group": "nim-page"}}, f"{NIM_BASE}/v1/infer"),
(
f"{NIM_BASE}/v1",
"nim-page/v1/infer",
{"litellm_metadata": {"model_group": "nim-page"}},
f"{NIM_BASE}/v1/infer",
),
(f"{NIM_BASE}/v1/", "/v1/infer", {}, f"{NIM_BASE}/v1/infer"),
(NIM_BASE, "v1/infer", {}, f"{NIM_BASE}/v1/infer"),
(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(
api_base, endpoint, litellm_params, expected
):
url, base = NvidiaNimPassthroughConfig().get_complete_url(
api_base=api_base,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint=endpoint,
request_query_params=None,
litellm_params=litellm_params,
)
assert str(url) == expected
assert base == expected.removesuffix("/v1/infer").removesuffix("/infer")
def test_query_params_are_forwarded_on_the_relay_url():
url, _ = NvidiaNimPassthroughConfig().get_complete_url(
api_base=NIM_BASE,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params={"timeout": "30"},
litellm_params={},
)
assert str(url) == f"{NIM_BASE}/v1/infer?timeout=30"
def test_env_api_base_is_used_when_the_deployment_has_none(monkeypatch):
monkeypatch.setenv("NVIDIA_NIM_API_BASE", f"{NIM_BASE}/v1")
url, _ = NvidiaNimPassthroughConfig().get_complete_url(
api_base=None,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params=None,
litellm_params={},
)
assert str(url) == f"{NIM_BASE}/v1/infer"
def test_missing_api_base_raises_instead_of_building_a_relative_url():
with pytest.raises(ValueError, match="NVIDIA_NIM_API_BASE"):
NvidiaNimPassthroughConfig().get_complete_url(
api_base=None,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params=None,
litellm_params={},
)
def test_deployment_key_becomes_a_bearer_token_and_caller_headers_are_kept():
caller_headers = MappingProxyType({"x-request-id": "abc"})
headers = NvidiaNimPassthroughConfig().validate_environment(
headers=caller_headers,
model="nvidia/nemoretriever-page-elements-v2",
messages=[],
optional_params={},
litellm_params={},
api_key="nvapi-secret",
)
assert headers == {"x-request-id": "abc", "Authorization": "Bearer nvapi-secret"}
def test_self_hosted_nim_without_a_key_sends_no_authorization_header():
headers = NvidiaNimPassthroughConfig().validate_environment(
headers={}, model="nvidia/x", messages=[], optional_params={}, litellm_params={}, api_key=None
)
assert "Authorization" not in headers
def test_env_api_key_fills_in_when_the_deployment_has_none(monkeypatch):
monkeypatch.setenv("NVIDIA_NIM_API_KEY", "nvapi-from-env")
assert NvidiaNimPassthroughConfig.get_api_key(None) == "nvapi-from-env"
assert NvidiaNimPassthroughConfig.get_api_key("nvapi-deployment") == "nvapi-deployment"
@pytest.mark.parametrize(
"endpoint, router_models, expected",
[
("nim-page/v1/infer", ("nim-page", "nim-table"), "nim-page"),
("/nim-page/v1/infer", ("nim-page",), "nim-page"),
(
"nvidia/nemoretriever-page-elements-v2/v1/infer",
("nvidia/nemoretriever-page-elements-v2",),
"nvidia/nemoretriever-page-elements-v2",
),
("nim/v1/infer", ("nim", "nim/v1"), "nim/v1"),
("v1/infer", ("nim-page",), None),
("nim-page-elements/v1/infer", ("nim-page",), None),
("", ("nim-page",), None),
],
)
def test_router_model_in_endpoint_takes_the_longest_leading_model_group(endpoint, router_models, expected):
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
def test_non_streaming_relay_logs_the_upstream_json_body():
response = httpx.Response(
200,
json={"data": [{"index": 0, "bounding_boxes": {}}]},
request=httpx.Request("POST", f"{NIM_BASE}/v1/infer"),
)
result = NvidiaNimPassthroughConfig().logging_non_streaming_response(
model="nvidia/nemoretriever-page-elements-v2",
custom_llm_provider="nvidia_nim",
httpx_response=response,
request_data=INFER_BODY,
logging_obj=None, # pyright: ignore[reportArgumentType] # not read for a plain passthrough body
endpoint="v1/infer",
)
assert result == {"response": {"data": [{"index": 0, "bounding_boxes": {}}]}}
@pytest.mark.asyncio
async def test_object_detection_relay_sends_the_native_body_unchanged_to_v1_infer():
upstream_requests: list[httpx.Request] = []
def nim(request: httpx.Request) -> httpx.Response:
upstream_requests.append(request)
return httpx.Response(200, json={"data": [{"index": 0}, {"index": 1}]}, headers={"x-nim": "1"})
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(nim))
response = await litellm.allm_passthrough_route(
model="nvidia_nim/nvidia/nemoretriever-page-elements-v2",
endpoint="nim-page/v1/infer",
method="POST",
api_base=f"{NIM_BASE}/v1",
api_key="nvapi-secret",
json=dict(INFER_BODY),
litellm_metadata={"model_group": "nim-page"},
client=client,
)
(sent,) = upstream_requests
assert str(sent.url) == f"{NIM_BASE}/v1/infer"
assert json.loads(sent.content) == INFER_BODY
assert sent.headers["authorization"] == "Bearer nvapi-secret"
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

View file

@ -853,6 +853,82 @@ def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_pa
assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected
def _nvidia_nim_relay_router():
from litellm.router import Router
return Router(
model_list=[
{
"model_name": "nim-page-elements",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"api_base": "http://nim-a.internal:8000",
"api_key": "k",
},
},
{
"model_name": "nvidia/nemoretriever-table-structure-v1",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-table-structure-v1",
"api_base": "http://nim-b.internal:8000",
"api_key": "k",
},
},
{
"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"},
},
]
)
NIM_INFER_BODY = {"input": [{"type": "image_url", "url": "data:image/png;base64,AAAA"}]}
@pytest.mark.parametrize(
"route, request_data, expected",
[
("/nvidia_nim/nim-page-elements/v1/infer", NIM_INFER_BODY, "nim-page-elements"),
(
"/nvidia_nim/nim-page-elements/v1/infer",
{"model": "nvidia/nemoretriever-table-structure-v1"},
"nim-page-elements",
),
(
"/nvidia_nim/nvidia/nemoretriever-table-structure-v1/v1/infer",
NIM_INFER_BODY,
"nvidia/nemoretriever-table-structure-v1",
),
("/nvidia_nim/v1/infer", NIM_INFER_BODY, None),
("/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):
assert (
get_model_from_request(request_data=request_data, route=route, llm_router=_nvidia_nim_relay_router())
== expected
)
def test_get_model_from_request_nvidia_nim_relay_without_a_router_has_no_model():
assert get_model_from_request(request_data=NIM_INFER_BODY, route="/nvidia_nim/nim-page-elements/v1/infer") is None
def test_get_model_from_request_includes_file_endpoint_header_model():
assert (
get_model_from_request(

View file

@ -693,6 +693,7 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied():
"/anthropic/v1/count_tokens",
"/gemini/v1/models",
"/gemini/countTokens",
"/nvidia_nim/nim-page-elements/v1/infer",
],
)
def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):

View file

@ -40,6 +40,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
llm_passthrough_factory_proxy_route,
milvus_proxy_route,
mistral_proxy_route,
relay_nvidia_nim_request,
openai_proxy_route,
vertex_discovery_proxy_route,
vertex_proxy_route,
@ -5375,6 +5376,186 @@ class TestRouterModelRelayUpstreamContract:
assert result.headers["x-ms-request-id"] == "req-1"
NIM_INFER_BODY = {
"input": [
{"type": "image_url", "url": "data:image/png;base64,AAAA"},
{"type": "image_url", "url": "data:image/png;base64,BBBB"},
]
}
class TestNvidiaNimProxyRoute:
def _request(self) -> MagicMock:
request = MagicMock(spec=Request)
request.method = "POST"
request.headers = {"content-type": "application/json"}
request.query_params = {}
return request
def _recording_router(self, captured: list[dict], deployments: dict[str, str]):
class RecordingRouter:
def get_model_list(self):
return [{"model_name": name, "litellm_params": {"model": model}} for name, model in deployments.items()]
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(
200, json={"data": [{"index": 0, "bounding_boxes": {}}]}, headers={"x-nim-request": "r1"}
)
return RecordingRouter()
async def _relay(self, llm_router, endpoint: str, body: dict, user_api_key_dict=None) -> Response:
return await relay_nvidia_nim_request(
llm_router=llm_router,
endpoint=endpoint,
request=self._request(),
request_body=dict(body),
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(api_key="hashed-token"),
)
@pytest.mark.asyncio
async def test_model_group_in_the_path_selects_the_deployment_and_the_body_stays_model_free(self):
captured: list[dict] = []
router = self._recording_router(
captured,
{
"nim-page-elements": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"nim-table": "nvidia_nim/nvidia/nemoretriever-table-structure-v1",
},
)
result = await self._relay(
router,
"nim-page-elements/v1/infer",
NIM_INFER_BODY,
UserAPIKeyAuth(api_key="hashed-token", team_id="team-1"),
)
(relay,) = captured
assert relay["model"] == "nim-page-elements"
assert relay["endpoint"] == "nim-page-elements/v1/infer"
assert relay["method"] == "POST"
assert relay["json"] == NIM_INFER_BODY
assert "model" not in relay["json"]
assert relay["litellm_metadata"]["user_api_key_team_id"] == "team-1"
assert result.status_code == 200
assert json.loads(result.body) == {"data": [{"index": 0, "bounding_boxes": {}}]}
assert result.headers["x-nim-request"] == "r1"
@pytest.mark.asyncio
async def test_model_group_with_a_slash_is_matched_as_the_longest_leading_path(self):
captured: list[dict] = []
router = self._recording_router(
captured, {"nvidia/nemoretriever-page-elements-v2": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"}
)
await self._relay(router, "nvidia/nemoretriever-page-elements-v2/v1/infer", NIM_INFER_BODY)
assert captured[0]["model"] == "nvidia/nemoretriever-page-elements-v2"
@pytest.mark.asyncio
async def test_custom_llm_provider_marks_a_deployment_as_nim_without_the_model_prefix(self):
captured: list[dict] = []
class ProviderRouter:
def get_model_list(self):
return [
{
"model_name": "page-elements",
"litellm_params": {
"model": "nvidia/nemoretriever-page-elements-v2",
"custom_llm_provider": "nvidia_nim",
},
}
]
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(200, json={"data": []})
await self._relay(ProviderRouter(), "page-elements/v1/infer", NIM_INFER_BODY)
assert captured[0]["model"] == "page-elements"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"endpoint",
["v1/infer", "unknown-group/v1/infer", "nim-page-elements-v2/v1/infer", "gpt-4o/v1/infer"],
)
async def test_path_without_a_nim_model_group_is_rejected_before_any_upstream_call(self, endpoint):
captured: list[dict] = []
router = self._recording_router(
captured,
{"nim-page-elements": "nvidia_nim/nvidia/nemoretriever-page-elements-v2", "gpt-4o": "openai/gpt-4o"},
)
with pytest.raises(HTTPException) as exc_info:
await self._relay(router, endpoint, NIM_INFER_BODY)
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):
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):
with pytest.raises(HTTPException) as exc_info:
await self._relay(None, "nim-page-elements/v1/infer", NIM_INFER_BODY)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_upstream_rejection_is_relayed_with_its_status_body_and_headers(self):
upstream_body = {"detail": "input[0].url must be a data URL"}
class RejectingRouter:
def get_model_list(self):
return [
{
"model_name": "nim-page-elements",
"litellm_params": {"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"},
}
]
async def allm_passthrough_route(self, **kwargs):
upstream_request = httpx.Request("POST", "http://nim.internal:8000/v1/infer")
upstream = httpx.Response(
422, json=upstream_body, headers={"x-nim-request": "r2"}, request=upstream_request
)
raise httpx.HTTPStatusError("422", request=upstream_request, response=upstream)
result = await self._relay(
RejectingRouter(), "nim-page-elements/v1/infer", {"input": [{"type": "image_url", "url": "x"}]}
)
assert result.status_code == 422
assert json.loads(result.body) == upstream_body
assert result.headers["x-nim-request"] == "r2"
@pytest.mark.asyncio
async def test_bedrock_count_tokens_error_forwards_provider_headers():
"""The count tokens route converts BedrockError into an HTTPException, and dropping the

View file

@ -9676,6 +9676,62 @@ export interface paths {
patch?: never;
trace?: never;
};
"/nvidia_nim/{endpoint}": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
/**
* Nvidia Nim Proxy Route
* @description Relay a native NVIDIA NIM request through a LiteLLM model group.
*
* `{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
* `api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
* virtual key auth, model access checks, and spend logging.
*/
get: operations["nvidia_nim_proxy_route_nvidia_nim__endpoint__get"];
/**
* Nvidia Nim Proxy Route
* @description Relay a native NVIDIA NIM request through a LiteLLM model group.
*
* `{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
* `api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
* virtual key auth, model access checks, and spend logging.
*/
put: operations["nvidia_nim_proxy_route_nvidia_nim__endpoint__put"];
/**
* Nvidia Nim Proxy Route
* @description Relay a native NVIDIA NIM request through a LiteLLM model group.
*
* `{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
* `api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
* virtual key auth, model access checks, and spend logging.
*/
post: operations["nvidia_nim_proxy_route_nvidia_nim__endpoint__post"];
/**
* Nvidia Nim Proxy Route
* @description Relay a native NVIDIA NIM request through a LiteLLM model group.
*
* `{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
* `api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
* virtual key auth, model access checks, and spend logging.
*/
delete: operations["nvidia_nim_proxy_route_nvidia_nim__endpoint__delete"];
options?: never;
head?: never;
/**
* Nvidia Nim Proxy Route
* @description Relay a native NVIDIA NIM request through a LiteLLM model group.
*
* `{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
* `api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
* virtual key auth, model access checks, and spend logging.
*/
patch: operations["nvidia_nim_proxy_route_nvidia_nim__endpoint__patch"];
trace?: never;
};
"/ocr": {
parameters: {
query?: never;
@ -53485,6 +53541,161 @@ export interface operations {
};
};
};
nvidia_nim_proxy_route_nvidia_nim__endpoint__get: {
parameters: {
query?: never;
header?: never;
path: {
endpoint: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
nvidia_nim_proxy_route_nvidia_nim__endpoint__put: {
parameters: {
query?: never;
header?: never;
path: {
endpoint: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
nvidia_nim_proxy_route_nvidia_nim__endpoint__post: {
parameters: {
query?: never;
header?: never;
path: {
endpoint: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
nvidia_nim_proxy_route_nvidia_nim__endpoint__delete: {
parameters: {
query?: never;
header?: never;
path: {
endpoint: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
nvidia_nim_proxy_route_nvidia_nim__endpoint__patch: {
parameters: {
query?: never;
header?: never;
path: {
endpoint: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
ocr_ocr_post: {
parameters: {
query?: never;