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.
- huggingface_api_haystack-0.2.0/CHANGELOG.md +9 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/PKG-INFO +1 -1
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/pydoc/config_docusaurus.yml +1 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/pyproject.toml +2 -1
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +11 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +11 -0
- {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
- huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +6 -0
- huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +297 -0
- huggingface_api_haystack-0.2.0/src/haystack_integrations/components/rankers/py.typed +0 -0
- huggingface_api_haystack-0.2.0/tests/conftest.py +27 -0
- huggingface_api_haystack-0.2.0/tests/test_ranker.py +351 -0
- huggingface_api_haystack-0.1.0/tests/conftest.py +0 -12
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/.gitignore +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/LICENSE.txt +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/README.md +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/huggingface_api/utils.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/common/py.typed +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/__init__.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_chat_generator.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_document_embedder.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_files/apple.jpg +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_text_embedder.py +0 -0
- {huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_utils.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: huggingface-api-haystack
|
|
3
|
-
Version: 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
|
{huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/pydoc/config_docusaurus.yml
RENAMED
|
@@ -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
|
|
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,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)
|
|
File without changes
|
|
@@ -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"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_chat_generator.py
RENAMED
|
File without changes
|
{huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_document_embedder.py
RENAMED
|
File without changes
|
{huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_files/apple.jpg
RENAMED
|
File without changes
|
{huggingface_api_haystack-0.1.0 → huggingface_api_haystack-0.2.0}/tests/test_text_embedder.py
RENAMED
|
File without changes
|
|
File without changes
|