From aa7292950ba50c9d9169a96547d0d36f082be4ce Mon Sep 17 00:00:00 2001
From: xiang <1984871009@qq.com>
Date: Tue, 13 Jun 2023 10:10:29 +0800
Subject: [PATCH] fix: some error

---
 index.ts           |  1 +
 model/base.ts      | 33 ++++++++++++++++++---------------
 model/you/index.ts |  1 +
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/index.ts b/index.ts
index 57913cc..b08df07 100644
--- a/index.ts
+++ b/index.ts
@@ -57,6 +57,7 @@ router.get('/ask/stream', async (ctx) => {
         "Connection": "keep-alive",
     });
     const es = new EventStream();
+    ctx.body = es.stream();
     if (!prompt) {
         es.write(Event.error, {error: 'need prompt in query'})
         return;
diff --git a/model/base.ts b/model/base.ts
index 8d49b4d..870fc13 100644
--- a/model/base.ts
+++ b/model/base.ts
@@ -23,23 +23,26 @@ export interface ChatRequest {
     model: ModelType;
 }
 
-export function PromptToString(prompt: string | Message[], limit: number): string {
-    if (typeof prompt === "string") {
+export function PromptToString(prompt: string, limit: number): string {
+    try {
+        const messages: Message[] = JSON.parse(prompt);
+        let result: Message[] = [];
+        let tokenSize = 0;
+        for (let i = messages.length - 1; i >= 0; i--) {
+            const item = messages[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')}\nassistant: `;
+    } catch (e) {
         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 {
diff --git a/model/you/index.ts b/model/you/index.ts
index 9d821ba..4557320 100644
--- a/model/you/index.ts
+++ b/model/you/index.ts
@@ -124,6 +124,7 @@ export class You extends Chat {
                     break;
                 case 'done':
                     stream.write(Event.done, {content: 'done'})
+                    stream.stream().end();
                     return;
                 default:
                     return;