phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 04:04:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

precommit fix

Browse source
This commit is contained in:
Kai Wu 2025-09-30 09:22:58 -07:00
parent 9adabb09cf
commit 774e4311ad
2 changed files with 88 additions and 210 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -4,11 +4,11 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import copy
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from enum import Enum from enum import Enum
from typing import Any, cast from typing import Any, cast
import copy
import httpx import httpx
from mcp import ClientSession, McpError from mcp import ClientSession, McpError
@ -32,10 +32,9 @@ logger = get_logger(__name__, category="tools")
protocol_cache = TTLDict(ttl_seconds=3600) protocol_cache = TTLDict(ttl_seconds=3600)
def resolve_json_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: def resolve_json_schema_refs(schema: Any) -> Any:
""" """
Resolve JSON Schema $ref references using $defs. Resolve JSON Schema $ref references using $defs.
This function takes a JSON schema that may contain $ref and $defs, This function takes a JSON schema that may contain $ref and $defs,
and returns a new schema with all $ref references resolved inline. and returns a new schema with all $ref references resolved inline.
""" """
@ -53,7 +52,7 @@ def resolve_json_schema_refs(schema: dict[str, Any]) -> dict[str, Any]:
# Extract the reference path # Extract the reference path
ref_path = obj["$ref"] ref_path = obj["$ref"]
if ref_path.startswith("#/$defs/"): if ref_path.startswith("#/$defs/"):
def_name = ref_path[len("#/$defs/"):] def_name = ref_path[len("#/$defs/") :]
if def_name in defs: if def_name in defs:
# Recursively resolve refs in the definition itself # Recursively resolve refs in the definition itself
resolved_def = resolve_refs(defs[def_name]) resolved_def = resolve_refs(defs[def_name])

191
tests/unit/providers/utils/tools/test_mcp.py
View file

