mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 01:43:27 +08:00
992 lines
36 KiB
Python
992 lines
36 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.
|
|
#
|
|
"""Shared helpers for the knowlege_compile pipelines (structure + wiki).
|
|
|
|
Both ``structure.py`` (compile_structure_from_text / merge_compiled_structures)
|
|
and ``wiki.py`` (the MAP→REDUCE→PLAN→REFINE artifact pipeline) need the same set
|
|
of plumbing: encode-through-LLMBundle, stable id minting, search-tokenizer
|
|
pairs, order-preserving chunk-id unions, defensive LLMBundle validation, the
|
|
``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead`` token-budget
|
|
calculation, and thin ES I/O wrappers.
|
|
|
|
Anything in this module is meant to be:
|
|
- LLMBundle-aware but provider-agnostic;
|
|
- Safe to import from either pipeline without circular references;
|
|
- Synchronous unless an awaitable behaviour is required.
|
|
|
|
Heavier shared logic that is conceptually identical but happens to differ in
|
|
shape between the two pipelines (e.g. pairwise-cosine dedup, LLM "are these
|
|
the same?" batching) intentionally stays in each pipeline file for now —
|
|
extract those only when their shapes converge.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import string
|
|
from typing import Any, Awaitable, Callable, Iterable, Optional
|
|
|
|
import xxhash
|
|
|
|
from common.misc_utils import thread_pool_exec
|
|
from common.token_utils import num_tokens_from_string
|
|
from rag.nlp import rag_tokenizer
|
|
from rag.prompts.generator import INPUT_UTILIZATION, gen_json, split_chunks
|
|
|
|
|
|
def knowledge_compile_gen_conf(chat_mdl, gen_conf: Optional[dict] = None) -> dict:
|
|
"""Add model-specific reasoning controls for knowledge compilation only."""
|
|
conf = dict(gen_conf or {})
|
|
model_config = getattr(chat_mdl, "model_config", None)
|
|
if not isinstance(model_config, dict):
|
|
model_config = {}
|
|
model_name = str(model_config.get("llm_name") or getattr(chat_mdl, "llm_name", "")).lower()
|
|
|
|
if "deepseek-v4" in model_name:
|
|
extra_body = conf.get("extra_body")
|
|
extra_body = dict(extra_body) if isinstance(extra_body, dict) else {}
|
|
extra_body["thinking"] = {"type": "disabled"}
|
|
conf["extra_body"] = extra_body
|
|
elif "qwen3" in model_name:
|
|
# chat_model.py maps this flag to the provider-specific request body.
|
|
conf["enable_thinking"] = False
|
|
else:
|
|
# LiteLLM maps this common control for providers that support it and
|
|
# drops it for providers that do not. Keep model-specific overrides
|
|
# above because some APIs use a different, explicit request field.
|
|
conf["reasoning_effort"] = "none"
|
|
|
|
return conf
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ID minting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def stable_row_id(*parts) -> str:
|
|
"""xxh64 hexdigest of ``":".join(parts)`` — stable per part tuple, used
|
|
as the ES row id when we want idempotent upserts.
|
|
|
|
``None`` parts become empty strings, everything else is ``str()``-ified.
|
|
"""
|
|
key = ":".join("" if p is None else str(p) for p in parts)
|
|
return xxhash.xxh64(key.encode("utf-8", "surrogatepass")).hexdigest()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Embedding
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def encode(embd_mdl, texts: list[str]) -> list:
|
|
"""``LLMBundle.encode`` wrapped in ``thread_pool_exec``.
|
|
|
|
Returns the embeddings list (drops the ``used_tokens`` count); empty
|
|
input returns ``[]``. Caller is responsible for ensuring ``embd_mdl``
|
|
is a real bundle — use :func:`ensure_llm_bundle` to validate at entry.
|
|
"""
|
|
if not texts:
|
|
return []
|
|
embeddings, _ = await thread_pool_exec(embd_mdl.encode, texts)
|
|
return list(embeddings)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tokenization for keyword search
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def tokenize_for_search(text: str) -> tuple[str, str]:
|
|
"""Returns ``(content_ltks, content_sm_ltks)`` for a piece of text.
|
|
|
|
Empty / non-string input returns ``("", "")``. Used wherever we write a
|
|
searchable ES row that needs both tokenizations.
|
|
"""
|
|
if not isinstance(text, str) or not text:
|
|
return "", ""
|
|
ltks = rag_tokenizer.tokenize(text)
|
|
if not ltks:
|
|
return "", ""
|
|
sm = rag_tokenizer.fine_grained_tokenize(ltks)
|
|
return ltks, sm
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Order-preserving union of string lists
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def union_ordered(*lists: Optional[Iterable]) -> list[str]:
|
|
"""Concatenate iterables and dedupe, preserving first-seen order.
|
|
Falsy values and non-strings are silently dropped.
|
|
"""
|
|
seen_set: set[str] = set()
|
|
seen: list[str] = []
|
|
for lst in lists:
|
|
if not lst:
|
|
continue
|
|
for v in lst:
|
|
if not v or not isinstance(v, str):
|
|
continue
|
|
if v in seen_set:
|
|
continue
|
|
seen_set.add(v)
|
|
seen.append(v)
|
|
return seen
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token-budget calculation for split_chunks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def make_input_budget(
|
|
chat_mdl,
|
|
*prompts: str,
|
|
floor: int = 1024,
|
|
utilization: float = INPUT_UTILIZATION,
|
|
) -> int:
|
|
"""``chat_mdl.max_length * utilization - num_tokens(sum of prompts)``,
|
|
floored at ``floor``.
|
|
|
|
Mirrors the budget idiom used by ``compile_structure_from_text`` and
|
|
``wiki_map_from_chunks``: caller passes the constant prompt scaffolding
|
|
(system prompt + user template) — ``split_chunks`` then sizes batches
|
|
to leave that much room.
|
|
"""
|
|
overhead = num_tokens_from_string("".join(p or "" for p in prompts))
|
|
budget = int(chat_mdl.max_length * utilization) - overhead
|
|
return max(budget, floor)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Defensive LLMBundle validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def ensure_llm_bundle(mdl, method: str, *, label: str = "model"):
|
|
"""Return ``mdl`` if it exposes ``method``; otherwise try to unwrap a
|
|
tuple, otherwise return ``None`` and log an error.
|
|
|
|
Common cause for tuple inputs at call sites: ``LLMBundle.encode()`` and
|
|
similar methods return ``(embeddings, used_tokens)``. If a caller stores
|
|
the *result* of ``encode()`` into a variable named like
|
|
``embedding_model`` and passes that in, we end up with a tuple here.
|
|
We unwrap with a warning so the pipeline keeps working while the caller
|
|
is fixed.
|
|
"""
|
|
if hasattr(mdl, method):
|
|
return mdl
|
|
if isinstance(mdl, tuple) and mdl and hasattr(mdl[0], method):
|
|
logging.warning(
|
|
"%s arrived as a %s; unwrapping to first element (check the call site — was %s()'s return value passed instead of the LLMBundle?)",
|
|
label,
|
|
type(mdl).__name__,
|
|
method,
|
|
)
|
|
return mdl[0]
|
|
logging.error(
|
|
"%s has no .%s method (type=%s); aborting",
|
|
label,
|
|
method,
|
|
type(mdl).__name__,
|
|
)
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ES I/O wrappers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def doc_storage_search(
|
|
select_fields: list[str],
|
|
condition: dict,
|
|
*,
|
|
tenant_id: str,
|
|
kb_ids: list[str],
|
|
match_expressions: list | None = None,
|
|
offset: int = 0,
|
|
limit: int = 1000,
|
|
label: str = "doc_storage_search",
|
|
) -> dict:
|
|
"""Thin wrapper around ``docStoreConn.search`` + ``get_fields``.
|
|
|
|
Returns ``{row_id: row_dict}``. Returns ``{}`` on failure (with a
|
|
logged exception). ``label`` is included in the failure log so each
|
|
call site is identifiable.
|
|
"""
|
|
from common import settings
|
|
from common.doc_store.doc_store_base import OrderByExpr
|
|
from rag.nlp import search as _rag_search
|
|
|
|
index = _rag_search.index_name(tenant_id)
|
|
try:
|
|
res = await thread_pool_exec(
|
|
settings.docStoreConn.search,
|
|
select_fields,
|
|
[],
|
|
condition,
|
|
match_expressions or [],
|
|
OrderByExpr(),
|
|
offset,
|
|
limit,
|
|
index,
|
|
kb_ids,
|
|
)
|
|
return settings.docStoreConn.get_fields(res, select_fields) or {}
|
|
except Exception:
|
|
logging.exception("%s failed (condition=%r)", label, condition)
|
|
return {}
|
|
|
|
|
|
async def doc_storage_insert(
|
|
rows: list[dict],
|
|
tenant_id: str,
|
|
kb_id: str,
|
|
*,
|
|
label: str = "doc_storage_insert",
|
|
) -> None:
|
|
"""Bulk insert wrapped in ``thread_pool_exec``. Logs on failure."""
|
|
if not rows:
|
|
return
|
|
from common import settings
|
|
from rag.nlp import search as _rag_search
|
|
|
|
index = _rag_search.index_name(tenant_id)
|
|
try:
|
|
await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id)
|
|
except Exception:
|
|
logging.exception("%s failed (%d row(s))", label, len(rows))
|
|
|
|
|
|
async def doc_storage_delete(
|
|
condition: dict,
|
|
tenant_id: str,
|
|
kb_id: str,
|
|
*,
|
|
label: str = "doc_storage_delete",
|
|
) -> None:
|
|
"""Bulk delete wrapped in ``thread_pool_exec``. Best-effort; logs on
|
|
failure (some callers rely on id-based upsert as a fallback)."""
|
|
from common import settings
|
|
from rag.nlp import search as _rag_search
|
|
|
|
index = _rag_search.index_name(tenant_id)
|
|
try:
|
|
await thread_pool_exec(settings.docStoreConn.delete, condition, index, kb_id)
|
|
except Exception:
|
|
logging.debug("%s failed (condition=%r); caller may rely on id-upsert", label, condition)
|
|
|
|
|
|
async def doc_storage_upsert_one(
|
|
filter_condition: dict,
|
|
row: dict,
|
|
tenant_id: str,
|
|
kb_id: str,
|
|
*,
|
|
label: str = "doc_storage_upsert_one",
|
|
) -> None:
|
|
"""Delete-by-filter then insert. Used when an in-place update would
|
|
require knowing the existing row's id and we'd rather drop+re-create.
|
|
|
|
Best-effort delete (failures are debug-logged) followed by the insert.
|
|
Set ``row["id"]`` to a stable value derived from the filter
|
|
(:func:`stable_row_id`) so id-based dedup at the connector catches any
|
|
race that bypasses the delete.
|
|
"""
|
|
await doc_storage_delete(filter_condition, tenant_id, kb_id, label=f"{label}.delete")
|
|
await doc_storage_insert([row], tenant_id, kb_id, label=f"{label}.insert")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Doc-vector field discovery
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def find_vec_field(doc: dict) -> tuple[Optional[str], Optional[list]]:
|
|
"""Locate the ``q_<dim>_vec`` field on an ES doc dict. Returns
|
|
``(field_name, vec)`` or ``(None, None)`` if the doc carries no
|
|
embedding."""
|
|
for k, v in doc.items():
|
|
if isinstance(k, str) and k.startswith("q_") and k.endswith("_vec"):
|
|
return k, v
|
|
return None, None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Chunked-LLM pipeline engine
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Both artifact MAP and compile_structure_from_text follow the same outer shape:
|
|
#
|
|
# 1. Filter chunks (drop empty text, optionally skip a "resume" set);
|
|
# 2. Pack remaining chunks into batches via ``split_chunks`` sized to leave
|
|
# room for the prompt scaffolding;
|
|
# 3. Run an LLM-driven ``process_batch`` over each batch in parallel under
|
|
# an ``asyncio.Semaphore(max_workers)``;
|
|
# 4. Aggregate the per-batch results into a single value.
|
|
#
|
|
# The inner LLM call shape diverges between the pipelines — artifact uses a
|
|
# single ``gen_json`` per batch with ``[CHUNK_ID Cn]``-labelled bodies,
|
|
# structure uses two ``gen_json`` calls (nodes then edges) with ``---``
|
|
# separators and no per-chunk attribution. That divergence lives in each
|
|
# pipeline's ``process_batch`` closure; this engine only owns the scaffold.
|
|
|
|
|
|
def _default_chunk_text(chunk: dict) -> str:
|
|
if not isinstance(chunk, dict):
|
|
return ""
|
|
text = chunk.get("text") or chunk.get("content_with_weight") or chunk.get("content") or ""
|
|
return text if isinstance(text, str) else ""
|
|
|
|
|
|
def _default_label(position_in_batch: int) -> str:
|
|
return f"C{position_in_batch + 1}"
|
|
|
|
|
|
def build_chunk_batches(
|
|
chunks: list[dict],
|
|
chat_mdl,
|
|
*,
|
|
prompt_overhead_tokens: int,
|
|
resume_chunk_ids: Optional[set[str]] = None,
|
|
scrub_text: Optional[Callable[[str], str]] = None,
|
|
label_fn: Callable[[int], str] = _default_label,
|
|
chunk_text_picker: Optional[Callable[[dict], str]] = None,
|
|
budget_floor: int = 1024,
|
|
batch_size_cap: Optional[int] = None,
|
|
window_fraction: Optional[float] = None,
|
|
) -> tuple[list[list[dict]], dict]:
|
|
"""Filter chunks, pack into batches, return per-batch entries.
|
|
|
|
Each batch entry is ``{"label": str, "chunk_id": str, "text": str}``
|
|
where ``label`` is per-batch positional (default ``C1``, ``C2``, …) and
|
|
``text`` is the post-scrub chunk body. Empty or resume-skipped chunks
|
|
are dropped.
|
|
|
|
Two packing modes:
|
|
- **Default (split_chunks)**: ``input_budget`` derived from
|
|
``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead_tokens``.
|
|
Used by ``structure.py`` and the legacy artifact MAP path.
|
|
- **Cap+fraction (greedy)**: when ``batch_size_cap`` is provided,
|
|
chunks are packed greedily with two cutoffs — chunk-count exceeds
|
|
``batch_size_cap`` OR accumulated tokens exceed
|
|
``chat_mdl.max_length * window_fraction``. This is the artifact
|
|
compilation rule (BS=8, window=0.5).
|
|
|
|
Returns ``(batches, info)`` where ``info`` is a small stats dict.
|
|
"""
|
|
if not chunks:
|
|
return [], {"total": 0, "kept": 0, "skipped_resume": 0, "skipped_empty": 0, "input_budget": 0, "n_batches": 0}
|
|
|
|
picker = chunk_text_picker or _default_chunk_text
|
|
resume_set = resume_chunk_ids or set()
|
|
|
|
chunk_ids: list[str] = []
|
|
chunk_texts: list[str] = []
|
|
skipped_resume = 0
|
|
skipped_empty = 0
|
|
|
|
for chunk in chunks:
|
|
cid = chunk.get("id") or chunk.get("chunk_id")
|
|
if not cid:
|
|
skipped_empty += 1
|
|
continue
|
|
if cid in resume_set:
|
|
skipped_resume += 1
|
|
continue
|
|
text = picker(chunk)
|
|
if not text or not text.strip():
|
|
skipped_empty += 1
|
|
continue
|
|
if scrub_text is not None:
|
|
text = scrub_text(text)
|
|
if not text or not text.strip():
|
|
skipped_empty += 1
|
|
continue
|
|
chunk_ids.append(cid)
|
|
chunk_texts.append(text)
|
|
|
|
if not chunk_texts:
|
|
return [], {
|
|
"total": len(chunks),
|
|
"kept": 0,
|
|
"skipped_resume": skipped_resume,
|
|
"skipped_empty": skipped_empty,
|
|
"input_budget": 0,
|
|
"n_batches": 0,
|
|
}
|
|
|
|
batches: list[list[dict]] = []
|
|
input_budget: int
|
|
|
|
if batch_size_cap is not None:
|
|
# Artifact mode — greedy bin-packing with chunk-count + token caps.
|
|
fraction = window_fraction if window_fraction is not None else 0.5
|
|
token_cap = max(int(chat_mdl.max_length * fraction), budget_floor)
|
|
input_budget = token_cap
|
|
|
|
current: list[dict] = []
|
|
current_tks = 0
|
|
for idx, text in enumerate(chunk_texts):
|
|
tks = num_tokens_from_string(text)
|
|
would_overflow_count = len(current) >= batch_size_cap
|
|
would_overflow_tokens = current and (current_tks + tks > token_cap)
|
|
if would_overflow_count or would_overflow_tokens:
|
|
batches.append(current)
|
|
current = []
|
|
current_tks = 0
|
|
current.append(
|
|
{
|
|
"label": label_fn(len(current)),
|
|
"chunk_id": chunk_ids[idx],
|
|
"text": text,
|
|
}
|
|
)
|
|
current_tks += tks
|
|
if current:
|
|
batches.append(current)
|
|
else:
|
|
input_budget = max(
|
|
int(chat_mdl.max_length * INPUT_UTILIZATION) - prompt_overhead_tokens,
|
|
budget_floor,
|
|
)
|
|
|
|
raw_batches = split_chunks(chunk_texts, input_budget) or []
|
|
for batch in raw_batches:
|
|
packed: list[dict] = []
|
|
for position, item in enumerate(batch):
|
|
for idx, text in item.items():
|
|
packed.append(
|
|
{
|
|
"label": label_fn(position),
|
|
"chunk_id": chunk_ids[idx],
|
|
"text": text,
|
|
}
|
|
)
|
|
if packed:
|
|
batches.append(packed)
|
|
|
|
info = {
|
|
"total": len(chunks),
|
|
"kept": len(chunk_texts),
|
|
"skipped_resume": skipped_resume,
|
|
"skipped_empty": skipped_empty,
|
|
"input_budget": input_budget,
|
|
"n_batches": len(batches),
|
|
}
|
|
return batches, info
|
|
|
|
|
|
async def run_chunked_pipeline(
|
|
batches: list[list[dict]],
|
|
*,
|
|
process_batch: Callable[..., Awaitable[Any]],
|
|
aggregate: Optional[Callable[[list[Any]], Any]] = None,
|
|
max_workers: int = 6,
|
|
callback: Optional[Callable] = None,
|
|
log_prefix: str = "chunked_pipeline",
|
|
) -> Any:
|
|
"""Run ``process_batch`` over each batch in parallel.
|
|
|
|
``process_batch`` is called as
|
|
``await process_batch(entries: list[dict], batch_idx: int, total: int)``
|
|
and may return anything; ``aggregate`` (if given) is called with the
|
|
list of per-batch results and its return value is the engine's return.
|
|
Without ``aggregate`` the raw per-batch results list is returned.
|
|
|
|
Cancel-on-error semantics: if any task raises, all sibling tasks are
|
|
cancelled and the exception propagates.
|
|
"""
|
|
if not batches:
|
|
return aggregate([]) if aggregate else []
|
|
|
|
total = len(batches)
|
|
worker_count = max_workers if max_workers and max_workers > 0 else 6
|
|
work_queue: asyncio.Queue[tuple[int, list[dict]] | None] = asyncio.Queue(maxsize=worker_count)
|
|
results: list[Any] = [None] * total
|
|
completed: list[bool] = [False] * total
|
|
|
|
async def _producer() -> None:
|
|
for idx, entries in enumerate(batches):
|
|
await work_queue.put((idx, entries))
|
|
for _ in range(worker_count):
|
|
await work_queue.put(None)
|
|
|
|
async def _worker() -> None:
|
|
while True:
|
|
item = await work_queue.get()
|
|
if item is None:
|
|
work_queue.task_done()
|
|
return
|
|
idx, entries = item
|
|
try:
|
|
results[idx] = await process_batch(entries, idx, total)
|
|
completed[idx] = True
|
|
finally:
|
|
work_queue.task_done()
|
|
|
|
producer = asyncio.create_task(_producer())
|
|
workers = [asyncio.create_task(_worker()) for _ in range(worker_count)]
|
|
try:
|
|
await asyncio.gather(producer, *workers)
|
|
except BaseException:
|
|
producer.cancel()
|
|
for t in workers:
|
|
t.cancel()
|
|
await asyncio.gather(producer, *workers, return_exceptions=True)
|
|
raise
|
|
|
|
if callback:
|
|
try:
|
|
callback(1.0, f"{log_prefix}: {total} batch(es) complete")
|
|
except Exception:
|
|
logging.debug("%s: completion callback failed", log_prefix, exc_info=True)
|
|
|
|
ordered_results = [results[idx] for idx in range(total) if completed[idx]]
|
|
return aggregate(ordered_results) if aggregate else ordered_results
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bulk dedup engine — exact + embedding + LLM disambiguation
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Replaces wiki's _wiki_exact_dedup_entities / _wiki_exact_dedup_concepts /
|
|
# _wiki_embedding_dedup_entities / _wiki_resolve_ambiguous_entities /
|
|
# _wiki_apply_merges with one parameterised engine. structure.py's
|
|
# merge_compiled_structures uses a different algorithm (incremental
|
|
# kept-set + per-pair LLM judgement) and stays as-is.
|
|
|
|
|
|
_PUNCT_TABLE = str.maketrans("", "", string.punctuation)
|
|
|
|
|
|
DEFAULT_DISAMBIGUATE_SYSTEM = "You are a named-entity resolution assistant. Return only JSON."
|
|
|
|
|
|
def normalize_key(name) -> str:
|
|
"""Lowercase + strip whitespace + strip ASCII punctuation. Used as the
|
|
bucket key for exact dedup."""
|
|
if not isinstance(name, str):
|
|
return ""
|
|
return name.lower().strip().translate(_PUNCT_TABLE)
|
|
|
|
|
|
def _exact_dedup_by_key(
|
|
items: list[dict],
|
|
*,
|
|
name_key: str,
|
|
type_key: Optional[str] = None,
|
|
aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
|
|
) -> list[dict]:
|
|
"""Group items by ``(normalize(item[name_key]), item.get(type_key))``.
|
|
|
|
Canonical record per group:
|
|
- ``<name_key>``: the most-common spelling across the group
|
|
- ``<type_key>`` (if given): the group's shared value
|
|
- ``aliases``: sorted union of every name + every input alias, minus
|
|
the canonical name
|
|
- ``mention_count``: sum of input ``mention_count`` values (defaults
|
|
to ``1`` per missing)
|
|
- ``chunk_ids``: order-preserving union
|
|
- ``_norm``: the normalized key (stripped by ``bulk_dedup_items``)
|
|
- any extras from ``aggregate_extra(group)``
|
|
"""
|
|
groups: dict[tuple, list[dict]] = {}
|
|
for it in items:
|
|
if not isinstance(it, dict):
|
|
continue
|
|
norm = normalize_key(it.get(name_key, ""))
|
|
if not norm:
|
|
continue
|
|
key = (norm, it.get(type_key) if type_key else None)
|
|
groups.setdefault(key, []).append(it)
|
|
|
|
canonical: list[dict] = []
|
|
for (norm, type_val), group in groups.items():
|
|
name_counts: dict[str, int] = {}
|
|
for it in group:
|
|
n = it.get(name_key, "")
|
|
if isinstance(n, str) and n:
|
|
name_counts[n] = name_counts.get(n, 0) + 1
|
|
best = max(name_counts, key=lambda k: name_counts[k]) if name_counts else ""
|
|
|
|
aliases: set[str] = set()
|
|
chunk_id_lists: list[list] = []
|
|
mention_count = 0
|
|
for it in group:
|
|
n = it.get(name_key, "")
|
|
if isinstance(n, str) and n:
|
|
aliases.add(n)
|
|
for a in it.get("aliases") or []:
|
|
if isinstance(a, str) and a:
|
|
aliases.add(a)
|
|
chunk_id_lists.append(it.get("chunk_ids") or [])
|
|
mention_count += int(it.get("mention_count") or 1)
|
|
aliases.discard(best)
|
|
|
|
record: dict = {
|
|
name_key: best,
|
|
"aliases": sorted(aliases),
|
|
"mention_count": mention_count,
|
|
"chunk_ids": union_ordered(*chunk_id_lists),
|
|
"_norm": norm,
|
|
}
|
|
if type_key:
|
|
record[type_key] = type_val
|
|
if aggregate_extra is not None:
|
|
try:
|
|
extras = aggregate_extra(group) or {}
|
|
if isinstance(extras, dict):
|
|
record.update(extras)
|
|
except Exception:
|
|
logging.exception("bulk_dedup: aggregate_extra failed for group %r", norm)
|
|
canonical.append(record)
|
|
|
|
return canonical
|
|
|
|
|
|
async def _embedding_dedup(
|
|
canonical: list[dict],
|
|
embd_mdl,
|
|
*,
|
|
name_key: str,
|
|
type_key: Optional[str] = None,
|
|
merge_threshold: float = 0.90,
|
|
ambiguous_low: float = 0.75,
|
|
) -> tuple[dict[int, int], list[tuple[int, int]], Optional[list]]:
|
|
"""Blockwise pairwise cosine; same-type-only when ``type_key`` given.
|
|
|
|
Returns ``(merged_into, ambiguous_pairs, vectors)``. ``merged_into``
|
|
is a union-find map ``index → parent_index``. ``ambiguous_pairs`` is the
|
|
[ambiguous_low, merge_threshold) bucket (after removing pairs already
|
|
linked by auto-merges). ``vectors`` is ``None`` on embedding failure
|
|
(caller should skip dedup).
|
|
"""
|
|
n = len(canonical)
|
|
if n <= 1:
|
|
return {}, [], []
|
|
|
|
names = [it.get(name_key, "") for it in canonical]
|
|
try:
|
|
vectors = await encode(embd_mdl, names)
|
|
except Exception:
|
|
logging.exception("bulk_dedup: embedding batch failed")
|
|
return {}, [], None
|
|
if vectors is None or len(vectors) != n:
|
|
return {}, [], None
|
|
|
|
try:
|
|
import numpy as np
|
|
|
|
matrix = np.asarray([list(v) for v in vectors], dtype=float)
|
|
if matrix.ndim != 2 or matrix.shape[0] != n:
|
|
raise ValueError(f"invalid embedding matrix shape: {matrix.shape}")
|
|
|
|
# Normalize once, then calculate only bounded blocks of the upper
|
|
# triangle. The old implementation materialized an N x N matrix,
|
|
# which becomes a significant memory spike for large KBs.
|
|
norms = np.linalg.norm(matrix, axis=1, keepdims=True)
|
|
matrix = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0)
|
|
except Exception:
|
|
logging.exception("bulk_dedup: pairwise cosine failed; skipping")
|
|
return {}, [], vectors
|
|
|
|
merged_into: dict[int, int] = {}
|
|
|
|
def _root(i: int) -> int:
|
|
while i in merged_into:
|
|
i = merged_into[i]
|
|
return i
|
|
|
|
auto_pairs: list[tuple[int, int]] = []
|
|
ambiguous_pairs: list[tuple[int, int]] = []
|
|
block_size = 1024
|
|
groups: dict[Any, list[int]] = {}
|
|
for index, item in enumerate(canonical):
|
|
groups.setdefault(item.get(type_key) if type_key else None, []).append(index)
|
|
|
|
for group_indices in groups.values():
|
|
for left_start in range(0, len(group_indices), block_size):
|
|
left_indices = group_indices[left_start : left_start + block_size]
|
|
left_vectors = matrix[left_indices]
|
|
for right_start in range(left_start, len(group_indices), block_size):
|
|
right_indices = group_indices[right_start : right_start + block_size]
|
|
sims = left_vectors @ matrix[right_indices].T
|
|
if right_start == left_start:
|
|
candidate_mask = np.triu(sims >= ambiguous_low, k=1)
|
|
else:
|
|
candidate_mask = sims >= ambiguous_low
|
|
rows, cols = np.nonzero(candidate_mask)
|
|
for row, col in zip(rows.tolist(), cols.tolist(), strict=True):
|
|
score = float(sims[row, col])
|
|
if score >= merge_threshold:
|
|
auto_pairs.append((left_indices[row], right_indices[col]))
|
|
else:
|
|
ambiguous_pairs.append((left_indices[row], right_indices[col]))
|
|
|
|
for i, j in auto_pairs:
|
|
ri, rj = _root(i), _root(j)
|
|
if ri == rj:
|
|
continue
|
|
if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
|
|
merged_into[rj] = ri
|
|
else:
|
|
merged_into[ri] = rj
|
|
|
|
still_ambiguous = [(i, j) for i, j in ambiguous_pairs if _root(i) != _root(j)]
|
|
return merged_into, still_ambiguous, vectors
|
|
|
|
|
|
async def _resolve_ambiguous_pairs(
|
|
canonical: list[dict],
|
|
ambiguous_pairs: list[tuple[int, int]],
|
|
merged_into: dict[int, int],
|
|
chat_mdl,
|
|
*,
|
|
name_key: str,
|
|
type_key: Optional[str] = None,
|
|
batch_size: int = 50,
|
|
llm_timeout: int = 60,
|
|
system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
|
|
) -> dict[int, int]:
|
|
"""LLM-judged disambiguation in bounded concurrent waves.
|
|
|
|
Each wave filters pairs using the latest union-find roots, then runs at
|
|
most one batch per available pooled LLM slot. Results are applied in one
|
|
serial pass before the next wave starts, avoiding redundant calls for
|
|
pairs merged by an earlier wave.
|
|
"""
|
|
if not ambiguous_pairs:
|
|
return merged_into
|
|
|
|
def _root(i: int) -> int:
|
|
while i in merged_into:
|
|
i = merged_into[i]
|
|
return i
|
|
|
|
async def _resolve_batch(batch: list[tuple[int, int]]) -> tuple[list[tuple[int, int]], Optional[list]]:
|
|
lines: list[str] = []
|
|
for k, (i, j) in enumerate(batch):
|
|
a_type = f" ({canonical[i].get(type_key, '')})" if type_key else ""
|
|
b_type = f" ({canonical[j].get(type_key, '')})" if type_key else ""
|
|
lines.append(f'{k + 1}. "{canonical[i].get(name_key, "")}"{a_type} vs "{canonical[j].get(name_key, "")}"{b_type}')
|
|
|
|
user_prompt = (
|
|
"For each pair below, determine if they refer to the same real-world entity.\n"
|
|
f"Return a JSON array of exactly {len(batch)} booleans "
|
|
"(true = same entity, false = different).\n"
|
|
"Return ONLY the JSON array.\n\n" + "\n".join(lines)
|
|
)
|
|
|
|
try:
|
|
res = await asyncio.wait_for(
|
|
gen_json(
|
|
system_prompt,
|
|
user_prompt,
|
|
chat_mdl,
|
|
gen_conf=knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}),
|
|
),
|
|
timeout=llm_timeout,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
logging.warning("bulk_dedup: disambiguation timed out (%d pairs)", len(batch))
|
|
return batch, None
|
|
except Exception:
|
|
logging.exception("bulk_dedup: disambiguation call failed (%d pairs)", len(batch))
|
|
return batch, None
|
|
|
|
decisions = None
|
|
if isinstance(res, list):
|
|
decisions = res
|
|
elif isinstance(res, dict):
|
|
for v in res.values():
|
|
if isinstance(v, list):
|
|
decisions = v
|
|
break
|
|
if not isinstance(decisions, list):
|
|
logging.warning("bulk_dedup: disambiguation returned unexpected shape: %r", type(res))
|
|
return batch, None
|
|
return batch, decisions
|
|
|
|
pool = getattr(chat_mdl, "_pool", None)
|
|
wave_batch_count = max(1, int(getattr(pool, "max_concurrency", 10)))
|
|
remaining = list(ambiguous_pairs)
|
|
while remaining:
|
|
eligible = [(i, j) for i, j in remaining if _root(i) != _root(j)]
|
|
if not eligible:
|
|
break
|
|
wave_pairs = eligible[: wave_batch_count * batch_size]
|
|
batches = [wave_pairs[start : start + batch_size] for start in range(0, len(wave_pairs), batch_size)]
|
|
results = await asyncio.gather(*(_resolve_batch(batch) for batch in batches))
|
|
for batch, decisions in results:
|
|
if decisions is None:
|
|
continue
|
|
for k, (i, j) in enumerate(batch):
|
|
verdict = decisions[k] if k < len(decisions) else False
|
|
if not verdict:
|
|
continue
|
|
ri, rj = _root(i), _root(j)
|
|
if ri == rj:
|
|
continue
|
|
if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
|
|
merged_into[rj] = ri
|
|
else:
|
|
merged_into[ri] = rj
|
|
remaining = eligible[len(wave_pairs) :]
|
|
|
|
return merged_into
|
|
|
|
|
|
def _apply_dedup_merges(
|
|
canonical: list[dict],
|
|
merged_into: dict[int, int],
|
|
*,
|
|
name_key: str,
|
|
) -> list[dict]:
|
|
"""Union-find collapse: sum ``mention_count``, union ``aliases`` and
|
|
``chunk_ids`` per canonical."""
|
|
|
|
def _root(i: int) -> int:
|
|
while i in merged_into:
|
|
i = merged_into[i]
|
|
return i
|
|
|
|
roots: set[int] = {_root(i) for i in range(len(canonical))}
|
|
out: list[dict] = []
|
|
for ri in roots:
|
|
base = dict(canonical[ri])
|
|
aliases: set[str] = set(base.get("aliases") or [])
|
|
chunk_id_lists: list[list] = [base.get("chunk_ids") or []]
|
|
mention_count = int(base.get("mention_count") or 0)
|
|
for i, it in enumerate(canonical):
|
|
if i == ri or _root(i) != ri:
|
|
continue
|
|
mention_count += int(it.get("mention_count") or 0)
|
|
aliases.update(it.get("aliases") or [])
|
|
n = it.get(name_key)
|
|
if isinstance(n, str) and n:
|
|
aliases.add(n)
|
|
chunk_id_lists.append(it.get("chunk_ids") or [])
|
|
aliases.discard(base.get(name_key) or "")
|
|
base["aliases"] = sorted(aliases)
|
|
base["mention_count"] = mention_count
|
|
base["chunk_ids"] = union_ordered(*chunk_id_lists)
|
|
out.append(base)
|
|
return out
|
|
|
|
|
|
async def bulk_dedup_items(
|
|
items: list[dict],
|
|
*,
|
|
name_key: str,
|
|
type_key: Optional[str] = None,
|
|
chat_mdl=None,
|
|
embd_mdl=None,
|
|
merge_threshold: float = 0.90,
|
|
ambiguous_low: float = 0.75,
|
|
ambiguous_batch_size: int = 50,
|
|
disambiguate_system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
|
|
llm_timeout: int = 60,
|
|
aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
|
|
strip_norm_key: bool = True,
|
|
) -> list[dict]:
|
|
"""Three-phase dedup → canonical items.
|
|
|
|
Phase 1 (always): exact dedup by ``(normalize(item[name_key]),
|
|
item.get(type_key))`` — groups by normalized key, sums mention_count,
|
|
unions aliases and chunk_ids, optionally adds extras via
|
|
``aggregate_extra(group)``.
|
|
|
|
Phase 2 (when ``embd_mdl`` is provided AND ``len(canonical) > 1``):
|
|
blockwise pairwise cosine over the canonical ``name_key`` values.
|
|
Pairs at similarity ≥ ``merge_threshold`` auto-merge; pairs in
|
|
``[ambiguous_low, merge_threshold)`` move to phase 3. When ``type_key``
|
|
is given, pairs are only considered when both endpoints share the same
|
|
type. Embedding failures cause this phase (and 3) to be skipped.
|
|
|
|
Phase 3 (when ``chat_mdl`` is provided AND ambiguous pairs remain):
|
|
batched LLM disambiguation via ``gen_json`` — each batch asks for a
|
|
JSON array of booleans. True verdicts join the union-find.
|
|
|
|
Apply: union-find collapse — sum mention_count, union aliases /
|
|
chunk_ids per canonical.
|
|
|
|
Setting both ``chat_mdl`` and ``embd_mdl`` to ``None`` makes this an
|
|
exact-dedup-only call (which is what artifact uses for concepts).
|
|
"""
|
|
canonical = _exact_dedup_by_key(
|
|
items,
|
|
name_key=name_key,
|
|
type_key=type_key,
|
|
aggregate_extra=aggregate_extra,
|
|
)
|
|
|
|
if len(canonical) > 1 and embd_mdl is not None:
|
|
merged_into, ambig, vectors = await _embedding_dedup(
|
|
canonical,
|
|
embd_mdl,
|
|
name_key=name_key,
|
|
type_key=type_key,
|
|
merge_threshold=merge_threshold,
|
|
ambiguous_low=ambiguous_low,
|
|
)
|
|
if vectors is None:
|
|
logging.warning("bulk_dedup: embedding phase skipped — keeping exact-dedup result")
|
|
else:
|
|
if chat_mdl is not None and ambig:
|
|
merged_into = await _resolve_ambiguous_pairs(
|
|
canonical,
|
|
ambig,
|
|
merged_into,
|
|
chat_mdl,
|
|
name_key=name_key,
|
|
type_key=type_key,
|
|
batch_size=ambiguous_batch_size,
|
|
llm_timeout=llm_timeout,
|
|
system_prompt=disambiguate_system_prompt,
|
|
)
|
|
canonical = _apply_dedup_merges(canonical, merged_into, name_key=name_key)
|
|
|
|
if strip_norm_key:
|
|
for it in canonical:
|
|
it.pop("_norm", None)
|
|
return canonical
|
|
|
|
|
|
__all__ = [
|
|
"stable_row_id",
|
|
"encode",
|
|
"tokenize_for_search",
|
|
"union_ordered",
|
|
"make_input_budget",
|
|
"ensure_llm_bundle",
|
|
"doc_storage_search",
|
|
"doc_storage_insert",
|
|
"doc_storage_delete",
|
|
"doc_storage_upsert_one",
|
|
"find_vec_field",
|
|
# New engines
|
|
"normalize_key",
|
|
"build_chunk_batches",
|
|
"run_chunked_pipeline",
|
|
"bulk_dedup_items",
|
|
"DEFAULT_DISAMBIGUATE_SYSTEM",
|
|
]
|