feat: add create react agent

Author: cotovanu-cristian

src/uipath_langchain/agent/react/__init__.py

@@ -0,0 +1,12 @@
+"""UiPath ReAct Agent implementation"""
+
+from .agent import create_agent
+from .state import AgentGraphNode, AgentGraphState
+from .utils import resolve_output_model
+
+__all__ = [
+    "create_agent",
+    "AgentGraphState",
+    "AgentGraphNode",
+    "resolve_output_model",
+]

src/uipath_langchain/agent/react/agent.py

@@ -0,0 +1,76 @@
+import os
+from typing import Sequence, Type
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.tools import BaseTool
+from langgraph.constants import END, START
+from langgraph.graph import StateGraph
+from pydantic import BaseModel
+
+from ..tools import create_tool_node
+from .init_node import (
+    create_init_node,
+)
+from .llm_node import (
+    create_llm_node,
+)
+from .router import (
+    route_agent,
+)
+from .state import AgentGraphNode, AgentGraphState
+from .terminate_node import (
+    create_terminate_node,
+)
+from .tools import create_flow_control_tools
+
+
+def create_agent(
+    model: BaseChatModel,
+    tools: Sequence[BaseTool],
+    messages: Sequence[SystemMessage | HumanMessage],
+    *,
+    state_schema: Type[AgentGraphState] = AgentGraphState,
+    response_format: type[BaseModel] | None = None,
+    recursion_limit: int = 50,
+) -> StateGraph[AgentGraphState]:
+    """Build agent graph with INIT -> AGENT <-> TOOLS loop, terminated by control flow tools.
+
+    Control flow tools (end_execution, raise_error) are auto-injected alongside regular tools.
+    """
+    os.environ["LANGCHAIN_RECURSION_LIMIT"] = str(recursion_limit)
+
+    agent_tools = list(tools)
+    flow_control_tools: list[BaseTool] = create_flow_control_tools(response_format)
+    llm_tools: list[BaseTool] = [*agent_tools, *flow_control_tools]
+
+    init_node = create_init_node(messages)
+    agent_node = create_llm_node(model, llm_tools)
+    tool_nodes = create_tool_node(agent_tools)
+    terminate_node = create_terminate_node(response_format)
+
+    builder: StateGraph[AgentGraphState] = StateGraph(state_schema)
+    builder.add_node(AgentGraphNode.INIT, init_node)
+    builder.add_node(AgentGraphNode.AGENT, agent_node)
+
+    for tool_name, tool_node in tool_nodes.items():
+        builder.add_node(tool_name, tool_node)
+
+    builder.add_node(AgentGraphNode.TERMINATE, terminate_node)
+
+    builder.add_edge(START, AgentGraphNode.INIT)
+    builder.add_edge(AgentGraphNode.INIT, AgentGraphNode.AGENT)
+
+    tool_node_names = list(tool_nodes.keys())
+    builder.add_conditional_edges(
+        AgentGraphNode.AGENT,
+        route_agent,
+        [AgentGraphNode.AGENT, *tool_node_names, AgentGraphNode.TERMINATE],
+    )
+
+    for tool_name in tool_node_names:
+        builder.add_edge(tool_name, AgentGraphNode.AGENT)
+
+    builder.add_edge(AgentGraphNode.TERMINATE, END)
+
+    return builder

src/uipath_langchain/agent/react/constants.py

@@ -0,0 +1,2 @@
+# Agent routing configuration
+MAX_SUCCESSIVE_COMPLETIONS = 1

src/uipath_langchain/agent/react/exceptions.py

@@ -0,0 +1,11 @@
+"""Exceptions for the basic agent loop."""
+
+from uipath._cli._runtime._contracts import UiPathRuntimeError
+
+
+class AgentNodeRoutingException(Exception):
+    pass
+
+
+class AgentTerminationException(UiPathRuntimeError):
+    pass

src/uipath_langchain/agent/react/init_node.py

