nvidia-haystack 1.4.0__tar.gz → 3.0.0__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 (39) hide show
  1. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/CHANGELOG.md +61 -12
  2. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/PKG-INFO +4 -5
  3. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/examples/chat_generator_with_structured_outputs.py +1 -1
  4. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/pydoc/config_docusaurus.yml +0 -1
  5. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/pyproject.toml +3 -4
  6. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +12 -25
  7. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +12 -25
  8. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/generators/nvidia/__init__.py +1 -2
  9. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py +1 -1
  10. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/rankers/nvidia/ranker.py +28 -37
  11. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/rankers/nvidia/truncate.py +1 -1
  12. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/nvidia/models.py +20 -3
  13. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/nvidia/nim_backend.py +6 -62
  14. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_base_url.py +5 -6
  15. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_document_embedder.py +47 -30
  16. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_nim_backend.py +9 -74
  17. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_nvidia_chat_generator.py +67 -94
  18. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_ranker.py +94 -82
  19. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_text_embedder.py +47 -24
  20. nvidia_haystack-1.4.0/src/haystack_integrations/components/generators/nvidia/generator.py +0 -199
  21. nvidia_haystack-1.4.0/tests/test_generator.py +0 -302
  22. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/.gitignore +0 -0
  23. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/LICENSE.txt +0 -0
  24. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/README.md +0 -0
  25. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/embedders/nvidia/__init__.py +0 -0
  26. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/embedders/nvidia/truncate.py +0 -0
  27. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  28. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/generators/nvidia/chat/__init__.py +0 -0
  29. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  30. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/rankers/nvidia/__init__.py +0 -0
  31. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  32. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/nvidia/__init__.py +0 -0
  33. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/nvidia/client.py +0 -0
  34. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/nvidia/utils.py +0 -0
  35. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/src/haystack_integrations/utils/py.typed +0 -0
  36. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/__init__.py +0 -0
  37. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/conftest.py +0 -0
  38. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_embedding_truncate_mode.py +0 -0
  39. {nvidia_haystack-1.4.0 → nvidia_haystack-3.0.0}/tests/test_utils.py +0 -0
