pitchfast/convex/usageEvents.ts

import type { Doc, Id } from "./_generated/dataModel";
import { internalMutation, query } from "./_generated/server";
import type { QueryCtx } from "./_generated/server";
import {
  normalizeListLimit,
  USAGE_EVENT_OPERATIONS,
  USAGE_EVENT_PROVIDERS,
} from "./domain";
import { v } from "convex/values";

const usageEventProvider = v.union(
  ...USAGE_EVENT_PROVIDERS.map((provider) => v.literal(provider)),
);
const usageEventOperation = v.union(
  ...USAGE_EVENT_OPERATIONS.map((operation) => v.literal(operation)),
);
const usageEventTokens = v.object({
  inputTokens: v.optional(v.number()),
  outputTokens: v.optional(v.number()),
  promptTokens: v.optional(v.number()),
  completionTokens: v.optional(v.number()),
  totalTokens: v.optional(v.number()),
  cacheReadTokens: v.optional(v.number()),
});
const usageEventCallCounts = v.object({
  requests: v.optional(v.number()),
  pages: v.optional(v.number()),
  screenshots: v.optional(v.number()),
  lookups: v.optional(v.number()),
});
const usageEventDoc = v.object({
  _id: v.id("usageEvents"),
  _creationTime: v.number(),
  provider: usageEventProvider,
  operation: usageEventOperation,
  runId: v.optional(v.id("agentRuns")),
  leadId: v.optional(v.id("leads")),
  auditId: v.optional(v.id("audits")),
  estimatedCostUsd: v.number(),
  tokens: v.optional(usageEventTokens),
  callCounts: v.optional(usageEventCallCounts),
  createdAt: v.number(),
});

type UsageEventTokens = {
  inputTokens?: number;
  outputTokens?: number;
  promptTokens?: number;
  completionTokens?: number;
  totalTokens?: number;
  cacheReadTokens?: number;
};

type UsageEventCallCounts = {
  requests?: number;
  pages?: number;
  screenshots?: number;
  lookups?: number;
};

type UsageEventNumberArgs = {
  estimatedCostUsd: number;
  tokens?: UsageEventTokens;
  callCounts?: UsageEventCallCounts;
};

const requireOperator = async (ctx: QueryCtx) => {
  const identity = await ctx.auth.getUserIdentity();
  if (!identity) {
    throw new Error("Nicht autorisiert.");
  }
};

function assertFiniteNonNegativeNumber(value: number, fieldName: string) {
  if (!Number.isFinite(value) || value < 0) {
    throw new Error(`${fieldName} must be a finite non-negative number.`);
  }
}

function assertFiniteNonNegativeInteger(
  value: number | undefined,
  fieldName: string,
) {
  if (value === undefined) {
    return;
  }

  if (!Number.isFinite(value) || value < 0 || !Number.isInteger(value)) {
    throw new Error(`${fieldName} must be a finite non-negative integer.`);
  }
}

function assertValidUsageEventNumbers(args: UsageEventNumberArgs) {
  assertFiniteNonNegativeNumber(args.estimatedCostUsd, "estimatedCostUsd");
  assertFiniteNonNegativeInteger(args.tokens?.inputTokens, "tokens.inputTokens");
  assertFiniteNonNegativeInteger(args.tokens?.outputTokens, "tokens.outputTokens");
  assertFiniteNonNegativeInteger(args.tokens?.promptTokens, "tokens.promptTokens");
  assertFiniteNonNegativeInteger(args.tokens?.completionTokens, "tokens.completionTokens");
  assertFiniteNonNegativeInteger(args.tokens?.totalTokens, "tokens.totalTokens");
  assertFiniteNonNegativeInteger(args.tokens?.cacheReadTokens, "tokens.cacheReadTokens");
  assertFiniteNonNegativeInteger(args.callCounts?.requests, "callCounts.requests");
  assertFiniteNonNegativeInteger(args.callCounts?.pages, "callCounts.pages");
  assertFiniteNonNegativeInteger(args.callCounts?.screenshots, "callCounts.screenshots");
  assertFiniteNonNegativeInteger(args.callCounts?.lookups, "callCounts.lookups");
}

