Merge pull request #359 from UiPath/feat/configurable-successive-optional-llms

feat: add config for agent allowed consecutive thinking messages

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.33"
+version = "0.1.34"
 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/agent.py

@@ -24,7 +24,7 @@
     create_llm_node,
 )
 from .router import (
-    route_agent,
+    create_route_agent,
 )
 from .terminate_node import (
     create_terminate_node,
@@ -58,7 +58,7 @@ def create_agent(
     config: AgentGraphConfig | None = None,
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None = None,
 ) -> StateGraph[AgentGraphState, None, InputT, OutputT]:
-    """Build agent graph with INIT -> AGENT(subgraph) <-> TOOLS loop, terminated by control flow tools.
+    """Build agent graph with INIT -> AGENT (subgraph) <-> TOOLS loop, terminated by control flow tools.
 
     The AGENT node is a subgraph that runs:
     - before-agent guardrail middlewares
@@ -107,14 +107,15 @@ def create_agent(
 
     builder.add_edge(START, AgentGraphNode.INIT)
 
-    llm_node = create_llm_node(model, llm_tools)
+    llm_node = create_llm_node(model, llm_tools, config.thinking_messages_limit)
     llm_with_guardrails_subgraph = create_llm_guardrails_subgraph(
         (AgentGraphNode.LLM, llm_node), guardrails
     )
     builder.add_node(AgentGraphNode.AGENT, llm_with_guardrails_subgraph)
     builder.add_edge(AgentGraphNode.INIT, AgentGraphNode.AGENT)
 
     tool_node_names = list(tool_nodes_with_guardrails.keys())
+    route_agent = create_route_agent(config.thinking_messages_limit)
     builder.add_conditional_edges(
         AgentGraphNode.AGENT,
         route_agent,

src/uipath_langchain/agent/react/constants.py

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

src/uipath_langchain/agent/react/llm_node.py

@@ -1,34 +1,65 @@
-"""LLM node implementation for LangGraph."""
+"""LLM node for ReAct Agent graph."""
 
-from typing import Sequence
+from typing import Literal, 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 .constants import MAX_CONSECUTIVE_THINKING_MESSAGES
 from .types import AgentGraphState
-from .utils import count_successive_completions
+from .utils import count_consecutive_thinking_messages
+
+OPENAI_COMPATIBLE_CHAT_MODELS = (
+    "UiPathChatOpenAI",
+    "AzureChatOpenAI",
+    "ChatOpenAI",
+    "UiPathChat",
+    "UiPathAzureChatOpenAI",
+)
+
+
+def _get_required_tool_choice_by_model(
+    model: BaseChatModel,
+) -> Literal["required", "any"]:
+    """Get the appropriate tool_choice value to enforce tool usage based on model type.
+
+    "required" - OpenAI compatible required tool_choice value
+    "any" - Vertex and Bedrock parameter for required tool_choice value
+    """
+    model_class_name = model.__class__.__name__
+    if model_class_name in OPENAI_COMPATIBLE_CHAT_MODELS:
+        return "required"
+    return "any"
 
 
 def create_llm_node(
     model: BaseChatModel,
     tools: Sequence[BaseTool] | None = None,
+    thinking_messages_limit: int = MAX_CONSECUTIVE_THINKING_MESSAGES,
 ):
-    """Invoke LLM with tools and dynamically control tool_choice based on successive completions.
+    """Create LLM node with dynamic tool_choice enforcement.
 
-    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.
+    Controls when to force tool usage based on consecutive thinking steps
+    to prevent infinite loops and ensure progress.
+
+    Args:
+        model: The chat model to use
+        tools: Available tools to bind
+        thinking_messages_limit: Max consecutive LLM responses without tool calls
+            before enforcing tool usage. 0 = force tools every time.
     """
     bindable_tools = list(tools) if tools else []
     base_llm = model.bind_tools(bindable_tools) if bindable_tools else model
+    tool_choice_required_value = _get_required_tool_choice_by_model(model)
 
     async def llm_node(state: AgentGraphState):
         messages: list[AnyMessage] = state.messages
 
-        successive_completions = count_successive_completions(messages)
-        if successive_completions >= MAX_SUCCESSIVE_COMPLETIONS:
-            llm = base_llm.bind(tool_choice="required")
+        consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
+
+        if bindable_tools and consecutive_thinking_messages >= thinking_messages_limit:
+            llm = base_llm.bind(tool_choice=tool_choice_required_value)
         else:
             llm = base_llm
 

src/uipath_langchain/agent/react/router.py

@@ -6,9 +6,8 @@
 from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
 
 from ..exceptions import AgentNodeRoutingException
-from .constants import MAX_SUCCESSIVE_COMPLETIONS
 from .types import AgentGraphNode, AgentGraphState
-from .utils import count_successive_completions
+from .utils import count_consecutive_thinking_messages
 
 FLOW_CONTROL_TOOLS = [END_EXECUTION_TOOL.name, RAISE_ERROR_TOOL.name]
 
@@ -48,50 +47,62 @@ def __validate_last_message_is_AI(messages: list[AnyMessage]) -> AIMessage:
     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.
+def create_route_agent(thinking_messages_limit: int = 0):
+    """Create a routing function configured with thinking_messages_limit.
 
-    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
+    Args:
+        thinking_messages_limit: Max consecutive thinking messages before error
 
     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)
+        Routing function for LangGraph conditional edges
     """
-    messages = state.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)
+    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 consecutive completions
+
+        Returns:
+            - list[str]: Tool node names for parallel execution
+            - AgentGraphNode.AGENT: For consecutive completions
+            - AgentGraphNode.TERMINATE: For control flow termination
+
+        Raises:
+            AgentNodeRoutingException: When encountering unexpected state (empty messages, non-AIMessage, or excessive completions)
+        """
+        messages = state.messages
+        last_message = __validate_last_message_is_AI(messages)
 
-    if tool_calls and __has_control_flow_tool(tool_calls):
-        return AgentGraphNode.TERMINATE
+        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:
-        return [tc["name"] for tc in tool_calls]
+        if tool_calls and __has_control_flow_tool(tool_calls):
+            return AgentGraphNode.TERMINATE
 
-    successive_completions = count_successive_completions(messages)
+        if tool_calls:
+            return [tc["name"] for tc in tool_calls]
+
+        consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
+
+        if consecutive_thinking_messages > thinking_messages_limit:
+            raise AgentNodeRoutingException(
+                f"Agent exceeded consecutive completions limit without producing tool calls "
+                f"(completions: {consecutive_thinking_messages}, max: {thinking_messages_limit}). "
+                f"This should not happen as tool_choice='required' is enforced at the limit."
+            )
+
+        if last_message.content:
+            return AgentGraphNode.AGENT
 
-    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."
+            f"Agent produced empty response without tool calls "
+            f"(completions: {consecutive_thinking_messages}, has_content: False)"
         )
 
-    if last_message.content:
-        return AgentGraphNode.AGENT
-
-    raise AgentNodeRoutingException(
-        f"Agent produced empty response without tool calls "
-        f"(completions: {successive_completions}, has_content: False)"
-    )
+    return route_agent

src/uipath_langchain/agent/react/types.py

@@ -46,3 +46,8 @@ class AgentGraphConfig(BaseModel):
     recursion_limit: int = Field(
         default=50, ge=1, description="Maximum recursion limit for the agent graph"
     )
+    thinking_messages_limit: int = Field(
+        default=0,
+        ge=0,
+        description="Max consecutive thinking messages before enforcing tool usage. 0 = force tools every time.",
+    )

src/uipath_langchain/agent/react/utils.py

@@ -28,7 +28,7 @@ def resolve_output_model(
     return END_EXECUTION_TOOL.args_schema
 
 
-def count_successive_completions(messages: Sequence[BaseMessage]) -> int:
+def count_consecutive_thinking_messages(messages: Sequence[BaseMessage]) -> int:
     """Count consecutive AIMessages without tool calls at end of message history."""
     if not messages:
         return 0

tests/agent/react/test_utils.py

@@ -2,28 +2,28 @@
 
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
-from uipath_langchain.agent.react.utils import count_successive_completions
+from uipath_langchain.agent.react.utils import count_consecutive_thinking_messages
 
 
 class TestCountSuccessiveCompletions:
     """Test successive completions calculation from message history."""
 
     def test_empty_messages(self):
         """Should return 0 for empty message list."""
-        assert count_successive_completions([]) == 0
+        assert count_consecutive_thinking_messages([]) == 0
 
     def test_no_ai_messages(self):
         """Should return 0 when no AI messages exist."""
         messages = [HumanMessage(content="test")]
-        assert count_successive_completions(messages) == 0
+        assert count_consecutive_thinking_messages(messages) == 0
 
     def test_last_message_not_ai(self):
         """Should return 0 when last message is not AI."""
         messages = [
             AIMessage(content="response"),
             HumanMessage(content="follow-up"),
         ]
-        assert count_successive_completions(messages) == 0
+        assert count_consecutive_thinking_messages(messages) == 0
 
     def test_ai_message_with_tool_calls(self):
         """Should return 0 when last AI message has tool calls."""
@@ -34,23 +34,23 @@ def test_ai_message_with_tool_calls(self):
                 tool_calls=[{"name": "test", "args": {}, "id": "call_1"}],
             ),
         ]
