RAGSieve

Reference implementation for “RAGSieve: Self-Referenced Local Contrast for Knowledge-Poison Detection in Retrieval-Augmented Generation.”

RAGSieve detects knowledge-poisoning documents at two control points in a retrieval-augmented generation (RAG) system. RAGSieve-Query (RSQ) is an online filter: it scores the five documents selected for generation against ranks 6--20 for the current query. RAGSieve-Graph (RSG) is an offline scanner: it scores every corpus document against its own semantic--lexical neighborhood. The package contains the implementations used in the preprint, exact full-corpus retrieval, Top-5 filtering and refill, document-level evaluation, and a compact runnable example.

Repository contents

data/datasets/       NQ, HotpotQA, and MS MARCO text knowledge bases
data/demo/           four target queries and 20 example poison documents
docs/                schemas, prompts, and RAGSieve method diagrams
src/ragsieve/        RSQ, RSG, retrieval, filtering, and metrics
install.sh           environment installation
data_preparation.sh  local construction of the nine dense indices
run_demo.sh          end-to-end RSQ and RSG functionality check

The three released knowledge bases contain 128,044 NQ documents, 9,961 HotpotQA documents, and 8,239 MS MARCO passages. They include raw text, 1,000 sampled queries, qrels, target answers, and the fixed 100-query attack subset. Dense vectors are generated locally and are never part of the repository.

The demo contains two PR-W and two CEM-C realizations from NQ, with five poison documents per query. These attacks give a short, inspectable path through both detectors without shipping the complete attack collection or attack-generation implementations.

Installation

Python 3.11 or 3.12, uv, and a CUDA-capable GPU are recommended.

bash install.sh

The detector downloads the Hugging Face models named in the paper on first use.

For end-to-end QA, copy the shared OpenAI-compatible configuration and fill in the three values:

cp .env.example .env

OPENAI_BASE_URL=https://api.openai.com/v1
OPENAI_MODEL=your-model-name
OPENAI_API_KEY=your-api-key

The same endpoint and model generate answers and perform semantic ASR judging.

Quick evaluation

bash run_demo.sh

This command builds a local BGE-M3 index for the 520-document RSG snapshot, runs RSQ on the prepared full-corpus retrieval contexts, filters and refills the generation Top-5, and evaluates both detectors. Results are written to outputs/demo/:

• rsq-metrics.json: RSQ AUROC and poison detection at 5% clean removal;
• rsg-metrics.json: RSG metrics on the query-free corpus snapshot;
• filtered-top5.jsonl: the five documents retained for generation after RSQ.
• rsg-*-top5.jsonl and joint-*-top5.jsonl: retained contexts for each demo attack after RSG alone and after joint filtering. QA is a separate, opt-in command.

With the released examples, RSQ reaches 100% AUROC and 100% poison detection at 5% clean removal for both PR-W and CEM-C. RSG reaches 100% AUROC and 100% poison detection at the same clean-removal budget; its fixed rule removes 85% of the poison documents and 0% of the clean documents. These values are a functionality check on the curated demo, not the paper's aggregate result.

Ground-truth fields are read only by the evaluation commands. They are not passed to RSQ or RSG.

Build the complete local indices

bash data_preparation.sh

By default, this builds the 3 datasets x 3 retrievers evaluated in the paper: BGE-M3, E5-large-v2, and all-MiniLM-L6-v2. The output is placed under the ignored data/indices/ directory. A subset can be selected without editing the script:

DATASETS="nq" ENCODERS="bge-m3" DEVICE="cuda:0" bash data_preparation.sh

Each index is produced from the complete released corpus with the paper's pooling, prefix, truncation, and cosine-normalization settings.

RSQ: online query-time filtering

Prepare a clean retrieval and a retrieval with the bundled attacks:

uv run ragsieve prepare-contexts \
  --dataset-dir data/datasets/nq \
  --index-dir data/indices/nq/bge-m3 \
  --subset data/demo/subset.json \
  --attacks data/demo/attacks.jsonl \
  --output outputs/demo/contexts.jsonl \
  --device cuda:0

prepare-contexts performs exact search over the full corpus once for each query, scores the injected documents separately, and saves the top 100 in their original order. RSQ scores candidates 1--5 against ranks 6--20; the remaining ranks are available for refill. Run RSQ, compute document-level metrics, and refill the generation context:

uv run ragsieve detect \
  --input outputs/demo/contexts.jsonl \
  --output outputs/demo/rsq-predictions.jsonl \
  --device cuda:0

uv run ragsieve evaluate \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --output outputs/demo/rsq-metrics.json

