huggingface-api-haystack 0.1.0__tar.gz → 0.2.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 (30) hide show
  1. huggingface_api_haystack-0.2.0/CHANGELOG.md +9 -0
  2. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/PKG-INFO +1 -1
  3. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/pydoc/config_docusaurus.yml +1 -0
  4. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/pyproject.toml +2 -1
  5. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +11 -0
  6. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +11 -0
  7. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/huggingface_api/chat/chat_generator.py +7 -0
  8. huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +6 -0
  9. huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +297 -0
  10. huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/py.typed +0 -0
  11. huggingface_api_haystack-0.2.0/tests/conftest.py +27 -0
  12. huggingface_api_haystack-0.2.0/tests/test_ranker.py +351 -0
  13. huggingface_api_haystack-0.1.0/tests/conftest.py +0 -12
  14. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/.gitignore +0 -0
  15. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/LICENSE.txt +0 -0
  16. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/README.md +0 -0
  17. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/huggingface_api/__init__.py +0 -0
  18. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/huggingface_api/utils.py +0 -0
  19. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/py.typed +0 -0
  20. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +0 -0
  21. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  22. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
  23. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
  24. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  25. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/__init__.py +0 -0
  26. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_chat_generator.py +0 -0
  27. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_document_embedder.py +0 -0
  28. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_files/apple.jpg +0 -0
  29. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_text_embedder.py +0 -0
  30. {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_utils.py +0 -0
@@ -0,0 +1,9 @@
1
+ # Changelog
2
+
3
+ ## [integrations/huggingface_api-v0.1.0] - 2026-06-05
4
+
5
+ ### 🚀 Features
6
+
7
+ - Move HF API components to core integrations (#3403)
8
+
9
+ <!-- generated by git-cliff -->
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: huggingface-api-haystack
3
- Version: 0.1.0
3
+ Version: 0.2.0
4
4
  Summary: Haystack integration for Hugging Face API
5
5
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/huggingface_api#readme
6
6
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -3,6 +3,7 @@ loaders:
3
3
  - haystack_integrations.components.embedders.huggingface_api.document_embedder
4
4
  - haystack_integrations.components.embedders.huggingface_api.text_embedder
5
5
  - haystack_integrations.components.generators.huggingface_api.chat.chat_generator
6
+ - haystack_integrations.components.rankers.huggingface_api.ranker
6
7
  search_path: [../src]
7
8
  processors:
8
9
  - type: filter
@@ -70,7 +70,8 @@ unit-cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 3
70
70
  integration-cov-append-retry = 'pytest --cov=haystack_integrations --cov-append --reruns 3 --reruns-delay 30 -x -m "integration" {args:tests}'
71
71
  types = """mypy -p haystack_integrations.components.common.huggingface_api \
72
72
  -p haystack_integrations.components.embedders.huggingface_api \
73
- -p haystack_integrations.components.generators.huggingface_api {args}"""
73
+ -p haystack_integrations.components.generators.huggingface_api \
74
+ -p haystack_integrations.components.rankers.huggingface_api {args}"""
74
75
 
75
76
  [tool.mypy]
76
77
  install_types = true
@@ -147,6 +147,9 @@ class HuggingFaceAPIDocumentEmbedder:
147
147
  :param concurrency_limit:
148
148
  The maximum number of requests that should be allowed to run concurrently.
149
149
  This parameter is only used in the `run_async` method.
150
+ :raises ValueError:
151
+ If the required `model` or `url` is missing from `api_params`, the `url` is invalid,
152
+ or the `api_type` is unknown.
150
153
  """
151
154
  if isinstance(api_type, str):
152
155
  api_type = HFEmbeddingAPIType.from_str(api_type)
@@ -331,6 +334,10 @@ class HuggingFaceAPIDocumentEmbedder:
331
334
  :param documents:
332
335
  Documents to embed.
333
336
 
337
+ :raises TypeError:
338
+ If `documents` is not a list of Documents.
339
+ :raises ValueError:
340
+ If the embeddings returned by the API have an unexpected shape.
334
341
  :returns:
335
342
  A dictionary with the following keys:
336
343
  - `documents`: A list of documents with embeddings.
@@ -360,6 +367,10 @@ class HuggingFaceAPIDocumentEmbedder:
360
367
  :param documents:
361
368
  Documents to embed.
362
369
 
370
+ :raises TypeError:
371
+ If `documents` is not a list of Documents.
372
+ :raises ValueError:
373
+ If the embeddings returned by the API have an unexpected shape.
363
374
  :returns:
364
375
  A dictionary with the following keys:
365
376
  - `documents`: A list of documents with embeddings.
@@ -113,6 +113,9 @@ class HuggingFaceAPITextEmbedder:
113
113
  Applicable when `api_type` is `TEXT_EMBEDDINGS_INFERENCE`, or `INFERENCE_ENDPOINTS`
114
114
  if the backend uses Text Embeddings Inference.
115
115
  If `api_type` is `SERVERLESS_INFERENCE_API`, this parameter is ignored.
116
+ :raises ValueError:
117
+ If the required `model` or `url` is missing from `api_params`, the `url` is invalid,
118
+ or the `api_type` is unknown.
116
119
  """
117
120
  if isinstance(api_type, str):
118
121
  api_type = HFEmbeddingAPIType.from_str(api_type)
@@ -213,6 +216,10 @@ class HuggingFaceAPITextEmbedder:
213
216
  :param text:
214
217
  Text to embed.
215
218
 
219
+ :raises TypeError:
220
+ If `text` is not a string.
221
+ :raises ValueError:
222
+ If the embedding returned by the API has an unexpected shape.
216
223
  :returns:
217
224
  A dictionary with the following keys:
218
225
  - `embedding`: The embedding of the input text.
@@ -241,6 +248,10 @@ class HuggingFaceAPITextEmbedder:
241
248
  :param text:
242
249
  Text to embed.
243
250
 
251
+ :raises TypeError:
252
+ If `text` is not a string.
253
+ :raises ValueError:
254
+ If the embedding returned by the API has an unexpected shape.
244
255
  :returns:
245
256
  A dictionary with the following keys:
246
257
  - `embedding`: The embedding of the input text.
@@ -392,6 +392,9 @@ class HuggingFaceAPIChatGenerator:
392
392
  The chosen model should support tool/function calling, according to the model card.
393
393
  Support for tools in the Hugging Face API and TGI is not yet fully refined and you may experience
394
394
  unexpected behavior.
395
+ :raises ValueError:
396
+ If the required `model` or `url` is missing from `api_params`, the `url` is invalid, the `api_type`
397
+ is unknown, `tools` and `streaming_callback` are used together, or duplicate tool names are provided.
395
398
  """
396
399
  if isinstance(api_type, str):
397
400
  api_type = HFGenerationAPIType.from_str(api_type)
@@ -510,6 +513,8 @@ class HuggingFaceAPIChatGenerator:
510
513
  :param streaming_callback:
511
514
  An optional callable for handling streaming responses. If set, it will override the `streaming_callback`
512
515
  parameter set during component initialization.
516
+ :raises ValueError:
517
+ If `tools` and a streaming callback are used together, or if duplicate tool names are provided.
513
518
  :returns: A dictionary with the following keys:
514
519
  - `replies`: A list containing the generated responses as ChatMessage objects.
515
520
  """
@@ -568,6 +573,8 @@ class HuggingFaceAPIChatGenerator:
568
573
  :param streaming_callback:
569
574
  An optional callable for handling streaming responses. If set, it will override the `streaming_callback`
570
575
  parameter set during component initialization.
576
+ :raises ValueError:
577
+ If `tools` and a streaming callback are used together, or if duplicate tool names are provided.
571
578
  :returns: A dictionary with the following keys:
572
579
  - `replies`: A list containing the generated responses as ChatMessage objects.
573
580
  """
@@ -0,0 +1,6 @@
1
+ # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+ from .ranker import HuggingFaceTEIRanker, TruncationDirection
5
+
6
+ __all__ = ["HuggingFaceTEIRanker", "TruncationDirection"]
@@ -0,0 +1,297 @@
1
+ # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from dataclasses import replace
6
+ from enum import Enum
7
+ from typing import Any
8
+ from urllib.parse import urljoin
9
+
10
+ import httpx
11
+ from haystack import Document, component, default_from_dict, default_to_dict
12
+ from haystack.utils import Secret
13
+ from haystack.utils.misc import _deduplicate_documents
14
+ from haystack.utils.requests_utils import async_request_with_retry, request_with_retry
15
+
16
+
17
+ class TruncationDirection(str, Enum):
18
+ """
19
+ Defines the direction to truncate text when input length exceeds the model's limit.
20
+
21
+ Attributes:
22
+ LEFT: Truncate text from the left side (start of text).
23
+ RIGHT: Truncate text from the right side (end of text).
24
+ """
25
+
26
+ LEFT = "Left"
27
+ RIGHT = "Right"
28
+
29
+
30
+ @component
31
+ class HuggingFaceTEIRanker:
32
+ """
33
+ Ranks documents based on their semantic similarity to the query.
34
+
35
+ It can be used with a Text Embeddings Inference (TEI) API endpoint:
36
+ - [Self-hosted Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference)
37
+ - [Hugging Face Inference Endpoints](https://huggingface.co/inference-endpoints)
38
+
39
+ Usage example:
40
+ ```python
41
+ from haystack import Document
42
+ from haystack.utils import Secret
43
+
44
+ from haystack_integrations.components.rankers.huggingface_api import HuggingFaceTEIRanker
45
+
46
+ reranker = HuggingFaceTEIRanker(
47
+ url="http://localhost:8080",
48
+ top_k=5,
49
+ timeout=30,
50
+ token=Secret.from_token("my_api_token")
51
+ )
52
+
53
+ docs = [Document(content="The capital of France is Paris"), Document(content="The capital of Germany is Berlin")]
54
+
55
+ result = reranker.run(query="What is the capital of France?", documents=docs)
56
+
57
+ ranked_docs = result["documents"]
58
+ print(ranked_docs)
59
+ # >> {'documents': [Document(id=..., content: 'the capital of France is Paris', score: 0.9979767),
60
+ # >> Document(id=..., content: 'the capital of Germany is Berlin', score: 0.13982213)]}
61
+ ```
62
+ """
63
+
64
+ def __init__(
65
+ self,
66
+ *,
67
+ url: str,
68
+ top_k: int = 10,
69
+ raw_scores: bool = False,
70
+ timeout: int | None = 30,
71
+ max_retries: int = 3,
72
+ retry_status_codes: list[int] | None = None,
73
+ token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
74
+ ) -> None:
75
+ """
76
+ Initializes the TEI reranker component.
77
+
78
+ :param url: Base URL of the TEI reranking service (for example, "https://api.example.com").
79
+ :param top_k: Maximum number of top documents to return.
80
+ :param raw_scores: If True, include raw relevance scores in the API payload.
81
+ :param timeout: Request timeout in seconds.
82
+ :param max_retries: Maximum number of retry attempts for failed requests.
83
+ :param retry_status_codes: List of HTTP status codes that will trigger a retry.
84
+ When None, HTTP 408, 418, 429 and 503 will be retried (default: None).
85
+ :param token: The Hugging Face token to use as HTTP bearer authorization. Not always required
86
+ depending on your TEI server configuration.
87
+ Check your HF token in your [account settings](https://huggingface.co/settings/tokens).
88
+ """
89
+ self.url = url
90
+ self.top_k = top_k
91
+ self.timeout = timeout
92
+ self.token = token
93
+ self.max_retries = max_retries
94
+ self.retry_status_codes = retry_status_codes
95
+ self.raw_scores = raw_scores
96
+
97
+ def to_dict(self) -> dict[str, Any]:
98
+ """
99
+ Serializes the component to a dictionary.
100
+
101
+ :returns:
102
+ Dictionary with serialized data.
103
+ """
104
+ return default_to_dict(
105
+ self,
106
+ url=self.url,
107
+ top_k=self.top_k,
108
+ timeout=self.timeout,
109
+ token=self.token,
110
+ max_retries=self.max_retries,
111
+ retry_status_codes=self.retry_status_codes,
112
+ )
113
+
114
+ @classmethod
115
+ def from_dict(cls, data: dict[str, Any]) -> "HuggingFaceTEIRanker":
116
+ """
117
+ Deserializes the component from a dictionary.
118
+
119
+ :param data:
120
+ Dictionary to deserialize from.
121
+ :returns:
122
+ Deserialized component.
123
+ """
124
+ return default_from_dict(cls, data)
125
+
126
+ def _compose_response(
127
+ self, result: dict[str, str] | list[dict[str, Any]], top_k: int | None, documents: list[Document]
128
+ ) -> dict[str, list[Document]]:
129
+ """
130
+ Processes the API response into a structured format.
131
+
132
+ :param result: The raw response from the API.
133
+
134
+ :returns: A dictionary with the following keys:
135
+ - `documents`: A list of reranked documents.
136
+
137
+ :raises RuntimeError:
138
+ - If the API request fails.
139
+
140
+ :raises RuntimeError:
141
+ - If the API returns an error response.
142
+
143
+ :raises TypeError:
144
+ - If the API response is not in the expected list format.
145
+ """
146
+ if isinstance(result, dict) and "error" in result:
147
+ error_type = result.get("error_type", "UnknownError")
148
+ error_msg = result.get("error", "No additional information.")
149
+ msg = f"HuggingFaceTEIRanker API call failed ({error_type}): {error_msg}"
150
+ raise RuntimeError(msg)
151
+
152
+ # Ensure we have a list of score dicts
153
+ if not isinstance(result, list):
154
+ # Expected list or dict, but encountered an unknown response format.
155
+ error_msg = f"Expected a list of score dictionaries, but got `{type(result).__name__}`. "
156
+ error_msg += f"Response content: {result}"
157
+ msg = f"Unexpected response format from text-embeddings-inference rerank API: {error_msg}"
158
+ raise TypeError(msg)
159
+
160
+ # Determine number of docs to return
161
+ final_k = min(top_k or self.top_k, len(result))
162
+
163
+ # Select and return the top_k documents
164
+ ranked_docs = []
165
+ for item in result[:final_k]:
166
+ index: int = item["index"]
167
+ ranked_docs.append(replace(documents[index], score=item["score"]))
168
+ return {"documents": ranked_docs}
169
+
170
+ @component.output_types(documents=list[Document])
171
+ def run(
172
+ self,
173
+ query: str,
174
+ documents: list[Document],
175
+ top_k: int | None = None,
176
+ truncation_direction: TruncationDirection | None = None,
177
+ ) -> dict[str, list[Document]]:
178
+ """
179
+ Reranks the provided documents by relevance to the query using the TEI API.
180
+
181
+ Before ranking, documents are deduplicated by their id, retaining only the document with the highest score
182
+ if a score is present.
183
+
184
+ :param query: The user query string to guide reranking.
185
+ :param documents: List of `Document` objects to rerank.
186
+ :param top_k: Optional override for the maximum number of documents to return.
187
+ :param truncation_direction: If set, enables text truncation in the specified direction.
188
+
189
+ :returns: A dictionary with the following keys:
190
+ - `documents`: A list of reranked documents.
191
+
192
+ :raises RuntimeError:
193
+ - If the API request fails.
194
+
195
+ :raises RuntimeError:
196
+ - If the API returns an error response.
197
+
198
+ :raises TypeError:
199
+ - If the API response is not in the expected list format.
200
+ """
201
+ # Return empty if no documents provided
202
+ if not documents:
203
+ return {"documents": []}
204
+
205
+ # Prepare the payload
206
+ deduplicated_documents = _deduplicate_documents(documents)
207
+ texts = [doc.content for doc in deduplicated_documents]
208
+ payload: dict[str, Any] = {"query": query, "texts": texts, "raw_scores": self.raw_scores}
209
+ if truncation_direction:
210
+ payload.update({"truncate": True, "truncation_direction": truncation_direction.value})
211
+
212
+ headers = {}
213
+ if self.token and self.token.resolve_value():
214
+ headers["Authorization"] = f"Bearer {self.token.resolve_value()}"
215
+
216
+ # Call the external service with retry
217
+ try:
218
+ response = request_with_retry(
219
+ method="POST",
220
+ url=urljoin(self.url, "/rerank"),
221
+ json=payload,
222
+ timeout=self.timeout,
223
+ headers=headers,
224
+ attempts=self.max_retries,
225
+ status_codes_to_retry=self.retry_status_codes,
226
+ )
227
+ except httpx.HTTPStatusError as e:
228
+ msg = f"HuggingFaceTEIRanker API call failed. Error: {e}, Response: {e.response.text}"
229
+ raise RuntimeError(msg) from e
230
+
231
+ result: dict[str, str] | list[dict[str, Any]] = response.json()
232
+
233
+ return self._compose_response(result, top_k, deduplicated_documents)
234
+
235
+ @component.output_types(documents=list[Document])
236
+ async def run_async(
237
+ self,
238
+ query: str,
239
+ documents: list[Document],
240
+ top_k: int | None = None,
241
+ truncation_direction: TruncationDirection | None = None,
242
+ ) -> dict[str, list[Document]]:
243
+ """
244
+ Asynchronously reranks the provided documents by relevance to the query using the TEI API.
245
+
246
+ Before ranking, documents are deduplicated by their id, retaining only the document with the highest score
247
+ if a score is present.
248
+
249
+ :param query: The user query string to guide reranking.
250
+ :param documents: List of `Document` objects to rerank.
251
+ :param top_k: Optional override for the maximum number of documents to return.
252
+ :param truncation_direction: If set, enables text truncation in the specified direction.
253
+
254
+ :returns: A dictionary with the following keys:
255
+ - `documents`: A list of reranked documents.
256
+
257
+ :raises httpx.RequestError:
258
+ - If the API request fails.
259
+ :raises RuntimeError:
260
+ - If the API returns an error response.
261
+
262
+ :raises TypeError:
263
+ - If the API response is not in the expected list format.
264
+ """
265
+ # Return empty if no documents provided
266
+ if not documents:
267
+ return {"documents": []}
268
+
269
+ # Prepare the payload
270
+ deduplicated_documents = _deduplicate_documents(documents)
271
+ texts = [doc.content for doc in deduplicated_documents]
272
+ payload: dict[str, Any] = {"query": query, "texts": texts, "raw_scores": self.raw_scores}
273
+ if truncation_direction:
274
+ payload.update({"truncate": True, "truncation_direction": truncation_direction.value})
275
+
276
+ headers = {}
277
+ if self.token and self.token.resolve_value():
278
+ headers["Authorization"] = f"Bearer {self.token.resolve_value()}"
279
+
280
+ # Call the external service with retry
281
+ try:
282
+ response = await async_request_with_retry(
283
+ method="POST",
284
+ url=urljoin(self.url, "/rerank"),
285
+ json=payload,
286
+ timeout=self.timeout,
287
+ headers=headers,
288
+ attempts=self.max_retries,
289
+ status_codes_to_retry=self.retry_status_codes,
290
+ )
291
+ except httpx.HTTPStatusError as e:
292
+ msg = f"HuggingFaceTEIRanker API call failed. Error: {e}, Response: {e.response.text}"
293
+ raise RuntimeError(msg) from e
294
+
295
+ result: dict[str, str] | list[dict[str, Any]] = response.json()
296
+
297
+ return self._compose_response(result, top_k, deduplicated_documents)
@@ -0,0 +1,27 @@
1
+ # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from pathlib import Path
6
+
7
+ import pytest
8
+
9
+
10
+ @pytest.fixture()
11
+ def test_files_path():
12
+ return Path(__file__).parent / "test_files"
13
+
14
+
15
+ @pytest.fixture()
16
+ def del_hf_env_vars(monkeypatch):
17
+ """
18
+ Delete Hugging Face environment variables for tests.
19
+
20
+ Prevents passing empty tokens to Hugging Face, which would cause API calls to fail.
21
+ This is particularly relevant for PRs opened from forks, where secrets are not available
22
+ and empty environment variables might be set instead of being removed.
23
+
24
+ See https://github.com/deepset-ai/haystack/issues/8811 for more details.
25
+ """
26
+ monkeypatch.delenv("HF_API_TOKEN", raising=False)
27
+ monkeypatch.delenv("HF_TOKEN", raising=False)
@@ -0,0 +1,351 @@
1
+ # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from unittest.mock import MagicMock, patch
6
+
7
+ import httpx
8
+ import pytest
9
+ from haystack import Document
10
+ from haystack.utils import Secret
11
+
12
+ from haystack_integrations.components.rankers.huggingface_api import HuggingFaceTEIRanker, TruncationDirection
13
+
14
+
15
+ class TestHuggingFaceTEIRanker:
16
+ def test_init(self, del_hf_env_vars):
17
+ """Test initialization with default and custom parameters"""
18
+ # Default parameters
19
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
20
+ assert ranker.url == "https://api.my-tei-service.com"
21
+ assert ranker.top_k == 10
22
+ assert ranker.timeout == 30
23
+ assert not ranker.token.resolve_value()
24
+ assert ranker.max_retries == 3
25
+ assert ranker.retry_status_codes is None
26
+
27
+ # Custom parameters
28
+ token = Secret.from_token("my_api_token")
29
+ ranker = HuggingFaceTEIRanker(
30
+ url="https://api.my-tei-service.com",
31
+ top_k=5,
32
+ timeout=60,
33
+ token=token,
34
+ max_retries=5,
35
+ retry_status_codes=[500, 502, 503],
36
+ )
37
+ assert ranker.url == "https://api.my-tei-service.com"
38
+ assert ranker.top_k == 5
39
+ assert ranker.timeout == 60
40
+ assert ranker.token == token
41
+ assert ranker.max_retries == 5
42
+ assert ranker.retry_status_codes == [500, 502, 503]
43
+
44
+ def test_to_dict(self, del_hf_env_vars):
45
+ """Test serialization to dict with Secret token"""
46
+ component = HuggingFaceTEIRanker(
47
+ url="https://api.my-tei-service.com", top_k=5, timeout=30, max_retries=4, retry_status_codes=[500, 502]
48
+ )
49
+ data = component.to_dict()
50
+
51
+ assert data["type"] == "haystack_integrations.components.rankers.huggingface_api.ranker.HuggingFaceTEIRanker"
52
+ assert data["init_parameters"]["url"] == "https://api.my-tei-service.com"
53
+ assert data["init_parameters"]["top_k"] == 5
54
+ assert data["init_parameters"]["timeout"] == 30
55
+ assert data["init_parameters"]["token"] == {
56
+ "env_vars": ["HF_API_TOKEN", "HF_TOKEN"],
57
+ "strict": False,
58
+ "type": "env_var",
59
+ }
60
+ assert data["init_parameters"]["max_retries"] == 4
61
+ assert data["init_parameters"]["retry_status_codes"] == [500, 502]
62
+
63
+ def test_from_dict(self, del_hf_env_vars):
64
+ """Test deserialization from dict with environment variable token"""
65
+ data = {
66
+ "type": "haystack_integrations.components.rankers.huggingface_api.ranker.HuggingFaceTEIRanker",
67
+ "init_parameters": {
68
+ "url": "https://api.my-tei-service.com",
69
+ "top_k": 5,
70
+ "timeout": 30,
71
+ "token": {"type": "env_var", "env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False},
72
+ "max_retries": 4,
73
+ "retry_status_codes": [500, 502],
74
+ },
75
+ }
76
+
77
+ component = HuggingFaceTEIRanker.from_dict(data)
78
+
79
+ assert component.url == "https://api.my-tei-service.com"
80
+ assert component.top_k == 5
81
+ assert component.timeout == 30
82
+ assert component.max_retries == 4
83
+ assert component.retry_status_codes == [500, 502]
84
+
85
+ def test_empty_documents(self, del_hf_env_vars):
86
+ """Test that empty documents list returns empty result"""
87
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
88
+ result = ranker.run(query="test query", documents=[])
89
+ assert result == {"documents": []}
90
+
91
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.request_with_retry")
92
+ def test_run_with_mock(self, mock_request, del_hf_env_vars):
93
+ """Test run method with mocked API response"""
94
+ # Setup mock response
95
+ mock_response = MagicMock(spec=httpx.Response)
96
+ mock_response.json.return_value = [
97
+ {"index": 2, "score": 0.95},
98
+ {"index": 1, "score": 0.85},
99
+ {"index": 0, "score": 0.75},
100
+ ]
101
+ mock_request.return_value = mock_response
102
+
103
+ # Create ranker and test documents
104
+ token = Secret.from_token("test_token")
105
+ ranker = HuggingFaceTEIRanker(
106
+ url="https://api.my-tei-service.com",
107
+ top_k=3,
108
+ timeout=30,
109
+ token=token,
110
+ max_retries=4,
111
+ retry_status_codes=[500, 502],
112
+ )
113
+
114
+ docs = [Document(content="Document A"), Document(content="Document B"), Document(content="Document C")]
115
+
116
+ # Run the ranker
117
+ result = ranker.run(query="test query", documents=docs)
118
+
119
+ # Check that request_with_retry was called with correct parameters
120
+ mock_request.assert_called_once_with(
121
+ method="POST",
122
+ url="https://api.my-tei-service.com/rerank",
123
+ json={"query": "test query", "texts": ["Document A", "Document B", "Document C"], "raw_scores": False},
124
+ timeout=30,
125
+ headers={"Authorization": "Bearer test_token"},
126
+ attempts=4,
127
+ status_codes_to_retry=[500, 502],
128
+ )
129
+
130
+ # Check that documents are ranked correctly
131
+ assert len(result["documents"]) == 3
132
+ assert result["documents"][0].content == "Document C"
133
+ assert result["documents"][0].score == 0.95
134
+ assert result["documents"][1].content == "Document B"
135
+ assert result["documents"][1].score == 0.85
136
+ assert result["documents"][2].content == "Document A"
137
+ assert result["documents"][2].score == 0.75
138
+
139
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.request_with_retry")
140
+ def test_run_with_truncation_direction(self, mock_request, del_hf_env_vars):
141
+ """Test run method with truncation direction parameter"""
142
+ # Setup mock response
143
+ mock_response = MagicMock(spec=httpx.Response)
144
+ mock_response.json.return_value = [{"index": 0, "score": 0.95}]
145
+ mock_request.return_value = mock_response
146
+
147
+ # Create ranker and test documents
148
+ token = Secret.from_token("test_token")
149
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com", token=token)
150
+ docs = [Document(content="Document A")]
151
+
152
+ # Run the ranker with truncation direction
153
+ ranker.run(query="test query", documents=docs, truncation_direction=TruncationDirection.LEFT)
154
+
155
+ # Check that request includes truncation parameters
156
+ mock_request.assert_called_once_with(
157
+ method="POST",
158
+ url="https://api.my-tei-service.com/rerank",
159
+ json={
160
+ "query": "test query",
161
+ "texts": ["Document A"],
162
+ "raw_scores": False,
163
+ "truncate": True,
164
+ "truncation_direction": "Left",
165
+ },
166
+ timeout=30,
167
+ headers={"Authorization": "Bearer test_token"},
168
+ attempts=3,
169
+ status_codes_to_retry=None,
170
+ )
171
+
172
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.request_with_retry")
173
+ def test_run_with_custom_top_k(self, mock_request, del_hf_env_vars):
174
+ """Test run method with custom top_k parameter"""
175
+ # Setup mock response with 5 documents
176
+ mock_response = MagicMock(spec=httpx.Response)
177
+ mock_response.json.return_value = [
178
+ {"index": 4, "score": 0.95},
179
+ {"index": 3, "score": 0.90},
180
+ {"index": 2, "score": 0.85},
181
+ {"index": 1, "score": 0.80},
182
+ {"index": 0, "score": 0.75},
183
+ ]
184
+ mock_request.return_value = mock_response
185
+
186
+ # Create ranker with top_k=3
187
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com", top_k=3)
188
+
189
+ # Create 5 test documents
190
+ docs = [Document(content=f"Document {i}") for i in range(5)]
191
+
192
+ # Run the ranker
193
+ result = ranker.run(query="test query", documents=docs)
194
+
195
+ # Check that only top 3 documents are returned
196
+ assert len(result["documents"]) == 3
197
+ assert result["documents"][0].content == "Document 4"
198
+ assert result["documents"][1].content == "Document 3"
199
+ assert result["documents"][2].content == "Document 2"
200
+
201
+ # Test with run-time top_k override
202
+ result = ranker.run(query="test query", documents=docs, top_k=2)
203
+
204
+ # Check that only top 2 documents are returned
205
+ assert len(result["documents"]) == 2
206
+ assert result["documents"][0].content == "Document 4"
207
+ assert result["documents"][1].content == "Document 3"
208
+
209
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.request_with_retry")
210
+ def test_run_deduplicates_documents(self, mock_request, del_hf_env_vars):
211
+ """Test that duplicate documents are removed before sending to the API."""
212
+ mock_response = MagicMock(spec=httpx.Response)
213
+ mock_response.json.return_value = [{"index": 1, "score": 0.9}, {"index": 0, "score": 0.2}]
214
+ mock_request.return_value = mock_response
215
+
216
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
217
+ # Document with duplicate id and lower score should be dropped
218
+ docs = [
219
+ Document(id="duplicate", content="keep me", score=0.9),
220
+ Document(id="duplicate", content="drop me", score=0.1),
221
+ Document(id="unique", content="unique"),
222
+ ]
223
+
224
+ result = ranker.run(query="test query", documents=docs)
225
+
226
+ mock_request.assert_called_once_with(
227
+ method="POST",
228
+ url="https://api.my-tei-service.com/rerank",
229
+ json={"query": "test query", "texts": ["keep me", "unique"], "raw_scores": False},
230
+ timeout=30,
231
+ headers={},
232
+ attempts=3,
233
+ status_codes_to_retry=None,
234
+ )
235
+ assert len(result["documents"]) == 2
236
+ assert result["documents"][0].content == "unique"
237
+ assert result["documents"][1].content == "keep me"
238
+
239
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.request_with_retry")
240
+ def test_error_handling(self, mock_request, del_hf_env_vars):
241
+ """Test error handling in the ranker"""
242
+ # Setup mock response with error
243
+ mock_response = MagicMock(spec=httpx.Response)
244
+ mock_response.json.return_value = {"error": "Some error occurred", "error_type": "TestError"}
245
+ mock_request.return_value = mock_response
246
+
247
+ # Create ranker and test documents
248
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
249
+ docs = [Document(content="Document A")]
250
+
251
+ # Test that RuntimeError is raised with the correct message
252
+ with pytest.raises(
253
+ RuntimeError, match=r"HuggingFaceTEIRanker API call failed \(TestError\): Some error occurred"
254
+ ):
255
+ ranker.run(query="test query", documents=docs)
256
+
257
+ # Test unexpected response format
258
+ mock_response.json.return_value = {"unexpected": "format"}
259
+ with pytest.raises(TypeError, match="Unexpected response format from text-embeddings-inference rerank API"):
260
+ ranker.run(query="test query", documents=docs)
261
+
262
+ @pytest.mark.asyncio
263
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.async_request_with_retry")
264
+ async def test_run_async_with_mock(self, mock_request, del_hf_env_vars):
265
+ """Test run_async method with mocked API response"""
266
+ # Setup mock response
267
+ mock_response = MagicMock(spec=httpx.Response)
268
+ mock_response.json.return_value = [
269
+ {"index": 2, "score": 0.95},
270
+ {"index": 1, "score": 0.85},
271
+ {"index": 0, "score": 0.75},
272
+ ]
273
+ mock_request.return_value = mock_response
274
+
275
+ # Create ranker and test documents
276
+ token = Secret.from_token("test_token")
277
+ ranker = HuggingFaceTEIRanker(
278
+ url="https://api.my-tei-service.com",
279
+ top_k=3,
280
+ timeout=30,
281
+ token=token,
282
+ max_retries=4,
283
+ retry_status_codes=[500, 502],
284
+ )
285
+
286
+ docs = [Document(content="Document A"), Document(content="Document B"), Document(content="Document C")]
287
+
288
+ # Run the ranker asynchronously
289
+ result = await ranker.run_async(query="test query", documents=docs)
290
+
291
+ # Check that async_request_with_retry was called with correct parameters
292
+ mock_request.assert_called_once_with(
293
+ method="POST",
294
+ url="https://api.my-tei-service.com/rerank",
295
+ json={"query": "test query", "texts": ["Document A", "Document B", "Document C"], "raw_scores": False},
296
+ timeout=30,
297
+ headers={"Authorization": "Bearer test_token"},
298
+ attempts=4,
299
+ status_codes_to_retry=[500, 502],
300
+ )
301
+
302
+ # Check that documents are ranked correctly
303
+ assert len(result["documents"]) == 3
304
+ assert result["documents"][0].content == "Document C"
305
+ assert result["documents"][0].score == 0.95
306
+ assert result["documents"][1].content == "Document B"
307
+ assert result["documents"][1].score == 0.85
308
+ assert result["documents"][2].content == "Document A"
309
+ assert result["documents"][2].score == 0.75
310
+
311
+ @pytest.mark.asyncio
312
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.async_request_with_retry")
313
+ async def test_run_async_deduplicates_documents(self, mock_request, del_hf_env_vars):
314
+ """Test that duplicate documents are removed before sending to the API."""
315
+ mock_response = MagicMock(spec=httpx.Response)
316
+ mock_response.json.return_value = [{"index": 1, "score": 0.9}, {"index": 0, "score": 0.2}]
317
+ mock_request.return_value = mock_response
318
+
319
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
320
+ # Document with duplicate id and lower score should be dropped
321
+ docs = [
322
+ Document(id="duplicate", content="keep me", score=0.9),
323
+ Document(id="duplicate", content="drop me", score=0.1),
324
+ Document(id="unique", content="unique"),
325
+ ]
326
+
327
+ result = await ranker.run_async(query="test query", documents=docs)
328
+
329
+ mock_request.assert_called_once_with(
330
+ method="POST",
331
+ url="https://api.my-tei-service.com/rerank",
332
+ json={"query": "test query", "texts": ["keep me", "unique"], "raw_scores": False},
333
+ timeout=30,
334
+ headers={},
335
+ attempts=3,
336
+ status_codes_to_retry=None,
337
+ )
338
+ assert len(result["documents"]) == 2
339
+ assert result["documents"][0].content == "unique"
340
+ assert result["documents"][1].content == "keep me"
341
+
342
+ @pytest.mark.asyncio
343
+ @patch("haystack_integrations.components.rankers.huggingface_api.ranker.async_request_with_retry")
344
+ async def test_run_async_empty_documents(self, mock_request, del_hf_env_vars):
345
+ """Test run_async with empty documents list"""
346
+ ranker = HuggingFaceTEIRanker(url="https://api.my-tei-service.com")
347
+ result = await ranker.run_async(query="test query", documents=[])
348
+
349
+ # Check that no API call was made
350
+ mock_request.assert_not_called()
351
+ assert result == {"documents": []}
@@ -1,12 +0,0 @@
1
- # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
- #
3
- # SPDX-License-Identifier: Apache-2.0
4
-
5
- from pathlib import Path
6
-
7
- import pytest
8
-
9
-
10
- @pytest.fixture()
11
- def test_files_path():
12
- return Path(__file__).parent / "test_files"