feat: AI block UI

examples/09-ai/01-minimal/App.tsx

@@ -1,11 +1,11 @@
 import {
   AIToolbarButton,
-  BlockNoteAIContextProvider,
   BlockNoteAIUI,
   locales as aiLocales,
+  createAIExtension,
   createBlockNoteAIClient,
+  getAIExtension,
   getAISlashMenuItems,
-  useBlockNoteAIContext,
 } from "@blocknote/xl-ai";
 
 import { createGroq } from "@ai-sdk/groq";
@@ -32,6 +32,7 @@ import {
 import "@blocknote/xl-ai/style.css";
 import { Fieldset, Switch } from "@mantine/core";
 import { useMemo, useState } from "react";
+import { useStore } from "zustand";
 import { BasicAutocomplete } from "./AutoComplete.js";
 import RadioGroupComponent from "./components/RadioGroupComponent.js";
 
@@ -51,16 +52,9 @@ const client = createBlockNoteAIClient({
 
 export default function App() {
   const [aiModelString, setAiModelString] = useState(
-    "openai/gpt-4o-2024-08-06"
+    // "openai/gpt-4o-2024-08-06"
+    "groq/llama-3.3-70b-specdec"
   );
-  // Creates a new editor instance.
-  const editor = useCreateBlockNote({
-    schema,
-    dictionary: {
-      ...en,
-      ai: aiLocales.en,
-    } as any,
-  });
 
   const model = useMemo(() => {
     const [provider, ...modelNameParts] = aiModelString.split("/");
@@ -84,13 +78,59 @@ export default function App() {
     throw new Error(`Unknown model: ${aiModelString}`);
   }, [aiModelString]);
 
-  const [dataFormat, setDataFormat] = useState<"json" | "markdown" | "html">(
-    "json"
-  );
+  // Creates a new editor instance.
+  const editor = useCreateBlockNote({
+    schema,
+    dictionary: {
+      ...en,
+      ai: aiLocales.en,
+    } as any,
+    _extensions: {
+      // TO DISCUSS
+      // we pass the extension here.
+      // I don't really like the fact that "ai" is hardcoded - I think we need an API that also allows us to pass in extensions without specifying the key
+      ai: createAIExtension({
+        model,
+      }),
+    },
+    initialContent: [
+      {
+        type: "heading",
+        props: {
+          level: 1,
+        },
+        content: "I love cats",
+      },
+      {
+        type: "paragraph",
+        content:
+          "Cats are one of the most beloved and fascinating animals in the world. Known for their agility, independence, and charm, cats have been companions to humans for thousands of years. Domesticated cats, scientifically named Felis catus, come in various breeds, colors, and personalities, making them a popular choice for pet owners everywhere. Their mysterious behavior, sharp reflexes, and quiet affection have earned them a special place in countless households.",
+      },
+      {
+        type: "paragraph",
+        content:
+          "Beyond their role as pets, cats have a rich history and cultural significance. In ancient Egypt, they were revered and even worshipped as symbols of protection and grace. Throughout history, they’ve appeared in folklore, art, and literature, often associated with curiosity, luck, and mystery. Despite superstitions surrounding black cats in some cultures, many societies around the world admire and cherish these sleek and graceful animals.",
+      },
+      {
+        type: "paragraph",
+        content:
+          "Cats also offer emotional and physical benefits to their owners. Studies have shown that interacting with cats can reduce stress, lower blood pressure, and improve mental well-being. Their gentle purring, playful antics, and warm companionship provide comfort to people of all ages. Whether lounging in the sun, chasing a toy, or curling up on a lap, cats bring joy, peace, and a bit of magic to the lives of those who welcome them into their homes.",
+      },
+    ],
+  });
+
+  const ai = getAIExtension(editor);
+  // TBD: we now derive this from the LLM extension options. desirable?
+
+  // const [dataFormat, setDataFormat] = useState<"json" | "markdown" | "html">(
+  //   "html"
+  // );
 
-  const [stream, setStream] = useState(true);
+  const dataFormat = useStore(ai.options, (state) => state.dataFormat);
 
-  // const stream = dataFormat === "markdown" ? false : streamStateValue;
+  // const [stream, setStream] = useState(true);
+
+  const stream = useStore(ai.options, (state) => state.stream);
 
   // Renders the editor instance using a React component.
   return (
@@ -110,12 +150,14 @@ export default function App() {
             },
           ]}
           value={dataFormat}
-          onChange={(value) => setDataFormat(value as "json" | "markdown")}
+          onChange={(value) =>
+            ai.options.setState({ dataFormat: value as "json" | "markdown" })
+          }
         />
 
         <Switch
           checked={stream}
-          onChange={(e) => setStream(e.target.checked)}
+          onChange={(e) => ai.options.setState({ stream: e.target.checked })}
           label="Streaming"
         />
       </Fieldset>
@@ -124,40 +166,34 @@ export default function App() {
         editor={editor}
         formattingToolbar={false}
         slashMenu={false}>
-        <BlockNoteAIContextProvider
-          model={model}
-          dataFormat={dataFormat}
-          stream={stream}>
-          <BlockNoteAIUI />
-          <FormattingToolbarController
-            formattingToolbar={() => (
-              <FormattingToolbar>
-                {...getFormattingToolbarItems([
-                  ...blockTypeSelectItems(editor.dictionary),
-                  // ...aiBlockTypeSelectItems(aiLocales.en),
-                ])}
-                <AIToolbarButton />
-              </FormattingToolbar>
-            )}
-          />
-          <SuggestionMenu editor={editor} />
-        </BlockNoteAIContextProvider>
+        <BlockNoteAIUI />
+        <FormattingToolbarController
+          formattingToolbar={() => (
+            <FormattingToolbar>
+              {...getFormattingToolbarItems([
+                ...blockTypeSelectItems(editor.dictionary),
+                // ...aiBlockTypeSelectItems(aiLocales.en),
+              ])}
+              <AIToolbarButton />
+            </FormattingToolbar>
+          )}
+        />
+        <SuggestionMenu editor={editor} />
         {/* TODO: Side Menu customization */}
       </BlockNoteView>
     </div>
   );
 }
 
 function SuggestionMenu(props: { editor: BlockNoteEditor<any, any, any> }) {
-  const ctx = useBlockNoteAIContext();
   return (
     <SuggestionMenuController
       triggerCharacter="/"
       getItems={async (query) =>
         filterSuggestionItems(
           [
             ...getDefaultReactSlashMenuItems(props.editor),
-            ...getAISlashMenuItems(props.editor, ctx),
+            ...getAISlashMenuItems(props.editor),
           ],
           query
         )

packages/core/src/editor/BlockNoteEditor.ts

@@ -107,16 +107,26 @@ import { CodeBlockOptions } from "../blocks/CodeBlockContent/CodeBlockContent.js
 import type { ThreadStore, User } from "../comments/index.js";
 import "../style.css";
 import { EventEmitter } from "../util/EventEmitter.js";
+import { BlockNoteExtension } from "./BlockNoteExtension.js";
 
+/**
+ * A factory function that returns a BlockNoteExtension
+ * This is useful so we can create extensions that require an editor instance
+ * in the constructor
+ */
 export type BlockNoteExtensionFactory = (
   editor: BlockNoteEditor<any, any, any>
 ) => BlockNoteExtension;
 
-export type BlockNoteExtension =
+/**
+ * We support Tiptap extensions and BlockNoteExtension based extensions
+ */
+export type SupportedExtension =
   | AnyExtension
   | {
-      plugin: Plugin;
-    };
+      plugin: Plugin; // TODO: deprecate this format and use BlockNoteExtension instead
+    }
+  | BlockNoteExtension;
 
 export type BlockNoteEditorOptions<
   BSchema extends BlockSchema,
@@ -355,7 +365,7 @@ export class BlockNoteEditor<
   /**
    * extensions that are added to the editor, can be tiptap extensions or prosemirror plugins
    */
-  public readonly extensions: Record<string, BlockNoteExtension> = {};
+  public readonly extensions: Record<string, SupportedExtension> = {};
 
   /**
    * Boolean indicating whether the editor is in headless mode.
@@ -630,6 +640,10 @@ export class BlockNoteEditor<
 
     const tiptapExtensions = [
       ...Object.entries(this.extensions).map(([key, ext]) => {
+        if (ext instanceof BlockNoteExtension) {
+          return ext.plugin;
+        }
+
         if (
           ext instanceof Extension ||
           ext instanceof TipTapNode ||
@@ -651,7 +665,8 @@ export class BlockNoteEditor<
           addProseMirrorPlugins: () => [ext.plugin],
         });
       }),
-    ];
+    ].filter((ext): ext is Extension => ext !== undefined);
+
     const tiptapOptions: BlockNoteTipTapEditorOptions = {
       ...blockNoteTipTapOptions,
       ...newOptions._tiptapOptions,
@@ -698,6 +713,26 @@ export class BlockNoteEditor<
     this._tiptapEditor.dispatch(tr);
   };
 
+  // TO DISCUSS
+  /**
+   * Shorthand to get a typed extension from the editor, by
+   * just passing in the extension class.
+   *
+   * @param ext - The extension class to get
+   * @param key - optional, the key of the extension in the extensions object (defaults to the extension name)
+   * @returns The extension instance
+   */
+  public extension<T extends BlockNoteExtension>(
+    ext: { new (...args: any[]): T } & typeof BlockNoteExtension,
+    key = ext.name()
+  ): T {
+    const extension = this.extensions[key] as T;
+    if (!extension) {
+      throw new Error(`Extension ${key} not found`);
+    }
+    return extension;
+  }
+
   /**
    * Mount the editor to a parent DOM element. Call mount(undefined) to clean up
    *

packages/core/src/editor/BlockNoteExtension.ts

@@ -0,0 +1,14 @@
+import { Plugin } from "prosemirror-state";
+
+export abstract class BlockNoteExtension {
+  public static name(): string {
+    throw new Error("You must implement the name method in your extension");
+  }
+
+  public plugin?: Plugin;
+
+  constructor(...args: any[]) {
+    // Allow subclasses to have constructors with parameters
+    // without this, we can't easily implement BlockNoteEditor.extension(MyExtension) pattern
+  }
+}

packages/core/src/index.ts

@@ -30,6 +30,7 @@ export * from "./blocks/defaultBlockTypeGuards.js";
 export * from "./blocks/defaultBlocks.js";
 export * from "./blocks/defaultProps.js";
 export * from "./editor/BlockNoteEditor.js";
+export * from "./editor/BlockNoteExtension.js";
 export * from "./editor/BlockNoteExtensions.js";
 export * from "./editor/BlockNoteSchema.js";
 export * from "./editor/defaultColors.js";

packages/xl-ai/package.json

@@ -73,7 +73,8 @@
     "react-icons": "^5.2.1",
     "remark-parse": "^10.0.1",
     "remark-stringify": "^10.0.2",
-    "unified": "^10.1.2"
+    "unified": "^10.1.2",
+    "zustand": "^5.0.3"
   },
   "devDependencies": {
     "@types/diff": "^6.0.0",