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.
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/CHANGELOG.md +8 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/PKG-INFO +5 -4
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/README.md +2 -1
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/pyproject.toml +2 -1
- vllm_haystack-2.0.0/src/haystack_integrations/common/vllm/utils.py +65 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/document_embedder.py +37 -17
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/text_embedder.py +35 -17
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/chat/chat_generator.py +51 -25
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/vllm/ranker.py +39 -20
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_chat_generator.py +99 -34
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_document_embedder.py +87 -30
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_ranker.py +103 -29
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/tests/test_text_embedder.py +87 -28
- vllm_haystack-2.0.0/tests/test_utils.py +41 -0
- vllm_haystack-1.3.0/src/haystack_integrations/common/vllm/utils.py +0 -38
- vllm_haystack-1.3.0/tests/test_utils.py +0 -29
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/.gitignore +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/LICENSE.txt +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/common/py.typed +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/common/vllm/__init__.py +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/embedders/vllm/__init__.py +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/__init__.py +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/generators/vllm/chat/__init__.py +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {vllm_haystack-1.3.0 → vllm_haystack-2.0.0}/src/haystack_integrations/components/rankers/vllm/__init__.py +0 -0
- {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.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: vllm-haystack
|
|
3
|
-
Version:
|
|
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>=
|
|
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
|
[](https://pypi.org/project/vllm-haystack)
|
|
31
31
|
[](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
|
[](https://pypi.org/project/vllm-haystack)
|
|
4
4
|
[](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>=
|
|
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
|
|
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
|
|
138
|
-
if self.
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
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
|
-
|
|
258
|
-
|
|
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
|
-
|
|
284
|
-
|
|
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
|
|
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
|
|
110
|
-
if self.
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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.
|
|
258
|
+
self._tools_warmed_up = False
|
|
259
259
|
|
|
260
|
-
def
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
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
|
-
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
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
|
|
455
|
-
|
|
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
|
|
513
|
-
|
|
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
|
|
121
|
-
"""
|
|
122
|
-
|
|
123
|
-
|
|
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
|
-
|
|
126
|
-
|
|
127
|
-
self._client
|
|
128
|
-
|
|
129
|
-
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
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
|
|
274
|
+
response = await self._async_client.post(url, json=body)
|
|
256
275
|
return self._parse_response(response.json(), documents, score_threshold)
|