73 lines
1.8 KiB
TypeScript
73 lines
1.8 KiB
TypeScript
import { ChatItemType } from '@/types/chat';
|
|
import { modelToolMap } from '@/utils/plugin';
|
|
import { ChatRoleEnum } from '@/constants/chat';
|
|
import { OpenAiChatEnum } from '@/constants/model';
|
|
import type { NextApiResponse } from 'next';
|
|
|
|
export type ChatCompletionResponseType = {
|
|
streamResponse: any;
|
|
responseMessages: ChatItemType[];
|
|
responseText: string;
|
|
totalTokens: number;
|
|
};
|
|
export type StreamResponseType = {
|
|
chatResponse: any;
|
|
prompts: ChatItemType[];
|
|
res: NextApiResponse;
|
|
model: `${OpenAiChatEnum}`;
|
|
[key: string]: any;
|
|
};
|
|
|
|
/* slice chat context by tokens */
|
|
export const ChatContextFilter = ({
|
|
model,
|
|
prompts = [],
|
|
maxTokens
|
|
}: {
|
|
model: string;
|
|
prompts: ChatItemType[];
|
|
maxTokens: number;
|
|
}) => {
|
|
if (!Array.isArray(prompts)) {
|
|
return [];
|
|
}
|
|
const rawTextLen = prompts.reduce((sum, item) => sum + item.value.length, 0);
|
|
|
|
// If the text length is less than half of the maximum token, no calculation is required
|
|
if (rawTextLen < maxTokens * 0.5) {
|
|
return prompts;
|
|
}
|
|
|
|
// filter startWith system prompt
|
|
const chatStartIndex = prompts.findIndex((item) => item.obj !== ChatRoleEnum.System);
|
|
const systemPrompts: ChatItemType[] = prompts.slice(0, chatStartIndex);
|
|
const chatPrompts: ChatItemType[] = prompts.slice(chatStartIndex);
|
|
|
|
// reduce token of systemPrompt
|
|
maxTokens -= modelToolMap.countTokens({
|
|
model,
|
|
messages: systemPrompts
|
|
});
|
|
|
|
// 根据 tokens 截断内容
|
|
const chats: ChatItemType[] = [];
|
|
|
|
// 从后往前截取对话内容
|
|
for (let i = chatPrompts.length - 1; i >= 0; i--) {
|
|
chats.unshift(chatPrompts[i]);
|
|
|
|
const tokens = modelToolMap.countTokens({
|
|
model,
|
|
messages: chats
|
|
});
|
|
|
|
/* 整体 tokens 超出范围, system必须保留 */
|
|
if (tokens >= maxTokens) {
|
|
chats.shift();
|
|
break;
|
|
}
|
|
}
|
|
|
|
return [...systemPrompts, ...chats];
|
|
};
|