feat(ee): AI menu (#1912)

* feat(ee): AI menu

* - Add insert below and copy option

* prebuild @editor-ext

* sanitize output

* clear existing output

* switch to menu component

* refactor directory

* separator

* refactor directory

* support more languages

* pass markdown to model

* fix: close AI menu on page change

* enhance text input and preview styling

* fix: Use absolute positioning for the AI menu

* make preview scrollable

* activation controls

* enhance bubble menu

* sync

* set width

* fix line break

* switch terminologies

* cloud

* buffer

---------

Co-authored-by: Philipinho <16838612+Philipinho@users.noreply.github.com>

apps/client/public/locales/en-US/translation.json

@@ -582,11 +582,15 @@
   "Ask AI": "Ask AI",
   "AI is thinking...": "AI is thinking...",
   "Ask a question...": "Ask a question...",
-  "AI-powered search (Ask AI)": "AI-powered search (Ask AI)",
+  "AI Answers": "AI Answers",
+  "AI-powered search (AI Answers)": "AI-powered search (AI Answers)",
   "AI search uses vector embeddings to provide semantic search capabilities across your workspace content.": "AI search uses vector embeddings to provide semantic search capabilities across your workspace content.",
   "Toggle AI search": "Toggle AI search",
+  "Generative AI (Ask AI)": "Generative AI (Ask AI)",
+  "Enable AI-powered content generation in the editor. Allows users to generate, improve, translate and transform text.": "Enable AI-powered content generation in the editor. Allows users to generate, improve, translate and transform text.",
+  "Toggle generative AI": "Toggle generative AI",
   "Sources": "Sources",
-  "Ask AI not available for attachments": "Ask AI not available for attachments",
+  "AI Answers not available for attachments": "AI Answers not available for attachments",
   "No answer available": "No answer available",
   "Background color": "Background color",
   "Highlight color": "Highlight color",

apps/client/src/components/settings/settings-sidebar.tsx

@@ -115,7 +115,6 @@ const groupedData: DataGroup[] = [
         icon: IconSparkles,
         path: "/settings/ai",
         isAdmin: true,
-        isSelfhosted: true,
       },
     ],
   },

apps/client/src/ee/ai/components/editor/ai-menu/ai-menu.module.css

@@ -0,0 +1,61 @@
+.aiMenu {
+  display: flex;
+  flex-direction: column;
+  width: 100%;
+  max-width: 600px;
+  min-height: 2.25rem;
+}
+
+.aiInput {
+  width: 100%;
+
+  & input {
+    height: 44px;
+    border-radius: 22px;
+    padding-left: 20px;
+    padding-right: 40px;
+    border: 1px solid
+      light-dark(var(--mantine-color-gray-3), var(--mantine-color-dark-4));
+    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
+    font-size: var(--mantine-font-size-sm);
+
+    &:focus {
+      border-color: light-dark(
+        var(--mantine-color-gray-4),
+        var(--mantine-color-dark-3)
+      );
+      box-shadow: 0 4px 12px rgba(0, 0, 0, 0.12);
+    }
+  }
+}
+.menuItemSelected {
+  background-color: var(--mantine-color-gray-1);
+
+  @mixin dark {
+    background-color: var(--mantine-color-dark-5);
+  }
+}
+
+.resultPreview {
+  background-color: light-dark(
+    var(--mantine-color-white),
+    var(--mantine-color-dark-6)
+  );
+  border: 1px solid
+    light-dark(var(--mantine-color-gray-3), var(--mantine-color-dark-4));
+  box-shadow: 0 4px 16px rgba(0, 0, 0, 0.1);
+  overflow: hidden;
+}
+
+.resultPreviewWrapper {
+  font-size: var(--mantine-font-size-md);
+  line-height: 1.6;
+  padding: var(--mantine-spacing-md);
+
+  *:first-child {
+    margin-top: 0;
+  }
+  *:last-child {
+    margin-bottom: 0;
+  }
+}

apps/client/src/ee/ai/components/editor/ai-menu/ai-menu.tsx

