From e9637c9f94b4b9875793f8a56320ed9a47d5cc06 Mon Sep 17 00:00:00 2001 From: Sbaaoui Idriss <112825897+6ba3i@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:15:34 +0800 Subject: [PATCH] fix: list model function for nvidia on python/go not returning current models (#17501) ### Summary fix the list model logic for nvidia models --- api/apps/restful_apis/provider_api.py | 4 +- api/apps/services/provider_api_service.py | 133 +++- api/db/services/tenant_model_service.py | 20 +- conf/llm_factories.json | 883 +++++----------------- conf/models/nvidia.json | 333 +++----- internal/entity/models/nvidia.go | 322 +++++++- internal/entity/models/nvidia_test.go | 230 ++++++ internal/service/model_service.go | 121 +++ internal/service/model_service_test.go | 140 ++++ rag/llm/chat_model.py | 11 +- rag/llm/cv_model.py | 82 +- rag/llm/embedding_model.py | 8 +- rag/llm/model_meta.py | 177 +++++ rag/utils/url_utils.py | 14 + 14 files changed, 1429 insertions(+), 1049 deletions(-) create mode 100644 internal/entity/models/nvidia_test.go diff --git a/api/apps/restful_apis/provider_api.py b/api/apps/restful_apis/provider_api.py index 422e4a2956..537e290e04 100644 --- a/api/apps/restful_apis/provider_api.py +++ b/api/apps/restful_apis/provider_api.py @@ -715,7 +715,7 @@ async def drop_provider_instances(tenant_id: str = None, provider_id_or_name: st @manager.route("/providers//instances//models", methods=["GET"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs -def list_instance_models(tenant_id: str = None, provider_id_or_name: str = None, instance_id_or_name: str = None): +async def list_instance_models(tenant_id: str = None, provider_id_or_name: str = None, instance_id_or_name: str = None): """ List models for a provider instance. --- @@ -757,7 +757,7 @@ def list_instance_models(tenant_id: str = None, provider_id_or_name: str = None, """ supported_only = request.args.get("supported", "").lower() == "true" try: - success, result = provider_api_service.list_instance_models(tenant_id, provider_id_or_name, instance_id_or_name, supported_only) + success, result = await provider_api_service.list_instance_models(tenant_id, provider_id_or_name, instance_id_or_name, supported_only) if success: return get_result(data=result) else: diff --git a/api/apps/services/provider_api_service.py b/api/apps/services/provider_api_service.py index ec128c2350..4da8b57927 100644 --- a/api/apps/services/provider_api_service.py +++ b/api/apps/services/provider_api_service.py @@ -20,6 +20,7 @@ import asyncio from common.constants import LLMType, ActiveStatusEnum, ModelVerifyStatusEnum from common.settings import FACTORY_LLM_INFOS +from api.db.db_models import DB from api.db.joint_services.tenant_model_service import resolve_model_config, delete_models_by_instance_ids, delete_instances_by_provider_ids from api.db.services.tenant_model_provider_service import TenantModelProviderService from api.db.services.tenant_model_instance_service import TenantModelInstanceService @@ -944,7 +945,100 @@ def drop_provider_instances(tenant_id: str, provider_id_or_name: str, instance_i return True, None -def list_instance_models(tenant_id: str, provider_id_or_name: str, instance_id_or_name: str, supported_only: bool = False): +def _public_factory_model(llm: dict) -> dict: + return { + "name": _factory_llm_name(llm), + "max_tokens": llm.get("max_tokens", 8192), + "model_types": _factory_model_types(llm), + "features": (llm.get("features") if llm.get("features") is not None else ((["is_tools"] if llm.get("is_tools") else []) + (["thinking"] if llm.get("thinking") else []))), + } + + +def _merge_nvidia_models(factory_info: dict, remote_models: list[dict]) -> list[dict]: + static_models = {_factory_llm_name(llm): _public_factory_model(llm) for llm in factory_info.get("llm", [])} + merged = [] + seen = set() + for remote in remote_models: + model_name = str(remote.get("name", "")).strip() + if not model_name or model_name in seen: + continue + seen.add(model_name) + model = dict(remote) + model["name"] = model_name + if preset := static_models.get(model_name): + model = {**model, **preset} + if not model.get("model_types"): + model["model_types"] = [LLMType.CHAT.value] + model["max_tokens"] = _to_int(model.get("max_tokens"), 8192) + model.setdefault("features", []) + merged.append(model) + merged.sort(key=lambda model: model["name"]) + return merged + + +def _set_discovered_model_metadata(extra: dict, model: dict): + extra["max_tokens"] = _to_int(model.get("max_tokens"), 8192) + if model.get("max_dimension") is not None: + extra["max_dimension"] = model["max_dimension"] + if model.get("dimensions"): + extra["dimensions"] = model["dimensions"] + features = model.get("features") or [] + extra["is_tools"] = "is_tools" in features + extra["thinking"] = "thinking" in features + + +def _reconcile_nvidia_instance_models(provider_obj, instance_obj, remote_models: list[dict]): + normalized = [] + seen = set() + for model in remote_models: + model_name = str(model.get("name", "")).strip() + if not model_name or model_name in seen: + continue + seen.add(model_name) + normalized.append({**model, "name": model_name}) + if not normalized: + raise ValueError("NVIDIA model discovery returned no usable models") + + with DB.atomic(): + existing_models = TenantModelService.get_models_by_instance_id(instance_obj.id) + existing_by_name = {model.model_name: model for model in existing_models} + + for model in normalized: + model_name = model["name"] + model_types = model.get("model_types") or [LLMType.CHAT.value] + model_type = calculate_model_type(model_types) + if existing := existing_by_name.pop(model_name, None): + extra = json.loads(existing.extra or "{}") + _set_discovered_model_metadata(extra, model) + TenantModelService.update_model(existing.id, {"model_type": model_type, "extra": json.dumps(extra)}) + continue + + extra = {"verify": ModelVerifyStatusEnum.UNKNOWN.value} + _set_discovered_model_metadata(extra, model) + TenantModelService.insert( + model_name=model_name, + provider_id=provider_obj.id, + instance_id=instance_obj.id, + model_type=model_type, + status=ActiveStatusEnum.ACTIVE.value, + extra=json.dumps(extra), + ) + + TenantModelService.delete_by_ids([model.id for model in existing_by_name.values()]) + + +def _get_provider_instance(provider_obj, instance_id_or_name: str): + instance_obj = None + if instance_id_or_name: + _, instance_obj = TenantModelInstanceService.get_by_id(instance_id_or_name) + if instance_obj and instance_obj.provider_id != provider_obj.id: + instance_obj = None + if not instance_obj: + instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_obj.id, instance_id_or_name) + return instance_obj + + +async def list_instance_models(tenant_id: str, provider_id_or_name: str, instance_id_or_name: str, supported_only: bool = False): """ List models for a provider instance. @@ -966,15 +1060,34 @@ def list_instance_models(tenant_id: str, provider_id_or_name: str, instance_id_o return False, f"No provider found for provider '{provider_id_or_name}'" if supported_only: - # List all models supported by this provider from the LLM dictionary - factory_info = [f for f in FACTORY_LLM_INFOS if f["name"] == provider_obj.provider_name] - if not factory_info: + # List all models supported by this provider from the LLM dictionary. + factory_infos = [f for f in FACTORY_LLM_INFOS if f["name"] == provider_obj.provider_name] + if not factory_infos: return False, f"Provider '{provider_id_or_name}' not found" - llms = factory_info[0].get("llm", []) + factory_info = factory_infos[0] + + if provider_obj.provider_name == "NVIDIA": + instance_obj = _get_provider_instance(provider_obj, instance_id_or_name) + if not instance_obj: + return False, f"No instance found for provider '{provider_id_or_name}' and instance '{instance_id_or_name}'" + instance_extra = json.loads(instance_obj.extra or "{}") + base_url = instance_extra.get("base_url") or factory_info.get("url", "") + remote_models = await ModelMeta["NVIDIA"](instance_obj.api_key, base_url).get_model_list() + models = _merge_nvidia_models(factory_info, remote_models) + if not models: + return False, "NVIDIA model discovery returned no usable models" + _reconcile_nvidia_instance_models(provider_obj, instance_obj, models) + return True, models + + llms = factory_info.get("llm", []) models = [{"name": llm["llm_name"], "rank": _to_int(llm.get("rank", 500))} for llm in llms] models.sort(key=lambda x: (-x["rank"], x["name"])) return True, models + instance_obj = _get_provider_instance(provider_obj, instance_id_or_name) + if not instance_obj: + return False, f"No instance found for provider '{provider_id_or_name}' and instance '{instance_id_or_name}'" + # Build rank mapping from LLM dictionary for the provider factory_info = [f for f in FACTORY_LLM_INFOS if f["name"] == provider_obj.provider_name] model_rank_map = {} @@ -982,16 +1095,6 @@ def list_instance_models(tenant_id: str, provider_id_or_name: str, instance_id_o for llm in factory_info[0].get("llm", []): model_rank_map[llm["llm_name"]] = _to_int(llm.get("rank", 500)) - # Get instance - instance_obj = None - if instance_id_or_name: - _, instance_obj = TenantModelInstanceService.get_by_id(instance_id_or_name) - if instance_obj and instance_obj.provider_id != provider_obj.id: - instance_obj = None - if not instance_obj: - instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_obj.id, instance_id_or_name) - if not instance_obj: - return False, f"No instance found for provider '{provider_id_or_name}' and instance '{instance_id_or_name}'" # Get models model_objs = TenantModelService.get_models_by_instance_id(instance_obj.id) model_list = [] diff --git a/api/db/services/tenant_model_service.py b/api/db/services/tenant_model_service.py index d8274d8ab2..f4af464e0a 100644 --- a/api/db/services/tenant_model_service.py +++ b/api/db/services/tenant_model_service.py @@ -31,23 +31,14 @@ class TenantModelService(CommonService): def get_by_provider_id_and_instance_id_and_model_type_and_model_name(cls, provider_id, instance_id, model_type, model_name): if isinstance(model_type, str): model_type = calculate_model_type([model_type]) - return cls.model.get_or_none( - cls.model.provider_id == provider_id, - cls.model.instance_id == instance_id, - cls.model.model_type.bin_and(model_type) > 0, - cls.model.model_name == model_name - ) + return cls.model.get_or_none(cls.model.provider_id == provider_id, cls.model.instance_id == instance_id, cls.model.model_type.bin_and(model_type) > 0, cls.model.model_name == model_name) @classmethod @DB.connection_context() def get_by_provider_id_and_instance_id_and_model_type(cls, provider_id, instance_id, model_type): if isinstance(model_type, str): model_type = calculate_model_type([model_type]) - return cls.model.get_or_none( - cls.model.provider_id == provider_id, - cls.model.instance_id == instance_id, - cls.model.model_type.bin_and(model_type) > 0 - ) + return cls.model.get_or_none(cls.model.provider_id == provider_id, cls.model.instance_id == instance_id, cls.model.model_type.bin_and(model_type) > 0) @classmethod @DB.connection_context() @@ -87,3 +78,10 @@ class TenantModelService(CommonService): @DB.connection_context() def delete_by_instance_ids(cls, instance_ids): return cls.model.delete().where(cls.model.instance_id.in_(instance_ids)).execute() + + @classmethod + @DB.connection_context() + def delete_by_ids(cls, model_ids): + if not model_ids: + return 0 + return cls.model.delete().where(cls.model.id.in_(model_ids)).execute() diff --git a/conf/llm_factories.json b/conf/llm_factories.json index aa1f75acc3..c7970ca634 100644 --- a/conf/llm_factories.json +++ b/conf/llm_factories.json @@ -2809,213 +2809,29 @@ "status": "1", "llm": [ { - "llm_name": "01-ai/yi-large", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "abacusai/dracarys-llama-3.1-70b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "ai21labs/jamba-1.5-large-instruct", - "tags": "LLM,CHAT,256K", - "max_tokens": 256000, - "model_type": "chat" - }, - { - "llm_name": "ai21labs/jamba-1.5-mini-instruct", - "tags": "LLM,CHAT,256K", - "max_tokens": 256000, - "model_type": "chat" - }, - { - "llm_name": "aisingapore/sea-lion-7b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "baichuan-inc/baichuan2-13b-chat", - "tags": "LLM,CHAT,192K", - "max_tokens": 196608, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "bigcode/starcoder2-7b", - "tags": "LLM,CHAT,16K", - "max_tokens": 16384, - "model_type": "chat" - }, - { - "llm_name": "bigcode/starcoder2-15b", - "tags": "LLM,CHAT,16K", - "max_tokens": 16384, - "model_type": "chat" - }, - { - "llm_name": "databricks/dbrx-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "deepseek-ai/deepseek-r1", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "google/gemma-2b", + "llm_name": "deepseek-ai/deepseek-v4-flash", "tags": "LLM,CHAT,8K", "max_tokens": 8192, "model_type": "chat" }, { - "llm_name": "google/gemma-7b", + "llm_name": "deepseek-ai/deepseek-v4-pro", "tags": "LLM,CHAT,8K", "max_tokens": 8192, "model_type": "chat" }, { - "llm_name": "google/gemma-2-2b-it", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "google/gemma-2-9b-it", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "google/gemma-2-27b-it", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "google/codegemma-1.1-7b", + "llm_name": "google/diffusiongemma-26b-a4b-it", "tags": "LLM,CHAT,8K", "max_tokens": 8192, "model_type": "chat" }, { - "llm_name": "google/codegemma-7b", + "llm_name": "google/gemma-4-31b-it", "tags": "LLM,CHAT,8K", "max_tokens": 8192, "model_type": "chat" }, - { - "llm_name": "google/recurrentgemma-2b", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "google/shieldgemma-9b", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "ibm/granite-3.0-3b-a800m-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "ibm/granite-3.0-8b-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "ibm/granite-34b-code-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "ibm/granite-8b-code-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "ibm/granite-guardian-3.0-8b", - "tags": "LLM,CHAT,128k", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "igenius / colosseum-355b_instruct_16k", - "tags": "LLM,CHAT,16K", - "max_tokens": 16384, - "model_type": "chat" - }, - { - "llm_name": "igenius / italia_10b_instruct_16k", - "tags": "LLM,CHAT,16K", - "max_tokens": 16384, - "model_type": "chat" - }, - { - "llm_name": "institute-of-science-tokyo/llama-3.1-swallow-70b-instruct-v01", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "institute-of-science-tokyo/llama-3.1-swallow-8b-instruct-v0.1", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "mediatek/breeze-7b-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "meta/codellama-70b", - "tags": "LLM,CHAT,100K", - "max_tokens": 100000, - "model_type": "chat" - }, - { - "llm_name": "meta/llama2-70b", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "meta/llama3-8b", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "meta/llama3-70b", - "tags": "LLM,CHAT,", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "meta/llama-3.1-8b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, { "llm_name": "meta/llama-3.1-70b-instruct", "tags": "LLM,CHAT,128K", @@ -3023,11 +2839,20 @@ "model_type": "chat" }, { - "llm_name": "meta/llama-3.1-405b-instruct", + "llm_name": "meta/llama-3.1-8b-instruct", "tags": "LLM,CHAT,128K", "max_tokens": 131072, "model_type": "chat" }, + { + "llm_name": "meta/llama-3.2-11b-vision-instruct", + "tags": "IMAGE2TEXT,128K", + "max_tokens": 131072, + "model_type": [ + "image2text", + "chat" + ] + }, { "llm_name": "meta/llama-3.2-1b-instruct", "tags": "LLM,CHAT,128K", @@ -3040,482 +2865,6 @@ "max_tokens": 131072, "model_type": "chat" }, - { - "llm_name": "meta/llama-3.3-70b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-medium-128k-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-medium-4k-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-mini-128k-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-mini-4k-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-small-128k-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3-small-8k-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3.5-mini", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "microsoft/phi-3.5-moe-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "mistralai/codestral-22b-instruct-v0.1", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "mistralai/mamba-codestral-7b-v0.1", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mistral-2-large-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mathstral-7b-v01", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mistral-7b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mistral-7b-instruct-v0.3", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mixtral-8x7b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mixtral-8x22b-instruct", - "tags": "LLM,CHAT,64K", - "max_tokens": 65536, - "model_type": "chat" - }, - { - "llm_name": "mistralai/mistral-large", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "mistralai/mistral-small-24b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama3-chatqa-1.5-8b", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama-3.1-nemoguard-8b-content-safety", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama-3.1-nemoguard-8b-topic-control", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama-3.1-nemotron-51b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama-3.1-nemotron-70b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama-3.1-nemotron-70b-reward", - "tags": "LLM,CHAT,128K", - "max_tokens": 128000, - "model_type": "chat" - }, - { - "llm_name": "nvidia/llama3-chatqa-1.5-70b", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "nvidia/mistral-nemo-minitron-8b-base", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "nvidia/mistral-nemo-minitron-8b-8k-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "nvidia/nemotron-4-340b-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "nvidia/nemotron-4-340b-reward", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "nvidia/nemotron-4-mini-hindi-4b-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "nvidia/nemotron-mini-4b-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "nv-mistralai/mistral-nemo-12b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "qwen/qwen2-7b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "qwen/qwen2.5-7b-instruct", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "qwen/qwen2.5-coder-7b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "rakuten/rakutenai-7b-chat", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "rakuten/rakutenai-7b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "seallms/seallm-7b-v2.5", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "snowflake/arctic", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "tokyotech-llm/llama-3-swallow-70b-instruct-v01", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "thudm/chatglm3-6b", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "tiiuae/falcon3-7b-instruct", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "upstage/solar-10.7b-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat", - "is_tools": true - }, - { - "llm_name": "writer/palmyra-creative-122b", - "tags": "LLM,CHAT,128K", - "max_tokens": 131072, - "model_type": "chat" - }, - { - "llm_name": "writer/palmyra-fin-70b-32k", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "writer/palmyra-med-70b-32k", - "tags": "LLM,CHAT,32K", - "max_tokens": 32768, - "model_type": "chat" - }, - { - "llm_name": "writer/palmyra-med-70b", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "yentinglin/llama-3-taiwan-70b-instruct", - "tags": "LLM,CHAT,8K", - "max_tokens": 8192, - "model_type": "chat" - }, - { - "llm_name": "zyphra/zamba2-7b-instruct", - "tags": "LLM,CHAT,4K", - "max_tokens": 4096, - "model_type": "chat" - }, - { - "llm_name": "BAAI/bge-m3", - "tags": "TEXT EMBEDDING", - "max_tokens": 8192, - "model_type": "embedding" - }, - { - "llm_name": "BAAI/bge-m3-unsupervised", - "tags": "TEXT EMBEDDING", - "max_tokens": 8192, - "model_type": "embedding" - }, - { - "llm_name": "BAAI/bge-m3-retromae", - "tags": "TEXT EMBEDDING", - "max_tokens": 8129, - "model_type": "embedding" - }, - { - "llm_name": "BAAI/bge-large-en-v1.5", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "BAAI/bge-base-en-v1.5", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "BAAI/bge-small-en-v1.5", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/embed-qa-4", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/llama-3.2-nv-embedqa-1b-v1", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/llama-3.2-nv-embedqa-1b-v2", - "tags": "TEXT EMBEDDING", - "max_tokens": 8192, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/llama-3.2-nv-rerankqa-1b-v1", - "tags": "RE-RANK,512", - "max_tokens": 512, - "model_type": "rerank" - }, - { - "llm_name": "nvidia/llama-3.2-nv-rerankqa-1b-v2", - "tags": "RE-RANK,8K", - "max_tokens": 8192, - "model_type": "rerank" - }, - { - "llm_name": "nvidia/nvclip", - "tags": "TEXT EMBEDDING", - "max_tokens": 1024, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/nv-embed-v1", - "tags": "TEXT EMBEDDING", - "max_tokens": 4096, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/nv-embedqa-e5-v5", - "tags": "TEXT EMBEDDING", - "max_tokens": 1024, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/nv-embedqa-mistral-7b-v2", - "tags": "TEXT EMBEDDING", - "max_tokens": 4096, - "model_type": "embedding" - }, - { - "llm_name": "nvidia/nv-rerankqa-mistral-4b-v3", - "tags": "RE-RANK,512", - "max_tokens": 512, - "model_type": "rerank" - }, - { - "llm_name": "nvidia/rerank-qa-mistral-4b", - "tags": "RE-RANK,512", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "snowflake-arctic-embed-xs", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "snowflake-arctic-embed-s", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "snowflake-arctic-embed-m", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "snowflake-arctic-embed-m-long", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "snowflake-arctic-embed-l", - "tags": "TEXT EMBEDDING", - "max_tokens": 512, - "model_type": "embedding" - }, - { - "llm_name": "adept/fuyu-8b", - "tags": "IMAGE2TEXT,1K", - "max_tokens": 1024, - "model_type": [ - "image2text", - "chat" - ] - }, - { - "llm_name": "google/deplot", - "tags": "IMAGE2TEXT,8K", - "max_tokens": 8192, - "model_type": [ - "image2text", - "chat" - ] - }, - { - "llm_name": "google/paligemma", - "tags": "IMAGE2TEXT,256K", - "max_tokens": 256000, - "model_type": [ - "image2text", - "chat" - ] - }, - { - "llm_name": "meta/llama-3.2-11b-vision-instruct", - "tags": "IMAGE2TEXT,128K", - "max_tokens": 131072, - "model_type": [ - "image2text", - "chat" - ] - }, { "llm_name": "meta/llama-3.2-90b-vision-instruct", "tags": "IMAGE2TEXT,128K", @@ -3526,51 +2875,193 @@ ] }, { - "llm_name": "microsoft/florence-2", - "tags": "IMAGE2TEXT,1K", - "max_tokens": 1024, + "llm_name": "meta/llama-3.3-70b-instruct", + "tags": "LLM,CHAT,128K", + "max_tokens": 131072, + "model_type": "chat" + }, + { + "llm_name": "meta/llama-guard-4-12b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "minimaxai/minimax-m3", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "mistralai/mistral-medium-3.5-128b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "mistralai/mistral-nemotron", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "moonshotai/kimi-k2.6", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/ising-calibration-1.5-31b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/llama-3.1-nemotron-nano-8b-v1", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/llama-3.1-nemotron-nano-vl-8b-v1", + "tags": "LLM,IMAGE2TEXT,8K", + "max_tokens": 8192, "model_type": [ - "image2text", - "chat" + "chat", + "vision" ] }, { - "llm_name": "microsoft/kosmos-2", - "tags": "IMAGE2TEXT,4K", + "llm_name": "nvidia/llama-3.1-nemotron-safety-guard-8b-v3", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/llama-3.3-nemotron-super-49b-v1", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/llama-3.3-nemotron-super-49b-v1.5", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-3-embed-1b", + "tags": "TEXT EMBEDDING,8K", + "max_tokens": 8192, + "model_type": "embedding" + }, + { + "llm_name": "nvidia/nemotron-3-nano-30b-a3b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-3-super-120b-a12b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-3-ultra-550b-a55b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-3.5-content-safety", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-mini-4b-instruct", + "tags": "LLM,CHAT,4K", "max_tokens": 4096, + "model_type": "chat" + }, + { + "llm_name": "nvidia/nemotron-nano-12b-v2-vl", + "tags": "LLM,IMAGE2TEXT,8K", + "max_tokens": 8192, "model_type": [ - "image2text", - "chat" + "chat", + "vision" ] }, { - "llm_name": "microsoft/phi-3-vision-128k-instruct", - "tags": "IMAGE2TEXT,128K", - "max_tokens": 131072, - "model_type": [ - "image2text", - "chat" - ] + "llm_name": "nvidia/nv-embed-v1", + "tags": "TEXT EMBEDDING", + "max_tokens": 4096, + "model_type": "embedding" }, { - "llm_name": "microsoft/phi-3.5-vision-instruct", - "tags": "IMAGE2TEXT,128K", - "max_tokens": 131072, - "model_type": [ - "image2text", - "chat" - ] + "llm_name": "nvidia/nv-embedcode-7b-v1", + "tags": "TEXT EMBEDDING,8K", + "max_tokens": 8192, + "model_type": "embedding" }, { - "llm_name": "nvidia/neva-22b", - "tags": "IMAGE2TEXT,1K", - "max_tokens": 1024, - "model_type": [ - "image2text", - "chat" - ] + "llm_name": "nvidia/nvidia-nemotron-nano-9b-v2", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "nvidia/riva-translate-4b-instruct-v2", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "openai/gpt-oss-120b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "openai/gpt-oss-20b", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "poolside/laguna-xs-2.1", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "stepfun-ai/step-3.7-flash", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "thinkingmachines/inkling", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" + }, + { + "llm_name": "z-ai/glm-5.2", + "tags": "LLM,CHAT,8K", + "max_tokens": 8192, + "model_type": "chat" } - ] + ], + "url": "https://integrate.api.nvidia.com/v1" }, { "name": "LM-Studio", diff --git a/conf/models/nvidia.json b/conf/models/nvidia.json index 9ef55bdbb7..00b1acbe99 100644 --- a/conf/models/nvidia.json +++ b/conf/models/nvidia.json @@ -11,26 +11,6 @@ }, "class": "nvidia", "models": [ - { - "name": "abacusai/dracarys-llama-3.1-70b-instruct", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "bytedance/seed-oss-36b-instruct", - "max_tokens": 32768, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, { "name": "deepseek-ai/deepseek-v4-flash", "max_tokens": 1048576, @@ -52,32 +32,12 @@ } }, { - "name": "nvidia/nv-embed-v1", + "name": "google/diffusiongemma-26b-a4b-it", "max_tokens": 8192, "model_types": [ - "embedding" + "chat" ] }, - { - "name": "google/codegemma-7b", - "max_tokens": 8192, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "google/gemma-2-2b-it", - "max_tokens": 8192, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, { "name": "google/gemma-4-31b-it", "max_tokens": 131072, @@ -88,6 +48,42 @@ "support": true } }, + { + "name": "meta/llama-3.1-70b-instruct", + "max_tokens": 8192, + "model_types": [ + "chat" + ] + }, + { + "name": "meta/llama-3.1-8b-instruct", + "max_tokens": 8192, + "model_types": [ + "chat" + ] + }, + { + "name": "meta/llama-3.2-11b-vision-instruct", + "max_tokens": 8192, + "model_types": [ + "chat", + "vision" + ] + }, + { + "name": "meta/llama-3.2-1b-instruct", + "max_tokens": 8192, + "model_types": [ + "chat" + ] + }, + { + "name": "meta/llama-3.2-3b-instruct", + "max_tokens": 8192, + "model_types": [ + "chat" + ] + }, { "name": "meta/llama-3.2-90b-vision-instruct", "max_tokens": 131072, @@ -100,54 +96,25 @@ } }, { - "name": "meta/llama-4-maverick-17b-128e-instruct", - "max_tokens": 1048576, + "name": "meta/llama-3.3-70b-instruct", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } + ] }, { - "name": "minimaxai/minimax-m2.5", - "max_tokens": 204800, + "name": "meta/llama-guard-4-12b", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } + ] }, { - "name": "minimaxai/minimax-m2.7", - "max_tokens": 204800, + "name": "minimaxai/minimax-m3", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "mistralai/mistral-7b-instruct-v0.3", - "max_tokens": 32768, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "mistralai/mistral-large-3-675b-instruct-2512", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } + ] }, { "name": "mistralai/mistral-medium-3.5-128b", @@ -182,58 +149,11 @@ } }, { - "name": "moonshotai/kimi-k2-instruct", - "max_tokens": 131072, + "name": "nvidia/ising-calibration-1.5-31b", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "moonshotai/kimi-k2-thinking", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } - }, - { - "name": "nvidia/gliner-pii", - "max_tokens": 4096, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "nvidia/llama-3.1-nemoguard-8b-content-safety", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } - }, - { - "name": "nvidia/llama-3.1-nemoguard-8b-topic-control", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "tools": { - "support": true - } + ] }, { "name": "nvidia/llama-3.1-nemotron-nano-8b-v1", @@ -245,6 +165,14 @@ "support": true } }, + { + "name": "nvidia/llama-3.1-nemotron-nano-vl-8b-v1", + "max_tokens": 8192, + "model_types": [ + "chat", + "vision" + ] + }, { "name": "nvidia/llama-3.1-nemotron-safety-guard-8b-v3", "max_tokens": 131072, @@ -255,27 +183,6 @@ "support": true } }, - { - "name": "nvidia/llama-3.1-nemotron-ultra-253b-v1", - "max_tokens": 131072, - "model_types": [ - "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } - }, - { - "name": "nvidia/llama-3.2-nemoretriever-1b-vlm-embed-v1", - "max_tokens": 8192, - "model_types": [ - "embedding" - ] - }, { "name": "nvidia/llama-3.3-nemotron-super-49b-v1", "max_tokens": 131072, @@ -300,6 +207,13 @@ "support": true } }, + { + "name": "nvidia/nemotron-3-embed-1b", + "max_tokens": 8192, + "model_types": [ + "embedding" + ] + }, { "name": "nvidia/nemotron-3-nano-30b-a3b", "max_tokens": 131072, @@ -336,14 +250,18 @@ } }, { - "name": "nvidia/nemotron-content-safety-reasoning-4b", + "name": "nvidia/nemotron-3-ultra-550b-a55b", "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } + ] + }, + { + "name": "nvidia/nemotron-3.5-content-safety", + "max_tokens": 8192, + "model_types": [ + "chat" + ] }, { "name": "nvidia/nemotron-mini-4b-instruct", @@ -355,6 +273,14 @@ "support": true } }, + { + "name": "nvidia/nemotron-nano-12b-v2-vl", + "max_tokens": 8192, + "model_types": [ + "chat", + "vision" + ] + }, { "name": "nvidia/nv-embed-v1", "max_tokens": 32768, @@ -363,33 +289,12 @@ ] }, { - "name": "nvidia/nv-embedqa-e5-v5", - "max_tokens": 512, + "name": "nvidia/nv-embedcode-7b-v1", + "max_tokens": 8192, "model_types": [ "embedding" ] }, - { - "name": "nvidia/nv-embedqa-mistral-7b-v2", - "max_tokens": 512, - "model_types": [ - "embedding" - ] - }, - { - "name": "nvidia/nv-rerankqa-mistral-4b-v3", - "max_tokens": 4096, - "model_types": [ - "rerank" - ] - }, - { - "name": "nvidia/llama-3.2-nv-rerankqa-1b-v2", - "max_tokens": 4096, - "model_types": [ - "rerank" - ] - }, { "name": "nvidia/nvidia-nemotron-nano-9b-v2", "max_tokens": 131072, @@ -401,14 +306,11 @@ } }, { - "name": "nvidia/riva-translate-4b-instruct-v1.1", - "max_tokens": 4096, + "name": "nvidia/riva-translate-4b-instruct-v2", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } + ] }, { "name": "openai/gpt-oss-120b", @@ -421,70 +323,39 @@ } }, { - "name": "qwen/qwen3.5-122b-a10b", - "max_tokens": 131072, + "name": "openai/gpt-oss-20b", + "max_tokens": 8192, "model_types": [ "chat" - ], - "tools": { - "support": true - } + ] }, { - "name": "qwen/qwen3-coder-480b-a35b-instruct", - "max_tokens": 262144, + "name": "poolside/laguna-xs-2.1", + "max_tokens": 8192, "model_types": [ "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } + ] }, { - "name": "z-ai/glm5", - "max_tokens": 131072, + "name": "stepfun-ai/step-3.7-flash", + "max_tokens": 8192, "model_types": [ "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } + ] }, { - "name": "z-ai/glm-5.1", - "max_tokens": 131072, + "name": "thinkingmachines/inkling", + "max_tokens": 8192, "model_types": [ "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } + ] }, { - "name": "z-ai/glm4.7", - "max_tokens": 131072, + "name": "z-ai/glm-5.2", + "max_tokens": 8192, "model_types": [ "chat" - ], - "thinking": { - "default_value": true, - "clear_thinking": true - }, - "tools": { - "support": true - } + ] } ] } diff --git a/internal/entity/models/nvidia.go b/internal/entity/models/nvidia.go index e469c60a1b..333b89da75 100644 --- a/internal/entity/models/nvidia.go +++ b/internal/entity/models/nvidia.go @@ -23,13 +23,24 @@ import ( "fmt" "io" "net/http" + "net/url" "ragflow/internal/common" + "sort" "strings" + "time" +) + +const ( + nvidiaHostedAPIHost = "integrate.api.nvidia.com" + nvidiaCatalogURL = "https://api.ngc.nvidia.com/v2/search/catalog/resources/ENDPOINT" + nvidiaCatalogPageSize = 500 ) // NvidiaModel implements ModelDriver for Nvidia type NvidiaModel struct { - baseModel BaseModel + baseModel BaseModel + catalogURL string + hostedAPIHost string } // NewNvidiaModel creates a new Nvidia model instance @@ -40,6 +51,8 @@ func NewNvidiaModel(baseURL map[string]string, urlSuffix URLSuffix) *NvidiaModel URLSuffix: urlSuffix, httpClient: NewDriverHTTPClient(), }, + catalogURL: nvidiaCatalogURL, + hostedAPIHost: nvidiaHostedAPIHost, } } @@ -521,22 +534,18 @@ func (n NvidiaModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]Li if err != nil { return nil, err } - baseURL := resolvedBaseURL - if baseURL == "" { - baseURL = resolvedBaseURL - } + url := fmt.Sprintf("%s/%s", strings.TrimRight(resolvedBaseURL, "/"), strings.TrimLeft(n.baseModel.URLSuffix.Models, "/")) - url := fmt.Sprintf("%s/%s", baseURL, n.baseModel.URLSuffix.Models) - - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + modelListCtx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) defer cancel() - req, err := http.NewRequestWithContext(ctx, "GET", url, nil) + req, err := http.NewRequestWithContext(modelListCtx, "GET", url, nil) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Accept", "application/json") resp, err := n.baseModel.httpClient.Do(req) if err != nil { @@ -562,7 +571,300 @@ func (n NvidiaModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]Li return nil, fmt.Errorf("invalid models list format") } - return ParseListModel(modelList), nil + var provider *Provider + if pm := GetProviderManager(); pm != nil { + provider = pm.FindProvider("NVIDIA") + } + models := parseNvidiaModelList(modelList, provider) + if n.usesHostedCatalog(resolvedBaseURL) { + catalogCtx, catalogCancel := context.WithTimeout(ctx, nonStreamCallTimeout) + catalog, catalogErr := n.fetchHostedCatalog(catalogCtx) + catalogCancel() + if catalogErr != nil { + return nil, catalogErr + } + models = filterNvidiaHostedModels(models, catalog, time.Now().UTC()) + } + if len(models) == 0 { + return nil, fmt.Errorf("Nvidia models API returned no usable models") + } + return models, nil +} + +type nvidiaCatalogLabel struct { + Key string `json:"key"` + Values []string `json:"values"` + UnresolvedValues []string `json:"unresolvedValues"` +} + +type nvidiaCatalogAttribute struct { + Key string `json:"key"` + Value string `json:"value"` +} + +type nvidiaCatalogResource struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + Labels []nvidiaCatalogLabel `json:"labels"` + Attributes []nvidiaCatalogAttribute `json:"attributes"` +} + +type nvidiaCatalogGroup struct { + GroupValue string `json:"groupValue"` + TotalCount int `json:"totalCount"` + Resources []nvidiaCatalogResource `json:"resources"` +} + +type nvidiaCatalogResponse struct { + Results []nvidiaCatalogGroup `json:"results"` +} + +func (n NvidiaModel) usesHostedCatalog(baseURL string) bool { + parsed, err := url.Parse(baseURL) + return err == nil && strings.EqualFold(parsed.Hostname(), n.hostedAPIHost) +} + +func (n NvidiaModel) fetchHostedCatalog(ctx context.Context) (*nvidiaCatalogResponse, error) { + parsed, err := url.Parse(n.catalogURL) + if err != nil { + return nil, fmt.Errorf("failed to parse Nvidia endpoint catalog URL: %w", err) + } + + var merged *nvidiaCatalogResponse + endpointGroupIndex := -1 + totalCount := -1 + for page := 0; ; page++ { + catalogQuery, err := json.Marshal(map[string]int{"page": page, "pageSize": nvidiaCatalogPageSize}) + if err != nil { + return nil, fmt.Errorf("failed to encode Nvidia endpoint catalog query: %w", err) + } + pageURL := *parsed + query := pageURL.Query() + query.Set("q", string(catalogQuery)) + query.Set("group-labels-by-labelset", "true") + pageURL.RawQuery = query.Encode() + + catalogPage, err := n.fetchHostedCatalogPage(ctx, pageURL.String()) + if err != nil { + return nil, err + } + + pageGroupIndex := -1 + for i := range catalogPage.Results { + if catalogPage.Results[i].GroupValue == "ENDPOINT" { + pageGroupIndex = i + break + } + } + if pageGroupIndex < 0 { + return nil, fmt.Errorf("Nvidia endpoint catalog response is missing the ENDPOINT group") + } + + pageGroup := catalogPage.Results[pageGroupIndex] + if pageGroup.TotalCount < 0 { + return nil, fmt.Errorf("Nvidia endpoint catalog returned invalid total count %d", pageGroup.TotalCount) + } + if merged == nil { + merged = catalogPage + endpointGroupIndex = pageGroupIndex + totalCount = pageGroup.TotalCount + merged.Results[endpointGroupIndex].Resources = nil + } else if pageGroup.TotalCount != totalCount { + return nil, fmt.Errorf("Nvidia endpoint catalog total count changed from %d to %d", totalCount, pageGroup.TotalCount) + } + + merged.Results[endpointGroupIndex].Resources = append(merged.Results[endpointGroupIndex].Resources, pageGroup.Resources...) + if len(pageGroup.Resources) < nvidiaCatalogPageSize { + break + } + } + + if merged == nil || endpointGroupIndex < 0 { + return nil, fmt.Errorf("Nvidia endpoint catalog returned no pages") + } + if len(merged.Results[endpointGroupIndex].Resources) < totalCount { + return nil, fmt.Errorf( + "Nvidia endpoint catalog returned %d of %d resources", + len(merged.Results[endpointGroupIndex].Resources), + totalCount, + ) + } + return merged, nil +} + +func (n NvidiaModel) fetchHostedCatalogPage(ctx context.Context, pageURL string) (*nvidiaCatalogResponse, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, pageURL, nil) + if err != nil { + return nil, fmt.Errorf("failed to create Nvidia endpoint catalog request: %w", err) + } + req.Header.Set("Accept", "application/json") + + resp, err := n.baseModel.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to request Nvidia endpoint catalog: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read Nvidia endpoint catalog: %w", err) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("Nvidia endpoint catalog error: %s, body: %s", resp.Status, string(body)) + } + + var catalog nvidiaCatalogResponse + if err = json.Unmarshal(body, &catalog); err != nil { + return nil, fmt.Errorf("failed to parse Nvidia endpoint catalog: %w", err) + } + return &catalog, nil +} + +func filterNvidiaHostedModels(models []ListModelResponse, catalog *nvidiaCatalogResponse, now time.Time) []ListModelResponse { + if catalog == nil { + return nil + } + + var resources []nvidiaCatalogResource + for _, group := range catalog.Results { + if group.GroupValue != "ENDPOINT" { + continue + } + if group.TotalCount < 0 || len(group.Resources) < group.TotalCount { + return nil + } + resources = group.Resources + break + } + if resources == nil { + return nil + } + + activeEndpoints := make(map[string]struct{}, len(resources)) + for _, resource := range resources { + if !nvidiaCatalogResourceIsActive(resource, now) { + continue + } + publishers := nvidiaCatalogLabelValues(resource, "publisher", true) + displayName := resource.DisplayName + if displayName == "" { + displayName = resource.Name + } + if len(publishers) == 0 || displayName == "" { + continue + } + activeEndpoints[nvidiaCatalogKey(publishers[0], displayName)] = struct{}{} + } + + filtered := make([]ListModelResponse, 0, len(models)) + for _, model := range models { + publisher, modelName, found := strings.Cut(model.Name, "/") + if !found { + continue + } + if _, ok := activeEndpoints[nvidiaCatalogKey(publisher, modelName)]; ok { + filtered = append(filtered, model) + } + } + return filtered +} + +func nvidiaCatalogResourceIsActive(resource nvidiaCatalogResource, now time.Time) bool { + if !containsNvidiaCatalogValue(nvidiaCatalogLabelValues(resource, "nimType", false), "Free Endpoint") { + return false + } + deprecation := "" + for _, attribute := range resource.Attributes { + if attribute.Key == "DEPRECATION" { + deprecation = attribute.Value + break + } + } + if deprecation == "" { + return true + } + cutoff, err := time.Parse("01/02/2006", deprecation) + if err != nil { + return false + } + today := time.Date(now.UTC().Year(), now.UTC().Month(), now.UTC().Day(), 0, 0, 0, 0, time.UTC) + return cutoff.After(today) +} + +func nvidiaCatalogLabelValues(resource nvidiaCatalogResource, key string, unresolved bool) []string { + for _, label := range resource.Labels { + if label.Key != key { + continue + } + if unresolved { + return label.UnresolvedValues + } + return label.Values + } + return nil +} + +func nvidiaCatalogKey(publisher, modelName string) string { + publisher = strings.ReplaceAll(strings.ToLower(strings.TrimSpace(publisher)), "_", "-") + return publisher + "\x00" + modelName +} + +func containsNvidiaCatalogValue(values []string, target string) bool { + for _, value := range values { + if value == target { + return true + } + } + return false +} + +func parseNvidiaModelList(modelList ModelList, provider *Provider) []ListModelResponse { + const defaultMaxTokens = 8192 + + models := make([]ListModelResponse, 0, len(modelList.Models)) + seen := make(map[string]struct{}, len(modelList.Models)) + for _, item := range modelList.Models { + modelName := strings.TrimSpace(item.ID) + if modelName == "" { + continue + } + if _, ok := seen[modelName]; ok { + continue + } + seen[modelName] = struct{}{} + + response := ListModelResponse{Name: modelName} + var preset *Model + if provider != nil { + for _, model := range provider.Models { + if strings.EqualFold(model.Name, modelName) { + preset = model + break + } + } + } + if preset != nil { + response.MaxTokens = preset.MaxTokens + response.ModelTypes = append([]string(nil), preset.ModelTypes...) + response.Thinking = preset.Thinking + response.MaxDimension = preset.MaxDimension + response.Dimensions = append([]int(nil), preset.Dimensions...) + } else { + maxTokens := defaultMaxTokens + response.MaxTokens = &maxTokens + response.ModelTypes = InferModelTypes(modelName) + } + if len(response.ModelTypes) == 0 { + response.ModelTypes = InferModelTypes(modelName) + } + + models = append(models, response) + } + + sort.Slice(models, func(i, j int) bool { + return models[i].Name < models[j].Name + }) + return models } func (n NvidiaModel) Balance(ctx context.Context, apiConfig *APIConfig) (map[string]interface{}, error) { diff --git a/internal/entity/models/nvidia_test.go b/internal/entity/models/nvidia_test.go new file mode 100644 index 0000000000..8896b88b42 --- /dev/null +++ b/internal/entity/models/nvidia_test.go @@ -0,0 +1,230 @@ +package models + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" +) + +func TestNvidiaListModelsUsesExactEndpointIDs(t *testing.T) { + const apiKey = "nvapi-test" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + t.Fatalf("method = %s, want GET", r.Method) + } + if r.URL.Path != "/v1/models" { + t.Fatalf("path = %s, want /v1/models", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer "+apiKey { + t.Fatalf("Authorization = %q, want Bearer token", got) + } + _ = json.NewEncoder(w).Encode(ModelList{ + Object: "list", + Models: []ModelListItem{ + {ID: "meta/llama-3.3-70b-instruct", Object: "model", OwnedBy: "meta"}, + {ID: " nvidia/nv-embed-v1 ", Object: "model", OwnedBy: "nvidia"}, + {ID: "meta/llama-3.3-70b-instruct", Object: "model", OwnedBy: "meta"}, + {ID: " ", Object: "model", OwnedBy: "nvidia"}, + }, + }) + })) + defer server.Close() + + driver := NewNvidiaModel( + map[string]string{"default": server.URL + "/v1"}, + URLSuffix{Models: "models"}, + ) + region := "default" + models, err := driver.ListModels(context.Background(), &APIConfig{ApiKey: ptr(apiKey), Region: ®ion}) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + + if got := joinModelNames(models, ","); got != "meta/llama-3.3-70b-instruct,nvidia/nv-embed-v1" { + t.Fatalf("model names = %q", got) + } + if got := models[0].ModelTypes; len(got) != 1 || got[0] != "chat" { + t.Fatalf("chat model types = %v, want [chat]", got) + } + if got := models[1].ModelTypes; len(got) != 1 || got[0] != "embedding" { + t.Fatalf("embedding model types = %v, want [embedding]", got) + } +} + +func TestParseNvidiaModelListPrefersPresetMetadata(t *testing.T) { + maxTokens := 131072 + provider := &Provider{Models: []*Model{ + { + Name: "nvidia/nemotron-3-super-120b-a12b", + MaxTokens: &maxTokens, + ModelTypes: []string{"chat"}, + Thinking: &ModelThinking{DefaultValue: true, ClearThinking: true}, + }, + }} + + models := parseNvidiaModelList(ModelList{Models: []ModelListItem{ + {ID: "nvidia/nemotron-3-super-120b-a12b", OwnedBy: "nvidia"}, + }}, provider) + if len(models) != 1 { + t.Fatalf("len(models) = %d, want 1", len(models)) + } + if models[0].MaxTokens == nil || *models[0].MaxTokens != maxTokens { + t.Fatalf("MaxTokens = %v, want %d", models[0].MaxTokens, maxTokens) + } + if models[0].Thinking == nil || !models[0].Thinking.DefaultValue { + t.Fatalf("Thinking = %#v, want preset metadata", models[0].Thinking) + } +} + +func TestParseNvidiaModelListInfersTypesForPresetWithoutTypes(t *testing.T) { + preset := &Model{Name: "nvidia/nv-embed-v1"} + models := parseNvidiaModelList(ModelList{Models: []ModelListItem{ + {ID: "nvidia/nv-embed-v1", OwnedBy: "nvidia"}, + }}, &Provider{Models: []*Model{preset}}) + + if len(models) != 1 { + t.Fatalf("len(models) = %d, want 1", len(models)) + } + if got := models[0].ModelTypes; len(got) != 1 || got[0] != "embedding" { + t.Fatalf("ModelTypes = %v, want [embedding]", got) + } + if preset.ModelTypes != nil { + t.Fatalf("preset ModelTypes mutated to %v", preset.ModelTypes) + } +} + +func TestNvidiaListModelsFiltersHostedCatalog(t *testing.T) { + const apiKey = "nvapi-test" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models": + _ = json.NewEncoder(w).Encode(ModelList{Models: []ModelListItem{ + {ID: "aisingapore/sea-lion-7b-instruct"}, + {ID: "meta/llama-3.1-8b-instruct"}, + {ID: "nvidia/nv-embed-v1"}, + {ID: "future/preserved-model"}, + }}) + case "/catalog": + _ = json.NewEncoder(w).Encode(nvidiaCatalogResponse{Results: []nvidiaCatalogGroup{{ + GroupValue: "ENDPOINT", + TotalCount: 4, + Resources: []nvidiaCatalogResource{ + {DisplayName: "sea-lion-7b-instruct", Labels: []nvidiaCatalogLabel{{Key: "publisher", UnresolvedValues: []string{"aisingapore"}}, {Key: "nimType", Values: []string{"Free Endpoint"}}}, Attributes: []nvidiaCatalogAttribute{{Key: "DEPRECATION", Value: "04/17/2026"}}}, + {DisplayName: "llama-3.1-8b-instruct", Labels: []nvidiaCatalogLabel{{Key: "publisher", UnresolvedValues: []string{"meta"}}, {Key: "nimType", Values: []string{"Free Endpoint"}}}}, + {DisplayName: "nv-embed-v1", Labels: []nvidiaCatalogLabel{{Key: "publisher", UnresolvedValues: []string{"nvidia"}}, {Key: "nimType", Values: []string{"Partner Endpoint"}}}}, + {DisplayName: "preserved-model", Labels: []nvidiaCatalogLabel{{Key: "publisher", UnresolvedValues: []string{"future"}}, {Key: "nimType", Values: []string{"Free Endpoint"}}}, Attributes: []nvidiaCatalogAttribute{{Key: "DEPRECATION", Value: "12/31/2099"}}}, + }, + }}}) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + driver := NewNvidiaModel(map[string]string{"default": server.URL + "/v1"}, URLSuffix{Models: "models"}) + driver.catalogURL = server.URL + "/catalog" + serverURL, _ := url.Parse(server.URL) + driver.hostedAPIHost = serverURL.Hostname() + region := "default" + models, err := driver.ListModels(context.Background(), &APIConfig{ApiKey: ptr(apiKey), Region: ®ion}) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + + if got := joinModelNames(models, ","); got != "future/preserved-model,meta/llama-3.1-8b-instruct" { + t.Fatalf("model names = %q, want active hosted models", got) + } +} + +func TestNvidiaFetchHostedCatalogPaginates(t *testing.T) { + resources := make([]nvidiaCatalogResource, nvidiaCatalogPageSize+1) + for i := range resources { + resources[i] = nvidiaCatalogResource{DisplayName: fmt.Sprintf("model-%d", i)} + } + + requestedPages := make([]int, 0, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var query struct { + Page int `json:"page"` + PageSize int `json:"pageSize"` + } + if err := json.Unmarshal([]byte(r.URL.Query().Get("q")), &query); err != nil { + t.Errorf("decode catalog query: %v", err) + return + } + requestedPages = append(requestedPages, query.Page) + if query.PageSize != nvidiaCatalogPageSize { + t.Errorf("pageSize = %d, want %d", query.PageSize, nvidiaCatalogPageSize) + return + } + + start := query.Page * query.PageSize + end := min(start+query.PageSize, len(resources)) + pageResources := []nvidiaCatalogResource{} + if start < len(resources) { + pageResources = resources[start:end] + } + _ = json.NewEncoder(w).Encode(nvidiaCatalogResponse{Results: []nvidiaCatalogGroup{{ + GroupValue: "ENDPOINT", + TotalCount: len(resources), + Resources: pageResources, + }}}) + })) + defer server.Close() + + driver := NewNvidiaModel(nil, URLSuffix{}) + driver.catalogURL = server.URL + catalog, err := driver.fetchHostedCatalog(t.Context()) + if err != nil { + t.Fatalf("fetchHostedCatalog() error = %v", err) + } + if got := requestedPages; len(got) != 2 || got[0] != 0 || got[1] != 1 { + t.Fatalf("requested pages = %v, want [0 1]", got) + } + if got := len(catalog.Results[0].Resources); got != len(resources) { + t.Fatalf("resources = %d, want %d", got, len(resources)) + } +} + +func TestNvidiaListModelsRejectsPartialHostedCatalog(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/models" { + _ = json.NewEncoder(w).Encode(ModelList{Models: []ModelListItem{{ID: "meta/llama-3.1-8b-instruct"}}}) + return + } + _ = json.NewEncoder(w).Encode(nvidiaCatalogResponse{Results: []nvidiaCatalogGroup{{ + GroupValue: "ENDPOINT", TotalCount: 2, + Resources: []nvidiaCatalogResource{{DisplayName: "llama-3.1-8b-instruct"}}, + }}}) + })) + defer server.Close() + + driver := NewNvidiaModel(map[string]string{"default": server.URL + "/v1"}, URLSuffix{Models: "models"}) + driver.catalogURL = server.URL + "/catalog" + serverURL, _ := url.Parse(server.URL) + driver.hostedAPIHost = serverURL.Hostname() + region := "default" + if _, err := driver.ListModels(context.Background(), &APIConfig{ApiKey: ptr("nvapi-test"), Region: ®ion}); err == nil { + t.Fatal("ListModels() error = nil, want partial catalog rejection") + } +} + +func TestNvidiaCatalogResourceRejectsMalformedDeprecation(t *testing.T) { + resource := nvidiaCatalogResource{ + Labels: []nvidiaCatalogLabel{{Key: "nimType", Values: []string{"Free Endpoint"}}}, + Attributes: []nvidiaCatalogAttribute{{Key: "DEPRECATION", Value: "not-a-date"}}, + } + + if nvidiaCatalogResourceIsActive(resource, time.Now()) { + t.Fatal("nvidiaCatalogResourceIsActive() = true, want false") + } +} + +func ptr[T any](value T) *T { + return &value +} diff --git a/internal/service/model_service.go b/internal/service/model_service.go index a2bfa8eb85..b7ea0954ce 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -385,6 +385,11 @@ func (m *ModelProviderService) ListSupportedModels(ctx context.Context, provider if err != nil { return nil, err } + if strings.EqualFold(provider.ProviderName, "NVIDIA") { + if err = m.reconcileNvidiaInstanceModels(ctx, dao.DB, provider, instance, modelList); err != nil { + return nil, err + } + } var result []map[string]interface{} for _, model := range modelList { @@ -400,6 +405,122 @@ func (m *ModelProviderService) ListSupportedModels(ctx context.Context, provider return result, nil } +func (m *ModelProviderService) reconcileNvidiaInstanceModels( + ctx context.Context, + db *gorm.DB, + provider *entity.TenantModelProvider, + instance *entity.TenantModelInstance, + remoteModels []modelModule.ListModelResponse, +) error { + if provider == nil || instance == nil || provider.ID == "" || instance.ID == "" || instance.ProviderID != provider.ID { + return errors.New("invalid NVIDIA provider instance scope") + } + + normalized := make([]modelModule.ListModelResponse, 0, len(remoteModels)) + seen := make(map[string]struct{}, len(remoteModels)) + for _, remote := range remoteModels { + remote.Name = strings.TrimSpace(remote.Name) + if remote.Name == "" { + continue + } + if _, ok := seen[remote.Name]; ok { + continue + } + seen[remote.Name] = struct{}{} + if len(remote.ModelTypes) == 0 { + remote.ModelTypes = modelModule.InferModelTypes(remote.Name) + } + normalized = append(normalized, remote) + } + if len(normalized) == 0 { + return errors.New("NVIDIA model discovery returned no usable models") + } + + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + existingModels, err := m.modelDAO.GetModelsByInstanceID(ctx, tx, instance.ID) + if err != nil { + return err + } + existingByName := make(map[string]*entity.TenantModel, len(existingModels)) + for _, existing := range existingModels { + existingByName[existing.ModelName] = existing + } + + for _, remote := range normalized { + maxTokens := 8192 + if remote.MaxTokens != nil && *remote.MaxTokens > 0 { + maxTokens = *remote.MaxTokens + } + modelType := int(entity.ModelTypeFromStrings(remote.ModelTypes)) + + if existing, ok := existingByName[remote.Name]; ok { + extra := make(map[string]interface{}) + if existing.Extra != "" { + if err := json.Unmarshal([]byte(existing.Extra), &extra); err != nil { + return fmt.Errorf("decode metadata for NVIDIA model %q: %w", remote.Name, err) + } + } + setDiscoveredModelMetadata(extra, remote, maxTokens) + extraBytes, err := json.Marshal(extra) + if err != nil { + return fmt.Errorf("encode metadata for NVIDIA model %q: %w", remote.Name, err) + } + if err = m.modelDAO.UpdateByID(ctx, tx, existing.ID, map[string]interface{}{ + "model_type": modelType, + "extra": string(extraBytes), + }); err != nil { + return err + } + delete(existingByName, remote.Name) + continue + } + + extra := map[string]interface{}{"verify": entity.ModelVerifyUnknown} + setDiscoveredModelMetadata(extra, remote, maxTokens) + extraBytes, err := json.Marshal(extra) + if err != nil { + return fmt.Errorf("encode metadata for NVIDIA model %q: %w", remote.Name, err) + } + if err = m.modelDAO.Create(ctx, tx, &entity.TenantModel{ + ID: utility.GenerateToken(), + ModelName: remote.Name, + ModelType: modelType, + ProviderID: provider.ID, + InstanceID: instance.ID, + Status: "active", + Extra: string(extraBytes), + }); err != nil { + return err + } + } + + staleIDs := make([]string, 0, len(existingByName)) + for _, stale := range existingByName { + staleIDs = append(staleIDs, stale.ID) + } + if len(staleIDs) > 0 { + if _, err := m.modelDAO.DeleteByIDs(ctx, tx, staleIDs); err != nil { + return err + } + } + return nil + }) +} + +func setDiscoveredModelMetadata(extra map[string]interface{}, model modelModule.ListModelResponse, maxTokens int) { + extra["max_tokens"] = maxTokens + if model.MaxDimension != nil { + extra["max_dimension"] = *model.MaxDimension + } + if len(model.Dimensions) > 0 { + extra["dimensions"] = model.Dimensions + } + if model.Thinking != nil { + extra["thinking"] = model.Thinking.DefaultValue + extra["clear_thinking"] = model.Thinking.ClearThinking + } +} + type CreateInstanceModelInfo struct { ModelName string `json:"model_name"` ModelTypes []string `json:"model_type"` diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index 6732480ee2..1011cbdacc 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -1,6 +1,8 @@ package service import ( + "context" + "encoding/json" "strings" "testing" @@ -245,6 +247,144 @@ func TestModelProviderServiceAlterModelRejectsWrongScopedModelID(t *testing.T) { } } +func TestReconcileNvidiaInstanceModelsAddsUpdatesAndDeletes(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + provider := &entity.TenantModelProvider{ID: "provider-nvidia", TenantID: "tenant-1", ProviderName: "NVIDIA"} + instance := &entity.TenantModelInstance{ID: "instance-nvidia", ProviderID: provider.ID, InstanceName: "default", APIKey: "nvapi-test", Status: "active", Extra: "{}"} + rows := []interface{}{ + provider, + instance, + &entity.TenantModel{ + ID: "keep-id", + ProviderID: provider.ID, + InstanceID: instance.ID, + ModelName: "nvidia/keep", + ModelType: int(entity.ModelTypeChat), + Status: "inactive", + Extra: `{"max_tokens":4096,"verify":"success","custom":"preserved"}`, + }, + &entity.TenantModel{ + ID: "stale-id", + ProviderID: provider.ID, + InstanceID: instance.ID, + ModelName: "nvidia/stale", + ModelType: int(entity.ModelTypeChat), + Status: "active", + Extra: `{}`, + }, + } + for _, row := range rows { + if err := db.Create(row).Error; err != nil { + t.Fatalf("seed %T: %v", row, err) + } + } + + maxTokens := 131072 + maxDimension := 2048 + remote := []modelModule.ListModelResponse{ + {Name: "nvidia/keep", MaxTokens: &maxTokens, ModelTypes: []string{"chat", "vision"}}, + {Name: "nvidia/new-embed", MaxTokens: ptrService(8192), MaxDimension: &maxDimension, Dimensions: []int{1024, 2048}, ModelTypes: []string{"embedding"}}, + } + + err := NewModelProviderService().reconcileNvidiaInstanceModels(context.Background(), db, provider, instance, remote) + if err != nil { + t.Fatalf("reconcileNvidiaInstanceModels() error = %v", err) + } + + var got []*entity.TenantModel + if err := db.Order("model_name").Find(&got).Error; err != nil { + t.Fatalf("list models: %v", err) + } + if len(got) != 2 || got[0].ModelName != "nvidia/keep" || got[1].ModelName != "nvidia/new-embed" { + t.Fatalf("models = %#v, want keep and new", got) + } + if got[0].ID != "keep-id" || got[0].Status != "inactive" { + t.Fatalf("retained model identity/status = %q/%q", got[0].ID, got[0].Status) + } + if got[0].ModelType != int(entity.ModelTypeChat|entity.ModelTypeImage2Text) { + t.Fatalf("retained model type = %d", got[0].ModelType) + } + var keepExtra map[string]interface{} + if err := json.Unmarshal([]byte(got[0].Extra), &keepExtra); err != nil { + t.Fatalf("decode retained extra: %v", err) + } + if keepExtra["custom"] != "preserved" || keepExtra["verify"] != "success" || int(keepExtra["max_tokens"].(float64)) != maxTokens { + t.Fatalf("retained extra = %#v", keepExtra) + } + var newExtra map[string]interface{} + if err := json.Unmarshal([]byte(got[1].Extra), &newExtra); err != nil { + t.Fatalf("decode new extra: %v", err) + } + if newExtra["verify"] != entity.ModelVerifyUnknown || int(newExtra["max_dimension"].(float64)) != maxDimension { + t.Fatalf("new extra = %#v", newExtra) + } +} + +func TestReconcileNvidiaInstanceModelsRejectsEmptyDiscoveryWithoutMutation(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + provider := &entity.TenantModelProvider{ID: "provider-nvidia", TenantID: "tenant-1", ProviderName: "NVIDIA"} + instance := &entity.TenantModelInstance{ID: "instance-nvidia", ProviderID: provider.ID, InstanceName: "default", Status: "active", Extra: "{}"} + existing := &entity.TenantModel{ID: "keep-id", ProviderID: provider.ID, InstanceID: instance.ID, ModelName: "nvidia/keep", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: "{}"} + for _, row := range []interface{}{provider, instance, existing} { + if err := db.Create(row).Error; err != nil { + t.Fatalf("seed %T: %v", row, err) + } + } + + err := NewModelProviderService().reconcileNvidiaInstanceModels(context.Background(), db, provider, instance, nil) + if err == nil { + t.Fatal("reconcileNvidiaInstanceModels() error = nil, want empty discovery error") + } + var count int64 + if err := db.Model(&entity.TenantModel{}).Where("id = ?", existing.ID).Count(&count).Error; err != nil { + t.Fatalf("count retained model: %v", err) + } + if count != 1 { + t.Fatalf("retained model count = %d, want 1", count) + } +} + +func TestReconcileNvidiaInstanceModelsRollsBackPartialRefresh(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + provider := &entity.TenantModelProvider{ID: "provider-nvidia", TenantID: "tenant-1", ProviderName: "NVIDIA"} + instance := &entity.TenantModelInstance{ID: "instance-nvidia", ProviderID: provider.ID, InstanceName: "default", Status: "active", Extra: "{}"} + existing := &entity.TenantModel{ + ID: "keep-id", + ProviderID: provider.ID, + InstanceID: instance.ID, + ModelName: "nvidia/keep", + ModelType: int(entity.ModelTypeChat), + Status: "active", + Extra: "{invalid-json", + } + for _, row := range []interface{}{provider, instance, existing} { + if err := db.Create(row).Error; err != nil { + t.Fatalf("seed %T: %v", row, err) + } + } + + remote := []modelModule.ListModelResponse{ + {Name: "nvidia/new", ModelTypes: []string{"chat"}}, + {Name: "nvidia/keep", ModelTypes: []string{"chat"}}, + } + err := NewModelProviderService().reconcileNvidiaInstanceModels(context.Background(), db, provider, instance, remote) + if err == nil { + t.Fatal("reconcileNvidiaInstanceModels() error = nil, want metadata error") + } + + var got []*entity.TenantModel + if err := db.Order("model_name").Find(&got).Error; err != nil { + t.Fatalf("list models: %v", err) + } + if len(got) != 1 || got[0].ID != existing.ID { + t.Fatalf("models after rollback = %#v, want only original model", got) + } +} + +func ptrService[T any](value T) *T { + return &value +} + func TestParseModelName(t *testing.T) { tests := []struct { name string diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 11274bb948..c3a917528d 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -1835,7 +1835,11 @@ class LiteLLMBase(ABC): logging.warning(f"Error: {error_code}. Retrying in {delay:.2f} seconds... (Attempt {attempt + 1}/{self.max_retries})") await asyncio.sleep(delay) return None - msg = f"{ERROR_PREFIX}: {error_code} - {str(e)}" + error_detail = str(e) + if self.provider == SupportedLiteLLMProvider.Nvidia and "function" in error_detail.lower() and "not found for account" in error_detail.lower(): + model_name = self.model_name.removeprefix(self.prefix) + error_detail = f"NVIDIA hosted endpoint '{model_name}' is unavailable or deprecated; refresh the provider model list and select an active Free Endpoint. Original error: {error_detail}" + msg = f"{ERROR_PREFIX}: {error_code} - {error_detail}" logging.error(f"async_chat_streamly giving up: {msg}") return msg @@ -2227,9 +2231,12 @@ class LiteLLMBase(ABC): "model": self.model_name, "messages": history, "api_key": self.api_key, - "num_retries": self.max_retries, **kwargs, } + if self.provider == SupportedLiteLLMProvider.Nvidia: + completion_args["num_retries"] = 0 + else: + completion_args.setdefault("num_retries", self.max_retries) # Forward the originating session/user as the OpenAI-standard `user` field so # providers (OpenAI, OpenRouter, ...) receive it in the request body and # upstream activity can be correlated back to the session. An explicit diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index c668186e10..30080854fc 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -24,7 +24,6 @@ from abc import ABC from copy import deepcopy from io import BytesIO from pathlib import Path -from urllib.parse import urljoin from json.decoder import JSONDecodeError import requests @@ -1101,86 +1100,13 @@ class GeminiCV(Base): tmp_path.unlink() -class NvidiaCV(Base): +class NvidiaCV(GptV4): _FACTORY_NAME = "NVIDIA" - def __init__(self, key, model_name, lang="Chinese", base_url="https://ai.api.nvidia.com/v1/vlm", **kwargs): + def __init__(self, key, model_name, lang="Chinese", base_url="https://integrate.api.nvidia.com/v1", **kwargs): if not base_url: - base_url = ("https://ai.api.nvidia.com/v1/vlm",) - self.lang = lang - factory, llm_name = model_name.split("/") - if factory != "liuhaotian": - self.base_url = urljoin(base_url, f"{factory}/{llm_name}") - else: - self.base_url = urljoin(f"{base_url}/community", llm_name.replace("-v1.6", "16")) - self.key = key - Base.__init__(self, **kwargs) - - def _image_prompt(self, text, images): - if not images: - return text - htmls = "" - for img in images: - htmls += ' '.format(f"data:image/jpeg;base64,{img}" if img[:4] != "data" else img) - return text + htmls - - def describe(self, image): - b64 = self.image2base64(image) - response = requests.post( - url=self.base_url, - headers={ - "accept": "application/json", - "content-type": "application/json", - "Authorization": f"Bearer {self.key}", - }, - json={"messages": self.prompt(b64)}, - timeout=60, - ) - response = response.json() - return ( - response["choices"][0]["message"]["content"].strip(), - total_token_count_from_response(response), - ) - - def _request(self, msg, gen_conf=None): - gen_conf = dict(gen_conf or {}) - response = requests.post( - url=self.base_url, - headers={ - "accept": "application/json", - "content-type": "application/json", - "Authorization": f"Bearer {self.key}", - }, - json={"messages": msg, **gen_conf}, - timeout=60, - ) - return response.json() - - def describe_with_prompt(self, image, prompt=None): - b64 = self.image2base64(image) - vision_prompt = self.vision_llm_prompt(b64, prompt) if prompt else self.vision_llm_prompt(b64) - response = self._request(vision_prompt) - return (response["choices"][0]["message"]["content"].strip(), total_token_count_from_response(response)) - - async def async_chat(self, system, history, gen_conf, images=None, **kwargs): - try: - response = await thread_pool_exec(self._request, self._form_history(system, history, images), gen_conf) - return (response["choices"][0]["message"]["content"].strip(), total_token_count_from_response(response)) - except Exception as e: - return "**ERROR**: " + str(e), 0 - - async def async_chat_streamly(self, system, history, gen_conf, images=None, **kwargs): - total_tokens = 0 - try: - response = await thread_pool_exec(self._request, self._form_history(system, history, images), gen_conf) - cnt = response["choices"][0]["message"]["content"] - total_tokens += total_token_count_from_response(response) - for resp in cnt: - yield resp - except Exception as e: - yield "\n**ERROR**: " + str(e) - - yield total_tokens + base_url = "https://integrate.api.nvidia.com/v1" + super().__init__(key, model_name, lang=lang, base_url=base_url, **kwargs) class AnthropicCV(Base): diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index c182631d50..560b324d6a 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -33,7 +33,7 @@ from common.aimlapi_utils import attribution_headers from common.exceptions import ModelException from common.token_utils import num_tokens_from_string, truncate, total_token_count_from_response from rag.llm.key_utils import _normalize_replicate_key -from rag.utils.url_utils import ensure_v1 +from rag.utils.url_utils import append_api_path, ensure_v1 import logging import base64 @@ -821,11 +821,11 @@ class GeminiEmbed(Base): class NvidiaEmbed(Base): _FACTORY_NAME = "NVIDIA" - def __init__(self, key, model_name, base_url="https://integrate.api.nvidia.com/v1/embeddings"): + def __init__(self, key, model_name, base_url="https://integrate.api.nvidia.com/v1"): if not base_url: - base_url = "https://integrate.api.nvidia.com/v1/embeddings" + base_url = "https://integrate.api.nvidia.com/v1" self.api_key = key - self.base_url = ensure_v1(base_url) + self.base_url = append_api_path(ensure_v1(base_url), "embeddings") self.headers = { "accept": "application/json", "Content-Type": "application/json", diff --git a/rag/llm/model_meta.py b/rag/llm/model_meta.py index cb518a3f36..8e5a06ae67 100644 --- a/rag/llm/model_meta.py +++ b/rag/llm/model_meta.py @@ -17,6 +17,7 @@ import json import logging import aiohttp from abc import ABC +from datetime import datetime, timezone from urllib.parse import urlparse from json.decoder import JSONDecodeError from typing import ClassVar @@ -461,6 +462,182 @@ class OpenAIAPICompatible(Base): return model_list +class NVIDIA(OpenAIAPICompatible): + _FACTORY_NAME = "NVIDIA" + _HOSTED_API_HOST = "integrate.api.nvidia.com" + _CATALOG_URL = "https://api.ngc.nvidia.com/v2/search/catalog/resources/ENDPOINT" + _CATALOG_PAGE_SIZE = 500 + _MODEL_LIST_TIMEOUT = aiohttp.ClientTimeout(total=30) + + @staticmethod + def _normalized_publisher(value): + return value.strip().lower().replace("_", "-") + + @classmethod + def _catalog_resources(cls, catalog): + catalogs = catalog if isinstance(catalog, list) else [catalog] + resources = [] + total_count = None + for catalog_page in catalogs: + page_resources, page_total_count = cls._catalog_page(catalog_page) + if page_resources is None: + return None + if total_count is None: + total_count = page_total_count + elif page_total_count != total_count: + return None + resources.extend(page_resources) + if total_count is None or len(resources) < total_count: + return None + return resources + + @staticmethod + def _catalog_page(catalog): + if not isinstance(catalog, dict): + return None, None + for group in catalog.get("results", []): + if not isinstance(group, dict) or group.get("groupValue") != "ENDPOINT": + continue + resources = group.get("resources") + total_count = group.get("totalCount") + if not isinstance(resources, list) or not isinstance(total_count, int) or total_count < 0: + return None, None + return resources, total_count + return None, None + + @staticmethod + def _label_values(resource, key): + for label in resource.get("labels", []): + if isinstance(label, dict) and label.get("key") == key: + values = label.get("values", []) + return values if isinstance(values, list) else [] + return [] + + @staticmethod + def _unresolved_label_values(resource, key): + for label in resource.get("labels", []): + if isinstance(label, dict) and label.get("key") == key: + values = label.get("unresolvedValues", []) + return values if isinstance(values, list) else [] + return [] + + @staticmethod + def _attribute_value(resource, key): + for attribute in resource.get("attributes", []): + if isinstance(attribute, dict) and attribute.get("key") == key: + return attribute.get("value") + return None + + @classmethod + def _is_active_free_endpoint(cls, resource, now): + if "Free Endpoint" not in cls._label_values(resource, "nimType"): + return False + deprecation = cls._attribute_value(resource, "DEPRECATION") + if not deprecation: + return True + try: + cutoff = datetime.strptime(deprecation, "%m/%d/%Y").replace(tzinfo=timezone.utc) + except (TypeError, ValueError): + return False + return cutoff.date() > now.date() + + @classmethod + def _filter_hosted_models(cls, models, catalog, now=None): + resources = cls._catalog_resources(catalog) + if resources is None: + return [] + + now = now or datetime.now(timezone.utc) + active_endpoints = cls._active_endpoint_keys(resources, now) + + filtered = [] + for model in models: + publisher, separator, model_name = model["name"].partition("/") + if separator and (cls._normalized_publisher(publisher), model_name) in active_endpoints: + filtered.append(model) + return filtered + + @classmethod + def _active_endpoint_keys(cls, resources, now): + active_endpoints = set() + for resource in resources: + if not isinstance(resource, dict) or not cls._is_active_free_endpoint(resource, now): + continue + publishers = cls._unresolved_label_values(resource, "publisher") + display_name = resource.get("displayName") or resource.get("name") + if not publishers or not isinstance(display_name, str) or not display_name: + continue + active_endpoints.add((cls._normalized_publisher(publishers[0]), display_name)) + return active_endpoints + + def _uses_hosted_catalog(self): + model_list_url = self._get_model_list_url() + return bool(model_list_url and urlparse(model_list_url).hostname == self._HOSTED_API_HOST) + + async def get_model_list(self): + if not self._uses_hosted_catalog(): + return await super().get_model_list() + + async with aiohttp.ClientSession(timeout=self._MODEL_LIST_TIMEOUT) as session: + async with session.get(self._get_model_list_url(), headers={"Authorization": f"Bearer {self._get_api_key()}"}) as response: + response.raise_for_status() + raw_models = await response.json() + + catalog_pages = [] + resource_count = 0 + total_count = None + page = 0 + while total_count is None or resource_count < total_count: + catalog_query = {"page": page, "pageSize": self._CATALOG_PAGE_SIZE} + async with session.get( + self._CATALOG_URL, + params={"q": json.dumps(catalog_query, separators=(",", ":")), "group-labels-by-labelset": "true"}, + ) as response: + response.raise_for_status() + catalog_page = await response.json() + + page_resources, page_total_count = self._catalog_page(catalog_page) + if page_resources is None: + raise ValueError("NVIDIA endpoint catalog response is missing a valid ENDPOINT group") + if total_count is None: + total_count = page_total_count + elif page_total_count != total_count: + raise ValueError(f"NVIDIA endpoint catalog total count changed from {total_count} to {page_total_count}") + + catalog_pages.append(catalog_page) + resource_count += len(page_resources) + if resource_count < total_count and len(page_resources) < self._CATALOG_PAGE_SIZE: + raise ValueError(f"NVIDIA endpoint catalog returned {resource_count} of {total_count} resources") + page += 1 + + formatted_models = self._format_model_list(raw_models) + resources = self._catalog_resources(catalog_pages) + if resources is None: + raise ValueError("NVIDIA endpoint catalog response is incomplete") + now = datetime.now(timezone.utc) + filtered_models = self._filter_hosted_models(formatted_models, catalog_pages, now) + raw_model_items = raw_models.get("data") if isinstance(raw_models, dict) else raw_models + raw_model_count = len(raw_model_items) if isinstance(raw_model_items, list) else 0 + logging.info( + "[NVIDIA] Hosted model discovery succeeded: raw_models=%d active_endpoints=%d filtered_models=%d", + raw_model_count, + len(self._active_endpoint_keys(resources, now)), + len(filtered_models), + ) + return filtered_models + + def _format_model_list(self, raw_model_list): + models = super()._format_model_list(raw_model_list) + unique_models = {} + for model in models: + model_name = model["name"].strip() + if not model_name or model_name in unique_models: + continue + model["name"] = model_name + unique_models[model_name] = model + return [unique_models[name] for name in sorted(unique_models)] + + class GreenPT(OpenAIAPICompatible): """Discover and classify GreenPT models from the live catalog.""" diff --git a/rag/utils/url_utils.py b/rag/utils/url_utils.py index 8fd6c7138b..f2c2bce19d 100644 --- a/rag/utils/url_utils.py +++ b/rag/utils/url_utils.py @@ -48,3 +48,17 @@ def ensure_v1(url: str) -> str: # No versioned segment found – append /v1 new_path = (path + "/v1") if path else "/v1" return urlunparse((parsed.scheme, parsed.netloc, new_path, parsed.params, parsed.query, parsed.fragment)) + + +def append_api_path(url: str, endpoint: str) -> str: + """Append an API endpoint path exactly once while preserving the base path.""" + if not url: + return url + + parsed = urlparse(url) + path = parsed.path.rstrip("/") + endpoint_path = f"/{endpoint.strip('/')}" + if not path.endswith(endpoint_path): + path = f"{path}{endpoint_path}" + + return urlunparse((parsed.scheme, parsed.netloc, path, parsed.params, parsed.query, parsed.fragment))