feat(ai): expand selections to contain words (#2304)

* feat(ai): expand selection to words

* ai unit test (wip)

* fix tests

Author: YousefED

packages/core/src/api/blockManipulation/selections/selection.ts

@@ -1,6 +1,5 @@
 import { TextSelection, type Transaction } from "prosemirror-state";
 import { TableMap } from "prosemirror-tables";
-
 import { Block } from "../../../blocks/defaultBlocks.js";
 import { Selection } from "../../../editor/selectionTypes.js";
 import {
@@ -9,6 +8,7 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "../../../schema/index.js";
+import { expandPMRangeToWords } from "../../../util/expandToWords.js";
 import { getBlockInfo, getNearestBlockPos } from "../../getBlockInfoFromPos.js";
 import {
   nodeToBlock,
@@ -220,12 +220,17 @@ export function setSelection(
   tr.setSelection(TextSelection.create(tr.doc, startPos, endPos));
 }
 
-export function getSelectionCutBlocks(tr: Transaction) {
+export function getSelectionCutBlocks(tr: Transaction, expandToWords = false) {
   // TODO: fix image node selection
 
   const pmSchema = getPmSchema(tr);
-  let start = tr.selection.$from;
-  let end = tr.selection.$to;
+
+  const range = expandToWords
+    ? expandPMRangeToWords(tr.doc, tr.selection)
+    : tr.selection;
+
+  let start = range.$from;
+  let end = range.$to;
 
   // the selection moves below are used to make sure `prosemirrorSliceToSlicedBlocks` returns
   // the correct information about whether content is cut at the start or end of a block

packages/core/src/editor/BlockNoteEditor.ts

@@ -7,7 +7,6 @@ import {
 } from "@tiptap/core";
 import { type Command, type Plugin, type Transaction } from "@tiptap/pm/state";
 import { Node, Schema } from "prosemirror-model";
-
 import type { BlocksChanged } from "../api/getBlocksChangedByTransaction.js";
 import { blockToNode } from "../api/nodeConversions/blockToNode.js";
 import {
@@ -18,6 +17,8 @@ import {
   DefaultStyleSchema,
   PartialBlock,
 } from "../blocks/index.js";
+import type { CollaborationOptions } from "../extensions/Collaboration/Collaboration.js";
+import { BlockChangeExtension } from "../extensions/index.js";
 import { UniqueID } from "../extensions/tiptap-extensions/UniqueID/UniqueID.js";
 import type { Dictionary } from "../i18n/dictionary.js";
 import { en } from "../i18n/locales/index.js";
@@ -51,8 +52,6 @@ import {
 } from "./managers/index.js";
 import type { Selection } from "./selectionTypes.js";
 import { transformPasted } from "./transformPasted.js";
-import { BlockChangeExtension } from "../extensions/index.js";
-import type { CollaborationOptions } from "../extensions/Collaboration/Collaboration.js";
 
 export type BlockCache<
   BSchema extends BlockSchema = any,
@@ -934,8 +933,8 @@ export class BlockNoteEditor<
    * If the selection starts / ends halfway through a block, the returned block will be
    * only the part of the block that is included in the selection.
    */
-  public getSelectionCutBlocks() {
-    return this._selectionManager.getSelectionCutBlocks();
+  public getSelectionCutBlocks(expandToWords = false) {
+    return this._selectionManager.getSelectionCutBlocks(expandToWords);
   }
 
   /**

packages/core/src/editor/managers/SelectionManager.ts

@@ -1,3 +1,4 @@
+import { isNodeSelection, posToDOMRect } from "@tiptap/core";
 import {
   getSelection,
   getSelectionCutBlocks,
@@ -7,21 +8,20 @@ import {
   getTextCursorPosition,
   setTextCursorPosition,
 } from "../../api/blockManipulation/selections/textCursorPosition.js";
-import { isNodeSelection, posToDOMRect } from "@tiptap/core";
+import {
+  DefaultBlockSchema,
+  DefaultInlineContentSchema,
+  DefaultStyleSchema,
+} from "../../blocks/defaultBlocks.js";
 import {
   BlockIdentifier,
   BlockSchema,
   InlineContentSchema,
   StyleSchema,
 } from "../../schema/index.js";
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-} from "../../blocks/defaultBlocks.js";
-import { Selection } from "../selectionTypes.js";
-import { TextCursorPosition } from "../cursorPositionTypes.js";
 import { BlockNoteEditor } from "../BlockNoteEditor.js";
+import { TextCursorPosition } from "../cursorPositionTypes.js";
+import { Selection } from "../selectionTypes.js";
 
 export class SelectionManager<
   BSchema extends BlockSchema = DefaultBlockSchema,
@@ -47,8 +47,8 @@ export class SelectionManager<
    * If the selection starts / ends halfway through a block, the returned block will be
    * only the part of the block that is included in the selection.
    */
-  public getSelectionCutBlocks() {
-    return this.editor.transact((tr) => getSelectionCutBlocks(tr));
+  public getSelectionCutBlocks(expandToWords = false) {
+    return this.editor.transact((tr) => getSelectionCutBlocks(tr, expandToWords));
   }
 
   /**

packages/core/src/index.ts

@@ -1,7 +1,7 @@
 export * from "./api/blockManipulation/commands/insertBlocks/insertBlocks.js";
 export * from "./api/blockManipulation/commands/replaceBlocks/replaceBlocks.js";
-export * from "./api/blockManipulation/commands/updateBlock/updateBlock.js";
 export * from "./api/blockManipulation/commands/replaceBlocks/util/fixColumnList.js";
+export * from "./api/blockManipulation/commands/updateBlock/updateBlock.js";
 export * from "./api/exporters/html/externalHTMLExporter.js";
 export * from "./api/exporters/html/internalHTMLSerializer.js";
 export * from "./api/getBlockInfoFromPos.js";
@@ -19,6 +19,7 @@ export * from "./i18n/dictionary.js";
 export * from "./schema/index.js";
 export * from "./util/browser.js";
 export * from "./util/combineByGroup.js";
+export * from "./util/expandToWords.js";
 export * from "./util/string.js";
 export * from "./util/table.js";
 export * from "./util/typescript.js";

packages/core/src/util/expandToWords.ts

@@ -0,0 +1,38 @@
+import type { Node, ResolvedPos } from "prosemirror-model";
+
+/**
+ * Expands a range (start to end) to include the rest of the word if it starts or ends within a word
+ */
+export function expandPMRangeToWords(
+  doc: Node,
+  range: { $from: ResolvedPos; $to: ResolvedPos },
+) {
+  let { $from, $to } = range;
+
+  // Expand Start
+  // If the selection starts with a word character or punctuation, check if we need to expand left to include the rest of the word
+  if ($from.pos > $from.start() && $from.pos < doc.content.size) {
+    const charAfterStart = doc.textBetween($from.pos, $from.pos + 1);
+    if (/^[\w\p{P}]$/u.test(charAfterStart)) {
+      const textBefore = doc.textBetween($from.start(), $from.pos);
+      const wordMatch = textBefore.match(/[\w\p{P}]+$/u);
+      if (wordMatch) {
+        $from = doc.resolve($from.pos - wordMatch[0].length);
+      }
+    }
+  }
+
+  // Expand End
+  // If the selection ends with a word characte or punctuation, check if we need to expand right to include the rest of the word
+  if ($to.pos < $to.end() && $to.pos > 0) {
+    const charBeforeEnd = doc.textBetween($to.pos - 1, $to.pos);
+    if (/^[\w\p{P}]$/u.test(charBeforeEnd)) {
+      const textAfter = doc.textBetween($to.pos, $to.end());
+      const wordMatch = textAfter.match(/^[\w\p{P}]+/u);
+      if (wordMatch) {
+        $to = doc.resolve($to.pos + wordMatch[0].length);
+      }
+    }
+  }
+  return { $from, $to, from: $from.pos, to: $to.pos };
+}

packages/xl-ai/src/api/aiRequest/builder.ts

@@ -51,7 +51,7 @@ export async function buildAIRequest(opts: {
       : undefined;
 
   const selectionInfo = useSelection
-    ? opts.editor.getSelectionCutBlocks()
+    ? opts.editor.getSelectionCutBlocks(true)
     : undefined;
 
   const streamTools = streamToolsProvider.getStreamTools(