From 9c3fe2af136821be5cb00bcebcc1875c3083a975 Mon Sep 17 00:00:00 2001
From: zyh <zyh320888@vip.qq.com>
Date: Tue, 22 Oct 2024 09:37:06 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=9E=E7=8E=B0=E8=81=8A=E5=A4=A9?=
 =?UTF-8?q?=E5=8A=9F=E8=83=BD=E5=B9=B6=E4=BC=98=E5=8C=96=E6=94=AF=E4=BB=98?=
 =?UTF-8?q?=E9=80=9A=E7=9F=A5=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/routes/api.chat.ts                        | 43 +++++++++++++++++++
 app/routes/api.payment-notify.ts              | 17 +++++++-
 ..._create_token_consumption_history_table.js | 17 ++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 db/migrations/20241022173400_create_token_consumption_history_table.js

diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index b685ac8..61fa951 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -3,12 +3,21 @@ import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants';
 import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts';
 import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
 import SwitchableStream from '~/lib/.server/llm/switchable-stream';
+import { db } from '~/utils/db.server';
+import { requireAuth } from '~/middleware/auth.server';
 
 export async function action(args: ActionFunctionArgs) {
   return chatAction(args);
 }
 
 async function chatAction({ context, request }: ActionFunctionArgs) {
+  let userId;
+  try {
+    userId = await requireAuth(request);
+  } catch (error) {
+    return error as Response;
+  }
+
   const { messages } = await request.json<{ messages: Messages }>();
 
   const stream = new SwitchableStream();
@@ -18,6 +27,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       toolChoice: 'none',
       onFinish: async ({ text: content, finishReason }) => {
         if (finishReason !== 'length') {
+          await recordTokenConsumption(userId, calculateTokensConsumed(messages, content));
           return stream.close();
         }
 
@@ -57,3 +67,36 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
     });
   }
 }
+
+async function recordTokenConsumption(userId: string, tokensConsumed: number) {
+  try {
+    await db.transaction(async (trx) => {
+      // 记录代币消耗历史
+      await trx('token_consumption_history').insert({
+        user_id: userId,
+        tokens_consumed: tokensConsumed,
+        timestamp: new Date(),
+      });
+
+      // 更新用户的代币余额
+      const updatedUser = await trx('users')
+        .where('_id', userId)
+        .decrement('token_balance', tokensConsumed)
+        .returning('token_balance');
+
+      if (updatedUser[0].token_balance < 0) {
+        throw new Error('Insufficient token balance');
+      }
+    });
+  } catch (error) {
+    console.error('Error recording token consumption:', error);
+    throw error;
+  }
+}
+
+function calculateTokensConsumed(messages: Messages, response: string): number {
+  // 这里的计算方法需要根据您的具体需求来实现
+  // 这只是一个简单的示例
+  const totalLength = messages.reduce((sum, message) => sum + message.content.length, 0) + response.length;
+  return Math.ceil(totalLength / 4); // 假设每4个字符消耗1个token
+}
diff --git a/app/routes/api.payment-notify.ts b/app/routes/api.payment-notify.ts
index cf6c468..8f9e1af 100644
--- a/app/routes/api.payment-notify.ts
+++ b/app/routes/api.payment-notify.ts
@@ -35,6 +35,21 @@ export async function action({ request }: { request: Request }) {
       await trx('users')
         .where('_id', transaction.user_id)
         .increment('token_balance', transaction.tokens);
+
+      // 如果是订阅计划,更新用户的订阅信息
+      if (transaction.type === 'subscription') {
+        const now = new Date();
+        const expirationDate = new Date(now.setMonth(now.getMonth() + (transaction.billing_cycle === 'yearly' ? 12 : 1)));
+
+        await trx('user_subscriptions').insert({
+          user_id: transaction.user_id,
+          plan_id: transaction.plan_id,
+          start_date: db.fn.now(),
+          expiration_date: expirationDate,
+          status: 'active',
+        }).onConflict('user_id')
+        .merge();
+      }
     });
 
     return json({ success: true });
@@ -42,4 +57,4 @@ export async function action({ request }: { request: Request }) {
     console.error('Error processing payment notification:', error);
     return json({ error: 'Failed to process payment notification' }, { status: 500 });
   }
-}
\ No newline at end of file
+}
diff --git a/db/migrations/20241022173400_create_token_consumption_history_table.js b/db/migrations/20241022173400_create_token_consumption_history_table.js
new file mode 100644
index 0000000..6b41111
--- /dev/null
+++ b/db/migrations/20241022173400_create_token_consumption_history_table.js
@@ -0,0 +1,17 @@
+export function up(knex) {
+  return knex.schema.createTable('token_consumption_history', function(table) {
+    table.increments('_id').primary().comment('主键ID');
+    table.integer('user_id').unsigned().notNullable().comment('用户ID');
+    table.foreign('user_id').references('users._id').onDelete('CASCADE');
+    table.integer('tokens_consumed').unsigned().notNullable().comment('消耗的代币数量');
+    table.string('session_id', 255).comment('会话ID');
+    table.timestamp('_create').defaultTo(knex.fn.now()).comment('消耗时间');
+    table.text('context').comment('消耗上下文');
+    table.index('user_id');
+    table.index('session_id');
+  });
+}
+
+export function down(knex) {
+  return knex.schema.dropTable('token_consumption_history');
+}