SAG Benchmark

Companion benchmark reproduction repository for the SAG paper. This repository is for reproducing the paper benchmark scores with the quick-start commands. General users, please see the SAG project.

English | 中文

Paper: https://arxiv.org/abs/2606.15971

SAG-How-It-Works-English.mp4

Benchmark Score Reproduction

This repository contains upload, retrieval, and evaluation scripts for SAG on HotpotQA, 2WikiMultiHopQA, and MuSiQue. The current quick-start workflow reproduces the SAG, BM25, and bge-large-en-v1.5 vector retrieval results, and supports the Triple indexing ablation through the atomic upload/search path.

Default paper setup:

Item	Value
Embedding	`bge-large-en-v1.5`
LLM	`qwen3.6-flash`
Main paper metrics	Recall@5 / F1
Main scripts	`scripts/run_upload.py`, `scripts/run_search_benchmark.py`, `scripts/run_qa_benchmark.py`

Main results:

Experiments on HotpotQA, 2WikiMultiHopQA, and MuSiQue show that SAG achieves the best retrieval and end-to-end QA performance on every benchmark.

Across the three datasets, SAG averages 90.07%/72.96% in Recall@5 and F1, outperforming the strongest baseline for each metric by 6.79/4.33 percentage points, respectively.
On the most challenging MuSiQue dataset, SAG outperforms the strongest baseline for each metric by 11.52/7.01 percentage points in Recall@5 and F1, respectively.

Note: The current CLI supports reproducing SAG2, BM25, and the bge-large-en-v1.5 vector retrieval results. Atomic indexing/search is available for the Triple indexing ablation.

Method Figures

SAG organizes text into lightweight chunk -> event, chunk -> entities, and event <-> entities indexes. It does not maintain a heavy global knowledge graph; it uses the event/entity index for SQL, vector search, full-text search, and multi-hop expansion.

Quick Start

1. Install Dependencies

Requirements:

Python 3.11+
uv
Docker Compose
Available LLM and embedding endpoints; a rerank endpoint is only needed for the final-selection reranker ablation

uv sync
cp .env.example .env

Activate the virtual environment when you want to run Python commands directly:

source .venv/bin/activate

On Windows PowerShell:

.\.venv\Scripts\Activate.ps1

Edit .env and fill in the storage backend, LLM, and embedding settings. Standard SAG reproduction only needs these three groups of settings. Add the rerank settings only when running the final-selection reranker ablation described below.

At minimum, check these .env values before running upload or search:

STORAGE_PROFILE=mysql_es

LLM_API_KEY=sk-...
LLM_BASE_URL=https://...
LLM_MODEL=qwen3.6-flash
LLM_LANGUAGE=en

EMBEDDING_API_KEY=...
EMBEDDING_BASE_URL=http://...
EMBEDDING_MODEL_NAME=text-embedding-bge-large-en-v1.5
EMBEDDING_DIMENSIONS=1024

Only add the following settings for the final-selection reranker ablation:

RERANK_BASE_URL=http://...
RERANK_MODEL_NAME=Qwen/Qwen3-Reranker-8B
RERANK_ENDPOINT=/rerank

Then fill the storage connection used by your profile:

# mysql_es
MYSQL_HOST=localhost
MYSQL_PORT=3306
MYSQL_USER=sag2
MYSQL_PASSWORD=sag2_pass
MYSQL_DATABASE=sag2

# oceanbase_es / oceanbase_full
OCEANBASE_HOST=localhost
OCEANBASE_PORT=2881
OCEANBASE_USER=sag2@sag2
OCEANBASE_PASSWORD=sag2_pass
OCEANBASE_DATABASE=sag2

# mysql_es / oceanbase_es
ES_HOST=localhost
ES_PORT=9200
ES_SCHEME=http

Storage is selected through STORAGE_PROFILE. The application code uses one storage facade, so upload and search callers do not need to know whether vectors are stored in Elasticsearch or OceanBase.

