require json output for tool results

Author: charleslien

.agents/types/util-types.ts

@@ -74,16 +74,15 @@ export type ToolCallPart = {
   providerExecuted?: boolean
 }
 
-export type ToolResultOutput =
-  | {
-      type: 'json'
-      value: JSONValue
-    }
-  | {
-      type: 'media'
-      data: string
-      mediaType: string
-    }
+export type MediaToolResultOutputSchema = {
+  data: string
+  mediaType: string
+}
+
+export type ToolResultOutput = {
+  value: JSONValue
+  media?: MediaToolResultOutputSchema[]
+}
 
 // ===== Message Types =====
 type AuxiliaryData = {

common/src/templates/initial-agents-dir/types/util-types.ts

@@ -74,16 +74,15 @@ export type ToolCallPart = {
   providerExecuted?: boolean
 }
 
-export type ToolResultOutput =
-  | {
-      type: 'json'
-      value: JSONValue
-    }
-  | {
-      type: 'media'
-      data: string
-      mediaType: string
-    }
+export type MediaToolResultOutputSchema = {
+  data: string
+  mediaType: string
+}
+
+export type ToolResultOutput = {
+  value: JSONValue
+  media?: MediaToolResultOutputSchema[]
+}
 
 // ===== Message Types =====
 type AuxiliaryData = {

common/src/types/messages/codebuff-message.ts

@@ -41,7 +41,7 @@ export type ToolMessage = {
   role: 'tool'
   toolCallId: string
   toolName: string
-  content: ToolResultOutput[]
+  content: ToolResultOutput
 } & AuxiliaryMessageData
 
 export type Message =

common/src/types/messages/content-part.ts

@@ -45,15 +45,16 @@ export const toolCallPartSchema = z.object({
 })
 export type ToolCallPart = z.infer<typeof toolCallPartSchema>
 
-export const toolResultOutputSchema = z.discriminatedUnion('type', [
-  z.object({
-    type: z.literal('json'),
-    value: jsonValueSchema,
-  }),
-  z.object({
-    type: z.literal('media'),
-    data: z.string(),
-    mediaType: z.string(),
-  }),
-])
+export const mediaToolResultOutputSchema = z.object({
+  data: z.string(),
+  mediaType: z.string(),
+})
+export type MediaToolResultOutputSchema = z.infer<
+  typeof mediaToolResultOutputSchema
+>
+
+export const toolResultOutputSchema = z.object({
+  value: jsonValueSchema,
+  media: mediaToolResultOutputSchema.array().optional(),
+})
 export type ToolResultOutput = z.infer<typeof toolResultOutputSchema>

common/src/util/messages.ts

@@ -1,6 +1,5 @@
 import { cloneDeep, has, isEqual } from 'lodash'
 
-import type { JSONValue } from '../types/json'
 import type {
   AssistantMessage,
   AuxiliaryMessageData,
@@ -9,7 +8,6 @@ import type {
   ToolMessage,
   UserMessage,
 } from '../types/messages/codebuff-message'
-import type { ToolResultOutput } from '../types/messages/content-part'
 import type { ProviderMetadata } from '../types/messages/provider-metadata'
 import type {
   AssistantModelMessage,
@@ -119,25 +117,31 @@ function assistantToCodebuffMessage(
 function convertToolResultMessage(
   message: ToolMessage,
 ): ModelMessageWithAuxiliaryData[] {
-  return message.content.map((c) => {
-    if (c.type === 'json') {
-      return cloneDeep<ToolModelMessage>({
-        ...message,
-        role: 'tool',
-        content: [{ ...message, output: c, type: 'tool-result' }],
-      })
-    }
-    if (c.type === 'media') {
-      return cloneDeep<UserMessage>({
+  const messages: ModelMessageWithAuxiliaryData[] = [
+    cloneDeep<ToolModelMessage>({
+      ...message,
+      role: 'tool',
+      content: [
+        {
+          ...message,
+          output: { type: 'json', value: message.content.value },
+          type: 'tool-result',
+        },
+      ],
+    }),
+  ]
+
+  for (const c of message.content.media ?? []) {
+    messages.push(
+      cloneDeep<UserMessage>({
         ...message,
         role: 'user',
         content: [{ type: 'file', data: c.data, mediaType: c.mediaType }],
-      })
-    }
-    c satisfies never
-    const cAny = c as any
-    throw new Error(`Invalid tool output type: ${cAny.type}`)
-  })
+      }),
+    )
+  }
+
+  return messages
 }
 
 function convertToolMessage(message: Message): ModelMessageWithAuxiliaryData[] {
@@ -411,32 +415,3 @@ export function assistantMessage(
     content: assistantContent(params),
   }
 }
-
-export function jsonToolResult<T extends JSONValue>(
-  value: T,
-): [
-  Extract<ToolResultOutput, { type: 'json' }> & {
-    value: T
-  },
-] {
-  return [
-    {
-      type: 'json',
-      value,
-    },
-  ]
-}
-
-export function mediaToolResult(params: {
-  data: string
-  mediaType: string
-}): [Extract<ToolResultOutput, { type: 'media' }>] {
-  const { data, mediaType } = params
-  return [
-    {
-      type: 'media',
-      data,
-      mediaType,
-    },
-  ]
-}