diff --git a/src/uipath_langchain/agent/react/forced_extraction.py b/src/uipath_langchain/agent/react/forced_extraction.py index ceadbb483..693c22438 100644 --- a/src/uipath_langchain/agent/react/forced_extraction.py +++ b/src/uipath_langchain/agent/react/forced_extraction.py @@ -39,6 +39,13 @@ def _ensure_trailing_user_turn(messages: list[AnyMessage]) -> list[AnyMessage]: ] +def build_nudge_messages(messages: list[AnyMessage]) -> list[AnyMessage]: + """Messages for re-asking a model that can't be forced: thinking and reasoning blocks + stay (its thinking can't be turned off), and the request ends on a user turn asking + for a tool call.""" + return _ensure_trailing_user_turn(messages) + + def build_extraction_call( model: BaseChatModel, messages: list[AnyMessage] ) -> tuple[BaseChatModel, list[AnyMessage]]: diff --git a/src/uipath_langchain/agent/react/llm_node.py b/src/uipath_langchain/agent/react/llm_node.py index e0184b542..44fb09380 100644 --- a/src/uipath_langchain/agent/react/llm_node.py +++ b/src/uipath_langchain/agent/react/llm_node.py @@ -16,7 +16,10 @@ from uipath.runtime.errors import UiPathErrorCategory from uipath_langchain.chat.handlers import get_payload_handler -from uipath_langchain.chat.thinking import thinking_rejects_forced_tool_choice +from uipath_langchain.chat.thinking import ( + model_rejects_forced_tool_choice, + thinking_rejects_forced_tool_choice, +) from ..exceptions import ( AgentRuntimeError, @@ -30,7 +33,7 @@ from ..messages.message_utils import replace_tool_calls from ..tools.static_args import StaticArgsHandler from .constants import DEFAULT_MAX_LLM_MESSAGES -from .forced_extraction import build_extraction_call +from .forced_extraction import build_extraction_call, build_nudge_messages from .types import FLOW_CONTROL_TOOLS, AgentGraphState from .utils import count_consecutive_tool_less_turns @@ -108,6 +111,7 @@ async def llm_node(state: StateT): current_tool_choice: Literal["auto", "any"] = tool_choice consecutive_tool_less = count_consecutive_tool_less_turns(messages) thinking_rejects_forcing = thinking_rejects_forced_tool_choice(model) + model_rejects_forcing = model_rejects_forced_tool_choice(model) call_model: BaseChatModel = model call_messages: list[AnyMessage] = messages handler = payload_handler @@ -122,10 +126,13 @@ async def llm_node(state: StateT): "configuration, verify your model deployment respects tool_choice.", category=UiPathErrorCategory.SYSTEM, ) - current_tool_choice = "any" - if thinking_rejects_forcing and consecutive_tool_less > 0: - call_model, call_messages = build_extraction_call(model, messages) - handler = get_payload_handler(call_model) + current_tool_choice = "auto" if model_rejects_forcing else "any" + if consecutive_tool_less > 0: + if model_rejects_forcing: + call_messages = build_nudge_messages(messages) + elif thinking_rejects_forcing: + call_model, call_messages = build_extraction_call(model, messages) + handler = get_payload_handler(call_model) binding_kwargs = handler.get_tool_binding_kwargs( tools=static_schema_tools, diff --git a/src/uipath_langchain/chat/thinking.py b/src/uipath_langchain/chat/thinking.py index 978f34935..a001e9548 100644 --- a/src/uipath_langchain/chat/thinking.py +++ b/src/uipath_langchain/chat/thinking.py @@ -43,6 +43,17 @@ def thinking_rejects_forced_tool_choice(model: Any) -> bool: return False +def model_rejects_forced_tool_choice(model: Any) -> bool: + """True if the gateway flags the model as rejecting any forced tool_choice. + + Discovery's ``modelDetails.shouldSkipForcedToolChoice`` (e.g. Claude Opus 5.5, which + 400s on tool_choice ``any`` / ``tool`` and whose thinking can't be turned off, so the + thinking-off extraction retry isn't an option either). + """ + details = getattr(model, "model_details", None) + return isinstance(details, dict) and bool(details.get("shouldSkipForcedToolChoice")) + + def strip_thinking(model: BaseChatModel) -> BaseChatModel: """Copy of the model with thinking config stripped, so forcing is honored. diff --git a/tests/agent/react/test_llm_node.py b/tests/agent/react/test_llm_node.py index 68279d300..4c2c8d7cd 100644 --- a/tests/agent/react/test_llm_node.py +++ b/tests/agent/react/test_llm_node.py @@ -682,3 +682,88 @@ async def test_plain_model_forces_from_first_turn(self) -> None: node = create_llm_node(model, [tool]) await node(AgentGraphState(messages=[HumanMessage(content="q")])) assert model.bind_tools.call_args.kwargs["tool_choice"] == "any" + + +class TestModelRejectsForcedToolChoice: + """Models flagged `shouldSkipForcedToolChoice` (e.g. Claude Opus 5.5) 400 on a forced + tool_choice and can't turn thinking off, so the node sends `auto` and re-asks a + stalled turn with thinking kept instead of running the thinking-off extraction.""" + + def _flagged_model(self) -> Any: + model: Any = _StubAzureChatOpenAI.model_construct() + model.model_details = {"shouldSkipForcedToolChoice": True} + model.thinking = {"type": "adaptive"} + model.bind_tools = Mock(return_value=model) + model.bind = Mock(return_value=model) + return model + + def _tool(self) -> Any: + tool = Mock(spec=BaseTool) + tool.name = "t" + return tool + + def _end_call(self) -> AIMessage: + return AIMessage( + content="", + tool_calls=[ + create_tool_call(name=END_EXECUTION_TOOL.name, args={}, id="c1") + ], + ) + + def _stalled_state(self, stalls: int = 1) -> AgentGraphState: + prior = AIMessage( + content=[ + {"type": "thinking", "thinking": "think", "signature": "s"}, + {"type": "text", "text": "answer"}, + ] + ) + return AgentGraphState( + messages=[HumanMessage(content="q"), *([prior] * stalls)] + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("configured", ["auto", "any"]) + async def test_first_turn_sends_auto( + self, configured: Literal["auto", "any"] + ) -> None: + model = self._flagged_model() + model.ainvoke = AsyncMock(return_value=self._end_call()) + + node = create_llm_node(model, [self._tool()], tool_choice=configured) + await node(AgentGraphState(messages=[HumanMessage(content="q")])) + + assert model.bind_tools.call_args.kwargs["tool_choice"] == "auto" + + @pytest.mark.asyncio + async def test_stall_is_nudged_with_thinking_kept(self) -> None: + model = self._flagged_model() + captured: dict[str, Any] = {} + + async def fake_ainvoke(msgs: Any) -> AIMessage: + captured["msgs"] = msgs + return self._end_call() + + model.ainvoke = AsyncMock(side_effect=fake_ainvoke) + with patch( + "uipath_langchain.agent.react.llm_node.build_extraction_call" + ) as spy: + await create_llm_node(model, [self._tool()])(self._stalled_state()) + + spy.assert_not_called() + assert model.bind_tools.call_args.kwargs["tool_choice"] == "auto" + msgs = captured["msgs"] + assert msgs[-2].content[0]["type"] == "thinking" + assert isinstance(msgs[-1], HumanMessage) + + @pytest.mark.asyncio + async def test_second_stall_raises_thinking_limit(self) -> None: + model = self._flagged_model() + model.ainvoke = AsyncMock(return_value=AIMessage(content="still stalling")) + + with pytest.raises(AgentRuntimeError) as exc_info: + await create_llm_node(model, [self._tool()])(self._stalled_state(stalls=2)) + + assert exc_info.value.error_info.code.endswith( + AgentRuntimeErrorCode.THINKING_LIMIT_EXCEEDED.value + ) + model.ainvoke.assert_not_awaited() diff --git a/tests/chat/test_thinking.py b/tests/chat/test_thinking.py new file mode 100644 index 000000000..1ffcedd08 --- /dev/null +++ b/tests/chat/test_thinking.py @@ -0,0 +1,28 @@ +from types import SimpleNamespace +from typing import Any + +import pytest + +from uipath_langchain.chat.thinking import model_rejects_forced_tool_choice + + +@pytest.mark.parametrize( + ("details", "expected"), + [ + ({"shouldSkipForcedToolChoice": True}, True), + ({"shouldSkipForcedToolChoice": False}, False), + ({}, False), + (None, False), + ], +) +def test_model_rejects_forced_tool_choice_reads_discovery_flag( + details: Any, expected: bool +) -> None: + assert ( + model_rejects_forced_tool_choice(SimpleNamespace(model_details=details)) + is expected + ) + + +def test_model_without_model_details_does_not_reject() -> None: + assert model_rejects_forced_tool_choice(object()) is False