Lexical で文章の修正提案機能を自作しよう！

はじめに

こんにちは。Gunosy新規事業開発室の上村 id:muemura です。

本記事は、Gunosy Advent Calendar 2024 10日目の記事です。 昨日の記事は、 小野 id:takuto_ono さんの「3年目iOSエンジニアがGunosyに転職して思ったこと」でした。

tech.gunosy.io

今回は、新規事業開発室で行っているエディタ開発において、文章の修正提案機能を自作した話を紹介します。

エディタ開発

現在、新規事業開発室ではWebサイト上での文章作成を支援するエディタの開発を行っています。 エディタの開発では、Meta社が開発したオープンソースのテキストエディタフレームワークであるLexicalを使用しています。

lexical.dev

Lexicalは、軽量で高速なパフォーマンスを持ち、プラグインを通じて機能を拡張できる柔軟性が特徴です。 公式でプレイグランドを公開しているため、簡単に試すことができます。

現在開発している環境では、TypeScript + React を使用しているため、本記事ではこの環境を前提として進めます。

単語の修正提案機能

どんなもの？

エディタ上で文章を編集する際、特定の単語を自動で検知して修正を提案する機能です。 具体的には、あらかじめ登録された単語辞書等を基に、文章中の特定の単語に対して修正案を表示し、ユーザーが必要に応じて修正を適用できます。 この文章校正機能により、より正確で一貫性のある文章作成をサポートします。

修正提案機能

Lexical の拡張機能について簡単に紹介

Lexical は、拡張機能を追加することで様々な機能を追加できます。 詳しくは公式ドキュメントを参照してください。

ざっくりとした流れ（一例）は以下の通りです。

1. 独自のLexicalNodeを作成
   • TextNode や ElementNode などのノードを継承して、持たせたいデータを持つノードを定義します。
2. 定義したNodeに変換、または新規追加するプラグインを作成
   • 正規表現などを用いて特定の文字列を検出し、独自定義したNodeに変換するプラグインを作成します。
3. プラグインとNodeを登録
   • 独自で追加したNodeとプラグインをエディタに登録して扱えるようにします。

修正内容を保持する LexicalNode の作成

修正内容を保持するノードを作成します。 今回は、ノードを SuggestionNode として定義します。

import type { EditorConfig, LexicalNode, TextNode } from "lexical";

export class SuggestionNode extends TextNode {
  /** @internal */
  __suggestText?: string;

  static getType(): string {
    return "suggestion";
  }

  createDOM(config: EditorConfig): HTMLElement {
    const dom = super.createDOM(config);
    dom.style.cursor = "default";
    dom.className = "suggestion";
    return dom;
  }

  getSuggestText(): string {
    return this.getLatest().__suggestText ?? "";
  }

  setSuggestText(suggestText: string): this {
    const self = this.getWritable();
    self.__suggestText = suggestText;
    return self;
  }

  /* その他 serialize, clone まわりの定義も必要 */
  ...
}

export function $createSuggestionNode(
  text: string,
  suggestText: string
): SuggestionNode {
  return new SuggestionNode(text, suggestText);
}

export function $isSuggestionNode(
  node: LexicalNode | null | undefined
): boolean {
  return node instanceof SuggestionNode;
}

SuggestionNode は、修正内容のテキストを保持するための __suggestText プロパティを持たせています。 また、createDOM メソッドをオーバーライドして、修正提案のノードを表示する際のスタイルを設定しています。

プラグインの作成

次に、修正提案機能を実装するプラグインを作成します。今回は以下の2つのプラグインを作成します。

1. 既存のテキストから特定の単語を検出して SuggestionNode に変換するプラグイン
2. SuggestionNode を hover した際に修正提案を表示するプラグイン

これらを順に実装していきます。

1. 既存のテキストから特定の単語を検出して SuggestionNode に変換するプラグイン

単純な正規表現で変換を行う場合は、useLexicalTextEntity *1 を使用することで簡単に実装できます。