-        assert count_successive_completions(messages) == 0
+        assert count_consecutive_thinking_messages(messages) == 0
 
     def test_ai_message_without_content(self):
         """Should return 0 when last AI message has no content."""
         messages = [
             HumanMessage(content="query"),
             AIMessage(content=""),
         ]
-        assert count_successive_completions(messages) == 0
+        assert count_consecutive_thinking_messages(messages) == 0
 
     def test_single_text_completion(self):
         """Should count single text-only AI message."""
         messages = [
             HumanMessage(content="query"),
             AIMessage(content="thinking"),
         ]
-        assert count_successive_completions(messages) == 1
+        assert count_consecutive_thinking_messages(messages) == 1
 
     def test_two_successive_completions(self):
         """Should count multiple consecutive text-only AI messages."""
@@ -59,7 +59,7 @@ def test_two_successive_completions(self):
             AIMessage(content="thinking 1"),
             AIMessage(content="thinking 2"),
         ]
-        assert count_successive_completions(messages) == 2
+        assert count_consecutive_thinking_messages(messages) == 2
 
     def test_three_successive_completions(self):
         """Should count all consecutive text-only AI messages at end."""
@@ -69,7 +69,7 @@ def test_three_successive_completions(self):
             AIMessage(content="thinking 2"),
             AIMessage(content="thinking 3"),
         ]
