mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-29 15:23:51 +00:00
cleanup: fix nits from review
Signed-off-by: Calum Murray <cmurray@redhat.com>
This commit is contained in:
parent
ffbea13e07
commit
033160b277
1 changed files with 3 additions and 4 deletions
|
@ -44,8 +44,7 @@ async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerat
|
||||||
if err.response.status_code == 401:
|
if err.response.status_code == 401:
|
||||||
raise AuthenticationRequiredError(exc) from exc
|
raise AuthenticationRequiredError(exc) from exc
|
||||||
except* McpError:
|
except* McpError:
|
||||||
logger.debug("failed to connect via streamable http, falling back to sse")
|
logger.warning("failed to connect via streamable http, falling back to sse")
|
||||||
# fallback to sse
|
|
||||||
try:
|
try:
|
||||||
async with sse_client(endpoint, headers=headers) as streams:
|
async with sse_client(endpoint, headers=headers) as streams:
|
||||||
async with ClientSession(*streams) as session:
|
async with ClientSession(*streams) as session:
|
||||||
|
@ -65,7 +64,7 @@ async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerat
|
||||||
async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefsResponse:
|
async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefsResponse:
|
||||||
tools = []
|
tools = []
|
||||||
async with client_wrapper(endpoint, headers) as session:
|
async with client_wrapper(endpoint, headers) as session:
|
||||||
logger.info("listing mcp tools...")
|
logger.debug("listing mcp tools")
|
||||||
tools_result = await session.list_tools()
|
tools_result = await session.list_tools()
|
||||||
for tool in tools_result.tools:
|
for tool in tools_result.tools:
|
||||||
parameters = []
|
parameters = []
|
||||||
|
@ -94,7 +93,7 @@ async def invoke_mcp_tool(
|
||||||
endpoint: str, headers: dict[str, str], tool_name: str, kwargs: dict[str, Any]
|
endpoint: str, headers: dict[str, str], tool_name: str, kwargs: dict[str, Any]
|
||||||
) -> ToolInvocationResult:
|
) -> ToolInvocationResult:
|
||||||
async with client_wrapper(endpoint, headers) as session:
|
async with client_wrapper(endpoint, headers) as session:
|
||||||
logger.info("invoking mcp tool")
|
logger.debug("invoking mcp tool")
|
||||||
result = await session.call_tool(tool_name, kwargs)
|
result = await session.call_tool(tool_name, kwargs)
|
||||||
|
|
||||||
content: list[InterleavedContentItem] = []
|
content: list[InterleavedContentItem] = []
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue