context.py 18.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_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
155
156
            # append_output is called only once before tool calling
            # in non-streaming case
            # so we can append all the parser messages to _messages
157
            output_msgs = self.parser.messages
158
159
            # The responses finish reason is set in the last message
            self.finish_reason = output.outputs[0].finish_reason
160
161
162
163
164
        else:
            # Tool output.
            output_msgs = output
        self._messages.extend(output_msgs)

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

168
169
170
171
172
        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
173

174
        Tool output tokens are calculated as:
175
        current_prompt_tokens - last_turn_prompt_tokens -
176
177
        last_turn_output_tokens
        This represents tokens added between turns (typically tool responses).
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
221
222
223
        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.
224

225
226
227
228
        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
229

230
231
        In streaming mode, this is called for each token generated.
        In non-streaming mode, this is called once with all output tokens.
232

233
234
        Args:
            output: The RequestOutput containing generated token information
235

236
237
238
239
240
241
242
243
244
245
246
247
        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

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

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

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

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

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

320
    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
321
322
                                 exit_stack: AsyncExitStack,
                                 request_id: str) -> None:
323
324
325
        if tool_server:
            for tool_name in self.available_tools:
                if tool_name not in self._tool_sessions:
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
375
376
377
                    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))
378

379
380
381
382
383
384
385
386
387
388

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

    @property
    def messages(self) -> list:
393
        return self._messages
394

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

    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