feat: add image token (#675)

* feat: add user statistics by model endpoint and export

* fix: lint

* feat: add image params

* feat: add image usage

* fix: type check

* feat: add image token

* feat: delete log

* feat: add image config

Author: xiaokaiwangdragonplus

service/src/chatgpt/index.ts

@@ -1,6 +1,6 @@
 import type { ClientOptions } from 'openai'
 import type { RequestInit } from 'undici'
-import type { APIMODEL, AuditConfig, Config, KeyConfig, SearchResult, UserInfo } from '../storage/model'
+import type { APIMODEL, AuditConfig, Config, ImageUsageItem, KeyConfig, SearchResult, UserInfo } from '../storage/model'
 import type { TextAuditService } from '../utils/textAudit'
 import type { ChatMessage, RequestOptions } from './types'
 import { tavily } from '@tavily/core'
@@ -19,6 +19,59 @@ function renderSystemMessage(template: string, currentTime: string): string {
   return template.replace(/\{current_time\}/g, currentTime)
 }
 
+/**
+ * 根据图片的size和quality计算token
+ * @param size 图片尺寸，如 '1024x1024', '1024x1536', '1536x1024'
+ * @param quality 图片质量，如 'low', 'medium', 'high'
+ * @returns token数量，如果无法匹配则返回0
+ */
+function calculateImageTokens(size: string | undefined, quality: string | undefined): number {
+  if (!size || !quality) {
+    return 0
+  }
+
+  // 标准化quality
+  const normalizedQuality = quality.toLowerCase()
+  // 标准化size，判断是square、portrait还是landscape
+  const sizeLower = size.toLowerCase()
+  let sizeType: 'square' | 'portrait' | 'landscape' | null = null
+
+  if (sizeLower === '1024x1024') {
+    sizeType = 'square'
+  }
+  else if (sizeLower === '1024x1536') {
+    sizeType = 'portrait'
+  }
+  else if (sizeLower === '1536x1024') {
+    sizeType = 'landscape'
+  }
+
+  if (!sizeType) {
+    return 0
+  }
+
+  // Token计算表
+  const tokenMap: Record<string, Record<string, number>> = {
+    low: {
+      square: 272,
+      portrait: 408,
+      landscape: 400,
+    },
+    medium: {
+      square: 1056,
+      portrait: 1584,
+      landscape: 1568,
+    },
+    high: {
+      square: 4160,
+      portrait: 6240,
+      landscape: 6208,
+    },
+  }
+
+  return tokenMap[normalizedQuality]?.[sizeType] || 0
+}
+
 const ErrorCodeMessage: Record<string, string> = {
   401: '[OpenAI] 提供错误的API密钥 | Incorrect API key provided',
   403: '[OpenAI] 服务器拒绝访问，请稍后再试 | Server refused to access, please try again later',
@@ -275,6 +328,27 @@ search result: <search_result>${searchResultContent}</search_result>`,
         }
       }
 
+      // 如果 tools 中有 image_generation，并且 keyConfig 中有配置，则使用 keyConfig 中的配置
+      let finalTools = tools
+      if (tools && tools.length > 0) {
+        finalTools = tools.map((tool: any) => {
+          if (tool.type === 'image_generation') {
+            // 从 keyConfig 读取配置，如果不存在则使用默认值
+            const inputFidelity = key.inputFidelity || tool.input_fidelity || 'high'
+            const quality = key.quality || tool.quality || 'high'
+            const model = key.imageModel || tool.model || 'gpt-image-1.5'
+
+            return {
+              type: 'image_generation',
+              input_fidelity: inputFidelity,
+              quality,
+              model,
+            }
+          }
+          return tool
+        })
+      }
+
       const stream = await openai.responses.create(
         {
           model,
@@ -283,7 +357,7 @@ search result: <search_result>${searchResultContent}</search_result>`,
           reasoning,
           store: options.room.toolsEnabled,
           stream: true,
-          ...(tools && tools.length > 0 && { tools: tools as any }),
+          ...(finalTools && finalTools.length > 0 && { tools: finalTools as any }),
           // 如果有图片代表是编辑当前图片，不传递该参数，否则传递该参数
           ...(previousResponseId && uploadFileKeys.length === 0 && { previous_response_id: previousResponseId }),
         },
@@ -299,6 +373,7 @@ search result: <search_result>${searchResultContent}</search_result>`,
       const usage = new UsageResponse()
       const toolCalls: Array<{ type: string, result?: any }> = []
       let editImageId: string | undefined
+      let imageUsageList: ImageUsageItem[] = []
 
       // 心跳机制：防止生图等长时间操作时连接超时
       let heartbeatInterval: NodeJS.Timeout | null = null
@@ -341,26 +416,52 @@ search result: <search_result>${searchResultContent}</search_result>`,
             usage.prompt_tokens = resp.usage.input_tokens
             usage.completion_tokens = resp.usage.output_tokens
             usage.total_tokens = resp.usage.total_tokens
+            // 重置图片使用列表
+            imageUsageList = []
+
+            // 获取主模型和图片生成模型
+            const mainModel = resp.model
+            const imageModel = resp.tools && Array.isArray(resp.tools) && resp.tools.length > 0
+              ? (resp.tools.find((tool: any) => tool.type === 'image_generation') as any)?.model
+              : undefined
 
             // Extract tool calls from response
             if (resp.output && Array.isArray(resp.output)) {
               editImageId = responseId
               for (const output of resp.output) {
-                if (output.type === 'image_generation_call' && output.result) {
-                  const base64Data = output.result
-                  const fileIdentifier = await saveBase64ToFile(base64Data)
-
-                  if (fileIdentifier) {
-                    toolCalls.push({
-                      type: 'image_generation',
-                      result: fileIdentifier, // 文件名或S3 URL，前端会自动处理
-                    })
-                  }
-                  else {
-                    toolCalls.push({
-                      type: 'image_generation',
-                      result: base64Data,
-                    })
+                if (output.type === 'image_generation_call') {
+                  // 记录图片使用信息：size、quality、模型
+                  const imageOutput = output as any
+                  const imageSize = imageOutput.size
+                  const imageQuality = imageOutput.quality
+                  const imageTokens = calculateImageTokens(imageSize, imageQuality)
+
+                  imageUsageList.push({
+                    size: imageSize,
+                    quality: imageQuality,
+                    model: imageModel,
+                    mainModel,
+                    tokens: imageTokens,
+                    data_type: 'output',
+                  })
+
+                  // 处理图片结果
+                  if (output.result) {
+                    const base64Data = output.result
+                    const fileIdentifier = await saveBase64ToFile(base64Data)
+
+                    if (fileIdentifier) {
+                      toolCalls.push({
+                        type: 'image_generation',
+                        result: fileIdentifier, // 文件名或S3 URL，前端会自动处理
+                      })
+                    }
+                    else {
+                      toolCalls.push({
+                        type: 'image_generation',
+                        result: base64Data,
+                      })
+                    }
                   }
                 }
               }
@@ -384,6 +485,7 @@ search result: <search_result>${searchResultContent}</search_result>`,
         detail: { usage },
         ...(toolCalls.length > 0 && { tool_calls: toolCalls }),
         ...(editImageId && { editImageId }),
+        ...(imageUsageList.length > 0 && { image_usage: imageUsageList }),
       }
       return sendResponse({ type: 'Success', data: response })
     }

service/src/chatgpt/types.ts

@@ -1,5 +1,6 @@
 import type OpenAI from 'openai'
-import type { ChatRoom, SearchResult, UserInfo } from 'src/storage/model'
+import type { ChatRoom, ImageUsageItem, SearchResult, UserInfo } from 'src/storage/model'
+import type { ImageGenerationTool } from '../types'
 
 export interface ChatMessage {
   id: string
@@ -35,14 +36,15 @@ export interface ResponseChunk {
   }>
   // 编辑图片时使用的文件 ID（用于后续作为 previousResponseId）
   editImageId?: string
+  image_usage?: ImageUsageItem[]
 }
 
 export interface RequestOptions {
   message: string
   uploadFileKeys?: string[]
   parentMessageId?: string
   previousResponseId?: string
-  tools?: Array<{ type: 'image_generation' }>
+  tools?: Array<ImageGenerationTool>
   process?: (chunk: ResponseChunk) => void
   systemMessage?: string
   temperature?: number

service/src/index.ts

@@ -25,6 +25,7 @@ import {
   getUserById,
   getUsers,
   getUserStatisticsByDay,
+  getUserStatisticsByModel,
   initializeMongoDB,
   updateApiKeyStatus,
   updateConfig,
@@ -1084,6 +1085,22 @@ router.post('/statistics/by-day', auth, async (req, res) => {
   }
 })
 
+router.post('/statistics/by-model', auth, async (req, res) => {
+  try {
+    const { start, end } = req.body as { start?: number, end?: number }
+
+    // 只有管理员可以查看所有用户的统计
+    if (!isAdmin(req.headers.userId as string))
+      throw new Error('无权限 | No permission')
+
+    const data = await getUserStatisticsByModel(start, end)
+    res.send({ status: 'Success', message: '', data })
+  }
+  catch (error) {
+    res.send(error)
+  }
+})
+
 app.use('', chatRouter)
 app.use('/api', chatRouter)
 

service/src/routes/chat.ts

@@ -1,5 +1,5 @@
 import type { ResponseChunk } from '../chatgpt/types'
-import type { ChatInfo, ChatOptions, UsageResponse, UserInfo } from '../storage/model'
+import type { ChatInfo, ChatOptions, ImageUsageItem, UsageResponse, UserInfo } from '../storage/model'
 import type { RequestProps } from '../types'
 import * as console from 'node:console'
 import * as process from 'node:process'
@@ -450,7 +450,8 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
               }
 
               if (result.data.detail?.usage) {
-                await insertChatUsage(new ObjectId(req.headers.userId), roomId, message._id, result.data.id, model, result.data.detail?.usage as UsageResponse)
+                const imageUsage = result.data.image_usage as ImageUsageItem[] | undefined
+                await insertChatUsage(new ObjectId(req.headers.userId), roomId, message._id, result.data.id, model, result.data.detail?.usage as UsageResponse, imageUsage)
               }
               // update personal useAmount moved here
               // if not fakeuserid, and has valid user info and valid useAmount set by admin nut null and limit is enabled

service/src/storage/model.ts

@@ -167,6 +167,15 @@ export class UsageResponse {
   estimated: boolean
 }
 
+export interface ImageUsageItem {
+  size?: string
+  quality?: string
+  model?: string
+  mainModel?: string
+  tokens?: number
+  data_type?: string
+}
+
 export class ChatUsage {
   _id: ObjectId
   userId: ObjectId
@@ -179,7 +188,8 @@ export class ChatUsage {
   totalTokens: number
   estimated: boolean
   dateTime: number
-  constructor(userId: ObjectId, roomId: number, chatId: ObjectId, messageId: string, model: string, usage?: UsageResponse) {
+  imageUsage?: ImageUsageItem[]
+  constructor(userId: ObjectId, roomId: number, chatId: ObjectId, messageId: string, model: string, usage?: UsageResponse, imageUsage?: ImageUsageItem[]) {
     this.userId = userId
     this.roomId = roomId
     this.chatId = chatId
@@ -191,6 +201,9 @@ export class ChatUsage {
       this.totalTokens = usage.total_tokens
       this.estimated = usage.estimated
     }
+    if (imageUsage && imageUsage.length > 0) {
+      this.imageUsage = imageUsage
+    }
     this.dateTime = new Date().getTime()
   }
 }
@@ -316,6 +329,9 @@ export class KeyConfig {
   baseUrl?: string
   toolsEnabled?: boolean
   imageUploadEnabled?: boolean
+  inputFidelity?: 'low' | 'medium' | 'high'
+  quality?: 'low' | 'medium' | 'high'
+  imageModel?: 'gpt-image-1' | 'gpt-image-1.5'
   constructor(key: string, keyModel: APIMODEL, chatModel: string, userRoles: UserRole[], remark: string) {
     this.key = key
     this.keyModel = keyModel