@@ -0,0 +1,16 @@
+"""State initialization node for the ReAct Agent graph."""
+
+from typing import Sequence
+
+from langchain_core.messages import HumanMessage, SystemMessage
+
+from .state import AgentGraphState
+
+
+def create_init_node(
+    messages: Sequence[SystemMessage | HumanMessage],
+):
+    def graph_state_init(_: AgentGraphState):
+        return {"messages": list(messages)}
+
+    return graph_state_init

src/uipath_langchain/agent/react/llm_node.py

@@ -0,0 +1,44 @@
+"""LLM node implementation for LangGraph."""
+
+from typing import Sequence
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, AnyMessage
+from langchain_core.tools import BaseTool
+
+from .constants import MAX_SUCCESSIVE_COMPLETIONS
+from .state import AgentGraphState
+from .utils import count_successive_completions
+
+
+def create_llm_node(
+    model: BaseChatModel,
+    tools: Sequence[BaseTool] | None = None,
+):
+    """Invoke LLM with tools and dynamically control tool_choice based on successive completions.
+
+    When successive completions reach the limit, tool_choice is set to "required" to force
+    the LLM to use a tool and prevent infinite reasoning loops.
+    """
+    bindable_tools = list(tools) if tools else []
+    base_llm = model.bind_tools(bindable_tools) if bindable_tools else model
+
+    async def llm_node(state: AgentGraphState):
+        messages: list[AnyMessage] = state["messages"]
+
+        successive_completions = count_successive_completions(messages)
+
+        if successive_completions >= MAX_SUCCESSIVE_COMPLETIONS and bindable_tools:
+            llm = base_llm.bind(tool_choice="required")
+        else:
+            llm = base_llm
+
+        response = await llm.ainvoke(messages)
+        if not isinstance(response, AIMessage):
+            raise TypeError(
+                f"LLM returned {type(response).__name__} instead of AIMessage"
+            )
+
+        return {"messages": [response]}
+
+    return llm_node

src/uipath_langchain/agent/react/router.py

