BubbleRAG

Official implementation of BubbleRAG: Evidence-Driven Retrieval-Augmented Generation for Black-Box Knowledge Graphs.

• Paper: arXiv:2603.20309

If you have any questions or would like to chat, feel free to contact me at azula_fire@163.com or dpan457@connect.hkust-gz.edu.cn. WeChat is also fine: Azula_Fire. I am also happy to discuss anything related to retrieval or GraphRAG.

Overview

Existing graph RAG methods often overlook a practical mismatch: during retrieval, LLMs tend to act according to their own parametric knowledge, but that behavior may not match how knowledge is actually expressed in the target KG. Without training or adaptation on a specific KG, an LLM often has insufficient knowledge of that KG's schema, entities, relations, and evidence patterns. This problem is especially common for private KGs and large-scale KGs, which are effectively black boxes to the LLM.

For example, consider the query: "find experts in machine learning". In a KG, valid answers may appear in many forms: a person who has published many machine learning papers, a senior engineer working at a well-known AI company, or a person whose students include a famous scholar. BubbleRAG focuses on this kind of evidence-driven retrieval problem over black-box KGs. For the full technical formulation and method details, please refer to the paper.

The implementation in this repository supports two usage modes:

1. Reproduction mode: run the provided benchmark configs on included datasets.
2. Custom-data mode: build an index from your own corpus and query it with the same BubbleRAG pipeline.

Repository Layout

BubbleRAG/
  bubblerag/
    api.py                    # High-level BubbleRAG class API
    bubble.py                 # BubbleRAG retrieval and query pipeline
    storage/                  # Milvus-backed KG/vector storage
    llm/                      # OpenAI-compatible LLM wrapper
    embedding/                # OpenAI-compatible embedding wrapper
  information_extraction/
    index.py                  # Corpus -> chunks/entities/relations/triples
  scripts/
    build_index.py            # Build index for any corpus
    query.py                  # Query an existing index
    evaluate.py               # Run batch QA evaluation
    reproduce.py              # Run reproduction configs
  configs/
    reproduce_hotpotqa.json
    reproduce_2wikimultihopqa.json
    reproduce_musique.json
    custom_template.json
  assets/
    main.png                  # README overview figure
    main.pdf                  # Source figure
  dataset/                    # Included benchmark subsets and corpora
    quickstart_corpus.json     # 3-document quick start corpus
  DATA_FORMAT.md              # Custom corpus/QA JSON formats

Installation

Use Python 3.9+.

pip install -e .

BubbleRAG uses OpenAI-compatible chat and embedding APIs. Set credentials through environment variables:

export LLM_API_KEY=your_llm_key
export EMB_API_KEY=your_embedding_key

Model names and base URLs are normal configuration values. For CLI runs, set them in the models section of the JSON config. For Python usage, pass them directly to BubbleRAG.from_openai(). No API keys are stored in config files or source code.

Quick Start: Python API

This is the recommended usage.

from bubblerag import BubbleRAG

WORKING_DIR = "runs/quickstart/kg"
CORPUS_FILE = "dataset/quickstart_corpus.json"
QUESTION = "Who leads the knowledge systems group at Arcadia AI?"

LLM_API_KEY = "your-llm-api-key"
LLM_MODEL_NAME = "your-chat-model"
LLM_BASE_URL = "https://your-llm-endpoint/v1"

EMB_API_KEY = "your-embedding-api-key"
EMB_MODEL_NAME = "your-embedding-model"
EMB_BASE_URL = "https://your-embedding-endpoint/v1"

rag = BubbleRAG.from_openai(
    working_dir=WORKING_DIR,
    llm_api_key=LLM_API_KEY,
    embedding_api_key=EMB_API_KEY,
    llm_model_name=LLM_MODEL_NAME,
    llm_base_url=LLM_BASE_URL,
    embedding_model_name=EMB_MODEL_NAME,
    embedding_base_url=EMB_BASE_URL,
)

summary = rag.insert(CORPUS_FILE, batch_size=3, max_concurrency=6)
print(summary)

result = rag.query(QUESTION)
print(result["answer"])

The same example is available as test.py. After the index has already been built, instantiate the class with the same working_dir and call query() directly:

rag = BubbleRAG.from_openai(
    working_dir=WORKING_DIR,
    llm_api_key=LLM_API_KEY,
    embedding_api_key=EMB_API_KEY,
    llm_model_name=LLM_MODEL_NAME,
    llm_base_url=LLM_BASE_URL,
    embedding_model_name=EMB_MODEL_NAME,
    embedding_base_url=EMB_BASE_URL,
)

result = rag.query(QUESTION)

Quick Start: CLI

Build an index from a corpus:

python -m scripts.build_index \
  --corpus-file dataset/quickstart_corpus.json \
  --working-dir runs/quickstart/kg

Query an existing index:

python -m scripts.query \
  --working-dir runs/quickstart/kg \
  --question "Who leads the knowledge systems group at Arcadia AI?"

Evaluate on a QA file:

python -m scripts.evaluate \
  --working-dir runs/custom/kg \
  --qa-file path/to/qa.json \
  --output-file runs/custom/results.json

Run only a small subset for a smoke test:

python -m scripts.evaluate \
  --working-dir runs/custom/kg \
  --qa-file path/to/qa.json \
  --output-file runs/custom/results.json \
  --limit 10

Reproducing Paper Experiments

HotpotQA:

