omniai-server/src/pricing.js

/**
 * Model pricing helpers backed by the database.
 *
 * All monetary values are stored as integers:
 *   - prices: mills (厘, 1/1000 CNY) — allows fine-grained pricing
 *   - costs calculated in mills, converted to cents only at billing layer
 *
 * Currency is unified to CNY; USD models are converted at retail rates.
 */

const { pool } = require("./db");

const LEGACY_MODEL_PRICING_MILLS = {
  "qwen-max": { input: 20, output: 60, currency: "CNY" },
  "qwen-max-latest": { input: 20, output: 60, currency: "CNY" },
  "qwen-plus": { input: 4, output: 12, currency: "CNY" },
  "qwen-plus-latest": { input: 4, output: 12, currency: "CNY" },
  "qwen-turbo": { input: 2, output: 6, currency: "CNY" },
  "qwen-turbo-latest": { input: 2, output: 6, currency: "CNY" },
  "qwen-long": { input: 0.5, output: 2, currency: "CNY" },
  "qwen3-235b-a22b": { input: 20, output: 60, currency: "CNY" },
  "qwen3-32b": { input: 4, output: 12, currency: "CNY" },
  "qwen3-14b": { input: 2, output: 6, currency: "CNY" },

  "gemini-3.1-pro": { input: 15, output: 45, currency: "CNY" },

  "nano-banana-pro": { flat: 200, currency: "CNY" },
  "nano-banana-2": { flat: 200, currency: "CNY" },
  "nano-banana-fast": { flat: 200, currency: "CNY" },
  "gpt-image-2": { flat: 200, currency: "CNY" },
  "gpt-image-2-vip": { flat: 200, currency: "CNY" },
  "wan2.7-image": { flat: 200, currency: "CNY" },
  "wan2.7-image-pro": { flat: 200, currency: "CNY" },

  "kling-v2-master": { flat: 100, currency: "CNY" },
  "kling-v2-std": { flat: 50, currency: "CNY" },
  "kling-v1-std": { flat: 50, currency: "CNY" },
  "kling-v1-pro": { flat: 100, currency: "CNY" },
  "kling-v3-omni": { flat: 100, currency: "CNY" },

  "seedance-2.0": { flat: 50, currency: "CNY" },
  "seedance-2.0-fast": { flat: 30, currency: "CNY" },
  "seedance-2.0-pro": { flat: 80, currency: "CNY" },
  "seedance-2.0-lite": { flat: 20, currency: "CNY" },
  "seedance-2.0-ark": { flat: 50, currency: "CNY" },
  "seedance-2.0-fast-ark": { flat: 30, currency: "CNY" },

  "gpt-4o": { input: 27, output: 108, currency: "CNY" },
  "gpt-4o-mini": { input: 2, output: 6, currency: "CNY" },
  "gpt-4.1": { input: 22, output: 86, currency: "CNY" },
  "gpt-4.1-mini": { input: 4, output: 17, currency: "CNY" },
  "gpt-4.1-nano": { input: 1, output: 4, currency: "CNY" },
};

function inferCategory(modelKey, pricing) {
  if (!pricing.flat) return "text";

  const normalized = modelKey.toLowerCase();
  if (
    normalized.includes("kling") ||
    normalized.includes("seedance") ||
    normalized.includes("video") ||
    normalized.includes("omni")
  ) {
    return "video";
  }

  return "image";
}

const DEFAULT_MODEL_PRICES = [
  ...Object.entries(LEGACY_MODEL_PRICING_MILLS).map(([modelKey, pricing]) => ({
    modelKey,
    displayName: modelKey,
    category: inferCategory(modelKey, pricing),
    pricingType: pricing.flat ? "flat" : "token",
    inputPriceMills: pricing.input != null ? Math.round(pricing.input) : null,
    outputPriceMills: pricing.output != null ? Math.round(pricing.output) : null,
    flatPriceMills: pricing.flat != null ? Math.round(pricing.flat) : null,
    inputPrice: pricing.input ?? null,
    outputPrice: pricing.output ?? null,
    flatPrice: pricing.flat ?? null,
    currency: pricing.currency || "CNY",
    enabled: 1,
  })),
  {
    modelKey: "image-generation-flat",
    displayName: "图片生成统一计费",
    category: "image",
    pricingType: "flat",
    inputPriceMills: null,
    outputPriceMills: null,
    flatPriceMills: 200,
    inputPrice: null,
    outputPrice: null,
    flatPrice: 200,
    currency: "CNY",
    enabled: 1,
  },
];

