uipath-langchain 0.1.34__py3-none-any.whl → 0.3.1__py3-none-any.whl

uipath_langchain/chat/mapper.py

@@ -41,6 +41,7 @@ class UiPathChatMessagesMapper:
     def __init__(self):
         """Initialize the mapper with empty state."""
         self.tool_call_to_ai_message: dict[str, str] = {}
+        self.current_message: AIMessageChunk
         self.seen_message_ids: set[str] = set()
 
     def _extract_text(self, content: Any) -> str:
@@ -141,7 +142,7 @@ class UiPathChatMessagesMapper:
     def map_event(
         self,
         message: BaseMessage,
-    ) -> UiPathConversationMessageEvent | None:
+    ) -> list[UiPathConversationMessageEvent] | None:
         """Convert LangGraph BaseMessage (chunk or full) into a UiPathConversationMessageEvent.
 
         Args:
@@ -168,16 +169,45 @@ class UiPathChatMessagesMapper:
 
             # Check if this is the last chunk by examining chunk_position
             if message.chunk_position == "last":
+                events: list[UiPathConversationMessageEvent] = []
+
+                # Loop through all content_blocks in current_message and create toolCallStart events for each tool_call_chunk
+                if self.current_message and self.current_message.content_blocks:
+                    for block in self.current_message.content_blocks:
+                        if block.get("type") == "tool_call_chunk":
+                            tool_chunk_block = cast(ToolCallChunk, block)
+                            tool_call_id = tool_chunk_block.get("id")
+                            tool_name = tool_chunk_block.get("name")
+                            tool_args = tool_chunk_block.get("args")
+
+                            if tool_call_id:
+                                tool_event = UiPathConversationMessageEvent(
+                                    message_id=message.id,
+                                    tool_call=UiPathConversationToolCallEvent(
+                                        tool_call_id=tool_call_id,
+                                        start=UiPathConversationToolCallStartEvent(
+                                            tool_name=tool_name,
+                                            timestamp=timestamp,
+                                            input=UiPathInlineValue(inline=tool_args),
+                                        ),
+                                    ),
+                                )
+                                events.append(tool_event)
+
+                # Create the final event for the message
                 msg_event.end = UiPathConversationMessageEndEvent(timestamp=timestamp)
                 msg_event.content_part = UiPathConversationContentPartEvent(
                     content_part_id=f"chunk-{message.id}-0",
                     end=UiPathConversationContentPartEndEvent(),
                 )
-                return msg_event
+                events.append(msg_event)
+
+                return events
 
             # For every new message_id, start a new message
             if message.id not in self.seen_message_ids:
                 self.seen_message_ids.add(message.id)
+                self.current_message = message
                 msg_event.start = UiPathConversationMessageStartEvent(
                     role="assistant", timestamp=timestamp
                 )
@@ -200,7 +230,6 @@ class UiPathChatMessagesMapper:
                             content_part_id=f"chunk-{message.id}-0",
                             chunk=UiPathConversationContentPartChunkEvent(
                                 data=text,
-                                content_part_sequence=0,
                             ),
                         )
 
@@ -210,19 +239,10 @@ class UiPathChatMessagesMapper:
                         tool_call_id = tool_chunk_block.get("id")
                         if tool_call_id:
                             # Track tool_call_id -> ai_message_id mapping
-                            self.tool_call_to_ai_message[str(tool_call_id)] = message.id
-
-                        args = tool_chunk_block.get("args") or ""
+                            self.tool_call_to_ai_message[tool_call_id] = message.id
 
-                        msg_event.content_part = UiPathConversationContentPartEvent(
-                            content_part_id=f"chunk-{message.id}-0",
-                            chunk=UiPathConversationContentPartChunkEvent(
-                                data=args,
-                                content_part_sequence=0,
-                            ),
-                        )
-                        # Continue so that multiple tool_call_chunks in the same block list
-                        # are handled correctly
+                        # Accumulate the message chunk
+                        self.current_message = self.current_message + message
                         continue
 
             # Fallback: raw string content on the chunk (rare when using content_blocks)
@@ -231,7 +251,6 @@ class UiPathChatMessagesMapper:
                     content_part_id=f"content-{message.id}",
                     chunk=UiPathConversationContentPartChunkEvent(
                         data=message.content,
-                        content_part_sequence=0,
                     ),
                 )
 
@@ -241,7 +260,7 @@ class UiPathChatMessagesMapper:
                 or msg_event.tool_call
                 or msg_event.end
             ):
-                return msg_event
+                return [msg_event]
 
             return None
 
@@ -275,35 +294,34 @@ class UiPathChatMessagesMapper:
                     # Keep as string if not valid JSON
                     pass
 
-            return UiPathConversationMessageEvent(
-                message_id=result_message_id or str(uuid4()),
-                tool_call=UiPathConversationToolCallEvent(
-                    tool_call_id=message.tool_call_id,
-                    start=UiPathConversationToolCallStartEvent(
-                        tool_name=message.name,
-                        arguments=None,
-                        timestamp=timestamp,
+            return [
+                UiPathConversationMessageEvent(
+                    message_id=result_message_id or str(uuid4()),
+                    tool_call=UiPathConversationToolCallEvent(
+                        tool_call_id=message.tool_call_id,
+                        end=UiPathConversationToolCallEndEvent(
+                            timestamp=timestamp,
+                            output=UiPathInlineValue(inline=content_value),
+                        ),
                     ),
-                    end=UiPathConversationToolCallEndEvent(
-                        timestamp=timestamp,
-                        output=UiPathInlineValue(inline=content_value),
-                    ),
-                ),
-            )
+                )
+            ]
 
         # --- Fallback for other BaseMessage types ---
         text_content = self._extract_text(message.content)
-        return UiPathConversationMessageEvent(
-            message_id=message.id,
-            start=UiPathConversationMessageStartEvent(
-                role="assistant", timestamp=timestamp
-            ),
-            content_part=UiPathConversationContentPartEvent(
-                content_part_id=f"cp-{message.id}",
-                chunk=UiPathConversationContentPartChunkEvent(data=text_content),
-            ),
-            end=UiPathConversationMessageEndEvent(),
-        )
+        return [
+            UiPathConversationMessageEvent(
+                message_id=message.id,
+                start=UiPathConversationMessageStartEvent(
+                    role="assistant", timestamp=timestamp
+                ),
+                content_part=UiPathConversationContentPartEvent(
+                    content_part_id=f"cp-{message.id}",
+                    chunk=UiPathConversationContentPartChunkEvent(data=text_content),
+                ),
+                end=UiPathConversationMessageEndEvent(),
+            )
+        ]
 
 
 __all__ = ["UiPathChatMessagesMapper"]

uipath_langchain/runtime/factory.py

@@ -92,7 +92,7 @@ class UiPathLangGraphRuntimeFactory:
         return self._config
 
     async def _load_graph(
-        self, entrypoint: str
+        self, entrypoint: str, **kwargs
     ) -> StateGraph[Any, Any, Any] | CompiledStateGraph[Any, Any, Any, Any]:
         """
         Load a graph for the given entrypoint.
