feat: 实现聊天功能并优化支付通知处理

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
+}

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 });
   }
-}
+}

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');
+}