context.py 33.8 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, 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
295
    ):
        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
        # 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,
296
            tool_parser_cls=tool_parser_cls,
297
        )
298
299
        self.tool_parser_cls = tool_parser_cls
        self.request = request
300

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

        self.tool_dicts = construct_tool_dicts(request.tools, request.tool_choice)
306
307
        self.chat_template = chat_template
        self.chat_template_content_format = chat_template_content_format
308

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

312
313
314
315
316
317
    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])

318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
        # 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,
                )
            )

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

    def need_builtin_tool_call(self) -> bool:
        """Return true if the last message is a MCP tool call"""
348
        last_message = self.parser.response_messages[-1]
349
350
351
352
353
354
355
356
        # TODO(qandrew): figure out which tools are MCP tools
        if last_message.type == "function_call":  # noqa: SIM102
            if last_message.name in (
                "code_interpreter",
                "python",
                "web_search_preview",
            ) or last_message.name.startswith("container"):
                return True
357

358
359
        return False

360
361
362
363
364
365
366
367
368
369
370
371
372
373
    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(
374
            id=f"mcpo_{random_uuid()}",
375
376
377
378
379
380
381
382
            type="function_call_output",
            call_id=f"call_{random_uuid()}",
            output=result_str,
            status="completed",
        )

        return [message]

383
384
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
    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]

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

    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],
    ):
479
480
481
482
483
484
485
486
487
488
489
490
491
492
        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)
493
494
495

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

        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
            )
        )
511
512


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

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

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

540
541
542
543
    def _update_num_reasoning_tokens(self):
        # Count all analysis and commentary channels as reasoning tokens
        if self.parser.current_channel in {"analysis", "commentary"}:
            self.num_reasoning_tokens += 1
544

545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
    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
567
568
        self._messages.extend(output_msgs)

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

572
573
574
575
576
        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
577

578
        Tool output tokens are calculated as:
579
        current_prompt_tokens - last_turn_prompt_tokens -
580
581
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
582

583
584
585
586
587
588
589
        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
590
            logger.error("RequestOutput appended contains no prompt_token_ids.")
591
592

        # Update current turn input tokens
593
        self.current_turn_metrics.input_tokens = this_turn_input_tokens
594
595
596
597
598
599
        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:
600
            previous_turn = self.all_turn_metrics[-1]
601
602
603
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
604
            this_turn_tool_tokens = (
605
606
607
                self.current_turn_metrics.input_tokens
                - previous_turn.input_tokens
                - previous_turn.output_tokens
608
            )
609
610
611
612
613
614
615
616

            # 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.",
617
                    this_turn_tool_tokens,
618
619
620
                    self.current_turn_metrics.input_tokens,
                    previous_turn.input_tokens,
                    previous_turn.output_tokens,
621
                )
622
623
624
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens
625
            self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
626
627

        # Update cached tokens
628
629
630
631
        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
632
633
634

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

636
637
638
639
        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
640

641
642
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
643

644
645
        Args:
            output: The RequestOutput containing generated token information
646

647
648
649
650
651
652
653
654
655
        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
656
            self.current_turn_metrics.output_tokens += updated_output_token_count
657
658
        return updated_output_token_count

659
660
661
662
663
664
665
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
666
667
668
669
670
        return recipient is not None and (
            recipient.startswith("browser.")
            or recipient.startswith("python")
            or recipient.startswith("container.")
        )
671
672
673
674
675
676
677
678
679

    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(
680
681
                    self._tool_sessions["browser"], last_msg
                )
682
683
            elif recipient.startswith("python"):
                return await self.call_python_tool(
684
685
                    self._tool_sessions["python"], last_msg
                )
686
687
            elif recipient.startswith("container."):
                return await self.call_container_tool(
688
689
                    self._tool_sessions["container"], last_msg
                )
690
691
692
693
694
        raise ValueError("No tool call found")

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

695
696
697
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
698
        self.called_tools.add("browser")
699
700
701
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1]
702
703
704
705
706
707
708
        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)
709
710
711
712
713
        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 [
714
715
716
717
718
719
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
720
721
        ]

722
723
724
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
725
        self.called_tools.add("python")
726
727
728
729
730
731
732
733
734
735
736
737
        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 [
738
739
740
741
742
743
            Message(
                author=author,
                content=[content],
                channel=last_msg.channel,
                recipient=Role.ASSISTANT,
            )
744
        ]
745

746
747
    async def init_tool_sessions(
        self,
748
        tool_server: ToolServer | None,
749
750
751
752
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
753
754
755
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
756
                    tool_type = _map_tool_name_to_tool_type(tool_name)
757
758
759
                    headers = (
                        mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                    )
760
                    tool_session = await exit_stack.enter_async_context(
761
762
                        tool_server.new_session(tool_name, request_id, headers)
                    )
763
764
765
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

766
767
768
    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
769
        """
770
771
772
773
774
775
776
777
778
779
780
781
782
783
        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",
                }
784
785
786
787
788
        """
        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]
789
790
791
792
793
794
795
        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)
796
797
798
799
800
        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 [
801
802
803
804
805
806
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
807
808
809
810
811
812
813
        ]

    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):
814
815
816
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
817
818
819
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

820
821
822
823
824
825
        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
826

827
828
829
830
831
832
833
834
835

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
836
        self.first_tok_of_message = True
837
        self.last_content_delta = None
838
839
840

    @property
    def messages(self) -> list:
841
        return self._messages
842

843
844
845
    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.
846
        self.last_content_delta = None
847
848
849
850
851
852
853
        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
854
        last_delta_text = ""
855
856
        for tok in output.outputs[0].token_ids:
            self.parser.process(tok)
857
858
859
            last_delta_text += self.parser.last_content_delta or ""
        if last_delta_text:
            self.last_content_delta = last_delta_text
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
        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
887
888
889
890
891

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

    def is_assistant_action_turn(self) -> bool:
892
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
893
894
895

    def render_for_completion(self) -> list[int]:
        # now this list of tokens as next turn's starting tokens
896
        # `<|start|>assistant`,
897
898
899
900
901
902
903
904
905
906
907
908
        # 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