mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-25 09:32:38 +08:00
2aaa6baf0c
### Summary 1. It changes the fallback semantics of the locate phase. When no chunks are found, the system stays in locate. If the same claim has two consecutive locate rounds with neither evidence chunks nor newly routed document scope, web_search is admitted to the candidate tool set on the next locate round as an external fallback. 2. It makes locate_empty_streak claim-scoped instead of shared in the global context. This prevents one claim’s empty locate rounds from affecting sibling claims running in parallel. 3. On the config side, it only raises max_parallel_agents for high / ultra to 4, without changing max_agent_cycles. This increases parallel claim execution without deepening per-claim search.
166 lines
7.5 KiB
Python
166 lines
7.5 KiB
Python
"""Pipeline — unified tool execution dispatcher."""
|
|
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
from rag.advanced_rag.harness.tools.registry import TOOL_REGISTRY
|
|
from rag.advanced_rag.harness.types import ToolResult
|
|
|
|
_LOG = logging.getLogger(__name__)
|
|
|
|
# Tools that retrieve *within* a set of documents. When a routing tool
|
|
# (``dataset_navigation_search``) has produced a relevant-document set, these
|
|
# inherit it as their ``doc_scope`` unless the caller passed one explicitly, so
|
|
# a follow-up search stays within the routed docs instead of re-scanning the KB.
|
|
_DOC_SCOPE_CONSUMERS = {"ontology_navigate", "mindmap_navigate", "graph_explore", "hybrid_search", "vector_search", "bm25_search", "structured_query", "dataset_navigation_search"}
|
|
|
|
|
|
class Pipeline:
|
|
"""Unified tool execution layer.
|
|
|
|
- execute(tool_name, **kwargs): dispatch to registered tool, normalize result
|
|
- available_tools(mode_tools): return LLM-visible tool definitions (compilation-filtered)
|
|
- get_chunks(evidence_ids): retrieve raw chunks for sufficiency cross-check
|
|
- trace: execution history for auditing
|
|
"""
|
|
|
|
def __init__(self, rag_tools, compilation_map: dict[str, set[str]] | None = None):
|
|
self.tools = rag_tools
|
|
self.compilation_map = compilation_map or {}
|
|
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."""
|
|
tool = TOOL_REGISTRY.get(tool_name)
|
|
if not tool:
|
|
return ToolResult(chunks=[], metadata={}, error=f"Unknown tool: {tool_name}")
|
|
|
|
fn = tool.get("fn")
|
|
if not fn:
|
|
return ToolResult(chunks=[], metadata={}, error=f"Tool {tool_name} has no executor")
|
|
|
|
# Downstream scoping: a within-document tool inherits the doc IDs a prior
|
|
# router (dataset_navigation_search) produced, unless the caller passed
|
|
# an explicit doc_scope.
|
|
if tool_name in _DOC_SCOPE_CONSUMERS and self._routed_docs and not kwargs.get("doc_scope"):
|
|
kwargs["doc_scope"] = list(self._routed_docs)
|
|
|
|
start = time.time()
|
|
try:
|
|
raw = await fn(self.tools, **kwargs)
|
|
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"):
|
|
new_routed_docs = self.tools.scoped_doc_ids(list(result.docs)) or []
|
|
else:
|
|
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.
|
|
self._merge_into_kbinfos(result)
|
|
return result
|
|
except Exception as e:
|
|
elapsed = time.time() - start
|
|
_LOG.exception("Pipeline.execute(%s) failed", tool_name)
|
|
self.trace.append({"tool": tool_name, "args": kwargs, "elapsed": elapsed, "success": False, "error": str(e)})
|
|
return ToolResult(chunks=[], metadata={}, error=str(e))
|
|
|
|
def available_tools(self, mode_tools: list[str]) -> list[dict]:
|
|
"""Return LLM-visible tool definitions, filtered by compilation availability."""
|
|
names = filter_available_tools(mode_tools, self.compilation_map)
|
|
defs = []
|
|
for name in names:
|
|
tool = TOOL_REGISTRY.get(name)
|
|
if tool and tool.get("function_schema"):
|
|
defs.append(tool["function_schema"])
|
|
return defs
|
|
|
|
def get_chunks(self, evidence_ids: list[int]) -> dict[int, dict]:
|
|
"""Retrieve raw chunks by ID from current kbinfos."""
|
|
result = {}
|
|
chunks = self.tools.kbinfos.get("chunks", [])
|
|
for eid in evidence_ids:
|
|
if 0 <= eid < len(chunks):
|
|
result[eid] = chunks[eid]
|
|
return result
|
|
|
|
def get_trace(self) -> list[dict]:
|
|
return list(self.trace)
|
|
|
|
# ── Private ──
|
|
|
|
def _merge_into_kbinfos(self, result: ToolResult) -> None:
|
|
"""Merge a tool result's chunks/doc_aggs into ``tools.kbinfos``, deduped."""
|
|
if not result or not result.chunks:
|
|
return
|
|
kb = self.tools.kbinfos
|
|
seen = {c.get("chunk_id") or c.get("id") or id(c) for c in kb.get("chunks", [])}
|
|
for c in result.chunks:
|
|
k = c.get("chunk_id") or c.get("id") or id(c)
|
|
if k in seen:
|
|
continue
|
|
seen.add(k)
|
|
kb.setdefault("chunks", []).append(c)
|
|
aggs = result.metadata.get("aggs") if isinstance(result.metadata, dict) else None
|
|
if aggs:
|
|
dseen = {d.get("doc_id") for d in kb.get("doc_aggs", [])}
|
|
for d in aggs:
|
|
if d.get("doc_id") in dseen:
|
|
continue
|
|
dseen.add(d.get("doc_id"))
|
|
kb.setdefault("doc_aggs", []).append(d)
|
|
|
|
@staticmethod
|
|
def _normalize(raw: Any) -> ToolResult:
|
|
if isinstance(raw, ToolResult):
|
|
return raw
|
|
if isinstance(raw, dict):
|
|
return ToolResult(
|
|
chunks=raw.get("chunks", []),
|
|
docs=raw.get("docs") or None,
|
|
metadata={"aggs": raw.get("doc_aggs", []), "answer": raw.get("answer", "")},
|
|
)
|
|
if isinstance(raw, list):
|
|
# A list of doc-id strings is a document-routing result (e.g.
|
|
# dataset_navigation_search); a list of dicts is chunks.
|
|
if raw and all(isinstance(x, str) for x in raw):
|
|
return ToolResult(docs=list(raw), metadata={})
|
|
return ToolResult(chunks=raw, metadata={})
|
|
return ToolResult(chunks=[], metadata={"raw": str(raw)})
|
|
|
|
|
|
def filter_available_tools(tool_names: list[str], compilation_map: dict[str, set[str]]) -> list[str]:
|
|
"""Filter tool list by compilation artifact availability."""
|
|
available = []
|
|
for name in tool_names:
|
|
tool = TOOL_REGISTRY.get(name)
|
|
if not tool:
|
|
continue
|
|
if tool.get("requires_compilation"):
|
|
comp_type = tool.get("compilation_type")
|
|
# ``compilation_type`` may name one artifact or several (a tool that
|
|
# reads either one is available when ANY of them is compiled).
|
|
wanted = {comp_type} if isinstance(comp_type, str) else set(comp_type or ())
|
|
if wanted and not any(wanted & comps for comps in compilation_map.values()):
|
|
continue
|
|
available.append(name)
|
|
return available
|