context.py 19.5 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 contextlib import AsyncExitStack
9
from typing import TYPE_CHECKING, Union
10

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

from vllm.entrypoints.harmony_utils import (
15
16
17
18
    get_encoding,
    get_streamable_parser_for_assistant,
    render_for_completion,
)
19
from vllm.entrypoints.tool import Tool
20
from vllm.entrypoints.tool_server import ToolServer
21
22
from vllm.outputs import RequestOutput

23
24
25
if TYPE_CHECKING:
    from mcp.client import ClientSession

26
27
logger = logging.getLogger(__name__)

28
29
30
31
32
33
34
35
36
37
38
39
# 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:
40
        available_tools = ", ".join(_TOOL_NAME_TO_TYPE_MAP.keys())
41
42
        raise ValueError(
            f"Built-in tool name '{tool_name}' not defined in mapping. "
43
44
            f"Available tools: {available_tools}"
        )
45
46
    return _TOOL_NAME_TO_TYPE_MAP[tool_name]

47

48
49
class TurnMetrics:
    """Tracks token and toolcall details for a single conversation turn."""
50

51
52
53
54
55
56
57
    def __init__(
        self,
        input_tokens=0,
        output_tokens=0,
        cached_input_tokens=0,
        tool_output_tokens=0,
    ):
58
59
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
60
61
        self.cached_input_tokens = cached_input_tokens
        self.tool_output_tokens = tool_output_tokens
62
63
64
65
66

    def reset(self):
        """Reset counters for a new turn."""
        self.input_tokens = 0
        self.output_tokens = 0
67
68
        self.cached_input_tokens = 0
        self.tool_output_tokens = 0
69
70
71

    def copy(self):
        """Create a copy of this turn's token counts."""
72
73
74
75
76
77
        return TurnMetrics(
            self.input_tokens,
            self.output_tokens,
            self.cached_input_tokens,
            self.tool_output_tokens,
        )
78
79


80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class ConversationContext(ABC):
    @abstractmethod
    def append_output(self, output) -> None:
        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

97
    @abstractmethod
98
99
    async def init_tool_sessions(
        self,
100
        tool_server: ToolServer | None,
101
102
103
104
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
105
106
        pass

107
108
109
110
    @abstractmethod
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

111
112
113
114

class SimpleContext(ConversationContext):
    def __init__(self):
        self.last_output = None
115
116
117
118
119
        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
120
121
        # not implemented yet for SimpleContext
        self.all_turn_metrics = []
122
123
124

    def append_output(self, output) -> None:
        self.last_output = output
125
126
127
128
129
        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 [])
130
131
132
133
134
135
136
137
138
139

    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.")

140
141
    async def init_tool_sessions(
        self,
142
        tool_server: ToolServer | None,
143
144
145
146
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ) -> None:
147
148
        pass

149
150
151
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

152
153
154
155
156

class HarmonyContext(ConversationContext):
    def __init__(
        self,
        messages: list,
157
        available_tools: list[str],
158
159
    ):
        self._messages = messages
160
        self.finish_reason: str | None = None
161
        self.available_tools = available_tools
162
        self._tool_sessions: dict[str, ClientSession | Tool] = {}
163
        self.called_tools: set[str] = set()
164
165
166
167
168

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
169
        self.num_cached_tokens = 0
170
        self.num_reasoning_tokens = 0
171
        self.num_tool_output_tokens = 0
172

173
        # Turn tracking - replaces multiple individual tracking variables
174
175
176
        self.current_turn_metrics = TurnMetrics()
        # Track metrics for all turns
        self.all_turn_metrics: list[TurnMetrics] = []
177
178
        self.is_first_turn = True
        self.first_tok_of_message = True  # For streaming support
179

180
181
182
183
    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
184

185
    def append_output(self, output: RequestOutput | list[Message]) -> None:
186
187
        if isinstance(output, RequestOutput):
            output_token_ids = output.outputs[0].token_ids
188
            self.parser = get_streamable_parser_for_assistant()
189
190
            for token_id in output_token_ids:
                self.parser.process(token_id)
191
                # Check if the current token is part of reasoning content
192
193
194
                self._update_num_reasoning_tokens()
            self._update_prefill_token_usage(output)
            self._update_decode_token_usage(output)
195
196
197
            # 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()
198
199
200
            # append_output is called only once before tool calling
            # in non-streaming case
            # so we can append all the parser messages to _messages
201
            output_msgs = self.parser.messages
202
203
            # The responses finish reason is set in the last message
            self.finish_reason = output.outputs[0].finish_reason
204
205
206
207
208
        else:
            # Tool output.
            output_msgs = output
        self._messages.extend(output_msgs)

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

212
213
214
215
216
        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
217

218
        Tool output tokens are calculated as:
219
        current_prompt_tokens - last_turn_prompt_tokens -
220
221
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
222

223
224
225
226
227
228
229
        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
230
            logger.error("RequestOutput appended contains no prompt_token_ids.")
231
232

        # Update current turn input tokens
233
        self.current_turn_metrics.input_tokens = this_turn_input_tokens
234
235
236
237
238
239
        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:
240
            previous_turn = self.all_turn_metrics[-1]
241
242
243
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
244
            this_turn_tool_tokens = (
245
246
247
                self.current_turn_metrics.input_tokens
                - previous_turn.input_tokens
                - previous_turn.output_tokens
248
            )
249
250
251
252
253
254
255
256

            # 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.",
257
                    this_turn_tool_tokens,
258
259
260
                    self.current_turn_metrics.input_tokens,
                    previous_turn.input_tokens,
                    previous_turn.output_tokens,
261
                )
262
263
264
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens
265
            self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
266
267

        # Update cached tokens
268
269
270
271
        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
272
273
274

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

276
277
278
279
        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
280

281
282
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
283

284
285
        Args:
            output: The RequestOutput containing generated token information
286

287
288
289
290
291
292
293
294
295
        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
296
            self.current_turn_metrics.output_tokens += updated_output_token_count
297
298
        return updated_output_token_count

299
300
301
302
303
304
305
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
306
307
308
309
310
        return recipient is not None and (
            recipient.startswith("browser.")
            or recipient.startswith("python")
            or recipient.startswith("container.")
        )
311
312
313
314
315
316
317
318
319

    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(
320
321
                    self._tool_sessions["browser"], last_msg
                )
322
323
            elif recipient.startswith("python"):
                return await self.call_python_tool(
324
325
                    self._tool_sessions["python"], last_msg
                )
326
327
            elif recipient.startswith("container."):
                return await self.call_container_tool(
328
329
                    self._tool_sessions["container"], last_msg
                )
330
331
332
333
334
        raise ValueError("No tool call found")

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

335
336
337
    async def call_search_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
338
        self.called_tools.add("browser")
339
340
341
342
343
344
345
346
347
        if isinstance(tool_session, Tool):
            return await tool_session.get_result(self)
        tool_name = last_msg.recipient.split(".")[1]
        args = json.loads(last_msg.content[0].text)
        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 [
348
349
350
351
352
353
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
354
355
        ]

356
357
358
    async def call_python_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
359
        self.called_tools.add("python")
360
361
362
363
364
365
366
367
368
369
370
371
        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 [
372
373
374
375
376
377
            Message(
                author=author,
                content=[content],
                channel=last_msg.channel,
                recipient=Role.ASSISTANT,
            )
378
        ]
379

380
381
    async def init_tool_sessions(
        self,
382
        tool_server: ToolServer | None,
383
384
385
386
        exit_stack: AsyncExitStack,
        request_id: str,
        mcp_tools: dict[str, Mcp],
    ):
387
388
389
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
390
                    tool_type = _map_tool_name_to_tool_type(tool_name)
391
392
393
                    headers = (
                        mcp_tools[tool_type].headers if tool_type in mcp_tools else None
                    )
394
                    tool_session = await exit_stack.enter_async_context(
395
396
                        tool_server.new_session(tool_name, request_id, headers)
                    )
397
398
399
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

400
401
402
    async def call_container_tool(
        self, tool_session: Union["ClientSession", Tool], last_msg: Message
    ) -> list[Message]:
403
        """
404
405
406
407
408
409
410
411
412
413
414
415
416
417
        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",
                }
418
419
420
421
422
423
424
425
426
427
428
        """
        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]
        args = json.loads(last_msg.content[0].text)
        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 [
429
430
431
432
433
434
            Message(
                author=author,
                content=[content],
                recipient=Role.ASSISTANT,
                channel=last_msg.channel,
            )
435
436
437
438
439
440
441
        ]

    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):
442
443
444
                logger.info(
                    "Cleaning up tool session for %s", tool_session._client_info
                )
445
446
447
                with contextlib.suppress(Exception):
                    await tool_session.call_tool("cleanup_session", {})

448
449
450
451
452
453
        await asyncio.gather(
            *(
                cleanup_tool_session(self._tool_sessions[tool])
                for tool in self.called_tools
            )
        )
454

455
456
457
458
459
460
461
462
463

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
464
        self.first_tok_of_message = True
465
466
467

    @property
    def messages(self) -> list:
468
        return self._messages
469

470
    def append_output(self, output: RequestOutput | list[Message]) -> None:
471
        if isinstance(output, RequestOutput):
472
473
474
            # 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:
475
                self._update_prefill_token_usage(output)
476
477
478
479
480
            # 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
481
482
            for tok in output.outputs[0].token_ids:
                self.parser.process(tok)
483
484
485
486
            self._update_decode_token_usage(output)

            # For streaming, update previous turn when message is complete
            if output.finished:
487
488
                self.all_turn_metrics.append(self.current_turn_metrics.copy())
                self.current_turn_metrics.reset()
489
            # Check if the current token is part of reasoning content
490
            self._update_num_reasoning_tokens()
491
            self.last_tok = tok
492
            if len(self._messages) - self.num_init_messages < len(self.parser.messages):
493
                self._messages.extend(
494
495
                    self.parser.messages[len(self._messages) - self.num_init_messages :]
                )
496
497
498
499
500
501
502
503
504
505
506
507
        else:
            # 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]
508
            # TODO: add tool_output messages to self._messages
509
510
511
512
513

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

    def is_assistant_action_turn(self) -> bool:
514
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530

    def render_for_completion(self) -> list[int]:
        # now this list of tokens as next turn's starting tokens
        # `<|start|>assistant``,
        # 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