From fbe3c9c588ecd1f7c1eef2fe2c4a21df6953ae08 Mon Sep 17 00:00:00 2001 From: Philip Date: Thu, 13 Mar 2025 16:15:11 -0400 Subject: [PATCH] Fix to issue 1586 --- .../model_context_protocol/model_context_protocol.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py b/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py index f7dc376f8..def69327d 100644 --- a/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py +++ b/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py @@ -36,7 +36,12 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime): raise ValueError("mcp_endpoint is required") tools = [] - async with sse_client(mcp_endpoint.uri) as streams: + if hasattr(mcp_endpoint, "__getitem__"): + if mcp_endpoint["uri"] is not None: + uri = mcp_endpoint["uri"] + else: + uri = mcp_endpoint.uri + async with sse_client(uri) as streams: async with ClientSession(*streams) as session: await session.initialize() tools_result = await session.list_tools() @@ -56,7 +61,7 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime): description=tool.description, parameters=parameters, metadata={ - "endpoint": mcp_endpoint.uri, + "endpoint": uri, }, ) )