From edaa10ab2d86365868493ea8590d620171495265 Mon Sep 17 00:00:00 2001 From: Arek Nawo Date: Sat, 20 Apr 2024 00:46:46 +0200 Subject: [PATCH] wip: Autocomplete --- apps/web/src/context/client.tsx | 2 +- .../web/src/lib/editor/extensions/document.ts | 178 +++++++++++++++++- apps/web/src/views/editor/editor.tsx | 4 +- .../src/routes/utils/handlers/autocomplete.ts | 54 ++++++ packages/backend/src/routes/utils/index.ts | 7 + 5 files changed, 239 insertions(+), 6 deletions(-) create mode 100644 packages/backend/src/routes/utils/handlers/autocomplete.ts diff --git a/apps/web/src/context/client.tsx b/apps/web/src/context/client.tsx index e7b71de3..9ba8173c 100644 --- a/apps/web/src/context/client.tsx +++ b/apps/web/src/context/client.tsx @@ -132,4 +132,4 @@ const useClient = (): Client => { }; export { ClientProvider, useClient }; -export type { App, RouterInput, RouterOutput }; +export type { App, Client, RouterInput, RouterOutput }; diff --git a/apps/web/src/lib/editor/extensions/document.ts b/apps/web/src/lib/editor/extensions/document.ts index dd8590f3..62d6d1d0 100644 --- a/apps/web/src/lib/editor/extensions/document.ts +++ b/apps/web/src/lib/editor/extensions/document.ts @@ -1,13 +1,149 @@ -import { markInputRule, markPasteRule } from "@tiptap/core"; import { Document as BaseDocument } from "@vrite/editor"; -import { AllSelection, PluginKey, Plugin, TextSelection } from "@tiptap/pm/state"; -import { useNotifications } from "#context"; +import { AllSelection, EditorState, Plugin, TextSelection, Transaction } from "@tiptap/pm/state"; +import { Decoration, DecorationSet } from "@tiptap/pm/view"; +import { Client, useNotifications } from "#context"; -const Document = BaseDocument.extend({ +interface DocumentOptions { + client: Client | null; +} +interface DocumentStorage { + autocompletion?: string; +} + +const autocompleteDecoration = ( + state: EditorState | Transaction, + client: Client | null, + storage: DocumentStorage +): { decoration: DecorationSet; cancel(): void } | null => { + const { selection } = state; + const { parent } = selection.$from; + const endOfParagraphPos = + selection.$from.pos - selection.$from.parentOffset + parent.nodeSize - 2; + + let element: HTMLElement | null = null; + + if ( + parent.type.name !== "paragraph" || + selection.$from.pos !== endOfParagraphPos || + parent.nodeSize === 2 + ) { + return null; + } + + const controller = new AbortController(); + const timeoutHandle = setTimeout(() => { + if (state instanceof Transaction && state.getMeta("autocompletion")) { + return; + } + + const sentences = (parent.textContent || "").trim().split(/(?<=\w)\. /g); + + if (sentences.at(-1)?.endsWith(".")) return; + + const paragraph = sentences.at(-1) || ""; + const context = sentences.slice(0, -1).join(". "); + + client?.utils.autocomplete + .query( + { + context, + paragraph + }, + { signal: controller.signal } + ) + .then((response) => { + storage.autocompletion = response.autocompletion; + + if (element) { + element.textContent = response.autocompletion; + } + }) + .catch(() => {}); + }, 250); + const cancel = (): void => { + controller.abort(); + clearTimeout(timeoutHandle); + }; + const decoration = DecorationSet.create(state.doc, [ + Decoration.widget(state.selection.from, () => { + const span = document.createElement("span"); + + span.setAttribute("class", "opacity-30"); + span.setAttribute("contenteditable", "false"); + + if (state instanceof Transaction && state.getMeta("autocompletion")) { + span.textContent = storage.autocompletion || ""; + } + + element = span; + + return span; + }) + ]); + + return { decoration, cancel }; +}; +const Document = BaseDocument.extend({ + addOptions() { + return { client: null }; + }, + addStorage() { + return { + autocompletion: "" + }; + }, addKeyboardShortcuts() { const { notify } = useNotifications() || {}; return { + "Tab": ({ editor }) => { + const autocompletion = this.storage.autocompletion || ""; + + if (autocompletion) { + const { tr } = editor.state; + const { dispatch } = editor.view; + const { $from } = editor.state.selection; + const endOfParagraphPos = $from.pos - $from.parentOffset + $from.parent.nodeSize - 2; + + this.storage.autocompletion = ""; + dispatch( + tr + .setMeta("autocompletion", true) + .insertText(autocompletion, endOfParagraphPos, endOfParagraphPos) + ); + + return true; + } + + // Handle tab when at the end of a paragraph. + return true; + }, + "Shift-Tab": ({ editor }) => { + const autocompletion = this.storage.autocompletion || ""; + + if (autocompletion) { + const { tr } = editor.state; + const { dispatch } = editor.view; + const { $from } = editor.state.selection; + const endOfParagraphPos = $from.pos - $from.parentOffset + $from.parent.nodeSize - 2; + const words = autocompletion.split(" "); + + this.storage.autocompletion = words.slice(1).join(" "); + dispatch( + tr + .setMeta("autocompletion", true) + .insertText( + `${words[0]}${words.length > 1 ? " " : ""}`, + endOfParagraphPos, + endOfParagraphPos + ) + ); + + return true; + } + + return false; + }, "Mod-s": () => { notify?.({ text: "Vrite autosaves your content", type: "success" }); @@ -47,6 +183,40 @@ const Document = BaseDocument.extend({ return true; } }; + }, + addProseMirrorPlugins() { + const { options, storage, editor } = this; + + return [ + new Plugin({ + state: { + init(_, state) { + return autocompleteDecoration(state, options.client, storage); + }, + apply(tr, previousValue, oldState) { + const selectionPosChanged = tr.selection.from !== oldState.selection.from; + + if ( + (!(tr.selectionSet && selectionPosChanged) && !tr.docChanged) || + tr.getMeta("tableColumnResizing$") || + !tr.selection.empty || + !editor.isFocused + ) { + return previousValue || null; + } + + previousValue?.cancel(); + + return autocompleteDecoration(tr, options.client, storage); + } + }, + props: { + decorations(state) { + return this.getState(state)?.decoration || null; + } + } + }) + ]; } }); diff --git a/apps/web/src/views/editor/editor.tsx b/apps/web/src/views/editor/editor.tsx index 7284226c..90094a28 100644 --- a/apps/web/src/views/editor/editor.tsx +++ b/apps/web/src/views/editor/editor.tsx @@ -51,6 +51,7 @@ import { App, hasPermission, useAuthenticatedUserData, + useClient, useContentData, useExtensions, useHostConfig, @@ -132,6 +133,7 @@ const Editor: Component = (props) => { const [showBlockBubbleMenu, setShowBlockBubbleMenu] = createSignal(false); const [isNodeSelection, setIsNodeSelection] = createSignal(false); const { workspaceSettings } = useAuthenticatedUserData(); + const client = useClient(); const extensionsContext = useExtensions(); const updateBubbleMenuPlacement = debounce(() => { bubbleMenuInstance()?.setProps({ placement: isNodeSelection() ? "top-start" : "top" }); @@ -156,7 +158,7 @@ const Editor: Component = (props) => { }, extensions: [ BlockPaste.configure({ workspaceSettings }), - Document, + Document.configure({ client }), Placeholder, Paragraph, Text, diff --git a/packages/backend/src/routes/utils/handlers/autocomplete.ts b/packages/backend/src/routes/utils/handlers/autocomplete.ts new file mode 100644 index 00000000..db995c96 --- /dev/null +++ b/packages/backend/src/routes/utils/handlers/autocomplete.ts @@ -0,0 +1,54 @@ +import { z } from "zod"; +import OpenAI from "openai"; +import { Context } from "#lib/context"; + +const inputSchema = z.object({ + context: z.string().optional(), + paragraph: z.string() +}); +const outputSchema = z.object({ + autocompletion: z.string() +}); +const handler = async ( + ctx: Context, + input: z.infer +): Promise> => { + const openai = new OpenAI({ + apiKey: ctx.fastify.config.OPENAI_API_KEY, + organization: ctx.fastify.config.OPENAI_ORGANIZATION + }); + const prompt = `Briefly and concisely autocomplete the sentence to fit the provided context. DO NOT INCLUDE THE SENTENCE IN THE OUTPUT.\nContext: ${input.context || ""}\nSentence: ${input.paragraph}\nAutocompletion:`; + const result = await openai.chat.completions.create({ + model: "gpt-3.5-turbo", + temperature: 0, + max_tokens: 32, + messages: [ + { + role: "user", + content: prompt + } + ] + }); + + let autocompletion = result.choices[0].message.content || ""; + + if (input.paragraph.endsWith(" ")) { + autocompletion = autocompletion.trim(); + } + + console.log( + JSON.stringify( + { + prompt: `Briefly and concisely autocomplete the sentence to fit the provided context. DO NOT INCLUDE THE SENTENCE IN THE OUTPUT.\nContext: ${input.context || ""}\nSentence: ${input.paragraph}\nAutocompletion:`, + autocompletion + }, + null, + 2 + ) + ); + autocompletion = autocompletion.replace(input.paragraph, ""); + + return { autocompletion }; +}; + +export { inputSchema, outputSchema, handler }; diff --git a/packages/backend/src/routes/utils/index.ts b/packages/backend/src/routes/utils/index.ts index 2fb9930a..acaa7749 100644 --- a/packages/backend/src/routes/utils/index.ts +++ b/packages/backend/src/routes/utils/index.ts @@ -2,6 +2,7 @@ import { PreviewData } from "./handlers/link-preview"; import * as getHostConfig from "./handlers/host-config"; import * as getLinkPreview from "./handlers/link-preview"; import * as generateCSS from "./handlers/generate-css"; +import * as autocomplete from "./handlers/autocomplete"; import type { HostConfig } from "#lib/host-config"; import { procedure, router } from "#lib/trpc"; import { isAuthenticated } from "#lib/middleware"; @@ -22,6 +23,12 @@ const utilsRouter = router({ .output(generateCSS.outputSchema) .query(async ({ ctx, input }) => { return generateCSS.handler(ctx, input); + }), + autocomplete: authenticatedProcedure + .input(autocomplete.inputSchema) + .output(autocomplete.outputSchema) + .query(async ({ ctx, input }) => { + return autocomplete.handler(ctx, input); }) });