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

import type {
  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 } from './qdrant-regulatory-index';

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 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.');
    await this.index.ensureCollection();
    await this.index.replaceDocuments(chunks.map((chunk, index) => ({ chunk, embedding: vectors[index] })));
    return { documentCount: documents.length, chunkCount: chunks.length, embeddingModel: this.embeddings.model };
  }
}

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;
}