Profile	SQL database	Vector/search backend	Notes
`mysql_es`	MySQL	Elasticsearch	Default paper-compatible local setup
`oceanbase_es`	OceanBase	Elasticsearch	Uses OceanBase for structured tables, keeps ES for vector search
`oceanbase_full`	OceanBase	OceanBase	Stores vectors in OceanBase table columns and searches through OceanBase vector indexes

DATABASE_BACKEND and VECTOR_BACKEND are advanced overrides. Leave them empty unless you need to bypass the profile mapping.

2. Start Local Services

All local services are managed by docker-compose.yml.

Service	Container	Default port	Notes
MySQL	`sag2_mysql`	`3306`	Default user `sag2`
Elasticsearch	`new_sag_elasticsearch`	`9200`	Security disabled
OceanBase	`oceanbase-ce`	`2881`	Optional backend for `oceanbase_es` / `oceanbase_full`
MLflow	`sag2_mlflow`	`5000`	Optional experiment tracking

Ports can be overridden in .env with MYSQL_PORT, ES_PORT, and MLFLOW_PORT. OceanBase is exposed on 2881 for SQL client traffic.

Choose one startup path for the selected STORAGE_PROFILE. Do not run all of them.

2.1 `mysql_es`

docker compose up -d mysql elasticsearch
docker compose ps

2.2 `oceanbase_es`

docker compose up -d oceanbase elasticsearch
docker compose ps

For oceanbase_es, wait until the OceanBase container log shows that tenant DDL is ready and init.sql has completed before running project initialization:

==> sag2 tenant ready.
==> Waiting for sag2 tenant DDL...
==> sag2 tenant DDL ready.
==> init.sql executed.
==> All done.

2.3 `oceanbase_full`

docker compose up -d oceanbase
docker compose ps

For oceanbase_full, also wait until the OceanBase container log shows that tenant DDL is ready and init.sql has completed:

==> sag2 tenant ready.
==> Waiting for sag2 tenant DDL...
==> sag2 tenant DDL ready.
==> init.sql executed.
==> All done.

Optional MLflow tracking can be started separately:

docker compose up -d mlflow

3. Initialize Database and Indexes

Choose one initialization path for the selected STORAGE_PROFILE. Do not run all of them.

3.1 `mysql_es`

uv run python scripts/init_database.py --fix-grants
uv run python scripts/init_elasticsearch.py

3.2 `oceanbase_es`

uv run python scripts/init_database.py
uv run python scripts/init_elasticsearch.py

3.3 `oceanbase_full`

uv run python scripts/init_database.py

scripts/init_database.py reads STORAGE_PROFILE and initializes the active SQL backend:

mysql_es: creates the normal MySQL structured tables.
oceanbase_es: creates the normal OceanBase structured tables and adds OceanBase-only compatibility columns such as source_event.entities.
oceanbase_full: does everything from oceanbase_es, then idempotently adds OceanBase vector columns and vector indexes for chunks, events, entities, and event-entity relations.

Run scripts/init_elasticsearch.py only when the active vector backend is Elasticsearch (mysql_es or oceanbase_es). It is not required for oceanbase_full. The --fix-grants option is only for local MySQL permission repair; do not use it for OceanBase profiles.

4. Upload Datasets

run_upload.py first converts pipeline/evaluation/dataset/<dataset>.json into a Markdown corpus, then writes structured rows and vectors through the configured storage facade. With mysql_es, data goes to MySQL plus Elasticsearch. With oceanbase_full, structured data and vectors are written into OceanBase. After upload, it generates:

pipeline/evaluation/source/SAG/<LLM_MODEL>/<dataset>/<timestamp>/source_info.json

The file contains the source_config_id used by benchmark runs.

uv run python scripts/run_upload.py --dataset hotpotqa
uv run python scripts/run_upload.py --dataset 2wikimultihopqa
uv run python scripts/run_upload.py --dataset musique

For a quick smoke test, use the smaller datasets first:

uv run python scripts/run_upload.py --dataset test_hotpotqa
uv run python scripts/run_upload.py --dataset sample

To reproduce the triplet (atomic event) mode — where each event contains exactly 2 entities (subject-relation-object) — add --atomic when uploading:

