2023-05-07 01:30:22 +09:00
|
|
|
import Koa, {Context, Next} from 'koa';
|
2023-05-04 22:32:21 +09:00
|
|
|
import Router from 'koa-router'
|
|
|
|
import bodyParser from 'koa-bodyparser';
|
2023-05-05 10:57:42 +09:00
|
|
|
import {ChatModelFactory, Model} from "./model";
|
2023-05-11 20:32:10 +09:00
|
|
|
import dotenv from 'dotenv';
|
2023-05-12 18:16:25 +09:00
|
|
|
import {freeBrowserPool} from "./pool/puppeteer";
|
2023-05-11 20:32:10 +09:00
|
|
|
|
|
|
|
dotenv.config();
|
2023-05-04 22:32:21 +09:00
|
|
|
|
|
|
|
const app = new Koa();
|
|
|
|
const router = new Router();
|
2023-05-07 01:30:22 +09:00
|
|
|
const errorHandler = async (ctx: Context, next: Next) => {
|
|
|
|
try {
|
|
|
|
await next();
|
2023-05-10 00:01:40 +09:00
|
|
|
} catch (err: any) {
|
2023-05-07 21:20:27 +09:00
|
|
|
console.error(err);
|
|
|
|
ctx.body = JSON.stringify(err);
|
|
|
|
ctx.res.end();
|
2023-05-07 01:30:22 +09:00
|
|
|
}
|
|
|
|
};
|
|
|
|
app.use(errorHandler);
|
2023-05-04 22:32:21 +09:00
|
|
|
app.use(bodyParser());
|
2023-05-10 00:01:40 +09:00
|
|
|
const chatModel = new ChatModelFactory();
|
2023-05-04 22:32:21 +09:00
|
|
|
|
|
|
|
interface AskReq {
|
|
|
|
prompt: string;
|
2023-05-05 10:57:42 +09:00
|
|
|
model: Model;
|
2023-05-04 22:32:21 +09:00
|
|
|
}
|
|
|
|
|
|
|
|
router.get('/ask', async (ctx) => {
|
2023-05-07 01:30:22 +09:00
|
|
|
const {prompt, model = Model.Forefront, ...options} = ctx.query as unknown as AskReq;
|
2023-05-04 22:32:21 +09:00
|
|
|
if (!prompt) {
|
|
|
|
ctx.body = 'please input prompt';
|
|
|
|
return;
|
|
|
|
}
|
2023-05-05 10:57:42 +09:00
|
|
|
const chat = chatModel.get(model);
|
|
|
|
if (!chat) {
|
|
|
|
ctx.body = 'Unsupported model';
|
|
|
|
return;
|
|
|
|
}
|
2023-05-07 01:30:22 +09:00
|
|
|
const res = await chat.ask({prompt: prompt as string, options});
|
2023-05-11 11:39:59 +09:00
|
|
|
ctx.body = res;
|
2023-05-04 22:32:21 +09:00
|
|
|
});
|
|
|
|
|
2023-05-04 23:01:55 +09:00
|
|
|
router.get('/ask/stream', async (ctx) => {
|
2023-05-07 01:30:22 +09:00
|
|
|
const {prompt, model = Model.Forefront, ...options} = ctx.query as unknown as AskReq;
|
2023-05-04 22:32:21 +09:00
|
|
|
if (!prompt) {
|
|
|
|
ctx.body = 'please input prompt';
|
|
|
|
return;
|
|
|
|
}
|
2023-05-05 10:57:42 +09:00
|
|
|
const chat = chatModel.get(model);
|
|
|
|
if (!chat) {
|
|
|
|
ctx.body = 'Unsupported model';
|
|
|
|
return;
|
|
|
|
}
|
2023-05-04 22:32:21 +09:00
|
|
|
ctx.set({
|
2023-05-07 02:10:52 +09:00
|
|
|
"Content-Type": "text/event-stream;charset=utf-8",
|
2023-05-04 22:32:21 +09:00
|
|
|
"Cache-Control": "no-cache",
|
|
|
|
"Connection": "keep-alive",
|
2023-05-05 10:57:42 +09:00
|
|
|
});
|
2023-05-07 01:30:22 +09:00
|
|
|
const res = await chat.askStream({prompt: prompt as string, options});
|
2023-05-07 21:20:27 +09:00
|
|
|
ctx.body = res?.text;
|
2023-05-04 22:32:21 +09:00
|
|
|
})
|
|
|
|
|
|
|
|
app.use(router.routes());
|
|
|
|
|
2023-05-12 18:16:25 +09:00
|
|
|
(async () => {
|
|
|
|
await freeBrowserPool.init(1, process.env.DEBUG === '1');
|
|
|
|
const server = app.listen(3000, () => {
|
|
|
|
console.log("Now listening: 127.0.0.1:3000");
|
2023-05-11 20:28:15 +09:00
|
|
|
});
|
2023-05-12 18:16:25 +09:00
|
|
|
process.on('SIGINT', () => {
|
|
|
|
server.close(() => {
|
|
|
|
process.exit(0);
|
|
|
|
});
|
|
|
|
});
|
|
|
|
})()
|
|
|
|
|