Skip to content

Commit b4f138d

Browse files
author
Kerwin
committed
chore: suport custom chat model
1 parent 948c3b2 commit b4f138d

File tree

11 files changed

+59
-57
lines changed

11 files changed

+59
-57
lines changed

service/src/chatgpt/index.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import httpsProxyAgent from 'https-proxy-agent'
77
import fetch from 'node-fetch'
88
import jwt_decode from 'jwt-decode'
99
import dayjs from 'dayjs'
10-
import type { AuditConfig, CHATMODEL, KeyConfig, UserInfo } from '../storage/model'
10+
import type { AuditConfig, KeyConfig, UserInfo } from '../storage/model'
1111
import { Status } from '../storage/model'
1212
import type { TextAuditService } from '../utils/textAudit'
1313
import { textAuditServices } from '../utils/textAudit'
@@ -34,7 +34,7 @@ const ErrorCodeMessage: Record<string, string> = {
3434
let auditService: TextAuditService
3535
const _lockedKeys: { key: string; lockedTime: number }[] = []
3636

37-
export async function initApi(key: KeyConfig, chatModel: CHATMODEL) {
37+
export async function initApi(key: KeyConfig, chatModel: string) {
3838
// More Info: https://github.com/transitive-bullshit/chatgpt-api
3939

4040
const config = await getCacheConfig()
@@ -400,7 +400,7 @@ async function randomKeyConfig(keys: KeyConfig[]): Promise<KeyConfig | null> {
400400
return thisKey
401401
}
402402

403-
async function getRandomApiKey(user: UserInfo, chatModel: CHATMODEL, accountId?: string): Promise<KeyConfig | undefined> {
403+
async function getRandomApiKey(user: UserInfo, chatModel: string, accountId?: string): Promise<KeyConfig | undefined> {
404404
let keys = (await getCacheApiKeys()).filter(d => hasAnyRole(d.userRoles, user.roles))
405405
.filter(d => d.chatModels.includes(chatModel))
406406
if (accountId)

service/src/index.ts

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ import type { ChatMessage } from './chatgpt'
99
import { abortChatProcess, chatConfig, chatReplyProcess, containsSensitiveWords, initAuditService } from './chatgpt'
1010
import { auth, getUserId } from './middleware/auth'
1111
import { clearApiKeyCache, clearConfigCache, getApiKeys, getCacheApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
12-
import type { AuditConfig, CHATMODEL, ChatInfo, ChatOptions, Config, KeyConfig, MailConfig, SiteConfig, UserConfig, UserInfo } from './storage/model'
13-
import { Status, UsageResponse, UserRole, chatModelOptions } from './storage/model'
12+
import type { AuditConfig, ChatInfo, ChatOptions, Config, KeyConfig, MailConfig, SiteConfig, UserConfig, UserInfo } from './storage/model'
13+
import { Status, UsageResponse, UserRole } from './storage/model'
1414
import {
1515
clearChat,
1616
createChatRoom,
@@ -136,7 +136,7 @@ router.post('/room-prompt', auth, async (req, res) => {
136136
router.post('/room-chatmodel', auth, async (req, res) => {
137137
try {
138138
const userId = req.headers.userId as string
139-
const { chatModel, roomId } = req.body as { chatModel: CHATMODEL; roomId: number }
139+
const { chatModel, roomId } = req.body as { chatModel: string; roomId: number }
140140
const success = await updateRoomChatModel(userId, roomId, chatModel)
141141
if (success)
142142
res.send({ status: 'Success', message: 'Saved successfully', data: null })
@@ -423,8 +423,8 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
423423
result.data.detail = {}
424424
result.data.detail.usage = new UsageResponse()
425425
// 因为 token 本身不计算, 所以这里默认以 gpt 3.5 的算做一个伪统计
426-
result.data.detail.usage.prompt_tokens = textTokens(prompt, 'gpt-3.5-turbo-0613')
427-
result.data.detail.usage.completion_tokens = textTokens(result.data.text, 'gpt-3.5-turbo-0613')
426+
result.data.detail.usage.prompt_tokens = textTokens(prompt, 'gpt-3.5-turbo')
427+
result.data.detail.usage.completion_tokens = textTokens(result.data.text, 'gpt-3.5-turbo')
428428
result.data.detail.usage.total_tokens = result.data.detail.usage.prompt_tokens + result.data.detail.usage.completion_tokens
429429
result.data.detail.usage.estimated = true
430430
}
@@ -579,6 +579,18 @@ router.post('/session', async (req, res) => {
579579
key: string
580580
value: string
581581
}[] = []
582+
583+
const chatModelOptions = config.siteConfig.chatModels.split(',').map((model: string) => {
584+
let label = model
585+
if (model === 'text-davinci-002-render-sha-mobile')
586+
label = 'gpt-3.5-mobile'
587+
return {
588+
label,
589+
key: model,
590+
value: model,
591+
}
592+
})
593+
582594
let userInfo: { name: string; description: string; avatar: string; userId: string; root: boolean; roles: UserRole[]; config: UserConfig }
583595
if (userId != null) {
584596
const user = await getUserById(userId)
@@ -740,7 +752,7 @@ router.post('/user-info', auth, async (req, res) => {
740752

741753
router.post('/user-chat-model', auth, async (req, res) => {
742754
try {
743-
const { chatModel } = req.body as { chatModel: CHATMODEL }
755+
const { chatModel } = req.body as { chatModel: string }
744756
const userId = req.headers.userId.toString()
745757

746758
const user = await getUserById(userId)

service/src/storage/config.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { ObjectId } from 'mongodb'
22
import * as dotenv from 'dotenv'
33
import type { TextAuditServiceProvider } from 'src/utils/textAudit'
44
import { isNotEmptyString, isTextAuditServiceProvider } from '../utils/is'
5-
import { AuditConfig, CHATMODELS, Config, KeyConfig, MailConfig, SiteConfig, TextAudioType, UserRole } from './model'
5+
import { AuditConfig, Config, KeyConfig, MailConfig, SiteConfig, TextAudioType, UserRole } from './model'
66
import { getConfig, getKeys, upsertKey } from './mongo'
77

88
dotenv.config()
@@ -93,6 +93,9 @@ export async function getOriginConfig() {
9393
'',
9494
)
9595
}
96+
if (!isNotEmptyString(config.siteConfig.chatModels))
97+
config.siteConfig.chatModels = 'gpt-3.5-turbo,gpt-3.5-turbo-1106,gpt-3.5-turbo-16k,gpt-3.5-turbo-16k-0613,gpt-4,gpt-4-0613,gpt-4-32k,gpt-4-32k-0613,text-davinci-002-render-sha-mobile,text-embedding-ada-002,gpt-4-mobile,gpt-4-browsing,gpt-4-1106-preview,gpt-4-vision-preview'
98+
9699
return config
97100
}
98101

@@ -140,8 +143,8 @@ export function clearApiKeyCache() {
140143

141144
export async function getApiKeys() {
142145
const result = await getKeys()
146+
const config = await getCacheConfig()
143147
if (result.keys.length <= 0) {
144-
const config = await getCacheConfig()
145148
if (config.apiModel === 'ChatGPTAPI')
146149
result.keys.push(await upsertKey(new KeyConfig(config.apiKey, 'ChatGPTAPI', [], [], '')))
147150

@@ -157,7 +160,7 @@ export async function getApiKeys() {
157160
key.userRoles.push(UserRole.Guest)
158161
}
159162
if (key.chatModels == null || key.chatModels.length <= 0) {
160-
CHATMODELS.forEach((chatModel) => {
163+
config.siteConfig.chatModels.split(',').forEach((chatModel) => {
161164
key.chatModels.push(chatModel)
162165
})
163166
}

service/src/storage/model.ts

Lines changed: 6 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -52,29 +52,8 @@ export class UserInfo {
5252
}
5353

5454
export class UserConfig {
55-
chatModel: CHATMODEL
56-
}
57-
58-
// https://platform.openai.com/docs/models/overview
59-
// 除此之外,gpt-4-0314、gpt-4-32k-0314、gpt-3.5-turbo-0301 模型将在 9 月 13 日被弃用。
60-
export type CHATMODEL = 'gpt-3.5-turbo' | 'gpt-3.5-turbo-0301' | 'gpt-3.5-turbo-0613' | 'gpt-3.5-turbo-16k' | 'gpt-3.5-turbo-16k-0613' | 'gpt-4' | 'gpt-4-0314' | 'gpt-4-32k' | 'gpt-4-32k-0314' | 'gpt-4-0613' | 'gpt-4-32k-0613' | 'text-davinci-002-render-sha-mobile' | 'text-embedding-ada-002' | 'gpt-4-mobile' | 'gpt-4-browsing' | 'gpt-4-1106-preview' | 'gpt-4-vision-preview'
61-
62-
export const CHATMODELS: CHATMODEL[] = [
63-
'gpt-3.5-turbo', 'gpt-3.5-turbo-0301', 'gpt-3.5-turbo-0613', 'gpt-3.5-turbo-16k', 'gpt-3.5-turbo-16k-0613', 'gpt-4', 'gpt-4-0314', 'gpt-4-32k', 'gpt-4-32k-0314', 'gpt-4-0613', 'gpt-4-32k-0613', 'text-davinci-002-render-sha-mobile', 'text-embedding-ada-002', 'gpt-4-mobile', 'gpt-4-browsing', 'gpt-4-1106-preview', 'gpt-4-vision-preview',
64-
]
65-
66-
export const chatModelOptions = [
67-
'gpt-3.5-turbo', 'gpt-3.5-turbo-0301', 'gpt-3.5-turbo-0613', 'gpt-3.5-turbo-16k', 'gpt-3.5-turbo-16k-0613', 'gpt-4', 'gpt-4-0314', 'gpt-4-32k', 'gpt-4-32k-0314', 'gpt-4-0613', 'gpt-4-32k-0613', 'text-davinci-002-render-sha-mobile', 'text-embedding-ada-002', 'gpt-4-mobile', 'gpt-4-browsing', 'gpt-4-1106-preview', 'gpt-4-vision-preview',
68-
].map((model: string) => {
69-
let label = model
70-
if (model === 'text-davinci-002-render-sha-mobile')
71-
label = 'gpt-3.5-mobile'
72-
return {
73-
label,
74-
key: model,
75-
value: model,
76-
}
77-
})
55+
chatModel: string
56+
}
7857

7958
export class ChatRoom {
8059
_id: ObjectId
@@ -86,7 +65,7 @@ export class ChatRoom {
8665
status: Status = Status.Normal
8766
// only access token used
8867
accountId?: string
89-
chatModel: CHATMODEL
68+
chatModel: string
9069
constructor(userId: string, title: string, roomId: number) {
9170
this.userId = userId
9271
this.title = title
@@ -198,6 +177,7 @@ export class SiteConfig {
198177
public registerReview?: boolean,
199178
public registerMails?: string,
200179
public siteDomain?: string,
180+
public chatModels?: string,
201181
) { }
202182
}
203183

@@ -233,11 +213,11 @@ export class KeyConfig {
233213
_id: ObjectId
234214
key: string
235215
keyModel: APIMODEL
236-
chatModels: CHATMODEL[]
216+
chatModels: string[]
237217
userRoles: UserRole[]
238218
status: Status
239219
remark: string
240-
constructor(key: string, keyModel: APIMODEL, chatModels: CHATMODEL[], userRoles: UserRole[], remark: string) {
220+
constructor(key: string, keyModel: APIMODEL, chatModels: string[], userRoles: UserRole[], remark: string) {
241221
this.key = key
242222
this.keyModel = keyModel
243223
this.chatModels = chatModels

service/src/storage/mongo.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import * as dotenv from 'dotenv'
33
import dayjs from 'dayjs'
44
import { md5 } from '../utils/security'
55
import { ChatInfo, ChatRoom, ChatUsage, Status, UserConfig, UserInfo, UserRole } from './model'
6-
import type { CHATMODEL, ChatOptions, Config, KeyConfig, UsageResponse } from './model'
6+
import type { ChatOptions, Config, KeyConfig, UsageResponse } from './model'
77

88
dotenv.config()
99

@@ -120,7 +120,7 @@ export async function updateRoomAccountId(userId: string, roomId: number, accoun
120120
return result.modifiedCount > 0
121121
}
122122

123-
export async function updateRoomChatModel(userId: string, roomId: number, chatModel: CHATMODEL) {
123+
export async function updateRoomChatModel(userId: string, roomId: number, chatModel: string) {
124124
const query = { userId, roomId }
125125
const update = {
126126
$set: {
@@ -217,7 +217,7 @@ export async function updateUserInfo(userId: string, user: UserInfo) {
217217
, { $set: { name: user.name, description: user.description, avatar: user.avatar } })
218218
}
219219

220-
export async function updateUserChatModel(userId: string, chatModel: CHATMODEL) {
220+
export async function updateUserChatModel(userId: string, chatModel: string) {
221221
return userCol.updateOne({ _id: new ObjectId(userId) }
222222
, { $set: { 'config.chatModel': chatModel } })
223223
}

src/api/index.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import type { AxiosProgressEvent, GenericAbortSignal } from 'axios'
22
import { get, post } from '@/utils/request'
3-
import type { AuditConfig, CHATMODEL, ConfigState, KeyConfig, MailConfig, SiteConfig, Status, UserInfo, UserPassword } from '@/components/common/Setting/model'
3+
import type { AuditConfig, ConfigState, KeyConfig, MailConfig, SiteConfig, Status, UserInfo, UserPassword } from '@/components/common/Setting/model'
44
import { useAuthStore, useSettingStore } from '@/store'
55

66
export function fetchChatConfig<T = any>() {
@@ -116,7 +116,7 @@ export function fetchUpdateUserInfo<T = any>(name: string, avatar: string, descr
116116
})
117117
}
118118

119-
export function fetchUpdateUserChatModel<T = any>(chatModel: CHATMODEL) {
119+
export function fetchUpdateUserChatModel<T = any>(chatModel: string) {
120120
return post<T>({
121121
url: '/user-chat-model',
122122
data: { chatModel },

src/components/common/Setting/Keys.vue

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
<script lang="ts" setup>
22
import { h, onMounted, reactive, ref } from 'vue'
33
import { NButton, NDataTable, NInput, NModal, NSelect, NSpace, NSwitch, NTag, useDialog, useMessage } from 'naive-ui'
4-
import type { CHATMODEL } from './model'
54
import { KeyConfig, Status, UserRole, apiModelOptions, userRoleOptions } from './model'
65
import { fetchGetKeys, fetchUpdateApiKeyStatus, fetchUpsertApiKey } from '@/api'
76
import { t } from '@/locales'
@@ -39,7 +38,7 @@ const columns = [
3938
key: 'chatModels',
4039
width: 320,
4140
render(row: any) {
42-
const tags = row.chatModels.map((chatModel: CHATMODEL) => {
41+
const tags = row.chatModels.map((chatModel: string) => {
4342
return h(
4443
NTag,
4544
{

src/components/common/Setting/Site.vue

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,18 @@ onMounted(() => {
117117
/>
118118
</div>
119119
</div>
120+
<div class="flex items-center space-x-4">
121+
<span class="flex-shrink-0 w-[100px]">{{ $t('setting.chatModels') }}</span>
122+
<div class="flex-1">
123+
<NInput
124+
:value="config && config.chatModels"
125+
placeholder="英文逗号分割 | English comma separated"
126+
type="textarea"
127+
:autosize="{ minRows: 1, maxRows: 4 }"
128+
@input="(val) => { if (config) config.chatModels = val }"
129+
/>
130+
</div>
131+
</div>
120132
<div class="flex items-center space-x-4">
121133
<span class="flex-shrink-0 w-[100px]" />
122134
<NButton :loading="saving" type="primary" @click="updateSiteInfo(config)">

src/components/common/Setting/model.ts

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,11 @@ export class ConfigState {
1515
auditConfig?: AuditConfig
1616
}
1717

18-
// https://platform.openai.com/docs/models/overview
19-
export type CHATMODEL = 'gpt-3.5-turbo' | 'gpt-3.5-turbo-0301' | 'gpt-3.5-turbo-0613' | 'gpt-3.5-turbo-16k' | 'gpt-3.5-turbo-16k-0613' | 'gpt-4' | 'gpt-4-0314' | 'gpt-4-32k' | 'gpt-4-32k-0314' | 'gpt-4-0613' | 'gpt-4-32k-0613' | 'text-davinci-002-render-sha-mobile' | 'text-embedding-ada-002' | 'gpt-4-mobile' | 'gpt-4-browsing' | 'gpt-4-1106-preview' | 'gpt-4-vision-preview'
20-
2118
export class UserConfig {
22-
chatModel?: CHATMODEL
19+
chatModel?: string
2320
}
2421

22+
// https://platform.openai.com/docs/models/overview
2523
export class SiteConfig {
2624
siteTitle?: string
2725
loginEnabled?: boolean
@@ -30,6 +28,7 @@ export class SiteConfig {
3028
registerReview?: boolean
3129
registerMails?: string
3230
siteDomain?: string
31+
chatModels?: string
3332
}
3433

3534
export class MailConfig {
@@ -88,11 +87,11 @@ export class KeyConfig {
8887
_id?: string
8988
key: string
9089
keyModel: APIMODEL
91-
chatModels: CHATMODEL[]
90+
chatModels: string[]
9291
userRoles: UserRole[]
9392
status: Status
9493
remark: string
95-
constructor(key: string, keyModel: APIMODEL, chatModels: CHATMODEL[], userRoles: UserRole[], remark: string) {
94+
constructor(key: string, keyModel: APIMODEL, chatModels: string[], userRoles: UserRole[], remark: string) {
9695
this.key = key
9796
this.keyModel = keyModel
9897
this.chatModels = chatModels

src/typings/chat.d.ts

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
declare namespace Chat {
2-
import { CHATMODEL } from "@/components/common/Setting/model"
3-
42
interface Chat {
53
uuid?: number
64
dateTime: string
@@ -27,7 +25,7 @@ declare namespace Chat {
2725
all?: boolean
2826
prompt?: string
2927
usingContext: boolean
30-
chatModel?: CHATMODEL
28+
chatModel?: string
3129
}
3230

3331
interface ChatState {

0 commit comments

Comments
 (0)