From f15f693a1c0a3ccffaa4c8c97485e382407de783 Mon Sep 17 00:00:00 2001
From: zyh <zyh320888@vip.qq.com>
Date: Tue, 22 Oct 2024 01:24:06 +0000
Subject: [PATCH] =?UTF-8?q?feat(api):=20=E9=87=8D=E6=9E=84API=E4=BB=A5?=
 =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=A8=A1=E6=8B=9F=E6=95=B0=E6=8D=AE=E5=92=8C?=
 =?UTF-8?q?AI=E6=9C=8D=E5=8A=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/routes/api.enhancer.ts | 111 ++++++++++++++++++++-----------------
 1 file changed, 60 insertions(+), 51 deletions(-)

diff --git a/app/routes/api.enhancer.ts b/app/routes/api.enhancer.ts
index daedbb8..d725b24 100644
--- a/app/routes/api.enhancer.ts
+++ b/app/routes/api.enhancer.ts
@@ -24,59 +24,68 @@ function* mockDataGenerator(message: string) {
 
 // 环境变量或配置来控制是否使用模拟数据
 // const USE_MOCK_DATA = process.env.USE_MOCK_DATA === 'true';
-const USE_MOCK_DATA = true;
+const USE_MOCK_DATA = false;
 
-export async function action({ context, request }: ActionFunctionArgs) {
-  const { message } = await request.json<{ message: string }>();
+export async function action({ request, context }: ActionFunctionArgs) {
+  const { message } = (await request.json()) as { message: string };
 
-  try {
-    if (USE_MOCK_DATA) {
-      // 使用模拟数据
-      const stream = new ReadableStream({
-        async start(controller) {
-          for (const chunk of mockDataGenerator(message)) {
-            controller.enqueue(encoder.encode(`0:${JSON.stringify(chunk)}\n`));
-            await new Promise(resolve => setTimeout(resolve, 50)); // 模拟延迟
-          }
-          controller.enqueue(encoder.encode('e:{"finishReason":"unknown","usage":{"promptTokens":null,"completionTokens":null},"isContinued":false}\n'));
-          controller.enqueue(encoder.encode('d:{"finishReason":"unknown","usage":{"promptTokens":null,"completionTokens":null}}\n'));
-          controller.close();
-        }
-      });
+  let stream: ReadableStream;
 
-      return new Response(stream, {
-        headers: { 'Content-Type': 'text/plain; charset=utf-8' }
-      });
-    } else {
-      // 使用实际 AI 服务
-      const result = await streamText(
-        [
-          {
-            role: 'user',
-            content: stripIndents`
-            I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
-
-            IMPORTANT: Only respond with the improved prompt and nothing else!
-
-            Also, please ensure your response is entirely in Chinese.
-
-            <original_prompt>
-              ${message}
-            </original_prompt>
-          `,
-          },
-        ],
-        context.cloudflare.env,
-      );
-
-      return result.toDataStreamResponse();
-    }
-  } catch (error) {
-    console.log(error);
-
-    throw new Response(null, {
-      status: 500,
-      statusText: 'Internal Server Error',
-    });
+  if (USE_MOCK_DATA) {
+    stream = createMockStream(message);
+  } else {
+    stream = await createAIStream(message, context.cloudflare.env);
   }
+
+  return new Response(stream, {
+    headers: {
+      'Content-Type': 'text/event-stream',
+      'Cache-Control': 'no-cache',
+      Connection: 'keep-alive',
+    },
+  });
+}
+
+function createMockStream(message: string): ReadableStream {
+  return new ReadableStream({
+    async start(controller) {
+      for (const chunk of mockDataGenerator(message)) {
+        controller.enqueue(encoder.encode(`0:${JSON.stringify(chunk)}\n`));
+        await new Promise((resolve) => setTimeout(resolve, 50));
+      }
+      controller.enqueue(
+        encoder.encode(
+          'e:{"finishReason":"unknown","usage":{"promptTokens":null,"completionTokens":null},"isContinued":false}\n',
+        ),
+      );
+      controller.enqueue(
+        encoder.encode('d:{"finishReason":"unknown","usage":{"promptTokens":null,"completionTokens":null}}\n'),
+      );
+      controller.close();
+    },
+  });
+}
+
+async function createAIStream(message: string, env: any): Promise<ReadableStream> {
+  const result = await streamText(
+    [
+      {
+        role: 'user',
+        content: stripIndents`
+          I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
+
+          IMPORTANT: Only respond with the improved prompt and nothing else!
+
+          Also, please ensure your response is entirely in Chinese.
+
+          <original_prompt>
+            ${message}
+          </original_prompt>
+        `,
+      },
+    ],
+    env,
+  );
+
+  return result.toDataStreamResponse().body as ReadableStream;
 }