From 74c5d2bada441b285399bb5b6aa8d9fad83bc22d Mon Sep 17 00:00:00 2001 From: C10H14N2O5 <100066858+C10H14N2O5@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:56:24 +0800 Subject: [PATCH 1/3] fix: update deprecated NVIDIA retrieval models --- astrbot/core/config/default.py | 6 +- .../sources/nvidia_embedding_source.py | 2 +- .../provider/sources/nvidia_rerank_source.py | 2 +- tests/test_nvidia_embedding_source.py | 61 ++++++++++++++++ tests/test_nvidia_rerank_source.py | 69 +++++++++++++++++++ 5 files changed, 135 insertions(+), 5 deletions(-) create mode 100644 tests/test_nvidia_embedding_source.py create mode 100644 tests/test_nvidia_rerank_source.py diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 780e070098..942bcda65f 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1904,9 +1904,9 @@ "enable": True, "embedding_api_key": "", "embedding_api_base": "https://integrate.api.nvidia.com/v1", - "embedding_model": "nvidia/llama-nemotron-embed-1b-v2", + "embedding_model": "nvidia/nemotron-3-embed-1b", "input_type": "passage", - "embedding_dimensions": 1024, + "embedding_dimensions": 2048, "timeout": 20, "proxy": "", }, @@ -1982,7 +1982,7 @@ "enable": True, "nvidia_rerank_api_key": "", "nvidia_rerank_api_base": "https://ai.api.nvidia.com/v1/retrieval", - "nvidia_rerank_model": "nv-rerank-qa-mistral-4b:1", + "nvidia_rerank_model": "nvidia/llama-nemotron-rerank-vl-1b-v2", "nvidia_rerank_model_endpoint": "/reranking", "timeout": 20, "nvidia_rerank_truncate": "", diff --git a/astrbot/core/provider/sources/nvidia_embedding_source.py b/astrbot/core/provider/sources/nvidia_embedding_source.py index b13ee0d201..475fadeecc 100644 --- a/astrbot/core/provider/sources/nvidia_embedding_source.py +++ b/astrbot/core/provider/sources/nvidia_embedding_source.py @@ -28,7 +28,7 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: ) self.timeout = int(provider_config.get("timeout", 20)) self.model = provider_config.get( - "embedding_model", "nvidia/llama-nemotron-embed-1b-v2" + "embedding_model", "nvidia/nemotron-3-embed-1b" ) self.input_type = provider_config.get("input_type", "passage") diff --git a/astrbot/core/provider/sources/nvidia_rerank_source.py b/astrbot/core/provider/sources/nvidia_rerank_source.py index c168da4a6e..6f3f92c1cd 100644 --- a/astrbot/core/provider/sources/nvidia_rerank_source.py +++ b/astrbot/core/provider/sources/nvidia_rerank_source.py @@ -19,7 +19,7 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: ).rstrip("/") self.timeout = provider_config.get("timeout", 20) self.model = provider_config.get( - "nvidia_rerank_model", "nv-rerank-qa-mistral-4b:1" + "nvidia_rerank_model", "nvidia/llama-nemotron-rerank-vl-1b-v2" ) self.model_endpoint = provider_config.get( "nvidia_rerank_model_endpoint", "/reranking" diff --git a/tests/test_nvidia_embedding_source.py b/tests/test_nvidia_embedding_source.py new file mode 100644 index 0000000000..76a9b65129 --- /dev/null +++ b/tests/test_nvidia_embedding_source.py @@ -0,0 +1,61 @@ +from astrbot.core.config.default import CONFIG_METADATA_2 +from astrbot.core.provider.sources.nvidia_embedding_source import ( + NvidiaEmbeddingProvider, +) + +NEW_MODEL = "nvidia/nemotron-3-embed-1b" +OLD_MODEL = "nvidia/llama-nemotron-embed-1b-v2" + + +def test_nvidia_embedding_config_template_uses_new_model_and_dimension(): + templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ + "config_template" + ] + + assert templates["NVIDIA Embedding"]["embedding_model"] == NEW_MODEL + assert templates["NVIDIA Embedding"]["embedding_dimensions"] == 2048 + + +def test_nvidia_embedding_provider_uses_new_fallback_model(): + provider = NvidiaEmbeddingProvider({}, {}) + + assert provider.model == NEW_MODEL + assert provider.get_model() == NEW_MODEL + + +def test_nvidia_embedding_provider_preserves_explicit_old_model(): + provider = NvidiaEmbeddingProvider( + { + "embedding_model": OLD_MODEL, + "embedding_dimensions": 1024, + }, + {}, + ) + + assert provider.model == OLD_MODEL + assert provider.get_dim() == 1024 + + +def test_nvidia_embedding_new_model_uses_existing_api_contract(): + provider = NvidiaEmbeddingProvider( + { + "embedding_model": NEW_MODEL, + "input_type": "passage", + }, + {}, + ) + + assert provider._build_payload(["first", "second"]) == { + "input": ["first", "second"], + "model": NEW_MODEL, + "input_type": "passage", + "encoding_format": "float", + } + assert provider._parse_response( + { + "data": [ + {"index": 0, "embedding": [0.1, 0.2]}, + {"index": 1, "embedding": [0.3, 0.4]}, + ] + } + ) == [[0.1, 0.2], [0.3, 0.4]] diff --git a/tests/test_nvidia_rerank_source.py b/tests/test_nvidia_rerank_source.py new file mode 100644 index 0000000000..02788314ee --- /dev/null +++ b/tests/test_nvidia_rerank_source.py @@ -0,0 +1,69 @@ +from astrbot.core.config.default import CONFIG_METADATA_2 +from astrbot.core.provider.sources.nvidia_rerank_source import NvidiaRerankProvider + +NEW_MODEL = "nvidia/llama-nemotron-rerank-vl-1b-v2" +OLD_MODEL = "nv-rerank-qa-mistral-4b:1" + + +def test_nvidia_rerank_config_template_uses_new_model(): + templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ + "config_template" + ] + + assert templates["NVIDIA Rerank"]["nvidia_rerank_model"] == NEW_MODEL + + +def test_nvidia_rerank_provider_uses_new_fallback_model(): + provider = NvidiaRerankProvider({}, {}) + + assert provider.model == NEW_MODEL + assert provider.get_model() == NEW_MODEL + + +def test_nvidia_rerank_provider_preserves_explicit_old_model(): + provider = NvidiaRerankProvider({"nvidia_rerank_model": OLD_MODEL}, {}) + + assert provider.model == OLD_MODEL + assert provider._get_endpoint() == ( + "https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking" + ) + + +def test_nvidia_rerank_new_model_uses_existing_api_contract(): + provider = NvidiaRerankProvider( + { + "nvidia_rerank_model": NEW_MODEL, + "nvidia_rerank_truncate": "END", + }, + {}, + ) + + assert provider._get_endpoint() == ( + "https://ai.api.nvidia.com/v1/retrieval/nvidia/" + "llama-nemotron-rerank-vl-1b-v2/reranking" + ) + assert provider._build_payload("query", ["first", "second"]) == { + "model": NEW_MODEL, + "query": {"text": "query"}, + "passages": [{"text": "first"}, {"text": "second"}], + "truncate": "END", + } + + +def test_nvidia_rerank_parses_official_rankings_response(): + provider = NvidiaRerankProvider({}, {}) + + results = provider._parse_results( + { + "rankings": [ + {"index": 1, "logit": -0.25}, + {"index": 0, "logit": 0.75}, + ] + }, + top_n=None, + ) + + assert [(result.index, result.relevance_score) for result in results] == [ + (0, 0.75), + (1, -0.25), + ] From 5cb0b6b966f9e39114cfda594cb24d36d6d10b80 Mon Sep 17 00:00:00 2001 From: Soulter <37870767+Soulter@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:16:28 +0800 Subject: [PATCH 2/3] Delete tests/test_nvidia_embedding_source.py --- tests/test_nvidia_embedding_source.py | 61 --------------------------- 1 file changed, 61 deletions(-) delete mode 100644 tests/test_nvidia_embedding_source.py diff --git a/tests/test_nvidia_embedding_source.py b/tests/test_nvidia_embedding_source.py deleted file mode 100644 index 76a9b65129..0000000000 --- a/tests/test_nvidia_embedding_source.py +++ /dev/null @@ -1,61 +0,0 @@ -from astrbot.core.config.default import CONFIG_METADATA_2 -from astrbot.core.provider.sources.nvidia_embedding_source import ( - NvidiaEmbeddingProvider, -) - -NEW_MODEL = "nvidia/nemotron-3-embed-1b" -OLD_MODEL = "nvidia/llama-nemotron-embed-1b-v2" - - -def test_nvidia_embedding_config_template_uses_new_model_and_dimension(): - templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ - "config_template" - ] - - assert templates["NVIDIA Embedding"]["embedding_model"] == NEW_MODEL - assert templates["NVIDIA Embedding"]["embedding_dimensions"] == 2048 - - -def test_nvidia_embedding_provider_uses_new_fallback_model(): - provider = NvidiaEmbeddingProvider({}, {}) - - assert provider.model == NEW_MODEL - assert provider.get_model() == NEW_MODEL - - -def test_nvidia_embedding_provider_preserves_explicit_old_model(): - provider = NvidiaEmbeddingProvider( - { - "embedding_model": OLD_MODEL, - "embedding_dimensions": 1024, - }, - {}, - ) - - assert provider.model == OLD_MODEL - assert provider.get_dim() == 1024 - - -def test_nvidia_embedding_new_model_uses_existing_api_contract(): - provider = NvidiaEmbeddingProvider( - { - "embedding_model": NEW_MODEL, - "input_type": "passage", - }, - {}, - ) - - assert provider._build_payload(["first", "second"]) == { - "input": ["first", "second"], - "model": NEW_MODEL, - "input_type": "passage", - "encoding_format": "float", - } - assert provider._parse_response( - { - "data": [ - {"index": 0, "embedding": [0.1, 0.2]}, - {"index": 1, "embedding": [0.3, 0.4]}, - ] - } - ) == [[0.1, 0.2], [0.3, 0.4]] From f62b8a5c3054f496bae3d0851791c43869acf049 Mon Sep 17 00:00:00 2001 From: Soulter <37870767+Soulter@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:16:40 +0800 Subject: [PATCH 3/3] Delete tests/test_nvidia_rerank_source.py --- tests/test_nvidia_rerank_source.py | 69 ------------------------------ 1 file changed, 69 deletions(-) delete mode 100644 tests/test_nvidia_rerank_source.py diff --git a/tests/test_nvidia_rerank_source.py b/tests/test_nvidia_rerank_source.py deleted file mode 100644 index 02788314ee..0000000000 --- a/tests/test_nvidia_rerank_source.py +++ /dev/null @@ -1,69 +0,0 @@ -from astrbot.core.config.default import CONFIG_METADATA_2 -from astrbot.core.provider.sources.nvidia_rerank_source import NvidiaRerankProvider - -NEW_MODEL = "nvidia/llama-nemotron-rerank-vl-1b-v2" -OLD_MODEL = "nv-rerank-qa-mistral-4b:1" - - -def test_nvidia_rerank_config_template_uses_new_model(): - templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ - "config_template" - ] - - assert templates["NVIDIA Rerank"]["nvidia_rerank_model"] == NEW_MODEL - - -def test_nvidia_rerank_provider_uses_new_fallback_model(): - provider = NvidiaRerankProvider({}, {}) - - assert provider.model == NEW_MODEL - assert provider.get_model() == NEW_MODEL - - -def test_nvidia_rerank_provider_preserves_explicit_old_model(): - provider = NvidiaRerankProvider({"nvidia_rerank_model": OLD_MODEL}, {}) - - assert provider.model == OLD_MODEL - assert provider._get_endpoint() == ( - "https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking" - ) - - -def test_nvidia_rerank_new_model_uses_existing_api_contract(): - provider = NvidiaRerankProvider( - { - "nvidia_rerank_model": NEW_MODEL, - "nvidia_rerank_truncate": "END", - }, - {}, - ) - - assert provider._get_endpoint() == ( - "https://ai.api.nvidia.com/v1/retrieval/nvidia/" - "llama-nemotron-rerank-vl-1b-v2/reranking" - ) - assert provider._build_payload("query", ["first", "second"]) == { - "model": NEW_MODEL, - "query": {"text": "query"}, - "passages": [{"text": "first"}, {"text": "second"}], - "truncate": "END", - } - - -def test_nvidia_rerank_parses_official_rankings_response(): - provider = NvidiaRerankProvider({}, {}) - - results = provider._parse_results( - { - "rankings": [ - {"index": 1, "logit": -0.25}, - {"index": 0, "logit": 0.75}, - ] - }, - top_n=None, - ) - - assert [(result.index, result.relevance_score) for result in results] == [ - (0, 0.75), - (1, -0.25), - ]