Backfill missing Langfuse generation costs with Narev

import { randomUUID } from "node:crypto";

const required = [
  "LANGFUSE_PUBLIC_KEY",
  "LANGFUSE_SECRET_KEY",
  "LANGFUSE_FROM",
  "LANGFUSE_TO",
  "USAGE_BUCKETS",
];

for (const name of required) {
  if (!process.env[name]) {
    throw new Error(`Missing required environment variable: ${name}`);
  }
}

const host = process.env.LANGFUSE_HOST ?? "https://cloud.langfuse.com";
const dryRun = process.env.DRY_RUN !== "false";
const auth = Buffer.from(
  `${process.env.LANGFUSE_PUBLIC_KEY}:${process.env.LANGFUSE_SECRET_KEY}`,
).toString("base64");

// Add aliases when a Langfuse model name does not exactly match one Narev row.
const modelOverrides = {
  // "my-direct-deepseek-alias": {
  //   modelId: "deepseek-v4-pro",
  //   providerId: "deepseek",
  // },
};

if (!["exclusive", "inclusive"].includes(process.env.USAGE_BUCKETS)) {
  throw new Error('USAGE_BUCKETS must be exclusive or inclusive');
}
const from = Date.parse(process.env.LANGFUSE_FROM);
const to = Date.parse(process.env.LANGFUSE_TO);
if (!Number.isFinite(from) || !Number.isFinite(to) || from >= to) {
  throw new Error('Provide a valid increasing UTC date range');
}
const priceCache = new Map();

async function requestJson(url, options = {}) {
  const response = await fetch(url, options);
  const text = await response.text();
  let body;

  try {
    body = text ? JSON.parse(text) : {};
  } catch {
    throw new Error(`${response.status} ${url}: expected JSON, received ${text}`);
  }

  if (!response.ok) {
    throw new Error(`${response.status} ${url}: ${JSON.stringify(body)}`);
  }

  return body;
}

async function* listUnpricedGenerations() {
  let cursor;

  do {
    const url = new URL("/api/public/v2/observations", host);
    url.searchParams.set("fields", "core,model,usage");
    url.searchParams.set("type", "GENERATION");
    url.searchParams.set("fromStartTime", process.env.LANGFUSE_FROM);
    url.searchParams.set("toStartTime", process.env.LANGFUSE_TO);
    url.searchParams.set("limit", "1000");
    if (cursor) url.searchParams.set("cursor", cursor);

    const page = await requestJson(url, {
      headers: { Authorization: `Basic ${auth}` },
    });

    for (const observation of page.data ?? []) {
      if (
        observation.costDetails?.total == null &&
        observation.totalCost == null
      ) {
        yield observation;
      }
    }

    cursor = page.meta?.cursor;
  } while (cursor);
}

const usageKeys = {
  prompt_tokens: ["input", "prompt_tokens"],
  completion_tokens: ["output", "completion_tokens"],
  cache_read_tokens: ["cache_read_input_tokens", "input_cached_tokens", "cache_read_tokens"],
  cache_write_tokens: ["cache_write_input_tokens", "input_cache_creation", "cache_write_tokens"],
  reasoning_tokens: ["output_reasoning_tokens", "reasoning_tokens"],
};

function mapUsage(details) {
  if (!details || typeof details !== "object") throw new Error("Missing usageDetails");
  const known = new Set([...Object.values(usageKeys).flat(), "total"]);
  for (const [key, value] of Object.entries(details)) {
    if (!known.has(key)) throw new Error(`Unmapped usage key: ${key}`);
    if (!Number.isSafeInteger(value) || value < 0) throw new Error(`Invalid usage: ${key}`);
  }
  const usage = {};
  for (const [field, aliases] of Object.entries(usageKeys)) {
    const present = aliases.filter((key) => details[key] != null);
    if (present.length > 1) throw new Error(`Ambiguous aliases for ${field}`);
    if (present.length === 0 && ["prompt_tokens", "completion_tokens"].includes(field)) {
      throw new Error(`Missing ${field}`);
    }
    usage[field] = present.length ? details[present[0]] : 0;
  }
  if (process.env.USAGE_BUCKETS === "exclusive") {
    usage.prompt_tokens += usage.cache_read_tokens + usage.cache_write_tokens;
    usage.completion_tokens += usage.reasoning_tokens;
  }
  if (usage.cache_read_tokens + usage.cache_write_tokens > usage.prompt_tokens ||
      usage.reasoning_tokens > usage.completion_tokens) {
    throw new Error("Subcategories exceed token totals; check USAGE_BUCKETS");
  }
  return usage;
}

async function resolveModel(providedModelName) {
  if (modelOverrides[providedModelName]) {
    return modelOverrides[providedModelName];
  }

  if (priceCache.has(providedModelName)) {
    return priceCache.get(providedModelName);
  }

  const url = new URL("https://api.narev.ai/v1/prices/search");
  url.searchParams.set("q", providedModelName);
  url.searchParams.set("page_size", "100");
  const exact = [];
  let page = 1;
  let totalPages;
  do {
    url.searchParams.set("page", String(page));
    const result = await requestJson(url);
    exact.push(...(result.data ?? []).filter(
      (row) => row.model_id === providedModelName,
    ));
    totalPages = result.meta.total_pages;
    page += 1;
  } while (page <= totalPages);

  if (exact.length !== 1 || !exact[0].pricing) {
    throw new Error(
      `Expected one Narev price for ${providedModelName}, found ${exact.length}. Add a modelOverrides entry.`,
    );
  }

  const resolved = {
    modelId: exact[0].model_id,
    providerId: exact[0].provider_id,
  };
  priceCache.set(providedModelName, resolved);
  return resolved;
}

async function calculateCost(modelId, providerId, usage) {
  return requestJson("https://api.narev.ai/v1/traces/cost", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      model_id: modelId,
      provider_id: providerId,
      usage,
    }),
  });
}

async function updateCost(observation, result) {
  const response = await requestJson(`${host}/api/public/ingestion`, {
    method: "POST",
    headers: {
      Authorization: `Basic ${auth}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify({
      batch: [
        {
          id: randomUUID(),
          type: "generation-update",
          timestamp: new Date().toISOString(),
          body: {
            id: observation.id,
            traceId: observation.traceId,
            costDetails: { total: result.cost_breakdown.total },
          },
        },
      ],
      metadata: { batch_size: 1, sdk_name: "narev-guide" },
    }),
  });

  if ((response.errors ?? []).length > 0) {
    throw new Error(`Langfuse rejected the update: ${JSON.stringify(response.errors)}`);
  }
}

let priced = 0;
let skipped = 0;

for await (const observation of listUnpricedGenerations()) {
  try {
    if (!observation.model) {
      throw new Error("Observation has no model");
    }

    const usage = mapUsage(observation.usageDetails);
    const { modelId, providerId } = await resolveModel(
      observation.model,
    );
    const result = await calculateCost(modelId, providerId, usage);
    const total = result.cost_breakdown?.total;

    if (!Number.isFinite(total) || total < 0) {
      throw new Error(`Invalid Narev cost response: ${JSON.stringify(result)}`);
    }

    if (dryRun) {
      console.log(`[dry run] ${observation.id}: $${total}`);
    } else {
      await updateCost(observation, result);
      console.log(`Updated ${observation.id}: $${total}`);
    }
    priced += 1;
  } catch (error) {
    skipped += 1;
    console.error(`Skipped ${observation.id}: ${error.message}`);
  }
}

console.log(`${dryRun ? "Checked" : "Updated"} ${priced}; skipped ${skipped}`);
if (skipped > 0) process.exitCode = 1;