mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-16 20:57:21 +08:00
983 lines
43 KiB
Python
983 lines
43 KiB
Python
#
|
|
# Copyright 2024 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.
|
|
#
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
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
|
|
from common.connection_utils import timeout
|
|
from common.exceptions import TaskCanceledException
|
|
from common.token_utils import truncate
|
|
from rag.graphrag.utils import (
|
|
chat_limiter,
|
|
get_embed_cache,
|
|
get_llm_cache,
|
|
set_embed_cache,
|
|
set_llm_cache,
|
|
)
|
|
from common.misc_utils import thread_pool_exec
|
|
from rag.utils.raptor_utils import (
|
|
AHC_CLUSTERING_METHOD,
|
|
GMM_CLUSTERING_METHOD,
|
|
PSI_TREE_BUILDER,
|
|
RAPTOR_TREE_BUILDER,
|
|
SUPPORTED_CLUSTERING_METHODS,
|
|
SUPPORTED_TREE_BUILDERS,
|
|
)
|
|
|
|
# Regularization added to GMM covariance diagonals; keeps components
|
|
# from collapsing on singleton/near-identical reduced points.
|
|
_GMM_REG_COVAR = 1e-4
|
|
|
|
|
|
@dataclass
|
|
class _PsiTreeNode:
|
|
"""Node used to represent the in-memory Psi merge tree."""
|
|
|
|
index: int
|
|
text: str = ""
|
|
embedding: np.ndarray | None = None
|
|
children: list["_PsiTreeNode"] = field(default_factory=list)
|
|
parent: "_PsiTreeNode | None" = None
|
|
# Original (leaf-level) chunk ids that contributed to this node. On
|
|
# a leaf this is a single-element list with the leaf's own id; on an
|
|
# internal node it's the order-preserving deduped union of its
|
|
# children's lists. Carried up through the merge tree so each
|
|
# produced summary knows which source chunks it covers.
|
|
source_chunk_ids: list[str] = field(default_factory=list)
|
|
|
|
|
|
class _PsiUnionFind:
|
|
"""Build parent links for the Psi merge tree from ranked leaf pairs."""
|
|
|
|
def __init__(self, n: int):
|
|
"""Initialize the union-find state for n leaf nodes."""
|
|
self._rank = [0 for _ in range(n)]
|
|
self._parent_chains = [[] for _ in range(n)]
|
|
self._node_ids = [[i] for i in range(n)]
|
|
self._tree = [-1 for _ in range(max(1, 2 * n - 1))]
|
|
self._next_id = n
|
|
|
|
@staticmethod
|
|
def _ordered_extend(target: list[int], values: list[int]):
|
|
"""Append unseen values while preserving their original order."""
|
|
for value in values:
|
|
if value not in target:
|
|
target.append(value)
|
|
|
|
def _find(self, i: int) -> list[int]:
|
|
"""Return the parent chain for a leaf, extending it lazily."""
|
|
chain = self._parent_chains[i]
|
|
if not chain or (len(chain) == 1 and chain[0] == i):
|
|
return [i]
|
|
if chain[0] == i:
|
|
self._ordered_extend(chain, self._find(chain[1]))
|
|
else:
|
|
self._ordered_extend(chain, self._find(chain[0]))
|
|
return chain
|
|
|
|
def _rank_bisect_right(self, chain: list[int], rank: int) -> int:
|
|
"""Return the first chain index whose rank is greater than rank."""
|
|
idx = 0
|
|
while idx < len(chain) and self._rank[chain[idx]] <= rank:
|
|
idx += 1
|
|
return idx
|
|
|
|
def _build(self, i: int, j: int, insert_point: int | None = None):
|
|
"""Record a merge edge in the compact parent array."""
|
|
if insert_point is not None:
|
|
parent_ids = self._node_ids[insert_point]
|
|
parent_rank_idx = self._rank[i] + 1
|
|
if parent_rank_idx >= len(parent_ids):
|
|
logging.warning(
|
|
"RAPTOR Psi union fallback: rank index %d is out of bounds for node %d with %d parent ids",
|
|
parent_rank_idx,
|
|
insert_point,
|
|
len(parent_ids),
|
|
)
|
|
parent_rank_idx = len(parent_ids) - 1
|
|
self._tree[self._node_ids[i][-1]] = parent_ids[parent_rank_idx]
|
|
return
|
|
self._tree[self._node_ids[i][-1]] = self._next_id
|
|
self._tree[self._node_ids[j][-1]] = self._next_id
|
|
self._node_ids[i].append(self._next_id)
|
|
self._next_id += 1
|
|
|
|
def union(self, i: int, j: int) -> bool:
|
|
"""Merge two ranked leaves and return whether a new edge was added."""
|
|
root_i = self._find(i)[-1]
|
|
root_j = self._find(j)[-1]
|
|
if root_i == root_j:
|
|
return False
|
|
|
|
if self._rank[root_i] < self._rank[root_j]:
|
|
if not self._parent_chains[root_j]:
|
|
self._parent_chains[root_j].append(root_j)
|
|
chain = self._parent_chains[j]
|
|
higher_rank_idx = self._rank_bisect_right(chain, self._rank[root_i])
|
|
if higher_rank_idx >= len(chain):
|
|
higher_rank_idx = len(chain) - 1
|
|
insert_point = chain[higher_rank_idx]
|
|
self._ordered_extend(self._parent_chains[root_i], chain[higher_rank_idx:])
|
|
self._build(root_i, root_j, insert_point=insert_point)
|
|
elif self._rank[root_i] > self._rank[root_j]:
|
|
if not self._parent_chains[root_i]:
|
|
self._parent_chains[root_i].append(root_i)
|
|
chain = self._parent_chains[i]
|
|
higher_rank_idx = self._rank_bisect_right(chain, self._rank[root_j])
|
|
if higher_rank_idx >= len(chain):
|
|
higher_rank_idx = len(chain) - 1
|
|
insert_point = chain[higher_rank_idx]
|
|
self._ordered_extend(self._parent_chains[root_j], chain[higher_rank_idx:])
|
|
self._build(root_j, root_i, insert_point=insert_point)
|
|
else:
|
|
if not self._parent_chains[root_i]:
|
|
self._parent_chains[root_i].append(root_i)
|
|
self._ordered_extend(self._parent_chains[root_j], self._parent_chains[i][-1:])
|
|
self._rank[root_i] += 1
|
|
self._build(root_i, root_j)
|
|
return True
|
|
|
|
@property
|
|
def tree(self) -> list[int]:
|
|
"""Return the compact child-to-parent array for constructed nodes."""
|
|
return self._tree[: self._next_id]
|
|
|
|
|
|
class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
|
|
"""Build RAPTOR summary layers with the classic or Psi tree strategy."""
|
|
|
|
def __init__(
|
|
self,
|
|
max_cluster,
|
|
llm_model,
|
|
embd_model,
|
|
prompt,
|
|
max_token=512,
|
|
threshold=0.1,
|
|
small_layer_collapse=8,
|
|
max_errors=3,
|
|
tree_builder=RAPTOR_TREE_BUILDER,
|
|
clustering_method=GMM_CLUSTERING_METHOD,
|
|
psi_exact_max_leaves=4096,
|
|
psi_bucket_size=1024,
|
|
):
|
|
"""Configure RAPTOR summarization, clustering, and Psi limits."""
|
|
self._max_cluster = max_cluster
|
|
self._small_layer_collapse = small_layer_collapse
|
|
self._llm_model = llm_model
|
|
self._embd_model = embd_model
|
|
self._threshold = threshold
|
|
self._prompt = prompt
|
|
self._max_token = max_token
|
|
self._max_errors = max(1, max_errors)
|
|
self._error_count = 0
|
|
self._tree_builder = tree_builder or RAPTOR_TREE_BUILDER
|
|
if self._tree_builder not in SUPPORTED_TREE_BUILDERS:
|
|
raise ValueError(f"Unsupported RAPTOR tree builder: {self._tree_builder}")
|
|
self._clustering_method = clustering_method or GMM_CLUSTERING_METHOD
|
|
if self._clustering_method not in SUPPORTED_CLUSTERING_METHODS:
|
|
raise ValueError(f"Unsupported RAPTOR clustering method: {self._clustering_method}")
|
|
self._psi_exact_max_leaves = max(2, int(psi_exact_max_leaves or 4096))
|
|
self._psi_bucket_size = min(max(2, int(psi_bucket_size or 1024)), self._psi_exact_max_leaves)
|
|
|
|
def _check_task_canceled(self, task_id: str, message: str = ""):
|
|
"""Raise if the current document task was canceled."""
|
|
if task_id and has_canceled(task_id):
|
|
log_msg = f"Task {task_id} cancelled during RAPTOR {message}."
|
|
logging.info(log_msg)
|
|
raise TaskCanceledException(f"Task {task_id} was cancelled")
|
|
|
|
@timeout(60 * 20)
|
|
async def _chat(self, system, history, gen_conf):
|
|
"""Call the configured LLM with caching and short retries."""
|
|
cached = await thread_pool_exec(get_llm_cache, self._llm_model.llm_name, system, history, gen_conf)
|
|
if cached:
|
|
return cached
|
|
|
|
last_exc = None
|
|
for attempt in range(3):
|
|
try:
|
|
response = await self._llm_model.async_chat(system, history, gen_conf)
|
|
response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
|
|
if response.find("**ERROR**") >= 0:
|
|
raise Exception(response)
|
|
await thread_pool_exec(set_llm_cache, self._llm_model.llm_name, system, response, history, gen_conf)
|
|
return response
|
|
except Exception as exc:
|
|
last_exc = exc
|
|
logging.warning("RAPTOR LLM call failed on attempt %d/3: %s", attempt + 1, exc)
|
|
if attempt < 2:
|
|
await asyncio.sleep(1 + attempt)
|
|
|
|
raise last_exc if last_exc else Exception("LLM chat failed without exception")
|
|
|
|
@timeout(20)
|
|
async def _embedding_encode(self, txt):
|
|
"""Encode text with the configured embedding model and cache result."""
|
|
response = await thread_pool_exec(get_embed_cache, self._embd_model.llm_name, txt)
|
|
if response is not None:
|
|
return response
|
|
embds, _ = await thread_pool_exec(self._embd_model.encode, [txt])
|
|
if len(embds) < 1 or len(embds[0]) < 1:
|
|
raise Exception("Embedding error: empty embeddings returned")
|
|
embds = embds[0]
|
|
await thread_pool_exec(set_embed_cache, self._embd_model.llm_name, txt, embds)
|
|
return embds
|
|
|
|
def _get_optimal_clusters(self, embeddings: np.ndarray, random_state: int, task_id: str = ""):
|
|
"""Choose the GMM cluster count with the lowest BIC score."""
|
|
max_clusters = min(self._max_cluster, len(embeddings))
|
|
if max_clusters <= 1:
|
|
logging.info(
|
|
"RAPTOR GMM: _get_optimal_clusters returning 1 (max_clusters=%s, embeddings=%d)",
|
|
max_clusters,
|
|
len(embeddings),
|
|
)
|
|
return 1
|
|
n_clusters = np.arange(1, max_clusters + 1)
|
|
bics = []
|
|
for n in n_clusters:
|
|
self._check_task_canceled(task_id, "get optimal clusters")
|
|
|
|
gm = GaussianMixture(n_components=n, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR)
|
|
gm.fit(embeddings)
|
|
bics.append(gm.bic(embeddings))
|
|
optimal_clusters = n_clusters[np.argmin(bics)]
|
|
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."""
|
|
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)
|
|
|
|
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)
|
|
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)
|
|
|
|
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])
|
|
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:
|
|
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)
|
|
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")
|
|
else:
|
|
n_clusters = int(self._get_optimal_clusters(reduced_embeddings, random_state, task_id=task_id))
|
|
if n_clusters <= 1:
|
|
labels = [0 for _ in range(len(reduced_embeddings))]
|
|
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)
|
|
labels = []
|
|
for prob in probs:
|
|
candidates = np.where(prob > self._threshold)[0]
|
|
labels.append(int(candidates[0]) if len(candidates) else int(np.argmax(prob)))
|
|
|
|
normalized_labels: list[int] = []
|
|
for label in labels:
|
|
if isinstance(label, np.ndarray):
|
|
normalized_labels.append(int(label[0]) if len(label) else 0)
|
|
else:
|
|
normalized_labels.append(int(label))
|
|
|
|
if len(normalized_labels) <= 0:
|
|
return 0, []
|
|
unique_labels = np.unique(normalized_labels)
|
|
if len(unique_labels) <= 1:
|
|
return 1, [0 for _ in normalized_labels]
|
|
label_map = {int(old): idx for idx, old in enumerate(unique_labels)}
|
|
return len(unique_labels), [label_map[label] for label in normalized_labels]
|
|
|
|
@timeout(60 * 20)
|
|
async def _summarize_texts(self, texts: list[str], callback=None, task_id: str = ""):
|
|
"""Summarize a cluster and return text plus embedding when successful."""
|
|
self._check_task_canceled(task_id, "summarization")
|
|
|
|
len_per_chunk = int((self._llm_model.max_length - self._max_token) / len(texts))
|
|
cluster_content = "\n".join([truncate(t, max(1, len_per_chunk)) for t in texts])
|
|
try:
|
|
async with chat_limiter:
|
|
self._check_task_canceled(task_id, "before LLM call")
|
|
|
|
cnt = await self._chat(
|
|
"You're a helpful assistant.\n\nHelp me with the following task.\n\n%s" % self._prompt.format(cluster_content=cluster_content),
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": "Beside the summarization, give a title at the first line of your summarization. Must be in the same language as the paragraphs.",
|
|
}
|
|
],
|
|
{"max_tokens": max(self._max_token, 512)}, # fix issue: #10235
|
|
)
|
|
cnt = re.sub(
|
|
"(······\n由于长度的原因，回答被截断了，要继续吗？|For the content length reason, it stopped, continue?)",
|
|
"",
|
|
cnt,
|
|
)
|
|
logging.debug(f"SUM: {cnt}")
|
|
|
|
self._check_task_canceled(task_id, "before embedding")
|
|
|
|
embds = await self._embedding_encode(cnt)
|
|
return cnt.split("\n")[0], cnt, embds
|
|
except TaskCanceledException:
|
|
raise
|
|
except Exception as exc:
|
|
self._error_count += 1
|
|
warn_msg = f"[RAPTOR] Skip cluster ({len(texts)} chunks) due to error: {exc}"
|
|
logging.warning(warn_msg)
|
|
if callback:
|
|
callback(msg=warn_msg)
|
|
if self._error_count >= self._max_errors:
|
|
raise RuntimeError(f"RAPTOR aborted after {self._error_count} errors. Last error: {exc}") from exc
|
|
return None
|
|
|
|
@staticmethod
|
|
def _root(node: _PsiTreeNode) -> _PsiTreeNode:
|
|
"""Return the current root for a Psi tree node."""
|
|
while node.parent is not None:
|
|
node = node.parent
|
|
return node
|
|
|
|
def _rank_leaf_pairs(self, leaves: list[_PsiTreeNode]) -> np.ndarray:
|
|
"""Rank all leaf pairs by original embedding-space cosine similarity."""
|
|
node_embeddings = np.asarray([leaf.embedding for leaf in leaves], dtype=np.float64)
|
|
node_embeddings = self._normalize_embeddings(node_embeddings)
|
|
similarities = node_embeddings @ node_embeddings.T
|
|
lower = np.tril_indices(len(leaves), -1)
|
|
ordered = np.argsort(similarities[lower], axis=0)[::-1]
|
|
return np.stack([lower[0][ordered], lower[1][ordered]], axis=-1)
|
|
|
|
@staticmethod
|
|
def _normalize_embeddings(node_embeddings: np.ndarray) -> np.ndarray:
|
|
"""Normalize embeddings for cosine operations while tolerating zero vectors."""
|
|
node_embeddings = np.asarray(node_embeddings, dtype=np.float64)
|
|
norms = np.linalg.norm(node_embeddings, axis=1, keepdims=True)
|
|
return node_embeddings / np.maximum(norms, 1e-12)
|
|
|
|
def _split_psi_buckets(self, nodes: list[_PsiTreeNode]) -> list[list[_PsiTreeNode]]:
|
|
"""Split large Psi inputs so exact pair ranking is bounded per bucket."""
|
|
if len(nodes) <= self._psi_bucket_size:
|
|
return [nodes]
|
|
|
|
node_embeddings = self._normalize_embeddings(np.asarray([node.embedding for node in nodes], dtype=np.float64))
|
|
groups = [np.arange(len(nodes), dtype=int)]
|
|
buckets = []
|
|
|
|
while groups:
|
|
group = np.asarray(groups.pop(), dtype=int)
|
|
if len(group) <= self._psi_bucket_size:
|
|
buckets.append(group.tolist())
|
|
continue
|
|
|
|
fanout = min(max(2, int(np.ceil(len(group) / self._psi_bucket_size))), len(group), 32)
|
|
group_embeddings = node_embeddings[group]
|
|
center_idx = np.linspace(0, len(group_embeddings) - 1, num=fanout, dtype=int)
|
|
centers = group_embeddings[center_idx].copy()
|
|
|
|
for _ in range(5):
|
|
labels = np.argmax(group_embeddings @ centers.T, axis=1)
|
|
for center_id in range(fanout):
|
|
mask = labels == center_id
|
|
if not np.any(mask):
|
|
continue
|
|
center = group_embeddings[mask].mean(axis=0)
|
|
norm = np.linalg.norm(center)
|
|
centers[center_id] = center / norm if norm > 0 else center
|
|
|
|
labels = np.argmax(group_embeddings @ centers.T, axis=1)
|
|
split_groups = [group[labels == center_id].tolist() for center_id in range(fanout)]
|
|
split_groups = [bucket for bucket in split_groups if bucket]
|
|
if len(split_groups) <= 1:
|
|
split_groups = [group[start : start + self._psi_bucket_size].tolist() for start in range(0, len(group), self._psi_bucket_size)]
|
|
groups.extend(split_groups)
|
|
|
|
buckets = [bucket for bucket in buckets if bucket]
|
|
buckets.sort(key=lambda bucket: (len(bucket), bucket[0]))
|
|
return [[nodes[idx] for idx in bucket] for bucket in buckets]
|
|
|
|
def _assign_prototype_embeddings(self, node: _PsiTreeNode) -> np.ndarray:
|
|
"""Assign mean child embeddings to internal Psi nodes for bucket-level ranking."""
|
|
if not node.children:
|
|
return np.asarray(node.embedding, dtype=np.float64)
|
|
embeddings = np.asarray([self._assign_prototype_embeddings(child) for child in node.children], dtype=np.float64)
|
|
node.embedding = embeddings.mean(axis=0)
|
|
return node.embedding
|
|
|
|
@staticmethod
|
|
def _iter_nodes(root: _PsiTreeNode):
|
|
"""Yield nodes in a Psi tree using a stack traversal."""
|
|
stack = [root]
|
|
while stack:
|
|
node = stack.pop()
|
|
yield node
|
|
stack.extend(node.children)
|
|
|
|
def _create_psi_parent(self, index: int, children: list[_PsiTreeNode]) -> _PsiTreeNode:
|
|
"""Create a parent node and attach the provided children to it."""
|
|
parent = _PsiTreeNode(index=index, children=children)
|
|
for child in children:
|
|
child.parent = parent
|
|
return parent
|
|
|
|
def _rebalance_psi_tree(self, root: _PsiTreeNode, next_index: int) -> tuple[_PsiTreeNode, int]:
|
|
"""Group oversized Psi tree nodes so fanout stays within max_cluster."""
|
|
max_children = max(2, int(self._max_cluster or 2))
|
|
|
|
def rebalance(node: _PsiTreeNode):
|
|
"""Recursively group children when a Psi node exceeds fanout."""
|
|
nonlocal next_index
|
|
|
|
for child in list(node.children):
|
|
rebalance(child)
|
|
|
|
while len(node.children) > max_children:
|
|
original_children = len(node.children)
|
|
grouped_children = []
|
|
for start in range(0, len(node.children), max_children):
|
|
batch = node.children[start : start + max_children]
|
|
if len(batch) == 1:
|
|
grouped_children.append(batch[0])
|
|
batch[0].parent = node
|
|
else:
|
|
grouped_children.append(self._create_psi_parent(next_index, batch))
|
|
grouped_children[-1].parent = node
|
|
next_index += 1
|
|
node.children = grouped_children
|
|
logging.info(
|
|
"RAPTOR Psi rebalance: node=%s children=%d grouped_to=%d max_cluster=%d",
|
|
node.index,
|
|
original_children,
|
|
len(grouped_children),
|
|
max_children,
|
|
)
|
|
|
|
rebalance(root)
|
|
return self._root(root), next_index
|
|
|
|
def _build_exact_psi_structure(
|
|
self,
|
|
nodes: list[_PsiTreeNode],
|
|
next_index: int,
|
|
task_id: str = "",
|
|
) -> tuple[_PsiTreeNode, int, int]:
|
|
"""Build an exact Psi subtree for a bounded node set."""
|
|
if len(nodes) == 1:
|
|
return nodes[0], next_index, 0
|
|
|
|
ranked_pairs = self._rank_leaf_pairs(nodes)
|
|
union_find = _PsiUnionFind(len(nodes))
|
|
merges = 0
|
|
for left_idx, right_idx in ranked_pairs:
|
|
self._check_task_canceled(task_id, "Psi tree construction")
|
|
if union_find.union(int(left_idx), int(right_idx)):
|
|
merges += 1
|
|
if merges == len(nodes) - 1:
|
|
break
|
|
|
|
local_nodes = {idx: node for idx, node in enumerate(nodes)}
|
|
tree = union_find.tree
|
|
children_by_parent = {}
|
|
for child_idx, parent_idx in enumerate(tree):
|
|
if child_idx not in local_nodes:
|
|
local_nodes[child_idx] = _PsiTreeNode(index=next_index)
|
|
next_index += 1
|
|
if parent_idx == -1:
|
|
continue
|
|
children_by_parent.setdefault(parent_idx, []).append(child_idx)
|
|
if parent_idx not in local_nodes:
|
|
local_nodes[parent_idx] = _PsiTreeNode(index=next_index)
|
|
next_index += 1
|
|
|
|
for parent_idx, child_indices in children_by_parent.items():
|
|
parent = local_nodes[parent_idx]
|
|
parent.children = [local_nodes[child_idx] for child_idx in child_indices]
|
|
for child in parent.children:
|
|
child.parent = parent
|
|
|
|
roots = [local_nodes[idx] for idx, parent_idx in enumerate(tree) if parent_idx == -1 and idx in local_nodes]
|
|
root = max(roots, key=lambda node: node.index)
|
|
return root, next_index, merges
|
|
|
|
def _build_bucketed_psi_structure(
|
|
self,
|
|
nodes: list[_PsiTreeNode],
|
|
next_index: int,
|
|
task_id: str = "",
|
|
) -> tuple[_PsiTreeNode, int, int]:
|
|
"""Build large Psi trees by exact-ranking bounded buckets, then bucket roots."""
|
|
buckets = self._split_psi_buckets(nodes)
|
|
logging.info(
|
|
"RAPTOR Psi bucketed build: nodes=%d buckets=%d bucket_size=%d exact_max_leaves=%d",
|
|
len(nodes),
|
|
len(buckets),
|
|
self._psi_bucket_size,
|
|
self._psi_exact_max_leaves,
|
|
)
|
|
|
|
bucket_roots = []
|
|
merges = 0
|
|
for bucket in buckets:
|
|
bucket_root, next_index, bucket_merges = self._build_psi_structure_from_nodes(bucket, next_index, task_id)
|
|
self._assign_prototype_embeddings(bucket_root)
|
|
bucket_roots.append(bucket_root)
|
|
merges += bucket_merges
|
|
|
|
if len(bucket_roots) == 1:
|
|
return bucket_roots[0], next_index, merges
|
|
|
|
root, next_index, root_merges = self._build_psi_structure_from_nodes(bucket_roots, next_index, task_id)
|
|
return root, next_index, merges + root_merges
|
|
|
|
def _build_psi_structure_from_nodes(
|
|
self,
|
|
nodes: list[_PsiTreeNode],
|
|
next_index: int,
|
|
task_id: str = "",
|
|
) -> tuple[_PsiTreeNode, int, int]:
|
|
"""Build Psi structure exactly for small sets and bucket large sets."""
|
|
if len(nodes) <= self._psi_exact_max_leaves:
|
|
return self._build_exact_psi_structure(nodes, next_index, task_id)
|
|
return self._build_bucketed_psi_structure(nodes, next_index, task_id)
|
|
|
|
def _build_psi_structure(self, chunks, task_id: str = "") -> tuple[_PsiTreeNode, list[_PsiTreeNode]]:
|
|
"""Build the Psi merge tree from original chunk embeddings.
|
|
|
|
``chunks`` is expected in the normalized 3-tuple shape
|
|
``(text, vec, source_chunk_ids)`` — leaves are seeded with
|
|
their own source ids, internal nodes get their ids set during
|
|
layer materialization in ``_build_psi_layers``.
|
|
"""
|
|
leaves = [
|
|
_PsiTreeNode(
|
|
index=i,
|
|
text=item[0],
|
|
embedding=np.asarray(item[1]),
|
|
source_chunk_ids=list(item[2] if len(item) > 2 else []),
|
|
)
|
|
for i, item in enumerate(chunks)
|
|
]
|
|
if len(leaves) == 1:
|
|
return leaves[0], leaves
|
|
|
|
root, next_index, merges = self._build_psi_structure_from_nodes(leaves, len(leaves), task_id)
|
|
root, _ = self._rebalance_psi_tree(root, next_index)
|
|
logging.info(
|
|
"RAPTOR Psi tree built: leaves=%d merges=%d root_fanout=%d",
|
|
len(leaves),
|
|
merges,
|
|
len(root.children),
|
|
)
|
|
return root, leaves
|
|
|
|
@staticmethod
|
|
def _psi_layers(root: _PsiTreeNode) -> dict[int, list[_PsiTreeNode]]:
|
|
"""Collect non-leaf Psi nodes by height for bottom-up summarization."""
|
|
layers = {}
|
|
|
|
def height(node: _PsiTreeNode) -> int:
|
|
"""Return node height while collecting internal nodes by layer."""
|
|
if not node.children:
|
|
return 0
|
|
node_height = max(height(child) for child in node.children) + 1
|
|
layers.setdefault(node_height, []).append(node)
|
|
return node_height
|
|
|
|
height(root)
|
|
return layers
|
|
|
|
async def _build_psi_layers(self, chunks, callback=None, task_id: str = ""):
|
|
"""Materialize Psi tree layers as summary chunks."""
|
|
layers = [(0, len(chunks))]
|
|
root, _ = self._build_psi_structure(chunks, task_id=task_id)
|
|
|
|
for layer_idx, (_, nodes) in enumerate(sorted(self._psi_layers(root).items()), start=1):
|
|
layer_start = len(chunks)
|
|
|
|
async def summarize_node(node: _PsiTreeNode):
|
|
"""Summarize one Psi internal node if its children have text.
|
|
|
|
Also propagates leaf provenance: the node's
|
|
``source_chunk_ids`` becomes the order-preserving deduped
|
|
union of every child's ``source_chunk_ids``. Because
|
|
children at this layer have already been processed (leaves
|
|
first, then bottom-up), each child carries the full set
|
|
of leaf ids underneath it — so the union here is the
|
|
complete leaf set this summary covers.
|
|
"""
|
|
texts = [child.text for child in node.children if child.text]
|
|
if not texts:
|
|
logging.warning("RAPTOR Psi node %s skipped because it has no child text to summarize", node.index)
|
|
return None
|
|
result = await self._summarize_texts(texts, callback, task_id)
|
|
if result is None:
|
|
logging.warning("RAPTOR Psi node %s skipped because summarization failed", node.index)
|
|
return None
|
|
_, node.text, node.embedding = result
|
|
merged_ids: list[str] = []
|
|
seen: set[str] = set()
|
|
for child in node.children:
|
|
for src in child.source_chunk_ids:
|
|
if src and src not in seen:
|
|
seen.add(src)
|
|
merged_ids.append(src)
|
|
node.source_chunk_ids = merged_ids
|
|
return node
|
|
|
|
tasks = [asyncio.create_task(summarize_node(node)) for node in nodes]
|
|
try:
|
|
summarized_nodes = await asyncio.gather(*tasks, return_exceptions=False)
|
|
except Exception as e:
|
|
logging.error(f"Error in RAPTOR Psi tree processing: {e}")
|
|
for task in tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
raise
|
|
|
|
summarized_nodes = [node for node in summarized_nodes if node is not None]
|
|
for node in summarized_nodes:
|
|
chunks.append((node.text, node.embedding, list(node.source_chunk_ids)))
|
|
|
|
if len(chunks) > layer_start:
|
|
layers.append((layer_start, len(chunks)))
|
|
logging.info(
|
|
"RAPTOR Psi layer materialized: layer=%d nodes=%d summaries=%d",
|
|
layer_idx,
|
|
len(nodes),
|
|
len(chunks) - layer_start,
|
|
)
|
|
if callback:
|
|
callback(msg="Build one Psi-RAG layer: {} -> {}".format(len(nodes), len(chunks) - layer_start))
|
|
else:
|
|
logging.warning("RAPTOR Psi layer %d produced no summaries; stopping materialization", layer_idx)
|
|
break
|
|
|
|
return chunks, layers
|
|
|
|
async def __call__(
|
|
self,
|
|
chunks,
|
|
random_state,
|
|
callback=None,
|
|
task_id: str = "",
|
|
is_tree: bool = False,
|
|
):
|
|
"""Build summary chunks and layer boundaries for RAPTOR retrieval.
|
|
|
|
``chunks`` accepts either the legacy 2-tuple shape
|
|
``(text, vec)`` or the provenance-carrying 3-tuple shape
|
|
``(text, vec, source_chunk_ids)`` where ``source_chunk_ids`` is
|
|
the list of original chunk ids that produced this entry. Output
|
|
always uses the 3-tuple shape so every appended summary carries
|
|
its leaves' ids. ``[]`` is left in the slot for a leaf whose id
|
|
was missing — see the caller for the normalization rules.
|
|
|
|
Return shapes:
|
|
* ``is_tree=False`` (default) — original behavior: returns
|
|
``(chunks, layers)`` where ``chunks`` is the flat list
|
|
(originals + summaries) and ``layers`` is the per-level
|
|
index range ``[(start, end), ...]``.
|
|
* ``is_tree=True`` — returns a hierarchical tree dict via
|
|
``_materialize_tree``. Supported for the classic builder
|
|
only; raises ``NotImplementedError`` for PSI_TREE_BUILDER
|
|
(PSI's hyperedge-driven summarization doesn't form a strict
|
|
parent-of relation). Returns ``None`` when there's nothing
|
|
to materialize.
|
|
"""
|
|
if len(chunks) <= 1:
|
|
return None if is_tree else ([], [])
|
|
|
|
# Normalize input to the 3-tuple shape. Reject empties / bad
|
|
# vectors at the same time the legacy path used to.
|
|
def _normalize(item):
|
|
if len(item) >= 3:
|
|
text, vec, src = item[0], item[1], item[2]
|
|
else:
|
|
text, vec = item[0], item[1]
|
|
src = []
|
|
if not text or vec is None or len(vec) <= 0:
|
|
return None
|
|
# Defensive: a leaf should carry a list of strings. Drop
|
|
# falsy entries so we don't propagate empty ids upward.
|
|
if isinstance(src, (list, tuple)):
|
|
src = [s for s in src if s]
|
|
else:
|
|
src = [src] if src else []
|
|
return (text, vec, list(src), "")
|
|
|
|
normalized = [t for t in (_normalize(c) for c in chunks) if t is not None]
|
|
if len(normalized) <= 1:
|
|
return None if is_tree else (normalized, [(0, len(normalized))])
|
|
chunks = normalized
|
|
|
|
if self._tree_builder == PSI_TREE_BUILDER:
|
|
if is_tree:
|
|
raise NotImplementedError(
|
|
"is_tree=True is not supported for PSI_TREE_BUILDER",
|
|
)
|
|
logging.info("RAPTOR: using %s tree builder for %d chunks", self._tree_builder, len(chunks))
|
|
return await self._build_psi_layers(chunks, callback, task_id)
|
|
|
|
# ``parent_child_map`` records each summary's immediate
|
|
# children so ``_materialize_tree`` can walk back into a tree
|
|
# when ``is_tree`` is set. Always populated (cheap) so the
|
|
# tree path is just a return-shape choice at the end.
|
|
parent_child_map: dict[int, list[int]] = {}
|
|
n_originals = len(chunks)
|
|
|
|
layers = [(0, len(chunks))]
|
|
start, end = 0, len(chunks)
|
|
|
|
@timeout(60 * 20)
|
|
async def summarize(ck_idx: list[int]):
|
|
"""Summarize one classic RAPTOR cluster into the chunk list.
|
|
|
|
On success appends ``(summary_text, summary_vec, src_ids)``
|
|
where ``src_ids`` is the order-preserving deduped union of
|
|
the ``source_chunk_ids`` of every chunk indexed in
|
|
``ck_idx`` — i.e. the full leaf set that contributed to
|
|
the cluster, even through nested summaries.
|
|
"""
|
|
nonlocal chunks
|
|
|
|
texts = [chunks[i][0] for i in ck_idx]
|
|
result = await self._summarize_texts(texts, callback, task_id)
|
|
if result is not None:
|
|
# ``dict.fromkeys`` is the cheapest way to de-dup a
|
|
# list of strings while preserving first-seen order.
|
|
merged_ids: list[str] = []
|
|
seen: set[str] = set()
|
|
for i in ck_idx:
|
|
for src in chunks[i][2]:
|
|
if src and src not in seen:
|
|
seen.add(src)
|
|
merged_ids.append(src)
|
|
summary_ti, summary_text, summary_vec = result
|
|
chunks.append((summary_text, summary_vec, merged_ids, summary_ti))
|
|
# Index of the just-appended summary; map it to its
|
|
# immediate children for the tree materializer below.
|
|
parent_child_map[len(chunks) - 1] = list(ck_idx)
|
|
|
|
while end - start > 1:
|
|
self._check_task_canceled(task_id, "layer processing")
|
|
|
|
# ``chunks`` is a mix of 3-tuples (layer-0 originals from
|
|
# _normalize) and 4-tuples (summaries appended by
|
|
# summarize). Vector is always at index 1 in both shapes,
|
|
# so use positional access — the older ``_, embd, _, _``
|
|
# form crashed on layer-0 entries.
|
|
embeddings = [entry[1] for entry in chunks[start:end]]
|
|
if end - start <= self._small_layer_collapse:
|
|
# Too few nodes for meaningful sub-clustering. Skip the
|
|
# clustering pass entirely and summarize the whole layer
|
|
# into one parent, so the upper tree doesn't descend one
|
|
# node per layer (N -> N-1 -> N-2 -> ... each a full
|
|
# clustering + summarize pass).
|
|
await summarize(list(range(start, end)))
|
|
produced = len(chunks) - end
|
|
if produced == 0:
|
|
logging.warning("RAPTOR layer produced no summaries; stopping materialization")
|
|
break
|
|
logging.info(
|
|
"RAPTOR small-N collapse: layer of %d node(s) [%d:%d] collapsed into %d summary; stopping at tree top",
|
|
end - start, start, end, produced,
|
|
)
|
|
layers.append((end, len(chunks)))
|
|
if callback:
|
|
callback(msg="Cluster one layer: {} -> {} (small-N collapse)".format(end - start, produced))
|
|
break
|
|
|
|
n_clusters, lbls = self.clustering(
|
|
embeddings,
|
|
random_state=random_state,
|
|
task_id=task_id,
|
|
)
|
|
|
|
# Loop-termination guarantee. The outer ``while end - start > 1``
|
|
# relies on each layer strictly shrinking the input count. If
|
|
# the clusterer degenerates and returns one cluster per input,
|
|
# every "cluster" is a single chunk, ``summarize()`` produces
|
|
# one summary per input, and ``produced == end - start`` —
|
|
# the same count carries into the next iteration and the loop
|
|
# spins forever, logging "Cluster one layer: N -> N".
|
|
#
|
|
# Collapse everything at this level into a single cluster so
|
|
# the layer produces exactly one summary. The tree gets a
|
|
# taller-than-usual "single trunk" segment at this depth
|
|
# instead of an infinite loop; downstream consumers only care
|
|
# that ``layers`` is monotonically shrinking.
|
|
if n_clusters >= len(embeddings):
|
|
logging.warning(
|
|
"RAPTOR clustering did not reduce input count (%d inputs → %d clusters); collapsing this layer into a single summary to prevent a non-terminating loop",
|
|
len(embeddings),
|
|
n_clusters,
|
|
)
|
|
n_clusters = 1
|
|
lbls = [0] * len(embeddings)
|
|
|
|
tasks = []
|
|
for c in range(n_clusters):
|
|
ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c]
|
|
assert len(ck_idx) > 0
|
|
self._check_task_canceled(task_id, "before cluster processing")
|
|
tasks.append(asyncio.create_task(summarize(ck_idx)))
|
|
try:
|
|
await asyncio.gather(*tasks, return_exceptions=False)
|
|
except Exception as e:
|
|
logging.error(f"Error in RAPTOR cluster processing: {e}")
|
|
for t in tasks:
|
|
t.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
raise
|
|
|
|
produced = len(chunks) - end
|
|
assert produced <= n_clusters, "{} vs. {}".format(produced, n_clusters)
|
|
if produced < n_clusters:
|
|
logging.warning(
|
|
"RAPTOR layer produced %d/%d cluster summaries; skipped %d cluster(s) due to errors",
|
|
produced,
|
|
n_clusters,
|
|
n_clusters - produced,
|
|
)
|
|
if produced == 0:
|
|
logging.warning("RAPTOR layer produced no summaries; stopping materialization")
|
|
break
|
|
layers.append((end, len(chunks)))
|
|
if callback:
|
|
callback(msg="Cluster one layer: {} -> {}".format(end - start, produced))
|
|
start = end
|
|
end = len(chunks)
|
|
|
|
if is_tree:
|
|
return self._materialize_tree(chunks, layers, parent_child_map, n_originals), []
|
|
return chunks, layers
|
|
|
|
@staticmethod
|
|
def _materialize_tree(chunks, layers, parent_child_map, n_originals):
|
|
"""Walk ``parent_child_map`` from the top layer down to layer-1
|
|
and emit the user-facing tree dict. See ``__call__``'s
|
|
``is_tree=True`` contract for the shape.
|
|
chunks: [(summary_text, summary_vec, merged_ids, summary_ti)]"""
|
|
if not layers or len(chunks) == 0:
|
|
return None
|
|
top_start, top_end = layers[-1]
|
|
if top_end <= top_start:
|
|
return None
|
|
|
|
def _title_at(idx: int) -> str:
|
|
# Summary tuples are (text, vec, merged_ids, summary_ti)
|
|
# — title is the 4th slot. Layer-0 originals are 3-tuples
|
|
# and don't appear as tree nodes themselves (they collapse
|
|
# into source_chunk_ids on their layer-1 parent).
|
|
return chunks[idx][3] if len(chunks[idx]) >= 4 else ""
|
|
|
|
def _desc_at(idx: int) -> str:
|
|
return chunks[idx][0] if chunks[idx] else ""
|
|
|
|
def _build_node(idx: int) -> dict:
|
|
children_idx = parent_child_map.get(idx, [])
|
|
# If every immediate child is a layer-0 original, this
|
|
# node is a "leaf" in the tree contract — collapse to
|
|
# source_chunk_ids.
|
|
if children_idx and all(c < n_originals for c in children_idx):
|
|
ids: list[str] = []
|
|
seen: set[str] = set()
|
|
for c in children_idx:
|
|
for s in chunks[c][2]:
|
|
if s and s not in seen:
|
|
seen.add(s)
|
|
ids.append(s)
|
|
return {"title": _title_at(idx), "source_chunk_ids": ids, "description": _desc_at(idx)}
|
|
return {"children": [_build_node(c) for c in children_idx], "title": _title_at(idx), "description": _desc_at(idx)}
|
|
|
|
top_nodes = [_build_node(i) for i in range(top_start, top_end)]
|
|
if len(top_nodes) == 1:
|
|
return top_nodes[0]
|
|
# Multiple top-layer summaries — clustering didn't collapse to
|
|
# a single root. Wrap in a synthetic root so the caller always
|
|
# sees one dict.
|
|
return {"title": "(root)", "children": top_nodes}
|