@@ -0,0 +1,97 @@
+"""Routing functions for conditional edges in the agent graph."""
+
+from typing import Literal
+
+from langchain_core.messages import AIMessage, AnyMessage, ToolCall
+from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
+
+from .constants import MAX_SUCCESSIVE_COMPLETIONS
+from .exceptions import AgentNodeRoutingException
+from .state import AgentGraphNode, AgentGraphState
+from .utils import count_successive_completions
+
+FLOW_CONTROL_TOOLS = [END_EXECUTION_TOOL.name, RAISE_ERROR_TOOL.name]
+
+
+def __filter_control_flow_tool_calls(
+    tool_calls: list[ToolCall],
+) -> list[ToolCall]:
+    """Remove control flow tools when multiple tool calls exist."""
+    if len(tool_calls) <= 1:
+        return tool_calls
+
+    return [tc for tc in tool_calls if tc.get("name") not in FLOW_CONTROL_TOOLS]
+
+
+def __has_control_flow_tool(tool_calls: list[ToolCall]) -> bool:
+    """Check if any tool call is of a control flow tool."""
+    return any(tc.get("name") in FLOW_CONTROL_TOOLS for tc in tool_calls)
+
+
+def __validate_last_message_is_AI(messages: list[AnyMessage]) -> AIMessage:
+    """Validate and return last message from state.
+
+    Raises:
+        AgentNodeRoutingException: If messages are empty or last message is not AIMessage
+    """
+    if not messages:
+        raise AgentNodeRoutingException(
+            "No messages in state - cannot route after agent"
+        )
+
+    last_message = messages[-1]
+    if not isinstance(last_message, AIMessage):
+        raise AgentNodeRoutingException(
+            f"Last message is not AIMessage (type: {type(last_message).__name__}) - cannot route after agent"
+        )
+
+    return last_message
+
+
+def route_agent(
+    state: AgentGraphState,
+) -> list[str] | Literal[AgentGraphNode.AGENT, AgentGraphNode.TERMINATE]:
+    """Route after agent: handles all routing logic including control flow detection.
+
+    Routing logic:
+    1. If multiple tool calls exist, filter out control flow tools (EndExecution, RaiseError)
+    2. If control flow tool(s) remain, route to TERMINATE
+    3. If regular tool calls remain, route to specific tool nodes (return list of tool names)
+    4. If no tool calls, handle successive completions
+
+    Returns:
+        - list[str]: Tool node names for parallel execution
+        - AgentGraphNode.AGENT: For successive completions
+        - AgentGraphNode.TERMINATE: For control flow termination
+
+    Raises:
+        AgentNodeRoutingException: When encountering unexpected state (empty messages, non-AIMessage, or excessive completions)
+    """
+    messages = state.get("messages", [])
+    last_message = __validate_last_message_is_AI(messages)
+
+    tool_calls = list(last_message.tool_calls) if last_message.tool_calls else []
+    tool_calls = __filter_control_flow_tool_calls(tool_calls)
+
+    if tool_calls and __has_control_flow_tool(tool_calls):
+        return AgentGraphNode.TERMINATE
+
+    if tool_calls:
+        return [tc["name"] for tc in tool_calls]
+
+    successive_completions = count_successive_completions(messages)
+
+    if successive_completions > MAX_SUCCESSIVE_COMPLETIONS:
+        raise AgentNodeRoutingException(
+            f"Agent exceeded successive completions limit without producing tool calls "
+            f"(completions: {successive_completions}, max: {MAX_SUCCESSIVE_COMPLETIONS}). "
+            f"This should not happen as tool_choice='required' is enforced at the limit."
+        )
+
+    if last_message.content:
+        return AgentGraphNode.AGENT
+
+    raise AgentNodeRoutingException(
+        f"Agent produced empty response without tool calls "
+        f"(completions: {successive_completions}, has_content: False)"
+    )

src/uipath_langchain/agent/react/state.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from enum import StrEnum
+
+from langgraph.graph import MessagesState
+
+
+class AgentGraphState(MessagesState):
+    """Agent Graph state for standard loop execution."""
+
+    pass
+
+
+class AgentGraphNode(StrEnum):
+    INIT = "init"
+    AGENT = "agent"
+    TOOLS = "tools"
+    TERMINATE = "terminate"

src/uipath_langchain/agent/react/terminate_node.py

@@ -0,0 +1,50 @@
+"""Termination node for the Agent graph."""
+
+from __future__ import annotations
+
+from langchain_core.messages import AIMessage
+from pydantic import BaseModel
+from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
+
+from .exceptions import (
+    AgentNodeRoutingException,
+    AgentTerminationException,
+)
+from .state import AgentGraphState
+
+
+def create_terminate_node(
+    response_schema: type[BaseModel] | None = None,
+):
+    """Validates and returns end_execution args, or raises AgentTerminationException for raise_error."""
+
+    def terminate_node(state: AgentGraphState):
+        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:
+                args = tool_call["args"]
+                output_schema = response_schema or END_EXECUTION_TOOL.args_schema
+                validated = output_schema.model_validate(args)
+                return validated.model_dump()
+
+            if tool_name == RAISE_ERROR_TOOL.name:
+                error_message = tool_call["args"].get(
+                    "message", "The LLM did not set the error message"
+                )
+                detail = tool_call["args"].get("details", "")
+                raise AgentTerminationException(
+                    code="400", title=error_message, detail=detail
+                )
+
+        raise AgentNodeRoutingException(
+            "No control flow tool call found in terminate node. Unexpected state."
+        )
+
+    return terminate_node

src/uipath_langchain/agent/react/tools/__init__.py

@@ -0,0 +1,7 @@
+from .tools import (
+    create_flow_control_tools,
+)
+
+__all__ = [
+    "create_flow_control_tools",
+]