plugin-ai-api 1.0.13 → 1.0.15

package/src/server/usage.ts

@@ -1,21 +1,110 @@
 import { Context } from '@nocobase/actions';
 
-export type Usage = { prompt_tokens?: number | null; completion_tokens?: number | null; total_tokens?: number | null };
+export type Usage = {
+  prompt_tokens: number | null;
+  completion_tokens: number | null;
+  total_tokens: number | null;
+};
 
-function normalizeUsage(value: unknown): Usage | undefined {
+export type AiApiAuthType = 'apiKey' | 'bearer' | 'oidc' | 'unknown';
+
+export interface AiApiUsageResult {
+  source: 'provider' | 'unavailable';
+  usage?: Usage;
+  gatewayResponseId?: string;
+  providerRequestId?: string;
+}
+
+export interface AiApiStreamResult {
+  succeeded: boolean;
+  id?: string;
+  errorCode?: string;
+}
+
+interface OAuthPrincipal {
+  clientId?: string;
+  subject?: string;
+  scopes?: string[];
+}
+
+interface AiApiContextState {
+  aiApiAuthType?: AiApiAuthType;
+  aiApiUsageResult?: AiApiUsageResult;
+  aiApiStreamResult?: AiApiStreamResult;
+  currentRole?: string;
+  currentRoles?: string[];
+  currentUser?: { id?: string | number | bigint };
+  oauthPrincipal?: OAuthPrincipal;
+}
+
+function getAiApiState(ctx: Context): AiApiContextState {
+  return ctx.state as AiApiContextState;
+}
+
+function normalizeTokenCount(value: unknown): number | null {
+  return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : null;
+}
+
+export function normalizeUsage(value: unknown): Usage | undefined {
   if (!value || typeof value !== 'object') return undefined;
   const source = value as Record<string, unknown>;
-  const prompt = source.prompt_tokens ?? source.input_tokens;
-  const completion = source.completion_tokens ?? source.output_tokens;
-  const total = source.total_tokens;
-  if (prompt == null && completion == null && total == null) return undefined;
+  const prompt = normalizeTokenCount(source.prompt_tokens ?? source.input_tokens);
+  const completion = normalizeTokenCount(source.completion_tokens ?? source.output_tokens);
+  let total = normalizeTokenCount(source.total_tokens);
+
+  if (total === null && prompt !== null && completion !== null) {
+    total = prompt + completion;
+  }
+
+  if (prompt === null && completion === null && total === null) return undefined;
+
   return {
-    prompt_tokens: typeof prompt === 'number' ? prompt : null,
-    completion_tokens: typeof completion === 'number' ? completion : null,
-    total_tokens: typeof total === 'number' ? total : null,
+    prompt_tokens: prompt,
+    completion_tokens: completion,
+    total_tokens: total,
+  };
+}
+
+export function setAiApiUsageResult(
+  ctx: Context,
+  value: unknown,
+  metadata: Pick<AiApiUsageResult, 'gatewayResponseId' | 'providerRequestId'> = {},
+): Usage | undefined {
+  const usage = normalizeUsage(value);
+  getAiApiState(ctx).aiApiUsageResult = usage
+    ? { source: 'provider', usage, ...metadata }
+    : { source: 'unavailable', ...metadata };
+  return usage;
+}
+
+export function setAiApiUsageUnavailable(ctx: Context, gatewayResponseId?: string): void {
+  getAiApiState(ctx).aiApiUsageResult = {
+    source: 'unavailable',
+    ...(gatewayResponseId ? { gatewayResponseId } : {}),
   };
 }
 
+export function extractProviderRequestId(value: unknown): string | undefined {
+  if (!value || typeof value !== 'object') return undefined;
+  const source = value as Record<string, unknown>;
+  const responseMetadata =
+    source.response_metadata && typeof source.response_metadata === 'object'
+      ? (source.response_metadata as Record<string, unknown>)
+      : undefined;
+  const headers =
+    responseMetadata?.headers && typeof responseMetadata.headers === 'object'
+      ? (responseMetadata.headers as Record<string, unknown>)
+      : undefined;
+  const candidate =
+    responseMetadata?.request_id ??
+    responseMetadata?.requestId ??
+    responseMetadata?.id ??
+    headers?.['x-request-id'] ??
+    headers?.['request-id'];
+
+  return typeof candidate === 'string' && candidate.length > 0 ? candidate : undefined;
+}
+
 export async function startUsageRecord(
   ctx: Context,
   requestId: string,
@@ -26,13 +115,21 @@ export async function startUsageRecord(
 ) {
   const body = (ctx.request.body || {}) as Record<string, unknown>;
   const messages = Array.isArray(body.messages) ? body.messages : undefined;
-  const oauth = ctx.state.oauthPrincipal as { clientId?: string; subject?: string; scopes?: string[] } | undefined;
+  const promptCount = Array.isArray(body.prompt) ? body.prompt.length : body.prompt === undefined ? undefined : 1;
+  const embeddingInputCount = Array.isArray(body.input) ? body.input.length : body.input === undefined ? undefined : 1;
+  const state = getAiApiState(ctx);
+  const userId = state.currentUser?.id;
+  if (userId === undefined || userId === null) {
+    throw new Error('AI API usage record requires an authenticated NocoBase user ID.');
+  }
+
+  const oauth = state.oauthPrincipal;
   const record = await ctx.db.getRepository('aiApiUsageRecords').create({
     values: {
       requestId,
-      userId: String(ctx.state.currentUser?.id),
-      roleName: ctx.state.currentRole || ctx.state.currentRoles?.[0] || 'unknown',
-      authType: ctx.state.aiApiAuthType || (oauth ? 'oidc' : 'session'),
+      userId,
+      roleName: state.currentRole || state.currentRoles?.[0] || 'unknown',
+      authType: state.aiApiAuthType || (oauth ? 'oidc' : 'unknown'),
       oauthClientId: oauth?.clientId,
       oauthSubject: oauth?.subject,
       oauthScopes: oauth?.scopes,
@@ -42,7 +139,12 @@ export async function startUsageRecord(
       status: 'pending',
       streaming,
       startedAt: new Date(),
-      requestMetadata: { messageCount: messages?.length, requestedMaxTokens: body.max_tokens },
+      requestMetadata: {
+        messageCount: messages?.length,
+        promptCount,
+        embeddingInputCount,
+        requestedMaxTokens: body.max_completion_tokens ?? body.max_tokens,
+      },
     },
   });
   return record.id;
@@ -50,29 +152,28 @@ export async function startUsageRecord(
 
 export async function finishUsageRecord(ctx: Context, id: unknown, startedAt: number, status: 'succeeded' | 'failed') {
   const response = (ctx.body || {}) as {
-    usage?: Usage;
-    response_metadata?: { usage?: unknown };
     id?: string;
     error?: { code?: string };
   };
-  const streamResult = ctx.state.aiApiStreamResult as
-    | { usage?: Usage; id?: string; errorCode?: string; succeeded: boolean }
-    | undefined;
-  const usage =
-    normalizeUsage(response.usage) ||
-    normalizeUsage(response.response_metadata?.usage) ||
-    normalizeUsage(streamResult?.usage);
+  const state = getAiApiState(ctx);
+  const streamResult = state.aiApiStreamResult;
+  const usageResult = state.aiApiUsageResult ?? { source: 'unavailable' as const };
+  const usage = usageResult.source === 'provider' ? usageResult.usage : undefined;
+  const gatewayResponseId = usageResult.gatewayResponseId || response.id || streamResult?.id;
   const values = {
     status: streamResult ? (streamResult.succeeded ? 'succeeded' : 'failed') : status,
     httpStatus: ctx.status,
     errorCode: response.error?.code || streamResult?.errorCode,
-    inputTokens: usage?.prompt_tokens,
-    outputTokens: usage?.completion_tokens,
-    totalTokens: usage?.total_tokens,
-    providerRequestId: response.id || streamResult?.id,
+    inputTokens: usage?.prompt_tokens ?? null,
+    outputTokens: usage?.completion_tokens ?? null,
+    totalTokens: usage?.total_tokens ?? null,
+    providerRequestId: usageResult.providerRequestId ?? null,
     completedAt: new Date(),
     durationMs: Date.now() - startedAt,
-    responseMetadata: { usageSource: usage ? 'provider' : 'unavailable' },
+    responseMetadata: {
+      usageSource: usageResult.source,
+      ...(gatewayResponseId ? { gatewayResponseId } : {}),
+    },
   };
   await ctx.db.getRepository('aiApiUsageRecords').update({ filterByTk: id, values });
 }

package/src/server/utils/openai-format.ts

@@ -58,7 +58,11 @@ export function toOpenAIResponse(options: {
   model: string;
   content: string;
   finishReason?: string;
-  usage?: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number };
+  usage?: {
+    prompt_tokens?: number | null;
+    completion_tokens?: number | null;
+    total_tokens?: number | null;
+  };
   toolCalls?: OpenAIToolCall[];
 }) {
   const {
@@ -156,7 +160,7 @@ export function toOpenAIEmbeddingsResponse(options: {
 /**
  * Format a streaming chunk as an SSE data line.
  */
-export function formatSSE(data: any): string {
+export function formatSSE(data: unknown): string {
   return `data: ${JSON.stringify(data)}\n\n`;
 }