1  import * as z from 'zod/v4'
  2  
  3  import {
  4    TextContentSchema as TextContentPartSchema,
  5    PromptMessage
  6  } from '@modelcontextprotocol/sdk/types'
  7  
  8  type Primitive = string | number | boolean | bigint | null | undefined
  9  type Flatten<T> = T extends Primitive
 10    ? T
 11    : T extends Array<infer U>
 12      ? Array<Flatten<U>>
 13      : T extends Set<infer U>
 14        ? Set<Flatten<U>>
 15        : T extends Map<infer K, infer V>
 16          ? Map<Flatten<K>, Flatten<V>>
 17          : T extends object
 18            ? {
 19                [K in keyof T]: Flatten<T[K]>
 20              }
 21            : T
 22  type Infer<Schema extends z.ZodTypeAny> = Flatten<z.infer<Schema>>
 23  
 24  const TextContent = z.string()
 25  
 26  // Either a URL of the image or the base64 encoded image data.
 27  const ImageURL = z.string()
 28  
 29  const RefusalContentPartSchema = z
 30    .object({
 31      type: z.literal('refusal'),
 32      text: z.string()
 33    })
 34    .passthrough()
 35  
 36  const ImageLevel = z.union([z.literal('auto'), z.literal('low'), z.literal('high')])
 37  
 38  const ImageContentPartSchema = z.object({
 39    type: z.literal('image_url'),
 40    image_url: z.object({
 41      url: ImageURL,
 42      detail: ImageLevel.optional()
 43    })
 44  })
 45  
 46  const ToolCallSchema = z.object({
 47    id: z.string(),
 48    type: z.literal('function'),
 49    function: z.object({
 50      name: z.string(),
 51      arguments: z.string()
 52    })
 53  })
 54  
 55  const SystemMessageSchema = z.object({
 56    content: z.union([TextContent, z.array(TextContentPartSchema)]),
 57    role: z.literal('system'),
 58    name: z.string().optional()
 59  })
 60  
 61  const UserMessageSchema = z.object({
 62    content: z.union([
 63      TextContent,
 64      z.array(z.union([TextContentPartSchema, ImageContentPartSchema]))
 65    ]),
 66    role: z.literal('user'),
 67    name: z.string().optional()
 68  })
 69  
 70  const AssistantMessageSchema = z.object({
 71    content: z.union([
 72      TextContent,
 73      z.array(z.union([TextContentPartSchema, RefusalContentPartSchema]))
 74    ]),
 75    reasoning_content: z.string().optional(),
 76    tool_calls: z.array(ToolCallSchema).optional(),
 77    role: z.literal('assistant'),
 78    name: z.string().optional()
 79  })
 80  
 81  const ToolMessageSchema = z.object({
 82    content: z.union([TextContent, z.array(TextContentPartSchema)]),
 83    role: z.literal('tool'),
 84    tool_call_id: z.string()
 85  })
 86  
 87  const ChatCompletionRequestMessageSchema = z.union([
 88    UserMessageSchema,
 89    AssistantMessageSchema,
 90    ToolMessageSchema
 91  ])
 92  
 93  const ChatCompletionMessageSchema = z.union([
 94    SystemMessageSchema,
 95    UserMessageSchema,
 96    AssistantMessageSchema,
 97    ToolMessageSchema
 98  ])
 99  
100  const ChatCompletionResponseMessageSchema = z.object({
101    content: z.string(),
102    refusal: z.string().optional(),
103    reasoning_content: z.string().optional(),
104  
105    tool_calls: z.array(ToolCallSchema).optional(),
106    role: z.literal('assistant')
107  })
108  
109  export type ToolCall = Infer<typeof ToolCallSchema>
110  // type SystemMessage = Infer<typeof SystemMessageSchema>
111  export type UserMessage = Infer<typeof UserMessageSchema>
112  export type AssistantMessage = Infer<typeof AssistantMessageSchema>
113  export type ToolMessage = Infer<typeof ToolMessageSchema>
114  
115  type TextContentPart = Infer<typeof TextContentPartSchema>
116  // type RefusalContentPart = Infer<typeof RefusalContentPartSchema>
117  type ImageContentPart = Infer<typeof ImageContentPartSchema>
118  
119  // ============================================================================
120  // ## OpenAI API compatible type
121  // ============================================================================
122  
123  // ## Supports text and images only
124  export type ChatCompletionRequestContent = TextContentPart | ImageContentPart
125  
126  export type ChatCompletionRequestMessage = Infer<typeof ChatCompletionRequestMessageSchema>
127  
128  export type ChatCompletionPromptMessage = PromptMessage
129  
130  export type ChatCompletionResponseMessage = Infer<typeof ChatCompletionResponseMessageSchema>
131  
132  export type ChatCompletionMessage = Infer<typeof ChatCompletionMessageSchema>
133  
134  export type ChatConversationMessage = ChatCompletionRequestMessage | ChatCompletionResponseMessage