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,18 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                         UiPathErrorCategory.USER,
                     ) from e
                 try:
-                    serialized_output = self._serialize_object(output)
+                    if isinstance(output, AgentOutput):
+                        if output.structured_response is not None:
+                            serialized_output = self._serialize_object(
+                                output.structured_response
+                            )
+                        else:
+                            serialized_output = output.response.content or str(
+                                output.response
+                            )
+                    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

@@ -4,6 +4,7 @@
 import uuid
 from typing import Any, Callable, Dict, overload
 
+from llama_index.core.agent.workflow import BaseWorkflowAgent
 from llama_index.core.workflow import (
     HumanResponseEvent,
     InputRequiredEvent,
@@ -71,18 +72,45 @@ 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) and getattr(workflow, 'output_cls', None) is not None:
+        # Use the output_cls schema for structured output
+        try:
+            output_schema = workflow.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",