ragflow/rag/advanced_rag/knowlege_compile/dataset_nav.py

#
#  Copyright 2026 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.
#

"""Incremental clustering for dataset-level navigation.

Replaces the old 128-doc markdown with a hierarchy of nav_cluster and nav_doc
ES rows.  Each new document is embedded and placed into the nearest cluster
via layered KNN search + threshold-based merge/create.

Storage: one ES/Infinity row per nav_cluster or nav_doc node.
Tree structure encoded via ``parent_kwd`` on each row — no full-tree JSON blob.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
from typing import Any

import xxhash

from common.misc_utils import thread_pool_exec
from rag.utils.redis_conn import RedisDistributedLock

from ._common import encode as _encode
from ._common import knowledge_compile_gen_conf as _knowledge_compile_gen_conf

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_COMPILE_KWD = "dataset_nav"

# Embedding dimension — inferred at runtime from the first encode() call.
# None until _embed_dim is set.
_EMBED_DIM: int | None = None

# Similarity thresholds
_MERGE_THRESHOLD = 0.80  # merge doc into cluster
_RECURSE_THRESHOLD = 0.65  # continue descending into children
_MIN_SIM = 0.50  # minimum similarity to be considered related

# Max child count before triggering rebalance
_MAX_FANOUT = 64

# Max docs per leaf cluster before triggering split
_MAX_DOCS_PER_CLUSTER = 50

# Concurrency lock TTL
_LOCK_TIMEOUT_S = 30
_LOCK_BLOCKING_TIMEOUT_S = 5

# Hard limit on how many sibling clusters we evaluate per KNN call
_KNN_TOP_K = 5

# Fields needed to shape a nav hit for hybrid scoring / rendering.
_NAV_SEARCH_FIELDS = [
    "id",
    "content_with_weight",
    "name",
    "doc_id",
    "type_kwd",
    "doc_ids_kwd",
    "doc_count_int",
]

# Weight of the dense leg in hybrid search (1 - this = BM25 leg weight).
_NAV_HYBRID_DENSE_W = 0.5

# Stop-words skipped when deriving routing tag-words from a summary.
_NAV_STOP_WORDS = {
    "the",
    "a",
    "an",
    "and",
    "or",
    "of",
    "to",
    "in",
    "on",
    "for",
    "with",
    "at",
    "is",
    "are",
    "was",
    "were",
    "be",
    "been",
    "being",
    "this",
    "that",
    "these",
    "those",
    "it",
    "its",
    "as",
    "by",
    "from",
    "about",
    "into",
}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _nav_doc_id(doc_id: str) -> str:
    """Stable row id for a nav_doc (deterministic by doc_id)."""
    return xxhash.xxh64(
        f"dataset_nav:doc:{doc_id}".encode("utf-8", "surrogatepass"),
    ).hexdigest()


def _nav_cluster_id(kb_id: str, name: str) -> str:
    """Stable row id for a nav_cluster (deterministic by kb_id + name)."""
    return xxhash.xxh64(
        f"dataset_nav:{kb_id}:cluster:{name}".encode("utf-8", "surrogatepass"),
    ).hexdigest()


def _nav_lock_key(kb_id: str) -> str:
    """Redis lock key for concurrency control on a KB's nav tree."""
    return f"dataset_nav:{kb_id}"


def _extract_root_summary_from_tree(tree: dict | None) -> str:
    """Extract the doc-level summary from a RAPTOR tree (or bare string)."""
    if not isinstance(tree, dict):
        return ""
    title = tree.get("title") or ""
    if isinstance(title, str) and title.strip():
        return title.strip()
    for alt in ("summary", "content_with_weight", "content"):
        v = tree.get(alt)
        if isinstance(v, str) and v.strip():
            return v.strip()
    return ""


def _index_name(tenant_id: str) -> str:
    from rag.nlp import search as _rag_search

    return _rag_search.index_name(tenant_id)


def _vec_field(dim: int) -> str:
    return f"q_{dim}_vec"


# ---------------------------------------------------------------------------
# Doc store I/O — works with any engine (ES, Infinity, …) via docStoreConn
# ---------------------------------------------------------------------------


async def _store_get(tenant_id: str, kb_id: str, row_id: str) -> dict | None:
    from common import settings

    index = _index_name(tenant_id)
    try:
        return (
            await thread_pool_exec(
                settings.docStoreConn.get,
                row_id,
                index,
                [kb_id],
            )
            or None
        )
    except Exception:
        return None


async def _store_search(
    tenant_id: str,
    kb_id: str,
    condition: dict,
    fields: list[str],
    limit: int = 10000,
) -> list[dict]:
    from common import settings
    from common.doc_store.doc_store_base import OrderByExpr

    index = _index_name(tenant_id)
    res = await thread_pool_exec(
        settings.docStoreConn.search,
        fields,
        [],
        condition,
        [],
        OrderByExpr(),
        0,
        limit,
        index,
        [kb_id],
    )
    rows = settings.docStoreConn.get_fields(res, fields) if res else {}
    return list(rows.values())


async def _store_text_search(
    tenant_id: str,
    kb_id: str,
    query: str,
    fields: list[str],
    limit: int = 100,
    *,
    compile_kwd: str = _COMPILE_KWD,
    type_kwd: str = "",
    extra_filter: dict | None = None,
) -> list[dict]:
    """Full-text (BM25) leg over the nav rows' tokenized fields.

    Recalls nav nodes whose ``content_ltks`` / ``content_sm_ltks`` match the
    query — the lexical half of hybrid search, complementing the KNN leg for
    exact/proper-noun recall (e.g. "关羽", "ImageNet 2012").

    Args:
        compile_kwd: Which compile partition to search within
            (default ``_COMPILE_KWD`` = ``"dataset_nav"``).
        type_kwd: Optional type filter (``"nav_doc"``, ``"nav_cluster"``, or ``""``).
        extra_filter: Additional filter conditions merged into the query.
    """
    from common import settings
    from common.doc_store.doc_store_base import MatchTextExpr, OrderByExpr

    # Pre-tokenize the query with the same tokenizer used to index content_ltks.
    # ES content_ltks uses the "whitespace" analyzer (no stemming), while our
    # Python tokenizer applies stemming (e.g. "Christie" → "christi").
    # Without this, raw query terms won't match pre-stemmed tokens in the index.
    tokenized_query = _tokenize(query)

    index = _index_name(tenant_id)
    filter_condition: dict = {"compile_kwd": [compile_kwd]}
    if type_kwd:
        filter_condition["type_kwd"] = type_kwd
    if extra_filter:
        filter_condition.update(extra_filter)
    res = await thread_pool_exec(
        settings.docStoreConn.search,
        fields,
        [],
        filter_condition,
        [MatchTextExpr(["content_ltks", "content_sm_ltks"], tokenized_query, limit)],
        OrderByExpr(),
        0,
        limit,
        index,
        [kb_id],
    )
    rows = settings.docStoreConn.get_fields(res, fields) if res else {}
    return list(rows.values())


