fix: remove namespace filtering from create_model (#662)

Author: andreitava-uip

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.3"
+version = "0.8.4"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/react/jsonschema_pydantic_converter.py

@@ -1,7 +1,7 @@
 import inspect
 import sys
 from types import ModuleType
-from typing import Any, Type, get_args, get_origin
+from typing import Any, Type
 
 from jsonschema_pydantic_converter import transform_with_modules
 from pydantic import BaseModel
@@ -37,53 +37,17 @@ def create_model(
         )
 
     model, namespace = transform_with_modules(schema)
-    corrected_namespace: dict[str, Any] = {}
-
-    def collect_types(annotation: Any) -> None:
-        """Recursively collect all BaseModel types from an annotation."""
-        # Unwrap generic types like List, Optional, etc.
-        origin = get_origin(annotation)
-        if origin is not None:
-            for arg in get_args(annotation):
-                collect_types(arg)
-
-        elif inspect.isclass(annotation) and issubclass(annotation, BaseModel):
-            # Find the original name for this type from the namespace
-            for type_name, type_def in namespace.items():
-                # Match by class name since rebuild may create new instances
-                if (
-                    hasattr(annotation, "__name__")
-                    and hasattr(type_def, "__name__")
-                    and annotation.__name__ == type_def.__name__
-                ):
-                    # Store the actual annotation type, not the old namespace one
-                    annotation.__name__ = type_name
-                    corrected_namespace[type_name] = annotation
-                    break
-
-    # Collect all types from field annotations
-    for field_info in model.model_fields.values():
-        collect_types(field_info.annotation)
-
-    # Get the shared pseudo-module and populate it with this schema's types
-    # This ensures that forward references can be resolved by get_type_hints()
-    # when the model is used with external libraries (e.g., LangGraph)
+
     pseudo_module = _get_or_create_dynamic_module()
 
-    # Populate the pseudo-module with all types from the namespace
-    # Use the original names so forward references resolve correctly
-    for type_name, type_def in corrected_namespace.items():
+    for type_name, type_def in namespace.items():
         setattr(pseudo_module, type_name, type_def)
+        if inspect.isclass(type_def) and issubclass(type_def, BaseModel):
+            type_def.__module__ = _DYNAMIC_MODULE_NAME
 
     setattr(pseudo_module, model.__name__, model)
-
-    # Update the model's __module__ to point to the shared pseudo-module
     model.__module__ = _DYNAMIC_MODULE_NAME
 
-    # Update the __module__ of all generated types in the namespace
-    for type_def in corrected_namespace.values():
-        if inspect.isclass(type_def) and issubclass(type_def, BaseModel):
-            type_def.__module__ = _DYNAMIC_MODULE_NAME
     return model
 
 

tests/agent/react/test_json_utils.py

@@ -1,11 +1,12 @@
-from typing import Optional
+from typing import Any, Optional
 
 from pydantic import BaseModel, RootModel
 
 from uipath_langchain.agent.react.json_utils import (
     extract_values_by_paths,
     get_json_paths_by_type,
 )
+from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 
 
 class Target(BaseModel):
@@ -151,3 +152,78 @@ def test_extract_no_paths(self):
     def test_extract_path_not_found(self):
         values = extract_values_by_paths({"a": 1}, ["$.missing"])
         assert values == []
+
+
+# -- get_json_paths_by_type: dynamic models from create_model ------------------
+
+
+class TestGetJsonPathsByTypeDynamic:
+    """Exercise _get_target_type pseudo-module lookup with dynamic models."""
+
+    def test_direct_ref(self) -> None:
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "attachment": {"$ref": "#/definitions/job-attachment"},
+                "name": {"type": "string"},
+            },
+            "definitions": {
+                "job-attachment": {
+                    "type": "object",
+                    "properties": {
+                        "ID": {"type": "string"},
+                        "FullName": {"type": "string"},
+                    },
+                }
+            },
+        }
+        model = create_model(schema)
+        paths = get_json_paths_by_type(model, "__Job_attachment")
+        assert paths == ["$.attachment"]
+
+    def test_array_ref(self) -> None:
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "items": {
+                    "type": "array",
+                    "items": {"$ref": "#/definitions/job-attachment"},
+                }
+            },
+            "definitions": {
+                "job-attachment": {
+                    "type": "object",
+                    "properties": {"ID": {"type": "string"}},
+                }
+            },
+        }
+        model = create_model(schema)
+        paths = get_json_paths_by_type(model, "__Job_attachment")
+        assert paths == ["$.items[*]"]
+
+    def test_nested_ref_only_in_defs(self) -> None:
+        """Type referenced only by another $def, not directly by root."""
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "order": {"$ref": "#/$defs/Order"},
+            },
+            "$defs": {
+                "Order": {
+                    "type": "object",
+                    "properties": {
+                        "item": {"$ref": "#/$defs/Item"},
+                        "quantity": {"type": "integer"},
+                    },
+                },
+                "Item": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                    },
+                },
+            },
+        }
+        model = create_model(schema)
+        paths = get_json_paths_by_type(model, "__Item")
+        assert paths == ["$.order.item"]

