Assessor_2.0/server/regulatory/qdrant-regulatory-index.ts

import { createHash } from 'node:crypto';
import type {
  RegulatoryChunk,
  RegulatoryRetrievalCandidate,
  RegulatoryRetrievalFilters
} from '../../src/domain/regulatory-rag';

export const REGULATORY_QDRANT_COLLECTION = 'assessor_regulatory_chunks_v1';
export const REGULATORY_DENSE_VECTOR = 'dense';
export const REGULATORY_SPARSE_VECTOR = 'sparse';

export interface QdrantEnvironment {
  ASSESSOR_QDRANT_URL?: string;
  ASSESSOR_QDRANT_API_KEY?: string;
  ASSESSOR_QDRANT_TIMEOUT_MS?: string;
}

export interface IndexedRegulatoryChunk {
  chunk: RegulatoryChunk;
  embedding: number[];
}

export interface RegulatoryChunkIndex {
  ensureCollection(): Promise<void>;
  replaceDocuments(chunks: IndexedRegulatoryChunk[]): Promise<void>;
  searchDense(vector: number[], limit: number, filters?: RegulatoryRetrievalFilters): Promise<RegulatoryRetrievalCandidate[]>;
}

export class QdrantClientError extends Error {
  constructor(readonly code: string, message: string) {
    super(message);
    this.name = 'QdrantClientError';
  }
}

interface QdrantRegulatoryIndexOptions {
  baseUrl: URL;
  apiKey: string;
  dimensions: number;
  embeddingModel: string;
  timeoutMs: number;
  fetcher?: typeof fetch;
}

const payloadIndexes = [
  ['schema_version', 'integer'],
  ['corpus_id', 'keyword'],
  ['corpus_version', 'keyword'],
  ['document_id', 'keyword'],
  ['framework', 'keyword'],
  ['framework_version', 'keyword'],
  ['control_id', 'keyword'],
  ['language', 'keyword'],
  ['embedding_model', 'keyword']
] as const;

export class QdrantRegulatoryIndex implements RegulatoryChunkIndex {
  private readonly fetcher: typeof fetch;

  constructor(private readonly options: QdrantRegulatoryIndexOptions) {
    this.fetcher = options.fetcher ?? globalThis.fetch.bind(globalThis);
  }

  async ensureCollection(): Promise<void> {
    const existing = await this.request('GET', `/collections/${REGULATORY_QDRANT_COLLECTION}`, undefined, [200, 404]);
    if (existing.status === 404) {
      await this.request('PUT', `/collections/${REGULATORY_QDRANT_COLLECTION}`, {
        vectors: { [REGULATORY_DENSE_VECTOR]: { size: this.options.dimensions, distance: 'Cosine' } },
        sparse_vectors: { [REGULATORY_SPARSE_VECTOR]: {} },
        on_disk_payload: true
      });
    } else {
      validateCollection(existing.body, this.options.dimensions);
    }
    const current = existing.status === 404
      ? await this.request('GET', `/collections/${REGULATORY_QDRANT_COLLECTION}`)
      : existing;
    const indexedFields = payloadSchemaFields(current.body);
    for (const [fieldName, fieldSchema] of payloadIndexes) {
      if (indexedFields.has(fieldName)) continue;
      await this.request('PUT', `/collections/${REGULATORY_QDRANT_COLLECTION}/index?wait=true`, {
        field_name: fieldName,
        field_schema: fieldSchema
      });
    }
  }

  async replaceDocuments(chunks: IndexedRegulatoryChunk[]): Promise<void> {
    if (!chunks.length) return;
    for (const entry of chunks) validateEmbedding(entry.embedding, this.options.dimensions);
    const documents = new Map<string, { corpusId: string; documentId: string }>();
    for (const { chunk } of chunks) documents.set(`${chunk.corpusId}\0${chunk.documentId}`, { corpusId: chunk.corpusId, documentId: chunk.documentId });
    for (const document of documents.values()) {
      await this.request('POST', `/collections/${REGULATORY_QDRANT_COLLECTION}/points/delete?wait=true`, {
        filter: { must: [keywordCondition('corpus_id', document.corpusId), keywordCondition('document_id', document.documentId)] }
      });
    }
    for (let offset = 0; offset < chunks.length; offset += 64) {
      const batch = chunks.slice(offset, offset + 64);
      await this.request('PUT', `/collections/${REGULATORY_QDRANT_COLLECTION}/points?wait=true`, {
        points: batch.map(({ chunk, embedding }) => ({
          id: pointId(chunk.chunkId),
          vector: { [REGULATORY_DENSE_VECTOR]: embedding },
          payload: chunkPayload(chunk, this.options.embeddingModel)
        }))
      });
    }
  }

