mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-18 05:37:24 +08:00
4391e03886
### What problem does this PR solve? Improve concurrency in the knowledge compilation pipeline: - Run Compile LLM requests concurrently while preserving ordered commits. - Run merge flush tasks concurrently while keeping ES writes ordered. - Improve concurrency for local deduplication, chain validation, and ES deduplication. - Remove temporary debugging instrumentation and unused timing variables. ### Type of change - [x] Refactor (no functional change)
460 lines
19 KiB
Python
460 lines
19 KiB
Python
#
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
"""Handler-free core for document-scoped knowledge (structure) compilation.
|
|
|
|
Extracted from
|
|
``rag/svr/task_executor_refactor/chunk_post_processor.py`` so the same
|
|
template resolution, batching, accumulate / merge-flush and synthesis logic
|
|
can be driven from two places:
|
|
|
|
* the chunking **task executor**, which streams a document's chunks out of
|
|
the doc store (``run_document_structure_compile``), and
|
|
* the ``rag.flow`` **Compiler** component, which receives chunks in-memory
|
|
from an upstream pipeline node.
|
|
|
|
Only the non-``tree`` template kinds are handled here. ``tree`` templates run
|
|
RAPTOR over the whole document and are still driven from the task executor
|
|
(``run_tree_templates``), which owns the doc-store reload + ``RaptorService``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import AsyncIterator, Callable
|
|
|
|
from api.db.services.compilation_template_service import CompilationTemplateService
|
|
from api.db.services.compilation_template_group_service import (
|
|
CompilationTemplateGroupService,
|
|
)
|
|
from api.db.services.llm_service import LLMBundle
|
|
from common.exceptions import TaskCanceledException
|
|
from rag.advanced_rag.knowlege_compile.structure import (
|
|
LLMCallPool,
|
|
compile_structure_from_text,
|
|
merge_compiled_structures,
|
|
)
|
|
|
|
|
|
# ----- tunables ------------------------------------------------------
|
|
# Bound how many source chunks are handed to a single
|
|
# ``compile_structure_from_text`` invocation.
|
|
DOC_STRUCTURE_COMPILE_BATCH_CHUNKS = 4
|
|
|
|
# Bound the number of batch/template extraction calls in flight. Results are
|
|
# committed in submission order so accumulator updates and merge flushes stay
|
|
# deterministic while the LLM calls run concurrently.
|
|
DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT = 15
|
|
|
|
# Total task-scoped Chat LLM capacity shared by compile, chain validation and
|
|
# merge decisions. A request waits in the priority queue when all slots are busy.
|
|
DOC_STRUCTURE_LLM_POOL_SIZE = 20
|
|
|
|
# Bound how many compiled ES-ready docs may accumulate before we flush
|
|
# them through ``merge_compiled_structures``. The merger does pairwise
|
|
# cosine + LLM duplicate-judging, so it's the more expensive step; we
|
|
# cap the per-flush set to keep the local-dedup buckets tractable.
|
|
DOC_STRUCTURE_MERGE_MAX_DOCS = 512
|
|
|
|
# Hard wall on the chain-validator LLM correction step.
|
|
STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S = 120.0
|
|
|
|
|
|
# ----- template resolution -------------------------------------------
|
|
|
|
|
|
def resolve_template_ids_from_groups(group_ids, tenant_id: str) -> list[str]:
|
|
"""Resolve an ordered, de-duplicated list of compilation-template ids
|
|
from a list of template-*group* ids.
|
|
|
|
Mirrors ``_parser_config_compilation_template_ids`` but takes the group
|
|
ids directly (the ``rag.flow`` Compiler carries them as a component
|
|
parameter rather than inside ``parser_config``).
|
|
"""
|
|
template_ids: list[str] = []
|
|
seen: set[str] = set()
|
|
for group_id in group_ids or []:
|
|
if not isinstance(group_id, str) or not group_id.strip():
|
|
continue
|
|
for template_id in CompilationTemplateGroupService.resolve_template_ids(
|
|
group_id.strip(),
|
|
tenant_id,
|
|
):
|
|
if template_id in seen:
|
|
continue
|
|
seen.add(template_id)
|
|
template_ids.append(template_id)
|
|
return template_ids
|
|
|
|
|
|
def load_active_templates(template_ids, tenant_id: str) -> list[tuple[str, dict]]:
|
|
"""Load each template's saved config and keep only the ones that drive a
|
|
real, non-``artifacts`` structure compilation.
|
|
|
|
Returns ``[(template_id, parser_cfg), ...]`` — templates that are missing,
|
|
have an invalid config, or resolve to no/``artifacts`` kind are dropped
|
|
(with a warning for the missing/invalid cases).
|
|
"""
|
|
from api.apps.restful_apis.chunk_api import _compilation_template_kind
|
|
|
|
active_templates: list[tuple[str, dict]] = []
|
|
for template_id in template_ids:
|
|
template = CompilationTemplateService.get_saved(template_id, tenant_id)
|
|
if not template:
|
|
logging.warning("document_structure_compile: template %s not found", template_id)
|
|
continue
|
|
parser_cfg = template.get("config") or {}
|
|
if not isinstance(parser_cfg, dict):
|
|
logging.warning("document_structure_compile: template %s config is invalid", template_id)
|
|
continue
|
|
kind = _compilation_template_kind(parser_cfg.get("kind"))
|
|
if not kind or kind == "artifacts":
|
|
continue
|
|
active_templates.append((template_id, parser_cfg))
|
|
return active_templates
|
|
|
|
|
|
def split_tree_templates(
|
|
active_templates: list[tuple[str, dict]],
|
|
) -> tuple[list[tuple[str, dict]], list[tuple[str, dict]]]:
|
|
"""Partition templates into ``(tree, non_tree)`` by kind."""
|
|
from api.apps.restful_apis.chunk_api import _compilation_template_kind
|
|
|
|
tree_templates: list[tuple[str, dict]] = []
|
|
non_tree_templates: list[tuple[str, dict]] = []
|
|
for tid, cfg in active_templates:
|
|
if _compilation_template_kind((cfg or {}).get("kind")) == "tree":
|
|
tree_templates.append((tid, cfg))
|
|
else:
|
|
non_tree_templates.append((tid, cfg))
|
|
return tree_templates, non_tree_templates
|
|
|
|
|
|
# ----- non-tree compilation core -------------------------------------
|
|
|
|
|
|
async def run_structure_compile_over_batches(
|
|
*,
|
|
active_templates: list[tuple[str, dict]],
|
|
chat_mdl_by_tid: dict[str, LLMBundle],
|
|
embedding_model: LLMBundle,
|
|
tenant_id: str,
|
|
kb_id: str,
|
|
doc_id: str,
|
|
language: str,
|
|
chunk_batches: AsyncIterator[list[dict]],
|
|
progress_cb: Callable[..., None],
|
|
cancel_check: Callable[[], bool] = lambda: False,
|
|
record: Callable[[str, dict], None] | None = None,
|
|
) -> dict[str, dict]:
|
|
"""Extract + merge structures for every non-``tree`` template over an
|
|
async stream of chunk batches, then run the optional synthesis phase.
|
|
|
|
``active_templates`` must already be the non-tree subset with a resolved
|
|
chat model in ``chat_mdl_by_tid``. Chunks arrive as an async iterator of
|
|
batches so callers can stream them from the doc store or hand over an
|
|
in-memory list; each ``dict`` must expose ``id`` and text
|
|
(``content_with_weight`` / ``text``).
|
|
|
|
Returns ``{template_id: {"inserted", "updated", "duplicates_dropped"}}``.
|
|
Raises :class:`TaskCanceledException` when ``cancel_check`` trips.
|
|
"""
|
|
from api.apps.restful_apis.chunk_api import _compilation_template_kind
|
|
|
|
if not active_templates:
|
|
return {}
|
|
|
|
total = len(active_templates)
|
|
llm_pool = LLMCallPool(DOC_STRUCTURE_LLM_POOL_SIZE)
|
|
|
|
accumulators: dict[str, list[dict]] = {tid: [] for tid, _ in active_templates}
|
|
template_kinds: dict[str, str] = {tid: _compilation_template_kind((cfg or {}).get("kind")) for tid, cfg in active_templates}
|
|
agg_infos: dict[str, dict] = {tid: {"inserted": 0, "updated": 0, "duplicates_dropped": 0} for tid, _ in active_templates}
|
|
chunks_by_id: dict[str, str] = {}
|
|
flush_sequence = 0
|
|
flush_tasks: set[asyncio.Task[None]] = set()
|
|
es_condition = asyncio.Condition()
|
|
next_es_sequence = 0
|
|
|
|
async def _flush(template_id: str) -> None:
|
|
nonlocal flush_sequence
|
|
acc = accumulators[template_id]
|
|
if not acc:
|
|
return
|
|
docs = list(acc)
|
|
acc.clear()
|
|
flush_sequence += 1
|
|
sequence = flush_sequence - 1
|
|
timing_context = f"{doc_id}:{template_id}:flush-{flush_sequence}"
|
|
|
|
async def _run_flush() -> None:
|
|
nonlocal next_es_sequence
|
|
es_acquired = False
|
|
es_released = False
|
|
|
|
async def _wait_for_es() -> None:
|
|
nonlocal es_acquired
|
|
async with es_condition:
|
|
await es_condition.wait_for(lambda: next_es_sequence == sequence)
|
|
es_acquired = True
|
|
|
|
async def _release_es() -> None:
|
|
nonlocal next_es_sequence, es_released
|
|
async with es_condition:
|
|
if next_es_sequence != sequence:
|
|
raise RuntimeError(f"ES sequence mismatch: expected {next_es_sequence}, releasing {sequence}")
|
|
next_es_sequence += 1
|
|
es_released = True
|
|
es_condition.notify_all()
|
|
|
|
kind = template_kinds.get(template_id, "")
|
|
merge_chat_mdl = llm_pool.wrap(
|
|
chat_mdl_by_tid[template_id],
|
|
priority=20,
|
|
label=f"merge:{template_id}",
|
|
context=timing_context,
|
|
)
|
|
try:
|
|
info = await merge_compiled_structures(
|
|
docs,
|
|
merge_chat_mdl,
|
|
embedding_model,
|
|
tenant_id,
|
|
kb_id,
|
|
compilation_template_id=template_id,
|
|
cancel_check=cancel_check,
|
|
timing_context=timing_context,
|
|
chunks_by_id=chunks_by_id,
|
|
chain_kind=kind,
|
|
chain_callback=progress_cb,
|
|
chain_timeout_seconds=STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S,
|
|
es_waiter=_wait_for_es,
|
|
es_releaser=_release_es,
|
|
)
|
|
finally:
|
|
if not es_released:
|
|
if not es_acquired:
|
|
await _wait_for_es()
|
|
await _release_es()
|
|
if isinstance(info, dict):
|
|
agg = agg_infos[template_id]
|
|
for k in ("inserted", "updated", "duplicates_dropped"):
|
|
agg[k] = agg.get(k, 0) + int(info.get(k, 0) or 0)
|
|
|
|
flush_tasks.add(asyncio.create_task(_run_flush()))
|
|
|
|
progress_cb(msg=f"Start document knowledge compilation ({total} template(s)) ...")
|
|
|
|
async def _compile_batch(batch_no: int, batch: list[dict], template_id: str, parser_cfg: dict) -> list[dict]:
|
|
context = f"{doc_id}:{template_id}:compile-batch-{batch_no}"
|
|
progress_cb(msg=f" compile batch {batch_no} ({len(batch)} chunks) for template ({template_ids_by_id[template_id]}/{total})")
|
|
compile_chat_mdl = llm_pool.wrap(
|
|
chat_mdl_by_tid[template_id],
|
|
priority=30,
|
|
label=f"compile:{template_id}:batch-{batch_no}",
|
|
context=context,
|
|
)
|
|
return await compile_structure_from_text(
|
|
batch,
|
|
parser_cfg,
|
|
compile_chat_mdl,
|
|
embedding_model,
|
|
doc_id,
|
|
language=language,
|
|
callback=progress_cb,
|
|
max_workers=3,
|
|
compilation_template_id=template_id,
|
|
)
|
|
|
|
async def _commit_result(batch_no: int, batch_len: int, template_id: str, docs: list[dict]) -> None:
|
|
if docs:
|
|
accumulators[template_id].extend(docs)
|
|
if len(accumulators[template_id]) >= DOC_STRUCTURE_MERGE_MAX_DOCS:
|
|
progress_cb(msg=f" merge flush ({len(accumulators[template_id])} docs) for batch {batch_no} ({batch_len} chunks) for template ({template_ids_by_id[template_id]}/{total})")
|
|
await _flush(template_id)
|
|
|
|
template_ids_by_id = {template_id: idx + 1 for idx, (template_id, _) in enumerate(active_templates)}
|
|
inflight: dict[asyncio.Task[list[dict]], tuple[int, int, int, str]] = {}
|
|
completed: dict[int, tuple[int, int, str, list[dict]]] = {}
|
|
submit_sequence = 0
|
|
commit_sequence = 0
|
|
|
|
async def _commit_ready() -> None:
|
|
nonlocal commit_sequence
|
|
while commit_sequence in completed:
|
|
batch_no, batch_len, template_id, docs = completed.pop(commit_sequence)
|
|
await _commit_result(batch_no, batch_len, template_id, docs)
|
|
commit_sequence += 1
|
|
|
|
async def _cancel_pending() -> None:
|
|
pending = [task for task in (*inflight, *flush_tasks) if not task.done()]
|
|
for task in pending:
|
|
task.cancel()
|
|
if pending:
|
|
await asyncio.gather(*pending, return_exceptions=True)
|
|
inflight.clear()
|
|
flush_tasks.clear()
|
|
|
|
async def _reap_one() -> None:
|
|
if not inflight:
|
|
return
|
|
try:
|
|
done, _ = await asyncio.wait(tuple(inflight), return_when=asyncio.FIRST_COMPLETED)
|
|
except BaseException:
|
|
await _cancel_pending()
|
|
raise
|
|
for task in done:
|
|
sequence, batch_no, batch_len, template_id = inflight.pop(task)
|
|
try:
|
|
docs = task.result()
|
|
except BaseException:
|
|
await _cancel_pending()
|
|
raise
|
|
completed[sequence] = (batch_no, batch_len, template_id, docs)
|
|
await _commit_ready()
|
|
|
|
async def _submit_batches() -> None:
|
|
nonlocal submit_sequence
|
|
batch_no = 0
|
|
try:
|
|
async for batch in chunk_batches:
|
|
batch_no += 1
|
|
for chunk in batch:
|
|
cid = chunk.get("id")
|
|
if isinstance(cid, str) and cid not in chunks_by_id:
|
|
text = chunk.get("content_with_weight") or chunk.get("text") or ""
|
|
chunks_by_id[cid] = text if isinstance(text, str) else ""
|
|
for template_id, parser_cfg in active_templates:
|
|
if cancel_check():
|
|
raise TaskCanceledException("Task was cancelled during document knowledge compilation")
|
|
task = asyncio.create_task(_compile_batch(batch_no, batch, template_id, parser_cfg))
|
|
inflight[task] = (submit_sequence, batch_no, len(batch), template_id)
|
|
submit_sequence += 1
|
|
if len(inflight) + len(completed) >= DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT:
|
|
await _reap_one()
|
|
except BaseException:
|
|
await _cancel_pending()
|
|
raise
|
|
|
|
await _submit_batches()
|
|
|
|
while inflight:
|
|
if cancel_check():
|
|
await _cancel_pending()
|
|
raise TaskCanceledException("Task was cancelled during document knowledge compilation")
|
|
await _reap_one()
|
|
await _commit_ready()
|
|
|
|
for template_id, _ in active_templates:
|
|
if cancel_check():
|
|
await _cancel_pending()
|
|
raise TaskCanceledException("Task was cancelled before merge flush")
|
|
await _flush(template_id)
|
|
if flush_tasks:
|
|
try:
|
|
await asyncio.gather(*flush_tasks)
|
|
except BaseException:
|
|
await _cancel_pending()
|
|
raise
|
|
finally:
|
|
flush_tasks.clear()
|
|
|
|
for idx, (template_id, parser_cfg) in enumerate(active_templates):
|
|
if cancel_check():
|
|
raise TaskCanceledException("Task was cancelled during document knowledge compilation")
|
|
agg = agg_infos[template_id]
|
|
if record:
|
|
record(f"document_structure_compile:{template_id}", agg)
|
|
progress_cb(msg=f"Document knowledge compilation done ({idx + 1}/{total}): {agg}")
|
|
|
|
# ── Synthesis phase ──────────────────────────────────────────────
|
|
# If the template has synthesis.enabled, run wiki PLAN+REFINE
|
|
# to generate output (wiki page, essence paragraph, etc.).
|
|
synthesis_cfg = (parser_cfg or {}).get("synthesis") or {}
|
|
if synthesis_cfg.get("enabled"):
|
|
example = synthesis_cfg.get("example")
|
|
compile_kwd = synthesis_cfg.get("compile_kwd", "artifact_page")
|
|
plan_cfg = synthesis_cfg.get("plan") or {}
|
|
|
|
# Reserved for future wiki_plan_from_reduction extension:
|
|
# entity_type_filter, mention_count_threshold, top_n
|
|
if plan_cfg:
|
|
logging.debug(
|
|
"synthesis: template %s plan config %r reserved for future use",
|
|
template_id,
|
|
plan_cfg,
|
|
)
|
|
|
|
if cancel_check():
|
|
raise TaskCanceledException("Task was cancelled before synthesis PLAN")
|
|
|
|
if not example:
|
|
logging.warning(
|
|
"synthesis: template %s has synthesis.enabled but no example; skipping",
|
|
template_id,
|
|
)
|
|
else:
|
|
try:
|
|
from rag.advanced_rag.knowlege_compile.wiki import (
|
|
wiki_plan_from_reduction,
|
|
wiki_refine_from_plan,
|
|
)
|
|
|
|
progress_cb(msg=f"Synthesis PLAN for template {template_id} (kind={compile_kwd}) ...")
|
|
plan = await wiki_plan_from_reduction(
|
|
chat_mdl=llm_pool.wrap(
|
|
chat_mdl_by_tid[template_id],
|
|
priority=20,
|
|
label=f"synthesis-plan:{template_id}",
|
|
context=f"{doc_id}:{template_id}:synthesis-plan",
|
|
),
|
|
embd_mdl=embedding_model,
|
|
tenant_id=tenant_id,
|
|
kb_id=kb_id,
|
|
callback=progress_cb,
|
|
)
|
|
if cancel_check():
|
|
raise TaskCanceledException("Task was cancelled after synthesis PLAN")
|
|
|
|
if not plan or not plan.get("pages"):
|
|
progress_cb(msg=f"Synthesis: no pages planned for template {template_id}.")
|
|
else:
|
|
progress_cb(msg=f"Synthesis REFINE for template {template_id} ({len(plan['pages'])} page(s)) ...")
|
|
pages = await wiki_refine_from_plan(
|
|
chat_mdl=llm_pool.wrap(
|
|
chat_mdl_by_tid[template_id],
|
|
priority=20,
|
|
label=f"synthesis-refine:{template_id}",
|
|
context=f"{doc_id}:{template_id}:synthesis-refine",
|
|
),
|
|
embd_mdl=embedding_model,
|
|
tenant_id=tenant_id,
|
|
kb_id=kb_id,
|
|
callback=progress_cb,
|
|
example=example,
|
|
)
|
|
# Overwrite compile_kwd on every output page so the
|
|
# synthesis type is tracked correctly in ES.
|
|
for p in pages or []:
|
|
p["compile_kwd"] = compile_kwd
|
|
progress_cb(msg=f"Synthesis done: {len(pages or [])} {compile_kwd} page(s) written.")
|
|
except TaskCanceledException:
|
|
raise
|
|
except Exception:
|
|
logging.exception("synthesis: failed for template %s", template_id)
|
|
|
|
return agg_infos
|