@@ -0,0 +1,325 @@
+import { Editor } from "@tiptap/react";
+import { ActionIcon, TextInput, Tooltip } from "@mantine/core";
+import { useDebouncedCallback, useMediaQuery } from "@mantine/hooks";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { createPortal } from "react-dom";
+import { useAtom } from "jotai";
+import { IconArrowUp } from "@tabler/icons-react";
+import { showAiMenuAtom } from "@/features/editor/atoms/editor-atoms.ts";
+import { useAiGenerateStreamMutation } from "@/ee/ai/queries/ai-query.ts";
+import { AiAction } from "@/ee/ai/types/ai.types.ts";
+import { CommandItem, commandItems, CommandSet } from "./command-items.ts";
+import { CommandSelector } from "./command-selector.tsx";
+import { ResultPreview } from "./result-preview.tsx";
+import classes from "./ai-menu.module.css";
+import { marked } from "marked";
+import { DOMSerializer } from "@tiptap/pm/model";
+import { htmlToMarkdown } from "@docmost/editor-ext";
+import { useLocation } from "react-router-dom";
+
+interface EditorAiMenuProps {
+  editor: Editor | null;
+}
+
+const EditorAiMenu = ({ editor }: EditorAiMenuProps): JSX.Element | null => {
+  const aiGenerateStreamMutation = useAiGenerateStreamMutation();
+  const location = useLocation();
+  const isSmBreakpoint = useMediaQuery("(max-width: 48em)");
+  const [showAiMenu, setShowAiMenu] = useAtom(showAiMenuAtom);
+  const containerRef = useRef<HTMLDivElement | null>(null);
+  const inputRef = useRef<HTMLInputElement | null>(null);
+  const [prompt, setPrompt] = useState("");
+  const [output, setOutput] = useState("");
+  const [isLoading, setIsLoading] = useState(false);
+  const [selectedIndex, setSelectedIndex] = useState(-1);
+  const [activeCommandSet, setActiveCommandSet] = useState<CommandSet>("main");
+  const [lastAction, setLastAction] = useState<CommandItem | null>(null);
+  const [menuPlacement, setMenuPlacement] = useState<{
+    top: number;
+    left: number;
+    width: number;
+  }>({
+    top: 0,
+    left: 0,
+    width: 0,
+  });
+  const currentItems = useMemo(() => {
+    return commandItems[activeCommandSet].filter((item) => {
+      return item.name.toLowerCase().includes(prompt.toLowerCase());
+    });
+  }, [prompt, output, activeCommandSet]);
+  const updateMenuPlacement = useCallback(() => {
+    if (!editor || !showAiMenu) return;
+
+    const { view } = editor;
+    const { to } = editor.state.selection;
+    const editorRect = view.dom.getBoundingClientRect();
+    const cursorCoords = view.coordsAtPos(to);
+    const topOffset = 8;
+    const editorPadding = isSmBreakpoint ? 16 : 48;
+
+    setMenuPlacement({
+      top: cursorCoords.bottom + topOffset + window.scrollY,
+      left: editorRect.left + editorPadding + window.scrollX,
+      width: editorRect.width - editorPadding * 2,
+    });
+  }, [editor, showAiMenu, isSmBreakpoint]);
+  const resetMenu = useCallback(() => {
+    setPrompt("");
+    setOutput("");
+    setActiveCommandSet("main");
+    setLastAction(null);
+    aiGenerateStreamMutation.reset();
+  }, [aiGenerateStreamMutation.reset]);
+  const debouncedUpdateMenuPlacement = useDebouncedCallback(
+    updateMenuPlacement,
+    60,
+  );
+  const handleGenerate = useCallback(
+    (item?: CommandItem) => {
+      if (!editor || isLoading) return;
+
+      let command: CommandItem | null = item || null;
+
+      if (!command) {
+        if (!prompt) return;
+
+        command = {
+          id: "custom",
+          name: "Custom",
+          action: AiAction.CUSTOM,
+          prompt,
+        };
+      }
+
+      const { from, to } = editor.state.selection;
+      const slice = editor.state.doc.slice(from, to);
+      const serializer = DOMSerializer.fromSchema(editor.schema);
+      const fragment = serializer.serializeFragment(slice.content);
+      const wrapper = document.createElement("div");
+      wrapper.appendChild(fragment);
+      const content = htmlToMarkdown(wrapper.innerHTML);
+
+      setOutput("");
+      setIsLoading(true);
+      aiGenerateStreamMutation.mutate({
+        action: command.action,
+        prompt: command.prompt,
+        content,
+        onChunk: (chunk) => {
+          setOutput((output) => output + chunk.content);
+        },
+        onComplete: () => {
+          setIsLoading(false);
+          setActiveCommandSet("result");
+        },
+        onError: () => {
+          setIsLoading(false);
+          resetMenu();
+        },
+      });
+      setLastAction(command);
+    },
+    [
+      editor,
+      prompt,
+      isLoading,
+      aiGenerateStreamMutation.mutateAsync,
+      resetMenu,
+    ],
+  );
+  const handleCommand = useCallback(
+    (item?: CommandItem) => {
+      setPrompt("");
+
+      if (!item) {
+        return handleGenerate();
+      }
+      if (item.id === "back") {
+        return setActiveCommandSet("main");
+      }
+      if (item.id === "result-replace") {
+        const chain = editor.chain().focus();
+
+        if (lastAction.action === AiAction.CONTINUE_WRITING) {
+          chain.setTextSelection(editor.state.selection.to);
+        }
+
+        const html = (marked.parse(output) as string).trim();
+        // Strip <p> wrapper for single-paragraph output to preserve inline context
+        const content =
+          html.startsWith("<p>") &&
+          html.endsWith("</p>") &&
+          html.lastIndexOf("<p>") === 0
+            ? html.slice(3, -4)
+            : html;
+
+        chain.insertContent(content).run();
+
+        return setShowAiMenu(false);
+      }
+      if (item.id === "result-insert-below") {
+        editor
+          .chain()
+          .focus()
+          .setTextSelection(editor.state.selection.to)
+          .insertContent(marked.parse(output))
+          .run();
+
+        return setShowAiMenu(false);
+      }
+      if (item.id === "result-copy") {
+        navigator.clipboard.writeText(output);
+
+        return setShowAiMenu(false);
+      }
+      if (item.id === "result-discard") {
+        setOutput("");
+
+        return resetMenu();
+      }
+      if (item.id === "result-try-again" && lastAction) {
+        return handleGenerate(lastAction);
+      }
+      if (item.subCommandSet) {
+        return setActiveCommandSet(item.subCommandSet);
+      }
+
+      return handleGenerate(item);
+    },
+    [editor, output, lastAction, handleGenerate, resetMenu],
+  );
+  const handleKeyDown = useCallback(
+    (event: React.KeyboardEvent<HTMLInputElement>) => {
+      const totalItems = currentItems.length;
+      const cycleSize = totalItems + 1;
+
+      if (event.key === "Escape") {
+        return setShowAiMenu(false);
+      }
+
+      if (event.key === "ArrowDown" || event.key === "ArrowUp") {
+        event.preventDefault();
+
+        return setSelectedIndex((selectedIndex) => {
+          const direction = event.key === "ArrowDown" ? 1 : -1;
+          const newIndex = selectedIndex + direction;
+
+          if (newIndex < -1) return cycleSize - 1;
+          if (newIndex >= cycleSize) return 0;
+
+          return newIndex;
+        });
+      }
+
+      if (event.key === "Enter") {
+        event.preventDefault();
+
+        return handleCommand(currentItems[selectedIndex]);
+      }
+    },
+    [currentItems, selectedIndex],
+  );
+
+  useEffect(() => {
+    if (!editor) return;
+
+    const handleClose = () => setShowAiMenu(false);
+    const observer = new ResizeObserver(() => {
+      debouncedUpdateMenuPlacement();
+    });
+
+    updateMenuPlacement();
+    editor.on("focus", handleClose);
+    editor.on("blur", handleClose);
+    window.addEventListener("resize", debouncedUpdateMenuPlacement);
+    window.addEventListener("scroll", debouncedUpdateMenuPlacement, true);
+    observer.observe(editor.view.dom);
+
+    return () => {
+      editor.off("focus", handleClose);
+      editor.off("blur", handleClose);
+      window.removeEventListener("resize", debouncedUpdateMenuPlacement);
+      window.removeEventListener("scroll", debouncedUpdateMenuPlacement, true);
+      observer.disconnect();
+    };
+  }, [editor, updateMenuPlacement, debouncedUpdateMenuPlacement]);
+
+  useEffect(() => {
+    setShowAiMenu(false);
+  }, [location]);
+  useEffect(() => {
+    if (showAiMenu) {
+      resetMenu();
+    }
+  }, [showAiMenu, resetMenu]);
+  useEffect(() => {
+    // Focus input when menu opens or command set changes
+    requestAnimationFrame(() => {
+      inputRef.current?.focus({ preventScroll: true });
+    });
+  }, [showAiMenu, isLoading, currentItems]);
+  useEffect(() => {
+    if (!currentItems.length) {
+      setSelectedIndex(-1);
+    }
+    setSelectedIndex(prompt || activeCommandSet !== "main" ? 0 : -1);
+  }, [prompt, activeCommandSet, currentItems]);
+
+  if (!showAiMenu) return null;
+
+  return createPortal(
+    <div
+      style={{
+        zIndex: 200,
+        position: "absolute",
+        top: menuPlacement.top,
+        left: menuPlacement.left,
+        width: menuPlacement.width,
+        pointerEvents: "none",
+      }}
+    >
+      <div
+        className={classes.aiMenu}
+        style={{ pointerEvents: "auto" }}
+        tabIndex={0}
+        ref={containerRef}
+      >
+        <ResultPreview output={output} isLoading={isLoading} />
+        <CommandSelector
+          selectedIndex={selectedIndex}
+          isLoading={isLoading}
+          output={output}
+          currentItems={currentItems}
+          handleCommand={handleCommand}
+        >
+          <TextInput
+            ref={inputRef}
+            className={classes.aiInput}
+            placeholder="Ask AI..."
+            data-autofocus
+            value={prompt}
+            disabled={isLoading}
+            onChange={(e) => setPrompt(e.currentTarget.value)}
+            rightSection={
+              <ActionIcon
+                disabled={!prompt || isLoading}
+                variant="filled"
+                color="blue"
+                radius="xl"
+                size="sm"
+                onClick={() => handleGenerate()}
+              >
+                <IconArrowUp size={14} stroke={2.5} />
+              </ActionIcon>
+            }
+            onKeyDown={handleKeyDown}
+          />
+        </CommandSelector>
+      </div>
+    </div>,
+    document.body,
+  );
+};
+
+export { EditorAiMenu };

apps/client/src/ee/ai/components/editor/ai-menu/command-items.ts

@@ -0,0 +1,219 @@
+import { AiAction } from "@/ee/ai/types/ai.types.ts";
+import {
+  IconSparkles,
+  IconArrowsMaximize,
+  IconArrowsMinimize,
+  IconWriting,
+  IconHelp,
+  IconList,
+  IconMoodSmile,
+  IconLanguage,
+  IconTrash,
+  IconRefresh,
+  IconChevronLeft,
+  IconCheck,
+  IconArrowDownLeft,
+  IconCopy,
+  IconTextPlus,
+  IconAlignJustified,
+} from "@tabler/icons-react";
+
+interface CommandItem {
+  name: string;
+  id: string;
+  icon?: typeof IconSparkles;
+  action?: AiAction;
+  prompt?: string;
+  subCommandSet?: CommandSet;
+}
+
+type CommandSet = "main" | "tone" | "translate" | "result";
+
+const mainItems: CommandItem[] = [
+  {
+    id: "improve-writing",
+    name: "Improve writing",
+    icon: IconSparkles,
+    action: AiAction.IMPROVE_WRITING,
+  },
+  {
+    id: "fix-spelling-grammar",
+    name: "Fix spelling & grammar",
+    icon: IconCheck,
+    action: AiAction.FIX_SPELLING_GRAMMAR,
+  },
+  {
+    id: "make-longer",
+    name: "Make longer",
+    icon: IconTextPlus,
+    action: AiAction.MAKE_LONGER,
+  },
+  {
+    id: "make-shorter",
+    name: "Make shorter",
+    icon: IconAlignJustified,
+    action: AiAction.MAKE_SHORTER,
+  },
+  {
+    id: "continue-writing",
+    name: "Continue writing",
+    icon: IconWriting,
+    action: AiAction.CONTINUE_WRITING,
+  },
+  {
+    id: "explain",
+    name: "Explain",
+    icon: IconHelp,
+    action: AiAction.EXPLAIN,
+  },
+  {
+    id: "summarize",
+    name: "Summarize",
+    icon: IconList,
+    action: AiAction.SUMMARIZE,
+  },
+  {
+    id: "change-tone",
+    name: "Change tone",
+    icon: IconMoodSmile,
+    subCommandSet: "tone",
+  },
+  {
+    id: "translate",
+    name: "Translate",
+    icon: IconLanguage,
+    subCommandSet: "translate",
+  },
+];
+const toneItems: CommandItem[] = [
+  {
+    id: "back",
+    name: "Back",
+    icon: IconChevronLeft,
+  },
+  {
+    id: "tone-professional",
+    name: "Professional",
+    icon: IconMoodSmile,
+    action: AiAction.CHANGE_TONE,
+    prompt: "Professional",
+  },
+  {
+    id: "tone-casual",
+    name: "Casual",
+    icon: IconMoodSmile,
+    action: AiAction.CHANGE_TONE,
+    prompt: "Casual",
+  },
+  {
+    id: "tone-friendly",
+    name: "Friendly",
+    icon: IconMoodSmile,
+    action: AiAction.CHANGE_TONE,
+    prompt: "Friendly",
+  },
+];
+const translateItems: CommandItem[] = [
+  {
+    id: "back",
+    name: "Back",
+    icon: IconChevronLeft,
+  },
+  {
+    id: "translate-english",
+    name: "English",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "English",
+  },
+  {
+    id: "translate-spanish",
+    name: "Spanish",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Spanish",
+  },
+  {
+    id: "translate-german",
+    name: "German",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "German",
+  },
+  {
+    id: "translate-french",
+    name: "French",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "French",
+  },
+  {
+    id: "translate-dutch",
+    name: "Dutch",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Dutch",
+  },
+  {
+    id: "translate-portuguese",
+    name: "Portuguese",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Portuguese",
+  },
+  {
+    id: "translate-italian",
+    name: "Italian",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Italian",
+  },
+  {
+    id: "translate-japanese",
+    name: "Japanese",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Japanese",
+  },
+  {
+    id: "translate-korean",
+    name: "Korean",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Korean",
+  },
+  {
+    id: "translate-swedish",
+    name: "Swedish",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Swedish",
+  },
+  {
+    id: "translate-chinese",
+    name: "Chinese (Simplified)",
+    icon: IconLanguage,
+    action: AiAction.TRANSLATE,
+    prompt: "Simplified Chinese",
+  },
+];
+const resultItems: CommandItem[] = [
+  { id: "result-replace", name: "Replace", icon: IconCheck },
+  { id: "result-insert-below", name: "Insert below", icon: IconArrowDownLeft },
+  { id: "result-copy", name: "Copy", icon: IconCopy },
+  { id: "result-discard", name: "Discard", icon: IconTrash },
+  {
+    id: "result-try-again",
+    name: "Try again",
+    icon: IconRefresh,
+  },
+];
+const commandItems: Record<CommandSet, CommandItem[]> = {
+  main: mainItems,
+  tone: toneItems,
+  translate: translateItems,
+  result: resultItems,
+};
+
+export type { CommandItem, CommandSet };
+export { commandItems };

apps/client/src/ee/ai/components/editor/ai-menu/command-selector.tsx

@@ -0,0 +1,72 @@
+import { Loader, Menu, ScrollArea } from "@mantine/core";
+import { IconChevronRight } from "@tabler/icons-react";
+import { ReactNode } from "react";
+import { CommandItem } from "./command-items.ts";
+import classes from "./ai-menu.module.css";
+
+interface CommandSelectorProps {
+  selectedIndex: number;
+
+  isLoading: boolean;
+  output: string;
+  currentItems: CommandItem[];
+  children: ReactNode;
+  handleCommand(item: CommandItem): void;
+}
+
+const CommandSelector = ({
+  selectedIndex,
+  children,
+  isLoading,
+  output,
+  currentItems,
+  handleCommand,
+}: CommandSelectorProps) => {
+  return (
+    <Menu
+      opened={!isLoading && currentItems.length > 0}
+      middlewares={{ flip: false }}
+      position="bottom-start"
+      offset={4}
+      width={250}
+      trapFocus={false}
+      shadow="lg"
+    >
+      <Menu.Target>{children}</Menu.Target>
+      <Menu.Dropdown>
+        <ScrollArea.Autosize type="scroll" scrollbarSize={5} mah={300}>
+          {currentItems.map((item, index) => {
+            const isSelected = selectedIndex === index;
+            const showLoader =
+              isLoading && output === "" && !item.subCommandSet;
+
+            return (
+              <Menu.Item
+                key={item.id}
+                className={isSelected ? classes.menuItemSelected : undefined}
+                leftSection={
+                  showLoader ? (
+                    <Loader size={14} />
+                  ) : item.icon ? (
+                    <item.icon size={16} />
+                  ) : undefined
+                }
+                rightSection={
+                  item.subCommandSet ? (
+                    <IconChevronRight size={14} />
+                  ) : undefined
+                }
+                onClick={() => handleCommand(item)}
+                disabled={isLoading}
+              >
+                {item.name}
+              </Menu.Item>
+            );
+          })}
+        </ScrollArea.Autosize>
+      </Menu.Dropdown>
+    </Menu>
+  );
+};
+
+export { CommandSelector };

apps/client/src/ee/ai/components/editor/ai-menu/result-preview.tsx

@@ -0,0 +1,32 @@
+import { Loader, Paper, ScrollArea } from "@mantine/core";
+import DOMPurify from "dompurify";
+import { marked } from "marked";
+import { memo } from "react";
+import classes from "./ai-menu.module.css";
+
+interface ResultPreviewProps {
+  output: string;
+  isLoading: boolean;
+}
+const ResultPreview = memo(({ output, isLoading }: ResultPreviewProps) => {
+  if (!output && !isLoading) return;
+
+  const parsedOutput = `${marked.parse(output)}`;
+
+  return (
+    <Paper mb={4} shadow="lg" radius="md" className={classes.resultPreview}>
+      <ScrollArea.Autosize mah={300} type="scroll" scrollbarSize={5}>
+        <div className={classes.resultPreviewWrapper}>
+          {parsedOutput && (
+            <div
+              dangerouslySetInnerHTML={{ __html: DOMPurify.sanitize(parsedOutput) }}
+            />
+          )}
+          {isLoading && <Loader size={12} ml="xs" display="inline-block" />}
+        </div>
+      </ScrollArea.Autosize>
+    </Paper>
+  );
+});
+
+export { ResultPreview };

apps/client/src/ee/ai/components/enable-ai-search.tsx