  async searchDense(
    vector: number[],
    limit: number,
    filters: RegulatoryRetrievalFilters = {}
  ): Promise<RegulatoryRetrievalCandidate[]> {
    validateEmbedding(vector, this.options.dimensions);
    if (!Number.isInteger(limit) || limit < 1 || limit > 100) throw new Error('Regulatory retrieval limit must be between 1 and 100.');
    const must = [keywordCondition('embedding_model', this.options.embeddingModel), ...filterConditions(filters)];
    const response = await this.request('POST', `/collections/${REGULATORY_QDRANT_COLLECTION}/points/query`, {
      query: vector,
      using: REGULATORY_DENSE_VECTOR,
      filter: { must },
      limit,
      with_payload: true,
      with_vector: false
    });
    const result = object(response.body)?.result;
    const points = Array.isArray(object(result)?.points) ? object(result)!.points as unknown[] : Array.isArray(result) ? result : undefined;
    if (!points) throw new QdrantClientError('qdrant_invalid_response', 'Qdrant returned an invalid query response.');
    return points.map(point => parseCandidate(point));
  }

  private async request(method: string, path: string, body?: unknown, acceptedStatuses: number[] = [200]): Promise<{ status: number; body: unknown }> {
    const url = new URL(path, withTrailingSlash(this.options.baseUrl));
    let response: Response;
    try {
      response = await this.fetcher(url, {
        method,
        redirect: 'error',
        signal: AbortSignal.timeout(this.options.timeoutMs),
        headers: {
          Accept: 'application/json',
          ...(body === undefined ? {} : { 'Content-Type': 'application/json' }),
          'api-key': this.options.apiKey
        },
        ...(body === undefined ? {} : { body: JSON.stringify(body) })
      });
    } catch {
      throw new QdrantClientError('qdrant_unavailable', 'Qdrant could not be reached.');
    }
    if (!acceptedStatuses.includes(response.status)) throw new QdrantClientError('qdrant_http_error', `Qdrant returned HTTP ${response.status}.`);
    if (response.status === 404) return { status: 404, body: undefined };
    const text = await readBoundedText(response, 16 * 1024 * 1024);
    try { return { status: response.status, body: text ? JSON.parse(text) : undefined }; }
    catch { throw new QdrantClientError('qdrant_invalid_response', 'Qdrant returned invalid JSON.'); }
  }
}

export function createQdrantRegulatoryIndexFromEnvironment(
  dimensions: number,
  embeddingModel: string,
  environment: QdrantEnvironment = {
    ASSESSOR_QDRANT_URL: process.env.ASSESSOR_QDRANT_URL,
    ASSESSOR_QDRANT_API_KEY: process.env.ASSESSOR_QDRANT_API_KEY,
    ASSESSOR_QDRANT_TIMEOUT_MS: process.env.ASSESSOR_QDRANT_TIMEOUT_MS
  },
  fetcher?: typeof fetch
): QdrantRegulatoryIndex {
  const urlValue = environment.ASSESSOR_QDRANT_URL?.trim();
  const apiKey = environment.ASSESSOR_QDRANT_API_KEY?.trim();
  if (!urlValue || !apiKey) throw new Error('Regulatory RAG requires ASSESSOR_QDRANT_URL and ASSESSOR_QDRANT_API_KEY.');
  let baseUrl: URL;
  try { baseUrl = new URL(urlValue); } catch { throw new Error('ASSESSOR_QDRANT_URL must be an absolute URL.'); }
  if (!['http:', 'https:'].includes(baseUrl.protocol) || baseUrl.username || baseUrl.password || baseUrl.search || baseUrl.hash) {
    throw new Error('ASSESSOR_QDRANT_URL must be an HTTP(S) URL without credentials, query parameters, or a fragment.');
  }
  if (baseUrl.pathname !== '/' && baseUrl.pathname !== '') throw new Error('ASSESSOR_QDRANT_URL must not contain a path.');
  if (baseUrl.protocol !== 'https:' && !['localhost', '127.0.0.1', '[::1]', 'qdrant'].includes(baseUrl.hostname.toLowerCase())) {
    throw new Error('ASSESSOR_QDRANT_URL must use HTTPS unless it targets loopback or the internal qdrant service.');
  }
  return new QdrantRegulatoryIndex({
    baseUrl,
    apiKey,
    dimensions,
    embeddingModel,
    timeoutMs: boundedInteger(environment.ASSESSOR_QDRANT_TIMEOUT_MS, 15_000, 1_000, 120_000, 'ASSESSOR_QDRANT_TIMEOUT_MS'),
    fetcher
  });
}

function validateCollection(value: unknown, dimensions: number): void {
  const result = object(object(value)?.result);
  const params = object(object(object(result?.config)?.params));
  const vectors = object(params?.vectors);
  const dense = object(vectors?.[REGULATORY_DENSE_VECTOR]);
  const sparse = object(params?.sparse_vectors);
  if (dense?.size !== dimensions || String(dense.distance).toLowerCase() !== 'cosine' || !object(sparse?.[REGULATORY_SPARSE_VECTOR])) {
    throw new QdrantClientError('qdrant_collection_incompatible', `Qdrant collection ${REGULATORY_QDRANT_COLLECTION} is incompatible and must be rebuilt.`);
  }
}

function payloadSchemaFields(value: unknown): Set<string> {
  const payloadSchema = object(object(value)?.result)?.payload_schema;
  return new Set(Object.keys(object(payloadSchema) ?? {}));
}

