context.py 34.4 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
import contextlib
5
import copy
6
import json
7
8
import logging
from abc import ABC, abstractmethod
9
from collections.abc import Callable
10
from contextlib import AsyncExitStack
11
from dataclasses import replace
12
from typing import TYPE_CHECKING, Final, Union
13

14
15
16
from openai.types.responses.response_function_tool_call_output_item import (
    ResponseFunctionToolCallOutputItem,
)
17
from openai.types.responses.tool import Mcp
18
from openai_harmony import Author, Message, Role, StreamState, TextContent
19

20
from vllm import envs
21
22
23
from vllm.entrypoints.chat_utils import (
    ChatTemplateContentFormatOption,
)
24
from vllm.entrypoints.constants import MCP_PREFIX
25
26
from vllm.entrypoints.mcp.tool import Tool
from vllm.entrypoints.mcp.tool_server import ToolServer
27
28
29
from vllm.entrypoints.openai.engine.protocol import (
    FunctionCall,
)
30
from vllm.entrypoints.openai.parser.harmony_utils import (
31
32
33
34
    get_encoding,
    get_streamable_parser_for_assistant,
    render_for_completion,
)
35
36
37
from vllm.entrypoints.openai.parser.responses_parser import (
    get_responses_parser_for_simple_context,
)
38
39
40
41
42
from vllm.entrypoints.openai.responses.protocol import (
    ResponseInputOutputItem,
    ResponseRawMessageAndToken,
    ResponsesRequest,
)
43
from vllm.entrypoints.openai.responses.utils import construct_tool_dicts
44
from vllm.outputs import RequestOutput
45
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
46
from vllm.tokenizers import TokenizerLike
47
from vllm.tool_parsers.abstract_tool_parser import ToolParser
48
from vllm.utils import random_uuid
49

50
51
52
if TYPE_CHECKING:
    from mcp.client import ClientSession

53
54
logger = logging.getLogger(__name__)

55
56
57
58
59
60
61
62
63
64
65
66
# This is currently needed as the tool type doesn't 1:1 match the
# tool namespace, which is what is used to look up the
# connection to the tool server
_TOOL_NAME_TO_TYPE_MAP = {
    "browser": "web_search_preview",
    "python": "code_interpreter",
    "container": "container",
}


def _map_tool_name_to_tool_type(tool_name: str) -> str:
    if tool_name not in _TOOL_NAME_TO_TYPE_MAP:
67
        available_tools = ", ".join(_TOOL_NAME_TO_TYPE_MAP.keys())
68
69
        raise ValueError(
            f"Built-in tool name '{tool_name}' not defined in mapping. "
70
71
            f"Available tools: {available_tools}"
        )
72
73
    return _TOOL_NAME_TO_TYPE_MAP[tool_name]

74

75
76
class TurnMetrics:
    """Tracks token and toolcall details for a single conversation turn."""
77

78
79
    def __init__(
        self,
80
81
82
83
84
        input_tokens: int = 0,
        output_tokens: int = 0,
        cached_input_tokens: int = 0,
        tool_output_tokens: int = 0,
    ) -> None:
85
86
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
87
88
        self.cached_input_tokens = cached_input_tokens
        self.tool_output_tokens = tool_output_tokens
89

90
    def reset(self) -> None:
91
92
93
        """Reset counters for a new turn."""
        self.input_tokens = 0
        self.output_tokens = 0
94
95
        self.cached_input_tokens = 0
        self.tool_output_tokens = 0
96

97
    def copy(self) -> "TurnMetrics":
98
        """Create a copy of this turn's token counts."""
99
100
101
102
103
104
        return TurnMetrics(
            self.input_tokens,
            self.output_tokens,
            self.cached_input_tokens,
            self.tool_output_tokens,
        )
105
106


107
108
class ConversationContext(ABC):
    @abstractmethod
109
110
111
112
113
    def append_output(self, output: RequestOutput) -> None:
        pass

    @abstractmethod
    def append_tool_output(self, output) -> None:
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        pass

    @abstractmethod
    async def call_tool(self) -> list[Message]:
        pass

    @abstractmethod
    def need_builtin_tool_call(self) -> bool:
        pass

    @abstractmethod
    def render_for_completion(self) -> list[int]:
        pass

