feat(tools): use { input_schema, output_schema } for ToolDefinition

tests/unit/tools/test_tools_json_schema.py

@@ -0,0 +1,297 @@
+# 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"