Tool streaming
Continuation from the AI agent chatbot guide
This guide continues the AI agent chatbot guide. Read it first.
Activate the AI Toolkit's tool streaming capabilities to update the document in real-time while the AI is generating content.
See the source code on GitHub.
Key changes
To add streaming to the AI agent chatbot we built in the previous guide, we need to replace the executeTool method with the streamTool method.
First, while the tool call is streaming, you can call streamTool repeatedly, every time a new streaming part is received. This will update the document incrementally.
const aiToolkit = getAiToolkit(editor) const result = aiToolkit.streamTool({ toolCallId: 'call_123', toolName: 'insertContent', // Content is still streaming, so we pass a partial object. input: { position: 'documentEnd', // The HTML content has not fully been received yet html: '<p>HTML cont', }, // This parameter indicates that the tool streaming has not finished yet hasFinished: false, })Then, when the tool call is complete, call the streamTool method again with hasFinished: true to indicate that the tool call has finished streaming. This will update the document with the final content.
const result = aiToolkit.streamTool({ toolCallId: 'call_123', toolName: 'insertContent', // Streaming is complete, so we can pass the full object input: { position: 'documentEnd', // The HTML content has fully been received html: '<p>HTML content</p>', }, // This parameter indicates that the tool streaming has finished hasFinished: true, })To implement this process in the AI agent chatbot we built in the previous guide, follow these steps:
1. Handle streaming updates
Add a useEffect hook to handle streaming updates while the tool call is in progress. Inside this hook, we call streamTool repeatedly, every time a new streaming part is received.
// While the tool streaming is in progress, we need to update the document // as the tool input changes useEffect(() => { if (!editor) return // Find the last message const lastMessage = messages[messages.length - 1] if (!lastMessage) return // Find the last tool that the AI has just called const toolCallParts = lastMessage.parts.filter((p) => p.type.startsWith('tool-')) ?? [] const lastToolCall = toolCallParts[toolCallParts.length - 1] if (!lastToolCall) return // Get the tool call data interface ToolStreamingPart { input: unknown state: string toolCallId: string type: string } const part = lastToolCall as ToolStreamingPart if (!(part.state === 'input-streaming')) return const toolName = part.type.replace('tool-', '') // Apply the tool call to the document, while it is streaming const toolkit = getAiToolkit(editor) toolkit.streamTool({ toolCallId: part.toolCallId, toolName, input: part.input, // This parameter indicates that the tool streaming has not finished yet hasFinished: false, }) }, [addToolResult, editor, messages])2. Handle streaming completion
In our demo, we use the useChat hook from the AI SDK by Vercel to implement the AI agent chatbot. This hook contains an onToolCall event handler that runs when the tool call finishes streaming.
Inside this handler, we call streamTool with hasFinished: true to indicate that the tool call has finished streaming.
const { messages, sendMessage, addToolResult } = useChat({ transport: new DefaultChatTransport({ api: '/api/chat' }), sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, async onToolCall({ toolCall }) { const editor = editorRef.current if (!editor) return const { toolName, input, toolCallId } = toolCall // Use the AI Toolkit to stream the tool const toolkit = getAiToolkit(editor) const result = toolkit.streamTool({ toolCallId, toolName, input, // This parameter indicates that the tool streaming is complete hasFinished: true, }) addToolResult({ tool: toolName, toolCallId, output: result.output }) }, })Pass the same attributes in all streamTool calls
You should pass the same attribute values in all streamTool calls. For example, in the first call to the streamTool method, if you pass {mode: 'preview'} to the reviewOptions parameter, you should pass the same value ({mode: 'preview'}) in all subsequent streamTool calls.
Complete implementation
Here's the complete updated component with tool streaming:
'use client' import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai' import { useChat } from '@ai-sdk/react' import { EditorContent, useEditor } from '@tiptap/react' import StarterKit from '@tiptap/starter-kit' import { useEffect, useRef, useState } from 'react' import { AiToolkit, getAiToolkit } from '@tiptap-pro/ai-toolkit' export default function Page() { const editor = useEditor({ immediatelyRender: false, extensions: [StarterKit, AiToolkit], content: `<h1>AI Agent Demo</h1><p>Ask the AI to improve this.</p>`, }) // Fixes issue: https://github.com/vercel/ai/issues/8148 const editorRef = useRef(editor) editorRef.current = editor const { messages, sendMessage, addToolResult } = useChat({ transport: new DefaultChatTransport({ api: '/api/chat' }), sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, async onToolCall({ toolCall }) { const editor = editorRef.current if (!editor) return const { toolName, input, toolCallId } = toolCall // When the tool streaming is complete, we need to apply the tool call to the document // Use the AI Toolkit to execute the tool const toolkit = getAiToolkit(editor) const result = toolkit.streamTool({ toolCallId, toolName, input, // This parameter indicates that the tool streaming is complete hasFinished: true, }) addToolResult({ tool: toolName, toolCallId, output: result.output }) }, }) const [input, setInput] = useState( 'Insert, at the end of the document, a long story with 10 paragraphs about Tiptap', ) // While the tool streaming is in progress, we need to update the document // as the tool input changes useEffect(() => { if (!editor) return // Find the last message const lastMessage = messages[messages.length - 1] if (!lastMessage) return // Find the last tool that the AI has just called const toolCallParts = lastMessage.parts.filter((p) => p.type.startsWith('tool-')) ?? [] const lastToolCall = toolCallParts[toolCallParts.length - 1] if (!lastToolCall) return // Get the tool call data interface ToolStreamingPart { input: unknown state: string toolCallId: string type: string } const part = lastToolCall as ToolStreamingPart if (!(part.state === 'input-streaming')) return const toolName = part.type.replace('tool-', '') // Apply the tool call to the document, while it is streaming const toolkit = getAiToolkit(editor) toolkit.streamTool({ toolCallId: part.toolCallId, toolName, input: part.input, // This parameter indicates that the tool streaming has not finished yet hasFinished: false, }) }, [addToolResult, editor, messages]) if (!editor) return null return ( <div> <EditorContent editor={editor} /> {messages?.map((message) => ( <div key={message.id} style={{ whiteSpace: 'pre-wrap' }}> <strong>{message.role}</strong> <br /> {message.parts .filter((p) => p.type === 'text') .map((p) => p.text) .join('\n')} </div> ))} <form onSubmit={(e) => { e.preventDefault() sendMessage({ text: input }) setInput('') }} > <input value={input} onChange={(e) => setInput(e.target.value)} /> </form> </div> ) }End result
With tool streaming, users can see changes happening in real-time as the AI generates them. Try it out:
See the source code on GitHub.
Next steps
Combine tool streaming with change review for the best user experience. See the review changes guide to learn how to let users preview and approve changes before they're applied.