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.