uv run python scripts/run_upload.py --dataset sample --atomic

Upload has three supported modes:

Mode	Command	Behavior
Compact (default)	`uv run python scripts/run_upload.py --dataset musique`	Uses the compact extraction prompt and creates merged events. No extra flag is required.
Atomic	`uv run python scripts/run_upload.py --dataset musique --atomic`	Uses the atomic-event template; each event is constrained to a subject-relation-object pair.
No extraction (Vector-only)	`uv run python scripts/run_upload.py --dataset musique --no-extraction`	Loads and indexes the corpus without event/entity extraction; use this when you only want the pure embedding `vector` search path.

Choose one mode per upload. The default compact mode is the recommended path for SAG2 and for experiments that compare SAG2 with vector or BM25. Use --no-extraction only for a vector-only run when SAG is not part of the experiment; SAG2 requires extracted events, so do not use --no-extraction if you plan to run SAG2 as well.

5. Run Paper Reproduction Benchmarks

Quick validation:

uv run python scripts/run_search_benchmark.py \
  --dataset-name test_hotpotqa \
  --strategy sag2 \
  --top-k 10 \
  --k-values "1,2,5,10" \
  --max-concurrency 5 \
  --limit 10

Main datasets:

uv run python scripts/run_search_benchmark.py \
  --dataset-name hotpotqa \
  --strategy sag2 \
  --top-k 10 \
  --k-values "1,2,5,10" \
  --max-concurrency 10 \
  --bench-size 20

uv run python scripts/run_search_benchmark.py \
  --dataset-name 2wikimultihopqa \
  --strategy sag2 \
  --top-k 10 \
  --k-values "1,2,5,10" \
  --max-concurrency 10 \
  --bench-size 20

uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy sag2 \
  --top-k 10 \
  --k-values "1,2,5,10" \
  --max-concurrency 10 \
  --bench-size 20

To run the SAG2 event-candidate-pool variant, enable the scope and set its pool size (k_pool) with --sag2-event-top-k:

uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy sag2 \
  --sag2-scope \
  --sag2-event-top-k 500 \
  --top-k 10

The candidate-pool variant is recorded as a separate search configuration at the same experiment level as the vector and atomic baselines. BM25 is also a first-class strategy:

uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy bm25 \
  --top-k 10

To reproduce the bge-large-en-v1.5 pure-vector baseline, use --no-extraction during upload only when Vector is the sole method being run:

uv run python scripts/run_upload.py --dataset musique --no-extraction
uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy vector \
  --top-k 10

If the experiment compares Vector with SAG2, upload with the default Compact mode instead so the same source contains the events required by SAG2.

To pin a specific uploaded source, pass the source_config_id generated during upload:

uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy sag2 \
  --source-config-id musique-20260512_213908 \
  --top-k 10 \
  --k-values "1,2,5,10" \
  --max-concurrency 10

Enable MLflow:

uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy sag2 \
  --use-mlflow \
  --mlflow-url http://localhost:5000 \
  --mlflow-experiment sag-benchmark

Default output directory:

output/<dataset>/<strategy>/<timestamp>/

Main output files:

File	Description
`search_results.json`	Per-question retrieval results
`benchmark_results.json`	Recall, Precision, F1, and summary metrics
`run.log`	Run log

6. Run QA on Retrieval Results

run_qa_benchmark.py consumes the search_results.json produced by a retrieval run, asks the configured LLM to answer each question using the retrieved passages, and writes EM/F1 results to a new qa_<timestamp>/ directory next to the input file.

uv run python scripts/run_qa_benchmark.py \
  --dataset-name musique \
  --input output/musique/sag2/20260730_172839/search_results.json \
  --qa-top-k 5

--qa-top-k controls how many retrieved passages are placed in each QA prompt. The main output is qa_results.json; use --output-dir to choose a different output directory, or --max-concurrency and --limit to control runtime and scope.

7. Reproduce SAG Ablation Results

The following experiments correspond to Table 6 of the paper. Run them on MuSiQue with the same dataset, source configuration, embedding model, and top-k values, changing only the indicated component.