@@ -181,7 +181,7 @@ class UiPathLangGraphRuntimeFactory:
         return builder.compile(checkpointer=memory)
 
     async def _resolve_and_compile_graph(
-        self, entrypoint: str, memory: AsyncSqliteSaver
+        self, entrypoint: str, memory: AsyncSqliteSaver, **kwargs
     ) -> CompiledStateGraph[Any, Any, Any, Any]:
         """
         Resolve a graph from configuration and compile it.
@@ -201,7 +201,7 @@ class UiPathLangGraphRuntimeFactory:
             if entrypoint in self._graph_cache:
                 return self._graph_cache[entrypoint]
 
-            loaded_graph = await self._load_graph(entrypoint)
+            loaded_graph = await self._load_graph(entrypoint, **kwargs)
 
             compiled_graph = await self._compile_graph(loaded_graph, memory)
 
@@ -249,6 +249,7 @@ class UiPathLangGraphRuntimeFactory:
         compiled_graph: CompiledStateGraph[Any, Any, Any, Any],
         runtime_id: str,
         entrypoint: str,
+        **kwargs,
     ) -> UiPathRuntimeProtocol:
         """
         Create a runtime instance from a compiled graph.
@@ -275,10 +276,11 @@ class UiPathLangGraphRuntimeFactory:
             delegate=base_runtime,
             storage=storage,
             trigger_manager=trigger_manager,
