| 1 | import { createZodDto, zodI18nString } from '@yikart/common' |
| 2 | import { z } from 'zod' |
| 3 | import { messageContentComplexSchema } from '../dtos/chat.dto' |
| 4 | import { AiLogChannel } from '../enums' |
| 5 | |
| 6 | const modalitiesTokenDetails = z.object({ |
| 7 | text: z.number().optional(), |
| 8 | image: z.number().optional(), |
| 9 | audio: z.number().optional(), |
| 10 | video: z.number().optional(), |
| 11 | document: z.number().optional(), |
| 12 | }) |
| 13 | |
| 14 | const chatInputTokenDetailsSchema = z.object({ |
| 15 | ...modalitiesTokenDetails.shape, |
| 16 | cache_read: z.number().optional(), |
| 17 | cache_creation_5m: z.number().optional(), |
| 18 | cache_creation_1h: z.number().optional(), |
| 19 | }).optional() |
| 20 | |
| 21 | const chatOutputTokenDetailsSchema = z.object({ |
| 22 | ...modalitiesTokenDetails.shape, |
| 23 | reasoning: z.number().optional(), |
| 24 | }).optional() |
| 25 | |
| 26 | const chatCompletionVoSchema = z.object({ |
| 27 | content: z.union([z.string(), z.array(messageContentComplexSchema)]).describe('生成内容'), |
| 28 | model: z.string().optional().describe('使用的模型'), |
| 29 | usage: z.object({ |
| 30 | input_tokens: z.number().optional().describe('输入token数'), |
| 31 | output_tokens: z.number().optional().describe('输出token数'), |
| 32 | total_tokens: z.number().optional().describe('总token数'), |
| 33 | input_token_details: chatInputTokenDetailsSchema, |
| 34 | output_token_details: chatOutputTokenDetailsSchema, |
| 35 | }).optional().describe('token 使用情况'), |
| 36 | }) |
| 37 | |
| 38 | export class ChatCompletionVo extends createZodDto(chatCompletionVoSchema) {} |
| 39 | |
| 40 | // 流式响应 Chunk VO |
| 41 | export const chatCompletionChunkVoSchema = z.union([ |
| 42 | z.object({ |
| 43 | type: z.literal('content'), |
| 44 | content: z.union([z.string(), z.array(messageContentComplexSchema)]).describe('流式内容'), |
| 45 | }), |
| 46 | z.object({ |
| 47 | type: z.literal('complete'), |
| 48 | content: z.union([z.string(), z.array(messageContentComplexSchema)]).describe('完整内容'), |
| 49 | usage: z.object({ |
| 50 | input_tokens: z.number().optional().describe('输入token数'), |
| 51 | output_tokens: z.number().optional().describe('输出token数'), |
| 52 | total_tokens: z.number().optional().describe('总token数'), |
| 53 | input_token_details: chatInputTokenDetailsSchema, |
| 54 | output_token_details: chatOutputTokenDetailsSchema, |
| 55 | }).describe('token 使用情况'), |
| 56 | }), |
| 57 | ]) |
| 58 | |
| 59 | export type ChatCompletionChunkVo = z.infer<typeof chatCompletionChunkVoSchema> |
| 60 | |
| 61 | // 对话模型参数 VO |
| 62 | export const chatModelSchema = z.object({ |
| 63 | name: z.string(), |
| 64 | description: z.string(), |
| 65 | summary: z.string().optional(), |
| 66 | logo: z.string().optional(), |
| 67 | tags: z.array(zodI18nString()).default([]), |
| 68 | mainTag: z.string().optional(), |
| 69 | channel: z.enum(AiLogChannel).describe('渠道'), |
| 70 | scenes: z.string().array().optional().describe('适用场景'), |
| 71 | inputModalities: z.array(z.enum(['text', 'image', 'video', 'audio'])), |
| 72 | outputModalities: z.array(z.enum(['text', 'image', 'video', 'audio'])), |
| 73 | }) |
| 74 | |
| 75 | export class ChatModelConfigVo extends createZodDto(chatModelSchema) {} |
| 76 |