128
    @abstractmethod
129
130
    async def init_tool_sessions(
        self,
131
        tool_server: ToolServer | None,
132
133
134
135
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
136
137
        pass

138
139
140
141
    @abstractmethod
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

142

143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def _create_json_parse_error_messages(
    last_msg: Message, e: json.JSONDecodeError
) -> list[Message]:
    """
    Creates an error message when json parse failed.
    """
    error_msg = (
        f"Error parsing tool arguments as JSON: {str(e)}. "
        "Please ensure the tool call arguments are valid JSON and try again."
    )
    content = TextContent(text=error_msg)
    author = Author(role=Role.TOOL, name=last_msg.recipient)
    return [
        Message(
            author=author,
            content=[content],
            recipient=Role.ASSISTANT,
            channel=last_msg.channel,
        )
    ]


165
class SimpleContext(ConversationContext):
166
167
    """This is a context that cannot handle MCP tool calls"""

168
169
    def __init__(self):
        self.last_output = None
170
171
172
173
174
175

        # Accumulated final output for streaming mode
        self._accumulated_text: str = ""
        self._accumulated_token_ids: list[int] = []
        self._accumulated_logprobs: list = []

176
177
178
179
180
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
        self.num_cached_tokens = 0
        # todo num_reasoning_tokens is not implemented yet.
        self.num_reasoning_tokens = 0
181
182
        # not implemented yet for SimpleContext
        self.all_turn_metrics = []
183

184
185
        self.input_messages: list[ResponseRawMessageAndToken] = []

186
187
    def append_output(self, output) -> None:
        self.last_output = output
188
189
190
191
192
        if not isinstance(output, RequestOutput):
            raise ValueError("SimpleContext only supports RequestOutput.")
        self.num_prompt_tokens = len(output.prompt_token_ids or [])
        self.num_cached_tokens = output.num_cached_tokens or 0
        self.num_output_tokens += len(output.outputs[0].token_ids or [])
193

194
195
196
197
198
199
200
        # Accumulate text, token_ids, and logprobs for streaming mode
        delta_output = output.outputs[0]
        self._accumulated_text += delta_output.text
        self._accumulated_token_ids.extend(delta_output.token_ids)
        if delta_output.logprobs is not None:
            self._accumulated_logprobs.extend(delta_output.logprobs)

201
202
203
204
205
206
207
208
209
        if len(self.input_messages) == 0:
            output_prompt = output.prompt or ""
            output_prompt_token_ids = output.prompt_token_ids or []
            self.input_messages.append(
                ResponseRawMessageAndToken(
                    message=output_prompt,
                    tokens=output_prompt_token_ids,
                )
            )
210
211
212
213
214
215
216
217
218
219
220

    @property
    def output_messages(self) -> list[ResponseRawMessageAndToken]:
        """Return consolidated output as a single message.

        In streaming mode, text and tokens are accumulated across many deltas.
        This property returns them as a single entry rather than one per delta.
        """
        if not self._accumulated_text and not self._accumulated_token_ids:
            return []
        return [
221
            ResponseRawMessageAndToken(
222
223
                message=self._accumulated_text,
                tokens=list(self._accumulated_token_ids),
224
            )
225
        ]
226

227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    @property
    def final_output(self) -> RequestOutput | None:
        """Return the final output, with complete text/token_ids/logprobs."""
        if self.last_output is not None and self.last_output.outputs:
            assert isinstance(self.last_output, RequestOutput)
            final_output = copy.copy(self.last_output)
            # copy inner item to avoid modify last_output
            final_output.outputs = [replace(item) for item in self.last_output.outputs]
            final_output.outputs[0].text = self._accumulated_text
            final_output.outputs[0].token_ids = tuple(self._accumulated_token_ids)
            if self._accumulated_logprobs:
                final_output.outputs[0].logprobs = self._accumulated_logprobs
            return final_output
        return self.last_output

242
243
244
    def append_tool_output(self, output) -> None:
        raise NotImplementedError("Should not be called.")

245
246
247
248
249
250
251
252
253
    def need_builtin_tool_call(self) -> bool:
        return False

    async def call_tool(self) -> list[Message]:
        raise NotImplementedError("Should not be called.")

    def render_for_completion(self) -> list[int]:
        raise NotImplementedError("Should not be called.")

254
255
    async def init_tool_sessions(
        self,
256
        tool_server: ToolServer | None,
257
258
259
260
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
261
262
        pass

263
264
265
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

266

267
268
269
270
271
class ParsableContext(ConversationContext):
    def __init__(
        self,
        *,
        response_messages: list[ResponseInputOutputItem],
272
        tokenizer: TokenizerLike,
273
        reasoning_parser_cls: Callable[[TokenizerLike], ReasoningParser] | None,
274
        request: ResponsesRequest,
275
276
277
278
        available_tools: list[str] | None,
        tool_parser_cls: Callable[[TokenizerLike], ToolParser] | None,
        chat_template: str | None,
        chat_template_content_format: ChatTemplateContentFormatOption,
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
    ):
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
        self.num_cached_tokens = 0
        self.num_reasoning_tokens = 0
        # not implemented yet for ParsableContext
        self.all_turn_metrics: list[TurnMetrics] = []

        if reasoning_parser_cls is None:
            raise ValueError("reasoning_parser_cls must be provided.")

        self.parser = get_responses_parser_for_simple_context(
            tokenizer=tokenizer,
            reasoning_parser_cls=reasoning_parser_cls,
            response_messages=response_messages,
            request=request,
295
            tool_parser_cls=tool_parser_cls,
296
        )
297
298
        self.tool_parser_cls = tool_parser_cls
        self.request = request
299

300
        self.available_tools = available_tools or []
301
302
303
304
        self._tool_sessions: dict[str, ClientSession | Tool] = {}
        self.called_tools: set[str] = set()

        self.tool_dicts = construct_tool_dicts(request.tools, request.tool_choice)
305
        self.chat_template = chat_template
306
        self.chat_template_content_format: Final = chat_template_content_format
307

308
309
        self.input_messages: list[ResponseRawMessageAndToken] = []
        self.output_messages: list[ResponseRawMessageAndToken] = []
310
        self._accumulated_token_ids: list[int] = []
311

312
313
314
315
316
    def append_output(self, output: RequestOutput) -> None:
        self.num_prompt_tokens = len(output.prompt_token_ids or [])
        self.num_cached_tokens = output.num_cached_tokens or 0
        self.num_output_tokens += len(output.outputs[0].token_ids or [])
        self.parser.process(output.outputs[0])
317
318
        output_token_ids = output.outputs[0].token_ids or []
        self._accumulated_token_ids.extend(output_token_ids)
319

320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
        # only store if enable_response_messages is True, save memory
        if self.request.enable_response_messages:
            output_prompt = output.prompt or ""
            output_prompt_token_ids = output.prompt_token_ids or []
            if len(self.input_messages) == 0:
                self.input_messages.append(
                    ResponseRawMessageAndToken(
                        message=output_prompt,
                        tokens=output_prompt_token_ids,
                    )
                )
            else:
                self.output_messages.append(
                    ResponseRawMessageAndToken(
                        message=output_prompt,
                        tokens=output_prompt_token_ids,
                    )
                )
            self.output_messages.append(
                ResponseRawMessageAndToken(
                    message=output.outputs[0].text,
                    tokens=output.outputs[0].token_ids,
                )
            )

345
    def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None:
346
        self.parser.response_messages.extend(output)
347
348

    def need_builtin_tool_call(self) -> bool:
349
350
        """Return true if the last message is a builtin tool call
        that the request has enabled."""
351
        last_message = self.parser.response_messages[-1]
352
353
354
355
356
357
358
359
        if last_message.type != "function_call":
            return False
        if last_message.name in ("code_interpreter", "python"):
            return "python" in self.available_tools
        if last_message.name == "web_search_preview":
            return "browser" in self.available_tools
        if last_message.name.startswith("container"):
            return "container" in self.available_tools
360
361
        return False

362
363
364
365
366
367
368
369
370
371
372
373
374
375
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: FunctionCall
    ) -> list[ResponseInputOutputItem]:
        self.called_tools.add("python")
        if isinstance(tool_session, Tool):
            return await tool_session.get_result_parsable_context(self)
        args = json.loads(last_msg.arguments)
        param = {
            "code": args["code"],
        }
        result = await tool_session.call_tool("python", param)
        result_str = result.content[0].text

        message = ResponseFunctionToolCallOutputItem(
376
            id=f"mcpo_{random_uuid()}",
377
378
379
380
381
382
383
384
            type="function_call_output",
            call_id=f"call_{random_uuid()}",
            output=result_str,
            status="completed",
        )

        return [message]

385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: FunctionCall
    ) -> list[ResponseInputOutputItem]:
        self.called_tools.add("browser")
        if isinstance(tool_session, Tool):
            return await tool_session.get_result_parsable_context(self)
        if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
            try:
                args = json.loads(last_msg.arguments)
            except json.JSONDecodeError as e:
                return _create_json_parse_error_messages(last_msg, e)
        else:
            args = json.loads(last_msg.arguments)
        result = await tool_session.call_tool("search", args)
        result_str = result.content[0].text

        message = ResponseFunctionToolCallOutputItem(
            id=f"fco_{random_uuid()}",
            type="function_call_output",
            call_id=f"call_{random_uuid()}",
            output=result_str,
            status="completed",
        )

        return [message]

    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
        """
        Call container tool. Expect this to be run in a stateful docker
        with command line terminal.
        The official container tool would at least
        expect the following format:
        - for tool name: exec
            - args:
                {
                    "cmd":List[str] "command to execute",
                    "workdir":optional[str] "current working directory",
                    "env":optional[object/dict] "environment variables",
                    "session_name":optional[str] "session name",
                    "timeout":optional[int] "timeout in seconds",
                    "user":optional[str] "user name",
                }
        """
        self.called_tools.add("container")
        if isinstance(tool_session, Tool):
            return await tool_session.get_result_parsable_context(self)
        # tool_name = last_msg.recipient.split(".")[1].split(" ")[0]
        if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
            try:
                args = json.loads(last_msg.arguments)
            except json.JSONDecodeError as e:
                return _create_json_parse_error_messages(last_msg, e)
        else:
            args = json.loads(last_msg.arguments)
        result = await tool_session.call_tool("exec", args)
        result_str = result.content[0].text

        message = ResponseFunctionToolCallOutputItem(
            id=f"fco_{random_uuid()}",
            type="function_call_output",
            call_id=f"call_{random_uuid()}",
            output=result_str,
            status="completed",
        )

        return [message]

