nolo-cli 0.1.13 → 0.1.14

package/ai/token/normalizeUsage.ts

@@ -0,0 +1,84 @@
+// 文件路径: ai/token/normalizeUsage.ts
+
+import { RawUsage, NormalizedUsage } from "./types";
+
+/**
+ * 规范化 usage 数据
+ */
+export const normalizeUsage = (usage: RawUsage): NormalizedUsage => {
+  const inputTokens =
+    "input_tokens" in usage
+      ? usage.input_tokens
+      : "prompt_tokens" in usage
+        ? usage.prompt_tokens
+        : 0;
+
+  const outputTokens =
+    "output_tokens" in usage
+      ? usage.output_tokens
+      : "completion_tokens" in usage
+        ? usage.completion_tokens
+        : 0;
+
+  const cacheCreationInputTokens =
+    "cache_creation_input_tokens" in usage
+      ? usage.cache_creation_input_tokens
+      : "prompt_cache_miss_tokens" in usage
+        ? usage.prompt_cache_miss_tokens
+        : 0;
+
+  const cacheReadInputTokens =
+    "cache_read_input_tokens" in usage
+      ? usage.cache_read_input_tokens
+      : "prompt_cache_hit_tokens" in usage
+        ? usage.prompt_cache_hit_tokens
+        : 0;
+
+  // ✅ 如果 provider 返回 usage.cost，则先原样透传；
+  // 后续由 calculatePrice 按各 provider 的账单单位再做换算。
+  let cost = 0;
+  if ("cost" in usage && typeof (usage as any).cost === "number") {
+    cost = (usage as any).cost;
+  }
+
+  const billingProvider =
+    "billing_provider" in usage && typeof usage.billing_provider === "string"
+      ? usage.billing_provider.trim() || undefined
+      : undefined;
+
+  const billingModel =
+    "billing_model" in usage && typeof usage.billing_model === "string"
+      ? usage.billing_model.trim() || undefined
+      : undefined;
+
+  const billingServiceTier =
+    "billing_service_tier" in usage &&
+    typeof usage.billing_service_tier === "string"
+      ? usage.billing_service_tier.trim() || undefined
+      : undefined;
+  const billingEstimated =
+    "billing_estimated" in usage && usage.billing_estimated === true;
+  const imageGenerationCount =
+    "image_generation_count" in usage &&
+    typeof usage.image_generation_count === "number" &&
+    Number.isFinite(usage.image_generation_count)
+      ? usage.image_generation_count
+      : undefined;
+
+  return {
+    input_tokens: inputTokens,
+    output_tokens: outputTokens,
+    cache_creation_input_tokens: cacheCreationInputTokens,
+    cache_read_input_tokens: cacheReadInputTokens,
+    cost,
+    ...(typeof imageGenerationCount === "number"
+      ? { image_generation_count: imageGenerationCount }
+      : {}),
+    ...(billingProvider ? { billing_provider: billingProvider } : {}),
+    ...(billingModel ? { billing_model: billingModel } : {}),
+    ...(billingServiceTier
+      ? { billing_service_tier: billingServiceTier }
+      : {}),
+    ...(billingEstimated ? { billing_estimated: true } : {}),
+  };
+};

package/ai/token/openaiImageGenerationUsage.ts

@@ -0,0 +1,56 @@
+import { getModelInfo } from "ai/llm/getModelContextWindow";
+import { countImageParts } from "chat/messages/messageContract";
+
+type AgentLike = {
+  provider?: string | null;
+  model?: string | null;
+  imageConfig?: { enabled?: boolean } | null;
+};
+
+type UsageLike = Record<string, any> | null | undefined;
+
+type TraceMessageLike = {
+  role?: string;
+  content?: string | any[] | null;
+};
+
+export const isOpenAIBuiltInImageGenerationAgent = (
+  agentConfig: AgentLike | null | undefined
+): boolean =>
+  String(agentConfig?.provider || "").toLowerCase() === "openai" &&
+  !getModelInfo(String(agentConfig?.model || ""))?.hasImageOutput &&
+  !!agentConfig?.imageConfig?.enabled;
+
+export const withImageGenerationCount = <T extends UsageLike>(
+  usage: T,
+  imageGenerationCount: number
+): T => {
+  if (!usage || !Number.isFinite(imageGenerationCount) || imageGenerationCount <= 0) {
+    return usage;
+  }
+
+  const existingCount =
+    typeof usage.image_generation_count === "number" &&
+    Number.isFinite(usage.image_generation_count)
+      ? usage.image_generation_count
+      : 0;
+
+  return {
+    ...usage,
+    image_generation_count: Math.max(existingCount, imageGenerationCount),
+  };
+};
+
+export const countAssistantImageGenerationOutputs = (
+  trace: TraceMessageLike[] | null | undefined
+): number => {
+  if (!Array.isArray(trace)) return 0;
+  return trace.reduce((total, message) => {
+    if (message?.role !== "assistant") return total;
+    return total + countImageParts(message.content ?? null);
+  }, 0);
+};
+
+export const countImageGenerationOutputsInContent = (
+  content: string | any[] | null | undefined
+): number => countImageParts(content);

