Merge pull request #227 from UiPath/fix/refactor_runtime

fix: refactor runtime execute

Author: cristipufu

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "uipath-langchain"
-version = "0.0.140"
+version = "0.0.141"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
-    "uipath>=2.1.76, <2.2.0",
+    "uipath>=2.1.96, <2.2.0",
     "langgraph>=0.5.0, <0.7.0",
     "langchain-core>=0.3.34",
     "langgraph-checkpoint-sqlite>=2.0.3",

src/uipath_langchain/_cli/_runtime/_context.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Optional
 
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from uipath._cli._runtime._contracts import UiPathRuntimeContext
@@ -7,10 +7,4 @@
 class LangGraphRuntimeContext(UiPathRuntimeContext):
     """Context information passed throughout the runtime execution."""
 
-    output: Optional[Any] = None
-    state: Optional[Any] = (
-        None  # TypedDict issue, the actual type is: Optional[langgraph.types.StateSnapshot]
-    )
     memory: Optional[AsyncSqliteSaver] = None
-    langsmith_tracing_enabled: Union[str, bool, None] = False
-    resume_triggers_table: str = "__uipath_resume_triggers"

src/uipath_langchain/_cli/_runtime/_input.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Any, Optional, cast
 
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.types import Command
 from uipath._cli._runtime._contracts import (
     UiPathApiTrigger,
@@ -17,123 +18,130 @@
 logger = logging.getLogger(__name__)
 
 
-class LangGraphInputProcessor:
+async def get_graph_input(
+    context: LangGraphRuntimeContext,
+    memory: AsyncSqliteSaver,
+    resume_triggers_table: str = "__uipath_resume_triggers",
+) -> Any:
     """
-    Handles input processing for graph execution, including resume scenarios
-    where it needs to fetch data from UiPath.
+    Process the input data for graph execution, handling both fresh starts and resume scenarios.
+
+    This method determines whether the graph is being executed fresh or resumed from a previous state.
+    For fresh executions, it returns the input JSON directly. For resume scenarios, it fetches
+    the latest trigger information from the database and constructs a Command object with the
+    appropriate resume data.
+
+    The method handles different types of resume triggers:
+    - API triggers: Creates an UiPathApiTrigger with inbox_id and request payload
+    - Other triggers: Uses the HitlReader to process the resume data
+
+    Args:
+        context: The runtime context for the graph execution.
+        memory: AsyncSqliteSaver. The async database saver used to fetch resume trigger data.
+        resume_triggers_table: str, optional. The name of the database table containing resume triggers (default: "__uipath_resume_triggers").
+
+    Returns:
+        Any: For fresh executions, returns the input JSON data directly.
+             For resume scenarios, returns a Command object containing the resume data
+             processed through the appropriate trigger handler.
+
+    Raises:
+        LangGraphRuntimeError: If there's an error fetching trigger data from the database
+            during resume processing.
     """
-
-    def __init__(self, context: LangGraphRuntimeContext):
-        """
-        Initialize the LangGraphInputProcessor.
-
-        Args:
-            context: The runtime context for the graph execution.
-        """
-        self.context = context
-
-    async def process(self) -> Any:
-        """
-        Process the input data for graph execution, handling both fresh starts and resume scenarios.
-
-        This method determines whether the graph is being executed fresh or resumed from a previous state.
-        For fresh executions, it returns the input JSON directly. For resume scenarios, it fetches
-        the latest trigger information from the database and constructs a Command object with the
-        appropriate resume data.
-
-        The method handles different types of resume triggers:
-        - API triggers: Creates an UiPathApiTrigger with inbox_id and request payload
-        - Other triggers: Uses the HitlReader to process the resume data
-
-        Returns:
-            Any: For fresh executions, returns the input JSON data directly.
-                 For resume scenarios, returns a Command object containing the resume data
-                 processed through the appropriate trigger handler.
-
-        Raises:
-            LangGraphRuntimeError: If there's an error fetching trigger data from the database
-                during resume processing.
-        """
-        logger.debug(f"Resumed: {self.context.resume} Input: {self.context.input_json}")
-
-        if not self.context.resume:
-            if self.context.input_message:
-                return {
-                    "messages": uipath_to_human_messages(self.context.input_message)
-                }
-            return self.context.input_json
-
-        if self.context.input_json:
-            return Command(resume=self.context.input_json)
-
-        trigger = await self._get_latest_trigger()
-        if not trigger:
-            return Command(resume=self.context.input_json)
-
-        trigger_type, key, folder_path, folder_key, payload = trigger
-        resume_trigger = UiPathResumeTrigger(
-            trigger_type=trigger_type,
-            item_key=key,
-            folder_path=folder_path,
-            folder_key=folder_key,
-            payload=payload,
+    logger.debug(f"Resumed: {context.resume} Input: {context.input_json}")
+
+    # Fresh execution - return input directly
+    if not context.resume:
+        if context.input_message:
+            return {"messages": uipath_to_human_messages(context.input_message)}
+        return context.input_json
+
+    # Resume with explicit input provided
+    if context.input_json:
+        return Command(resume=context.input_json)
+
+    # Resume from database trigger
+    trigger = await _get_latest_trigger(
+        memory, resume_triggers_table=resume_triggers_table
+    )
+    if not trigger:
+        return Command(resume=context.input_json)
+
+    trigger_type, key, folder_path, folder_key, payload = trigger
+    resume_trigger = UiPathResumeTrigger(
+        trigger_type=trigger_type,
+        item_key=key,
+        folder_path=folder_path,
+        folder_key=folder_key,
+        payload=payload,
+    )
+    logger.debug(f"ResumeTrigger: {trigger_type} {key}")
+
+    # Populate back expected fields for api_triggers
+    if resume_trigger.trigger_type == UiPathResumeTriggerType.API:
+        resume_trigger.api_resume = UiPathApiTrigger(
+            inbox_id=resume_trigger.item_key, request=resume_trigger.payload
         )
-        logger.debug(f"ResumeTrigger: {trigger_type} {key}")
-
-        # populate back expected fields for api_triggers
-        if resume_trigger.trigger_type == UiPathResumeTriggerType.API:
-            resume_trigger.api_resume = UiPathApiTrigger(
-                inbox_id=resume_trigger.item_key, request=resume_trigger.payload
-            )
-        return Command(resume=await HitlReader.read(resume_trigger))
-
-    async def _get_latest_trigger(self) -> Optional[tuple[str, str, str, str, str]]:
-        """
-        Fetch the most recent resume trigger from the database.
-
-        This private method queries the resume triggers table to retrieve the latest trigger
-        information based on timestamp. It handles database connection setup and executes
-        a SQL query to fetch trigger data needed for resume operations.
-
-        The method returns trigger information as a tuple containing:
-        - type: The type of trigger (e.g., 'API', 'MANUAL', etc.)
-        - key: The unique identifier for the trigger/item
-        - folder_path: The path to the folder containing the trigger
-        - folder_key: The unique identifier for the folder
-        - payload: The serialized payload data associated with the trigger
-
-        Returns:
-            Optional[tuple[str, str, str, str, str]]: A tuple containing (type, key, folder_path,
-                folder_key, payload) for the most recent trigger, or None if no triggers are found
-                or if the memory context is not available.
-
-        Raises:
-            LangGraphRuntimeError: If there's an error during database connection setup, query
-                execution, or result fetching. The original exception is wrapped with context
-                about the database operation failure.
-        """
-        if self.context.memory is None:
-            return None
-        try:
-            await self.context.memory.setup()
-            async with (
-                self.context.memory.lock,
-                self.context.memory.conn.cursor() as cur,
-            ):
-                await cur.execute(f"""
-                    SELECT type, key, folder_path, folder_key, payload
-                    FROM {self.context.resume_triggers_table}
-                    ORDER BY timestamp DESC
-                    LIMIT 1
-                """)
-                result = await cur.fetchone()
-                if result is None:
-                    return None
-                return cast(tuple[str, str, str, str, str], tuple(result))
-        except Exception as e:
-            raise LangGraphRuntimeError(
-                "DB_QUERY_FAILED",
-                "Database query failed",
-                f"Error querying resume trigger information: {str(e)}",
-                UiPathErrorCategory.SYSTEM,
-            ) from e
+
+    return Command(resume=await HitlReader.read(resume_trigger))
+
+
+async def _get_latest_trigger(
+    memory: AsyncSqliteSaver,
+    resume_triggers_table: str = "__uipath_resume_triggers",
+) -> Optional[tuple[str, str, str, str, str]]:
+    """
+    Fetch the most recent resume trigger from the database.
+
+    This private method queries the resume triggers table to retrieve the latest trigger
+    information based on timestamp. It handles database connection setup and executes
+    a SQL query to fetch trigger data needed for resume operations.
+
+    The method returns trigger information as a tuple containing:
+    - type: The type of trigger (e.g., 'API', 'MANUAL', etc.)
+    - key: The unique identifier for the trigger/item
+    - folder_path: The path to the folder containing the trigger
+    - folder_key: The unique identifier for the folder
+    - payload: The serialized payload data associated with the trigger
+
+    Args:
+        memory: The AsyncSqliteSaver instance used to access the database.
+        resume_triggers_table: The name of the table containing resume triggers (default: "__uipath_resume_triggers").
+
+    Returns:
+        Optional[tuple[str, str, str, str, str]]: A tuple containing (type, key, folder_path,
+            folder_key, payload) for the most recent trigger, or None if no triggers are found
+            or if the memory context is not available.
+
+    Raises:
+        LangGraphRuntimeError: If there's an error during database connection setup, query
+            execution, or result fetching. The original exception is wrapped with context
+            about the database operation failure.
+    """
+    if memory is None:
+        return None
+
+    try:
+        await memory.setup()
+        async with (
+            memory.lock,
+            memory.conn.cursor() as cur,
+        ):
+            await cur.execute(f"""
+                SELECT type, key, folder_path, folder_key, payload
+                FROM {resume_triggers_table}
+                ORDER BY timestamp DESC
+                LIMIT 1
+            """)
+            result = await cur.fetchone()
+            if result is None:
+                return None
+            return cast(tuple[str, str, str, str, str], tuple(result))
+    except Exception as e:
+        raise LangGraphRuntimeError(
+            "DB_QUERY_FAILED",
+            "Database query failed",
+            f"Error querying resume trigger information: {str(e)}",
+            UiPathErrorCategory.SYSTEM,
+        ) from e