uipath-langchain 0.0.140__py3-none-any.whl → 0.0.142__py3-none-any.whl

uipath_langchain/_cli/_runtime/_runtime.py

@@ -1,24 +1,32 @@
 import logging
 import os
-from typing import Any, List, Optional, Sequence
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, AsyncIterator, Optional, Sequence
 
-from langchain_core.callbacks.base import BaseCallbackHandler
-from langchain_core.messages import BaseMessage
 from langchain_core.runnables.config import RunnableConfig
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.errors import EmptyInputError, GraphRecursionError, InvalidUpdateError
+from langgraph.graph.state import CompiledStateGraph, StateGraph
+from langgraph.types import Interrupt, StateSnapshot
 from uipath._cli._runtime._contracts import (
     UiPathBaseRuntime,
+    UiPathBreakpointResult,
     UiPathErrorCategory,
+    UiPathResumeTrigger,
     UiPathRuntimeResult,
+    UiPathRuntimeStatus,
+)
+from uipath._events._events import (
+    UiPathAgentMessageEvent,
+    UiPathAgentStateEvent,
+    UiPathRuntimeEvent,
 )
 
 from ._context import LangGraphRuntimeContext
-from ._conversation import map_message
 from ._exception import LangGraphRuntimeError
 from ._graph_resolver import AsyncResolver, LangGraphJsonResolver
-from ._input import LangGraphInputProcessor
-from ._output import LangGraphOutputProcessor
+from ._input import get_graph_input
+from ._output import create_and_save_resume_trigger, serialize_output
 
 logger = logging.getLogger(__name__)
 
@@ -33,171 +41,221 @@ class LangGraphRuntime(UiPathBaseRuntime):
         super().__init__(context)
         self.context: LangGraphRuntimeContext = context
         self.graph_resolver: AsyncResolver = graph_resolver
+        self.resume_triggers_table: str = "__uipath_resume_triggers"
 
-    async def execute(self) -> Optional[UiPathRuntimeResult]:
+    @asynccontextmanager
+    async def _get_or_create_memory(self) -> AsyncIterator[AsyncSqliteSaver]:
         """
-        Execute the graph with the provided input and configuration.
+        Get existing memory from context or create a new one.
 
-        Returns:
-            Dictionary with execution results
+        If memory is created, it will be automatically disposed at the end.
+        If memory already exists in context, it will be reused without disposal.
 
-        Raises:
-            LangGraphRuntimeError: If execution fails
+        Yields:
+            AsyncSqliteSaver instance
         """
+        # Check if memory already exists in context
+        if self.context.memory is not None:
+            # Use existing memory, don't dispose
+            yield self.context.memory
+        else:
+            # Create new memory and dispose at the end
+            async with AsyncSqliteSaver.from_conn_string(
+                self.state_file_path
+            ) as memory:
+                yield memory
+                # Memory is automatically disposed by the context manager
 
+    async def execute(self) -> Optional[UiPathRuntimeResult]:
+        """Execute the graph with the provided input and configuration."""
         graph = await self.graph_resolver()
         if not graph:
             return None
 
         try:
-            async with AsyncSqliteSaver.from_conn_string(
-                self.state_file_path
-            ) as memory:
-                self.context.memory = memory
+            async with self._get_or_create_memory() as memory:
+                compiled_graph = await self._setup_graph(memory, graph)
+                graph_input = await self._get_graph_input(memory)
+                graph_config = self._get_graph_config()
 
-                # Compile the graph with the checkpointer
-                compiled_graph = graph.compile(checkpointer=self.context.memory)
+                # Execute without streaming
+                graph_output = await compiled_graph.ainvoke(graph_input, graph_config)
 
-                # Process input, handling resume if needed
-                input_processor = LangGraphInputProcessor(context=self.context)
+                # Get final state and create result
+                self.context.result = await self._create_runtime_result(
+                    compiled_graph, graph_config, memory, graph_output
+                )
 
-                processed_input = await input_processor.process()
+            return self.context.result
 
-                callbacks: List[BaseCallbackHandler] = []
+        except Exception as e:
+            raise self._create_runtime_error(e) from e
 
