[gpt-oss] Harmony changes with container tool support (#23386)

Signed-off-by: zhiweiz <zhiweiz@fb.com>
Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
Signed-off-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Co-authored-by: zhiweiz <zhiweiz@fb.com>
Co-authored-by: Aaron Pham <contact@aarnphm.xyz>
Co-authored-by: Simon Mo <simon.mo@hey.com>
Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
This commit is contained in:
zhiweiz
2025-09-08 19:03:50 -07:00
committed by GitHub
parent 955c624915
commit 170129eb28
5 changed files with 170 additions and 27 deletions

View File

@@ -86,7 +86,8 @@ class ToolServer(ABC):
pass
@abstractmethod
def new_session(self, tool_name: str) -> AbstractAsyncContextManager[Any]:
def new_session(self, tool_name: str,
session_id: str) -> AbstractAsyncContextManager[Any]:
"""
Create a session for the tool.
"""
@@ -124,7 +125,8 @@ class MCPToolServer(ToolServer):
description=tool.description,
parameters=tool.inputSchema)
for tool in list_tools_response.tools
])
],
)
self.harmony_tool_descriptions[tool_from_mcp.name] = tool_from_mcp
if tool_from_mcp.name not in self.urls:
self.urls[tool_from_mcp.name] = url
@@ -142,14 +144,16 @@ class MCPToolServer(ToolServer):
return self.harmony_tool_descriptions.get(tool_name)
@asynccontextmanager
async def new_session(self, tool_name: str):
async def new_session(self, tool_name: str, session_id: str):
from mcp import ClientSession
from mcp.client.sse import sse_client
url = self.urls.get(tool_name)
headers = {"x-session-id": session_id}
if not url:
raise KeyError(f"Tool '{tool_name}' is not supported")
async with sse_client(url=url) as streams, ClientSession(
*streams) as session:
async with sse_client(url=url,
headers=headers) as streams, ClientSession(
*streams) as session:
await session.initialize()
yield session
@@ -182,7 +186,7 @@ class DemoToolServer(ToolServer):
raise ValueError(f"Unknown tool {tool_name}")
@asynccontextmanager
async def new_session(self, tool_name: str):
async def new_session(self, tool_name: str, session_id: str):
if tool_name not in self.tools:
raise KeyError(f"Tool '{tool_name}' is not supported")
yield self.tools[tool_name]