context.py 24.1 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 json
6
7
import logging
from abc import ABC, abstractmethod
8
from collections.abc import Callable
9
from contextlib import AsyncExitStack
10
from typing import TYPE_CHECKING, Union
11

12
from openai.types.responses.tool import Mcp
13
from openai_harmony import Author, Message, Role, StreamState, TextContent
14

15
from vllm import envs
16
from vllm.entrypoints.harmony_utils import (
17
18
19
20
    get_encoding,
    get_streamable_parser_for_assistant,
    render_for_completion,
)
21
22
23
24
25
from vllm.entrypoints.openai.parser.responses_parser import (
    get_responses_parser_for_simple_context,
)
from vllm.entrypoints.openai.protocol import (
    ResponseInputOutputItem,
26
    ResponseRawMessageAndToken,
27
28
29
    ResponsesRequest,
)
from vllm.entrypoints.responses_utils import construct_tool_dicts
30
from vllm.entrypoints.tool import Tool
31
from vllm.entrypoints.tool_server import ToolServer
32
from vllm.outputs import RequestOutput
33
34
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
from vllm.transformers_utils.tokenizer import AnyTokenizer
35

36
37
38
if TYPE_CHECKING:
    from mcp.client import ClientSession

39
40
logger = logging.getLogger(__name__)

41
42
43
44
45
46
47
48
49
50
51
52
# 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:
53
        available_tools = ", ".join(_TOOL_NAME_TO_TYPE_MAP.keys())
54
55
        raise ValueError(
            f"Built-in tool name '{tool_name}' not defined in mapping. "
56
57
            f"Available tools: {available_tools}"
        )
58
59
    return _TOOL_NAME_TO_TYPE_MAP[tool_name]

60

61
62
class TurnMetrics:
    """Tracks token and toolcall details for a single conversation turn."""
63

64
65
66
67
68
69
70
    def __init__(
        self,
        input_tokens=0,
        output_tokens=0,
        cached_input_tokens=0,
        tool_output_tokens=0,
    ):
71
72
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
73
74
        self.cached_input_tokens = cached_input_tokens
        self.tool_output_tokens = tool_output_tokens
75
76
77
78
79

    def reset(self):
        """Reset counters for a new turn."""
        self.input_tokens = 0
        self.output_tokens = 0
80
81
        self.cached_input_tokens = 0
        self.tool_output_tokens = 0
82
83
84

    def copy(self):
        """Create a copy of this turn's token counts."""
85
86
87
88
89
90
        return TurnMetrics(
            self.input_tokens,
            self.output_tokens,
            self.cached_input_tokens,
            self.tool_output_tokens,
        )
91
92


93
94
class ConversationContext(ABC):
    @abstractmethod
95
96
97
98
99
    def append_output(self, output: RequestOutput) -> None:
        pass

    @abstractmethod
    def append_tool_output(self, output) -> None:
100
101
102
103
104
105
106
107
108
109
110
111
112
113
        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

114
    @abstractmethod
115
116
    async def init_tool_sessions(
        self,
117
        tool_server: ToolServer | None,
118
119
120
121
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
122
123
        pass

124
125
126
127
    @abstractmethod
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

128

129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
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,
        )
    ]


151
class SimpleContext(ConversationContext):
152
153
    """This is a context that cannot handle MCP tool calls"""

154
155
    def __init__(self):
        self.last_output = None
156
157
158
159
160
        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
161
162
        # not implemented yet for SimpleContext
        self.all_turn_metrics = []
163

164
165
166
        self.input_messages: list[ResponseRawMessageAndToken] = []
        self.output_messages: list[ResponseRawMessageAndToken] = []

167
168
    def append_output(self, output) -> None:
        self.last_output = output
169
170
171
172
173
        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 [])
174

175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        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,
                )
            )
        self.output_messages.append(
            ResponseRawMessageAndToken(
                message=output.outputs[0].text,
                tokens=output.outputs[0].token_ids,
            )
        )

191
192
193
    def append_tool_output(self, output) -> None:
        raise NotImplementedError("Should not be called.")

194
195
196
197
198
199
200
201
202
    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.")

203
204
    async def init_tool_sessions(
        self,
205
        tool_server: ToolServer | None,
206
207
208
209
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
210
211
        pass

212
213
214
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

215

216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
class ParsableContext(ConversationContext):
    def __init__(
        self,
        *,
        response_messages: list[ResponseInputOutputItem],
        tokenizer: AnyTokenizer,
        reasoning_parser_cls: Callable[[AnyTokenizer], ReasoningParser] | None,
        request: ResponsesRequest,
    ):
        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,
        )

        self._tool_sessions: dict[str, ClientSession | Tool] = {}
        self.called_tools: set[str] = set()

        self.tool_dicts = construct_tool_dicts(request.tools, request.tool_choice)

    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])

    def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None:
        raise NotImplementedError("Should not be called.")

    def need_builtin_tool_call(self) -> bool:
        """Return true if the last message is a MCP tool call"""
        return False

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

    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],
    ):
        pass

    async def cleanup_session(self, *args, **kwargs) -> None:
        """Can be used as coro to used in __aexit__"""
        raise NotImplementedError("Should not be called.")


