diff --git a/examples/09-ai/01-minimal/src/App.tsx b/examples/09-ai/01-minimal/src/App.tsx index 8c0ce912a2..f85818afa3 100644 --- a/examples/09-ai/01-minimal/src/App.tsx +++ b/examples/09-ai/01-minimal/src/App.tsx @@ -14,6 +14,7 @@ import { import { AIMenuController, AIToolbarButton, + createAIAutoCompleteExtension, createAIExtension, getAISlashMenuItems, } from "@blocknote/xl-ai"; @@ -26,6 +27,49 @@ import { getEnv } from "./getEnv"; const BASE_URL = getEnv("BLOCKNOTE_AI_SERVER_BASE_URL") || "https://localhost:3000/ai"; +/** + * Fetches suggestions for the auto complete plugin from our backend API. + */ +async function autoCompleteProvider( + editor: BlockNoteEditor, + signal: AbortSignal, +) { + // TODO: + // - API is very prosemirror-based, make something more BlockNote-native + // - Add simple method to retrieve relevant context (e.g. block content / json until selection) + + const state = editor.prosemirrorState; + const text = state.doc.textBetween( + state.selection.from - 300, + state.selection.from, + ); + + const response = await fetch( + "https://blocknote-pr-2191.onrender.com/ai/autocomplete/generateText", + // `https://localhost:3000/ai/autocomplete/generateText`, + { + method: "POST", + body: JSON.stringify({ text }), + signal, + }, + ); + const data = await response.json(); + return data.suggestions.map((suggestion: string) => ({ + position: state.selection.from, + suggestion: suggestion, + })); + // return [ + // { + // position: state.selection.from, + // suggestion: "Hello World", + // }, + // { + // position: state.selection.from, + // suggestion: "Hello Planet", + // }, + // ]; +} + export default function App() { // Creates a new editor instance. const editor = useCreateBlockNote({ @@ -41,6 +85,7 @@ export default function App() { api: `${BASE_URL}/regular/streamText`, }), }), + createAIAutoCompleteExtension({ autoCompleteProvider }), ], // We set some initial content for demo purposes initialContent: [ diff --git a/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts b/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts index 0c4bb7f4a6..10b5ee8818 100644 --- a/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts +++ b/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts @@ -290,7 +290,7 @@ export class SuggestionMenuProseMirrorPlugin< }, props: { - handleTextInput(view, from, to, text) { + handleTextInput(view, from, to, text, deflt) { // only on insert if (from === to) { const doc = view.state.doc; @@ -301,18 +301,21 @@ export class SuggestionMenuProseMirrorPlugin< : text; if (str === snippet) { - view.dispatch(view.state.tr.insertText(text)); view.dispatch( - view.state.tr - .setMeta(suggestionMenuPluginKey, { - triggerCharacter: snippet, - }) - .scrollIntoView(), + deflt().setMeta(suggestionMenuPluginKey, { + triggerCharacter: snippet, + }), ); return true; } } } + if (this.getState(view.state)) { + // when menu is open, we dispatch the default transaction + // and return true so that other event handlers (i.e.: AI AutoComplete) are not triggered + view.dispatch(deflt()); + return true; + } return false; }, diff --git a/packages/xl-ai-server/src/index.ts b/packages/xl-ai-server/src/index.ts index f66ef6738e..375252e1e3 100644 --- a/packages/xl-ai-server/src/index.ts +++ b/packages/xl-ai-server/src/index.ts @@ -5,6 +5,7 @@ import { cors } from "hono/cors"; import { existsSync, readFileSync } from "node:fs"; import { createSecureServer } from "node:http2"; import { Agent, setGlobalDispatcher } from "undici"; +import { autocompleteRoute } from "./routes/autocomplete.js"; import { modelPlaygroundRoute } from "./routes/model-playground/index.js"; import { objectGenerationRoute } from "./routes/objectGeneration.js"; import { proxyRoute } from "./routes/proxy.js"; @@ -37,6 +38,7 @@ app.route("/ai/proxy", proxyRoute); app.route("/ai/object-generation", objectGenerationRoute); app.route("/ai/server-promptbuilder", serverPromptbuilderRoute); app.route("/ai/model-playground", modelPlaygroundRoute); +app.route("/ai/autocomplete", autocompleteRoute); const http2 = existsSync("localhost.pem"); serve( diff --git a/packages/xl-ai-server/src/routes/autocomplete.ts b/packages/xl-ai-server/src/routes/autocomplete.ts new file mode 100644 index 0000000000..b5a47a00a1 --- /dev/null +++ b/packages/xl-ai-server/src/routes/autocomplete.ts @@ -0,0 +1,46 @@ +import { createGroq } from "@ai-sdk/groq"; +import { generateText } from "ai"; +import { Hono } from "hono"; + +export const autocompleteRoute = new Hono(); + +// Setup your model +// const model = createOpenAI({ +// apiKey: process.env.OPENAI_API_KEY, +// })("gpt-4.1-nano"); + +const model = createGroq({ + apiKey: process.env.GROQ_API_KEY, +})("openai/gpt-oss-20b"); + +// Use `streamText` to stream text responses from the LLM +autocompleteRoute.post("/generateText", async (c) => { + const { text } = await c.req.json(); + + const result = await generateText({ + model, + system: `You are a writing assistant. Predict and generate the most likely next part of the text. +- separate suggestions by newlines +- max 3 suggestions +- keep it short, max 5 words per suggestion +- don't include other text (or explanations) +- YOU MUST ONLY return the text to be appended. Your suggestion will EXACTLY replace [SUGGESTION_HERE]. +- YOU MUST NOT include the original text / characters (prefix) in your suggestion. +- YOU MUST add a space (or other relevant punctuation) before the suggestion IF starting a new word (the suggestion will be directly concatenated to the text)`, + messages: [ + { + role: "user", + content: `Complete the following text: + ${text}[SUGGESTION_HERE]`, + }, + ], + abortSignal: c.req.raw.signal, + }); + + return c.json({ + suggestions: result.text + .split("\n") + .map((suggestion) => suggestion.trimEnd()) + .filter((suggestion) => suggestion.trim().length > 0), + }); +}); diff --git a/packages/xl-ai/src/index.ts b/packages/xl-ai/src/index.ts index 843bb17d03..7db8f0650d 100644 --- a/packages/xl-ai/src/index.ts +++ b/packages/xl-ai/src/index.ts @@ -11,4 +11,5 @@ export * from "./components/AIMenu/PromptSuggestionMenu.js"; export * from "./components/FormattingToolbar/AIToolbarButton.js"; export * from "./components/SuggestionMenu/getAISlashMenuItems.js"; export * from "./i18n/dictionary.js"; +export * from "./plugins/AutoCompletePlugin.js"; export * from "./streamTool/index.js"; diff --git a/packages/xl-ai/src/plugins/AutoCompletePlugin.ts b/packages/xl-ai/src/plugins/AutoCompletePlugin.ts new file mode 100644 index 0000000000..3d332d4a05 --- /dev/null +++ b/packages/xl-ai/src/plugins/AutoCompletePlugin.ts @@ -0,0 +1,336 @@ +import { + BlockNoteEditor, + BlockNoteExtension, + BlockSchema, + InlineContentSchema, + StyleSchema, +} from "@blocknote/core"; +import { EditorState, Plugin, PluginKey } from "prosemirror-state"; +import { Decoration, DecorationSet } from "prosemirror-view"; + +export type AutoCompleteState = + | { + autoCompleteSuggestion: AutoCompleteSuggestion; + } + | undefined; + +const autoCompletePluginKey = new PluginKey<{ isUserInput: boolean }>( + "AutoCompletePlugin", +); + +type AutoCompleteSuggestion = { + position: number; + suggestion: string; +}; + +type AutoCompleteProvider = ( + editor: BlockNoteEditor, + signal: AbortSignal, +) => Promise; + +function getMatchingSuggestions( + autoCompleteSuggestions: AutoCompleteSuggestion[], + state: EditorState, +): AutoCompleteSuggestion[] { + return autoCompleteSuggestions + .map((suggestion) => { + if (suggestion.position > state.selection.from) { + return false; + } + + if ( + !state.doc + .resolve(suggestion.position) + .sameParent(state.selection.$from) + ) { + return false; + } + + const text = state.doc.textBetween( + suggestion.position, + state.selection.from, + ); + if ( + suggestion.suggestion.startsWith(text) && + suggestion.suggestion.length > text.length + ) { + return { + position: suggestion.position, + suggestion: suggestion.suggestion.slice(text.length), + }; + } + return false; + }) + .filter((suggestion) => suggestion !== false); +} + +export class AutoCompleteProseMirrorPlugin< + BSchema extends BlockSchema, + I extends InlineContentSchema, + S extends StyleSchema, +> extends BlockNoteExtension { + public static key() { + return "suggestionMenu"; + } + + public get priority(): number | undefined { + return 1000000; // should be lower (e.g.: -1000 to be below suggestion menu, but that currently breaks Tab) + } + + // private view: AutoCompleteView | undefined; + + // private view: EditorView | undefined; + private autoCompleteSuggestions: AutoCompleteSuggestion[] = []; + + private debounceFetchSuggestions = debounceWithAbort( + async (editor: BlockNoteEditor, signal: AbortSignal) => { + // fetch suggestions + const autoCompleteSuggestions = await this.options.autoCompleteProvider( + editor, + signal, + ); + + // TODO: map positions? + + if (signal.aborted) { + return; + } + + this.autoCompleteSuggestions = autoCompleteSuggestions; + this.editor.transact((tr) => { + tr.setMeta(autoCompletePluginKey, { + autoCompleteSuggestions, + }); + }); + }, + ); + + constructor( + private readonly editor: BlockNoteEditor, + private readonly options: { + autoCompleteProvider: AutoCompleteProvider; + }, + ) { + super(); + + // eslint-disable-next-line @typescript-eslint/no-this-alias + const self = this; + this.addProsemirrorPlugin( + new Plugin({ + key: autoCompletePluginKey, + + // view: (view) => { + // this.view = new AutoCompleteView(editor, view); + // return this.view; + // }, + + state: { + // Initialize the plugin's internal state. + init(): AutoCompleteState { + return undefined; + }, + + // Apply changes to the plugin state from an editor transaction. + apply: ( + transaction, + _prev, + _oldState, + newState, + ): AutoCompleteState => { + // selection is active, no autocomplete + if (newState.selection.from !== newState.selection.to) { + this.debounceFetchSuggestions.cancel(); + return undefined; + } + + // Are there matching suggestions? + const matchingSuggestions = getMatchingSuggestions( + this.autoCompleteSuggestions, + newState, + ); + + if (matchingSuggestions.length > 0) { + this.debounceFetchSuggestions.cancel(); + return { + autoCompleteSuggestion: matchingSuggestions[0], + }; + } + + // No matching suggestions, if isUserInput is true, debounce fetch suggestions + if (transaction.getMeta(autoCompletePluginKey)?.isUserInput) { + // TODO: this queueMicrotask is a workaround to ensure the transaction is applied before the debounceFetchSuggestions is called + // (discuss with Nick what ideal architecture would be) + queueMicrotask(() => { + this.debounceFetchSuggestions(self.editor).catch((error) => { + /* eslint-disable-next-line no-console */ + console.error(error); + }); + }); + } else { + // clear suggestions + this.autoCompleteSuggestions = []; + } + return undefined; + }, + }, + + props: { + handleKeyDown(view, event) { + if (event.key === "Tab") { + // TODO (discuss with Nick): + // Plugin priority needs to be below suggestion menu, so no auto complete is triggered when the suggestion menu is open + // However, Plugin priority needs to be above other Tab handlers (because now indentation will be wrongly prioritized over auto complete) + const autoCompleteState = this.getState(view.state); + + if (autoCompleteState) { + // insert suggestion + view.dispatch( + view.state.tr + .insertText( + autoCompleteState.autoCompleteSuggestion.suggestion, + ) + .setMeta(autoCompletePluginKey, { isUserInput: true }), // isUserInput true to trigger new fetch + ); + return true; + } + + // if tab to suggest is enabled (TODO: make configurable) + view.dispatch( + view.state.tr.setMeta(autoCompletePluginKey, { + isUserInput: true, + }), + ); + return true; + } + + if (event.key === "Escape") { + self.autoCompleteSuggestions = []; + self.debounceFetchSuggestions.cancel(); + view.dispatch(view.state.tr.setMeta(autoCompletePluginKey, {})); + return true; + } + + return false; + }, + handleTextInput(view, _from, _to, _text, deflt) { + const tr = deflt(); + tr.setMeta(autoCompletePluginKey, { + isUserInput: true, + }); + view.dispatch(tr); + return true; + }, + + // Setup decorator on the currently active suggestion. + decorations(state) { + const autoCompleteState: AutoCompleteState = this.getState(state); + + if (!autoCompleteState) { + return null; + } + + // console.log(autoCompleteState); + // Creates an inline decoration around the trigger character. + return DecorationSet.create(state.doc, [ + Decoration.widget( + state.selection.from, + renderAutoCompleteSuggestion( + autoCompleteState.autoCompleteSuggestion.suggestion, + ), + {}, + ), + ]); + }, + }, + }), + ); + } +} + +function renderAutoCompleteSuggestion(suggestion: string) { + const element = document.createElement("span"); + element.classList.add("bn-autocomplete-decorator"); + element.textContent = suggestion; + return element; +} + +export function debounceWithAbort( + fn: (...args: [...T, AbortSignal]) => Promise | R, + delay = 300, // TODO: configurable +) { + let timeoutId: ReturnType | null = null; + let controller: AbortController | null = null; + + const debounced = (...args: T): Promise => { + // Clear pending timeout + if (timeoutId) { + clearTimeout(timeoutId); + } + + // Abort any in-flight execution + if (controller) { + controller.abort(); + } + + controller = new AbortController(); + const signal = controller.signal; + + return new Promise((resolve, reject) => { + timeoutId = setTimeout(async () => { + try { + const result = await fn(...args, signal); + resolve(result); + } catch (err) { + reject(err); + } + }, delay); + }); + }; + + // External cancel method + debounced.cancel = () => { + if (timeoutId) { + clearTimeout(timeoutId); + } + timeoutId = null; + + if (controller) { + controller.abort(); + } + controller = null; + }; + + return debounced; +} + +// Add a type for the cancel method +export interface DebouncedFunction { + (...args: T): Promise; + cancel(): void; +} + +/** + * Create a new AIExtension instance, this can be passed to the BlockNote editor via the `extensions` option + */ +export function createAIAutoCompleteExtension( + options: ConstructorParameters[1], +) { + return (editor: BlockNoteEditor) => { + return new AutoCompleteProseMirrorPlugin(editor, options); + }; +} + +/** + * Return the AIExtension instance from the editor + */ +export function getAIAutoCompleteExtension( + editor: BlockNoteEditor, +) { + return editor.extension(AutoCompleteProseMirrorPlugin); +} + +// TODO: move more to blocknote API? +// TODO: test with Collaboration edits +// TODO: compare kilocode / cline etc +// TODO: think about advanced scenarios (e.g.: multiple suggestions, etc.) +// TODO: double tap -> insert extra long suggestion diff --git a/packages/xl-ai/src/style.css b/packages/xl-ai/src/style.css index 4b7558d518..5547d181b3 100644 --- a/packages/xl-ai/src/style.css +++ b/packages/xl-ai/src/style.css @@ -31,3 +31,9 @@ del, text-decoration: line-through; text-decoration-thickness: 1px; } + +.bn-autocomplete-decorator { + width: 50px; + height: 50px; + color: var(--bn-colors-side-menu); +}