込み入った変換を行う場合は、TextNode から変換する処理 ($textNodeTransform) と TextNode に戻す変換をする処理 ($reverseNodeTransform) を自前で定義して registerNodeTransform で処理を登録しましょう。
今回は簡単に実装するため、useLexicalTextEntity を使用する例を紹介します。

import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
import { useLexicalTextEntity } from '@lexical/react/useLexicalTextEntity';
import { useCallback, useEffect } from "react";

const SUGGESTION_WORDS_MAP: Map<string, string> = ([
    // 修正を提案したい単語と修正後の単語のマッピング
    // 例: ["修正を提案したい単語1", "修正後の単語1"]
    ["修正を提案したい単語1","修正後の単語1"],
    ["修正を提案したい単語2","修正後の単語2"],
]);
const SUGGESTION_REGEX = new RegExp(
  Array.from(SUGGESTION_WORDS_MAP.keys()).join("|"),
  "g"
);

export default function SuggestionPlugin(): JSX.Element | null {
  const [editor] = useLexicalComposerContext();

  // エディタにSuggestionNodeが登録されているか確認
  useEffect(() => {
    if (!editor.hasNodes([SuggestionNode])) {
      throw new Error(
        "SuggestionPlugin: SuggestionNode not registered on editor"
      );
    }
  }, [editor]);

  // TextNode から SuggestionNode に変換する処理
  const $createSuggestionNode_ = useCallback(
    (textNode: TextNode): SuggestionNode => {
      // 登録された単語を検出し、対応する修正案に置換
      const targetText = textNode.getTextContent();
      const suggestText = targetText.replace(SUGGESTION_REGEX, (match) =>
        SUGGESTION_WORDS_MAP[match]
      );
      // 新しい SuggestionNode を作成して返す
      return $createSuggestionNode(text, suggestText);
    },
    []
  );

  // Matcher (どこからどこまでを SuggestionNode に変換するか)の定義
  const getSuggestionMatch = useCallback((text: string) => {
    const matchArr = SUGGESTION_REGEX.exec(text);
    if (matchArr === null) {
      return null;
    }

    return {
      end: matchArr.index + matchArr[0].length,
      start: matchArr.index,
    };
  }, []);

  // プラグインの登録
  useLexicalTextEntity<SuggestionNode>(
    getSuggestionMatch,
    SuggestionNode,
    $createSuggestionNode_,
  );

  return null;
}

SUGGESTION_WORDS_MAP に登録された単語を検出して、それぞれの修正後の単語を保持する SuggestionNode に変換するプラグインとして SuggestionsPlugin を定義しています。 useLexicalTextEntity は、マッチするテキストを特定のノードに変換することができる便利なフックです。

2. SuggestionNode を hover した際に修正提案を表示するプラグイン

SuggestionNode だけでは修正提案を表示することができません。 今回は、エディタ上で SuggestionNode を hover した際に修正提案を表示し、修正するためのプラグインを作成します。

import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
import {
  $getNearestNodeFromDOMNode,
  $getNodeByKey,
  LexicalNode,
  NodeKey,
} from "lexical";
import {useState, useEffect, useCallback, useRef} from "react";
import { createPortal } from "react-dom";
import { mergeRegister } from "@lexical/utils";
// https://github.com/facebook/lexical/blob/05fa244bd0f6043114ffb8feab2922d8e4de7e6f/packages/lexical-playground/src/plugins/CodeActionMenuPlugin/utils.ts
import {useDebounce} from "./utils";
import IconClose from "@/components/svg/IconClose";

