diff --git a/rag/advanced_rag/knowlege_compile/raptor.py b/rag/advanced_rag/knowlege_compile/raptor.py index 9dff4a749c..92618dfa49 100644 --- a/rag/advanced_rag/knowlege_compile/raptor.py +++ b/rag/advanced_rag/knowlege_compile/raptor.py @@ -19,7 +19,7 @@ import logging import re import numpy as np -from sklearn.cluster import AgglomerativeClustering + from sklearn.mixture import GaussianMixture from api.db.services.task_service import has_canceled @@ -42,7 +42,6 @@ from rag.utils.raptor_utils import ( SUPPORTED_CLUSTERING_METHODS, SUPPORTED_TREE_BUILDERS, ) -from ._common import knowledge_compile_gen_conf # Regularization added to GMM covariance diagonals; keeps components # from collapsing on singleton/near-identical reduced points. @@ -180,10 +179,20 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: clustering_method=GMM_CLUSTERING_METHOD, psi_exact_max_leaves=4096, psi_bucket_size=1024, + cluster_percentile=30, ): - """Configure RAPTOR summarization, clustering, and Psi limits.""" + """Configure RAPTOR summarization, clustering, and Psi limits. + + Args: + cluster_percentile: AHC distance threshold is set to this + percentile of all pairwise cosine distances in each + layer. A lower value produces finer (more) clusters. + Default 30 means the threshold excludes the top 70% + most dissimilar pairs. + """ self._max_cluster = max_cluster self._small_layer_collapse = small_layer_collapse + self._cluster_percentile = cluster_percentile self._llm_model = llm_model self._embd_model = embd_model self._threshold = threshold @@ -217,7 +226,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: last_exc = None for attempt in range(3): try: - response = await self._llm_model.async_chat(system, history, knowledge_compile_gen_conf(self._llm_model, gen_conf)) + response = await self._llm_model.async_chat(system, history, gen_conf) response = re.sub(r"^.*", "", response, flags=re.DOTALL) if response.find("**ERROR**") >= 0: raise Exception(response) @@ -266,92 +275,92 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: return int(optimal_clusters) def _get_clusters_ahc(self, embeddings: np.ndarray, task_id: str = "") -> np.ndarray: - """Cluster embeddings with Ward-linkage AHC and a dendrogram gap heuristic.""" + """Sequential clustering of adjacent embeddings (cosine similarity). + + Only compares **adjacent** pairs (chunk i vs chunk i+1), not all + pairwise — ``O(N)`` complexity per layer instead of ``O(N²)``. + + The similarity threshold is the ``p``-th percentile of all + adjacent-pair similarities in the current layer, so it adapts + to each layer's data distribution automatically. + + Returns an array of cluster labels (contiguous 0..K-1). + """ n = len(embeddings) if n <= 1: return np.zeros(n, dtype=int) if n == 2: return np.arange(n) - self._check_task_canceled(task_id, "_get_clusters_ahc dendrogram") - full_clust = AgglomerativeClustering( - n_clusters=None, - distance_threshold=0, - compute_distances=True, - linkage="ward", - ) - full_clust.fit(embeddings) + self._check_task_canceled(task_id, "_get_clusters_ahc") - distances = full_clust.distances_ - if len(distances) > 1: - gaps = np.diff(distances) - max_gap_idx = int(np.argmax(gaps)) - n_clusters = max(1, min(n - max_gap_idx - 1, self._max_cluster)) - else: - n_clusters = max(1, min(n, self._max_cluster)) - if n_clusters <= 1: - logging.info("RAPTOR AHC: _get_clusters_ahc selected one cluster for %d embeddings", n) + # L2-normalize embeddings so dot product = cosine similarity + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + norms = np.where(norms == 0, 1.0, norms) + normalized = embeddings / norms + + # Adjacent cosine similarities (n-1 pairs) + adj_sims = np.sum(normalized[:-1] * normalized[1:], axis=1) + if len(adj_sims) == 0: return np.zeros(n, dtype=int) - logging.info("RAPTOR AHC: _get_clusters_ahc selected n_clusters=%d for %d embeddings", n_clusters, n) - self._check_task_canceled(task_id, "_get_clusters_ahc fit") - clustering = AgglomerativeClustering(n_clusters=n_clusters, linkage="ward") - return clustering.fit_predict(embeddings) + # Adaptive threshold from adjacent distribution + threshold = float(np.percentile(adj_sims, self._cluster_percentile)) + labels = np.zeros(n, dtype=int) + cluster_id = 0 + for i in range(1, n): + if adj_sims[i - 1] >= threshold: + labels[i] = cluster_id + else: + cluster_id += 1 + labels[i] = cluster_id - def _adjust_tree_nodes(self, embeddings: np.ndarray, labels: np.ndarray, max_iter: int = 5) -> np.ndarray: - """Refine AHC assignments by reassigning nodes to nearest centroids.""" - labels = labels.copy() - for _ in range(max_iter): - unique_labels = np.unique(labels) - if len(unique_labels) <= 1: - return labels - centroids = np.stack([embeddings[labels == lbl].mean(axis=0) for lbl in unique_labels]) - diffs = embeddings[:, np.newaxis, :] - centroids[np.newaxis, :, :] - sq_dists = (diffs**2).sum(axis=2) - new_label_indices = np.argmin(sq_dists, axis=1) - new_labels = unique_labels[new_label_indices] - if np.array_equal(new_labels, labels): - break - unique_new = np.unique(new_labels) - remap = {old: new for new, old in enumerate(unique_new)} - labels = np.array([remap[int(lbl)] for lbl in new_labels]) + logging.info( + "RAPTOR seq-clus: p=%d threshold=%.4f n_clusters=%d for %d embeddings (adj pairs=%d)", + self._cluster_percentile, + threshold, + int(np.unique(labels).size), + n, + len(adj_sims), + ) return labels def clustering(self, embeddings, random_state: int, task_id: str = "") -> tuple[int, list[int]]: """Cluster one RAPTOR layer and return contiguous labels.""" - reduced_embeddings = np.asarray(embeddings, dtype=np.float64) - if len(reduced_embeddings) == 0: + if len(embeddings) == 0: return 0, [] - # Degrade too much ?? - n_neighbors = min(int((len(embeddings) - 1) ** 0.8), 100) - import umap - - reduced_embeddings = umap.UMAP( - n_neighbors=max(2, n_neighbors), - n_components=min(12, len(embeddings) - 2), - metric="cosine", - ).fit_transform(embeddings) if self._clustering_method == AHC_CLUSTERING_METHOD: - logging.info("RAPTOR: using clustering_method=%s before _get_clusters_ahc", self._clustering_method) - raw_labels = self._get_clusters_ahc(reduced_embeddings, task_id=task_id) + # AHC: cluster on raw embeddings with cosine distance. + # UMAP is skipped because it discards semantic information + # that average-linkage + cosine can leverage directly. + logging.info("RAPTOR: using clustering_method=%s on raw embeddings (dim=%d)", self._clustering_method, len(embeddings[0]) if hasattr(embeddings[0], "__len__") else "?") + asarray = np.asarray(embeddings, dtype=np.float64) + raw_labels = self._get_clusters_ahc(asarray, task_id=task_id) raw_cluster_count = np.unique(raw_labels).size logging.info("RAPTOR AHC: _get_clusters_ahc produced n_clusters=%d", raw_cluster_count) - if raw_cluster_count > 1: - labels = self._adjust_tree_nodes(reduced_embeddings, raw_labels) - adjusted_cluster_count = np.unique(labels).size - logging.info("RAPTOR AHC: _adjust_tree_nodes adjusted n_clusters=%d", adjusted_cluster_count) - else: - labels = raw_labels - logging.warning("RAPTOR AHC: _adjust_tree_nodes skipped because _get_clusters_ahc returned one cluster") + labels = raw_labels else: - n_clusters = int(self._get_optimal_clusters(reduced_embeddings, random_state, task_id=task_id)) + # GMM: reduce dimensionality first (UMAP, 12D) so the + # Gaussian mixture can find meaningful clusters. + if len(embeddings) == 0: + return 0, [] + reduced = np.asarray(embeddings, dtype=np.float64) + n_neighbors = min(int((len(embeddings) - 1) ** 0.8), 100) + import umap + + reduced = umap.UMAP( + n_neighbors=max(2, n_neighbors), + n_components=min(12, len(embeddings) - 2), + metric="cosine", + ).fit_transform(embeddings) + n_clusters = int(self._get_optimal_clusters(reduced, random_state, task_id=task_id)) if n_clusters <= 1: - labels = [0 for _ in range(len(reduced_embeddings))] + labels = [0 for _ in range(len(reduced))] else: gm = GaussianMixture(n_components=n_clusters, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR) - gm.fit(reduced_embeddings) - probs = gm.predict_proba(reduced_embeddings) + gm.fit(reduced) + probs = gm.predict_proba(reduced) labels = [] for prob in probs: candidates = np.where(prob > self._threshold)[0] diff --git a/rag/flow/compiler/compiler.py b/rag/flow/compiler/compiler.py index 081ac53345..82bf990359 100644 --- a/rag/flow/compiler/compiler.py +++ b/rag/flow/compiler/compiler.py @@ -133,7 +133,7 @@ class Compiler(ProcessBase, LLM): chat_mdl=chat_mdl_by_tid[template_id], embd_mdl=embedding_model, tree_builder="raptor", - clustering_method="gmm", + clustering_method="ahc", max_errors=3, ) except Exception: diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py index b97f667cdf..61956ee841 100644 --- a/rag/svr/task_executor_refactor/chunk_post_processor.py +++ b/rag/svr/task_executor_refactor/chunk_post_processor.py @@ -476,6 +476,39 @@ def raptor_tree_to_graph(tree: Dict) -> Dict: entities: list[dict] = [] relations: list[dict] = [] + def _collapse_unary(node: dict) -> dict: + """Collapse tree nodes that only wrap one child.""" + collapsed = dict(node) + collapsed["children"] = [_collapse_unary(child) for child in node.get("children") or [] if isinstance(child, dict)] + + while len(collapsed["children"]) == 1: + child = collapsed["children"][0] + parent_title = collapsed.get("title") or "" + child_title = child.get("title") or "" + parent_description = collapsed.get("description") or parent_title + child_description = child.get("description") or child_title + + descriptions = [str(parent_description)] + if child_title and child_title != parent_title and child_title not in child_description: + descriptions.append(str(child_title)) + if child_description and child_description not in descriptions: + descriptions.append(str(child_description)) + + source_chunk_ids = [] + for source in (collapsed.get("source_chunk_ids") or [], child.get("source_chunk_ids") or []): + for chunk_id in source: + if isinstance(chunk_id, str) and chunk_id and chunk_id not in source_chunk_ids: + source_chunk_ids.append(chunk_id) + + collapsed["description"] = "\n\n".join(descriptions) + if source_chunk_ids: + collapsed["source_chunk_ids"] = source_chunk_ids + collapsed["children"] = child.get("children") or [] + + return collapsed + + tree = _collapse_unary(tree) if isinstance(tree, dict) else tree + def _walk(node: dict, parent_id: Optional[str]) -> None: if not isinstance(node, dict): return @@ -491,7 +524,10 @@ def raptor_tree_to_graph(tree: Dict) -> Dict: if isinstance(src_ids, list) and src_ids: ent["source_chunk_ids"] = [s for s in src_ids if isinstance(s, str) and s] entities.append(ent) - if parent_id is not None: + # A summary and its child can occasionally receive the same LLM-generated + # title. They are still valid tree nodes, but must not become a self-loop + # when the tree is projected to graph relations. + if parent_id is not None and parent_id != node_id: relations.append({"from": parent_id, "to": node_id, "type": "child"}) for child in node.get("children") or []: _walk(child, node_id) @@ -500,6 +536,77 @@ def raptor_tree_to_graph(tree: Dict) -> Dict: return {"entities": entities, "relations": relations} +async def rewrite_duplicate_tree_names(tree: Dict, chat_mdl) -> None: + """Rewrite only duplicate tree titles whose descriptions differ.""" + from rag.advanced_rag.knowlege_compile._common import knowledge_compile_gen_conf + from rag.prompts.generator import gen_json + + groups: dict[str, list[tuple[dict, str, str]]] = {} + + def _walk(node: dict, path: tuple[int, ...]) -> None: + if not isinstance(node, dict): + return + title = str(node.get("title") or "").strip() + if title: + description = str(node.get("description") or title).strip() + node_key = ".".join(str(index) for index in path) + groups.setdefault(title, []).append((node, node_key, description)) + for index, child in enumerate(node.get("children") or []): + _walk(child, (*path, index)) + + _walk(tree, (0,)) + for title, candidates in groups.items(): + descriptions = {description for _, _, description in candidates} + if len(candidates) < 2 or len(descriptions) < 2: + continue + + items = [{"id": node_key, "description": description} for _, node_key, description in candidates] + prompt = ( + "The following tree nodes currently have the same title but describe different content. " + "Give each node a concise, distinct human-readable title. Preserve the original language, " + "do not add numbering unless necessary, and return only a JSON array of objects with the " + "same ids and a name field.\n\n" + f"Current title: {title}\n" + f"Nodes: {json.dumps(items, ensure_ascii=False)}" + ) + try: + result = await gen_json( + "You rename duplicate tree node titles for display.", + prompt, + chat_mdl, + gen_conf=knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}), + ) + except Exception: + logging.exception("tree-template: duplicate title rewrite failed for title=%s", title) + continue + + rewrites = {} + if isinstance(result, list): + rewrites = {str(item.get("id")): str(item.get("name")).strip() for item in result if isinstance(item, dict) and item.get("id") and str(item.get("name") or "").strip()} + for node, node_key, _ in candidates: + new_title = rewrites.get(node_key) + if new_title: + node["title"] = new_title + + # The LLM is asked to produce distinct names, but enforce that contract + # deterministically before the graph uses titles as relation endpoints. + used_names: dict[str, int] = {} + + def _ensure_unique(node: dict) -> None: + if not isinstance(node, dict): + return + title = str(node.get("title") or "").strip() + if title: + occurrence = used_names.get(title, 0) + 1 + used_names[title] = occurrence + if occurrence > 1: + node["title"] = f"{title} ({occurrence})" + for child in node.get("children") or []: + _ensure_unique(child) + + _ensure_unique(tree) + + async def load_chunks_with_vec( tenant_id: str, kb_id: str, @@ -840,7 +947,7 @@ async def run_tree_templates( chat_mdl=chat_mdl_by_tid[template_id], embd_mdl=embedding_model, tree_builder="raptor", - clustering_method="gmm", + clustering_method="ahc", max_errors=3, ) except Exception: @@ -873,6 +980,7 @@ async def run_tree_templates( doc_id, ) + await rewrite_duplicate_tree_names(tree, chat_mdl_by_tid[template_id]) graph = raptor_tree_to_graph(tree) try: await _struct_upsert_graph_json( diff --git a/test/unit_test/rag/test_raptor_psi_tree_builder.py b/test/unit_test/rag/test_raptor_psi_tree_builder.py index de8465b3ea..a2c2178bda 100644 --- a/test/unit_test/rag/test_raptor_psi_tree_builder.py +++ b/test/unit_test/rag/test_raptor_psi_tree_builder.py @@ -18,8 +18,6 @@ import importlib import os import sys import types -from unittest.mock import MagicMock - import pytest np = pytest.importorskip("numpy") @@ -54,11 +52,12 @@ def raptor_module(monkeypatch): return np.ones((len(embeddings), 1)) class DummyAgglomerativeClustering: - def __init__(self, n_clusters=None, distance_threshold=None, compute_distances=False, linkage="ward"): + def __init__(self, n_clusters=None, distance_threshold=None, compute_distances=False, linkage="ward", metric="euclidean"): self.n_clusters = n_clusters self.distance_threshold = distance_threshold self.compute_distances = compute_distances self.linkage = linkage + self.metric = metric self.distances_ = np.array([0.1, 0.2, 1.0]) def fit(self, embeddings): diff --git a/web/src/components/ui/tree-view.tsx b/web/src/components/ui/tree-view.tsx index 04568725b7..661694e521 100644 --- a/web/src/components/ui/tree-view.tsx +++ b/web/src/components/ui/tree-view.tsx @@ -142,7 +142,7 @@ const TreeItem = React.forwardRef(