function chunkPayload(chunk: RegulatoryChunk, embeddingModel: string): Record<string, unknown> {
  return {
    schema_version: chunk.schemaVersion,
    corpus_id: chunk.corpusId,
    corpus_version: chunk.corpusVersion,
    document_id: chunk.documentId,
    framework: chunk.framework,
    framework_version: chunk.frameworkVersion,
    control_id: chunk.controlId,
    chunk_id: chunk.chunkId,
    chunk_index: chunk.chunkIndex,
    language: chunk.language,
    title: chunk.title,
    text: chunk.text,
    text_sha256: chunk.textSha256,
    embedding_model: embeddingModel,
    ...(chunk.sourceUri ? { source_uri: chunk.sourceUri } : {})
  };
}

function parseCandidate(value: unknown): RegulatoryRetrievalCandidate {
  const point = object(value);
  const payload = object(point?.payload);
  const score = point?.score;
  if (!payload || typeof score !== 'number' || !Number.isFinite(score)) throw new QdrantClientError('qdrant_invalid_response', 'Qdrant returned an invalid point.');
  const stringField = (name: string): string => {
    const field = payload[name];
    if (typeof field !== 'string' || !field) throw new QdrantClientError('qdrant_invalid_response', `Qdrant point is missing ${name}.`);
    return field;
  };
  if (payload.schema_version !== 1 || !Number.isInteger(payload.chunk_index)) throw new QdrantClientError('qdrant_invalid_response', 'Qdrant point has an unsupported regulatory schema.');
  return {
    score,
    chunk: {
      schemaVersion: 1,
      corpusId: stringField('corpus_id'),
      corpusVersion: stringField('corpus_version'),
      documentId: stringField('document_id'),
      framework: stringField('framework'),
      frameworkVersion: stringField('framework_version'),
      controlId: stringField('control_id'),
      chunkId: stringField('chunk_id'),
      chunkIndex: Number(payload.chunk_index),
      language: stringField('language'),
      title: stringField('title'),
      text: stringField('text'),
      textSha256: stringField('text_sha256'),
      ...(typeof payload.source_uri === 'string' && payload.source_uri ? { sourceUri: payload.source_uri } : {})
    }
  };
}

function filterConditions(filters: RegulatoryRetrievalFilters): Array<Record<string, unknown>> {
  const values: Array<[string, string | undefined]> = [
    ['corpus_id', filters.corpusId],
    ['corpus_version', filters.corpusVersion],
    ['framework', filters.framework],
    ['framework_version', filters.frameworkVersion],
    ['control_id', filters.controlId],
    ['language', filters.language?.toLowerCase()]
  ];
  return values.filter((entry): entry is [string, string] => Boolean(entry[1])).map(([key, value]) => keywordCondition(key, value));
}

function keywordCondition(key: string, value: string): Record<string, unknown> {
  return { key, match: { value } };
}

function pointId(chunkId: string): string {
  const hash = createHash('sha256').update(chunkId, 'utf8').digest('hex').slice(0, 32);
  return `${hash.slice(0, 8)}-${hash.slice(8, 12)}-${hash.slice(12, 16)}-${hash.slice(16, 20)}-${hash.slice(20)}`;
}

function validateEmbedding(value: number[], dimensions: number): void {
  if (!Array.isArray(value) || value.length !== dimensions || value.some(entry => typeof entry !== 'number' || !Number.isFinite(entry))) {
    throw new Error(`Regulatory embeddings must contain ${dimensions} finite values.`);
  }
}

function boundedInteger(value: string | undefined, fallback: number, minimum: number, maximum: number, name: string): number {
  if (!value?.trim()) return fallback;
  const parsed = Number(value);
  if (!Number.isInteger(parsed) || parsed < minimum || parsed > maximum) throw new Error(`${name} must be an integer between ${minimum} and ${maximum}.`);
  return parsed;
}

function withTrailingSlash(url: URL): URL {
  const copy = new URL(url.href);
  if (!copy.pathname.endsWith('/')) copy.pathname += '/';
  return copy;
}

async function readBoundedText(response: Response, maximumBytes: number): Promise<string> {
  const declared = Number(response.headers.get('content-length') ?? 0);
  if (declared > maximumBytes) throw new QdrantClientError('qdrant_response_too_large', 'Qdrant response exceeded the size limit.');
  if (!response.body) return '';
  const reader = response.body.getReader();
  const chunks: Uint8Array[] = [];
  let size = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    size += value.byteLength;
    if (size > maximumBytes) {
      await reader.cancel();
      throw new QdrantClientError('qdrant_response_too_large', 'Qdrant response exceeded the size limit.');
    }
    chunks.push(value);
  }
  const merged = new Uint8Array(size);
  let offset = 0;
  for (const chunk of chunks) { merged.set(chunk, offset); offset += chunk.byteLength; }
  return new TextDecoder().decode(merged);
}

function object(value: unknown): Record<string, any> | undefined {
  return value && typeof value === 'object' && !Array.isArray(value) ? value as Record<string, any> : undefined;
}