From 4064317a1a24cfbe85ec6a8873a14ca595d3170b Mon Sep 17 00:00:00 2001
From: zyh <zyh320888@vip.qq.com>
Date: Tue, 22 Oct 2024 10:02:57 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E4=BC=9A=E8=AF=9DID?=
 =?UTF-8?q?=E5=92=8C=E6=B6=88=E6=81=AF=E8=AE=B0=E5=BD=95=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/routes/api.chat.ts | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index df20854..efe25c2 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -18,7 +18,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
     return error as Response;
   }
 
-  const { messages } = await request.json<{ messages: Messages }>();
+  const { messages, sessionId } = await request.json<{ messages: Messages; sessionId: string }>();
 
   const stream = new SwitchableStream();
 
@@ -27,7 +27,11 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       toolChoice: 'none',
       onFinish: async ({ text: content, finishReason }) => {
         if (finishReason !== 'length') {
-          await recordTokenConsumption(userId, calculateTokensConsumed(messages, content));
+          const tokensConsumed = calculateTokensConsumed(messages, content);
+          await Promise.all([
+            recordTokenConsumption(userId, tokensConsumed),
+            recordChatHistory(userId, content, 'assistant', sessionId, tokensConsumed)
+          ]);
           return stream.close();
         }
 
@@ -48,6 +52,9 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       },
     };
 
+    // 记录用户的消息
+    await recordChatHistory(userId, messages[messages.length - 1].content, 'user', sessionId, 0);
+
     const result = await streamText(messages, context.cloudflare.env, options);
 
     stream.switchSource(result.toAIStream());
@@ -68,6 +75,21 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
   }
 }
 
+async function recordChatHistory(userId: string, message: string, role: 'user' | 'assistant', sessionId: string, tokensUsed: number) {
+  try {
+    await db('chat_histories').insert({
+      user_id: userId,
+      message,
+      role,
+      session_id: sessionId,
+      tokens_used: tokensUsed,
+    });
+  } catch (error) {
+    console.error('Error recording chat history:', error);
+    // 这里我们只记录错误,不中断流程
+  }
+}
+
 async function recordTokenConsumption(userId: string, tokensConsumed: number) {
   try {
     await db.transaction(async (trx) => {