A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[gpt-oss] Add Tool/ConversationContext classes and harmony_utils (#22340)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Woosuk Kwon committed
Aug 6, 2025 at 01:08 UTC
178d03fbd64e18999647b349623cd1489f816c8c
3 files changed
+375
vllm/entrypoints/context.py
new
+177
@@ -0,0 +1,177 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+from abc import ABC, abstractmethod
5
+
6
+from openai_harmony import Message, Role, StreamState
7
+
8
+from vllm.entrypoints.harmony_utils import (
9
+ get_encoding, get_streamable_parser_for_assistant, render_for_completion)
10
+from vllm.entrypoints.tool import Tool
11
+from vllm.outputs import RequestOutput
12
+
13
+logger = logging.getLogger(__name__)
14
+
15
+
16
+class ConversationContext(ABC):
17
+
18
+ @abstractmethod
19
+ def append_output(self, output) -> None:
20
+ pass
21
+
22
+ @abstractmethod
23
+ async def call_tool(self) -> list[Message]:
24
+ pass
25
+
26
+ @abstractmethod
27
+ def need_builtin_tool_call(self) -> bool:
28
+ pass
29
+
30
+ @abstractmethod
31
+ def render_for_completion(self) -> list[int]:
32
+ pass
33
+
34
+
35
+class SimpleContext(ConversationContext):
36
+
37
+ def __init__(self):
38
+ self.last_output = None
39
+
40
+ def append_output(self, output) -> None:
41
+ self.last_output = output
42
+
43
+ def need_builtin_tool_call(self) -> bool:
44
+ return False
45
+
46
+ async def call_tool(self) -> list[Message]:
47
+ raise NotImplementedError("Should not be called.")
48
+
49
+ def render_for_completion(self) -> list[int]:
50
+ raise NotImplementedError("Should not be called.")
51
+
52
+
53
+class HarmonyContext(ConversationContext):
54
+
55
+ def __init__(
56
+ self,
57
+ messages: list,
58
+ tool_sessions: dict[str, Tool],
59
+ ):
60
+ self._messages = messages
61
+ self.tool_sessions = tool_sessions
62
+
63
+ self.parser = get_streamable_parser_for_assistant()
64
+ self.num_init_messages = len(messages)
65
+ # TODO(woosuk): Implement the following fields.
66
+ self.num_prompt_tokens = 0
67
+ self.num_cached_tokens = 0
68
+ self.num_output_tokens = 0
69
+ self.num_reasoning_tokens = 0
70
+
71
+ def append_output(self, output) -> None:
72
+ if isinstance(output, RequestOutput):
73
+ output_token_ids = output.outputs[0].token_ids
74
+ for token_id in output_token_ids:
75
+ self.parser.process(token_id)
76
+ output_msgs = self.parser.messages
77
+ else:
78
+ # Tool output.
79
+ output_msgs = output
80
+ self._messages.extend(output_msgs)
81
+
82
+ @property
83
+ def messages(self) -> list:
84
+ return self._messages
85
+
86
+ def need_builtin_tool_call(self) -> bool:
87
+ last_msg = self.messages[-1]
88
+ recipient = last_msg.recipient
89
+ return recipient is not None and (recipient.startswith("browser.")
90
+ or recipient.startswith("python"))
91
+
92
+ async def call_tool(self) -> list[Message]:
93
+ if not self.messages:
94
+ return []
95
+ last_msg = self.messages[-1]
96
+ recipient = last_msg.recipient
97
+ if recipient is not None:
98
+ if recipient.startswith("browser."):
99
+ return await self.call_search_tool(
100
+ self.tool_sessions["browser"], last_msg)
101
+ elif recipient.startswith("python"):
102
+ return await self.call_python_tool(
103
+ self.tool_sessions["python"], last_msg)
104
+ raise ValueError("No tool call found")
105
+
106
+ def render_for_completion(self) -> list[int]:
107
+ return render_for_completion(self.messages)
108
+
109
+ async def call_search_tool(
110
+ self,
111
+ tool_session: Tool,
112
+ last_msg: Message,
113
+ ) -> list[Message]:
114
+ return await tool_session.get_result(self)
115
+
116
+ async def call_python_tool(
117
+ self,
118
+ tool_session: Tool,
119
+ last_msg: Message,
120
+ ) -> list[Message]:
121
+ return await tool_session.get_result(self)
122
+
123
+
124
+class StreamingHarmonyContext(HarmonyContext):
125
+
126
+ def __init__(self, *args, **kwargs):
127
+ super().__init__(*args, **kwargs)
128
+ self.last_output = None
129
+
130
+ self.parser = get_streamable_parser_for_assistant()
131
+ self.encoding = get_encoding()
132
+ self.last_tok = None
133
+
134
+ @property
135
+ def messages(self) -> list:
136
+ return self.parser.messages
137
+
138
+ def append_output(self, output) -> None:
139
+ if isinstance(output, RequestOutput):
140
+ tok = output.outputs[0].token_ids[0]
141
+ self.parser.process(tok)
142
+ self.last_tok = tok
143
+ else:
144
+ # Handle the case of tool output in direct message format
145
+ assert len(output) == 1, "Tool output should be a single message"
146
+ msg = output[0]
147
+ # Sometimes the recipient is not set for tool messages,
148
+ # so we set it to "assistant"
149
+ if msg.author.role == Role.TOOL and msg.recipient is None:
150
+ msg.recipient = "assistant"
151
+ toks = self.encoding.render(msg)
152
+ for tok in toks:
153
+ self.parser.process(tok)
154
+ self.last_tok = toks[-1]
155
+
156
+ def is_expecting_start(self) -> bool:
157
+ return self.parser.state == StreamState.EXPECT_START
158
+
159
+ def is_assistant_action_turn(self) -> bool:
160
+ return self.last_tok in self.encoding.stop_tokens_for_assistant_actions(
161
+ )
162
+
163
+ def render_for_completion(self) -> list[int]:
164
+ # now this list of tokens as next turn's starting tokens
165
+ # `<|start|>assistant``,
166
+ # we need to process them in parser.
167
+ rendered_tokens = super().render_for_completion()
168
+
169
+ last_n = -1
170
+ to_process = []
171
+ while rendered_tokens[last_n] != self.last_tok:
172
+ to_process.append(rendered_tokens[last_n])
173
+ last_n -= 1
174
+ for tok in reversed(to_process):
175
+ self.parser.process(tok)
176
+
177
+ return rendered_tokens
vllm/entrypoints/harmony_utils.py
new
+111
@@ -0,0 +1,111 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import datetime
4
+from typing import Literal, Optional
5
+
6
+from openai.types.responses.tool import Tool
7
+from openai_harmony import (Conversation, DeveloperContent,
8
+ HarmonyEncodingName, Message, ReasoningEffort,
9
+ Role, StreamableParser, SystemContent, TextContent,
10
+ ToolDescription, load_harmony_encoding)
11
+
12
+REASONING_EFFORT = {
13
+ "high": ReasoningEffort.HIGH,
14
+ "medium": ReasoningEffort.MEDIUM,
15
+ "low": ReasoningEffort.LOW,
16
+}
17
+
18
+_harmony_encoding = None
19
+
20
+
21
+def get_encoding():
22
+ global _harmony_encoding
23
+ if _harmony_encoding is None:
24
+ _harmony_encoding = load_harmony_encoding(
25
+ HarmonyEncodingName.HARMONY_GPT_OSS)
26
+ return _harmony_encoding
27
+
28
+
29
+def get_system_message(
30
+ model_identity: Optional[str] = None,
31
+ reasoning_effort: Optional[Literal["high", "medium", "low"]] = None,
32
+ start_date: Optional[str] = None,
33
+ browser_description: Optional[str] = None,
34
+ python_description: Optional[str] = None,
35
+) -> Message:
36
+ sys_msg_content = SystemContent.new()
37
+ if model_identity is not None:
38
+ sys_msg_content = sys_msg_content.with_model_identity(model_identity)
39
+ if reasoning_effort is not None:
40
+ sys_msg_content = sys_msg_content.with_reasoning_effort(
41
+ REASONING_EFFORT[reasoning_effort])
42
+ if start_date is None:
43
+ # NOTE(woosuk): This brings non-determinism in vLLM. Be careful.
44
+ start_date = datetime.datetime.now().strftime("%Y-%m-%d")
45
+ sys_msg_content = sys_msg_content.with_conversation_start_date(start_date)
46
+ if browser_description is not None:
47
+ sys_msg_content = sys_msg_content.with_tools(browser_description)
48
+ if python_description is not None:
49
+ sys_msg_content = sys_msg_content.with_tools(python_description)
50
+ sys_msg = Message.from_role_and_content(Role.SYSTEM, sys_msg_content)
51
+ return sys_msg
52
+
53
+
54
+def get_developer_message(instructions: Optional[str] = None,
55
+ tools: Optional[list[Tool]] = None) -> Message:
56
+ dev_msg_content = DeveloperContent.new()
57
+ if instructions is not None:
58
+ dev_msg_content = dev_msg_content.with_instructions(instructions)
59
+ if tools is not None:
60
+ function_tools = []
61
+ for tool in tools:
62
+ if tool.type in ("web_search_preview", "code_interpreter"):
63
+ # These are built-in tools that are added to the system message.
64
+ pass
65
+ elif tool.type == "function":
66
+ function_tools.append(tool)
67
+ else:
68
+ raise ValueError(f"tool type {tool.type} not supported")
69
+ if function_tools:
70
+ function_tool_descriptions = [
71
+ ToolDescription.new(
72
+ name=tool.name,
73
+ description=tool.description,
74
+ parameters=tool.parameters,
75
+ ) for tool in function_tools
76
+ ]
77
+ dev_msg_content = dev_msg_content.with_function_tools(
78
+ function_tool_descriptions)
79
+ dev_msg = Message.from_role_and_content(Role.DEVELOPER, dev_msg_content)
80
+ return dev_msg
81
+
82
+
83
+def get_user_message(content: str) -> Message:
84
+ return Message.from_role_and_content(Role.USER, content)
85
+
86
+
87
+def parse_chat_input(chat_msg) -> Message:
88
+ role = chat_msg["role"]
89
+ content = chat_msg["content"]
90
+ if isinstance(content, str):
91
+ contents = [TextContent(text=content)]
92
+ else:
93
+ # TODO: Support refusal.
94
+ contents = [TextContent(text=c["text"]) for c in content]
95
+ msg = Message.from_role_and_contents(role, contents)
96
+ return msg
97
+
98
+
99
+def render_for_completion(messages: list[Message]) -> list[int]:
100
+ conversation = Conversation.from_messages(messages)
101
+ token_ids = get_encoding().render_conversation_for_completion(
102
+ conversation, Role.ASSISTANT)
103
+ return token_ids
104
+
105
+
106
+def get_stop_tokens_for_assistant_actions() -> list[int]:
107
+ return get_encoding().stop_tokens_for_assistant_actions()
108
+
109
+
110
+def get_streamable_parser_for_assistant() -> StreamableParser:
111
+ return StreamableParser(get_encoding(), role=Role.ASSISTANT)
vllm/entrypoints/tool.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+from abc import ABC, abstractmethod
5
+from typing import TYPE_CHECKING, Any
6
+
7
+from vllm.logger import init_logger
8
+
9
+if TYPE_CHECKING:
10
+ # Avoid circular import.
11
+ from vllm.entrypoints.context import ConversationContext
12
+
13
+logger = init_logger(__name__)
14
+
15
+
16
+class Tool(ABC):
17
+
18
+ @abstractmethod
19
+ async def get_result(self, context: "ConversationContext") -> Any:
20
+ pass
21
+
22
+
23
+class HarmonyBrowserTool(Tool):
24
+
25
+ def __init__(self):
26
+ self.enabled = True
27
+ exa_api_key = os.getenv("EXA_API_KEY")
28
+ if not exa_api_key:
29
+ self.enabled = False
30
+ logger.warning_once("EXA_API_KEY is not set, browsing is disabled")
31
+ return
32
+
33
+ try:
34
+ from gpt_oss.tools.simple_browser import SimpleBrowserTool
35
+ from gpt_oss.tools.simple_browser.backend import ExaBackend
36
+ except ImportError:
37
+ self.enabled = False
38
+ logger.warning_once(
39
+ "gpt_oss is not installed, browsing is disabled")
40
+ return
41
+
42
+ browser_backend = ExaBackend(source="web", api_key=exa_api_key)
43
+ self.browser_tool = SimpleBrowserTool(backend=browser_backend)
44
+ logger.info_once("Browser tool initialized")
45
+
46
+ async def get_result(self, context: "ConversationContext") -> Any:
47
+ from vllm.entrypoints.context import HarmonyContext
48
+ assert isinstance(context, HarmonyContext)
49
+ last_msg = context.messages[-1]
50
+ tool_output_msgs = []
51
+ async for msg in self.browser_tool.process(last_msg):
52
+ tool_output_msgs.append(msg)
53
+ return tool_output_msgs
54
+
55
+ @property
56
+ def tool_config(self) -> Any:
57
+ return self.browser_tool.tool_config
58
+
59
+
60
+class HarmonyPythonTool(Tool):
61
+
62
+ def __init__(self):
63
+ self.enabled = True
64
+
65
+ try:
66
+ from gpt_oss.tools.python_docker.docker_tool import PythonTool
67
+ except ImportError:
68
+ self.enabled = False
69
+ logger.warning_once(
70
+ "gpt_oss is not installed, code interpreter is disabled")
71
+ return
72
+
73
+ self.python_tool = PythonTool()
74
+ logger.info_once("Code interpreter tool initialized")
75
+
76
+ async def get_result(self, context: "ConversationContext") -> Any:
77
+ from vllm.entrypoints.context import HarmonyContext
78
+ assert isinstance(context, HarmonyContext)
79
+ last_msg = context.messages[-1]
80
+ tool_output_msgs = []
81
+ async for msg in self.python_tool.process(last_msg):
82
+ tool_output_msgs.append(msg)
83
+ return tool_output_msgs
84
+
85
+ @property
86
+ def tool_config(self) -> Any:
87
+ return self.python_tool.tool_config