Doctor-RAG (DR-RAG): Diagnose-and-Repair for Agentic RAG

A diagnose-and-repair framework that corrects failures in Agentic Retrieval-Augmented Generation systems via explicit error localization and prefix reuse.

The framework figure is also available as PDF (figs/main_fig.pdf) for higher resolution.

Overview

Agentic RAG interleaves retrieval and reasoning for multi-hop QA. As reasoning trajectories lengthen, failures become frequent. Existing methods either stop at diagnosis or rely on coarse replanning, incurring high computational cost.

DR-RAG addresses this with a two-stage framework:

Diagnosis — A distilled diagnosis model inspects the failed trajectory, classifies the failure into one of four error types, and localizes the earliest failure point (culprit step).
Repair — Based on the diagnosis, a targeted repair strategy is applied only at the failure point, reusing all valid prefixes and previously retrieved evidence. This avoids costly full reruns.

Error Taxonomy & Repair Strategies

Error Type	When it Occurs	Repair Strategy
Format-Invalid Answer Error	Model answered correctly but string format doesn't match EM	Generate all possible EM-valid answer variants
Reasoning Logic Error	Retrieved evidence is sufficient but reasoning went wrong	Compress evidence + re-reason from scratch
Retriever Failure	Correct queries issued but retriever returned irrelevant docs	Rewrite queries + re-retrieve + generate answer
Reasoning-Induced Missing Retrieval	Reasoning errors caused the model to search for wrong things	Truncate at culprit → multi-turn search with planning

Key Design Choices

Prefix Reuse: The repair system truncates the trajectory at the diagnosed failure point and reuses all valid preceding steps, avoiding redundant computation.
Distilled Diagnosis: A small student model (e.g., 4B/8B params) is distilled from a strong teacher's chain-of-thought annotations, enabling efficient inference-time diagnosis.
Token-Budget-Aware Prompts: All prompts enforce a configurable token limit (default 30K) with intelligent truncation to fit within model context windows.

Main Results

All methods operate on the same set of failed trajectories (EM=0). Repair Rate is averaged across three datasets; ΔEM, ΔF1, and ΔR-L denote per-dataset improvements over the original failed outputs.

Qwen3-8B Backbone

Baseline	Method	Repair Rate	HotpotQA ΔEM	HotpotQA ΔF1	2Wiki ΔEM	2Wiki ΔF1	MuSiQue ΔEM	MuSiQue ΔF1
ReAct	Rerun	12.2	17.6	16.3	15.1	15.3	3.8	3.8
	Step-wise	8.1	13.0	13.0	8.3	12.8	3.2	1.1
	RAG-Critic	13.9	18.1	17.3	15.6	13.1	6.9	5.7
	DR-RAG	21.8	31.6	27.3	26.0	24.4	7.7	6.2
Search-o1	Rerun	9.8	14.7	11.8	7.3	5.2	7.3	3.5
	Step-wise	7.3	12.7	9.1	5.6	1.0	3.5	0.5
	RAG-Critic	9.9	13.2	12.0	5.6	2.0	10.8	9.8
	DR-RAG	22.5	29.8	22.7	19.5	15.3	18.1	15.2
Search-R1	Rerun	16.6	19.4	20.8	23.3	23.5	7.2	8.8
	Step-wise	7.4	7.6	9.0	12.9	14.8	1.7	2.0
	RAG-Critic	12.2	11.8	12.8	15.3	13.8	7.3	8.2
	DR-RAG	19.8	25.4	26.1	25.8	24.4	8.3	9.0

LLaMA-3.1-8B-Instruct Backbone

Baseline	Method	Repair Rate	HotpotQA ΔEM	HotpotQA ΔF1	2Wiki ΔEM	2Wiki ΔF1	MuSiQue ΔEM	MuSiQue ΔF1
ReAct	Rerun	9.6	16.4	13.9	10.1	8.4	2.2	0.3
	Step-wise	8.4	14.3	12.5	8.3	9.2	2.5	0.8
	RAG-Critic	5.3	7.1	5.4	6.9	3.7	1.9	-1.7
	DR-RAG	15.5	27.8	23.4	13.3	9.9	5.4	1.5
Search-o1	Rerun	7.6	12.7	7.8	5.6	0.8	4.5	2.3
	Step-wise	6.8	9.3	4.0	6.7	1.7	4.5	1.4
	RAG-Critic	5.7	9.8	7.2	5.6	2.4	1.7	0.1
	DR-RAG	17.8	28.4	17.8	12.3	4.0	12.5	6.3