-                graph_config: RunnableConfig = {
-                    "configurable": {
-                        "thread_id": (
-                            self.context.execution_id
-                            or self.context.job_id
-                            or "default"
-                        )
-                    },
-                    "callbacks": callbacks,
-                }
+    async def stream(
+        self,
+    ) -> AsyncGenerator[UiPathRuntimeEvent | UiPathRuntimeResult, None]:
+        """
+        Stream graph execution events in real-time.
+
+        Yields UiPath UiPathRuntimeEvent instances (thin wrappers around framework data),
+        then yields the final UiPathRuntimeResult as the last item.
+        The result is also stored in self.context.result.
+
+        Yields:
+            - UiPathAgentMessageEvent: Wraps framework messages (BaseMessage, chunks, etc.)
+            - UiPathAgentStateEvent: Wraps framework state updates
+            - Final event: UiPathRuntimeResult or UiPathBreakpointResult
+
+        Example:
+            async for event in runtime.stream():
+                if isinstance(event, UiPathRuntimeResult):
+                    # Last event is the result
+                    print(f"Final result: {event}")
+                elif isinstance(event, UiPathAgentMessageEvent):
+                    # Access framework-specific message
+                    message = event.payload  # BaseMessage or AIMessageChunk
+                    print(f"Message: {message.content}")
+                elif isinstance(event, UiPathAgentStateEvent):
+                    # Access framework-specific state
+                    state = event.payload
+                    print(f"Node {event.node_name} updated: {state}")
 
-                recursion_limit = os.environ.get("LANGCHAIN_RECURSION_LIMIT", None)
-                max_concurrency = os.environ.get("LANGCHAIN_MAX_CONCURRENCY", None)
-
-                if recursion_limit is not None:
-                    graph_config["recursion_limit"] = int(recursion_limit)
-                if max_concurrency is not None:
-                    graph_config["max_concurrency"] = int(max_concurrency)
-
-                if self.context.chat_handler or self.is_debug_run():
-                    final_chunk: Optional[dict[Any, Any]] = None
-                    async for stream_chunk in compiled_graph.astream(
-                        processed_input,
-                        graph_config,
-                        stream_mode=["messages", "updates"],
-                        subgraphs=True,
-                    ):
-                        _, chunk_type, data = stream_chunk
-                        if chunk_type == "messages":
-                            if self.context.chat_handler:
-                                if isinstance(data, tuple):
-                                    message, _ = data
-                                    event = map_message(
-                                        message=message,
-                                        conversation_id=self.context.execution_id,
-                                        exchange_id=self.context.execution_id,
+        Raises:
+            LangGraphRuntimeError: If execution fails
+        """
+        graph = await self.graph_resolver()
+        if not graph:
+            return
+
+        try:
+            async with self._get_or_create_memory() as memory:
+                compiled_graph = await self._setup_graph(memory, graph)
+                graph_input = await self._get_graph_input(memory)
+                graph_config = self._get_graph_config()
+
+                # Track final chunk for result creation
+                final_chunk: Optional[dict[Any, Any]] = None
+
+                # Stream events from graph
+                async for stream_chunk in compiled_graph.astream(
+                    graph_input,
+                    graph_config,
+                    stream_mode=["messages", "updates"],
+                    subgraphs=True,
+                ):
+                    _, chunk_type, data = stream_chunk
+
+                    # Emit UiPathAgentMessageEvent for messages
+                    if chunk_type == "messages":
+                        if isinstance(data, tuple):
+                            message, _ = data
+                            event = UiPathAgentMessageEvent(
+                                payload=message,
+                                execution_id=self.context.execution_id,
+                            )
+                            yield event
+
+                    # Emit UiPathAgentStateEvent for state updates
+                    elif chunk_type == "updates":
+                        if isinstance(data, dict):
+                            final_chunk = data
+
+                            # Emit state update event for each node
+                            for node_name, agent_data in data.items():
+                                if isinstance(agent_data, dict):
+                                    state_event = UiPathAgentStateEvent(
+                                        payload=agent_data,
+                                        node_name=node_name,
+                                        execution_id=self.context.execution_id,
                                     )
