🚀 Quick Start

# 1. Install
git clone https://github.com/Ayanami0730/arag.git && cd arag
uv sync --extra full                  # or: pip install -e ".[full]"

# 2. Download benchmark datasets from HuggingFace
git clone https://huggingface.co/datasets/Ayanami0730/rag_test data --depth 1
rm -rf data/.git data/README.md

# 3. Build embedding index
#    We use Qwen3-Embedding-0.6B in our paper (https://huggingface.co/Qwen/Qwen3-Embedding-0.6B)
#    You can also use a local path: --model /path/to/Qwen3-Embedding-0.6B
uv run python scripts/build_index.py \
    --chunks data/musique/chunks.json \
    --output data/musique/index \
    --model Qwen/Qwen3-Embedding-0.6B \
    --device cuda:0

# 4. Set environment variables
export ARAG_API_KEY="your-api-key"
export ARAG_BASE_URL="https://api.openai.com/v1"
export ARAG_MODEL="gpt-5-mini"

# 5. Run A-RAG agent
uv run python scripts/batch_runner.py \
    --config configs/example.yaml \
    --questions data/musique/questions.json \
    --output results/musique \
    --limit 10 --workers 5

# 6. Evaluate results
uv run python scripts/eval.py \
    --predictions results/musique/predictions.jsonl \
    --workers 5

Note: Datasets hosted on HuggingFace 🤗, reformatted from Zly0523/linear-rag and GraphRAG-Bench into a unified format.

Don't have uv? Install it: curl -LsSf https://astral.sh/uv/install.sh | sh

✨ News

• [Feb 2026] 📄 Paper released on arXiv
• [Feb 2026] 🚀 Initial code and evaluation suite released

📖 Overview

Frontier language models have demonstrated strong reasoning and long-horizon tool-use capabilities. However, existing RAG systems fail to leverage these capabilities. They still rely on two paradigms:

1. Graph RAG: Designing an algorithm that retrieves passages in a single shot and concatenates them into the model's input
2. Workflow RAG: Predefining a workflow and prompting the model to execute it step-by-step

Neither paradigm allows the model to participate in retrieval decisions, preventing efficient scaling with model improvements.

Three Principles of Agentic RAG

We identify three key principles that define true agentic autonomy:

• Autonomous Strategy: The agent dynamically chooses retrieval strategies based on task characteristics
• Iterative Execution: The agent supports multi-round execution, adapting based on intermediate results
• Interleaved Tool Use: The agent follows a ReAct-like action→observation→reasoning loop

Comparison of three RAG paradigms. Only A-RAG satisfies all three principles, making it a truly agentic framework.

Our Solution: A-RAG

A-RAG is an Agentic RAG framework that exposes hierarchical retrieval interfaces directly to the model. A-RAG provides three retrieval tools: keyword_search, semantic_search, and chunk_read, enabling the agent to adaptively search and retrieve information across multiple granularities.

Overview of A-RAG framework. The agent iteratively uses hierarchical retrieval tools to gather information from the corpus and autonomously decides when to provide the final answer.

Key Features

• 🔍 Hierarchical Retrieval: Keyword-level, sentence-level, and chunk-level information access
• 🤖 True Agentic Autonomy: Autonomous strategy, iterative execution, and interleaved tool use
• 📈 Test-Time Scaling: Performance improves with increased compute resources
• ⚡ Context Efficient: Achieves superior accuracy with comparable or fewer retrieved tokens

📊 Main Results

Results (%) of baselines and A-RAG on benchmark datasets in terms of LLM-Evaluation Accuracy (LLM-Acc) and Contain-Match Accuracy (Cont-Acc). Best results are in bold, second best are underlined.

GPT-4o-mini Backbone

GPT-5-mini Backbone

📁 Project Structure

arag/
├── src/arag/              # Main package
│   ├── core/              # Core modules
│   │   ├── config.py      # Configuration management
│   │   ├── context.py     # Agent context & state tracking
│   │   └── llm.py         # LLM client with cost tracking
│   ├── agent/             # Agent implementations
│   │   ├── base.py        # BaseAgent with ReAct loop
│   │   └── prompts/       # System prompts
│   └── tools/             # Retrieval tools
│       ├── keyword_search.py
│       ├── semantic_search.py
│       └── read_chunk.py
├── scripts/               # CLI scripts
│   ├── build_index.py     # Build embedding index
│   ├── batch_runner.py    # Batch processing
│   └── eval.py            # Evaluation
├── configs/               # Configuration examples
├── tests/                 # Test suite (gitignored, add your own tests)
├── .github/               # Issue templates
└── CITATION.cff           # Citation metadata