Search-R1	Rerun	6.9	8.9	11.6	7.7	9.6	4.1	3.1
	Step-wise	5.4	8.3	9.5	5.2	8.8	2.6	2.7
	RAG-Critic	2.7	3.2	1.2	3.7	2.0	1.2	-0.7
	DR-RAG	12.4	17.2	14.9	13.5	10.0	6.5	4.5

DR-RAG consistently outperforms all baselines across all backbone × baseline × dataset combinations, achieving up to 2.3× higher repair rates than the best baseline.

Project Structure

dr_rag_open/
├── src/                        # All source code
│   ├── diagnose.py             # Stage 1: Unified fault localization (vLLM)
│   ├── repair.py               # Stage 2: Targeted repair with error routing
│   ├── prompt.py               # Token-budget-aware prompt templates
│   ├── metric.py               # Evaluation (EM, F1, ROUGE, BERTScore, SBERT)
│   ├── 1_split_data.py         # Distillation Step 1: Filter EM=0 samples
│   ├── 2_annotate_diagnosis.py # Distillation Step 2: Teacher annotation
│   ├── 3_build_sft_data.py     # Distillation Step 3: Build SFT data
│   ├── train_config.yaml       # LLaMA-Factory training config
│   └── dataset_info.json       # Dataset registration for LLaMA-Factory
├── scripts/                    # Batch execution scripts
│   ├── run_diagnosis.sh
│   ├── run_repair.sh
│   └── run_train.sh
├── figs/                       # Figures
│   └── main_fig_v4.pdf         # Framework overview
├── LLaMA-Factory/              # Training framework (third-party)
├── requirements.txt
└── README.md

Installation

conda create -n drrag python=3.11 -y
conda activate drrag

pip install torch>=2.4.0
pip install vllm>=0.11.0
pip install -r requirements.txt

# Install LLaMA-Factory for training
cd LLaMA-Factory && pip install -e . && cd ..

# NLTK data (for METEOR evaluation)
python -c "import nltk; nltk.download('punkt'); nltk.download('wordnet'); nltk.download('omw-1.4')"

Data Format

Input data (original_data.json) schema:

{
  "Question": "Which film had a younger director, Film A or Film B?",
  "response": "<search>Film A director</search><information>Doc 1 ...</information><reason>...</reason><answer>Film A</answer>",
  "Predicted Answer": "Film A",
  "ground_truth": { "target": ["Film B"] },
  "metric": { "exact_match": 0 },
  "score": { "label_in": 1.0 },
  "retrieved_documents": ["doc text 1", "doc text 2"]
}

Key fields:

response: Full agentic RAG trajectory with XML tags (<search>, <information>, <answer>, <reason>)
metric.exact_match: 0 for failed samples (input to our pipeline)
score.label_in: Evidence coverage (1.0 = all evidence retrieved, <1.0 = partial)
retrieved_documents: List of retrieved document texts

Reproducing the Full Pipeline

Prerequisites

Retrieval Service at http://localhost:8000:

POST /retrieve
Body: {"queries": ["q1", "q2"], "topk": 5, "return_scores": true}
Response: {"results": [[{"document": {"id": "...", "contents": "..."}, "score": 0.9}]]}

Evaluation Models: all-MiniLM-L6-v2, roberta-large
Base LLM: Any instruction-tuned model (e.g., Qwen3-8B, LLaMA-3.1-8B-Instruct)

Step 1: Diagnosis Model Distillation

1a. Prepare data

python src/1_split_data.py --data_root ./data

1b. Annotate with teacher model

python src/2_annotate_diagnosis.py \
    --model_name Qwen3-8B \
    --data_root ./data \
    --api_base https://your-api-endpoint/v1 \
    --api_keys YOUR_KEY_1,YOUR_KEY_2 \
    --teacher_model your-teacher-model \
    --rpm 30 --concurrency 30

1c. Build SFT dataset

python src/3_build_sft_data.py \
    --data_root ./data \
    --model_name Qwen3-8B \
    --output LLaMA-Factory/data/sft_train_distill.json \
    --oversample_factor 2

1d. Train diagnosis model

# Edit src/train_config.yaml: set model_name_or_path and output_dir
bash scripts/run_train.sh 8 src/train_config.yaml

Training hyperparameters:

Parameter	Value
Finetuning	Full parameter
Thinking mode	Enabled (CoT distillation)
Sequence length	16384
Learning rate	5e-6 (cosine, 10% warmup)
Epochs	3
Batch size	1 × 2 grad accum
Precision	bf16
DeepSpeed	ZeRO-2