python -m scripts.reproduce --config configs/reproduce_hotpotqa.json

2WikiMultihopQA:

python -m scripts.reproduce --config configs/reproduce_2wikimultihopqa.json

MuSiQue:

python -m scripts.reproduce --config configs/reproduce_musique.json

If the index already exists and you only want to rerun QA evaluation:

python -m scripts.reproduce \
  --config configs/reproduce_hotpotqa.json \
  --skip-index

Reproduction configs store dataset paths, output paths, runtime parameters, and retrieval parameters. API keys remain environment variables.

Reproducing Ablation Experiments

The ablations are exposed through scripts.reproduce --ablation:

# Full BubbleRAG
python -m scripts.reproduce --config configs/reproduce_hotpotqa.json

# Without anchor specialization
python -m scripts.reproduce \
  --config configs/reproduce_hotpotqa.json \
  --skip-index \
  --ablation no-anchor-specialization

# Without schema relaxation
python -m scripts.reproduce \
  --config configs/reproduce_hotpotqa.json \
  --skip-index \
  --ablation no-schema-relaxation

# Without CEG ranking
python -m scripts.reproduce \
  --config configs/reproduce_hotpotqa.json \
  --skip-index \
  --ablation no-ceg-ranking

Use the matching reproduction config for 2WikiMultihopQA or MuSiQue runs. These ablations only change query-time retrieval behavior, so the same built index can be reused with --skip-index. By default, each run writes a timestamped result file, for example runs/hotpotqa/results.no-schema-relaxation.20260601-153000.json, and each record includes an ablation field. Reproduction outputs are stored as:

{
  "metadata": {
    "run_id": "20260601-153000",
    "ablation": "no-schema-relaxation",
    "runtime_config": {
      "batch_size": 16,
      "max_concurrency": 6,
      "limit": 10
    },
    "retrieval_config": {
      "bfs_step": 6,
      "bfs_min_k": 2,
      "use_schema_relaxation": false
    },
    "elapsed_seconds": 0.0
  },
  "results": []
}

The ablation switches correspond to the following config fields:

• use_anchor_specialization=false: use query-only keyword extraction instead of context-specialized keyword extraction.
• use_schema_relaxation=false: skip the initial chunk/community retrieval and early answer check before graph retrieval.
• use_ceg_ranking=false: use the first generated Steiner candidate instead of ranking candidate evidence graphs by the BubbleRAG scoring function.

Custom Data

BubbleRAG expects a corpus JSON file with title and text fields:

[
  {
    "title": "Example Document",
    "text": "Document text used for entity and triple extraction."
  }
]

For evaluation, the QA file should contain question and answer fields:

[
  {
    "question": "Who founded Example Document?",
    "answer": "Alice",
    "answer_aliases": ["Alice Smith"]
  }
]

See DATA_FORMAT.md for details.

Configuration

Use configs/custom_template.json as a starting point for custom experiments:

python -m scripts.build_index --config configs/custom_template.json
python -m scripts.evaluate --config configs/custom_template.json

Important fields:

• models.llm_model_name: OpenAI-compatible chat model name
• models.llm_base_url: chat completion API base URL
• models.embedding_model_name: embedding model name
• models.embedding_base_url: embedding API base URL
• paths.corpus_file: corpus used for indexing
• paths.qa_file: QA file used for evaluation
• paths.working_dir: Milvus Lite database directory
• paths.output_file: evaluation result path
• runtime.batch_size: embedding batch size
• runtime.max_concurrency: LLM extraction concurrency
• runtime.corpus_limit: number of corpus records to index; 0 means all
• runtime.limit: number of QA records to evaluate; 0 means all
• retrieval_config.*: retrieval and expansion parameters
• retrieval_config.use_anchor_specialization: enable context-specialized keyword extraction
• retrieval_config.use_schema_relaxation: enable initial chunk retrieval and early answer checking
• retrieval_config.use_ceg_ranking: enable BubbleRAG candidate evidence graph ranking

The CLI also accepts optional environment overrides named LLM_MODEL_NAME, LLM_BASE_URL, EMB_MODEL_NAME, and EMB_BASE_URL; config files are usually clearer for reproducible runs.

Outputs

By default, experiment artifacts should be written under runs/:

runs/
  hotpotqa/
    kg/              # Milvus Lite database
    results.20260601-153000.json
    results.no-schema-relaxation.20260601-153100.json

Each evaluation record contains the question, ground-truth answer(s), generated answer, LLM reasoning text, F1 score, runtime, and error information if the query failed. The top-level metadata records the run id, output path, model names and base URLs, runtime settings, full effective retrieval config, total elapsed time, and completion progress. Passing --output-file disables automatic timestamping and writes exactly to the requested path.

Notes

• Index construction calls both an LLM and an embedding model. Start with a small corpus when testing custom data.
• The quality of the KG depends on entity/triple extraction quality.
• working_dir identifies the database. Reusing the same directory lets you query an existing index without rebuilding it.
• This is research code; configs are intended to make experiments transparent and easy to modify.

Citation

@misc{pan2026bubblerag,
  title={BubbleRAG: Evidence-Driven Retrieval-Augmented Generation for Black-Box Knowledge Graphs},
  author={Pan, Duyi and Lou, Tianao and Li, Xin and Song, Haoze and Wu, Yiwen and Deng, Mengyi and Yang, Mingyu and Wang, Wei},
  year={2026},
  eprint={2603.20309},
  archivePrefix={arXiv}
}