export const recordUsageEvent = internalMutation({
  args: {
    provider: usageEventProvider,
    operation: usageEventOperation,
    runId: v.optional(v.id("agentRuns")),
    leadId: v.optional(v.id("leads")),
    auditId: v.optional(v.id("audits")),
    estimatedCostUsd: v.number(),
    tokens: v.optional(usageEventTokens),
    callCounts: v.optional(usageEventCallCounts),
    createdAt: v.optional(v.number()),
  },
  returns: v.id("usageEvents"),
  handler: async (ctx, args): Promise<Id<"usageEvents">> => {
    assertValidUsageEventNumbers(args);

    const now = args.createdAt ?? Date.now();

    return await ctx.db.insert("usageEvents", {
      provider: args.provider,
      operation: args.operation,
      ...(args.runId ? { runId: args.runId } : {}),
      ...(args.leadId ? { leadId: args.leadId } : {}),
      ...(args.auditId ? { auditId: args.auditId } : {}),
      estimatedCostUsd: args.estimatedCostUsd,
      ...(args.tokens ? { tokens: args.tokens } : {}),
      ...(args.callCounts ? { callCounts: args.callCounts } : {}),
      createdAt: now,
    });
  },
});

export const listLatestUsageEvents = query({
  args: {
    limit: v.optional(v.number()),
  },
  returns: v.array(usageEventDoc),
  handler: async (ctx, args): Promise<Doc<"usageEvents">[]> => {
    await requireOperator(ctx);

    return await ctx.db
      .query("usageEvents")
      .withIndex("by_createdAt")
      .order("desc")
      .take(normalizeListLimit(args.limit));
  },
});

export const listUsageEventsByRun = query({
  args: {
    runId: v.id("agentRuns"),
    limit: v.optional(v.number()),
  },
  returns: v.array(usageEventDoc),
  handler: async (ctx, args): Promise<Doc<"usageEvents">[]> => {
    await requireOperator(ctx);

    return await ctx.db
      .query("usageEvents")
      .withIndex("by_runId_and_createdAt", (q) => q.eq("runId", args.runId))
      .order("desc")
      .take(normalizeListLimit(args.limit));
  },
});

export const listUsageEventsByLead = query({
  args: {
    leadId: v.id("leads"),
    limit: v.optional(v.number()),
  },
  returns: v.array(usageEventDoc),
  handler: async (ctx, args): Promise<Doc<"usageEvents">[]> => {
    await requireOperator(ctx);

    return await ctx.db
      .query("usageEvents")
      .withIndex("by_leadId_and_createdAt", (q) => q.eq("leadId", args.leadId))
      .order("desc")
      .take(normalizeListLimit(args.limit));
  },
});

export const listUsageEventsByAudit = query({
  args: {
    auditId: v.id("audits"),
    limit: v.optional(v.number()),
  },
  returns: v.array(usageEventDoc),
  handler: async (ctx, args): Promise<Doc<"usageEvents">[]> => {
    await requireOperator(ctx);

    return await ctx.db
      .query("usageEvents")
      .withIndex("by_auditId_and_createdAt", (q) => q.eq("auditId", args.auditId))
      .order("desc")
      .take(normalizeListLimit(args.limit));
  },
});

export const listUsageEventsByProvider = query({
  args: {
    provider: usageEventProvider,
    limit: v.optional(v.number()),
  },
  returns: v.array(usageEventDoc),
  handler: async (ctx, args): Promise<Doc<"usageEvents">[]> => {
    await requireOperator(ctx);

    return await ctx.db
      .query("usageEvents")
      .withIndex("by_provider_and_createdAt", (q) =>
        q.eq("provider", args.provider),
      )
      .order("desc")
      .take(normalizeListLimit(args.limit));
  },
});