@blocknote/xl-ai 0.42.2 → 0.43.0

package/src/AIExtension.ts

@@ -1,10 +1,15 @@
 import { Chat } from "@ai-sdk/react";
 import {
-  BlockNoteEditor,
-  BlockNoteExtension,
+  createExtension,
+  createStore,
+  ExtensionOptions,
   getNodeById,
   UnreachableCaseError,
 } from "@blocknote/core";
+import {
+  ForkYDocExtension,
+  ShowSelectionExtension,
+} from "@blocknote/core/extensions";
 import {
   applySuggestions,
   revertSuggestions,
@@ -14,7 +19,6 @@ import { UIMessage } from "ai";
 import { Fragment, Slice } from "prosemirror-model";
 import { Plugin, PluginKey } from "prosemirror-state";
 import { fixTablesKey } from "prosemirror-tables";
-import { createStore, StoreApi } from "zustand/vanilla";
 
 import {
   aiDocumentFormats,
@@ -25,18 +29,7 @@ import {
 import { createAgentCursorPlugin } from "./plugins/AgentCursorPlugin.js";
 import { AIRequestHelpers, InvokeAIOptions } from "./types.js";
 
-type ReadonlyStoreApi<T> = Pick<
-  StoreApi<T>,
-  "getState" | "getInitialState" | "subscribe"
->;
-
 type AIPluginState = {
-  // zustand design considerations:
-  //  - moved this to a nested object to have better typescript typing
-  //  - if we'd do this without a nested object, then we could easily set "wrong" values,
-  //    because "setState" takes a partial object (unless the second parameter "replace" = true),
-  //    and thus we'd lose typescript's typing help
-
   aiMenuState:
     | ({
         /**
@@ -58,228 +51,211 @@ type AIPluginState = {
 
 const PLUGIN_KEY = new PluginKey(`blocknote-ai-plugin`);
 
-export class AIExtension extends BlockNoteExtension {
-  private chatSession:
-    | {
-        previousRequestOptions: InvokeAIOptions;
-        chat: Chat<UIMessage>;
-      }
-    | undefined;
-
-  private scrollInProgress = false;
-  private autoScroll = false;
-
-  public static key(): string {
-    return "ai";
-  }
-
-  // internal store including setters
-  private readonly _store = createStore<AIPluginState>()((_set) => ({
-    aiMenuState: "closed",
-  }));
-
-  /**
-   * Returns a read-only zustand store with the state of the AI Extension
-   */
-  public get store() {
-    // externally, don't expose setters (if we want to do so, expose `set` methods seperately)
-    return this._store as ReadonlyStoreApi<AIPluginState>;
-  }
-
-  /**
-   * Returns a zustand store with the global configuration of the AI Extension.
-   * These options are used by default across all LLM calls when calling {@link executeLLMRequest}
-   */
-  public readonly options: ReturnType<
-    ReturnType<typeof createStore<AIRequestHelpers>>
-  >;
-
-  /**
-   * @internal use `createAIExtension` instead
-   */
-  constructor(
-    public readonly editor: BlockNoteEditor<any, any, any>,
-    options: AIRequestHelpers & {
-      /**
-       * The name and color of the agent cursor
-       *
-       * @default { name: "AI", color: "#8bc6ff" }
-       */
-      agentCursor?: { name: string; color: string };
-    },
-  ) {
-    super();
-
-    this.options = createStore<AIRequestHelpers>()((_set) => ({
-      ...options,
-    }));
-
-    this.addProsemirrorPlugin(
-      new Plugin({
-        key: PLUGIN_KEY,
-        filterTransaction: (tr) => {
-          const menuState = this.store.getState().aiMenuState;
-
-          if (menuState !== "closed" && menuState.status === "ai-writing") {
-            if (tr.getMeta(fixTablesKey)?.fixTables) {
-              // the fixtables plugin causes the steps between of the AI Agent to become invalid
-              // so we need to prevent it from running
-              // (we might need to filter out other / or maybe any transactions during the writing phase)
-              return false;
-            }
-          }
-          return true;
-        },
-      }),
-    );
-    this.addProsemirrorPlugin(suggestChanges());
-    this.addProsemirrorPlugin(
-      createAgentCursorPlugin(
-        options.agentCursor || { name: "AI", color: "#8bc6ff" },
-      ),
-    );
-
-    // Listens for `scroll` and `scrollend` events to see if a new scroll was
-    // started before an existing one ended. This is the most reliable way we
-    // have of checking if a scroll event was caused by the user and not by
-    // `scrollIntoView`, as the events are otherwise indistinguishable. If a
-    // scroll was started before an existing one finished (meaning the user has
-    // scrolled), auto scrolling is disabled.
-    document.addEventListener(
-      "scroll",
-      () => {
-        if (this.scrollInProgress) {
-          this.autoScroll = false;
+export const AIExtension = createExtension(
+  ({
+    editor,
+    options: editorOptions,
+  }: ExtensionOptions<
+    | (AIRequestHelpers & {
+        /**
+         * The name and color of the agent cursor
+         *
+         * @default { name: "AI", color: "#8bc6ff" }
+         */
+        agentCursor?: { name: string; color: string };
+      })
+    | undefined
+  >) => {
+    // TODO should we really expose it like this?
+    const options = createStore<AIRequestHelpers>(editorOptions ?? {});
+    const store = createStore<AIPluginState>({
+      aiMenuState: "closed",
+    });
+    let chatSession:
+      | {
+          previousRequestOptions: InvokeAIOptions;
+          chat: Chat<UIMessage>;
         }
+      | undefined;
+    let autoScroll = false;
+
+    return {
+      key: "ai",
+      options,
+      store,
+      mount({ signal }: { signal: AbortSignal }) {
+        let scrollInProgress = false;
+        // Listens for `scroll` and `scrollend` events to see if a new scroll was
+        // started before an existing one ended. This is the most reliable way we
+        // have of checking if a scroll event was caused by the user and not by
+        // `scrollIntoView`, as the events are otherwise indistinguishable. If a
+        // scroll was started before an existing one finished (meaning the user has
+        // scrolled), auto scrolling is disabled.
+        document.addEventListener(
+          "scroll",
+          () => {
+            if (scrollInProgress) {
+              autoScroll = false;
+            }
 
-        this.scrollInProgress = true;
+            scrollInProgress = true;
+          },
+          {
+            capture: true,
+            signal,
+          },
+        );
+        document.addEventListener(
+          "scrollend",
+          () => {
+            scrollInProgress = false;
+          },
+          {
+            capture: true,
+            signal,
+          },
+        );
       },
-      true,
-    );
-    document.addEventListener(
-      "scrollend",
-      () => {
-        this.scrollInProgress = false;
+      prosemirrorPlugins: [
+        new Plugin({
+          key: PLUGIN_KEY,
+          filterTransaction: (tr) => {
+            const menuState = store.state.aiMenuState;
+
+            if (menuState !== "closed" && menuState.status === "ai-writing") {
+              if (tr.getMeta(fixTablesKey)?.fixTables) {
+                // the fixtables plugin causes the steps between of the AI Agent to become invalid
+                // so we need to prevent it from running
+                // (we might need to filter out other / or maybe any transactions during the writing phase)
+                return false;
+              }
+            }
+            return true;
+          },
+        }),
+        suggestChanges(),
+        createAgentCursorPlugin(
+          editorOptions?.agentCursor || { name: "AI", color: "#8bc6ff" },
+        ),
+      ],
+
+      /**
+       * Open the AI menu at a specific block
+       */
+      openAIMenuAtBlock(blockID: string) {
+        editor.getExtension(ShowSelectionExtension)?.showSelection(true);
+        editor.isEditable = false;
+        store.setState({
+          aiMenuState: {
+            blockId: blockID,
+            status: "user-input",
+          },
+        });
+
+        // Scrolls to the block when the menu opens.
+        const blockElement = editor.domElement?.querySelector(
+          `[data-node-type="blockContainer"][data-id="${blockID}"]`,
+        );
+        blockElement?.scrollIntoView({ block: "center" });
       },
-      true,
-    );
-  }
-
-  /**
-   * Open the AI menu at a specific block
-   */
-  public openAIMenuAtBlock(blockID: string) {
-    this.editor.setForceSelectionVisible(true);
-    this.editor.isEditable = false;
-    this._store.setState({
-      aiMenuState: {
-        blockId: blockID,
-        status: "user-input",
+
+      /**
+       * Close the AI menu
+       */
+      closeAIMenu() {
+        store.setState({
+          aiMenuState: "closed",
+        });
+        chatSession = undefined;
+        editor.getExtension(ShowSelectionExtension)?.showSelection(false);
+        editor.isEditable = true;
+        editor.focus();
       },
-    });
 
-    // Scrolls to the block when the menu opens.
-    const blockElement = this.editor.domElement?.querySelector(
-      `[data-node-type="blockContainer"][data-id="${blockID}"]`,
-    );
-    blockElement?.scrollIntoView({ block: "center" });
-  }
-
-  /**
-   * Close the AI menu
-   */
-  public closeAIMenu() {
-    this._store.setState({
-      aiMenuState: "closed",
-    });
-    this.chatSession = undefined;
-    this.editor.setForceSelectionVisible(false);
-    this.editor.isEditable = true;
-    this.editor.focus();
-  }
-
-  /**
-   * Accept the changes made by the LLM
-   */
-  public acceptChanges() {
-    // This is slightly convoluted, to try to maintain the undo history as much as possible
-    // The idea is that the LLM call has appended a number of updates to the document, moving the document from state `A` to state `C`
-    // But we want the undo history to skip all of the intermediate states and go straight from `C` to `A`
-    // To do this, we capture the document state `C` (post-LLM call), and then reject the suggestions to recover the original document state `A`
-    // Then we create an intermediate state `B` that captures the diff between `A` and `C`
-    // Then we apply the suggestions to `B` to get the final state `C`
-    // This causes the undo history to skip `B` and go straight from `C` back to `A`
-
-    // Capture the document state `C'` (post-LLM call with all suggestions still in the document)
-    const markedUpDocument = this.editor.prosemirrorState.doc;
-
-    // revert the suggestions to get back to the original document state `A`
-    this.editor.exec((state, dispatch) => {
-      return revertSuggestions(state, (tr) => {
-        dispatch?.(tr.setMeta("addToHistory", false));
-      });
-    });
+      /**
+       * Accept the changes made by the LLM
+       */
+      acceptChanges() {
+        // This is slightly convoluted, to try to maintain the undo history as much as possible
+        // The idea is that the LLM call has appended a number of updates to the document, moving the document from state `A` to state `C`
+        // But we want the undo history to skip all of the intermediate states and go straight from `C` to `A`
+        // To do this, we capture the document state `C` (post-LLM call), and then reject the suggestions to recover the original document state `A`
+        // Then we create an intermediate state `B` that captures the diff between `A` and `C`
+        // Then we apply the suggestions to `B` to get the final state `C`
+        // This causes the undo history to skip `B` and go straight from `C` back to `A`
+
+        // Capture the document state `C'` (post-LLM call with all suggestions still in the document)
+        const markedUpDocument = editor.prosemirrorState.doc;
+
+        // revert the suggestions to get back to the original document state `A`
+        editor.exec((state, dispatch) => {
+          return revertSuggestions(state, (tr) => {
+            dispatch?.(tr.setMeta("addToHistory", false));
+          });
+        });
 
-    // Create an intermediate state `B` that captures the diff between the original document and the marked up document
-    this.editor.exec((state, dispatch) => {
-      const tr = state.tr;
-      tr.replace(
-        0,
-        tr.doc.content.size,
-        new Slice(Fragment.from(markedUpDocument), 0, 0),
-      );
-      const nextState = state.apply(tr);
-      // Apply the suggestions to the intermediate state `B` to get the final state `C`
-      return applySuggestions(nextState, (resultTr) => {
-        dispatch?.(
+        // Create an intermediate state `B` that captures the diff between the original document and the marked up document
+        editor.exec((state, dispatch) => {
+          const tr = state.tr;
           tr.replace(
             0,
             tr.doc.content.size,
-            new Slice(Fragment.from(resultTr.doc), 0, 0),
-          ),
-        );
-      });
-    });
+            new Slice(Fragment.from(markedUpDocument), 0, 0),
+          );
+          const nextState = state.apply(tr);
+          // Apply the suggestions to the intermediate state `B` to get the final state `C`
+          return applySuggestions(nextState, (resultTr) => {
+            dispatch?.(
+              tr.replace(
+                0,
+                tr.doc.content.size,
+                new Slice(Fragment.from(resultTr.doc), 0, 0),
+              ),
+            );
+          });
+        });
 
-    // If in collaboration mode, merge the changes back into the original yDoc
-    this.editor.forkYDocPlugin?.merge({ keepChanges: true });
-
-    this.closeAIMenu();
-  }
-
-  /**
-   * Reject the changes made by the LLM
-   */
-  public rejectChanges() {
-    // Revert the suggestions to get back to the original document
-    this.editor.exec((state, dispatch) => {
-      return revertSuggestions(state, (tr) => {
-        // Do so without adding to history (so the last undo step is just prior to the LLM call)
-        dispatch?.(tr.setMeta("addToHistory", false));
-      });
-    });
+        // If in collaboration mode, merge the changes back into the original yDoc
+        editor.getExtension(ForkYDocExtension)?.merge({ keepChanges: true });
+
+        this.closeAIMenu();
+      },
 
-    // If in collaboration mode, discard the changes and revert to the original yDoc
-    this.editor.forkYDocPlugin?.merge({ keepChanges: false });
-    this.closeAIMenu();
-  }
-
-  /**
-   * Retry the previous LLM call.
-   *
-   * Only valid if the current status is "error"
-   */
-  public async retry() {
-    const { aiMenuState } = this.store.getState();
-    if (
-      aiMenuState === "closed" ||
-      aiMenuState.status !== "error" ||
-      !this.chatSession
-    ) {
-      throw new Error("retry() is only valid when a previous response failed");
-    }
-
-    /*
+      /**
+       * Reject the changes made by the LLM
+       */
+      rejectChanges() {
+        // Revert the suggestions to get back to the original document
+        editor.exec((state, dispatch) => {
+          return revertSuggestions(state, (tr) => {
+            // Do so without adding to history (so the last undo step is just prior to the LLM call)
+            dispatch?.(tr.setMeta("addToHistory", false));
+          });
+        });
+
+        // If in collaboration mode, discard the changes and revert to the original yDoc
+        editor.getExtension(ForkYDocExtension)?.merge({ keepChanges: false });
+        this.closeAIMenu();
+      },
+
+      /**
+       * Retry the previous LLM call.
+       *
+       * Only valid if the current status is "error"
+       */
+      async retry() {
+        const { aiMenuState } = store.state;
+        if (
+          aiMenuState === "closed" ||
+          aiMenuState.status !== "error" ||
+          !chatSession
+        ) {
+          throw new Error(
+            "retry() is only valid when a previous response failed",
+          );
+        }
+
+        /*
     Design decisions:
     - we cannot use chat.regenerate() because the document might have been updated already by toolcalls that failed mid-way
     - we also cannot revert the document changes and use chat.regenerate(), 
@@ -289,198 +265,179 @@ export class AIExtension extends BlockNoteExtension {
       when a call fails
     */
 
-    if (this.chatSession?.chat.status === "error") {
-      // the LLM call failed (i.e. a network error)
-      // console.log("retry failed LLM call", this.chatSession.chat.error);
-      return this.invokeAI({
-        ...this.chatSession.previousRequestOptions,
-        userPrompt: `An error occured in the previous request. Please retry to accomplish the last user prompt.`,
-      });
-    } else {
-      // an error occurred while parsing / executing the previous LLM call
-      // give the LLM a chance to fix the error
-      // console.log("retry failed tool execution");
-      return this.invokeAI({
-        ...this.chatSession.previousRequestOptions,
-        userPrompt: `An error occured while executing the previous tool call. Please retry to accomplish the last user prompt.`,
-      });
-    }
-  }
-
-  /**
-   * Update the status of a call to an LLM
-   *
-   * @warning This method should usually only be used for advanced use-cases
-   * if you want to implement how an LLM call is executed. Usually, you should
-   * use {@link executeLLMRequest} instead which will handle the status updates for you.
-   */
-  public setAIResponseStatus(
-    status:
-      | "user-input"
-      | "thinking"
-      | "ai-writing"
-      | "user-reviewing"
-      | {
-          status: "error";
-          error: any;
-        },
-  ) {
-    const aiMenuState = this.store.getState().aiMenuState;
-    if (aiMenuState === "closed") {
-      return;
-    }
-
-    if (status === "ai-writing") {
-      this.editor.setForceSelectionVisible(false);
-    }
-
-    if (typeof status === "object") {
-      if (status.status !== "error") {
-        throw new UnreachableCaseError(status.status);
-      }
-      this._store.setState({
-        aiMenuState: {
-          status: status.status,
-          error: status.error,
-          blockId: aiMenuState.blockId,
-        },
-      });
-    } else {
-      this._store.setState({
-        aiMenuState: {
-          status: status,
-          blockId: aiMenuState.blockId,
-        },
-      });
-    }
-  }
-
-  /**
-   * @deprecated Use {@link invokeAI} instead
-   */
-  public async callLLM(opts: InvokeAIOptions) {
-    return this.invokeAI(opts);
-  }
-
-  /**
-   * Execute a call to an LLM and apply the result to the editor
-   */
-  public async invokeAI(opts: InvokeAIOptions) {
-    this.setAIResponseStatus("thinking");
-    this.editor.forkYDocPlugin?.fork();
-
-    try {
-      if (!this.chatSession) {
-        // note: in the current implementation opts.transport is only used when creating a new chat
-        // (so changing transport for a subsequent call in the same chat-session is not supported)
-        this.chatSession = {
-          previousRequestOptions: opts,
-          chat: new Chat<UIMessage>({
-            sendAutomaticallyWhen: () => false,
-            transport: opts.transport || this.options.getState().transport,
-          }),
-        };
-      } else {
-        this.chatSession.previousRequestOptions = opts;
-      }
-      const chat = this.chatSession.chat;
-
-      // merge the global options with the local options
-      const globalOpts = this.options.getState();
-      opts = {
-        ...globalOpts,
-        ...opts,
-      } as InvokeAIOptions;
-
-      const sender =
-        opts.aiRequestSender ??
-        defaultAIRequestSender(
-          aiDocumentFormats.html.defaultPromptBuilder,
-          aiDocumentFormats.html.defaultPromptInputDataBuilder,
-        );
+        if (chatSession?.chat.status === "error") {
+          // the LLM call failed (i.e. a network error)
+          // console.log("retry failed LLM call", this.chatSession.chat.error);
+          return this.invokeAI({
+            ...chatSession.previousRequestOptions,
+            userPrompt: `An error occured in the previous request. Please retry to accomplish the last user prompt.`,
+          });
+        } else {
+          // an error occurred while parsing / executing the previous LLM call
+          // give the LLM a chance to fix the error
+          // console.log("retry failed tool execution");
+          return this.invokeAI({
+            ...chatSession.previousRequestOptions,
+            userPrompt: `An error occured while executing the previous tool call. Please retry to accomplish the last user prompt.`,
+          });
+        }
+      },
+
+      /**
+       * Update the status of a call to an LLM
+       *
+       * @warning This method should usually only be used for advanced use-cases
+       * if you want to implement how an LLM call is executed. Usually, you should
+       * use {@link executeLLMRequest} instead which will handle the status updates for you.
+       */
+      setAIResponseStatus(
+        status:
+          | "user-input"
+          | "thinking"
+          | "ai-writing"
+          | "user-reviewing"
+          | {
+              status: "error";
+              error: any;
+            },
+      ) {
+        const aiMenuState = store.state.aiMenuState;
+        if (aiMenuState === "closed") {
+          return;
+        }
+
+        if (status === "ai-writing") {
+          editor.getExtension(ShowSelectionExtension)?.showSelection(false);
+        }
 
-      const aiRequest = buildAIRequest({
-        editor: this.editor,
-        chat,
-        userPrompt: opts.userPrompt,
-        useSelection: opts.useSelection,
-        deleteEmptyCursorBlock: opts.deleteEmptyCursorBlock,
-        streamToolsProvider: opts.streamToolsProvider,
-        onBlockUpdated: (blockId) => {
-          // NOTE: does this setState with an anon object trigger unnecessary re-renders?
-          this._store.setState({
+        if (typeof status === "object") {
+          if (status.status !== "error") {
+            throw new UnreachableCaseError(status.status);
+          }
+          store.setState({
+            aiMenuState: {
+              status: status.status,
+              error: status.error,
+              blockId: aiMenuState.blockId,
+            },
+          });
+        } else {
+          store.setState({
             aiMenuState: {
-              blockId,
-              status: "ai-writing",
+              status: status,
+              blockId: aiMenuState.blockId,
             },
           });
+        }
+      },
 
-          // Scrolls to the block being edited by the AI while auto scrolling is
-          // enabled.
-          if (!this.autoScroll) {
-            return;
-          }
+      /**
+       * @deprecated Use {@link invokeAI} instead
+       */
+      async callLLM(opts: InvokeAIOptions) {
+        return this.invokeAI(opts);
+      },
 
-          const aiMenuState = this._store.getState().aiMenuState;
-          const aiMenuOpenState =
-            aiMenuState === "closed" ? undefined : aiMenuState;
-          if (!aiMenuOpenState || aiMenuOpenState.status !== "ai-writing") {
-            return;
+      /**
+       * Execute a call to an LLM and apply the result to the editor
+       */
+      async invokeAI(opts: InvokeAIOptions) {
+        this.setAIResponseStatus("thinking");
+        editor.getExtension(ForkYDocExtension)?.fork();
+
+        try {
+          if (!chatSession) {
+            // note: in the current implementation opts.transport is only used when creating a new chat
+            // (so changing transport for a subsequent call in the same chat-session is not supported)
+            chatSession = {
+              previousRequestOptions: opts,
+              chat: new Chat<UIMessage>({
+                sendAutomaticallyWhen: () => false,
+                transport: opts.transport || options.state.transport,
+              }),
+            };
+          } else {
+            chatSession.previousRequestOptions = opts;
           }
+          const chat = chatSession.chat;
+
+          // merge the global options with the local options
+          const globalOpts = options.state;
+          opts = {
+            ...globalOpts,
+            ...opts,
+          } as InvokeAIOptions;
+
+          const sender =
+            opts.aiRequestSender ??
+            defaultAIRequestSender(
+              aiDocumentFormats.html.defaultPromptBuilder,
+              aiDocumentFormats.html.defaultPromptInputDataBuilder,
+            );
+
+          const aiRequest = buildAIRequest({
+            editor,
+            chat,
+            userPrompt: opts.userPrompt,
+            useSelection: opts.useSelection,
+            deleteEmptyCursorBlock: opts.deleteEmptyCursorBlock,
+            streamToolsProvider: opts.streamToolsProvider,
+            onBlockUpdated: (blockId) => {
+              const aiMenuState = store.state.aiMenuState;
+              const aiMenuOpenState =
+                aiMenuState === "closed" ? undefined : aiMenuState;
+              if (!aiMenuOpenState || aiMenuOpenState.status !== "ai-writing") {
+                return;
+              }
+
+              // TODO: Sometimes, the updated block doesn't actually exist in
+              // the editor. I don't know why this happens, seems like a bug?
+              const nodeInfo = getNodeById(
+                blockId,
+                editor.prosemirrorState.doc,
+              );
+              if (!nodeInfo) {
+                return;
+              }
+
+              store.setState({
+                aiMenuState: {
+                  blockId,
+                  status: "ai-writing",
+                },
+              });
+
+              if (autoScroll) {
+                const blockElement = editor.prosemirrorView.domAtPos(
+                  nodeInfo.posBeforeNode + 1,
+                );
+                (blockElement.node as HTMLElement).scrollIntoView({
+                  block: "center",
+                });
+              }
+            },
+          });
 
-          const nodeInfo = getNodeById(
-            aiMenuOpenState.blockId,
-            this.editor.prosemirrorState.doc,
-          );
-          if (!nodeInfo) {
-            return;
-          }
+          await executeAIRequest({
+            aiRequest,
+            sender,
+            chatRequestOptions: opts.chatRequestOptions,
+            onStart: () => {
+              autoScroll = true;
+              this.setAIResponseStatus("ai-writing");
+            },
+          });
 
-          const blockElement = this.editor.prosemirrorView.domAtPos(
-            nodeInfo.posBeforeNode + 1,
-          );
-          (blockElement.node as HTMLElement).scrollIntoView({
-            block: "center",
+          this.setAIResponseStatus("user-reviewing");
+        } catch (e) {
+          this.setAIResponseStatus({
+            status: "error",
+            error: e,
           });
-        },
-      });
-
-      await executeAIRequest({
-        aiRequest,
-        sender,
-        chatRequestOptions: opts.chatRequestOptions,
-        onStart: () => {
-          this.autoScroll = true;
-          this.setAIResponseStatus("ai-writing");
-        },
-      });
-
-      this.setAIResponseStatus("user-reviewing");
-    } catch (e) {
-      this.setAIResponseStatus({
-        status: "error",
-        error: e,
-      });
-      // eslint-disable-next-line no-console
-      console.warn("Error calling LLM", e, this.chatSession?.chat.messages);
-    }
-  }
-}
-
-/**
- * Create a new AIExtension instance, this can be passed to the BlockNote editor via the `extensions` option
- */
-export function createAIExtension(
-  options: ConstructorParameters<typeof AIExtension>[1],
-) {
-  return (editor: BlockNoteEditor<any, any, any>) => {
-    return new AIExtension(editor, options);
-  };
-}
-
-/**
- * Return the AIExtension instance from the editor
- */
-export function getAIExtension(editor: BlockNoteEditor<any, any, any>) {
-  return editor.extension(AIExtension);
-}
+          // eslint-disable-next-line no-console
+          console.warn("Error calling LLM", e, chatSession?.chat.messages);
+        }
+      },
+    } as const;
+  },
+);