232 lines
10 KiB
TypeScript
232 lines
10 KiB
TypeScript
import type {
|
|
RegulatoryChunk,
|
|
RegulatoryControlDocument,
|
|
RegulatoryReranker,
|
|
RegulatoryRetrievalCandidate,
|
|
RegulatoryRetrievalRequest,
|
|
RegulatoryRetriever
|
|
} from '../../src/domain/regulatory-rag';
|
|
import type { TextEmbeddingClient } from './embedding-client';
|
|
import { chunkRegulatoryControl, regulatoryChunkEmbeddingText, validateRegulatoryControlDocument } from './regulatory-chunker';
|
|
import type { RegulatoryChunkIndex, RegulatorySparseVector } from './qdrant-regulatory-index';
|
|
|
|
const HYBRID_RRF_K = 60;
|
|
const BM25_K1 = 1.2;
|
|
const BM25_B = 0.75;
|
|
|
|
export interface RegulatoryCandidateRetriever {
|
|
retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]>;
|
|
}
|
|
|
|
export class DenseRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
|
|
constructor(
|
|
private readonly embeddings: TextEmbeddingClient,
|
|
private readonly index: RegulatoryChunkIndex
|
|
) {}
|
|
|
|
async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
|
|
const [vector] = await this.embeddings.embed([validatedQuery(request.query)]);
|
|
return this.index.searchDense(vector, validatedLimit(request.limit), request.filters);
|
|
}
|
|
}
|
|
|
|
export class SparseRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
|
|
constructor(private readonly index: RegulatoryChunkIndex) {}
|
|
|
|
async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
|
|
const vector = regulatoryQuerySparseVector(validatedQuery(request.query));
|
|
if (!vector) return [];
|
|
return this.index.searchSparse(vector, validatedLimit(request.limit), request.filters);
|
|
}
|
|
}
|
|
|
|
export class HybridRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
|
|
constructor(
|
|
private readonly dense: RegulatoryCandidateRetriever,
|
|
private readonly sparse: RegulatoryCandidateRetriever
|
|
) {}
|
|
|
|
async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
|
|
const limit = validatedLimit(request.limit);
|
|
const [dense, sparse] = await Promise.all([
|
|
this.dense.retrieveCandidates({ ...request, limit }),
|
|
this.sparse.retrieveCandidates({ ...request, limit })
|
|
]);
|
|
return reciprocalRankFuseCandidates([dense, sparse], limit);
|
|
}
|
|
}
|
|
|
|
export class RegulatoryRetrievalPipeline implements RegulatoryRetriever {
|
|
constructor(
|
|
private readonly candidates: RegulatoryCandidateRetriever,
|
|
private readonly reranker?: RegulatoryReranker,
|
|
private readonly rerankCandidateMultiplier = 4
|
|
) {
|
|
if (!Number.isInteger(rerankCandidateMultiplier) || rerankCandidateMultiplier < 1 || rerankCandidateMultiplier > 10) {
|
|
throw new Error('rerankCandidateMultiplier must be an integer between 1 and 10.');
|
|
}
|
|
}
|
|
|
|
async retrieve(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
|
|
const limit = validatedLimit(request.limit);
|
|
const query = validatedQuery(request.query);
|
|
const candidateLimit = this.reranker ? Math.min(100, limit * this.rerankCandidateMultiplier) : limit;
|
|
const candidates = await this.candidates.retrieveCandidates({ ...request, query, limit: candidateLimit });
|
|
if (!this.reranker) return candidates.slice(0, limit);
|
|
return (await this.reranker.rerank(query, candidates)).slice(0, limit);
|
|
}
|
|
}
|
|
|
|
export interface RegulatoryIngestionOptions {
|
|
maximumChunkCharacters?: number;
|
|
embeddingBatchSize?: number;
|
|
}
|
|
|
|
export interface RegulatoryIngestionResult {
|
|
documentCount: number;
|
|
chunkCount: number;
|
|
embeddingModel: string;
|
|
}
|
|
|
|
export class RegulatoryIngestionService {
|
|
private readonly maximumChunkCharacters: number;
|
|
private readonly embeddingBatchSize: number;
|
|
|
|
constructor(
|
|
private readonly embeddings: TextEmbeddingClient,
|
|
private readonly index: RegulatoryChunkIndex,
|
|
options: RegulatoryIngestionOptions = {}
|
|
) {
|
|
this.maximumChunkCharacters = options.maximumChunkCharacters ?? 1_600;
|
|
this.embeddingBatchSize = options.embeddingBatchSize ?? 32;
|
|
if (!Number.isInteger(this.embeddingBatchSize) || this.embeddingBatchSize < 1 || this.embeddingBatchSize > 128) {
|
|
throw new Error('embeddingBatchSize must be an integer between 1 and 128.');
|
|
}
|
|
}
|
|
|
|
async ingest(documents: RegulatoryControlDocument[], onProgress?: (completedChunks: number, totalChunks: number) => Promise<void> | void): Promise<RegulatoryIngestionResult> {
|
|
if (!documents.length) throw new Error('At least one regulatory control document is required.');
|
|
const documentKeys = new Set<string>();
|
|
const chunks = documents.flatMap(input => {
|
|
const document = validateRegulatoryControlDocument(input);
|
|
const key = `${document.corpusId}\0${document.documentId}`;
|
|
if (documentKeys.has(key)) throw new Error(`Duplicate regulatory document identity: ${document.documentId}.`);
|
|
documentKeys.add(key);
|
|
return chunkRegulatoryControl(document, { maximumCharacters: this.maximumChunkCharacters });
|
|
});
|
|
const vectors: number[][] = [];
|
|
await onProgress?.(0, chunks.length);
|
|
for (let offset = 0; offset < chunks.length; offset += this.embeddingBatchSize) {
|
|
const batch = chunks.slice(offset, offset + this.embeddingBatchSize).map(regulatoryChunkEmbeddingText);
|
|
vectors.push(...await this.embeddings.embed(batch));
|
|
await onProgress?.(Math.min(offset + batch.length, chunks.length), chunks.length);
|
|
}
|
|
if (vectors.length !== chunks.length) throw new Error('Embedding provider returned an unexpected vector count.');
|
|
const sparseVectors = buildRegulatorySparseVectors(chunks);
|
|
await this.index.ensureCollection();
|
|
await this.index.replaceDocuments(chunks.map((chunk, index) => ({
|
|
chunk,
|
|
embedding: vectors[index],
|
|
...(sparseVectors.get(chunk.chunkId) ? { sparseVector: sparseVectors.get(chunk.chunkId) } : {})
|
|
})));
|
|
return { documentCount: documents.length, chunkCount: chunks.length, embeddingModel: this.embeddings.model };
|
|
}
|
|
}
|
|
|
|
/** Builds deterministic hashed BM25 document vectors without a scanner-specific vocabulary. */
|
|
export function buildRegulatorySparseVectors(chunks: RegulatoryChunk[]): Map<string, RegulatorySparseVector> {
|
|
const documents = chunks.map(chunk => tokenFrequencies(`${chunk.title} ${chunk.title} ${chunk.text}`));
|
|
const documentFrequency = new Map<number, number>();
|
|
for (const document of documents) {
|
|
for (const index of document.frequencies.keys()) documentFrequency.set(index, (documentFrequency.get(index) ?? 0) + 1);
|
|
}
|
|
const averageLength = documents.reduce((sum, document) => sum + document.length, 0) / Math.max(documents.length, 1);
|
|
const vectors = new Map<string, RegulatorySparseVector>();
|
|
documents.forEach((document, documentIndex) => {
|
|
if (!document.frequencies.size) return;
|
|
const entries = [...document.frequencies.entries()].map(([index, frequency]) => {
|
|
const matchingDocuments = documentFrequency.get(index) ?? 1;
|
|
const inverseDocumentFrequency = Math.log(1 + (documents.length - matchingDocuments + 0.5) / (matchingDocuments + 0.5));
|
|
const normalizedFrequency = frequency * (BM25_K1 + 1)
|
|
/ (frequency + BM25_K1 * (1 - BM25_B + BM25_B * document.length / Math.max(averageLength, 1)));
|
|
return [index, inverseDocumentFrequency * normalizedFrequency] as const;
|
|
}).sort(([left], [right]) => left - right);
|
|
vectors.set(chunks[documentIndex].chunkId, {
|
|
indices: entries.map(([index]) => index),
|
|
values: entries.map(([, score]) => score)
|
|
});
|
|
});
|
|
return vectors;
|
|
}
|
|
|
|
export function regulatoryQuerySparseVector(query: string): RegulatorySparseVector | undefined {
|
|
const indices = [...new Set(regulatoryTokens(query).map(sparseIndex))].sort((left, right) => left - right);
|
|
return indices.length ? { indices, values: indices.map(() => 1) } : undefined;
|
|
}
|
|
|
|
export function reciprocalRankFuseCandidates(
|
|
rankings: RegulatoryRetrievalCandidate[][],
|
|
limit: number
|
|
): RegulatoryRetrievalCandidate[] {
|
|
const fused = new Map<string, { score: number; bestRank: number; candidate: RegulatoryRetrievalCandidate }>();
|
|
for (const ranking of rankings) {
|
|
const seen = new Set<string>();
|
|
let rank = 0;
|
|
for (const candidate of ranking) {
|
|
const controlId = candidate.chunk.controlId;
|
|
if (seen.has(controlId)) continue;
|
|
seen.add(controlId);
|
|
rank += 1;
|
|
const existing = fused.get(controlId);
|
|
const entry = existing ?? { score: 0, bestRank: Number.MAX_SAFE_INTEGER, candidate };
|
|
entry.score += 1 / (HYBRID_RRF_K + rank);
|
|
if (rank < entry.bestRank) {
|
|
entry.bestRank = rank;
|
|
entry.candidate = candidate;
|
|
}
|
|
fused.set(controlId, entry);
|
|
}
|
|
}
|
|
return [...fused.values()]
|
|
.sort((left, right) => right.score - left.score
|
|
|| left.bestRank - right.bestRank
|
|
|| left.candidate.chunk.controlId.localeCompare(right.candidate.chunk.controlId, 'en', { numeric: true }))
|
|
.slice(0, limit)
|
|
.map(entry => ({ ...entry.candidate, score: entry.score }));
|
|
}
|
|
|
|
function tokenFrequencies(value: string): { frequencies: Map<number, number>; length: number } {
|
|
const frequencies = new Map<number, number>();
|
|
const tokens = regulatoryTokens(value);
|
|
for (const token of tokens) {
|
|
const index = sparseIndex(token);
|
|
frequencies.set(index, (frequencies.get(index) ?? 0) + 1);
|
|
}
|
|
return { frequencies, length: tokens.length };
|
|
}
|
|
|
|
function regulatoryTokens(value: string): string[] {
|
|
return (value.normalize('NFKD').toLocaleLowerCase('und').replace(/\p{M}/gu, '').match(/[\p{L}\p{N}]+/gu) ?? [])
|
|
.filter(token => token.length > 1);
|
|
}
|
|
|
|
function sparseIndex(token: string): number {
|
|
let hash = 0x811c9dc5;
|
|
for (let index = 0; index < token.length; index += 1) {
|
|
hash ^= token.charCodeAt(index);
|
|
hash = Math.imul(hash, 0x01000193);
|
|
}
|
|
return hash >>> 0;
|
|
}
|
|
|
|
function validatedQuery(value: string): string {
|
|
if (typeof value !== 'string' || !value.trim()) throw new Error('Regulatory retrieval query must be non-empty.');
|
|
if (value.length > 20_000) throw new Error('Regulatory retrieval query exceeds 20000 characters.');
|
|
return value.trim();
|
|
}
|
|
|
|
function validatedLimit(value: number): number {
|
|
if (!Number.isInteger(value) || value < 1 || value > 100) throw new Error('Regulatory retrieval limit must be between 1 and 100.');
|
|
return value;
|
|
}
|