-                                    if event:
-                                        self.context.chat_handler.on_event(event)
-                        elif chunk_type == "updates":
-                            if isinstance(data, dict):
-                                # data is a dict, e.g. {'agent': {'messages': [...]}}
-                                for agent_data in data.values():
-                                    if isinstance(agent_data, dict):
-                                        messages = agent_data.get("messages", [])
-                                        if isinstance(messages, list):
-                                            for message in messages:
-                                                if isinstance(message, BaseMessage):
-                                                    message.pretty_print()
-                                final_chunk = data
-
-                    self.context.output = self._extract_graph_result(
-                        final_chunk, compiled_graph.output_channels
-                    )
-                else:
-                    # Execute the graph normally at runtime or eval
-                    self.context.output = await compiled_graph.ainvoke(
-                        processed_input, graph_config
-                    )
-
-                # Get the state if available
-                try:
-                    self.context.state = await compiled_graph.aget_state(graph_config)
-                except Exception:
-                    pass
-
-                output_processor = await LangGraphOutputProcessor.create(self.context)
-
-                self.context.result = await output_processor.process()
-
-                return self.context.result
+                                    yield state_event
 
-        except Exception as e:
-            if isinstance(e, LangGraphRuntimeError):
-                raise
-
-            detail = f"Error: {str(e)}"
-
-            if isinstance(e, GraphRecursionError):
-                raise LangGraphRuntimeError(
-                    "GRAPH_RECURSION_ERROR",
-                    "Graph recursion limit exceeded",
-                    detail,
-                    UiPathErrorCategory.USER,
-                ) from e
-
-            if isinstance(e, InvalidUpdateError):
-                raise LangGraphRuntimeError(
-                    "GRAPH_INVALID_UPDATE",
-                    str(e),
-                    detail,
-                    UiPathErrorCategory.USER,
-                ) from e
-
-            if isinstance(e, EmptyInputError):
-                raise LangGraphRuntimeError(
-                    "GRAPH_EMPTY_INPUT",
-                    "The input data is empty",
-                    detail,
-                    UiPathErrorCategory.USER,
-                ) from e
-
-            raise LangGraphRuntimeError(
-                "EXECUTION_ERROR",
-                "Graph execution failed",
-                detail,
-                UiPathErrorCategory.USER,
-            ) from e
-        finally:
-            pass
+                # Extract output from final chunk
+                graph_output = self._extract_graph_result(
+                    final_chunk, compiled_graph.output_channels
+                )
 
-    async def validate(self) -> None:
-        pass
+                # Get final state and create result
+                self.context.result = await self._create_runtime_result(
+                    compiled_graph, graph_config, memory, graph_output
+                )
 
-    async def cleanup(self):
-        pass
+                # Yield the final result as last event
+                yield self.context.result
 
-    def _extract_graph_result(self, final_chunk, output_channels: str | Sequence[str]):
+        except Exception as e:
+            raise self._create_runtime_error(e) from e
+
+    async def _setup_graph(
+        self, memory: AsyncSqliteSaver, graph: StateGraph[Any, Any, Any]
+    ) -> CompiledStateGraph[Any, Any, Any]:
+        """Setup and compile the graph with memory and interrupts."""
+        interrupt_before: list[str] = []
+        interrupt_after: list[str] = []
+
+        return graph.compile(
+            checkpointer=memory,
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
+        )
+
+    def _get_graph_config(self) -> RunnableConfig:
+        """Build graph execution configuration."""
+        graph_config: RunnableConfig = {
+            "configurable": {
+                "thread_id": (
+                    self.context.execution_id or self.context.job_id or "default"
+                )
+            },
+            "callbacks": [],
+        }
+
+        # Add optional config from environment
+        recursion_limit = os.environ.get("LANGCHAIN_RECURSION_LIMIT")
+        max_concurrency = os.environ.get("LANGCHAIN_MAX_CONCURRENCY")
+
+        if recursion_limit is not None:
+            graph_config["recursion_limit"] = int(recursion_limit)
+        if max_concurrency is not None:
+            graph_config["max_concurrency"] = int(max_concurrency)
+
+        return graph_config
+
+    async def _get_graph_input(self, memory: AsyncSqliteSaver) -> Any:
+        """Process and return graph input."""
+        return await get_graph_input(
+            context=self.context,
+            memory=memory,
+            resume_triggers_table=self.resume_triggers_table,
+        )
+
+    async def _get_graph_state(
+        self,
+        compiled_graph: CompiledStateGraph[Any, Any, Any],
+        graph_config: RunnableConfig,
+    ) -> Optional[StateSnapshot]:
+        """Get final graph state."""
+        try:
+            return await compiled_graph.aget_state(graph_config)
+        except Exception:
+            return None
+
+    def _extract_graph_result(
+        self, final_chunk: Any, output_channels: str | Sequence[str]
+    ) -> Any:
         """
         Extract the result from a LangGraph output chunk according to the graph's output channels.
 
         Args:
             final_chunk: The final chunk from graph.astream()
-            graph: The LangGraph instance
+            output_channels: The graph's output channel configuration
 
         Returns:
             The extracted result according to the graph's output_channels configuration
         """
         # Unwrap from subgraph tuple format if needed
         if isinstance(final_chunk, tuple) and len(final_chunk) == 2:
