support selections

TypeCellOS · YousefED · May 20, 2025 · Sep 26, 2024 · Sep 26, 2024 · Sep 26, 2024
commit 78b8f98cdf17d8ad84fd2eae2a55756a961488d8
diff --git a/packages/core/src/api/nodeConversions/nodeToBlock.ts b/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,

diff --git a/packages/core/src/api/nodeConversions/selectionMarkers.test.ts b/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";
    });

diff --git a/packages/core/src/editor/BlockNoteEditor.ts b/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,

diff --git a/packages/xl-ai/src/api/formats/json/json.ts b/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,
  };


diff --git a/packages/xl-ai/src/api/formats/markdown/markdown.ts b/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
  };