package/ai/token/prepareTokenUsageData.ts

@@ -0,0 +1,88 @@
+import { extractUserId } from "core/prefix";
+import { calculatePrice } from "./calculatePrice";
+import { normalizeUsage } from "./normalizeUsage";
+import { resolveBillingTarget } from "./resolveBillingTarget";
+import type { RawUsage, TokenUsageData } from "./types";
+
+interface BillingAgentConfig {
+  model: string;
+  provider?: string;
+  apiSource?: string;
+  inputPrice?: number;
+  outputPrice?: number;
+  id?: string;
+}
+
+interface PrepareTokenUsageDataParams {
+  rawUsage: RawUsage;
+  agentConfig: BillingAgentConfig;
+  userId?: string;
+  username?: string;
+  cybotId: string;
+  dialogId: string;
+  timestamp?: number;
+}
+
+export interface PreparedTokenUsageData {
+  usage: ReturnType<typeof normalizeUsage>;
+  billedProvider?: string;
+  billedModel: string;
+  billedServiceTier?: string;
+  recordProvider: string;
+  tokenData: TokenUsageData;
+}
+
+export const prepareTokenUsageData = ({
+  rawUsage,
+  agentConfig,
+  userId,
+  username,
+  cybotId,
+  dialogId,
+  timestamp = Date.now(),
+}: PrepareTokenUsageDataParams): PreparedTokenUsageData => {
+  const usage = normalizeUsage(rawUsage);
+  const billingTarget = resolveBillingTarget({
+    usage,
+    fallbackProvider: agentConfig.provider,
+    fallbackModel: agentConfig.model,
+  });
+  const billedProvider = billingTarget.provider;
+  const billedModel = billingTarget.model;
+  const billedServiceTier = billingTarget.serviceTier;
+  const recordProvider = billedProvider ?? agentConfig.provider ?? "unknown";
+  const { cost, pay } = calculatePrice({
+    provider: billedProvider,
+    modelName: billedModel,
+    billingServiceTier: billedServiceTier,
+    usage,
+    externalPrice: agentConfig.id
+      ? {
+          input: agentConfig.inputPrice ?? 0,
+          output: agentConfig.outputPrice ?? 0,
+          creatorId: extractUserId(agentConfig.id),
+        }
+      : undefined,
+  });
+
+  return {
+    usage,
+    billedProvider,
+    billedModel,
+    billedServiceTier,
+    recordProvider,
+    tokenData: {
+      ...usage,
+      userId,
+      username,
+      cybotId,
+      model: billedModel,
+      provider: recordProvider,
+      billing_service_tier: billedServiceTier,
+      dialogId,
+      cost,
+      pay,
+      timestamp,
+    },
+  };
+};

package/ai/token/query.ts

@@ -0,0 +1,88 @@
+import { createTokenStatsKey } from "database/keys";
+
+
+
+
+
+export interface QueryParams {
+  userId: string;
+  startTime?: number; // UTC timestamp
+  endTime?: number; // UTC timestamp
+  model?: string;
+  limit?: number;
+  offset?: number;
+}
+
+export interface TokenStats {
+  total: number;
+  date: string; // UTC YYYY-MM-DD
+  inputTokens: number;
+  outputTokens: number;
+  cost: number;
+}
+
+export interface StatsParams {
+  userId: string;
+  period: "day";
+  startDate: string; // UTC YYYY-MM-DD
+  endDate: string; // UTC YYYY-MM-DD
+}
+
+// 通用数据库迭代器
+const iterateDb = async <T>(
+  db: any,
+  options: any,
+  filter: (v: T) => boolean
+): Promise<T[]> => {
+  const records: T[] = [];
+  const { offset = 0, limit } = options;
+  let count = 0;
+
+  try {
+    for await (const [_, value] of db.iterator(options)) {
+      if (filter(value)) {
+        if (count >= offset) {
+          records.push(value);
+          if (limit && records.length >= limit) {
+            break;
+          }
+        }
+        count++;
+      }
+    }
+    return records;
+  } catch (err) {
+    throw err;
+  }
+};
+
+/**
+ * 获取Token统计数据
+ * @param db - 数据库实例
+ * @param params.userId - 用户ID
+ * @param params.startDate - 开始日期 YYYY-MM-DD (UTC)
+ * @param params.endDate - 结束日期 YYYY-MM-DD (UTC)
+ * @param params.period - 统计周期,目前支持 "day"
+ */
+export const getTokenStats = async (
+  db: any,
+  params: StatsParams
+): Promise<TokenStats[]> => {
+  const { userId, startDate, endDate } = params;
+
+  // 统计数据按UTC日期(00:00:00)存储
+  const startKey = createTokenStatsKey(userId, startDate);
+  const endKey = createTokenStatsKey(userId, endDate);
+
+  const records = await iterateDb<TokenStats>(
+    db,
+    {
+      gte: startKey,
+      lte: endKey,
+    },
+    (record) => Boolean(record?.total)
+  );
+
+  return records;
+};
+