454
    async def call_tool(self) -> list[ResponseInputOutputItem]:
455
456
457
        if not self.parser.response_messages:
            return []
        last_msg = self.parser.response_messages[-1]
458
459
460
        # change this to a mcp_ function call
        last_msg.id = f"{MCP_PREFIX}{random_uuid()}"
        self.parser.response_messages[-1] = last_msg
461
462
        if last_msg.name == "code_interpreter":
            return await self.call_python_tool(self._tool_sessions["python"], last_msg)
463
464
465
466
467
468
        elif last_msg.name == "web_search_preview":
            return await self.call_search_tool(self._tool_sessions["browser"], last_msg)
        elif last_msg.name.startswith("container"):
            return await self.call_container_tool(
                self._tool_sessions["container"], last_msg
            )
469
        return []
470
471
472
473
474
475
476
477
478
479
480

    def render_for_completion(self):
        raise NotImplementedError("Should not be called.")

    async def init_tool_sessions(
        self,
        tool_server: ToolServer | None,
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
481
482
483
484
485
486
487
488
489
490
491
492
493
494
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name in self._tool_sessions:
                    continue

                tool_type = _map_tool_name_to_tool_type(tool_name)
                headers = (
                    mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                )
                tool_session = await exit_stack.enter_async_context(
                    tool_server.new_session(tool_name, request_id, headers)
                )
                self._tool_sessions[tool_name] = tool_session
                exit_stack.push_async_exit(self.cleanup_session)
495
496
497

    async def cleanup_session(self, *args, **kwargs) -> None:
        """Can be used as coro to used in __aexit__"""
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512

        async def cleanup_tool_session(tool_session):
            if not isinstance(tool_session, Tool):
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
513
514


515
516
517
518
class HarmonyContext(ConversationContext):
    def __init__(
        self,
        messages: list,
519
        available_tools: list[str],
520
521
    ):
        self._messages = messages
522
        self.finish_reason: str | None = None
523
        self.available_tools = available_tools
524
        self._tool_sessions: dict[str, ClientSession | Tool] = {}
525
        self.called_tools: set[str] = set()
526
527
528
529
530

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
531
        self.num_cached_tokens = 0
532
        self.num_reasoning_tokens = 0
533
        self.num_tool_output_tokens = 0
534

535
        # Turn tracking - replaces multiple individual tracking variables
536
537
538
        self.current_turn_metrics = TurnMetrics()
        # Track metrics for all turns
        self.all_turn_metrics: list[TurnMetrics] = []
539
540
        self.is_first_turn = True
        self.first_tok_of_message = True  # For streaming support
541

542
    def _update_num_reasoning_tokens(self):
543
544
545
546
547
548
        channel = self.parser.current_channel
        if channel == "analysis":
            self.num_reasoning_tokens += 1
        elif channel == "commentary" and self.parser.current_recipient is not None:
            # Tool interactions (python/browser/container) are hidden.
            # Preambles (recipient=None) are visible user text.
549
            self.num_reasoning_tokens += 1
550

551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
    def append_output(self, output: RequestOutput) -> None:
        output_token_ids = output.outputs[0].token_ids
        self.parser = get_streamable_parser_for_assistant()
        for token_id in output_token_ids:
            self.parser.process(token_id)
            # Check if the current token is part of reasoning content
            self._update_num_reasoning_tokens()
        self._update_prefill_token_usage(output)
        self._update_decode_token_usage(output)
        # Append current turn to all turn list for next turn's calculations
        self.all_turn_metrics.append(self.current_turn_metrics.copy())
        self.current_turn_metrics.reset()
        # append_output is called only once before tool calling
        # in non-streaming case
        # so we can append all the parser messages to _messages
        output_msgs = self.parser.messages
        # The responses finish reason is set in the last message
        self.finish_reason = output.outputs[0].finish_reason
        self._messages.extend(output_msgs)

    def append_tool_output(self, output: list[Message]) -> None:
        output_msgs = output
573
574
        self._messages.extend(output_msgs)

575
576
    def _update_prefill_token_usage(self, output: RequestOutput) -> None:
        """Update token usage statistics for the prefill phase of generation.
577

578
579
580
581
582
        The prefill phase processes the input prompt tokens. This method:
        1. Counts the prompt tokens for this turn
        2. Calculates tool output tokens for multi-turn conversations
        3. Updates cached token counts
        4. Tracks state for next turn calculations
583

584
        Tool output tokens are calculated as:
585
        current_prompt_tokens - last_turn_prompt_tokens -
586
587
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
588

589
590
591
592
593
594
595
        Args:
            output: The RequestOutput containing prompt token information
        """
        if output.prompt_token_ids is not None:
            this_turn_input_tokens = len(output.prompt_token_ids)
        else:
            this_turn_input_tokens = 0
596
            logger.error("RequestOutput appended contains no prompt_token_ids.")
597
598

        # Update current turn input tokens
599
        self.current_turn_metrics.input_tokens = this_turn_input_tokens
600
601
602
603
604
605
        self.num_prompt_tokens += this_turn_input_tokens

        # Calculate tool tokens (except on first turn)
        if self.is_first_turn:
            self.is_first_turn = False
        else:
606
            previous_turn = self.all_turn_metrics[-1]
607
608
609
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
610
            this_turn_tool_tokens = (
611
612
613
                self.current_turn_metrics.input_tokens
                - previous_turn.input_tokens
                - previous_turn.output_tokens
614
            )
615
616
617
618
619
620
621
622

            # Handle negative tool token counts (shouldn't happen in normal
            # cases)
            if this_turn_tool_tokens < 0:
                logger.error(
                    "Negative tool output tokens calculated: %d "
                    "(current_input=%d, previous_input=%d, "
                    "previous_output=%d). Setting to 0.",
623
                    this_turn_tool_tokens,
624
625
626
                    self.current_turn_metrics.input_tokens,
                    previous_turn.input_tokens,
                    previous_turn.output_tokens,
627
                )
628
629
630
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens
631
            self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
632
633

        # Update cached tokens
634
635
636
637
        num_cached_token = output.num_cached_tokens
        if num_cached_token is not None:
            self.num_cached_tokens += num_cached_token
            self.current_turn_metrics.cached_input_tokens = num_cached_token
638
639
640

    def _update_decode_token_usage(self, output: RequestOutput) -> int:
        """Update token usage statistics for the decode phase of generation.
641

642
643
644
645
        The decode phase processes the generated output tokens. This method:
        1. Counts output tokens from all completion outputs
        2. Updates the total output token count
        3. Tracks tokens generated in the current turn
646

647
648
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
649

650
651
        Args:
            output: The RequestOutput containing generated token information
652

653
654
655
656
657
658
659
660
661
        Returns:
            int: Number of output tokens processed in this call
        """
        updated_output_token_count = 0
        if output.outputs:
            for completion_output in output.outputs:
                # only keep last round
                updated_output_token_count += len(completion_output.token_ids)
            self.num_output_tokens += updated_output_token_count
662
            self.current_turn_metrics.output_tokens += updated_output_token_count
663
664
        return updated_output_token_count

665
666
667
668
669
670
671
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
672
673
674
675
676
677
678
679
680
        if recipient is None:
            return False
        if recipient.startswith("browser."):
            return "browser" in self.available_tools
        if recipient.startswith("python"):
            return "python" in self.available_tools
        if recipient.startswith("container."):
            return "container" in self.available_tools
        return False
681
682
683
684
685
686
687
688
689

    async def call_tool(self) -> list[Message]:
        if not self.messages:
            return []
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
        if recipient is not None:
            if recipient.startswith("browser."):
                return await self.call_search_tool(
690
691
                    self._tool_sessions["browser"], last_msg
                )
692
693
            elif recipient.startswith("python"):
                return await self.call_python_tool(
694
695
                    self._tool_sessions["python"], last_msg
                )
696
697
            elif recipient.startswith("container."):
                return await self.call_container_tool(
698
699
                    self._tool_sessions["container"], last_msg
                )
700
701
702
703
704
        raise ValueError("No tool call found")

    def render_for_completion(self) -> list[int]:
        return render_for_completion(self.messages)

705
706
707
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
708
        self.called_tools.add("browser")
709
710
711
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1]
712
713
714
715
716
717
718
        if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
            try:
                args = json.loads(last_msg.content[0].text)
            except json.JSONDecodeError as e:
                return _create_json_parse_error_messages(last_msg, e)
        else:
            args = json.loads(last_msg.content[0].text)
