From 73963771fbfcd18a45b8bf61d0f0c3f8fdb56ee8 Mon Sep 17 00:00:00 2001 From: liangzhechen655 <186018027+liangzhechen655@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:33:05 +0800 Subject: [PATCH 1/4] fix: preserve scalar MCP enum and const constraints --- .../connectors/mcp/schema_transformer.py | 12 +++- .../connectors/mcp/test_schema_transformer.py | 71 +++++++++++++++++++ 2 files changed, 82 insertions(+), 1 deletion(-) diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index ad8cad80..830430d1 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -43,7 +43,10 @@ def _resolve_ref(ref_path: str, root_schema: Dict[str, Any], model_cache: Dict[s model_name = ref_schema.get("title", ref_name) # Avoid infinite recursion by adding placeholder first model_cache[ref_name] = Any - model = SchemaTransformer._create_nested_model(ref_schema, model_name, root_schema, model_cache) + if ref_schema.get("type") == "object" or "properties" in ref_schema: + model = SchemaTransformer._create_nested_model(ref_schema, model_name, root_schema, model_cache) + else: + model, _ = SchemaTransformer.json_to_pydantic_field(ref_schema, True, root_schema, model_cache) model_cache[ref_name] = model return model @@ -96,6 +99,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 = Literal[prop_schema["const"]] + elif prop_schema.get("enum") and all( + isinstance(value, (str, int, float, bool, type(None))) for value in prop_schema["enum"] + ): + python_type = Literal[tuple(prop_schema["enum"])] # 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 781beef8..36240618 100644 --- a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py +++ b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py @@ -71,6 +71,77 @@ def test_no_type(self): class TestCreateModelFromSchema: + @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("definitions_key", ["$defs", "definitions"]) @pytest.mark.parametrize( "name, reference_token", From 6dbfbe4d22838c5e0126014a5676297f7e88cf41 Mon Sep 17 00:00:00 2001 From: liangzhechen655 <186018027+liangzhechen655@users.noreply.github.com> Date: Sat, 10 Oct 2026 17:16:20 +0800 Subject: [PATCH 2/4] fix: intersect MCP scalar choices with declared types --- .../connectors/mcp/schema_transformer.py | 37 ++++++++++++++-- .../connectors/mcp/test_schema_transformer.py | 43 +++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index 5dc846a8..b94e2df9 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -1,10 +1,10 @@ """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 pydantic import BeforeValidator, Field, create_model from atomic_agents.base.base_io_schema import BaseIOSchema @@ -25,6 +25,33 @@ 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]) -> Type: + """Intersect scalar choices with the declared type for validation and schema output.""" + declared_type = prop_schema.get("type") + if declared_type is None: + return Literal[tuple(choices)] + json_types = declared_type if isinstance(declared_type, list) else [declared_type] + choices = [value for value in choices if SchemaTransformer._matches_scalar_type(value, json_types)] + if not choices: + raise ValueError("Scalar choices do not match the declared JSON Schema type") + + def validate_type(value: Any) -> Any: + if not SchemaTransformer._matches_scalar_type(value, json_types): + raise ValueError("Value does not match the declared JSON Schema type") + return value + + return Annotated[Literal[tuple(choices)], BeforeValidator(validate_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,11 +127,11 @@ def json_to_pydantic_field( 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 = Literal[prop_schema["const"]] + python_type = SchemaTransformer._scalar_choice_type(prop_schema, [prop_schema["const"]]) elif prop_schema.get("enum") and all( isinstance(value, (str, int, float, bool, type(None))) for value in prop_schema["enum"] ): - python_type = Literal[tuple(prop_schema["enum"])] + python_type = SchemaTransformer._scalar_choice_type(prop_schema, prop_schema["enum"]) # Handle oneOf/anyOf (unions) elif "oneOf" in prop_schema or "anyOf" in prop_schema: union_schemas = prop_schema.get("oneOf", prop_schema.get("anyOf", [])) @@ -144,6 +171,8 @@ def json_to_pydantic_field( python_type = Dict[str, Any] field_kwargs = {"description": description} + if ("const" in prop_schema or "enum" in prop_schema) and "type" in prop_schema: + field_kwargs["json_schema_extra"] = {"type": prop_schema["type"]} if required: field_kwargs["default"] = ... elif default is not None: diff --git a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py index 36240618..bbf1192e 100644 --- a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py +++ b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py @@ -71,6 +71,49 @@ def test_no_type(self): class TestCreateModelFromSchema: + @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", [ From 06c30e61f06f219f7f8e6dfb43c76108df5ec615 Mon Sep 17 00:00:00 2001 From: liangzhechen655 <186018027+liangzhechen655@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:06:34 +0800 Subject: [PATCH 3/4] fix: preserve scalar choice identity and optional schema types --- .../connectors/mcp/schema_transformer.py | 25 +++++--- .../connectors/mcp/test_schema_transformer.py | 64 +++++++++++++++++++ 2 files changed, 79 insertions(+), 10 deletions(-) diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index 3f8a0621..b381ce0c 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -38,19 +38,26 @@ def _matches_scalar_type(value: Any, json_types: List[str]) -> bool: def _scalar_choice_type(prop_schema: Dict[str, Any], choices: List[Any]) -> Type: """Intersect scalar choices with the declared type for validation and schema output.""" declared_type = prop_schema.get("type") - if declared_type is None: - return Literal[tuple(choices)] json_types = declared_type if isinstance(declared_type, list) else [declared_type] - choices = [value for value in choices if SchemaTransformer._matches_scalar_type(value, json_types)] + if declared_type is not None: + choices = [value for value in choices if SchemaTransformer._matches_scalar_type(value, json_types)] if not choices: raise ValueError("Scalar choices do not match the declared JSON Schema type") - def validate_type(value: Any) -> Any: - if not SchemaTransformer._matches_scalar_type(value, json_types): + 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") - return value - - return Annotated[Literal[tuple(choices)], BeforeValidator(validate_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: @@ -177,8 +184,6 @@ def json_to_pydantic_field( python_type = Dict[str, value_type] field_kwargs = {"description": description} - if ("const" in prop_schema or "enum" in prop_schema) and "type" in prop_schema: - field_kwargs["json_schema_extra"] = {"type": prop_schema["type"]} if required: field_kwargs["default"] = ... elif default is not None: diff --git a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py index 7d02f566..1091e1ed 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,69 @@ def test_no_type(self): class TestCreateModelFromSchema: + @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", [ From 7105d62c4265a584f7aa9368982462d20641a16e Mon Sep 17 00:00:00 2001 From: liangzhechen655 <186018027+liangzhechen655@users.noreply.github.com> Date: Sun, 25 Oct 2026 17:19:39 +0800 Subject: [PATCH 4/4] fix: intersect MCP scalar choices with composition constraints --- .../connectors/mcp/schema_transformer.py | 30 +++++-- .../connectors/mcp/test_schema_transformer.py | 86 +++++++++++++++++++ pyproject.toml | 2 + uv.lock | 4 + 4 files changed, 115 insertions(+), 7 deletions(-) diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index b381ce0c..826b70c2 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -4,12 +4,26 @@ from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, Annotated, cast from atomic_agents.connectors.mcp.mcp_definition_service import MCPAttributeType +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, @@ -35,14 +49,16 @@ def _matches_scalar_type(value: Any, json_types: List[str]) -> bool: 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]) -> Type: - """Intersect scalar choices with the declared type for validation and schema output.""" + 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] - if declared_type is not None: - choices = [value for value in choices if SchemaTransformer._matches_scalar_type(value, json_types)] + # 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 type") + 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): @@ -136,11 +152,11 @@ def json_to_pydantic_field( 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"]]) + 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"]) + 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 1091e1ed..fce7906b 100644 --- a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py +++ b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py @@ -76,6 +76,92 @@ 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", [ 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" },