diff --git a/rag/advanced_rag/knowlege_compile/runner.py b/rag/advanced_rag/knowlege_compile/runner.py index b1146a72e1..0bfd4450a6 100644 --- a/rag/advanced_rag/knowlege_compile/runner.py +++ b/rag/advanced_rag/knowlege_compile/runner.py @@ -43,30 +43,33 @@ from api.db.services.compilation_template_group_service import ( from api.db.services.llm_service import LLMBundle from common.exceptions import TaskCanceledException from rag.advanced_rag.knowlege_compile.structure import ( - CHAIN_KINDS, + LLMCallPool, compile_structure_from_text, merge_compiled_structures, - validate_and_correct_chain, ) # ----- tunables ------------------------------------------------------ # Bound how many source chunks are handed to a single -# ``compile_structure_from_text`` invocation. The call fans them out -# across max_workers internally, so a moderate window keeps memory + -# LLM-context pressure predictable for long docs. +# ``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. ``list`` and -# ``timeline`` kinds run this just before each merge flush; anything -# longer than this is treated as a blocked LLM and the uncorrected -# docs are flushed instead. +# Hard wall on the chain-validator LLM correction step. STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S = 120.0 @@ -175,88 +178,203 @@ async def run_structure_compile_over_batches( 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 - kind = template_kinds.get(template_id, "") - if kind in CHAIN_KINDS: - try: - acc = await asyncio.wait_for( - validate_and_correct_chain( - acc, - chunks_by_id, - chat_mdl_by_tid[template_id], - kind, - callback=progress_cb, - ), - timeout=STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S, - ) - accumulators[template_id] = acc - except asyncio.TimeoutError: - logging.warning( - "chain validate: timed out after %ss for template %s; using uncorrected docs", - STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S, - template_id, - ) - except Exception: - logging.exception( - "chain validate: unexpected failure for template %s; using uncorrected docs", - template_id, - ) - info = await merge_compiled_structures( - acc, - chat_mdl_by_tid[template_id], - embedding_model, - tenant_id, - kb_id, - compilation_template_id=template_id, - cancel_check=cancel_check, - ) + docs = list(acc) acc.clear() - 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_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)) ...") - batch_no = 0 - 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 idx, (template_id, parser_cfg) in enumerate(active_templates): - progress_cb(msg=f" compile batch {batch_no} ({len(batch)} chunks) for template ({idx + 1}/{total})") - docs = await compile_structure_from_text( - batch, - parser_cfg, - chat_mdl_by_tid[template_id], - embedding_model, - doc_id, - language=language, - callback=progress_cb, - compilation_template_id=template_id, - ) - 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 template ({idx + 1}/{total})") - await _flush(template_id) + 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") - await _flush(template_id) agg = agg_infos[template_id] if record: record(f"document_structure_compile:{template_id}", agg) @@ -276,7 +394,8 @@ async def run_structure_compile_over_batches( if plan_cfg: logging.debug( "synthesis: template %s plan config %r reserved for future use", - template_id, plan_cfg, + template_id, + plan_cfg, ) if cancel_check(): @@ -296,7 +415,12 @@ async def run_structure_compile_over_batches( progress_cb(msg=f"Synthesis PLAN for template {template_id} (kind={compile_kwd}) ...") plan = await wiki_plan_from_reduction( - chat_mdl=chat_mdl_by_tid[template_id], + 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, @@ -310,7 +434,12 @@ async def run_structure_compile_over_batches( else: progress_cb(msg=f"Synthesis REFINE for template {template_id} ({len(plan['pages'])} page(s)) ...") pages = await wiki_refine_from_plan( - chat_mdl=chat_mdl_by_tid[template_id], + 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, diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index e60f82e178..32f62ba4a6 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -14,9 +14,11 @@ # limitations under the License. # import datetime +import asyncio +import heapq import json import logging -from typing import Callable, Tuple +from typing import Awaitable, Callable, Tuple import xxhash @@ -38,6 +40,72 @@ from ._common import ( _STRUCT_TYPES = ("list", "set", "hypergraph") +_ES_DEDUP_KNN_CONCURRENCY = 8 +_ES_DEDUP_LLM_CONCURRENCY = 16 +_ES_DEDUP_LLM_BATCH_SIZE = 16 +_ES_DEDUP_EMBED_BATCH_SIZE = 64 +_ES_DEDUP_INSERT_BATCH_SIZE = 256 + + +class LLMCallPool: + """Task-scoped priority scheduler for actual chat-model calls.""" + + def __init__(self, max_concurrency: int = 10): + self.max_concurrency = max(1, int(max_concurrency)) + self._active = 0 + self._ticket = 0 + self._waiting: list[tuple[int, int]] = [] + self._condition = asyncio.Condition() + + def wrap(self, chat_mdl, *, priority: int, label: str, context: str | None = None): + return PooledChatModel(self, chat_mdl, priority=priority, label=label, context=context) + + async def call(self, fn, *, priority: int, label: str, context: str | None = None): + async with self._condition: + ticket = (int(priority), self._ticket) + self._ticket += 1 + heapq.heappush(self._waiting, ticket) + try: + while self._active >= self.max_concurrency or self._waiting[0] != ticket: + await self._condition.wait() + except BaseException: + if ticket in self._waiting: + self._waiting.remove(ticket) + heapq.heapify(self._waiting) + self._condition.notify_all() + raise + heapq.heappop(self._waiting) + self._active += 1 + try: + result = await fn() + return result + except BaseException: + raise + finally: + async with self._condition: + self._active -= 1 + self._condition.notify_all() + + +class PooledChatModel: + def __init__(self, pool: LLMCallPool, chat_mdl, *, priority: int, label: str, context: str | None): + self._pool = pool + self._chat_mdl = chat_mdl + self._priority = priority + self._label = label + self._context = context + + def __getattr__(self, name): + return getattr(self._chat_mdl, name) + + async def async_chat(self, system, history, gen_conf=None, **kwargs): + return await self._pool.call( + lambda: self._chat_mdl.async_chat(system, history, gen_conf=gen_conf, **kwargs), + priority=self._priority, + label=self._label, + context=self._context, + ) + def _struct_normalize_kind(kind) -> str: if not isinstance(kind, str): @@ -781,6 +849,58 @@ def _struct_union_chunk_ids(*chunk_id_lists) -> list: return _union_ordered(*normalized) +def _struct_entity_name(doc_or_payload: dict) -> str: + value = doc_or_payload.get("name") if isinstance(doc_or_payload, dict) else None + if value is None and isinstance(doc_or_payload, dict): + try: + value = json.loads(doc_or_payload.get("content_with_weight") or "{}").get("name") + except Exception: + value = None + return str(value).strip() if value is not None else "" + + +def _struct_resolve_entity_alias(name: str, aliases: dict[str, str]) -> str: + current = str(name).strip() + seen = set() + while current in aliases and current not in seen: + seen.add(current) + current = aliases[current] + return current + + +def _struct_rewrite_relation_payload(payload: dict, aliases: dict[str, str]) -> bool: + changed = False + for fields in (("source", "src", "from"), ("target", "tgt", "to")): + for field in fields: + if field not in payload or payload[field] is None: + continue + old = str(payload[field]).strip() + new = _struct_resolve_entity_alias(old, aliases) + if new != old: + payload[field] = new + changed = True + return changed + + +async def _struct_rewrite_relation_doc(doc: dict, aliases: dict[str, str], embd_mdl) -> dict: + if doc.get("knowledge_graph_kwd") != "relation" or not aliases: + return doc + try: + payload = json.loads(doc.get("content_with_weight") or "{}") + except Exception: + return doc + if not isinstance(payload, dict) or not _struct_rewrite_relation_payload(payload, aliases): + return doc + vecs = await _struct_embed(embd_mdl, [_struct_payload_description(payload)]) + if not vecs: + return doc + base = dict(doc) + base["content_with_weight"] = json.dumps(payload, ensure_ascii=False) + base["from_entity_kwd"] = _struct_resolve_entity_alias(base.get("from_entity_kwd", ""), aliases) + base["to_entity_kwd"] = _struct_resolve_entity_alias(base.get("to_entity_kwd", ""), aliases) + return _struct_rebuild_es_doc(payload, base, vecs[0], doc.get("source_chunk_ids") or [], preserve_id=True) + + async def _struct_merge_pair(existing: dict, incoming: dict, chat_mdl) -> dict | None: """LLM-judged merge. Returns merged payload dict if duplicate, else None. @@ -864,6 +984,8 @@ def _struct_rebuild_es_doc( kind=kind, src_field=src_field, target_field=target_field, + compilation_template_id=_struct_doc_template_id(base_doc), + compilation_template_kind=base_doc.get("compilation_template_kind_kwd"), ) if preserve_id and base_doc.get("id"): new_doc["id"] = base_doc["id"] @@ -881,11 +1003,615 @@ async def _struct_reembed_payload(payload: dict, embd_mdl): return vecs[0] if vecs else None +def _struct_es_dedup_condition(doc: dict) -> dict: + condition = { + "compile_kwd": [doc["compile_kwd"]], + "doc_id": [doc["doc_id"]], + } + if doc.get("knowledge_graph_kwd"): + condition["knowledge_graph_kwd"] = [doc["knowledge_graph_kwd"]] + if doc.get("from_entity_kwd"): + condition["from_entity_kwd"] = [doc["from_entity_kwd"]] + if doc.get("to_entity_kwd"): + condition["to_entity_kwd"] = [doc["to_entity_kwd"]] + template_id = _struct_doc_template_id(doc) + if template_id: + condition["compilation_template_ids"] = [template_id] + return condition + + +async def _struct_es_knn_candidate( + doc: dict, + tenant_id: str, + kb_id: str, + similarity_threshold: float, + index: str, + select_fields: list[str], + timing_context: str | None, + item_index: int, +) -> dict | None: + """Run one KNN lookup; the caller controls concurrency.""" + from common import settings + from common.doc_store.doc_store_base import MatchDenseExpr, OrderByExpr + + vec_field, vec = _struct_doc_vec(doc) + if not vec_field or vec is None: + return None + match_expr = MatchDenseExpr( + vector_column_name=vec_field, + embedding_data=list(vec), + embedding_data_type="float", + distance_type="cosine", + topn=1, + extra_options={"similarity": similarity_threshold}, + ) + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + _struct_es_dedup_condition(doc), + [match_expr], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) + if not field_map: + return None + old_id, old_doc = next(iter(field_map.items())) + old_doc = dict(old_doc) + old_doc.setdefault("id", old_id) + return old_doc + except Exception: + logging.exception("merge_compiled_structures: ES KNN search failed; treating doc as new") + return None + + +ES_GROUP_MERGE_PROMPT = """Existing item: +{existing} + +Incoming items: +{incoming} + +Decide which incoming items refer to the same logical entity or relation as +the existing item. Merge all duplicated incoming items with the existing item. +Incoming items that are not duplicates must remain separate. Do not invent +data and do not merge unrelated incoming items with each other. + +Return ONLY JSON with this exact shape: +{{ + "duplicate_indices": [, ...], + "merged": +}} +""" + +ES_GROUP_BATCH_MERGE_PROMPT = """You are judging multiple independent ES deduplication groups. + +For every group, compare every incoming item with that group's existing item. +You must make a separate duplicated decision for every incoming item. Only +incoming items marked duplicated=true may contribute to that group's merged +payload. Incoming items marked duplicated=false must remain separate. Do not +merge items from different groups and do not invent data. + +Return ONLY JSON with this exact shape: +{{ + "groups": [ + {{ + "group_id": "", + "decisions": [ + {{"incoming_index": 0, "duplicated": true}}, + {{"incoming_index": 1, "duplicated": false}} + ], + "merged": + }} + ] +}} + +Groups: +{groups} +""" + +ES_GROUP_DECISION_BATCH_PROMPT = """You are judging multiple independent ES deduplication groups. + +For every incoming item, independently decide whether it is a duplicate of +the existing item in the same group. Do not merge anything and do not judge +items from different groups against each other. + +Return ONLY JSON with this exact shape: +{{ + "groups": [ + {{ + "group_id": "", + "decisions": [ + {{"incoming_index": 0, "duplicated": true}}, + {{"incoming_index": 1, "duplicated": false}} + ] + }} + ] +}} + +Groups: +{groups} +""" + + +async def _struct_judge_es_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, set[int]]: + """Judge every incoming item independently without generating a merge.""" + prompt_groups = [] + for spec in group_specs: + try: + existing_payload = json.loads(spec["old_doc"].get("content_with_weight") or "{}") + incoming_payloads = [json.loads(d.get("content_with_weight") or "{}") for d in spec["incoming_docs"]] + except Exception: + logging.exception("merge: failed to parse ES decision group") + continue + if not isinstance(existing_payload, dict) or not all(isinstance(p, dict) for p in incoming_payloads): + continue + prompt_groups.append( + { + "group_id": spec["request_group_id"], + "existing": existing_payload, + "incoming": [{"index": i, "item": payload} for i, payload in enumerate(incoming_payloads)], + } + ) + if not prompt_groups: + return {spec["request_group_id"]: set() for spec in group_specs} + + user_prompt = ES_GROUP_DECISION_BATCH_PROMPT.format(groups=json.dumps(prompt_groups, ensure_ascii=False)) + system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + ES_GROUP_DECISION_BATCH_PROMPT.split("Groups:", 1)[0] + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + raw_groups = res.get("groups") if isinstance(res, dict) else None + if not isinstance(raw_groups, list): + return {spec["request_group_id"]: set() for spec in group_specs} + + by_id = {spec["request_group_id"]: spec for spec in group_specs} + result: dict[str, set[int]] = {} + for raw in raw_groups: + if not isinstance(raw, dict) or raw.get("group_id") not in by_id: + continue + spec = by_id[raw["group_id"]] + decisions = raw.get("decisions") + if not isinstance(decisions, list): + result[spec["request_group_id"]] = set() + continue + result[spec["request_group_id"]] = { + item["incoming_index"] + for item in decisions + if isinstance(item, dict) and item.get("duplicated") is True and isinstance(item.get("incoming_index"), int) and 0 <= item["incoming_index"] < len(spec["incoming_docs"]) + } + for spec in group_specs: + result.setdefault(spec["request_group_id"], set()) + return result + + +async def _struct_merge_es_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, tuple[list[dict], dict | None]]: + """Judge multiple old_id groups in one LLM request.""" + prompt_groups = [] + for spec in group_specs: + old_doc = spec["old_doc"] + incoming_docs = spec["incoming_docs"] + try: + existing_payload = json.loads(old_doc.get("content_with_weight") or "{}") + incoming_payloads = [json.loads(d.get("content_with_weight") or "{}") for d in incoming_docs] + except Exception: + logging.exception("merge: failed to parse grouped content_with_weight") + continue + if not isinstance(existing_payload, dict) or not all(isinstance(p, dict) for p in incoming_payloads): + continue + prompt_groups.append( + { + "group_id": spec["old_id"], + "existing": existing_payload, + "incoming": [{"index": i, "item": payload} for i, payload in enumerate(incoming_payloads)], + } + ) + if not prompt_groups: + return {spec["old_id"]: (list(spec["incoming_docs"]), None) for spec in group_specs} + + user_prompt = ES_GROUP_BATCH_MERGE_PROMPT.format(groups=json.dumps(prompt_groups, ensure_ascii=False)) + system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + ES_GROUP_BATCH_MERGE_PROMPT.split("Groups:", 1)[0] + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + raw_groups = res.get("groups") if isinstance(res, dict) else None + if not isinstance(raw_groups, list): + return {spec["old_id"]: (list(spec["incoming_docs"]), None) for spec in group_specs} + + result = {} + by_id = {spec["old_id"]: spec for spec in group_specs} + for raw in raw_groups: + if not isinstance(raw, dict) or raw.get("group_id") not in by_id: + continue + spec = by_id[raw["group_id"]] + decisions = raw.get("decisions") + merged = raw.get("merged") + if not isinstance(decisions, list): + result[spec["old_id"]] = (list(spec["incoming_docs"]), None) + continue + duplicate_indices = {item.get("incoming_index") for item in decisions if isinstance(item, dict) and item.get("duplicated") is True and isinstance(item.get("incoming_index"), int)} + duplicate_indices = {i for i in duplicate_indices if 0 <= i < len(spec["incoming_docs"])} + if not duplicate_indices or not isinstance(merged, dict): + result[spec["old_id"]] = (list(spec["incoming_docs"]), None) + continue + separate = [d for i, d in enumerate(spec["incoming_docs"]) if i not in duplicate_indices] + result[spec["old_id"]] = (separate, merged) + + for spec in group_specs: + result.setdefault(spec["old_id"], (list(spec["incoming_docs"]), None)) + return result + + +async def _struct_merge_es_group(old_doc: dict, incoming_docs: list[dict], chat_mdl) -> tuple[list[dict], dict | None]: + """Judge one ES candidate group with one LLM request. + + Returns ``(non_duplicate_docs, merged_payload)``. The existing ES row is + updated only when ``merged_payload`` is a dict. + """ + if len(incoming_docs) == 1: + merged = await _struct_merge_pair(old_doc, incoming_docs[0], chat_mdl) + return ([] if merged is not None else list(incoming_docs), merged) + + try: + existing_payload = json.loads(old_doc.get("content_with_weight") or "{}") + incoming_payloads = [json.loads(d.get("content_with_weight") or "{}") for d in incoming_docs] + except Exception: + logging.exception("merge: failed to parse grouped content_with_weight") + return list(incoming_docs), None + if not isinstance(existing_payload, dict) or not all(isinstance(p, dict) for p in incoming_payloads): + return list(incoming_docs), None + + system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + ES_GROUP_MERGE_PROMPT + user_prompt = ES_GROUP_MERGE_PROMPT.format( + existing=json.dumps(existing_payload, ensure_ascii=False), + incoming=json.dumps( + [{"index": i, "item": payload} for i, payload in enumerate(incoming_payloads)], + ensure_ascii=False, + ), + ) + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + if not isinstance(res, dict): + return list(incoming_docs), None + indices = res.get("duplicate_indices") + merged = res.get("merged") + if not isinstance(indices, list) or not isinstance(merged, dict): + return list(incoming_docs), None + duplicate_indices = {i for i in indices if isinstance(i, int) and 0 <= i < len(incoming_docs)} + if not duplicate_indices: + return list(incoming_docs), None + separate = [d for i, d in enumerate(incoming_docs) if i not in duplicate_indices] + return separate, merged + + +async def _struct_es_dedup_batch( + docs: list[dict], + chat_mdl, + embd_mdl, + tenant_id: str, + kb_id: str, + similarity_threshold: float, + timing_context: str | None = None, + cancel_check: Callable[[], bool] | None = None, +) -> tuple[int, int]: + """Batch ES dedup: concurrent KNN, parallel decisions, then grouped merges.""" + from common import settings + from rag.nlp import search as _rag_search + + index = _rag_search.index_name(tenant_id) + + def _raise_if_canceled() -> None: + if callable(cancel_check) and cancel_check(): + raise TaskCanceledException("Task was cancelled during ES dedup") + + select_fields = [ + "id", + "content_with_weight", + "source_chunk_ids", + "knowledge_graph_kwd", + "compile_kwd", + "doc_id", + "from_entity_kwd", + "to_entity_kwd", + "compilation_template_ids", + "compilation_template_kind_kwd", + ] + + # One semaphore shared by all tasks; constructing it per task would not + # limit concurrency. + knn_semaphore = asyncio.Semaphore(_ES_DEDUP_KNN_CONCURRENCY) + + async def run_knn_shared(item_index: int, doc: dict): + _raise_if_canceled() + async with knn_semaphore: + return doc, await _struct_es_knn_candidate( + doc, + tenant_id, + kb_id, + similarity_threshold, + index, + select_fields, + timing_context, + item_index, + ) + + _raise_if_canceled() + knn_results = await asyncio.gather(*(run_knn_shared(i, d) for i, d in enumerate(docs))) + _raise_if_canceled() + groups: dict[str, tuple[dict, list[dict]]] = {} + inserts: list[dict] = [] + for doc, old_doc in knn_results: + if old_doc is None: + inserts.append(doc) + continue + old_id = str(old_doc["id"]) + if old_id not in groups: + groups[old_id] = (old_doc, []) + groups[old_id][1].append(doc) + + # Stage 1 is deliberately read-only: every decision for an old_id uses + # the same KNN snapshot. This lets sub-batches of one large group run in + # parallel without one completed request changing the input of another. + states = { + old_id: { + "old_doc": old_doc, + "incoming_docs": incoming, + "separate": [], + "duplicate_docs": [], + "merged": None, + "chunk_ids": list(old_doc.get("source_chunk_ids") or []), + "entity_aliases": {}, + } + for old_id, (old_doc, incoming) in groups.items() + } + entity_aliases: dict[str, str] = {} + llm_semaphore = asyncio.Semaphore(_ES_DEDUP_LLM_CONCURRENCY) + + decision_specs = [] + for old_id, state in states.items(): + incoming_docs = state["incoming_docs"] + for part, start in enumerate(range(0, len(incoming_docs), _ES_DEDUP_LLM_BATCH_SIZE)): + decision_specs.append( + { + "old_id": old_id, + "request_group_id": f"{old_id}:part-{part}", + "old_doc": state["old_doc"], + "incoming_docs": incoming_docs[start : start + _ES_DEDUP_LLM_BATCH_SIZE], + } + ) + + decision_batches = [] + current_batch = [] + current_size = 0 + for spec in decision_specs: + size = len(spec["incoming_docs"]) + if current_batch and current_size + size > _ES_DEDUP_LLM_BATCH_SIZE: + decision_batches.append(current_batch) + current_batch = [] + current_size = 0 + current_batch.append(spec) + current_size += size + if current_batch: + decision_batches.append(current_batch) + + async def run_decision_batch(batch_no: int, batch_specs: list[dict]): + _raise_if_canceled() + async with llm_semaphore: + try: + result = await _struct_judge_es_group_batch(batch_specs, chat_mdl) + except Exception: + logging.exception("merge_compiled_structures: ES decision batch failed") + result = {spec["request_group_id"]: set() for spec in batch_specs} + return result + + decision_results = await asyncio.gather( + *(run_decision_batch(i, batch) for i, batch in enumerate(decision_batches)), + ) + _raise_if_canceled() + for batch, result in zip(decision_batches, decision_results): + for spec in batch: + state = states[spec["old_id"]] + duplicate_indices = result.get(spec["request_group_id"], set()) + for incoming_index, doc in enumerate(spec["incoming_docs"]): + if incoming_index in duplicate_indices: + state["duplicate_docs"].append(doc) + else: + state["separate"].append(doc) + + async def merge_one_group(old_id: str, state: dict): + _raise_if_canceled() + duplicate_docs = state["duplicate_docs"] + if not duplicate_docs: + return + old_doc = state["old_doc"] + current_doc = dict(old_doc) + current_chunk_ids = list(state["chunk_ids"]) + merged_payload = None + # A normal group gets exactly one merge request. Only pathological + # groups are folded in <= batch-sized sequential pieces. + for start in range(0, len(duplicate_docs), _ES_DEDUP_LLM_BATCH_SIZE): + _raise_if_canceled() + candidate_docs = duplicate_docs[start : start + _ES_DEDUP_LLM_BATCH_SIZE] + separate, candidate_merged = await _struct_merge_es_group(current_doc, candidate_docs, chat_mdl) + state["separate"].extend(separate) + if candidate_merged is None: + continue + candidate_merged = _struct_apply_merge_invariants(current_doc, candidate_merged) + if old_doc.get("knowledge_graph_kwd") == "entity": + old_name = _struct_entity_name(current_doc) + canonical_name = _struct_entity_name(candidate_merged) or old_name + for candidate in candidate_docs: + candidate_name = _struct_entity_name(candidate) + if candidate_name and candidate_name != canonical_name: + state["entity_aliases"][candidate_name] = canonical_name + if old_name and old_name != canonical_name: + state["entity_aliases"][old_name] = canonical_name + separate_ids = {id(doc) for doc in separate} + current_chunk_ids = _struct_union_chunk_ids( + current_chunk_ids, + *(d.get("source_chunk_ids") for d in candidate_docs if id(d) not in separate_ids), + ) + current_doc["content_with_weight"] = json.dumps(candidate_merged, ensure_ascii=False) + current_doc["source_chunk_ids"] = current_chunk_ids + merged_payload = candidate_merged + if merged_payload is not None: + state["merged"] = merged_payload + state["chunk_ids"] = current_chunk_ids + + merge_jobs = [merge_one_group(old_id, state) for old_id, state in states.items() if state["duplicate_docs"]] + await asyncio.gather(*merge_jobs) + _raise_if_canceled() + + existing_relation_updates = 0 + + merged_jobs = [] + for old_id, state in states.items(): + separate_docs = state["separate"] + inserts.extend(separate_docs) + if state["merged"] is None: + continue + merged_jobs.append( + { + "old_id": old_id, + "old_doc": state["old_doc"], + "payload": state["merged"], + "chunk_ids": state["chunk_ids"], + "entity_aliases": dict(state.get("entity_aliases") or {}), + } + ) + + # Encode all merged groups in batches, independent of the LLM grouping. + for start in range(0, len(merged_jobs), _ES_DEDUP_EMBED_BATCH_SIZE): + batch = merged_jobs[start : start + _ES_DEDUP_EMBED_BATCH_SIZE] + texts = [_struct_payload_description(job["payload"]) for job in batch] + try: + vectors = await _struct_embed(embd_mdl, texts) + except Exception: + logging.exception("merge_compiled_structures: grouped embedding failed for %d docs", len(batch)) + vectors = [] + for job, vec in zip(batch, vectors): + job["rebuilt"] = _struct_rebuild_es_doc( + job["payload"], + job["old_doc"], + vec, + job["chunk_ids"], + preserve_id=True, + ) + + updated_jobs = [job for job in merged_jobs if job.get("rebuilt")] + writes = inserts + [job["rebuilt"] for job in updated_jobs] + inserted = 0 + updated = 0 + successful_entity_aliases: dict[str, str] = {} + for start in range(0, len(writes), _ES_DEDUP_INSERT_BATCH_SIZE): + _raise_if_canceled() + batch = writes[start : start + _ES_DEDUP_INSERT_BATCH_SIZE] + if not batch: + continue + try: + await thread_pool_exec(settings.docStoreConn.insert, batch, index, kb_id) + updated_in_batch = sum(1 for doc in batch if any(doc is job.get("rebuilt") for job in updated_jobs)) + updated += updated_in_batch + inserted += len(batch) - updated_in_batch + for job in updated_jobs: + if job.get("rebuilt") not in batch: + continue + if job["old_doc"].get("knowledge_graph_kwd") == "entity": + successful_entity_aliases.update(job.get("entity_aliases") or {}) + except Exception: + logging.exception("merge_compiled_structures: bulk insert failed for %d docs", len(batch)) + + # Only publish aliases after the canonical entity writes have completed. + entity_aliases.update(successful_entity_aliases) + if entity_aliases: + relation_fields = [ + "id", + "content_with_weight", + "source_chunk_ids", + "knowledge_graph_kwd", + "compile_kwd", + "doc_id", + "from_entity_kwd", + "to_entity_kwd", + "compilation_template_ids", + "compilation_template_kind_kwd", + ] + from common.doc_store.doc_store_base import OrderByExpr + + scopes = { + ( + state["old_doc"].get("doc_id"), + state["old_doc"].get("compile_kwd"), + _struct_doc_template_id(state["old_doc"]), + ) + for state in states.values() + if state["old_doc"].get("knowledge_graph_kwd") == "entity" + } + for doc_id, compile_kwd, template_id in scopes: + condition = { + "doc_id": [doc_id], + "compile_kwd": [compile_kwd], + "knowledge_graph_kwd": ["relation"], + } + if template_id: + condition["compilation_template_ids"] = [template_id] + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + relation_fields, + [], + condition, + [], + OrderByExpr(), + 0, + 10000, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, relation_fields) + except Exception: + logging.exception("merge_compiled_structures: relation reference search failed") + continue + rewrite_batch = [] + for row_id, row in rows.items(): + payload = _struct_load_payload(row) + if not isinstance(payload, dict) or not _struct_rewrite_relation_payload(payload, entity_aliases): + continue + base = dict(row) + base["id"] = row_id + base["content_with_weight"] = json.dumps(payload, ensure_ascii=False) + base["from_entity_kwd"] = _struct_resolve_entity_alias(base.get("from_entity_kwd", ""), entity_aliases) + base["to_entity_kwd"] = _struct_resolve_entity_alias(base.get("to_entity_kwd", ""), entity_aliases) + rewrite_batch.append((base, payload)) + for start in range(0, len(rewrite_batch), _ES_DEDUP_EMBED_BATCH_SIZE): + batch = rewrite_batch[start : start + _ES_DEDUP_EMBED_BATCH_SIZE] + vectors = await _struct_embed(embd_mdl, [_struct_payload_description(payload) for _, payload in batch]) + rewritten = [_struct_rebuild_es_doc(payload, base, vector, base.get("source_chunk_ids") or [], preserve_id=True) for (base, payload), vector in zip(batch, vectors)] + if rewritten: + await thread_pool_exec(settings.docStoreConn.insert, rewritten, index, kb_id) + existing_relation_updates += len(rewritten) + rewritten_inserts = [await _struct_rewrite_relation_doc(doc, entity_aliases, embd_mdl) if doc.get("knowledge_graph_kwd") == "relation" else doc for doc in inserts] + if rewritten_inserts != inserts: + await thread_pool_exec(settings.docStoreConn.insert, rewritten_inserts, index, kb_id) + for job in merged_jobs: + if job["old_doc"].get("knowledge_graph_kwd") != "relation": + continue + if not _struct_rewrite_relation_payload(job["payload"], entity_aliases): + continue + vector = await _struct_reembed_payload(job["payload"], embd_mdl) + if vector is not None: + rewritten = _struct_rebuild_es_doc(job["payload"], job["old_doc"], vector, job["chunk_ids"], preserve_id=True) + await thread_pool_exec(settings.docStoreConn.insert, [rewritten], index, kb_id) + return inserted, updated + existing_relation_updates + + async def _struct_local_dedup( docs: list[dict], chat_mdl, embd_mdl, similarity_threshold: float, + timing_context: str | None = None, + rewrite_relations: bool = True, + return_aliases: bool = False, ) -> tuple[list[dict], int]: """Single-pass dedup inside ``docs``. Returns (deduped, dropped_count).""" from sklearn.metrics.pairwise import cosine_similarity @@ -901,11 +1627,12 @@ async def _struct_local_dedup( dropped = 0 deduped: list[dict] = [] + entity_aliases: dict[str, str] = {} - for key in order: + for group_index, key in enumerate(order): kept: list[dict] = [] - for incoming in groups[key]: - inc_field, inc_vec = _struct_doc_vec(incoming) + for incoming_index, incoming in enumerate(groups[key]): + _, inc_vec = _struct_doc_vec(incoming) if not inc_vec or not kept: kept.append(incoming) continue @@ -920,14 +1647,22 @@ async def _struct_local_dedup( sims = cosine_similarity([list(inc_vec)], [list(v) for _, v in kept_with_vecs])[0] sims_list = sims.tolist() if hasattr(sims, "tolist") else list(sims) best_idx = max(range(len(sims_list)), key=lambda i: sims_list[i]) - if sims_list[best_idx] < similarity_threshold: + best_score = float(sims_list[best_idx]) + existing = kept_with_vecs[best_idx][0] + if best_score < similarity_threshold: kept.append(incoming) continue - existing = kept_with_vecs[best_idx][0] merged_payload = await _struct_merge_pair(existing, incoming, chat_mdl) if merged_payload is None: kept.append(incoming) continue + if existing.get("knowledge_graph_kwd") == "entity": + old_name = _struct_entity_name(existing) + incoming_name = _struct_entity_name(incoming) + canonical_name = _struct_entity_name(merged_payload) or old_name + for alias in (old_name, incoming_name): + if alias and alias != canonical_name: + entity_aliases[alias] = canonical_name merged_payload = _struct_apply_merge_invariants(existing, merged_payload) merged_chunk_ids = _struct_union_chunk_ids( existing.get("source_chunk_ids"), @@ -953,128 +1688,135 @@ async def _struct_local_dedup( dropped += 1 deduped.extend(kept) + if rewrite_relations and entity_aliases: + entity_docs = [d for d in deduped if d.get("knowledge_graph_kwd") != "relation"] + relation_docs = [d for d in deduped if d.get("knowledge_graph_kwd") == "relation"] + rewritten_relations = [await _struct_rewrite_relation_doc(d, entity_aliases, embd_mdl) for d in relation_docs] + relation_deduped, relation_dropped = await _struct_local_dedup( + rewritten_relations, + chat_mdl, + embd_mdl, + similarity_threshold, + timing_context=timing_context, + rewrite_relations=False, + ) + deduped = entity_docs + relation_deduped + dropped += relation_dropped + + if return_aliases: + return deduped, dropped, entity_aliases return deduped, dropped -async def _struct_es_dedup_one( - doc: dict, +_LOCAL_DEDUP_GROUP_CONCURRENCY = 8 + + +def _struct_entity_candidate_groups(docs: list[dict], similarity_threshold: float) -> list[list[dict]]: + """Partition entity candidates into independent cosine-connected groups.""" + from sklearn.metrics.pairwise import cosine_similarity + + buckets: dict[tuple, list[dict]] = {} + order: list[tuple] = [] + for doc in docs: + key = _struct_filter_key(doc) + if key not in buckets: + buckets[key] = [] + order.append(key) + buckets[key].append(doc) + + result: list[list[dict]] = [] + for key in order: + bucket = buckets[key] + vectors = [_struct_doc_vec(doc)[1] for doc in bucket] + valid = [i for i, vector in enumerate(vectors) if vector] + parent = list(range(len(bucket))) + + def find(index: int) -> int: + while parent[index] != index: + parent[index] = parent[parent[index]] + index = parent[index] + return index + + def union(left: int, right: int) -> None: + left_root, right_root = find(left), find(right) + if left_root != right_root: + parent[right_root] = left_root + + if len(valid) > 1: + matrix = cosine_similarity([list(vectors[i]) for i in valid]) + for left_offset, left in enumerate(valid): + for right_offset in range(left_offset + 1, len(valid)): + right = valid[right_offset] + if float(matrix[left_offset, right_offset]) >= similarity_threshold: + union(left, right) + + components: dict[int, list[dict]] = {} + component_order: list[int] = [] + for index, doc in enumerate(bucket): + root = find(index) if index in valid else index + if root not in components: + components[root] = [] + component_order.append(root) + components[root].append(doc) + result.extend(components[root] for root in component_order) + return result + + +async def _struct_local_dedup_parallel( + docs: list[dict], chat_mdl, embd_mdl, - tenant_id: str, - kb_id: str, similarity_threshold: float, -) -> str: - """Persist a single doc into ES with merge-or-insert semantics. + timing_context: str | None = None, +) -> tuple[list[dict], int]: + """Deduplicate entities and relations in dependency order with group concurrency.""" + if not docs: + return [], 0 - Returns one of: 'inserted', 'updated', 'skipped'. - """ - from common import settings - from rag.nlp import search as _rag_search - from common.doc_store.doc_store_base import MatchDenseExpr, OrderByExpr + entity_docs = [doc for doc in docs if doc.get("knowledge_graph_kwd") != "relation"] + relation_docs = [doc for doc in docs if doc.get("knowledge_graph_kwd") == "relation"] + entity_groups = _struct_entity_candidate_groups(entity_docs, similarity_threshold) + group_semaphore = asyncio.Semaphore(_LOCAL_DEDUP_GROUP_CONCURRENCY) - index = _rag_search.index_name(tenant_id) + async def dedup_group(group: list[dict]): + async with group_semaphore: + return await _struct_local_dedup( + group, + chat_mdl, + embd_mdl, + similarity_threshold, + timing_context=timing_context, + rewrite_relations=False, + return_aliases=True, + ) - condition = { - "compile_kwd": [doc["compile_kwd"]], - "doc_id": [doc["doc_id"]], - } - if doc.get("knowledge_graph_kwd"): - condition["knowledge_graph_kwd"] = [doc["knowledge_graph_kwd"]] - if doc.get("from_entity_kwd"): - condition["from_entity_kwd"] = [doc["from_entity_kwd"]] - if doc.get("to_entity_kwd"): - condition["to_entity_kwd"] = [doc["to_entity_kwd"]] - # KNN dedup must stay within the same template — two templates can - # produce identical-looking entities (e.g. two "list" kinds extracting - # the same headline) but they live on independent tabs in the UI. - incoming_template = _struct_doc_template_id(doc) - if incoming_template: - condition["compilation_template_ids"] = [incoming_template] + entity_results = await asyncio.gather(*(dedup_group(group) for group in entity_groups)) + deduped_entities: list[dict] = [] + entity_aliases: dict[str, str] = {} + dropped = 0 + for entity_result, group in zip(entity_results, entity_groups): + group_docs, group_dropped, group_aliases = entity_result + deduped_entities.extend(group_docs) + dropped += group_dropped + entity_aliases.update(group_aliases) - vec_field, vec = _struct_doc_vec(doc) - if not vec_field or vec is None: - await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id) - return "inserted" + rewritten_relations = await asyncio.gather(*(_struct_rewrite_relation_doc(doc, entity_aliases, embd_mdl) for doc in relation_docs)) + relation_buckets: dict[tuple, list[dict]] = {} + relation_order: list[tuple] = [] + for doc in rewritten_relations: + key = _struct_filter_key(doc) + if key not in relation_buckets: + relation_buckets[key] = [] + relation_order.append(key) + relation_buckets[key].append(doc) - match_expr = MatchDenseExpr( - vector_column_name=vec_field, - embedding_data=list(vec), - embedding_data_type="float", - distance_type="cosine", - topn=1, - extra_options={"similarity": similarity_threshold}, - ) - select_fields = [ - "id", - "content_with_weight", - "source_chunk_ids", - "knowledge_graph_kwd", - "compile_kwd", - "doc_id", - "from_entity_kwd", - "to_entity_kwd", - ] - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - condition, - [match_expr], - OrderByExpr(), - 0, - 1, - index, - [kb_id], - ) - field_map = settings.docStoreConn.get_fields(res, select_fields) - except Exception: - logging.exception("merge_compiled_structures: ES KNN search failed; inserting as new") - await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id) - return "inserted" - - if not field_map: - await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id) - return "inserted" - - old_id, old_doc = next(iter(field_map.items())) - old_doc = dict(old_doc) - old_doc.setdefault("id", old_id) - - merged_payload = await _struct_merge_pair(old_doc, doc, chat_mdl) - if merged_payload is None: - await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id) - return "inserted" - - merged_payload = _struct_apply_merge_invariants(old_doc, merged_payload) - merged_chunk_ids = _struct_union_chunk_ids( - old_doc.get("source_chunk_ids"), - doc.get("source_chunk_ids"), - ) - new_vec = await _struct_reembed_payload(merged_payload, embd_mdl) - if new_vec is None: - return "skipped" - - rebuilt = _struct_rebuild_es_doc( - merged_payload, - old_doc, - new_vec, - merged_chunk_ids, - preserve_id=True, - ) - update_fields = {k: v for k, v in rebuilt.items() if k != "id"} - try: - await thread_pool_exec( - settings.docStoreConn.update, - {"id": old_id}, - update_fields, - index, - kb_id, - ) - return "updated" - except Exception: - logging.exception("merge_compiled_structures: ES update failed for id %s", old_id) - return "skipped" + relation_results = await asyncio.gather(*(dedup_group(relation_buckets[key]) for key in relation_order)) + deduped_relations: list[dict] = [] + for relation_result in relation_results: + group_docs, group_dropped, _ = relation_result + deduped_relations.extend(group_docs) + dropped += group_dropped + return deduped_entities + deduped_relations, dropped def _struct_graph_row_id( @@ -1228,6 +1970,8 @@ CHAIN_KINDS: tuple[str, ...] = ("list", "timeline") # Max source-chunk text length passed to the LLM in the correction prompt. _CHAIN_CORRECTION_MAX_CHUNK_CHARS = 8196 _CHAIN_CORRECTION_MAX_CHUNKS = 12 +_CHAIN_CORRECTION_MAX_RELATIONS = 16 +_CHAIN_CORRECTION_CONCURRENCY = 10 CHAIN_CORRECTION_PROMPT = """You are correcting an extracted {kind}-kind structure. @@ -1440,55 +2184,70 @@ async def validate_and_correct_chain( return docs bad_edges = list(violations.keys()) - bad_docs: list[dict] = [] - for e in bad_edges: - bad_docs.extend(edge_to_docs.get(e, ())) - bad_relations_repr = [{"from": e[0], "to": e[1], "issue": "; ".join(reasons)} for e, reasons in violations.items()] - chunk_pairs = _chain_gather_chunk_text(bad_docs, chunks_by_id) - source_chunks_text = "\n\n".join(f"[{cid}]\n{text}" for cid, text in chunk_pairs) or "(no source chunks available)" - prompt = CHAIN_CORRECTION_PROMPT.format( - kind=kind, - bad_relations_json=json.dumps(bad_relations_repr, ensure_ascii=False), - source_chunks_text=source_chunks_text, - ) if callable(callback): try: callback(msg=f"chain validation: {len(bad_edges)} flagged for LLM correction") except Exception: pass - res = await gen_json( - "You correct extracted graph relations to satisfy a strict-chain constraint.", - prompt, - chat_mdl, - gen_conf={"temperature": 0.0}, - ) except Exception: - logging.exception("chain validate: detection / LLM call failed; skipping correction") - return docs - - if not isinstance(res, dict): - return docs - keep_raw = res.get("keep") - if not isinstance(keep_raw, list): + logging.exception("chain validate: detection failed; skipping correction") return docs bad_edge_set = set(bad_edges) keep_set: set[tuple[str, str]] = set() - for item in keep_raw: - if not isinstance(item, dict): - continue - s = item.get("from") - t = item.get("to") - if not isinstance(s, str) or not isinstance(t, str): - continue - edge = (s.strip(), t.strip()) - # Reject anything that wasn't in the bad set — we don't invent - # new relations and we don't allow the LLM to "rescue" a - # never-extracted edge. - if edge in bad_edge_set: - keep_set.add(edge) + correction_batches = [bad_edges[i : i + _CHAIN_CORRECTION_MAX_RELATIONS] for i in range(0, len(bad_edges), _CHAIN_CORRECTION_MAX_RELATIONS)] + correction_semaphore = asyncio.Semaphore(_CHAIN_CORRECTION_CONCURRENCY) + + async def correct_batch(batch_no: int, batch_edges: list[tuple[str, str]]) -> set[tuple[str, str]]: + # Fail open for a failed or malformed batch: retain its relations. + batch_keep = set(batch_edges) + batch_docs = [doc for edge in batch_edges for doc in edge_to_docs.get(edge, ())] + batch_relations = [{"from": e[0], "to": e[1], "issue": "; ".join(violations.get(e, ("cross-batch conflict",)))} for e in batch_edges] + chunk_pairs = _chain_gather_chunk_text(batch_docs, chunks_by_id) + source_chunks_text = "\n\n".join(f"[{cid}]\n{text}" for cid, text in chunk_pairs) or "(no source chunks available)" + prompt = CHAIN_CORRECTION_PROMPT.format( + kind=kind, + bad_relations_json=json.dumps(batch_relations, ensure_ascii=False), + source_chunks_text=source_chunks_text, + ) + try: + async with correction_semaphore: + res = await gen_json( + "You correct extracted graph relations to satisfy a strict-chain constraint.", + prompt, + chat_mdl, + gen_conf={"temperature": 0.0}, + ) + keep_raw = res.get("keep") if isinstance(res, dict) else None + if isinstance(keep_raw, list): + batch_keep = set() + batch_edge_set = set(batch_edges) + for item in keep_raw: + if not isinstance(item, dict): + continue + s, t = item.get("from"), item.get("to") + edge = (s.strip(), t.strip()) if isinstance(s, str) and isinstance(t, str) else None + if edge in batch_edge_set: + batch_keep.add(edge) + except Exception: + logging.exception("chain validate: correction batch %d failed; retaining its relations", batch_no) + return batch_keep + + batch_keeps = await asyncio.gather(*(correct_batch(i, batch) for i, batch in enumerate(correction_batches))) + for batch_keep in batch_keeps: + keep_set.update(batch_keep) + + # Independent corrections can be valid inside each request but conflict + # after their results are combined. Re-check the combined keep set and + # give the model one final decision over the remaining conflicts. + combined_violations = _chain_detect_violations(list(keep_set)) + if combined_violations: + conflict_edges = list(combined_violations) + final_keep = await correct_batch(-1, conflict_edges) + keep_set.difference_update(conflict_edges) + keep_set.update(final_keep) if keep_set == bad_edge_set: # LLM kept everything → no correction applied. @@ -1523,6 +2282,13 @@ async def merge_compiled_structures( similarity_threshold: float = 0.99, compilation_template_id: str | None = None, cancel_check: Callable[[], bool] | None = None, + timing_context: str | None = None, + chunks_by_id: dict[str, str] | None = None, + chain_kind: str = "", + chain_callback=None, + chain_timeout_seconds: float = 120.0, + es_waiter: Callable[[], Awaitable[None]] | None = None, + es_releaser: Callable[[], Awaitable[None]] | None = None, ) -> dict: """Merge ``docs`` (the output of ``compile_structure_from_text``) before inserting them into ES. @@ -1560,13 +2326,37 @@ async def merge_compiled_structures( if not docs: return {"inserted": 0, "updated": 0, "duplicates_dropped": 0} - deduped, dropped = await _struct_local_dedup( + if callable(cancel_check) and cancel_check(): + raise TaskCanceledException("Task was cancelled before local dedup") + deduped, dropped = await _struct_local_dedup_parallel( docs, chat_mdl, embd_mdl, similarity_threshold, + timing_context=timing_context, ) + if callable(cancel_check) and cancel_check(): + raise TaskCanceledException("Task was cancelled after local dedup") + if chain_kind in CHAIN_KINDS: + try: + deduped = await asyncio.wait_for( + validate_and_correct_chain( + deduped, + chunks_by_id or {}, + chat_mdl, + chain_kind, + callback=chain_callback, + ), + timeout=chain_timeout_seconds, + ) + except asyncio.TimeoutError: + logging.warning("chain validate: timed out after %ss; using local-deduped docs", chain_timeout_seconds) + except Exception: + logging.exception("chain validate: unexpected failure; using local-deduped docs") + + if callable(cancel_check) and cancel_check(): + raise TaskCanceledException("Task was cancelled after chain validation") graph_keys = { ( str(d.get("doc_id")), @@ -1581,29 +2371,28 @@ async def merge_compiled_structures( if callable(cancel_check) and cancel_check(): raise TaskCanceledException("Task was cancelled during structure ES dedup merge") - inserted = 0 - updated = 0 - for d in deduped: - _raise_if_canceled() - try: - result = await _struct_es_dedup_one( - d, - chat_mdl, - embd_mdl, - tenant_id, - kb_id, - similarity_threshold, - ) - except Exception: - logging.exception("merge_compiled_structures: per-doc dedup failed") - continue - if result == "inserted": - inserted += 1 - elif result == "updated": - updated += 1 + if es_waiter is not None: + await es_waiter() + _raise_if_canceled() + try: + inserted, updated = await _struct_es_dedup_batch( + deduped, + chat_mdl, + embd_mdl, + tenant_id, + kb_id, + similarity_threshold, + timing_context=timing_context, + cancel_check=cancel_check, + ) + except Exception: + logging.exception("merge_compiled_structures: batched ES dedup failed") + inserted = updated = 0 + if es_releaser is not None: + await es_releaser() graphs = 0 - for doc_id, compile_kwd, template_id in graph_keys: + for graph_index, (doc_id, compile_kwd, template_id) in enumerate(graph_keys): _raise_if_canceled() try: await rebuild_structure_graph_json( @@ -1622,12 +2411,13 @@ async def merge_compiled_structures( template_id, ) - return { + info = { "inserted": inserted, "updated": updated, "duplicates_dropped": dropped, "graphs": graphs, } + return info __all__ = [ diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py index 555656f4ba..b97f667cdf 100644 --- a/rag/svr/task_executor_refactor/chunk_post_processor.py +++ b/rag/svr/task_executor_refactor/chunk_post_processor.py @@ -515,7 +515,7 @@ async def load_chunks_with_vec( index_nm = search.index_name(tenant_id) if not settings.docStoreConn.index_exist(index_nm, kb_id): return [] - select_fields = ["id", "doc_id", "content_with_weight", vctr_nm] + select_fields = ["id", "doc_id", "content_with_weight", "compile_kwd", vctr_nm] order_by = OrderByExpr() order_by.asc("page_num_int") order_by.asc("top_int") @@ -529,7 +529,11 @@ async def load_chunks_with_vec( settings.docStoreConn.search, select_fields, [], - {"doc_id": [doc_id], "available_int": 1}, + { + "doc_id": [doc_id], + "available_int": 1, + "must_not": {"exists": "compile_kwd"}, + }, [], order_by, offset, diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index 983cd01cb1..3d6b507f4e 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -730,6 +730,7 @@ class TaskHandler: "content_with_weight", "page_num_int", "top_int", + "compile_kwd", ] order_by = OrderByExpr() order_by.asc("page_num_int") @@ -742,7 +743,15 @@ class TaskHandler: settings.docStoreConn.search, select_fields, [], - {"doc_id": [doc_id], "available_int": 1}, + { + "doc_id": [doc_id], + "available_int": 1, + # Compilation writes its output back to the same + # document index. Exclude those rows in the query so + # they cannot change offset pagination while this + # task is still streaming source chunks. + "must_not": {"exists": "compile_kwd"}, + }, [], order_by, offset,