+            runtime_id=runtime_id,
         )
 
     async def new_runtime(
-        self, entrypoint: str, runtime_id: str
+        self, entrypoint: str, runtime_id: str, **kwargs
     ) -> UiPathRuntimeProtocol:
         """
         Create a new LangGraph runtime instance.
@@ -293,12 +295,15 @@ class UiPathLangGraphRuntimeFactory:
         # Get shared memory instance
         memory = await self._get_memory()
 
-        compiled_graph = await self._resolve_and_compile_graph(entrypoint, memory)
+        compiled_graph = await self._resolve_and_compile_graph(
+            entrypoint, memory, **kwargs
+        )
 
         return await self._create_runtime_instance(
             compiled_graph=compiled_graph,
             runtime_id=runtime_id,
             entrypoint=entrypoint,
+            **kwargs,
         )
 
     async def dispose(self) -> None:

uipath_langchain/runtime/runtime.py

@@ -3,6 +3,7 @@ import os
 from typing import Any, AsyncGenerator
 from uuid import uuid4
 
+from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.runnables.config import RunnableConfig
 from langgraph.errors import EmptyInputError, GraphRecursionError, InvalidUpdateError
 from langgraph.graph.state import CompiledStateGraph
@@ -41,6 +42,7 @@ class UiPathLangGraphRuntime:
         graph: CompiledStateGraph[Any, Any, Any, Any],
         runtime_id: str | None = None,
         entrypoint: str | None = None,
+        callbacks: list[BaseCallbackHandler] | None = None,
     ):
         """
         Initialize the runtime.
@@ -53,6 +55,7 @@ class UiPathLangGraphRuntime:
         self.graph: CompiledStateGraph[Any, Any, Any, Any] = graph
         self.runtime_id: str = runtime_id or "default"
         self.entrypoint: str | None = entrypoint
+        self.callbacks: list[BaseCallbackHandler] = callbacks or []
         self.chat = UiPathChatMessagesMapper()
         self._middleware_node_names: set[str] = self._detect_middleware_nodes()
 
@@ -135,10 +138,17 @@ class UiPathLangGraphRuntime:
                 if chunk_type == "messages":
                     if isinstance(data, tuple):
                         message, _ = data
-                        event = UiPathRuntimeMessageEvent(
-                            payload=self.chat.map_event(message),
-                        )
-                        yield event
+                        try:
+                            events = self.chat.map_event(message)
+                        except Exception as e:
+                            logger.warning(f"Error mapping message event: {e}")
+                            events = None
+                        if events:
+                            for mapped_event in events:
+                                event = UiPathRuntimeMessageEvent(
+                                    payload=mapped_event,
+                                )
+                                yield event
 
                 # Emit UiPathRuntimeStateEvent for state updates
                 elif chunk_type == "updates":
@@ -153,6 +163,8 @@ class UiPathLangGraphRuntime:
 
                         # Emit state update event for each node
                         for node_name, agent_data in data.items():
+                            if node_name in ("__metadata__",):
+                                continue
                             if isinstance(agent_data, dict):
                                 state_event = UiPathRuntimeStateEvent(
                                     payload=serialize_output(agent_data),
@@ -189,7 +201,7 @@ class UiPathLangGraphRuntime:
         """Build graph execution configuration."""
         graph_config: RunnableConfig = {
             "configurable": {"thread_id": self.runtime_id},
-            "callbacks": [],
+            "callbacks": self.callbacks,
         }
 
         # Add optional config from environment
@@ -283,29 +295,9 @@ class UiPathLangGraphRuntime:
 
     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) -> Interrupt | None:
