wip: Autocomplete

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 };

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<DocumentOptions, DocumentStorage>({
+  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;
+          }
+        }
+      })
+    ];
   }
 });
 

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<EditorProps> = (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<EditorProps> = (props) => {
     },
     extensions: [
       BlockPaste.configure({ workspaceSettings }),
-      Document,
+      Document.configure({ client }),
       Placeholder,
       Paragraph,
       Text,

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<typeof inputSchema>
+): Promise<z.infer<typeof outputSchema>> => {
+  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 };

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);
     })
 });