-            final_chunk = final_chunk[
-                1
-            ]  # Extract data part from (namespace, data) tuple
+            final_chunk = final_chunk[1]
 
         # If the result isn't a dict or graph doesn't define output channels, return as is
         if not isinstance(final_chunk, dict):
@@ -205,10 +263,7 @@ class LangGraphRuntime(UiPathBaseRuntime):
 
         # Case 1: Single output channel as string
         if isinstance(output_channels, str):
-            if output_channels in final_chunk:
-                return final_chunk[output_channels]
-            else:
-                return final_chunk
+            return final_chunk.get(output_channels, final_chunk)
 
         # Case 2: Multiple output channels as sequence
         elif hasattr(output_channels, "__iter__") and not isinstance(
@@ -217,12 +272,12 @@ class LangGraphRuntime(UiPathBaseRuntime):
             # 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
+            # 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
-                ):
+            if not available_channels and len(final_chunk) == 1:
+                potential_unwrap = next(iter(final_chunk.values()))
+                if isinstance(potential_unwrap, dict):
+                    unwrapped_final_chunk = potential_unwrap
                     available_channels = [
                         ch for ch in output_channels if ch in unwrapped_final_chunk
                     ]
@@ -230,14 +285,184 @@ class LangGraphRuntime(UiPathBaseRuntime):
             if available_channels:
                 # Create a dict with the available channels
                 return {
-                    channel: final_chunk.get(channel, None)
-                    or unwrapped_final_chunk[channel]
+                    channel: final_chunk.get(channel)
+                    or unwrapped_final_chunk.get(channel)
                     for channel in available_channels
                 }
 
         # Fallback for any other case
         return final_chunk
 
