diff --git a/src/aiq_agent/agents/clarifier/agent.py b/src/aiq_agent/agents/clarifier/agent.py index ebcaf1170..11c70a8ce 100644 --- a/src/aiq_agent/agents/clarifier/agent.py +++ b/src/aiq_agent/agents/clarifier/agent.py @@ -105,6 +105,19 @@ ) """Reminder prompt added after tool results to reinforce JSON-only output.""" +FORCE_SEARCH_GUIDANCE = ( + "You attempted to ask the user for clarification before gathering any context. " + "Before asking the user a question, you MUST first use the available search tools " + "to look up unfamiliar entities, acronyms, products, or terms in their request. " + "Issue one focused tool call now with a query derived from the user's request. " + "Only after reviewing the tool results should you decide whether clarification is still needed." +) +"""Guidance prompt injected when the LLM tries to clarify without having searched first.""" + +SKIPPED_CLARIFICATION_SENTINEL = "[skipped clarification]" +"""Stand-in user turn persisted when a skip reply is blank, so an empty +HumanMessage is never written to state (some chat APIs reject empty content).""" + class ClarifierAgent: """ @@ -485,6 +498,45 @@ def _get_fallback_clarification(self, query: str | None = None) -> str: SKIP_COMMANDS = {"skip", "done", "exit", "quit", "proceed", "continue", "no", "n", ""} """Set of commands that indicate the user wants to skip clarification.""" + @staticmethod + def _has_tool_invocations(messages: Sequence[Any]) -> bool: + """ + Check whether any prior assistant message in the conversation issued tool calls. + + Args: + messages: The conversation message history. + + Returns: + True if any AIMessage in the history carries non-empty tool_calls, + False otherwise. + """ + for msg in messages: + tool_calls = getattr(msg, "tool_calls", None) + if tool_calls: + return True + return False + + @staticmethod + def _searched_since_last_user_turn(messages: Sequence[Any]) -> bool: + """Check whether a tool call has occurred since the latest user message. + + Scopes the search-before-clarify guard to the *current* request: tool + calls from earlier conversation turns must not suppress the nudge for a + fresh user query. + + Args: + messages: The conversation message history. + + Returns: + True if any message after the most recent HumanMessage carries tool + calls, False otherwise. + """ + last_user_idx = -1 + for idx, msg in enumerate(messages): + if isinstance(msg, HumanMessage): + last_user_idx = idx + return ClarifierAgent._has_tool_invocations(messages[last_user_idx + 1 :]) + def _is_skip_command(self, user_reply: str) -> bool: """ Check if the user's reply indicates they want to skip clarification. @@ -504,16 +556,21 @@ def _build_graph(self) -> CompiledStateGraph: """ Build the LangGraph StateGraph for the clarification workflow. - Creates a graph with three nodes: - - agent: Generates clarification questions using the LLM + Creates a graph with the following nodes: + - agent: Generates clarification questions using the LLM. On the first + turn it also enforces search-before-clarify (issue #234): if the model + asks for clarification without using its bound search tools, it nudges + the model once and retries inline. - tools: Executes tool calls (e.g., web search) for context - ask_for_clarification: Prompts user and processes response + - plan_preview: Optional plan approval flow The graph flow: - 1. agent generates a response (question, tool call, or completion) + 1. agent generates a response (question, tool call, or completion); + on turn 0 it may force one search-and-retry before yielding 2. If tool call → tools node → back to agent - 3. If question → ask_for_clarification → back to agent - 4. If complete → end + 3. If complete → end (or plan_preview if enabled) + 4. Otherwise → ask_for_clarification → back to agent Returns: Compiled LangGraph StateGraph ready for execution. @@ -526,9 +583,38 @@ def _build_graph(self) -> CompiledStateGraph: graph = StateGraph(ClarifierAgentState) async def agent_node(state: ClarifierAgentState): + """Run the LLM for one turn and apply the search-before-clarify nudge. + + Emits a completion when the clarification budget is exhausted, and on + the first turn forces one search-and-retry if the model tries to ask + for clarification without using its bound tools. + """ if state.remaining_questions <= 0: - complete_response = ClarificationResponse(needs_clarification=False, clarification_question=None) - return {"messages": [AIMessage(content=complete_response.model_dump_json())]} + # Clarification budget is exhausted — emit a completion signal, + # but never create an invalid history (two adjacent assistant + # messages, or a pending tool call left unresolved), since it can + # reach plan_preview's planner. + last_message = state.messages[-1] if state.messages else None + if isinstance(last_message, AIMessage) and getattr(last_message, "tool_calls", None): + # A pending tool call must be resolved before we complete; + # let decide_route route to the tools node, after which the + # tool result re-enters here as the last turn. + return {} + if isinstance(last_message, AIMessage) and self._is_complete(getattr(last_message, "content", "")): + # A prior node (e.g. the skip-command branch) already emitted + # the completion; don't duplicate it. Let decide_route end. + return {} + complete = AIMessage( + content=ClarificationResponse( + needs_clarification=False, clarification_question=None + ).model_dump_json() + ) + if isinstance(last_message, AIMessage): + # The last turn is a non-complete assistant message (e.g. an + # unanswered clarification at exhaustion). Interleave a + # sentinel user turn so the completion is not adjacent to it. + return {"messages": [HumanMessage(content=SKIPPED_CLARIFICATION_SENTINEL), complete]} + return {"messages": [complete]} tools_info = [ {"name": getattr(t, "name", ""), "description": getattr(t, "description", "")} for t in self.tools ] @@ -554,9 +640,37 @@ async def agent_node(state: ClarifierAgentState): messages.append(HumanMessage(content=JSON_REMINDER_AFTER_TOOLS)) response = await bound_llm.ainvoke(messages) + + # Search-before-clarify (issue #234): on the first turn, if the model + # asks for clarification without searching, nudge it once (guidance as + # ephemeral scaffolding, never persisted to state) and retry inline. + # The guard is one-shot: iteration == 0 and no tool call yet for the + # current request (scoped to the latest user turn). Return only the + # retry so the skipped first response can't leave two adjacent + # assistant messages in history. The guidance is folded into the + # leading system prompt (a trailing SystemMessage is rejected by + # providers that only accept a leading one). + if ( + self.tools + and state.iteration == 0 + and not self._searched_since_last_user_turn(state.messages) + and not getattr(response, "tool_calls", None) + and self._is_needed(response.content) + ): + logger.info("Clarifier: model skipped search before clarifying; injecting guidance and retrying once") + retry_system = SystemMessage(content=f"{rendered_system_prompt}\n\n{FORCE_SEARCH_GUIDANCE}") + retry_messages = [retry_system, *messages[1:], response] + retry_response = await bound_llm.ainvoke(retry_messages) + return {"messages": [retry_response]} + return {"messages": [response]} async def ask_clarification(state: ClarifierAgentState): + """Prompt the user with the pending question and record their reply. + + Handles skip commands (substituting a non-empty sentinel for blank + replies) and the max-turns cutoff, advancing the clarification log. + """ iteration = state.iteration max_turns = state.max_turns clarifier_log = state.clarifier_log @@ -582,8 +696,23 @@ async def ask_clarification(state: ClarifierAgentState): logger.info("Clarifier: User requested to skip clarification") complete_response = ClarificationResponse(needs_clarification=False, clarification_question=None) clarifier_log = f"{clarifier_log}\n**Turn {iteration + 1} - User:** [Skipped clarification]" + # Persist the user's reply as a HumanMessage before the + # completion AIMessage. The prior turn already left an + # AIMessage(clarification) in history; without an interleaving + # human message the two assistant turns would be adjacent, which + # the OpenAI/Anthropic APIs reject. (The duplicate completion on + # graph re-entry is suppressed by the guard in agent_node.) + # + # A blank/whitespace reply also counts as skip (see + # SKIP_COMMANDS), but an empty HumanMessage must not be persisted + # -- it would flow into plan generation, and some chat APIs reject + # empty message content. Substitute a non-empty sentinel. + skip_reply = user_reply if user_reply.strip() else SKIPPED_CLARIFICATION_SENTINEL return { - "messages": [AIMessage(content=complete_response.model_dump_json())], + "messages": [ + HumanMessage(content=skip_reply), + AIMessage(content=complete_response.model_dump_json()), + ], "iteration": max_turns, # Force end of clarification "clarifier_log": clarifier_log, } @@ -596,6 +725,7 @@ async def ask_clarification(state: ClarifierAgentState): } def decide_route(state: ClarifierAgentState | dict): + """Route after agent_node: to tools, plan preview, end, or the user.""" if isinstance(state, dict): messages = state.get("messages", []) elif hasattr(state, "messages"): @@ -616,6 +746,11 @@ def decide_route(state: ClarifierAgentState | dict): if self.enable_plan_approval: return "plan_preview" return "__end__" + + # The search-before-clarify nudge (issue #234) is handled inline in + # agent_node, not here — see the retry block there. By the time a + # clarification response reaches this router, any forced search has + # already happened, so we route straight to the user. return "ask_for_clarification" async def plan_preview_node(state: ClarifierAgentState): diff --git a/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 b/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 index 8cbcccd8d..ae719a242 100644 --- a/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 +++ b/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 @@ -29,9 +29,11 @@ Your ONLY responsibility is to determine whether a research request requires cla ## Tool Usage -- You may use search tools ONLY to understand unfamiliar domains -- Use at most 1-2 high-value searches -- Searches are for your internal understanding only +- **Search first, ask second.** If the user's request contains any unfamiliar entity, acronym, project, person, product, or technical term that you cannot fully define from your training data, you MUST issue a search tool call before deciding clarification is needed. Do not ask the user to define terms that a quick search would resolve. +- On the first turn, prefer at least one search to ground the topic in current context whenever search tools are available. +- Use at most 1-2 high-value searches per turn — keep queries focused on the specific unknown. +- Searches are for your internal understanding only. Do not summarize or report search results to the user. +- After reviewing search results, re-evaluate: if the request is now sufficiently specified, return `needs_clarification: false`. Only ask a clarification question if a genuine ambiguity remains. --- diff --git a/tests/aiq_agent/agents/clarifier/test_agent.py b/tests/aiq_agent/agents/clarifier/test_agent.py index 38f53b018..d6ca78f8d 100644 --- a/tests/aiq_agent/agents/clarifier/test_agent.py +++ b/tests/aiq_agent/agents/clarifier/test_agent.py @@ -22,10 +22,14 @@ import pytest from langchain_core.messages import AIMessage +from langchain_core.messages import BaseMessage from langchain_core.messages import HumanMessage +from langchain_core.messages import SystemMessage +from langchain_core.messages import ToolMessage from langchain_core.tools import tool from aiq_agent.agents.clarifier.agent import DEFAULT_CLARIFICATION_PROMPT +from aiq_agent.agents.clarifier.agent import FORCE_SEARCH_GUIDANCE from aiq_agent.agents.clarifier.agent import ClarifierAgent from aiq_agent.agents.clarifier.models import ClarificationResponse from aiq_agent.agents.clarifier.models import ClarifierAgentState @@ -40,6 +44,15 @@ def web_search_tool(query: str) -> str: return f"Results for: {query}" +def _adjacent_assistant_pairs(messages: list[BaseMessage]) -> list[tuple[int, int]]: + """Return index pairs of any two consecutive AIMessages (the invalid shape).""" + return [ + (i, i + 1) + for i in range(len(messages) - 1) + if isinstance(messages[i], AIMessage) and isinstance(messages[i + 1], AIMessage) + ] + + class TestClarifierAgentInit: """Tests for ClarifierAgent initialization.""" @@ -872,3 +885,677 @@ def test_init_with_planner_llm(self, mock_llm_provider): ) assert agent.planner_llm == planner_llm + + +class TestHasToolInvocations: + """Tests for the _has_tool_invocations helper.""" + + def test_empty_messages(self): + """No messages -> no invocations.""" + assert ClarifierAgent._has_tool_invocations([]) is False + + def test_only_human_messages(self): + """Human-only history has no invocations.""" + messages = [HumanMessage(content="hi"), HumanMessage(content="more")] + assert ClarifierAgent._has_tool_invocations(messages) is False + + def test_ai_message_without_tool_calls(self): + """An AIMessage without tool_calls counts as no invocation.""" + messages = [HumanMessage(content="hi"), AIMessage(content="hello")] + assert ClarifierAgent._has_tool_invocations(messages) is False + + def test_ai_message_with_tool_calls(self): + """An AIMessage with tool_calls counts as an invocation.""" + ai = AIMessage( + content="", + tool_calls=[{"name": "web_search_tool", "args": {"query": "x"}, "id": "call_1"}], + ) + assert ClarifierAgent._has_tool_invocations([HumanMessage(content="hi"), ai]) is True + + def test_with_tool_message(self): + """A ToolMessage by itself does not count - we look at the assistant turn.""" + msg = ToolMessage(content="result", tool_call_id="call_1") + assert ClarifierAgent._has_tool_invocations([msg]) is False + + +class TestSearchedSinceLastUserTurn: + """Tests for _searched_since_last_user_turn (current-request scoping).""" + + def test_no_messages(self): + """An empty history has no search this turn.""" + assert ClarifierAgent._searched_since_last_user_turn([]) is False + + def test_only_current_query_no_search(self): + """Just the user query, no tool calls yet -> not searched.""" + msgs = [HumanMessage(content="Research AI")] + assert ClarifierAgent._searched_since_last_user_turn(msgs) is False + + def test_tool_call_after_query_counts(self): + """A tool call after the latest user turn counts as searched.""" + ai = AIMessage(content="", tool_calls=[{"name": "web_search_tool", "args": {}, "id": "c1"}]) + msgs = [HumanMessage(content="Research AI"), ai] + assert ClarifierAgent._searched_since_last_user_turn(msgs) is True + + def test_tool_call_before_latest_user_turn_is_ignored(self): + """A tool call from an earlier turn must not suppress the nudge for a new query.""" + prior_tool = AIMessage(content="", tool_calls=[{"name": "web_search_tool", "args": {}, "id": "c0"}]) + msgs = [ + HumanMessage(content="earlier question"), + prior_tool, + ToolMessage(content="result", tool_call_id="c0"), + HumanMessage(content="a fresh research query"), # latest user turn, no search after it + ] + assert ClarifierAgent._searched_since_last_user_turn(msgs) is False + + +class TestClarifierForceSearch: + """Tests for the search-before-clarify behavior (issue #234).""" + + @pytest.fixture + def mock_llm(self): + """Create a mock LLM.""" + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + return llm + + @pytest.fixture + def mock_llm_provider(self, mock_llm): + """Create a mock LLM provider.""" + provider = MagicMock(spec=LLMProvider) + provider.get = MagicMock(return_value=mock_llm) + return provider + + @pytest.mark.asyncio + async def test_force_search_fires_when_llm_skips_tools(self, mock_llm_provider, mock_llm): + """When tools are configured and the LLM tries to clarify without searching, + the agent must first nudge the LLM with a force-search guidance message, + then route to tools when the LLM complies.""" + + # 1st LLM call: skip tools, ask for clarification. + # 2nd LLM call (after force_search): produce a tool call. + # 3rd LLM call (after tool result): return complete. + clarif_msg = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + tool_call_msg = AIMessage( + content="", + tool_calls=[{"name": "web_search_tool", "args": {"query": "AI"}, "id": "call_1"}], + ) + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg, tool_call_msg, complete_msg]) + + user_callback = AsyncMock() + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research Foo Project XYZ")]) + result = await agent.run(state) + + assert result is not None + assert isinstance(result, ClarifierResult) + # The LLM was invoked three times: clarify-attempt, tool-call, finalize. + assert mock_llm.ainvoke.call_count == 3 + # The user was never prompted because the search-then-complete path was taken. + user_callback.assert_not_called() + # The 2nd LLM call (after force_search guidance) must carry the guidance + # folded into a single LEADING SystemMessage — providers that only accept + # a leading system message reject a trailing one. + second_call_messages = mock_llm.ainvoke.call_args_list[1].args[0] + system_messages = [m for m in second_call_messages if isinstance(m, SystemMessage)] + assert len(system_messages) == 1, "retry must contain exactly one system message" + assert isinstance(second_call_messages[0], SystemMessage), "system message must lead the list" + assert FORCE_SEARCH_GUIDANCE in second_call_messages[0].content + + @pytest.mark.asyncio + async def test_force_search_skipped_when_no_tools(self, mock_llm_provider, mock_llm): + """When no tools are configured, force_search must NOT fire; the agent + should fall back to asking the user immediately.""" + clarif_msg = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg, complete_msg]) + + user_callback = AsyncMock(return_value="technical deep dive") + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[], # no tools available + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + result = await agent.run(state) + + assert result is not None + # User callback is called once for clarification (no force_search detour). + user_callback.assert_called_once() + + @pytest.mark.asyncio + async def test_force_search_fires_at_most_once(self, mock_llm_provider, mock_llm): + """Even if the LLM stubbornly refuses to call a tool after the force_search + nudge, the agent must not loop forever - it should proceed to asking the + user after the single nudge attempt.""" + clarif_msg_1 = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + # After the force_search nudge, the model still refuses to call a tool + # and returns another clarification request. + clarif_msg_2 = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + # After the user replies, the model completes. + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg_1, clarif_msg_2, complete_msg]) + + user_callback = AsyncMock(return_value="technical") + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + result = await agent.run(state) + + assert result is not None + # The LLM was invoked three times max - the nudge fired once, then we + # fell through to ask_for_clarification, and the user reply produced + # the final completion. + assert mock_llm.ainvoke.call_count == 3 + # The user was prompted exactly once (no infinite loop of nudges). + user_callback.assert_called_once() + + @pytest.mark.asyncio + async def test_force_search_not_triggered_when_llm_searches_first(self, mock_llm_provider, mock_llm): + """When the LLM voluntarily issues a tool call on the first turn, the + force_search path is never entered - normal flow is preserved.""" + tool_call_msg = AIMessage( + content="", + tool_calls=[{"name": "web_search_tool", "args": {"query": "AI"}, "id": "call_1"}], + ) + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[tool_call_msg, complete_msg]) + + user_callback = AsyncMock() + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + result = await agent.run(state) + + assert result is not None + assert mock_llm.ainvoke.call_count == 2 + user_callback.assert_not_called() + # The guidance must not appear in ANY message of ANY call — the model + # searched voluntarily, so the nudge path must never have fired. + for call in mock_llm.ainvoke.call_args_list: + assert not any(FORCE_SEARCH_GUIDANCE in str(m.content) for m in call.args[0]) + + @pytest.mark.asyncio + async def test_force_search_guidance_not_in_state_messages(self, mock_llm_provider, mock_llm): + """The force_search guidance must be injected ephemerally only; it must + never end up in state.messages, otherwise helpers like + get_latest_user_query would surface internal scaffolding back to the + user in fallback text.""" + # The LLM ignores the nudge and returns invalid JSON, triggering the + # invalid-format fallback path inside ask_clarification. The user then + # replies "skip", which forces completion - so only two LLM calls + # actually happen in this run. + clarif_msg_1 = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + clarif_invalid = AIMessage(content="not valid JSON at all") + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg_1, clarif_invalid]) + + # Capture what gets sent to the user. + prompts_received: list[str] = [] + + async def user_callback(question: str) -> str: + """Return the canned user reply for this test.""" + prompts_received.append(question) + return "skip" + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=user_callback, + ) + + original_query = "Research Project Foo at Acme" + state = ClarifierAgentState(messages=[HumanMessage(content=original_query)]) + final = await agent.graph.ainvoke(state, config={"callbacks": []}) + + # The guidance must never be persisted into state.messages. + final_messages = final["messages"] if isinstance(final, dict) else final.messages + assert not any(FORCE_SEARCH_GUIDANCE in str(m.content) for m in final_messages), ( + "force_search guidance leaked into persisted state" + ) + # The user was prompted exactly once - with a fallback derived from + # their actual query (the full topic survives), never from the + # force-search guidance. + assert len(prompts_received) == 1 + prompt_text = prompts_received[0] + assert "Project Foo" in prompt_text and "Acme" in prompt_text + assert FORCE_SEARCH_GUIDANCE not in prompt_text + # The internal force-search guidance must never be visible in any + # message the user-facing fallback would draw from. + assert "You attempted to ask the user" not in prompt_text + + @pytest.mark.asyncio + async def test_force_search_guidance_not_injected_after_user_reply(self, mock_llm_provider, mock_llm): + """After the user has actually replied (iteration > 0), the agent must + NOT re-inject the search-first nudge on the next LLM call. Otherwise + the model would receive 'issue a tool call now' immediately after the + user provided clarifying answer, causing a gratuitous search instead + of synthesizing the answer.""" + clarif_msg_1 = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + # After the nudge, the model still refuses to call a tool. + clarif_msg_2 = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="Which area?" + ).model_dump_json() + ) + # After the user replies, the model completes. + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg_1, clarif_msg_2, complete_msg]) + + user_callback = AsyncMock(return_value="technical deep dive") + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + await agent.run(state) + + # The 3rd LLM call happens AFTER the user reply (iteration moves from + # 0 to 1), so the inline search-before-clarify guard (gated on + # iteration == 0) must not fire again and the nudge must not appear in + # that call's message list. + assert mock_llm.ainvoke.call_count == 3 + third_call_messages = mock_llm.ainvoke.call_args_list[2].args[0] + assert not any(FORCE_SEARCH_GUIDANCE in str(m.content) for m in third_call_messages), ( + "force_search guidance must not be re-injected after the user replies" + ) + + @pytest.mark.asyncio + async def test_forced_retry_does_not_emit_consecutive_assistant_messages(self, mock_llm_provider, mock_llm): + """After a forced search-retry whose retry produces a tool call, the + message list fed to the LLM on the next turn must NOT contain two + consecutive assistant (AIMessage) turns. Two adjacent assistant + messages with no interleaved user/tool message are rejected with a 400 + by the OpenAI Chat Completions and Anthropic Messages APIs; mocked LLMs + don't enforce this, so we assert the invariant explicitly.""" + clarif_msg = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + tool_call_msg = AIMessage( + content="", + tool_calls=[{"name": "web_search_tool", "args": {"query": "AI"}, "id": "call_1"}], + ) + complete_msg = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif_msg, tool_call_msg, complete_msg]) + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=AsyncMock(), + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research Foo Project XYZ")]) + await agent.run(state) + + # Inspect every message list that was actually sent to the LLM and assert + # no two consecutive AIMessages appear (the API-invalid shape). + for call_idx, call in enumerate(mock_llm.ainvoke.call_args_list): + sent_messages = call.args[0] + offenders = _adjacent_assistant_pairs(sent_messages) + assert not offenders, ( + f"LLM call #{call_idx} contained consecutive assistant messages at {offenders}; " + "this is rejected by OpenAI/Anthropic APIs" + ) + + # Specifically: the forced retry must not persist the skipped + # clarification, so the post-tool history is [..., AIMessage(tool_call), + # ToolMessage, ...] with no stale AIMessage before the tool call. + third_call_messages = mock_llm.ainvoke.call_args_list[2].args[0] + ai_then_tool = any( + isinstance(third_call_messages[i], AIMessage) + and getattr(third_call_messages[i], "tool_calls", None) + and isinstance(third_call_messages[i + 1], ToolMessage) + for i in range(len(third_call_messages) - 1) + ) + assert ai_then_tool, "expected a tool-call AIMessage immediately followed by its ToolMessage in history" + + +class TestClarifierSkipMessageOrdering: + """Skip-command branch must not produce consecutive assistant messages. + + Regression tests for the ordering bug surfaced during the PR #245 audit: + the skip branch returned an AIMessage(complete) without persisting the + user's reply, leaving it adjacent to the prior clarification AIMessage; the + graph then re-entered agent_node and (with the budget exhausted) appended a + third AIMessage. Two/three consecutive assistant turns are rejected by the + OpenAI/Anthropic APIs and corrupt the planner call when plan approval is on. + """ + + @pytest.fixture + def mock_llm(self): + """Create a mock LLM.""" + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + return llm + + @pytest.fixture + def mock_llm_provider(self, mock_llm): + """Create a mock LLM provider returning the mock LLM.""" + provider = MagicMock(spec=LLMProvider) + provider.get = MagicMock(return_value=mock_llm) + return provider + + @pytest.mark.asyncio + async def test_skip_persists_reply_and_no_consecutive_assistants(self, mock_llm_provider, mock_llm): + """User skips after one clarification: final history must interleave the + skip reply (HumanMessage) and contain no consecutive AIMessages.""" + clarif = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + # Only one real LLM response is needed; after skip the graph completes + # without another model call (the early-complete guard returns {}). + mock_llm.ainvoke = AsyncMock(side_effect=[clarif]) + + async def user_callback(question: str) -> str: + """Always skip the clarification.""" + return "skip" + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[], # no tools → no forced search; isolate the skip path + user_prompt_callback=user_callback, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + + # Inspect persisted message ordering from the final graph state. + final = await agent.graph.ainvoke(state, config={"callbacks": []}) + msgs = final["messages"] if isinstance(final, dict) else final.messages + offenders = _adjacent_assistant_pairs(msgs) + assert not offenders, f"persisted history has consecutive assistant messages at {offenders}: {msgs}" + # The skip reply must be persisted as a HumanMessage. + assert any(isinstance(m, HumanMessage) and m.content == "skip" for m in msgs), ( + "skip reply was not persisted as a HumanMessage" + ) + # Exactly one *terminal completion* AIMessage should end the dialog: a + # message that parses to needs_clarification=false. Matching any non-tool + # AIMessage (e.g. the earlier clarification prompt) or allowing >= 1 would + # not catch a missing or duplicated completion. + completion_ais = [ + m + for m in msgs + if isinstance(m, AIMessage) + and not ClarifierAgent._has_tool_invocations([m]) + and agent._is_complete(m.content) + ] + got = len(completion_ais) + assert got == 1, f"expected exactly one terminal completion AIMessage, got {got}" + + @pytest.mark.asyncio + async def test_skip_with_plan_approval_planner_gets_valid_sequence(self, mock_llm_provider, mock_llm): + """With plan approval on, the corrupted skip history previously flowed + into the planner's ainvoke. Assert the planner never receives two + consecutive assistant messages.""" + clarif = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif]) + + planner_llm = MagicMock() + plan_json = '{"title": "Plan", "sections": ["Intro", "Analysis"]}' + planner_llm.ainvoke = AsyncMock(return_value=AIMessage(content=plan_json)) + + replies = iter(["skip", "approve"]) + + async def user_callback(question: str) -> str: + """Return the canned user reply for this test.""" + return next(replies) + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[], + user_prompt_callback=user_callback, + enable_plan_approval=True, + planner_llm=planner_llm, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + result = await agent.run(state) + + assert result is not None + assert result.plan_approved is True + # The planner was called; none of its input message lists may contain + # consecutive assistant messages. + assert planner_llm.ainvoke.call_count == 1 + for call_idx, call in enumerate(planner_llm.ainvoke.call_args_list): + sent = call.args[0] + offenders = _adjacent_assistant_pairs(sent) + assert not offenders, f"planner ainvoke #{call_idx} had consecutive assistant messages at {offenders}" + + @pytest.mark.asyncio + async def test_blank_skip_reply_does_not_persist_empty_human_message(self, mock_llm_provider, mock_llm): + """A blank skip reply must not put an empty HumanMessage into history. + + ``_is_skip_command`` treats ``""``/whitespace as skip, so the skip branch + substitutes a non-empty sentinel. With plan approval enabled the history + flows into the planner's ainvoke; an empty-content message there is + rejected by some chat APIs. Assert no empty HumanMessage reaches the + planner. Regression test for the blank-skip-reply edge case. + """ + clarif = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif]) + + planner_llm = MagicMock() + planner_llm.ainvoke = AsyncMock( + return_value=AIMessage(content='{"title": "Plan", "sections": ["Intro", "Analysis"]}') + ) + + # First reply is blank (still a skip command); then approve the plan. + replies = iter(["", "approve"]) + + async def user_callback(question: str) -> str: + """Return the canned user reply for this test.""" + return next(replies) + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[], + user_prompt_callback=user_callback, + enable_plan_approval=True, + planner_llm=planner_llm, + ) + + state = ClarifierAgentState(messages=[HumanMessage(content="Research AI")]) + result = await agent.run(state) + + assert result is not None + assert planner_llm.ainvoke.call_count == 1 + for call_idx, call in enumerate(planner_llm.ainvoke.call_args_list): + sent = call.args[0] + empty = [i for i, m in enumerate(sent) if isinstance(m, HumanMessage) and not str(m.content).strip()] + assert not empty, f"planner ainvoke #{call_idx} received empty HumanMessage(s) at {empty}" + + +class TestClarifierReviewRegressions: + """Regression tests for edge cases surfaced in deep review of #245.""" + + @pytest.fixture + def mock_llm(self): + """Create a mock LLM.""" + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + return llm + + @pytest.fixture + def mock_llm_provider(self, mock_llm): + """Create a mock LLM provider returning the mock LLM.""" + provider = MagicMock(spec=LLMProvider) + provider.get = MagicMock(return_value=mock_llm) + return provider + + @pytest.mark.asyncio + async def test_force_search_fires_despite_prior_turn_tool_calls(self, mock_llm_provider, mock_llm): + """A tool call from an earlier conversation turn must not suppress the + search-before-clarify nudge for a fresh user query. The guard is scoped + to tool activity since the latest user turn.""" + # First model call (for the new query): clarify without a tool call. + # Second call (after the nudge): emit a tool call. Third: complete. + clarif = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + tool_call = AIMessage(content="", tool_calls=[{"name": "web_search_tool", "args": {"query": "x"}, "id": "c1"}]) + complete = AIMessage( + content=ClarificationResponse(needs_clarification=False, clarification_question=None).model_dump_json() + ) + mock_llm.ainvoke = AsyncMock(side_effect=[clarif, tool_call, complete]) + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=AsyncMock(), + ) + + # History carries a tool call from an EARLIER turn, then the fresh query. + prior_tool = AIMessage(content="", tool_calls=[{"name": "web_search_tool", "args": {}, "id": "c0"}]) + state = ClarifierAgentState( + messages=[ + HumanMessage(content="an earlier question"), + prior_tool, + ToolMessage(content="old result", tool_call_id="c0"), + HumanMessage(content="a fresh research query"), + ] + ) + await agent.run(state) + + # The nudge must still have fired: 3 model calls, and the 2nd received + # the FORCE_SEARCH_GUIDANCE (it is not suppressed by the prior tool call). + assert mock_llm.ainvoke.call_count == 3 + second_call_messages = mock_llm.ainvoke.call_args_list[1].args[0] + assert any(isinstance(m, SystemMessage) and FORCE_SEARCH_GUIDANCE in m.content for m in second_call_messages) + + @pytest.mark.asyncio + async def test_exhausted_entry_with_clarification_last_no_adjacent_assistants(self, mock_llm_provider, mock_llm): + """If the budget is already exhausted and the last message is a + non-complete clarification AIMessage, the completion must be interleaved + with a sentinel user turn so no two assistant messages are adjacent + (which would 400 the planner under plan approval).""" + planner_llm = MagicMock() + planner_llm.ainvoke = AsyncMock( + return_value=AIMessage(content='{"title": "Plan", "sections": ["Intro", "Analysis"]}') + ) + + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[], + user_prompt_callback=AsyncMock(return_value="approve"), + enable_plan_approval=True, + planner_llm=planner_llm, + ) + + # Externally-provided state at exhaustion (max_turns=0) whose last turn + # is a non-complete clarification AIMessage. + clarif = AIMessage( + content=ClarificationResponse( + needs_clarification=True, clarification_question="What aspect?" + ).model_dump_json() + ) + state = ClarifierAgentState( + messages=[HumanMessage(content="Research AI"), clarif], + max_turns=0, + ) + result = await agent.run(state) + + assert result is not None + # The planner must never receive two consecutive assistant messages. + assert planner_llm.ainvoke.call_count == 1 + for call_idx, call in enumerate(planner_llm.ainvoke.call_args_list): + sent = call.args[0] + offenders = _adjacent_assistant_pairs(sent) + assert not offenders, f"planner ainvoke #{call_idx} had consecutive assistant messages at {offenders}" + + @pytest.mark.asyncio + async def test_exhausted_entry_with_pending_tool_call_is_not_completed_directly(self, mock_llm_provider, mock_llm): + """At exhaustion, if the last message is an AIMessage with pending tool + calls, agent_node must not stack a completion on top of it (which would + leave the tool call unresolved / produce invalid history). It returns no + new message so the graph can route the tool call to the tools node.""" + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=AsyncMock(), + ) + + pending_tool = AIMessage( + content="", tool_calls=[{"name": "web_search_tool", "args": {"query": "x"}, "id": "c1"}] + ) + state = ClarifierAgentState( + messages=[HumanMessage(content="Research AI"), pending_tool], + max_turns=0, + ) + + # Run only the agent node against this state. + result = await agent.graph.nodes["agent"].ainvoke(state) + assert isinstance(result, dict), f"agent node must return a state-update dict, got {type(result)}" + new_messages = result.get("messages", []) + # No completion AIMessage may be appended on top of the pending tool call. + assert not any(isinstance(m, AIMessage) and not getattr(m, "tool_calls", None) for m in new_messages), ( + "must not append a completion on top of a pending tool call" + )