nvidia-haystack 0.1.0__tar.gz → 0.1.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/CHANGELOG.md +11 -1
  2. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/PKG-INFO +3 -2
  3. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/README.md +1 -1
  4. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/pyproject.toml +1 -1
  5. nvidia_haystack-0.1.2/src/haystack_integrations/components/__init__.py +3 -0
  6. nvidia_haystack-0.1.2/src/haystack_integrations/components/embedders/__init__.py +3 -0
  7. nvidia_haystack-0.1.2/src/haystack_integrations/components/embedders/nvidia/__init__.py +9 -0
  8. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +25 -8
  9. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +29 -5
  10. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/truncate.py +4 -0
  11. nvidia_haystack-0.1.2/src/haystack_integrations/components/generators/__init__.py +3 -0
  12. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/generators/nvidia/__init__.py +1 -0
  13. nvidia_haystack-0.1.2/src/haystack_integrations/components/generators/nvidia/chat/__init__.py +3 -0
  14. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/generators/nvidia/generator.py +11 -1
  15. nvidia_haystack-0.1.2/src/haystack_integrations/components/rankers/__init__.py +3 -0
  16. nvidia_haystack-0.1.2/src/haystack_integrations/components/rankers/nvidia/__init__.py +7 -0
  17. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/rankers/nvidia/ranker.py +91 -33
  18. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/rankers/nvidia/truncate.py +4 -0
  19. nvidia_haystack-0.1.2/src/haystack_integrations/utils/__init__.py +3 -0
  20. nvidia_haystack-0.1.2/src/haystack_integrations/utils/nvidia/__init__.py +8 -0
  21. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/utils/nvidia/nim_backend.py +29 -11
  22. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/utils/nvidia/utils.py +6 -2
  23. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/__init__.py +1 -0
  24. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/conftest.py +4 -0
  25. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_base_url.py +4 -0
  26. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_document_embedder.py +47 -5
  27. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_embedding_truncate_mode.py +4 -0
  28. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_generator.py +14 -0
  29. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_ranker.py +112 -7
  30. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_text_embedder.py +37 -2
  31. nvidia_haystack-0.1.0/src/haystack_integrations/components/embedders/nvidia/__init__.py +0 -5
  32. nvidia_haystack-0.1.0/src/haystack_integrations/components/rankers/nvidia/__init__.py +0 -3
  33. nvidia_haystack-0.1.0/src/haystack_integrations/utils/nvidia/__init__.py +0 -4
  34. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/.gitignore +0 -0
  35. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/LICENSE.txt +0 -0
  36. {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/pydoc/config.yml +0 -0
  37. {nvidia_haystack-0.1.0/src/haystack_integrations/components/generators/nvidia/chat → nvidia_haystack-0.1.2/src/haystack_integrations}/__init__.py +0 -0
@@ -1,15 +1,23 @@
1
1
  # Changelog
2
2
 
3
- ## [unreleased]
3
+ ## [integrations/nvidia-v0.1.1] - 2024-11-14
4
+
5
+ ### 🐛 Bug Fixes
6
+
7
+ - Fixes to NvidiaRanker (#1191)
8
+
9
+ ## [integrations/nvidia-v0.1.0] - 2024-11-13
4
10
 
5
11
  ### 🚀 Features
6
12
 
7
13
  - Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
8
14
  - Add NVIDIA NIM ranker support (#1023)
15
+ - Raise error when attempting to embed empty documents/strings with Nvidia embedders (#1118)
9
16
 
10
17
  ### 🐛 Bug Fixes
11
18
 
12
19
  - Lints in `nvidia-haystack` (#993)
20
+ - Missing Nvidia embedding truncate mode (#1043)
13
21
 
14
22
  ### 🚜 Refactor
15
23
 
@@ -27,6 +35,8 @@
27
35
 
28
36
  - Retry tests to reduce flakyness (#836)
29
37
  - Update ruff invocation to include check parameter (#853)
38
+ - Update ruff linting scripts and settings (#1105)
39
+ - Adopt uv as installer (#1142)
30
40
 
31
41
  ### Docs
32
42
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: nvidia-haystack
3
- Version: 0.1.0
3
+ Version: 0.1.2
4
4
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme
5
5
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
6
6
  Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
19
19
  Requires-Python: >=3.8
20
20
  Requires-Dist: haystack-ai
21
21
  Requires-Dist: requests
22
+ Requires-Dist: tqdm
22
23
  Description-Content-Type: text/markdown
23
24
 
24
25
  # nvidia-haystack
@@ -61,7 +62,7 @@ hatch run test
61
62
  To only run unit tests:
62
63
 
63
64
  ```
64
- hatch run test -m"not integration"
65
+ hatch run test -m "not integration"
65
66
  ```
66
67
 
67
68
  To run the linters `ruff` and `mypy`:
@@ -38,7 +38,7 @@ hatch run test
38
38
  To only run unit tests:
39
39
 
40
40
  ```
41
- hatch run test -m"not integration"
41
+ hatch run test -m "not integration"
42
42
  ```
43
43
 
44
44
  To run the linters `ruff` and `mypy`:
@@ -23,7 +23,7 @@ classifiers = [
23
23
  "Programming Language :: Python :: Implementation :: CPython",
24
24
  "Programming Language :: Python :: Implementation :: PyPy",
25
25
  ]
26
- dependencies = ["haystack-ai", "requests"]
26
+ dependencies = ["haystack-ai", "requests", "tqdm"]
27
27
 
28
28
  [project.urls]
29
29
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme"
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,9 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from .document_embedder import NvidiaDocumentEmbedder
6
+ from .text_embedder import NvidiaTextEmbedder
7
+ from .truncate import EmbeddingTruncateMode
8
+
9
+ __all__ = ["EmbeddingTruncateMode", "NvidiaDocumentEmbedder", "NvidiaTextEmbedder"]
@@ -1,13 +1,19 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
1
6
  import warnings
2
7
  from typing import Any, Dict, List, Optional, Tuple, Union
3
8
 
4
- from haystack import Document, component, default_from_dict, default_to_dict
9
+ from haystack import Document, component, default_from_dict, default_to_dict, logging
5
10
  from haystack.utils import Secret, deserialize_secrets_inplace
6
11
  from tqdm import tqdm
7
12
 
13
+ from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
8
14
  from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
9
15
 
10
- from .truncate import EmbeddingTruncateMode
16
+ logger = logging.getLogger(__name__)
11
17
 
12
18
  _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
13
19
 
@@ -44,6 +50,7 @@ class NvidiaDocumentEmbedder:
44
50
  meta_fields_to_embed: Optional[List[str]] = None,
45
51
  embedding_separator: str = "\n",
46
52
  truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
53
+ timeout: Optional[float] = None,
47
54
  ):
48
55
  """
49
56
  Create a NvidiaTextEmbedder component.
@@ -71,8 +78,11 @@ class NvidiaDocumentEmbedder:
71
78
  :param embedding_separator:
72
79
  Separator used to concatenate the meta fields to the Document text.
73
80
  :param truncate:
74
- Specifies how inputs longer that the maximum token length should be truncated.
81
+ Specifies how inputs longer than the maximum token length should be truncated.
75
82
  If None the behavior is model-dependent, see the official documentation for more information.
83
+ :param timeout:
84
+ Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
85
+ or set to 60 by default.
76
86
  """
77
87
 
78
88
  self.api_key = api_key
@@ -95,6 +105,10 @@ class NvidiaDocumentEmbedder:
95
105
  if is_hosted(api_url) and not self.model: # manually set default model
96
106
  self.model = "nvidia/nv-embedqa-e5-v5"
97
107
 
108
+ if timeout is None:
109
+ timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
110
+ self.timeout = timeout
111
+
98
112
  def default_model(self):
99
113
  """Set default model in local NIM mode."""
100
114
  valid_models = [
@@ -125,10 +139,11 @@ class NvidiaDocumentEmbedder:
125
139
  if self.truncate is not None:
126
140
  model_kwargs["truncate"] = str(self.truncate)
127
141
  self.backend = NimBackend(
128
- self.model,
142
+ model=self.model,
129
143
  api_url=self.api_url,
130
144
  api_key=self.api_key,
131
145
  model_kwargs=model_kwargs,
146
+ timeout=self.timeout,
132
147
  )
133
148
 
134
149
  self._initialized = True
@@ -155,6 +170,7 @@ class NvidiaDocumentEmbedder:
155
170
  meta_fields_to_embed=self.meta_fields_to_embed,
156
171
  embedding_separator=self.embedding_separator,
157
172
  truncate=str(self.truncate) if self.truncate is not None else None,
173
+ timeout=self.timeout,
158
174
  )
159
175
 
160
176
  @classmethod
@@ -167,7 +183,9 @@ class NvidiaDocumentEmbedder:
167
183
  :returns:
168
184
  The deserialized component.
169
185
  """
170
- deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
186
+ init_parameters = data.get("init_parameters", {})
187
+ if init_parameters:
188
+ deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
171
189
  return default_from_dict(cls, data)
172
190
 
173
191
  def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
@@ -224,7 +242,7 @@ class NvidiaDocumentEmbedder:
224
242
  if not self._initialized:
225
243
  msg = "The embedding model has not been loaded. Please call warm_up() before running."
226
244
  raise RuntimeError(msg)
227
- elif not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
245
+ elif not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
228
246
  msg = (
229
247
  "NvidiaDocumentEmbedder expects a list of Documents as input."
230
248
  "In case you want to embed a string, please use the NvidiaTextEmbedder."
@@ -233,8 +251,7 @@ class NvidiaDocumentEmbedder:
233
251
 
234
252
  for doc in documents:
235
253
  if not doc.content:
236
- msg = f"Document '{doc.id}' has no content to embed."
237
- raise ValueError(msg)
254
+ logger.warning(f"Document '{doc.id}' has no content to embed.")
238
255
 
239
256
  texts_to_embed = self._prepare_texts_to_embed(documents)
240
257
  embeddings, metadata = self._embed_batch(texts_to_embed, self.batch_size)
@@ -1,12 +1,18 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
1
6
  import warnings
2
7
  from typing import Any, Dict, List, Optional, Union
3
8
 
4
- from haystack import component, default_from_dict, default_to_dict
9
+ from haystack import component, default_from_dict, default_to_dict, logging
5
10
  from haystack.utils import Secret, deserialize_secrets_inplace
6
11
 
12
+ from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
7
13
  from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
8
14
 
9
- from .truncate import EmbeddingTruncateMode
15
+ logger = logging.getLogger(__name__)
10
16
 
11
17
  _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
12
18
 
@@ -41,6 +47,7 @@ class NvidiaTextEmbedder:
41
47
  prefix: str = "",
42
48
  suffix: str = "",
43
49
  truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
50
+ timeout: Optional[float] = None,
44
51
  ):
45
52
  """
46
53
  Create a NvidiaTextEmbedder component.
@@ -61,6 +68,9 @@ class NvidiaTextEmbedder:
61
68
  :param truncate:
62
69
  Specifies how inputs longer that the maximum token length should be truncated.
63
70
  If None the behavior is model-dependent, see the official documentation for more information.
71
+ :param timeout:
72
+ Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
73
+ or set to 60 by default.
64
74
  """
65
75
 
66
76
  self.api_key = api_key
@@ -79,6 +89,10 @@ class NvidiaTextEmbedder:
79
89
  if is_hosted(api_url) and not self.model: # manually set default model
80
90
  self.model = "nvidia/nv-embedqa-e5-v5"
81
91
 
92
+ if timeout is None:
93
+ timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
94
+ self.timeout = timeout
95
+
82
96
  def default_model(self):
83
97
  """Set default model in local NIM mode."""
84
98
  valid_models = [
@@ -86,6 +100,12 @@ class NvidiaTextEmbedder:
86
100
  ]
87
101
  name = next(iter(valid_models), None)
88
102
  if name:
103
+ logger.warning(
104
+ "Default model is set as: {model_name}. \n"
105
+ "Set model using model parameter. \n"
106
+ "To get available models use available_models property.",
107
+ model_name=name,
108
+ )
89
109
  warnings.warn(
90
110
  f"Default model is set as: {name}. \n"
91
111
  "Set model using model parameter. \n"
@@ -109,10 +129,11 @@ class NvidiaTextEmbedder:
109
129
  if self.truncate is not None:
110
130
  model_kwargs["truncate"] = str(self.truncate)
111
131
  self.backend = NimBackend(
112
- self.model,
132
+ model=self.model,
113
133
  api_url=self.api_url,
114
134
  api_key=self.api_key,
115
135
  model_kwargs=model_kwargs,
136
+ timeout=self.timeout,
116
137
  )
117
138
 
118
139
  self._initialized = True
@@ -135,6 +156,7 @@ class NvidiaTextEmbedder:
135
156
  prefix=self.prefix,
136
157
  suffix=self.suffix,
137
158
  truncate=str(self.truncate) if self.truncate is not None else None,
159
+ timeout=self.timeout,
138
160
  )
139
161
 
140
162
  @classmethod
@@ -147,7 +169,9 @@ class NvidiaTextEmbedder:
147
169
  :returns:
148
170
  The deserialized component.
149
171
  """
150
- deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
172
+ init_parameters = data.get("init_parameters", {})
173
+ if init_parameters:
174
+ deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
151
175
  return default_from_dict(cls, data)
152
176
 
153
177
  @component.output_types(embedding=List[float], meta=Dict[str, Any])
@@ -159,7 +183,7 @@ class NvidiaTextEmbedder:
159
183
  The text to embed.
160
184
  :returns:
161
185
  A dictionary with the following keys and values:
162
- - `embedding` - Embeddng of the text.
186
+ - `embedding` - Embedding of the text.
163
187
  - `meta` - Metadata on usage statistics, etc.
164
188
  :raises RuntimeError:
165
189
  If the component was not initialized.
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  from enum import Enum
2
6
 
3
7
 
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -1,6 +1,7 @@
1
1
  # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
+
4
5
  from .generator import NvidiaGenerator
5
6
 
6
7
  __all__ = ["NvidiaGenerator"]
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -1,6 +1,8 @@
1
1
  # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
4
6
  import warnings
5
7
  from typing import Any, Dict, List, Optional
6
8
 
@@ -48,6 +50,7 @@ class NvidiaGenerator:
48
50
  api_url: str = _DEFAULT_API_URL,
49
51
  api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
50
52
  model_arguments: Optional[Dict[str, Any]] = None,
53
+ timeout: Optional[float] = None,
51
54
  ):
52
55
  """
53
56
  Create a NvidiaGenerator component.
@@ -69,6 +72,9 @@ class NvidiaGenerator:
69
72
  specific to a model.
70
73
  Search your model in the [NVIDIA NIM](https://ai.nvidia.com)
71
74
  to find the arguments it accepts.
75
+ :param timeout:
76
+ Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
77
+ or set to 60 by default.
72
78
  """
73
79
  self._model = model
74
80
  self._api_url = url_validation(api_url, _DEFAULT_API_URL, ["v1/chat/completions"])
@@ -78,6 +84,9 @@ class NvidiaGenerator:
78
84
  self._backend: Optional[Any] = None
79
85
 
80
86
  self.is_hosted = is_hosted(api_url)
87
+ if timeout is None:
88
+ timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
89
+ self.timeout = timeout
81
90
 
82
91
  def default_model(self):
83
92
  """Set default model in local NIM mode."""
@@ -109,10 +118,11 @@ class NvidiaGenerator:
109
118
  msg = "API key is required for hosted NVIDIA NIMs."
110
119
  raise ValueError(msg)
111
120
  self._backend = NimBackend(
112
- self._model,
121
+ model=self._model,
113
122
  api_url=self._api_url,
114
123
  api_key=self._api_key,
115
124
  model_kwargs=self._model_arguments,
125
+ timeout=self.timeout,
116
126
  )
117
127
 
118
128
  if not self.is_hosted and not self._model:
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from .ranker import NvidiaRanker
6
+
7
+ __all__ = ["NvidiaRanker"]
@@ -1,12 +1,18 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
1
6
  import warnings
2
7
  from typing import Any, Dict, List, Optional, Union
3
8
 
4
- from haystack import Document, component, default_from_dict, default_to_dict
9
+ from haystack import Document, component, default_from_dict, default_to_dict, logging
5
10
  from haystack.utils import Secret, deserialize_secrets_inplace
6
11
 
12
+ from haystack_integrations.components.rankers.nvidia.truncate import RankerTruncateMode
7
13
  from haystack_integrations.utils.nvidia import NimBackend, url_validation
8
14
 
9
- from .truncate import RankerTruncateMode
15
+ logger = logging.getLogger(__name__)
10
16
 
11
17
  _DEFAULT_MODEL = "nvidia/nv-rerankqa-mistral-4b-v3"
12
18
 
@@ -51,8 +57,13 @@ class NvidiaRanker:
51
57
  model: Optional[str] = None,
52
58
  truncate: Optional[Union[RankerTruncateMode, str]] = None,
53
59
  api_url: Optional[str] = None,
54
- api_key: Optional[Secret] = None,
60
+ api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
55
61
  top_k: int = 5,
62
+ query_prefix: str = "",
63
+ document_prefix: str = "",
64
+ meta_fields_to_embed: Optional[List[str]] = None,
65
+ embedding_separator: str = "\n",
66
+ timeout: Optional[float] = None,
56
67
  ):
57
68
  """
58
69
  Create a NvidiaRanker component.
@@ -67,6 +78,19 @@ class NvidiaRanker:
67
78
  Custom API URL for the NVIDIA NIM.
68
79
  :param top_k:
69
80
  Number of documents to return.
81
+ :param query_prefix:
82
+ A string to add at the beginning of the query text before ranking.
83
+ Use it to prepend the text with an instruction, as required by reranking models like `bge`.
84
+ :param document_prefix:
85
+ A string to add at the beginning of each document before ranking. You can use it to prepend the document
86
+ with an instruction, as required by embedding models like `bge`.
87
+ :param meta_fields_to_embed:
88
+ List of metadata fields to embed with the document.
89
+ :param embedding_separator:
90
+ Separator to concatenate metadata fields to the document.
91
+ :param timeout:
92
+ Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
93
+ or set to 60 by default.
70
94
  """
71
95
  if model is not None and not isinstance(model, str):
72
96
  msg = "Ranker expects the `model` parameter to be a string."
@@ -81,25 +105,34 @@ class NvidiaRanker:
81
105
  raise TypeError(msg)
82
106
 
83
107
  # todo: detect default in non-hosted case (when api_url is provided)
84
- self._model = model or _DEFAULT_MODEL
85
- self._truncate = truncate
86
- self._api_key = api_key
108
+ self.model = model or _DEFAULT_MODEL
109
+ self.truncate = truncate
110
+ self.api_key = api_key
87
111
  # if no api_url is provided, we're using a hosted model and can
88
112
  # - assume the default url will work, because there's only one model
89
113
  # - assume we won't call backend.models()
90
114
  if api_url is not None:
91
- self._api_url = url_validation(api_url, None, ["v1/ranking"])
92
- self._endpoint = None # we let backend.rank() handle the endpoint
115
+ self.api_url = url_validation(api_url, None, ["v1/ranking"])
116
+ self.endpoint = None # we let backend.rank() handle the endpoint
93
117
  else:
94
- if self._model not in _MODEL_ENDPOINT_MAP:
118
+ if self.model not in _MODEL_ENDPOINT_MAP:
95
119
  msg = f"Model '{model}' is unknown. Please provide an api_url to access it."
96
120
  raise ValueError(msg)
97
- self._api_url = None # we handle the endpoint
98
- self._endpoint = _MODEL_ENDPOINT_MAP[self._model]
121
+ self.api_url = None # we handle the endpoint
122
+ self.endpoint = _MODEL_ENDPOINT_MAP[self.model]
99
123
  if api_key is None:
100
124
  self._api_key = Secret.from_env_var("NVIDIA_API_KEY")
101
- self._top_k = top_k
125
+ self.top_k = top_k
102
126
  self._initialized = False
127
+ self._backend: Optional[Any] = None
128
+
129
+ self.query_prefix = query_prefix
130
+ self.document_prefix = document_prefix
131
+ self.meta_fields_to_embed = meta_fields_to_embed or []
132
+ self.embedding_separator = embedding_separator
133
+ if timeout is None:
134
+ timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
135
+ self.timeout = timeout
103
136
 
104
137
  def to_dict(self) -> Dict[str, Any]:
105
138
  """
@@ -109,11 +142,16 @@ class NvidiaRanker:
109
142
  """
110
143
  return default_to_dict(
111
144
  self,
112
- model=self._model,
113
- top_k=self._top_k,
114
- truncate=self._truncate,
115
- api_url=self._api_url,
116
- api_key=self._api_key,
145
+ model=self.model,
146
+ top_k=self.top_k,
147
+ truncate=self.truncate,
148
+ api_url=self.api_url,
149
+ api_key=self.api_key.to_dict() if self.api_key else None,
150
+ query_prefix=self.query_prefix,
151
+ document_prefix=self.document_prefix,
152
+ meta_fields_to_embed=self.meta_fields_to_embed,
153
+ embedding_separator=self.embedding_separator,
154
+ timeout=self.timeout,
117
155
  )
118
156
 
119
157
  @classmethod
@@ -124,7 +162,9 @@ class NvidiaRanker:
124
162
  :param data: A dictionary containing the ranker's attributes.
125
163
  :returns: The deserialized ranker.
126
164
  """
127
- deserialize_secrets_inplace(data, keys=["api_key"])
165
+ init_parameters = data.get("init_parameters", {})
166
+ if init_parameters:
167
+ deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
128
168
  return default_from_dict(cls, data)
129
169
 
130
170
  def warm_up(self):
@@ -135,18 +175,31 @@ class NvidiaRanker:
135
175
  """
136
176
  if not self._initialized:
137
177
  model_kwargs = {}
138
- if self._truncate is not None:
139
- model_kwargs.update(truncate=str(self._truncate))
178
+ if self.truncate is not None:
179
+ model_kwargs.update(truncate=str(self.truncate))
140
180
  self._backend = NimBackend(
141
- self._model,
142
- api_url=self._api_url,
143
- api_key=self._api_key,
181
+ model=self.model,
182
+ api_url=self.api_url,
183
+ api_key=self.api_key,
144
184
  model_kwargs=model_kwargs,
185
+ timeout=self.timeout,
145
186
  )
146
- if not self._model:
147
- self._model = _DEFAULT_MODEL
187
+ if not self.model:
188
+ self.model = _DEFAULT_MODEL
148
189
  self._initialized = True
149
190
 
191
+ def _prepare_documents_to_embed(self, documents: List[Document]) -> List[str]:
192
+ document_texts = []
193
+ for doc in documents:
194
+ meta_values_to_embed = [
195
+ str(doc.meta[key])
196
+ for key in self.meta_fields_to_embed
197
+ if key in doc.meta and doc.meta[key] # noqa: RUF019
198
+ ]
199
+ text_to_embed = self.embedding_separator.join([*meta_values_to_embed, doc.content or ""])
200
+ document_texts.append(self.document_prefix + text_to_embed)
201
+ return document_texts
202
+
150
203
  @component.output_types(documents=List[Document])
151
204
  def run(
152
205
  self,
@@ -170,32 +223,37 @@ class NvidiaRanker:
170
223
  msg = "The ranker has not been loaded. Please call warm_up() before running."
171
224
  raise RuntimeError(msg)
172
225
  if not isinstance(query, str):
173
- msg = "Ranker expects the `query` parameter to be a string."
226
+ msg = "NvidiaRanker expects the `query` parameter to be a string."
174
227
  raise TypeError(msg)
175
228
  if not isinstance(documents, list):
176
- msg = "Ranker expects the `documents` parameter to be a list."
229
+ msg = "NvidiaRanker expects the `documents` parameter to be a list."
177
230
  raise TypeError(msg)
178
231
  if not all(isinstance(doc, Document) for doc in documents):
179
- msg = "Ranker expects the `documents` parameter to be a list of Document objects."
232
+ msg = "NvidiaRanker expects the `documents` parameter to be a list of Document objects."
180
233
  raise TypeError(msg)
181
234
  if top_k is not None and not isinstance(top_k, int):
182
- msg = "Ranker expects the `top_k` parameter to be an integer."
235
+ msg = "NvidiaRanker expects the `top_k` parameter to be an integer."
183
236
  raise TypeError(msg)
184
237
 
185
238
  if len(documents) == 0:
186
239
  return {"documents": []}
187
240
 
188
- top_k = top_k if top_k is not None else self._top_k
241
+ top_k = top_k if top_k is not None else self.top_k
189
242
  if top_k < 1:
243
+ logger.warning("top_k should be at least 1, returning nothing")
190
244
  warnings.warn("top_k should be at least 1, returning nothing", stacklevel=2)
191
245
  return {"documents": []}
192
246
 
193
247
  assert self._backend is not None
248
+
249
+ query_text = self.query_prefix + query
250
+ document_texts = self._prepare_documents_to_embed(documents=documents)
251
+
194
252
  # rank result is list[{index: int, logit: float}] sorted by logit
195
253
  sorted_indexes_and_scores = self._backend.rank(
196
- query,
197
- documents,
198
- endpoint=self._endpoint,
254
+ query_text=query_text,
255
+ document_texts=document_texts,
256
+ endpoint=self.endpoint,
199
257
  )
200
258
  sorted_documents = []
201
259
  for item in sorted_indexes_and_scores[:top_k]:
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  from enum import Enum
2
6
 
3
7
 
@@ -0,0 +1,3 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,8 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from .nim_backend import Model, NimBackend
6
+ from .utils import is_hosted, url_validation
7
+
8
+ __all__ = ["Model", "NimBackend", "is_hosted", "url_validation"]
@@ -1,11 +1,18 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
1
6
  from dataclasses import dataclass, field
2
7
  from typing import Any, Dict, List, Optional, Tuple
3
8
 
4
9
  import requests
5
- from haystack import Document
10
+ from haystack import logging
6
11
  from haystack.utils import Secret
7
12
 
8
- REQUEST_TIMEOUT = 60
13
+ logger = logging.getLogger(__name__)
14
+
15
+ REQUEST_TIMEOUT = 60.0
9
16
 
10
17
 
11
18
  @dataclass
@@ -31,6 +38,7 @@ class NimBackend:
31
38
  api_url: str,
32
39
  api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
33
40
  model_kwargs: Optional[Dict[str, Any]] = None,
41
+ timeout: Optional[float] = None,
34
42
  ):
35
43
  headers = {
36
44
  "Content-Type": "application/json",
@@ -46,6 +54,9 @@ class NimBackend:
46
54
  self.model = model
47
55
  self.api_url = api_url
48
56
  self.model_kwargs = model_kwargs or {}
57
+ if timeout is None:
58
+ timeout = float(os.environ.get("NVIDIA_TIMEOUT", REQUEST_TIMEOUT))
59
+ self.timeout = timeout
49
60
 
50
61
  def embed(self, texts: List[str]) -> Tuple[List[List[float]], Dict[str, Any]]:
51
62
  url = f"{self.api_url}/embeddings"
@@ -58,10 +69,11 @@ class NimBackend:
58
69
  "input": texts,
59
70
  **self.model_kwargs,
60
71
  },
61
- timeout=REQUEST_TIMEOUT,
72
+ timeout=self.timeout,
62
73
  )
63
74
  res.raise_for_status()
64
75
  except requests.HTTPError as e:
76
+ logger.error("Error when calling NIM embedding endpoint: Error - {error}", error=e.response.text)
65
77
  msg = f"Failed to query embedding endpoint: Error - {e.response.text}"
66
78
  raise ValueError(msg) from e
67
79
 
@@ -90,10 +102,11 @@ class NimBackend:
90
102
  ],
91
103
  **self.model_kwargs,
92
104
  },
93
- timeout=REQUEST_TIMEOUT,
105
+ timeout=self.timeout,
94
106
  )
95
107
  res.raise_for_status()
96
108
  except requests.HTTPError as e:
109
+ logger.error("Error when calling NIM chat completion endpoint: Error - {error}", error=e.response.text)
97
110
  msg = f"Failed to query chat completion endpoint: Error - {e.response.text}"
98
111
  raise ValueError(msg) from e
99
112
 
@@ -128,21 +141,22 @@ class NimBackend:
128
141
 
129
142
  res = self.session.get(
130
143
  url,
131
- timeout=REQUEST_TIMEOUT,
144
+ timeout=self.timeout,
132
145
  )
133
146
  res.raise_for_status()
134
147
 
135
148
  data = res.json()["data"]
136
149
  models = [Model(element["id"]) for element in data if "id" in element]
137
150
  if not models:
151
+ logger.error("No hosted model were found at URL '{u}'.", u=url)
138
152
  msg = f"No hosted model were found at URL '{url}'."
139
153
  raise ValueError(msg)
140
154
  return models
141
155
 
142
156
  def rank(
143
157
  self,
144
- query: str,
145
- documents: List[Document],
158
+ query_text: str,
159
+ document_texts: List[str],
146
160
  endpoint: Optional[str] = None,
147
161
  ) -> List[Dict[str, Any]]:
148
162
  url = endpoint or f"{self.api_url}/ranking"
@@ -152,18 +166,22 @@ class NimBackend:
152
166
  url,
153
167
  json={
154
168
  "model": self.model,
155
- "query": {"text": query},
156
- "passages": [{"text": doc.content} for doc in documents],
169
+ "query": {"text": query_text},
170
+ "passages": [{"text": text} for text in document_texts],
157
171
  **self.model_kwargs,
158
172
  },
159
- timeout=REQUEST_TIMEOUT,
173
+ timeout=self.timeout,
160
174
  )
161
175
  res.raise_for_status()
162
176
  except requests.HTTPError as e:
177
+ logger.error("Error when calling NIM ranking endpoint: Error - {error}", error=e.response.text)
163
178
  msg = f"Failed to rank endpoint: Error - {e.response.text}"
164
179
  raise ValueError(msg) from e
165
180
 
166
181
  data = res.json()
167
- assert "rankings" in data, f"Expected 'rankings' in response, got {data}"
182
+ if "rankings" not in data:
183
+ logger.error("Expected 'rankings' in response, got {d}", d=data)
184
+ msg = f"Expected 'rankings' in response, got {data}"
185
+ raise ValueError(msg)
168
186
 
169
187
  return data["rankings"]
@@ -1,9 +1,13 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import warnings
2
- from typing import List
6
+ from typing import List, Optional
3
7
  from urllib.parse import urlparse, urlunparse
4
8
 
5
9
 
6
- def url_validation(api_url: str, default_api_url: str, allowed_paths: List[str]) -> str:
10
+ def url_validation(api_url: str, default_api_url: Optional[str], allowed_paths: List[str]) -> str:
7
11
  """
8
12
  Validate and normalize an API URL.
9
13
 
@@ -1,6 +1,7 @@
1
1
  # SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
+
4
5
  from .conftest import MockBackend
5
6
 
6
7
  __all__ = ["MockBackend"]
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  from typing import Any, Dict, List, Optional, Tuple
2
6
 
3
7
  import pytest
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import pytest
2
6
 
3
7
  from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaTextEmbedder
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import os
2
6
 
3
7
  import pytest
@@ -71,6 +75,7 @@ class TestNvidiaDocumentEmbedder:
71
75
  "meta_fields_to_embed": [],
72
76
  "embedding_separator": "\n",
73
77
  "truncate": None,
78
+ "timeout": 60.0,
74
79
  },
75
80
  }
76
81
 
@@ -86,6 +91,7 @@ class TestNvidiaDocumentEmbedder:
86
91
  meta_fields_to_embed=["test_field"],
87
92
  embedding_separator=" | ",
88
93
  truncate=EmbeddingTruncateMode.END,
94
+ timeout=45.0,
89
95
  )
90
96
  data = component.to_dict()
91
97
  assert data == {
@@ -101,10 +107,11 @@ class TestNvidiaDocumentEmbedder:
101
107
  "meta_fields_to_embed": ["test_field"],
102
108
  "embedding_separator": " | ",
103
109
  "truncate": "END",
110
+ "timeout": 45.0,
104
111
  },
105
112
  }
106
113
 
107
- def from_dict(self, monkeypatch):
114
+ def test_from_dict(self, monkeypatch):
108
115
  monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
109
116
  data = {
110
117
  "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder",
@@ -119,18 +126,38 @@ class TestNvidiaDocumentEmbedder:
119
126
  "meta_fields_to_embed": ["test_field"],
120
127
  "embedding_separator": " | ",
121
128
  "truncate": "START",
129
+ "timeout": 45.0,
122
130
  },
123
131
  }
124
132
  component = NvidiaDocumentEmbedder.from_dict(data)
125
- assert component.model == "nvolveqa_40k"
133
+ assert component.model == "playground_nvolveqa_40k"
126
134
  assert component.api_url == "https://example.com/v1"
127
135
  assert component.prefix == "prefix"
128
136
  assert component.suffix == "suffix"
137
+ assert component.batch_size == 10
138
+ assert component.progress_bar is False
139
+ assert component.meta_fields_to_embed == ["test_field"]
140
+ assert component.embedding_separator == " | "
141
+ assert component.truncate == EmbeddingTruncateMode.START
142
+ assert component.timeout == 45.0
143
+
144
+ def test_from_dict_defaults(self, monkeypatch):
145
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
146
+ data = {
147
+ "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder",
148
+ "init_parameters": {},
149
+ }
150
+ component = NvidiaDocumentEmbedder.from_dict(data)
151
+ assert component.model == "nvidia/nv-embedqa-e5-v5"
152
+ assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia"
153
+ assert component.prefix == ""
154
+ assert component.suffix == ""
129
155
  assert component.batch_size == 32
130
156
  assert component.progress_bar
131
157
  assert component.meta_fields_to_embed == []
132
158
  assert component.embedding_separator == "\n"
133
- assert component.truncate == EmbeddingTruncateMode.START
159
+ assert component.truncate is None
160
+ assert component.timeout == 60.0
134
161
 
135
162
  def test_prepare_texts_to_embed_w_metadata(self):
136
163
  documents = [
@@ -326,7 +353,7 @@ class TestNvidiaDocumentEmbedder:
326
353
  with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"):
327
354
  embedder.run(documents=list_integers_input)
328
355
 
329
- def test_run_empty_document(self):
356
+ def test_run_empty_document(self, caplog):
330
357
  model = "playground_nvolveqa_40k"
331
358
  api_key = Secret.from_token("fake-api-key")
332
359
  embedder = NvidiaDocumentEmbedder(model, api_key=api_key)
@@ -334,8 +361,10 @@ class TestNvidiaDocumentEmbedder:
334
361
  embedder.warm_up()
335
362
  embedder.backend = MockBackend(model=model, api_key=api_key)
336
363
 
337
- with pytest.raises(ValueError, match="no content to embed"):
364
+ # Write check using caplog that a logger.warning is raised
365
+ with caplog.at_level("WARNING"):
338
366
  embedder.run(documents=[Document(content="")])
367
+ assert "has no content to embed." in caplog.text
339
368
 
340
369
  def test_run_on_empty_list(self):
341
370
  model = "playground_nvolveqa_40k"
@@ -351,6 +380,19 @@ class TestNvidiaDocumentEmbedder:
351
380
  assert result["documents"] is not None
352
381
  assert not result["documents"] # empty list
353
382
 
383
+ def test_setting_timeout(self, monkeypatch):
384
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
385
+ embedder = NvidiaDocumentEmbedder(timeout=10.0)
386
+ embedder.warm_up()
387
+ assert embedder.backend.timeout == 10.0
388
+
389
+ def test_setting_timeout_env(self, monkeypatch):
390
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
391
+ monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
392
+ embedder = NvidiaDocumentEmbedder()
393
+ embedder.warm_up()
394
+ assert embedder.backend.timeout == 45.0
395
+
354
396
  @pytest.mark.skipif(
355
397
  not os.environ.get("NVIDIA_API_KEY", None),
356
398
  reason="Export an env var called NVIDIA_API_KEY containing the Nvidia API key to run this test.",
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import pytest
2
6
 
3
7
  from haystack_integrations.components.embedders.nvidia import EmbeddingTruncateMode
@@ -1,6 +1,7 @@
1
1
  # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
+
4
5
  import os
5
6
 
6
7
  import pytest
@@ -123,6 +124,19 @@ class TestNvidiaGenerator:
123
124
  },
124
125
  }
125
126
 
127
+ def test_setting_timeout(self, monkeypatch):
128
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
129
+ generator = NvidiaGenerator(timeout=10.0)
130
+ generator.warm_up()
131
+ assert generator._backend.timeout == 10.0
132
+
133
+ def test_setting_timeout_env(self, monkeypatch):
134
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
135
+ monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
136
+ generator = NvidiaGenerator()
137
+ generator.warm_up()
138
+ assert generator._backend.timeout == 45.0
139
+
126
140
  @pytest.mark.skipif(
127
141
  not os.environ.get("NVIDIA_NIM_GENERATOR_MODEL", None) or not os.environ.get("NVIDIA_NIM_ENDPOINT_URL", None),
128
142
  reason="Export an env var called NVIDIA_NIM_GENERATOR_MODEL containing the hosted model name and "
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import os
2
6
  import re
3
7
  from typing import Any, Optional, Union
@@ -15,8 +19,8 @@ class TestNvidiaRanker:
15
19
  def test_init_default(self, monkeypatch):
16
20
  monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
17
21
  client = NvidiaRanker()
18
- assert client._model == _DEFAULT_MODEL
19
- assert client._api_key == Secret.from_env_var("NVIDIA_API_KEY")
22
+ assert client.model == _DEFAULT_MODEL
23
+ assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
20
24
 
21
25
  def test_init_with_parameters(self):
22
26
  client = NvidiaRanker(
@@ -25,10 +29,10 @@ class TestNvidiaRanker:
25
29
  top_k=3,
26
30
  truncate="END",
27
31
  )
28
- assert client._api_key == Secret.from_token("fake-api-key")
29
- assert client._model == _DEFAULT_MODEL
30
- assert client._top_k == 3
31
- assert client._truncate == RankerTruncateMode.END
32
+ assert client.api_key == Secret.from_token("fake-api-key")
33
+ assert client.model == _DEFAULT_MODEL
34
+ assert client.top_k == 3
35
+ assert client.truncate == RankerTruncateMode.END
32
36
 
33
37
  def test_init_fail_wo_api_key(self, monkeypatch):
34
38
  monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
@@ -39,7 +43,7 @@ class TestNvidiaRanker:
39
43
  def test_init_pass_wo_api_key_w_api_url(self):
40
44
  url = "https://url.bogus/v1"
41
45
  client = NvidiaRanker(api_url=url)
42
- assert client._api_url == url
46
+ assert client.api_url == url
43
47
 
44
48
  def test_warm_up_required(self):
45
49
  client = NvidiaRanker()
@@ -256,3 +260,104 @@ class TestNvidiaRanker:
256
260
  backend = client._backend
257
261
  client.warm_up()
258
262
  assert backend == client._backend
263
+
264
+ def test_to_dict(self) -> None:
265
+ client = NvidiaRanker()
266
+ assert client.to_dict() == {
267
+ "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
268
+ "init_parameters": {
269
+ "model": "nvidia/nv-rerankqa-mistral-4b-v3",
270
+ "top_k": 5,
271
+ "truncate": None,
272
+ "api_url": None,
273
+ "api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True},
274
+ "query_prefix": "",
275
+ "document_prefix": "",
276
+ "meta_fields_to_embed": [],
277
+ "embedding_separator": "\n",
278
+ "timeout": 60.0,
279
+ },
280
+ }
281
+
282
+ def test_from_dict(self) -> None:
283
+ client = NvidiaRanker.from_dict(
284
+ {
285
+ "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
286
+ "init_parameters": {
287
+ "model": "nvidia/nv-rerankqa-mistral-4b-v3",
288
+ "top_k": 5,
289
+ "truncate": None,
290
+ "api_url": None,
291
+ "api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True},
292
+ "query_prefix": "",
293
+ "document_prefix": "",
294
+ "meta_fields_to_embed": [],
295
+ "embedding_separator": "\n",
296
+ "timeout": 45.0,
297
+ },
298
+ }
299
+ )
300
+ assert client.model == "nvidia/nv-rerankqa-mistral-4b-v3"
301
+ assert client.top_k == 5
302
+ assert client.truncate is None
303
+ assert client.api_url is None
304
+ assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
305
+ assert client.query_prefix == ""
306
+ assert client.document_prefix == ""
307
+ assert client.meta_fields_to_embed == []
308
+ assert client.embedding_separator == "\n"
309
+ assert client.timeout == 45.0
310
+
311
+ def test_from_dict_defaults(self) -> None:
312
+ client = NvidiaRanker.from_dict(
313
+ {
314
+ "type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
315
+ "init_parameters": {},
316
+ }
317
+ )
318
+ assert client.model == "nvidia/nv-rerankqa-mistral-4b-v3"
319
+ assert client.top_k == 5
320
+ assert client.truncate is None
321
+ assert client.api_url is None
322
+ assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
323
+ assert client.query_prefix == ""
324
+ assert client.document_prefix == ""
325
+ assert client.meta_fields_to_embed == []
326
+ assert client.embedding_separator == "\n"
327
+ assert client.timeout == 60.0
328
+
329
+ def test_setting_timeout(self, monkeypatch):
330
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
331
+ client = NvidiaRanker(timeout=10.0)
332
+ client.warm_up()
333
+ assert client._backend.timeout == 10.0
334
+
335
+ def test_setting_timeout_env(self, monkeypatch):
336
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
337
+ monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
338
+ client = NvidiaRanker()
339
+ client.warm_up()
340
+ assert client._backend.timeout == 45.0
341
+
342
+ def test_prepare_texts_to_embed_w_metadata(self):
343
+ documents = [
344
+ Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5)
345
+ ]
346
+
347
+ ranker = NvidiaRanker(
348
+ model=None,
349
+ api_key=Secret.from_token("fake-api-key"),
350
+ meta_fields_to_embed=["meta_field"],
351
+ embedding_separator=" | ",
352
+ )
353
+
354
+ prepared_texts = ranker._prepare_documents_to_embed(documents)
355
+
356
+ # note that newline is replaced by space
357
+ assert prepared_texts == [
358
+ "meta_value 0 | document number 0:\ncontent",
359
+ "meta_value 1 | document number 1:\ncontent",
360
+ "meta_value 2 | document number 2:\ncontent",
361
+ "meta_value 3 | document number 3:\ncontent",
362
+ "meta_value 4 | document number 4:\ncontent",
363
+ ]
@@ -1,3 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
1
5
  import os
2
6
 
3
7
  import pytest
@@ -52,6 +56,7 @@ class TestNvidiaTextEmbedder:
52
56
  "prefix": "",
53
57
  "suffix": "",
54
58
  "truncate": None,
59
+ "timeout": 60.0,
55
60
  },
56
61
  }
57
62
 
@@ -63,6 +68,7 @@ class TestNvidiaTextEmbedder:
63
68
  prefix="prefix",
64
69
  suffix="suffix",
65
70
  truncate=EmbeddingTruncateMode.START,
71
+ timeout=10.0,
66
72
  )
67
73
  data = component.to_dict()
68
74
  assert data == {
@@ -74,10 +80,11 @@ class TestNvidiaTextEmbedder:
74
80
  "prefix": "prefix",
75
81
  "suffix": "suffix",
76
82
  "truncate": "START",
83
+ "timeout": 10.0,
77
84
  },
78
85
  }
79
86
 
80
- def from_dict(self, monkeypatch):
87
+ def test_from_dict(self, monkeypatch):
81
88
  monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
82
89
  data = {
83
90
  "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder",
@@ -88,6 +95,7 @@ class TestNvidiaTextEmbedder:
88
95
  "prefix": "prefix",
89
96
  "suffix": "suffix",
90
97
  "truncate": "START",
98
+ "timeout": 10.0,
91
99
  },
92
100
  }
93
101
  component = NvidiaTextEmbedder.from_dict(data)
@@ -95,7 +103,21 @@ class TestNvidiaTextEmbedder:
95
103
  assert component.api_url == "https://example.com/v1"
96
104
  assert component.prefix == "prefix"
97
105
  assert component.suffix == "suffix"
98
- assert component.truncate == "START"
106
+ assert component.truncate == EmbeddingTruncateMode.START
107
+ assert component.timeout == 10.0
108
+
109
+ def test_from_dict_defaults(self, monkeypatch):
110
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
111
+ data = {
112
+ "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder",
113
+ "init_parameters": {},
114
+ }
115
+ component = NvidiaTextEmbedder.from_dict(data)
116
+ assert component.model == "nvidia/nv-embedqa-e5-v5"
117
+ assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia"
118
+ assert component.prefix == ""
119
+ assert component.suffix == ""
120
+ assert component.truncate is None
99
121
 
100
122
  @pytest.mark.usefixtures("mock_local_models")
101
123
  def test_run_default_model(self):
@@ -158,6 +180,19 @@ class TestNvidiaTextEmbedder:
158
180
  with pytest.raises(ValueError, match="empty string"):
159
181
  embedder.run(text="")
160
182
 
183
+ def test_setting_timeout(self, monkeypatch):
184
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
185
+ embedder = NvidiaTextEmbedder(timeout=10.0)
186
+ embedder.warm_up()
187
+ assert embedder.backend.timeout == 10.0
188
+
189
+ def test_setting_timeout_env(self, monkeypatch):
190
+ monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
191
+ monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
192
+ embedder = NvidiaTextEmbedder()
193
+ embedder.warm_up()
194
+ assert embedder.backend.timeout == 45.0
195
+
161
196
  @pytest.mark.skipif(
162
197
  not os.environ.get("NVIDIA_NIM_EMBEDDER_MODEL", None) or not os.environ.get("NVIDIA_NIM_ENDPOINT_URL", None),
163
198
  reason="Export an env var called NVIDIA_NIM_EMBEDDER_MODEL containing the hosted model name and "
@@ -1,5 +0,0 @@
1
- from .document_embedder import NvidiaDocumentEmbedder
2
- from .text_embedder import NvidiaTextEmbedder
3
- from .truncate import EmbeddingTruncateMode
4
-
5
- __all__ = ["NvidiaDocumentEmbedder", "NvidiaTextEmbedder", "EmbeddingTruncateMode"]
@@ -1,3 +0,0 @@
1
- from .ranker import NvidiaRanker
2
-
3
- __all__ = ["NvidiaRanker"]
@@ -1,4 +0,0 @@
1
- from .nim_backend import Model, NimBackend
2
- from .utils import is_hosted, url_validation
3
-
4
- __all__ = ["NimBackend", "Model", "is_hosted", "url_validation"]