# -*- coding: utf-8 -*-
"""RAPTORRAG standalone test."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from bdirag.rag_methods.raptor_rag import RAPTORRAG
from examples.rag_test_utils import build_vector_rag, install_sklearn_fallback, run_retrieval_test

install_sklearn_fallback()


def build_tree(rag, documents):
    rag.build_tree(documents)


if __name__ == "__main__":
    rag = build_vector_rag(RAPTORRAG, max_tree_depth=2, cluster_size=2)
    run_retrieval_test("RAPTORRAG", rag, index_func=build_tree)