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

import type { ExpectResult } from "../../expect/types";
import type { LlmCheckObservation, LlmExpectConfig, LlmUsageExpect } from "./types";

import { checkContentRules } from "../../expect/content";
import { mismatchFailure } from "../../expect/failure";
import { checkValueMatcher } from "../../expect/matcher";
import { checkHeaders, checkStatus } from "../http/expect";

export function checkStreamExpect(observation: LlmCheckObservation, expect: LlmExpectConfig): ExpectResult {
  if (!observation.stream || !expect.stream) return { failure: null, matched: true };

  const expectedCompleted = expect.stream.completed ?? true;
  if (observation.stream.completed !== expectedCompleted) {
    return {
      failure: mismatchFailure(
        "stream",
        "stream.completed",
        expectedCompleted,
        observation.stream.completed,
        "stream.completed mismatch",
      ),
      matched: false,
    };
  }

  if (expect.stream.firstTokenMs && observation.stream.firstTokenMs !== null) {
    return checkValueMatcher(observation.stream.firstTokenMs, expect.stream.firstTokenMs, {
      message: "stream.firstTokenMs mismatch",
      path: "stream.firstTokenMs",
      phase: "stream",
    });
  } else if (expect.stream.firstTokenMs && observation.stream.firstTokenMs === null) {
    return {
      failure: mismatchFailure(
        "stream",
        "stream.firstTokenMs",
        expect.stream.firstTokenMs,
        null,
        "stream.firstTokenMs missing",
      ),
      matched: false,
    };
  }

  return { failure: null, matched: true };
}

export function runExpects(observation: LlmCheckObservation, expect: LlmExpectConfig | undefined): ExpectResult {
  if (!expect) {
    const defaultStatus = checkStatus(observation.http?.status ?? 0, [200]);
    if (!defaultStatus.matched) return defaultStatus;
    return { failure: null, matched: true };
  }

  const http = observation.http;

  const statusResult = checkStatus(http?.status ?? 0, expect.status ?? [200]);
  if (!statusResult.matched) return statusResult;

  if (http && expect.headers) {
    const headersResult = checkHeaders(http.headers, expect.headers);
    if (!headersResult.matched) return headersResult;
  }

  if (observation.stream !== null) {
    const streamResult = checkStreamExpect(observation, expect);
    if (!streamResult.matched) return streamResult;
  }

  const outputResult = checkContentRules(observation.outputText, expect.output, { path: "output", phase: "output" });
  if (!outputResult.matched) return outputResult;

  if (expect.finishReason !== undefined) {
    const result = checkValueMatcher(observation.finishReason, expect.finishReason, {
      message: "finishReason mismatch",
      path: "finishReason",
      phase: "finishReason",
    });
    if (!result.matched) return result;
  }

  if (expect.rawFinishReason !== undefined) {
    const result = checkValueMatcher(observation.rawFinishReason, expect.rawFinishReason, {
      message: "rawFinishReason mismatch",
      path: "rawFinishReason",
      phase: "rawFinishReason",
    });
    if (!result.matched) return result;
  }

  if (expect.usage && observation.usage) {
    const usageResult = checkUsageExpect(observation.usage, expect.usage);
    if (!usageResult.matched) return usageResult;
  }

  return { failure: null, matched: true };
}

function checkUsageExpect(
  usage: { inputTokens: number; outputTokens: number; totalTokens: number },
  expectUsage: LlmUsageExpect,
): ExpectResult {
  if (expectUsage.inputTokens !== undefined) {
    const result = checkValueMatcher(usage.inputTokens, expectUsage.inputTokens, {
      message: "usage.inputTokens mismatch",
      path: "usage.inputTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  if (expectUsage.outputTokens !== undefined) {
    const result = checkValueMatcher(usage.outputTokens, expectUsage.outputTokens, {
      message: "usage.outputTokens mismatch",
      path: "usage.outputTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  if (expectUsage.totalTokens !== undefined) {
    const result = checkValueMatcher(usage.totalTokens, expectUsage.totalTokens, {
      message: "usage.totalTokens mismatch",
      path: "usage.totalTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  return { failure: null, matched: true };
}