function SuggestionQuickFixContainer({anchorElem}: {anchorElem: HTMLElement | null}): JSX.Element {
  const [editor] = useLexicalComposerContext();
  // 修正提案のテキスト
  const [suggestionText, setSuggestText] = useState<string>("");
  // 修正対象のSuggestionNodeのキー
  const [suggestionNodeKey, setSuggestionNodeKey] = useState<NodeKey | null>(
    null
  );
  // 修正提案のダイアログを表示するかどうか
  const [isShown, setIsShown] = useState<boolean>(false);
  // マウスの動きを監視するかどうか
  const [shouldListenMouseMove, setShouldListenMouseMove] =
    useState<boolean>(false);
  // ダイアログの表示位置
  const [position, setPosition] = useState({});
  // 現在エディタ上に存在するSuggestionNodeのキーを保持するSet
  const codeSetRef = useRef<Set<NodeKey>>(new Set());
  // 現在ホバーしているSuggestionNodeのDOM要素への参照
  const sugestionDOMNodeRef = useRef<HTMLElement | null>(null);

  // マウスの移動を監視するため、debounce して処理を行う
  const debouncedOnMouseMove = useDebounce(
    (event: MouseEvent) => {
      // マウスの位置情報を取得して、SuggestionNode に hover しているか判定
      const { isOutside, sugestionDOMNode } = getMouseInfo(event);
      if (isOutside) {
        setIsShown(false);
        return;
      }

      if (!sugestionDOMNode) {
        return;
      }

      sugestionDOMNodeRef.current = sugestionDOMNode;

      let hoveredNode: LexicalNode | null = null;
      let sugestionDOMElement: HTMLElement | null = null;

      editor.update(() => {
        // DOM から取得した Node が SuggestionNode かどうか判定
        const maybeSuggestionNode =
          $getNearestNodeFromDOMNode(sugestionDOMNode);

        // SuggestionNode であれば、NodeKey を元に修正提案のテキストを取得
        if (maybeSuggestionNode && $isSuggestionNode(maybeSuggestionNode)) {
          sugestionDOMElement = editor.getElementByKey(
            maybeSuggestionNode?.getKey()
          );
          setSuggestionNodeKey(maybeSuggestionNode.getKey());
          setSuggestText(
            (maybeSuggestionNode as SuggestionNode).getSuggestText()
          );

          if (sugestionDOMElement) {
            hoveredNode = maybeSuggestionNode;
          }
        }
      });

      // 修正提案のダイアログを表示する位置を設定 (SuggestionNode の右下あたりに表示されるように調整)
      if (sugestionDOMElement) {
        const {
          height: suggestionElemHeight,
          top: suggestionElemTop,
          right: suggestionElemRight,
        } = (sugestionDOMElement as HTMLSpanElement).getBoundingClientRect();
        const { y: editorElemY } = anchorElem.getBoundingClientRect();

        if (hoveredNode) {
          setIsShown(true);
          setPosition({
            left: suggestionElemRight,
            top: suggestionElemTop + suggestionElemHeight - editorElemY,
          });
        }
      }
    },
    50,
    250
  );

  // mousemove にイベントリスナーを登録
  useEffect(() => {
    if (!shouldListenMouseMove) {
      return;
    }
    document.addEventListener("mousemove", debouncedOnMouseMove);

    return () => {
      setIsShown(false);
      debouncedOnMouseMove.cancel();
      document.removeEventListener("mousemove", debouncedOnMouseMove);
    };
  }, [shouldListenMouseMove, debouncedOnMouseMove]);

  // SuggestionNode の変更を監視
  useEffect(() => {
    return mergeRegister(
      editor.registerMutationListener(SuggestionNode, (mutations) => {
        editor.getEditorState().read(() => {
          for (const [key, type] of mutations) {
            switch (type) {
              case "created":
                codeSetRef.current.add(key);
                setShouldListenMouseMove(codeSetRef.current.size > 0);
                break;
              case "updated":
                if (!codeSetRef.current.has(key)) {
                  codeSetRef.current.add(key);
                }
                setShouldListenMouseMove(codeSetRef.current.size > 0);
                break;
              case "destroyed":
                codeSetRef.current.delete(key);
                setShouldListenMouseMove(codeSetRef.current.size > 0);
                break;

              default:
                break;
            }
          }
        });
      })
    );
  }, [editor]);

  // 修正提案を適用する処理
  // setTextContent で内容が変更されると、上記で定義した正規表現に引っ掛からなくなるため、自動的に TextNode に変換される
  const quickFixHandler = useCallback(() => {
    editor.update(() => {
      if (suggestionNodeKey) {
        const maybeSuggestionNode = $getNodeByKey(suggestionNodeKey);
        if (maybeSuggestionNode && $isSuggestionNode(maybeSuggestionNode)) {
          (maybeSuggestionNode as SuggestionNode).setTextContent(
            suggestionText
          );
        }
      }
    });

    setIsShown(false);
  }, [editor, suggestionText, suggestionNodeKey]);

  return (
    <>
      {isShown && (
        <div
          className="suggestionQuickFix"
          style={{ ...position }}
        >
          <button
            className="suggestionQuickFixButton"
            onClick={quickFixHandler}
          >
            {suggestionText}
          </button>
          <button
            className="suggestionQuickFixCloseButton"
            onClick={setIsShown(false)}
          >
            <IconClose />
          </button>
        </div>
      )}
    </>
  );
}

