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 b4c4e00aeb..8ad004f097 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 @@ -4,6 +4,7 @@ import os import warnings +from dataclasses import replace from typing import Any, Optional, Union from haystack import Document, component, default_from_dict, default_to_dict, logging @@ -49,7 +50,7 @@ def __init__( embedding_separator: str = "\n", truncate: Optional[Union[EmbeddingTruncateMode, str]] = None, timeout: Optional[float] = None, - ): + ) -> None: """ Create a NvidiaTextEmbedder component. @@ -108,7 +109,7 @@ def __init__( def class_name(cls) -> str: return "NvidiaDocumentEmbedder" - def default_model(self): + def default_model(self) -> None: """Set default model in local NIM mode.""" valid_models = [ model.id for model in self.available_models if not model.base_model or model.base_model == model.id @@ -129,7 +130,7 @@ def default_model(self): error_message = "No locally hosted model was found." raise ValueError(error_message) - def warm_up(self): + def warm_up(self) -> None: """ Initializes the component. """ @@ -246,14 +247,12 @@ def run(self, documents: list[Document]) -> dict[str, Union[list[Document], dict A dictionary with the following keys and values: - `documents` - List of processed Documents with embeddings. - `meta` - Metadata on usage statistics, etc. - :raises RuntimeError: - If the component was not initialized. :raises TypeError: - If the input is not a string. + If the input is not a list of Documents. """ if not self._initialized: - msg = "The embedding model has not been loaded. Please call warm_up() before running." - raise RuntimeError(msg) + self.warm_up() + elif not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)): msg = ( "NvidiaDocumentEmbedder expects a list of Documents as input." @@ -267,7 +266,9 @@ def run(self, documents: list[Document]) -> dict[str, Union[list[Document], dict texts_to_embed = self._prepare_texts_to_embed(documents) embeddings, metadata = self._embed_batch(texts_to_embed, self.batch_size) + + new_documents = [] for doc, emb in zip(documents, embeddings): - doc.embedding = emb + new_documents.append(replace(doc, embedding=emb)) - return {"documents": documents, "meta": metadata} + return {"documents": new_documents, "meta": metadata} 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 076a1d3020..ab7b3fc77d 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 @@ -197,14 +197,14 @@ def run(self, text: str) -> dict[str, Union[list[float], dict[str, Any]]]: A dictionary with the following keys and values: - `embedding` - Embedding of the text. - `meta` - Metadata on usage statistics, etc. - :raises RuntimeError: - If the component was not initialized. :raises TypeError: If the input is not a string. + :raises ValueError: + If the input string is empty. """ if not self._initialized: - msg = "The embedding model has not been loaded. Please call warm_up() before running." - raise RuntimeError(msg) + self.warm_up() + elif not isinstance(text, str): msg = ( "NvidiaTextEmbedder expects a string as an input." diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py index 9f623e7a19..14ef8ee7a7 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py @@ -60,7 +60,7 @@ def __init__( timeout: Optional[float] = None, max_retries: Optional[int] = None, http_client_kwargs: Optional[dict[str, Any]] = None, - ): + ) -> None: """ Creates an instance of NvidiaChatGenerator. 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 312431cdad..3c608f66cc 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py @@ -49,7 +49,7 @@ def __init__( api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"), model_arguments: Optional[dict[str, Any]] = None, timeout: Optional[float] = None, - ): + ) -> None: """ Create a NvidiaGenerator component. @@ -90,7 +90,7 @@ def __init__( def class_name(cls) -> str: return "NvidiaGenerator" - def default_model(self): + def default_model(self) -> None: """Set default model in local NIM mode.""" valid_models = [ model.id for model in self.available_models if not model.base_model or model.base_model == model.id @@ -111,7 +111,7 @@ def default_model(self): error_message = "No locally hosted model was found." raise ValueError(error_message) - def warm_up(self): + def warm_up(self) -> None: """ Initializes the component. """ @@ -183,8 +183,7 @@ def run(self, prompt: str) -> dict[str, Union[list[str], list[dict[str, Any]]]]: - `meta` - Metadata for each reply. """ if self.backend is None: - msg = "The generation model has not been loaded. Call warm_up() before running." - raise RuntimeError(msg) + self.warm_up() assert self.backend is not None replies, meta = self.backend.generate(prompt=prompt) 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 ef931212e8..1b27cca6e6 100644 --- a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py @@ -57,7 +57,7 @@ def __init__( meta_fields_to_embed: Optional[list[str]] = None, embedding_separator: str = "\n", timeout: Optional[float] = None, - ): + ) -> None: """ Create a NvidiaRanker component. @@ -155,7 +155,7 @@ def from_dict(cls, data: dict[str, Any]) -> "NvidiaRanker": deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) return default_from_dict(cls, data) - def warm_up(self): + def warm_up(self) -> None: """ Initialize the ranker. @@ -192,12 +192,7 @@ def _prepare_documents_to_embed(self, documents: list[Document]) -> list[str]: return document_texts @component.output_types(documents=list[Document]) - def run( - self, - query: str, - documents: list[Document], - top_k: Optional[int] = None, - ) -> dict[str, list[Document]]: + def run(self, query: str, documents: list[Document], top_k: Optional[int] = None) -> dict[str, list[Document]]: """ Rank a list of documents based on a given query. @@ -205,14 +200,13 @@ def run( :param documents: The list of documents to rank. :param top_k: The number of documents to return. - :raises RuntimeError: If the ranker has not been loaded. :raises TypeError: If the arguments are of the wrong type. :returns: A dictionary containing the ranked documents. """ if not self._initialized: - msg = "The ranker has not been loaded. Please call warm_up() before running." - raise RuntimeError(msg) + self.warm_up() + if not isinstance(query, str): msg = "NvidiaRanker expects the `query` parameter to be a string." raise TypeError(msg) diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py index 04a0bc5f15..f6dacf9a4f 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.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/nvidia/models.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/models.py index 5478dbac7d..081a8ea899 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/models.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/models.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: 2024-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + from dataclasses import dataclass from typing import Literal, Optional, Union diff --git a/integrations/nvidia/tests/test_nim_backend.py b/integrations/nvidia/tests/test_nim_backend.py index cb55294a2e..16e4ececb0 100644 --- a/integrations/nvidia/tests/test_nim_backend.py +++ b/integrations/nvidia/tests/test_nim_backend.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2024-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 + import json from unittest.mock import patch diff --git a/integrations/nvidia/tests/test_ranker.py b/integrations/nvidia/tests/test_ranker.py index e684844cdf..f12c79f2bf 100644 --- a/integrations/nvidia/tests/test_ranker.py +++ b/integrations/nvidia/tests/test_ranker.py @@ -46,12 +46,6 @@ def test_init_pass_wo_api_key_w_api_url(self): client = NvidiaRanker(api_url=url) assert client.api_url == url - def test_warm_up_required(self): - client = NvidiaRanker() - with pytest.raises(RuntimeError) as e: - client.run("query", [Document(content="doc")]) - assert "not been loaded" in str(e.value) - @pytest.mark.parametrize( "truncate", [