Assessor_2.0/server/regulatory/regulatory-finding-retrieval.ts

import type {
  RegulatoryFindingQuery,
  RegulatoryRetrievalCandidate,
  RegulatoryRetrievalFilters,
  RegulatoryRetriever
} from '../../src/domain/regulatory-rag';
import type { UnifiedResultProjection } from '../../src/domain/unified-result';
import { buildRegulatoryFindingQuery } from '../../src/profiles/public-search/regulatory-query';

const RECIPROCAL_RANK_FUSION_K = 60;
const FINDING_QUERY_SEARCH_LIMIT = 64;
const SECTION_SEARCH_LIMIT = 32;
const FULL_CONTEXT_CONTROL_LIMIT = 100;
const CANDIDATE_SECTION_KINDS = ['control', 'purpose', 'guidance'] as const;

export interface RegulatoryFindingRetrievalResult {
  sourceResultId: string;
  status: 'candidates' | 'insufficient-public-context';
  preparedQuery?: RegulatoryFindingQuery;
  candidates: RegulatoryRetrievalCandidate[];
}

export class RegulatoryFindingRetrievalService {
  constructor(private readonly retriever: RegulatoryRetriever) {}

  async retrieve(
    finding: UnifiedResultProjection,
    limit: number,
    filters?: RegulatoryRetrievalFilters
  ): Promise<RegulatoryFindingRetrievalResult> {
    const preparedQuery = buildRegulatoryFindingQuery(finding);
    if (!preparedQuery) {
      return { sourceResultId: finding.sourceResultId, status: 'insufficient-public-context', candidates: [] };
    }
    const query = preparedQuery.retrievalQuery ?? preparedQuery.query;
    if (filters?.sectionKind) {
      const candidates = await this.retriever.retrieve({ query, limit, filters });
      return { sourceResultId: finding.sourceResultId, status: 'candidates', preparedQuery, candidates };
    }
    const selectionQueries = [...new Set([
      preparedQuery.problemQuery ?? query,
      preparedQuery.remediationQuery
    ].filter((value): value is string => Boolean(value)))];
    const selectionRankings = await Promise.all(selectionQueries.map(selectionQuery => this.retriever.retrieve({
      query: selectionQuery,
      limit: Math.max(limit, FINDING_QUERY_SEARCH_LIMIT),
      filters
    })));
    const sectionRankings = await Promise.all(CANDIDATE_SECTION_KINDS.map(sectionKind => this.retriever.retrieve({
      query,
      limit: Math.max(limit, SECTION_SEARCH_LIMIT),
      filters: { ...filters, sectionKind }
    })));
    const candidates = mergeFindingQueryRankings(selectionRankings, sectionRankings, limit);
    return { sourceResultId: finding.sourceResultId, status: 'candidates', preparedQuery, candidates };
  }
}

/**
 * Collapses the full-context section searches into one control ranking, then
 * fuses that family with the independently phrased problem and remediation
 * rankings. A control can therefore receive at most one final vote from the
 * section family while its section coverage still determines that family's
 * internal order.
 */
export function mergeFindingQueryRankings(
  selectionRankings: RegulatoryRetrievalCandidate[][],
  evidenceRankings: RegulatoryRetrievalCandidate[][],
  limit: number
): RegulatoryRetrievalCandidate[] {
  if (!Number.isInteger(limit) || limit < 1 || limit > 100) throw new Error('Regulatory retrieval limit must be between 1 and 100.');
  const fullContextRanking = reciprocalRankFuseByControl(evidenceRankings, FULL_CONTEXT_CONTROL_LIMIT);
  const selected = uniqueControlIds(reciprocalRankFuseByControl(
    [...selectionRankings, fullContextRanking],
    limit
  ));
  const allEvidenceRankings = [...evidenceRankings, ...selectionRankings];
  return selected.flatMap((controlId, index) => {
    const seenChunks = new Set<string>();
    return allEvidenceRankings.flatMap(ranking => {
      const hit = ranking.find(candidate => candidate.chunk.controlId === controlId && !seenChunks.has(candidate.chunk.chunkId));
      if (!hit) return [];
      seenChunks.add(hit.chunk.chunkId);
      return [{ ...hit, score: 1 / (RECIPROCAL_RANK_FUSION_K + index + 1) }];
    });
  });
}

/** Fuses independent rankings at control level, allowing one vote per control and ranking. */
export function reciprocalRankFuseByControl(
  rankings: RegulatoryRetrievalCandidate[][],
  limit: number
): RegulatoryRetrievalCandidate[] {
  if (!Number.isInteger(limit) || limit < 1 || limit > 100) throw new Error('Regulatory retrieval limit must be between 1 and 100.');
  const controls = new Map<string, {
    score: number;
    bestRank: number;
    hits: Map<number, RegulatoryRetrievalCandidate>;
  }>();
  rankings.forEach((ranking, rankingIndex) => {
    const seen = new Set<string>();
    for (const hit of ranking) {
      const controlId = hit.chunk.controlId;
      if (seen.has(controlId)) continue;
      seen.add(controlId);
      const controlRank = seen.size;
      const entry = controls.get(controlId) ?? { score: 0, bestRank: Number.MAX_SAFE_INTEGER, hits: new Map() };
      entry.score += 1 / (RECIPROCAL_RANK_FUSION_K + controlRank);
      entry.bestRank = Math.min(entry.bestRank, controlRank);
      entry.hits.set(rankingIndex, hit);
      controls.set(controlId, entry);
    }
  });
  return [...controls.entries()]
    .sort(([leftId, left], [rightId, right]) => right.score - left.score
      || left.bestRank - right.bestRank
      || leftId.localeCompare(rightId, 'en', { numeric: true }))
    .slice(0, limit)
    .flatMap(([, entry]) => [...entry.hits.entries()]
      .sort(([left], [right]) => left - right)
      .map(([, hit]) => ({ ...hit, score: entry.score })));
}

function uniqueControlIds(candidates: RegulatoryRetrievalCandidate[]): string[] {
  const seen = new Set<string>();
  return candidates.flatMap(candidate => {
    const controlId = candidate.chunk.controlId;
    if (seen.has(controlId)) return [];
    seen.add(controlId);
    return [controlId];
  });
}