🔧 Hierarchical Retrieval Tools

A-RAG provides three retrieval tools that operate at different granularities:

Keyword Search

• Method: Exact lexical matching (case-insensitive)
• Best for: Known entities, names, technical terms
• Score: Score(chunk, keywords) = Σ count(k, chunk) × |k|
• No pre-indexing required

Semantic Search

• Method: Dense retrieval using sentence-level embeddings
• Best for: Conceptual queries, when exact wording is unknown
• Score: Cosine similarity between query and sentence embeddings
• Requires pre-built index

Chunk Read

• Method: Retrieve full content of specified chunks
• Strategy: Read promising chunks identified by search, read adjacent chunks (±1) for context
• Context Tracker: Prevents redundant reading of already-accessed chunks

📚 Benchmarks & Datasets

Supported Datasets

Custom Data Format

Prepare your own corpus as a JSON file:

["0:Document chunk content here...", "1:Another chunk..."]

Full Evaluation Example

# Using HuggingFace model (auto-download)
uv run python scripts/build_index.py \
    --chunks data/musique/chunks.json \
    --output data/musique/index \
    --model Qwen/Qwen3-Embedding-0.6B \
    --device cuda:0

# Or using a local model path
uv run python scripts/build_index.py \
    --chunks data/musique/chunks.json \
    --output data/musique/index \
    --model /path/to/Qwen3-Embedding-0.6B \
    --device cuda:0

llm:
  temperature: 0.0
  max_tokens: 16384
  reasoning_effort: "medium"

embedding:
  model: "Qwen/Qwen3-Embedding-0.6B"  # or local path
  device: "cuda:0"
  batch_size: 16

agent:
  max_loops: 15
  max_token_budget: 128000
  verbose: false

data:
  chunks_file: "data/musique/chunks.json"
  index_dir: "data/musique/index"

export ARAG_API_KEY="your-api-key"
export ARAG_BASE_URL="https://api.openai.com/v1"
export ARAG_MODEL="gpt-5-mini"
export CUDA_VISIBLE_DEVICES=0

# Run all questions
uv run python scripts/batch_runner.py \
    --config configs/test_musique.yaml \
    --questions data/musique/questions.json \
    --output results/musique \
    --workers 10

# Evaluate
uv run python scripts/eval.py \
    --predictions results/musique/predictions.jsonl \
    --workers 10

🐍 Python API

from arag import LLMClient, BaseAgent, ToolRegistry
from arag.tools.keyword_search import KeywordSearchTool
from arag.tools.semantic_search import SemanticSearchTool
from arag.tools.read_chunk import ReadChunkTool

# Initialize LLM client
client = LLMClient(
    model="gpt-5-mini",
    api_key="your-api-key",
    base_url="https://api.openai.com/v1"
)

# Setup tools
tools = ToolRegistry()
tools.register(KeywordSearchTool(chunks_file="data/chunks.json"))
tools.register(SemanticSearchTool(
    chunks_file="data/chunks.json",
    index_dir="data/index",
    embedding_model="Qwen/Qwen3-Embedding-0.6B"
))
tools.register(ReadChunkTool(chunks_file="data/chunks.json"))

# Create agent
agent = BaseAgent(
    llm_client=client,
    tools=tools,
    max_loops=15,
    max_token_budget=128000
)

# Run query
result = agent.run("What is the capital of France?")
print(f"Answer: {result['answer']}")
print(f"Cost: ${result['total_cost']:.6f}")
print(f"Loops: {result['loops']}")

🗺️ Roadmap

• Baseline Scripts: Compatible scripts for all baseline methods (GraphRAG, HippoRAG2, LinearRAG, etc.)
• Ablation Interfaces: Complete interfaces for ablation studies (w/o keyword search, w/o semantic search, w/o chunk read)
• Multi-Provider Support: Native API support for Anthropic Claude and Google Gemini (currently only OpenAI-compatible APIs)
• Additional Benchmarks: Scripts for HotpotQA, 2WikiMQA, and GraphRAG-Bench evaluation
• Visualization Tools: Trajectory visualization and analysis tools

Contributions and feedback are welcome!

📝 Citation

If you use A-RAG in your research, please cite our paper:

@misc{du2026aragscalingagenticretrievalaugmented,
      title={A-RAG: Scaling Agentic Retrieval-Augmented Generation via Hierarchical Retrieval Interfaces}, 
      author={Mingxuan Du and Benfeng Xu and Chiwei Zhu and Shaohan Wang and Pengyu Wang and Xiaorui Wang and Zhendong Mao},
      year={2026},
      eprint={2602.03442},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2602.03442}, 
}

📄 License

MIT License