feat: Update Nvidia components to auto call run warm_up and don't modify Documents in place (#2680)

* Add missing license headers, auto call warmup, don't modify docs in place

* Updates

* fix docstring

---------

Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: sjrl

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}

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."

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.
 

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)

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,27 +192,21 @@ 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.
 
         :param query: The query to rank the documents against.
         :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)

integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from enum import Enum
 
 

integrations/nvidia/src/haystack_integrations/utils/nvidia/models.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from dataclasses import dataclass
 from typing import Literal, Optional, Union
 

integrations/nvidia/tests/test_nim_backend.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
 #
 # SPDX-License-Identifier: Apache-2.0
+
 import json
 from unittest.mock import patch
 

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",
         [