Adding one unit test which works locally, fixing type to name translation issues

Signed-off-by: Alec Solder <alecs@fb.com>

Author: Alec Solder

tests/entrypoints/openai/test_response_api_with_harmony.py

@@ -373,6 +373,24 @@ async def test_code_interpreter(client: OpenAI, model_name: str):
     assert response.status == "completed"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.skip(reason="Code interpreter tool is not available in CI yet.")
+async def test_mcp_tool(client: OpenAI, model_name: str):
+    response = await client.responses.create(
+        model=model_name,
+        input="Please multiply 123 and 456 using the available python tool.",
+        tools=[{
+            "type": "mcp",
+            "server_label": "code_interpreter",
+            # URL unused for DemoToolServer
+            "server_url": "http://localhost:8888"
+        }],
+    )
+    assert response is not None
+    assert response.status == "completed"
+
+
 def get_weather(latitude, longitude):
     response = requests.get(
         f"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m"  # noqa

vllm/entrypoints/context.py

@@ -8,6 +8,7 @@
 from contextlib import AsyncExitStack
 from typing import TYPE_CHECKING, Optional, Union
 
+from openai.types.responses.tool import Mcp
 from openai_harmony import Author, Message, Role, StreamState, TextContent
 
 from vllm.entrypoints.harmony_utils import (
@@ -22,6 +23,21 @@
 logger = logging.getLogger(__name__)
 
 
+# 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
+def _map_tool_name_to_tool_type(tool_name: str) -> str:
+    if tool_name == "browser":
+        return "web_search_preview"
+    elif tool_name == "python":
+        return "code_interpreter"
+    elif tool_name == "container":
+        return "container"
+    else:
+        raise ValueError(
+            f"Built in tool name not defined in mapping: {tool_name}")
+
+
 class TurnTokens:
     """Tracks token counts for a single conversation turn."""
 
@@ -59,8 +75,8 @@ def render_for_completion(self) -> list[int]:
 
     @abstractmethod
     async def init_tool_sessions(self, tool_server: Optional[ToolServer],
-                                 exit_stack: AsyncExitStack,
-                                 request_id: str) -> None:
+                                 exit_stack: AsyncExitStack, request_id: str,
+                                 mcp_tools: dict[str, Mcp]) -> None:
         pass
 
     @abstractmethod
@@ -96,8 +112,8 @@ def render_for_completion(self) -> list[int]:
         raise NotImplementedError("Should not be called.")
 
     async def init_tool_sessions(self, tool_server: Optional[ToolServer],
-                                 exit_stack: AsyncExitStack,
-                                 request_id: str) -> None:
+                                 exit_stack: AsyncExitStack, request_id: str,
+                                 mcp_tools: dict[str, Mcp]) -> None:
         pass
 
     async def cleanup_session(self) -> None:
@@ -310,20 +326,18 @@ async def call_python_tool(self, tool_session: Union["ClientSession",
                     recipient=Role.ASSISTANT)
         ]
 
-    async def init_tool_sessions(
-            self,
-            tool_server: Optional[ToolServer],
-            exit_stack: AsyncExitStack,
-            request_id: str,
-            mcp_tool_headers: dict[str, dict[str, str]] = None) -> None:
+    async def init_tool_sessions(self, tool_server: Optional[ToolServer],
+                                 exit_stack: AsyncExitStack, request_id: str,
+                                 mcp_tools: dict[str, Mcp]):
         if tool_server:
             for tool_name in self.available_tools:
                 if tool_name not in self._tool_sessions:
+                    tool_type = _map_tool_name_to_tool_type(tool_name)
+                    headers = mcp_tools[
+                        tool_type].headers if tool_type in mcp_tools else None
                     tool_session = await exit_stack.enter_async_context(
-                        tool_server.new_session(
-                            tool_name, request_id,
-                            mcp_tool_headers.get(tool_name)
-                            if tool_name in mcp_tool_headers else {}))
+                        tool_server.new_session(tool_name, request_id,
+                                                headers))
                     logger.info("Created new session for %s", tool_name)
                     self._tool_sessions[tool_name] = tool_session
                     exit_stack.push_async_exit(self.cleanup_session)

vllm/entrypoints/harmony_utils.py

@@ -126,8 +126,10 @@ def get_developer_message(
         function_tools: list[Union[Tool, ChatCompletionToolsParam]] = []
         for tool in tools:
             if tool.type in ("web_search_preview", "code_interpreter",
-                             "container"):
+                             "container", "mcp"):
                 # These are built-in tools that are added to the system message.
+                # Adding in MCP for now until we support MCP tools executed
+                # server side
                 pass
 
             elif tool.type == "function":

vllm/entrypoints/openai/serving_responses.py

@@ -451,8 +451,12 @@ async def responses_full_generator(
 
         async with AsyncExitStack() as exit_stack:
             try:
+                mcp_tools = {
+                    tool.server_label: tool
+                    for tool in request.tools if tool.type == "mcp"
+                }
                 await context.init_tool_sessions(self.tool_server, exit_stack,
-                                                 request.request_id)
+                                                 request.request_id, mcp_tools)
                 async for _ in result_generator:
                     pass
             except asyncio.CancelledError:
@@ -1635,8 +1639,12 @@ def _send_event(event: BaseModel):
         async with AsyncExitStack() as exit_stack:
             processer = None
             if self.use_harmony:
+                mcp_tools = {
+                    tool.server_label: tool
+                    for tool in request.tools if tool.type == "mcp"
+                }
                 await context.init_tool_sessions(self.tool_server, exit_stack,
-                                                 request.request_id)
+                                                 request.request_id, mcp_tools)
                 processer = self._process_harmony_streaming_events
             else:
                 processer = self._process_simple_streaming_events

vllm/entrypoints/tool_server.py

@@ -18,7 +18,6 @@
 async def list_server_and_tools(server_url: str):
     from mcp import ClientSession
     from mcp.client.sse import sse_client
-
     async with sse_client(url=server_url) as streams, ClientSession(
             *streams) as session:
         initialize_response = await session.initialize()
@@ -155,14 +154,14 @@ async def new_session(self,
         from mcp import ClientSession
         from mcp.client.sse import sse_client
         url = self.urls.get(tool_name)
-        if headers is None:
-            headers = {}
-        headers = {"x-session-id": session_id, **headers}
+        request_headers = {"x-session-id": session_id}
+        if headers is not None:
+            request_headers.update(headers)
         if not url:
             raise KeyError(f"Tool '{tool_name}' is not supported")
-        async with sse_client(url=url,
-                              headers=headers) as streams, ClientSession(
-                                  *streams) as session:
+        async with sse_client(
+                url=url, headers=request_headers) as streams, ClientSession(
+                    *streams) as session:
             await session.initialize()
             yield session
 
@@ -198,7 +197,10 @@ def get_tool_description(self,
             raise ValueError(f"Unknown tool {tool_name}")
 
     @asynccontextmanager
-    async def new_session(self, tool_name: str, session_id: str):
+    async def new_session(self,
+                          tool_name: str,
+                          session_id: str,
+                          headers: Optional[dict[str, str]] = None):
         if tool_name not in self.tools:
             raise KeyError(f"Tool '{tool_name}' is not supported")
         yield self.tools[tool_name]