package/ai/token/queryUserTokens.ts

@@ -0,0 +1,59 @@
+import { TokenRecord } from "./types";
+import { createTokenKey } from "database/keys";
+
+
+
+export interface QueryParams {
+  userId: string;
+  startTime?: number; // 不传则查询当天
+  model?: string;
+  pageSize?: number;
+  offset?: number;
+}
+
+export interface QueryResult {
+  records: TokenRecord[];
+  total: number;
+}
+
+/**
+ * 查询用户token使用记录
+ * @param db 数据库实例
+ * @param params 查询参数，包含用户ID、起始时间（默认当天）、模型、分页参数
+ * @returns 包含记录列表和总数的结果
+ */
+export const queryUserTokens = async (
+  db: any,
+  params: QueryParams
+): Promise<QueryResult> => {
+  const { userId, startTime, model, pageSize = 100, offset = 0 } = params;
+  const { start, end } = createTokenKey.range(userId, startTime || Date.now());
+  const records: TokenRecord[] = [];
+
+  try {
+    // 单次扫描：同时计总数与收集分页数据
+    let count = 0;
+    for await (const [_, value] of db.iterator({
+      gte: start,
+      lte: end,
+      reverse: true,
+    })) {
+      if (!model || value.model === model) {
+        if (count >= offset && records.length < pageSize) {
+          records.push({
+            ...value,
+            createdAt:
+              typeof value.createdAt === "number"
+                ? value.createdAt
+                : value.timestamp,
+          });
+        }
+        count++;
+      }
+    }
+
+    return { records, total: count };
+  } catch (err) {
+    throw err;
+  }
+};

package/ai/token/resolveBillingTarget.ts

@@ -0,0 +1,52 @@
+import { resolveOpenAIServiceTier } from "integrations/openai/flexTier";
+import type { NormalizedUsage, RawUsage } from "./types";
+
+type BillingUsage = Pick<
+  RawUsage | NormalizedUsage,
+  "billing_provider" | "billing_model" | "billing_service_tier"
+>;
+
+const normalizeString = (value?: string) => {
+  const trimmed = typeof value === "string" ? value.trim() : "";
+  return trimmed || undefined;
+};
+
+const resolveBillingModel = (
+  usageModel?: string,
+  fallbackModel?: string
+) => {
+  const model = normalizeString(usageModel) ?? normalizeString(fallbackModel);
+  if (!model) {
+    throw new Error("Billing model is required");
+  }
+  return model;
+};
+
+export const resolveBillingTarget = ({
+  usage,
+  fallbackProvider,
+  fallbackModel,
+}: {
+  usage?: BillingUsage | null;
+  fallbackProvider?: string;
+  fallbackModel: string;
+}) => {
+  const provider =
+    normalizeString(usage?.billing_provider) ?? normalizeString(fallbackProvider);
+  const model = resolveBillingModel(usage?.billing_model, fallbackModel);
+  const requestedServiceTier = normalizeString(usage?.billing_service_tier);
+  const serviceTier =
+    provider === "google"
+      ? requestedServiceTier
+      : resolveOpenAIServiceTier({
+          providerName: provider,
+          model,
+          requestedServiceTier,
+        });
+
+  return {
+    provider,
+    model,
+    serviceTier,
+  };
+};

package/ai/token/saveTokenRecord.ts

