@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 185 lines 6.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4 import os
5 from abc import ABC, abstractmethod
6 from typing import TYPE_CHECKING, Any
7
8 from openai.types.responses.response_function_tool_call_output_item import (
9 ResponseFunctionToolCallOutputItem,
10 )
11 from openai_harmony import Author, Message, Role, TextContent
12
13 from vllm.logger import init_logger
14 from vllm.utils import random_uuid
15
16 if TYPE_CHECKING:
17 # Avoid circular import.
18 from vllm.entrypoints.openai.responses.context import ConversationContext
19
20 logger = init_logger(__name__)
21
22 MIN_GPT_OSS_VERSION = "0.0.7"
23
24
25 def validate_gpt_oss_install():
26 """Check if the gpt-oss is installed and its version is at least 0.0.7.
27 If not, raise an ImportError.
28 """
29 from importlib.metadata import PackageNotFoundError, version
30
31 from packaging.version import InvalidVersion, Version
32
33 try:
34 pkg_version_str = version("gpt_oss")
35 pkg_version = Version(pkg_version_str)
36 except PackageNotFoundError:
37 raise ImportError("Package 'gpt_oss' is not installed.") from None
38 except InvalidVersion as e:
39 raise ImportError(f"Invalid version string for 'gpt_oss': {e}") from None
40
41 if pkg_version < Version(MIN_GPT_OSS_VERSION):
42 raise ImportError(
43 f"gpt_oss >= {MIN_GPT_OSS_VERSION} is required, "
44 f"but {pkg_version} is installed."
45 ) from None
46
47
48 class Tool(ABC):
49 @abstractmethod
50 async def get_result(self, context: "ConversationContext") -> Any:
51 pass
52
53 @abstractmethod
54 async def get_result_parsable_context(self, context: "ConversationContext") -> Any:
55 pass
56
57
58 class HarmonyBrowserTool(Tool):
59 def __init__(self):
60 self.enabled = True
61 exa_api_key = os.getenv("EXA_API_KEY")
62 if not exa_api_key:
63 self.enabled = False
64 logger.warning_once("EXA_API_KEY is not set, browsing is disabled")
65 return
66
67 try:
68 validate_gpt_oss_install()
69 from gpt_oss.tools.simple_browser import SimpleBrowserTool
70 from gpt_oss.tools.simple_browser.backend import ExaBackend
71 except ImportError as e:
72 self.enabled = False
73 logger.warning_once(
74 "gpt_oss is not installed properly (%s), browsing is disabled", e
75 )
76 return
77
78 browser_backend = ExaBackend(source="web", api_key=exa_api_key)
79 self.browser_tool = SimpleBrowserTool(backend=browser_backend)
80 logger.info_once("Browser tool initialized")
81
82 async def get_result(self, context: "ConversationContext") -> Any:
83 from vllm.entrypoints.openai.responses.context import HarmonyContext
84
85 assert isinstance(context, HarmonyContext)
86 last_msg = context.messages[-1]
87 tool_output_msgs = []
88 async for msg in self.browser_tool.process(last_msg):
89 tool_output_msgs.append(msg)
90 return tool_output_msgs
91
92 async def get_result_parsable_context(self, context: "ConversationContext") -> Any:
93 raise NotImplementedError("Not implemented yet")
94
95 @property
96 def tool_config(self) -> Any:
97 return self.browser_tool.tool_config
98
99
100 class HarmonyPythonTool(Tool):
101 def __init__(self):
102 self.enabled = True
103
104 try:
105 validate_gpt_oss_install()
106 from gpt_oss.tools.python_docker.docker_tool import PythonTool
107 except ImportError as e:
108 self.enabled = False
109 logger.warning_once(
110 "gpt_oss is not installed properly (%s), code interpreter is disabled",
111 e,
112 )
113 return
114
115 self.python_tool = PythonTool()
116
117 async def validate(self):
118 if not self.enabled:
119 return
120 try:
121 message = Message(
122 author=Author(role=Role.ASSISTANT),
123 content=[TextContent(text="print('Hello, world!')")],
124 channel="analysis",
125 recipient="python",
126 content_type="code",
127 )
128 msgs = []
129 async for msg in self.python_tool.process(message):
130 msgs.append(msg)
131 assert msgs[0].content[0].text == "Hello, world!\n"
132 except Exception as e:
133 self.enabled = False
134 logger.warning_once(
135 "Code interpreter tool failed to initialize (%s), code "
136 "interpreter is disabled",
137 e,
138 )
139 return
140 logger.info_once("Code interpreter tool initialized")
141
142 async def get_result(self, context: "ConversationContext") -> Any:
143 from vllm.entrypoints.openai.responses.context import HarmonyContext
144
145 assert isinstance(context, HarmonyContext)
146 last_msg = context.messages[-1]
147 tool_output_msgs = []
148 async for msg in self.python_tool.process(last_msg):
149 tool_output_msgs.append(msg)
150 return tool_output_msgs
151
152 async def get_result_parsable_context(self, context: "ConversationContext") -> Any:
153 """This function converts parsable context types to harmony and
154 back so we can use GPTOSS demo python tool
155 """
156 from vllm.entrypoints.openai.responses.context import ParsableContext
157
158 assert isinstance(context, ParsableContext)
159
160 last_msg = context.response_messages[-1]
161 args = json.loads(last_msg.arguments)
162
163 last_msg_harmony = Message(
164 author=Author(role="assistant", name=None),
165 content=[TextContent(text=args["code"])],
166 channel="analysis",
167 recipient="python",
168 content_type="code",
169 )
170
171 tool_output_msgs = []
172 async for msg in self.python_tool.process(last_msg_harmony):
173 processed = ResponseFunctionToolCallOutputItem(
174 id=f"fco_{random_uuid()}",
175 type="function_call_output",
176 call_id=last_msg.call_id,
177 output=msg.content[0].text,
178 status="completed",
179 )
180 tool_output_msgs.append(processed)
181 return tool_output_msgs
182
183 @property
184 def tool_config(self) -> Any:
185 return self.python_tool.tool_config