feat: Uniform Request and Response Parameters

shishehgar · Jun 13, 2023 · b4c8b52 · b4c8b52
1 parent 749c9b1
commit b4c8b52
Show file tree

Hide file tree

Showing 11 changed files with 265 additions and 305 deletions.
diff --git a/index.ts b/index.ts
@@ -1,8 +1,10 @@
 import Koa, {Context, Next} from 'koa';
 import Router from 'koa-router'
 import bodyParser from 'koa-bodyparser';
-import {ChatModelFactory, Model} from "./model";
+import {ChatModelFactory, Site} from "./model";
 import dotenv from 'dotenv';
+import {ChatRequest, ChatResponse, ModelType, PromptToString} from "./model/base";
+import {Event, EventStream} from "./utils";
 
 dotenv.config();
 
@@ -21,44 +23,56 @@ app.use(errorHandler);
 app.use(bodyParser());
 const chatModel = new ChatModelFactory();
 
-interface AskReq {
-    prompt: string;
-    model: Model;
+interface AskReq extends ChatRequest {
+    site: Site;
+}
+
+interface AskRes extends ChatResponse {
 }
 
 router.get('/ask', async (ctx) => {
-    const {prompt, model = Model.Mcbbs, ...options} = ctx.query as unknown as AskReq;
+    const {prompt, model = ModelType.GPT3p5, site = Site.You} = ctx.query as unknown as AskReq;
     if (!prompt) {
-        ctx.body = 'please input prompt';
+        ctx.body = {error: `need prompt in query`} as AskRes;
         return;
     }
-    const chat = chatModel.get(model);
+    const chat = chatModel.get(site);
     if (!chat) {
-        ctx.body = 'Unsupported  model';
+        ctx.body = {error: `not support site: ${site} `} as AskRes;
+        return;
+    }
+    const tokenLimit = chat.support(model);
+    if (!tokenLimit) {
+        ctx.body = {error: `${site} not support model ${model}`} as AskRes;
         return;
     }
-    const res = await chat.ask({prompt: prompt as string, options});
-    ctx.body = res.text;
+    ctx.body = await chat.ask({prompt: PromptToString(prompt, tokenLimit), model});
 });
 
 router.get('/ask/stream', async (ctx) => {
-    const {prompt, model = Model.Mcbbs, ...options} = ctx.query as unknown as AskReq;
+    const {prompt, model = ModelType.GPT3p5, site = Site.You} = ctx.query as unknown as AskReq;
+    ctx.set({
+        "Content-Type": "text/event-stream;charset=utf-8",
+        "Cache-Control": "no-cache",
+        "Connection": "keep-alive",
+    });
+    const es = new EventStream();
     if (!prompt) {
-        ctx.body = 'please input prompt';
+        es.write(Event.error, {error: 'need prompt in query'})
         return;
     }
-    const chat = chatModel.get(model);
+    const chat = chatModel.get(site);
     if (!chat) {
-        ctx.body = 'Unsupported  model';
+        es.write(Event.error, {error: `not support site: ${site} `})
         return;
     }
-    ctx.set({
-        "Content-Type": "text/event-stream;charset=utf-8",
-        "Cache-Control": "no-cache",
-        "Connection": "keep-alive",
-    });
-    const res = await chat.askStream({prompt: prompt as string, options});
-    ctx.body = res?.text;
+    const tokenLimit = chat.support(model);
+    if (!tokenLimit) {
+        es.write(Event.error, {error: `${site} not support model ${model}`})
+        return;
+    }
+    await chat.askStream({prompt: PromptToString(prompt, tokenLimit), model}, es);
+    ctx.body = es.stream();
 })
 
 app.use(router.routes());

diff --git a/model/aidream/index.ts b/model/aidream/index.ts
diff --git a/model/base.ts b/model/base.ts
@@ -1,21 +1,45 @@
-import {Stream} from "stream";
+import {EventStream, getTokenSize} from "../utils";
 
 export interface ChatOptions {
 }
 
-export interface Response {
-    text: string | null;
-    other?: any;
+export interface ChatResponse {
+    content?: string;
+    error?: string;
 }
 
-export interface ResponseStream {
-    text: Stream;
-    other?: any;
+export type Message = {
+    role: string;
+    content: string;
 }
 
-export interface Request {
+export enum ModelType {
+    GPT3p5 = 'gpt3.5',
+    GPT4 = 'gpt4',
+}
+
+export interface ChatRequest {
     prompt: string;
-    options?: any;
+    model: ModelType;
+}
+
+export function PromptToString(prompt: string | Message[], limit: number): string {
+    if (typeof prompt === "string") {
+        return prompt;
+    }
+    let result: Message[] = [];
+
+    let tokenSize = 0;
+    for (let i = prompt.length - 1; i >= 0; i--) {
+        const item = prompt[i];
+        const {role, content} = item;
+        tokenSize += getTokenSize(content);
+        if (tokenSize > limit) {
+            break;
+        }
+        result.push(item);
+    }
+    return result.reverse().map(item => `${item.role}: ${item.content}`).join('\n');
 }
 
 export abstract class Chat {
@@ -25,7 +49,9 @@ export abstract class Chat {
         this.options = options;
     }
 
-    public abstract ask(req: Request): Promise<Response>
+    public abstract support(model: ModelType): number
+
+    public abstract ask(req: ChatRequest): Promise<ChatResponse>
 
-    public abstract askStream(req: Request): Promise<ResponseStream>
+    public abstract askStream(req: ChatRequest, stream: EventStream): Promise<void>
 }