fix: support non-object output schemas (#637)

Author: andreitava-uip

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.0"
+version = "0.8.1"
 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"
@@ -18,7 +18,7 @@ dependencies = [
     "python-dotenv>=1.0.1",
     "httpx>=0.27.0",
     "openinference-instrumentation-langchain>=0.1.56",
-    "jsonschema-pydantic-converter>=0.1.9",
+    "jsonschema-pydantic-converter>=0.2.0",
     "jsonpath-ng>=1.7.0",
     "mcp==1.26.0",
     "langchain-mcp-adapters==0.2.1",

src/uipath_langchain/agent/react/json_utils.py

@@ -2,7 +2,7 @@
 from typing import Any, ForwardRef, Union, get_args, get_origin
 
 from jsonpath_ng import parse  # type: ignore[import-untyped]
-from pydantic import BaseModel
+from pydantic import BaseModel, RootModel
 
 
 def get_json_paths_by_type(model: type[BaseModel], type_name: str) -> list[str]:
@@ -63,26 +63,35 @@ def _recursive_search(
                 continue
 
             if origin is list:
-                args = get_args(annotation)
-                if args:
-                    list_item_type = args[0]
-                    if matches_type(list_item_type):
-                        json_paths.append(f"{field_path}[*]")
-                        continue
-
-                    if _is_pydantic_model(list_item_type):
-                        nested_paths = _recursive_search(
-                            list_item_type, f"{field_path}[*]"
-                        )
-                        json_paths.extend(nested_paths)
-                        continue
+                inner_type, suffix = _unwrap_lists(annotation)
+                inner_path = f"{field_path}{suffix}"
+                if matches_type(inner_type):
+                    json_paths.append(inner_path)
+                    continue
+                if _is_pydantic_model(inner_type):
+                    nested_paths = _recursive_search(inner_type, inner_path)
+                    json_paths.extend(nested_paths)
+                    continue
 
             if _is_pydantic_model(annotation):
                 nested_paths = _recursive_search(annotation, field_path)
                 json_paths.extend(nested_paths)
 
         return json_paths
 
+    # RootModel serializes without the "root" wrapper — e.g. RootModel[list[X]]
+    # dumps as [...], not {"root": [...]}.  Iterating model_fields directly would
+    # produce wrong paths like "$.root.field".  Instead we peel off the RootModel
+    # envelope (and any Optional/list layers) so _recursive_search only ever sees
+    # a plain BaseModel with correct JSONPath prefixes (e.g. "$[*].field").
+    if issubclass(model, RootModel):
+        inner = _unwrap_optional(model.model_fields["root"].annotation)
+        inner, suffix = _unwrap_lists(inner)
+        # Primitive or non-model root types can't contain nested typed fields.
+        if not _is_pydantic_model(inner):
+            return []
+        return _recursive_search(inner, f"${suffix}" if suffix else "")
+
     return _recursive_search(model, "")
 
 
@@ -179,5 +188,21 @@ def _unwrap_optional(annotation: Any) -> Any:
     return annotation
 
 
+def _unwrap_lists(annotation: Any) -> tuple[Any, str]:
+    """Unwrap nested list types, returning (inner_type, jsonpath_suffix).
+
+    Each list layer adds a "[*]" wildcard so the resulting suffix maps directly
+    to JSONPath:  list[list[X]] → (X, "[*][*]").
+    """
+    suffix = ""
+    while get_origin(annotation) is list:
+        args = get_args(annotation)
+        if not args:
+            break
+        annotation = args[0]
+        suffix += "[*]"
+    return annotation, suffix
+
+
 def _is_pydantic_model(annotation: Any) -> bool:
     return isinstance(annotation, type) and issubclass(annotation, BaseModel)

tests/agent/react/test_json_utils.py

@@ -0,0 +1,153 @@
+from typing import Optional
+
+from pydantic import BaseModel, RootModel
+
+from uipath_langchain.agent.react.json_utils import (
+    extract_values_by_paths,
+    get_json_paths_by_type,
+)
+
+
+class Target(BaseModel):
+    id: str
+
+
+class Nested(BaseModel):
+    target: Target
+    name: str
+
+
+class WithList(BaseModel):
+    items: list[Target]
+
+
+class WithOptional(BaseModel):
+    target: Optional[Target] = None
+
+
+class WithNestedList(BaseModel):
+    matrix: list[list[Target]]
+
+
+class Mixed(BaseModel):
+    direct: Target
+    items: list[Target]
+    nested: Nested
+    plain: str
+
+
+# -- get_json_paths_by_type: regular BaseModel ---------------------------------
+
+
+class TestGetJsonPathsByType:
+    def test_direct_field(self):
+        paths = get_json_paths_by_type(Nested, "Target")
+        assert paths == ["$.target"]
+
+    def test_list_field(self):
+        paths = get_json_paths_by_type(WithList, "Target")
+        assert paths == ["$.items[*]"]
+
+    def test_optional_field(self):
+        paths = get_json_paths_by_type(WithOptional, "Target")
+        assert paths == ["$.target"]
+
+    def test_nested_list_of_lists(self):
+        paths = get_json_paths_by_type(WithNestedList, "Target")
+        assert paths == ["$.matrix[*][*]"]
+
+    def test_mixed_fields(self):
+        paths = get_json_paths_by_type(Mixed, "Target")
+        assert set(paths) == {"$.direct", "$.items[*]", "$.nested.target"}
+
+    def test_no_match(self):
+        class Unrelated(BaseModel):
+            name: str
+            value: int
+
+        paths = get_json_paths_by_type(Unrelated, "Target")
+        assert paths == []
+
+    def test_nested_model_in_list(self):
+        class Outer(BaseModel):
+            groups: list[Nested]
+
+        paths = get_json_paths_by_type(Outer, "Target")
+        assert paths == ["$.groups[*].target"]
+
+
+# -- get_json_paths_by_type: RootModel -----------------------------------------
+
+
+class TestGetJsonPathsByTypeRootModel:
+    def test_root_model_single_object(self):
+        Model = RootModel[Target]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == []
+
+    def test_root_model_list_of_target(self):
+        """Target is the leaf type itself — no nested fields of type Target within it."""
+        Model = RootModel[list[Target]]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == []
+
+    def test_root_model_list_of_nested(self):
+        Model = RootModel[list[Nested]]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == ["$[*].target"]
+
+    def test_root_model_list_of_list(self):
+        Model = RootModel[list[list[Nested]]]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == ["$[*][*].target"]
+
+    def test_root_model_primitive(self):
+        Model = RootModel[str]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == []
+
+    def test_root_model_list_of_primitives(self):
+        Model = RootModel[list[str]]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == []
+
+    def test_root_model_optional_list(self):
+        Model = RootModel[Optional[list[Nested]]]
+        paths = get_json_paths_by_type(Model, "Target")
+        assert paths == ["$[*].target"]
+
+
+# -- extract_values_by_paths ---------------------------------------------------
+
+
+class TestExtractValuesByPaths:
+    def test_extract_from_dict(self):
+        data = {"target": {"id": "1"}, "name": "test"}
+        values = extract_values_by_paths(data, ["$.target"])
+        assert values == [{"id": "1"}]
+
+    def test_extract_from_list(self):
+        data = {"items": [{"id": "1"}, {"id": "2"}]}
+        values = extract_values_by_paths(data, ["$.items[*]"])
+        assert values == [{"id": "1"}, {"id": "2"}]
+
+    def test_extract_from_pydantic_model(self):
+        obj = Nested(target=Target(id="1"), name="test")
+        values = extract_values_by_paths(obj, ["$.target"])
+        assert values == [{"id": "1"}]
+
+    def test_extract_multiple_paths(self):
+        data = {
+            "direct": {"id": "1"},
+            "items": [{"id": "2"}, {"id": "3"}],
+        }
+        values = extract_values_by_paths(data, ["$.direct", "$.items[*]"])
+        assert values == [{"id": "1"}, {"id": "2"}, {"id": "3"}]
+
+    def test_extract_no_paths(self):
+        values = extract_values_by_paths({"a": 1}, [])
+        assert values == []
+
+    def test_extract_path_not_found(self):
+        values = extract_values_by_paths({"a": 1}, ["$.missing"])
+        assert values == []