Assessor_2.0/server/regmap/regulation-retrieval-service.ts

// Regulation retrieval service: for a finding definition, embed a definition-level
// query, cosine-search the Qdrant catalog, group hits by atomic unit, and return
// the best candidate units. No LLM selection, no persistence.
//
// Never logs query text, catalog text or vectors — only counts / codes.

import {
  RegulationMappingUnavailableError,
  type RegulationCapabilityReport,
  type RegulationMapping,
  type RegulationMappingConfig
} from './regulation-mapping';
import type { RegulationDefinitionInput, RegulationRetrievalResult, RegulationCandidateUnit } from './regulation-retrieval-types';
import { createEmbeddingClient, type EmbeddingClient } from './embedding-client';
import { QdrantClient, type QdrantSearchHit } from './qdrant-client';

export interface RegulationRetrievalDependencies {
  embeddingClient?: EmbeddingClient;
  qdrantClient?: Pick<QdrantClient, 'collectionInfo' | 'search'>;
}

function stringField(payload: Record<string, unknown>, key: string): string | undefined {
  const value = payload[key];
  return typeof value === 'string' && value.trim() ? value : undefined;
}

export class RegulationRetrievalService implements RegulationMapping {
  private readonly embedding: EmbeddingClient;
  private readonly qdrant: Pick<QdrantClient, 'collectionInfo' | 'search'>;
  private report: RegulationCapabilityReport;

  constructor(
    private readonly config: RegulationMappingConfig,
    dependencies: RegulationRetrievalDependencies = {}
  ) {
    this.embedding = dependencies.embeddingClient ?? createEmbeddingClient({
      baseUrl: config.embeddingBaseUrl,
      model: config.embeddingModel,
      apiKey: config.embeddingApiKey,
      dimensions: config.embeddingDimensions,
      queryPrefix: config.queryPrefix,
      passagePrefix: config.passagePrefix,
      requestTimeoutMs: config.requestTimeoutMs
    });
    this.qdrant = dependencies.qdrantClient ?? new QdrantClient({
      url: config.qdrantUrl,
      apiKey: config.qdrantApiKey,
      collection: config.collection,
      requestTimeoutMs: config.requestTimeoutMs
    });
    // Not yet validated against the live backend.
    this.report = this.reportWith('misconfigured', 'not_validated');
  }

  private reportWith(capability: RegulationCapabilityReport['capability'], detailCode?: string): RegulationCapabilityReport {
    return {
      capability,
      catalogVersion: this.config.catalogVersion,
      embeddingModel: this.config.embeddingModel,
      collection: this.config.collection,
      detailCode
    };
  }

  cachedCapability(): RegulationCapabilityReport {
    return this.report;
  }

  async refreshCapability(): Promise<RegulationCapabilityReport> {
    try {
      const info = await this.qdrant.collectionInfo();
      if (!info.exists) {
        this.report = this.reportWith('misconfigured', 'collection_missing');
      } else if (info.vectorSize !== undefined && info.vectorSize !== this.config.embeddingDimensions) {
        this.report = this.reportWith('misconfigured', 'dimension_mismatch');
      } else if (info.distance && info.distance !== 'Cosine') {
        this.report = this.reportWith('misconfigured', 'distance_mismatch');
      } else {
        // Verify the embedding endpoint/model end to end with a synthetic probe.
        await this.embedding.embed(['probe'], 'query');
        this.report = this.reportWith('ready');
      }
    } catch (error) {
      this.report = this.reportWith('misconfigured', safeErrorCode(error));
    }
    return this.report;
  }

  async retrieve(input: RegulationDefinitionInput): Promise<RegulationRetrievalResult> {
    if (this.report.capability !== 'ready') {
      const refreshed = await this.refreshCapability();
      if (refreshed.capability !== 'ready') throw new RegulationMappingUnavailableError(refreshed);
    }

    const query = this.buildQuery(input);
    let hits: QdrantSearchHit[];
    try {
      const [vector] = await this.embedding.embed([query], 'query');
      hits = await this.qdrant.search(vector, this.config.retrievalChunkLimit);
    } catch (error) {
      this.report = this.reportWith('misconfigured', safeErrorCode(error));
      throw new RegulationMappingUnavailableError(this.report);
    }

    return {
      catalogVersion: this.config.catalogVersion,
      embeddingModel: this.config.embeddingModel,
      candidates: groupHitsByUnit(hits, this.config.candidateUnitLimit)
    };
  }

  private buildQuery(input: RegulationDefinitionInput): string {
    const parts = [
      input.title,
      input.description,
      input.recommendation ? `Remediation: ${input.recommendation}` : undefined,
      input.category ? `Category: ${input.category}` : undefined
    ].filter((part): part is string => typeof part === 'string' && part.trim().length > 0);
    return parts.join('. ').slice(0, this.config.maximumQueryCharacters);
  }
}

function safeErrorCode(error: unknown): string {
  const code = (error as { code?: unknown })?.code;
  return typeof code === 'string' && code ? code : 'validation_failed';
}

/** Group chunk hits back to their atomic unit, scoring each unit by its best
 * chunk, then return the top-N units in deterministic order. */
export function groupHitsByUnit(hits: QdrantSearchHit[], limit: number): RegulationCandidateUnit[] {
  const byUnit = new Map<string, RegulationCandidateUnit>();
  for (const hit of hits) {
    const unitId = stringField(hit.payload, 'unit_id');
    if (!unitId) continue;
    const existing = byUnit.get(unitId);
    if (existing && existing.score >= hit.score) continue;
    byUnit.set(unitId, {
      unitId,
      regulationId: stringField(hit.payload, 'regulation_id'),
      regulationTitle: stringField(hit.payload, 'regulation_title'),
      controlId: stringField(hit.payload, 'control_id'),
      controlTitle: stringField(hit.payload, 'control_title'),
      title: stringField(hit.payload, 'control_title') ?? stringField(hit.payload, 'title'),
      score: hit.score,
      bestChunkId: stringField(hit.payload, 'chunk_id')
    });
  }
  return [...byUnit.values()]
    .sort((a, b) => (b.score - a.score) || a.unitId.localeCompare(b.unitId))
    .slice(0, limit);
}