RAGSieve

Код для RAGSieve: Detecting Knowledge-Poisoned Documents in RAG Without a Trusted Reference.

RAGSieve обнаруживает документы с отравлением знаний в двух точках системы retrieval-augmented generation (RAG). RSQ — это онлайн-фильтр: он оценивает пять документов, выбранных для генерации, относительно рангов 6--20 для текущего запроса. RSG — это офлайн-сканер: он оценивает каждый документ корпуса относительно его собственного семантико-лексического окружения. Пакет содержит реализации RSQ и RSG, использованные в статье, точный полнотекстовый поиск по корпусу, фильтрацию Top-5 и пополнение, оценку на уровне документов и компактный запускаемый пример.

Содержимое репозитория

data/datasets/       NQ, HotpotQA, and MS MARCO text knowledge bases
data/demo/           four target queries and 20 example poison documents
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

Три выпущенные базы знаний содержат 128 044 документа NQ, 9 961 документ HotpotQA и 8 239 пассажей MS MARCO. Они включают необработанный текст, 1 000 выборочных запросов, qrels, целевые ответы и фиксированное подмножество из 100 запросов для атаки. Плотные векторы генерируются локально и никогда не являются частью репозитория.

Демо содержит две реализации PR-W и две CEM-C из NQ, с пятью отравленными документами на запрос. Эти атаки дают короткий, проверяемый путь через оба детектора без поставки полной коллекции атак или реализаций генерации атак.

Установка

Рекомендуются Python 3.11 или 3.12, uv и GPU с поддержкой CUDA.

bash install.sh

Детектор загружает модели Hugging Face, указанные в статье, при первом использовании.

Для сквозного QA скопируйте общую конфигурацию, совместимую с OpenAI, и заполните три значения:

cp .env.example .env

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

Тот же endpoint и модель генерируют ответы и выполняют семантическое судейство ASR.

Быстрая оценка

bash run_demo.sh

Эта команда строит локальный индекс BGE-M3 для снимка RSG из 520 документов, запускает RSQ на подготовленных контекстах полнотекстового поиска, фильтрует и пополняет Top-5 для генерации и оценивает оба детектора. Результаты записываются в outputs/demo/:

• rsq-metrics.json: AUROC RSQ и обнаружение отравы при 5% удалении чистых документов;
• rsg-metrics.json: метрики RSG на снимке корпуса без запросов;
• filtered-top5.jsonl: пять документов, сохранённых для генерации после RSQ.
• rsg-*-top5.jsonl и joint-*-top5.jsonl: сохранённые контексты для каждой демо-атаки после одного RSG и после совместной фильтрации. QA — отдельная команда, запускаемая по желанию.

С выпущенными примерами RSQ достигает 100% AUROC и 100% обнаружения отравы при 5% удалении чистых документов как для PR-W, так и для CEM-C. RSG достигает 100% AUROC и 100% обнаружения отравы при том же бюджете удаления чистых документов; его фиксированное правило удаляет 85% отравленных документов и 0% чистых документов. Эти значения являются проверкой функциональности на подготовленном демо, а не агрегированным результатом статьи.

Поля ground-truth читаются только командами оценки. Они не передаются в RSQ или RSG.

Построение полных локальных индексов

bash data_preparation.sh

По умолчанию это строит 3 набора данных x 3 ретривера, оценённых в статье: BGE-M3, E5-large-v2 и all-MiniLM-L6-v2. Вывод размещается в игнорируемом каталоге data/indices/. Подмножество можно выбрать без редактирования скрипта:

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

Каждый индекс создаётся из полного выпущенного корпуса с настройками пулинга, префикса, усечения и косинусной нормализации из статьи.

RSQ: онлайн-фильтрация во время запроса

Подготовьте чистый поиск и поиск с включёнными атаками:

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 выполняет точный поиск по полному корпусу один раз для каждого запроса, оценивает внедрённые документы отдельно и сохраняет top 100 в их исходном порядке. RSQ оценивает кандидатов 1--5 относительно рангов 6--20; оставшиеся ранги доступны для пополнения. Запустите RSQ, вычислите метрики на уровне документов и пополните контекст генерации:

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 объединяет четыре источника доказательств: концентрацию якорей ответа, целостность системы письма, локальную неожиданность языковой модели и переход выравнивания запроса. Все настройки в значениях по умолчанию CLI являются настройками статьи.

RSG: офлайн-инспекция корпуса

RSG использует текст документов и векторы, уже поддерживаемые индексом RAG:

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 строит точный граф косинусного top-16, сохраняет семантически близких и лексически далёких соседей, измеряет рост плотности каждого документа относительно его собственного минимума окружения и объединяет эмпирический хвост корпуса с целостностью системы письма.

Совместное развёртывание

Совместное развёртывание последовательное: сначала RSG исключает документы, затем RSQ оценивает новые top five относительно выживших рангов 6--20. Помеченные кандидаты удаляются, и контекст пополняется один раз из выжившего ранжирования; документы-замены не переоцениваются. После bash run_demo.sh пример PR-W выглядит так:

uv run ragsieve quarantine-contexts \
  --contexts data/demo/contexts.jsonl \
  --graph-predictions outputs/demo/rsg-predictions.jsonl \
  --condition pr_w \
  --output outputs/demo/serial-pr_w-contexts.jsonl
uv run ragsieve detect \
  --input outputs/demo/serial-pr_w-contexts.jsonl \
  --output outputs/demo/serial-pr_w-rsq.jsonl --device cuda:0
uv run ragsieve filter-contexts \
  --contexts outputs/demo/serial-pr_w-contexts.jsonl \
  --predictions outputs/demo/serial-pr_w-rsq.jsonl \
  --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

Для одного RSG запустите filter-contexts на исходных контекстах только с --graph-predictions и --condition. Для CEM-C используйте --condition ipi_cem_c и отдельные выходные пути. --condition привязывает сканирование RSG к оцениваемому условию поиска; неотравленный QA использует отдельное сканирование RSG чистого корпуса. Входящий в комплект снимок RSG содержит обе демо-атаки. Шаги для полного корпуса и чистого QA описаны в руководстве по артефакту.

Абляции компонентов