719
720
721
722
723
        result = await tool_session.call_tool(tool_name, args)
        result_str = result.content[0].text
        content = TextContent(text=result_str)
        author = Author(role=Role.TOOL, name=last_msg.recipient)
        return [
724
725
726
727
728
729
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
730
731
        ]

732
733
734
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
735
        self.called_tools.add("python")
736
737
738
739
740
741
742
743
744
745
746
747
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        param = {
            "code": last_msg.content[0].text,
        }
        result = await tool_session.call_tool("python", param)
        result_str = result.content[0].text

        content = TextContent(text=result_str)
        author = Author(role=Role.TOOL, name="python")

        return [
748
749
750
751
752
753
            Message(
                author=author,
                content=[content],
                channel=last_msg.channel,
                recipient=Role.ASSISTANT,
            )
754
        ]
755

756
757
    async def init_tool_sessions(
        self,
758
        tool_server: ToolServer | None,
759
760
761
762
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
763
764
765
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
766
                    tool_type = _map_tool_name_to_tool_type(tool_name)
767
768
769
                    headers = (
                        mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                    )
770
                    tool_session = await exit_stack.enter_async_context(
771
772
                        tool_server.new_session(tool_name, request_id, headers)
                    )
773
774
775
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

776
777
778
    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
779
        """
780
781
782
783
784
785
786
787
788
789
790
791
792
793
        Call container tool. Expect this to be run in a stateful docker
        with command line terminal.
        The official container tool would at least
        expect the following format:
        - for tool name: exec
            - args:
                {
                    "cmd":List[str] "command to execute",
                    "workdir":optional[str] "current working directory",
                    "env":optional[object/dict] "environment variables",
                    "session_name":optional[str] "session name",
                    "timeout":optional[int] "timeout in seconds",
                    "user":optional[str] "user name",
                }
794
795
796
797
798
        """
        self.called_tools.add("container")
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1].split(" ")[0]
799
800
801
802
803
804
805
        if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
            try:
                args = json.loads(last_msg.content[0].text)
            except json.JSONDecodeError as e:
                return _create_json_parse_error_messages(last_msg, e)
        else:
            args = json.loads(last_msg.content[0].text)
