Assessor_2.0/tests/server/regulatory-rag-service.test.ts

import { describe, expect, it, vi } from 'vitest';
import type { RegulatoryControlDocument, RegulatoryRetrievalCandidate } from '../../src/domain/regulatory-rag';
import type { TextEmbeddingClient } from '../../server/regulatory/embedding-client';
import type { IndexedRegulatoryChunk, RegulatoryChunkIndex } from '../../server/regulatory/qdrant-regulatory-index';
import {
  buildRegulatorySparseVectors,
  DenseRegulatoryCandidateRetriever,
  HybridRegulatoryCandidateRetriever,
  RegulatoryIngestionService,
  RegulatoryRetrievalPipeline,
  SparseRegulatoryCandidateRetriever
} from '../../server/regulatory/regulatory-rag-service';

const document: RegulatoryControlDocument = {
  schemaVersion: 1, corpusId: 'nis2', corpusVersion: '2024-1', documentId: 'nis2-21-2-a-de',
  framework: 'NIS2', frameworkVersion: '2022/2555', controlId: '21(2)(a)', language: 'de',
  title: 'Risikoanalyse', text: 'Synthetischer Kontrolltext. '.repeat(30)
};

describe('regulatory RAG services', () => {
  it('embeds every chunk before replacing documents in the derived index', async () => {
    const embeddingBatches: string[][] = [];
    const embeddings: TextEmbeddingClient = {
      model: 'multilingual-model', dimensions: 2,
      embed: async texts => {
        embeddingBatches.push(texts);
        return texts.map((_, index) => [index, 1]);
      }
    };
    const ensureCollection = vi.fn(async () => {});
    const replaceDocuments = vi.fn(async (_chunks: IndexedRegulatoryChunk[]) => {});
    const index: RegulatoryChunkIndex = { ensureCollection, replaceDocuments, searchDense: vi.fn(), searchSparse: vi.fn(), list: vi.fn() };
    const service = new RegulatoryIngestionService(embeddings, index, { maximumChunkCharacters: 220, embeddingBatchSize: 2 });

    const result = await service.ingest([document]);

    expect(result.documentCount).toBe(1);
    expect(result.chunkCount).toBeGreaterThan(1);
    expect(embeddingBatches.every(batch => batch.length <= 2)).toBe(true);
    expect(ensureCollection).toHaveBeenCalledOnce();
    expect(replaceDocuments).toHaveBeenCalledOnce();
    expect(replaceDocuments.mock.calls[0][0]).toHaveLength(result.chunkCount);
    expect(replaceDocuments.mock.calls[0][0].every(entry => entry.sparseVector?.indices.length)).toBe(true);
  });

  it('supports multilingual dense candidates and an optional later reranker stage', async () => {
    const embed = vi.fn(async () => [[0.5, 0.5]]);
    const candidates: RegulatoryRetrievalCandidate[] = Array.from({ length: 8 }, (_, index) => ({
      score: 1 - index / 10,
      chunk: { ...document, chunkId: `chunk-${index}`, chunkIndex: index, textSha256: String(index).padStart(64, '0') }
    }));
    const searchDense = vi.fn(async () => candidates);
    const dense = new DenseRegulatoryCandidateRetriever(
      { model: 'multilingual-model', dimensions: 2, embed },
      { ensureCollection: vi.fn(), replaceDocuments: vi.fn(), searchDense, searchSparse: vi.fn(), list: vi.fn() }
    );
    const rerank = vi.fn(async (_query: string, values: RegulatoryRetrievalCandidate[]) => values.slice().reverse());
    const retrieval = new RegulatoryRetrievalPipeline(dense, { rerank }, 4);

    const result = await retrieval.retrieve({ query: 'How is the German control relevant?', limit: 2, filters: { language: 'de' } });

    expect(embed).toHaveBeenCalledWith(['How is the German control relevant?']);
    expect(searchDense).toHaveBeenCalledWith([0.5, 0.5], 8, { language: 'de' });
    expect(rerank).toHaveBeenCalledOnce();
    expect(result.map(item => item.chunk.chunkId)).toEqual(['chunk-7', 'chunk-6']);
  });

  it('combines dense and lexical rankings once per control with reciprocal-rank fusion', async () => {
    const candidate = (controlId: string, chunkId: string, score: number): RegulatoryRetrievalCandidate => ({
      score,
      chunk: { ...document, controlId, chunkId, chunkIndex: 0, textSha256: chunkId.padEnd(64, '0') }
    });
    const dense = { retrieveCandidates: vi.fn(async () => [
      candidate('C-1', 'dense-c1-a', 0.9),
      candidate('C-1', 'dense-c1-b', 0.89),
      candidate('C-2', 'dense-c2', 0.8)
    ]) };
    const sparse = { retrieveCandidates: vi.fn(async () => [
      candidate('C-2', 'sparse-c2', 4.2),
      candidate('C-3', 'sparse-c3', 3.1)
    ]) };
    const hybrid = new HybridRegulatoryCandidateRetriever(dense, sparse);

    const result = await hybrid.retrieveCandidates({ query: 'synthetic access control', limit: 3 });

    expect(result.map(item => item.chunk.controlId)).toEqual(['C-2', 'C-1', 'C-3']);
    expect(result.map(item => item.chunk.chunkId)).toEqual(['sparse-c2', 'dense-c1-a', 'sparse-c3']);
    expect(dense.retrieveCandidates).toHaveBeenCalledOnce();
    expect(sparse.retrieveCandidates).toHaveBeenCalledOnce();
  });

  it('builds deterministic BM25 vectors and uses the sparse Qdrant slot for lexical retrieval', async () => {
    const chunks = [
      { ...document, chunkId: 'chunk-a', chunkIndex: 0, textSha256: 'a'.repeat(64), title: 'Privileged access', text: 'Restrict privileged access rights.' },
      { ...document, chunkId: 'chunk-b', chunkIndex: 1, textSha256: 'b'.repeat(64), title: 'Backup', text: 'Protect backup copies.' }
    ];
    const vectors = buildRegulatorySparseVectors(chunks);
    expect(vectors.get('chunk-a')?.indices).toEqual(vectors.get('chunk-a')?.indices.slice().sort((left, right) => left - right));
    expect(vectors.get('chunk-a')?.values.every(value => value > 0)).toBe(true);

    const searchSparse = vi.fn(async () => [{ score: 2.4, chunk: chunks[0] }]);
    const sparse = new SparseRegulatoryCandidateRetriever({
      ensureCollection: vi.fn(), replaceDocuments: vi.fn(), searchDense: vi.fn(), searchSparse, list: vi.fn()
    });
    await expect(sparse.retrieveCandidates({ query: 'privileged access', limit: 5, filters: { language: 'de' } }))
      .resolves.toEqual([{ score: 2.4, chunk: chunks[0] }]);
    expect(searchSparse).toHaveBeenCalledWith(expect.objectContaining({ indices: expect.any(Array), values: expect.any(Array) }), 5, { language: 'de' });
  });
});