tests/agent/react/test_schemas.py

@@ -275,6 +275,148 @@ def test_returns_false(self, schema: dict[str, Any]) -> None:
         assert has_underscore_fields(schema) is False
 
 
+class TestCreateModelJsonSchemaRoundtrip:
+    """Verify that model_json_schema() works on models produced by create_model().
+
+    Pydantic re-resolves forward references via sys.modules[cls.__module__],
+    so all types must be registered in the pseudo-module.
+    """
+
+    def test_single_ref(self) -> None:
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "address": {"$ref": "#/$defs/Address"},
+            },
+            "$defs": {
+                "Address": {
+                    "type": "object",
+                    "properties": {
+                        "street": {"type": "string"},
+                        "city": {"type": "string"},
+                    },
+                },
+            },
+        }
+        model = create_model(schema)
+        result = model.model_json_schema()
+        defs = result.get("$defs") or result.get("definitions") or {}
+        assert len(defs) == 1
+
+    def test_deeply_nested_defs_chain(self) -> None:
+        """Type B is only referenced by type A (not by root). Old code missed B."""
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "order": {"$ref": "#/$defs/Order"},
+            },
+            "$defs": {
+                "Order": {
+                    "type": "object",
+                    "properties": {
+                        "item": {"$ref": "#/$defs/Item"},
+                        "quantity": {"type": "integer"},
+                    },
+                },
+                "Item": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "price": {"type": "number"},
+                    },
+                },
+            },
+        }
+        model = create_model(schema)
+        result = model.model_json_schema()
+        defs = result.get("$defs") or result.get("definitions") or {}
+        assert len(defs) == 2
+
+    def test_three_level_defs_chain(self) -> None:
+        """Root -> A -> B -> C. C is two levels removed from root fields."""
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "company": {"$ref": "#/$defs/Company"},
+            },
+            "$defs": {
+                "Company": {
+                    "type": "object",
+                    "properties": {
+                        "department": {"$ref": "#/$defs/Department"},
+                    },
+                },
+                "Department": {
+                    "type": "object",
+                    "properties": {
+                        "manager": {"$ref": "#/$defs/Employee"},
+                    },
+                },
+                "Employee": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "role": {"type": "string"},
+                    },
+                },
+            },
+        }
+        model = create_model(schema)
+        result = model.model_json_schema()
+        defs = result.get("$defs") or result.get("definitions") or {}
+        assert len(defs) == 3
+
+    def test_enum_in_defs(self) -> None:
+        """Non-BaseModel types (enums) in $defs must also be registered."""
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "status": {"$ref": "#/$defs/Status"},
+            },
+            "$defs": {
+                "Status": {
+                    "type": "string",
+                    "enum": ["active", "inactive", "pending"],
+                },
+            },
+        }
+        model = create_model(schema)
+        result = model.model_json_schema()
+        assert result is not None
+
+    def test_array_of_nested_refs(self) -> None:
+        """Array field referencing a $def type that itself has a $ref."""
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": {
+                "tasks": {
+                    "type": "array",
+                    "items": {"$ref": "#/$defs/Task"},
+                },
+            },
+            "$defs": {
+                "Task": {
+                    "type": "object",
+                    "properties": {
+                        "assignee": {"$ref": "#/$defs/Person"},
+                        "title": {"type": "string"},
+                    },
+                },
+                "Person": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "email": {"type": "string"},
+                    },
+                },
+            },
+        }
+        model = create_model(schema)
+        result = model.model_json_schema()
+        defs = result.get("$defs") or result.get("definitions") or {}
+        assert len(defs) == 2
+
+
 class TestCreateModelRejectsUnderscoreFields:
     def test_top_level_underscore_field(self) -> None:
         schema = {