Merge pull request #180 from UiPath/fix/llamaindex-use-shared-runtime-utils

fix: use shared serialization and schema utils from uipath core

Author: cristipufu

.gitignore

@@ -191,3 +191,4 @@ playground.py
 **/samples/**/bindings.json
 **/samples/**/entry-points.json
 **/samples/**/uv.lock
+**/**/.claude/settings.local.json

packages/uipath-llamaindex/pyproject.toml

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

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/__init__.py

@@ -27,8 +27,6 @@ def create_factory(
     )
 
 
-register_runtime_factory()
-
 __all__ = [
     "register_runtime_factory",
     "get_entrypoints_schema",

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/_serialize.py

@@ -12,6 +12,7 @@
     ToolCall,
     ToolCallResult,
 )
+from uipath.core.serialization import serialize_defaults
 from uipath.runtime import (
     UiPathExecuteOptions,
     UiPathRuntimeResult,
@@ -49,8 +50,6 @@
 from uipath_llamaindex.runtime.schema import get_entrypoints_schema, get_workflow_schema
 from uipath_llamaindex.runtime.storage import SqliteResumableStorage
 
-from ._serialize import serialize_output
-
 
 class UiPathLlamaIndexRuntime:
     """
@@ -179,14 +178,14 @@ async def _run_workflow(
                     (AgentOutput, AgentInput, AgentStream, ToolCall, ToolCallResult),
                 ):
                     message_event = UiPathRuntimeMessageEvent(
-                        payload=serialize_output(event),
+                        payload=serialize_defaults(event),
                         node_name=node_name,
                         execution_id=self.runtime_id,
                     )
                     yield message_event
                 elif not isinstance(event, BreakpointEvent):
                     state_event = UiPathRuntimeStateEvent(
-                        payload=serialize_output(event),
+                        payload=serialize_defaults(event),
                         node_name=node_name,
                         execution_id=self.runtime_id,
                     )
@@ -256,7 +255,7 @@ def _create_breakpoint_result(
         return UiPathBreakpointResult(
             breakpoint_node=self._get_node_name(event),
             breakpoint_type="before",
-            current_state=serialize_output(event),
+            current_state=serialize_defaults(event),
             next_nodes=[],  # We don't know what's next in the stream
         )
 
@@ -285,11 +284,11 @@ def _create_success_result(self, output: Any) -> UiPathRuntimeResult:
         """Create result for successful completion."""
         if isinstance(output, AgentOutput):
             if output.structured_response is not None:
-                serialized_output = serialize_output(output.structured_response)
+                serialized_output = serialize_defaults(output.structured_response)
             else:
-                serialized_output = serialize_output(output)
+                serialized_output = serialize_defaults(output)
         else:
-            serialized_output = serialize_output(output)
+            serialized_output = serialize_defaults(output)
 
         if isinstance(serialized_output, str):
             serialized_output = {"result": serialized_output}

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/runtime.py

@@ -13,6 +13,8 @@
     UiPathRuntimeEdge,
     UiPathRuntimeGraph,
     UiPathRuntimeNode,
+    transform_nullable_types,
+    transform_references,
 )
 from workflows import Workflow
 from workflows.decorators import StepConfig
@@ -60,8 +62,8 @@ def get_entrypoints_schema(workflow: Workflow) -> dict[str, Any]:
         else:
             input_schema = start_event_class.model_json_schema()
             # Resolve references and handle nullable types
-            unpacked_input = _resolve_refs(input_schema)
-            schema["input"]["properties"] = _process_nullable_types(
+            unpacked_input, _ = transform_references(input_schema)
+            schema["input"]["properties"] = transform_nullable_types(
                 unpacked_input.get("properties", {})
             )
             schema["input"]["required"] = unpacked_input.get("required", [])
@@ -75,8 +77,8 @@ def get_entrypoints_schema(workflow: Workflow) -> dict[str, Any]:
             try:
                 output_schema = output_cls.model_json_schema()
                 # Resolve references and handle nullable types
-                unpacked_output = _resolve_refs(output_schema)
-                schema["output"]["properties"] = _process_nullable_types(
+                unpacked_output, _ = transform_references(output_schema)
+                schema["output"]["properties"] = transform_nullable_types(
                     unpacked_output.get("properties", {})
                 )
                 schema["output"]["required"] = unpacked_output.get("required", [])
@@ -100,8 +102,8 @@ def get_entrypoints_schema(workflow: Workflow) -> dict[str, Any]:
         try:
             output_schema = stop_event_class.model_json_schema()
             # Resolve references and handle nullable types
-            unpacked_output = _resolve_refs(output_schema)
-            schema["output"]["properties"] = _process_nullable_types(
+            unpacked_output, _ = transform_references(output_schema)
+            schema["output"]["properties"] = transform_nullable_types(
                 unpacked_output.get("properties", {})
             )
             schema["output"]["required"] = unpacked_output.get("required", [])
@@ -305,92 +307,6 @@ def get_step_config(step_name: str, step_func: Any) -> StepConfig | None:
     )
 
 
-def _resolve_refs(
-    schema: dict[str, Any],
-    root: dict[str, Any] | None = None,
-    visited: set[str] | None = None,
-) -> dict[str, Any]:
-    """
-    Recursively resolves $ref references in a JSON schema.
-
-    Args:
-        schema: The schema dictionary to resolve
-        root: The root schema for reference resolution
-        visited: Set of visited references to detect circular dependencies
-
-    Returns:
-        Resolved schema dictionary
-    """
-    if root is None:
-        root = schema
-
-    if visited is None:
-        visited = set()
-
-    if isinstance(schema, dict):
-        if "$ref" in schema:
-            ref_path = schema["$ref"]
-
-            if ref_path in visited:
-                # Circular dependency detected
-                return {
-                    "type": "object",
-                    "description": f"Circular reference to {ref_path}",
-                }
-
-            visited.add(ref_path)
-
-            # Resolve the reference - handle both #/definitions/ and #/$defs/ formats
-            ref_parts = ref_path.lstrip("#/").split("/")
-            ref_schema = root
-            for part in ref_parts:
-                ref_schema = ref_schema.get(part, {})
-
-            result = _resolve_refs(ref_schema, root, visited)
-
-            # Remove from visited after resolution
-            visited.discard(ref_path)
-
-            return result
-
-        return {k: _resolve_refs(v, root, visited) for k, v in schema.items()}
-
-    elif isinstance(schema, list):
-        return [_resolve_refs(item, root, visited) for item in schema]
-
-    return schema
-
-
-def _process_nullable_types(properties: dict[str, Any]) -> dict[str, Any]:
-    """
-    Process properties to handle nullable types correctly.
-
-    This matches the original implementation that adds "nullable": True
-    instead of simplifying the schema structure.
-
-    Args:
-        properties: The properties dictionary from a schema
-
-    Returns:
-        Processed properties with nullable types marked
-    """
-    result = {}
-    for name, prop in properties.items():
-        if "anyOf" in prop:
-            types = [item.get("type") for item in prop["anyOf"] if "type" in item]
-            if "null" in types:
-                non_null_types = [t for t in types if t != "null"]
-                if len(non_null_types) == 1:
-                    result[name] = {"type": non_null_types[0], "nullable": True}
-                else:
-                    result[name] = {"type": non_null_types, "nullable": True}
-            else:
-                result[name] = prop
-        else:
-            result[name] = prop
-    return result
-
-
 __all__ = [
     "get_entrypoints_schema",
     "get_workflow_schema",

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/schema.py

@@ -2,6 +2,11 @@
 
 import pytest
 
+from uipath_llamaindex.runtime import register_runtime_factory
+
+# Register the LlamaIndex runtime factory for CLI tests
+register_runtime_factory()
+
 # Get the tests directory path relative to this conftest file
 TESTS_DIR = Path(__file__).parent.parent
 MOCKS_DIR = TESTS_DIR / "mocks"

packages/uipath-llamaindex/tests/cli/conftest.py

@@ -104,7 +104,6 @@ def test_init_custom_config_generation(
 
                 assert "param" in props
                 assert props["param"]["type"] == "string"
-                assert props["param"]["nullable"]
 
                 # Verify required fields in input
                 assert input_schema["required"] == ["topic"]
@@ -126,7 +125,6 @@ def test_init_custom_config_generation(
 
                 assert "param" in out_props
                 assert out_props["param"]["type"] == "string"
-                assert out_props["param"]["nullable"]
 
                 # Verify required fields in output
                 assert output_schema["required"] == ["joke", "critique"]