@@ -0,0 +1,53 @@
+import { TokenUsageData, TokenRecord } from "ai/token/types";
+import { DataType } from "create/types";
+import { createTokenKey } from "database/keys";
+import { write } from "database/dbSlice";
+import { toast } from "app/utils/toast";
+import { pino } from "pino";
+
+const logger = pino({ name: "token-record", level: "info" });
+
+type TokenCount = { input: number; output: number };
+
+export interface ModelStats {
+  count: number;
+  tokens: TokenCount;
+  cost: number;
+}
+
+export const createTokenRecord = (
+  data: TokenUsageData,
+  { cost, inputPrice, outputPrice }: Partial<TokenRecord> = {}
+): TokenRecord => ({
+  ...data,
+  cost: cost || data.cost,
+  inputPrice,
+  outputPrice,
+});
+
+export const saveTokenRecord = async (
+  tokenData: TokenUsageData,
+  record: TokenRecord,
+  thunkApi
+) => {
+  const key = createTokenKey.record(tokenData.userId, tokenData.timestamp);
+  try {
+    await thunkApi.dispatch(
+      write({
+        data: { ...record, id: key, type: DataType.TOKEN },
+        customKey: key,
+      })
+    );
+  } catch (error) {
+    logger.error(
+      {
+        key,
+        userId: tokenData.userId,
+        error: error.message,
+      },
+      "Failed to save token record"
+    );
+    toast.error("Failed to save token record");
+    throw error;
+  }
+};

package/ai/token/serverDialogProjection.ts

@@ -0,0 +1,78 @@
+import pino from "pino";
+
+import serverDb from "database/server/db";
+import { createKey } from "database/keys";
+import { DataType } from "create/types";
+import { isLevelNotFoundError } from "share/helpers";
+
+const logger = pino({ name: "server-dialog-projection" });
+const dialogProjectionQueue = new Map<string, Promise<void>>();
+
+interface ApplyServerDialogProjectionDeltaOpts {
+    userId: string;
+    dialogId: string;
+    inputTokensDelta?: number;
+    outputTokensDelta?: number;
+    costDelta?: number;
+}
+
+export async function applyServerDialogProjectionDelta(
+    opts: ApplyServerDialogProjectionDeltaOpts,
+): Promise<void> {
+    const {
+        userId,
+        dialogId,
+        inputTokensDelta = 0,
+        outputTokensDelta = 0,
+        costDelta = 0,
+    } = opts;
+
+    if (!dialogId) return;
+    if (inputTokensDelta === 0 && outputTokensDelta === 0 && costDelta === 0) return;
+
+    const dialogKey = createKey(DataType.DIALOG, userId, dialogId);
+    const previousTask = dialogProjectionQueue.get(dialogKey) ?? Promise.resolve();
+    const nextTask = previousTask.catch(() => undefined).then(async () => {
+        let existingDialog: any;
+        try {
+            existingDialog = await serverDb.get(dialogKey);
+        } catch (error) {
+            if (isLevelNotFoundError(error)) {
+                logger.warn(
+                    {
+                        userId,
+                        dialogId,
+                        dialogKey,
+                        inputTokensDelta,
+                        outputTokensDelta,
+                        costDelta,
+                    },
+                    "Skipping dialog projection update because dialog record is missing",
+                );
+                return;
+            }
+            throw error;
+        }
+
+        await serverDb.put(dialogKey, {
+            ...existingDialog,
+            inputTokens: (existingDialog.inputTokens ?? 0) + inputTokensDelta,
+            outputTokens: (existingDialog.outputTokens ?? 0) + outputTokensDelta,
+            totalCost: Number(((existingDialog.totalCost ?? 0) + costDelta).toFixed(6)),
+        });
+    });
+    const queueEntry = nextTask.then(
+        () => undefined,
+        () => undefined,
+    );
+
+    dialogProjectionQueue.set(dialogKey, queueEntry);
+
+    try {
+        await nextTask;
+    } finally {
+        if (dialogProjectionQueue.get(dialogKey) === queueEntry) {
+            dialogProjectionQueue.delete(dialogKey);
+        }
+    }
+}

package/ai/token/serverTokenWriter.ts