Default SAG

Use the default Compact upload and the sag2 search strategy. The default final selection is llm_rank.

uv run python scripts/run_upload.py --dataset musique
uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy sag2 \
  --top-k 10 \
  --k-values "1,2,5,10"

Indexing with Triple indexing

Upload with the atomic prompt and select the atomic search strategy:

uv run python scripts/run_upload.py --dataset musique --atomic
uv run python scripts/run_search_benchmark.py \
  --dataset-name musique \
  --strategy atomic \
  --top-k 10 \
  --k-values "1,2,5,10"

Expansion without expansion (`L=0`)

In pipeline/modules/search/config.py, temporarily change the SAG2 expansion switch from:

enabled: bool = Field(default=True, description="是否启用扩展")

to:

enabled: bool = Field(default=False, description="是否启用扩展")

Then upload with the default Compact mode and run --strategy sag2. Restore default=True after the ablation so normal SAG2 runs keep one-hop expansion.

Final selection with Qwen3-Reranker-8B

In pipeline/modules/search/config.py, temporarily change SAG2RerankConfig.strategy from the default LLM selection:

strategy: Literal["rerank", "llm_rank", "rrf"] = Field(
    default="llm_rank", description="排序策略"
)

to:

strategy: Literal["rerank", "llm_rank", "rrf"] = Field(
    default="rerank", description="排序策略"
)

Configure RERANK_BASE_URL, RERANK_MODEL_NAME=Qwen/Qwen3-Reranker-8B, and RERANK_ENDPOINT in .env, then run the default Compact upload and --strategy sag2. Restore default="llm_rank" after the experiment.

Paper Ablation Results on MuSiQue

Stage	Configuration	R@1	R@2	R@5	R@10
Default SAG (Ours)	Default configuration	36.82	63.62	80.36	83.37
Indexing	Triple indexing	35.66	61.83	77.61	81.54
Expansion	w/o Expansion (`L=0`)	35.70	57.75	69.41	74.76
Final selection	Qwen3-Reranker-8B	32.12	48.97	67.11	76.51

The default row uses hyperedge indexing, one-hop expansion (L=1), and Qwen3.6-Flash final selection. Keep all non-ablated settings fixed when comparing rows.

Datasets

Name	Description
`hotpotqa`	HotpotQA multi-hop QA
`2wikimultihopqa`	2WikiMultiHopQA
`musique`	MuSiQue multi-hop QA
`test_hotpotqa`	Small HotpotQA test set
`sample`	Tiny sample set for pipeline debugging

Dataset files are under pipeline/evaluation/dataset/.

Retrieval Strategies

Strategy	Description
`sag2`	SAG2 graph recall, expansion, and LLM reranking
`sag2` + `--sag2-scope`	SAG2 event-candidate-pool variant; `--sag2-event-top-k` sets `k_pool`
`atomic`	Entity-first atomic retrieval with step-by-step hop expansion
`vector`	Pure vector retrieval baseline
`bm25`	Elasticsearch BM25 keyword retrieval baseline

See docs/search.md for full arguments.

Common Scripts

Search Only

uv run python scripts/run_search.py \
  --dataset-name test_hotpotqa \
  --strategy sag2 \
  --output-dir output/manual-search

Evaluate Existing Results

uv run python scripts/run_benchmark.py \
  --results output/<dataset>/<strategy>/<timestamp>/search_results.json \
  --dataset musique

Repository Layout

SAG-Benchmark/
├── assets/                         # README figures and logo
├── pipeline/
│   ├── core/                       # Config, AI clients, storage layer
│   ├── db/                         # SQLAlchemy ORM
│   ├── evaluation/
│   │   ├── dataset/                # Evaluation datasets
│   │   ├── metrics/                # Recall and related metrics
│   │   └── utils/                  # Data loading, MLflow, token tracking
│   ├── modules/
│   │   ├── extract/                # Event/entity extraction
│   │   ├── load/                   # Document loading and chunking
│   │   └── search/                 # Retrieval strategies
│   ├── storage/                    # Storage facade and backend providers
│   └── utils/
├── scripts/
│   ├── init_database.py
│   ├── init_oceanbase_vectors.py
│   ├── init_elasticsearch.py
│   ├── run_upload.py
│   ├── run_search_benchmark.py
│   ├── run_qa_benchmark.py
│   ├── run_search.py
│   └── run_benchmark.py
├── docs/
├── docker-compose.yml
├── .env.example
├── README.md
└── README-CN.md

