Skip to content

Commit

Permalink
Merge pull request #15 from zhujunsan/record_usage
Browse files Browse the repository at this point in the history
记录每个请求的 Token 消耗(预估)
  • Loading branch information
Kerwin1202 authored Apr 9, 2023
2 parents dcc40e8 + fa47199 commit 0315da4
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 25 deletions.
2 changes: 1 addition & 1 deletion service/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
},
"dependencies": {
"axios": "^1.3.4",
"chatgpt": "^5.1.2",
"chatgpt": "^5.2.2",
"dotenv": "^16.0.3",
"esno": "^0.16.3",
"express": "^4.18.2",
Expand Down
22 changes: 12 additions & 10 deletions service/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

51 changes: 42 additions & 9 deletions service/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import type { ChatContext, ChatMessage } from './chatgpt'
import { chatConfig, chatReplyProcess, currentModel, initApi } from './chatgpt'
import { auth } from './middleware/auth'
import { clearConfigCache, getCacheConfig, getOriginConfig } from './storage/config'
import type { ChatOptions, Config, MailConfig, SiteConfig, UserInfo } from './storage/model'
import type { ChatOptions, Config, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
import { Status } from './storage/model'
import {
clearChat,
Expand All @@ -22,6 +22,7 @@ import {
getUser,
getUserById,
insertChat,
insertChatUsage,
renameChatRoom,
updateChat,
updateConfig,
Expand Down Expand Up @@ -229,10 +230,26 @@ router.post('/chat', auth, async (req, res) => {
if (response.status === 'Success') {
if (regenerate && message.options.messageId) {
const previousResponse = message.previousResponse || []
previousResponse.push({ response: message.response, messageId: message.options.messageId })
await updateChat(message._id, response.data.text, response.data.id, previousResponse)
} else {
await updateChat(message._id, response.data.text, response.data.id)
previousResponse.push({ response: message.response, options: message.options })
await updateChat(message._id as unknown as string,
response.data.text,
response.data.id,
response.data.detail.usage as UsageResponse,
previousResponse)
}
else {
await updateChat(message._id as unknown as string,
response.data.text,
response.data.id,
response.data.detail.usage as UsageResponse)
}

if (response.data.usage) {
await insertChatUsage(req.headers.userId as string,
roomId,
message._id,
response.data.id,
response.data.detail.usage as UsageResponse)
}
}
res.send(response)
Expand Down Expand Up @@ -263,10 +280,26 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
if (result.status === 'Success') {
if (regenerate && message.options.messageId) {
const previousResponse = message.previousResponse || []
previousResponse.push({ response: message.response, messageId: message.options.messageId })
await updateChat(message._id, result.data.text, result.data.id, previousResponse)
} else {
await updateChat(message._id, result.data.text, result.data.id)
previousResponse.push({ response: message.response, options: message.options })
await updateChat(message._id as unknown as string,
result.data.text,
result.data.id,
result.data.detail.usage as UsageResponse,
previousResponse)
}
else {
await updateChat(message._id as unknown as string,
result.data.text,
result.data.id,
result.data.detail.usage as UsageResponse)
}

if (result.data.detail.usage) {
await insertChatUsage(req.headers.userId as string,
roomId,
message._id,
result.data.id,
result.data.detail.usage as UsageResponse)
}
}
}
Expand Down
37 changes: 36 additions & 1 deletion service/src/storage/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ export class ChatOptions {
parentMessageId?: string
messageId?: string
conversationId?: string
promptTokens?: number
completionTokens?: number
totalTokens?: number
estimated?: boolean
constructor(parentMessageId?: string, messageId?: string, conversationId?: string) {
this.parentMessageId = parentMessageId
this.messageId = messageId
Expand All @@ -55,7 +59,7 @@ export class ChatOptions {

export class previousResponse {
response: string
messageId: string
options: ChatOptions
}

export class ChatInfo {
Expand All @@ -77,6 +81,37 @@ export class ChatInfo {
}
}

export class UsageResponse {
prompt_tokens: number
completion_tokens: number
total_tokens: number
estimated: boolean
}

export class ChatUsage {
_id: ObjectId
userId: string
roomId: number
chatId: ObjectId
messageId: string
promptTokens: number
completionTokens: number
totalTokens: number
estimated: boolean
dateTime: number
constructor(userId: string, roomId: number, chatId: ObjectId, messageId: string, usage: UsageResponse) {
this.userId = userId
this.roomId = roomId
this.chatId = chatId
this.messageId = messageId
this.promptTokens = usage.prompt_tokens
this.completionTokens = usage.completion_tokens
this.totalTokens = usage.total_tokens
this.estimated = usage.estimated
this.dateTime = new Date().getTime()
}
}

export class Config {
constructor(
public _id: ObjectId,
Expand Down
22 changes: 18 additions & 4 deletions service/src/storage/mongo.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { MongoClient, ObjectId } from 'mongodb'
import * as dotenv from 'dotenv'
import { ChatInfo, ChatRoom, Status, UserInfo } from './model'
import type { ChatOptions, Config } from './model'
import { ChatInfo, ChatRoom, ChatUsage, Status, UserInfo } from './model'
import type { ChatOptions, Config, UsageResponse } from './model'

dotenv.config()

Expand All @@ -11,6 +11,7 @@ const chatCol = client.db('chatgpt').collection('chat')
const roomCol = client.db('chatgpt').collection('chat_room')
const userCol = client.db('chatgpt').collection('user')
const configCol = client.db('chatgpt').collection('config')
const usageCol = client.db('chatgpt').collection('chat_usage')

/**
* 插入聊天信息
Expand All @@ -29,10 +30,17 @@ export async function getChat(roomId: number, uuid: number) {
return await chatCol.findOne({ roomId, uuid })
}

export async function updateChat(chatId: string, response: string, messageId: string, previousResponse?: []) {
export async function updateChat(chatId: string, response: string, messageId: string, usage: UsageResponse, previousResponse?: []) {
const query = { _id: new ObjectId(chatId) }
const update = {
$set: { 'response': response, 'options.messageId': messageId },
$set: {
'response': response,
'options.messageId': messageId,
'options.prompt_tokens': usage.prompt_tokens,
'options.completion_tokens': usage.completion_tokens,
'options.total_tokens': usage.total_tokens,
'options.estimated': usage.estimated,
},
}

if (previousResponse)
Expand All @@ -41,6 +49,12 @@ export async function updateChat(chatId: string, response: string, messageId: st
await chatCol.updateOne(query, update)
}

export async function insertChatUsage(userId: string, roomId: number, chatId: ObjectId, messageId: string, usage: UsageResponse) {
const chatUsage = new ChatUsage(userId, roomId, chatId, messageId, usage)
await usageCol.insertOne(chatUsage)
return chatUsage
}

export async function createChatRoom(userId: string, title: string, roomId: number) {
const room = new ChatRoom(userId, title, roomId)
await roomCol.insertOne(room)
Expand Down

0 comments on commit 0315da4

Please sign in to comment.