@ -4,25 +4,14 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import pytest from unittest.mock import AsyncMock, Mock, patch
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator
import httpx
from mcp import McpError
from mcp import types as mcp_types
from llama_stack.apis.tools import ListToolDefsResponse
from llama_stack.providers.utils.tools.mcp import ( from llama_stack.providers.utils.tools.mcp import (
resolve_json_schema_refs,
MCPProtol, MCPProtol,
client_wrapper,
list_mcp_tools, list_mcp_tools,
invoke_mcp_tool, resolve_json_schema_refs,
protocol_cache,
) )
from llama_stack.apis.tools import ListToolDefsResponse, ToolDef, ToolParameter, ToolInvocationResult
from llama_stack.apis.common.content_types import TextContentItem, ImageContentItem
from llama_stack.core.datatypes import AuthenticationRequiredError
class TestResolveJsonSchemaRefs: class TestResolveJsonSchemaRefs:
@ -32,18 +21,10 @@ class TestResolveJsonSchemaRefs:
"""Test resolving a simple $ref reference.""" """Test resolving a simple $ref reference."""
schema = { schema = {
"type": "object", "type": "object",
"properties": { "properties": {"user": {"$ref": "#/$defs/User"}},
"user": {"$ref": "#/$defs/User"}
},
"$defs": { "$defs": {
"User": { "User": {"type": "object", "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}}
"type": "object", },
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"}
}
}
}
} }
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
@ -51,14 +32,8 @@ class TestResolveJsonSchemaRefs:
expected = { expected = {
"type": "object", "type": "object",
"properties": { "properties": {
"user": { "user": {"type": "object", "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}}
"type": "object", },
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"}
}
}
}
} }
assert result == expected assert result == expected
@ -67,23 +42,11 @@ class TestResolveJsonSchemaRefs:
"""Test resolving nested $ref references.""" """Test resolving nested $ref references."""
schema = { schema = {
"type": "object", "type": "object",
"properties": { "properties": {"data": {"$ref": "#/$defs/Container"}},
"data": {"$ref": "#/$defs/Container"}
},
"$defs": { "$defs": {
"Container": { "Container": {"type": "object", "properties": {"user": {"$ref": "#/$defs/User"}}},
"type": "object", "User": {"type": "object", "properties": {"name": {"type": "string"}}},
"properties": {
"user": {"$ref": "#/$defs/User"}
}
}, },
"User": {
"type": "object",
"properties": {
"name": {"type": "string"}
}
}
}
} }
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
@ -93,16 +56,9 @@ class TestResolveJsonSchemaRefs:
"properties": { "properties": {
"data": { "data": {
"type": "object", "type": "object",
"properties": { "properties": {"user": {"type": "object", "properties": {"name": {"type": "string"}}}},
"user": {
"type": "object",
"properties": {
"name": {"type": "string"}
}
}
}
}
} }
},
} }
assert result == expected assert result == expected
@ -111,20 +67,8 @@ class TestResolveJsonSchemaRefs:
"""Test resolving $ref references within arrays.""" """Test resolving $ref references within arrays."""
schema = { schema = {
"type": "object", "type": "object",
"properties": { "properties": {"items": {"type": "array", "items": {"$ref": "#/$defs/Item"}}},
"items": { "$defs": {"Item": {"type": "object", "properties": {"id": {"type": "string"}}}},
"type": "array",
"items": {"$ref": "#/$defs/Item"}
}
},
"$defs": {
"Item": {
"type": "object",
"properties": {
"id": {"type": "string"}
}
}
}
} }
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
@ -132,16 +76,8 @@ class TestResolveJsonSchemaRefs:
expected = { expected = {
"type": "object", "type": "object",
"properties": { "properties": {
"items": { "items": {"type": "array", "items": {"type": "object", "properties": {"id": {"type": "string"}}}}
"type": "array", },
"items": {
"type": "object",
"properties": {
"id": {"type": "string"}
}
}
}
}
} }
assert result == expected assert result == expected
@ -150,65 +86,35 @@ class TestResolveJsonSchemaRefs:
"""Test handling of missing $ref definition.""" """Test handling of missing $ref definition."""
schema = { schema = {
"type": "object", "type": "object",
"properties": { "properties": {"user": {"$ref": "#/$defs/MissingUser"}},
"user": {"$ref": "#/$defs/MissingUser"} "$defs": {"User": {"type": "object", "properties": {"name": {"type": "string"}}}},
},
"$defs": {
"User": {
"type": "object",
"properties": {
"name": {"type": "string"}
}
}
}
} }
with patch('llama_stack.providers.utils.tools.mcp.logger') as mock_logger: with patch("llama_stack.providers.utils.tools.mcp.logger") as mock_logger:
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
mock_logger.warning.assert_called_once_with("Referenced definition 'MissingUser' not found in $defs") mock_logger.warning.assert_called_once_with("Referenced definition 'MissingUser' not found in $defs")
# Should return the original $ref unchanged # Should return the original $ref unchanged
expected = { expected = {"type": "object", "properties": {"user": {"$ref": "#/$defs/MissingUser"}}}
"type": "object",
"properties": {
"user": {"$ref": "#/$defs/MissingUser"}
}
}
assert result == expected assert result == expected
def test_resolve_unsupported_ref_format(self): def test_resolve_unsupported_ref_format(self):
"""Test handling of unsupported $ref format.""" """Test handling of unsupported $ref format."""
schema = { schema = {"type": "object", "properties": {"user": {"$ref": "http://example.com/schema"}}, "$defs": {}}
"type": "object",
"properties": {
"user": {"$ref": "http://example.com/schema"}
},
"$defs": {}
}
with patch('llama_stack.providers.utils.tools.mcp.logger') as mock_logger: with patch("llama_stack.providers.utils.tools.mcp.logger") as mock_logger:
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
mock_logger.warning.assert_called_once_with("Unsupported $ref format: http://example.com/schema") mock_logger.warning.assert_called_once_with("Unsupported $ref format: http://example.com/schema")
# Should return the original $ref unchanged # Should return the original $ref unchanged
expected = { expected = {"type": "object", "properties": {"user": {"$ref": "http://example.com/schema"}}}
"type": "object",
"properties": {
"user": {"$ref": "http://example.com/schema"}
}
}
assert result == expected assert result == expected
def test_resolve_no_defs(self): def test_resolve_no_defs(self):
"""Test schema without $defs section.""" """Test schema without $defs section."""
schema = { schema = {"type": "object", "properties": {"name": {"type": "string"}}}
"type": "object",
"properties": {
"name": {"type": "string"}
}
}
result = resolve_json_schema_refs(schema) result = resolve_json_schema_refs(schema)
@ -225,20 +131,10 @@ class TestResolveJsonSchemaRefs:
"""Test that original schema is not modified.""" """Test that original schema is not modified."""
original_schema = { original_schema = {
"type": "object", "type": "object",
"properties": { "properties": {"user": {"$ref": "#/$defs/User"}},
"user": {"$ref": "#/$defs/User"} "$defs": {"User": {"type": "object", "properties": {"name": {"type": "string"}}}},
},
"$defs": {
"User": {
"type": "object",
"properties": {
"name": {"type": "string"}
}
}
}
} }
original_copy = original_schema.copy()
resolve_json_schema_refs(original_schema) resolve_json_schema_refs(original_schema)
# Original should be unchanged (but this is a shallow comparison) # Original should be unchanged (but this is a shallow comparison)
@ -259,7 +155,6 @@ class TestMCPProtocol:
class TestListMcpTools: class TestListMcpTools:
"""Test cases for list_mcp_tools function.""" """Test cases for list_mcp_tools function."""
@pytest.mark.asyncio
async def test_list_tools_success(self): async def test_list_tools_success(self):
"""Test successful listing of MCP tools.""" """Test successful listing of MCP tools."""
endpoint = "http://example.com/mcp" endpoint = "http://example.com/mcp"
@ -272,16 +167,9 @@ class TestListMcpTools:
mock_tool.inputSchema = { mock_tool.inputSchema = {
"type": "object", "type": "object",
"properties": { "properties": {
"param1": { "param1": {"type": "string", "description": "First parameter", "default": "default_value"},
"type": "string", "param2": {"type": "integer", "description": "Second parameter"},
"description": "First parameter",
"default": "default_value"
}, },
"param2": {
"type": "integer",
"description": "Second parameter"
}
}
} }
mock_tools_result = Mock() mock_tools_result = Mock()
@ -290,7 +178,7 @@ class TestListMcpTools:
mock_session = Mock() mock_session = Mock()
mock_session.list_tools = AsyncMock(return_value=mock_tools_result) mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
with patch('llama_stack.providers.utils.tools.mcp.client_wrapper') as mock_wrapper: with patch("llama_stack.providers.utils.tools.mcp.client_wrapper") as mock_wrapper:
mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session) mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_wrapper.return_value.__aexit__ = AsyncMock() mock_wrapper.return_value.__aexit__ = AsyncMock()
@ -318,7 +206,6 @@ class TestListMcpTools:
assert param2.description == "Second parameter" assert param2.description == "Second parameter"
assert param2.required is True # No default value assert param2.required is True # No default value
@pytest.mark.asyncio
async def test_list_tools_with_schema_refs(self): async def test_list_tools_with_schema_refs(self):
"""Test listing tools with JSON Schema $refs.""" """Test listing tools with JSON Schema $refs."""
endpoint = "http://example.com/mcp" endpoint = "http://example.com/mcp"
@ -330,17 +217,10 @@ class TestListMcpTools:
mock_tool.description = "Tool with refs" mock_tool.description = "Tool with refs"
mock_tool.inputSchema = { mock_tool.inputSchema = {
"type": "object", "type": "object",
"properties": { "properties": {"user": {"$ref": "#/$defs/User"}},
"user": {"$ref": "#/$defs/User"}
},
"$defs": { "$defs": {
"User": { "User": {"type": "object", "properties": {"name": {"type": "string", "description": "User name"}}}
"type": "object", },
"properties": {
"name": {"type": "string", "description": "User name"}
}
}
}
} }
mock_tools_result = Mock() mock_tools_result = Mock()
@ -349,7 +229,7 @@ class TestListMcpTools:
mock_session = Mock() mock_session = Mock()
mock_session.list_tools = AsyncMock(return_value=mock_tools_result) mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
with patch('llama_stack.providers.utils.tools.mcp.client_wrapper') as mock_wrapper: with patch("llama_stack.providers.utils.tools.mcp.client_wrapper") as mock_wrapper:
mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session) mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_wrapper.return_value.__aexit__ = AsyncMock() mock_wrapper.return_value.__aexit__ = AsyncMock()
@ -362,7 +242,6 @@ class TestListMcpTools:
# The user parameter should be flattened from the resolved $ref # The user parameter should be flattened from the resolved $ref
# Note: This depends on how the schema resolution works with nested objects # Note: This depends on how the schema resolution works with nested objects
@pytest.mark.asyncio
async def test_list_tools_empty_result(self): async def test_list_tools_empty_result(self):
"""Test listing tools when no tools are available.""" """Test listing tools when no tools are available."""
endpoint = "http://example.com/mcp" endpoint = "http://example.com/mcp"
@ -374,7 +253,7 @@ class TestListMcpTools:
mock_session = Mock() mock_session = Mock()
mock_session.list_tools = AsyncMock(return_value=mock_tools_result) mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
with patch('llama_stack.providers.utils.tools.mcp.client_wrapper') as mock_wrapper: with patch("llama_stack.providers.utils.tools.mcp.client_wrapper") as mock_wrapper:
mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session) mock_wrapper.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_wrapper.return_value.__aexit__ = AsyncMock() mock_wrapper.return_value.__aexit__ = AsyncMock()