Multi-document AI agent

npx create-next-app@latest multi-document-ai-agent

npm install @tiptap/react @tiptap/starter-kit ai @ai-sdk/react @ai-sdk/openai

npm install @tiptap-pro/ai-toolkit @tiptap-pro/ai-toolkit-ai-sdk

// app/api/chat/route.ts
import { openai } from '@ai-sdk/openai'
import { toolDefinitions } from '@tiptap-pro/ai-toolkit-ai-sdk'
import { convertToModelMessages, streamText, tool, UIMessage } from 'ai'
import { z } from 'zod'

export async function POST(req: Request) {
  const { messages }: { messages: UIMessage[] } = await req.json()

  const result = streamText({
    model: openai('gpt-5'),
    system: `You are an assistant that can edit rich text documents. 
    You have access to multiple documents and can switch between them. 
    At any point in time, the "active document" is the document that is open in the editor.
    When you call the tools to read and edit the document, they will read and edit the active document.
    To read and edit another document, you should use the tools to switch to that document and then read and edit it.
    Before making any edits, you should always list the documents and see which is the active document.
    `,
    messages: convertToModelMessages(messages),
    tools: {
      ...toolDefinitions(),
      createDocument: tool({
        description: 'Create a new document',
        inputSchema: z.object({
          documentName: z.string(),
        }),
      }),
      listDocuments: tool({
        description:
          'See a list of all the documents you have access to, and see which is the active document',
        inputSchema: z.object({}),
      }),
      setActiveDocument: tool({
        description: 'Switch to a specific document, so that it becomes the active document',
        inputSchema: z.object({
          documentName: z.string(),
        }),
      }),
      deleteDocument: tool({
        description: 'Delete a document',
        inputSchema: z.object({
          documentName: z.string(),
        }),
      }),
    },
  })

  return result.toUIMessageStreamResponse()
}

// app/page.tsx
'use client'

