phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

cleanup: simplify connection code

Browse source 
Signed-off-by: Calum Murray <cmurray@redhat.com>
This commit is contained in:
Calum Murray 2025-07-10 14:15:29 -04:00
parent 873bf8d95a
commit caecaa8b31
No known key found for this signature in database
GPG key ID: B67F01AEB13FE187
1 changed files with 14 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

20
llama_stack/providers/utils/tools/mcp.py
View file

@ -46,17 +46,19 @@ async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerat
for i, strategy in enumerate(connection_strategies): for i, strategy in enumerate(connection_strategies):
try: try:
if strategy == PROTOCOL_STREAMABLEHTTP: if strategy == PROTOCOL_STREAMABLEHTTP:
async with streamablehttp_client(endpoint, headers=headers) as (read_stream, write_stream, _): client = streamablehttp_client
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
protocol_cache[endpoint] = PROTOCOL_STREAMABLEHTTP
yield session
return
elif strategy == PROTOCOL_SSE: elif strategy == PROTOCOL_SSE:
async with sse_client(endpoint, headers=headers) as streams: client = sse_client
async with ClientSession(*streams) as session: else:
# this should not happen
logger.warning(
"tried to establish MCP connection with unknown protocol, defaulting to try with streamable_http"
)
client = streamablehttp_client
async with client(endpoint, headers=headers) as client_streams:
async with ClientSession(read_stream=client_streams[0], write_stream=client_streams[1]) as session:
await session.initialize() await session.initialize()
protocol_cache[endpoint] = PROTOCOL_SSE protocol_cache[endpoint] = strategy
yield session yield session
return return
except* httpx.HTTPStatusError as eg: except* httpx.HTTPStatusError as eg: