handle ai completion markdown formatting

ibastawisi · Jun 17, 2024 · 1a0e83c · 1a0e83c
1 parent e18a5ea
commit 1a0e83c
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 26 deletions.
diff --git a/src/app/api/completion/route.ts b/src/app/api/completion/route.ts
@@ -24,6 +24,8 @@ export async function POST(req: Request) {
           "You are an AI writing assistant for the text editor application 'Math Editor'. " +
           "You are asked to continue writing more text following user's " +
           "Limit your response to no more than 200 characters, but make sure to construct complete sentences." +
+          "Use Markdown to format your response. " +
+          "Use latex to write math formulas. " +
           "Respond directly without any conversation starters.",
       },
       {
@@ -38,6 +40,8 @@ export async function POST(req: Request) {
           "You are an AI writing assistant for the text editor application 'Math Editor'. " +
           "You are asked to rewrite what user writes in another way. " +
           "Limit your response to no more than 200 characters, but make sure to construct complete sentences." +
+          "Use Markdown to format your response. " +
+          "Use latex to write math formulas. " +
           "Respond directly without any conversation starters.",
       },
       {
@@ -51,6 +55,8 @@ export async function POST(req: Request) {
         content:
           "You are an AI writing assistant for the text editor application 'Math Editor'. " +
           "You are asked to rewrite what user writes in a shorter form. " +
+          "Use Markdown to format your response. " +
+          "Use latex to write math formulas. " +
           "Respond directly without any conversation starters.",
       },
       {
@@ -64,6 +70,8 @@ export async function POST(req: Request) {
         content:
           "You are an AI writing assistant for the text editor application 'Math Editor'. " +
           "You are asked to rewrite what user writes in a longer form. " +
+          "Use Markdown to format your response. " +
+          "Use latex to write math formulas. " +
           "Respond directly without any conversation starters.",
       },
       {
@@ -77,6 +85,8 @@ export async function POST(req: Request) {
         content:
           "You are an AI writing assistant for the text editor application 'Math Editor'. " +
           "You are asked to help the user with his document. " +
+          "Use Markdown to format your response. " +
+          "Use latex to write math formulas. " +
           "Respond directly without any conversation starters.",
       },
       {

diff --git a/src/editor/plugins/ToolbarPlugin/Tools/AITools.tsx b/src/editor/plugins/ToolbarPlugin/Tools/AITools.tsx
@@ -1,5 +1,5 @@
 "use client"
-import { $getSelection, $isRangeSelection, BLUR_COMMAND, CLICK_COMMAND, COMMAND_PRIORITY_CRITICAL, INSERT_PARAGRAPH_COMMAND, KEY_DOWN_COMMAND, LexicalEditor, LexicalNode, SELECTION_CHANGE_COMMAND, } from "lexical";
+import { $createTextNode, $getSelection, $isRangeSelection, BLUR_COMMAND, CLICK_COMMAND, COMMAND_PRIORITY_CRITICAL, INSERT_PARAGRAPH_COMMAND, KEY_DOWN_COMMAND, LexicalEditor, LexicalNode, SELECTION_CHANGE_COMMAND, } from "lexical";
 import { mergeRegister } from "@lexical/utils";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { Menu, Button, MenuItem, ListItemIcon, ListItemText, Typography, TextField, CircularProgress } from "@mui/material";
@@ -9,6 +9,8 @@ import { useCompletion } from "ai/react";
 import { SET_DIALOGS_COMMAND } from "../Dialogs/commands";
 import { ANNOUNCE_COMMAND, UPDATE_DOCUMENT_COMMAND } from "@/editor/commands";
 import { Announcement } from "@/types";
+import { $isCodeNode } from "@/editor/nodes/CodeNode";
+import { $isListNode } from "@/editor/nodes/ListNode";
 
 export default function AITools({ editor, sx }: { editor: LexicalEditor, sx?: SxProps<Theme> }): JSX.Element {
   const [anchorEl, setAnchorEl] = useState<null | HTMLElement>(null);
@@ -123,18 +125,53 @@ export default function AITools({ editor, sx }: { editor: LexicalEditor, sx?: Sx
     const hasCompletion = completion.length > 0;
     if (!hasCompletion) return;
     const isStarting = offset.current === 0;
+    let shouldInsertNewlineOnUpdate = false;
     editor.update(() => {
       const selection = $getSelection();
       if (!$isRangeSelection(selection)) return;
       const delta = completion.slice(offset.current);
       offset.current = completion.length;
+      const anchorNode = selection.anchor.getNode();
+      const elementNode = anchorNode.getTopLevelElement();
+      const isCodeNode = $isCodeNode(elementNode);
+      const isListNode = $isListNode(elementNode);
       const isCollapsed = selection.isCollapsed();
       const isAtNewline = selection.anchor.offset === 0 && selection.focus.offset === 0;
-      const shouldInsertNewline = isStarting && isCollapsed && !isAtNewline;
+      const shouldInsertNewline = isStarting && isCollapsed && !isAtNewline && !isCodeNode && !isListNode;
       const isEndinginNewline = delta.endsWith("\n\n") || delta.endsWith("\n") || delta.endsWith("<eos>");
-      if (shouldInsertNewline || isEndinginNewline) selection.insertParagraph();
-      else selection.insertText(delta);
-    }, { tag: !isStarting ? "history-merge" : undefined });
+      shouldInsertNewlineOnUpdate = isEndinginNewline && !isCodeNode && !isListNode;
+      if (shouldInsertNewline && !isCodeNode) selection.insertParagraph();
+      let newDelta = delta;
+      if (isEndinginNewline && !isCodeNode) newDelta = delta.trimEnd();
+      if (isCodeNode) {
+        const language = completion.match(/```(\w+)$/)?.[1];
+        if (language) return elementNode.setLanguage(language);
+        if (elementNode.getTextContent() === "\n") elementNode.getFirstChild()?.remove();
+        if (elementNode.getTextContentSize() === 0 && newDelta === "\n") newDelta = "";
+        const textNode = $createTextNode(newDelta);
+        elementNode.append(textNode).selectEnd();
+        const isEnding = elementNode.getTextContent().trimEnd().endsWith("```");
+        if (isEnding) {
+          elementNode.getLastChild()?.remove();
+          elementNode.getLastChild()?.remove();
+          elementNode.getLastChild()?.remove();
+          shouldInsertNewlineOnUpdate = true;
+          elementNode.selectNext().insertParagraph();
+        }
+      }
+      else selection.insertText(newDelta);
+      if (isListNode && isEndinginNewline) elementNode.selectNext().insertParagraph();
+    }, {
+      tag: !isStarting ? "history-merge" : undefined,
+      onUpdate() {
+        if (!shouldInsertNewlineOnUpdate) return;
+        editor.update(() => {
+          const selection = $getSelection();
+          if (!$isRangeSelection(selection)) return;
+          selection.insertParagraph();
+        }, { tag: "history-merge" });
+      },
+    });
   }, [completion]);
 
   useEffect(() => {

diff --git a/src/editor/plugins/ToolbarPlugin/index.tsx b/src/editor/plugins/ToolbarPlugin/index.tsx
@@ -312,13 +312,15 @@ function ToolbarPlugin() {
   const showImageTools = $isImageNode(selectedNode);
   const showTableTools = !!selectedTable;
   const showTextTools = (!showMathTools && !showImageTools) || $isStickyNode(selectedNode);
+  const showCodeTools = blockType === 'code';
+  const showTextFormatTools = showTextTools && !showCodeTools;
   const showAITools = !!isOnline;
 
   return (
     <>
       <AppBar elevation={toolbarTrigger ? 4 : 0} position={toolbarTrigger ? 'fixed' : 'static'}>
         <Toolbar className="editor-toolbar" sx={{
-          position: "relative", 
+          position: "relative",
           displayPrint: 'none', px: `${(toolbarTrigger ? 1 : 0)}!important`,
           justifyContent: "space-between", alignItems: "start", gap: 0.5, py: 1,
         }}>
@@ -334,28 +336,22 @@ function ToolbarPlugin() {
           </Box>
           <Box sx={{ display: "flex", gap: 0.5, mx: 'auto', flexWrap: "wrap", justifyContent: "center" }}>
             {showMathTools && <MathTools editor={activeEditor} node={selectedNode} />}
-            {(showImageTools) && <ImageTools editor={activeEditor} node={selectedNode} />}
+            {showImageTools && <ImageTools editor={activeEditor} node={selectedNode} />}
             {showTextTools && <>
               {blockType in blockTypeToBlockName && <BlockFormatSelect blockType={blockType} editor={activeEditor} />}
-              {blockType === 'code' ? (
-                <Select size='small' onChange={onCodeLanguageSelect} value={codeLanguage}>
-                  {CODE_LANGUAGE_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option}>{text}</MenuItem>)}
-                </Select>
-              ) : (
-                <>
-                  <Select size='small' sx={{ width: { xs: 68, md: "auto" } }} onChange={onFontFamilySelect} value={fontFamily}>
-                    {FONT_FAMILY_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option}>{text}</MenuItem>)}
-                  </Select>
-                  <Select size='small' sx={{ width: { xs: 68, md: "auto" } }} onChange={onFontSizeSelect} value={fontSize}>
-                    {FONT_SIZE_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option} sx={{ justifyContent: "center" }}>{text}</MenuItem>)}
-                  </Select>
-                  {showAITools && <AITools editor={activeEditor} />}
-                  {showTableTools && <TableTools editor={activeEditor} node={selectedTable} />}
-                  <TextFormatToggles editor={activeEditor} sx={{ display: { xs: "none", sm: "none", md: "none", lg: "flex" } }} />
-                </>
-              )}
-            </>
-            }
+              {showCodeTools && <Select size='small' onChange={onCodeLanguageSelect} value={codeLanguage}>
+                {CODE_LANGUAGE_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option}>{text}</MenuItem>)}
+              </Select>}
+              {showTextFormatTools && <Select size='small' sx={{ width: { xs: 68, md: "auto" } }} onChange={onFontFamilySelect} value={fontFamily}>
+                {FONT_FAMILY_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option}>{text}</MenuItem>)}
+              </Select>}
+              {showTextFormatTools && <Select size='small' sx={{ width: { xs: 68, md: "auto" } }} onChange={onFontSizeSelect} value={fontSize}>
+                {FONT_SIZE_OPTIONS.map(([option, text]) => <MenuItem key={option} value={option} sx={{ justifyContent: "center" }}>{text}</MenuItem>)}
+              </Select>}
+              {showAITools && <AITools editor={activeEditor} />}
+              {showTableTools && <TableTools editor={activeEditor} node={selectedTable} />}
+              {showTextFormatTools && <TextFormatToggles editor={activeEditor} sx={{ display: { xs: "none", sm: "none", md: "none", lg: "flex" } }} />}
+            </>}
           </Box>
           <Box sx={{ display: "flex", gridColumn: "3/-1" }}>
             <InsertToolMenu editor={activeEditor} />