mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-03 19:57:35 +00:00
384 lines
12 KiB
Python
384 lines
12 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import pytest
|
|
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.providers.utils.tools.mcp import (
|
|
resolve_json_schema_refs,
|
|
MCPProtol,
|
|
client_wrapper,
|
|
list_mcp_tools,
|
|
invoke_mcp_tool,
|
|
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:
|
|
"""Test cases for resolve_json_schema_refs function."""
|
|
|
|
def test_resolve_simple_ref(self):
|
|
"""Test resolving a simple $ref reference."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/User"}
|
|
},
|
|
"$defs": {
|
|
"User": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"age": {"type": "integer"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
result = resolve_json_schema_refs(schema)
|
|
|
|
expected = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"age": {"type": "integer"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert result == expected
|
|
|
|
def test_resolve_nested_refs(self):
|
|
"""Test resolving nested $ref references."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"data": {"$ref": "#/$defs/Container"}
|
|
},
|
|
"$defs": {
|
|
"Container": {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/User"}
|
|
}
|
|
},
|
|
"User": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
result = resolve_json_schema_refs(schema)
|
|
|
|
expected = {
|
|
"type": "object",
|
|
"properties": {
|
|
"data": {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert result == expected
|
|
|
|
def test_resolve_refs_in_array(self):
|
|
"""Test resolving $ref references within arrays."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"items": {
|
|
"type": "array",
|
|
"items": {"$ref": "#/$defs/Item"}
|
|
}
|
|
},
|
|
"$defs": {
|
|
"Item": {
|
|
"type": "object",
|
|
"properties": {
|
|
"id": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
result = resolve_json_schema_refs(schema)
|
|
|
|
expected = {
|
|
"type": "object",
|
|
"properties": {
|
|
"items": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"id": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert result == expected
|
|
|
|
def test_resolve_missing_ref(self):
|
|
"""Test handling of missing $ref definition."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/MissingUser"}
|
|
},
|
|
"$defs": {
|
|
"User": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
with patch('llama_stack.providers.utils.tools.mcp.logger') as mock_logger:
|
|
result = resolve_json_schema_refs(schema)
|
|
mock_logger.warning.assert_called_once_with("Referenced definition 'MissingUser' not found in $defs")
|
|
|
|
# Should return the original $ref unchanged
|
|
expected = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/MissingUser"}
|
|
}
|
|
}
|
|
|
|
assert result == expected
|
|
|
|
def test_resolve_unsupported_ref_format(self):
|
|
"""Test handling of unsupported $ref format."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "http://example.com/schema"}
|
|
},
|
|
"$defs": {}
|
|
}
|
|
|
|
with patch('llama_stack.providers.utils.tools.mcp.logger') as mock_logger:
|
|
result = resolve_json_schema_refs(schema)
|
|
mock_logger.warning.assert_called_once_with("Unsupported $ref format: http://example.com/schema")
|
|
|
|
# Should return the original $ref unchanged
|
|
expected = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "http://example.com/schema"}
|
|
}
|
|
}
|
|
|
|
assert result == expected
|
|
|
|
def test_resolve_no_defs(self):
|
|
"""Test schema without $defs section."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"}
|
|
}
|
|
}
|
|
|
|
result = resolve_json_schema_refs(schema)
|
|
|
|
assert result == schema
|
|
|
|
def test_resolve_non_dict_input(self):
|
|
"""Test with non-dictionary input."""
|
|
assert resolve_json_schema_refs("string") == "string"
|
|
assert resolve_json_schema_refs(123) == 123
|
|
assert resolve_json_schema_refs(["list"]) == ["list"]
|
|
assert resolve_json_schema_refs(None) is None
|
|
|
|
def test_resolve_preserves_original(self):
|
|
"""Test that original schema is not modified."""
|
|
original_schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/User"}
|
|
},
|
|
"$defs": {
|
|
"User": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
original_copy = original_schema.copy()
|
|
resolve_json_schema_refs(original_schema)
|
|
|
|
# Original should be unchanged (but this is a shallow comparison)
|
|
assert "$ref" in original_schema["properties"]["user"]
|
|
assert "$defs" in original_schema
|
|
|
|
|
|
class TestMCPProtocol:
|
|
"""Test cases for MCPProtol enum."""
|
|
|
|
def test_protocol_values(self):
|
|
"""Test enum values are correct."""
|
|
assert MCPProtol.UNKNOWN.value == 0
|
|
assert MCPProtol.STREAMABLE_HTTP.value == 1
|
|
assert MCPProtol.SSE.value == 2
|
|
|
|
|
|
class TestListMcpTools:
|
|
"""Test cases for list_mcp_tools function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_success(self):
|
|
"""Test successful listing of MCP tools."""
|
|
endpoint = "http://example.com/mcp"
|
|
headers = {"Authorization": "Bearer token"}
|
|
|
|
# Mock tool from MCP
|
|
mock_tool = Mock()
|
|
mock_tool.name = "test_tool"
|
|
mock_tool.description = "A test tool"
|
|
mock_tool.inputSchema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"param1": {
|
|
"type": "string",
|
|
"description": "First parameter",
|
|
"default": "default_value"
|
|
},
|
|
"param2": {
|
|
"type": "integer",
|
|
"description": "Second parameter"
|
|
}
|
|
}
|
|
}
|
|
|
|
mock_tools_result = Mock()
|
|
mock_tools_result.tools = [mock_tool]
|
|
|
|
mock_session = Mock()
|
|
mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
|
|
|
|
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.__aexit__ = AsyncMock()
|
|
|
|
result = await list_mcp_tools(endpoint, headers)
|
|
|
|
assert isinstance(result, ListToolDefsResponse)
|
|
assert len(result.data) == 1
|
|
|
|
tool_def = result.data[0]
|
|
assert tool_def.name == "test_tool"
|
|
assert tool_def.description == "A test tool"
|
|
assert tool_def.metadata["endpoint"] == endpoint
|
|
|
|
# Check parameters
|
|
assert len(tool_def.parameters) == 2
|
|
|
|
param1 = next(p for p in tool_def.parameters if p.name == "param1")
|
|
assert param1.parameter_type == "string"
|
|
assert param1.description == "First parameter"
|
|
assert param1.required is False # Has default value
|
|
assert param1.default == "default_value"
|
|
|
|
param2 = next(p for p in tool_def.parameters if p.name == "param2")
|
|
assert param2.parameter_type == "integer"
|
|
assert param2.description == "Second parameter"
|
|
assert param2.required is True # No default value
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_with_schema_refs(self):
|
|
"""Test listing tools with JSON Schema $refs."""
|
|
endpoint = "http://example.com/mcp"
|
|
headers = {}
|
|
|
|
# Mock tool with $ref in schema
|
|
mock_tool = Mock()
|
|
mock_tool.name = "ref_tool"
|
|
mock_tool.description = "Tool with refs"
|
|
mock_tool.inputSchema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"user": {"$ref": "#/$defs/User"}
|
|
},
|
|
"$defs": {
|
|
"User": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string", "description": "User name"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
mock_tools_result = Mock()
|
|
mock_tools_result.tools = [mock_tool]
|
|
|
|
mock_session = Mock()
|
|
mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
|
|
|
|
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.__aexit__ = AsyncMock()
|
|
|
|
result = await list_mcp_tools(endpoint, headers)
|
|
|
|
# Should have resolved the $ref
|
|
tool_def = result.data[0]
|
|
assert len(tool_def.parameters) == 1
|
|
|
|
# The user parameter should be flattened from the resolved $ref
|
|
# Note: This depends on how the schema resolution works with nested objects
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_empty_result(self):
|
|
"""Test listing tools when no tools are available."""
|
|
endpoint = "http://example.com/mcp"
|
|
headers = {}
|
|
|
|
mock_tools_result = Mock()
|
|
mock_tools_result.tools = []
|
|
|
|
mock_session = Mock()
|
|
mock_session.list_tools = AsyncMock(return_value=mock_tools_result)
|
|
|
|
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.__aexit__ = AsyncMock()
|
|
|
|
result = await list_mcp_tools(endpoint, headers)
|
|
|
|
assert isinstance(result, ListToolDefsResponse)
|
|
assert len(result.data) == 0
|