message.ts 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. import z from "zod"
  2. import { Bus } from "../bus"
  3. import { Provider } from "../provider/provider"
  4. import { NamedError } from "../util/error"
  5. export namespace Message {
  6. export const ToolCall = z
  7. .object({
  8. state: z.literal("call"),
  9. step: z.number().optional(),
  10. toolCallId: z.string(),
  11. toolName: z.string(),
  12. args: z.custom<Required<unknown>>(),
  13. })
  14. .openapi({
  15. ref: "Message.ToolInvocation.ToolCall",
  16. })
  17. export type ToolCall = z.infer<typeof ToolCall>
  18. export const ToolPartialCall = z
  19. .object({
  20. state: z.literal("partial-call"),
  21. step: z.number().optional(),
  22. toolCallId: z.string(),
  23. toolName: z.string(),
  24. args: z.custom<Required<unknown>>(),
  25. })
  26. .openapi({
  27. ref: "Message.ToolInvocation.ToolPartialCall",
  28. })
  29. export type ToolPartialCall = z.infer<typeof ToolPartialCall>
  30. export const ToolResult = z
  31. .object({
  32. state: z.literal("result"),
  33. step: z.number().optional(),
  34. toolCallId: z.string(),
  35. toolName: z.string(),
  36. args: z.custom<Required<unknown>>(),
  37. result: z.string(),
  38. })
  39. .openapi({
  40. ref: "Message.ToolInvocation.ToolResult",
  41. })
  42. export type ToolResult = z.infer<typeof ToolResult>
  43. export const ToolInvocation = z
  44. .discriminatedUnion("state", [ToolCall, ToolPartialCall, ToolResult])
  45. .openapi({
  46. ref: "Message.ToolInvocation",
  47. })
  48. export type ToolInvocation = z.infer<typeof ToolInvocation>
  49. export const TextPart = z
  50. .object({
  51. type: z.literal("text"),
  52. text: z.string(),
  53. })
  54. .openapi({
  55. ref: "Message.Part.Text",
  56. })
  57. export type TextPart = z.infer<typeof TextPart>
  58. export const ReasoningPart = z
  59. .object({
  60. type: z.literal("reasoning"),
  61. text: z.string(),
  62. providerMetadata: z.record(z.any()).optional(),
  63. })
  64. .openapi({
  65. ref: "Message.Part.Reasoning",
  66. })
  67. export type ReasoningPart = z.infer<typeof ReasoningPart>
  68. export const ToolInvocationPart = z
  69. .object({
  70. type: z.literal("tool-invocation"),
  71. toolInvocation: ToolInvocation,
  72. })
  73. .openapi({
  74. ref: "Message.Part.ToolInvocation",
  75. })
  76. export type ToolInvocationPart = z.infer<typeof ToolInvocationPart>
  77. export const SourceUrlPart = z
  78. .object({
  79. type: z.literal("source-url"),
  80. sourceId: z.string(),
  81. url: z.string(),
  82. title: z.string().optional(),
  83. providerMetadata: z.record(z.any()).optional(),
  84. })
  85. .openapi({
  86. ref: "Message.Part.SourceUrl",
  87. })
  88. export type SourceUrlPart = z.infer<typeof SourceUrlPart>
  89. export const FilePart = z
  90. .object({
  91. type: z.literal("file"),
  92. mediaType: z.string(),
  93. filename: z.string().optional(),
  94. url: z.string(),
  95. })
  96. .openapi({
  97. ref: "Message.Part.File",
  98. })
  99. export type FilePart = z.infer<typeof FilePart>
  100. export const StepStartPart = z
  101. .object({
  102. type: z.literal("step-start"),
  103. })
  104. .openapi({
  105. ref: "Message.Part.StepStart",
  106. })
  107. export type StepStartPart = z.infer<typeof StepStartPart>
  108. export const Part = z
  109. .discriminatedUnion("type", [
  110. TextPart,
  111. ReasoningPart,
  112. ToolInvocationPart,
  113. SourceUrlPart,
  114. FilePart,
  115. StepStartPart,
  116. ])
  117. .openapi({
  118. ref: "Message.Part",
  119. })
  120. export type Part = z.infer<typeof Part>
  121. export const Info = z
  122. .object({
  123. id: z.string(),
  124. role: z.enum(["system", "user", "assistant"]),
  125. parts: z.array(Part),
  126. metadata: z.object({
  127. time: z.object({
  128. created: z.number(),
  129. completed: z.number().optional(),
  130. }),
  131. error: z
  132. .discriminatedUnion("name", [
  133. Provider.AuthError.Schema,
  134. NamedError.Unknown.Schema,
  135. ])
  136. .optional(),
  137. sessionID: z.string(),
  138. tool: z.record(
  139. z.string(),
  140. z
  141. .object({
  142. title: z.string(),
  143. time: z.object({
  144. start: z.number(),
  145. end: z.number(),
  146. }),
  147. })
  148. .catchall(z.any()),
  149. ),
  150. assistant: z
  151. .object({
  152. modelID: z.string(),
  153. providerID: z.string(),
  154. cost: z.number(),
  155. summary: z.boolean().optional(),
  156. tokens: z.object({
  157. input: z.number(),
  158. output: z.number(),
  159. reasoning: z.number(),
  160. }),
  161. })
  162. .optional(),
  163. }),
  164. })
  165. .openapi({
  166. ref: "Message.Info",
  167. })
  168. export type Info = z.infer<typeof Info>
  169. export const Event = {
  170. Updated: Bus.event(
  171. "message.updated",
  172. z.object({
  173. info: Info,
  174. }),
  175. ),
  176. PartUpdated: Bus.event(
  177. "message.part.updated",
  178. z.object({ part: Part, sessionID: z.string(), messageID: z.string() }),
  179. ),
  180. }
  181. }