@@ -1,5 +1,27 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/nvidia-v2.0.0] - 2026-09-08
4
+
5
+ ### 🚀 Features
6
+
7
+ - [**breaking**] Nvidia - add `close` method; change the default models (#3930)
8
+
9
+ ### ⚙️ CI
10
+
11
+ - Improve changelog generation; fix existing changelogs (#3883)
12
+
13
+ ### 🧹 Chores
14
+
15
+ - [**breaking**] Remove deprecated `NvidiaGenerator` (use `NvidiaChatGenerator` instead) (#3891)
16
+
17
+
18
+ ## [integrations/nvidia-v1.4.0] - 2026-08-11
19
+
20
+ ### 🐛 Bug Fixes
21
+
22
+ - NvidiaRanker reordered the top_k fallback to happen after validation (#3749)
23
+
24
+
3
25
  ## [integrations/nvidia-v1.3.0] - 2026-08-04
4
26
 
5
27
  ### 🐛 Bug Fixes
@@ -185,9 +207,14 @@
185
207
 
186
208
  ### 🚀 Features
187
209
 
188
- - Improvements to NvidiaRanker and adding user input timeout (#1193)
189
210
  - Add model `nvidia/llama-3.2-nv-rerankqa-1b-v2` to `_MODEL_ENDPOINT_MAP` (#1260)
190
211
 
212
+ ## [integrations/nvidia-v0.1.2] - 2024-11-25
213
+
214
+ ### 🚀 Features
215
+
216
+ - Improvements to NvidiaRanker and adding user input timeout (#1193)
217
+
191
218
  ### 🧹 Chores
192
219
 
193
220
  - Fix linting/isort (#1215)
@@ -204,18 +231,34 @@
204
231
 
205
232
  ### 🚀 Features
206
233
 
207
- - Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
208
- - Add NVIDIA NIM ranker support (#1023)
209
234
  - Raise error when attempting to embed empty documents/strings with Nvidia embedders (#1118)
210
235
 
211
236
  ### 🐛 Bug Fixes
212
237
 
213
- - Lints in `nvidia-haystack` (#993)
214
238
  - Missing Nvidia embedding truncate mode (#1043)
215
239
 
216
- ### 🚜 Refactor
240
+ ### ⚙️ CI
217
241
 
218
- - Remove deprecated Nvidia Cloud Functions backend and related code. (#803)
242
+ - Adopt uv as installer (#1142)
243
+
244
+ ### 🧹 Chores
245
+
246
+ - Update ruff linting scripts and settings (#1105)
247
+
248
+ ### 🌀 Miscellaneous
249
+
250
+ - Feat: add nvidia/llama-3.2-nv-rerankqa-1b-v1 to set of known ranking models (#1183)
251
+
252
+ ## [integrations/nvidia-v0.0.5] - 2024-08-28
253
+
254
+ ### 🚀 Features
255
+
256
+ - Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
257
+ - Add NVIDIA NIM ranker support (#1023)
258
+
259
+ ### 🐛 Bug Fixes
260
+
261
+ - Lints in `nvidia-haystack` (#993)
219
262
 
220
263
  ### 📚 Documentation
221
264
 
@@ -225,25 +268,31 @@
225
268
 
226
269
  - Do not retry tests in `hatch run test` command (#954)
227
270
 
271
+ ### 🌀 Miscellaneous
272
+
273
+ - Raise warning for base_url ../embeddings .../completions .../rankings (#922)
274
+ - Update NvidiaGenerator docstrings (#966)
275
+ - Add default model for NVIDIA HayStack local NIM endpoints (#915)
276
+
277
+ ## [integrations/nvidia-v0.0.4] - 2024-07-24
278
+
279
+ ### 🚜 Refactor
280
+
281
+ - Remove deprecated Nvidia Cloud Functions backend and related code. (#803)
282
+
228
283
  ### ⚙️ CI
229
284
 
230
285
  - Retry tests to reduce flakyness (#836)
231
- - Adopt uv as installer (#1142)
232
286
 
233
287
  ### 🧹 Chores
234
288
 
235
289
  - Update ruff invocation to include check parameter (#853)
236
- - Update ruff linting scripts and settings (#1105)
237
290
 
238
291
  ### 🌀 Miscellaneous
239
292
 
240
293
  - Fix: make hosted nim default (#734)
241
294
  - Fix: align tests and docs on NVIDIA_API_KEY (instead of NVIDIA_CATALOG_API_KEY) (#731)
242
295
  - Ci: install `pytest-rerunfailures` where needed; add retry config to `test-cov` script (#845)
243
- - Raise warning for base_url ../embeddings .../completions .../rankings (#922)
244
- - Update NvidiaGenerator docstrings (#966)
245
- - Add default model for NVIDIA HayStack local NIM endpoints (#915)
246
- - Feat: add nvidia/llama-3.2-nv-rerankqa-1b-v1 to set of known ranking models (#1183)
247
296
 
248
297
  ## [integrations/nvidia-v0.0.3] - 2024-05-22
249
298
 
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: nvidia-haystack
3
- Version: 1.4.0
3
+ Version: 3.0.0
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
@@ -10,15 +10,14 @@ License-File: LICENSE.txt
10
10
  Classifier: Development Status :: 4 - Beta
11
11
  Classifier: License :: OSI Approved :: Apache Software License
12
12
  Classifier: Programming Language :: Python
13
- Classifier: Programming Language :: Python :: 3.10
14
13
  Classifier: Programming Language :: Python :: 3.11
15
14
  Classifier: Programming Language :: Python :: 3.12
16
15
  Classifier: Programming Language :: Python :: 3.13
17
16
  Classifier: Programming Language :: Python :: 3.14
18
17
  Classifier: Programming Language :: Python :: Implementation :: CPython
19
18
  Classifier: Programming Language :: Python :: Implementation :: PyPy
20
- Requires-Python: >=3.10
21
- Requires-Dist: haystack-ai>=2.22.0
19
+ Requires-Python: >=3.11
20
+ Requires-Dist: haystack-ai>=3.0.0
22
21
  Requires-Dist: requests>=2.25.0
23
22
  Requires-Dist: tqdm>=4.21.0
24
23
  Description-Content-Type: text/markdown
@@ -27,7 +27,7 @@ chat_messages = [
27
27
  ]
28
28
 
29
29
  component = NvidiaChatGenerator(
30
- model="meta/llama-3.1-70b-instruct",
30
+ model="nvidia/nemotron-3.5-lightning-30b-a3b",
31
31
  generation_kwargs={"extra_body": {"nvext": {"guided_json": json_schema}}},
32
32
  )
33
33
  results = component.run(chat_messages)
@@ -3,7 +3,6 @@ loaders:
3
3
  - haystack_integrations.components.embedders.nvidia.document_embedder
4
4
  - haystack_integrations.components.embedders.nvidia.text_embedder
5
5
  - haystack_integrations.components.embedders.nvidia.truncate
6
- - haystack_integrations.components.generators.nvidia.generator
7
6
  - haystack_integrations.components.generators.nvidia.chat.chat_generator
8
7
  - haystack_integrations.components.rankers.nvidia.ranker
9
8
  - haystack_integrations.components.rankers.nvidia.truncate
@@ -7,7 +7,7 @@ name = "nvidia-haystack"
7
7
  dynamic = ["version"]
8
8
  description = ''
9
9
  readme = "README.md"
10
- requires-python = ">=3.10"
10
+ requires-python = ">=3.11"
11
11
  license = "Apache-2.0"
12
12
  keywords = []
13
13
  authors = [{ name = "deepset GmbH", email = "info@deepset.ai" }]
@@ -15,7 +15,6 @@ classifiers = [
15
15
  "License :: OSI Approved :: Apache Software License",
16
16
  "Development Status :: 4 - Beta",
17
17
  "Programming Language :: Python",
18
- "Programming Language :: Python :: 3.10",
19
18
  "Programming Language :: Python :: 3.11",
20
19
  "Programming Language :: Python :: 3.12",
21
20
  "Programming Language :: Python :: 3.13",
@@ -23,7 +22,7 @@ classifiers = [
23
22
  "Programming Language :: Python :: Implementation :: CPython",
24
23
  "Programming Language :: Python :: Implementation :: PyPy",
25
24
  ]
26
- dependencies = ["haystack-ai>=2.22.0", "requests>=2.25.0", "tqdm>=4.21.0"]
25
+ dependencies = ["haystack-ai>=3.0.0", "requests>=2.25.0", "tqdm>=4.21.0"]
27
26
 
28
27
  [project.urls]
29
28
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme"
@@ -168,4 +167,4 @@ addopts = "--strict-markers"
168
167
  markers = [
169
168
  "integration: integration tests",
170
169
  ]
171
- log_cli = true
170
+ log_cli = true
@@ -7,8 +7,8 @@ import warnings
7
7
  from dataclasses import replace
8
8
  from typing import Any
9
9
 
10
- from haystack import Document, component, default_from_dict, default_to_dict, logging
11
- from haystack.utils import Secret, deserialize_secrets_inplace
10
+ from haystack import Document, component, default_to_dict, logging
11
+ from haystack.utils import Secret
12
12
  from tqdm import tqdm
13
13
 
14
14
  from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
@@ -28,7 +28,7 @@ class NvidiaDocumentEmbedder:
28
28
 
29
29
  doc = Document(content="I love pizza!")
30
30
 
31
- text_embedder = NvidiaDocumentEmbedder(model="nvidia/nv-embedqa-e5-v5", api_url="https://integrate.api.nvidia.com/v1")
31
+ text_embedder = NvidiaDocumentEmbedder(model="nvidia/nemotron-3-embed-1b", api_url="https://integrate.api.nvidia.com/v1")
32
32
  # Components warm up automatically on first run.
33
33
 
34
34
  result = document_embedder.run([doc])
@@ -97,8 +97,7 @@ class NvidiaDocumentEmbedder:
97
97
  truncate = EmbeddingTruncateMode.from_str(truncate)
98
98
  self.truncate = truncate
99
99
 
100
- self.backend: Any | None = None
101
- self._initialized = False
100
+ self.backend: NimBackend | None = None
102
101
 
103
102
  if timeout is None:
104
103
  timeout = float(os.environ.get("NVIDIA_TIMEOUT", "60.0"))
@@ -134,7 +133,7 @@ class NvidiaDocumentEmbedder:
134
133
  """
135
134
  Initializes the component.
136
135
  """
137
- if self._initialized:
136
+ if self.backend is not None:
138
137
  return
139
138
 
140
139
  model_kwargs = {"input_type": "passage"}
@@ -152,11 +151,15 @@ class NvidiaDocumentEmbedder:
152
151
  if not self.model and self.backend.model:
153
152
  self.model = self.backend.model
154
153
 
155
- self._initialized = True
156
-
157
154
  if not self.model:
158
155
  self.default_model()
159
156
 
157
+ def close(self) -> None:
158
+ """Close the backend and release its resources."""
159
+ if self.backend is not None:
160
+ self.backend.close()
161
+ self.backend = None
162
+
160
163
  def to_dict(self) -> dict[str, Any]:
161
164
  """
162
165
  Serializes the component to a dictionary.
@@ -186,21 +189,6 @@ class NvidiaDocumentEmbedder:
186
189
  """
187
190
  return self.backend.models() if self.backend else []
188
191
 
189
- @classmethod
190
- def from_dict(cls, data: dict[str, Any]) -> "NvidiaDocumentEmbedder":
191
- """
192
- Deserializes the component from a dictionary.
193
-
194
- :param data:
195
- The dictionary to deserialize from.
196
- :returns:
197
- The deserialized component.
198
- """
199
- init_parameters = data.get("init_parameters", {})
200
- if init_parameters:
201
- deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
202
- return default_from_dict(cls, data)
203
-
204
192
  def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
205
193
  texts_to_embed = []
206
194
  for doc in documents:
@@ -250,8 +238,7 @@ class NvidiaDocumentEmbedder:
250
238
  :raises TypeError:
251
239
  If the input is not a list of Documents.
252
240
  """
253
- if not self._initialized:
254
- self.warm_up()
241
+ self.warm_up()
255
242
 
256
243
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
257
244
  msg = (
@@ -6,8 +6,8 @@ import os
6
6
  import warnings
7
7
  from typing import Any
8
8
 
9
- from haystack import component, default_from_dict, default_to_dict, logging
10
- from haystack.utils import Secret, deserialize_secrets_inplace
9
+ from haystack import component, default_to_dict, logging
10
+ from haystack.utils import Secret
11
11
 
12
12
  from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
13
13
  from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, Model, NimBackend, url_validation
@@ -29,7 +29,7 @@ class NvidiaTextEmbedder:
29
29
 
30
30
  text_to_embed = "I love pizza!"
31
31
 
32
- text_embedder = NvidiaTextEmbedder(model="nvidia/nv-embedqa-e5-v5", api_url="https://integrate.api.nvidia.com/v1")
32
+ text_embedder = NvidiaTextEmbedder(model="nvidia/nemotron-3-embed-1b", api_url="https://integrate.api.nvidia.com/v1")
33
33
  # Components warm up automatically on first run.
34
34
 
35
35
  print(text_embedder.run(text_to_embed))
@@ -80,8 +80,7 @@ class NvidiaTextEmbedder:
80
80
  truncate = EmbeddingTruncateMode.from_str(truncate)
81
81
  self.truncate = truncate
82
82
 
83
- self.backend: Any | None = None
84
- self._initialized = False
83
+ self.backend: NimBackend | None = None
85
84
 
86
85
  if timeout is None:
87
86
  timeout = float(os.environ.get("NVIDIA_TIMEOUT", "60.0"))
@@ -123,7 +122,7 @@ class NvidiaTextEmbedder:
123
122
  """
124
123
  Initializes the component.
125
124
  """
126
- if self._initialized:
125
+ if self.backend is not None:
127
126
  return
128
127
 
129
128
  model_kwargs = {"input_type": "query"}
@@ -138,14 +137,18 @@ class NvidiaTextEmbedder:
138
137
  timeout=self.timeout,
139
138
  client=Client.NVIDIA_TEXT_EMBEDDER,
140
139
  )
141
- self._initialized = True
142
-
143
140
  if not self.model:
144
141
  if self.backend.model:
145
142
  self.model = self.backend.model
146
143
  else:
147
144
  self.default_model()
148
145
 
146
+ def close(self) -> None:
147
+ """Close the backend and release its resources."""
148
+ if self.backend is not None:
149
+ self.backend.close()
150
+ self.backend = None
151
+
149
152
  def to_dict(self) -> dict[str, Any]:
150
153
  """
151
154
  Serializes the component to a dictionary.
@@ -171,21 +174,6 @@ class NvidiaTextEmbedder:
171
174
  """
172
175
  return self.backend.models() if self.backend else []
173
176
 
174
- @classmethod
175
- def from_dict(cls, data: dict[str, Any]) -> "NvidiaTextEmbedder":
176
- """
177
- Deserializes the component from a dictionary.
178
-
179
- :param data:
180
- The dictionary to deserialize from.
181
- :returns:
182
- The deserialized component.
183
- """
184
- init_parameters = data.get("init_parameters", {})
185
- if init_parameters:
186
- deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
187
- return default_from_dict(cls, data)
188
-
189
177
  @component.output_types(embedding=list[float], meta=dict[str, Any])
190
178
  def run(self, text: str) -> dict[str, list[float] | dict[str, Any]]:
191
179
  """
@@ -202,8 +190,7 @@ class NvidiaTextEmbedder:
202
190
  :raises ValueError:
203
191
  If the input string is empty.
204
192
  """
205
- if not self._initialized:
206
- self.warm_up()
193
+ self.warm_up()
207
194
 
208
195
  if not isinstance(text, str):
209
196
  msg = (
@@ -3,6 +3,5 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  from .chat.chat_generator import NvidiaChatGenerator
6
- from .generator import NvidiaGenerator
7
6
 
8
- __all__ = ["NvidiaChatGenerator", "NvidiaGenerator"]
7
+ __all__ = ["NvidiaChatGenerator"]
@@ -53,7 +53,7 @@ class NvidiaChatGenerator(OpenAIChatGenerator):
53
53
  self,
54
54
  *,
55
55
  api_key: Secret = Secret.from_env_var("NVIDIA_API_KEY"),
56
- model: str = "meta/llama-3.1-8b-instruct",
56
+ model: str = "nvidia/nemotron-3.5-lightning-30b-a3b",
57
57
  streaming_callback: StreamingCallbackT | None = None,
58
58
  api_base_url: str | None = os.getenv("NVIDIA_API_URL", DEFAULT_API_URL),
59
59
  generation_kwargs: dict[str, Any] | None = None,
@@ -6,8 +6,8 @@ import os
6
6
  from dataclasses import replace
7
7
  from typing import Any
8
8
 
9
- from haystack import Document, component, default_from_dict, default_to_dict
10
- from haystack.utils import Secret, deserialize_secrets_inplace
9
+ from haystack import Document, component, default_to_dict
10
+ from haystack.utils import Secret
11
11
 
12
12
  from haystack_integrations.components.rankers.nvidia.truncate import RankerTruncateMode
13
13
  from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, NimBackend, is_hosted, url_validation
@@ -25,7 +25,7 @@ class NvidiaRanker:
25
25
  from haystack.utils import Secret
26
26
 
27
27
  ranker = NvidiaRanker(
28
- model="nvidia/nv-rerankqa-mistral-4b-v3",
28
+ model="nvidia/llama-nemotron-rerank-vl-1b-v2",
29
29
  api_key=Secret.from_env_var("NVIDIA_API_KEY"),
30
30
  )
31
31
  # Components warm up automatically on first run.
@@ -108,8 +108,7 @@ class NvidiaRanker:
108
108
  # - assume we won't call backend.models()
109
109
  self.api_url = url_validation(api_url)
110
110
  self.top_k = top_k
111
- self._initialized = False
112
- self.backend: Any | None = None
111
+ self.backend: NimBackend | None = None
113
112
  self.is_hosted = is_hosted(api_url)
114
113
 
115
114
  self.query_prefix = query_prefix
@@ -145,42 +144,35 @@ class NvidiaRanker:
145
144
  timeout=self.timeout,
146
145
  )
147
146
 
148
- @classmethod
149
- def from_dict(cls, data: dict[str, Any]) -> "NvidiaRanker":
150
- """
151
- Deserialize the ranker from a dictionary.
152
-
153
- :param data: A dictionary containing the ranker's attributes.
154
- :returns: The deserialized ranker.
155
- """
156
- init_parameters = data.get("init_parameters", {})
157
- if init_parameters:
158
- deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
159
- return default_from_dict(cls, data)
160
-
161
147
  def warm_up(self) -> None:
162
148
  """
163
149
  Initialize the ranker.
164
150
 
165
151
  :raises ValueError: If the API key is required for hosted NVIDIA NIMs.
166
152
  """
167
- if not self._initialized:
168
- model_kwargs: dict[str, Any] = {}
169
- if self.truncate is not None:
170
- model_kwargs.update(truncate=str(self.truncate))
171
- self.backend = NimBackend(
172
- model=self.model,
173
- model_type="ranking",
174
- api_url=self.api_url,
175
- api_key=self.api_key,
176
- model_kwargs=model_kwargs,
177
- timeout=self.timeout,
178
- client=Client.NVIDIA_RANKER,
179
- )
180
- if not self.is_hosted and not self.model:
181
- if self.backend.model:
182
- self.model = self.backend.model
183
- self._initialized = True
153
+ if self.backend is not None:
154
+ return
155
+
156
+ model_kwargs: dict[str, Any] = {}
157
+ if self.truncate is not None:
158
+ model_kwargs.update(truncate=str(self.truncate))
159
+ self.backend = NimBackend(
160
+ model=self.model,
161
+ model_type="ranking",
162
+ api_url=self.api_url,
163
+ api_key=self.api_key,
164
+ model_kwargs=model_kwargs,
165
+ timeout=self.timeout,
166
+ client=Client.NVIDIA_RANKER,
167
+ )
168
+ if not self.is_hosted and not self.model and self.backend.model:
169
+ self.model = self.backend.model
170
+
171
+ def close(self) -> None:
172
+ """Close the backend and release its resources."""
173
+ if self.backend is not None:
174
+ self.backend.close()
175
+ self.backend = None
184
176
 
185
177
  def _prepare_documents_to_embed(self, documents: list[Document]) -> list[str]:
186
178
  document_texts = []
@@ -208,8 +200,7 @@ class NvidiaRanker:
208
200
 
209
201
  :returns: A dictionary containing the ranked documents.
210
202
  """
211
- if not self._initialized:
212
- self.warm_up()
203
+ self.warm_up()
213
204
 
214
205
  if not isinstance(query, str):
215
206
  msg = "NvidiaRanker expects the `query` parameter to be a string."
@@ -5,7 +5,7 @@
5
5
  from enum import Enum
6
6
 
7
7
 
8
- class RankerTruncateMode(str, Enum):
8
+ class RankerTruncateMode(str, Enum): # noqa: UP042 # StrEnum would change str() of this public enum
9
9
  """
10
10
  Specifies how inputs to the NVIDIA ranker components are truncated.
11
11
 
@@ -308,6 +308,13 @@ CHAT_MODEL_TABLE = {
308
308
  supports_tools=True,
309
309
  supports_structured_output=True,
310
310
  ),
311
+ "nvidia/nemotron-3.5-lightning-30b-a3b": Model(
312
+ id="nvidia/nemotron-3.5-lightning-30b-a3b",
313
+ model_type="chat",
314
+ client=Client.NVIDIA_GENERATOR,
315
+ supports_tools=True,
316
+ supports_structured_output=True,
317
+ ),
311
318
  "meta/llama-3.1-405b-instruct": Model(
312
319
  id="meta/llama-3.1-405b-instruct",
313
320
  model_type="chat",
@@ -430,6 +437,10 @@ EMBEDDING_MODEL_TABLE = {
430
437
  id="nvidia/nv-embedqa-mistral-7b-v2",
431
438
  model_type="embedding",
432
439
  ),
440
+ "nvidia/nemotron-3-embed-1b": Model(
441
+ id="nvidia/nemotron-3-embed-1b",
442
+ model_type="embedding",
443
+ ),
433
444
  "nvidia/nv-embedqa-e5-v5": Model(
434
445
  id="nvidia/nv-embedqa-e5-v5",
435
446
  model_type="embedding",
@@ -452,6 +463,12 @@ EMBEDDING_MODEL_TABLE = {
452
463
  ),
453
464
  }
454
465
  RANKING_MODEL_TABLE = {
466
+ "nvidia/llama-nemotron-rerank-vl-1b-v2": Model(
467
+ id="nvidia/llama-nemotron-rerank-vl-1b-v2",
468
+ model_type="ranking",
469
+ client=Client.NVIDIA_RANKER,
470
+ endpoint="https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-vl-1b-v2/reranking",
471
+ ),
455
472
  "nv-rerank-qa-mistral-4b:1": Model(
456
473
  id="nv-rerank-qa-mistral-4b:1",
457
474
  model_type="ranking",
@@ -480,9 +497,9 @@ RANKING_MODEL_TABLE = {
480
497
  }
481
498
 
482
499
  DEFAULT_MODELS = {
483
- "embedding": "nvidia/nv-embedqa-e5-v5",
484
- "ranking": "nv-rerank-qa-mistral-4b:1",
485
- "chat": "meta/llama3-8b-instruct",
500
+ "embedding": "nvidia/nemotron-3-embed-1b",
501
+ "ranking": "nvidia/llama-nemotron-rerank-vl-1b-v2",
502
+ "chat": "nvidia/nemotron-3.5-lightning-30b-a3b",
486
503
  }
487
504
 
488
505
 
@@ -38,9 +38,6 @@ class NimBackend:
38
38
  if api_key:
39
39
  headers["authorization"] = f"Bearer {api_key.resolve_value()}"
40
40
 
41
- self.session = requests.Session()
42
- self.session.headers.update(headers)
43
-
44
41
  self.api_url = api_url
45
42
  # Hosted ranking models are assigned a per-model endpoint that already includes the
46
43
  # full path (e.g. ".../nv-rerankqa-mistral-4b-v3/reranking"), unlike other hosted
@@ -79,6 +76,12 @@ class NimBackend:
79
76
  if timeout is None:
80
77
  timeout = float(os.environ.get("NVIDIA_TIMEOUT", REQUEST_TIMEOUT))
81
78
  self.timeout = timeout
79
+ self.session = requests.Session()
80
+ self.session.headers.update(headers)
81
+
82
+ def close(self) -> None:
83
+ """Close the HTTP session owned by this backend."""
84
+ self.session.close()
82
85
 
83
86
  def embed(self, texts: list[str]) -> tuple[list[list[float]], dict[str, Any]]:
84
87
  """Compute embeddings for a list of texts via the NIM API."""
@@ -111,65 +114,6 @@ class NimBackend:
111
114
 
112
115
  return embeddings, {"usage": data["usage"]}
113
116
 
114
- def generate(self, prompt: str) -> tuple[list[str], list[dict[str, Any]]]:
115
- """Generate text completions for a prompt via the NIM API."""
116
- # We're using the chat completion endpoint as the NIM API doesn't support
117
- # the /completions endpoint. So both the non-chat and chat generator will use this.
118
- # This is the same for local containers and the cloud API.
119
- url = f"{self.api_url}/chat/completions"
120
-
121
- try:
122
- res = self.session.post(
123
- url,
124
- json={
125
- "model": self.model,
126
- "messages": [
127
- {
128
- "role": "user",
129
- "content": prompt,
130
- },
131
- ],
132
- **self.model_kwargs,
133
- },
134
- timeout=self.timeout,
135
- )
136
- res.raise_for_status()
137
- except requests.HTTPError as e:
138
- response_text = e.response.text.strip() if e.response is not None else ""
139
- error_details = f": {response_text}" if response_text else ""
140
- logger.error(
141
- "Error when calling NIM chat completion endpoint{error}",
142
- error=error_details,
143
- )
144
- msg = f"Failed to query chat completion endpoint{error_details}"
145
- raise ValueError(msg) from e
146
-
147
- completions = res.json()
148
- choices = completions["choices"]
149
- # Sort the choices by index, we don't know whether they're out of order or not
150
- choices.sort(key=lambda c: c["index"])
151
- replies = []
152
- meta = []
153
- for choice in choices:
154
- message = choice["message"]
155
- replies.append(message["content"])
156
- choice_meta = {
157
- "role": message["role"],
158
- "usage": {
159
- "prompt_tokens": completions["usage"]["prompt_tokens"],
160
- "total_tokens": completions["usage"]["total_tokens"],
161
- },
162
- }
163
- # These fields could be null, the others will always be present
164
- if "finish_reason" in choice:
165
- choice_meta["finish_reason"] = choice["finish_reason"]
166
- if "completion_tokens" in completions["usage"]:
167
- choice_meta["usage"]["completion_tokens"] = completions["usage"]["completion_tokens"]
168
-
169
- meta.append(choice_meta)
170
-
171
- return replies, meta
172
-
173
117
  def models(self) -> list[Model]:
174
118
  """Retrieve available models from the NIM API."""
175
119
  url = f"{self.api_url}/models"