DiAL/src/server/checker/runner/llm/observation.ts

import type { APICallError } from "ai";

import type {
  LlmCheckObservation,
  LlmHttpMetadata,
  LlmMode,
  LlmPersistedHttpMetadata,
  LlmPersistedObservation,
  LlmProvider,
  LlmStreamObservation,
  LlmUsageObservation,
} from "./types";

import { LLM_HEADERS_MAX, LLM_OUTPUT_PREVIEW_MAX } from "./types";

export function buildObservationFromApiCallError(
  error: APICallError,
  provider: LlmProvider,
  model: string,
  mode: LlmMode,
): LlmCheckObservation {
  const http: LlmHttpMetadata | null =
    error.statusCode !== undefined
      ? {
          headers: error.responseHeaders ?? {},
          status: error.statusCode,
          statusText: "",
        }
      : null;

  return {
    finishReason: null,
    http,
    mode,
    model,
    outputText: null,
    provider,
    rawFinishReason: null,
    stream: null,
    usage: null,
    warnings: [],
  };
}

export function buildObservationFromGenerateText(
  provider: LlmProvider,
  model: string,
  mode: LlmMode,
  result: {
    finishReason: string;
    rawFinishReason: null | string | undefined;
    text: string;
    usage: { inputTokens: number; outputTokens: number; totalTokens?: number | undefined };
    warnings?: string[];
  },
  http: LlmHttpMetadata | null,
): LlmCheckObservation {
  return {
    finishReason: result.finishReason,
    http,
    mode,
    model,
    outputText: result.text,
    provider,
    rawFinishReason: result.rawFinishReason ?? null,
    stream: null,
    usage: {
      inputTokens: result.usage.inputTokens ?? 0,
      outputTokens: result.usage.outputTokens ?? 0,
      totalTokens: result.usage.totalTokens ?? (result.usage.inputTokens ?? 0) + (result.usage.outputTokens ?? 0),
    },
    warnings: result.warnings ?? [],
  };
}

export async function buildObservationFromStreamText(
  provider: LlmProvider,
  model: string,
  mode: LlmMode,
  fullStream: AsyncIterable<unknown>,
  http: LlmHttpMetadata | null,
  startMs: number,
): Promise<LlmCheckObservation> {
  let outputText = "";
  let firstTokenMs: null | number = null;
  let completed = false;
  let finishReason: null | string = null;
  let rawFinishReason: null | string = null;
  let usage: LlmUsageObservation | null = null;
  const warnings: string[] = [];

  for await (const part of fullStream) {
    const p = part as Record<string, unknown>;
    const type = p["type"] as string;

    if (type === "text-delta") {
      const delta = p["textDelta"] as string;
      if (delta !== "") {
        firstTokenMs ??= Math.round(performance.now() - startMs);
        outputText += delta;
      }
    } else if (type === "finish") {
      completed = true;
      finishReason = (p["finishReason"] as string) ?? null;
      rawFinishReason = (p["rawFinishReason"] as string | undefined) ?? null;

      const totalUsage = p["totalUsage"] as
        | undefined
        | { inputTokens: number; outputTokens: number; totalTokens: number };
      const partUsage = p["usage"] as undefined | { inputTokens: number; outputTokens: number };

      usage = {
        inputTokens: totalUsage?.inputTokens ?? partUsage?.inputTokens ?? 0,
        outputTokens: totalUsage?.outputTokens ?? partUsage?.outputTokens ?? 0,
        totalTokens: totalUsage?.totalTokens ?? (partUsage?.inputTokens ?? 0) + (partUsage?.outputTokens ?? 0),
      };
    } else if (type === "error") {
      const err = p["error"] as Error | undefined;
      warnings.push(err?.message ?? "stream error");
    }
  }

  return {
    finishReason,
    http,
    mode,
    model,
    outputText: outputText || null,
    provider,
    rawFinishReason,
    stream: { completed, firstTokenMs } satisfies LlmStreamObservation,
    usage,
    warnings,
  };
}

export function toPersistedObservation(obs: LlmCheckObservation): LlmPersistedObservation {
  const outputText = obs.outputText;
  const outputPreview =
    outputText !== null
      ? outputText.length <= LLM_OUTPUT_PREVIEW_MAX
        ? outputText
        : outputText.slice(0, LLM_OUTPUT_PREVIEW_MAX)
      : null;
  const outputLength = outputText !== null ? outputText.length : null;

  const http: LlmPersistedHttpMetadata | null = obs.http
    ? {
        headers: truncateHeaders(obs.http.headers, LLM_HEADERS_MAX),
        status: obs.http.status,
        statusText: obs.http.statusText,
      }
    : null;

  return {
    finishReason: obs.finishReason,
    http,
    mode: obs.mode,
    model: obs.model,
    outputLength,
    outputPreview,
    provider: obs.provider,
    rawFinishReason: obs.rawFinishReason,
    stream: obs.stream,
    usage: obs.usage,
    warnings: obs.warnings,
  };
}

function truncateHeaders(headers: Record<string, string>, maxCount: number): Record<string, string> {
  const entries = Object.entries(headers);
  if (entries.length <= maxCount) return headers;
  return Object.fromEntries(entries.slice(0, maxCount));
}