llama-gpt/pages/api/chat.ts

61 lines
1.7 KiB
TypeScript

import { ChatBody, Message } from '@/types/chat';
import { DEFAULT_SYSTEM_PROMPT } from '@/utils/app/const';
import { OpenAIError, OpenAIStream } from '@/utils/server';
import tiktokenModel from '@dqbd/tiktoken/encoders/cl100k_base.json';
import { init, Tiktoken } from '@dqbd/tiktoken/lite/init';
// @ts-expect-error
import wasm from '../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module';
export const config = {
runtime: 'edge',
};
const handler = async (req: Request): Promise<Response> => {
try {
const { model, messages, key, prompt } = (await req.json()) as ChatBody;
await init((imports) => WebAssembly.instantiate(wasm, imports));
const encoding = new Tiktoken(
tiktokenModel.bpe_ranks,
tiktokenModel.special_tokens,
tiktokenModel.pat_str,
);
let promptToSend = prompt;
if (!promptToSend) {
promptToSend = DEFAULT_SYSTEM_PROMPT;
}
const prompt_tokens = encoding.encode(promptToSend);
let tokenCount = prompt_tokens.length;
let messagesToSend: Message[] = [];
for (let i = messages.length - 1; i >= 0; i--) {
const message = messages[i];
const tokens = encoding.encode(message.content);
if (tokenCount + tokens.length + 1000 > model.tokenLimit) {
break;
}
tokenCount += tokens.length;
messagesToSend = [message, ...messagesToSend];
}
encoding.free();
const stream = await OpenAIStream(model, promptToSend, key, messagesToSend);
return new Response(stream);
} catch (error) {
console.error(error);
if (error instanceof OpenAIError) {
return new Response('Error', { status: 500, statusText: error.message });
} else {
return new Response('Error', { status: 500 });
}
}
};
export default handler;