async def _store_knn(
    tenant_id: str,
    kb_id: str,
    vec: list[float],
    vec_dim: int,
    filter_condition: dict,
    top_k: int = _KNN_TOP_K,
) -> list[dict]:
    """KNN search with dense vector and filter, returning top_k hits."""
    from common import settings
    from common.doc_store.doc_store_base import MatchDenseExpr, OrderByExpr

    index = _index_name(tenant_id)
    vf = _vec_field(vec_dim)
    fields = [
        "content_with_weight",
        "name",
        "doc_id",
        "compile_kwd",
        "type_kwd",
        "parent_kwd",
        "depth_int",
        "doc_count_int",
        "doc_ids_kwd",
        vf,
    ]
    match_expr = MatchDenseExpr(
        vector_column_name=vf,
        embedding_data=list(vec),
        embedding_data_type="float",
        distance_type="cosine",
        topn=top_k,
        extra_options={},
    )
    res = await thread_pool_exec(
        settings.docStoreConn.search,
        fields,
        [],
        filter_condition,
        [match_expr],
        OrderByExpr(),
        0,
        top_k,
        index,
        [kb_id],
    )
    results = settings.docStoreConn.get_fields(res, fields) if res else {}
    rows = list(results.values())
    if filter_condition and any(not _matches_condition(row, filter_condition) for row in rows):
        scanned = await _store_search(tenant_id, kb_id, filter_condition, fields, limit=10000)
        rows = [row for row in scanned if _vector_len(row.get(vf)) == vec_dim and _matches_condition(row, filter_condition)]
        rows.sort(key=lambda row: _cosine_sim(vec, row.get(vf)), reverse=True)
        rows = rows[:top_k]
    return rows


async def _store_upsert(tenant_id: str, kb_id: str, doc: dict) -> None:
    from common import settings

    index = _index_name(tenant_id)
    row_id = doc.get("id", "")
    existing = await thread_pool_exec(
        settings.docStoreConn.get,
        row_id,
        index,
        [kb_id],
    )
    if existing:
        upd = {k: v for k, v in doc.items() if k != "id"}
        await thread_pool_exec(
            settings.docStoreConn.update,
            {"id": row_id},
            upd,
            index,
            kb_id,
        )
    else:
        await thread_pool_exec(
            settings.docStoreConn.insert,
            [doc],
            index,
            kb_id,
        )


async def _store_delete(tenant_id: str, kb_id: str, row_id: str) -> None:
    from common import settings

    index = _index_name(tenant_id)
    try:
        await thread_pool_exec(
            settings.docStoreConn.delete,
            {"id": [row_id]},
            index,
            kb_id,
        )
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Embedding helpers
# ---------------------------------------------------------------------------


async def _embed(embd_mdl, text: str) -> list[float]:
    """Encode a single text string and return its embedding vector."""
    global _EMBED_DIM
    vecs = await _encode(embd_mdl, [text])
    if vecs and _vector_len(vecs[0]) > 0:
        dim = len(vecs[0])
        if _EMBED_DIM is None:
            _EMBED_DIM = dim
        return vecs[0]
    return []


def _vector_len(vec) -> int:
    if vec is None:
        return 0
    try:
        return len(vec)
    except TypeError:
        return 0


def _cosine_sim(a, b) -> float:
    """Compute cosine similarity between two vectors."""
    a_len = _vector_len(a)
    b_len = _vector_len(b)
    if a_len == 0 or b_len == 0 or a_len != b_len:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(x * x for x in b) ** 0.5
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


# ---------------------------------------------------------------------------
# Fabrication of nav rows
# ---------------------------------------------------------------------------


def _make_nav_doc_row(
    kb_id: str,
    doc_id: str,
    summary: str,
    parent_kwd: str,
    depth_int: int,
    embd_mdl=None,
    embedding: list[float] | None = None,
    *,
    graph_content: str = "",
) -> dict:
    """Build a nav_doc ES/Infinity row dict for a single document leaf node.

    Args:
        summary: Short human-readable title / description for display.
        graph_content: Optional full-graph text used for richer keyword /
            entity extraction.  When set, keywords and entities are derived
            from graph_content instead of summary, and the text is stored in
            the ``graph_content`` payload field.
    """
    kw_text = graph_content or summary
    payload = {
        "type": "nav_doc",
        "description": summary,
        "keywords": _nav_keywords(kw_text),
        "entities": _nav_entities(kw_text),
    }
    if graph_content:
        payload["graph_content"] = graph_content
    row: dict = {
        "id": _nav_doc_id(doc_id),
        "kb_id": kb_id,
        "doc_id": doc_id,
        "compile_kwd": _COMPILE_KWD,
        "knowledge_graph_kwd": "entity",
        "type_kwd": "nav_doc",
        "name": f"{parent_kwd}_{xxhash.xxh64(summary.encode()).hexdigest()[:12]}",
        "parent_kwd": parent_kwd,
        "depth_int": depth_int,
        "available_int": 0,
    }
    row["content_with_weight"] = json.dumps(payload, ensure_ascii=False)
    ltks = _tokenize(kw_text)
    row["content_ltks"] = ltks
    row["content_sm_ltks"] = _fine_tokenize(ltks)
    if _vector_len(embedding) > 0:
        dim = len(embedding)
        row[_vec_field(dim)] = embedding
    return row


def _make_nav_cluster_row(
    kb_id: str,
    name: str,
    description: str,
    parent_kwd: str,
    depth_int: int,
    doc_ids: list[str],
    embedding: list[float] | None = None,
) -> dict:
    """Build a nav_cluster ES/Infinity row dict for an internal tree node."""
    cluster_id = _nav_cluster_id(kb_id, name)
    row: dict = {
        "id": cluster_id,
        "kb_id": kb_id,
        "doc_id": kb_id,
        "compile_kwd": _COMPILE_KWD,
        "knowledge_graph_kwd": "entity",
        "type_kwd": "nav_cluster",
        "name": name,
        "parent_kwd": parent_kwd,
        "depth_int": depth_int,
        "doc_ids_kwd": doc_ids,
        "doc_count_int": len(doc_ids),
        "available_int": 0,
    }
    payload = {
        "type": "nav_cluster",
        "description": description,
        "keywords": _nav_keywords(description),
        "entities": _nav_entities(description),
    }
    row["content_with_weight"] = json.dumps(payload, ensure_ascii=False)
    ltks = _tokenize(description)
    row["content_ltks"] = ltks
    row["content_sm_ltks"] = _fine_tokenize(ltks)
    if _vector_len(embedding) > 0:
        dim = len(embedding)
        row[_vec_field(dim)] = embedding
    return row


