nvidia-haystack 1.4.0__tar.gz → 2.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.
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/CHANGELOG.md +46 -12
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/PKG-INFO +3 -3
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/examples/chat_generator_with_structured_outputs.py +1 -1
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -1
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/pyproject.toml +2 -2
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +10 -8
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +10 -8
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/generators/nvidia/__init__.py +1 -2
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py +1 -1
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/rankers/nvidia/ranker.py +26 -22
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/models.py +20 -3
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/nim_backend.py +6 -62
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_base_url.py +5 -6
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_document_embedder.py +45 -10
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_nim_backend.py +9 -74
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_nvidia_chat_generator.py +67 -94
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_ranker.py +109 -82
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_text_embedder.py +45 -9
- nvidia_haystack-1.4.0/src/haystack_integrations/components/generators/nvidia/generator.py +0 -199
- nvidia_haystack-1.4.0/tests/test_generator.py +0 -302
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/.gitignore +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/LICENSE.txt +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/README.md +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/embedders/nvidia/__init__.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/embedders/nvidia/truncate.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/generators/nvidia/chat/__init__.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/rankers/nvidia/__init__.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/rankers/nvidia/truncate.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/__init__.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/client.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/utils.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/py.typed +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/__init__.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/conftest.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_embedding_truncate_mode.py +0 -0
- {nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/tests/test_utils.py +0 -0
|
@@ -1,5 +1,12 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [integrations/nvidia-v1.4.0] - 2026-08-11
|
|
4
|
+
|
|
5
|
+
### 🐛 Bug Fixes
|
|
6
|
+
|
|
7
|
+
- NvidiaRanker reordered the top_k fallback to happen after validation (#3749)
|
|
8
|
+
|
|
9
|
+
|
|
3
10
|
## [integrations/nvidia-v1.3.0] - 2026-08-04
|
|
4
11
|
|
|
5
12
|
### 🐛 Bug Fixes
|
|
@@ -185,9 +192,14 @@
|
|
|
185
192
|
|
|
186
193
|
### 🚀 Features
|
|
187
194
|
|
|
188
|
-
- Improvements to NvidiaRanker and adding user input timeout (#1193)
|
|
189
195
|
- Add model `nvidia/llama-3.2-nv-rerankqa-1b-v2` to `_MODEL_ENDPOINT_MAP` (#1260)
|
|
190
196
|
|
|
197
|
+
## [integrations/nvidia-v0.1.2] - 2024-11-25
|
|
198
|
+
|
|
199
|
+
### 🚀 Features
|
|
200
|
+
|
|
201
|
+
- Improvements to NvidiaRanker and adding user input timeout (#1193)
|
|
202
|
+
|
|
191
203
|
### 🧹 Chores
|
|
192
204
|
|
|
193
205
|
- Fix linting/isort (#1215)
|
|
@@ -204,18 +216,34 @@
|
|
|
204
216
|
|
|
205
217
|
### 🚀 Features
|
|
206
218
|
|
|
207
|
-
- Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
|
|
208
|
-
- Add NVIDIA NIM ranker support (#1023)
|
|
209
219
|
- Raise error when attempting to embed empty documents/strings with Nvidia embedders (#1118)
|
|
210
220
|
|
|
211
221
|
### 🐛 Bug Fixes
|
|
212
222
|
|
|
213
|
-
- Lints in `nvidia-haystack` (#993)
|
|
214
223
|
- Missing Nvidia embedding truncate mode (#1043)
|
|
215
224
|
|
|
216
|
-
###
|
|
225
|
+
### ⚙️ CI
|
|
217
226
|
|
|
218
|
-
-
|
|
227
|
+
- Adopt uv as installer (#1142)
|
|
228
|
+
|
|
229
|
+
### 🧹 Chores
|
|
230
|
+
|
|
231
|
+
- Update ruff linting scripts and settings (#1105)
|
|
232
|
+
|
|
233
|
+
### 🌀 Miscellaneous
|
|
234
|
+
|
|
235
|
+
- Feat: add nvidia/llama-3.2-nv-rerankqa-1b-v1 to set of known ranking models (#1183)
|
|
236
|
+
|
|
237
|
+
## [integrations/nvidia-v0.0.5] - 2024-08-28
|
|
238
|
+
|
|
239
|
+
### 🚀 Features
|
|
240
|
+
|
|
241
|
+
- Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
|
|
242
|
+
- Add NVIDIA NIM ranker support (#1023)
|
|
243
|
+
|
|
244
|
+
### 🐛 Bug Fixes
|
|
245
|
+
|
|
246
|
+
- Lints in `nvidia-haystack` (#993)
|
|
219
247
|
|
|
220
248
|
### 📚 Documentation
|
|
221
249
|
|
|
@@ -225,25 +253,31 @@
|
|
|
225
253
|
|
|
226
254
|
- Do not retry tests in `hatch run test` command (#954)
|
|
227
255
|
|
|
256
|
+
### 🌀 Miscellaneous
|
|
257
|
+
|
|
258
|
+
- Raise warning for base_url ../embeddings .../completions .../rankings (#922)
|
|
259
|
+
- Update NvidiaGenerator docstrings (#966)
|
|
260
|
+
- Add default model for NVIDIA HayStack local NIM endpoints (#915)
|
|
261
|
+
|
|
262
|
+
## [integrations/nvidia-v0.0.4] - 2024-07-24
|
|
263
|
+
|
|
264
|
+
### 🚜 Refactor
|
|
265
|
+
|
|
266
|
+
- Remove deprecated Nvidia Cloud Functions backend and related code. (#803)
|
|
267
|
+
|
|
228
268
|
### ⚙️ CI
|
|
229
269
|
|
|
230
270
|
- Retry tests to reduce flakyness (#836)
|
|
231
|
-
- Adopt uv as installer (#1142)
|
|
232
271
|
|
|
233
272
|
### 🧹 Chores
|
|
234
273
|
|
|
235
274
|
- Update ruff invocation to include check parameter (#853)
|
|
236
|
-
- Update ruff linting scripts and settings (#1105)
|
|
237
275
|
|
|
238
276
|
### 🌀 Miscellaneous
|
|
239
277
|
|
|
240
278
|
- Fix: make hosted nim default (#734)
|
|
241
279
|
- Fix: align tests and docs on NVIDIA_API_KEY (instead of NVIDIA_CATALOG_API_KEY) (#731)
|
|
242
280
|
- 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
281
|
|
|
248
282
|
## [integrations/nvidia-v0.0.3] - 2024-05-22
|
|
249
283
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: nvidia-haystack
|
|
3
|
-
Version:
|
|
3
|
+
Version: 2.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
|
|
@@ -18,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.14
|
|
|
18
18
|
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
19
19
|
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
20
20
|
Requires-Python: >=3.10
|
|
21
|
-
Requires-Dist: haystack-ai>=
|
|
21
|
+
Requires-Dist: haystack-ai>=3.0.0
|
|
22
22
|
Requires-Dist: requests>=2.25.0
|
|
23
23
|
Requires-Dist: tqdm>=4.21.0
|
|
24
24
|
Description-Content-Type: text/markdown
|
{nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/examples/chat_generator_with_structured_outputs.py
RENAMED
|
@@ -27,7 +27,7 @@ chat_messages = [
|
|
|
27
27
|
]
|
|
28
28
|
|
|
29
29
|
component = NvidiaChatGenerator(
|
|
30
|
-
model="
|
|
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
|
|
@@ -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>=
|
|
26
|
+
dependencies = ["haystack-ai>=3.0.0", "requests>=2.25.0", "tqdm>=4.21.0"]
|
|
27
27
|
|
|
28
28
|
[project.urls]
|
|
29
29
|
Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme"
|
|
@@ -168,4 +168,4 @@ addopts = "--strict-markers"
|
|
|
168
168
|
markers = [
|
|
169
169
|
"integration: integration tests",
|
|
170
170
|
]
|
|
171
|
-
log_cli = true
|
|
171
|
+
log_cli = true
|
|
@@ -28,7 +28,7 @@ class NvidiaDocumentEmbedder:
|
|
|
28
28
|
|
|
29
29
|
doc = Document(content="I love pizza!")
|
|
30
30
|
|
|
31
|
-
text_embedder = NvidiaDocumentEmbedder(model="nvidia/
|
|
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:
|
|
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.
|
|
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.
|
|
@@ -250,8 +253,7 @@ class NvidiaDocumentEmbedder:
|
|
|
250
253
|
:raises TypeError:
|
|
251
254
|
If the input is not a list of Documents.
|
|
252
255
|
"""
|
|
253
|
-
|
|
254
|
-
self.warm_up()
|
|
256
|
+
self.warm_up()
|
|
255
257
|
|
|
256
258
|
if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
|
|
257
259
|
msg = (
|
|
@@ -29,7 +29,7 @@ class NvidiaTextEmbedder:
|
|
|
29
29
|
|
|
30
30
|
text_to_embed = "I love pizza!"
|
|
31
31
|
|
|
32
|
-
text_embedder = NvidiaTextEmbedder(model="nvidia/
|
|
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:
|
|
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.
|
|
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.
|
|
@@ -202,8 +205,7 @@ class NvidiaTextEmbedder:
|
|
|
202
205
|
:raises ValueError:
|
|
203
206
|
If the input string is empty.
|
|
204
207
|
"""
|
|
205
|
-
|
|
206
|
-
self.warm_up()
|
|
208
|
+
self.warm_up()
|
|
207
209
|
|
|
208
210
|
if not isinstance(text, str):
|
|
209
211
|
msg = (
|
|
@@ -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 = "
|
|
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,
|
|
@@ -25,7 +25,7 @@ class NvidiaRanker:
|
|
|
25
25
|
from haystack.utils import Secret
|
|
26
26
|
|
|
27
27
|
ranker = NvidiaRanker(
|
|
28
|
-
model="nvidia/
|
|
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.
|
|
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
|
|
@@ -164,23 +163,29 @@ class NvidiaRanker:
|
|
|
164
163
|
|
|
165
164
|
:raises ValueError: If the API key is required for hosted NVIDIA NIMs.
|
|
166
165
|
"""
|
|
167
|
-
if not
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
self.
|
|
166
|
+
if self.backend is not None:
|
|
167
|
+
return
|
|
168
|
+
|
|
169
|
+
model_kwargs: dict[str, Any] = {}
|
|
170
|
+
if self.truncate is not None:
|
|
171
|
+
model_kwargs.update(truncate=str(self.truncate))
|
|
172
|
+
self.backend = NimBackend(
|
|
173
|
+
model=self.model,
|
|
174
|
+
model_type="ranking",
|
|
175
|
+
api_url=self.api_url,
|
|
176
|
+
api_key=self.api_key,
|
|
177
|
+
model_kwargs=model_kwargs,
|
|
178
|
+
timeout=self.timeout,
|
|
179
|
+
client=Client.NVIDIA_RANKER,
|
|
180
|
+
)
|
|
181
|
+
if not self.is_hosted and not self.model and self.backend.model:
|
|
182
|
+
self.model = self.backend.model
|
|
183
|
+
|
|
184
|
+
def close(self) -> None:
|
|
185
|
+
"""Close the backend and release its resources."""
|
|
186
|
+
if self.backend is not None:
|
|
187
|
+
self.backend.close()
|
|
188
|
+
self.backend = None
|
|
184
189
|
|
|
185
190
|
def _prepare_documents_to_embed(self, documents: list[Document]) -> list[str]:
|
|
186
191
|
document_texts = []
|
|
@@ -208,8 +213,7 @@ class NvidiaRanker:
|
|
|
208
213
|
|
|
209
214
|
:returns: A dictionary containing the ranked documents.
|
|
210
215
|
"""
|
|
211
|
-
|
|
212
|
-
self.warm_up()
|
|
216
|
+
self.warm_up()
|
|
213
217
|
|
|
214
218
|
if not isinstance(query, str):
|
|
215
219
|
msg = "NvidiaRanker expects the `query` parameter to be a string."
|
{nvidia_haystack-1.4.0 → nvidia_haystack-2.0.0}/src/haystack_integrations/utils/nvidia/models.py
RENAMED
|
@@ -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/
|
|
484
|
-
"ranking": "
|
|
485
|
-
"chat": "
|
|
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"
|
|
@@ -7,7 +7,6 @@ from typing import Any
|
|
|
7
7
|
import pytest
|
|
8
8
|
|
|
9
9
|
from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaTextEmbedder
|
|
10
|
-
from haystack_integrations.components.generators.nvidia import NvidiaGenerator
|
|
11
10
|
from haystack_integrations.components.rankers.nvidia import NvidiaRanker
|
|
12
11
|
|
|
13
12
|
|
|
@@ -23,7 +22,7 @@ from haystack_integrations.components.rankers.nvidia import NvidiaRanker
|
|
|
23
22
|
)
|
|
24
23
|
@pytest.mark.parametrize(
|
|
25
24
|
"component",
|
|
26
|
-
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker
|
|
25
|
+
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker],
|
|
27
26
|
)
|
|
28
27
|
def test_base_url_invalid_not_hosted(base_url: str, component) -> None:
|
|
29
28
|
with pytest.warns(UserWarning) as msg:
|
|
@@ -34,7 +33,7 @@ def test_base_url_invalid_not_hosted(base_url: str, component) -> None:
|
|
|
34
33
|
|
|
35
34
|
@pytest.mark.parametrize(
|
|
36
35
|
"component",
|
|
37
|
-
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker
|
|
36
|
+
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker],
|
|
38
37
|
)
|
|
39
38
|
def test_create_without_base_url(component: type, monkeypatch) -> None:
|
|
40
39
|
monkeypatch.setenv("NVIDIA_API_KEY", "valid_api_key")
|
|
@@ -45,7 +44,7 @@ def test_create_without_base_url(component: type, monkeypatch) -> None:
|
|
|
45
44
|
|
|
46
45
|
@pytest.mark.parametrize(
|
|
47
46
|
"component",
|
|
48
|
-
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker
|
|
47
|
+
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker],
|
|
49
48
|
)
|
|
50
49
|
def test_base_url_priority(component: type) -> None:
|
|
51
50
|
param_url = "https://PARAM/v1"
|
|
@@ -59,7 +58,7 @@ def test_base_url_priority(component: type) -> None:
|
|
|
59
58
|
|
|
60
59
|
@pytest.mark.parametrize(
|
|
61
60
|
"component",
|
|
62
|
-
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker
|
|
61
|
+
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker],
|
|
63
62
|
)
|
|
64
63
|
@pytest.mark.parametrize(
|
|
65
64
|
"api_url",
|
|
@@ -78,7 +77,7 @@ def test_param_api_url_negative(component: type, api_url: str) -> None:
|
|
|
78
77
|
|
|
79
78
|
@pytest.mark.parametrize(
|
|
80
79
|
"component",
|
|
81
|
-
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker
|
|
80
|
+
[NvidiaDocumentEmbedder, NvidiaTextEmbedder, NvidiaRanker],
|
|
82
81
|
)
|
|
83
82
|
@pytest.mark.parametrize(
|
|
84
83
|
"api_url",
|
|
@@ -3,6 +3,7 @@
|
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
|
|
5
5
|
import os
|
|
6
|
+
from unittest.mock import patch
|
|
6
7
|
|
|
7
8
|
import pytest
|
|
8
9
|
from haystack import Document
|
|
@@ -14,20 +15,20 @@ from haystack_integrations.utils.nvidia import DEFAULT_API_URL
|
|
|
14
15
|
from . import MockBackend
|
|
15
16
|
|
|
16
17
|
|
|
17
|
-
class
|
|
18
|
+
class TestInitialization:
|
|
18
19
|
def test_init_default(self, monkeypatch):
|
|
19
20
|
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
20
21
|
embedder = NvidiaDocumentEmbedder()
|
|
21
|
-
embedder.warm_up()
|
|
22
22
|
|
|
23
23
|
assert embedder.api_key == Secret.from_env_var("NVIDIA_API_KEY")
|
|
24
|
-
assert embedder.model
|
|
24
|
+
assert embedder.model is None
|
|
25
25
|
assert embedder.prefix == ""
|
|
26
26
|
assert embedder.suffix == ""
|
|
27
27
|
assert embedder.batch_size == 32
|
|
28
28
|
assert embedder.progress_bar is True
|
|
29
29
|
assert embedder.meta_fields_to_embed == []
|
|
30
30
|
assert embedder.embedding_separator == "\n"
|
|
31
|
+
assert embedder.backend is None
|
|
31
32
|
|
|
32
33
|
def test_init_with_parameters(self):
|
|
33
34
|
embedder = NvidiaDocumentEmbedder(
|
|
@@ -51,12 +52,8 @@ class TestNvidiaDocumentEmbedder:
|
|
|
51
52
|
assert embedder.meta_fields_to_embed == ["test_field"]
|
|
52
53
|
assert embedder.embedding_separator == " | "
|
|
53
54
|
|
|
54
|
-
def test_init_fail_wo_api_key(self, monkeypatch):
|
|
55
|
-
monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
|
|
56
|
-
embedder = NvidiaDocumentEmbedder("nvolveqa_40k")
|
|
57
|
-
with pytest.raises(ValueError):
|
|
58
|
-
embedder.warm_up()
|
|
59
55
|
|
|
56
|
+
class TestSerialization:
|
|
60
57
|
def test_to_dict(self, monkeypatch):
|
|
61
58
|
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
62
59
|
component = NvidiaDocumentEmbedder("nvidia/nv-embedqa-e5-v5")
|
|
@@ -160,6 +157,44 @@ class TestNvidiaDocumentEmbedder:
|
|
|
160
157
|
assert component.truncate is None
|
|
161
158
|
assert component.timeout == 60.0
|
|
162
159
|
|
|
160
|
+
|
|
161
|
+
class TestComponentLifecycle:
|
|
162
|
+
def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
|
|
163
|
+
monkeypatch.delenv("MISSING_NVIDIA_API_KEY", raising=False)
|
|
164
|
+
embedder = NvidiaDocumentEmbedder("nvolveqa_40k", api_key=Secret.from_env_var("MISSING_NVIDIA_API_KEY"))
|
|
165
|
+
|
|
166
|
+
with pytest.raises(ValueError, match="MISSING_NVIDIA_API_KEY"):
|
|
167
|
+
embedder.warm_up()
|
|
168
|
+
|
|
169
|
+
@patch("haystack_integrations.components.embedders.nvidia.document_embedder.NimBackend")
|
|
170
|
+
def test_sync_lifecycle(self, mock_backend):
|
|
171
|
+
embedder = NvidiaDocumentEmbedder("nvolveqa_40k", api_key=Secret.from_token("test-api-key"))
|
|
172
|
+
backend = mock_backend.return_value
|
|
173
|
+
|
|
174
|
+
embedder.warm_up()
|
|
175
|
+
assert embedder.backend is backend
|
|
176
|
+
|
|
177
|
+
embedder.close()
|
|
178
|
+
backend.close.assert_called_once_with()
|
|
179
|
+
assert embedder.backend is None
|
|
180
|
+
|
|
181
|
+
embedder.warm_up()
|
|
182
|
+
assert mock_backend.call_count == 2
|
|
183
|
+
|
|
184
|
+
@patch("haystack_integrations.components.embedders.nvidia.document_embedder.NimBackend")
|
|
185
|
+
def test_warm_up_is_idempotent(self, mock_backend):
|
|
186
|
+
embedder = NvidiaDocumentEmbedder("nvolveqa_40k", api_key=Secret.from_token("test-api-key"))
|
|
187
|
+
embedder.warm_up()
|
|
188
|
+
embedder.warm_up()
|
|
189
|
+
mock_backend.assert_called_once()
|
|
190
|
+
|
|
191
|
+
def test_close_is_safe_without_warm_up(self):
|
|
192
|
+
embedder = NvidiaDocumentEmbedder("nvolveqa_40k", api_key=Secret.from_token("test-api-key"))
|
|
193
|
+
embedder.close()
|
|
194
|
+
assert embedder.backend is None
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
class TestRun:
|
|
163
198
|
def test_prepare_texts_to_embed_w_metadata(self):
|
|
164
199
|
documents = [
|
|
165
200
|
Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5)
|
|
@@ -437,10 +472,10 @@ class TestNvidiaDocumentEmbedder:
|
|
|
437
472
|
@pytest.mark.parametrize(
|
|
438
473
|
"model, api_url",
|
|
439
474
|
[
|
|
440
|
-
("nvidia/
|
|
475
|
+
("nvidia/nemotron-3-embed-1b", "https://integrate.api.nvidia.com/v1"),
|
|
441
476
|
],
|
|
442
477
|
ids=[
|
|
443
|
-
"nvidia/
|
|
478
|
+
"nvidia/nemotron-3-embed-1b",
|
|
444
479
|
],
|
|
445
480
|
)
|
|
446
481
|
@pytest.mark.skipif(
|