function getMouseInfo(event: MouseEvent): {
  sugestionDOMNode: HTMLElement | null;
  isOutside: boolean;
} {
  const target = event.target;
  if (
    target &&
    (target instanceof HTMLElement || target instanceof SVGElement)
  ) {
    // suggestionNode は span.suggestion で定義されているため、それを元に判定
    const sugestionDOMNode = target.closest<HTMLElement>("span.suggestion");
    // 修正提案のダイアログ外をクリックしたかどうかの判定
    const isOutside = !(
      sugestionDOMNode || target.closest<HTMLElement>("div.suggestionQuickFix")
    );

    return { isOutside, sugestionDOMNode };
  } else {
    return { isOutside: true, sugestionDOMNode: null };
  }
}

export default function SuggestionQuickFixPlugin({
  isReadOnly = false,
  anchorElem = document.body,
}: {
  isReadOnly?: boolean;
  anchorElem?: HTMLElement;
}): React.ReactPortal | null {
  if (isReadOnly) {
    return null;
  }

  return createPortal(
    <SuggestionQuickFixContainer anchorElem={anchorElem} />,
    anchorElem
  );
}

SuggestionQuickFixContainer は、SuggestionNode を hover した際に修正提案を表示するためのコンポーネントを定義しています。 具体的には、マウスの位置から Node を特定し、それが SuggestionNode だった場合に修正提案を表示を行い、提案を適用するための処理を仕込むような形で実装しています。

プラグインの登録

最後に、作成したプラグインをエディタに登録します。詳細は、公式ドキュメントのページを参照してください。

// config に作成した Node を登録
const initialConfig = {
  namespace: 'MyEditor',
  nodes: [SuggestionNode],
};

// Pluginは、LexicalComposer の子要素として登録
 <LexicalComposer initialConfig={initialConfig}>
    <SuggestionPlugin />
    <SuggestionQuickFixPlugin />
 </LexicalComposer>

実際の挙動

以下のように、エディタ上で特定の単語を検知して、修正提案を表示し、修正を適用することができるようになりました。

修正提案機能の挙動

今回は簡単に実装するため、正規表現で単語を検出して修正提案を表示するような形で実装しましたが、LLMを使用してより高度な修正提案機能を実装したり、文章の続きを生成するような執筆のサポートを行ったりすることもできるので、興味があれば試してみてください。

終わりに

Lexical を使用して文章の修正提案機能を自作する方法を紹介しました。

実際 Lexical の開発をしてみると、まだ Lexical 自体のバグに当たることが多く、ある程度ソースコードを見ながら開発する必要があるため、最初はとっつきにくく感じることもあるかもしれません。 しかし、柔軟な拡張を行うことができる Lexical を使いこなすことで、様々な機能を追加することができるため、触ってみる価値はあると思います！ 今回の記事で、Lexical を使った自由なエディタ開発に興味を持っていただければ幸いです。

明日は、 UT id:yuutookun さんの「EM（エンジニアリングマネージャ）を１年半やって」です。お楽しみに！

*1:https://lexical.dev/docs/api/modules/lexical_react_useLexicalTextEntity