mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
refa: simplify RAPTOR tree clustering and configuration (#17614)
This commit is contained in:
buua436
@@ -188,12 +188,11 @@ class TestRaptorServiceRunRaptorForKb:
|
||||
"""RAPTOR config with file-level scope."""
|
||||
return {
|
||||
"raptor": {
|
||||
"tree_builder": "raptor",
|
||||
"clustering_method": "gmm",
|
||||
"scope": "file",
|
||||
"prompt": "summarize",
|
||||
"max_token": 512,
|
||||
"threshold": 0.5,
|
||||
"clustering_threshold": 0.5,
|
||||
"clustering_ratio": 0.5,
|
||||
"max_cluster": 64,
|
||||
"random_seed": 42,
|
||||
}
|
||||
@@ -204,12 +203,11 @@ class TestRaptorServiceRunRaptorForKb:
|
||||
"""RAPTOR config with dataset-level scope."""
|
||||
return {
|
||||
"raptor": {
|
||||
"tree_builder": "raptor",
|
||||
"clustering_method": "gmm",
|
||||
"scope": "dataset",
|
||||
"prompt": "summarize",
|
||||
"max_token": 512,
|
||||
"threshold": 0.5,
|
||||
"clustering_threshold": 0.5,
|
||||
"clustering_ratio": 0.5,
|
||||
"max_cluster": 64,
|
||||
"random_seed": 42,
|
||||
}
|
||||
@@ -321,7 +319,7 @@ class TestRaptorServiceRunRaptorForKb:
|
||||
):
|
||||
|
||||
async def mock_run_file(*args, **kwargs):
|
||||
cleanup_list = args[11]
|
||||
cleanup_list = args[9]
|
||||
cleanup_list.append(("doc_1", "tree_builder_a"))
|
||||
return [{"id": "c1"}], 10
|
||||
|
||||
@@ -383,10 +381,10 @@ class TestRaptorServiceRunRaptorForKb:
|
||||
await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
|
||||
|
||||
# Verify _run_file_level_raptor received vctr_nm with the correct vector size
|
||||
# Positional args: 0=raptor_config, 1=tree_builder, 2=clustering_method,
|
||||
# 3=chat_mdl, 4=embd_mdl, 5=vctr_nm
|
||||
# Positional args: 0=raptor_config, 1=chat_mdl, 2=embd_mdl,
|
||||
# 3=vctr_nm, 4=doc_ids, 5=doc_info_by_id
|
||||
positional_args = mock_file.call_args[0]
|
||||
assert positional_args[5] == "q_256_vec"
|
||||
assert positional_args[3] == "q_256_vec"
|
||||
|
||||
# ---- Document info collection through public API ----
|
||||
|
||||
@@ -405,9 +403,9 @@ class TestRaptorServiceRunRaptorForKb:
|
||||
await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
|
||||
|
||||
mock_collect.assert_called_once_with(doc_ids)
|
||||
# Verify doc_info_by_id was passed as positional arg[7] to _run_file_level_raptor
|
||||
# Verify doc_info_by_id was passed as positional arg[5] to _run_file_level_raptor
|
||||
positional_args = mock_file.call_args[0]
|
||||
assert positional_args[7] == expected_info
|
||||
assert positional_args[5] == expected_info
|
||||
|
||||
|
||||
class TestRaptorServiceFileLevelRaptorCheckpoint:
|
||||
@@ -429,12 +427,10 @@ class TestRaptorServiceFileLevelRaptorCheckpoint:
|
||||
"scope": "file",
|
||||
"max_cluster": 64,
|
||||
"prompt": "test prompt",
|
||||
"max_token": 256,
|
||||
"threshold": 0.1,
|
||||
"max_token": 512,
|
||||
"clustering_threshold": 0.3,
|
||||
"clustering_ratio": 0.5,
|
||||
"random_seed": 0,
|
||||
"clustering_method": "gmm",
|
||||
"tree_builder": "raptor",
|
||||
"ext": {},
|
||||
}
|
||||
|
||||
with patch.object(svc, "_get_raptor_chunk_methods", new_callable=AsyncMock) as mock_methods, patch.object(svc, "_should_skip_raptor", return_value=False):
|
||||
@@ -442,8 +438,6 @@ class TestRaptorServiceFileLevelRaptorCheckpoint:
|
||||
|
||||
result = await svc._run_file_level_raptor(
|
||||
raptor_config=raptor_config,
|
||||
tree_builder="raptor",
|
||||
clustering_method="gmm",
|
||||
chat_mdl=MagicMock(),
|
||||
embd_mdl=MagicMock(),
|
||||
vctr_nm="q_128_vec",
|
||||
|
||||
@@ -1,434 +0,0 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import pytest
|
||||
|
||||
np = pytest.importorskip("numpy")
|
||||
|
||||
from api.utils.validation_utils import RaptorConfig
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def raptor_module(monkeypatch):
|
||||
class TaskCanceledException(Exception):
|
||||
pass
|
||||
|
||||
class DummyLimiter:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
class DummyGaussianMixture:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def fit(self, embeddings):
|
||||
return self
|
||||
|
||||
def bic(self, embeddings):
|
||||
return 0
|
||||
|
||||
def predict_proba(self, embeddings):
|
||||
return np.ones((len(embeddings), 1))
|
||||
|
||||
class DummyAgglomerativeClustering:
|
||||
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):
|
||||
self.labels_ = self.fit_predict(embeddings)
|
||||
return self
|
||||
|
||||
def fit_predict(self, embeddings):
|
||||
if self.n_clusters is None:
|
||||
return np.zeros(len(embeddings), dtype=int)
|
||||
return np.array([idx % self.n_clusters for idx in range(len(embeddings))])
|
||||
|
||||
class DummyUMAP:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def fit_transform(self, embeddings):
|
||||
raise AssertionError("Psi tree builder must use original embeddings, not UMAP")
|
||||
|
||||
sklearn_module = types.ModuleType("sklearn")
|
||||
mixture_module = types.ModuleType("sklearn.mixture")
|
||||
mixture_module.GaussianMixture = DummyGaussianMixture
|
||||
cluster_module = types.ModuleType("sklearn.cluster")
|
||||
cluster_module.AgglomerativeClustering = DummyAgglomerativeClustering
|
||||
umap_module = types.ModuleType("umap")
|
||||
umap_module.UMAP = DummyUMAP
|
||||
task_service_module = types.ModuleType("api.db.services.task_service")
|
||||
task_service_module.has_canceled = lambda task_id: False
|
||||
connection_utils_module = types.ModuleType("common.connection_utils")
|
||||
connection_utils_module.timeout = lambda seconds: lambda fn: fn
|
||||
exceptions_module = types.ModuleType("common.exceptions")
|
||||
exceptions_module.TaskCanceledException = TaskCanceledException
|
||||
token_utils_module = types.ModuleType("common.token_utils")
|
||||
token_utils_module.truncate = lambda text, max_len: text[:max_len]
|
||||
graphrag_utils_module = types.ModuleType("rag.graphrag.utils")
|
||||
graphrag_utils_module.chat_limiter = DummyLimiter()
|
||||
graphrag_utils_module.get_embed_cache = lambda *args, **kwargs: None
|
||||
graphrag_utils_module.get_llm_cache = lambda *args, **kwargs: None
|
||||
graphrag_utils_module.set_embed_cache = lambda *args, **kwargs: None
|
||||
graphrag_utils_module.set_llm_cache = lambda *args, **kwargs: None
|
||||
|
||||
async def thread_pool_exec(fn, *args, **kwargs):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
misc_utils_module = types.ModuleType("common.misc_utils")
|
||||
misc_utils_module.thread_pool_exec = thread_pool_exec
|
||||
|
||||
monkeypatch.setitem(sys.modules, "sklearn", sklearn_module)
|
||||
monkeypatch.setitem(sys.modules, "sklearn.mixture", mixture_module)
|
||||
monkeypatch.setitem(sys.modules, "sklearn.cluster", cluster_module)
|
||||
monkeypatch.setitem(sys.modules, "umap", umap_module)
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service_module)
|
||||
monkeypatch.setitem(sys.modules, "common.connection_utils", connection_utils_module)
|
||||
monkeypatch.setitem(sys.modules, "common.exceptions", exceptions_module)
|
||||
monkeypatch.setitem(sys.modules, "common.token_utils", token_utils_module)
|
||||
monkeypatch.setitem(sys.modules, "rag.graphrag.utils", graphrag_utils_module)
|
||||
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_module)
|
||||
# Create stub parent packages and load raptor directly via spec_from_file_location
|
||||
# to bypass rag/advanced_rag/__init__.py (which triggers ES connection etc.).
|
||||
_test_dir = os.path.dirname(__file__)
|
||||
_rag_adv_kc_dir = os.path.normpath(os.path.join(_test_dir, "../../../rag/advanced_rag/knowlege_compile"))
|
||||
_rag_adv = types.ModuleType("rag.advanced_rag")
|
||||
_rag_adv.__path__ = [os.path.normpath(os.path.join(_test_dir, "../../../rag/advanced_rag"))]
|
||||
_rag_adv.__package__ = "rag.advanced_rag"
|
||||
monkeypatch.setitem(sys.modules, "rag.advanced_rag", _rag_adv)
|
||||
_rag_adv_kc = types.ModuleType("rag.advanced_rag.knowlege_compile")
|
||||
_rag_adv_kc.__path__ = [_rag_adv_kc_dir]
|
||||
_rag_adv_kc.__package__ = "rag.advanced_rag.knowlege_compile"
|
||||
monkeypatch.setitem(sys.modules, "rag.advanced_rag.knowlege_compile", _rag_adv_kc)
|
||||
monkeypatch.delitem(sys.modules, "rag.advanced_rag.knowlege_compile.raptor", raising=False)
|
||||
_raptor_spec = importlib.util.spec_from_file_location(
|
||||
"rag.advanced_rag.knowlege_compile.raptor",
|
||||
os.path.join(_rag_adv_kc_dir, "raptor.py"),
|
||||
)
|
||||
module = importlib.util.module_from_spec(_raptor_spec)
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.raptor"] = module
|
||||
_raptor_spec.loader.exec_module(module)
|
||||
yield module
|
||||
monkeypatch.delitem(sys.modules, "rag.advanced_rag.knowlege_compile.raptor", raising=False)
|
||||
|
||||
|
||||
class FakeChatModel:
|
||||
llm_name = "fake-chat"
|
||||
max_length = 4096
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
async def async_chat(self, system, history, gen_conf):
|
||||
self.calls.append(history[0]["content"])
|
||||
return f"summary-{len(self.calls)}"
|
||||
|
||||
|
||||
class FakeEmbeddingModel:
|
||||
llm_name = "fake-embedding"
|
||||
|
||||
def encode(self, texts):
|
||||
embeddings = []
|
||||
for text in texts:
|
||||
checksum = sum(ord(ch) for ch in text)
|
||||
embeddings.append(np.array([len(text), checksum % 17 + 1], dtype=float))
|
||||
return embeddings, len(texts)
|
||||
|
||||
|
||||
_DEFAULT_TREE_BUILDER = object()
|
||||
|
||||
|
||||
def _make_raptor(raptor_module, max_cluster=64, tree_builder=_DEFAULT_TREE_BUILDER, **kwargs):
|
||||
if tree_builder is _DEFAULT_TREE_BUILDER:
|
||||
kwargs["tree_builder"] = raptor_module.PSI_TREE_BUILDER
|
||||
else:
|
||||
kwargs["tree_builder"] = tree_builder
|
||||
return raptor_module.RecursiveAbstractiveProcessing4TreeOrganizedRetrieval(
|
||||
max_cluster,
|
||||
FakeChatModel(),
|
||||
FakeEmbeddingModel(),
|
||||
"{cluster_content}",
|
||||
max_token=32,
|
||||
threshold=0.1,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _chunks():
|
||||
return [
|
||||
("alpha first", np.array([1.0, 0.0])),
|
||||
("alpha second", np.array([0.99, 0.01])),
|
||||
("alpha third", np.array([0.98, 0.02])),
|
||||
]
|
||||
|
||||
|
||||
def test_default_tree_builder_remains_original_raptor(raptor_module):
|
||||
raptor = _make_raptor(raptor_module, tree_builder=None)
|
||||
|
||||
assert raptor._tree_builder == raptor_module.RAPTOR_TREE_BUILDER
|
||||
|
||||
|
||||
def test_unknown_tree_builder_is_rejected(raptor_module):
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
|
||||
_make_raptor(raptor_module, tree_builder="ahc")
|
||||
|
||||
|
||||
def test_raptor_config_accepts_hidden_psi_tree_builder():
|
||||
assert RaptorConfig().tree_builder == "raptor"
|
||||
assert RaptorConfig().clustering_method == "gmm"
|
||||
assert RaptorConfig(clustering_method="ahc").clustering_method == "ahc"
|
||||
assert RaptorConfig(tree_builder="psi").tree_builder == "psi"
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
RaptorConfig(tree_builder="ahc")
|
||||
with pytest.raises(ValidationError):
|
||||
RaptorConfig(clustering_method="psi")
|
||||
|
||||
|
||||
def test_ahc_clustering_method_is_supported_in_original_tree_builder(raptor_module):
|
||||
raptor = _make_raptor(raptor_module, tree_builder=raptor_module.RAPTOR_TREE_BUILDER, clustering_method="ahc")
|
||||
|
||||
labels = raptor._get_clusters_ahc(np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1, 10.0]]))
|
||||
|
||||
assert raptor._tree_builder == raptor_module.RAPTOR_TREE_BUILDER
|
||||
assert raptor._clustering_method == "ahc"
|
||||
assert len(labels) == 4
|
||||
|
||||
|
||||
def test_unknown_clustering_method_is_rejected(raptor_module):
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
|
||||
_make_raptor(raptor_module, clustering_method="psi")
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_get_optimal_clusters_handles_max_cluster_equal_one(raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=1)
|
||||
|
||||
optimal = raptor._get_optimal_clusters(
|
||||
np.array([[0.0, 0.0], [0.1, 0.0], [1.0, 1.0], [1.1, 1.1]]),
|
||||
random_state=0,
|
||||
)
|
||||
|
||||
assert optimal == 1
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_get_optimal_clusters_evaluates_upper_bound_candidate(monkeypatch, raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=3)
|
||||
evaluated = []
|
||||
|
||||
class RecordingGaussianMixture:
|
||||
def __init__(self, n_components, random_state=None, **kwargs):
|
||||
self.n_components = n_components
|
||||
evaluated.append(n_components)
|
||||
|
||||
def fit(self, embeddings):
|
||||
return self
|
||||
|
||||
def bic(self, embeddings):
|
||||
scores = {1: 30.0, 2: 20.0, 3: 10.0}
|
||||
return scores[self.n_components]
|
||||
|
||||
monkeypatch.setattr(raptor_module, "GaussianMixture", RecordingGaussianMixture)
|
||||
|
||||
optimal = raptor._get_optimal_clusters(
|
||||
np.array([[0.0, 0.0], [0.1, 0.0], [1.0, 1.0], [1.1, 1.1]]),
|
||||
random_state=0,
|
||||
)
|
||||
|
||||
assert optimal == 3
|
||||
assert evaluated == [1, 2, 3]
|
||||
|
||||
|
||||
def test_psi_tree_builder_ranks_all_leaf_pairs_by_original_cosine_similarity(raptor_module):
|
||||
raptor = _make_raptor(raptor_module)
|
||||
leaves = [
|
||||
raptor_module._PsiTreeNode(index=0, embedding=np.array([1.0, 0.0])),
|
||||
raptor_module._PsiTreeNode(index=1, embedding=np.array([0.0, 1.0])),
|
||||
raptor_module._PsiTreeNode(index=2, embedding=np.array([0.99, 0.01])),
|
||||
raptor_module._PsiTreeNode(index=3, embedding=np.array([-1.0, 0.0])),
|
||||
]
|
||||
|
||||
ranked_pairs = raptor._rank_leaf_pairs(leaves)
|
||||
|
||||
assert len(ranked_pairs) == 6
|
||||
assert tuple(ranked_pairs[0]) == (2, 0)
|
||||
|
||||
|
||||
def test_psi_tree_builder_uses_cosine_similarity_not_vector_magnitude(raptor_module):
|
||||
raptor = _make_raptor(raptor_module)
|
||||
leaves = [
|
||||
raptor_module._PsiTreeNode(index=0, embedding=np.array([100.0, 0.0])),
|
||||
raptor_module._PsiTreeNode(index=1, embedding=np.array([1.0, 1.0])),
|
||||
raptor_module._PsiTreeNode(index=2, embedding=np.array([0.1, 0.0])),
|
||||
]
|
||||
|
||||
ranked_pairs = raptor._rank_leaf_pairs(leaves)
|
||||
|
||||
assert tuple(ranked_pairs[0]) == (2, 0)
|
||||
|
||||
|
||||
def test_psi_tree_builder_handles_zero_vectors_in_cosine_ranking(raptor_module):
|
||||
raptor = _make_raptor(raptor_module)
|
||||
leaves = [
|
||||
raptor_module._PsiTreeNode(index=0, embedding=np.array([0.0, 0.0])),
|
||||
raptor_module._PsiTreeNode(index=1, embedding=np.array([1.0, 0.0])),
|
||||
raptor_module._PsiTreeNode(index=2, embedding=np.array([0.9, 0.1])),
|
||||
]
|
||||
|
||||
ranked_pairs = raptor._rank_leaf_pairs(leaves)
|
||||
|
||||
assert tuple(ranked_pairs[0]) == (2, 1)
|
||||
|
||||
|
||||
def test_psi_tree_builder_collapses_leaf_into_ranked_pair_parent(raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=64)
|
||||
|
||||
root, leaves = raptor._build_psi_structure(_chunks())
|
||||
|
||||
assert len(root.children) == 3
|
||||
assert {child.index for child in root.children} == {0, 1, 2}
|
||||
assert all(leaf.parent is root for leaf in leaves)
|
||||
|
||||
|
||||
def test_psi_tree_builder_collapses_leaf_at_matching_rank(monkeypatch, raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=64)
|
||||
chunks = [
|
||||
("node 0", np.array([1.0, 0.0])),
|
||||
("node 1", np.array([0.9, 0.1])),
|
||||
("node 2", np.array([-1.0, 0.0])),
|
||||
("node 3", np.array([-0.9, -0.1])),
|
||||
("node 4", np.array([0.8, 0.2])),
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
raptor,
|
||||
"_rank_leaf_pairs",
|
||||
lambda _leaves: np.array([[0, 1], [2, 3], [0, 2], [4, 0]]),
|
||||
)
|
||||
|
||||
root, leaves = raptor._build_psi_structure(chunks)
|
||||
|
||||
assert leaves[4].parent is leaves[0].parent
|
||||
assert leaves[4].parent is not root
|
||||
assert len(root.children) == 2
|
||||
|
||||
|
||||
def test_psi_union_find_clamps_out_of_bounds_parent_rank(caplog, raptor_module):
|
||||
union_find = raptor_module._PsiUnionFind(2)
|
||||
union_find._node_ids[1] = [1]
|
||||
union_find._rank[0] = 2
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
union_find._build(0, 1, insert_point=1)
|
||||
|
||||
assert union_find.tree[0] == 1
|
||||
assert "rank index" in caplog.text
|
||||
|
||||
|
||||
def test_psi_tree_builder_rebalances_nodes_over_max_children(raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=2)
|
||||
|
||||
root, _ = raptor._build_psi_structure(_chunks())
|
||||
|
||||
assert all(len(node.children) <= 2 for node in raptor._iter_nodes(root))
|
||||
assert len(root.children) == 2
|
||||
assert any(child.children for child in root.children)
|
||||
|
||||
|
||||
def test_psi_tree_builder_uses_bucketed_structure_for_large_inputs(monkeypatch, raptor_module):
|
||||
chunks = [(f"node {idx}", np.array([float(idx), float(idx % 3 + 1)])) for idx in range(8)]
|
||||
raptor = _make_raptor(
|
||||
raptor_module,
|
||||
max_cluster=3,
|
||||
psi_exact_max_leaves=3,
|
||||
psi_bucket_size=2,
|
||||
)
|
||||
ranked_sizes = []
|
||||
original_rank = raptor._rank_leaf_pairs
|
||||
|
||||
def track_rank(nodes):
|
||||
ranked_sizes.append(len(nodes))
|
||||
return original_rank(nodes)
|
||||
|
||||
monkeypatch.setattr(raptor, "_rank_leaf_pairs", track_rank)
|
||||
|
||||
root, leaves = raptor._build_psi_structure(chunks)
|
||||
|
||||
assert len(leaves) == len(chunks)
|
||||
assert all(leaf.parent is not None for leaf in leaves)
|
||||
assert all(len(node.children) <= 3 for node in raptor._iter_nodes(root))
|
||||
assert max(ranked_sizes) <= 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_psi_tree_builder_materializes_rebalanced_summary_layers_without_umap(monkeypatch, raptor_module):
|
||||
def fail_umap(*args, **kwargs):
|
||||
raise AssertionError("Psi tree builder must use original embeddings, not UMAP")
|
||||
|
||||
monkeypatch.setattr("umap.UMAP", fail_umap)
|
||||
raptor = _make_raptor(raptor_module, max_cluster=2)
|
||||
|
||||
chunks, layers = await raptor(_chunks(), random_state=0)
|
||||
|
||||
assert len(chunks) == 5
|
||||
assert layers == [(0, 3), (3, 4), (4, 5)]
|
||||
assert [chunk[0] for chunk in chunks[3:]] == ["summary-1", "summary-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_psi_tree_builder_skips_failed_node_summary(monkeypatch, raptor_module):
|
||||
raptor = _make_raptor(raptor_module, max_cluster=2)
|
||||
|
||||
async def fail_summary(*args, **kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(raptor, "_summarize_texts", fail_summary)
|
||||
|
||||
chunks, layers = await raptor(_chunks(), random_state=0)
|
||||
|
||||
assert len(chunks) == 3
|
||||
assert [chunk[0] for chunk in chunks] == [chunk[0] for chunk in _chunks()]
|
||||
assert layers == [(0, 3)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_original_raptor_stops_when_transient_summary_fails(monkeypatch, raptor_module):
|
||||
raptor = _make_raptor(raptor_module, tree_builder=raptor_module.RAPTOR_TREE_BUILDER)
|
||||
|
||||
async def fail_summary(*args, **kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(raptor, "_summarize_texts", fail_summary)
|
||||
|
||||
input_chunks = _chunks()[:2]
|
||||
chunks, layers = await raptor(input_chunks, random_state=0)
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert [chunk[0] for chunk in chunks] == [chunk[0] for chunk in input_chunks]
|
||||
assert layers == [(0, 2)]
|
||||
@@ -17,14 +17,8 @@
|
||||
Unit tests for rag/utils/raptor_utils.py module.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from rag.utils.raptor_utils import (
|
||||
RAPTOR_TREE_BUILDER,
|
||||
PSI_TREE_BUILDER,
|
||||
GMM_CLUSTERING_METHOD,
|
||||
AHC_CLUSTERING_METHOD,
|
||||
get_raptor_tree_builder,
|
||||
get_raptor_clustering_method,
|
||||
_as_extra_dict,
|
||||
_has_raptor_marker,
|
||||
_raptor_methods_from_fields,
|
||||
@@ -38,65 +32,6 @@ from rag.utils.raptor_utils import (
|
||||
)
|
||||
|
||||
|
||||
class TestGetRaptorTreeBuilder:
|
||||
"""Tests for get_raptor_tree_builder function."""
|
||||
|
||||
def test_returns_default_raptor_tree_builder(self):
|
||||
"""Test that default tree builder is 'raptor'."""
|
||||
result = get_raptor_tree_builder(None)
|
||||
assert result == RAPTOR_TREE_BUILDER
|
||||
|
||||
def test_returns_default_with_empty_config(self):
|
||||
"""Test that empty config returns default."""
|
||||
result = get_raptor_tree_builder({})
|
||||
assert result == RAPTOR_TREE_BUILDER
|
||||
|
||||
def test_returns_configured_tree_builder(self):
|
||||
"""Test that configured tree builder is returned."""
|
||||
config = {"tree_builder": PSI_TREE_BUILDER}
|
||||
result = get_raptor_tree_builder(config)
|
||||
assert result == PSI_TREE_BUILDER
|
||||
|
||||
def test_returns_ext_tree_builder(self):
|
||||
"""Test that ext.tree_builder takes precedence."""
|
||||
config = {"tree_builder": "old", "ext": {"tree_builder": PSI_TREE_BUILDER}}
|
||||
result = get_raptor_tree_builder(config)
|
||||
assert result == PSI_TREE_BUILDER
|
||||
|
||||
def test_raises_error_for_unsupported_tree_builder(self):
|
||||
"""Test that unsupported tree builder raises ValueError."""
|
||||
config = {"tree_builder": "unknown"}
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
|
||||
get_raptor_tree_builder(config)
|
||||
|
||||
|
||||
class TestGetRaptorClusteringMethod:
|
||||
"""Tests for get_raptor_clustering_method function."""
|
||||
|
||||
def test_returns_default_gmm(self):
|
||||
"""Test that default clustering method is 'gmm'."""
|
||||
result = get_raptor_clustering_method(None)
|
||||
assert result == GMM_CLUSTERING_METHOD
|
||||
|
||||
def test_returns_configured_clustering_method(self):
|
||||
"""Test that configured clustering method is returned."""
|
||||
config = {"clustering_method": AHC_CLUSTERING_METHOD}
|
||||
result = get_raptor_clustering_method(config)
|
||||
assert result == AHC_CLUSTERING_METHOD
|
||||
|
||||
def test_returns_ext_clustering_method(self):
|
||||
"""Test that ext.clustering_method takes precedence."""
|
||||
config = {"clustering_method": "old", "ext": {"clustering_method": AHC_CLUSTERING_METHOD}}
|
||||
result = get_raptor_clustering_method(config)
|
||||
assert result == AHC_CLUSTERING_METHOD
|
||||
|
||||
def test_raises_error_for_unsupported_clustering_method(self):
|
||||
"""Test that unsupported clustering method raises ValueError."""
|
||||
config = {"clustering_method": "unknown"}
|
||||
with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
|
||||
get_raptor_clustering_method(config)
|
||||
|
||||
|
||||
class TestAsExtraDict:
|
||||
"""Tests for _as_extra_dict function."""
|
||||
|
||||
@@ -129,14 +64,14 @@ class TestAsExtraDict:
|
||||
assert result == {}
|
||||
|
||||
def test_parses_python_dict_literal(self):
|
||||
"""Test that Python dict literal is parsed."""
|
||||
"""Test that Python dict literal string is parsed correctly."""
|
||||
input_str = "{'key': 'value'}"
|
||||
result = _as_extra_dict(input_str)
|
||||
assert result == {"key": "value"}
|
||||
|
||||
def test_returns_empty_dict_for_malformed_string(self):
|
||||
"""Test that malformed string returns empty dict."""
|
||||
input_str = "{invalid json}"
|
||||
input_str = "not a dict at all"
|
||||
result = _as_extra_dict(input_str)
|
||||
assert result == {}
|
||||
|
||||
@@ -162,7 +97,7 @@ class TestHasRaptorMarker:
|
||||
|
||||
def test_returns_false_for_list_without_raptor(self):
|
||||
"""Test that list without 'raptor' returns False."""
|
||||
assert _has_raptor_marker(["psi", "other"]) is False
|
||||
assert _has_raptor_marker(["other", "unknown"]) is False
|
||||
|
||||
|
||||
class TestRaptorMethodsFromFields:
|
||||
@@ -173,23 +108,23 @@ class TestRaptorMethodsFromFields:
|
||||
result = _raptor_methods_from_fields({})
|
||||
assert result == {RAPTOR_TREE_BUILDER}
|
||||
|
||||
def test_returns_method_from_extra_dict(self):
|
||||
"""Test that method is extracted from extra dict."""
|
||||
fields = {"extra": {"raptor_method": PSI_TREE_BUILDER}}
|
||||
def test_returns_raptor_method_from_extra_dict(self):
|
||||
"""Test that the RAPTOR method is extracted from extra dict."""
|
||||
fields = {"extra": {"raptor_method": RAPTOR_TREE_BUILDER}}
|
||||
result = _raptor_methods_from_fields(fields)
|
||||
assert result == {PSI_TREE_BUILDER}
|
||||
assert result == {RAPTOR_TREE_BUILDER}
|
||||
|
||||
def test_returns_method_from_extra_field(self):
|
||||
"""Test that method is extracted from extra field directly."""
|
||||
fields = {"extra": "{'raptor_method': 'psi'}"}
|
||||
fields = {"extra": "{'raptor_method': 'raptor'}"}
|
||||
result = _raptor_methods_from_fields(fields)
|
||||
assert result == {PSI_TREE_BUILDER}
|
||||
assert result == {RAPTOR_TREE_BUILDER}
|
||||
|
||||
def test_handles_list_method(self):
|
||||
"""Test that list method is converted to set."""
|
||||
fields = {"extra": {"raptor_method": ["raptor", "psi"]}}
|
||||
fields = {"extra": {"raptor_method": ["raptor", "other"]}}
|
||||
result = _raptor_methods_from_fields(fields)
|
||||
assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
|
||||
assert result == {RAPTOR_TREE_BUILDER, "other"}
|
||||
|
||||
def test_handles_empty_method(self):
|
||||
"""Test that empty method returns default."""
|
||||
@@ -208,21 +143,21 @@ class TestCollectRaptorMethods:
|
||||
|
||||
def test_collects_methods_from_raptor_chunks(self):
|
||||
"""Test that methods are collected from RAPTOR chunks."""
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": PSI_TREE_BUILDER}}}
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": RAPTOR_TREE_BUILDER}}}
|
||||
result = collect_raptor_methods(field_map)
|
||||
assert result == {PSI_TREE_BUILDER}
|
||||
assert result == {RAPTOR_TREE_BUILDER}
|
||||
|
||||
def test_skips_non_raptor_chunks(self):
|
||||
"""Test that non-RAPTOR chunks are skipped."""
|
||||
field_map = {"chunk_1": {"raptor_kwd": "other", "extra": {"raptor_method": PSI_TREE_BUILDER}}}
|
||||
field_map = {"chunk_1": {"raptor_kwd": "other", "extra": {"raptor_method": RAPTOR_TREE_BUILDER}}}
|
||||
result = collect_raptor_methods(field_map)
|
||||
assert result == set()
|
||||
|
||||
def test_collects_multiple_methods(self):
|
||||
"""Test that multiple methods are collected."""
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}, "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}}
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}, "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "other"}}}
|
||||
result = collect_raptor_methods(field_map)
|
||||
assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
|
||||
assert result == {RAPTOR_TREE_BUILDER, "other"}
|
||||
|
||||
|
||||
class TestCollectRaptorChunkIds:
|
||||
@@ -241,7 +176,7 @@ class TestCollectRaptorChunkIds:
|
||||
|
||||
def test_excludes_specified_methods(self):
|
||||
"""Test that specified methods are excluded."""
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}, "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}}
|
||||
field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}, "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "other"}}}
|
||||
result = collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"})
|
||||
assert result == {"chunk_2"}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user