Assessor_2.0/server/compliance/embedding-client.ts

export interface TextEmbeddingClient {
  readonly model: string;
  readonly dimensions: number;
  embed(texts: string[]): Promise<number[][]>;
}

export interface EmbeddingEnvironment {
  ASSESSOR_RAG_EMBEDDING_BASE_URL?: string;
  ASSESSOR_RAG_EMBEDDING_MODEL?: string;
  ASSESSOR_RAG_EMBEDDING_API_KEY?: string;
  ASSESSOR_RAG_EMBEDDING_DIMENSIONS?: string;
  ASSESSOR_RAG_EMBEDDING_REQUEST_DIMENSIONS?: string;
  ASSESSOR_RAG_EMBEDDING_TIMEOUT_MS?: string;
}

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

interface OpenAiCompatibleEmbeddingClientOptions {
  endpoint: URL;
  model: string;
  apiKey: string;
  dimensions: number;
  timeoutMs: number;
  requestDimensions?: boolean;
  fetcher?: typeof fetch;
}

export class OpenAiCompatibleEmbeddingClient implements TextEmbeddingClient {
  readonly model: string;
  readonly dimensions: number;
  private readonly endpoint: URL;
  private readonly apiKey: string;
  private readonly timeoutMs: number;
  private readonly requestDimensions: boolean;
  private readonly fetcher: typeof fetch;

  constructor(options: OpenAiCompatibleEmbeddingClientOptions) {
    this.endpoint = new URL(options.endpoint.href);
    this.model = options.model;
    this.apiKey = options.apiKey;
    this.dimensions = options.dimensions;
    this.timeoutMs = options.timeoutMs;
    this.requestDimensions = options.requestDimensions ?? false;
    this.fetcher = options.fetcher ?? globalThis.fetch.bind(globalThis);
  }

  async embed(texts: string[]): Promise<number[][]> {
    if (!texts.length || texts.some(text => typeof text !== 'string' || !text.trim())) {
      throw new EmbeddingClientError('invalid_embedding_input', 'At least one non-empty embedding input is required.');
    }
    let response: Response;
    try {
      response = await this.fetcher(this.endpoint, {
        method: 'POST',
        redirect: 'error',
        signal: AbortSignal.timeout(this.timeoutMs),
        headers: {
          Accept: 'application/json',
          'Content-Type': 'application/json',
          Authorization: `Bearer ${this.apiKey}`
        },
        body: JSON.stringify({
          model: this.model,
          input: texts,
          encoding_format: 'float',
          ...(this.requestDimensions ? { dimensions: this.dimensions } : {})
        })
      });
    } catch {
      throw new EmbeddingClientError('embedding_unavailable', 'The configured embedding service could not be reached.');
    }
    if (!response.ok) throw new EmbeddingClientError('embedding_http_error', `The embedding service returned HTTP ${response.status}.`);
    const body = await readBoundedJson(response, 32 * 1024 * 1024);
    if (!isObject(body) || !Array.isArray(body.data) || body.data.length !== texts.length) {
      throw new EmbeddingClientError('embedding_invalid_response', 'The embedding service returned an invalid result count.');
    }
    const ordered = body.data.map((entry, responseIndex) => ({ entry, responseIndex })).sort((left, right) => {
      const leftIndex = isObject(left.entry) && Number.isInteger(left.entry.index) ? Number(left.entry.index) : left.responseIndex;
      const rightIndex = isObject(right.entry) && Number.isInteger(right.entry.index) ? Number(right.entry.index) : right.responseIndex;
      return leftIndex - rightIndex;
    });
    return ordered.map(({ entry }) => {
      const vector = isObject(entry) && Array.isArray(entry.embedding) ? entry.embedding : undefined;
      if (!vector || vector.length !== this.dimensions || vector.some(value => typeof value !== 'number' || !Number.isFinite(value))) {
        throw new EmbeddingClientError('embedding_invalid_dimensions', `The embedding service must return ${this.dimensions} finite dimensions per input.`);
      }
      return vector as number[];
    });
  }
}

