diff --git a/rag/advanced_rag/harness/agent.py b/rag/advanced_rag/harness/agent.py index 5621e60201..684291716d 100644 --- a/rag/advanced_rag/harness/agent.py +++ b/rag/advanced_rag/harness/agent.py @@ -193,7 +193,7 @@ async def research_agent_loop( per-claim would race: the first claim to execute would clear it, starving the rest). """ - phase = determine_current_phase(context) + phase = determine_current_phase(context, claim=claim) phase_config = SEARCH_PHASES.get(phase, {}) gated_defs = get_gated_tools( phase=phase, @@ -202,15 +202,32 @@ async def research_agent_loop( context=context, has_routed_scope=bool(getattr(pipeline, "_routed_docs", None)), web_enabled=bool(getattr(tools, "has_web", lambda: False)()), + claim=claim, ) + pipeline._active_phase = phase + pipeline._round_had_evidence = False + pipeline._round_had_routed_scope_progress = False + # Clone so binding tools never leaks onto the shared chat model. agent_mdl = tools.chat_mdl.clone() if getattr(agent_mdl, "is_tools", False): - return await _research_native(claim, agent_mdl, pipeline, phase, phase_config, gated_defs, mode, followups) + result = await _research_native(claim, agent_mdl, pipeline, phase, phase_config, gated_defs, mode, followups) + else: + _LOG.info("research_agent: model lacks native tool support; falling back to text-based tool selection") + result = await _research_text(claim, tools, pipeline, phase, phase_config, gated_defs, mode, followups) - _LOG.info("research_agent: model lacks native tool support; falling back to text-based tool selection") - return await _research_text(claim, tools, pipeline, phase, phase_config, gated_defs, mode, followups) + # Bookkeeping for locate-phase web fallback: a locate round that produced + # evidence chunks or newly routed document scope counts as progress and + # resets this claim's streak. Pre-existing request doc_scope does not + # count: only scope produced by this round should prevent the fallback. + # The phase stays `locate`; gating.py decides when repeated locate + # failures should admit `web_search` into the candidate tool set. + if pipeline._round_had_evidence or pipeline._round_had_routed_scope_progress: + claim.locate_empty_streak = 0 + elif phase == "locate": + claim.locate_empty_streak += 1 + return result async def _research_native( @@ -228,8 +245,9 @@ async def _research_native( session = ResearchToolSession(pipeline, phase, claim) agent_mdl.bind_tools(session, schemas) # Bound the model's internal tool loop to the mode's agent-cycle budget. + base_rounds = max(1, mode.max_agent_cycles) if hasattr(agent_mdl, "mdl") and hasattr(agent_mdl.mdl, "max_rounds"): - agent_mdl.mdl.max_rounds = max(1, mode.max_agent_cycles) + agent_mdl.mdl.max_rounds = base_rounds system = RESEARCH_AGENT_PROMPT.format( claim_description=claim.description, @@ -283,12 +301,15 @@ async def _research_text( followups: list[str] | None = None, ) -> dict: """Fallback: prompt-based tool selection for models without native tools.""" + # Mirror the native path's cycle budget; the locate→web_search fallback is + # handled by gating.get_gated_tools injecting web_search into the tool set. + text_max_cycles = mode.max_agent_cycles system = RESEARCH_AGENT_TEXT_PROMPT.format( claim_description=claim.description, phase=phase, phase_hint=phase_config.get("tool_hint", ""), tool_list=_fmt_tool_list(gated_defs), - max_cycles=mode.max_agent_cycles, + max_cycles=text_max_cycles, ) history: list[dict] = [] @@ -304,7 +325,7 @@ async def _research_text( # normalized against this claim's recorded chunks (same as the native path). session = ResearchToolSession(pipeline, phase, claim) - for cycle in range(mode.max_agent_cycles): + for cycle in range(text_max_cycles): try: ans = await tools.chat_mdl.async_chat(system, history, {"temperature": 0.3}) if isinstance(ans, tuple): @@ -325,13 +346,15 @@ async def _research_text( if not isinstance(args, dict): _LOG.warning("generate_report: arguments not a dict (%s); using empty", type(args).__name__) args = {} - return session._normalize_report(args) + report = session._normalize_report(args) + return report if tool_call.get("name") == "think_tool": history.append({"role": "user", "content": "[continue]"}) continue args = tool_call.get("arguments", {}) + # Text fallback bypasses ResearchToolSession.tool_call_async(), so it result = await execute_with_fallback(pipeline, tool_call["name"], phase, **args) if result.chunks: session._record_evidence_ids(result.chunks) @@ -418,7 +441,8 @@ async def _force_generate_report( report = json_repair.loads(text) if isinstance(report, dict) and session is not None: - return session._normalize_report(report) + normalized = session._normalize_report(report) + return normalized return report if isinstance(report, dict) else {"report": str(report)} except Exception: _LOG.exception("force_generate_report failed") diff --git a/rag/advanced_rag/harness/config.py b/rag/advanced_rag/harness/config.py index 5dec62b581..651009998f 100644 --- a/rag/advanced_rag/harness/config.py +++ b/rag/advanced_rag/harness/config.py @@ -50,7 +50,7 @@ THINKING_MODES: dict[str, ExecutionStrategy] = { allows_replan=False, max_orchestrator_cycles=3, max_agent_cycles=2, - max_parallel_agents=2, + max_parallel_agents=4, available_tools=[ "hybrid_search", "web_search", @@ -79,7 +79,7 @@ THINKING_MODES: dict[str, ExecutionStrategy] = { allows_replan=True, max_orchestrator_cycles=4, max_agent_cycles=2, - max_parallel_agents=3, + max_parallel_agents=4, available_tools=[ "hybrid_search", "bm25_search", diff --git a/rag/advanced_rag/harness/orchestrator/agentic.py b/rag/advanced_rag/harness/orchestrator/agentic.py index f5b54708de..e3a5d7fdf6 100644 --- a/rag/advanced_rag/harness/orchestrator/agentic.py +++ b/rag/advanced_rag/harness/orchestrator/agentic.py @@ -324,6 +324,17 @@ async def _run_claim_research( except asyncio.CancelledError: raise except TimeoutError: + if getattr(pipeline, "_active_phase", None) == "locate": + if pipeline._round_had_evidence or pipeline._round_had_routed_scope_progress: + claim.locate_empty_streak = 0 + else: + claim.locate_empty_streak += 1 + _LOG.warning( + "[Agentic research] claim=%s timed out in locate (progress=%s, locate_empty_streak=%d).", + claim.claim_id, + pipeline._round_had_evidence or pipeline._round_had_routed_scope_progress, + claim.locate_empty_streak, + ) _LOG.warning( '[Agentic research] Gave up on "%s" — it took longer than %ss.', _snip(claim.description), diff --git a/rag/advanced_rag/harness/pipeline.py b/rag/advanced_rag/harness/pipeline.py index 3da9f154c0..c30afbe5ec 100644 --- a/rag/advanced_rag/harness/pipeline.py +++ b/rag/advanced_rag/harness/pipeline.py @@ -31,6 +31,11 @@ class Pipeline: self.trace: list[dict] = [] # Latest relevant-document set produced by a routing tool this run. self._routed_docs: list[str] = list(getattr(rag_tools, "doc_scope", None) or []) + # Per-claim round state, read by both normal completion and timeout handling. + self._active_phase: str | None = None + self._round_initial_routed_docs: tuple[str, ...] = tuple(self._routed_docs) + self._round_had_evidence = False + self._round_had_routed_scope_progress = False async def execute(self, tool_name: str, **kwargs) -> ToolResult: """Execute a registered tool by name.""" @@ -54,14 +59,19 @@ class Pipeline: elapsed = time.time() - start self.trace.append({"tool": tool_name, "args": kwargs, "elapsed": elapsed, "success": True}) result = self._normalize(raw) + if result.chunks: + self._round_had_evidence = True # A routing tool (e.g. dataset_navigation_search) yields the relevant # document IDs; remember them so the scope-consuming tools above can # inherit them on later turns. if result.docs: if hasattr(self.tools, "scoped_doc_ids"): - self._routed_docs = self.tools.scoped_doc_ids(list(result.docs)) or [] + new_routed_docs = self.tools.scoped_doc_ids(list(result.docs)) or [] else: - self._routed_docs = list(result.docs) + new_routed_docs = list(result.docs) + if new_routed_docs and tuple(new_routed_docs) != self._round_initial_routed_docs: + self._round_had_routed_scope_progress = True + self._routed_docs = new_routed_docs # Feed the shared citation pool: agent searches go through the # pipeline, so without this their evidence never reaches kbinfos and # the final answer has nothing to cite. diff --git a/rag/advanced_rag/harness/tools/gating.py b/rag/advanced_rag/harness/tools/gating.py index 7b8c1afff2..6ec6f37b1a 100644 --- a/rag/advanced_rag/harness/tools/gating.py +++ b/rag/advanced_rag/harness/tools/gating.py @@ -1,7 +1,11 @@ """Tool selection gating: phase-based filtering and fallback chain.""" +import logging + from rag.advanced_rag.harness.tools.registry import TOOL_REGISTRY -from rag.advanced_rag.harness.types import OrchestratorContext +from rag.advanced_rag.harness.types import ClaimTarget, OrchestratorContext + +_LOG = logging.getLogger(__name__) # Search phase definitions @@ -90,6 +94,7 @@ def get_gated_tools( context: OrchestratorContext, has_routed_scope: bool = False, web_enabled: bool = True, + claim: ClaimTarget | None = None, ) -> list[dict]: """Filter, sort, and gate tools by phase priority and context.""" phase_config = SEARCH_PHASES.get(phase) @@ -110,6 +115,8 @@ def get_gated_tools( sorted_tools.append(tool_name) selected = sorted_tools[: phase_config["max_returned"]] + if phase == "locate": + selected = _inject_locate_fallback_tools(selected, available_tools, claim, web_enabled) # Copy the registry schemas before annotating — the registry dicts are # shared process-wide, mutating them would leak phase hints across # concurrent requests. @@ -124,10 +131,46 @@ def _default_defs(tool_names: list[str], web_enabled: bool = True) -> list[dict] return [TOOL_REGISTRY[n]["function_schema"] for n in tool_names if n in TOOL_REGISTRY and (web_enabled or n != "web_search")] -def determine_current_phase(context: OrchestratorContext) -> str: +# After this many consecutive `locate` rounds for a single claim that produced +# zero evidence chunks, keep the phase in `locate` but admit `web_search` into +# the gated tool set. This preserves the phase semantics ("still trying to +# locate an answer") while giving the agent an external-search escape hatch +# when the knowledge base lacks the fact entirely. +LOCATE_EMPTY_ADVANCE_THRESHOLD = 2 + + +def determine_current_phase( + context: OrchestratorContext, + claim: ClaimTarget | None = None, +) -> str: """Determine the current search phase based on context.""" - if not context.has_any_chunks(): + if claim is not None: + evidence_ids = claim.agent_result.evidence_ids if claim.agent_result else None + if not evidence_ids: + return "locate" + elif not context.has_any_chunks(): return "locate" if context.verdict and context.verdict.has_conflicts: return "verify" return "explore" + + +def _inject_locate_fallback_tools( + selected: list[str], + available_tools: list[str], + claim: ClaimTarget | None, + web_enabled: bool, +) -> list[str]: + """Keep `locate` semantics but add external fallback when KB locate fails.""" + if not web_enabled or claim is None: + return selected + if claim.locate_empty_streak < LOCATE_EMPTY_ADVANCE_THRESHOLD: + return selected + if "web_search" not in available_tools or "web_search" in selected: + return selected + _LOG.info( + "[Tool gating] claim=%s: injecting web_search after repeated locate misses (locate_empty_streak=%d)", + claim.claim_id, + claim.locate_empty_streak, + ) + return [*selected, "web_search"] diff --git a/rag/advanced_rag/harness/types.py b/rag/advanced_rag/harness/types.py index bc6dfa27c5..60d735c84d 100644 --- a/rag/advanced_rag/harness/types.py +++ b/rag/advanced_rag/harness/types.py @@ -72,6 +72,14 @@ class ClaimTarget: confidence: float = 0.0 suggested_tools: list[str] = field(default_factory=list) agent_result: AgentResult | None = None + # Consecutive research rounds for THIS claim that stayed in `locate` and + # produced neither evidence chunks nor routed document scope. Claim-scoped + # on purpose: claims are researched in parallel over one shared + # OrchestratorContext, so a global counter would let one claim's empty + # locate rounds influence a sibling. Once it reaches + # LOCATE_EMPTY_ADVANCE_THRESHOLD, the phase still stays `locate`, but + # gating may admit `web_search` for facts the corpus lacks. + locate_empty_streak: int = 0 @dataclass @@ -170,7 +178,6 @@ class OrchestratorContext: claims: list[ClaimTarget] mode: str iteration: int = 0 - current_phase: str = "locate" verdict: SufficiencyVerdict | None = None history: list[dict] = field(default_factory=list) # Follow-up search queries produced by the Phase-2 LLM Sufficient Context