From 3c04cfec2f71bca26007612803e174e0b2eb3cd9 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Thu, 14 Nov 2024 14:53:54 +0100 Subject: [PATCH] fix: Fixes to NvidiaRanker (#1191) * Fixes to NvidiaRanker * Add inits and headers * More headers * updates * Reactivate test * Fix tests * Reenable test and add test --- integrations/nvidia/pyproject.toml | 2 +- .../src/haystack_integrations/__init__.py | 3 ++ .../components/__init__.py | 3 ++ .../components/embedders/__init__.py | 3 ++ .../components/embedders/nvidia/__init__.py | 4 ++ .../embedders/nvidia/document_embedder.py | 11 +++-- .../embedders/nvidia/text_embedder.py | 7 ++- .../components/embedders/nvidia/truncate.py | 4 ++ .../components/generators/__init__.py | 3 ++ .../components/generators/nvidia/__init__.py | 1 + .../components/generators/nvidia/generator.py | 1 + .../components/rankers/__init__.py | 3 ++ .../components/rankers/nvidia/__init__.py | 4 ++ .../components/rankers/nvidia/ranker.py | 27 ++++++---- .../components/rankers/nvidia/truncate.py | 4 ++ .../haystack_integrations/utils/__init__.py | 3 ++ .../utils/nvidia/__init__.py | 4 ++ .../utils/nvidia/nim_backend.py | 4 ++ .../utils/nvidia/utils.py | 8 ++- integrations/nvidia/tests/__init__.py | 1 + integrations/nvidia/tests/conftest.py | 4 ++ integrations/nvidia/tests/test_base_url.py | 4 ++ .../nvidia/tests/test_document_embedder.py | 27 ++++++++-- .../tests/test_embedding_truncate_mode.py | 4 ++ integrations/nvidia/tests/test_generator.py | 1 + integrations/nvidia/tests/test_ranker.py | 49 +++++++++++++++++++ .../nvidia/tests/test_text_embedder.py | 21 +++++++- 27 files changed, 188 insertions(+), 22 deletions(-) create mode 100644 integrations/nvidia/src/haystack_integrations/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/generators/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/rankers/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/utils/__init__.py diff --git a/integrations/nvidia/pyproject.toml b/integrations/nvidia/pyproject.toml index 7f0048c1b..586b50848 100644 --- a/integrations/nvidia/pyproject.toml +++ b/integrations/nvidia/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["haystack-ai", "requests"] +dependencies = ["haystack-ai", "requests", "tqdm"] [project.urls] Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme" diff --git a/integrations/nvidia/src/haystack_integrations/__init__.py b/integrations/nvidia/src/haystack_integrations/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/components/__init__.py b/integrations/nvidia/src/haystack_integrations/components/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/__init__.py b/integrations/nvidia/src/haystack_integrations/components/embedders/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py index bc2d9372c..827ad7dc6 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from .document_embedder import NvidiaDocumentEmbedder from .text_embedder import NvidiaTextEmbedder from .truncate import EmbeddingTruncateMode diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py index d746a75f4..606ec78fd 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import warnings from typing import Any, Dict, List, Optional, Tuple, Union @@ -5,10 +9,9 @@ from haystack.utils import Secret, deserialize_secrets_inplace from tqdm import tqdm +from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation -from .truncate import EmbeddingTruncateMode - _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia" @@ -167,7 +170,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "NvidiaDocumentEmbedder": :returns: The deserialized component. """ - deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) + init_parameters = data.get("init_parameters", {}) + if init_parameters: + deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) return default_from_dict(cls, data) def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]: diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py index 22bed8197..4b7072f33 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py @@ -1,13 +1,16 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import warnings from typing import Any, Dict, List, Optional, Union from haystack import component, default_from_dict, default_to_dict from haystack.utils import Secret, deserialize_secrets_inplace +from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation -from .truncate import EmbeddingTruncateMode - _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia" diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py index 3a8eb9d07..931c3cce3 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from enum import Enum diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/__init__.py b/integrations/nvidia/src/haystack_integrations/components/generators/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/generators/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py index 18354ea17..b809d83b9 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2024-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 + from .generator import NvidiaGenerator __all__ = ["NvidiaGenerator"] diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py index 3eadcc5df..5bf71a9e1 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2024-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 + import warnings from typing import Any, Dict, List, Optional diff --git a/integrations/nvidia/src/haystack_integrations/components/rankers/__init__.py b/integrations/nvidia/src/haystack_integrations/components/rankers/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/__init__.py index 29cb2f7f5..05daa1c54 100644 --- a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/__init__.py +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/__init__.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from .ranker import NvidiaRanker __all__ = ["NvidiaRanker"] diff --git a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py index 1553d1ac3..9938b37d1 100644 --- a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py @@ -1,12 +1,17 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import warnings from typing import Any, Dict, List, Optional, Union -from haystack import Document, component, default_from_dict, default_to_dict +from haystack import Document, component, default_from_dict, default_to_dict, logging from haystack.utils import Secret, deserialize_secrets_inplace +from haystack_integrations.components.rankers.nvidia.truncate import RankerTruncateMode from haystack_integrations.utils.nvidia import NimBackend, url_validation -from .truncate import RankerTruncateMode +logger = logging.getLogger(__name__) _DEFAULT_MODEL = "nvidia/nv-rerankqa-mistral-4b-v3" @@ -51,7 +56,7 @@ def __init__( model: Optional[str] = None, truncate: Optional[Union[RankerTruncateMode, str]] = None, api_url: Optional[str] = None, - api_key: Optional[Secret] = None, + api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"), top_k: int = 5, ): """ @@ -100,6 +105,7 @@ def __init__( self._api_key = Secret.from_env_var("NVIDIA_API_KEY") self._top_k = top_k self._initialized = False + self._backend: Optional[Any] = None def to_dict(self) -> Dict[str, Any]: """ @@ -113,7 +119,7 @@ def to_dict(self) -> Dict[str, Any]: top_k=self._top_k, truncate=self._truncate, api_url=self._api_url, - api_key=self._api_key, + api_key=self._api_key.to_dict() if self._api_key else None, ) @classmethod @@ -124,7 +130,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "NvidiaRanker": :param data: A dictionary containing the ranker's attributes. :returns: The deserialized ranker. """ - deserialize_secrets_inplace(data, keys=["api_key"]) + init_parameters = data.get("init_parameters", {}) + if init_parameters: + deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) return default_from_dict(cls, data) def warm_up(self): @@ -170,16 +178,16 @@ def run( msg = "The ranker has not been loaded. Please call warm_up() before running." raise RuntimeError(msg) if not isinstance(query, str): - msg = "Ranker expects the `query` parameter to be a string." + msg = "NvidiaRanker expects the `query` parameter to be a string." raise TypeError(msg) if not isinstance(documents, list): - msg = "Ranker expects the `documents` parameter to be a list." + msg = "NvidiaRanker expects the `documents` parameter to be a list." raise TypeError(msg) if not all(isinstance(doc, Document) for doc in documents): - msg = "Ranker expects the `documents` parameter to be a list of Document objects." + msg = "NvidiaRanker expects the `documents` parameter to be a list of Document objects." raise TypeError(msg) if top_k is not None and not isinstance(top_k, int): - msg = "Ranker expects the `top_k` parameter to be an integer." + msg = "NvidiaRanker expects the `top_k` parameter to be an integer." raise TypeError(msg) if len(documents) == 0: @@ -187,6 +195,7 @@ def run( top_k = top_k if top_k is not None else self._top_k if top_k < 1: + logger.warning("top_k should be at least 1, returning nothing") warnings.warn("top_k should be at least 1, returning nothing", stacklevel=2) return {"documents": []} diff --git a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/truncate.py b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/truncate.py index 3b5d7f40a..649ceaf9d 100644 --- a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/truncate.py +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/truncate.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from enum import Enum diff --git a/integrations/nvidia/src/haystack_integrations/utils/__init__.py b/integrations/nvidia/src/haystack_integrations/utils/__init__.py new file mode 100644 index 000000000..6b5e14dc1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/utils/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py index da301d29d..f08cda6cd 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from .nim_backend import Model, NimBackend from .utils import is_hosted, url_validation diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/nim_backend.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/nim_backend.py index cbb6b7c3f..0279cf608 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/nim_backend.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/nim_backend.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py index 7d4dfc3b4..f07989405 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py @@ -1,9 +1,13 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import warnings -from typing import List +from typing import List, Optional from urllib.parse import urlparse, urlunparse -def url_validation(api_url: str, default_api_url: str, allowed_paths: List[str]) -> str: +def url_validation(api_url: str, default_api_url: Optional[str], allowed_paths: List[str]) -> str: """ Validate and normalize an API URL. diff --git a/integrations/nvidia/tests/__init__.py b/integrations/nvidia/tests/__init__.py index 47611e0b9..38adc654d 100644 --- a/integrations/nvidia/tests/__init__.py +++ b/integrations/nvidia/tests/__init__.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2023-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 + from .conftest import MockBackend __all__ = ["MockBackend"] diff --git a/integrations/nvidia/tests/conftest.py b/integrations/nvidia/tests/conftest.py index a6c78ba4e..b6346c672 100644 --- a/integrations/nvidia/tests/conftest.py +++ b/integrations/nvidia/tests/conftest.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from typing import Any, Dict, List, Optional, Tuple import pytest diff --git a/integrations/nvidia/tests/test_base_url.py b/integrations/nvidia/tests/test_base_url.py index 426bacc25..506fbc385 100644 --- a/integrations/nvidia/tests/test_base_url.py +++ b/integrations/nvidia/tests/test_base_url.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import pytest from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaTextEmbedder diff --git a/integrations/nvidia/tests/test_document_embedder.py b/integrations/nvidia/tests/test_document_embedder.py index db69053e7..7e0e02f3d 100644 --- a/integrations/nvidia/tests/test_document_embedder.py +++ b/integrations/nvidia/tests/test_document_embedder.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import os import pytest @@ -104,7 +108,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): }, } - def from_dict(self, monkeypatch): + def test_from_dict(self, monkeypatch): monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") data = { "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder", @@ -122,15 +126,32 @@ def from_dict(self, monkeypatch): }, } component = NvidiaDocumentEmbedder.from_dict(data) - assert component.model == "nvolveqa_40k" + assert component.model == "playground_nvolveqa_40k" assert component.api_url == "https://example.com/v1" assert component.prefix == "prefix" assert component.suffix == "suffix" + assert component.batch_size == 10 + assert component.progress_bar is False + assert component.meta_fields_to_embed == ["test_field"] + assert component.embedding_separator == " | " + assert component.truncate == EmbeddingTruncateMode.START + + def test_from_dict_defaults(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + data = { + "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder", + "init_parameters": {}, + } + component = NvidiaDocumentEmbedder.from_dict(data) + assert component.model == "nvidia/nv-embedqa-e5-v5" + assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia" + assert component.prefix == "" + assert component.suffix == "" assert component.batch_size == 32 assert component.progress_bar assert component.meta_fields_to_embed == [] assert component.embedding_separator == "\n" - assert component.truncate == EmbeddingTruncateMode.START + assert component.truncate is None def test_prepare_texts_to_embed_w_metadata(self): documents = [ diff --git a/integrations/nvidia/tests/test_embedding_truncate_mode.py b/integrations/nvidia/tests/test_embedding_truncate_mode.py index e74d0308c..16f9112ea 100644 --- a/integrations/nvidia/tests/test_embedding_truncate_mode.py +++ b/integrations/nvidia/tests/test_embedding_truncate_mode.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import pytest from haystack_integrations.components.embedders.nvidia import EmbeddingTruncateMode diff --git a/integrations/nvidia/tests/test_generator.py b/integrations/nvidia/tests/test_generator.py index 0bd8b1fc6..055830ae5 100644 --- a/integrations/nvidia/tests/test_generator.py +++ b/integrations/nvidia/tests/test_generator.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2024-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 + import os import pytest diff --git a/integrations/nvidia/tests/test_ranker.py b/integrations/nvidia/tests/test_ranker.py index 566fd18a8..d66bb0f65 100644 --- a/integrations/nvidia/tests/test_ranker.py +++ b/integrations/nvidia/tests/test_ranker.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import os import re from typing import Any, Optional, Union @@ -256,3 +260,48 @@ def test_warm_up_once(self, monkeypatch) -> None: backend = client._backend client.warm_up() assert backend == client._backend + + def test_to_dict(self) -> None: + client = NvidiaRanker() + assert client.to_dict() == { + "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker", + "init_parameters": { + "model": "nvidia/nv-rerankqa-mistral-4b-v3", + "top_k": 5, + "truncate": None, + "api_url": None, + "api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True}, + }, + } + + def test_from_dict(self) -> None: + client = NvidiaRanker.from_dict( + { + "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker", + "init_parameters": { + "model": "nvidia/nv-rerankqa-mistral-4b-v3", + "top_k": 5, + "truncate": None, + "api_url": None, + "api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True}, + }, + } + ) + assert client._model == "nvidia/nv-rerankqa-mistral-4b-v3" + assert client._top_k == 5 + assert client._truncate is None + assert client._api_url is None + assert client._api_key == Secret.from_env_var("NVIDIA_API_KEY") + + def test_from_dict_defaults(self) -> None: + client = NvidiaRanker.from_dict( + { + "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker", + "init_parameters": {}, + } + ) + assert client._model == "nvidia/nv-rerankqa-mistral-4b-v3" + assert client._top_k == 5 + assert client._truncate is None + assert client._api_url is None + assert client._api_key == Secret.from_env_var("NVIDIA_API_KEY") diff --git a/integrations/nvidia/tests/test_text_embedder.py b/integrations/nvidia/tests/test_text_embedder.py index 8690de6b1..278fa5191 100644 --- a/integrations/nvidia/tests/test_text_embedder.py +++ b/integrations/nvidia/tests/test_text_embedder.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + import os import pytest @@ -77,7 +81,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): }, } - def from_dict(self, monkeypatch): + def test_from_dict(self, monkeypatch): monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") data = { "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder", @@ -95,7 +99,20 @@ def from_dict(self, monkeypatch): assert component.api_url == "https://example.com/v1" assert component.prefix == "prefix" assert component.suffix == "suffix" - assert component.truncate == "START" + assert component.truncate == EmbeddingTruncateMode.START + + def test_from_dict_defaults(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + data = { + "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder", + "init_parameters": {}, + } + component = NvidiaTextEmbedder.from_dict(data) + assert component.model == "nvidia/nv-embedqa-e5-v5" + assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia" + assert component.prefix == "" + assert component.suffix == "" + assert component.truncate is None @pytest.mark.usefixtures("mock_local_models") def test_run_default_model(self):