def _tokenize(text: str) -> str:
    """Coarse-grained tokenization for ES/Infinity full-text search."""
    from rag.nlp import rag_tokenizer

    return rag_tokenizer.tokenize(text)


def _fine_tokenize(text: str) -> str:
    """Fine-grained tokenization for ES/Infinity sub-word search."""
    from rag.nlp import rag_tokenizer

    return rag_tokenizer.fine_grained_tokenize(text)


def _nav_keywords(summary: str, max_kwds: int = 6) -> list[str]:
    """Derive routing tags from a nav summary (no LLM call, zero cost).

    These are the tokenized non-stop terms the model uses to fast-judge node
    relevance before reading the full description.
    """
    from rag.nlp import rag_tokenizer

    tokens = (rag_tokenizer.tokenize(summary or "") or "").split()
    seen: set[str] = set()
    out: list[str] = []
    for t in tokens:
        t = t.strip()
        if len(t) < 2 or t.isdigit() or t.lower() in _NAV_STOP_WORDS or t.lower() in seen:
            continue
        seen.add(t.lower())
        out.append(t)
        if len(out) >= max_kwds:
            break
    return out


def _nav_entities(summary: str, max_entities: int = 6) -> list[str]:
    """Extract likely named entities from a nav summary (no LLM call, zero cost).

    Uses a two-pass heuristic:
    1. Regex for English capitalized multi-word sequences (proper nouns).
    2. Tokenizer-based extraction for CJK and other non-Latin text.

    Returns up to ``max_entities`` deduplicated strings.
    """
    text = summary or ""
    entities: list[str] = []
    seen: set[str] = set()

    # Pass 1: English-like capitalized sequences (e.g. "New York", "Machine Learning")
    for m in re.finditer(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b", text):
        ent = m.group(1).strip()
        key = ent.lower()
        if key not in _NAV_STOP_WORDS and key not in seen:
            seen.add(key)
            entities.append(ent)
            if len(entities) >= max_entities:
                return entities

    # Pass 2: tokenizer-based for CJK / other scripts
    from rag.nlp import rag_tokenizer

    tokens = (rag_tokenizer.tokenize(text) or "").split()
    for t in tokens:
        t = t.strip()
        if len(t) < 3 or t.isdigit():
            continue
        if t.isascii() and t[0].islower():
            continue  # skip lowercase English fragments
        key = t.lower()
        if key in _NAV_STOP_WORDS or key in seen:
            continue
        seen.add(key)
        entities.append(t)
        if len(entities) >= max_entities:
            break

    return entities


def _matches_condition(row: dict, condition: dict) -> bool:
    """Check the simple equality filters used by dataset navigation."""
    for field, expected in condition.items():
        if not expected or field == "kb_id":
            continue
        actual = row.get(field)
        actual_values = actual if isinstance(actual, (list, tuple, set)) else [actual]
        expected_values = expected if isinstance(expected, (list, tuple, set)) else [expected]
        if not any(str(value) == str(item) for value in actual_values for item in expected_values):
            return False
    return True


# ---------------------------------------------------------------------------
# Incremental clustering core
# ---------------------------------------------------------------------------


async def _find_best_cluster(
    tenant_id: str,
    kb_id: str,
    doc_embedding: list[float],
    vec_dim: int,
) -> tuple[str | None, str | None, float]:
    """Locate the nearest cluster for a document via layered KNN descent.

    Starts from the root cluster (depth_int=0) and recursively descends into
    the best-matching child as long as similarity stays above
    ``_RECURSE_THRESHOLD``.  Returns the deepest cluster whose children are
    all less similar, along with the similarity score.

    Returns:
        (best_cluster_name, best_cluster_parent_name, similarity)
    """
    # Step 1: find the root cluster (depth_int=0)
    root_cond = {
        "kb_id": [kb_id],
        "compile_kwd": [_COMPILE_KWD],
        "type_kwd": ["nav_cluster"],
        "depth_int": [0],
    }
    roots = await _store_knn(tenant_id, kb_id, doc_embedding, vec_dim, root_cond, top_k=1)
    if not roots:
        return None, None, 0.0

    best = roots[0]
    best_name = best.get("name", "")
    best_parent = best.get("parent_kwd", "")
    # compute actual similarity to root
    stored = best.get(_vec_field(vec_dim))
    sim = _cosine_sim(doc_embedding, stored)
    visited_names = {best_name}

    # Step 2: recursively descend into children
    while sim >= _RECURSE_THRESHOLD:
        child_cond = {
            "kb_id": [kb_id],
            "compile_kwd": [_COMPILE_KWD],
            "type_kwd": ["nav_cluster"],
            "parent_kwd": [best_name],
        }
        children = await _store_knn(tenant_id, kb_id, doc_embedding, vec_dim, child_cond, top_k=1)
        if not children:
            break
        child = children[0]
        child_name = child.get("name", "")
        if not child_name or child_name in visited_names:
            break
        stored = child.get(_vec_field(vec_dim))
        child_sim = _cosine_sim(doc_embedding, stored)
        if child_sim < _RECURSE_THRESHOLD:
            break
        best_name = child_name
        best_parent = best.get("parent_kwd", best_parent)
        sim = child_sim
        best = child
        visited_names.add(best_name)

    return best_name, best_parent, sim


async def _llm_merge(chat_mdl, cluster_desc: str, doc_summary: str) -> str:
    """LLM merge: fuse existing cluster description with new doc summary."""
    if not chat_mdl:
        return cluster_desc  # no LLM available, keep old summary
    from rag.prompts.generator import gen_json

    prompt = (
        "Merge the following two descriptions of the same topic into "
        "a single concise summary (1-3 sentences):\n\n"
        f"Existing: {cluster_desc}\n\n"
        f"New: {doc_summary}\n\n"
        "Return ONLY the merged text, no commentary."
    )
    try:
        resp = await gen_json("", prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}))
        if isinstance(resp, dict):
            return str(resp.get("merged", resp.get("result", cluster_desc)))
        if isinstance(resp, str) and resp.strip():
            return resp.strip()
    except Exception:
        logging.exception("dataset_nav: LLM merge failed, keeping original")
    return cluster_desc


