| 1 | import { createZodDto, UserType } from '@yikart/common' |
| 2 | import { z } from 'zod' |
| 3 | import { AiLogChannel } from '../enums' |
| 4 | |
| 5 | export const messageContentTextSchema = z.object({ |
| 6 | type: z.literal('text'), |
| 7 | text: z.string(), |
| 8 | }) |
| 9 | |
| 10 | export const messageContentImageUrlSchema = z.object({ |
| 11 | type: z.literal('image_url'), |
| 12 | image_url: z.object({ |
| 13 | url: z.url(), |
| 14 | detail: z.enum(['auto', 'low', 'high']).optional(), |
| 15 | }), |
| 16 | }) |
| 17 | |
| 18 | const complexObjectSchema = z.record(z.string(), z.any()).and(z.object({ |
| 19 | type: z.string(), |
| 20 | })) |
| 21 | |
| 22 | export const messageContentComplexSchema = z.union([ |
| 23 | messageContentTextSchema, |
| 24 | messageContentImageUrlSchema, |
| 25 | complexObjectSchema, |
| 26 | ]) |
| 27 | |
| 28 | export type MessageContent = string | z.infer<typeof messageContentComplexSchema>[] |
| 29 | |
| 30 | export const chatMessageSchema = z.object({ |
| 31 | role: z.string().describe('消息角色'), |
| 32 | content: z.union([z.string(), z.array(messageContentComplexSchema)]).describe('消息内容'), |
| 33 | }) |
| 34 | export class ChatMessageDto extends createZodDto(chatMessageSchema) {} |
| 35 | |
| 36 | export type ChatMessage = z.infer<typeof chatMessageSchema> |
| 37 | |
| 38 | export const chatCompletionDtoSchema = z.object({ |
| 39 | messages: z.array(chatMessageSchema).min(1).describe('消息列表'), |
| 40 | model: z.string().describe('模型'), |
| 41 | temperature: z.number().min(0).max(2).optional().describe('温度参数'), |
| 42 | maxTokens: z.number().int().min(1).optional().describe('最大输出token数'), |
| 43 | maxCompletionTokens: z.number().optional(), |
| 44 | modalities: z.enum(['text', 'audio', 'image', 'video']).array().optional(), |
| 45 | topP: z.number().optional(), |
| 46 | modelKwargs: z.record(z.string(), z.any()).optional(), |
| 47 | }) |
| 48 | |
| 49 | export class ChatCompletionDto extends createZodDto(chatCompletionDtoSchema) {} |
| 50 | |
| 51 | const userChatCompletionDtoSchema = z.object({ |
| 52 | userId: z.string(), |
| 53 | userType: z.enum(UserType), |
| 54 | ...chatCompletionDtoSchema.shape, |
| 55 | }) |
| 56 | |
| 57 | export class UserChatCompletionDto extends createZodDto(userChatCompletionDtoSchema) {} |
| 58 | |
| 59 | const chatModelsQuerySchema = z.object({ |
| 60 | userId: z.string().optional().describe('用户ID'), |
| 61 | userType: z.enum(UserType).optional().describe('用户类型'), |
| 62 | channel: z.enum(AiLogChannel).optional().describe('渠道筛选'), |
| 63 | scene: z.string().optional().describe('场景筛选'), |
| 64 | }) |
| 65 | |
| 66 | export class ChatModelsQueryDto extends createZodDto(chatModelsQuerySchema) {} |
| 67 |