import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai'
import { useChat } from '@ai-sdk/react'
import { Editor, 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'

/**
 * One of the documents available in the app. Each
 * document has a name and a content.
 */
interface Document {
  /** The name/title of the document */
  name: string
  /** The HTML content of the document */
  content: string
}

/**
 * Initial set of documents available when the application starts
 */
const initialDocuments: Document[] = [
  {
    name: 'Document 1',
    content: '<h1>Document 1</h1><p>This is the content of the first document.</p>',
  },
]

export default function Page() {
  // The editor instance of the current active document
  const editorRef = useRef<Editor | null>(null)
  // The list of documents
  const [documents, setDocuments] = useState(initialDocuments)
  // The name of the active document
  const [activeDocumentName, setActiveDocumentName] = useState('Document 1')

  /**
   * Find a document by name
   * @param documentName
   * @returns
   */
  const findDocument = (documentName: string) => {
    return documents.find((doc) => doc.name === documentName)
  }

  /**
   * The active document is the one that is open in the editor
   */
  const activeDocument = findDocument(activeDocumentName)

  /**
   * Call this function before switching to a new document, to save
   * the content of the active document
   */
  const saveActiveDocument = () => {
    const editor = editorRef.current
    if (!editor) return
    const content = editor.getHTML()
    setDocuments((documents) =>
      documents.map((doc) => (doc.name === activeDocumentName ? { ...doc, content } : doc)),
    )
  }

  /**
   * Create a new document
   * @param documentName The name of the new document
   * @returns A message indicating the result of the operation
   */
  const createDocument = (documentName: string) => {
    saveActiveDocument()
    const existingDocument = findDocument(documentName)
    if (existingDocument) {
      setActiveDocumentName(documentName)
      return `Document already exists. Active document is now "${existingDocument.name}"`
    }
    const newDocument = {
      name: documentName,
      content: '<p></p>',
    }
    setDocuments((documents) => [...documents, newDocument])
    setActiveDocumentName(documentName)
    return `Document created. Active document is now "${documentName}"`
  }

  /**
   * Delete a document
   * @param documentName The name of the document to delete
   * @returns A message indicating the result of the operation
   */
  const deleteDocument = (documentName: string) => {
    if (documentName === activeDocumentName) {
      return `Cannot delete the active document. Active document is "${activeDocumentName}". Switch to another document first.`
    }
    const existingDocument = findDocument(documentName)
    if (!existingDocument) {
      return `Document does not exist. Active document is still "${activeDocumentName}"`
    }
    setDocuments((documents) => documents.filter((doc) => doc.name !== documentName))
    return `Deleted document "${documentName}".`
  }

  /**
   * Switch to a different document as the active document
   * @param documentName The name of the document to switch to
   * @returns A message indicating the result of the operation
   */
  const setActiveDocument = (documentName: string) => {
    const existingDocument = findDocument(documentName)
    if (!existingDocument) {
      return `Document does not exist.`
    }
    saveActiveDocument()
    setActiveDocumentName(documentName)
    return `Switched to document "${documentName}".`
  }

  /**
   * Get a formatted list of all documents and the currently active document
   * @returns A string containing the list of documents and active document name
   */
  const listDocuments = () => {
    return `Documents: ${documents
      .map((doc) => `"${doc.name}"`)
      .join(', ')}. Active document is "${activeDocumentName}".`
  }

  /**
   * Handle tool calls from the AI agent, routing them to appropriate functions
   * @param toolCall The tool call object containing tool name, input, and ID
   * @returns The result of the tool execution with tool name, ID, and output
   */
  const handleToolCall = (toolCall: { toolName: string; input: unknown }): string => {
    const editor = editorRef.current
    if (!editor) return ''

    const { toolName, input } = toolCall

    if (toolName === 'createDocument') {
      return createDocument((input as { documentName: string }).documentName)
    } else if (toolName === 'listDocuments') {
      return listDocuments()
    } else if (toolName === 'setActiveDocument') {
      return setActiveDocument((input as { documentName: string }).documentName)
    } else if (toolName === 'deleteDocument') {
      return deleteDocument((input as { documentName: string }).documentName)
    }

    // Use the AI Toolkit to execute the tool
    const toolkit = getAiToolkit(editor)
    const result = toolkit.executeTool({
      toolName,
      input,
    })

    return result.output
  }

  /**
   * Reference to the handleToolCall function to avoid stale closure issues
   * Fixes issue: https://github.com/vercel/ai/issues/8148
   */
  const handleToolCallRef = useRef(handleToolCall)
  handleToolCallRef.current = handleToolCall

  const { messages, sendMessage, addToolResult } = useChat({
    transport: new DefaultChatTransport({ api: '/api/multi-document' }),
    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
    async onToolCall({ toolCall }) {
      const output = handleToolCallRef.current(toolCall)
      if (output) {
        addToolResult({ tool: toolCall.toolName, toolCallId: toolCall.toolCallId, output })
      }
    },
  })

  const [input, setInput] = useState(
    'Create two documents, one with a short poem and another with a short story about Tiptap',
  )

  return (
    <div>
      {documents.map((doc) => (
        <button
          key={doc.name}
          disabled={doc.name === activeDocumentName}
          onClick={() => setActiveDocument(doc.name)}
        >
          {doc.name}
        </button>
      ))}

      {activeDocument && (
        <EditorComponent
          key={activeDocument.name}
          initialContent={activeDocument.content}
          onEditorInitialized={(editor) => (editorRef.current = 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>
  )
}

/**
 * Editor component that wraps the Tiptap editor with AI toolkit support
 * @param initialContent The initial HTML content to display in the editor
 * @param onEditorInitialized Callback function called when the editor is initialized
 */
function EditorComponent({
  initialContent,
  onEditorInitialized,
}: {
  /** The initial HTML content to display in the editor */
  initialContent: string
  /** Callback function called when the editor is initialized */
  onEditorInitialized: (editor: Editor) => void
}) {
  const editor = useEditor({
    immediatelyRender: false,
    extensions: [StarterKit, AiToolkit],
    content: initialContent,
  })

  useEffect(() => {
    if (editor) {
      onEditorInitialized(editor)
    }
  }, [editor, onEditorInitialized])

  return <EditorContent editor={editor} content={initialContent} />
}