uv run ragsieve filter-contexts \
  --contexts outputs/demo/contexts.jsonl \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --output outputs/demo/filtered-top5.jsonl

uv run ragsieve qa \
  --input outputs/demo/filtered-top5.jsonl \
  --output outputs/demo/qa-records.jsonl \
  --summary outputs/demo/qa-summary.json \
  --env-file .env

RSQ combines four evidence sources: answer-anchor concentration, writing-system integrity, local language-model surprisal, and a query-alignment transition. All settings in the CLI defaults are the paper settings.

RSG: offline corpus inspection

RSG consumes document text and the vectors already maintained by the RAG index:

uv run ragsieve detect-graph \
  --documents data/demo/corpus.jsonl \
  --embeddings data/indices/demo/bge-m3/embeddings.npy \
  --output outputs/demo/rsg-predictions.jsonl \
  --device cuda:0

uv run ragsieve evaluate-graph \
  --predictions outputs/demo/rsg-predictions.jsonl \
  --labels data/demo/labels.jsonl \
  --output outputs/demo/rsg-metrics.json

RSG builds the exact cosine top-16 graph, retains semantic-near and lexical-far neighbors, measures each document's density rise over its own neighborhood floor, and combines the empirical corpus tail with writing-system integrity.

Joint deployment

The paper's joint evaluation computes RSQ on the original retrieval result, combines its flags with RSG's corpus-level flags, and refills Top-5 from the unchanged ranking. RSQ is not rescored after RSG filtering. After bash run_demo.sh, the PR-W example is:

uv run ragsieve filter-contexts \
  --contexts data/demo/contexts.jsonl \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --graph-predictions outputs/demo/rsg-predictions.jsonl \
  --condition pr_w \
  --output outputs/demo/joint-pr_w-top5.jsonl

uv run ragsieve qa \
  --input outputs/demo/joint-pr_w-top5.jsonl \
  --output outputs/demo/joint-pr_w-qa.jsonl \
  --summary outputs/demo/joint-pr_w-qa-summary.json \
  --env-file .env

For RSG alone, omit --predictions. For CEM-C, use --condition ipi_cem_c and separate output paths. --condition binds a RSG scan to the retrieval condition being evaluated; unpoisoned QA uses a separate RSG scan of the clean corpus. The bundled RSG snapshot contains both demo attacks. Full-corpus and clean-QA steps are in the artifact guide.

Component ablations

Saved evidence is sufficient to run the published ablations without loading models again:

uv run ragsieve ablate \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --variant without-answer-anchor \
  --output outputs/demo/rsq-without-answer-anchor.jsonl
uv run ragsieve evaluate \
  --predictions outputs/demo/rsq-without-answer-anchor.jsonl \
  --output outputs/demo/rsq-without-answer-anchor-metrics.json

uv run ragsieve ablate \
  --predictions outputs/demo/rsg-predictions.jsonl \
  --variant corpus-local-only \
  --output outputs/demo/rsg-corpus-local-only.jsonl
uv run ragsieve evaluate-graph \
  --predictions outputs/demo/rsg-corpus-local-only.jsonl \
  --labels data/demo/labels.jsonl \
  --output outputs/demo/rsg-corpus-local-only-metrics.json

RSQ also supports without-script-integrity, without-surprisal, and without-query-alignment; RSG also supports script-integrity-only. The corpus-local ablation assigns the full 5% alert budget to the graph branch. Ablated predictions can be passed to the same filtering and QA commands.

Detection cost

uv run ragsieve benchmark --mode rsq \
  --input data/demo/contexts.jsonl \
  --output outputs/demo/rsq-cost.json --device cuda:0
uv run ragsieve benchmark --mode rsg \
  --input data/demo/corpus.jsonl \
  --embeddings data/indices/demo/bge-m3/embeddings.npy \
  --output outputs/demo/rsg-cost.json --device cuda:0

Measurements exclude model/data loading, retrieval, and QA. Models and embeddings remain resident; CUDA work is synchronized around each measurement. The command reports mean, median, and P95 execution time, throughput, and peak GPU allocation. Use the full NQ index for the paper's corpus-scan workload; the small demo is a functionality check. Workload and environment details are in the artifact guide.

Mapping to the paper

The release provides the RAGSieve results and ablations through the following commands. Baseline entries are obtained from the implementations cited in the paper.

Figures 1 and 2 are explanatory system diagrams rather than generated data figures. Their SVG sources are included as docs/ragsieve-overview.svg and docs/ragsieve-method.svg. Detailed schemas and the expected evaluation protocol are in docs/ARTIFACT.md.

License

The RAGSieve implementation is released under the MIT License. The released datasets and model weights retain their original licenses.