806
807
808
809
810
        result = await tool_session.call_tool(tool_name, args)
        result_str = result.content[0].text
        content = TextContent(text=result_str)
        author = Author(role=Role.TOOL, name=last_msg.recipient)
        return [
811
812
813
814
815
816
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
817
818
819
820
821
822
823
        ]

    async def cleanup_session(self, *args, **kwargs) -> None:
        """Can be used as coro to used in __aexit__"""

        async def cleanup_tool_session(tool_session):
            if not isinstance(tool_session, Tool):
824
825
826
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
827
828
829
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

830
831
832
833
834
835
        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
836

837
838
839
840
841
842
843
844
845

class StreamingHarmonyContext(HarmonyContext):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.last_output = None

        self.parser = get_streamable_parser_for_assistant()
        self.encoding = get_encoding()
        self.last_tok = None
846
        self.first_tok_of_message = True
847
        self.last_content_delta = None
848
849
850

    @property
    def messages(self) -> list:
851
        return self._messages
852

853
854
855
    def append_output(self, output: RequestOutput) -> None:
        # append_output is called for each output token in streaming case,
        # so we only want to add the prompt tokens once for each message.
856
        self.last_content_delta = None
857
858
859
860
861
862
863
        if self.first_tok_of_message:
            self._update_prefill_token_usage(output)
        # Reset self.first_tok_of_message if needed:
        # if the current token is the last one of the current message
        # (finished=True), then the next token processed will mark the
        # beginning of a new message
        self.first_tok_of_message = output.finished
