context.py 17.6 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_harmony import Author, Message, Role, StreamState, TextContent
12
13
14
15

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

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

22
23
24
logger = logging.getLogger(__name__)


25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
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)


42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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

60
61
    @abstractmethod
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
62
63
                                 exit_stack: AsyncExitStack,
                                 request_id: str) -> None:
64
65
        pass

66
67
68
69
    @abstractmethod
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

70
71
72
73
74

class SimpleContext(ConversationContext):

    def __init__(self):
        self.last_output = None
75
76
77
78
79
        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
80
81
82

    def append_output(self, output) -> None:
        self.last_output = output
83
84
85
86
87
        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 [])
88
89
90
91
92
93
94
95
96
97

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

98
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
99
100
                                 exit_stack: AsyncExitStack,
                                 request_id: str) -> None:
101
102
        pass

103
104
105
    async def cleanup_session(self) -> None:
        raise NotImplementedError("Should not be called.")

106
107
108
109
110
111

class HarmonyContext(ConversationContext):

    def __init__(
        self,
        messages: list,
112
        available_tools: list[str],
113
114
    ):
        self._messages = messages
115
        self.finish_reason: Optional[str] = None
116
117
        self.available_tools = available_tools
        self._tool_sessions: dict[str, Union[ClientSession, Tool]] = {}
118
        self.called_tools: set[str] = set()
119
120
121
122
123

        self.parser = get_streamable_parser_for_assistant()
        self.num_init_messages = len(messages)
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
124
        self.num_cached_tokens = 0
125
        self.num_reasoning_tokens = 0
126
        self.num_tool_output_tokens = 0
127

128
129
130
131
132
        # 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
133

134
135
136
137
    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
138

139
140
    def append_output(self, output: Union[RequestOutput,
                                          list[Message]]) -> None:
141
142
        if isinstance(output, RequestOutput):
            output_token_ids = output.outputs[0].token_ids
143
            self.parser = get_streamable_parser_for_assistant()
144
145
            for token_id in output_token_ids:
                self.parser.process(token_id)
146
                # Check if the current token is part of reasoning content
147
148
149
150
151
152
153
                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()
154
            output_msgs = self.parser.messages
155
156
            # The responses finish reason is set in the last message
            self.finish_reason = output.outputs[0].finish_reason
157
158
159
160
161
        else:
            # Tool output.
            output_msgs = output
        self._messages.extend(output_msgs)

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

165
166
167
168
169
        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
170

171
        Tool output tokens are calculated as:
172
        current_prompt_tokens - last_turn_prompt_tokens -
173
174
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
175

176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
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
        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.
221

222
223
224
225
        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
226

227
228
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
229

230
231
        Args:
            output: The RequestOutput containing generated token information
232

233
234
235
236
237
238
239
240
241
242
243
244
        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

245
246
247
248
249
250
251
252
    @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.")
253
254
                                          or recipient.startswith("python") or
                                          recipient.startswith("container."))
255
256
257
258
259
260
261
262
263

    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(
264
                    self._tool_sessions["browser"], last_msg)
265
266
            elif recipient.startswith("python"):
                return await self.call_python_tool(
267
                    self._tool_sessions["python"], last_msg)
268
269
270
            elif recipient.startswith("container."):
                return await self.call_container_tool(
                    self._tool_sessions["container"], last_msg)
271
272
273
274
275
        raise ValueError("No tool call found")

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

276
277
278
    async def call_search_tool(self, tool_session: Union["ClientSession",
                                                         Tool],
                               last_msg: Message) -> list[Message]:
279
        self.called_tools.add("browser")
280
281
282
283
284
285
286
287
288
        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 [
289
290
291
292
            Message(author=author,
                    content=[content],
                    recipient=Role.ASSISTANT,
                    channel=last_msg.channel)
293
294
295
296
297
        ]

    async def call_python_tool(self, tool_session: Union["ClientSession",
                                                         Tool],
                               last_msg: Message) -> list[Message]:
298
        self.called_tools.add("python")
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        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)
        ]
316

317
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
318
319
                                 exit_stack: AsyncExitStack,
                                 request_id: str) -> None:
320
321
322
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
                    tool_session = await exit_stack.enter_async_context(
                        tool_server.new_session(tool_name, request_id))
                    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))
375

376
377
378
379
380
381
382
383
384
385

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
386
        self.first_tok_of_message = True
387
388
389
390
391

    @property
    def messages(self) -> list:
        return self.parser.messages

392
393
    def append_output(self, output: Union[RequestOutput,
                                          list[Message]]) -> None:
394
        if isinstance(output, RequestOutput):
395
396
397
            # 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:
398
399
                self._update_prefill_token_usage(output)
                self.current_turn.output_tokens = 0
400
401
402
403
404
            # 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
405
406
            for tok in output.outputs[0].token_ids:
                self.parser.process(tok)
407
408
409
410
411
            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()
412
            # Check if the current token is part of reasoning content
413
            self._update_num_reasoning_tokens()
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
            self.last_tok = tok
        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]

    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