-        """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
+        # An execution is considered interrupted if there are any next nodes (static interrupt)
+        # or if there are any dynamic interrupts present
+        return bool(state.next) or bool(state.interrupts)
 
     async def _create_runtime_result(
         self,
@@ -334,13 +326,24 @@ class UiPathLangGraphRuntime:
         graph_state: StateSnapshot,
     ) -> UiPathRuntimeResult:
         """Create result for suspended execution."""
-        # Check if it's a dynamic interrupt
-        dynamic_interrupt = self._get_dynamic_interrupt(graph_state)
-
-        if dynamic_interrupt:
-            # Dynamic interrupt - should create and save resume trigger
+        interrupt_map: dict[str, Any] = {}
+
+        if graph_state.interrupts:
+            for interrupt in graph_state.interrupts:
+                if isinstance(interrupt, Interrupt):
+                    # Find which task this interrupt belongs to
+                    for task in graph_state.tasks:
+                        if task.interrupts and interrupt in task.interrupts:
+                            # Only include if this task is still waiting for interrupt resolution
+                            if task.interrupts and not task.result:
+                                interrupt_map[interrupt.id] = interrupt.value
+                            break
+
+        # If we have dynamic interrupts, return suspended with interrupt map
+        # The output is used to create the resume triggers
+        if interrupt_map:
             return UiPathRuntimeResult(
-                output=dynamic_interrupt.value,
+                output=interrupt_map,
                 status=UiPathRuntimeStatus.SUSPENDED,
             )
         else:
@@ -360,7 +363,7 @@ class UiPathLangGraphRuntime:
         if next_nodes:
             # Breakpoint is BEFORE these nodes (interrupt_before)
             breakpoint_type = "before"
-            breakpoint_node = next_nodes[0]
+            breakpoint_node = ", ".join(next_nodes)
         else:
             # Breakpoint is AFTER the last executed node (interrupt_after)
             # Get the last executed node from tasks

uipath_langchain/runtime/storage.py

@@ -1,115 +1,222 @@
 """SQLite implementation of UiPathResumableStorageProtocol."""
 
 import json
-from typing import cast
+from typing import Any, cast
 
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from pydantic import BaseModel
-from uipath.runtime import (
-    UiPathApiTrigger,
-    UiPathResumeTrigger,
-    UiPathResumeTriggerName,
-    UiPathResumeTriggerType,
-)
+from uipath.runtime import UiPathResumeTrigger
 
 
 class SqliteResumableStorage:
-    """SQLite storage for resume triggers."""
+    """SQLite storage for resume triggers and arbitrary kv pairs."""
 
     def __init__(
-        self, memory: AsyncSqliteSaver, table_name: str = "__uipath_resume_triggers"
+        self,
+        memory: AsyncSqliteSaver,
     ):
         self.memory = memory
-        self.table_name = table_name
+        self.rs_table_name = "__uipath_resume_triggers"
+        self.kv_table_name = "__uipath_runtime_kv"
         self._initialized = False
 
     async def _ensure_table(self) -> None:
-        """Create table if needed."""
+        """Create tables if needed."""
         if self._initialized:
             return
 
         await self.memory.setup()
         async with self.memory.lock, self.memory.conn.cursor() as cur:
-            await cur.execute(f"""
-                CREATE TABLE IF NOT EXISTS {self.table_name} (
+            # Enable WAL mode for high concurrency
+            await cur.execute("PRAGMA journal_mode=WAL")
+
+            await cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {self.rs_table_name} (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    type TEXT NOT NULL,
-                    name TEXT NOT NULL,
-                    key TEXT,
-                    folder_key TEXT,
-                    folder_path TEXT,
-                    payload TEXT,
+                    runtime_id TEXT NOT NULL,
+                    interrupt_id TEXT NOT NULL,
+                    data TEXT NOT NULL,
                     timestamp DATETIME DEFAULT (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc'))
                 )
-            """)
-            await self.memory.conn.commit()
-            self._initialized = True
+                """
+            )
 
-    async def save_trigger(self, trigger: UiPathResumeTrigger) -> None:
-        """Save resume trigger to database."""
-        await self._ensure_table()
+            await cur.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS idx_{self.rs_table_name}_runtime_id
+                ON {self.rs_table_name}(runtime_id)
+                """
+            )
 
-        trigger_key = (
-            trigger.api_resume.inbox_id if trigger.api_resume else trigger.item_key
-        )
-        payload = trigger.payload
-        if payload:
-            payload = (
-                (
-                    payload.model_dump()
-                    if isinstance(payload, BaseModel)
-                    else json.dumps(payload)
+            await cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {self.kv_table_name} (
+                    runtime_id TEXT NOT NULL,
+                    namespace TEXT NOT NULL,
+                    key TEXT NOT NULL,
+                    value TEXT,
+                    timestamp DATETIME DEFAULT (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc')),
+                    PRIMARY KEY (runtime_id, namespace, key)
                 )
-                if isinstance(payload, dict)
-                else str(payload)
+                """
             )
 
+            await self.memory.conn.commit()
+
+        self._initialized = True
+
+    async def save_triggers(
+        self, runtime_id: str, triggers: list[UiPathResumeTrigger]
+    ) -> None:
+        """Save resume triggers to database, replacing all existing triggers for this runtime_id."""
+        await self._ensure_table()
+
         async with self.memory.lock, self.memory.conn.cursor() as cur:
+            # Delete all existing triggers for this runtime_id
             await cur.execute(
-                f"INSERT INTO {self.table_name} (type, key, name, payload, folder_path, folder_key) VALUES (?, ?, ?, ?, ?, ?)",
-                (
-                    trigger.trigger_type.value,
-                    trigger_key,
-                    trigger.trigger_name.value,
-                    payload,
-                    trigger.folder_path,
-                    trigger.folder_key,
-                ),
+                f"""
+                DELETE FROM {self.rs_table_name}
+                WHERE runtime_id = ?
+                """,
+                (runtime_id,),
             )
+
+            # Insert new triggers
+            for trigger in triggers:
+                trigger_data = trigger.model_dump()
+                trigger_data["payload"] = trigger.payload
+                trigger_data["trigger_name"] = trigger.trigger_name
+
+                await cur.execute(
+                    f"""
+                    INSERT INTO {self.rs_table_name}
+                        (runtime_id, interrupt_id, data)
+                    VALUES (?, ?, ?)
+                    """,
+                    (
+                        runtime_id,
+                        trigger.interrupt_id,
+                        json.dumps(trigger_data),
+                    ),
+                )
             await self.memory.conn.commit()
 