864
        last_delta_text = ""
865
866
        for tok in output.outputs[0].token_ids:
            self.parser.process(tok)
867
868
869
            last_delta_text += self.parser.last_content_delta or ""
        if last_delta_text:
            self.last_content_delta = last_delta_text
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
        self._update_decode_token_usage(output)

        # For streaming, update previous turn when message is complete
        if output.finished:
            self.all_turn_metrics.append(self.current_turn_metrics.copy())
            self.current_turn_metrics.reset()
        # Check if the current token is part of reasoning content
        self._update_num_reasoning_tokens()
        self.last_tok = tok
        if len(self._messages) - self.num_init_messages < len(self.parser.messages):
            self._messages.extend(
                self.parser.messages[len(self._messages) - self.num_init_messages :]
            )

    def append_tool_output(self, output: list[Message]) -> None:
        # Handle the case of tool output in direct message format
        assert len(output) == 1, "Tool output should be a single message"
        msg = output[0]
        # Sometimes the recipient is not set for tool messages,
        # so we set it to "assistant"
        if msg.author.role == Role.TOOL and msg.recipient is None:
            msg.recipient = "assistant"
        toks = self.encoding.render(msg)
        for tok in toks:
            self.parser.process(tok)
        self.last_tok = toks[-1]
        # TODO: add tool_output messages to self._messages
897
898
899
900
901

    def is_expecting_start(self) -> bool:
        return self.parser.state == StreamState.EXPECT_START

    def is_assistant_action_turn(self) -> bool:
902
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
903
904
905

    def render_for_completion(self) -> list[int]:
        # now this list of tokens as next turn's starting tokens
906
        # `<|start|>assistant`,
907
908
909
910
911
912
913
914
915
916
917
918
        # we need to process them in parser.
        rendered_tokens = super().render_for_completion()

        last_n = -1
        to_process = []
        while rendered_tokens[last_n] != self.last_tok:
            to_process.append(rendered_tokens[last_n])
            last_n -= 1
        for tok in reversed(to_process):
            self.parser.process(tok)

        return rendered_tokens