Reproduction Notes

Results depend on the actual LLM, embedding, and rerank services configured in .env; changing models, embedding dimensions, or rerank settings can change the metrics.
Storage behavior depends on STORAGE_PROFILE. Keep the same profile for initialization, upload, and search unless you intentionally migrate data.
OceanBase vector search uses COSINE distance and returns an ES-compatible _score computed from the returned cosine distance. Approximate ANN search uses OceanBase APPROX LIMIT ... PARAMETERS (ef_search=...).
When --source-config-id is omitted, run_search_benchmark.py looks up the latest uploaded source based on LLM_MODEL in .env.
Full dataset upload and benchmark runs call external model services. Check quota, concurrency, and timeout settings before running.
Stop local services with docker compose down. To delete local database volumes, use docker compose down -v; this removes uploaded data.

Name		Name	Last commit message	Last commit date
.github		.github
assets		assets
docs		docs
pipeline		pipeline
prompts		prompts
scripts		scripts
tests		tests
.env.example		.env.example
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
CHANGELOG.md		CHANGELOG.md
CODE_OF_CONDUCT.md		CODE_OF_CONDUCT.md
CONTRIBUTING.md		CONTRIBUTING.md
LICENSE		LICENSE
README-CN.md		README-CN.md
README.md		README.md
SECURITY.md		SECURITY.md
SUPPORT.md		SUPPORT.md
docker-compose.yml		docker-compose.yml
init.sql		init.sql
ob-start.sh		ob-start.sh
pyproject.toml		pyproject.toml
requirements.txt		requirements.txt
uv.lock		uv.lock

Navigation Menu

Uh oh!

Repository files navigation

SAG Benchmark

Benchmark Score Reproduction

Method Figures

Quick Start

1. Install Dependencies

2. Start Local Services

2.1 `mysql_es`

2.2 `oceanbase_es`

2.3 `oceanbase_full`

3. Initialize Database and Indexes

3.1 `mysql_es`

3.2 `oceanbase_es`

3.3 `oceanbase_full`

4. Upload Datasets

5. Run Paper Reproduction Benchmarks

6. Run QA on Retrieval Results

7. Reproduce SAG Ablation Results

Default SAG

Indexing with Triple indexing

Expansion without expansion (`L=0`)

Final selection with Qwen3-Reranker-8B

Paper Ablation Results on MuSiQue

Datasets

Retrieval Strategies

Common Scripts

Search Only

Evaluate Existing Results

Repository Layout

Reproduction Notes

About

Releases

Packages

Contributors

Languages

Uh oh!

Latest commit

History

Folders and files

Repository files navigation

SAG Benchmark

Benchmark Score Reproduction

Method Figures

Quick Start

1. Install Dependencies

2. Start Local Services

2.1 mysql_es

2.2 oceanbase_es

2.3 oceanbase_full

3. Initialize Database and Indexes

3.1 mysql_es

3.2 oceanbase_es

3.3 oceanbase_full

4. Upload Datasets

5. Run Paper Reproduction Benchmarks

6. Run QA on Retrieval Results

7. Reproduce SAG Ablation Results

Default SAG

Indexing with Triple indexing

Expansion without expansion (L=0)

Final selection with Qwen3-Reranker-8B

Paper Ablation Results on MuSiQue

Datasets

Retrieval Strategies

Common Scripts

Search Only

Evaluate Existing Results

Repository Layout

Reproduction Notes

About

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages

2.1 `mysql_es`

2.2 `oceanbase_es`

2.3 `oceanbase_full`

3.1 `mysql_es`

3.2 `oceanbase_es`

3.3 `oceanbase_full`

Expansion without expansion (`L=0`)