diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index 441658a0..826b70c2 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -1,15 +1,29 @@ """Module for transforming JSON schemas to Pydantic models.""" import logging -from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, cast +from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, Annotated, cast from atomic_agents.connectors.mcp.mcp_definition_service import MCPAttributeType -from pydantic import Field, create_model +from jsonschema import Draft202012Validator, validators +from pydantic import BeforeValidator, Field, create_model +from referencing import Registry +from referencing.exceptions import Unresolvable from atomic_agents.base.base_io_schema import BaseIOSchema logger = logging.getLogger(__name__) + +def _validate_ref(validator: Any, reference: str, value: Any, schema: Any) -> Any: + try: + yield from Draft202012Validator.VALIDATORS["$ref"](validator, reference, value, schema) + except Unresolvable: + # Match the transformer's Any fallback only for the unresolved reference. + return + + +_ScalarChoiceValidator = validators.extend(Draft202012Validator, {"$ref": _validate_ref}) + # JSON type mapping JSON_TYPE_MAP = { "string": str, @@ -25,6 +39,42 @@ class SchemaTransformer: """Class for transforming JSON schemas to Pydantic models.""" + @staticmethod + def _matches_scalar_type(value: Any, json_types: List[str]) -> bool: + """Match JSON scalar types without treating Python bools as numbers.""" + if isinstance(value, bool): + return "boolean" in json_types + if isinstance(value, (int, float)): + return "number" in json_types or ("integer" in json_types and (isinstance(value, int) or value.is_integer())) + return (value is None and "null" in json_types) or (isinstance(value, str) and "string" in json_types) + + @staticmethod + def _scalar_choice_type(prop_schema: Dict[str, Any], choices: List[Any], root_schema: Dict[str, Any]) -> Type: + """Keep only choices satisfying all advertised constraints, including compositions.""" + declared_type = prop_schema.get("type") + json_types = declared_type if isinstance(declared_type, list) else [declared_type] + # The finite set can be intersected once, without coercing values through union types. + # Bind local references to the original root, and never retrieve remote schemas. + validator = _ScalarChoiceValidator(root_schema, registry=Registry()).evolve(schema=prop_schema) + choices = [value for value in choices if validator.is_valid(value)] + if not choices: + raise ValueError("Scalar choices do not match the declared JSON Schema constraints") + + def validate_choice(value: Any) -> Any: + if declared_type is not None and not SchemaTransformer._matches_scalar_type(value, json_types): + raise ValueError("Value does not match the declared JSON Schema type") + for choice in choices: + if isinstance(value, bool) == isinstance(choice, bool) and value == choice: + # Give Literal the matching choice so 1.0 cannot resolve to True. + return choice + raise ValueError("Value does not match the JSON Schema choices") + + choice_type = Annotated[Literal[tuple(choices)], BeforeValidator(validate_choice)] + if declared_type is not None: + # Keep type metadata inside the choice branch when Optional wraps it. + choice_type = Annotated[choice_type, Field(json_schema_extra={"type": declared_type})] + return choice_type + @staticmethod def _resolve_ref(ref_path: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type]) -> Type: """Resolve a $ref to its Pydantic-compatible value type.""" @@ -100,6 +150,13 @@ def json_to_pydantic_field( # Handle $ref if "$ref" in prop_schema: python_type = SchemaTransformer._resolve_ref(prop_schema["$ref"], root_schema, model_cache) + # Scalar value constraints must also appear in the model-facing schema. + elif "const" in prop_schema and isinstance(prop_schema["const"], (str, int, float, bool, type(None))): + python_type = SchemaTransformer._scalar_choice_type(prop_schema, [prop_schema["const"]], root_schema) + elif prop_schema.get("enum") and all( + isinstance(value, (str, int, float, bool, type(None))) for value in prop_schema["enum"] + ): + python_type = SchemaTransformer._scalar_choice_type(prop_schema, prop_schema["enum"], root_schema) # Handle oneOf/anyOf (unions) elif "oneOf" in prop_schema or "anyOf" in prop_schema: union_schemas = prop_schema.get("oneOf", prop_schema.get("anyOf", [])) diff --git a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py index ff03d65a..fce7906b 100644 --- a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py +++ b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py @@ -1,6 +1,7 @@ import logging import pytest +from jsonschema import Draft202012Validator from typing import Any, Dict, List, Optional, Union from pydantic import ValidationError @@ -75,6 +76,269 @@ def test_no_type(self): class TestCreateModelFromSchema: + @pytest.mark.parametrize("required", [True, False]) + @pytest.mark.parametrize( + "prop_schema", + [ + {"enum": [1, "bad"], "anyOf": [{"type": "integer"}]}, + {"enum": [1, 2.5, "bad"], "oneOf": [{"type": "integer"}, {"type": "number"}]}, + {"enum": [1, 2, 3, "bad"], "allOf": [{"type": "number"}, {"minimum": 2}]}, + {"const": 1, "anyOf": [{"type": "integer"}]}, + {"const": 2.5, "oneOf": [{"type": "integer"}, {"type": "number"}]}, + {"const": 3, "allOf": [{"type": "number"}, {"minimum": 2}]}, + {"enum": [True, 1, "bad"], "anyOf": [{"type": "boolean"}, {"type": "integer"}]}, + {"enum": [1, 2], "const": 2}, + {"enum": [1, "bad"], "anyOf": [{"$ref": "#/$defs/Integer"}]}, + { + "enum": [0, 1, 2, 3], + "anyOf": [{"type": "integer"}, {"type": "string"}], + "oneOf": [{"const": 1}, {"const": 2}], + "allOf": [{"minimum": 2}], + }, + {"enum": [1, "bad"], "anyOf": [False, {"type": "integer"}]}, + {"enum": ["brief", None, 1], "anyOf": [{"type": "string"}, {"type": "null"}]}, + {"const": None, "anyOf": [{"type": "string"}, {"type": "null"}]}, + ], + ) + def test_scalar_choices_intersect_compositions_in_validation_and_emitted_schema(self, prop_schema, required): + schema = { + "type": "object", + "properties": {"value": prop_schema}, + "required": ["value"] if required else [], + "$defs": {"Integer": {"type": "integer"}}, + } + model = SchemaTransformer.create_model_from_schema(schema, "ComposedChoiceInput", "choose") + advertised = Draft202012Validator(schema) + emitted = Draft202012Validator(model.model_json_schema()) + for value in [None, True, False, 0, 1, 1.0, 2, 2.5, 3, "brief", "bad", "unknown", [], {}]: + payload = {"tool_name": "choose", "value": value} + # Optional fields retain the transformer's existing local None allowance. + allowed = advertised.is_valid(payload) or (not required and value is None) + assert emitted.is_valid(payload) == allowed + if allowed: + assert model.model_validate(payload).value == value + else: + with pytest.raises(ValidationError): + model.model_validate(payload) + + @pytest.mark.parametrize( + "prop_schema", + [ + {"const": "bad", "anyOf": [{"type": "integer"}]}, + {"const": 1, "oneOf": [{"type": "integer"}, {"type": "number"}]}, + {"enum": [1], "allOf": [{"minimum": 2}]}, + ], + ) + def test_composition_with_no_admissible_scalar_choice_rejects_definition(self, prop_schema): + schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} + with pytest.raises(ValueError, match="choices"): + SchemaTransformer.create_model_from_schema(schema, "ImpossibleChoiceInput", "choose") + + @pytest.mark.parametrize("ref_first", [True, False]) + def test_scalar_choice_composition_unresolved_reference_keeps_known_sibling_constraint(self, ref_first): + branch = {"$ref": "https://example.invalid/schema.json", "type": "integer"} + if not ref_first: + branch = dict(reversed(branch.items())) + schema = { + "type": "object", + "properties": {"value": {"enum": [1, "bad"], "anyOf": [branch]}}, + "required": ["value"], + } + model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose") + assert model(tool_name="choose", value=1).value == 1 + with pytest.raises(ValidationError): + model(tool_name="choose", value="bad") + assert not Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "choose", "value": "bad"}) + + def test_scalar_choice_composition_unresolved_reference_alone_preserves_fallback(self): + schema = { + "type": "object", + "properties": {"value": {"enum": [1, "brief"], "anyOf": [{"$ref": "#/$defs/Missing"}]}}, + "required": ["value"], + } + model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose") + assert model(tool_name="choose", value=1).value == 1 + assert model(tool_name="choose", value="brief").value == "brief" + with pytest.raises(ValidationError): + model(tool_name="choose", value="other") + + @pytest.mark.parametrize( + "prop_schema, invalid", + [ + ({"type": ["boolean", "integer"], "enum": [True]}, 1), + ({"type": ["boolean", "integer"], "enum": [1]}, True), + ({"type": ["boolean", "number"], "const": False}, 0.0), + ({"type": ["boolean", "number"], "const": 0}, False), + ({"enum": [True]}, 1.0), + ({"enum": [0]}, False), + ({"const": True}, 1), + ({"const": 1}, True), + ], + ) + def test_scalar_choices_preserve_json_identity(self, prop_schema, invalid): + schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} + model = SchemaTransformer.create_model_from_schema(schema, "ChoiceInput", "choose") + payload = {"tool_name": "choose", "value": invalid} + + assert not Draft202012Validator(schema).is_valid(payload) + with pytest.raises(ValidationError): + model.model_validate(payload) + assert not Draft202012Validator(model.model_json_schema()).is_valid(payload) + + @pytest.mark.parametrize("choices", [[True, 1], [1, True], [False, 0], [0, False]]) + def test_mixed_boolean_and_numeric_choices_keep_the_matching_choice(self, choices): + schema = {"type": "object", "properties": {"value": {"enum": choices}}, "required": ["value"]} + model = SchemaTransformer.create_model_from_schema(schema, "MixedChoiceInput", "choose") + numeric = next(value for value in choices if not isinstance(value, bool)) + + assert type(model(tool_name="choose", value=float(numeric)).value) is int + boolean = next(value for value in choices if isinstance(value, bool)) + assert model(tool_name="choose", value=boolean).value is boolean + + @pytest.mark.parametrize("prop_schema", [{"type": "number", "enum": [1]}, {"const": 1.0}, {"enum": [0, 1]}]) + def test_scalar_choices_retain_json_numeric_equality(self, prop_schema): + schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} + model = SchemaTransformer.create_model_from_schema(schema, "NumericChoiceInput", "choose") + for value in (1, 1.0): + payload = {"tool_name": "choose", "value": value} + assert Draft202012Validator(schema).is_valid(payload) + assert model.model_validate(payload).value == value + assert Draft202012Validator(model.model_json_schema()).is_valid(payload) + + @pytest.mark.parametrize( + "prop_schema", + [ + {"type": "string", "enum": ["brief", "detailed"]}, + {"type": "string", "const": "brief"}, + {"type": ["string", "null"], "enum": ["brief", None]}, + {"type": "number", "enum": [1, 2]}, + ], + ) + def test_optional_scalar_choice_schema_matches_none_validation(self, prop_schema): + schema = {"type": "object", "properties": {"value": prop_schema}} + model = SchemaTransformer.create_model_from_schema(schema, "OptionalChoiceInput", "choose") + payload = {"tool_name": "choose", "value": None} + + assert model.model_validate(payload).value is None + emitted = model.model_json_schema() + assert Draft202012Validator(emitted).is_valid(payload) + assert Draft202012Validator(emitted).is_valid({"tool_name": "choose"}) + assert not Draft202012Validator(emitted).is_valid({"tool_name": "choose", "value": "unknown"}) + + @pytest.mark.parametrize( + "prop_schema, valid, invalid, emitted_type, choices", + [ + ({"type": "string", "enum": ["brief", None]}, "brief", None, "string", ["brief"]), + ({"type": "integer", "enum": [1, "two"]}, 1, "two", "integer", [1]), + ({"type": "integer", "enum": [1, True, 2.5]}, 1, True, "integer", [1]), + ({"type": "number", "enum": [1, 2.5, False]}, 2.5, False, "number", [1, 2.5]), + ({"type": "boolean", "enum": [False, 0, "false"]}, False, 0, "boolean", [False]), + ({"type": "integer", "enum": [1.0, 2.5]}, 1, 2.5, "integer", [1.0]), + ({"type": "boolean", "const": False}, False, 0, "boolean", [False]), + ], + ) + def test_typed_scalar_choices(self, prop_schema, valid, invalid, emitted_type, choices): + """Choice constraints intersect the declared JSON type, including bool/number boundaries.""" + schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} + model = SchemaTransformer.create_model_from_schema(schema, "TypedModeInput", "configure") + assert model(tool_name="configure", mode=valid).mode == valid + with pytest.raises(ValueError): + model(tool_name="configure", mode=invalid) + emitted = model.model_json_schema()["properties"]["mode"] + assert emitted["type"] == emitted_type + assert emitted.get("enum", [emitted.get("const")]) == choices + + @pytest.mark.parametrize("prop_schema", [{"type": "string", "const": None}, {"type": "integer", "enum": ["two"]}]) + def test_incompatible_scalar_choices_reject_definition(self, prop_schema): + """An impossible type/choice combination never becomes an unrestricted input field.""" + schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} + with pytest.raises(ValueError, match="choices.*declared"): + SchemaTransformer.create_model_from_schema(schema, "InvalidModeInput", "configure") + + def test_nullable_type_array_with_mixed_enum(self): + schema = { + "type": "object", + "properties": {"mode": {"type": ["string", "null"], "enum": ["brief", None, 1]}}, + "required": ["mode"], + } + model = SchemaTransformer.create_model_from_schema(schema, "NullableModeInput", "configure") + assert model(tool_name="configure", mode="brief").mode == "brief" + assert model(tool_name="configure", mode=None).mode is None + with pytest.raises(ValueError): + model(tool_name="configure", mode=1) + assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", None] + + @pytest.mark.parametrize( + "prop_schema, valid, invalid", + [ + ({"type": "string", "enum": ["brief", "detailed"]}, "brief", "unknown"), + ({"type": "integer", "enum": [1, 2]}, 1, 3), + ({"enum": ["brief", None]}, None, "unknown"), + ({"type": "string", "const": "brief"}, "brief", "detailed"), + ({"type": "integer", "const": 0}, 0, 1), + ({"type": "boolean", "const": False}, False, True), + ({"const": None}, None, "unknown"), + ], + ) + def test_scalar_value_constraints(self, prop_schema, valid, invalid): + """MCP parameter choices constrain both validation and the model-facing schema.""" + schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} + model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") + + assert model(tool_name="configure", mode=valid).mode == valid + with pytest.raises(ValueError): + model(tool_name="configure", mode=invalid) + emitted = model.model_json_schema()["properties"]["mode"] + choices = prop_schema.get("enum", [prop_schema.get("const")]) + if len(choices) == 1: + assert emitted["const"] == choices[0] + else: + assert emitted["enum"] == choices + + def test_array_item_enum(self): + """Enum constraints survive recursive conversion of tool parameter arrays.""" + schema = { + "type": "object", + "properties": {"modes": {"type": "array", "items": {"type": "string", "enum": ["brief", "detailed"]}}}, + "required": ["modes"], + } + model = SchemaTransformer.create_model_from_schema(schema, "ModesInput", "configure") + + assert model(tool_name="configure", modes=["brief", "detailed"]).modes == ["brief", "detailed"] + with pytest.raises(ValueError): + model(tool_name="configure", modes=["unknown"]) + assert model.model_json_schema()["properties"]["modes"]["items"]["enum"] == ["brief", "detailed"] + + @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"]) + def test_referenced_enum(self, definitions_key): + """Enum definitions emitted by MCP servers remain scalar parameter choices.""" + schema = { + "type": "object", + "properties": {"mode": {"$ref": f"#/{definitions_key}/Mode"}}, + "required": ["mode"], + definitions_key: {"Mode": {"type": "string", "enum": ["brief", "detailed"]}}, + } + model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") + + assert model(tool_name="configure", mode="brief").mode == "brief" + with pytest.raises(ValueError): + model(tool_name="configure", mode="unknown") + assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", "detailed"] + + def test_nullable_enum_with_default(self): + schema = { + "type": "object", + "properties": { + "mode": {"anyOf": [{"type": "string", "enum": ["brief", "detailed"]}, {"type": "null"}], "default": "brief"} + }, + } + model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") + + assert model(tool_name="configure").mode == "brief" + assert model(tool_name="configure", mode=None).mode is None + with pytest.raises(ValueError): + model(tool_name="configure", mode="unknown") + @pytest.mark.parametrize("invalid_field", ["first", "second"]) def test_repeated_reference_reuses_model_and_validates_both_fields(self, invalid_field: str) -> None: reference = {"$ref": "#/$defs/Item"} diff --git a/pyproject.toml b/pyproject.toml index fa9a959d..299a0f36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,8 @@ dependencies = [ "pyyaml>=6.0.2,<7.0.0", "requests>=2.32.3,<3.0.0", "mcp[cli]>=1.8.0,<2", + "jsonschema>=4.20.0,<5", + "referencing>=0.28.4,<1", "litellm>=1.50.0,<2.0.0", ] diff --git a/uv.lock b/uv.lock index bfd13e43..64365c4b 100644 --- a/uv.lock +++ b/uv.lock @@ -272,11 +272,13 @@ source = { editable = "." } dependencies = [ { name = "gitpython" }, { name = "instructor" }, + { name = "jsonschema" }, { name = "litellm" }, { name = "mcp", extra = ["cli"] }, { name = "pydantic" }, { name = "pyfiglet" }, { name = "pyyaml" }, + { name = "referencing" }, { name = "requests" }, { name = "rich" }, { name = "textual" }, @@ -321,6 +323,7 @@ dev = [ requires-dist = [ { name = "gitpython", specifier = ">=3.1.43,<4.0.0" }, { name = "instructor", specifier = ">=1.17.0,<2" }, + { name = "jsonschema", specifier = ">=4.20.0,<5" }, { name = "litellm", specifier = ">=1.50.0,<2.0.0" }, { name = "mcp", extras = ["cli"], specifier = ">=1.8.0,<2" }, { name = "openinference-instrumentation", marker = "extra == 'phoenix'", specifier = ">=0.1.38,<0.2.0" }, @@ -330,6 +333,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.11.4,<3.0.0" }, { name = "pyfiglet", specifier = ">=1.0.2,<2.0.0" }, { name = "pyyaml", specifier = ">=6.0.2,<7.0.0" }, + { name = "referencing", specifier = ">=0.28.4,<1" }, { name = "requests", specifier = ">=2.32.3,<3.0.0" }, { name = "rich", specifier = ">=13.7.1,<14.0.0" }, { name = "textual", specifier = ">=5.3.0,<6.0.0" },