function normalizeModelPriceRow(row) {
  if (!row) return null;

  return {
    id: row.id,
    modelKey: row.model_key,
    displayName: row.display_name,
    category: row.category,
    pricingType: row.pricing_type,
    inputPriceMills: row.input_price_mills ?? null,
    outputPriceMills: row.output_price_mills ?? null,
    flatPriceMills: row.flat_price_mills ?? null,
    currency: row.currency || "CNY",
    enabled: !!row.enabled,
    createdAt: row.created_at,
  };
}

function getModelPrice(modelKey, _options = {}) {
  // getModelPrice is still synchronous for use in calculateCostMills
  // We cache prices in a module-level map
  return _priceCache.get(modelKey) || null;
}

const _priceCache = new Map();

async function loadPriceCache() {
  const { rows } = await pool.query("SELECT * FROM model_prices WHERE enabled = 1");
  _priceCache.clear();
  for (const row of rows) {
    _priceCache.set(row.model_key, normalizeModelPriceRow(row));
  }
}

async function getModelPriceAsync(modelKey, options = {}) {
  const { includeDisabled = false } = options;
  const sql = includeDisabled
    ? "SELECT * FROM model_prices WHERE model_key = $1 LIMIT 1"
    : "SELECT * FROM model_prices WHERE model_key = $1 AND enabled = 1 LIMIT 1";
  const { rows } = await pool.query(sql, [modelKey]);
  return normalizeModelPriceRow(rows[0]);
}

async function listModelPrices(options = {}) {
  const { enabledOnly = false } = options;
  const sql = enabledOnly
    ? "SELECT * FROM model_prices WHERE enabled = 1 ORDER BY category, model_key"
    : "SELECT * FROM model_prices ORDER BY category, model_key";
  const { rows } = await pool.query(sql);
  return rows.map(normalizeModelPriceRow);
}

function calculateCostMills(model, promptTokens, completionTokens) {
  const pricing = getModelPrice(model);
  if (!pricing) return null;

  if (pricing.pricingType === "flat") return pricing.flatPriceMills;

  const inputCost = Math.round(((promptTokens || 0) / 1000) * (pricing.inputPriceMills || 0));
  const outputCost = Math.round(((completionTokens || 0) / 1000) * (pricing.outputPriceMills || 0));
  return inputCost + outputCost;
}

function calculateCost(model, promptTokens, completionTokens) {
  const mills = calculateCostMills(model, promptTokens, completionTokens);
  if (mills == null) return null;
  return Math.round(mills) / 1000;
}

function getCurrency(_model) {
  return "CNY";
}

async function getAverageCostCents(provider) {
  const { rows } = await pool.query(
    `
    SELECT CAST(ROUND(AVG(CASE
      WHEN cost_estimate IS NOT NULL THEN cost_estimate * 10000
      ELSE 0
    END)::numeric) AS INTEGER) AS avg_cents
    FROM api_call_logs
    WHERE provider = $1 AND status = 'success' AND created_at >= NOW() - INTERVAL '7 days'
  `,
    [provider],
  );
  return rows[0]?.avg_cents || 0;
}

module.exports = {
  DEFAULT_MODEL_PRICES,
  LEGACY_MODEL_PRICING_MILLS,
  normalizeModelPriceRow,
  getModelPrice,
  getModelPriceAsync,
  listModelPrices,
  calculateCostMills,
  calculateCost,
  getCurrency,
  getAverageCostCents,
  loadPriceCache,
};