A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Feature][Responses API] Support MCP tool in background mode (#23494)
Signed-off-by: wuhang <wuhang6@huawei.com>
wuhang committed
Aug 27, 2025 at 09:06 UTC
6891205b161e78ea6e255da194a4470e06997a3b
2 files changed
+164
-136
vllm/entrypoints/context.py
+26
-5
@@ -4,13 +4,15 @@ import json
4
import logging
5
from abc import ABC, abstractmethod
6
from collections.abc import Sequence
7
-from typing import TYPE_CHECKING, Union
7
+from contextlib import AsyncExitStack
8
+from typing import TYPE_CHECKING, Optional, Union
9
10
from openai_harmony import Author, Message, Role, StreamState, TextContent
11
12
from vllm.entrypoints.harmony_utils import (
13
get_encoding, get_streamable_parser_for_assistant, render_for_completion)
14
from vllm.entrypoints.tool import Tool
15
+from vllm.entrypoints.tool_server import ToolServer
16
from vllm.outputs import RequestOutput
17
18
if TYPE_CHECKING:
@@ -37,6 +39,11 @@ class ConversationContext(ABC):
39
def render_for_completion(self) -> list[int]:
40
pass
41
42
+ @abstractmethod
43
+ async def init_tool_sessions(self, tool_server: Optional[ToolServer],
44
+ exit_stack: AsyncExitStack) -> None:
45
+ pass
46
+
47
48
class SimpleContext(ConversationContext):
49
@@ -55,16 +62,21 @@ class SimpleContext(ConversationContext):
62
def render_for_completion(self) -> list[int]:
63
raise NotImplementedError("Should not be called.")
64
65
+ async def init_tool_sessions(self, tool_server: Optional[ToolServer],
66
+ exit_stack: AsyncExitStack) -> None:
67
+ pass
68
+
69
70
class HarmonyContext(ConversationContext):
71
72
def __init__(
73
self,
74
messages: list,
64
- tool_sessions: dict[str, Tool],
75
+ available_tools: list[str],
76
):
77
self._messages = messages
67
- self.tool_sessions = tool_sessions
78
+ self.available_tools = available_tools
79
+ self._tool_sessions: dict[str, Union[ClientSession, Tool]] = {}
80
81
self.parser = get_streamable_parser_for_assistant()
82
self.num_init_messages = len(messages)
@@ -116,10 +128,10 @@ class HarmonyContext(ConversationContext):
128
if recipient is not None:
129
if recipient.startswith("browser."):
130
return await self.call_search_tool(
119
- self.tool_sessions["browser"], last_msg)
131
+ self._tool_sessions["browser"], last_msg)
132
elif recipient.startswith("python"):
133
return await self.call_python_tool(
122
- self.tool_sessions["python"], last_msg)
134
+ self._tool_sessions["python"], last_msg)
135
raise ValueError("No tool call found")
136
137
def render_for_completion(self) -> list[int]:
@@ -161,6 +173,15 @@ class HarmonyContext(ConversationContext):
173
recipient=Role.ASSISTANT)
174
]
175
176
+ async def init_tool_sessions(self, tool_server: Optional[ToolServer],
177
+ exit_stack: AsyncExitStack) -> None:
178
+ if tool_server:
179
+ for tool_name in self.available_tools:
180
+ if tool_name not in self._tool_sessions:
181
+ self._tool_sessions[
182
+ tool_name] = await exit_stack.enter_async_context(
183
+ tool_server.new_session(tool_name))
184
+
185
186
class StreamingHarmonyContext(HarmonyContext):
187
vllm/entrypoints/openai/serving_responses.py
+138
-131
@@ -8,7 +8,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Sequence
8
from contextlib import AsyncExitStack
9
from copy import copy
10
from http import HTTPStatus
11
-from typing import Any, Callable, Final, Optional, Union
11
+from typing import Callable, Final, Optional, Union
12
13
import jinja2
14
import openai.types.responses as openai_responses_types
@@ -248,10 +248,10 @@ class OpenAIServingResponses(OpenAIServing):
248
raw_request.state.request_metadata = request_metadata
249
250
if self.tool_server is not None and isinstance(
251
- self.tool_server, MCPToolServer
252
- ) and (request.background or request.stream) and request.tools and any(
253
- tool.type in ["web_search_preview", "code_interpreter"]
254
- for tool in request.tools):
251
+ self.tool_server,
252
+ MCPToolServer) and request.stream and request.tools and any(
253
+ tool.type in ["web_search_preview", "code_interpreter"]
254
+ for tool in request.tools):
255
return self.create_error_response(
256
"MCP tool server is not supported in background mode and "
257
"streaming mode")
@@ -265,103 +265,70 @@ class OpenAIServingResponses(OpenAIServing):
265
builtin_tool_list.append("browser")
266
if self.tool_server.has_tool("python"):
267
builtin_tool_list.append("python")
268
- async with AsyncExitStack() as exit_stack:
269
- try:
270
- if self.tool_server is not None:
271
- # TODO: initialize tool sessions lazily when the session
272
- # is actually used.
273
- tool_session_ctxs: dict[str, Any] = {
274
- tool_name:
275
- exit_stack.enter_async_context(
276
- self.tool_server.new_session(tool_name))
277
- for tool_name in builtin_tool_list
278
- }
279
- tool_sessions = {}
280
- for tool_name in builtin_tool_list:
281
- tool_sessions[tool_name] = (
282
- await tool_session_ctxs[tool_name])
283
- else:
284
- assert len(builtin_tool_list) == 0
285
- tool_sessions = {}
286
- for i, engine_prompt in enumerate(engine_prompts):
287
- default_max_tokens = self.max_model_len - len(
288
- engine_prompt["prompt_token_ids"])
289
- sampling_params = request.to_sampling_params(
290
- default_max_tokens, self.default_sampling_params)
291
-
292
- trace_headers = (None if raw_request is None else await
293
- self._get_trace_headers(
294
- raw_request.headers))
295
-
296
- context: ConversationContext
297
- if self.use_harmony:
298
- if request.stream:
299
- context = StreamingHarmonyContext(
300
- messages, tool_sessions)
301
- else:
302
- context = HarmonyContext(messages, tool_sessions)
303
- else:
304
- context = SimpleContext()
305
- generator = self._generate_with_builtin_tools(
306
- request_id=request.request_id,
307
- request_prompt=request_prompts[i],
308
- engine_prompt=engine_prompt,
309
- sampling_params=sampling_params,
310
- context=context,
311
- lora_request=lora_request,
312
- priority=request.priority,
313
- trace_headers=trace_headers,
314
- )
315
- generators.append(generator)
316
- except ValueError as e:
317
- # TODO: Use a vllm-specific Validation Error
318
- return self.create_error_response(str(e))
319
-
320
- assert len(generators) == 1
321
- result_generator, = generators
322
-
323
- # Store the input messages.
324
- if request.store:
325
- self.msg_store[request.request_id] = messages
268
327
- if request.background:
328
- created_time = int(time.time())
329
- response = ResponsesResponse.from_request(
330
- request,
331
- sampling_params,
332
- model_name=model_name,
333
- created_time=created_time,
334
- output=[],
335
- status="queued",
336
- usage=None,
269
+ if self.tool_server is not None:
270
+ available_tools = builtin_tool_list
271
+ else:
272
+ assert len(builtin_tool_list) == 0
273
+ available_tools = []
274
+ try:
275
+ for i, engine_prompt in enumerate(engine_prompts):
276
+ default_max_tokens = self.max_model_len - len(
277
+ engine_prompt["prompt_token_ids"])
278
+ sampling_params = request.to_sampling_params(
279
+ default_max_tokens, self.default_sampling_params)
280
+
281
+ trace_headers = (None if raw_request is None else await
282
+ self._get_trace_headers(raw_request.headers))
283
+
284
+ context: ConversationContext
285
+ if self.use_harmony:
286
+ if request.stream:
287
+ context = StreamingHarmonyContext(
288
+ messages, available_tools)
289
+ else:
290
+ context = HarmonyContext(messages, available_tools)
291
+ else:
292
+ context = SimpleContext()
293
+ generator = self._generate_with_builtin_tools(
294
+ request_id=request.request_id,
295
+ request_prompt=request_prompts[i],
296
+ engine_prompt=engine_prompt,
297
+ sampling_params=sampling_params,
298
+ context=context,
299
+ lora_request=lora_request,
300
+ priority=request.priority,
301
+ trace_headers=trace_headers,
302
)
338
- async with self.response_store_lock:
339
- self.response_store[response.id] = response
303
+ generators.append(generator)
304
+ except ValueError as e:
305
+ # TODO: Use a vllm-specific Validation Error
306
+ return self.create_error_response(str(e))
307
341
- # Run the request in the background.
342
- task = asyncio.create_task(
343
- self._run_background_request(
344
- request,
345
- sampling_params,
346
- result_generator,
347
- context,
348
- model_name,
349
- tokenizer,
350
- request_metadata,
351
- created_time,
352
- ),
353
- name=f"create_{response.id}",
354
- )
308
+ assert len(generators) == 1
309
+ result_generator, = generators
310
+
311
+ # Store the input messages.
312
+ if request.store:
313
+ self.msg_store[request.request_id] = messages
314
356
- # For cleanup.
357
- response_id = response.id
358
- self.background_tasks[response_id] = task
359
- task.add_done_callback(
360
- lambda _: self.background_tasks.pop(response_id, None))
361
- return response
315
+ if request.background:
316
+ created_time = int(time.time())
317
+ response = ResponsesResponse.from_request(
318
+ request,
319
+ sampling_params,
320
+ model_name=model_name,
321
+ created_time=created_time,
322
+ output=[],
323
+ status="queued",
324
+ usage=None,
325
+ )
326
+ async with self.response_store_lock:
327
+ self.response_store[response.id] = response
328
363
- if request.stream:
364
- return self.responses_stream_generator(
329
+ # Run the request in the background.
330
+ task = asyncio.create_task(
331
+ self._run_background_request(
332
request,
333
sampling_params,
334
result_generator,
@@ -369,21 +336,41 @@ class OpenAIServingResponses(OpenAIServing):
336
model_name,
337
tokenizer,
338
request_metadata,
372
- )
339
+ created_time,
340
+ ),
341
+ name=f"create_{response.id}",
342
+ )
343
374
- try:
375
- return await self.responses_full_generator(
376
- request,
377
- sampling_params,
378
- result_generator,
379
- context,
380
- model_name,
381
- tokenizer,
382
- request_metadata,
383
- )
384
- except Exception as e:
385
- return self.create_error_response(str(e))
386
- return self.create_error_response("Should not reach here")
344
+ # For cleanup.
345
+ response_id = response.id
346
+ self.background_tasks[response_id] = task
347
+ task.add_done_callback(
348
+ lambda _: self.background_tasks.pop(response_id, None))
349
+ return response
350
+
351
+ if request.stream:
352
+ return self.responses_stream_generator(
353
+ request,
354
+ sampling_params,
355
+ result_generator,
356
+ context,
357
+ model_name,
358
+ tokenizer,
359
+ request_metadata,
360
+ )
361
+
362
+ try:
363
+ return await self.responses_full_generator(
364
+ request,
365
+ sampling_params,
366
+ result_generator,
367
+ context,
368
+ model_name,
369
+ tokenizer,
370
+ request_metadata,
371
+ )
372
+ except Exception as e:
373
+ return self.create_error_response(str(e))
374
375
async def _make_request(
376
self,
@@ -439,14 +426,16 @@ class OpenAIServingResponses(OpenAIServing):
426
if created_time is None:
427
created_time = int(time.time())
428
442
- try:
443
- async for _ in result_generator:
444
- pass
445
- except asyncio.CancelledError:
446
- return self.create_error_response("Client disconnected")
447
- except ValueError as e:
448
- # TODO: Use a vllm-specific Validation Error
449
- return self.create_error_response(str(e))
429
+ async with AsyncExitStack() as exit_stack:
430
+ try:
431
+ await context.init_tool_sessions(self.tool_server, exit_stack)
432
+ async for _ in result_generator:
433
+ pass
434
+ except asyncio.CancelledError:
435
+ return self.create_error_response("Client disconnected")
436
+ except ValueError as e:
437
+ # TODO: Use a vllm-specific Validation Error
438
+ return self.create_error_response(str(e))
439
440
if self.use_harmony:
441
assert isinstance(context, HarmonyContext)
@@ -838,7 +827,7 @@ class OpenAIServingResponses(OpenAIServing):
827
status_code=HTTPStatus.BAD_REQUEST,
828
)
829
841
- async def responses_stream_generator(
830
+ async def _process_streaming_events(
831
self,
832
request: ResponsesRequest,
833
sampling_params: SamplingParams,
@@ -847,18 +836,8 @@ class OpenAIServingResponses(OpenAIServing):
836
model_name: str,
837
tokenizer: AnyTokenizer,
838
request_metadata: RequestResponseMetadata,
850
- created_time: Optional[int] = None,
839
+ created_time: int,
840
) -> AsyncGenerator[str, None]:
852
- # TODO:
853
- # 1. Handle disconnect
854
-
855
- if not isinstance(context, StreamingHarmonyContext):
856
- raise NotImplementedError(
857
- "Streaming is not supported for responses API without Harmony."
858
- )
859
-
860
- created_time = created_time or int(time.time())
861
-
841
sequence_number = 0
842
843
def _send_event(event: BaseModel):
@@ -1270,3 +1249,31 @@ class OpenAIServingResponses(OpenAIServing):
1249
sequence_number=-1,
1250
response=final_response.model_dump(),
1251
))
1252
+
1253
+ async def responses_stream_generator(
1254
+ self,
1255
+ request: ResponsesRequest,
1256
+ sampling_params: SamplingParams,
1257
+ result_generator: AsyncIterator[Optional[ConversationContext]],
1258
+ context: ConversationContext,
1259
+ model_name: str,
1260
+ tokenizer: AnyTokenizer,
1261
+ request_metadata: RequestResponseMetadata,
1262
+ created_time: Optional[int] = None,
1263
+ ) -> AsyncGenerator[str, None]:
1264
+ # TODO:
1265
+ # 1. Handle disconnect
1266
+
1267
+ if not isinstance(context, StreamingHarmonyContext):
1268
+ raise NotImplementedError(
1269
+ "Streaming is not supported for responses API without Harmony."
1270
+ )
1271
+
1272
+ created_time = created_time or int(time.time())
1273
+
1274
+ async with AsyncExitStack() as exit_stack:
1275
+ await context.init_tool_sessions(self.tool_server, exit_stack)
1276
+ async for event_data in self._process_streaming_events(
1277
+ request, sampling_params, result_generator, context,
1278
+ model_name, tokenizer, request_metadata, created_time):
1279
+ yield event_data