fix: io inference for agent workflows

Author: cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.0.34"
+version = "0.0.35"
 description = "UiPath LlamaIndex SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
@@ -9,7 +9,7 @@ dependencies = [
     "llama-index-embeddings-azure-openai>=0.3.8",
     "llama-index-llms-azure-openai>=0.3.2",
     "openinference-instrumentation-llama-index>=4.3.0",
-    "uipath>=2.1.35, <2.2.0",
+    "uipath>=2.1.36, <2.2.0",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",

src/uipath_llamaindex/_cli/_runtime/_runtime.py

@@ -4,6 +4,7 @@
 import pickle
 from typing import Optional, cast
 
+from llama_index.core.agent.workflow.workflow_events import AgentOutput
 from llama_index.core.workflow import (
     Context,
     HumanResponseEvent,
@@ -102,7 +103,19 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                         UiPathErrorCategory.USER,
                     ) from e
                 try:
-                    serialized_output = self._serialize_object(output)
+                    if isinstance(output, AgentOutput):
+                        structured_response = getattr(
+                            output, "structured_response", None
+                        )
+                        if structured_response is not None:
+                            serialized_output = self._serialize_object(
+                                structured_response
+                            )
+                        else:
+                            serialized_output = self._serialize_object(output)
+                    else:
+                        serialized_output = self._serialize_object(output)
+
                     # create simple kvp from string
                     if type(serialized_output) is str:
                         serialized_output = {"result": serialized_output}

src/uipath_llamaindex/_cli/cli_init.py

@@ -2,8 +2,9 @@
 import json
 import os
 import uuid
-from typing import Any, Callable, Dict, overload
+from typing import Any, Callable, Dict, Optional, Type, overload
 
+from llama_index.core.agent.workflow import BaseWorkflowAgent
 from llama_index.core.workflow import (
     HumanResponseEvent,
     InputRequiredEvent,
@@ -15,6 +16,7 @@
     get_steps_from_class,
     get_steps_from_instance,
 )
+from pydantic import BaseModel
 from uipath._cli._utils._console import ConsoleLogger
 from uipath._cli._utils._parse_ast import generate_bindings_json  # type: ignore
 from uipath._cli.middlewares import MiddlewareResult
@@ -71,18 +73,46 @@ def generate_schema_from_workflow(workflow: Workflow) -> Dict[str, Any]:
 
     # Generate input schema from StartEvent using Pydantic's schema method
     try:
-        input_schema = start_event_class.model_json_schema()
-        # Resolve references and handle nullable types
-        input_schema = resolve_refs(input_schema)
-        schema["input"]["properties"] = process_nullable_types(
-            input_schema.get("properties", {})
-        )
-        schema["input"]["required"] = input_schema.get("required", [])
+        if isinstance(workflow, BaseWorkflowAgent):
+            # For workflow agents, define a simple schema with just user_msg
+            schema["input"] = {
+                "type": "object",
+                "properties": {
+                    "user_msg": {
+                        "type": "string",
+                        "title": "User Message",
+                        "description": "The user's question or request",
+                    }
+                },
+                "required": ["user_msg"],
+            }
+        else:
+            input_schema = start_event_class.model_json_schema()
+            # Resolve references and handle nullable types
+            input_schema = resolve_refs(input_schema)
+            schema["input"]["properties"] = process_nullable_types(
+                input_schema.get("properties", {})
+            )
+            schema["input"]["required"] = input_schema.get("required", [])
     except (AttributeError, Exception):
         pass
 
-    # For output schema, check if it's the base StopEvent or a custom subclass
-    if stop_event_class is StopEvent:
+    # Handle output schema - check if it's a workflow agent with output_cls first
+    if isinstance(workflow, BaseWorkflowAgent):
+        output_cls: Optional[Type[BaseModel]] = getattr(workflow, "output_cls", None)
+        if output_cls is not None:
+            try:
+                output_schema = output_cls.model_json_schema()
+                # Resolve references and handle nullable types
+                output_schema = resolve_refs(output_schema)
+                schema["output"]["properties"] = process_nullable_types(
+                    output_schema.get("properties", {})
+                )
+                schema["output"]["required"] = output_schema.get("required", [])
+            except (AttributeError, Exception):
+                pass
+    # Check if it's the base StopEvent or a custom subclass
+    elif stop_event_class is StopEvent:
         # base StopEvent
         schema["output"] = {
             "type": "object",