-        assert count_successive_completions(messages) == 3
+        assert count_consecutive_thinking_messages(messages) == 3
 
     def test_tool_call_resets_count(self):
         """Should only count completions after last tool call."""
@@ -84,7 +84,7 @@ def test_tool_call_resets_count(self):
             AIMessage(content="thinking 2"),
             AIMessage(content="thinking 3"),
         ]
-        assert count_successive_completions(messages) == 2
+        assert count_consecutive_thinking_messages(messages) == 2
 
     def test_mixed_message_types(self):
         """Should handle complex message patterns correctly."""
@@ -100,7 +100,7 @@ def test_mixed_message_types(self):
             HumanMessage(content="user follow-up"),
             AIMessage(content="responding to follow-up"),
         ]
-        assert count_successive_completions(messages) == 1
+        assert count_consecutive_thinking_messages(messages) == 1
 
     def test_multiple_tool_calls_in_message(self):
         """Should reset count even with multiple tool calls."""
@@ -115,15 +115,15 @@ def test_multiple_tool_calls_in_message(self):
                 ],
             ),
         ]
-        assert count_successive_completions(messages) == 0
+        assert count_consecutive_thinking_messages(messages) == 0
 
     def test_ai_message_with_empty_tool_calls_list(self):
         """Should handle AI message with empty tool_calls list."""
         messages = [
             HumanMessage(content="query"),
             AIMessage(content="thinking", tool_calls=[]),
         ]
-        assert count_successive_completions(messages) == 1
+        assert count_consecutive_thinking_messages(messages) == 1
 
     def test_only_ai_messages_all_text(self):
         """Should count all AI messages when all are text-only."""
@@ -132,4 +132,4 @@ def test_only_ai_messages_all_text(self):
             AIMessage(content="thought 2"),
             AIMessage(content="thought 3"),
         ]
-        assert count_successive_completions(messages) == 3
+        assert count_consecutive_thinking_messages(messages) == 3