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
[gpt-oss] Add loop for built-in tool call (#22374)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Woosuk Kwon committed
Aug 6, 2025 at 10:32 UTC
ec7cb1922478015b4e7eae73c6acde8b598a05a8
2 files changed
+73
-16
vllm/entrypoints/openai/serving_engine.py
+56
@@ -35,6 +35,7 @@ from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
35
apply_mistral_chat_template,
36
parse_chat_messages_futures,
37
resolve_chat_template_content_format)
38
+from vllm.entrypoints.context import ConversationContext
39
from vllm.entrypoints.logger import RequestLogger
40
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
41
ChatCompletionResponse,
@@ -948,6 +949,61 @@ class OpenAIServing:
949
950
return conversation, [request_prompt], [engine_prompt]
951
952
+ async def _generate_with_builtin_tools(
953
+ self,
954
+ request_id: str,
955
+ request_prompt: RequestPrompt,
956
+ engine_prompt: EngineTokensPrompt,
957
+ sampling_params: SamplingParams,
958
+ context: ConversationContext,
959
+ lora_request: Optional[LoRARequest] = None,
960
+ priority: int = 0,
961
+ **kwargs,
962
+ ):
963
+ orig_priority = priority
964
+ while True:
965
+ self._log_inputs(
966
+ request_id,
967
+ request_prompt,
968
+ params=sampling_params,
969
+ lora_request=lora_request,
970
+ )
971
+ generator = self.engine_client.generate(
972
+ engine_prompt,
973
+ sampling_params,
974
+ request_id,
975
+ lora_request=lora_request,
976
+ priority=priority,
977
+ **kwargs,
978
+ )
979
+ async for res in generator:
980
+ context.append_output(res)
981
+ # NOTE(woosuk): The stop condition is handled by the engine.
982
+ yield context
983
+
984
+ if not context.need_builtin_tool_call():
985
+ # The model did not ask for a tool call, so we're done.
986
+ break
987
+
988
+ # Call the tool and update the context with the result.
989
+ tool_output = await context.call_tool()
990
+ context.append_output(tool_output)
991
+
992
+ # TODO: uncomment this and enable tool output streaming
993
+ # yield context
994
+
995
+ # Create inputs for the next turn.
996
+ # Render the next prompt token ids.
997
+ prompt_token_ids = context.render_for_completion()
998
+ engine_prompt = EngineTokensPrompt(
999
+ prompt_token_ids=prompt_token_ids)
1000
+ request_prompt = prompt_token_ids
1001
+ # Update the sampling params.
1002
+ sampling_params.max_tokens = (self.max_model_len -
1003
+ len(prompt_token_ids))
1004
+ # OPTIMIZATION
1005
+ priority = orig_priority - 1
1006
+
1007
def _load_prompt_embeds(
1008
self,
1009
prompt_embeds: Optional[Union[bytes, list[bytes]]],
vllm/entrypoints/openai/serving_responses.py
+17
-16
@@ -16,6 +16,7 @@ from vllm.config import ModelConfig
16
from vllm.engine.protocol import EngineClient
17
from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
18
ChatTemplateContentFormatOption)
19
+from vllm.entrypoints.context import ConversationContext, SimpleContext
20
from vllm.entrypoints.logger import RequestLogger
21
# yapf conflicts with isort for this block
22
# yapf: disable
@@ -29,7 +30,6 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse,
30
from vllm.entrypoints.openai.serving_engine import OpenAIServing
31
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
32
from vllm.logger import init_logger
32
-from vllm.outputs import RequestOutput
33
from vllm.reasoning import ReasoningParser, ReasoningParserManager
34
from vllm.sampling_params import SamplingParams
35
from vllm.transformers_utils.tokenizer import AnyTokenizer
@@ -187,7 +187,7 @@ class OpenAIServingResponses(OpenAIServing):
187
raw_request.state.request_metadata = request_metadata
188
189
# Schedule the request and get the result generator.
190
- generators: list[AsyncGenerator[RequestOutput, None]] = []
190
+ generators: list[AsyncGenerator[ConversationContext, None]] = []
191
try:
192
for i, engine_prompt in enumerate(engine_prompts):
193
default_max_tokens = self.max_model_len - len(
@@ -195,21 +195,19 @@ class OpenAIServingResponses(OpenAIServing):
195
sampling_params = request.to_sampling_params(
196
default_max_tokens, self.default_sampling_params)
197
198
- self._log_inputs(request.request_id,
199
- request_prompts[i],
200
- params=sampling_params,
201
- lora_request=lora_request)
202
-
198
trace_headers = (None if raw_request is None else await
199
self._get_trace_headers(raw_request.headers))
200
206
- generator = self.engine_client.generate(
207
- engine_prompt,
208
- sampling_params,
209
- request.request_id,
201
+ context = SimpleContext()
202
+ generator = self._generate_with_builtin_tools(
203
+ request_id=request.request_id,
204
+ request_prompt=request_prompts[i],
205
+ engine_prompt=engine_prompt,
206
+ sampling_params=sampling_params,
207
+ context=context,
208
lora_request=lora_request,
211
- trace_headers=trace_headers,
209
priority=request.priority,
210
+ trace_headers=trace_headers,
211
)
212
generators.append(generator)
213
except ValueError as e:
@@ -277,7 +275,7 @@ class OpenAIServingResponses(OpenAIServing):
275
self,
276
request: ResponsesRequest,
277
sampling_params: SamplingParams,
280
- result_generator: AsyncIterator[RequestOutput],
278
+ result_generator: AsyncIterator[ConversationContext],
279
model_name: str,
280
tokenizer: AnyTokenizer,
281
request_metadata: RequestResponseMetadata,
@@ -285,17 +283,20 @@ class OpenAIServingResponses(OpenAIServing):
283
) -> Union[ErrorResponse, ResponsesResponse]:
284
if created_time is None:
285
created_time = int(time.time())
288
- final_res: Optional[RequestOutput] = None
286
287
+ context: Optional[ConversationContext] = None
288
try:
291
- async for res in result_generator:
292
- final_res = res
289
+ async for context in result_generator:
290
+ pass
291
except asyncio.CancelledError:
292
return self.create_error_response("Client disconnected")
293
except ValueError as e:
294
# TODO: Use a vllm-specific Validation Error
295
return self.create_error_response(str(e))
296
297
+ assert context is not None
298
+ assert isinstance(context, SimpleContext)
299
+ final_res = context.last_output
300
assert final_res is not None
301
assert len(final_res.outputs) == 1
302
final_output = final_res.outputs[0]