diff --git a/src/mcpadapt/utils/modeling.py b/src/mcpadapt/utils/modeling.py index 6ac704c..e5db088 100644 --- a/src/mcpadapt/utils/modeling.py +++ b/src/mcpadapt/utils/modeling.py @@ -1,7 +1,7 @@ -from copy import deepcopy from typing import Any, Dict, ForwardRef, List, Optional, Type, Union + +import jsonref from pydantic import BaseModel, Field, create_model -import re json_type_mapping: dict[str, Type] = { "string": str, @@ -13,31 +13,10 @@ } -def resolve_refs_and_remove_defs(json_obj): - # Extract $defs - defs = json_obj.get("$defs", {}) - - # Function to recursively resolve $ref - def _resolve(obj): - if isinstance(obj, dict): - if "$ref" in obj: - ref_path = obj["$ref"] - match = re.match(r"#/\$defs/(\w+)", ref_path) - if match: - def_key = match.group(1) - return _resolve(deepcopy(defs.get(def_key, {}))) - return {k: _resolve(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [_resolve(i) for i in obj] - else: - return obj - - json_obj = _resolve(json_obj) - - # Remove $defs - json_obj.pop("$defs", None) - - return json_obj +def resolve_refs_and_remove_defs(json_obj: dict[str, Any]) -> dict[str, Any]: + """Resolve JSON Schema $ref pointers (including in-document paths) and strip $defs.""" + resolved = jsonref.replace_refs(json_obj, proxies=False) + return {key: value for key, value in resolved.items() if key != "$defs"} def create_model_from_json_schema( diff --git a/tests/utils/test_modeling.py b/tests/utils/test_modeling.py index 5173418..91a479f 100644 --- a/tests/utils/test_modeling.py +++ b/tests/utils/test_modeling.py @@ -2,7 +2,10 @@ Tests for the modeling module, specifically focused on JSON Schema handling. """ -from mcpadapt.utils.modeling import create_model_from_json_schema +from mcpadapt.utils.modeling import ( + create_model_from_json_schema, + resolve_refs_and_remove_defs, +) def test_direct_modeling_with_list_type(): @@ -35,3 +38,38 @@ def test_direct_modeling_with_list_type(): # Verify the model works as expected with number instance = model(multi_type_field=42) assert instance.multi_type_field == 42 + + +def test_create_model_with_property_path_ref(): + """Regression for #75: in-document $ref paths must be resolved before modeling.""" + schema = { + "type": "object", + "properties": { + "requestResponse": { + "type": "object", + "properties": { + "request": { + "type": "object", + "properties": { + "filename": {"type": "string"}, + "body": {"type": "string"}, + }, + "required": ["body"], + "additionalProperties": False, + }, + "response": { + "$ref": "#/properties/requestResponse/properties/request" + }, + }, + "required": ["request", "response"], + "additionalProperties": False, + } + }, + "additionalProperties": False, + } + + model = create_model_from_json_schema(resolve_refs_and_remove_defs(schema)) + + generated_schema = model.model_json_schema() + assert "requestResponse" in generated_schema["properties"] + assert "$ref" not in str(generated_schema)