amazon-science/state-aware-rag

amazon-science/state-aware-rag

Language: Python

License: NOASSERTION

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-04-17T21:50:35Z

Pushed: 2026-06-30T19:55:22Z

Default branch: main

Fork: no

Archived: no

README: State-Aware RAG =================================

Official implementation of the paper [Reasoning with Memory: Adaptive Information Management for Retrieval-Augmented Generation](https://github.com/amazon-science/state-aware-rag).

> Hieu Man, Ro-ee Tal, Abhishek Kumar, JJ Cho, Benjamin Hsu. > University of Oregon, AWS AI.

State-Aware RAG is a modular multi-hop Retrieval-Augmented Generation (RAG) framework featuring two complementary search paradigms:

1. MCTS (Monte Carlo Tree Search) planner for exploratory, branching multi-step reasoning 2. CoT (Chain-of-Thought) linear planner for lightweight sequential reasoning

The system performs iterative (decompose → retrieve → extract → evaluate → synthesize) cycles while maintaining an explicit evolving reasoning state ("memory"). Trees and intermediate artifacts can be persisted, re-loaded, and post-analyzed.

> Core design goals: transparency, reproducibility, pluggability (LLMs & retrievers), and scalability (multi-processing + concurrency + caching).

Table of Contents -----------------

• Features
• Architecture Overview
• Quick Start
• Environment
• Minimal Run (Single Question)
• Dataset Inference
• Configuration System
• Components & Roles
• Search Modes (MCTS vs CoT)
• Caching & Performance Tips
• Evaluation & Metrics
• Directory Structure
• Extending the Framework
• Troubleshooting
• Roadmap / Ideas
• License & Citation

Features --------

• Multi-hop reasoning with explicit search state (reasoning nodes + memory)
• Two planners: MCTS (branching) and CoT (linear)
• Modular role agents (Generator / Retriever / Extractor / Evaluator)
• Pluggable LLM backends via OpenAI-compatible or LiteLLM interface
• Online (HTTP API) and optional offline (FlashRAG) retrieval
• Structured output extraction with fallback text parsing
• Concurrency + multi-processing + deterministic caching
• Reasoning tree export & reload (resumable search)
• Evaluation suite: F1, Exact Match, Sub EM, Retrieval Recall, LLM Judge
• Hydra + YAML config for reproducible experiment control

--- Architecture Overview ---------------------

High-level loop per question:

1. Generator proposes sub-questions, answers, synthesis attempts, rephrasings, or self-corrections using prompt templates under state_aware_rag/agents/prompts/. 2. Retriever fetches context passages (HTTP retriever API and/or local index). 3. Extractor pulls fine-grained answer spans / structured fields from retrieved context. 4. Evaluator scores candidate answers, ranks or synthesizes final answers, may perform majority vote or reasoning synthesis. 5. Planner (MCTS or CoT) orchestrates iterative expansion until termination conditions (depth, rollouts, convergence). 6. Final answer + reasoning chain saved to a Hugging Face dataset directory under results/ with optional result_trees/ JSONL.

Key files:

• inference.py – main entry (single question or dataset map)
• evaluate.py – metric computation and judged scoring
• state_aware_rag/agents/ – agent abstractions & role implementations
• state_aware_rag/planners/ – search logic (MCTS and CoT)
• configs/infer/ – Hydra configs for experiments
• cache/ – per-model LLM call cache

Each node logs: node type, memory snapshot, confidence, content (sub-question, answer, synthesis, final answer), plus lineage. Trees can be reloaded to continue or analyze.

Quick Start -----------

1. Environment

Python >= 3.10

Install (editable):

pip install -e .

You will need:

• An OpenAI-compatible endpoint (e.g. local server or provider) OR models accessible via LiteLLM routing
• A retriever service URL (see scripts/deploy_retriever_server.sh) or FlashRAG index
• API keys exported (e.g. OPENAI_API_KEY) if required by your LLM backend

Optional (FlashRAG offline retrieval):

pip install flashrag

Prerequisites: deploy the backing servers first

> Important: State-Aware RAG does not run standalone. Before any inference or > evaluation command below, you must have three services reachable (the default > endpoints are configured in configs/infer/base.yaml): > > - an LLM server (OpenAI-compatible) on localhost:30000 — used by the extractor/evaluator-metric models, > - an embedding server on localhost:8000 — used by the generator/evaluator and the retriever's encoder, > - the retriever server on localhost:5000. > > These require GPUs and model weights. See docs/server_deployment.md for the full > guide (LLM + embedding + retriever + evaluation). Minimal local stack:

# Start LLM
python -m sglang.launch_server --model-path Qwen/Qwen3-8B --host 0.0.0.0 --port 30000 --dtype bfloat16 &
# Start embedding server
python -m sglang.launch_server --model-path Qwen/Qwen3-Embedding-4B --is-embedding --host 0.0.0.0 --port 8000 --dtype bfloat16 &
# Start retriever
python -m state_aware_rag.servers.retriever --config configs/servers/retriever-Qwen3-4B-wiki-23.yaml --port 5000 --workers 2 --mmap_index &

2. Minimal Single-Question Inference

All agent settings (models, endpoints, generation params) are defined inline in configs/infer/base.yaml, so a single-question run needs no extra overrides:

python -m inference \
mode=mcts \
question="Who founded the company that created the iPhone?"

Hydra merges any overrides you pass; mode=cot switches planner, and search.save_tree=true persists the reasoning tree. To point an agent at a different model without editing base.yaml, either override a single field (e.g. agents.generator.client_kwargs.model_name=...) or swap in a standalone agent YAML, e.g. agents.generator=configs/generator.yaml (sample per-agent configs are provided under configs/).

3. Dataset Inference

Example (2Wiki dev subset, MCTS):

python -m inference \
mode=mcts \
data.name=2wiki \
data.limit=32 \
num_proc=4 \
search.max_depth=6 \
search.num_rollouts=8 \
search.save_tree=true

Results saved under:

results/mcts/Generator_/Extractor_/Evaluator_//

4. Evaluation

python -m evaluate \
mode=mcts \
data.name=2wiki \
data.metrics='["all"]'

The judge model used for llm_judge is the agents.evaluator_metric block in configs/infer/base.yaml. By default evaluate scores the results...