281
282
283
284
class HarmonyContext(ConversationContext):
    def __init__(
        self,
        messages: list,
285
        available_tools: list[str],
286
287
    ):
        self._messages = messages
288
        self.finish_reason: str | None = None
289
        self.available_tools = available_tools
290
        self._tool_sessions: dict[str, ClientSession | Tool] = {}
291
        self.called_tools: set[str] = set()
292
293
294
295
296

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
297
        self.num_cached_tokens = 0
298
        self.num_reasoning_tokens = 0
299
        self.num_tool_output_tokens = 0
300

301
        # Turn tracking - replaces multiple individual tracking variables
302
303
304
        self.current_turn_metrics = TurnMetrics()
        # Track metrics for all turns
        self.all_turn_metrics: list[TurnMetrics] = []
305
306
        self.is_first_turn = True
        self.first_tok_of_message = True  # For streaming support
307

308
309
310
311
    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
312

313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
    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
335
336
        self._messages.extend(output_msgs)

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

340
341
342
343
344
        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
345

346
        Tool output tokens are calculated as:
347
        current_prompt_tokens - last_turn_prompt_tokens -
348
349
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
350

351
352
353
354
355
356
357
        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
358
            logger.error("RequestOutput appended contains no prompt_token_ids.")
359
360

        # Update current turn input tokens
361
        self.current_turn_metrics.input_tokens = this_turn_input_tokens
362
363
364
365
366
367
        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:
368
            previous_turn = self.all_turn_metrics[-1]
369
370
371
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
372
            this_turn_tool_tokens = (
373
374
375
                self.current_turn_metrics.input_tokens
                - previous_turn.input_tokens
                - previous_turn.output_tokens
376
            )
377
378
379
380
381
382
383
384

            # 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.",
385
                    this_turn_tool_tokens,
386
387
388
                    self.current_turn_metrics.input_tokens,
                    previous_turn.input_tokens,
                    previous_turn.output_tokens,
389
                )
390
391
392
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens
393
            self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
394
395

        # Update cached tokens
396
397
398
399
        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
400
401
402

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

404
405
406
407
        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
408

409
410
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
411

412
413
        Args:
            output: The RequestOutput containing generated token information
414

415
416
417
418
419
420
421
422
423
        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
424
            self.current_turn_metrics.output_tokens += updated_output_token_count
425
426
        return updated_output_token_count

427
428
429
430
431
432
433
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
434
435
436
437
438
        return recipient is not None and (
            recipient.startswith("browser.")
            or recipient.startswith("python")
            or recipient.startswith("container.")
        )
439
440
441
442
443
444
445
446
447

    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(
448
449
                    self._tool_sessions["browser"], last_msg
                )
450
451
            elif recipient.startswith("python"):
                return await self.call_python_tool(
452
453
                    self._tool_sessions["python"], last_msg
                )
454
455
            elif recipient.startswith("container."):
                return await self.call_container_tool(
456
457
                    self._tool_sessions["container"], last_msg
                )
458
459
460
461
462
        raise ValueError("No tool call found")

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

463
464
465
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
466
        self.called_tools.add("browser")
467
468
469
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1]
470
471
472
473
474
475
476
        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)
477
478
479
480
481
        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 [
482
483
484
485
486
487
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
488
489
        ]

490
491
492
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
493
        self.called_tools.add("python")
494
495
496
497
498
499
500
501
502
503
504
505
        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 [
506
507
508
509
510
511
            Message(
                author=author,
                content=[content],
                channel=last_msg.channel,
                recipient=Role.ASSISTANT,
            )
512
        ]
513

514
515
    async def init_tool_sessions(
        self,
516
        tool_server: ToolServer | None,
517
518
519
520
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
521
522
523
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
524
                    tool_type = _map_tool_name_to_tool_type(tool_name)
525
526
527
                    headers = (
                        mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                    )
528
                    tool_session = await exit_stack.enter_async_context(
529
530
                        tool_server.new_session(tool_name, request_id, headers)
                    )
531
532
533
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

534
535
536
    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
537
        """
538
539
540
541
542
543
544
545
546
547
548
549
550
551
        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",
                }
552
553
554
555
556
        """
        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]
557
558
559
560
561
562
563
        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)
564
565
566
567
568
        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 [
569
570
571
572
573
574
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
575
576
577
578
579
580
581
        ]

    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):
582
583
584
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
585
586
587
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

588
589
590
591
592
593
        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
594

595
596
597
598
599
600
601
602
603

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
604
        self.first_tok_of_message = True
605
606
607

    @property
    def messages(self) -> list:
608
        return self._messages
609

610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
    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.
        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
        for tok in output.outputs[0].token_ids:
            self.parser.process(tok)
        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
649
650
651
652
653

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

    def is_assistant_action_turn(self) -> bool:
654
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
655
656
657

    def render_for_completion(self) -> list[int]:
        # now this list of tokens as next turn's starting tokens
658
        # `<|start|>assistant`,
659
660
661
662
663
664
665
666
667
668
669
670
        # 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