mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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:
commit
7eeba69016
16 changed files with 1247 additions and 23 deletions
|
|
@ -96,6 +96,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/langfuse/",
|
||||
"/vllm/",
|
||||
"/mistral/",
|
||||
"/nvidia_nim/",
|
||||
"/groq/",
|
||||
"/voyage/",
|
||||
"/cursor/",
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
0
litellm/llms/nvidia_nim/passthrough/__init__.py
Normal file
0
litellm/llms/nvidia_nim/passthrough/__init__.py
Normal file
139
litellm/llms/nvidia_nim/passthrough/transformation.py
Normal file
139
litellm/llms/nvidia_nim/passthrough/transformation.py
Normal 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))
|
||||
|
|
@ -205,6 +205,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/gigachat/",
|
||||
"/milvus/",
|
||||
"/mistral/",
|
||||
"/nvidia_nim/",
|
||||
"/openai/",
|
||||
"/openai_passthrough/",
|
||||
"/vertex-ai/",
|
||||
|
|
|
|||
|
|
@ -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```",
|
||||
|
|
|
|||
|
|
@ -486,6 +486,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/milvus",
|
||||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
211
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
211
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue