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

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

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

20
21
22
if TYPE_CHECKING:
    from mcp.client import ClientSession

23
24
logger = logging.getLogger(__name__)

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

43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class TurnTokens:
    """Tracks token counts for a single conversation turn."""

    def __init__(self, input_tokens=0, output_tokens=0):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens

    def reset(self):
        """Reset counters for a new turn."""
        self.input_tokens = 0
        self.output_tokens = 0

    def copy(self):
        """Create a copy of this turn's token counts."""
        return TurnTokens(self.input_tokens, self.output_tokens)


61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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

79
80
    @abstractmethod
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
81
82
                                 exit_stack: AsyncExitStack, request_id: str,
                                 mcp_tools: dict[str, Mcp]) -> None:
83
84
        pass

85
86
87
88
    @abstractmethod
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

89
90
91
92
93

class SimpleContext(ConversationContext):

    def __init__(self):
        self.last_output = None
94
95
96
97
98
        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
99
100
101

    def append_output(self, output) -> None:
        self.last_output = output
102
103
104
105
106
        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 [])
107
108
109
110
111
112
113
114
115
116

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

117
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
118
119
                                 exit_stack: AsyncExitStack, request_id: str,
                                 mcp_tools: dict[str, Mcp]) -> None:
120
121
        pass

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

125
126
127
128
129
130

class HarmonyContext(ConversationContext):

    def __init__(
        self,
        messages: list,
131
        available_tools: list[str],
132
133
    ):
        self._messages = messages
134
        self.finish_reason: Optional[str] = None
135
136
        self.available_tools = available_tools
        self._tool_sessions: dict[str, Union[ClientSession, Tool]] = {}
137
        self.called_tools: set[str] = set()
138
139
140
141
142

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
143
        self.num_cached_tokens = 0
144
        self.num_reasoning_tokens = 0
145
        self.num_tool_output_tokens = 0
146

147
148
149
150
151
        # Turn tracking - replaces multiple individual tracking variables
        self.current_turn = TurnTokens()
        self.previous_turn = TurnTokens()
        self.is_first_turn = True
        self.first_tok_of_message = True  # For streaming support
152

153
154
155
156
    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
157

158
159
    def append_output(self, output: Union[RequestOutput,
                                          list[Message]]) -> None:
160
161
        if isinstance(output, RequestOutput):
            output_token_ids = output.outputs[0].token_ids
162
            self.parser = get_streamable_parser_for_assistant()
163
164
            for token_id in output_token_ids:
                self.parser.process(token_id)
165
                # Check if the current token is part of reasoning content
166
167
168
169
170
171
172
                self._update_num_reasoning_tokens()
            self._update_prefill_token_usage(output)
            # Reset current turn output tokens for this turn
            self.current_turn.output_tokens = 0
            self._update_decode_token_usage(output)
            # Move current turn to previous turn for next turn's calculations
            self.previous_turn = self.current_turn.copy()
173
174
175
            # append_output is called only once before tool calling
            # in non-streaming case
            # so we can append all the parser messages to _messages
176
            output_msgs = self.parser.messages
177
178
            # The responses finish reason is set in the last message
            self.finish_reason = output.outputs[0].finish_reason
179
180
181
182
183
        else:
            # Tool output.
            output_msgs = output
        self._messages.extend(output_msgs)

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

187
188
189
190
191
        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
192

193
        Tool output tokens are calculated as:
194
        current_prompt_tokens - last_turn_prompt_tokens -
195
196
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
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
        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
            logger.error(
                "RequestOutput appended contains no prompt_token_ids.")

        # Update current turn input tokens
        self.current_turn.input_tokens = this_turn_input_tokens
        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:
            # start counting tool after first turn
            # tool tokens = this turn prefill - last turn prefill -
            # last turn decode
            this_turn_tool_tokens = (self.current_turn.input_tokens -
                                     self.previous_turn.input_tokens -
                                     self.previous_turn.output_tokens)

            # 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.",
                    this_turn_tool_tokens, self.current_turn.input_tokens,
                    self.previous_turn.input_tokens,
                    self.previous_turn.output_tokens)
                this_turn_tool_tokens = 0

            self.num_tool_output_tokens += this_turn_tool_tokens

        # Update cached tokens
        if output.num_cached_tokens is not None:
            self.num_cached_tokens += output.num_cached_tokens

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

244
245
246
247
        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
248

249
250
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
251

252
253
        Args:
            output: The RequestOutput containing generated token information
254

255
256
257
258
259
260
261
262
263
264
265
266
        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
            self.current_turn.output_tokens += updated_output_token_count
        return updated_output_token_count

267
268
269
270
271
272
273
274
    @property
    def messages(self) -> list:
        return self._messages

    def need_builtin_tool_call(self) -> bool:
        last_msg = self.messages[-1]
        recipient = last_msg.recipient
        return recipient is not None and (recipient.startswith("browser.")
275
276
                                          or recipient.startswith("python") or
                                          recipient.startswith("container."))
277
278
279
280
281
282
283
284
285

    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(
286
                    self._tool_sessions["browser"], last_msg)
287
288
            elif recipient.startswith("python"):
                return await self.call_python_tool(
289
                    self._tool_sessions["python"], last_msg)
290
291
292
            elif recipient.startswith("container."):
                return await self.call_container_tool(
                    self._tool_sessions["container"], last_msg)
293
294
295
296
297
        raise ValueError("No tool call found")

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

298
299
300
    async def call_search_tool(self, tool_session: Union["ClientSession",
                                                         Tool],
                               last_msg: Message) -> list[Message]:
301
        self.called_tools.add("browser")
302
303
304
305
306
307
308
309
310
        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 [
311
312
313
314
            Message(author=author,
                    content=[content],
                    recipient=Role.ASSISTANT,
                    channel=last_msg.channel)
315
316
317
318
319
        ]

    async def call_python_tool(self, tool_session: Union["ClientSession",
                                                         Tool],
                               last_msg: Message) -> list[Message]:
320
        self.called_tools.add("python")
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
        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 [
            Message(author=author,
                    content=[content],
                    channel=last_msg.channel,
                    recipient=Role.ASSISTANT)
        ]
338

339
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
340
341
                                 exit_stack: AsyncExitStack, request_id: str,
                                 mcp_tools: dict[str, Mcp]):
342
343
344
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
345
346
347
                    tool_type = _map_tool_name_to_tool_type(tool_name)
                    headers = mcp_tools[
                        tool_type].headers if tool_type in mcp_tools else None
348
                    tool_session = await exit_stack.enter_async_context(
349
350
                        tool_server.new_session(tool_name, request_id,
                                                headers))
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
                    self._tool_sessions[tool_name] = tool_session
                    exit_stack.push_async_exit(self.cleanup_session)

    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(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 [
            Message(author=author,
                    content=[content],
                    recipient=Role.ASSISTANT,
                    channel=last_msg.channel)
        ]

    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):
                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))
401

402
403
404
405
406
407
408
409
410
411

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
412
        self.first_tok_of_message = True
413
414
415

    @property
    def messages(self) -> list:
416
        return self._messages
417

418
419
    def append_output(self, output: Union[RequestOutput,
                                          list[Message]]) -> None:
420
        if isinstance(output, RequestOutput):
421
422
423
            # 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:
424
425
                self._update_prefill_token_usage(output)
                self.current_turn.output_tokens = 0
426
427
428
429
430
            # 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
431
432
            for tok in output.outputs[0].token_ids:
                self.parser.process(tok)
433
434
435
436
437
            self._update_decode_token_usage(output)

            # For streaming, update previous turn when message is complete
            if output.finished:
                self.previous_turn = self.current_turn.copy()
438
            # Check if the current token is part of reasoning content
439
            self._update_num_reasoning_tokens()
440
            self.last_tok = tok
441
442
443
444
445
            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:])
446
447
448
449
450
451
452
453
454
455
456
457
        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]
458
            # TODO: add tool_output messages to self._messages
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481

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

    def is_assistant_action_turn(self) -> bool:
        return self.last_tok in self.encoding.stop_tokens_for_assistant_actions(
        )

    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