support selections

Author: YousefED

packages/core/src/api/nodeConversions/nodeToBlock.ts

@@ -425,6 +425,36 @@ export function nodeToBlock<
   return block;
 }
 
+/**
+ * Convert a Prosemirror document to a BlockNote document (array of blocks)
+ */
+export function docToBlocks<
+  BSchema extends BlockSchema,
+  I extends InlineContentSchema,
+  S extends StyleSchema
+>(
+  doc: Node,
+  blockSchema: BSchema,
+  inlineContentSchema: I,
+  styleSchema: S,
+  blockCache?: WeakMap<Node, Block<BSchema, I, S>>
+) {
+  const blocks: Block<BSchema, I, S>[] = [];
+  doc.firstChild!.descendants((node) => {
+    blocks.push(
+      nodeToBlock(
+        node,
+        blockSchema,
+        inlineContentSchema,
+        styleSchema,
+        blockCache
+      )
+    );
+    return false;
+  });
+  return blocks;
+}
+
 export type SlicedBlock<
   BSchema extends BlockSchema,
   I extends InlineContentSchema,
@@ -462,18 +492,18 @@ export function selectionToInsertionEnd(tr: Transaction, startLen: number) {
   return end;
 }
 
-export function withSelectionMarkers<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema
->(
+/**
+ * Create a transaction that adds selection markers to the document at the given positions.
+ *
+ * @param state - The editor state.
+ * @param from - The start position of the selection.
+ * @param to - The end position of the selection.
+ * @returns The transaction and the new end position.
+ */
+export function addSelectionMarkersTr(
   state: EditorState,
   from: number,
-  to: number,
-  blockSchema: BSchema,
-  inlineContentSchema: I,
-  styleSchema: S,
-  blockCache?: WeakMap<Node, Block<BSchema, I, S>>
+  to: number
 ) {
   if (from >= to) {
     throw new Error("from must be less than to");
@@ -490,6 +520,52 @@ export function withSelectionMarkers<
       "tr.docChanged is false or insertText was not applied. Was a valid textselection passed?"
     );
   }
+  return {
+    tr,
+    newEnd,
+  };
+}
+
+export function getDocumentWithSelectionMarkers<
+  BSchema extends BlockSchema,
+  I extends InlineContentSchema,
+  S extends StyleSchema
+>(
+  state: EditorState,
+  from: number,
+  to: number,
+  blockSchema: BSchema,
+  inlineContentSchema: I,
+  styleSchema: S,
+  blockCache?: WeakMap<Node, Block<BSchema, I, S>>
+) {
+  const { tr } = addSelectionMarkersTr(state, from, to);
+  return docToBlocks(
+    tr.doc,
+    blockSchema,
+    inlineContentSchema,
+    styleSchema,
+    blockCache
+  );
+}
+
+/**
+ * Add selection markers to the document at the given positions and return the blocks that span the selection.
+ */
+export function getSelectedBlocksWithSelectionMarkers<
+  BSchema extends BlockSchema,
+  I extends InlineContentSchema,
+  S extends StyleSchema
+>(
+  state: EditorState,
+  from: number,
+  to: number,
+  blockSchema: BSchema,
+  inlineContentSchema: I,
+  styleSchema: S,
+  blockCache?: WeakMap<Node, Block<BSchema, I, S>>
+) {
+  const { tr, newEnd } = addSelectionMarkersTr(state, from, to);
 
   return getBlocksBetween(
     from,
@@ -597,6 +673,8 @@ export function withoutSliceMetadata<
  *     childrenCutAtEnd: true,
  *   },
  * ]
+ *
+ * TODO: do we actually need / use this?
  */
 export function prosemirrorSliceToSlicedBlocks<
   BSchema extends BlockSchema,

packages/core/src/api/nodeConversions/selectionMarkers.test.ts

@@ -160,7 +160,7 @@ describe("Test ProseMirror selection HTML conversion", () => {
 
     // const slice = editor._tiptapEditor.state.selection.content();
 
-    const blockNoteSelection = editor.getSelectionWithMarkers();
+    const blockNoteSelection = editor.getSelectedBlocksWithSelectionMarkers();
 
     expect(
       JSON.stringify(blockNoteSelection, undefined, 2)
@@ -282,7 +282,7 @@ describe("Test ProseMirror selection HTML conversion", () => {
       }
       editor.dispatch(editor._tiptapEditor.state.tr.setSelection(selection));
 
-      const blockNoteSelection = editor.getSelectionWithMarkers();
+      const blockNoteSelection = editor.getSelectedBlocksWithSelectionMarkers();
       const JSONString = JSON.stringify(blockNoteSelection);
       ret += JSONString + "\n";
     });
@@ -301,7 +301,7 @@ describe("Test ProseMirror selection HTML conversion", () => {
       }
       editor.dispatch(editor._tiptapEditor.state.tr.setSelection(selection));
 
-      const blockNoteSelection = editor.getSelectionWithMarkers();
+      const blockNoteSelection = editor.getSelectedBlocksWithSelectionMarkers();
       const JSONString = JSON.stringify(blockNoteSelection);
       ret += JSONString + "\n";
     });

packages/core/src/editor/BlockNoteEditor.ts

@@ -9,12 +9,6 @@ import {
 import { Node, Schema } from "prosemirror-model";
 // import "./blocknote.css";
 import * as Y from "yjs";
-import {
-  getBlock,
-  getNextBlock,
-  getParentBlock,
-  getPrevBlock,
-} from "../api/blockManipulation/getBlock/getBlock.js";
 import { insertBlocks } from "../api/blockManipulation/commands/insertBlocks/insertBlocks.js";
 import {
   moveBlocksDown,
@@ -29,15 +23,21 @@ import {
 import { removeBlocks } from "../api/blockManipulation/commands/removeBlocks/removeBlocks.js";
 import { replaceBlocks } from "../api/blockManipulation/commands/replaceBlocks/replaceBlocks.js";
 import { updateBlock } from "../api/blockManipulation/commands/updateBlock/updateBlock.js";
-import { insertContentAt } from "../api/blockManipulation/insertContentAt.js";
 import {
-  getTextCursorPosition,
-  setTextCursorPosition,
-} from "../api/blockManipulation/selections/textCursorPosition/textCursorPosition.js";
+  getBlock,
+  getNextBlock,
+  getParentBlock,
+  getPrevBlock,
+} from "../api/blockManipulation/getBlock/getBlock.js";
+import { insertContentAt } from "../api/blockManipulation/insertContentAt.js";
 import {
   getSelection,
   setSelection,
 } from "../api/blockManipulation/selections/selection.js";
+import {
+  getTextCursorPosition,
+  setTextCursorPosition,
+} from "../api/blockManipulation/selections/textCursorPosition/textCursorPosition.js";
 import { createExternalHTMLExporter } from "../api/exporters/html/externalHTMLExporter.js";
 import { blocksToMarkdown } from "../api/exporters/markdown/markdownExporter.js";
 import { HTMLToBlocks } from "../api/parsers/html/parseHTML.js";
@@ -89,15 +89,16 @@ import { en } from "../i18n/locales/index.js";
 
 import { Plugin, Transaction } from "@tiptap/pm/state";
 import { dropCursor } from "prosemirror-dropcursor";
+import { EditorView } from "prosemirror-view";
 import { createInternalHTMLSerializer } from "../api/exporters/html/internalHTMLSerializer.js";
 import { inlineContentToNodes } from "../api/nodeConversions/blockToNode.js";
 import {
-  nodeToBlock,
+  docToBlocks,
+  getDocumentWithSelectionMarkers,
+  getSelectedBlocksWithSelectionMarkers,
   prosemirrorSliceToSlicedBlocks,
-  withSelectionMarkers,
 } from "../api/nodeConversions/nodeToBlock.js";
 import "../style.css";
-import { EditorView } from "prosemirror-view";
 
 export type BlockNoteExtension =
   | AnyExtension
@@ -639,23 +640,13 @@ export class BlockNoteEditor<
    * @returns A snapshot of all top-level (non-nested) blocks in the editor.
    */
   public get document(): Block<BSchema, ISchema, SSchema>[] {
-    const blocks: Block<BSchema, ISchema, SSchema>[] = [];
-
-    this._tiptapEditor.state.doc.firstChild!.descendants((node) => {
-      blocks.push(
-        nodeToBlock(
-          node,
-          this.schema.blockSchema,
-          this.schema.inlineContentSchema,
-          this.schema.styleSchema,
-          this.blockCache
-        )
-      );
-
-      return false;
-    });
-
-    return blocks;
+    return docToBlocks(
+      this._tiptapEditor.state.doc,
+      this.schema.blockSchema,
+      this.schema.inlineContentSchema,
+      this.schema.styleSchema,
+      this.blockCache
+    );
   }
 
   /**
@@ -792,8 +783,19 @@ export class BlockNoteEditor<
     setTextCursorPosition(this, targetBlock, placement);
   }
 
+  public getDocumentWithSelectionMarkers() {
+    return getDocumentWithSelectionMarkers(
+      this._tiptapEditor.state,
+      this._tiptapEditor.state.selection.from,
+      this._tiptapEditor.state.selection.to,
+      this.schema.blockSchema,
+      this.schema.inlineContentSchema,
+      this.schema.styleSchema
+    );
+  }
+
   // TODO: what about node selections?
-  public getSelectionWithMarkers() {
+  public getSelectedBlocksWithSelectionMarkers() {
     const start = this._tiptapEditor.state.selection.$from;
     const end = this._tiptapEditor.state.selection.$to;
 
@@ -817,7 +819,7 @@ export class BlockNoteEditor<
     //   start = this._tiptapEditor.state.doc.resolve(start.pos + 1);
     // }
 
-    return withSelectionMarkers(
+    return getSelectedBlocksWithSelectionMarkers(
       this._tiptapEditor.state,
       start.pos,
       end.pos,

packages/xl-ai/src/api/formats/json/json.ts

@@ -1,5 +1,11 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { LanguageModel, generateObject, jsonSchema, streamObject } from "ai";
+import {
+  CoreMessage,
+  LanguageModel,
+  generateObject,
+  jsonSchema,
+  streamObject,
+} from "ai";
 
 import {
   executeAIOperation,
@@ -10,7 +16,10 @@ import { deleteFunction } from "../../functions/delete.js";
 import { AIFunction } from "../../functions/index.js";
 import { updateFunction } from "../../functions/update.js";
 import type { PromptOrMessages } from "../../index.js";
-import { promptManipulateDocumentUseJSONSchema } from "../../prompts/jsonSchemaPrompts.js";
+import {
+  promptManipulateDocumentUseJSONSchema,
+  promptManipulateSelectionJSONSchema,
+} from "../../prompts/jsonSchemaPrompts.js";
 import { createOperationsArraySchema } from "../../schema/operations.js";
 import { blockNoteSchemaToJSONSchema } from "../../schema/schemaToJSONSchema.js";
 
@@ -43,19 +52,30 @@ export async function callLLM(
   editor: BlockNoteEditor<any, any, any>,
   opts: CallLLMOptionsWithOptional
 ) {
-  const { prompt, ...rest } = opts;
+  const { prompt, useSelection, ...rest } = opts;
+
+  let messages: CoreMessage[];
+
+  if ("messages" in opts && opts.messages) {
+    messages = opts.messages;
+  } else if (useSelection) {
+    messages = promptManipulateSelectionJSONSchema({
+      editor,
+      userPrompt: opts.prompt!,
+      document: editor.getDocumentWithSelectionMarkers(),
+    });
+  } else {
+    messages = promptManipulateDocumentUseJSONSchema({
+      editor,
+      userPrompt: opts.prompt!,
+      document: editor.document,
+    });
+  }
 
   const options: CallLLMOptions = {
     functions: [updateFunction, addFunction, deleteFunction],
     stream: true,
-    messages:
-      "messages" in opts && opts.messages !== undefined
-        ? opts.messages
-        : promptManipulateDocumentUseJSONSchema({
-            editor,
-            userPrompt: opts.prompt!,
-            document: editor.document,
-          }),
+    messages,
     ...rest,
   };
 

packages/xl-ai/src/api/formats/markdown/markdown.ts

@@ -7,7 +7,10 @@ import { addFunction } from "../../functions/add.js";
 import { deleteFunction } from "../../functions/delete.js";
 import { updateFunction } from "../../functions/update.js";
 import type { PromptOrMessages } from "../../index.js";
-import { promptManipulateDocumentUseMarkdown } from "../../prompts/markdownPrompts.js";
+import {
+  promptManipulateDocumentUseMarkdown,
+  promptManipulateDocumentUseMarkdownWithSelection,
+} from "../../prompts/markdownPrompts.js";
 import { trimArray } from "../../util/trimArray.js";
 
 type BasicLLMRequestOptions = {
@@ -22,19 +25,31 @@ export async function callLLM(
   editor: BlockNoteEditor<any, any, any>,
   options: MarkdownLLMRequestOptions
 ) {
+  let messages: CoreMessage[];
   const markdown = await editor.blocksToMarkdownLossy();
 
+  if ("messages" in options && options.messages) {
+    messages = options.messages;
+  } else if (options.useSelection) {
+    const selection = editor.getDocumentWithSelectionMarkers();
+    const markdown = await editor.blocksToMarkdownLossy(selection);
+    messages = promptManipulateDocumentUseMarkdownWithSelection({
+      editor,
+      markdown,
+      userPrompt: (options as any).prompt,
+    });
+  } else {
+    messages = promptManipulateDocumentUseMarkdown({
+      editor,
+      markdown,
+      userPrompt: (options as any).prompt,
+    });
+  }
+
   const withDefaults: Required<
     Omit<MarkdownLLMRequestOptions & { messages: CoreMessage[] }, "prompt">
   > = {
-    messages:
-      "messages" in options
-        ? options.messages
-        : promptManipulateDocumentUseMarkdown({
-            editor,
-            markdown,
-            userPrompt: (options as any).prompt,
-          }),
+    messages,
     ...(options as any), // TODO
   };