@@ -0,0 +1,143 @@
+// ai/token/serverTokenWriter.ts
+// 服务端直接写 LevelDB 的 token 记录（对应客户端 saveTokenRecord + saveTokenUsage）
+// 用于 agentRun 等纯服务端调用场景，无 Redux 依赖
+
+import { ulid } from "ulid";
+import { format } from "date-fns";
+import serverDb from "database/server/db";
+import { createKey, createTokenKey, createTokenStatsKey } from "database/keys";
+import { DataType } from "create/types";
+import type { RawUsage } from "./types";
+import pino from "pino";
+import { prepareTokenUsageData } from "./prepareTokenUsageData";
+import { applyServerDialogProjectionDelta } from "./serverDialogProjection";
+
+const logger = pino({ name: "server-token-writer" });
+
+interface WriteTokenOpts {
+    userId: string;
+    username?: string;
+    agentKey: string;   // 对应 cybotId
+    agentConfig: {
+        model: string;
+        // apiSource="custom" 的 agent 没有 provider 字段，允许 undefined
+        provider?: string;
+        apiSource?: string;
+        inputPrice?: number;
+        outputPrice?: number;
+        id?: string;
+    };
+    runId: string;      // 对应 dialogId
+    rawUsage: RawUsage;
+}
+
+export async function writeServerTokenRecord(opts: WriteTokenOpts): Promise<{ cost: number }> {
+    const { userId, username = "", agentKey, agentConfig, runId, rawUsage } = opts;
+
+    const prepared = prepareTokenUsageData({
+        rawUsage,
+        agentConfig,
+        userId,
+        username,
+        cybotId: agentKey,
+        dialogId: runId,
+    });
+    const {
+        usage,
+        billedModel,
+        billedServiceTier,
+        recordProvider,
+        tokenData,
+    } = prepared;
+    const { cost, pay } = tokenData;
+
+    const timestamp = Date.now();
+    const dateKey = format(timestamp, "yyyy-MM-dd");
+    const tokenId = ulid(timestamp);
+
+    // 1. 写单条 token 明细记录（key: token-{userId}-{timestamp}）
+    const recordKey = createTokenKey.record(userId, timestamp);
+    const tokenRecord = {
+        id: tokenId,
+        type: DataType.TOKEN,
+        userId,
+        username,
+        cybotId: agentKey,
+        model: billedModel,
+        provider: recordProvider,
+        billing_service_tier: billedServiceTier,
+        // 记录 apiSource 便于后续分析 custom/platform 的收益分布
+        apiSource: agentConfig.apiSource,
+        dialogId: runId,    // runId 作为 dialogId，便于与 RunRecord 关联
+        input_tokens: usage.input_tokens,
+        output_tokens: usage.output_tokens,
+        cache_creation_input_tokens: usage.cache_creation_input_tokens,
+        cache_read_input_tokens: usage.cache_read_input_tokens,
+        cost,
+        pay,
+        createdAt: timestamp,
+        timestamp,
+        dateKey,
+    };
+
+    // 2. 读/更新每日统计（key: token-stats-day-user-{userId}-{date}）
+    const statsKey = createTokenStatsKey(userId, dateKey);
+    let stats: any;
+    try {
+        stats = await serverDb.get(statsKey);
+    } catch {
+        stats = null;
+    }
+
+    const modelName = billedModel || "unknown";
+    const providerName = recordProvider;
+    const inc = (s: any = { count: 0, tokens: { input: 0, output: 0 }, cost: 0 }) => ({
+        count: s.count + 1,
+        tokens: {
+            input: s.tokens.input + usage.input_tokens,
+            output: s.tokens.output + usage.output_tokens,
+        },
+        cost: s.cost + cost,
+    });
+
+    const updatedStats = stats ?? {
+        userId,
+        period: "day",
+        timeKey: dateKey,
+        total: { count: 0, tokens: { input: 0, output: 0 }, cost: 0 },
+        models: {},
+        providers: {},
+    };
+
+    const newStats = {
+        ...updatedStats,
+        id: statsKey,
+        type: DataType.TOKEN,
+        total: inc(updatedStats.total),
+        models: { ...updatedStats.models, [modelName]: inc(updatedStats.models?.[modelName]) },
+        providers: { ...updatedStats.providers, [providerName]: inc(updatedStats.providers?.[providerName]) },
+    };
+
+    await serverDb.batch([
+        { type: "put", key: recordKey, value: tokenRecord },
+        { type: "put", key: statsKey, value: newStats },
+    ]);
+
+    try {
+        await applyServerDialogProjectionDelta({
+            userId,
+            dialogId: runId,
+            inputTokensDelta: usage.input_tokens,
+            outputTokensDelta: usage.output_tokens,
+            costDelta: cost,
+        });
+    } catch (error) {
+        logger.warn(
+            { err: error, userId, runId, cost },
+            "Server token record persisted but dialog projection update failed",
+        );
+    }
+
+    logger.info({ userId, agentKey, cost, model: billedModel, provider: recordProvider }, "Token record written");
+    return { cost };
+}