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 { ApiTags, ApiOperation, ApiBearerAuth } from '@nestjs/swagger';
|
||||||
import type { Response } from 'express';
|
import type { Response } from 'express';
|
||||||
import { RagChatService } from './rag-chat.service';
|
import { RagChatService } from './rag-chat.service';
|
||||||
import type { CreateSessionResult } 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 { CurrentUser } from '../../common/decorators/current-user.decorator';
|
||||||
import type { UserPayload } from '../../common/types';
|
import type { UserPayload } from '../../common/types';
|
||||||
|
|
||||||
@ -10,7 +11,10 @@ import type { UserPayload } from '../../common/types';
|
|||||||
@Controller('rag-chat')
|
@Controller('rag-chat')
|
||||||
@ApiBearerAuth()
|
@ApiBearerAuth()
|
||||||
export class RagChatController {
|
export class RagChatController {
|
||||||
constructor(private readonly svc: RagChatService) {}
|
constructor(
|
||||||
|
private readonly svc: RagChatService,
|
||||||
|
private readonly quotaGuard: QuotaGuardService,
|
||||||
|
) {}
|
||||||
|
|
||||||
@Post('sessions')
|
@Post('sessions')
|
||||||
@ApiOperation({ summary: '创建/打开对话 (open-or-create)' })
|
@ApiOperation({ summary: '创建/打开对话 (open-or-create)' })
|
||||||
@ -67,23 +71,45 @@ export class RagChatController {
|
|||||||
@Post('sessions/:id/messages')
|
@Post('sessions/:id/messages')
|
||||||
@ApiOperation({ summary: '发送消息(同步)' })
|
@ApiOperation({ summary: '发送消息(同步)' })
|
||||||
async sendMessage(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }) {
|
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')
|
@Post('sessions/:id/stream')
|
||||||
@ApiOperation({ summary: '发送消息(SSE 流式,支持思考过程)' })
|
@ApiOperation({ summary: '发送消息(SSE 流式,支持思考过程)' })
|
||||||
async sendMessageStream(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }, @Res() res: Response) {
|
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('Content-Type', 'text/event-stream; charset=utf-8');
|
||||||
res.setHeader('Cache-Control', 'no-cache');
|
res.setHeader('Cache-Control', 'no-cache');
|
||||||
res.setHeader('Connection', 'keep-alive');
|
res.setHeader('Connection', 'keep-alive');
|
||||||
res.setHeader('X-Accel-Buffering', 'no');
|
res.setHeader('X-Accel-Buffering', 'no');
|
||||||
|
|
||||||
const userId = String(user.id);
|
try {
|
||||||
for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) {
|
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.write(`data: ${JSON.stringify(chunk)}\n\n`);
|
||||||
|
}
|
||||||
|
if (!res.destroyed) { res.end(); this.quotaGuard.confirm(opId).catch(() => {}); }
|
||||||
|
} catch {
|
||||||
|
if (!res.destroyed) res.end();
|
||||||
|
this.quotaGuard.cancel(opId).catch(() => {});
|
||||||
}
|
}
|
||||||
res.end();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
@Patch('sessions/:id')
|
@Patch('sessions/:id')
|
||||||
|
|||||||
@ -4,9 +4,10 @@ import { AdminRagChatController } from './admin-rag-chat.controller';
|
|||||||
import { RagChatService } from './rag-chat.service';
|
import { RagChatService } from './rag-chat.service';
|
||||||
import { PrismaService } from '../../infrastructure/database/prisma.service';
|
import { PrismaService } from '../../infrastructure/database/prisma.service';
|
||||||
import { AiModule } from '../ai/ai.module';
|
import { AiModule } from '../ai/ai.module';
|
||||||
|
import { MembershipModule } from '../membership/membership.module';
|
||||||
|
|
||||||
@Module({
|
@Module({
|
||||||
imports: [AiModule],
|
imports: [AiModule, MembershipModule],
|
||||||
controllers: [RagChatController, AdminRagChatController],
|
controllers: [RagChatController, AdminRagChatController],
|
||||||
providers: [RagChatService, PrismaService],
|
providers: [RagChatService, PrismaService],
|
||||||
exports: [RagChatService],
|
exports: [RagChatService],
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user