export function createEmbeddingClientFromEnvironment(
  environment: EmbeddingEnvironment = {
    ASSESSOR_RAG_EMBEDDING_BASE_URL: process.env.ASSESSOR_RAG_EMBEDDING_BASE_URL,
    ASSESSOR_RAG_EMBEDDING_MODEL: process.env.ASSESSOR_RAG_EMBEDDING_MODEL,
    ASSESSOR_RAG_EMBEDDING_API_KEY: process.env.ASSESSOR_RAG_EMBEDDING_API_KEY,
    ASSESSOR_RAG_EMBEDDING_DIMENSIONS: process.env.ASSESSOR_RAG_EMBEDDING_DIMENSIONS,
    ASSESSOR_RAG_EMBEDDING_REQUEST_DIMENSIONS: process.env.ASSESSOR_RAG_EMBEDDING_REQUEST_DIMENSIONS,
    ASSESSOR_RAG_EMBEDDING_TIMEOUT_MS: process.env.ASSESSOR_RAG_EMBEDDING_TIMEOUT_MS
  },
  fetcher?: typeof fetch
): OpenAiCompatibleEmbeddingClient {
  const baseUrl = environment.ASSESSOR_RAG_EMBEDDING_BASE_URL?.trim();
  const model = environment.ASSESSOR_RAG_EMBEDDING_MODEL?.trim();
  const apiKey = environment.ASSESSOR_RAG_EMBEDDING_API_KEY?.trim();
  if (!baseUrl || !model || !apiKey) {
    throw new Error('Regulatory RAG requires ASSESSOR_RAG_EMBEDDING_BASE_URL, ASSESSOR_RAG_EMBEDDING_MODEL, and ASSESSOR_RAG_EMBEDDING_API_KEY.');
  }
  return new OpenAiCompatibleEmbeddingClient({
    endpoint: embeddingEndpoint(baseUrl),
    model,
    apiKey,
    dimensions: boundedInteger(environment.ASSESSOR_RAG_EMBEDDING_DIMENSIONS, undefined, 64, 16_384, 'ASSESSOR_RAG_EMBEDDING_DIMENSIONS'),
    requestDimensions: booleanFlag(environment.ASSESSOR_RAG_EMBEDDING_REQUEST_DIMENSIONS, 'ASSESSOR_RAG_EMBEDDING_REQUEST_DIMENSIONS'),
    timeoutMs: boundedInteger(environment.ASSESSOR_RAG_EMBEDDING_TIMEOUT_MS, 60_000, 1_000, 300_000, 'ASSESSOR_RAG_EMBEDDING_TIMEOUT_MS'),
    fetcher
  });
}

export function embeddingEndpoint(baseUrl: string): URL {
  let url: URL;
  try { url = new URL(baseUrl); } catch { throw new Error('ASSESSOR_RAG_EMBEDDING_BASE_URL must be an absolute HTTP(S) URL.'); }
  const loopback = ['localhost', '127.0.0.1', '[::1]'].includes(url.hostname.toLowerCase());
  if (url.protocol !== 'https:' && !(url.protocol === 'http:' && loopback)) {
    throw new Error('ASSESSOR_RAG_EMBEDDING_BASE_URL must use HTTPS unless it targets loopback.');
  }
  if (url.username || url.password || url.search || url.hash) {
    throw new Error('ASSESSOR_RAG_EMBEDDING_BASE_URL must not contain credentials, query parameters, or a fragment.');
  }
  const path = url.pathname.replace(/\/+$/, '');
  url.pathname = path.endsWith('/embeddings') ? path : path ? `${path}/embeddings` : '/v1/embeddings';
  return url;
}

function booleanFlag(value: string | undefined, name: string): boolean {
  if (!value?.trim()) return false;
  if (value === 'true') return true;
  if (value === 'false') return false;
  throw new Error(`${name} must be true or false.`);
}

function boundedInteger(value: string | undefined, fallback: number | undefined, minimum: number, maximum: number, name: string): number {
  if (!value?.trim()) {
    if (fallback !== undefined) return fallback;
    throw new Error(`${name} is required.`);
  }
  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;
}

async function readBoundedJson(response: Response, maximumBytes: number): Promise<unknown> {
  const declared = Number(response.headers.get('content-length') ?? 0);
  if (declared > maximumBytes) throw new EmbeddingClientError('embedding_response_too_large', 'The embedding response exceeded the size limit.');
  if (!response.body) throw new EmbeddingClientError('embedding_invalid_response', 'The embedding service returned an empty response.');
  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 EmbeddingClientError('embedding_response_too_large', 'The embedding 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; }
  const text = new TextDecoder().decode(merged);
  try { return JSON.parse(text); } catch { throw new EmbeddingClientError('embedding_invalid_response', 'The embedding service returned invalid JSON.'); }
}

function isObject(value: unknown): value is Record<string, any> {
  return Boolean(value) && typeof value === 'object' && !Array.isArray(value);
}