RAGSieve

Código para RAGSieve: Detecting Knowledge-Poisoned Documents in RAG Without a Trusted Reference.

O RAGSieve detecta documentos com envenenamento de conhecimento em dois pontos de um sistema de geração aumentada por recuperação (RAG). RSQ é um filtro online: pontua os cinco documentos selecionados para geração em relação aos ranks 6--20 para a consulta atual. RSG é um scanner offline: pontua cada documento do corpus em relação à sua própria vizinhança semântico-lexical. O pacote contém as implementações de RSQ e RSG usadas no artigo, recuperação exata de corpus completo, filtragem e reabastecimento do Top-5, avaliação em nível de documento e um exemplo compacto executável.

Conteúdo do repositório

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

As três bases de conhecimento disponibilizadas contêm 128.044 documentos NQ, 9.961 documentos HotpotQA e 8.239 passagens MS MARCO. Elas incluem texto bruto, 1.000 consultas amostradas, qrels, respostas-alvo e o subconjunto fixo de 100 consultas de ataque. Os vetores densos são gerados localmente e nunca fazem parte do repositório.

A demonstração contém duas realizações PR-W e duas CEM-C do NQ, com cinco documentos envenenados por consulta. Esses ataques oferecem um caminho curto e inspecionável por ambos os detectores, sem enviar a coleção completa de ataques ou as implementações de geração de ataques.

Instalação

Recomenda-se Python 3.11 ou 3.12, uv e uma GPU compatível com CUDA.

bash install.sh

O detector baixa os modelos do Hugging Face nomeados no artigo no primeiro uso.

Para QA ponta a ponta, copie a configuração compartilhada compatível com OpenAI e preencha os três valores:

cp .env.example .env

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

O mesmo endpoint e modelo geram respostas e realizam o julgamento semântico de ASR.

Avaliação rápida

bash run_demo.sh

Este comando constrói um índice BGE-M3 local para o snapshot RSG de 520 documentos, executa o RSQ nos contextos de recuperação de corpus completo preparados, filtra e reabastece o Top-5 de geração e avalia ambos os detectores. Os resultados são gravados em outputs/demo/:

• rsq-metrics.json: AUROC do RSQ e detecção de envenenamento com 5% de remoção limpa;
• rsg-metrics.json: métricas do RSG no snapshot de corpus sem consulta;
• filtered-top5.jsonl: os cinco documentos retidos para geração após o RSQ.
• rsg-*-top5.jsonl e joint-*-top5.jsonl: contextos retidos para cada ataque de demonstração após o RSG isolado e após a filtragem conjunta. O QA é um comando separado e opcional.

Com os exemplos disponibilizados, o RSQ atinge 100% de AUROC e 100% de detecção de envenenamento com 5% de remoção limpa tanto para PR-W quanto para CEM-C. O RSG atinge 100% de AUROC e 100% de detecção de envenenamento com o mesmo orçamento de remoção limpa; sua regra fixa remove 85% dos documentos envenenados e 0% dos documentos limpos. Esses valores são uma verificação de funcionalidade na demonstração curada, não o resultado agregado do artigo.

Os campos de verdade fundamental são lidos apenas pelos comandos de avaliação. Eles não são passados para o RSQ ou o RSG.

Construir os índices locais completos

bash data_preparation.sh

Por padrão, isso constrói os 3 conjuntos de dados x 3 recuperadores avaliados no artigo: BGE-M3, E5-large-v2 e all-MiniLM-L6-v2. A saída é colocada no diretório ignorado data/indices/. Um subconjunto pode ser selecionado sem editar o script:

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

Cada índice é produzido a partir do corpus completo disponibilizado com as configurações de pooling, prefixo, truncamento e normalização por cosseno do artigo.

RSQ: filtragem online no momento da consulta

Prepare uma recuperação limpa e uma recuperação com os ataques incluídos:

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 realiza busca exata sobre o corpus completo uma vez para cada consulta, pontua os documentos injetados separadamente e salva os 100 primeiros em sua ordem original. O RSQ pontua os candidatos 1--5 em relação aos ranks 6--20; os ranks restantes ficam disponíveis para reabastecimento. Execute o RSQ, calcule as métricas em nível de documento e reabasteça o contexto de geração:

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

O RSQ combina quatro fontes de evidência: concentração de âncoras de resposta, integridade do sistema de escrita, surpresa de modelo de linguagem local e uma transição de alinhamento de consulta. Todas as configurações nos padrões da CLI são as configurações do artigo.

RSG: inspeção offline do corpus

O RSG consome o texto do documento e os vetores já mantidos pelo índice 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

O RSG constrói o grafo exato de cosseno top-16, retém vizinhos semanticamente próximos e lexicalmente distantes, mede o aumento de densidade de cada documento em relação ao piso de sua própria vizinhança e combina a cauda empírica do corpus com a integridade do sistema de escrita.

Implantação conjunta

A implantação conjunta é serial: o RSG exclui documentos primeiro, depois o RSQ pontua os novos cinco primeiros em relação aos ranks 6--20 sobreviventes. Os candidatos sinalizados são removidos e o contexto é reabastecido uma vez a partir do ranking sobrevivente; os documentos de substituição não são repontuados. Após bash run_demo.sh, o exemplo 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