homarr/apps/nextjs/src/components/assistant/assistant-message-metadata.ts

import type { UIMessage } from "ai";

export type AssistantUsage = {
  inputTokens?: number;
  outputTokens?: number;
  totalTokens?: number;
  reasoningTokens?: number;
  cachedInputTokens?: number;
  cacheWriteTokens?: number;
};

export type AssistantRequestStep = {
  index: number;
  durationMs: number;
  modelDurationMs: number;
  toolDurationMs: number;
  timeToFirstOutputMs?: number;
  outputTokensPerSecond?: number;
  inputTokens?: number;
  outputTokens?: number;
  cost?: number;
  upstreamCost?: number;
  generationId?: string;
  routedProvider?: string;
};

export type AssistantRequestTelemetry = {
  requestId: string;
  provider: string;
  modelId: string;
  startedAt: string;
  completedAt?: string;
  durationMs?: number;
  timeToFirstOutputMs?: number;
  outputTokensPerSecond?: number;
  contextLength?: number;
  contextUsed?: number;
  contextUtilization?: number;
  cost?: number;
  upstreamCost?: number;
  costType?: "reported" | "estimated";
  finishReason?: string;
  steps: AssistantRequestStep[];
};

export type AssistantMessageMetadata = {
  usage?: AssistantUsage;
  custom: {
    telemetry: AssistantRequestTelemetry;
  };
};

export type AssistantUIMessage = UIMessage<AssistantMessageMetadata>;

const getFiniteNonNegativeNumber = (value: unknown) =>
  typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;

export const getAssistantUsage = (metadata: unknown): AssistantUsage | null => {
  if (!metadata || typeof metadata !== "object" || Array.isArray(metadata) || !("usage" in metadata)) return null;
  const usage = metadata.usage;
  if (!usage || typeof usage !== "object" || Array.isArray(usage)) return null;
  const value = usage as Record<string, unknown>;
  return {
    ...(getFiniteNonNegativeNumber(value.inputTokens) !== undefined
      ? { inputTokens: getFiniteNonNegativeNumber(value.inputTokens) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.outputTokens) !== undefined
      ? { outputTokens: getFiniteNonNegativeNumber(value.outputTokens) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.totalTokens) !== undefined
      ? { totalTokens: getFiniteNonNegativeNumber(value.totalTokens) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.reasoningTokens) !== undefined
      ? { reasoningTokens: getFiniteNonNegativeNumber(value.reasoningTokens) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.cachedInputTokens) !== undefined
      ? { cachedInputTokens: getFiniteNonNegativeNumber(value.cachedInputTokens) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.cacheWriteTokens) !== undefined
      ? { cacheWriteTokens: getFiniteNonNegativeNumber(value.cacheWriteTokens) }
      : {}),
  };
};

export const getAssistantTelemetry = (metadata: unknown): AssistantRequestTelemetry | null => {
  if (!metadata || typeof metadata !== "object" || Array.isArray(metadata)) return null;
  const custom = (metadata as { custom?: unknown }).custom;
  if (!custom || typeof custom !== "object" || Array.isArray(custom)) return null;
  const telemetry = (custom as { telemetry?: unknown }).telemetry;
  if (!telemetry || typeof telemetry !== "object" || Array.isArray(telemetry)) return null;
  const value = telemetry as Partial<AssistantRequestTelemetry>;
  if (
    typeof value.requestId !== "string" ||
    typeof value.provider !== "string" ||
    typeof value.modelId !== "string" ||
    typeof value.startedAt !== "string"
  ) {
    return null;
  }
  const steps = Array.isArray(value.steps)
    ? value.steps.flatMap((step): AssistantRequestStep[] => {
        if (!step || typeof step !== "object" || Array.isArray(step)) return [];
        const candidate = step as Record<string, unknown>;
        const index = getFiniteNonNegativeNumber(candidate.index);
        const durationMs = getFiniteNonNegativeNumber(candidate.durationMs);
        const modelDurationMs = getFiniteNonNegativeNumber(candidate.modelDurationMs);
        const toolDurationMs = getFiniteNonNegativeNumber(candidate.toolDurationMs);
        if (
          index === undefined ||
          durationMs === undefined ||
          modelDurationMs === undefined ||
          toolDurationMs === undefined
        )
          return [];
        return [
          {
            index,
            durationMs,
            modelDurationMs,
            toolDurationMs,
            ...(getFiniteNonNegativeNumber(candidate.timeToFirstOutputMs) !== undefined
              ? { timeToFirstOutputMs: getFiniteNonNegativeNumber(candidate.timeToFirstOutputMs) }
              : {}),
            ...(getFiniteNonNegativeNumber(candidate.outputTokensPerSecond) !== undefined
              ? { outputTokensPerSecond: getFiniteNonNegativeNumber(candidate.outputTokensPerSecond) }
              : {}),
            ...(getFiniteNonNegativeNumber(candidate.inputTokens) !== undefined
              ? { inputTokens: getFiniteNonNegativeNumber(candidate.inputTokens) }
              : {}),
            ...(getFiniteNonNegativeNumber(candidate.outputTokens) !== undefined
              ? { outputTokens: getFiniteNonNegativeNumber(candidate.outputTokens) }
              : {}),
            ...(getFiniteNonNegativeNumber(candidate.cost) !== undefined
              ? { cost: getFiniteNonNegativeNumber(candidate.cost) }
              : {}),
            ...(getFiniteNonNegativeNumber(candidate.upstreamCost) !== undefined
              ? { upstreamCost: getFiniteNonNegativeNumber(candidate.upstreamCost) }
              : {}),
            ...(typeof candidate.generationId === "string" ? { generationId: candidate.generationId } : {}),
            ...(typeof candidate.routedProvider === "string" ? { routedProvider: candidate.routedProvider } : {}),
          },
        ];
      })
    : [];
  return {
    requestId: value.requestId,
    provider: value.provider,
    modelId: value.modelId,
    startedAt: value.startedAt,
    steps,
    ...(typeof value.completedAt === "string" ? { completedAt: value.completedAt } : {}),
    ...(getFiniteNonNegativeNumber(value.durationMs) !== undefined
      ? { durationMs: getFiniteNonNegativeNumber(value.durationMs) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.timeToFirstOutputMs) !== undefined
      ? { timeToFirstOutputMs: getFiniteNonNegativeNumber(value.timeToFirstOutputMs) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.outputTokensPerSecond) !== undefined
      ? { outputTokensPerSecond: getFiniteNonNegativeNumber(value.outputTokensPerSecond) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.contextLength) !== undefined
      ? { contextLength: getFiniteNonNegativeNumber(value.contextLength) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.contextUsed) !== undefined
      ? { contextUsed: getFiniteNonNegativeNumber(value.contextUsed) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.contextUtilization) !== undefined
      ? { contextUtilization: Math.min(getFiniteNonNegativeNumber(value.contextUtilization) ?? 0, 1) }
      : {}),
    ...(getFiniteNonNegativeNumber(value.cost) !== undefined ? { cost: getFiniteNonNegativeNumber(value.cost) } : {}),
    ...(getFiniteNonNegativeNumber(value.upstreamCost) !== undefined
      ? { upstreamCost: getFiniteNonNegativeNumber(value.upstreamCost) }
      : {}),
    ...(value.costType === "reported" || value.costType === "estimated" ? { costType: value.costType } : {}),
    ...(typeof value.finishReason === "string" ? { finishReason: value.finishReason } : {}),
  };
};