fix messages when resuming

Author: mike-deem-uipath

src/uipath_langchain/agent/react/init_node.py

@@ -25,9 +25,16 @@ def graph_state_init(state: Any) -> Any:
             resolved_messages = list(messages)
         if is_conversational:
             # For conversational agents we need to reorder the messages so that the system message is first, followed by
-            # the initial user message. The initial user message is put in the state by UiPathLangGraphRuntime. The add
-            # reducer is used for the messages property in the state, so by default new messages are appended to the end.
-            resolved_messages = Overwrite([*resolved_messages, *state.messages])
+            # the initial user message. When resuming the conversation, the state will have the entire message history,
+            # including the system message. In this case, we need to replace the system message from the state with the
+            # newly generated one. It will have the current date/time and reflect any changes to user settings. The add
+            # reducer is used for the messages property in the state, so by default new messages are appended to the end
+            # and using Overwrite will cause LangGraph to replace the entire array instead.
+            if len(state.messages) > 0 and isinstance(state.messages[0], SystemMessage):
+                preserved_messages = state.messages[1:]
+            else:
+                preserved_messages = state.messages
+            resolved_messages = Overwrite([*resolved_messages, *preserved_messages])
 
         schema = input_schema if input_schema is not None else BaseModel
         job_attachments = get_job_attachments(schema, state)

src/uipath_langchain/agent/react/router_conversational.py

@@ -23,7 +23,7 @@ def create_route_agent_conversational():
 
     def route_agent_conversational(
         state: AgentGraphState,
-    ) -> list[str] | Literal[AgentGraphNode.USER_MESSAGE_WAIT]:
+    ) -> list[str] | Literal[AgentGraphNode.TERMINATE]:
         """Route after agent
 
         Routing logic:
@@ -41,6 +41,6 @@ def route_agent_conversational(
         if last_message.tool_calls:
             return [tc["name"] for tc in last_message.tool_calls]
         else:
-            return AgentGraphNode.USER_MESSAGE_WAIT
+            return AgentGraphNode.TERMINATE
 
     return route_agent_conversational

src/uipath_langchain/agent/react/terminate_node.py

@@ -0,0 +1,83 @@
+"""Termination node for the Agent graph."""
+
+from __future__ import annotations
+
+from typing import Any, NoReturn
+
+from langchain_core.messages import AIMessage
+from pydantic import BaseModel
+from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
+from uipath.runtime.errors import UiPathErrorCode
+
+from ..exceptions import (
+    AgentNodeRoutingException,
+    AgentTerminationException,
+)
+from .types import AgentGraphState, AgentTermination
+
+
+def _handle_end_execution(
+    args: dict[str, Any], response_schema: type[BaseModel] | None
+) -> dict[str, Any]:
+    """Handle LLM-initiated termination via END_EXECUTION_TOOL."""
+    output_schema = response_schema or END_EXECUTION_TOOL.args_schema
+    validated = output_schema.model_validate(args)
+    return validated.model_dump()
+
+
+def _handle_raise_error(args: dict[str, Any]) -> NoReturn:
+    """Handle LLM-initiated error via RAISE_ERROR_TOOL."""
+    error_message = args.get("message", "The LLM did not set the error message")
+    detail = args.get("details", "")
+    raise AgentTerminationException(
+        code=UiPathErrorCode.EXECUTION_ERROR,
+        title=error_message,
+        detail=detail,
+    )
+
+
+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)
+    """
+
+    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):
+                raise AgentNodeRoutingException(
+                    f"Expected last message to be AIMessage, got {type(last_message).__name__}"
+                )
+
+            for tool_call in last_message.tool_calls:
+                tool_name = tool_call["name"]
+
+                if tool_name == END_EXECUTION_TOOL.name:
+                    return _handle_end_execution(tool_call["args"], response_schema)
+
+                if tool_name == RAISE_ERROR_TOOL.name:
+                    _handle_raise_error(tool_call["args"])
+
+            raise AgentNodeRoutingException(
+                "No control flow tool call found in terminate node. Unexpected state."
+            )
+
+    return terminate_node