clean up

Author: JoshParkSJ

src/uipath_langchain/agent/tools/tool_node.py

@@ -274,11 +274,7 @@ async def _afunc(state: AgentGraphState) -> OutputType:
                 raise
             return result
 
-    wrapped = RunnableCallable(func=_func, afunc=_afunc, name=tool_name)
-    # Preserve .tool so _get_tool_confirmation_info can find metadata
-    if hasattr(tool_node, "tool"):
-        wrapped.tool = tool_node.tool  # type: ignore[attr-defined]
-    return wrapped
+    return RunnableCallable(func=_func, afunc=_afunc, name=tool_name)
 
 
 class ToolWrapperMixin:

src/uipath_langchain/chat/hitl.py

@@ -108,11 +108,6 @@ def request_approval(
     """Interrupt the graph to request user approval for a tool call.
 
     Returns the (possibly edited) tool arguments if approved, or None if rejected.
-
-    Confirmation data (inputSchema, inputValue) is included in the startToolCall
-    event emitted by the mapper. The @durable_interrupt pauses the graph; the
-    resume payload is the confirmToolCall event body from CAS:
-      {"approved": bool, "input": <edited args | None>}
     """
     tool_call_id: str = tool_args.pop("tool_call_id")
 

src/uipath_langchain/runtime/messages.py

@@ -59,7 +59,6 @@ def __init__(self, runtime_id: str, storage: UiPathRuntimeStorageProtocol | None
         self.runtime_id = runtime_id
         self.storage = storage
         self.current_message: AIMessageChunk | AIMessage
-        self.tool_names_requiring_confirmation: set[str] = set()
         self.tool_confirmation_schemas: dict[str, Any] = {}
         self.seen_message_ids: set[str] = set()
         self._storage_lock = asyncio.Lock()
@@ -321,7 +320,6 @@ async def map_ai_message_chunk_to_events(
 
         events: list[UiPathConversationMessageEvent] = []
 
-        # For every new message_id, start a new message
         if message.id not in self.seen_message_ids:
             self.current_message = AIMessageChunk(content="", id=message.id)
             self.seen_message_ids.add(message.id)
@@ -428,7 +426,7 @@ async def map_current_message_to_start_tool_call_events(self):
 
                         tool_name = tool_call["name"]
                         require_confirmation = (
-                            tool_name in self.tool_names_requiring_confirmation
+                            tool_name in self.tool_confirmation_schemas
                         )
                         input_schema = self.tool_confirmation_schemas.get(tool_name)
                         events.append(

src/uipath_langchain/runtime/runtime.py

@@ -29,7 +29,6 @@
 )
 from uipath.runtime.schema import UiPathRuntimeSchema
 
-from uipath_langchain.chat.hitl import REQUIRE_CONVERSATIONAL_CONFIRMATION
 from uipath_langchain.runtime.errors import LangGraphErrorCode, LangGraphRuntimeError
 from uipath_langchain.runtime.messages import UiPathChatMessagesMapper
 from uipath_langchain.runtime.schema import get_entrypoints_schema, get_graph_schema
@@ -51,6 +50,7 @@ def __init__(
         entrypoint: str | None = None,
         callbacks: list[BaseCallbackHandler] | None = None,
         storage: UiPathRuntimeStorageProtocol | None = None,
+        tool_confirmation_schemas: dict[str, Any] | None = None,
     ):
         """
         Initialize the runtime.
@@ -59,15 +59,14 @@ def __init__(
             graph: The CompiledStateGraph to execute
             runtime_id: Unique identifier for this runtime instance
             entrypoint: Optional entrypoint name (for schema generation)
+            tool_confirmation_schemas: Map of tool name to JSON schema for tool confirmation
         """
         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.runtime_id, storage)
-        confirmation_names, confirmation_schemas = self._get_tool_confirmation_info()
-        self.chat.tool_names_requiring_confirmation = confirmation_names
-        self.chat.tool_confirmation_schemas = confirmation_schemas
+        self.chat.tool_confirmation_schemas = tool_confirmation_schemas or {}
         self._middleware_node_names: set[str] = self._detect_middleware_nodes()
 
     async def execute(
@@ -490,40 +489,6 @@ def _detect_middleware_nodes(self) -> set[str]:
 
         return middleware_nodes
 
-    def _get_tool_confirmation_info(self) -> tuple[set[str], dict[str, Any]]:
-        """Single pass over graph nodes to collect confirmation tool names and schemas."""
-        names: set[str] = set()
-        schemas: dict[str, Any] = {}
-
-        def _record(tool: Any, fallback_name: str) -> None:
-            metadata = getattr(tool, "metadata", None) or {}
-            if not metadata.get(REQUIRE_CONVERSATIONAL_CONFIRMATION):
-                return
-            tool_name = getattr(tool, "name", fallback_name)
-            names.add(tool_name)
-            tool_call_schema = getattr(tool, "tool_call_schema", None)
-            if tool_call_schema is not None:
-                schemas[tool_name] = tool_call_schema.model_json_schema()
-            else:
-                schemas[tool_name] = {}
-
-        for node_name, node_spec in self.graph.nodes.items():
-            bound = getattr(node_spec, "bound", None)
-            if bound is None:
-                continue
-            # Low-code: single-tool UiPathToolNode
-            tool = getattr(bound, "tool", None)
-            if tool is not None:
-                _record(tool, node_name)
-                continue
-            # Coded: multi-tool ToolNode (create_agent)
-            tools_by_name = getattr(bound, "tools_by_name", None)
-            if isinstance(tools_by_name, dict):
-                for fallback_name, tool in tools_by_name.items():
-                    _record(tool, fallback_name)
-
-        return names, schemas
-
     def _is_middleware_node(self, node_name: str) -> bool:
         """Check if a node name represents a middleware node."""
         return node_name in self._middleware_node_names

tests/runtime/test_chat_message_mapper.py

@@ -2019,7 +2019,7 @@ async def test_confirmation_tool_has_requires_confirmation_metadata(self):
         storage = create_mock_storage()
         storage.get_value.return_value = {}
         mapper = UiPathChatMessagesMapper("test-runtime", storage)
-        mapper.tool_names_requiring_confirmation = {"confirm_tool"}
+        mapper.tool_confirmation_schemas = {"confirm_tool": {}}
 
         first_chunk = AIMessageChunk(
             content="",
@@ -2053,7 +2053,7 @@ async def test_normal_tool_has_no_confirmation_metadata(self):
         storage = create_mock_storage()
         storage.get_value.return_value = {}
         mapper = UiPathChatMessagesMapper("test-runtime", storage)
-        mapper.tool_names_requiring_confirmation = {"other_tool"}
+        mapper.tool_confirmation_schemas = {"other_tool": {}}
 
         first_chunk = AIMessageChunk(
             content="",
@@ -2086,7 +2086,7 @@ async def test_mixed_tools_only_confirmation_has_metadata(self):
         storage = create_mock_storage()
         storage.get_value.return_value = {}
         mapper = UiPathChatMessagesMapper("test-runtime", storage)
-        mapper.tool_names_requiring_confirmation = {"confirm_tool"}
+        mapper.tool_confirmation_schemas = {"confirm_tool": {}}
 
         first_chunk = AIMessageChunk(
             content="",