def _clean_title(title: str) -> str:
    """Normalize an LLM title into a one-line, length-capped display name."""
    return " ".join((title or "").split())[:48].strip()


def _fallback_title(summary: str) -> str:
    """Derive a short readable title from a summary when the LLM gives none."""
    words = " ".join((summary or "").split()).split(" ")
    return " ".join(words[:6]).strip() or "Cluster"


def _readable_cluster_name(title: str, seed: str) -> str:
    """A readable yet unique nav-cluster key: ``"<title> <8-hex>"``.

    The title makes the node name human-readable; the short hash of ``seed``
    (the cluster description) preserves the per-KB uniqueness the tree keying
    relies on (``_nav_cluster_id`` / ``parent_kwd``).
    """
    suffix = xxhash.xxh64((seed or "").encode("utf-8")).hexdigest()[:8]
    return f"{_clean_title(title) or 'Cluster'} {suffix}"


async def _llm_create_summary(chat_mdl, doc_summaries: list[str]) -> tuple[str, str]:
    """LLM-derive a cluster's readable ``(name, summary)`` from doc summaries.

    ``name`` is a short human-readable topic title used to make the nav node
    name readable (the caller still appends a short hash to keep the tree key
    unique); ``summary`` is the 1-3 sentence description.
    """
    fallback_summary = doc_summaries[0] if doc_summaries else ""
    if not chat_mdl:
        return _fallback_title(fallback_summary), fallback_summary

    from rag.prompts.generator import gen_json

    texts = "\n---\n".join(doc_summaries)
    prompt = (
        "Given the document excerpts below, produce a short human-readable topic "
        "name and a concise description of their common topic.\n\n"
        f"{texts}\n\n"
        'Return ONLY JSON: {"name": "<2-6 word topic title>", "summary": "<1-3 sentence description>"}'
    )
    try:
        resp = await gen_json("", prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}))
        if isinstance(resp, dict):
            summary = str(resp.get("summary") or resp.get("result") or fallback_summary).strip()
            name = _clean_title(str(resp.get("name") or "")) or _fallback_title(summary)
            return name, (summary or fallback_summary)
        if isinstance(resp, str) and resp.strip():
            return _fallback_title(resp), resp.strip()
    except Exception:
        logging.exception("dataset_nav: LLM summary failed")
    return _fallback_title(fallback_summary), fallback_summary


# ---------------------------------------------------------------------------
# Public surface
# ---------------------------------------------------------------------------


async def upsert_dataset_nav_doc(
    tenant_id: str,
    kb_id: str,
    doc_id: str,
    summary_or_tree: Any,
    embd_mdl=None,
    chat_mdl=None,
) -> None:
    """Upsert a document into the nav clustering tree.

    Args:
        tenant_id: Tenant owning the KB.
        kb_id: Knowledge base id.
        doc_id: Document id.
        summary_or_tree:
            - A plain summary string (fallback, no RAPTOR graph).
            - A RAPTOR tree dict (backwards-compat, root summary extracted).
            - A dict with keys ``title`` (short display label) and
              ``graph_text`` (full graph entities + relations for richer
              embedding / keyword extraction).
        embd_mdl: LLMBundle for embedding (required for clustering).
        chat_mdl: LLMBundle for chat (required for LLM merge/summary).
    """
    if not doc_id or not kb_id:
        return

    # 1. Extract display summary and optional full graph content.
    graph_content = ""
    if isinstance(summary_or_tree, dict):
        if "title" in summary_or_tree and "graph_text" in summary_or_tree:
            # New extended format from generate_nav's RAPTOR graph path.
            summary = (summary_or_tree.get("title") or "").strip()
            graph_content = (summary_or_tree.get("graph_text") or "").strip()
        else:
            summary = _extract_root_summary_from_tree(summary_or_tree)
    elif isinstance(summary_or_tree, str):
        summary = summary_or_tree
    else:
        summary = ""
    if not summary:
        logging.info("dataset_nav: skipping doc=%s (kb=%s) — no summary", doc_id, kb_id)
        return

    # 2. Embed with the richest available text so that KNN placement
    # benefits from the full tree structure when present.
    embed_text = graph_content or summary
    doc_embedding = await _embed(embd_mdl, embed_text) if embd_mdl else []
    vec_dim = len(doc_embedding)

    lock = RedisDistributedLock(
        _nav_lock_key(kb_id),
        timeout=_LOCK_TIMEOUT_S,
        blocking_timeout=_LOCK_BLOCKING_TIMEOUT_S,
    )
    try:
        await lock.spin_acquire()
    except Exception:
        logging.exception("dataset_nav: lock acquire failed for kb=%s", kb_id)
        return

    try:
        # 3. Check and replace the existing nav_doc while holding the same
        # lock as placement. Otherwise concurrent updates can both observe
        # the old row and race while deleting/rebuilding its parent cluster.
        existing_doc = await _store_get(tenant_id, kb_id, _nav_doc_id(doc_id))
        if existing_doc:
            old_payload = json.loads(existing_doc.get("content_with_weight") or "{}")
            if old_payload.get("description") == summary:
                return
            await _remove_dataset_nav_doc_locked(tenant_id, kb_id, doc_id)

        # 4. Layered KNN search for nearest cluster
        if _vector_len(doc_embedding) > 0:
            best_name, best_parent, sim = await _find_best_cluster(
                tenant_id,
                kb_id,
                doc_embedding,
                vec_dim,
            )
        else:
            logging.warning("dataset_nav: embedding unavailable for doc=%s, skipping KNN placement", doc_id)
            best_name, best_parent, sim = None, None, 0.0

        if best_name and sim >= _MERGE_THRESHOLD:
            # ── Merge into best cluster ──
            cluster_id = _nav_cluster_id(kb_id, best_name)
            cluster_row = await _store_get(tenant_id, kb_id, cluster_id)
            if cluster_row:
                payload = json.loads(cluster_row.get("content_with_weight") or "{}")
                old_desc = payload.get("description", "")
                new_desc = await _llm_merge(chat_mdl, old_desc, summary)
                payload["description"] = new_desc
                cluster_row["content_with_weight"] = json.dumps(payload, ensure_ascii=False)
                doc_ids = cluster_row.get("doc_ids_kwd") or []
                if doc_id not in doc_ids:
                    doc_ids.append(doc_id)
                cluster_row["doc_ids_kwd"] = doc_ids
                cluster_row["doc_count_int"] = len(doc_ids)
                # Re-compute embedding for the new summary
                if embd_mdl and new_desc != old_desc:
                    new_emb = await _embed(embd_mdl, new_desc)
                    if _vector_len(new_emb) > 0:
                        cluster_row[_vec_field(len(new_emb))] = new_emb
                await _store_upsert(tenant_id, kb_id, cluster_row)

            # Upsert nav_doc under the cluster
            depth = cluster_row.get("depth_int", 1) + 1 if cluster_row else 2
            nav_doc_row = _make_nav_doc_row(
                kb_id,
                doc_id,
                summary,
                best_name,
                depth,
                embd_mdl,
                doc_embedding,
                graph_content=graph_content,
            )
            await _store_upsert(tenant_id, kb_id, nav_doc_row)

            # Check fanout — if the cluster now has too many children, trigger split
            await _maybe_split_cluster(
                tenant_id,
                kb_id,
                best_name,
                embd_mdl,
                chat_mdl,
            )

        elif best_name and sim >= _MIN_SIM:
            # ── Create new cluster as sibling/child ──
            parent_for_new = best_parent if best_parent else best_name
            depth_of_parent = 1  # default
            parent_row = await _store_get(
                tenant_id,
                kb_id,
                _nav_cluster_id(kb_id, parent_for_new),
            )
            if parent_row:
                depth_of_parent = parent_row.get("depth_int", 1)
            new_depth = depth_of_parent + 1
            new_title, new_desc = await _llm_create_summary(chat_mdl, [summary])
            new_name = _readable_cluster_name(new_title, summary)
            new_cluster = _make_nav_cluster_row(
                kb_id,
                new_name,
                new_desc,
                parent_for_new,
                depth_of_parent,
                [doc_id],
                doc_embedding,
            )
            if embd_mdl and _vector_len(doc_embedding) > 0:
                new_cluster[_vec_field(len(doc_embedding))] = doc_embedding
            await _store_upsert(tenant_id, kb_id, new_cluster)

            nav_doc_row = _make_nav_doc_row(
                kb_id,
                doc_id,
                summary,
                new_name,
                new_depth,
                embd_mdl,
                doc_embedding,
                graph_content=graph_content,
            )
            await _store_upsert(tenant_id, kb_id, nav_doc_row)
        else:
            # ── Create root-level new cluster ──
            root_title, new_desc = await _llm_create_summary(chat_mdl, [summary])
            new_name = _readable_cluster_name(root_title, summary)
            new_cluster = _make_nav_cluster_row(
                kb_id,
                new_name,
                new_desc,
                "root",
                0,
                [doc_id],
                doc_embedding,
            )
            if embd_mdl and _vector_len(doc_embedding) > 0:
                new_cluster[_vec_field(len(doc_embedding))] = doc_embedding
            await _store_upsert(tenant_id, kb_id, new_cluster)

            nav_doc_row = _make_nav_doc_row(
                kb_id,
                doc_id,
                summary,
                new_name,
                1,
                embd_mdl,
                doc_embedding,
                graph_content=graph_content,
            )
            await _store_upsert(tenant_id, kb_id, nav_doc_row)

    except Exception:
        logging.exception(
            "dataset_nav: upsert failed for kb=%s doc=%s",
            kb_id,
            doc_id,
        )
    finally:
        try:
            lock.release()
        except Exception:
            logging.exception("dataset_nav: lock release failed for kb=%s", kb_id)


