From 469c8f2aed6565ef92caec7cb066988bfd18bce0 Mon Sep 17 00:00:00 2001 From: Daniel Pruessner Date: Tue, 13 Jan 2026 13:20:42 +1300 Subject: [PATCH] Tools can return DocumentBlock, ImageBlock, VideoBlock in addition to string and JSONValue --- README.md | 1 + src/models/bedrock.ts | 27 +++++++++++++++++++++++ src/tools/__tests__/zod-tool.test.ts | 22 +++++++++++++++++++ src/tools/function-tool.ts | 10 +++++++++ src/tools/zod-tool.ts | 13 ++++++++--- src/types/__tests__/messages.test.ts | 32 ++++++++++++++++++++++++++++ src/types/messages.ts | 17 ++++++++++++--- 7 files changed, 116 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index e16612c4..3ad9e009 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,7 @@ const weatherTool = tool({ }), callback: (input) => { // input is fully typed based on the Zod schema + // Tools can return strings, JSON values, or media blocks (DocumentBlock, ImageBlock, VideoBlock) return `The weather in ${input.location} is 72°F and sunny.` }, }) diff --git a/src/models/bedrock.ts b/src/models/bedrock.ts index 6966ffbf..8e2bbbc2 100644 --- a/src/models/bedrock.ts +++ b/src/models/bedrock.ts @@ -558,6 +558,33 @@ export class BedrockModel extends Model { return { text: content.text } case 'jsonBlock': return { json: content.json } + case 'documentBlock': + return { + document: { + name: content.name, + format: content.format as BedrockContentBlock.DocumentMember['document']['format'], + source: this._formatDocumentSource(content.source), + ...(content.citations && { citations: content.citations }), + ...(content.context && { context: content.context }), + }, + } + case 'imageBlock': + return { + image: { + format: content.format as BedrockContentBlock.ImageMember['image']['format'], + source: this._formatMediaSource(content.source), + }, + } + case 'videoBlock': + return { + video: { + format: content.format as BedrockContentBlock.VideoMember['video']['format'], + source: this._formatMediaSource(content.source), + }, + } + default: + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return { text: `[Unsupported content type: ${(content as any).type}]` } } }) diff --git a/src/tools/__tests__/zod-tool.test.ts b/src/tools/__tests__/zod-tool.test.ts index ee8e95a9..166abc55 100644 --- a/src/tools/__tests__/zod-tool.test.ts +++ b/src/tools/__tests__/zod-tool.test.ts @@ -98,6 +98,28 @@ describe('tool', () => { const result = await myTool.invoke({ count: 3 }) expect(result).toBe(3) }) + + it('handles DocumentBlock return', async () => { + const { DocumentBlock } = await import('../../types/media.js') + + const docTool = tool({ + name: 'create_document', + description: 'Creates a document', + inputSchema: z.object({ content: z.string() }), + callback: (input) => { + return new DocumentBlock({ + name: 'RESULT', + format: 'md', + source: { bytes: new TextEncoder().encode(input.content) }, + }) + }, + }) + + const result = await docTool.invoke({ content: 'Hello World!' }) + expect(result.type).toBe('documentBlock') + expect(result.name).toBe('RESULT') + expect(result.format).toBe('md') + }) }) describe('validation', () => { diff --git a/src/tools/function-tool.ts b/src/tools/function-tool.ts index db347731..bcffd0e4 100644 --- a/src/tools/function-tool.ts +++ b/src/tools/function-tool.ts @@ -5,6 +5,7 @@ import type { ToolSpec } from './types.js' import type { JSONSchema, JSONValue } from '../types/json.js' import { deepCopy } from '../types/json.js' import { JsonBlock, TextBlock, ToolResultBlock } from '../types/messages.js' +import { DocumentBlock, ImageBlock, VideoBlock } from '../types/media.js' /** * Callback function for FunctionTool implementations. @@ -218,6 +219,15 @@ export class FunctionTool extends Tool { */ private _wrapInToolResult(value: unknown, toolUseId: string): ToolResultBlock { try { + // Handle DocumentBlock, ImageBlock, VideoBlock directly + if (value instanceof DocumentBlock || value instanceof ImageBlock || value instanceof VideoBlock) { + return new ToolResultBlock({ + toolUseId, + status: 'success', + content: [value], + }) + } + // Handle null with special string representation as text content if (value === null) { return new ToolResultBlock({ diff --git a/src/tools/zod-tool.ts b/src/tools/zod-tool.ts index 0e68e051..f1f35b33 100644 --- a/src/tools/zod-tool.ts +++ b/src/tools/zod-tool.ts @@ -4,6 +4,13 @@ import type { ToolSpec } from './types.js' import type { JSONSchema, JSONValue } from '../types/json.js' import { FunctionTool } from './function-tool.js' import { z, ZodVoid } from 'zod' +import { DocumentBlock, ImageBlock, VideoBlock } from '../types/media.js' + +/** + * Valid return types for tool callbacks. + * Includes JSON-serializable values and media blocks. + */ +export type ToolReturnValue = JSONValue | DocumentBlock | ImageBlock | VideoBlock /** * Helper type to infer input type from Zod schema or default to never. @@ -16,7 +23,7 @@ type ZodInferred = TInput extends z.ZodType ? z.infer : never * @typeParam TInput - Zod schema type for input validation * @typeParam TReturn - Return type of the callback function */ -export interface ToolConfig { +export interface ToolConfig { /** The name of the tool */ name: string @@ -46,7 +53,7 @@ export interface ToolConfig +class ZodTool extends Tool implements InvokableTool, TReturn> { @@ -231,7 +238,7 @@ class ZodTool( +export function tool( config: ToolConfig ): InvokableTool, TReturn> { return new ZodTool(config) diff --git a/src/types/__tests__/messages.test.ts b/src/types/__tests__/messages.test.ts index 4e6ac5bb..f1c6ed28 100644 --- a/src/types/__tests__/messages.test.ts +++ b/src/types/__tests__/messages.test.ts @@ -176,6 +176,38 @@ describe('Message.fromMessageData', () => { expect(toolResultBlock.content[0]).toBeInstanceOf(JsonBlock) }) + it('converts tool result block data to ToolResultBlock with document content', () => { + const messageData: MessageData = { + role: 'user', + content: [ + { + toolResult: { + toolUseId: 'tooluse_HUxQGMooooooooooooooeV', + status: 'success', + content: [ + { + document: { + format: 'md', + name: 'DOCUMENT', + source: { + bytes: Uint8Array.from( + 'CiMgUHJvZHVjdCBSZXF1aXJlbWVudHMgRG9jdW1lbnQKCiMjIE92ZXJ2aWV3ClRoaXMgZG9jdW1lbnQgb3V0bGluZXMgdGhlIHJlcXVpcmVtZW50cyBmb3IgYSBuZXcgZmVhdHVyZSBpbiBvdXIgYXBwbGljYXRpb24uCgojIyBSZXF1aXJlbWVudHMKMS4gVXNlciBhdXRoZW50aWNhdGlvbiBtdXN0IHN1cHBvcnQgT0F1dGggMi4wCjIuIEFQSSByZXNwb25zZSB0aW1lIG11c3QgYmUgdW5kZXIgMjAwbXMKMy4gU3lzdGVtIG11c3QgaGFuZGxlIDEwLDAwMCBjb25jdXJyZW50IHVzZXJzCjQuIERhdGEgbXVzdCBiZSBlbmNyeXB0ZWQgYXQgcmVzdCBhbmQgaW4gdHJhbnNpdAoKIyMgVGltZWxpbmUKLSBQaGFzZSAxOiBRMSAyMDI2Ci0gUGhhc2UgMjogUTIgMjAyNgo=' + ), + }, + }, + }, + ], + }, + }, + ], + } + const message = Message.fromMessageData(messageData) + expect(message.content).toHaveLength(1) + const toolResultBlock = message.content[0] as ToolResultBlock + expect(toolResultBlock.content).toHaveLength(1) + expect(toolResultBlock.content[0]).toBeInstanceOf(DocumentBlock) + }) + it('converts reasoning block data to ReasoningBlock', () => { const messageData: MessageData = { role: 'assistant', diff --git a/src/types/messages.ts b/src/types/messages.ts index ddb93c0a..5f682d0d 100644 --- a/src/types/messages.ts +++ b/src/types/messages.ts @@ -192,9 +192,14 @@ export class ToolUseBlock implements ToolUseBlockData { * * This is a discriminated union where the object key determines the content format. */ -export type ToolResultContentData = TextBlockData | JsonBlockData +export type ToolResultContentData = + | TextBlockData + | JsonBlockData + | { document: DocumentBlockData } + | { image: ImageBlockData } + | { video: VideoBlockData } -export type ToolResultContent = TextBlock | JsonBlock +export type ToolResultContent = TextBlock | JsonBlock | DocumentBlock | ImageBlock | VideoBlock /** * Data for a tool result block. @@ -226,7 +231,7 @@ export interface ToolResultBlockData { /** * Tool result content block. */ -export class ToolResultBlock implements ToolResultBlockData { +export class ToolResultBlock { /** * Discriminator for tool result content. */ @@ -596,6 +601,12 @@ export function contentBlockFromData(data: ContentBlockData): ContentBlock { return new TextBlock(contentItem.text) } else if ('json' in contentItem) { return new JsonBlock(contentItem) + } else if ('document' in contentItem) { + return new DocumentBlock(contentItem.document as DocumentBlockData) + } else if ('image' in contentItem) { + return new ImageBlock(contentItem.image as ImageBlockData) + } else if ('video' in contentItem) { + return new VideoBlock(contentItem.video as VideoBlockData) } else { throw new Error('Unknown ToolResultContentData type') }