-    async def get_latest_trigger(self) -> UiPathResumeTrigger | None:
-        """Get most recent trigger from database."""
+    async def get_triggers(self, runtime_id: str) -> list[UiPathResumeTrigger] | None:
+        """Get all triggers for runtime_id from database."""
         await self._ensure_table()
 
         async with self.memory.lock, self.memory.conn.cursor() as cur:
-            await cur.execute(f"""
-                SELECT type, key, name, folder_path, folder_key, payload
-                FROM {self.table_name}
-                ORDER BY timestamp DESC
-                LIMIT 1
-            """)
-            result = await cur.fetchone()
+            await cur.execute(
+                f"""
+                SELECT data
+                FROM {self.rs_table_name}
+                WHERE runtime_id = ?
+                ORDER BY timestamp ASC
+                """,
+                (runtime_id,),
+            )
+            results = await cur.fetchall()
+
+        if not results:
+            return None
 
-            if not result:
-                return None
+        triggers = []
+        for result in results:
+            data_text = cast(str, result[0])
+            trigger = UiPathResumeTrigger.model_validate_json(data_text)
+            triggers.append(trigger)
 
-            trigger_type, key, name, folder_path, folder_key, payload = cast(
-                tuple[str, str, str, str, str, str], tuple(result)
+        return triggers
+
+    async def delete_trigger(
+        self, runtime_id: str, trigger: UiPathResumeTrigger
+    ) -> None:
+        """Delete resume trigger from storage."""
+        await self._ensure_table()
+
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(
+                f"""
+                DELETE FROM {self.rs_table_name}
+                WHERE runtime_id = ? AND interrupt_id = ?
+                """,
+                (
+                    runtime_id,
+                    trigger.interrupt_id,
+                ),
             )
+            await self.memory.conn.commit()
+
+    async def set_value(
+        self,
+        runtime_id: str,
+        namespace: str,
+        key: str,
+        value: Any,
+    ) -> None:
+        """Save arbitrary key-value pair to database."""
+        if not (
+            isinstance(value, str)
+            or isinstance(value, dict)
+            or isinstance(value, BaseModel)
+            or value is None
+        ):
+            raise TypeError("Value must be str, dict, BaseModel or None.")
+
+        await self._ensure_table()
 
-            resume_trigger = UiPathResumeTrigger(
-                trigger_type=UiPathResumeTriggerType(trigger_type),
-                trigger_name=UiPathResumeTriggerName(name),
-                item_key=key,
-                folder_path=folder_path,
-                folder_key=folder_key,
-                payload=payload,
+        value_text = self._dump_value(value)
+
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(
+                f"""
+                INSERT INTO {self.kv_table_name} (runtime_id, namespace, key, value)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(runtime_id, namespace, key)
+                DO UPDATE SET
+                    value = excluded.value,
+                    timestamp = (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc'))
+                """,
+                (runtime_id, namespace, key, value_text),
             )
+            await self.memory.conn.commit()
 
-            if resume_trigger.trigger_type == UiPathResumeTriggerType.API:
-                resume_trigger.api_resume = UiPathApiTrigger(
-                    inbox_id=resume_trigger.item_key, request=resume_trigger.payload
-                )
+    async def get_value(self, runtime_id: str, namespace: str, key: str) -> Any:
+        """Get arbitrary key-value pair from database (scoped by runtime_id + namespace)."""
+        await self._ensure_table()
 
-            return resume_trigger
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(
+                f"""
+                SELECT value
+                FROM {self.kv_table_name}
+                WHERE runtime_id = ? AND namespace = ? AND key = ?
+                LIMIT 1
+                """,
+                (runtime_id, namespace, key),
+            )
+            row = await cur.fetchone()
+
+        if not row:
+            return None
+
+        return self._load_value(cast(str | None, row[0]))
+
+    def _dump_value(self, value: str | dict[str, Any] | BaseModel | None) -> str | None:
+        if value is None:
+            return None
+        if isinstance(value, BaseModel):
+            return "j:" + json.dumps(value.model_dump())
+        if isinstance(value, dict):
+            return "j:" + json.dumps(value)
+        return "s:" + value
+
+    def _load_value(self, raw: str | None) -> Any:
+        if raw is None:
+            return None
+        if raw.startswith("s:"):
+            return raw[2:]
+        if raw.startswith("j:"):
+            return json.loads(raw[2:])
+        return raw