async def _remove_dataset_nav_doc_locked(
    tenant_id: str,
    kb_id: str,
    doc_id: str,
) -> None:
    """Remove a document nav row; the caller must hold the KB nav lock."""
    # 1. Find and delete the nav_doc row
    doc_row_id = _nav_doc_id(doc_id)
    doc_row = await _store_get(tenant_id, kb_id, doc_row_id)
    if not doc_row:
        return
    parent_name = doc_row.get("parent_kwd", "")
    await _store_delete(tenant_id, kb_id, doc_row_id)

    # 2. Remove doc_id from the parent cluster's doc_ids_kwd
    if parent_name and parent_name != "root":
        cluster_id = _nav_cluster_id(kb_id, parent_name)
        cluster_row = await _store_get(tenant_id, kb_id, cluster_id)
        if cluster_row:
            doc_ids = cluster_row.get("doc_ids_kwd") or []
            if doc_id in doc_ids:
                doc_ids.remove(doc_id)
            if not doc_ids:
                # Cluster is empty — delete it
                await _store_delete(tenant_id, kb_id, cluster_id)
                # Recurse: check grandparent
                grandparent = cluster_row.get("parent_kwd", "")
                if grandparent and grandparent != "root":
                    await _cleanup_empty_cluster(
                        tenant_id,
                        kb_id,
                        grandparent,
                    )
            else:
                cluster_row["doc_ids_kwd"] = doc_ids
                cluster_row["doc_count_int"] = len(doc_ids)
                await _store_upsert(tenant_id, kb_id, cluster_row)


async def remove_dataset_nav_doc(
    tenant_id: str,
    kb_id: str,
    doc_id: str,
) -> None:
    """Remove a document from the nav clustering tree.

    Cascades: if the parent cluster becomes empty after removal, the cluster
    itself is also removed.
    """
    if not doc_id or not kb_id:
        return

    lock = RedisDistributedLock(
        _nav_lock_key(kb_id),
        timeout=_LOCK_TIMEOUT_S,
        blocking_timeout=_LOCK_BLOCKING_TIMEOUT_S,
    )
    try:
        await lock.spin_acquire()
    except Exception:
        logging.exception("dataset_nav: lock acquire failed for kb=%s", kb_id)
        return

    try:
        await _remove_dataset_nav_doc_locked(tenant_id, kb_id, doc_id)
    except Exception:
        logging.exception(
            "dataset_nav: remove failed for kb=%s doc=%s",
            kb_id,
            doc_id,
        )
    finally:
        try:
            lock.release()
        except Exception:
            logging.exception("dataset_nav: lock release failed for kb=%s", kb_id)


async def _cleanup_empty_cluster(
    tenant_id: str,
    kb_id: str,
    cluster_name: str,
) -> None:
    """Recursively remove a cluster if it has no doc children and no direct doc descendants."""
    cluster_id = _nav_cluster_id(kb_id, cluster_name)
    cluster = await _store_get(tenant_id, kb_id, cluster_id)
    if not cluster:
        return
    # Check direct children (nav_cluster)
    from common import settings
    from common.doc_store.doc_store_base import OrderByExpr

    index = _index_name(tenant_id)
    child_cond = {
        "kb_id": [kb_id],
        "compile_kwd": [_COMPILE_KWD],
        "parent_kwd": [cluster_name],
    }
    res = await thread_pool_exec(
        settings.docStoreConn.search,
        ["id"],
        [],
        child_cond,
        [],
        OrderByExpr(),
        0,
        100,
        index,
        [kb_id],
    )
    children = settings.docStoreConn.get_fields(res, ["id"]) if res else {}
    if not children and not cluster.get("doc_ids_kwd"):
        grandparent = cluster.get("parent_kwd", "")
        await _store_delete(tenant_id, kb_id, cluster_id)
        if grandparent and grandparent != "root":
            await _cleanup_empty_cluster(tenant_id, kb_id, grandparent)


