| 1 | import { createSdkMcpServer, McpSdkServerConfigWithInstance } from '@anthropic-ai/claude-agent-sdk' |
| 2 | import { Injectable, Logger } from '@nestjs/common' |
| 3 | import { FileUtil, UserType } from '@yikart/common' |
| 4 | import { AiLogStatus } from '@yikart/mongodb' |
| 5 | import dayjs from 'dayjs' |
| 6 | import { z } from 'zod' |
| 7 | import { AiAvailabilityService } from '../../ai-availability' |
| 8 | import { ImageService } from '../../ai/image' |
| 9 | import { GrokVideoService, OpenAIVideoService } from '../../ai/video' |
| 10 | import { McpServerName } from '../agent.constants' |
| 11 | import { errorResult, successResult, wrapTool } from './mcp.utils' |
| 12 | |
| 13 | const generateMediaSchema = z.object({ |
| 14 | prompt: z.string(), |
| 15 | imageUrls: z.array(z.string()).optional(), |
| 16 | imageSize: z.enum(['1K', '2K', '4K']).optional(), |
| 17 | aspectRatio: z.enum(['1:1', '2:3', '3:2', '3:4', '4:3', '4:5', '5:4', '9:16', '16:9', '21:9']).optional(), |
| 18 | model: z.enum(['gemini-3.1-flash-image-preview', 'gemini-3-pro-image-preview']).optional(), |
| 19 | }) |
| 20 | |
| 21 | const generateVideoSchema = z.object({ |
| 22 | prompt: z.string(), |
| 23 | input_reference: z.string().optional(), |
| 24 | model: z.enum(['sora-2', 'sora-2-pro']).default('sora-2'), |
| 25 | seconds: z.enum(['8', '10', '25']).optional(), |
| 26 | size: z.enum(['720x1280', '1280x720', '1024x1792', '1792x1024']).optional(), |
| 27 | }) |
| 28 | |
| 29 | const getMediaStatusSchema = z.object({ |
| 30 | taskId: z.string(), |
| 31 | }) |
| 32 | |
| 33 | const createSoraCharacterSchema = z.object({ |
| 34 | prompt: z.string(), |
| 35 | videoUrl: z.string().optional(), |
| 36 | taskId: z.string().optional(), |
| 37 | timestamps: z.string().regex(/^\d+,\d+$/), |
| 38 | }) |
| 39 | |
| 40 | const getSoraCharacterSchema = z.object({ |
| 41 | characterId: z.string(), |
| 42 | }) |
| 43 | |
| 44 | const generateVideoWithGrokSchema = z.object({ |
| 45 | prompt: z.string().describe('Video description prompt'), |
| 46 | model: z.string().default('grok-imagine-video').describe('Grok video model name'), |
| 47 | aspectRatio: z.enum(['1:1', '16:9', '9:16', '4:3', '3:4', '3:2', '2:3']).default('9:16').describe('Video aspect ratio'), |
| 48 | resolution: z.enum(['480p', '720p']).default('720p').describe('Video resolution'), |
| 49 | duration: z.number().int().min(1).max(15).optional().describe('Video duration in seconds'), |
| 50 | imageUrl: z.string().optional().describe('Reference image URL for image-to-video generation'), |
| 51 | }) |
| 52 | |
| 53 | const getGrokVideoStatusSchema = z.object({ |
| 54 | taskId: z.string().describe('Task ID returned from generateVideoWithGrok'), |
| 55 | }) |
| 56 | |
| 57 | export enum MediaToolName { |
| 58 | GenerateImage = 'generateImage', |
| 59 | GenerateVideo = 'generateVideo', |
| 60 | GetVideoStatus = 'getVideoStatus', |
| 61 | CreateSoraCharacter = 'createSoraCharacter', |
| 62 | GetSoraCharacter = 'getSoraCharacter', |
| 63 | GenerateVideoWithGrok = 'generateVideoWithGrok', |
| 64 | GetGrokVideoStatus = 'getGrokVideoStatus', |
| 65 | } |
| 66 | |
| 67 | @Injectable() |
| 68 | export class MediaMcp { |
| 69 | private readonly logger = new Logger(MediaMcp.name) |
| 70 | |
| 71 | constructor( |
| 72 | private readonly openaiVideoService: OpenAIVideoService, |
| 73 | private readonly imageService: ImageService, |
| 74 | private readonly aiAvailability: AiAvailabilityService, |
| 75 | private readonly grokVideoService: GrokVideoService, |
| 76 | ) { } |
| 77 | |
| 78 | createGenerateImageTool(userId: string, userType: UserType) { |
| 79 | return wrapTool( |
| 80 | this.logger, |
| 81 | MediaToolName.GenerateImage, |
| 82 | `Generate an image using Gemini model. |
| 83 | |
| 84 | Parameters: |
| 85 | - prompt: Text description of the image |
| 86 | - imageUrls (optional): Reference image URLs for editing |
| 87 | - imageSize (optional): "1K", "2K", or "4K" |
| 88 | - aspectRatio (optional): "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9" |
| 89 | - model (optional): "gemini-3.1-flash-image-preview" (default) or "gemini-3-pro-image-preview" |
| 90 | |
| 91 | Returns the generated image URL(s).`, |
| 92 | generateMediaSchema.shape, |
| 93 | async ({ prompt, imageUrls = [], imageSize, aspectRatio, model }) => { |
| 94 | this.logger.debug(`[generateImage] Starting image generation for user ${userId}`) |
| 95 | const startTime = Date.now() |
| 96 | |
| 97 | const result = await this.imageService.userGeminiGeneration({ |
| 98 | userId, |
| 99 | userType, |
| 100 | prompt, |
| 101 | imageUrls, |
| 102 | imageSize, |
| 103 | aspectRatio, |
| 104 | ...(model ? { model } : {}), |
| 105 | }) |
| 106 | |
| 107 | const generatedImages = result.images.map(img => ({ ...img, url: FileUtil.buildUrl(img.url!) })) |
| 108 | |
| 109 | const duration = Date.now() - startTime |
| 110 | this.logger.debug(`[generateImage] Image generation completed in ${duration}ms for user ${userId}, count: ${generatedImages.length}`) |
| 111 | |
| 112 | const text = generatedImages.map((img, index) => `Image URL ${index + 1}: ${img.url}`).join('\n') |
| 113 | |
| 114 | return { |
| 115 | content: [ |
| 116 | ...generatedImages.map((img, index) => ({ |
| 117 | type: 'resource_link', |
| 118 | uri: img.url, |
| 119 | name: `Image ${index + 1}`, |
| 120 | } as const)), |
| 121 | { type: 'text', text }, |
| 122 | ], |
| 123 | } |
| 124 | }, |
| 125 | this.aiAvailability, |
| 126 | ) |
| 127 | } |
| 128 | |
| 129 | createGenerateVideoTool(userId: string, userType: UserType) { |
| 130 | return wrapTool( |
| 131 | this.logger, |
| 132 | MediaToolName.GenerateVideo, |
| 133 | `Generate a video using OpenAI Sora model. Follow the Sora Prompting Guide for best results. |
| 134 | |
| 135 | **IMPORTANT: Use the same language as the user's request for the video prompt.** |
| 136 | |
| 137 | **Model Selection**: |
| 138 | - sora-2 (default): 10s 720p video, supports audio/voiceover |
| 139 | - sora-2-pro: 25s 1024p video - use when user requests high quality/HD or long shots |
| 140 | |
| 141 | **Sora Prompt Engineering Guidelines**: |
| 142 | |
| 143 | 1. **Specificity over Abstraction**: Use concrete details instead of vague descriptions. |
| 144 | - Bad: "a beautiful street" / "美丽的街道" |
| 145 | - Good: "wet asphalt reflecting neon lights, zebra crossing, steam rising from manholes" / "湿润的沥青反射着霓虹灯光，斑马线清晰可见，井盖冒着蒸汽" |
| 146 | |
| 147 | 2. **Establish Visual Style First**: Set the visual tone at the beginning. |
| 148 | - Example: "1970s film grain aesthetic" / "1970年代胶片质感", "cinematic documentary style" / "电影纪录片风格" |
| 149 | |
| 150 | 3. **Camera Composition**: Specify framing and position. |
| 151 | - Framing: wide shot (全景), medium close-up (中近景), extreme close-up (特写) |
| 152 | - Position: eye-level (平视), low angle (低角度), aerial view (鸟瞰) |
| 153 | |
| 154 | 4. **Depth and Focus**: Describe depth of field. |
| 155 | - "shallow depth of field with blurred background" / "浅景深，背景虚化" |
| 156 | |
| 157 | 5. **Lighting and Color**: Be specific about light sources and color palette. |
| 158 | - "golden hour sunlight casting long shadows" / "黄金时段的阳光投下长长的影子" |
| 159 | |
| 160 | 6. **Action and Timing**: One camera movement + one subject action per shot. |
| 161 | - "Camera slowly pans left as the woman walks towards the window" / "镜头缓慢向左平移，女子走向窗边" |
| 162 | |
| 163 | 7. **Character Consistency**: Include detailed character descriptions for recurring characters. |
| 164 | |
| 165 | **Audio/Voiceover Guidelines (口播/旁白)**: |
| 166 | - Sora supports native audio generation with voiceover/narration |
| 167 | - Supports: Chinese (中文), English, Japanese (日本語), etc. |
| 168 | - Describe narration as a separate element at the end of the prompt |
| 169 | |
| 170 | **Voiceover Prompt Format**: |
| 171 | [Scene description] + [Camera] + [Atmosphere] + [Action] + [Voiceover content] |
| 172 | |
| 173 | **Examples**: |
| 174 | Chinese: "纪录片场景：海浪拍打礁石。全景，稳定镜头。平静氛围。男声旁白：'大自然的力量无处不在'" |
| 175 | English: "Product showcase of smartphone on white surface. Close-up, slow rotation. Minimalist aesthetic. Female voiceover: 'The new design features a seamless display'" |
| 176 | |
| 177 | Returns task ID for status tracking.`, |
| 178 | generateVideoSchema.shape, |
| 179 | async ({ prompt, input_reference, model, seconds, size }) => { |
| 180 | const finalSeconds = seconds || (model === 'sora-2-pro' ? '25' : '10') |
| 181 | const finalSize = size || (model === 'sora-2-pro' ? '1024x1792' : '720x1280') |
| 182 | |
| 183 | const response = await this.openaiVideoService.createVideo({ |
| 184 | userId, |
| 185 | userType, |
| 186 | prompt, |
| 187 | input_reference, |
| 188 | model, |
| 189 | seconds: finalSeconds, |
| 190 | size: finalSize, |
| 191 | }) |
| 192 | const { id, error, status } = response |
| 193 | |
| 194 | if (status === AiLogStatus.Failed) { |
| 195 | return errorResult(`Failed to generate video with OpenAI ${model}, Error: ${error || 'Unknown error'}`) |
| 196 | } |
| 197 | |
| 198 | return successResult(`Video is generating with OpenAI ${model}, task id: ${id}`) |
| 199 | }, |
| 200 | this.aiAvailability, |
| 201 | ) |
| 202 | } |
| 203 | |
| 204 | createGetVideoStatusTool(userId: string, userType: UserType) { |
| 205 | return wrapTool( |
| 206 | this.logger, |
| 207 | MediaToolName.GetVideoStatus, |
| 208 | 'Get video generation task status. Provide taskId. Returns task status, progress percentage, and video URL when completed. Response includes start time, current time, and elapsed time for tracking generation progress.', |
| 209 | getMediaStatusSchema.shape, |
| 210 | async ({ taskId }) => { |
| 211 | const result = await this.openaiVideoService.getVideo(userId, userType, taskId) |
| 212 | |
| 213 | const startTime = result.created_at || 0 |
| 214 | let timeInfo = '' |
| 215 | if (startTime > 0) { |
| 216 | const start = dayjs.unix(startTime) |
| 217 | const current = dayjs() |
| 218 | const elapsedSeconds = current.diff(start, 'second') |
| 219 | const elapsedMinutes = current.diff(start, 'minute') |
| 220 | const remainingSeconds = elapsedSeconds % 60 |
| 221 | timeInfo = `\nStart time: ${start.toISOString()}\nCurrent time: ${current.toISOString()}\nElapsed time: ${elapsedMinutes} minutes ${remainingSeconds} seconds` |
| 222 | } |
| 223 | |
| 224 | if (result.status === 'completed' && (result.url || result.video_url)) { |
| 225 | const videoUrl = result.url || result.video_url || '' |
| 226 | const fullVideoUrl = FileUtil.buildUrl(videoUrl) |
| 227 | return successResult(`Video is completed, task id: ${taskId} and video url is ${fullVideoUrl}${timeInfo}`) |
| 228 | } |
| 229 | if (result.status === 'failed') { |
| 230 | return errorResult(`Video is failed, task id: ${taskId} and error message is ${result.error?.message || 'Unknown error'}${timeInfo}`) |
| 231 | } |
| 232 | return successResult(`Video is ${result.status}, progress: ${result.progress}%, task id: ${taskId}${timeInfo}`) |
| 233 | }, |
| 234 | this.aiAvailability, |
| 235 | ) |
| 236 | } |
| 237 | |
| 238 | createSoraCharacterTool(userId: string, userType: UserType) { |
| 239 | return wrapTool( |
| 240 | this.logger, |
| 241 | MediaToolName.CreateSoraCharacter, |
| 242 | `Create a reusable character from video for Sora generation. |
| 243 | |
| 244 | **Purpose**: Maintain consistent character appearance across multiple video generations. Once created, the same character can be referenced in different videos with identical visual features. |
| 245 | |
| 246 | **Creation Methods**: |
| 247 | 1. From Video URL: Provide videoUrl parameter |
| 248 | 2. From Task ID: Provide taskId parameter (existing Sora task) |
| 249 | |
| 250 | **timestamps**: Two numbers separated by comma (e.g., "1,3"), gap must be ≤ 3 seconds |
| 251 | |
| 252 | **Processing Time**: Character creation takes approximately 400 seconds (~7 minutes) |
| 253 | |
| 254 | **Usage**: After creation, use @{username} in prompts to reference character |
| 255 | Example: "@character1 walks through a garden"`, |
| 256 | createSoraCharacterSchema.shape, |
| 257 | async ({ prompt, videoUrl, taskId, timestamps }) => { |
| 258 | const response = await this.openaiVideoService.createCharacter({ |
| 259 | userId, |
| 260 | userType, |
| 261 | prompt, |
| 262 | videoUrl, |
| 263 | taskId, |
| 264 | timestamps, |
| 265 | }) |
| 266 | |
| 267 | if (response.status === 'failed') { |
| 268 | return errorResult(`Failed to create character, Error: ${response.error?.message || 'Unknown error'}`) |
| 269 | } |
| 270 | |
| 271 | return successResult(`Character is creating, character id: ${response.id}, use @${response.username} to reference in prompts`) |
| 272 | }, |
| 273 | this.aiAvailability, |
| 274 | ) |
| 275 | } |
| 276 | |
| 277 | createGetSoraCharacterTool(userId: string, userType: UserType) { |
| 278 | return wrapTool( |
| 279 | this.logger, |
| 280 | MediaToolName.GetSoraCharacter, |
| 281 | 'Get Sora character creation status. Provide characterId. Returns character status and username for referencing in prompts.', |
| 282 | getSoraCharacterSchema.shape, |
| 283 | async ({ characterId }) => { |
| 284 | const result = await this.openaiVideoService.getCharacter(userId, userType, characterId) |
| 285 | |
| 286 | if (result.status === 'completed') { |
| 287 | return successResult(`Character is ready, character id: ${result.id}, username: @${result.username}. Use @${result.username} in video prompts to reference this character.`) |
| 288 | } |
| 289 | if (result.status === 'failed') { |
| 290 | return errorResult(`Character creation failed, character id: ${result.id}, error: ${result.error?.message || 'Unknown error'}`) |
| 291 | } |
| 292 | return successResult(`Character is ${result.status}, character id: ${result.id}`) |
| 293 | }, |
| 294 | this.aiAvailability, |
| 295 | ) |
| 296 | } |
| 297 | |
| 298 | createGenerateVideoWithGrokTool(userId: string, userType: UserType) { |
| 299 | return wrapTool( |
| 300 | this.logger, |
| 301 | MediaToolName.GenerateVideoWithGrok, |
| 302 | `Generate a video using Grok video model. |
| 303 | |
| 304 | **IMPORTANT: Use the same language as the user's request for the video prompt.** |
| 305 | |
| 306 | **Model**: grok-imagine-video (default) |
| 307 | |
| 308 | **Generation Modes**: |
| 309 | - Text-to-video: Only prompt |
| 310 | - Image-to-video: prompt + imageUrl |
| 311 | |
| 312 | **Parameters**: |
| 313 | - prompt: Detailed video description |
| 314 | - model: Model name (default: grok-imagine-video) |
| 315 | - aspectRatio: "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3" (default: 9:16) |
| 316 | - resolution: "480p", "720p" (default: 720p) |
| 317 | - duration: Video duration in seconds (1-15, optional) |
| 318 | - imageUrl: Reference image URL for image-to-video generation (optional) |
| 319 | |
| 320 | Returns task ID for status tracking.`, |
| 321 | generateVideoWithGrokSchema.shape, |
| 322 | async ({ prompt, model, aspectRatio, resolution, duration, imageUrl }) => { |
| 323 | const response = await this.grokVideoService.createVideo({ |
| 324 | userId, |
| 325 | userType, |
| 326 | prompt, |
| 327 | model, |
| 328 | aspectRatio, |
| 329 | resolution, |
| 330 | duration, |
| 331 | image: imageUrl, |
| 332 | }) |
| 333 | |
| 334 | return successResult(`Video is generating with Grok ${model}, task id: ${response.id}`) |
| 335 | }, |
| 336 | this.aiAvailability, |
| 337 | ) |
| 338 | } |
| 339 | |
| 340 | createGetGrokVideoStatusTool(userId: string, userType: UserType) { |
| 341 | return wrapTool( |
| 342 | this.logger, |
| 343 | MediaToolName.GetGrokVideoStatus, |
| 344 | 'Get Grok video generation task status. Provide taskId. Returns task status and video URL when completed.', |
| 345 | getGrokVideoStatusSchema.shape, |
| 346 | async ({ taskId }) => { |
| 347 | const result = await this.grokVideoService.getTask(userId, userType, taskId) |
| 348 | |
| 349 | if (result.videoUrl) { |
| 350 | return successResult(`Video is completed, task id: ${taskId} and video url is ${FileUtil.buildUrl(result.videoUrl)}`) |
| 351 | } |
| 352 | if (result.error) { |
| 353 | return errorResult(`Video failed, task id: ${taskId} and error message is ${result.error}`) |
| 354 | } |
| 355 | return successResult(`Video is ${result.status}, task id: ${taskId}`) |
| 356 | }, |
| 357 | this.aiAvailability, |
| 358 | ) |
| 359 | } |
| 360 | |
| 361 | createServer(userId: string, userType: UserType): McpSdkServerConfigWithInstance { |
| 362 | return createSdkMcpServer({ |
| 363 | name: McpServerName.MediaGeneration, |
| 364 | version: '1.0.0', |
| 365 | tools: [ |
| 366 | this.createGenerateImageTool(userId, userType), |
| 367 | // this.createGenerateVideoTool(userId, userType), |
| 368 | // this.createGetVideoStatusTool(userId, userType), |
| 369 | // this.createSoraCharacterTool(userId, userType), |
| 370 | // this.createGetSoraCharacterTool(userId, userType), |
| 371 | this.createGenerateVideoWithGrokTool(userId, userType), |
| 372 | this.createGetGrokVideoStatusTool(userId, userType), |
| 373 | ], |
| 374 | }) |
| 375 | } |
| 376 | } |
| 377 |