perf: chat framwork

This commit is contained in:
archer
2023-05-03 15:28:25 +08:00
parent 91decc3683
commit 00a99261ae
23 changed files with 811 additions and 1011 deletions

View File

@@ -1,5 +0,0 @@
export enum OpenAiTuneStatusEnum {
cancelled = 'cancelled',
succeeded = 'succeeded',
pending = 'pending'
}

View File

@@ -1,14 +1,13 @@
import { SplitData } from '@/service/mongo';
import { getOpenAIApi } from '@/service/utils/auth';
import { axiosConfig } from '@/service/utils/tools';
import { getOpenApiKey } from '../utils/openai';
import type { ChatCompletionRequestMessage } from 'openai';
import { getApiKey } from '../utils/auth';
import { OpenAiChatEnum } from '@/constants/model';
import { pushSplitDataBill } from '@/service/events/pushBill';
import { generateVector } from './generateVector';
import { openaiError2 } from '../errorCode';
import { PgClient } from '@/service/pg';
import { ModelSplitDataSchema } from '@/types/mongoSchema';
import { modelServiceToolMap } from '../utils/chat';
import { ChatRoleEnum } from '@/constants/chat';
export async function generateQA(next = false): Promise<any> {
if (process.env.queueTask !== '1') {
@@ -47,11 +46,11 @@ export async function generateQA(next = false): Promise<any> {
// 获取 openapi Key
let userApiKey = '',
systemKey = '';
systemApiKey = '';
try {
const key = await getOpenApiKey(dataItem.userId);
const key = await getApiKey({ model: OpenAiChatEnum.GPT35, userId: dataItem.userId });
userApiKey = key.userApiKey;
systemKey = key.systemKey;
systemApiKey = key.systemApiKey;
} catch (error: any) {
if (error?.code === 501) {
// 余额不够了, 清空该记录
@@ -69,55 +68,44 @@ export async function generateQA(next = false): Promise<any> {
const startTime = Date.now();
// 获取 openai 请求实例
const chatAPI = getOpenAIApi(userApiKey || systemKey);
const systemPrompt: ChatCompletionRequestMessage = {
role: 'system',
content: `你是出题人
${dataItem.prompt || '下面是"一段长文本"'}
从中选出5至20个题目和答案.答案详细.按格式返回: Q1:
A1:
Q2:
A2:
...`
};
// 请求 chatgpt 获取回答
const response = await Promise.allSettled(
textList.map((text) =>
chatAPI
.createChatCompletion(
{
model: OpenAiChatEnum.GPT35,
temperature: 0.8,
n: 1,
messages: [
systemPrompt,
{
role: 'user',
content: text
}
]
},
{
timeout: 180000,
...axiosConfig()
}
)
.then((res) => {
const rawContent = res?.data.choices[0].message?.content || ''; // chatgpt 原本的回复
const result = formatSplitText(res?.data.choices[0].message?.content || ''); // 格式化后的QA对
modelServiceToolMap[OpenAiChatEnum.GPT35]
.chatCompletion({
apiKey: userApiKey || systemApiKey,
temperature: 0.8,
messages: [
{
obj: ChatRoleEnum.System,
value: `你是出题人
${dataItem.prompt || '下面是"一段长文本"'}
从中选出5至20个题目和答案.答案详细.按格式返回: Q1:
A1:
Q2:
A2:
...`
},
{
obj: 'Human',
value: text
}
],
stream: false
})
.then(({ totalTokens, responseText, responseMessages }) => {
const result = formatSplitText(responseText); // 格式化后的QA对
console.log(`split result length: `, result.length);
// 计费
pushSplitDataBill({
isPay: !userApiKey && result.length > 0,
userId: dataItem.userId,
type: 'QA',
text: systemPrompt.content + text + rawContent,
tokenLen: res.data.usage?.total_tokens || 0
textLen: responseMessages.map((item) => item.value).join('').length,
totalTokens
});
return {
rawContent,
rawContent: responseText,
result
};
})

View File

@@ -1,6 +1,8 @@
import { openaiCreateEmbedding, getOpenApiKey } from '../utils/openai';
import { openaiCreateEmbedding } from '../utils/chat/openai';
import { getApiKey } from '../utils/auth';
import { openaiError2 } from '../errorCode';
import { PgClient } from '@/service/pg';
import { embeddingModel } from '@/constants/model';
export async function generateVector(next = false): Promise<any> {
if (process.env.queueTask !== '1') {
@@ -40,11 +42,11 @@ export async function generateVector(next = false): Promise<any> {
dataId = dataItem.id;
// 获取 openapi Key
let userApiKey, systemKey;
let userApiKey, systemApiKey;
try {
const res = await getOpenApiKey(dataItem.userId);
const res = await getApiKey({ model: embeddingModel, userId: dataItem.userId });
userApiKey = res.userApiKey;
systemKey = res.systemKey;
systemApiKey = res.systemApiKey;
} catch (error: any) {
if (error?.code === 501) {
await PgClient.delete('modelData', {
@@ -61,8 +63,8 @@ export async function generateVector(next = false): Promise<any> {
const { vector } = await openaiCreateEmbedding({
text: dataItem.q,
userId: dataItem.userId,
isPay: !userApiKey,
apiKey: userApiKey || systemKey
userApiKey,
systemApiKey
});
// 更新 pg 向量和状态数据

View File

@@ -1,60 +1,54 @@
import { connectToDatabase, Bill, User } from '../mongo';
import { ChatModelMap, OpenAiChatEnum, ChatModelType, embeddingModel } from '@/constants/model';
import { BillTypeEnum } from '@/constants/user';
import { countChatTokens } from '@/utils/tools';
export const pushChatBill = async ({
isPay,
chatModel,
userId,
chatId,
messages
textLen,
tokens
}: {
isPay: boolean;
chatModel: ChatModelType;
userId: string;
chatId?: '' | string;
messages: { role: 'system' | 'user' | 'assistant'; content: string }[];
textLen: number;
tokens: number;
}) => {
console.log(`chat generate success. text len: ${textLen}. token len: ${tokens}. pay:${isPay}`);
if (!isPay) return;
let billId = '';
try {
// 计算 token 数量
const tokens = countChatTokens({ model: chatModel, messages });
const text = messages.map((item) => item.content).join('');
await connectToDatabase();
console.log(
`chat generate success. text len: ${text.length}. token len: ${tokens}. pay:${isPay}`
);
// 计算价格
const unitPrice = ChatModelMap[chatModel]?.price || 5;
const price = unitPrice * tokens;
if (isPay) {
await connectToDatabase();
try {
// 插入 Bill 记录
const res = await Bill.create({
userId,
type: 'chat',
modelName: chatModel,
chatId: chatId ? chatId : undefined,
textLen,
tokenLen: tokens,
price
});
billId = res._id;
// 计算价格
const unitPrice = ChatModelMap[chatModel]?.price || 5;
const price = unitPrice * tokens;
try {
// 插入 Bill 记录
const res = await Bill.create({
userId,
type: 'chat',
modelName: chatModel,
chatId: chatId ? chatId : undefined,
textLen: text.length,
tokenLen: tokens,
price
});
billId = res._id;
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
} catch (error) {
console.log(error);
@@ -64,54 +58,49 @@ export const pushChatBill = async ({
export const pushSplitDataBill = async ({
isPay,
userId,
tokenLen,
text,
totalTokens,
textLen,
type
}: {
isPay: boolean;
userId: string;
tokenLen: number;
text: string;
totalTokens: number;
textLen: number;
type: `${BillTypeEnum}`;
}) => {
await connectToDatabase();
console.log(
`splitData generate success. text len: ${textLen}. token len: ${totalTokens}. pay:${isPay}`
);
if (!isPay) return;
let billId;
try {
console.log(
`splitData generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
);
await connectToDatabase();
if (isPay) {
try {
// 获取模型单价格, 都是用 gpt35 拆分
const unitPrice = ChatModelMap[OpenAiChatEnum.GPT35]?.price || 3;
// 计算价格
const price = unitPrice * tokenLen;
// 获取模型单价格, 都是用 gpt35 拆分
const unitPrice = ChatModelMap[OpenAiChatEnum.GPT35].price || 3;
// 计算价格
const price = unitPrice * totalTokens;
// 插入 Bill 记录
const res = await Bill.create({
userId,
type,
modelName: OpenAiChatEnum.GPT35,
textLen: text.length,
tokenLen,
price
});
billId = res._id;
// 插入 Bill 记录
const res = await Bill.create({
userId,
type,
modelName: OpenAiChatEnum.GPT35,
textLen,
tokenLen: totalTokens,
price
});
billId = res._id;
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
}
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log(error);
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
};
@@ -126,41 +115,40 @@ export const pushGenerateVectorBill = async ({
text: string;
tokenLen: number;
}) => {
await connectToDatabase();
console.log(
`vector generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
);
if (!isPay) return;
let billId;
try {
console.log(
`vector generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
);
await connectToDatabase();
if (isPay) {
try {
const unitPrice = 0.4;
// 计算价格. 至少为1
let price = unitPrice * tokenLen;
price = price > 1 ? price : 1;
try {
const unitPrice = 0.4;
// 计算价格. 至少为1
let price = unitPrice * tokenLen;
price = price > 1 ? price : 1;
// 插入 Bill 记录
const res = await Bill.create({
userId,
type: BillTypeEnum.vector,
modelName: embeddingModel,
textLen: text.length,
tokenLen,
price
});
billId = res._id;
// 插入 Bill 记录
const res = await Bill.create({
userId,
type: BillTypeEnum.vector,
modelName: embeddingModel,
textLen: text.length,
tokenLen,
price
});
billId = res._id;
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
// 账号扣费
await User.findByIdAndUpdate(userId, {
$inc: { balance: -price }
});
} catch (error) {
console.log('创建账单失败:', error);
billId && Bill.findByIdAndDelete(billId);
}
} catch (error) {
console.log(error);

View File

@@ -1,5 +1,6 @@
import { Schema, model, models, Model } from 'mongoose';
import { ChatSchema as ChatType } from '@/types/mongoSchema';
import { ChatRoleMap } from '@/constants/chat';
const ChatSchema = new Schema({
userId: {
@@ -36,7 +37,7 @@ const ChatSchema = new Schema({
obj: {
type: String,
required: true,
enum: ['Human', 'AI', 'SYSTEM']
enum: Object.keys(ChatRoleMap)
},
value: {
type: String,

View File

@@ -1,22 +1,23 @@
import { openaiCreateEmbedding } from '../utils/openai';
import { PgClient } from '@/service/pg';
import { ModelDataStatusEnum, ModelVectorSearchModeEnum, ChatModelMap } from '@/constants/model';
import { ModelSchema } from '@/types/mongoSchema';
import { systemPromptFilter } from '../utils/tools';
import { openaiCreateEmbedding } from '../utils/chat/openai';
import { ChatRoleEnum } from '@/constants/chat';
import { sliceTextByToken } from '@/utils/chat';
/**
* use openai embedding search kb
*/
export const searchKb_openai = async ({
apiKey,
isPay = true,
export const searchKb = async ({
userApiKey,
systemApiKey,
text,
similarity = 0.2,
model,
userId
}: {
apiKey: string;
isPay: boolean;
userApiKey?: string;
systemApiKey: string;
text: string;
model: ModelSchema;
userId: string;
@@ -24,7 +25,7 @@ export const searchKb_openai = async ({
}): Promise<{
code: 200 | 201;
searchPrompt?: {
obj: 'Human' | 'AI' | 'SYSTEM';
obj: `${ChatRoleEnum}`;
value: string;
};
}> => {
@@ -32,8 +33,8 @@ export const searchKb_openai = async ({
// 获取提示词的向量
const { vector: promptVector } = await openaiCreateEmbedding({
isPay,
apiKey,
userApiKey,
systemApiKey,
userId,
text
});
@@ -61,7 +62,7 @@ export const searchKb_openai = async ({
return {
code: 201,
searchPrompt: {
obj: 'AI',
obj: ChatRoleEnum.AI,
value: '对不起,你的问题不在知识库中。'
}
};
@@ -72,7 +73,7 @@ export const searchKb_openai = async ({
code: 200,
searchPrompt: model.chat.systemPrompt
? {
obj: 'SYSTEM',
obj: ChatRoleEnum.System,
value: model.chat.systemPrompt
}
: undefined
@@ -81,16 +82,16 @@ export const searchKb_openai = async ({
// 有匹配情况下system 添加知识库内容。
// 系统提示词过滤,最多 65% tokens
const filterSystemPrompt = systemPromptFilter({
const filterSystemPrompt = sliceTextByToken({
model: model.chat.chatModel,
prompts: systemPrompts,
maxTokens: Math.floor(modelConstantsData.contextMaxToken * 0.65)
text: systemPrompts.join('\n'),
length: Math.floor(modelConstantsData.contextMaxToken * 0.65)
});
return {
code: 200,
searchPrompt: {
obj: 'SYSTEM',
obj: ChatRoleEnum.System,
value: `
${model.chat.systemPrompt}
${

View File

@@ -1,14 +1,18 @@
import { Configuration, OpenAIApi } from 'openai';
import type { NextApiRequest } from 'next';
import jwt from 'jsonwebtoken';
import { Chat, Model, OpenApi, User } from '../mongo';
import type { ModelSchema } from '@/types/mongoSchema';
import { getOpenApiKey } from './openai';
import type { ChatItemSimpleType } from '@/types/chat';
import mongoose from 'mongoose';
import { defaultModel } from '@/constants/model';
import { formatPrice } from '@/utils/user';
import { ERROR_ENUM } from '../errorCode';
import {
ChatModelType,
OpenAiChatEnum,
embeddingModel,
EmbeddingModelType
} from '@/constants/model';
/* 校验 token */
export const authToken = (token?: string): Promise<string> => {
@@ -29,13 +33,63 @@ export const authToken = (token?: string): Promise<string> => {
});
};
export const getOpenAIApi = (apiKey: string) => {
const configuration = new Configuration({
apiKey,
basePath: process.env.OPENAI_BASE_URL
});
/* 获取 api 请求的 key */
export const getApiKey = async ({
model,
userId
}: {
model: ChatModelType | EmbeddingModelType;
userId: string;
}) => {
const user = await User.findById(userId);
if (!user) {
return Promise.reject({
code: 501,
message: '找不到用户'
});
}
return new OpenAIApi(configuration);
const keyMap = {
[OpenAiChatEnum.GPT35]: {
userApiKey: user.openaiKey || '',
systemApiKey: process.env.OPENAIKEY as string
},
[OpenAiChatEnum.GPT4]: {
userApiKey: user.openaiKey || '',
systemApiKey: process.env.OPENAIKEY as string
},
[OpenAiChatEnum.GPT432k]: {
userApiKey: user.openaiKey || '',
systemApiKey: process.env.OPENAIKEY as string
},
[embeddingModel]: {
userApiKey: user.openaiKey || '',
systemApiKey: process.env.OPENAIKEY as string
}
};
// 有自己的key
if (keyMap[model].userApiKey) {
return {
user,
userApiKey: keyMap[model].userApiKey,
systemApiKey: ''
};
}
// 平台账号余额校验
if (formatPrice(user.balance) <= 0) {
return Promise.reject({
code: 501,
message: '账号余额不足'
});
}
return {
user,
userApiKey: '',
systemApiKey: keyMap[model].systemApiKey
};
};
// 模型使用权校验
@@ -122,11 +176,11 @@ export const authChat = async ({
]);
}
// 获取 user 的 apiKey
const { userApiKey, systemKey } = await getOpenApiKey(userId);
const { userApiKey, systemApiKey } = await getApiKey({ model: model.chat.chatModel, userId });
return {
userApiKey,
systemKey,
systemApiKey,
content,
userId,
model,

View File

@@ -0,0 +1,155 @@
import { ChatItemSimpleType } from '@/types/chat';
import { modelToolMap } from '@/utils/chat';
import type { ChatModelType } from '@/constants/model';
import { ChatRoleEnum, SYSTEM_PROMPT_PREFIX } from '@/constants/chat';
import { OpenAiChatEnum } from '@/constants/model';
import { chatResponse, openAiStreamResponse } from './openai';
import type { NextApiResponse } from 'next';
import type { PassThrough } from 'stream';
export type ChatCompletionType = {
apiKey: string;
temperature: number;
messages: ChatItemSimpleType[];
stream: boolean;
};
export type StreamResponseType = {
stream: PassThrough;
chatResponse: any;
prompts: ChatItemSimpleType[];
};
export const modelServiceToolMap = {
[OpenAiChatEnum.GPT35]: {
chatCompletion: (data: ChatCompletionType) =>
chatResponse({ model: OpenAiChatEnum.GPT35, ...data }),
streamResponse: (data: StreamResponseType) =>
openAiStreamResponse({
model: OpenAiChatEnum.GPT35,
...data
})
},
[OpenAiChatEnum.GPT4]: {
chatCompletion: (data: ChatCompletionType) =>
chatResponse({ model: OpenAiChatEnum.GPT4, ...data }),
streamResponse: (data: StreamResponseType) =>
openAiStreamResponse({
model: OpenAiChatEnum.GPT4,
...data
})
},
[OpenAiChatEnum.GPT432k]: {
chatCompletion: (data: ChatCompletionType) =>
chatResponse({ model: OpenAiChatEnum.GPT432k, ...data }),
streamResponse: (data: StreamResponseType) =>
openAiStreamResponse({
model: OpenAiChatEnum.GPT432k,
...data
})
}
};
/* delete invalid symbol */
const simplifyStr = (str: string) =>
str
.replace(/\n+/g, '\n') // 连续空行
.replace(/[^\S\r\n]+/g, ' ') // 连续空白内容
.trim();
/* 聊天上下文 tokens 截断 */
export const ChatContextFilter = ({
model,
prompts,
maxTokens
}: {
model: ChatModelType;
prompts: ChatItemSimpleType[];
maxTokens: number;
}) => {
let rawTextLen = 0;
const formatPrompts = prompts.map<ChatItemSimpleType>((item) => {
const val = simplifyStr(item.value);
rawTextLen += val.length;
return {
obj: item.obj,
value: val
};
});
// 长度太小时,不需要进行 token 截断
if (formatPrompts.length <= 2 || rawTextLen < maxTokens * 0.5) {
return formatPrompts;
}
// 根据 tokens 截断内容
const chats: ChatItemSimpleType[] = [];
let systemPrompt: ChatItemSimpleType | null = null;
// System 词保留
if (formatPrompts[0].obj === ChatRoleEnum.System) {
const prompt = formatPrompts.shift();
if (prompt) {
systemPrompt = prompt;
}
}
let messages: ChatItemSimpleType[] = [];
// 从后往前截取对话内容
for (let i = formatPrompts.length - 1; i >= 0; i--) {
chats.unshift(formatPrompts[i]);
messages = systemPrompt ? [systemPrompt, ...chats] : chats;
const tokens = modelToolMap[model].countTokens({
messages
});
/* 整体 tokens 超出范围 */
if (tokens >= maxTokens) {
return systemPrompt ? [systemPrompt, ...chats.slice(1)] : chats.slice(1);
}
}
return messages;
};
/* stream response */
export const resStreamResponse = async ({
model,
res,
stream,
chatResponse,
systemPrompt,
prompts
}: StreamResponseType & {
model: ChatModelType;
res: NextApiResponse;
systemPrompt?: string;
}) => {
// 创建响应流
res.setHeader('Content-Type', 'text/event-stream;charset-utf-8');
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('X-Accel-Buffering', 'no');
res.setHeader('Cache-Control', 'no-cache, no-transform');
stream.pipe(res);
const { responseContent, totalTokens, finishMessages } = await modelServiceToolMap[
model
].streamResponse({
chatResponse,
stream,
prompts
});
// push system prompt
!stream.destroyed &&
systemPrompt &&
stream.push(`${SYSTEM_PROMPT_PREFIX}${systemPrompt.replace(/\n/g, '<br/>')}`);
// close stream
!stream.destroyed && stream.push(null);
stream.destroy();
return { responseContent, totalTokens, finishMessages };
};

View File

@@ -0,0 +1,174 @@
import { Configuration, OpenAIApi } from 'openai';
import { createParser, ParsedEvent, ReconnectInterval } from 'eventsource-parser';
import { axiosConfig } from '../tools';
import { ChatModelMap, embeddingModel, OpenAiChatEnum } from '@/constants/model';
import { pushGenerateVectorBill } from '../../events/pushBill';
import { adaptChatItem_openAI } from '@/utils/chat/openai';
import { modelToolMap } from '@/utils/chat';
import { ChatCompletionType, ChatContextFilter, StreamResponseType } from './index';
import { ChatRoleEnum } from '@/constants/chat';
export const getOpenAIApi = (apiKey: string) => {
const configuration = new Configuration({
apiKey,
basePath: process.env.OPENAI_BASE_URL
});
return new OpenAIApi(configuration);
};
/* 获取向量 */
export const openaiCreateEmbedding = async ({
userApiKey,
systemApiKey,
userId,
text
}: {
userApiKey?: string;
systemApiKey: string;
userId: string;
text: string;
}) => {
// 获取 chatAPI
const chatAPI = getOpenAIApi(userApiKey || systemApiKey);
// 把输入的内容转成向量
const res = await chatAPI
.createEmbedding(
{
model: embeddingModel,
input: text
},
{
timeout: 60000,
...axiosConfig()
}
)
.then((res) => ({
tokenLen: res.data.usage.total_tokens || 0,
vector: res.data.data?.[0]?.embedding || []
}));
pushGenerateVectorBill({
isPay: !userApiKey,
userId,
text,
tokenLen: res.tokenLen
});
return {
vector: res.vector,
chatAPI
};
};
/* 模型对话 */
export const chatResponse = async ({
model,
apiKey,
temperature,
messages,
stream
}: ChatCompletionType & { model: `${OpenAiChatEnum}` }) => {
const filterMessages = ChatContextFilter({
model,
prompts: messages,
maxTokens: Math.ceil(ChatModelMap[model].contextMaxToken * 0.9)
});
const adaptMessages = adaptChatItem_openAI({ messages: filterMessages });
const chatAPI = getOpenAIApi(apiKey);
const response = await chatAPI.createChatCompletion(
{
model,
temperature: Number(temperature) || 0,
messages: adaptMessages,
frequency_penalty: 0.5, // 越大,重复内容越少
presence_penalty: -0.5, // 越大,越容易出现新内容
stream,
stop: ['.!?。']
},
{
timeout: stream ? 40000 : 240000,
responseType: stream ? 'stream' : 'json',
...axiosConfig()
}
);
let responseText = '';
let totalTokens = 0;
// adapt data
if (!stream) {
responseText = response.data.choices[0].message?.content || '';
totalTokens = response.data.usage?.total_tokens || 0;
}
return {
streamResponse: response,
responseMessages: filterMessages.concat({ obj: 'AI', value: responseText }),
responseText,
totalTokens
};
};
/* openai stream response */
export const openAiStreamResponse = async ({
model,
stream,
chatResponse,
prompts
}: StreamResponseType & {
model: `${OpenAiChatEnum}`;
}) => {
try {
let responseContent = '';
const onParse = async (event: ParsedEvent | ReconnectInterval) => {
if (event.type !== 'event') return;
const data = event.data;
if (data === '[DONE]') return;
try {
const json = JSON.parse(data);
const content: string = json?.choices?.[0].delta.content || '';
responseContent += content;
!stream.destroyed && content && stream.push(content.replace(/\n/g, '<br/>'));
} catch (error) {
error;
}
};
try {
const decoder = new TextDecoder();
const parser = createParser(onParse);
for await (const chunk of chatResponse.data as any) {
if (stream.destroyed) {
// 流被中断了,直接忽略后面的内容
break;
}
parser.feed(decoder.decode(chunk, { stream: true }));
}
} catch (error) {
console.log('pipe error', error);
}
// count tokens
const finishMessages = prompts.concat({
obj: ChatRoleEnum.AI,
value: responseContent
});
const totalTokens = modelToolMap[model].countTokens({
messages: finishMessages
});
return {
responseContent,
totalTokens,
finishMessages
};
} catch (error) {
return Promise.reject(error);
}
};

View File

@@ -1,179 +0,0 @@
import type { NextApiResponse } from 'next';
import type { PassThrough } from 'stream';
import { createParser, ParsedEvent, ReconnectInterval } from 'eventsource-parser';
import { getOpenAIApi } from '@/service/utils/auth';
import { axiosConfig } from './tools';
import { User } from '../models/user';
import { formatPrice } from '@/utils/user';
import { embeddingModel } from '@/constants/model';
import { pushGenerateVectorBill } from '../events/pushBill';
import { SYSTEM_PROMPT_PREFIX } from '@/constants/chat';
/* 获取用户 api 的 openai 信息 */
export const getUserApiOpenai = async (userId: string) => {
const user = await User.findById(userId);
const userApiKey = user?.openaiKey;
if (!userApiKey) {
return Promise.reject('缺少ApiKey, 无法请求');
}
return {
user,
openai: getOpenAIApi(userApiKey),
apiKey: userApiKey
};
};
/* 获取 open api key如果用户没有自己的key就用平台的用平台记得加账单 */
export const getOpenApiKey = async (userId: string) => {
const user = await User.findById(userId);
if (!user) {
return Promise.reject({
code: 501,
message: '找不到用户'
});
}
const userApiKey = user?.openaiKey;
// 有自己的key
if (userApiKey) {
return {
user,
userApiKey,
systemKey: ''
};
}
// 平台账号余额校验
if (formatPrice(user.balance) <= 0) {
return Promise.reject({
code: 501,
message: '账号余额不足'
});
}
return {
user,
userApiKey: '',
systemKey: process.env.OPENAIKEY as string
};
};
/* 获取向量 */
export const openaiCreateEmbedding = async ({
isPay,
userId,
apiKey,
text
}: {
isPay: boolean;
userId: string;
apiKey: string;
text: string;
}) => {
// 获取 chatAPI
const chatAPI = getOpenAIApi(apiKey);
// 把输入的内容转成向量
const res = await chatAPI
.createEmbedding(
{
model: embeddingModel,
input: text
},
{
timeout: 60000,
...axiosConfig()
}
)
.then((res) => ({
tokenLen: res.data.usage.total_tokens || 0,
vector: res.data.data?.[0]?.embedding || []
}));
pushGenerateVectorBill({
isPay,
userId,
text,
tokenLen: res.tokenLen
});
return {
vector: res.vector,
chatAPI
};
};
/* gpt35 响应 */
export const gpt35StreamResponse = ({
res,
stream,
chatResponse,
systemPrompt = ''
}: {
res: NextApiResponse;
stream: PassThrough;
chatResponse: any;
systemPrompt?: string;
}) =>
new Promise<{ responseContent: string }>(async (resolve, reject) => {
try {
// 创建响应流
res.setHeader('Content-Type', 'text/event-stream;charset-utf-8');
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('X-Accel-Buffering', 'no');
res.setHeader('Cache-Control', 'no-cache, no-transform');
stream.pipe(res);
let responseContent = '';
const onParse = async (event: ParsedEvent | ReconnectInterval) => {
if (event.type !== 'event') return;
const data = event.data;
if (data === '[DONE]') return;
try {
const json = JSON.parse(data);
const content: string = json?.choices?.[0].delta.content || '';
responseContent += content;
if (!stream.destroyed && content) {
stream.push(content.replace(/\n/g, '<br/>'));
}
} catch (error) {
error;
}
};
try {
const decoder = new TextDecoder();
const parser = createParser(onParse);
for await (const chunk of chatResponse.data as any) {
if (stream.destroyed) {
// 流被中断了,直接忽略后面的内容
break;
}
parser.feed(decoder.decode(chunk, { stream: true }));
}
} catch (error) {
console.log('pipe error', error);
}
// push system prompt
!stream.destroyed &&
systemPrompt &&
stream.push(`${SYSTEM_PROMPT_PREFIX}${systemPrompt.replace(/\n/g, '<br/>')}`);
// close stream
!stream.destroyed && stream.push(null);
stream.destroy();
resolve({
responseContent
});
} catch (error) {
reject(error);
}
});

View File

@@ -1,9 +1,5 @@
import crypto from 'crypto';
import jwt from 'jsonwebtoken';
import { ChatItemSimpleType } from '@/types/chat';
import { countChatTokens, sliceTextByToken } from '@/utils/tools';
import { ChatCompletionRequestMessageRoleEnum, ChatCompletionRequestMessage } from 'openai';
import type { ChatModelType } from '@/constants/model';
/* 密码加密 */
export const hashPassword = (psw: string) => {
@@ -30,92 +26,3 @@ export const axiosConfig = () => ({
auth: process.env.OPENAI_BASE_URL_AUTH || ''
}
});
/* delete invalid symbol */
const simplifyStr = (str: string) =>
str
.replace(/\n+/g, '\n') // 连续空行
.replace(/[^\S\r\n]+/g, ' ') // 连续空白内容
.trim();
/* 聊天内容 tokens 截断 */
export const openaiChatFilter = ({
model,
prompts,
maxTokens
}: {
model: ChatModelType;
prompts: ChatItemSimpleType[];
maxTokens: number;
}) => {
// role map
const map = {
Human: ChatCompletionRequestMessageRoleEnum.User,
AI: ChatCompletionRequestMessageRoleEnum.Assistant,
SYSTEM: ChatCompletionRequestMessageRoleEnum.System
};
let rawTextLen = 0;
const formatPrompts = prompts.map((item) => {
const val = simplifyStr(item.value);
rawTextLen += val.length;
return {
role: map[item.obj],
content: val
};
});
// 长度太小时,不需要进行 token 截断
if (rawTextLen < maxTokens * 0.5) {
return formatPrompts;
}
// 根据 tokens 截断内容
const chats: ChatCompletionRequestMessage[] = [];
let systemPrompt: ChatCompletionRequestMessage | null = null;
// System 词保留
if (formatPrompts[0]?.role === 'system') {
systemPrompt = formatPrompts.shift() as ChatCompletionRequestMessage;
}
let messages: { role: ChatCompletionRequestMessageRoleEnum; content: string }[] = [];
// 从后往前截取对话内容
for (let i = formatPrompts.length - 1; i >= 0; i--) {
chats.unshift(formatPrompts[i]);
messages = systemPrompt ? [systemPrompt, ...chats] : chats;
const tokens = countChatTokens({
model,
messages
});
/* 整体 tokens 超出范围 */
if (tokens >= maxTokens) {
return systemPrompt ? [systemPrompt, ...chats.slice(1)] : chats.slice(1);
}
}
return messages;
};
/* system 内容截断. 相似度从高到低 */
export const systemPromptFilter = ({
model,
prompts,
maxTokens
}: {
model: 'gpt-4' | 'gpt-4-32k' | 'gpt-3.5-turbo';
prompts: string[];
maxTokens: number;
}) => {
const systemPrompt = prompts.join('\n');
return sliceTextByToken({
model,
text: systemPrompt,
length: maxTokens
});
};