mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-06 19:38:36 +08:00
feat(raptor): add Psi tree builder with original-space ranking and safe migration (#14679)
### 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:
@@ -18,15 +18,22 @@
|
||||
Unit tests for Raptor utility functions.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
from rag.utils.raptor_utils import (
|
||||
CSV_EXTENSIONS,
|
||||
EXCEL_EXTENSIONS,
|
||||
STRUCTURED_EXTENSIONS,
|
||||
collect_raptor_chunk_ids,
|
||||
collect_raptor_methods,
|
||||
get_raptor_clustering_method,
|
||||
get_raptor_tree_builder,
|
||||
get_skip_reason,
|
||||
is_structured_file_type,
|
||||
is_tabular_pdf,
|
||||
make_raptor_summary_chunk_id,
|
||||
should_skip_raptor,
|
||||
get_skip_reason,
|
||||
EXCEL_EXTENSIONS,
|
||||
CSV_EXTENSIONS,
|
||||
STRUCTURED_EXTENSIONS
|
||||
)
|
||||
|
||||
|
||||
@@ -283,5 +290,117 @@ class TestIntegrationScenarios:
|
||||
assert should_skip_raptor(file_type, raptor_config=raptor_config) is False
|
||||
|
||||
|
||||
class TestRaptorTreeBuilderConfig:
|
||||
"""Test RAPTOR tree builder config resolution"""
|
||||
|
||||
def test_defaults_to_original_raptor_builder(self):
|
||||
assert get_raptor_tree_builder({}) == "raptor"
|
||||
assert get_raptor_tree_builder(None) == "raptor"
|
||||
|
||||
def test_reads_top_level_tree_builder(self):
|
||||
assert get_raptor_tree_builder({"tree_builder": "psi"}) == "psi"
|
||||
|
||||
def test_reads_legacy_ext_tree_builder(self):
|
||||
assert get_raptor_tree_builder({"ext": {"tree_builder": "psi"}}) == "psi"
|
||||
|
||||
def test_ext_tree_builder_overrides_stale_top_level_value(self):
|
||||
assert get_raptor_tree_builder({"tree_builder": "psi", "ext": {"tree_builder": "raptor"}}) == "raptor"
|
||||
|
||||
def test_rejects_unknown_tree_builder(self):
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
|
||||
get_raptor_tree_builder({"tree_builder": "ahc"})
|
||||
|
||||
|
||||
class TestRaptorClusteringMethodConfig:
|
||||
"""Test RAPTOR clustering method config resolution"""
|
||||
|
||||
def test_defaults_to_gmm(self):
|
||||
assert get_raptor_clustering_method({}) == "gmm"
|
||||
assert get_raptor_clustering_method(None) == "gmm"
|
||||
|
||||
def test_reads_top_level_clustering_method(self):
|
||||
assert get_raptor_clustering_method({"clustering_method": "gmm"}) == "gmm"
|
||||
assert get_raptor_clustering_method({"clustering_method": "ahc"}) == "ahc"
|
||||
|
||||
def test_reads_legacy_ext_clustering_method(self):
|
||||
assert get_raptor_clustering_method({"ext": {"clustering_method": "ahc"}}) == "ahc"
|
||||
|
||||
def test_ext_clustering_method_overrides_stale_top_level_value(self):
|
||||
assert get_raptor_clustering_method({"clustering_method": "gmm", "ext": {"clustering_method": "ahc"}}) == "ahc"
|
||||
|
||||
def test_rejects_unknown_clustering_method(self):
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
|
||||
get_raptor_clustering_method({"clustering_method": "unknown"})
|
||||
|
||||
|
||||
class TestRaptorMethodCollection:
|
||||
"""Test RAPTOR summary method extraction from doc-store fields"""
|
||||
|
||||
def test_legacy_summary_without_method_is_original_raptor(self):
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor"}}
|
||||
|
||||
assert collect_raptor_methods(field_map) == {"raptor"}
|
||||
assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
|
||||
|
||||
def test_extra_method_is_preserved(self):
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}}
|
||||
|
||||
assert collect_raptor_methods(field_map) == {"psi"}
|
||||
assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
|
||||
|
||||
def test_extra_field_supports_oceanbase_legacy_rows(self):
|
||||
field_map = {
|
||||
"chunk_1": {
|
||||
"extra": {
|
||||
"raptor_kwd": "raptor",
|
||||
"raptor_method": "psi",
|
||||
}
|
||||
},
|
||||
"chunk_2": {
|
||||
"extra": "{\"raptor_kwd\": \"raptor\"}",
|
||||
},
|
||||
"chunk_3": {
|
||||
"extra": {"raptor_kwd": ""},
|
||||
},
|
||||
}
|
||||
|
||||
assert collect_raptor_methods(field_map) == {"psi", "raptor"}
|
||||
assert collect_raptor_chunk_ids(field_map) == {"chunk_1", "chunk_2"}
|
||||
|
||||
def test_non_raptor_rows_are_ignored(self):
|
||||
field_map = {
|
||||
"chunk_1": {"raptor_kwd": ""},
|
||||
"chunk_2": {"extra": {"raptor_kwd": "graph"}},
|
||||
"chunk_3": {},
|
||||
}
|
||||
|
||||
assert collect_raptor_methods(field_map) == set()
|
||||
assert collect_raptor_chunk_ids(field_map) == set()
|
||||
|
||||
def test_malformed_extra_payload_is_logged_and_ignored(self, caplog):
|
||||
field_map = {"chunk_1": {"extra": "{bad json"}}
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
assert collect_raptor_methods(field_map) == set()
|
||||
assert collect_raptor_chunk_ids(field_map) == set()
|
||||
|
||||
assert "Ignoring malformed RAPTOR extra payload" in caplog.text
|
||||
|
||||
def test_chunk_id_collection_can_preserve_current_method(self):
|
||||
field_map = {
|
||||
"legacy": {"raptor_kwd": "raptor"},
|
||||
"old": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
|
||||
"current": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}},
|
||||
}
|
||||
|
||||
assert collect_raptor_chunk_ids(field_map, exclude_methods={"psi"}) == {"legacy", "old"}
|
||||
assert collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"}) == {"current"}
|
||||
|
||||
def test_summary_chunk_ids_include_real_document_id(self):
|
||||
content = "same generated summary"
|
||||
|
||||
assert make_raptor_summary_chunk_id(content, "doc-a") != make_raptor_summary_chunk_id(content, "doc-b")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
Reference in New Issue
Block a user