mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-04 20:14:13 +00:00
ef0736527d
This is a sweeping change to clean up some gunk around our "Tool" definitions. First, we had two types `Tool` and `ToolDef`. The first of these was a "Resource" type for the registry but we had stopped registering tools inside the Registry long back (and only registered ToolGroups.) The latter was for specifying tools for the Agents API. This PR removes the former and adds an optional `toolgroup_id` field to the latter. Secondly, as pointed out by @bbrowning in https://github.com/llamastack/llama-stack/pull/3003#issuecomment-3245270132, we were doing a lossy conversion from a full JSON schema from the MCP tool specification into our ToolDefinition to send it to the model. There is no necessity to do this -- we ourselves aren't doing any execution at all but merely passing it to the chat completions API which supports this. By doing this (and by doing it poorly), we encountered limitations like not supporting array items, or not resolving $refs, etc. To fix this, we replaced the `parameters` field by `{ input_schema, output_schema }` which can be full blown JSON schemas. Finally, there were some types in our llama-related chat format conversion which needed some cleanup. We are taking this opportunity to clean those up. This PR is a substantial breaking change to the API. However, given our window for introducing breaking changes, this suits us just fine. I will be landing a concurrent `llama-stack-client` change as well since API shapes are changing.
297 lines
11 KiB
Python
297 lines
11 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.
|
|
|
|
"""
|
|
Unit tests for JSON Schema-based tool definitions.
|
|
Tests the new input_schema and output_schema fields.
|
|
"""
|
|
|
|
from pydantic import ValidationError
|
|
|
|
from llama_stack.apis.tools import ToolDef
|
|
from llama_stack.models.llama.datatypes import BuiltinTool, ToolDefinition
|
|
|
|
|
|
class TestToolDefValidation:
|
|
"""Test ToolDef validation with JSON Schema."""
|
|
|
|
def test_simple_input_schema(self):
|
|
"""Test ToolDef with simple input schema."""
|
|
tool = ToolDef(
|
|
name="get_weather",
|
|
description="Get weather information",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"location": {"type": "string", "description": "City name"}},
|
|
"required": ["location"],
|
|
},
|
|
)
|
|
|
|
assert tool.name == "get_weather"
|
|
assert tool.input_schema["type"] == "object"
|
|
assert "location" in tool.input_schema["properties"]
|
|
assert tool.output_schema is None
|
|
|
|
def test_input_and_output_schema(self):
|
|
"""Test ToolDef with both input and output schemas."""
|
|
tool = ToolDef(
|
|
name="calculate",
|
|
description="Perform calculation",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"x": {"type": "number"}, "y": {"type": "number"}},
|
|
"required": ["x", "y"],
|
|
},
|
|
output_schema={"type": "object", "properties": {"result": {"type": "number"}}, "required": ["result"]},
|
|
)
|
|
|
|
assert tool.input_schema is not None
|
|
assert tool.output_schema is not None
|
|
assert "result" in tool.output_schema["properties"]
|
|
|
|
def test_schema_with_refs_and_defs(self):
|
|
"""Test that $ref and $defs are preserved in schemas."""
|
|
tool = ToolDef(
|
|
name="book_flight",
|
|
description="Book a flight",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"flight": {"$ref": "#/$defs/FlightInfo"},
|
|
"passengers": {"type": "array", "items": {"$ref": "#/$defs/Passenger"}},
|
|
},
|
|
"$defs": {
|
|
"FlightInfo": {
|
|
"type": "object",
|
|
"properties": {"from": {"type": "string"}, "to": {"type": "string"}},
|
|
},
|
|
"Passenger": {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
|
|
},
|
|
},
|
|
},
|
|
)
|
|
|
|
# Verify $defs are preserved
|
|
assert "$defs" in tool.input_schema
|
|
assert "FlightInfo" in tool.input_schema["$defs"]
|
|
assert "Passenger" in tool.input_schema["$defs"]
|
|
|
|
# Verify $ref are preserved
|
|
assert tool.input_schema["properties"]["flight"]["$ref"] == "#/$defs/FlightInfo"
|
|
assert tool.input_schema["properties"]["passengers"]["items"]["$ref"] == "#/$defs/Passenger"
|
|
|
|
def test_output_schema_with_refs(self):
|
|
"""Test that output_schema also supports $ref and $defs."""
|
|
tool = ToolDef(
|
|
name="search",
|
|
description="Search for items",
|
|
input_schema={"type": "object", "properties": {"query": {"type": "string"}}},
|
|
output_schema={
|
|
"type": "object",
|
|
"properties": {"results": {"type": "array", "items": {"$ref": "#/$defs/SearchResult"}}},
|
|
"$defs": {
|
|
"SearchResult": {
|
|
"type": "object",
|
|
"properties": {"title": {"type": "string"}, "score": {"type": "number"}},
|
|
}
|
|
},
|
|
},
|
|
)
|
|
|
|
assert "$defs" in tool.output_schema
|
|
assert "SearchResult" in tool.output_schema["$defs"]
|
|
|
|
def test_complex_json_schema_features(self):
|
|
"""Test various JSON Schema features are preserved."""
|
|
tool = ToolDef(
|
|
name="complex_tool",
|
|
description="Tool with complex schema",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {
|
|
# anyOf
|
|
"contact": {
|
|
"anyOf": [
|
|
{"type": "string", "format": "email"},
|
|
{"type": "string", "pattern": "^\\+?[0-9]{10,15}$"},
|
|
]
|
|
},
|
|
# enum
|
|
"status": {"type": "string", "enum": ["pending", "approved", "rejected"]},
|
|
# nested objects
|
|
"address": {
|
|
"type": "object",
|
|
"properties": {
|
|
"street": {"type": "string"},
|
|
"city": {"type": "string"},
|
|
"zipcode": {"type": "string", "pattern": "^[0-9]{5}$"},
|
|
},
|
|
"required": ["street", "city"],
|
|
},
|
|
# array with constraints
|
|
"tags": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"minItems": 1,
|
|
"maxItems": 10,
|
|
"uniqueItems": True,
|
|
},
|
|
},
|
|
},
|
|
)
|
|
|
|
# Verify anyOf
|
|
assert "anyOf" in tool.input_schema["properties"]["contact"]
|
|
|
|
# Verify enum
|
|
assert tool.input_schema["properties"]["status"]["enum"] == ["pending", "approved", "rejected"]
|
|
|
|
# Verify nested object
|
|
assert tool.input_schema["properties"]["address"]["type"] == "object"
|
|
assert "zipcode" in tool.input_schema["properties"]["address"]["properties"]
|
|
|
|
# Verify array constraints
|
|
tags_schema = tool.input_schema["properties"]["tags"]
|
|
assert tags_schema["minItems"] == 1
|
|
assert tags_schema["maxItems"] == 10
|
|
assert tags_schema["uniqueItems"] is True
|
|
|
|
def test_invalid_json_schema_raises_error(self):
|
|
"""Test that invalid JSON Schema raises validation error."""
|
|
# TODO: This test will pass once we add schema validation
|
|
# For now, Pydantic accepts any dict, so this is a placeholder
|
|
|
|
# This should eventually raise an error due to invalid schema
|
|
try:
|
|
ToolDef(
|
|
name="bad_tool",
|
|
input_schema={
|
|
"type": "invalid_type", # Not a valid JSON Schema type
|
|
"properties": "not_an_object", # Should be an object
|
|
},
|
|
)
|
|
# For now this passes, but shouldn't after we add validation
|
|
except ValidationError:
|
|
pass # Expected once validation is added
|
|
|
|
|
|
class TestToolDefinitionValidation:
|
|
"""Test ToolDefinition (internal) validation with JSON Schema."""
|
|
|
|
def test_simple_tool_definition(self):
|
|
"""Test ToolDefinition with simple schema."""
|
|
tool = ToolDefinition(
|
|
tool_name="get_time",
|
|
description="Get current time",
|
|
input_schema={"type": "object", "properties": {"timezone": {"type": "string"}}},
|
|
)
|
|
|
|
assert tool.tool_name == "get_time"
|
|
assert tool.input_schema is not None
|
|
|
|
def test_builtin_tool_with_schema(self):
|
|
"""Test ToolDefinition with BuiltinTool enum."""
|
|
tool = ToolDefinition(
|
|
tool_name=BuiltinTool.code_interpreter,
|
|
description="Run Python code",
|
|
input_schema={"type": "object", "properties": {"code": {"type": "string"}}, "required": ["code"]},
|
|
output_schema={"type": "object", "properties": {"output": {"type": "string"}, "error": {"type": "string"}}},
|
|
)
|
|
|
|
assert isinstance(tool.tool_name, BuiltinTool)
|
|
assert tool.input_schema is not None
|
|
assert tool.output_schema is not None
|
|
|
|
def test_tool_definition_with_refs(self):
|
|
"""Test ToolDefinition preserves $ref/$defs."""
|
|
tool = ToolDefinition(
|
|
tool_name="process_data",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"data": {"$ref": "#/$defs/DataObject"}},
|
|
"$defs": {
|
|
"DataObject": {
|
|
"type": "object",
|
|
"properties": {
|
|
"id": {"type": "integer"},
|
|
"values": {"type": "array", "items": {"type": "number"}},
|
|
},
|
|
}
|
|
},
|
|
},
|
|
)
|
|
|
|
assert "$defs" in tool.input_schema
|
|
assert tool.input_schema["properties"]["data"]["$ref"] == "#/$defs/DataObject"
|
|
|
|
|
|
class TestSchemaEquivalence:
|
|
"""Test that schemas remain unchanged through serialization."""
|
|
|
|
def test_schema_roundtrip(self):
|
|
"""Test that schemas survive model_dump/model_validate roundtrip."""
|
|
original = ToolDef(
|
|
name="test",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"x": {"$ref": "#/$defs/X"}},
|
|
"$defs": {"X": {"type": "string"}},
|
|
},
|
|
)
|
|
|
|
# Serialize and deserialize
|
|
dumped = original.model_dump()
|
|
restored = ToolDef(**dumped)
|
|
|
|
# Schemas should be identical
|
|
assert restored.input_schema == original.input_schema
|
|
assert "$defs" in restored.input_schema
|
|
assert restored.input_schema["properties"]["x"]["$ref"] == "#/$defs/X"
|
|
|
|
def test_json_serialization(self):
|
|
"""Test JSON serialization preserves schema."""
|
|
import json
|
|
|
|
tool = ToolDef(
|
|
name="test",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "string"}},
|
|
"$defs": {"T": {"type": "number"}},
|
|
},
|
|
output_schema={"type": "object", "properties": {"b": {"$ref": "#/$defs/T"}}},
|
|
)
|
|
|
|
# Serialize to JSON and back
|
|
json_str = tool.model_dump_json()
|
|
parsed = json.loads(json_str)
|
|
restored = ToolDef(**parsed)
|
|
|
|
assert restored.input_schema == tool.input_schema
|
|
assert restored.output_schema == tool.output_schema
|
|
assert "$defs" in restored.input_schema
|
|
|
|
|
|
class TestBackwardsCompatibility:
|
|
"""Test handling of legacy code patterns."""
|
|
|
|
def test_none_schemas(self):
|
|
"""Test tools with no schemas (legacy case)."""
|
|
tool = ToolDef(name="legacy_tool", description="Tool without schemas", input_schema=None, output_schema=None)
|
|
|
|
assert tool.input_schema is None
|
|
assert tool.output_schema is None
|
|
|
|
def test_metadata_preserved(self):
|
|
"""Test that metadata field still works."""
|
|
tool = ToolDef(
|
|
name="test", input_schema={"type": "object"}, metadata={"endpoint": "http://example.com", "version": "1.0"}
|
|
)
|
|
|
|
assert tool.metadata["endpoint"] == "http://example.com"
|
|
assert tool.metadata["version"] == "1.0"
|