Backfill missing Langfuse generation costs with Narev

import { randomUUID } from "node:crypto";

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

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 = {
  // "openrouter/anthropic/claude-sonnet-4-6": {
  //   modelId: "claude-sonnet-4-6",
  //   providerId: "anthropic",
  // },
};

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);
}

function readCount(details, keys) {
  for (const key of keys) {
    if (Number.isInteger(details?.[key]) && details[key] >= 0) {
      return details[key];
    }
  }
  return 0;
}

function mapUsage(details) {
  const usage = {
    prompt_tokens: readCount(details, ["input", "prompt_tokens"]),
    completion_tokens: readCount(details, ["output", "completion_tokens"]),
    cache_read_tokens: readCount(details, [
      "cache_read_input_tokens",
      "input_cached_tokens",
      "cache_read_tokens",
    ]),
    cache_write_tokens: readCount(details, [
      "cache_write_input_tokens",
      "input_cache_creation",
      "cache_write_tokens",
    ]),
    reasoning_tokens: readCount(details, [
      "output_reasoning_tokens",
      "reasoning_tokens",
    ]),
  };

  if (
    usage.prompt_tokens === 0 &&
    usage.completion_tokens === 0 &&
    usage.cache_read_tokens === 0 &&
    usage.cache_write_tokens === 0 &&
    usage.reasoning_tokens === 0
  ) {
    throw new Error(`Unsupported usageDetails: ${JSON.stringify(details)}`);
  }

  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 result = await requestJson(url);
  const exact = (result.data ?? []).filter(
    (row) => row.model_id === providedModelName && row.pricing,
  );

  if (exact.length !== 1) {
    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: {
              narev_model_id: result.model_id,
              narev_provider_id: result.provider_id,
            },
          },
        },
      ],
      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.providedModelName) {
      throw new Error("Observation has no providedModelName");
    }

    const usage = mapUsage(observation.usageDetails);
    const { modelId, providerId } = await resolveModel(
      observation.providedModelName,
    );
    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;