huggingface-api-haystack 0.6.0__tar.gz → 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/CHANGELOG.md +20 -1
  2. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/PKG-INFO +3 -3
  3. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/pyproject.toml +1 -1
  4. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/common/huggingface_api/utils.py +41 -8
  5. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +54 -8
  6. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +40 -3
  7. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/chat_generator.py +55 -22
  8. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_chat_generator.py +197 -78
  9. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_document_embedder.py +161 -30
  10. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_text_embedder.py +153 -28
  11. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_utils.py +35 -1
  12. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/.gitignore +0 -0
  13. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/LICENSE.txt +0 -0
  14. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/README.md +0 -0
  15. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/docker-compose.yml +0 -0
  16. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/pydoc/config_docusaurus.yml +0 -0
  17. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/common/huggingface_api/__init__.py +0 -0
  18. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/common/py.typed +0 -0
  19. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +0 -0
  20. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_document_embedder.py +0 -0
  21. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_embedding_utils.py +0 -0
  22. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +0 -0
  23. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  24. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
  25. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
  26. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  27. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +0 -0
  28. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +0 -0
  29. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  30. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/__init__.py +0 -0
  31. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/conftest.py +0 -0
  32. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_files/apple.jpg +0 -0
  33. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_ranker.py +0 -0
  34. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_sparse_document_embedder.py +0 -0
  35. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-1.0.0}/tests/test_sparse_text_embedder.py +0 -0
@@ -1,6 +1,25 @@
1
1
  # Changelog
2
2
 
