refactor: remove agent termination graph specific handling (#423)

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.4.13"
+version = "0.4.14"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/react/terminate_node.py

@@ -13,7 +13,7 @@
     AgentNodeRoutingException,
     AgentTerminationException,
 )
-from .types import AgentGraphState, AgentTermination
+from .types import AgentGraphState
 
 
 def _handle_end_execution(
@@ -36,30 +36,17 @@ def _handle_raise_error(args: dict[str, Any]) -> NoReturn:
     )
 
 
-def _handle_agent_termination(termination: AgentTermination) -> NoReturn:
-    """Handle Command-based termination."""
-    raise AgentTerminationException(
-        code=UiPathErrorCode.EXECUTION_ERROR,
-        title=termination.title,
-        detail=termination.detail,
-    )
-
-
 def create_terminate_node(
     response_schema: type[BaseModel] | None = None, is_conversational: bool = False
 ):
     """Handles Agent Graph termination for multiple sources and output or error propagation to Orchestrator.
 
     Termination scenarios:
-    1. Command based termination with information in state (e.g: escalation)
-    2. LLM-initiated termination (END_EXECUTION_TOOL)
-    3. LLM-initiated error (RAISE_ERROR_TOOL)
+    1. LLM-initiated termination (END_EXECUTION_TOOL)
+    2. LLM-initiated error (RAISE_ERROR_TOOL)
     """
 
     def terminate_node(state: AgentGraphState):
-        if state.inner_state.termination:
-            _handle_agent_termination(state.inner_state.termination)
-
         if not is_conversational:
             last_message = state.messages[-1]
             if not isinstance(last_message, AIMessage):

src/uipath_langchain/agent/react/types.py

@@ -12,21 +12,8 @@
 FLOW_CONTROL_TOOLS = [END_EXECUTION_TOOL.name, RAISE_ERROR_TOOL.name]
 
 
-class AgentTerminationSource(StrEnum):
-    ESCALATION = "escalation"
-
-
-class AgentTermination(BaseModel):
-    """Agent Graph Termination model."""
-
-    source: AgentTerminationSource
-    title: str
-    detail: str = ""
-
-
 class InnerAgentGraphState(BaseModel):
     job_attachments: Annotated[dict[str, Attachment], add_job_attachments] = {}
-    termination: AgentTermination | None = None
 
 
 class InnerAgentGuardrailsGraphState(InnerAgentGraphState):

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -3,10 +3,9 @@
 from enum import Enum
 from typing import Any
 
-from langchain_core.messages import ToolMessage
 from langchain_core.messages.tool import ToolCall
 from langchain_core.tools import BaseTool, StructuredTool
-from langgraph.types import Command, interrupt
+from langgraph.types import interrupt
 from uipath.agent.models.agent import (
     AgentEscalationChannel,
     AgentEscalationRecipient,
@@ -17,10 +16,11 @@
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
 from uipath.platform.common import CreateEscalation
+from uipath.runtime.errors import UiPathErrorCode
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 
-from ..react.types import AgentGraphNode, AgentGraphState, AgentTerminationSource
+from ..exceptions import AgentTerminationException
 from .tool_node import ToolWrapperMixin
 from .utils import sanitize_tool_name
 
@@ -132,8 +132,7 @@ async def escalation_tool_fn(**kwargs: Any) -> dict[str, Any]:
     async def escalation_wrapper(
         tool: BaseTool,
         call: ToolCall,
-        state: AgentGraphState,
-    ) -> dict[str, Any] | Command[Any]:
+    ) -> dict[str, Any] | None:
         result = await tool.ainvoke(call["args"])
 
         if result["action"] == EscalationAction.END:
@@ -143,23 +142,10 @@ async def escalation_wrapper(
                 f"with directive {result['escalation_action']}"
             )
 