+    def _is_interrupted(self, state: StateSnapshot) -> bool:
+        """Check if execution was interrupted (static or dynamic)."""
+        # Check for static interrupts (interrupt_before/after)
+        if hasattr(state, "next") and state.next:
+            return True
+
+        # Check for dynamic interrupts (interrupt() inside node)
+        if hasattr(state, "tasks"):
+            for task in state.tasks:
+                if hasattr(task, "interrupts") and task.interrupts:
+                    return True
+
+        return False
+
+    def _get_dynamic_interrupt(self, state: StateSnapshot) -> Optional[Interrupt]:
+        """Get the first dynamic interrupt if any."""
+        if not hasattr(state, "tasks"):
+            return None
+
+        for task in state.tasks:
+            if hasattr(task, "interrupts") and task.interrupts:
+                for interrupt in task.interrupts:
+                    if isinstance(interrupt, Interrupt):
+                        return interrupt
+        return None
+
+    async def _create_runtime_result(
+        self,
+        compiled_graph: CompiledStateGraph[Any, Any, Any],
+        graph_config: RunnableConfig,
+        memory: AsyncSqliteSaver,
+        graph_output: Optional[Any],
+    ) -> UiPathRuntimeResult:
+        """
+        Get final graph state and create the execution result.
+
+        Stores the result in self.context.result and self.context.state.
+
+        Args:
+            compiled_graph: The compiled graph instance
+            graph_config: The graph execution configuration
+            memory: The SQLite memory instance
+            graph_output: The graph execution output
+        """
+        # Get the final state
+        graph_state = await self._get_graph_state(compiled_graph, graph_config)
+
+        # Check if execution was interrupted (static or dynamic)
+        if graph_state and self._is_interrupted(graph_state):
+            return await self._create_suspended_result(
+                graph_state, memory, graph_output
+            )
+        else:
+            # Normal completion
+            return self._create_success_result(graph_output)
+
+    async def _create_suspended_result(
+        self,
+        graph_state: StateSnapshot,
+        graph_memory: AsyncSqliteSaver,
+        graph_output: Optional[Any],
+    ) -> UiPathRuntimeResult:
+        """Create result for suspended execution."""
+        # Check if it's a dynamic interrupt
+        dynamic_interrupt = self._get_dynamic_interrupt(graph_state)
+        resume_trigger: Optional[UiPathResumeTrigger] = None
+
+        if dynamic_interrupt:
+            # Dynamic interrupt - create and save resume trigger
+            resume_trigger = await create_and_save_resume_trigger(
+                interrupt_value=dynamic_interrupt.value,
+                memory=graph_memory,
+                resume_triggers_table=self.resume_triggers_table,
+            )
+            output = serialize_output(graph_output)
+            return UiPathRuntimeResult(
+                output=output,
+                status=UiPathRuntimeStatus.SUSPENDED,
+                resume=resume_trigger,
+            )
+        else:
+            # Static interrupt (breakpoint)
+            return self._create_breakpoint_result(graph_state)
+
+    def _create_breakpoint_result(
+        self,
+        graph_state: StateSnapshot,
+    ) -> UiPathBreakpointResult:
+        """Create result for execution paused at a breakpoint."""
+
+        # Get next nodes - these are the nodes that will execute when resumed
+        next_nodes = list(graph_state.next)
+
+        # Determine breakpoint type and node
+        if next_nodes:
+            # Breakpoint is BEFORE these nodes (interrupt_before)
+            breakpoint_type = "before"
+            breakpoint_node = next_nodes[0]
+        else:
+            # Breakpoint is AFTER the last executed node (interrupt_after)
+            # Get the last executed node from tasks
+            breakpoint_type = "after"
+            if graph_state.tasks:
+                # Tasks contain the nodes that just executed
+                # Get the last task's name
+                breakpoint_node = graph_state.tasks[-1].name
+            else:
+                # Fallback if no tasks (shouldn't happen)
+                breakpoint_node = "unknown"
+
+        return UiPathBreakpointResult(
+            breakpoint_node=breakpoint_node,
+            breakpoint_type=breakpoint_type,
+            current_state=graph_state.values,
+            next_nodes=next_nodes,
+        )
+
+    def _create_success_result(self, output: Optional[Any]) -> UiPathRuntimeResult:
+        """Create result for successful completion."""
+        return UiPathRuntimeResult(
+            output=serialize_output(output),
+            status=UiPathRuntimeStatus.SUCCESSFUL,
+        )
+
+    def _create_runtime_error(self, e: Exception) -> LangGraphRuntimeError:
+        """Handle execution errors and create appropriate LangGraphRuntimeError."""
+        if isinstance(e, LangGraphRuntimeError):
+            return e
+
+        detail = f"Error: {str(e)}"
+
+        if isinstance(e, GraphRecursionError):
+            return LangGraphRuntimeError(
+                "GRAPH_RECURSION_ERROR",
+                "Graph recursion limit exceeded",
+                detail,
+                UiPathErrorCategory.USER,
+            )
+
+        if isinstance(e, InvalidUpdateError):
+            return LangGraphRuntimeError(
+                "GRAPH_INVALID_UPDATE",
+                str(e),
+                detail,
+                UiPathErrorCategory.USER,
+            )
+
+        if isinstance(e, EmptyInputError):
+            return LangGraphRuntimeError(
+                "GRAPH_EMPTY_INPUT",
+                "The input data is empty",
+                detail,
+                UiPathErrorCategory.USER,
+            )
+
+        return LangGraphRuntimeError(
+            "EXECUTION_ERROR",
+            "Graph execution failed",
+            detail,
+            UiPathErrorCategory.USER,
+        )
+
+    async def validate(self) -> None:
+        """Validate runtime inputs."""
+        pass
+
+    async def cleanup(self) -> None:
+        """Cleanup runtime resources."""
+        pass
+
 
 class LangGraphScriptRuntime(LangGraphRuntime):
     """
@@ -250,6 +475,7 @@ class LangGraphScriptRuntime(LangGraphRuntime):
         self.resolver = LangGraphJsonResolver(entrypoint=entrypoint)
         super().__init__(context, self.resolver)
 
-    async def cleanup(self):
+    async def cleanup(self) -> None:
+        """Cleanup runtime resources including resolver."""
         await super().cleanup()
         await self.resolver.cleanup()