feat: chat content use tiktoken count

This commit is contained in:
archer
2023-04-24 16:46:39 +08:00
parent adbaa8b37b
commit 1f112f7715
23 changed files with 182 additions and 836 deletions

View File

@@ -4,7 +4,7 @@ import { httpsAgent } from '@/service/utils/tools';
import { getOpenApiKey } from '../utils/openai';
import type { ChatCompletionRequestMessage } from 'openai';
import { DataItemSchema } from '@/types/mongoSchema';
import { ChatModelNameEnum } from '@/constants/model';
import { ChatModelEnum } from '@/constants/model';
import { pushSplitDataBill } from '@/service/events/pushBill';
export async function generateAbstract(next = false): Promise<any> {
@@ -68,7 +68,7 @@ export async function generateAbstract(next = false): Promise<any> {
// 请求 chatgpt 获取摘要
const abstractResponse = await chatAPI.createChatCompletion(
{
model: ChatModelNameEnum.GPT35,
model: ChatModelEnum.GPT35,
temperature: 0.8,
n: 1,
messages: [

View File

@@ -3,7 +3,7 @@ import { getOpenAIApi } from '@/service/utils/auth';
import { httpsAgent } from '@/service/utils/tools';
import { getOpenApiKey } from '../utils/openai';
import type { ChatCompletionRequestMessage } from 'openai';
import { ChatModelNameEnum } from '@/constants/model';
import { ChatModelEnum } from '@/constants/model';
import { pushSplitDataBill } from '@/service/events/pushBill';
import { generateVector } from './generateVector';
import { openaiError2 } from '../errorCode';
@@ -84,7 +84,7 @@ A2:
chatAPI
.createChatCompletion(
{
model: ChatModelNameEnum.GPT35,
model: ChatModelEnum.GPT35,
temperature: 0.8,
n: 1,
messages: [

View File

@@ -1,27 +1,34 @@
import { connectToDatabase, Bill, User } from '../mongo';
import { modelList, ChatModelNameEnum } from '@/constants/model';
import { encode } from 'gpt-token-utils';
import {
modelList,
ChatModelEnum,
ModelNameEnum,
Model2ChatModelMap,
embeddingModel
} from '@/constants/model';
import { BillTypeEnum } from '@/constants/user';
import type { DataType } from '@/types/data';
import { countChatTokens } from '@/utils/tools';
export const pushChatBill = async ({
isPay,
modelName,
userId,
chatId,
text
messages
}: {
isPay: boolean;
modelName: string;
modelName: `${ModelNameEnum}`;
userId: string;
chatId?: '' | string;
text: string;
messages: { role: 'system' | 'user' | 'assistant'; content: string }[];
}) => {
let billId;
let billId = '';
try {
// 计算 token 数量
const tokens = Math.floor(encode(text).length * 0.75);
const tokens = countChatTokens({ model: Model2ChatModelMap[modelName] as any, messages });
const text = messages.map((item) => item.content).join('');
console.log(
`chat generate success. text len: ${text.length}. token len: ${tokens}. pay:${isPay}`
@@ -88,7 +95,7 @@ export const pushSplitDataBill = async ({
if (isPay) {
try {
// 获取模型单价格, 都是用 gpt35 拆分
const modelItem = modelList.find((item) => item.model === ChatModelNameEnum.GPT35);
const modelItem = modelList.find((item) => item.model === ChatModelEnum.GPT35);
const unitPrice = modelItem?.price || 3;
// 计算价格
const price = unitPrice * tokenLen;
@@ -97,7 +104,7 @@ export const pushSplitDataBill = async ({
const res = await Bill.create({
userId,
type,
modelName: ChatModelNameEnum.GPT35,
modelName: ChatModelEnum.GPT35,
textLen: text.length,
tokenLen,
price
@@ -149,7 +156,7 @@ export const pushGenerateVectorBill = async ({
const res = await Bill.create({
userId,
type: BillTypeEnum.vector,
modelName: ChatModelNameEnum.VECTOR,
modelName: embeddingModel,
textLen: text.length,
tokenLen,
price