-            return Command(
-                update={
-                    "messages": [
-                        ToolMessage(
-                            content=f"{termination_title}. {output_detail}",
-                            tool_call_id=call["id"],
-                        )
-                    ],
-                    "inner_state": {
-                        "termination": {
-                            "source": AgentTerminationSource.ESCALATION,
-                            "title": termination_title,
-                            "detail": output_detail,
-                        }
-                    },
-                },
-                goto=AgentGraphNode.TERMINATE,
+            raise AgentTerminationException(
+                code=UiPathErrorCode.EXECUTION_ERROR,
+                title=termination_title,
+                detail=output_detail,
             )
 
         return result["output"]

src/uipath_langchain/agent/tools/tool_node.py

@@ -19,13 +19,28 @@
 
 # the type safety can be improved with generics
 ToolWrapperReturnType = dict[str, Any] | Command[Any] | None
-ToolWrapperType = Callable[[BaseTool, ToolCall, Any], ToolWrapperReturnType]
-AsyncToolWrapperType = Callable[
+
+ToolWrapperWithoutState = Callable[[BaseTool, ToolCall], ToolWrapperReturnType]
+ToolWrapperWithState = Callable[[BaseTool, ToolCall, Any], ToolWrapperReturnType]
+ToolWrapperType = ToolWrapperWithoutState | ToolWrapperWithState
+
+AsyncToolWrapperWithoutState = Callable[
+    [BaseTool, ToolCall], Awaitable[ToolWrapperReturnType]
+]
+AsyncToolWrapperWithState = Callable[
     [BaseTool, ToolCall, Any], Awaitable[ToolWrapperReturnType]
 ]
+AsyncToolWrapperType = AsyncToolWrapperWithoutState | AsyncToolWrapperWithState
+
 OutputType = dict[Literal["messages"], list[ToolMessage]] | Command[Any] | None
 
 
+def _wrapper_needs_state(wrapper: ToolWrapperType | AsyncToolWrapperType) -> bool:
+    """Check if wrapper function expects a state parameter."""
+    params = list(signature(wrapper).parameters.values())
+    return len(params) >= 3
+
+
 class UiPathToolNode(RunnableCallable):
     """
     A ToolNode that can be used in a React agent graph.
@@ -57,8 +72,8 @@ def _func(self, state: AgentGraphState) -> OutputType:
         if call is None:
             return None
         if self.wrapper:
-            filtered_state = self._filter_state(state, self.wrapper)
-            result = self.wrapper(self.tool, call, filtered_state)
+            inputs = self._prepare_wrapper_inputs(self.wrapper, self.tool, call, state)
+            result = self.wrapper(*inputs)
         else:
             result = self.tool.invoke(call["args"])
         return self._process_result(call, result)
@@ -68,8 +83,8 @@ async def _afunc(self, state: AgentGraphState) -> OutputType:
         if call is None:
             return None
         if self.awrapper:
-            filtered_state = self._filter_state(state, self.awrapper)
-            result = await self.awrapper(self.tool, call, filtered_state)
+            inputs = self._prepare_wrapper_inputs(self.awrapper, self.tool, call, state)
+            result = await self.awrapper(*inputs)
         else:
             result = await self.tool.ainvoke(call["args"])
         return self._process_result(call, result)
@@ -106,6 +121,19 @@ def _process_result(
             )
             return {"messages": [message]}
 
+    def _prepare_wrapper_inputs(
+        self,
+        wrapper: ToolWrapperType | AsyncToolWrapperType,
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> Sequence[Any]:
+        """Prepare inputs for wrapper invocation based on its signature."""
+        if _wrapper_needs_state(wrapper):
+            filtered_state = self._filter_state(state, wrapper)
+            return tool, call, filtered_state
+        return tool, call
+
     def _filter_state(
         self, state: Any, wrapper: ToolWrapperType | AsyncToolWrapperType
     ) -> BaseModel:
@@ -137,7 +165,6 @@ def create_tool_node(tools: Sequence[BaseTool]) -> dict[str, UiPathToolNode]:
 
     Args:
         tools: Sequence of tools to create nodes for.
-        agentState: The type of the agent state model.
 
     Returns:
         Dict mapping tool.name -> ReactToolNode([tool]).

src/uipath_langchain/agent/wrappers/job_attachment_wrapper.py

@@ -11,10 +11,12 @@
     replace_job_attachment_ids,
 )
 from uipath_langchain.agent.react.types import AgentGraphState
-from uipath_langchain.agent.tools.tool_node import AsyncToolWrapperType
+from uipath_langchain.agent.tools.tool_node import AsyncToolWrapperWithState
 
 
-def get_job_attachment_wrapper(output_type: Any | None = None) -> AsyncToolWrapperType:
+def get_job_attachment_wrapper(
+    output_type: Any | None = None,
+) -> AsyncToolWrapperWithState:
     """Create a tool wrapper that handles job attachments in both tool inputs and outputs.
 
     This wrapper performs two main functions:

tests/agent/react/test_terminate_node.py

@@ -12,16 +12,11 @@
     AgentTerminationException,
 )
 from uipath_langchain.agent.react.terminate_node import create_terminate_node
-from uipath_langchain.agent.react.types import (
-    AgentTermination,
-    AgentTerminationSource,
-)
 
 
 class MockInnerState(BaseModel):
     """Mock inner state for testing."""
 
-    termination: AgentTermination | None = None
     job_attachments: dict[str, Any] = {}
 
 
@@ -52,19 +47,6 @@ def state_with_human_message(self):
         """Fixture for state with human message as last."""
         return MockAgentGraphState(messages=[HumanMessage(content="User message")])
 
-    @pytest.fixture
-    def state_with_termination(self):
-        """Fixture for state with termination set (e.g., escalation)."""
-        termination = AgentTermination(
-            source=AgentTerminationSource.ESCALATION,
-            title="Escalation required",
-            detail="User requested human assistance",
-        )
-        return MockAgentGraphState(
-            messages=[AIMessage(content="response")],
-            inner_state=MockInnerState(termination=termination),
-        )
-
     def test_conversational_returns_none_no_tool_calls(
         self, terminate_node, state_with_ai_message
     ):
@@ -82,15 +64,6 @@ def test_conversational_skips_ai_message_validation(
 
         assert result is None
 
-    def test_conversational_handles_termination(
-        self, terminate_node, state_with_termination
-    ):
-        """Conversational mode should still handle state.inner_state.termination."""
-        with pytest.raises(AgentTerminationException) as exc_info:
-            terminate_node(state_with_termination)
-
-        assert "Escalation required" in exc_info.value.error_info.title
-
     def test_conversational_ignores_end_execution_tool(self):
         """Conversational mode should ignore END_EXECUTION tool calls."""
         terminate_node = create_terminate_node(
@@ -169,19 +142,6 @@ def state_with_no_control_flow_tool(self):
         )
         return MockAgentGraphState(messages=[ai_message])
 
-    @pytest.fixture
-    def state_with_termination(self):
-        """Fixture for state with escalation termination."""
-        termination = AgentTermination(
-            source=AgentTerminationSource.ESCALATION,
-            title="Needs human review",
-            detail="Complex issue",
-        )
-        return MockAgentGraphState(
-            messages=[AIMessage(content="response")],
-            inner_state=MockInnerState(termination=termination),
-        )
-
     def test_non_conversational_handles_end_execution(
         self, terminate_node, state_with_end_execution
     ):
@@ -223,15 +183,6 @@ def test_non_conversational_raises_on_no_control_flow_tool(
         ):
             terminate_node(state_with_no_control_flow_tool)
 
-    def test_non_conversational_handles_termination_first(
-        self, terminate_node, state_with_termination
-    ):
-        """Non-conversational mode should check termination before tool calls."""
-        with pytest.raises(AgentTerminationException) as exc_info:
-            terminate_node(state_with_termination)
-
-        assert "Needs human review" in exc_info.value.error_info.title
-
 
 class TestTerminateNodeWithResponseSchema:
     """Test cases for terminate node with custom response schema."""