| 1 | import { createSdkMcpServer, McpSdkServerConfigWithInstance } from '@anthropic-ai/claude-agent-sdk' |
| 2 | import { Injectable, Logger } from '@nestjs/common' |
| 3 | import { UserType } from '@yikart/common' |
| 4 | import { z } from 'zod' |
| 5 | import { AiAvailabilityService } from '../../../ai-availability' |
| 6 | import { VideoStyleTransferService } from '../../../ai/aideo' |
| 7 | import { AideoTaskStatus } from '../../../ai/libs/volcengine' |
| 8 | import { McpServerName } from '../../agent.constants' |
| 9 | import { errorResult, successResult, wrapTool } from '../mcp.utils' |
| 10 | |
| 11 | const submitVideoStyleTransferSchema = z.object({ |
| 12 | videoInput: z.string(), |
| 13 | style: z.string().optional(), |
| 14 | resolution: z.enum(['480p', '720p', '1080p']), |
| 15 | }) |
| 16 | |
| 17 | const getVideoStyleTransferStatusSchema = z.object({ |
| 18 | taskId: z.string(), |
| 19 | }) |
| 20 | |
| 21 | export enum StyleTransferToolName { |
| 22 | SubmitVideoStyleTransfer = 'submitVideoStyleTransfer', |
| 23 | GetVideoStyleTransferStatus = 'getVideoStyleTransferStatus', |
| 24 | } |
| 25 | |
| 26 | @Injectable() |
| 27 | export class StyleTransferMcp { |
| 28 | private readonly logger = new Logger(StyleTransferMcp.name) |
| 29 | |
| 30 | constructor( |
| 31 | private readonly videoStyleTransferService: VideoStyleTransferService, |
| 32 | private readonly aiAvailability: AiAvailabilityService, |
| 33 | ) { } |
| 34 | |
| 35 | createSubmitVideoStyleTransferTool(userId: string, userType: UserType) { |
| 36 | return wrapTool( |
| 37 | this.logger, |
| 38 | StyleTransferToolName.SubmitVideoStyleTransfer, |
| 39 | `Submit a video style transfer task to convert live-action videos into artistic styles. |
| 40 | |
| 41 | **Parameters**: |
| 42 | - videoInput: Video URL or VID (vid://xxx format); URLs auto-uploaded |
| 43 | - style: Target style - "漫画风" (Comic, default), "3D卡通风格" (3D Cartoon), "日漫风格" (Japanese Anime) |
| 44 | - resolution: Output resolution - 480p (0.7x price), 720p (1.0x price), 1080p (1.8x price) |
| 45 | |
| 46 | Processing time: ~10 minutes per 1-minute video. Returns taskId for status tracking.`, |
| 47 | submitVideoStyleTransferSchema.shape, |
| 48 | async ({ videoInput, style, resolution }) => { |
| 49 | const result = await this.videoStyleTransferService.submitVideoStyleTransferTask({ |
| 50 | userId, |
| 51 | userType, |
| 52 | videoInput, |
| 53 | style, |
| 54 | resolution, |
| 55 | }) |
| 56 | return successResult(`Video style transfer task submitted successfully. TaskId: ${result.taskId}`) |
| 57 | }, |
| 58 | this.aiAvailability, |
| 59 | ) |
| 60 | } |
| 61 | |
| 62 | createGetVideoStyleTransferStatusTool(userId: string, userType: UserType) { |
| 63 | return wrapTool( |
| 64 | this.logger, |
| 65 | StyleTransferToolName.GetVideoStyleTransferStatus, |
| 66 | 'Get video style transfer task status. Returns status (Processing/Completed/Failed), and when completed, returns outputVid and outputUrl.', |
| 67 | getVideoStyleTransferStatusSchema.shape, |
| 68 | async ({ taskId }) => { |
| 69 | const result = await this.videoStyleTransferService.getVideoStyleTransferTask({ |
| 70 | userId, |
| 71 | userType, |
| 72 | taskId, |
| 73 | }) |
| 74 | |
| 75 | if (result.status === AideoTaskStatus.Completed) { |
| 76 | return successResult(`Task completed successfully! Output video URL: ${result.outputUrl}, VID: ${result.outputVid}`) |
| 77 | } |
| 78 | else if (result.status === AideoTaskStatus.Processing) { |
| 79 | return successResult(`Task is still processing. Please continue to wait...`) |
| 80 | } |
| 81 | else { |
| 82 | return errorResult(`Task failed: ${result.errorMessage}`) |
| 83 | } |
| 84 | }, |
| 85 | this.aiAvailability, |
| 86 | ) |
| 87 | } |
| 88 | |
| 89 | createServer(userId: string, userType: UserType): McpSdkServerConfigWithInstance { |
| 90 | return createSdkMcpServer({ |
| 91 | name: McpServerName.StyleTransfer, |
| 92 | version: '1.0.0', |
| 93 | tools: [ |
| 94 | this.createSubmitVideoStyleTransferTool(userId, userType), |
| 95 | this.createGetVideoStyleTransferStatusTool(userId, userType), |
| 96 | ], |
| 97 | }) |
| 98 | } |
| 99 | } |
| 100 |