Step 2: Run Diagnosis

python src/diagnose.py \
    --model qwen3-4b \
    --baseline react \
    --dataset hotpotqa \
    --data_root ./data \
    --model_path /path/to/distilled-diagnosis-model

# Or batch all:
bash scripts/run_diagnosis.sh /path/to/diagnosis-model ./data qwen3-4b

Step 3: Split by Coverage

import json

with open("data/qwen3-4b/react/hotpotqa/unified_diagnosis_result.json") as f:
    data = json.load(f)

cov1 = [d for d in data if d.get('label_in', d.get('score',{}).get('label_in',0)) == 1.0]
cov0 = [d for d in data if d.get('label_in', d.get('score',{}).get('label_in',0)) != 1.0]

with open("data/qwen3-4b/react/hotpotqa/efr_diagnosis_result.json", "w") as f:
    json.dump(cov1, f, indent=2, ensure_ascii=False)
with open("data/qwen3-4b/react/hotpotqa/pcefr_diagnosis_result.json", "w") as f:
    json.dump(cov0, f, indent=2, ensure_ascii=False)

Step 4: Run Repair

cd src && python repair.py \
    --model_name_or_path /path/to/base-model \
    --cov1_file ../data/qwen3-4b/react/hotpotqa/efr_diagnosis_result.json \
    --cov0_file ../data/qwen3-4b/react/hotpotqa/pcefr_diagnosis_result.json \
    --retriever_url http://localhost:8000 \
    --sample_n 0

# Or batch all:
bash scripts/run_repair.sh /path/to/base-model ./data qwen3-4b

Key Hyperparameters

Parameter	Value	Description
`temperature`	0.6	Sampling temperature for repair
`top_p`	0.95	Nucleus sampling
`top_k`	20	Top-k sampling
`max_tokens`	10240	Max generation length
`repetition_penalty`	1.2	Repetition penalty
`max_token_limit`	30000	Input prompt token budget
`max_turns`	4	Max turns for multi-turn search repair
`retriever_topk`	5–10	Documents per retrieval call

Supported Baselines

ReAct: <search>, <information>, <answer> (directly supported)
SearchO1: <|begin_search_query|>...<|end_search_query|> (preprocess to unified format)
SearchR1: \boxed{answer} (preprocess to unified format)

Acknowledgements

We gratefully acknowledge the LLaMA-Factory project for providing an excellent unified fine-tuning framework that powers our diagnosis model training.

Citation

@article{jiao2026doctor,
  title={Doctor-RAG: Failure-Aware Repair for Agentic Retrieval-Augmented Generation},
  author={Shuguang Jiao, Chengkai Huang, Shuhan Qi, Xuan Wang, Yifan Li, Quanchi Weng, Lingchuan Liu, Xunliang Cai, Lina Yao},
  journal={arXiv preprint arXiv:2604.00865},
  year={2026}
}

License

MIT License.

Navigation Menu

Repository files navigation

Doctor-RAG (DR-RAG): Diagnose-and-Repair for Agentic RAG

Overview

Error Taxonomy & Repair Strategies

Key Design Choices

Main Results

Qwen3-8B Backbone

LLaMA-3.1-8B-Instruct Backbone

Project Structure

Installation

Data Format

Reproducing the Full Pipeline

Prerequisites

Step 1: Diagnosis Model Distillation

1a. Prepare data

1b. Annotate with teacher model

1c. Build SFT dataset

1d. Train diagnosis model

Step 2: Run Diagnosis

Step 3: Split by Coverage

Step 4: Run Repair

Key Hyperparameters

Supported Baselines

Acknowledgements

Citation

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
LLaMA-Factory		LLaMA-Factory
figs		figs
scripts		scripts
src		src
.gitignore		.gitignore
README.md		README.md
requirements.txt		requirements.txt

Latest commit

History

Folders and files

Repository files navigation

Doctor-RAG (DR-RAG): Diagnose-and-Repair for Agentic RAG

Overview

Error Taxonomy & Repair Strategies

Key Design Choices

Main Results

Qwen3-8B Backbone

LLaMA-3.1-8B-Instruct Backbone

Project Structure

Installation

Data Format

Reproducing the Full Pipeline

Prerequisites

Step 1: Diagnosis Model Distillation

1a. Prepare data

1b. Annotate with teacher model

1c. Build SFT dataset

1d. Train diagnosis model

Step 2: Run Diagnosis

Step 3: Split by Coverage

Step 4: Run Repair

Key Hyperparameters

Supported Baselines

Acknowledgements

Citation

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages