fix: extract output from final chunk

Author: radu-mocanu

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "uipath-langchain"
-version = "0.0.120"
+version = "0.0.121"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
-    "uipath>=2.1.0, <2.2.0",
-    "langgraph>=0.2.70",
+    "uipath>=2.1.9, <2.2.0",
+    "langgraph>=0.5.0, <0.7.0",
     "langchain-core>=0.3.34",
     "langgraph-checkpoint-sqlite>=2.0.3",
     "langchain-community>=0.3.21",

src/uipath_langchain/_cli/_runtime/_runtime.py

@@ -101,7 +101,7 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                     graph_config["max_concurrency"] = int(max_concurrency)
 
                 # Stream the output at debug time
-                if self.context.job_id is None:
+                if self.is_debug_run():
                     # Get final chunk while streaming
                     final_chunk = None
                     async for stream_chunk in graph.astream(
@@ -115,7 +115,7 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
 
                     self.context.output = self._extract_graph_result(final_chunk, graph)
                 else:
-                    # Execute the graph normally at runtime
+                    # Execute the graph normally at runtime or eval
                     self.context.output = await graph.ainvoke(
                         processed_input, graph_config
                     )
@@ -310,9 +310,23 @@ def _extract_graph_result(
             # Check which channels are present
             available_channels = [ch for ch in output_channels if ch in final_chunk]
 
+            # if no available channels, output may contain the last_node name as key
+            unwrapped_final_chunk = {}
+            if not available_channels:
+                if len(final_chunk) == 1 and isinstance(
+                    unwrapped_final_chunk := next(iter(final_chunk.values())), dict
+                ):
+                    available_channels = [
+                        ch for ch in output_channels if ch in unwrapped_final_chunk
+                    ]
+
             if available_channels:
                 # Create a dict with the available channels
-                return {channel: final_chunk[channel] for channel in available_channels}
+                return {
+                    channel: final_chunk.get(channel, None)
+                    or unwrapped_final_chunk[channel]
+                    for channel in available_channels
+                }
 
         # Fallback for any other case
         return final_chunk

src/uipath_langchain/_cli/cli_run.py

@@ -42,7 +42,7 @@ async def execute():
             context.logs_min_level = env.get("LOG_LEVEL", "INFO")
             context.job_id = env.get("UIPATH_JOB_KEY")
             context.trace_id = env.get("UIPATH_TRACE_ID")
-            context.eval_run = kwargs.get("eval_run", False)
+            context.is_eval_run = kwargs.get("is_eval_run", False)
             context.tracing_enabled = tracing
             context.input_file = kwargs.get("input_file", None)
             context.execution_output_file = kwargs.get("execution_output_file", None)