zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-19 22:21:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(raptor): add Psi tree builder with original-space ranking and safe migration (#14679)

Browse Source 
### What problem does this PR solve?

Closes #14674.

This PR improves RAPTOR configuration and tree construction while
preserving the existing RAPTOR behavior as the default.

RAPTOR currently builds summary layers with the original UMAP + GMM
clustering path. This PR keeps that default path, and adds:

- A hidden backend tree-builder option:
  - `tree_builder="raptor"`: default, existing RAPTOR behavior.
- `tree_builder="psi"`: rank-aware Psi-style tree builder using original
embedding-space cosine ranking.
- A user-facing clustering method option for the default RAPTOR builder:
  - `clustering_method="gmm"`: existing default.
- `clustering_method="ahc"`: agglomerative hierarchical clustering path.
- A RAPTOR UI setting for `Clustering method` and `Max cluster`.

### What changed

#### Backend

- Added `tree_builder` support for RAPTOR/Psi.
- Added `clustering_method` support for GMM/AHC.
- Kept existing RAPTOR + GMM as the default.
- Added Psi tree building from original-space cosine similarity.
- Added bucketed Psi building controls for large inputs:
  - `raptor.ext.psi_exact_max_leaves`
  - `raptor.ext.psi_bucket_size`
- Added method-aware RAPTOR summary metadata using existing
`extra.raptor_method`.
- Avoided adding a dedicated DB schema field for experimental method
tracking.
- Added cleanup/migration logic to avoid mixing stale RAPTOR summary
trees.
- Added defensive checks for Psi tree construction and summary failures.

#### Frontend/UI

- Added `Clustering method` in RAPTOR settings with `GMM` and `AHC`.
- Added/kept `Max cluster` in RAPTOR settings.
- Enlarged max cluster UI limit to `1024`, matching backend validation.
- Kept AHC editable even when a RAPTOR task has already finished.
- Fixed the UI save payload so `clustering_method` and `tree_builder`
are serialized through `parser_config.raptor.ext`, avoiding backend
validation errors for extra top-level RAPTOR fields.

Example saved RAPTOR config:

```json
{
  "raptor": {
    "max_cluster": 317,
    "ext": {
      "clustering_method": "ahc",
      "tree_builder": "raptor"
    }
  }
}

Co-authored-by: CaptainTimon <CaptainTimon@users.noreply.github.com>
This commit is contained in:
CaptainTimon
2026-05-11 15:42:31 -10:00
committed by GitHub
parent 415169d497
commit 2717ee283f
21 changed files with 1722 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

295
rag/svr/task_executor.py
View File

@@ -36,7 +36,15 @@ from api.db.joint_services.memory_message_service import handle_save_to_memory_t
from common.connection_utils import timeout
from common.metadata_utils import turn2jsonschema, update_metadata_to
from rag.utils.base64_image import image2id
from rag.utils.raptor_utils import should_skip_raptor, get_skip_reason
from rag.utils.raptor_utils import (
collect_raptor_chunk_ids,
collect_raptor_methods,
get_raptor_clustering_method,
get_raptor_tree_builder,
get_skip_reason,
make_raptor_summary_chunk_id,
should_skip_raptor,
)
from common.log_utils import init_root_logger
from common.config_utils import show_configs
from rag.graphrag.general.index import run_graphrag_for_kb
@@ -70,7 +78,10 @@ from api.db.db_models import close_connection
from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, \
email, tag
from rag.nlp import search, rag_tokenizer, add_positions
from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
from rag.raptor import (
RAPTOR_TREE_BUILDER,
RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor,
)
from common.token_utils import num_tokens_from_string, truncate
from rag.utils.redis_conn import REDIS_CONN, RedisDistributedLock
from rag.graphrag.utils import chat_limiter
@@ -817,61 +828,160 @@ async def run_dataflow(task: dict):
dsl=str(pipeline))
async def has_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str) -> bool:
"""Return True if RAPTOR chunks already exist for doc_id in the doc store.
RAPTOR_METHOD_SEARCH_LIMIT = 10000
Queries directly for raptor_kwd="raptor" rows so a non-RAPTOR leading
chunk cannot produce a false-negative result. Uses thread_pool_exec so
the blocking doc-store call does not stall the event loop.
"""
async def get_raptor_chunk_field_map(doc_id: str, tenant_id: str, kb_id: str) -> dict:
"""Return stored RAPTOR marker fields for a document."""
from common.doc_store.doc_store_base import OrderByExpr
from rag.nlp import search as nlp_search
try:
condition = {"doc_id": doc_id, "raptor_kwd": ["raptor"]}
async def search_fields(fields: list[str], condition: dict, order_by=None):
"""Search chunk fields in the current knowledge base."""
res = await thread_pool_exec(
settings.docStoreConn.search,
["raptor_kwd"], [], condition, [], OrderByExpr(),
0, 1, nlp_search.index_name(tenant_id), [kb_id]
fields, [], condition, [], order_by or OrderByExpr(),
0, RAPTOR_METHOD_SEARCH_LIMIT, nlp_search.index_name(tenant_id), [kb_id]
)
field_map = settings.docStoreConn.get_fields(res, ["raptor_kwd"])
found = bool(field_map)
if found:
return settings.docStoreConn.get_fields(res, fields)
primary = await search_fields(["raptor_kwd", "extra"], {"doc_id": doc_id, "raptor_kwd": ["raptor"]})
if collect_raptor_chunk_ids(primary):
return primary
try:
return await search_fields(
["raptor_kwd", "extra"],
{"doc_id": doc_id},
OrderByExpr().desc("create_timestamp_flt"),
)
except Exception:
logging.debug("RAPTOR fallback method lookup with extra field failed for doc %s", doc_id, exc_info=True)
return primary
async def get_raptor_chunk_methods(doc_id: str, tenant_id: str, kb_id: str) -> set[str]:
"""Return the RAPTOR tree builders already stored for doc_id.
Queries directly for raptor_kwd="raptor" rows so a non-RAPTOR leading
chunk cannot produce a false-negative result. Legacy summary chunks that
do not have method metadata are treated as the original RAPTOR builder.
"""
try:
field_map = await get_raptor_chunk_field_map(doc_id, tenant_id, kb_id)
methods = collect_raptor_methods(field_map)
if methods:
logging.info(
"Checkpoint hit: RAPTOR chunks for doc %s (tenant=%s kb=%s) already exist",
doc_id, tenant_id, kb_id,
"Checkpoint hit: RAPTOR chunks for doc %s (tenant=%s kb=%s methods=%s) already exist",
doc_id, tenant_id, kb_id, sorted(methods),
)
else:
logging.info(
"Checkpoint miss: no RAPTOR chunks for doc %s (tenant=%s kb=%s)",
doc_id, tenant_id, kb_id,
)
return found
return methods
except Exception:
logging.exception("Failed to check RAPTOR chunks for doc %s", doc_id)
return False
raise
async def has_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, tree_builder: str = RAPTOR_TREE_BUILDER) -> bool:
"""Return whether doc_id already has summaries for tree_builder."""
methods = await get_raptor_chunk_methods(doc_id, tenant_id, kb_id)
return tree_builder in methods
async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_method: str | None = None):
"""Delete RAPTOR summaries for doc_id, optionally preserving one method."""
from rag.nlp import search as nlp_search
if keep_method is None:
logging.info(
"delete_raptor_chunks: removing all RAPTOR summaries (doc=%s tenant=%s kb=%s)",
doc_id, tenant_id, kb_id,
)
await thread_pool_exec(
settings.docStoreConn.delete,
{"doc_id": doc_id, "raptor_kwd": ["raptor"]},
nlp_search.index_name(tenant_id),
kb_id,
)
return 0
field_map = await get_raptor_chunk_field_map(doc_id, tenant_id, kb_id)
chunk_ids = collect_raptor_chunk_ids(field_map, exclude_methods={keep_method})
if not chunk_ids:
logging.debug(
"delete_raptor_chunks: no stale RAPTOR chunks to remove (doc=%s tenant=%s kb=%s keep=%s)",
doc_id, tenant_id, kb_id, keep_method,
)
return 0
logging.info(
"delete_raptor_chunks: removing %d stale RAPTOR chunks (doc=%s tenant=%s kb=%s keep=%s)",
len(chunk_ids), doc_id, tenant_id, kb_id, keep_method,
)
await thread_pool_exec(
settings.docStoreConn.delete,
{"id": list(chunk_ids)},
nlp_search.index_name(tenant_id),
kb_id,
)
return len(chunk_ids)
@timeout(3600)
async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_size, callback=None, doc_ids=[]):
"""Generate RAPTOR summaries for selected documents in a knowledge base."""
fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID
raptor_config = kb_parser_config.get("raptor", {})
raptor_ext_config = raptor_config.get("ext") or {}
tree_builder = get_raptor_tree_builder(raptor_config)
clustering_method = get_raptor_clustering_method(raptor_config)
vctr_nm = "q_%d_vec" % vector_size
res = []
tk_count = 0
cleanup_raptor_chunks = []
max_errors = int(os.environ.get("RAPTOR_MAX_ERRORS", 3))
doc_name_by_id = {}
doc_info_by_id = {}
for doc_id in set(doc_ids):
ok, source_doc = DocumentService.get_by_id(doc_id)
if not ok or not source_doc:
continue
source_name = getattr(source_doc, "name", "")
if source_name:
doc_name_by_id[doc_id] = source_name
doc_info_by_id[doc_id] = {
"name": getattr(source_doc, "name", ""),
"type": getattr(source_doc, "type", ""),
"parser_id": getattr(source_doc, "parser_id", ""),
"parser_config": getattr(source_doc, "parser_config", {}) or {},
}
def schedule_raptor_cleanup(doc_id: str, keep_method: str | None = None):
"""Queue stale RAPTOR summaries for deletion after successful insert."""
cleanup_plan = (doc_id, keep_method)
if cleanup_plan not in cleanup_raptor_chunks:
cleanup_raptor_chunks.append(cleanup_plan)
def skip_raptor_doc(doc_id: str) -> bool:
"""Return whether RAPTOR should be skipped for this source document."""
doc_info = doc_info_by_id.get(doc_id, {})
file_type = doc_info.get("type") or row.get("type", "")
parser_id = doc_info.get("parser_id") or row.get("parser_id", "")
parser_config = doc_info.get("parser_config") or row.get("parser_config", {})
if should_skip_raptor(file_type, parser_id, parser_config, raptor_config):
skip_reason = get_skip_reason(file_type, parser_id, parser_config)
doc_name = doc_info.get("name") or doc_id
logging.info("Skipping Raptor for document %s: %s", doc_name, skip_reason)
callback(msg=f"[RAPTOR] doc:{doc_id} skipped: {skip_reason}")
return True
return False
async def generate(chunks, did):
"""Run RAPTOR and append generated summary chunks for one doc id."""
nonlocal tk_count, res
logging.info("RAPTOR: using tree_builder=%s clustering_method=%s for doc %s", tree_builder, clustering_method, did)
raptor = Raptor(
raptor_config.get("max_cluster", 64),
chat_mdl,
@@ -880,16 +990,21 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
raptor_config["max_token"],
raptor_config["threshold"],
max_errors=max_errors,
tree_builder=tree_builder,
clustering_method=clustering_method,
psi_exact_max_leaves=raptor_ext_config.get("psi_exact_max_leaves", 4096),
psi_bucket_size=raptor_ext_config.get("psi_bucket_size", 1024),
)
original_length = len(chunks)
chunks, layers = await raptor(chunks, kb_parser_config["raptor"]["random_seed"], callback, row["id"])
effective_doc_name = row["name"] if did == fake_doc_id else doc_name_by_id.get(did, row["name"])
effective_doc_name = row["name"] if did == fake_doc_id else doc_info_by_id.get(did, {}).get("name") or row["name"]
doc = {
"doc_id": did,
"kb_id": [str(row["kb_id"])],
"docnm_kwd": effective_doc_name,
"title_tks": rag_tokenizer.tokenize(effective_doc_name),
"raptor_kwd": "raptor"
"raptor_kwd": "raptor",
"extra": {"raptor_method": tree_builder},
}
if row["pagerank"]:
doc[PAGERANK_FLD] = int(row["pagerank"])
@@ -906,7 +1021,7 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
for idx, (content, vctr) in enumerate(chunks[original_length:], start=original_length):
d = copy.deepcopy(doc)
d["id"] = xxhash.xxh64((content + str(fake_doc_id)).encode("utf-8")).hexdigest()
d["id"] = make_raptor_summary_chunk_id(content, did)
d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
d["create_timestamp_flt"] = datetime.now().timestamp()
d[vctr_nm] = vctr.tolist()
@@ -918,12 +1033,28 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
tk_count += num_tokens_from_string(content)
if raptor_config.get("scope", "file") == "file":
dataset_methods = await get_raptor_chunk_methods(fake_doc_id, row["tenant_id"], row["kb_id"])
remove_dataset_summaries = bool(dataset_methods)
has_file_level_target = False
if dataset_methods:
callback(msg="[RAPTOR] will remove dataset-level summaries after file-level summaries are available.")
for x, doc_id in enumerate(doc_ids):
# CHECKPOINT: skip docs that already have RAPTOR chunks in the doc store
if await has_raptor_chunks(doc_id, row["tenant_id"], row["kb_id"]):
callback(msg=f"[RAPTOR] doc:{doc_id} already has RAPTOR chunks, skipping.")
if skip_raptor_doc(doc_id):
callback(prog=(x + 1.) / len(doc_ids))
continue
# CHECKPOINT: skip docs that already have RAPTOR chunks in the doc store
existing_methods = await get_raptor_chunk_methods(doc_id, row["tenant_id"], row["kb_id"])
if tree_builder in existing_methods:
has_file_level_target = True
if existing_methods != {tree_builder}:
schedule_raptor_cleanup(doc_id, tree_builder)
callback(msg=f"[RAPTOR] doc:{doc_id} will remove old RAPTOR summaries after insert.")
callback(msg=f"[RAPTOR] doc:{doc_id} already has {tree_builder} RAPTOR chunks, skipping.")
callback(prog=(x + 1.) / len(doc_ids))
continue
if existing_methods:
callback(msg=f"[RAPTOR] doc:{doc_id} will migrate RAPTOR summaries to {tree_builder} after insert.")
chunks = []
skipped_chunks = 0
@@ -945,12 +1076,52 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
callback(msg=f"[WARN] No valid chunks with vectors found for doc {doc_id}, skipping")
continue
before_generate = len(res)
await generate(chunks, doc_id)
if len(res) > before_generate:
has_file_level_target = True
if existing_methods:
schedule_raptor_cleanup(doc_id, tree_builder)
callback(prog=(x + 1.) / len(doc_ids))
if remove_dataset_summaries:
if has_file_level_target:
schedule_raptor_cleanup(fake_doc_id)
else:
callback(msg="[RAPTOR] kept dataset-level summaries because no file-level summaries were built.")
else:
migrated_file_docs = 0
file_cleanup_doc_ids = []
skipped_doc_ids = set()
for doc_id in set(doc_ids):
if skip_raptor_doc(doc_id):
skipped_doc_ids.add(doc_id)
continue
existing_methods = await get_raptor_chunk_methods(doc_id, row["tenant_id"], row["kb_id"])
if existing_methods:
file_cleanup_doc_ids.append(doc_id)
migrated_file_docs += 1
if migrated_file_docs:
callback(msg=f"[RAPTOR] will remove file-level summaries for {migrated_file_docs} docs after dataset-level build succeeds.")
existing_methods = await get_raptor_chunk_methods(fake_doc_id, row["tenant_id"], row["kb_id"])
if tree_builder in existing_methods:
if existing_methods != {tree_builder}:
schedule_raptor_cleanup(fake_doc_id, tree_builder)
callback(msg="[RAPTOR] will remove old dataset-level RAPTOR summaries after insert.")
for doc_id in file_cleanup_doc_ids:
schedule_raptor_cleanup(doc_id)
callback(msg=f"[RAPTOR] dataset-level {tree_builder} summaries already exist, skipping.")
return res, tk_count, cleanup_raptor_chunks
migrate_dataset_summaries = bool(existing_methods)
if migrate_dataset_summaries:
callback(msg=f"[RAPTOR] will migrate dataset-level RAPTOR summaries to {tree_builder} after insert.")
chunks = []
skipped_chunks = 0
for doc_id in doc_ids:
if doc_id in skipped_doc_ids:
continue
for d in settings.retriever.chunk_list(doc_id, row["tenant_id"], [str(row["kb_id"])],
fields=["content_with_weight", vctr_nm],
sort_by_position=True):
@@ -965,13 +1136,22 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
callback(msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}'. Consider re-parsing documents with the current embedding model.")
if not chunks:
if skipped_doc_ids and len(skipped_doc_ids) == len(set(doc_ids)):
callback(msg="[RAPTOR] all documents were skipped by RAPTOR auto-disable rules.")
return res, tk_count, cleanup_raptor_chunks
logging.error(f"RAPTOR: No valid chunks with vectors found in any document for kb {row['kb_id']}")
callback(msg=f"[ERROR] No valid chunks with vectors found. Please ensure documents are parsed with the current embedding model (vector size: {vector_size}).")
return res, tk_count
return res, tk_count, cleanup_raptor_chunks
before_generate = len(res)
await generate(chunks, fake_doc_id)
if len(res) > before_generate:
for doc_id in file_cleanup_doc_ids:
schedule_raptor_cleanup(doc_id)
if migrate_dataset_summaries:
schedule_raptor_cleanup(fake_doc_id, tree_builder)
return res, tk_count
return res, tk_count, cleanup_raptor_chunks
async def delete_image(kb_id, chunk_id):
@@ -1029,6 +1209,29 @@ async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progre
search.index_name(task_tenant_id), task_dataset_id, )
task_canceled = has_canceled(task_id)
if task_canceled:
# Roll back partial RAPTOR summary inserts so the next run is not
# mistaken for a completed checkpoint by get_raptor_chunk_methods.
raptor_ids_to_rollback = [
c["id"] for c in chunks[:b + settings.DOC_BULK_SIZE]
if c.get("raptor_kwd") == "raptor"
]
if raptor_ids_to_rollback:
try:
await thread_pool_exec(
settings.docStoreConn.delete,
{"id": raptor_ids_to_rollback},
search.index_name(task_tenant_id),
task_dataset_id,
)
logging.info(
"insert_chunks: rolled back %d partial RAPTOR chunks after cancellation (task=%s)",
len(raptor_ids_to_rollback), task_id,
)
except Exception:
logging.exception(
"insert_chunks: failed to roll back partial RAPTOR chunks after cancellation (task=%s)",
task_id,
)
progress_callback(-1, msg="Task has been canceled.")
return False
if b % 128 == 0:
@@ -1088,6 +1291,7 @@ async def do_handle_task(task):
task_parser_config = task["parser_config"]
task_start_ts = timer()
toc_thread = None
raptor_cleanup_chunks = []
# prepare the progress callback function
progress_callback = partial(set_progress, task_id, task_from_page, task_to_page)
@@ -1135,7 +1339,9 @@ async def do_handle_task(task):
"threshold": 0.1,
"max_cluster": 64,
"random_seed": 0,
"scope": "file"
"scope": "file",
"clustering_method": "gmm",
"tree_builder": "raptor",
},
}
)
@@ -1143,23 +1349,12 @@ async def do_handle_task(task):
progress_callback(prog=-1.0, msg="Internal error: Invalid RAPTOR configuration")
return
# Check if Raptor should be skipped for structured data
file_type = task.get("type", "")
parser_id = task.get("parser_id", "")
raptor_config = kb_parser_config.get("raptor", {})
if should_skip_raptor(file_type, parser_id, task_parser_config, raptor_config):
skip_reason = get_skip_reason(file_type, parser_id, task_parser_config)
logging.info(f"Skipping Raptor for document {task_document_name}: {skip_reason}")
progress_callback(prog=1.0, msg=f"Raptor skipped: {skip_reason}")
return
# bind LLM for raptor
chat_model_config = get_model_config_by_type_and_name(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
chat_model = LLMBundle(task_tenant_id, chat_model_config, lang=task_language)
# run RAPTOR
async with kg_limiter:
chunks, token_count = await run_raptor_for_kb(
chunks, token_count, raptor_cleanup_chunks = await run_raptor_for_kb(
row=task,
kb_parser_config=kb_parser_config,
chat_mdl=chat_model,
@@ -1268,6 +1463,18 @@ async def do_handle_task(task):
progress_callback(-1, msg="Task has been canceled.")
return
if raptor_cleanup_chunks:
cleaned_chunks = 0
for cleanup_doc_id, keep_method in raptor_cleanup_chunks:
cleaned_chunks += await delete_raptor_chunks(
cleanup_doc_id,
task_tenant_id,
task_dataset_id,
keep_method=keep_method,
)
if cleaned_chunks:
progress_callback(msg=f"Cleaned up {cleaned_chunks} stale RAPTOR chunks.")
logging.info(
"Indexing doc({}), page({}-{}), chunks({}), elapsed: {:.2f}".format(
task_document_name, task_from_page, task_to_page, len(chunks), timer() - start_ts