vllm-haystack 1.3.0__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/CHANGELOG.md +8 -0
  2. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/PKG-INFO +5 -4
  3. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/README.md +2 -1
  4. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/pyproject.toml +2 -1
  5. vllm_haystack-2.0.0/src/haystack_integrations/common/vllm/utils.py +65 -0
  6. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/document_embedder.py +37 -17
  7. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/text_embedder.py +35 -17
  8. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/chat/chat_generator.py +51 -25
  9. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/vllm/ranker.py +39 -20
  10. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_chat_generator.py +99 -34
  11. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_document_embedder.py +87 -30
  12. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_ranker.py +103 -29
  13. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_text_embedder.py +87 -28
  14. vllm_haystack-2.0.0/tests/test_utils.py +41 -0
  15. vllm_haystack-1.3.0/src/haystack_integrations/common/vllm/utils.py +0 -38
  16. vllm_haystack-1.3.0/tests/test_utils.py +0 -29
  17. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/.gitignore +0 -0
  18. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/LICENSE.txt +0 -0
  19. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -0
  20. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/common/py.typed +0 -0
  21. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/common/vllm/__init__.py +0 -0
  22. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  23. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/__init__.py +0 -0
  24. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  25. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/__init__.py +0 -0
  26. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/chat/__init__.py +0 -0
  27. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  28. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/vllm/__init__.py +0 -0
  29. {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/__init__.py +0 -0
@@ -1,5 +1,13 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/vllm-v1.3.0] - 2026-07-07
4
+
5
+ ### 🧹 Chores
6
+
7
+ - Support sync streaming callbacks in async contexts for Haystack 2.x/3.x compatibility (#3534)
8
+ - *(vllm)* Unpin fastapi (#3557)
9
+
10
+
3
11
  ## [integrations/vllm-v1.2.1] - 2026-07-02
4
12
 
5
13
  ### 🚀 Features
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: vllm-haystack
3
- Version: 1.3.0
3
+ Version: 2.0.0
4
4
  Summary: Haystack integration for vllm
5
5
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/vllm#readme
6
6
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -19,7 +19,7 @@ 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: more-itertools>=9.0.0
24
24
  Requires-Dist: openai
25
25
  Requires-Dist: tqdm>=4.48.0
@@ -30,6 +30,7 @@ Description-Content-Type: text/markdown
30
30
  [![PyPI - Version](https://img.shields.io/pypi/v/vllm-haystack.svg)](https://pypi.org/project/vllm-haystack)
31
31
  [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/vllm-haystack.svg)](https://pypi.org/project/vllm-haystack)
32
32
 
33
+ - [Integration page](https://haystack.deepset.ai/integrations/vllm)
33
34
  - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/vllm/CHANGELOG.md)
34
35
 
35
36
  ---
@@ -60,4 +61,4 @@ To run the ranker server, use CPU Docker image:
60
61
  # ranker server (port 8002)
61
62
  docker run --rm -p 8002:8000 -e VLLM_CPU_OMP_THREADS_BIND=0-3 vllm/vllm-openai-cpu:latest \
62
63
  --model BAAI/bge-reranker-base --enforce-eager
63
- ```
64
+ ```
@@ -3,6 +3,7 @@
3
3
  [![PyPI - Version](https://img.shields.io/pypi/v/vllm-haystack.svg)](https://pypi.org/project/vllm-haystack)
4
4
  [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/vllm-haystack.svg)](https://pypi.org/project/vllm-haystack)
5
5
 
6
+ - [Integration page](https://haystack.deepset.ai/integrations/vllm)
6
7
  - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/vllm/CHANGELOG.md)
7
8
 
8
9
  ---
@@ -33,4 +34,4 @@ To run the ranker server, use CPU Docker image:
33
34
  # ranker server (port 8002)
34
35
  docker run --rm -p 8002:8000 -e VLLM_CPU_OMP_THREADS_BIND=0-3 vllm/vllm-openai-cpu:latest \
35
36
  --model BAAI/bge-reranker-base --enforce-eager
36
- ```
37
+ ```
@@ -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", "openai", "more_itertools>=9.0.0", "tqdm>=4.48.0"]
26
+ dependencies = ["haystack-ai>=3.0.0", "openai", "more_itertools>=9.0.0", "tqdm>=4.48.0"]
27
27
 
28
28
  [project.urls]
29
29
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/vllm#readme"
@@ -127,6 +127,7 @@ ignore = [
127
127
  "PLR0912",
128
128
  "PLR0913",
129
129
  "PLR0915",
130
+ "PLR0917",
130
131
  # Allow `Any` type - used legitimately for dynamic types and SDK boundaries
131
132
  "ANN401",
132
133
  ]
@@ -0,0 +1,65 @@
1
+ # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from typing import Any
6
+
7
+ from haystack.utils import Secret
8
+ from haystack.utils.http_client import init_http_client
9
+ from openai import AsyncOpenAI, OpenAI
10
+
11
+
12
+ def _openai_client_kwargs(
13
+ api_key: Secret | None,
14
+ api_base_url: str,
15
+ timeout: float | None,
16
+ max_retries: int | None,
17
+ ) -> dict[str, Any]:
18
+ """
19
+ Build the common keyword arguments for OpenAI clients pointing at a vLLM server.
20
+
21
+ A placeholder api key is used when the user did not supply one and no `VLLM_API_KEY` env var is set, because the
22
+ OpenAI client requires a non-empty value.
23
+ `timeout` and `max_retries` are only forwarded when provided: when None, the OpenAI client's own defaults apply.
24
+ """
25
+ resolved_api_key = "placeholder-api-key"
26
+ if api_key is not None and (value := api_key.resolve_value()):
27
+ resolved_api_key = value
28
+
29
+ client_kwargs: dict[str, Any] = {"api_key": resolved_api_key, "base_url": api_base_url}
30
+ if timeout is not None:
31
+ client_kwargs["timeout"] = timeout
32
+ if max_retries is not None:
33
+ client_kwargs["max_retries"] = max_retries
34
+
35
+ return client_kwargs
36
+
37
+
38
+ def _create_openai_client(
39
+ api_key: Secret | None,
40
+ api_base_url: str,
41
+ timeout: float | None,
42
+ max_retries: int | None,
43
+ http_client_kwargs: dict[str, Any] | None,
44
+ ) -> OpenAI:
45
+ """Build a synchronous OpenAI client pointing at a vLLM server."""
46
+ client_kwargs = _openai_client_kwargs(api_key, api_base_url, timeout, max_retries)
47
+ sync_http_client = init_http_client(http_client_kwargs, async_client=False)
48
+ # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
49
+ # https://github.com/openai/openai-python/blob/main/httpx2.md
50
+ return OpenAI(http_client=sync_http_client, **client_kwargs) # type: ignore[arg-type]
51
+
52
+
53
+ def _create_async_openai_client(
54
+ api_key: Secret | None,
55
+ api_base_url: str,
56
+ timeout: float | None,
57
+ max_retries: int | None,
58
+ http_client_kwargs: dict[str, Any] | None,
59
+ ) -> AsyncOpenAI:
60
+ """Build an asynchronous OpenAI client pointing at a vLLM server."""
61
+ client_kwargs = _openai_client_kwargs(api_key, api_base_url, timeout, max_retries)
62
+ async_http_client = init_http_client(http_client_kwargs, async_client=True)
63
+ # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
64
+ # https://github.com/openai/openai-python/blob/main/httpx2.md
65
+ return AsyncOpenAI(http_client=async_http_client, **client_kwargs) # type: ignore[arg-type]
@@ -12,7 +12,7 @@ from openai import APIError, AsyncOpenAI, OpenAI
12
12
  from tqdm import tqdm
13
13
  from tqdm.asyncio import tqdm as async_tqdm
14
14
 
15
- from haystack_integrations.common.vllm.utils import _create_openai_clients
15
+ from haystack_integrations.common.vllm.utils import _create_async_openai_client, _create_openai_client
16
16
 
17
17
  logger = logging.getLogger(__name__)
18
18
 
@@ -131,20 +131,40 @@ class VLLMDocumentEmbedder:
131
131
 
132
132
  self._client: OpenAI | None = None
133
133
  self._async_client: AsyncOpenAI | None = None
134
- self._is_warmed_up = False
135
134
 
136
135
  def warm_up(self) -> None:
137
- """Create the OpenAI clients."""
138
- if self._is_warmed_up:
139
- return
140
- self._client, self._async_client = _create_openai_clients(
141
- api_key=self.api_key,
142
- api_base_url=self.api_base_url,
143
- timeout=self.timeout,
144
- max_retries=self.max_retries,
145
- http_client_kwargs=self.http_client_kwargs,
146
- )
147
- self._is_warmed_up = True
136
+ """Create the synchronous OpenAI client."""
137
+ if self._client is None:
138
+ self._client = _create_openai_client(
139
+ api_key=self.api_key,
140
+ api_base_url=self.api_base_url,
141
+ timeout=self.timeout,
142
+ max_retries=self.max_retries,
143
+ http_client_kwargs=self.http_client_kwargs,
144
+ )
145
+
146
+ async def warm_up_async(self) -> None:
147
+ """Create the asynchronous OpenAI client."""
148
+ if self._async_client is None:
149
+ self._async_client = _create_async_openai_client(
150
+ api_key=self.api_key,
151
+ api_base_url=self.api_base_url,
152
+ timeout=self.timeout,
153
+ max_retries=self.max_retries,
154
+ http_client_kwargs=self.http_client_kwargs,
155
+ )
156
+
157
+ def close(self) -> None:
158
+ """Close the synchronous OpenAI client."""
159
+ if self._client is not None:
160
+ self._client.close()
161
+ self._client = None
162
+
163
+ async def close_async(self) -> None:
164
+ """Close the asynchronous OpenAI client."""
165
+ if self._async_client is not None:
166
+ await self._async_client.close()
167
+ self._async_client = None
148
168
 
149
169
  def _prepare_texts_to_embed(self, documents: list[Document]) -> dict[str, str]:
150
170
  """Concatenate each Document's text with the selected meta fields."""
@@ -254,8 +274,8 @@ class VLLMDocumentEmbedder:
254
274
  if not documents:
255
275
  return {"documents": [], "meta": {}}
256
276
 
257
- if not self._is_warmed_up:
258
- self.warm_up()
277
+ self.warm_up()
278
+ assert self._client is not None # noqa: S101
259
279
 
260
280
  texts_to_embed = self._prepare_texts_to_embed(documents)
261
281
  doc_ids_to_embeddings, meta = self._embed_batch(texts_to_embed, self.batch_size)
@@ -280,8 +300,8 @@ class VLLMDocumentEmbedder:
280
300
  if not documents:
281
301
  return {"documents": [], "meta": {}}
282
302
 
283
- if not self._is_warmed_up:
284
- self.warm_up()
303
+ await self.warm_up_async()
304
+ assert self._async_client is not None # noqa: S101
285
305
 
286
306
  texts_to_embed = self._prepare_texts_to_embed(documents)
287
307
  doc_ids_to_embeddings, meta = await self._embed_batch_async(texts_to_embed, self.batch_size)
@@ -9,7 +9,7 @@ from haystack.utils import Secret
9
9
  from openai import AsyncOpenAI, OpenAI
10
10
  from openai.types import CreateEmbeddingResponse
11
11
 
12
- from haystack_integrations.common.vllm.utils import _create_openai_clients
12
+ from haystack_integrations.common.vllm.utils import _create_async_openai_client, _create_openai_client
13
13
 
14
14
 
15
15
  @component
@@ -103,20 +103,40 @@ class VLLMTextEmbedder:
103
103
 
104
104
  self._client: OpenAI | None = None
105
105
  self._async_client: AsyncOpenAI | None = None
106
- self._is_warmed_up = False
107
106
 
108
107
  def warm_up(self) -> None:
109
- """Create the OpenAI clients."""
110
- if self._is_warmed_up:
111
- return
112
- self._client, self._async_client = _create_openai_clients(
113
- api_key=self.api_key,
114
- api_base_url=self.api_base_url,
115
- timeout=self.timeout,
116
- max_retries=self.max_retries,
117
- http_client_kwargs=self.http_client_kwargs,
118
- )
119
- self._is_warmed_up = True
108
+ """Create the synchronous OpenAI client."""
109
+ if self._client is None:
110
+ self._client = _create_openai_client(
111
+ api_key=self.api_key,
112
+ api_base_url=self.api_base_url,
113
+ timeout=self.timeout,
114
+ max_retries=self.max_retries,
115
+ http_client_kwargs=self.http_client_kwargs,
116
+ )
117
+
118
+ async def warm_up_async(self) -> None:
119
+ """Create the asynchronous OpenAI client."""
120
+ if self._async_client is None:
121
+ self._async_client = _create_async_openai_client(
122
+ api_key=self.api_key,
123
+ api_base_url=self.api_base_url,
124
+ timeout=self.timeout,
125
+ max_retries=self.max_retries,
126
+ http_client_kwargs=self.http_client_kwargs,
127
+ )
128
+
129
+ def close(self) -> None:
130
+ """Close the synchronous OpenAI client."""
131
+ if self._client is not None:
132
+ self._client.close()
133
+ self._client = None
134
+
135
+ async def close_async(self) -> None:
136
+ """Close the asynchronous OpenAI client."""
137
+ if self._async_client is not None:
138
+ await self._async_client.close()
139
+ self._async_client = None
120
140
 
121
141
  def _prepare_input(self, text: str) -> dict[str, Any]:
122
142
  if not isinstance(text, str):
@@ -155,8 +175,7 @@ class VLLMTextEmbedder:
155
175
  - `meta`: Information about the usage of the model.
156
176
  """
157
177
  kwargs = self._prepare_input(text)
158
- if not self._is_warmed_up:
159
- self.warm_up()
178
+ self.warm_up()
160
179
  assert self._client is not None # noqa: S101
161
180
  response = self._client.embeddings.create(**kwargs)
162
181
  return self._prepare_output(response)
@@ -172,8 +191,7 @@ class VLLMTextEmbedder:
172
191
  - `meta`: Information about the usage of the model.
173
192
  """
174
193
  kwargs = self._prepare_input(text)
175
- if not self._is_warmed_up:
176
- self.warm_up()
194
+ await self.warm_up_async()
177
195
  assert self._async_client is not None # noqa: S101
178
196
  response = await self._async_client.embeddings.create(**kwargs)
179
197
  return self._prepare_output(response)
@@ -37,7 +37,7 @@ from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
37
37
  from openai.types.chat import ChatCompletion, ChatCompletionChunk
38
38
  from openai.types.chat.chat_completion import Choice
39
39
 
40
- from haystack_integrations.common.vllm.utils import _create_openai_clients
40
+ from haystack_integrations.common.vllm.utils import _create_async_openai_client, _create_openai_client
41
41
 
42
42
  logger = logging.getLogger(__name__)
43
43
 
@@ -255,22 +255,48 @@ class VLLMChatGenerator:
255
255
 
256
256
  self._client: OpenAI | None = None
257
257
  self._async_client: AsyncOpenAI | None = None
258
- self._is_warmed_up = False
258
+ self._tools_warmed_up = False
259
259
 
260
- def warm_up(self) -> None:
261
- """Create the OpenAI clients and warm up tools."""
262
- if self._is_warmed_up:
263
- return
260
+ def _warm_up_tools(self) -> None:
261
+ if not self._tools_warmed_up:
262
+ warm_up_tools(self.tools)
263
+ self._tools_warmed_up = True
264
264
 
265
- self._client, self._async_client = _create_openai_clients(
266
- api_key=self.api_key,
267
- api_base_url=self.api_base_url,
268
- timeout=self.timeout,
269
- max_retries=self.max_retries,
270
- http_client_kwargs=self.http_client_kwargs,
271
- )
272
- warm_up_tools(self.tools)
273
- self._is_warmed_up = True
265
+ def warm_up(self) -> None:
266
+ """Create the synchronous OpenAI client and warm up tools."""
267
+ self._warm_up_tools()
268
+ if self._client is None:
269
+ self._client = _create_openai_client(
270
+ api_key=self.api_key,
271
+ api_base_url=self.api_base_url,
272
+ timeout=self.timeout,
273
+ max_retries=self.max_retries,
274
+ http_client_kwargs=self.http_client_kwargs,
275
+ )
276
+
277
+ async def warm_up_async(self) -> None:
278
+ """Create the asynchronous OpenAI client and warm up tools."""
279
+ self._warm_up_tools()
280
+ if self._async_client is None:
281
+ self._async_client = _create_async_openai_client(
282
+ api_key=self.api_key,
283
+ api_base_url=self.api_base_url,
284
+ timeout=self.timeout,
285
+ max_retries=self.max_retries,
286
+ http_client_kwargs=self.http_client_kwargs,
287
+ )
288
+
289
+ def close(self) -> None:
290
+ """Close the synchronous OpenAI client."""
291
+ if self._client is not None:
292
+ self._client.close()
293
+ self._client = None
294
+
295
+ async def close_async(self) -> None:
296
+ """Close the asynchronous OpenAI client."""
297
+ if self._async_client is not None:
298
+ await self._async_client.close()
299
+ self._async_client = None
274
300
 
275
301
  def to_dict(self) -> dict[str, Any]:
276
302
  """
@@ -451,8 +477,9 @@ class VLLMChatGenerator:
451
477
  :param streaming_callback:
452
478
  A callback function that is called when a new token is received from the stream.
453
479
  :param generation_kwargs:
454
- Additional keyword arguments for text generation. These parameters will
455
- override the parameters passed during component initialization.
480
+ Additional keyword arguments for text generation. These are merged per key with the
481
+ `generation_kwargs` passed at initialization: keys provided here take precedence, keys set only
482
+ at initialization are kept.
456
483
  For details on vLLM API parameters, see
457
484
  [vLLM documentation](https://docs.vllm.ai/en/stable/serving/openai_compatible_server/).
458
485
  :param tools:
@@ -463,9 +490,9 @@ class VLLMChatGenerator:
463
490
  A dictionary with the following key:
464
491
  - `replies`: A list containing the generated responses as ChatMessage instances.
465
492
  """
493
+ self.warm_up()
494
+ assert self._client is not None # noqa: S101
466
495
  messages = _normalize_messages(messages)
467
- if not self._is_warmed_up:
468
- self.warm_up()
469
496
 
470
497
  if len(messages) == 0:
471
498
  return {"replies": []}
@@ -475,7 +502,6 @@ class VLLMChatGenerator:
475
502
  )
476
503
 
477
504
  api_kwargs = self._prepare_api_call(messages, streaming_callback, generation_kwargs, tools)
478
- assert self._client is not None # noqa: S101
479
505
  chat_completion = self._client.chat.completions.create(**api_kwargs)
480
506
 
481
507
  if streaming_callback is not None:
@@ -509,8 +535,9 @@ class VLLMChatGenerator:
509
535
  A callback function that is called when a new token is received from the stream.
510
536
  Must be a coroutine.
511
537
  :param generation_kwargs:
512
- Additional keyword arguments for text generation. These parameters will
513
- override the parameters passed during component initialization.
538
+ Additional keyword arguments for text generation. These are merged per key with the
539
+ `generation_kwargs` passed at initialization: keys provided here take precedence, keys set only
540
+ at initialization are kept.
514
541
  For details on vLLM API parameters, see
515
542
  [vLLM documentation](https://docs.vllm.ai/en/stable/serving/openai_compatible_server/).
516
543
  :param tools:
@@ -521,9 +548,9 @@ class VLLMChatGenerator:
521
548
  A dictionary with the following key:
522
549
  - `replies`: A list containing the generated responses as ChatMessage instances.
523
550
  """
551
+ await self.warm_up_async()
552
+ assert self._async_client is not None # noqa: S101
524
553
  messages = _normalize_messages(messages)
525
- if not self._is_warmed_up:
526
- self.warm_up()
527
554
 
528
555
  if len(messages) == 0:
529
556
  return {"replies": []}
@@ -533,7 +560,6 @@ class VLLMChatGenerator:
533
560
  )
534
561
 
535
562
  api_kwargs = self._prepare_api_call(messages, streaming_callback, generation_kwargs, tools)
536
- assert self._async_client is not None # noqa: S101
537
563
  chat_completion = await self._async_client.chat.completions.create(**api_kwargs)
538
564
 
539
565
  if streaming_callback is not None:
@@ -109,24 +109,45 @@ class VLLMRanker:
109
109
  self.http_client_kwargs = http_client_kwargs
110
110
  self.extra_parameters = extra_parameters
111
111
 
112
- self._headers = {"Content-Type": "application/json"}
113
- if self.api_key is not None and (resolved_key := self.api_key.resolve_value()):
114
- self._headers["Authorization"] = f"Bearer {resolved_key}"
115
-
116
112
  self._client: httpx.Client | None = None
117
113
  self._async_client: httpx.AsyncClient | None = None
118
- self._is_warmed_up = False
119
114
 
120
- def warm_up(self) -> None:
121
- """Create the httpx clients."""
122
- if self._is_warmed_up:
123
- return
115
+ def _client_kwargs(self) -> dict[str, Any]:
116
+ """Build the keyword arguments used to create HTTP clients."""
117
+ headers = httpx.Headers((self.http_client_kwargs or {}).get("headers"))
118
+ headers["Content-Type"] = "application/json"
119
+ if self.api_key is not None and (resolved_key := self.api_key.resolve_value()):
120
+ headers["Authorization"] = f"Bearer {resolved_key}"
121
+
122
+ client_kwargs = self.http_client_kwargs.copy() if self.http_client_kwargs else {}
123
+ client_kwargs["headers"] = headers
124
+ return client_kwargs
124
125
 
125
- client = init_http_client(self.http_client_kwargs, async_client=False)
126
- async_client = init_http_client(self.http_client_kwargs, async_client=True)
127
- self._client = client if client is not None else httpx.Client()
128
- self._async_client = async_client if async_client is not None else httpx.AsyncClient()
129
- self._is_warmed_up = True
126
+ def warm_up(self) -> None:
127
+ """Create the synchronous HTTP client."""
128
+ if self._client is None:
129
+ client = init_http_client(http_client_kwargs=self._client_kwargs(), async_client=False)
130
+ assert client is not None # noqa: S101
131
+ self._client = client
132
+
133
+ async def warm_up_async(self) -> None:
134
+ """Create the asynchronous HTTP client."""
135
+ if self._async_client is None:
136
+ client = init_http_client(http_client_kwargs=self._client_kwargs(), async_client=True)
137
+ assert client is not None # noqa: S101
138
+ self._async_client = client
139
+
140
+ def close(self) -> None:
141
+ """Close the synchronous HTTP client."""
142
+ if self._client is not None:
143
+ self._client.close()
144
+ self._client = None
145
+
146
+ async def close_async(self) -> None:
147
+ """Close the asynchronous HTTP client."""
148
+ if self._async_client is not None:
149
+ await self._async_client.aclose()
150
+ self._async_client = None
130
151
 
131
152
  def _prepare_texts(self, documents: list[Document]) -> list[str]:
132
153
  """Concatenate each Document's text with the selected meta fields."""
@@ -210,13 +231,12 @@ class VLLMRanker:
210
231
 
211
232
  top_k, score_threshold = self._resolve_run_params(top_k, score_threshold)
212
233
 
213
- if not self._is_warmed_up:
214
- self.warm_up()
234
+ self.warm_up()
215
235
  assert self._client is not None # noqa: S101
216
236
 
217
237
  body = self._prepare_request(query, documents, top_k)
218
238
  url = f"{self.api_base_url.rstrip('/')}/rerank"
219
- response = self._client.post(url, json=body, headers=self._headers)
239
+ response = self._client.post(url, json=body)
220
240
  return self._parse_response(response.json(), documents, score_threshold)
221
241
 
222
242
  @component.output_types(documents=list[Document], meta=dict[str, Any])
@@ -246,11 +266,10 @@ class VLLMRanker:
246
266
 
247
267
  top_k, score_threshold = self._resolve_run_params(top_k, score_threshold)
248
268
 
249
- if not self._is_warmed_up:
250
- self.warm_up()
269
+ await self.warm_up_async()
251
270
  assert self._async_client is not None # noqa: S101
252
271
 
253
272
  body = self._prepare_request(query, documents, top_k)
254
273
  url = f"{self.api_base_url.rstrip('/')}/rerank"
255
- response = await self._async_client.post(url, json=body, headers=self._headers)
274
+ response = await self._async_client.post(url, json=body)
256
275
  return self._parse_response(response.json(), documents, score_threshold)