feat(quota): wire ai_chat_messages quota into chat controller
- Add QuotaGuard check to sendMessage and sendMessageStream - Stream disconnect → cancel reservation - Import MembershipModule in RagChatModule Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
6a5ab31ddd
commit
807488a929
@ -1,8 +1,9 @@
|
||||
import { Controller, Get, Post, Patch, Delete, Body, Param, Res, Query, HttpCode } from '@nestjs/common';
|
||||
import { Controller, Get, Post, Patch, Delete, Body, Param, Res, Query, HttpCode, BadRequestException } from '@nestjs/common';
|
||||
import { ApiTags, ApiOperation, ApiBearerAuth } from '@nestjs/swagger';
|
||||
import type { Response } from 'express';
|
||||
import { RagChatService } from './rag-chat.service';
|
||||
import type { CreateSessionResult } from './rag-chat.service';
|
||||
import { QuotaGuardService } from '../membership/quota-guard.service';
|
||||
import { CurrentUser } from '../../common/decorators/current-user.decorator';
|
||||
import type { UserPayload } from '../../common/types';
|
||||
|
||||
@ -10,7 +11,10 @@ import type { UserPayload } from '../../common/types';
|
||||
@Controller('rag-chat')
|
||||
@ApiBearerAuth()
|
||||
export class RagChatController {
|
||||
constructor(private readonly svc: RagChatService) {}
|
||||
constructor(
|
||||
private readonly svc: RagChatService,
|
||||
private readonly quotaGuard: QuotaGuardService,
|
||||
) {}
|
||||
|
||||
@Post('sessions')
|
||||
@ApiOperation({ summary: '创建/打开对话 (open-or-create)' })
|
||||
@ -67,23 +71,45 @@ export class RagChatController {
|
||||
@Post('sessions/:id/messages')
|
||||
@ApiOperation({ summary: '发送消息(同步)' })
|
||||
async sendMessage(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }) {
|
||||
return this.svc.sendMessage(String(user.id), id, dto.content);
|
||||
const userId = String(user.id);
|
||||
const opId = `chat:${id}:${Date.now()}`;
|
||||
const check = await this.quotaGuard.check(userId, 'ai_chat_messages', opId);
|
||||
if (!check.allowed) throw new BadRequestException(this.quotaGuard.buildError(check));
|
||||
try {
|
||||
const result = await this.svc.sendMessage(userId, id, dto.content);
|
||||
this.quotaGuard.confirm(opId).catch(() => {});
|
||||
return result;
|
||||
} catch {
|
||||
this.quotaGuard.cancel(opId).catch(() => {});
|
||||
throw new BadRequestException('AI Chat failed');
|
||||
}
|
||||
}
|
||||
|
||||
@Post('sessions/:id/stream')
|
||||
@ApiOperation({ summary: '发送消息(SSE 流式,支持思考过程)' })
|
||||
async sendMessageStream(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }, @Res() res: Response) {
|
||||
const userId = String(user.id);
|
||||
const opId = `chat_stream:${id}:${Date.now()}`;
|
||||
const check = await this.quotaGuard.check(userId, 'ai_chat_messages', opId);
|
||||
if (!check.allowed) {
|
||||
res.status(429).json(this.quotaGuard.buildError(check));
|
||||
return;
|
||||
}
|
||||
res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
|
||||
res.setHeader('Cache-Control', 'no-cache');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
res.setHeader('X-Accel-Buffering', 'no');
|
||||
|
||||
const userId = String(user.id);
|
||||
try {
|
||||
for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) {
|
||||
if (res.destroyed) break;
|
||||
if (res.destroyed) { this.quotaGuard.cancel(opId).catch(() => {}); break; }
|
||||
res.write(`data: ${JSON.stringify(chunk)}\n\n`);
|
||||
}
|
||||
res.end();
|
||||
if (!res.destroyed) { res.end(); this.quotaGuard.confirm(opId).catch(() => {}); }
|
||||
} catch {
|
||||
if (!res.destroyed) res.end();
|
||||
this.quotaGuard.cancel(opId).catch(() => {});
|
||||
}
|
||||
}
|
||||
|
||||
@Patch('sessions/:id')
|
||||
|
||||
@ -4,9 +4,10 @@ import { AdminRagChatController } from './admin-rag-chat.controller';
|
||||
import { RagChatService } from './rag-chat.service';
|
||||
import { PrismaService } from '../../infrastructure/database/prisma.service';
|
||||
import { AiModule } from '../ai/ai.module';
|
||||
import { MembershipModule } from '../membership/membership.module';
|
||||
|
||||
@Module({
|
||||
imports: [AiModule],
|
||||
imports: [AiModule, MembershipModule],
|
||||
controllers: [RagChatController, AdminRagChatController],
|
||||
providers: [RagChatService, PrismaService],
|
||||
exports: [RagChatService],
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user