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 so that the document updates in real-time while the AI is generating content.
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 })
},
})
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.