| 1 | import { |
| 2 | Body, |
| 3 | Controller, |
| 4 | Delete, |
| 5 | Get, |
| 6 | Header, |
| 7 | Logger, |
| 8 | Param, |
| 9 | Patch, |
| 10 | Post, |
| 11 | Query, |
| 12 | Req, |
| 13 | Res, |
| 14 | SetMetadata, |
| 15 | } from '@nestjs/common' |
| 16 | import { SSE_METADATA } from '@nestjs/common/constants' |
| 17 | import { ApiTags } from '@nestjs/swagger' |
| 18 | import { GetToken, Public, TokenInfo } from '@yikart/aitoearn-auth' |
| 19 | import { ApiDoc, AppException, ParseObjectIdPipe, ResponseCode, UserType } from '@yikart/common' |
| 20 | import { ContentGenerationTaskStatus } from '@yikart/mongodb' |
| 21 | import { RedisPubSubService } from '@yikart/redis' |
| 22 | import { Request, Response } from 'express' |
| 23 | import { AGENT_TASK_ABORT_CHANNEL } from './agent.constants' |
| 24 | import { |
| 25 | CreateContentGenerationTaskDto, |
| 26 | CreateContentGenerationTaskRatingDto, |
| 27 | CreateContentGenerationTaskRatingDtoSchema, |
| 28 | CreateContentGenerationTaskSchema, |
| 29 | CreatePublicShareDto, |
| 30 | GetTaskMessagesQueryDto, |
| 31 | GetTaskMessagesQueryDtoSchema, |
| 32 | ListContentGenerationTaskDto, |
| 33 | ListContentGenerationTaskDtoSchema, |
| 34 | UpdateContentGenerationTaskTitleDto, |
| 35 | UpdateContentGenerationTaskTitleDtoSchema, |
| 36 | } from './agent.dto' |
| 37 | import { AgentService } from './agent.service' |
| 38 | import { |
| 39 | ContentGenerationTaskChunkVoSchema, |
| 40 | ContentGenerationTaskListVo, |
| 41 | ContentGenerationTaskVo, |
| 42 | PublicShareVo, |
| 43 | TaskMessagesVo, |
| 44 | } from './agent.vo' |
| 45 | |
| 46 | @ApiTags('Me/Agent') |
| 47 | @Controller('agent') |
| 48 | export class AgentController { |
| 49 | private readonly logger = new Logger(AgentController.name) |
| 50 | constructor( |
| 51 | private readonly agentService: AgentService, |
| 52 | private readonly redisPubSubService: RedisPubSubService, |
| 53 | ) { } |
| 54 | |
| 55 | @ApiDoc({ |
| 56 | summary: 'Create Content Generation Task', |
| 57 | description: 'Create a content generation task.', |
| 58 | body: CreateContentGenerationTaskSchema, |
| 59 | response: ContentGenerationTaskChunkVoSchema, |
| 60 | }) |
| 61 | @SetMetadata(SSE_METADATA, true) |
| 62 | @Header('Cache-Control', 'no-cache, no-transform') |
| 63 | @Header('Connection', 'keep-alive') |
| 64 | @Header('X-Accel-Buffering', 'no') |
| 65 | @Header('Content-Encoding', 'none') |
| 66 | @Post('tasks') |
| 67 | createContentGenerationTask( |
| 68 | @GetToken() token: TokenInfo, |
| 69 | @Body() body: CreateContentGenerationTaskDto, |
| 70 | @Req() req: Request, |
| 71 | @Res({ passthrough: true }) res: Response, |
| 72 | ) { |
| 73 | const abortController = new AbortController() |
| 74 | |
| 75 | res.on('close', () => { |
| 76 | this.logger.debug(`User ${token.id} closed connection`) |
| 77 | }) |
| 78 | |
| 79 | return this.agentService.createContentGenerationTask(token.id, UserType.User, body, abortController, req, res) |
| 80 | } |
| 81 | |
| 82 | @ApiDoc({ |
| 83 | summary: 'Get Content Generation Task List', |
| 84 | description: 'Get a paginated list of content generation tasks.', |
| 85 | query: ListContentGenerationTaskDtoSchema, |
| 86 | response: ContentGenerationTaskListVo, |
| 87 | }) |
| 88 | @Get('tasks') |
| 89 | async getContentGenerationTaskListWithPagination( |
| 90 | @GetToken() token: TokenInfo, |
| 91 | @Query() query: ListContentGenerationTaskDto, |
| 92 | ) { |
| 93 | const [tasks, total] = await this.agentService.getTaskList(token.id, query) |
| 94 | return new ContentGenerationTaskListVo(tasks, total, query) |
| 95 | } |
| 96 | |
| 97 | @ApiDoc({ |
| 98 | summary: 'Get Content Generation Task', |
| 99 | description: 'Get a content generation task.', |
| 100 | response: ContentGenerationTaskVo, |
| 101 | }) |
| 102 | @Get('tasks/:taskId') |
| 103 | async getContentGenerationTask( |
| 104 | @GetToken() token: TokenInfo, |
| 105 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 106 | ) { |
| 107 | const task = await this.agentService.getTask(token.id, taskId) |
| 108 | return ContentGenerationTaskVo.create(task) |
| 109 | } |
| 110 | |
| 111 | @ApiDoc({ |
| 112 | summary: 'Get Content Generation Task Messages', |
| 113 | description: 'Get messages after specified message ID. Used for polling after SSE disconnection.', |
| 114 | query: GetTaskMessagesQueryDtoSchema, |
| 115 | response: TaskMessagesVo, |
| 116 | }) |
| 117 | @Get('/tasks/:taskId/messages') |
| 118 | async getTaskMessages( |
| 119 | @GetToken() token: TokenInfo, |
| 120 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 121 | @Query() query: GetTaskMessagesQueryDto, |
| 122 | ) { |
| 123 | const result = await this.agentService.getTaskMessages(token.id, taskId, query.lastMessageId) |
| 124 | return TaskMessagesVo.create(result) |
| 125 | } |
| 126 | |
| 127 | @ApiDoc({ |
| 128 | summary: 'Delete Content Generation Task', |
| 129 | description: 'Delete a content generation task.', |
| 130 | }) |
| 131 | @Delete('tasks/:taskId') |
| 132 | async deleteTask( |
| 133 | @GetToken() token: TokenInfo, |
| 134 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 135 | ) { |
| 136 | await this.agentService.deleteTask(token.id, taskId) |
| 137 | } |
| 138 | |
| 139 | @ApiDoc({ |
| 140 | summary: 'Update Content Generation Task', |
| 141 | description: 'Update content generation task.', |
| 142 | body: UpdateContentGenerationTaskTitleDtoSchema, |
| 143 | }) |
| 144 | @Patch('tasks/:taskId') |
| 145 | async updateTask( |
| 146 | @GetToken() token: TokenInfo, |
| 147 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 148 | @Body() body: UpdateContentGenerationTaskTitleDto, |
| 149 | ) { |
| 150 | await this.agentService.updateTask(token.id, taskId, body) |
| 151 | } |
| 152 | |
| 153 | @ApiDoc({ |
| 154 | summary: 'Create or Update Content Generation Task Rating', |
| 155 | description: 'Create or update a rating for a content generation task.', |
| 156 | body: CreateContentGenerationTaskRatingDtoSchema, |
| 157 | }) |
| 158 | @Post('tasks/:taskId/rating') |
| 159 | async createRating( |
| 160 | @GetToken() token: TokenInfo, |
| 161 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 162 | @Body() body: CreateContentGenerationTaskRatingDto, |
| 163 | ) { |
| 164 | await this.agentService.createRating(token.id, taskId, body) |
| 165 | } |
| 166 | |
| 167 | @ApiDoc({ |
| 168 | summary: 'Create a public share link for a task', |
| 169 | description: 'Generate a token so the task can be accessed publicly. Only owner can create.', |
| 170 | body: CreatePublicShareDto.schema, |
| 171 | response: PublicShareVo, |
| 172 | }) |
| 173 | @Post('tasks/:taskId/share') |
| 174 | async createPublicShare( |
| 175 | @GetToken() token: TokenInfo, |
| 176 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 177 | @Body() body: CreatePublicShareDto, |
| 178 | ) { |
| 179 | const ttlSeconds = body.ttlSeconds |
| 180 | const res = await this.agentService.createPublicShare(token.id, taskId, ttlSeconds) |
| 181 | return PublicShareVo.create(res) |
| 182 | } |
| 183 | |
| 184 | @ApiDoc({ |
| 185 | summary: 'Get public shared task by token', |
| 186 | description: 'Retrieve a task by public share token. No authentication required.', |
| 187 | response: ContentGenerationTaskVo, |
| 188 | }) |
| 189 | @Public() |
| 190 | @Get('tasks/shared/:token') |
| 191 | async getTaskByShareToken( |
| 192 | @Param('token') token: string, |
| 193 | ) { |
| 194 | const task = await this.agentService.getTaskByShareToken(token) |
| 195 | return ContentGenerationTaskVo.create(task) |
| 196 | } |
| 197 | |
| 198 | @ApiDoc({ |
| 199 | summary: 'Abort Content Generation Task', |
| 200 | description: 'Abort a running content generation task.', |
| 201 | }) |
| 202 | @Post('/tasks/:taskId/abort') |
| 203 | async abortTask( |
| 204 | @GetToken() token: TokenInfo, |
| 205 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 206 | ) { |
| 207 | const task = await this.agentService.getTask(token.id, taskId) |
| 208 | if (task.status !== ContentGenerationTaskStatus.Running) |
| 209 | throw new AppException(ResponseCode.AgentTaskNotRunning) |
| 210 | |
| 211 | await this.redisPubSubService.emit(AGENT_TASK_ABORT_CHANNEL, taskId) |
| 212 | } |
| 213 | |
| 214 | @ApiDoc({ |
| 215 | summary: 'Favorite Content Generation Task', |
| 216 | description: 'Add a task to favorites.', |
| 217 | }) |
| 218 | @Post('/tasks/:taskId/favorite') |
| 219 | async favoriteTask( |
| 220 | @GetToken() token: TokenInfo, |
| 221 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 222 | ) { |
| 223 | await this.agentService.favoriteTask(token.id, taskId) |
| 224 | } |
| 225 | |
| 226 | @ApiDoc({ |
| 227 | summary: 'Unfavorite Content Generation Task', |
| 228 | description: 'Remove a task from favorites.', |
| 229 | }) |
| 230 | @Delete('/tasks/:taskId/favorite') |
| 231 | async unfavoriteTask( |
| 232 | @GetToken() token: TokenInfo, |
| 233 | @Param('taskId', ParseObjectIdPipe) taskId: string, |
| 234 | ) { |
| 235 | await this.agentService.unfavoriteTask(token.id, taskId) |
| 236 | } |
| 237 | } |
| 238 |