@@ -15,7 +15,7 @@ export default function EnableAiSearch() {
     <>
       <Group justify="space-between" wrap="nowrap" gap="xl">
         <div>
-          <Text size="md">{t("AI-powered search (Ask AI)")}</Text>
+          <Text size="md">{t("AI-powered search (AI Answers)")}</Text>
           <Text size="sm" c="dimmed">
             {t(
               "AI search uses vector embeddings to provide semantic search capabilities across your workspace content.",

apps/client/src/ee/ai/components/enable-generative-ai.tsx

@@ -0,0 +1,48 @@
+import { Group, Text, Switch } from "@mantine/core";
+import { useAtom } from "jotai";
+import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
+import React, { useState } from "react";
+import { useTranslation } from "react-i18next";
+import { updateWorkspace } from "@/features/workspace/services/workspace-service.ts";
+import { notifications } from "@mantine/notifications";
+import { useIsCloudEE } from "@/hooks/use-is-cloud-ee.tsx";
+
+export default function EnableGenerativeAi() {
+  const { t } = useTranslation();
+  const [workspace, setWorkspace] = useAtom(workspaceAtom);
+  const [checked, setChecked] = useState(workspace?.settings?.ai?.generative);
+  const hasAccess = useIsCloudEE();
+
+  const handleChange = async (event: React.ChangeEvent<HTMLInputElement>) => {
+    const value = event.currentTarget.checked;
+    try {
+      const updatedWorkspace = await updateWorkspace({ generativeAi: value });
+      setChecked(value);
+      setWorkspace(updatedWorkspace);
+    } catch (err) {
+      notifications.show({
+        message: err?.response?.data?.message,
+        color: "red",
+      });
+    }
+  };
+
+  return (
+    <Group justify="space-between" wrap="nowrap" gap="xl">
+      <div>
+        <Text size="md">{t("Generative AI (Ask AI)")}</Text>
+        <Text size="sm" c="dimmed">
+          {t(
+            "Enable AI-powered content generation in the editor. Allows users to generate, improve, translate and transform text.",
+          )}
+        </Text>
+      </div>
+
+      <Switch
+        defaultChecked={checked}
+        onChange={handleChange}
+        disabled={!hasAccess}
+      />
+    </Group>
+  );
+}

apps/client/src/ee/ai/hooks/use-ai-search.ts

@@ -1,6 +1,6 @@
 import { useMutation, UseMutationResult } from "@tanstack/react-query";
 import { useState, useCallback } from "react";
-import { askAi, IAiSearchResponse } from "@/ee/ai/services/ai-search-service.ts";
+import { aiAnswers, IAiSearchResponse } from "@/ee/ai/services/ai-search-service.ts";
 import { IPageSearchParams } from "@/features/search/types/search.types.ts";
 
 // @ts-ignore
@@ -26,7 +26,7 @@ export function useAiSearch(): UseAiSearchResult {
 
       const { contentType, ...apiParams } = params;
 
-      return await askAi(apiParams, (chunk) => {
+      return await aiAnswers(apiParams, (chunk) => {
         if (chunk.content) {
           setStreamingAnswer((prev) => prev + chunk.content);
         }

apps/client/src/ee/ai/pages/ai-settings.tsx

@@ -1,25 +1,25 @@
 import { Helmet } from "react-helmet-async";
-import { getAppName, isCloud } from "@/lib/config.ts";
+import { getAppName } from "@/lib/config.ts";
 import SettingsTitle from "@/components/settings/settings-title.tsx";
 import React from "react";
 import useUserRole from "@/hooks/use-user-role.tsx";
 import { useTranslation } from "react-i18next";
-import useLicense from "@/ee/hooks/use-license.tsx";
 import EnableAiSearch from "@/ee/ai/components/enable-ai-search.tsx";
-import { Alert } from "@mantine/core";
+import EnableGenerativeAi from "@/ee/ai/components/enable-generative-ai.tsx";
+import { Alert, Stack } from "@mantine/core";
 import { IconInfoCircle } from "@tabler/icons-react";
+import { useIsCloudEE } from "@/hooks/use-is-cloud-ee.tsx";
+import { isCloud } from "@/lib/config.ts";
 
 export default function AiSettings() {
   const { t } = useTranslation();
   const { isAdmin } = useUserRole();
-  const { hasLicenseKey } = useLicense();
+  const hasAccess = useIsCloudEE();
 
   if (!isAdmin) {
     return null;
   }
 
-  const hasAccess = isCloud() || (!isCloud() && hasLicenseKey);
-
   return (
     <>
       <Helmet>
@@ -40,7 +40,10 @@ export default function AiSettings() {
         </Alert>
       )}
 
-      <EnableAiSearch />
+      <Stack gap="md">
+        {!isCloud() && <EnableAiSearch />}
+        <EnableGenerativeAi />
+      </Stack>
     </>
   );
 }

apps/client/src/ee/ai/services/ai-search-service.ts

@@ -15,11 +15,11 @@ export interface IAiSearchResponse {
   }>;
 }
 
-export async function askAi(
+export async function aiAnswers(
   params: IPageSearchParams,
   onChunk?: (chunk: { content?: string; sources?: any[] }) => void,
 ): Promise<IAiSearchResponse> {
-  const response = await fetch("/api/ai/ask", {
+  const response = await fetch("/api/ai/answers", {
     method: "POST",
     headers: {
       "Content-Type": "application/json",

apps/client/src/ee/ai/services/ai-service.ts

@@ -43,13 +43,16 @@ export async function generateAiContentStream(
     }
 
     const processStream = async () => {
+      let buffer = "";
       try {
         while (true) {
           const { done, value } = await reader.read();
           if (done) break;
 
-          const chunk = decoder.decode(value, { stream: true });
+          buffer += decoder.decode(value, { stream: true });
-          const lines = chunk.split("\n");
+          const lines = buffer.split("\n");
+
+          buffer = lines.pop() || "";
 
           for (const line of lines) {
             if (line.startsWith("data: ")) {
@@ -66,7 +69,7 @@ export async function generateAiContentStream(
                   onChunk(parsed);
                 }
               } catch (e) {
-                // Ignore parse errors for incomplete chunks
+                // Skip invalid JSON
               }
             }
           }

apps/client/src/ee/ai/types/ai.types.ts

@@ -6,6 +6,7 @@ export enum AiAction {
   SIMPLIFY = "simplify",
   CHANGE_TONE = "change_tone",
   SUMMARIZE = "summarize",
+  EXPLAIN = "explain",
   CONTINUE_WRITING = "continue_writing",
   TRANSLATE = "translate",
   CUSTOM = "custom",

apps/client/src/features/editor/atoms/editor-atoms.ts

@@ -8,3 +8,5 @@ export const titleEditorAtom = atom<Editor | null>(null);
 export const readOnlyEditorAtom = atom<Editor | null>(null);
 
 export const yjsConnectionStatusAtom = atom<string>("");
+
+export const showAiMenuAtom = atom(false);

apps/client/src/features/editor/components/bubble-menu/bubble-menu.module.css

@@ -1,11 +1,38 @@
 .bubbleMenu {
   display: flex;
+  flex-wrap: nowrap;
+  overflow-x: auto;
+  max-width: 100vw;
   width: fit-content;
-  border-radius: 2px;
+  box-shadow: 0 4px 12px light-dark(#cfcfcf, #0f0f0f);
+  border-radius: 6px;
   border: 1px solid
     light-dark(var(--mantine-color-gray-3), var(--mantine-color-gray-8));
 
+  > * {
+    flex-shrink: 0;
+  }
+
+  scrollbar-width: none;
+
+  &::-webkit-scrollbar {
+    display: none;
+  }
+
   .active {
     color: light-dark(var(--mantine-color-blue-8), var(--mantine-color-gray-5));
   }
 }
+
+.buttonRoot {
+  height: 34px;
+  padding-left: rem(8);
+  padding-right: rem(4);
+  border: none;
+  border-radius: 6px;
+}
+
+.buttonSeparator {
+  border-right: 1px solid
+    light-dark(var(--mantine-color-gray-3), var(--mantine-color-dark-4)) !important;
+}

apps/client/src/features/editor/components/bubble-menu/bubble-menu.tsx

@@ -9,10 +9,11 @@ import {
   IconStrikethrough,
   IconUnderline,
   IconMessage,
+  IconSparkles,
 } from "@tabler/icons-react";
 import clsx from "clsx";
 import classes from "./bubble-menu.module.css";
-import { ActionIcon, rem, Tooltip } from "@mantine/core";
+import { ActionIcon, Button, rem, Tooltip } from "@mantine/core";
 import { ColorSelector } from "./color-selector";
 import { NodeSelector } from "./node-selector";
 import { TextAlignmentSelector } from "./text-alignment-selector";
@@ -20,11 +21,13 @@ import {
   draftCommentIdAtom,
   showCommentPopupAtom,
 } from "@/features/comment/atoms/comment-atom";
-import { useAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { v7 as uuid7 } from "uuid";
 import { isCellSelection, isTextSelected } from "@docmost/editor-ext";
 import { LinkSelector } from "@/features/editor/components/bubble-menu/link-selector.tsx";
 import { useTranslation } from "react-i18next";
+import { showAiMenuAtom } from "@/features/editor/atoms/editor-atoms";
+import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
 
 export interface BubbleMenuItem {
   name: string;
@@ -39,14 +42,22 @@ type EditorBubbleMenuProps = Omit<BubbleMenuProps, "children" | "editor"> & {
 
 export const EditorBubbleMenu: FC<EditorBubbleMenuProps> = (props) => {
   const { t } = useTranslation();
+  const [showAiMenu, setShowAiMenu] = useAtom(showAiMenuAtom);
   const [showCommentPopup, setShowCommentPopup] = useAtom(showCommentPopupAtom);
+  const workspace = useAtomValue(workspaceAtom);
+  const isGenerativeAiEnabled = workspace?.settings?.ai?.generative === true;
   const [, setDraftCommentId] = useAtom(draftCommentIdAtom);
   const showCommentPopupRef = useRef(showCommentPopup);
+  const showAiMenuRef = useRef(showAiMenu);
 
   useEffect(() => {
     showCommentPopupRef.current = showCommentPopup;
   }, [showCommentPopup]);
 
+  useEffect(() => {
+    showAiMenuRef.current = showAiMenu;
+  }, [showAiMenu]);
+
   const editorState = useEditorState({
     editor: props.editor,
     selector: (ctx) => {
@@ -123,6 +134,7 @@ export const EditorBubbleMenu: FC<EditorBubbleMenuProps> = (props) => {
         empty ||
         isNodeSelection(selection) ||
         isCellSelection(selection) ||
+        showAiMenuRef.current ||
         showCommentPopupRef?.current
       ) {
         return false;
@@ -146,9 +158,28 @@ export const EditorBubbleMenu: FC<EditorBubbleMenuProps> = (props) => {
   const [isLinkSelectorOpen, setIsLinkSelectorOpen] = useState(false);
   const [isColorSelectorOpen, setIsColorSelectorOpen] = useState(false);
 
+  // Hide the bubble menu immediately when AI menu is shown
+  if (showAiMenu) return;
+
   return (
-    <BubbleMenu {...bubbleMenuProps} style={{ zIndex: 200, position: "relative"}}>
+    <BubbleMenu
+      {...bubbleMenuProps}
+      style={{ zIndex: 200, position: "relative" }}
+    >
       <div className={classes.bubbleMenu}>
+        {isGenerativeAiEnabled && (
+          <Button
+            variant="default"
+            className={clsx(classes.buttonRoot)}
+            radius="0"
+            leftSection={<IconSparkles size={16} />}
+            onClick={() => {
+              setShowAiMenu(true);
+            }}
+          >
+            {t("Ask AI")}
+          </Button>
+        )}
         <NodeSelector
           editor={props.editor}
           isOpen={isNodeSelectorOpen}
@@ -215,7 +246,7 @@ export const EditorBubbleMenu: FC<EditorBubbleMenuProps> = (props) => {
         <ActionIcon
           variant="default"
           size="lg"
-          radius="0"
+          radius="6px"
           aria-label={t(commentItem.name)}
           style={{ border: "none" }}
           onClick={commentItem.command}

apps/client/src/features/editor/components/bubble-menu/color-selector.tsx

@@ -1,7 +1,6 @@
 import React, { Dispatch, FC, SetStateAction } from "react";
-import { IconCheck, IconChevronDown, IconPalette } from "@tabler/icons-react";
+import { IconCheck, IconChevronDown } from "@tabler/icons-react";
 import {
-  ActionIcon,
   Button,
   Popover,
   rem,
@@ -15,6 +14,8 @@ import {
 import type { Editor } from "@tiptap/react";
 import { useEditorState } from "@tiptap/react";
 import { useTranslation } from "react-i18next";
+import clsx from "clsx";
+import classes from "./bubble-menu.module.css";
 
 export interface BubbleColorMenuItem {
   name: string;
@@ -166,14 +167,10 @@ export const ColorSelector: FC<ColorSelectorProps> = ({
             onClick={() => setIsOpen(!isOpen)}
             data-text-color={activeColorItem?.color || ""}
             data-highlight-color={activeHighlightItem?.color || ""}
-            className="color-selector-trigger"
+            className={clsx(["color-selector-trigger", classes.buttonRoot])}
             style={{
-              height: "34px",
-              border: "none",
               fontWeight: 500,
               fontSize: rem(16),
-              paddingLeft: rem(8),
-              paddingRight: rem(4),
             }}
           >
             A

apps/client/src/features/editor/components/bubble-menu/node-selector.tsx

@@ -16,6 +16,7 @@ import { Popover, Button, ScrollArea } from "@mantine/core";
 import type { Editor } from "@tiptap/react";
 import { useEditorState } from "@tiptap/react";
 import { useTranslation } from "react-i18next";
+import classes from "./bubble-menu.module.css";
 
 interface NodeSelectorProps {
   editor: Editor | null;
@@ -133,6 +134,7 @@ export const NodeSelector: FC<NodeSelectorProps> = ({
     <Popover opened={isOpen} withArrow>
       <Popover.Target>
         <Button
+          className={classes.buttonRoot}
           variant="default"
           style={{ border: "none", height: "34px" }}
           radius="0"

apps/client/src/features/editor/page-editor.tsx

@@ -66,6 +66,7 @@ import { PageEditMode } from "@/features/user/types/user.types.ts";
 import { jwtDecode } from "jwt-decode";
 import { searchSpotlight } from "@/features/search/constants.ts";
 import { useEditorScroll } from "./hooks/use-editor-scroll";
+import { EditorAiMenu } from "@/ee/ai/components/editor/ai-menu/ai-menu";
 
 interface PageEditorProps {
   pageId: string;
@@ -405,6 +406,7 @@ export default function PageEditor({
 
         {editor && editorIsEditable && (
           <div>
+            <EditorAiMenu editor={editor} />
             <EditorBubbleMenu editor={editor} />
             <TableMenu editor={editor} />
             <TableCellMenu editor={editor} appendTo={menuContainerRef} />

apps/client/src/features/search/components/search-spotlight-filters.tsx

@@ -140,7 +140,7 @@ export function SearchSpotlightFilters({
           <Switch
             checked={isAiMode}
             onChange={(event) => onAskClick()}
-            label={t("Ask AI")}
+            label={t("AI Answers")}
             size="sm"
             color="blue"
             labelPosition="left"
@@ -279,7 +279,7 @@ export function SearchSpotlightFilters({
                     isAiMode &&
                     option.value === "attachment" && (
                       <Text size="xs" mt={4}>
-                        {t("Ask AI not available for attachments")}
+                        {t("AI Answers not available for attachments")}
                       </Text>
                     )}
                 </div>

apps/client/src/features/workspace/types/workspace.types.ts

@@ -23,6 +23,7 @@ export interface IWorkspace {
   hasLicenseKey?: boolean;
   enforceMfa?: boolean;
   aiSearch?: boolean;
+  generativeAi?: boolean;
   disablePublicSharing?: boolean;
 }
 
@@ -33,6 +34,7 @@ export interface IWorkspaceSettings {
 
 export interface IWorkspaceAiSettings {
   search?: boolean;
+  generative?: boolean;
 }
 
 export interface IWorkspaceSharingSettings {

apps/server/src/core/workspace/services/workspace.service.ts

@@ -118,6 +118,7 @@ export class WorkspaceService {
         let status = undefined;
         let plan = undefined;
         let billingEmail = undefined;
+        let settings = undefined;
 
         if (this.environmentService.isCloud()) {
           // generate unique hostname
@@ -131,6 +132,7 @@ export class WorkspaceService {
           status = WorkspaceStatus.Active;
           plan = 'standard';
           billingEmail = user.email;
+          settings = { ai: { generative: true } };
         }
 
         // create workspace
@@ -143,6 +145,7 @@ export class WorkspaceService {
             trialEndAt,
             plan,
             billingEmail,
+            settings,
           },
           trx,
         );

nx.json

@@ -6,6 +6,11 @@
       ],
       "cache": true
     },
+    "start:dev": {
+      "dependsOn": [
+        "^build"
+      ]
+    },
     "lint": {
       "cache": true
     }