async def _maybe_split_cluster(
    tenant_id: str,
    kb_id: str,
    cluster_name: str,
    embd_mdl,
    chat_mdl,
) -> None:
    """If a cluster exceeds fanout or doc count, split children via AHC."""
    from common import settings
    from common.doc_store.doc_store_base import OrderByExpr

    index = _index_name(tenant_id)

    # Count children (nav_cluster)
    child_cond = {
        "kb_id": [kb_id],
        "compile_kwd": [_COMPILE_KWD],
        "parent_kwd": [cluster_name],
    }
    res = await thread_pool_exec(
        settings.docStoreConn.search,
        ["id", "name", "type_kwd"],
        [],
        child_cond,
        [],
        OrderByExpr(),
        0,
        200,
        index,
        [kb_id],
    )
    children = settings.docStoreConn.get_fields(res, ["id", "name", "type_kwd"]) if res else {}
    if not children:
        return

    nav_cluster_kids = [c for c in children.values() if c.get("type_kwd") == "nav_cluster"]
    nav_doc_kids = [c for c in children.values() if c.get("type_kwd") == "nav_doc"]

    should_split = len(nav_cluster_kids) + len(nav_doc_kids) > _MAX_FANOUT or len(nav_doc_kids) > _MAX_DOCS_PER_CLUSTER
    if not should_split:
        return

    # Load embeddings for all children
    vf = _vec_field(_EMBED_DIM) if _EMBED_DIM else "q_768_vec"
    child_details = await _store_search(
        tenant_id,
        kb_id,
        child_cond,
        ["id", "name", "type_kwd", "content_with_weight", vf],
        limit=200,
    )
    embeddings = []
    names = []
    name_to_type: dict[str, str] = {}
    for c in child_details:
        stored = c.get(vf)
        if stored:
            embeddings.append(stored)
            names.append(c.get("name", ""))
        cn = c.get("name", "")
        if cn:
            name_to_type[cn] = c.get("type_kwd", "nav_cluster")

    if len(embeddings) < 4:
        return

    # Simple k-means-like split into 2 groups (no scikit dependency at runtime)
    # Use the first two embeddings as initial centroids
    centroids = [embeddings[0][:], embeddings[len(embeddings) // 2][:]]
    for _ in range(10):
        groups = [[], []]
        for emb in embeddings:
            d0 = sum((a - b) ** 2 for a, b in zip(emb, centroids[0]))
            d1 = sum((a - b) ** 2 for a, b in zip(emb, centroids[1]))
            groups[0 if d0 < d1 else 1].append(emb)
        for gi in (0, 1):
            if groups[gi]:
                avg = [sum(c) / len(groups[gi]) for c in zip(*groups[gi])]
                centroids[gi] = avg

    # Relabel
    labels = []
    for emb in embeddings:
        d0 = sum((a - b) ** 2 for a, b in zip(emb, centroids[0]))
        d1 = sum((a - b) ** 2 for a, b in zip(emb, centroids[1]))
        labels.append(0 if d0 < d1 else 1)

    # Create sub-clusters
    cluster_row = await _store_get(tenant_id, kb_id, _nav_cluster_id(kb_id, cluster_name))
    depth = (cluster_row.get("depth_int", 0) if cluster_row else 0) + 1

    for gi in (0, 1):
        kid_names = [names[i] for i in range(len(names)) if labels[i] == gi]
        if not kid_names:
            continue
        # Collect doc_ids from all children
        doc_ids: list[str] = []
        descs: list[str] = []
        for kn in kid_names:
            is_doc = name_to_type.get(kn) == "nav_doc"
            cid = _nav_doc_id(kn) if is_doc else _nav_cluster_id(kb_id, kn)
            row = await _store_get(tenant_id, kb_id, cid)
            if row:
                payload = json.loads(row.get("content_with_weight") or "{}")
                descs.append(payload.get("description", ""))
                dids = row.get("doc_ids_kwd") or []
                for d in dids:
                    if d not in doc_ids:
                        doc_ids.append(d)
        if descs:
            group_title, group_desc = await _llm_create_summary(chat_mdl, descs)
        else:
            group_title = group_desc = f"Group {gi + 1}"
        group_name = _readable_cluster_name(group_title, group_desc)
        group_emb = await _embed(embd_mdl, group_desc) if embd_mdl else []
        new_cluster = _make_nav_cluster_row(
            kb_id,
            group_name,
            group_desc,
            cluster_name,
            depth,
            doc_ids,
            group_emb,
        )
        await _store_upsert(tenant_id, kb_id, new_cluster)

        # Reparent children to new split cluster
        for kn in kid_names:
            is_doc = name_to_type.get(kn) == "nav_doc"
            cid = _nav_doc_id(kn) if is_doc else _nav_cluster_id(kb_id, kn)
            row = await _store_get(tenant_id, kb_id, cid)
            if row:
                row["parent_kwd"] = group_name
                row["depth_int"] = depth + 1
                await _store_upsert(tenant_id, kb_id, row)


async def search_dataset_nav(
    tenant_id: str,
    kb_id: str,
    query: str,
    embd_mdl=None,
    top_k: int | None = None,
    *,
    type_kwd: str = "",
    compile_kwd: str = _COMPILE_KWD,
) -> list[dict]:
    """Find the nav-tree nodes most relevant to ``query`` for one KB.

    The nav rows are ``available_int=0`` (invisible to the normal retriever), so
    this is the sanctioned read seam: a caller uses the returned document ids to
    route a scoped chunk retrieval.

    Args:
        type_kwd: Optional type filter — ``"nav_doc"`` to restrict to document
            leaves, ``"nav_cluster"`` to restrict to clusters, or ``""`` for all.
        compile_kwd: Which compile partition to search within
            (default ``_COMPILE_KWD`` = ``"dataset_nav"``).

    Returns items shaped as::

        {"type": "nav_doc" | "nav_cluster",
         "doc_id": str | None,      # the document, for a leaf
         "doc_ids": [str],          # the documents a node covers
         "name": str, "description": str, "score": float}

    Ranked by hybrid search: KNN over the node-summary vectors fused with a
    BM25 (full-text) leg over the tokenized fields. When ``embd_mdl`` is absent
    the lexical leg alone ranks the results.
    """
    query = (query or "").strip()
    if not query:
        return []

    condition: dict = {"compile_kwd": [compile_kwd]}
    if type_kwd:
        condition["type_kwd"] = type_kwd
    # name -> [row, fused_score]; `name` uniquely identifies a nav node, so it
    # is the dedup key when the dense and lexical legs return the same node.
    fused: dict[str, list] = {}
    dense_w = _NAV_HYBRID_DENSE_W if embd_mdl is not None else 0.0

    # ── Dense leg: KNN over the node-summary vectors ──
    if embd_mdl is not None:
        try:
            vec = await _embed(embd_mdl, query)
        except Exception:
            logging.exception("search_dataset_nav: embed failed for kb=%s", kb_id)
            vec = []
        if _vector_len(vec) > 0:
            try:
                rows = await _store_knn(tenant_id, kb_id, vec, len(vec), condition, top_k=top_k or 10000)
                vf = _vec_field(len(vec))
                for r in rows:
                    rk = r.get("name") or r.get("doc_id") or ""
                    if not rk:
                        continue
                    fused.setdefault(rk, [r, 0.0])[1] += dense_w * _cosine_sim(vec, r.get(vf))
            except Exception:
                logging.exception("search_dataset_nav: knn failed for kb=%s", kb_id)

    # ── Lexical leg: engine BM25 over the tokenized fields ──
    text_w = 1.0 - dense_w
    try:
        text_rows = await _store_text_search(tenant_id, kb_id, query, _NAV_SEARCH_FIELDS, limit=max((top_k or 0) * 3, 20) if top_k else 10000, compile_kwd=compile_kwd, type_kwd=type_kwd)
    except Exception:
        logging.exception("search_dataset_nav: text search failed for kb=%s", kb_id)
        text_rows = []
    for r in text_rows:
        rk = r.get("name") or r.get("doc_id") or ""
        if not rk:
            continue
        ts = _nav_text_score(query, r)
        if ts <= 0:
            continue
        fused.setdefault(rk, [r, 0.0])[1] += text_w * ts

    rows_with_scores = [(r, s) for r, s in fused.values() if s > 0]
    rows_with_scores.sort(key=lambda item: item[1], reverse=True)
    if top_k is not None:
        rows_with_scores = rows_with_scores[:top_k]

    out: list[dict] = []
    for r, score in rows_with_scores:
        try:
            payload = json.loads(r.get("content_with_weight") or "{}")
        except Exception:
            payload = {}
        typ = payload.get("type") or r.get("type_kwd") or ("nav_cluster" if r.get("doc_ids_kwd") else "nav_doc")
        name = r.get("name") or ""
        if typ == "nav_cluster":
            doc_id = None
            doc_ids = _as_str_list(r.get("doc_ids_kwd"))
        else:
            # Leaf: ``name`` == the document id (see ``_make_nav_doc_row``).
            doc_id = r.get("doc_id") or name
            doc_ids = [doc_id] if doc_id else []
        out.append(
            {
                "type": typ,
                "doc_id": doc_id,
                "doc_ids": doc_ids,
                "name": name,
                "description": payload.get("description") or "",
                "keywords": _as_str_list(payload.get("keywords")),
                "entities": _as_str_list(payload.get("entities")),
                "graph_content": payload.get("graph_content") or "",
                "doc_title": payload.get("doc_title") or "",
                "source_type": payload.get("source_type") or "",
                "doc_count": int(r.get("doc_count_int") or len(doc_ids) or 0),
                "score": float(score or 0.0),
            }
        )
    return out


async def search_nav_tree_descent(
    tenant_id: str,
    kb_id: str,
    query: str,
    embd_mdl,
    top_k: int | None = None,
) -> list[dict]:
    """Tree-structured hybrid search: descend from root into the most relevant branches.

    Unlike ``search_dataset_nav`` (flat hybrid search over all nav nodes),
    this respects the parent-child hierarchy: it starts at root clusters,
    then at each level descends into the semantically closest children,
    collecting document identifiers from ``nav_doc`` leaves along the way.

    Each level uses hybrid search (KNN + BM25 text) fused with the same
    weights as ``search_dataset_nav``, so both vector similarity *and*
    lexical matching drive the descent.

    The search uses BFS with beam pruning — at each depth, only the
    *beam_width* most similar clusters are expanded further.

    Returns items shaped as ``{"doc_id": str, "score": float}``.
    """
    query = (query or "").strip()
    if not query:
        return []
    if embd_mdl is None:
        logging.warning(
            "search_nav_tree_descent: embd_mdl is None — falling back to text-only flat search for kb=%s query=%.80s",
            kb_id,
            query,
        )
        raw = await search_dataset_nav(tenant_id, kb_id, query, embd_mdl=None, top_k=top_k, type_kwd="nav_doc")
        return [{"doc_id": r.get("doc_id", ""), "score": r.get("score", 0.0)} for r in raw if r.get("doc_id")]

    vec = await _embed(embd_mdl, query)
    vec_dim = _vector_len(vec)
    if vec_dim == 0:
        return []

    beam_width = 5
    dense_w = _NAV_HYBRID_DENSE_W  # same weight as search_dataset_nav
    vf = _vec_field(vec_dim)

    fields = [
        "content_with_weight",
        "name",
        "doc_id",
        "compile_kwd",
        "type_kwd",
        "parent_kwd",
        "depth_int",
        "doc_count_int",
        "doc_ids_kwd",
        vf,
    ]

    collected: list[dict] = []
    seen_docs: set[str] = set()
    seen_nodes: set[str] = set()

    # ── Root clusters (KNN only — no text) ──
    # Root clusters at depth=0 are LLM-generated broad-topic summaries.
    # Specific query terms (e.g. "British") almost never appear in their
    # name / description, so we route purely by vector similarity at this
    # level.  Text matching takes over from depth ≥ 1 where cluster
    # descriptions contain concrete terms.
    root_cond = {
        "kb_id": [kb_id],
        "compile_kwd": [_COMPILE_KWD],
        "type_kwd": ["nav_cluster"],
        "depth_int": [0],
    }
    roots_knn = await _store_knn(tenant_id, kb_id, vec, vec_dim, root_cond, top_k=beam_width * 3)

    # If no root cluster (depth=0) exists — the dataset may have been
    # compiled without one — scan all nav_clusters to find the lowest
    # available depth and start beam search there.
    if not roots_knn:
        all_cond = {
            "kb_id": [kb_id],
            "compile_kwd": [_COMPILE_KWD],
            "type_kwd": ["nav_cluster"],
        }
        all_clusters = await _store_search(tenant_id, kb_id, all_cond, fields, limit=10000)
        if not all_clusters:
            return []

        min_depth = min(
            (r.get("depth_int", 0) for r in all_clusters if r.get("depth_int") is not None),
            default=0,
        )
        logging.warning(
            "search_nav_tree_descent: no root cluster at depth=0, starting beam search from depth=%d",
            min_depth,
        )

        starters = [r for r in all_clusters if r.get("depth_int") == min_depth]
        starters.sort(key=lambda r: _cosine_sim(vec, r.get(vf)), reverse=True)
        current_level = starters[:beam_width]
        for r in current_level:
            r["_score"] = _cosine_sim(vec, r.get(vf))
    else:
        # Route down to beam_width semantically closest root clusters.
        for r in roots_knn:
            r["_score"] = _cosine_sim(vec, r.get(vf))
        roots_knn.sort(key=lambda r: r["_score"], reverse=True)
        current_level = roots_knn[:beam_width]

    if not current_level:
        return []

    while current_level and (top_k is None or len(collected) < top_k):
        next_level: list[dict] = []

        for node in current_level:
            node_name = node.get("name", "")
            if node_name in seen_nodes:
                continue
            seen_nodes.add(node_name)
            parent_score = node.get("_score", 0.0)

            child_cond: dict = {
                "kb_id": [kb_id],
                "compile_kwd": [_COMPILE_KWD],
                "parent_kwd": [node_name],
            }
            children_knn = await _store_knn(tenant_id, kb_id, vec, vec_dim, child_cond, top_k=beam_width * 3)
            children_text = await _store_text_search(
                tenant_id,
                kb_id,
                query,
                fields,
                limit=beam_width * 3,
                extra_filter={"parent_kwd": [node_name], "kb_id": [kb_id]},
            )
            candidates = _hybrid_fuse(vec, vf, query, children_knn, children_text, dense_w, beam_width)

            for c in candidates:
                if top_k is not None and len(collected) >= top_k:
                    break
                if c.get("type_kwd") == "nav_doc":
                    doc_id = (c.get("doc_id") or "").strip()
                    if doc_id and doc_id not in seen_docs:
                        seen_docs.add(doc_id)
                        collected.append({"doc_id": doc_id, "score": round(c["_score"] or parent_score, 4)})
                else:
                    next_level.append(c)

            if top_k is not None and len(collected) >= top_k:
                break

        if top_k is not None and len(collected) >= top_k:
            break

        next_level.sort(key=lambda c: c.get("_score", 0.0), reverse=True)
        current_level = next_level[:beam_width]

    # Fallback: collect doc_ids from terminal cluster nodes.
    if not collected and current_level:
        for node in current_level:
            for did in node.get("doc_ids_kwd") or []:
                did_str = str(did).strip()
                if did_str and did_str not in seen_docs:
                    seen_docs.add(did_str)
                    collected.append({"doc_id": did_str, "score": round(node.get("_score", 0.0), 4)})
                    if top_k is not None and len(collected) >= top_k:
                        break
            if top_k is not None and len(collected) >= top_k:
                break

    return collected


def _hybrid_fuse(
    vec: list[float],
    vf: str,
    query: str,
    knn_rows: list[dict],
    text_rows: list[dict],
    dense_w: float,
    top_k: int,
) -> list[dict]:
    """Fuse KNN and text results into a scored, deduplicated, sorted list.

    ``_store_knn`` already enforces filter conditions on the KNN leg;
    text rows come from ``_store_text_search`` which applies its own
    filters.  Here we merge both legs by ``name``/``doc_id``.
    """
    text_w = 1.0 - dense_w
    fused: dict[str, tuple[dict, float]] = {}  # key → (row, score)

    # KNN leg: raw cosine * dense_w, matching search_dataset_nav's scale so
    # both search paths score the dense leg identically.
    for r in knn_rows:
        rk = r.get("name") or r.get("doc_id") or ""
        if not rk:
            continue
        key = f"knn:{rk}"
        fused[key] = (r, _cosine_sim(vec, r.get(vf, [])) * dense_w)

    # Text leg.
    for r in text_rows:
        rk = r.get("name") or r.get("doc_id") or ""
        if not rk:
            continue
        ts = _nav_text_score(query, r)
        if ts <= 0:
            continue
        key = f"text:{rk}"
        if key in fused:
            fused[key] = (fused[key][0], fused[key][1] + text_w * ts)
        else:
            key_knn = f"knn:{rk}"
            if key_knn in fused:
                fused[key_knn] = (fused[key_knn][0], fused[key_knn][1] + text_w * ts)
            else:
                fused[key] = (r, text_w * ts)

    rows_with_scores = [(r, s) for r, s in fused.values() if s > 0]
    rows_with_scores.sort(key=lambda item: item[1], reverse=True)
    result = rows_with_scores[:top_k]
    for r, s in result:
        r["_score"] = s
    return [r for r, _ in result]


def _as_str_list(value) -> list[str]:
    if isinstance(value, list):
        return [str(v) for v in value if v]
    if isinstance(value, str) and value:
        return [value]
    return []


def _nav_text_score(query: str, row: dict) -> float:
    try:
        payload = json.loads(row.get("content_with_weight") or "{}")
    except Exception:
        payload = {}
    keywords = payload.get("keywords") or []
    entities = payload.get("entities") or []
    graph_content = payload.get("graph_content") or ""
    haystack = " ".join(
        str(x or "")
        for x in (
            row.get("name"),
            payload.get("description"),
            graph_content,
            *keywords,
            *entities,
        )
    ).lower()
    q_terms = set(re.findall(r"[\w]+", query.lower()))
    if not q_terms:
        return 0.0
    hits = sum(1 for term in q_terms if term in haystack)
    return hits / max(len(q_terms), 1)


def remove_dataset_nav_doc_sync(
    tenant_id: str,
    kb_id: str,
    doc_id: str,
) -> None:
    """Sync wrapper around ``remove_dataset_nav_doc``."""
    try:
        loop = asyncio.new_event_loop()
        try:
            loop.run_until_complete(
                remove_dataset_nav_doc(tenant_id, kb_id, doc_id),
            )
        finally:
            loop.close()
    except Exception:
        logging.exception(
            "dataset_nav: sync remove failed for kb=%s doc=%s",
            kb_id,
            doc_id,
        )