context.py 23.3 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
26
27
28
from vllm.entrypoints.openai.parser.responses_parser import (
    get_responses_parser_for_simple_context,
)
from vllm.entrypoints.openai.protocol import (
    ResponseInputOutputItem,
    ResponsesRequest,
)
from vllm.entrypoints.responses_utils import construct_tool_dicts
29
from vllm.entrypoints.tool import Tool
30
from vllm.entrypoints.tool_server import ToolServer
31
from vllm.outputs import RequestOutput
32
33
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
from vllm.transformers_utils.tokenizer import AnyTokenizer
34

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

38
39
logger = logging.getLogger(__name__)

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

59

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

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

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

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


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

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

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

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

127

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


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

    def append_output(self, output) -> None:
        self.last_output = output
163
164
165
166
167
        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 [])
168

169
170
171
    def append_tool_output(self, output) -> None:
        raise NotImplementedError("Should not be called.")

172
173
174
175
176
177
178
179
180
    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.")

181
182
    async def init_tool_sessions(
        self,
183
        tool_server: ToolServer | None,
184
185
186
187
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
188
189
        pass

190
191
192
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

193

194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
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
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.")


259
260
261
262
class HarmonyContext(ConversationContext):
    def __init__(
        self,
        messages: list,
263
        available_tools: list[str],
264
265
    ):
        self._messages = messages
266
        self.finish_reason: str | None = None
267
        self.available_tools = available_tools
268
        self._tool_sessions: dict[str, ClientSession | Tool] = {}
269
        self.called_tools: set[str] = set()
270
271
272
273
274

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
275
        self.num_cached_tokens = 0
276
        self.num_reasoning_tokens = 0
277
        self.num_tool_output_tokens = 0
278

279
        # Turn tracking - replaces multiple individual tracking variables
280
281
282
        self.current_turn_metrics = TurnMetrics()
        # Track metrics for all turns
        self.all_turn_metrics: list[TurnMetrics] = []
283
284
        self.is_first_turn = True
        self.first_tok_of_message = True  # For streaming support
285

286
287
288
289
    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
290

291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
    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
313
314
        self._messages.extend(output_msgs)

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

318
319
320
321
322
        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
323

324
        Tool output tokens are calculated as:
325
        current_prompt_tokens - last_turn_prompt_tokens -
326
327
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
328

329
330
331
332
333
334
335
        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
336
            logger.error("RequestOutput appended contains no prompt_token_ids.")
337
338

        # Update current turn input tokens
339
        self.current_turn_metrics.input_tokens = this_turn_input_tokens
340
341
342
343
344
345
        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:
346
            previous_turn = self.all_turn_metrics[-1]
347
348
349
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
350
            this_turn_tool_tokens = (
351
352
353
                self.current_turn_metrics.input_tokens
                - previous_turn.input_tokens
                - previous_turn.output_tokens
354
            )
355
356
357
358
359
360
361
362

            # 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.",
363
                    this_turn_tool_tokens,
364
365
366
                    self.current_turn_metrics.input_tokens,
                    previous_turn.input_tokens,
                    previous_turn.output_tokens,
367
                )
368
369
370
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens
371
            self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
372
373

        # Update cached tokens
374
375
376
377
        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
378
379
380

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

382
383
384
385
        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
386

387
388
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
389

390
391
        Args:
            output: The RequestOutput containing generated token information
392

393
394
395
396
397
398
399
400
401
        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
402
            self.current_turn_metrics.output_tokens += updated_output_token_count
403
404
        return updated_output_token_count

405
406
407
408
409
410
411
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
412
413
414
415
416
        return recipient is not None and (
            recipient.startswith("browser.")
            or recipient.startswith("python")
            or recipient.startswith("container.")
        )
417
418
419
420
421
422
423
424
425

    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(
426
427
                    self._tool_sessions["browser"], last_msg
                )
428
429
            elif recipient.startswith("python"):
                return await self.call_python_tool(
430
431
                    self._tool_sessions["python"], last_msg
                )
432
433
            elif recipient.startswith("container."):
                return await self.call_container_tool(
434
435
                    self._tool_sessions["container"], last_msg
                )
436
437
438
439
440
        raise ValueError("No tool call found")

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

441
442
443
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
444
        self.called_tools.add("browser")
445
446
447
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1]
448
449
450
451
452
453
454
        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)
455
456
457
458
459
        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 [
460
461
462
463
464
465
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
466
467
        ]

468
469
470
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
471
        self.called_tools.add("python")
472
473
474
475
476
477
478
479
480
481
482
483
        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 [
484
485
486
487
488
489
            Message(
                author=author,
                content=[content],
                channel=last_msg.channel,
                recipient=Role.ASSISTANT,
            )
490
        ]
491

492
493
    async def init_tool_sessions(
        self,
494
        tool_server: ToolServer | None,
495
496
497
498
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
499
500
501
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
502
                    tool_type = _map_tool_name_to_tool_type(tool_name)
503
504
505
                    headers = (
                        mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                    )
506
                    tool_session = await exit_stack.enter_async_context(
507
508
                        tool_server.new_session(tool_name, request_id, headers)
                    )
509
510
511
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

512
513
514
    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
515
        """
516
517
518
519
520
521
522
523
524
525
526
527
528
529
        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",
                }
530
531
532
533
534
        """
        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]
535
536
537
538
539
540
541
        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)
542
543
544
545
546
        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 [
547
548
549
550
551
552
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
553
554
555
556
557
558
559
        ]

    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):
560
561
562
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
563
564
565
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

566
567
568
569
570
571
        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
572

573
574
575
576
577
578
579
580
581

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
582
        self.first_tok_of_message = True
583
584
585

    @property
    def messages(self) -> list:
586
        return self._messages
587

588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
    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
627
628
629
630
631

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

    def is_assistant_action_turn(self) -> bool:
632
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
633
634
635

    def render_for_completion(self) -> list[int]:
        # now this list of tokens as next turn's starting tokens
636
        # `<|start|>assistant`,
637
638
639
640
641
642
643
644
645
646
647
648
        # 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