3
- ## [unreleased]
3
+ ## [integrations/huggingface_api-v0.6.0] - 2026-08-27
4
+
5
+ ### 🚀 Features
6
+
7
+ - *(huggingface_api)* Add SPLADE sparse embedder components (#3835)
8
+
9
+ ## [integrations/huggingface_api-v0.5.0] - 2026-08-24
10
+
11
+ ### 🐛 Bug Fixes
12
+
13
+ - Fix new issues raised by ruff 0.16.0 (#3670)
14
+ - Serialize missing init params in to_dict (transformers, huggingface_api, ragas) (#3808)
15
+
16
+ ### 🧹 Chores
17
+
18
+ - Re-standardize READMEs (#3768)
19
+ - Clarify how generation_kwargs passed in run are handled (#3805)
20
+
21
+
22
+ ## [integrations/huggingface_api-v0.4.0] - 2026-07-09
4
23
 
5
24
  ### 🧪 Testing
6
25
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: huggingface-api-haystack
3
- Version: 0.6.0
3
+ Version: 1.0.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
@@ -19,9 +19,9 @@ Classifier: Programming Language :: Python :: 3.14
19
19
  Classifier: Programming Language :: Python :: Implementation :: CPython
20
20
  Classifier: Programming Language :: Python :: Implementation :: PyPy
21
21
  Requires-Python: >=3.10
22
- Requires-Dist: haystack-ai>=2.30.0
22
+ Requires-Dist: haystack-ai>=3.0.0
23
23
  Requires-Dist: httpx>=0.27.0
24
- Requires-Dist: huggingface-hub[inference]>=0.35.0
24
+ Requires-Dist: huggingface-hub[inference]>=1.0.0
25
25
  Description-Content-Type: text/markdown
26
26
 
27
27
  # huggingface-api-haystack
@@ -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>=2.30.0", "httpx>=0.27.0", "huggingface_hub[inference]>=0.35.0"]
26
+ dependencies = ["haystack-ai>=3.0.0", "httpx>=0.27.0", "huggingface_hub[inference]>=1.0.0"]
27
27
 
28
28
  [project.urls]
29
29
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/huggingface_api#readme"
@@ -7,6 +7,7 @@ from enum import Enum
7
7
  from haystack.utils import Secret
8
8
  from huggingface_hub import HfApi
9
9
  from huggingface_hub.errors import RepositoryNotFoundError
10
+ from huggingface_hub.utils import build_hf_headers, get_async_session, hf_raise_for_status
10
11
 
11
12
 
12
13
  class HFGenerationAPIType(Enum):
@@ -82,14 +83,15 @@ class HFModelType(Enum):
82
83
 
83
84
  def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | None) -> None:
84
85
  """
85
- Check if the provided model ID corresponds to a valid model on HuggingFace Hub.
86
+ Check if the provided model ID corresponds to a valid model on Hugging Face Hub.
86
87
 
87
- Also check if the model is an embedding or generation model.
88
+ Also check if the model supports the requested embedding or generation task.
88
89
 
89
- :param model_id: A string representing the HuggingFace model ID.
90
- :param model_type: the model type, HFModelType.EMBEDDING or HFModelType.GENERATION
91
- :param token: The optional authentication token.
92
- :raises ValueError: If the model is not found or is not a embedding model.
90
+ :param model_id: A string representing the Hugging Face model ID.
91
+ :param model_type: The model type, `HFModelType.EMBEDDING` or `HFModelType.GENERATION`.
92
+ :param token: The optional authentication secret.
93
+ :raises ValueError: If the secret cannot be resolved, the model is not found, or it does not match the
94
+ requested type.
93
95
  """
94
96
  api = HfApi()
95
97
  try:
@@ -98,11 +100,42 @@ def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | N
98
100
  msg = f"Model {model_id} not found on HuggingFace Hub. Please provide a valid HuggingFace model_id."
99
101
  raise ValueError(msg) from e
100
102
 
103
+ _validate_model_type(model_id, model_type, model_info.pipeline_tag)
104
+
105
+
106
+ async def _check_valid_model_async(model_id: str, model_type: HFModelType, token: Secret | None) -> None:
107
+ """
108
+ Asynchronously check if the provided model ID corresponds to a valid model on Hugging Face Hub.
109
+
110
+ Also check if the model supports the requested embedding or generation task.
111
+
112
+ :param model_id: A string representing the Hugging Face model ID.
113
+ :param model_type: The model type, `HFModelType.EMBEDDING` or `HFModelType.GENERATION`.
114
+ :param token: The optional authentication secret.
115
+ :raises ValueError: If the secret cannot be resolved, the model is not found, or it does not match the
116
+ requested type.
117
+ """
118
+ api = HfApi()
119
+ try:
120
+ async with get_async_session() as client:
121
+ response = await client.get(
122
+ f"{api.endpoint}/api/models/{model_id}",
123
+ headers=build_hf_headers(token=token.resolve_value() if token else None),
124
+ )
125
+ hf_raise_for_status(response)
126
+ except RepositoryNotFoundError as e:
127
+ msg = f"Model {model_id} not found on HuggingFace Hub. Please provide a valid HuggingFace model_id."
128
+ raise ValueError(msg) from e
129
+
130
+ _validate_model_type(model_id, model_type, response.json().get("pipeline_tag"))
131
+
132
+
133
+ def _validate_model_type(model_id: str, model_type: HFModelType, pipeline_tag: str | None) -> None:
101
134
  if model_type == HFModelType.EMBEDDING:
102
- allowed_model = model_info.pipeline_tag in ["sentence-similarity", "feature-extraction"]
135
+ allowed_model = pipeline_tag in ["sentence-similarity", "feature-extraction"]
103
136
  error_msg = f"Model {model_id} is not a embedding model. Please provide a embedding model."
104
137
  elif model_type == HFModelType.GENERATION:
105
- allowed_model = model_info.pipeline_tag in ["text-generation", "text2text-generation", "image-text-to-text"]
138
+ allowed_model = pipeline_tag in ["text-generation", "text2text-generation", "image-text-to-text"]
106
139
  error_msg = f"Model {model_id} is not a text generation model. Please provide a text generation model."
107
140
  else:
108
141
  allowed_model = False
@@ -18,6 +18,7 @@ from haystack_integrations.common.huggingface_api.utils import (
18
18
  HFEmbeddingAPIType,
19
19
  HFModelType,
20
20
  _check_valid_model,
21
+ _check_valid_model_async,
21
22
  )
22
23
 
23
24
  logger = logging.getLogger(__name__)
@@ -161,7 +162,6 @@ class HuggingFaceAPIDocumentEmbedder:
161
162
  if model is None:
162
163
  msg = "To use the Serverless Inference API, you need to specify the `model` parameter in `api_params`."
163
164
  raise ValueError(msg)
164
- _check_valid_model(model, HFModelType.EMBEDDING, token)
165
165
  model_or_url = model
166
166
  elif api_type in [HFEmbeddingAPIType.INFERENCE_ENDPOINTS, HFEmbeddingAPIType.TEXT_EMBEDDINGS_INFERENCE]:
167
167
  url = api_params.get("url")
@@ -179,8 +179,6 @@ class HuggingFaceAPIDocumentEmbedder:
179
179
  msg = f"Unknown api_type {api_type}"
180
180
  raise ValueError(msg)
181
181
 
182
- client_args: dict[str, Any] = {"model": model_or_url, "token": token.resolve_value() if token else None}
183
-
184
182
  self.api_type = api_type
185
183
  self.api_params = api_params
186
184
  self.token = token
@@ -193,8 +191,39 @@ class HuggingFaceAPIDocumentEmbedder:
193
191
  self.meta_fields_to_embed = meta_fields_to_embed or []
194
192
  self.embedding_separator = embedding_separator
195
193
  self.concurrency_limit = concurrency_limit
196
- self._client = InferenceClient(**client_args)
197
- self._async_client = AsyncInferenceClient(**client_args)
194
+ self._model_or_url = model_or_url
195
+ self._client: InferenceClient | None = None
196
+ self._async_client: AsyncInferenceClient | None = None
197
+
198
+ def _client_kwargs(self) -> dict[str, Any]:
199
+ """Build the keyword arguments used to create Hugging Face clients."""
200
+ return {"model": self._model_or_url, "token": self.token.resolve_value() if self.token else None}
201
+
202
+ def warm_up(self) -> None:
203
+ """Create the synchronous Hugging Face client."""
204
+ if self._client is None:
205
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
206
+ _check_valid_model(self._model_or_url, HFModelType.EMBEDDING, self.token)
207
+ self._client = InferenceClient(**self._client_kwargs())
208
+
209
+ async def warm_up_async(self) -> None:
210
+ """Create the asynchronous Hugging Face client."""
211
+ if self._async_client is None:
212
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
213
+ await _check_valid_model_async(self._model_or_url, HFModelType.EMBEDDING, self.token)
214
+ self._async_client = AsyncInferenceClient(**self._client_kwargs())
215
+
216
+ def close(self) -> None:
217
+ """Close the synchronous Hugging Face client."""
218
+ if self._client is not None:
219
+ self._client.close()
220
+ self._client = None
221
+
222
+ async def close_async(self) -> None:
223
+ """Close the asynchronous Hugging Face client."""
224
+ if self._async_client is not None:
225
+ await self._async_client.close()
226
+ self._async_client = None
198
227
 
199
228
  def to_dict(self) -> dict[str, Any]:
200
229
  """
@@ -270,6 +299,7 @@ class HuggingFaceAPIDocumentEmbedder:
270
299
  """
271
300
  Embed a list of texts in batches.
272
301
  """
302
+ assert self._client is not None # noqa: S101
273
303
  truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
274
304
 
275
305
  all_embeddings: list = []
@@ -278,7 +308,12 @@ class HuggingFaceAPIDocumentEmbedder:
278
308
  ):
279
309
  batch = texts_to_embed[i : i + batch_size]
280
310
 
281
- np_embeddings = self._client.feature_extraction(text=batch, truncate=truncate, normalize=normalize)
311
+ # a batch is accepted but text is typed as str
312
+ np_embeddings = self._client.feature_extraction(
313
+ text=batch, # type: ignore[arg-type]
314
+ truncate=truncate,
315
+ normalize=normalize,
316
+ )
282
317
 
283
318
  if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
284
319
  msg = f"Expected embedding shape ({batch_size}, embedding_dim), got {np_embeddings.shape}"
@@ -292,6 +327,8 @@ class HuggingFaceAPIDocumentEmbedder:
292
327
  """
293
328
  Embed a list of texts in batches asynchronously.
294
329
  """
330
+ assert self._async_client is not None # noqa: S101
331
+ async_client = self._async_client
295
332
  truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
296
333
  sem = Semaphore(max(1, self.concurrency_limit))
297
334
  num_batches = (len(texts_to_embed) + batch_size - 1) // batch_size
@@ -299,8 +336,11 @@ class HuggingFaceAPIDocumentEmbedder:
299
336
 
300
337
  async def _runner(batch: list[str]) -> list[list[float]]:
301
338
  async with sem:
302
- np_embeddings = await self._async_client.feature_extraction(
303
- text=batch, truncate=truncate, normalize=normalize
339
+ # a batch is accepted but text is typed as str
340
+ np_embeddings = await async_client.feature_extraction(
341
+ text=batch, # type: ignore[arg-type]
342
+ truncate=truncate,
343
+ normalize=normalize,
304
344
  )
305
345
 
306
346
  if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
@@ -342,6 +382,9 @@ class HuggingFaceAPIDocumentEmbedder:
342
382
  A dictionary with the following keys:
343
383
  - `documents`: A list of documents with embeddings.
344
384
  """
385
+ self.warm_up()
386
+ assert self._client is not None # noqa: S101
387
+
345
388
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
346
389
  msg = (
347
390
  "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
@@ -375,6 +418,9 @@ class HuggingFaceAPIDocumentEmbedder:
375
418
  A dictionary with the following keys:
376
419
  - `documents`: A list of documents with embeddings.
377
420
  """
421
+ await self.warm_up_async()
422
+ assert self._async_client is not None # noqa: S101
423
+
378
424
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
379
425
  msg = (
380
426
  "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
@@ -13,6 +13,7 @@ from haystack_integrations.common.huggingface_api.utils import (
13
13
  HFEmbeddingAPIType,
14
14
  HFModelType,
15
15
  _check_valid_model,
16
+ _check_valid_model_async,
16
17
  )
17
18
 
18
19
  logger = logging.getLogger(__name__)
@@ -125,7 +126,6 @@ class HuggingFaceAPITextEmbedder:
125
126
  if model is None:
126
127
  msg = "To use the Serverless Inference API, you need to specify the `model` parameter in `api_params`."
127
128
  raise ValueError(msg)
128
- _check_valid_model(model, HFModelType.EMBEDDING, token)
129
129
  model_or_url = model
130
130
  elif api_type in [HFEmbeddingAPIType.INFERENCE_ENDPOINTS, HFEmbeddingAPIType.TEXT_EMBEDDINGS_INFERENCE]:
131
131
  url = api_params.get("url")
@@ -150,8 +150,39 @@ class HuggingFaceAPITextEmbedder:
150
150
  self.suffix = suffix
151
151
  self.truncate = truncate
152
152
  self.normalize = normalize
153
- self._client = InferenceClient(model_or_url, token=token.resolve_value() if token else None)
154
- self._async_client = AsyncInferenceClient(model_or_url, token=token.resolve_value() if token else None)
153
+ self._model_or_url = model_or_url
154
+ self._client: InferenceClient | None = None
155
+ self._async_client: AsyncInferenceClient | None = None
156
+
157
+ def _client_kwargs(self) -> dict[str, Any]:
158
+ """Build the keyword arguments used to create Hugging Face clients."""
159
+ return {"model": self._model_or_url, "token": self.token.resolve_value() if self.token else None}
160
+
161
+ def warm_up(self) -> None:
162
+ """Create the synchronous Hugging Face client."""
163
+ if self._client is None:
164
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
165
+ _check_valid_model(self._model_or_url, HFModelType.EMBEDDING, self.token)
166
+ self._client = InferenceClient(**self._client_kwargs())
167
+
168
+ async def warm_up_async(self) -> None:
169
+ """Create the asynchronous Hugging Face client."""
170
+ if self._async_client is None:
171
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
172
+ await _check_valid_model_async(self._model_or_url, HFModelType.EMBEDDING, self.token)
173
+ self._async_client = AsyncInferenceClient(**self._client_kwargs())
174
+
175
+ def close(self) -> None:
176
+ """Close the synchronous Hugging Face client."""
177
+ if self._client is not None:
178
+ self._client.close()
179
+ self._client = None
180
+
181
+ async def close_async(self) -> None:
182
+ """Close the asynchronous Hugging Face client."""
183
+ if self._async_client is not None:
184
+ await self._async_client.close()
185
+ self._async_client = None
155
186
 
156
187
  def _prepare_input(self, text: str) -> tuple[str, bool | None, bool | None]:
157
188
  if not isinstance(text, str):
@@ -224,6 +255,9 @@ class HuggingFaceAPITextEmbedder:
224
255
  A dictionary with the following keys:
225
256
  - `embedding`: The embedding of the input text.
226
257
  """
258
+ self.warm_up()
259
+ assert self._client is not None # noqa: S101
260
+
227
261
  text_to_embed, truncate_val, normalize_val = self._prepare_input(text)
228
262
 
229
263
  np_embedding = self._client.feature_extraction(
@@ -256,6 +290,9 @@ class HuggingFaceAPITextEmbedder:
256
290
  A dictionary with the following keys:
257
291
  - `embedding`: The embedding of the input text.
258
292
  """
293
+ await self.warm_up_async()
294
+ assert self._async_client is not None # noqa: S101
295
+
259
296
  text_to_embed, truncate_val, normalize_val = self._prepare_input(text)
260
297
 
261
298
  np_embedding = await self._async_client.feature_extraction(
@@ -49,6 +49,7 @@ from haystack_integrations.common.huggingface_api.utils import (
49
49
  HFGenerationAPIType,
50
50
  HFModelType,
51
51
  _check_valid_model,
52
+ _check_valid_model_async,
52
53
  )
53
54
 
54
55
  logger = logging.getLogger(__name__)
@@ -280,7 +281,7 @@ class HuggingFaceAPIChatGenerator:
280
281
  api_type = "serverless_inference_api" # this is equivalent to the above
281
282
 
282
283
  generator = HuggingFaceAPIChatGenerator(api_type=api_type,
283
- api_params={"model": "Qwen/Qwen2.5-7B-Instruct",
284
+ api_params={"model": "Qwen/Qwen3.5-9B",
284
285
  "provider": "together"},
285
286
  token=Secret.from_token("<your-api-key>"))
286
287
 
@@ -305,8 +306,8 @@ class HuggingFaceAPIChatGenerator:
305
306
  generator = HuggingFaceAPIChatGenerator(
306
307
  api_type=HFGenerationAPIType.SERVERLESS_INFERENCE_API,
307
308
  api_params={
308
- "model": "Qwen/Qwen2.5-VL-7B-Instruct", # Vision Language Model
309
- "provider": "hyperbolic"
309
+ "model": "Qwen/Qwen3.5-9B", # Vision Language Model
310
+ "provider": "together"
310
311
  },
311
312
  token=Secret.from_token("<your-api-key>")
312
313
  )
@@ -404,7 +405,6 @@ class HuggingFaceAPIChatGenerator:
404
405
  if model is None:
405
406
  msg = "To use the Serverless Inference API, you need to specify the `model` parameter in `api_params`."
406
407
  raise ValueError(msg)
407
- _check_valid_model(model, HFModelType.GENERATION, token)
408
408
  model_or_url = model
409
409
  elif api_type in [HFGenerationAPIType.INFERENCE_ENDPOINTS, HFGenerationAPIType.TEXT_GENERATION_INFERENCE]:
410
410
  url = api_params.get("url")
@@ -438,27 +438,56 @@ class HuggingFaceAPIChatGenerator:
438
438
  self.token = token
439
439
  self.generation_kwargs = generation_kwargs
440
440
  self.streaming_callback = streaming_callback
441
-
442
- resolved_api_params: dict[str, Any] = {k: v for k, v in api_params.items() if k not in ("model", "url")}
443
- self._client = InferenceClient(
444
- model_or_url, token=token.resolve_value() if token else None, **resolved_api_params
445
- )
446
- self._async_client = AsyncInferenceClient(
447
- model_or_url, token=token.resolve_value() if token else None, **resolved_api_params
448
- )
449
441
  self.tools = tools
450
- self._is_warmed_up = False
442
+ self._model_or_url = model_or_url
443
+ self._client: InferenceClient | None = None
444
+ self._async_client: AsyncInferenceClient | None = None
445
+ self._tools_warmed_up = False
446
+
447
+ def _client_kwargs(self) -> dict[str, Any]:
448
+ """Build the keyword arguments used to create Hugging Face clients."""
449
+ return {
450
+ "model": self._model_or_url,
451
+ "token": self.token.resolve_value() if self.token else None,
452
+ **{k: v for k, v in self.api_params.items() if k not in ("model", "url")},
453
+ }
454
+
455
+ def _warm_up_tools(self) -> None:
456
+ if not self._tools_warmed_up:
457
+ warm_up_tools(self.tools)
458
+ self._tools_warmed_up = True
451
459
 
452
460
  def warm_up(self) -> None:
453
461
  """
454
462
  Warm up the Hugging Face API chat generator.
455
463
 
456
- This will warm up the tools registered in the chat generator.
457
- This method is idempotent and will only warm up the tools once.
464
+ This creates the synchronous client and warms up the configured tools.
458
465
  """
459
- if not self._is_warmed_up:
460
- warm_up_tools(self.tools)
461
- self._is_warmed_up = True
466
+ self._warm_up_tools()
467
+ if self._client is None:
468
+ if self.api_type == HFGenerationAPIType.SERVERLESS_INFERENCE_API:
469
+ _check_valid_model(self._model_or_url, HFModelType.GENERATION, self.token)
470
+ self._client = InferenceClient(**self._client_kwargs())
471
+
472
+ async def warm_up_async(self) -> None:
473
+ """Create the asynchronous Hugging Face client and warm up the configured tools."""
474
+ self._warm_up_tools()
475
+ if self._async_client is None:
476
+ if self.api_type == HFGenerationAPIType.SERVERLESS_INFERENCE_API:
477
+ await _check_valid_model_async(self._model_or_url, HFModelType.GENERATION, self.token)
478
+ self._async_client = AsyncInferenceClient(**self._client_kwargs())
479
+
480
+ def close(self) -> None:
481
+ """Close the synchronous Hugging Face client."""
482
+ if self._client is not None:
483
+ self._client.close()
484
+ self._client = None
485
+
486
+ async def close_async(self) -> None:
487
+ """Close the asynchronous Hugging Face client."""
488
+ if self._async_client is not None:
489
+ await self._async_client.close()
490
+ self._async_client = None
462
491
 
463
492
  def to_dict(self) -> dict[str, Any]:
464
493
  """
@@ -520,8 +549,8 @@ class HuggingFaceAPIChatGenerator:
520
549
  :returns: A dictionary with the following keys:
521
550
  - `replies`: A list containing the generated responses as ChatMessage objects.
522
551
  """
523
- if not self._is_warmed_up:
524
- self.warm_up()
552
+ self.warm_up()
553
+ assert self._client is not None # noqa: S101
525
554
 
526
555
  messages = _normalize_messages(messages)
527
556
 
@@ -582,8 +611,8 @@ class HuggingFaceAPIChatGenerator:
582
611
  :returns: A dictionary with the following keys:
583
612
  - `replies`: A list containing the generated responses as ChatMessage objects.
584
613
  """
585
- if not self._is_warmed_up:
586
- self.warm_up()
614
+ await self.warm_up_async()
615
+ assert self._async_client is not None # noqa: S101
587
616
 
588
617
  messages = _normalize_messages(messages)
589
618
 
@@ -615,6 +644,7 @@ class HuggingFaceAPIChatGenerator:
615
644
  generation_kwargs: dict[str, Any],
616
645
  streaming_callback: SyncStreamingCallbackT,
617
646
  ) -> dict[str, list[ChatMessage]]:
647
+ assert self._client is not None # noqa: S101
618
648
  api_output: Iterable[ChatCompletionStreamOutput] = self._client.chat_completion(
619
649
  messages,
620
650
  stream=True,
@@ -643,6 +673,7 @@ class HuggingFaceAPIChatGenerator:
643
673
  generation_kwargs: dict[str, Any],
644
674
  tools: list["ChatCompletionInputTool"] | None = None,
645
675
  ) -> dict[str, list[ChatMessage]]:
676
+ assert self._client is not None # noqa: S101
646
677
  api_chat_output: ChatCompletionOutput = self._client.chat_completion(
647
678
  messages=messages, tools=tools, **generation_kwargs
648
679
  )
@@ -686,6 +717,7 @@ class HuggingFaceAPIChatGenerator:
686
717
  generation_kwargs: dict[str, Any],
687
718
  streaming_callback: StreamingCallbackT,
688
719
  ) -> dict[str, list[ChatMessage]]:
720
+ assert self._async_client is not None # noqa: S101
689
721
  api_output: AsyncIterable[ChatCompletionStreamOutput] = await self._async_client.chat_completion(
690
722
  messages,
691
723
  stream=True,
@@ -717,6 +749,7 @@ class HuggingFaceAPIChatGenerator:
717
749
  generation_kwargs: dict[str, Any],
718
750
  tools: list["ChatCompletionInputTool"] | None = None,
719
751
  ) -> dict[str, list[ChatMessage]]:
752
+ assert self._async_client is not None # noqa: S101
720
753
  api_chat_output: ChatCompletionOutput = await self._async_client.chat_completion(
721
754
  messages=messages, tools=tools, **generation_kwargs
722
755
  )