64 lines
3.3 KiB
TypeScript
64 lines
3.3 KiB
TypeScript
import { describe, expect, it, vi } from 'vitest';
|
|
import type { RegulatoryControlDocument, RegulatoryRetrievalCandidate } from '../../src/compliance/regulatory-rag';
|
|
import type { TextEmbeddingClient } from '../../server/compliance/embedding-client';
|
|
import type { IndexedRegulatoryChunk, RegulatoryChunkIndex } from '../../server/compliance/qdrant-regulatory-index';
|
|
import {
|
|
DenseRegulatoryCandidateRetriever,
|
|
RegulatoryIngestionService,
|
|
RegulatoryRetrievalPipeline
|
|
} from '../../server/compliance/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(), 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);
|
|
});
|
|
|
|
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, 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']);
|
|
});
|
|
});
|