feat(core): impl AI switch (#8018)

close PD-1658

https://github.com/user-attachments/assets/2f3d1d26-1879-4d95-b80c-7c0965cefbd0

packages/frontend/core/src/blocksuite/presets/ai/ai-spec.ts

@@ -26,15 +26,15 @@ import { assertInstanceOf } from '@blocksuite/global/utils';
 import { literal, unsafeStatic } from 'lit/static-html.js';
 
 import { buildAIPanelConfig } from './ai-panel';
-import { setupCodeToolbarEntry } from './entries/code-toolbar/setup-code-toolbar';
+import { setupCodeToolbarAIEntry } from './entries/code-toolbar/setup-code-toolbar';
 import {
   setupEdgelessCopilot,
-  setupEdgelessElementToolbarEntry,
+  setupEdgelessElementToolbarAIEntry,
 } from './entries/edgeless/index';
-import { setupFormatBarEntry } from './entries/format-bar/setup-format-bar';
-import { setupImageToolbarEntry } from './entries/image-toolbar/setup-image-toolbar';
-import { setupSlashMenuEntry } from './entries/slash-menu/setup-slash-menu';
-import { setupSpaceEntry } from './entries/space/setup-space';
+import { setupFormatBarAIEntry } from './entries/format-bar/setup-format-bar';
+import { setupImageToolbarAIEntry } from './entries/image-toolbar/setup-image-toolbar';
+import { setupSlashMenuAIEntry } from './entries/slash-menu/setup-slash-menu';
+import { setupSpaceAIEntry } from './entries/space/setup-space';
 
 class AIPageRootWatcher extends BlockServiceWatcher {
   static override readonly flavour = 'affine:page';
@@ -45,15 +45,15 @@ class AIPageRootWatcher extends BlockServiceWatcher {
       if (view.component instanceof AffineAIPanelWidget) {
         view.component.style.width = '630px';
         view.component.config = buildAIPanelConfig(view.component);
-        setupSpaceEntry(view.component);
+        setupSpaceAIEntry(view.component);
       }
 
       if (view.component instanceof AffineFormatBarWidget) {
-        setupFormatBarEntry(view.component);
+        setupFormatBarAIEntry(view.component);
       }
 
       if (view.component instanceof AffineSlashMenuWidget) {
-        setupSlashMenuEntry(view.component);
+        setupSlashMenuAIEntry(view.component);
       }
     });
   }
@@ -85,7 +85,7 @@ class AIEdgelessRootWatcher extends BlockServiceWatcher {
       if (view.component instanceof AffineAIPanelWidget) {
         view.component.style.width = '430px';
         view.component.config = buildAIPanelConfig(view.component);
-        setupSpaceEntry(view.component);
+        setupSpaceAIEntry(view.component);
       }
 
       if (view.component instanceof EdgelessCopilotWidget) {
@@ -93,15 +93,15 @@ class AIEdgelessRootWatcher extends BlockServiceWatcher {
       }
 
       if (view.component instanceof EdgelessElementToolbarWidget) {
-        setupEdgelessElementToolbarEntry(view.component);
+        setupEdgelessElementToolbarAIEntry(view.component);
       }
 
       if (view.component instanceof AffineFormatBarWidget) {
-        setupFormatBarEntry(view.component);
+        setupFormatBarAIEntry(view.component);
       }
 
       if (view.component instanceof AffineSlashMenuWidget) {
-        setupSlashMenuEntry(view.component);
+        setupSlashMenuAIEntry(view.component);
       }
     });
   }
@@ -166,7 +166,7 @@ class AICodeBlockWatcher extends BlockServiceWatcher {
     const service = this.blockService;
     service.specSlots.widgetConnected.on(view => {
       if (view.component instanceof AffineCodeToolbarWidget) {
-        setupCodeToolbarEntry(view.component);
+        setupCodeToolbarAIEntry(view.component);
       }
     });
   }
@@ -184,7 +184,7 @@ class AIImageBlockWatcher extends BlockServiceWatcher {
     super.mounted();
     this.blockService.specSlots.widgetConnected.on(view => {
       if (view.component instanceof AffineImageToolbarWidget) {
-        setupImageToolbarEntry(view.component);
+        setupImageToolbarAIEntry(view.component);
       }
     });
   }

packages/frontend/core/src/blocksuite/presets/ai/entries/code-toolbar/setup-code-toolbar.ts

@@ -12,7 +12,7 @@ const buttonOptions: AskAIButtonOptions = {
 import type { AskAIButtonOptions } from '../../_common/components/ask-ai-button';
 import { buildAICodeItemGroups } from '../../_common/config';
 
-export function setupCodeToolbarEntry(codeToolbar: AffineCodeToolbarWidget) {
+export function setupCodeToolbarAIEntry(codeToolbar: AffineCodeToolbarWidget) {
   codeToolbar.addPrimaryItems([
     {
       type: 'ask-ai',

packages/frontend/core/src/blocksuite/presets/ai/entries/edgeless/index.ts

@@ -17,7 +17,7 @@ export function setupEdgelessCopilot(widget: EdgelessCopilotWidget) {
   widget.groups = edgelessActionGroups;
 }
 
-export function setupEdgelessElementToolbarEntry(
+export function setupEdgelessElementToolbarAIEntry(
   widget: EdgelessElementToolbarWidget
 ) {
   widget.registerEntry({

packages/frontend/core/src/blocksuite/presets/ai/entries/format-bar/setup-format-bar.ts

@@ -8,7 +8,7 @@ import { html, type TemplateResult } from 'lit';
 
 import { AIItemGroups } from '../../_common/config';
 
-export function setupFormatBarEntry(formatBar: AffineFormatBarWidget) {
+export function setupFormatBarAIEntry(formatBar: AffineFormatBarWidget) {
   toolbarDefaultConfig(formatBar);
   formatBar.addRawConfigItems(
     [

packages/frontend/core/src/blocksuite/presets/ai/entries/image-toolbar/setup-image-toolbar.ts

@@ -13,7 +13,9 @@ const buttonOptions: AskAIButtonOptions = {
   panelWidth: 300,
 };
 
-export function setupImageToolbarEntry(imageToolbar: AffineImageToolbarWidget) {
+export function setupImageToolbarAIEntry(
+  imageToolbar: AffineImageToolbarWidget
+) {
   imageToolbar.addPrimaryItems(
     [
       {

packages/frontend/core/src/blocksuite/presets/ai/entries/slash-menu/setup-slash-menu.ts

@@ -20,7 +20,7 @@ import { AIItemGroups } from '../../_common/config';
 import { handleInlineAskAIAction } from '../../actions/doc-handler';
 import { AIProvider } from '../../provider';
 
-export function setupSlashMenuEntry(slashMenu: AffineSlashMenuWidget) {
+export function setupSlashMenuAIEntry(slashMenu: AffineSlashMenuWidget) {
   const AIItems = AIItemGroups.map(group => group.items).flat();
 
   const iconWrapper = (icon: AIItemConfig['icon']) => {

packages/frontend/core/src/blocksuite/presets/ai/entries/space/setup-space.ts

@@ -3,7 +3,7 @@ import type { AffineAIPanelWidget } from '@blocksuite/blocks';
 import { handleInlineAskAIAction } from '../../actions/doc-handler';
 import { AIProvider } from '../../provider';
 
-export function setupSpaceEntry(panel: AffineAIPanelWidget) {
+export function setupSpaceAIEntry(panel: AffineAIPanelWidget) {
   panel.handleEvent('keyDown', ctx => {
     const host = panel.host;
     const keyboardState = ctx.get('keyboardState');

packages/frontend/core/src/components/affine/ai-onboarding/index.tsx

@@ -1,3 +1,5 @@
+import { EditorSettingService } from '@affine/core/modules/editor-settting';
+import { useLiveData, useService } from '@toeverything/infra';
 import { Suspense, useCallback, useEffect, useState } from 'react';
 
 import { AIOnboardingEdgeless } from './edgeless.dialog';
@@ -28,19 +30,29 @@ const useDismiss = (key: AIOnboardingType) => {
 export const WorkspaceAIOnboarding = () => {
   const [dismissGeneral] = useDismiss(AIOnboardingType.GENERAL);
   const [dismissLocal] = useDismiss(AIOnboardingType.LOCAL);
+  const editorSettingService = useService(EditorSettingService);
+  const enableAI = useLiveData(
+    editorSettingService.editorSetting.settings$.map(s => s.enableAI)
+  );
 
   return (
     <Suspense>
-      {dismissGeneral ? null : <AIOnboardingGeneral />}
-      {dismissLocal ? null : <AIOnboardingLocal />}
+      {!enableAI || dismissGeneral ? null : <AIOnboardingGeneral />}
+      {!enableAI || dismissLocal ? null : <AIOnboardingLocal />}
     </Suspense>
   );
 };
 
 export const PageAIOnboarding = () => {
   const [dismissEdgeless] = useDismiss(AIOnboardingType.EDGELESS);
+  const editorSettingService = useService(EditorSettingService);
+  const enableAI = useLiveData(
+    editorSettingService.editorSetting.settings$.map(s => s.enableAI)
+  );
 
   return (
-    <Suspense>{dismissEdgeless ? null : <AIOnboardingEdgeless />}</Suspense>
+    <Suspense>
+      {!enableAI || dismissEdgeless ? null : <AIOnboardingEdgeless />}
+    </Suspense>
   );
 };

packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx

@@ -8,6 +8,7 @@ import {
   type RadioItem,
   Scrollable,
   Switch,
+  useConfirmModal,
 } from '@affine/component';
 import {
   SettingRow,
@@ -393,16 +394,51 @@ export const SpellCheckSettings = () => {
   );
 };
 
+const AISettings = () => {
+  const t = useI18n();
+  const { openConfirmModal } = useConfirmModal();
+  const { editorSettingService } = useServices({ EditorSettingService });
+
+  const settings = useLiveData(editorSettingService.editorSetting.settings$);
+
+  const onAIChange = useCallback(
+    (checked: boolean) => {
+      editorSettingService.editorSetting.set('enableAI', checked);
+    },
+    [editorSettingService]
+  );
+  const onToggleAI = useCallback(
+    (checked: boolean) => {
+      openConfirmModal({
+        title: checked ? 'Enable AI?' : 'Disable AI?',
+        description: `Are you sure you want to ${checked ? 'enable' : 'disable'} AI?`,
+        confirmText: checked ? 'Enable' : 'Disable',
+        cancelText: 'Cancel',
+        onConfirm: () => onAIChange(checked),
+        confirmButtonOptions: {
+          variant: checked ? 'primary' : 'error',
+        },
+      });
+    },
+    [openConfirmModal, onAIChange]
+  );
+
+  return (
+    <SettingRow
+      name={t['com.affine.settings.editorSettings.general.ai.title']()}
+      desc={t['com.affine.settings.editorSettings.general.ai.description']()}
+    >
+      <Switch checked={settings.enableAI} onChange={onToggleAI} />
+    </SettingRow>
+  );
+};
+
 export const General = () => {
   const t = useI18n();
+
   return (
     <SettingWrapper title={t['com.affine.settings.editorSettings.general']()}>
-      <SettingRow
-        name={t['com.affine.settings.editorSettings.general.ai.title']()}
-        desc={t['com.affine.settings.editorSettings.general.ai.description']()}
-      >
-        <Switch />
-      </SettingRow>
+      <AISettings />
       <FontFamilySettings />
       <CustomFontFamilySettings />
       <NewDocDefaultModeSettings />

packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx

@@ -15,6 +15,7 @@ import {
   useFramework,
   useLiveData,
   useService,
+  useServices,
 } from '@toeverything/infra';
 import React, {
   forwardRef,
@@ -65,9 +66,13 @@ interface BlocksuiteEditorProps {
 
 const usePatchSpecs = (page: Doc, shared: boolean, mode: DocMode) => {
   const [reactToLit, portals] = useLitPortalFactory();
-  const peekViewService = useService(PeekViewService);
-  const docService = useService(DocService);
-  const docsService = useService(DocsService);
+  const { peekViewService, docService, docsService, editorSettingService } =
+    useServices({
+      PeekViewService,
+      DocService,
+      DocsService,
+      EditorSettingService,
+    });
   const framework = useFramework();
   const referenceRenderer: ReferenceReactRenderer = useMemo(() => {
     return function customReference(reference) {
@@ -89,10 +94,12 @@ const usePatchSpecs = (page: Doc, shared: boolean, mode: DocMode) => {
   }, [mode, page.collection]);
 
   const specs = useMemo(() => {
+    const enableAI =
+      editorSettingService.editorSetting.settings$.value.enableAI;
     return mode === 'edgeless'
-      ? createEdgelessModeSpecs(framework)
-      : createPageModeSpecs(framework);
-  }, [mode, framework]);
+      ? createEdgelessModeSpecs(framework, enableAI)
+      : createPageModeSpecs(framework, enableAI);
+  }, [editorSettingService, mode, framework]);
 
   const confirmModal = useConfirmModal();
   const patchedSpecs = useMemo(() => {

packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts

@@ -6,6 +6,7 @@ import {
 import type { ExtensionType } from '@blocksuite/block-std';
 import {
   BookmarkBlockSpec,
+  CodeBlockSpec,
   DatabaseBlockSpec,
   DataViewBlockSpec,
   DividerBlockSpec,
@@ -16,12 +17,15 @@ import {
   EmbedLoomBlockSpec,
   EmbedSyncedDocBlockSpec,
   EmbedYoutubeBlockSpec,
+  ImageBlockSpec,
   ListBlockSpec,
+  ParagraphBlockSpec,
 } from '@blocksuite/blocks';
+import { AIChatBlockSpec } from '@blocksuite/presets';
 
 import { CustomAttachmentBlockSpec } from './custom/attachment-block';
 
-export const CommonBlockSpecs: ExtensionType[] = [
+const CommonBlockSpecs: ExtensionType[] = [
   ListBlockSpec,
   DatabaseBlockSpec,
   DataViewBlockSpec,
@@ -36,7 +40,19 @@ export const CommonBlockSpecs: ExtensionType[] = [
   EmbedLinkedDocBlockSpec,
   // special
   CustomAttachmentBlockSpec,
+].flat();
+
+export const DefaultBlockSpecs: ExtensionType[] = [
+  CodeBlockSpec,
+  ImageBlockSpec,
+  ParagraphBlockSpec,
+  ...CommonBlockSpecs,
+].flat();
+
+export const AIBlockSpecs: ExtensionType[] = [
   AICodeBlockSpec,
   AIImageBlockSpec,
   AIParagraphBlockSpec,
+  AIChatBlockSpec,
+  ...CommonBlockSpecs,
 ].flat();

packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts

@@ -14,7 +14,9 @@ import {
 import type { RootService, TelemetryEventMap } from '@blocksuite/blocks';
 import {
   AffineCanvasTextFonts,
+  EdgelessRootBlockSpec,
   EdgelessRootService,
+  PageRootBlockSpec,
   PageRootService,
 } from '@blocksuite/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
@@ -55,11 +57,12 @@ function withAffineRootService(Service: typeof PageRootService) {
 }
 
 export function createPageRootBlockSpec(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   const editorSettingService = framework.get(EditorSettingService);
   return [
-    ...AIPageRootBlockSpec,
+    ...(enableAI ? AIPageRootBlockSpec : PageRootBlockSpec),
     {
       setup: di => {
         di.override(
@@ -79,11 +82,12 @@ export function createPageRootBlockSpec(
 }
 
 export function createEdgelessRootBlockSpec(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   const editorSettingService = framework.get(EditorSettingService);
   return [
-    ...AIEdgelessRootBlockSpec,
+    ...(enableAI ? AIEdgelessRootBlockSpec : EdgelessRootBlockSpec),
     {
       setup: di => {
         di.override(

packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts

@@ -6,24 +6,23 @@ import {
   EdgelessTextBlockSpec,
   FrameBlockSpec,
 } from '@blocksuite/blocks';
-import { AIChatBlockSpec } from '@blocksuite/presets';
 import type { FrameworkProvider } from '@toeverything/infra';
 
-import { CommonBlockSpecs } from './common';
+import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createEdgelessRootBlockSpec } from './custom/root-block';
 
 export function createEdgelessModeSpecs(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   return [
-    ...CommonBlockSpecs,
+    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     EdgelessSurfaceBlockSpec,
     EdgelessSurfaceRefBlockSpec,
     FrameBlockSpec,
     EdgelessTextBlockSpec,
     EdgelessNoteBlockSpec,
-    AIChatBlockSpec,
     // special
-    createEdgelessRootBlockSpec(framework),
+    createEdgelessRootBlockSpec(framework, enableAI),
   ].flat();
 }

packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts

@@ -4,22 +4,21 @@ import {
   PageSurfaceBlockSpec,
   PageSurfaceRefBlockSpec,
 } from '@blocksuite/blocks';
-import { AIChatBlockSpec } from '@blocksuite/presets';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { CommonBlockSpecs } from './common';
+import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createPageRootBlockSpec } from './custom/root-block';
 
 export function createPageModeSpecs(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   return [
-    ...CommonBlockSpecs,
+    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     PageSurfaceBlockSpec,
     PageSurfaceRefBlockSpec,
     NoteBlockSpec,
-    AIChatBlockSpec,
     // special
-    createPageRootBlockSpec(framework),
+    createPageRootBlockSpec(framework, enableAI),
   ].flat();
 }

packages/frontend/core/src/modules/editor-settting/schema.ts

@@ -29,6 +29,7 @@ export const fontStyleOptions = [
 }[];
 
 const AffineEditorSettingSchema = z.object({
+  enableAI: z.boolean().default(true),
   fontFamily: z.enum(['Sans', 'Serif', 'Mono', 'Custom']).default('Sans'),
   customFontFamily: z.string().default(''),
   newDocDefaultMode: z.enum(['edgeless', 'page']).default('page'),

packages/frontend/core/src/pages/workspace/detail-page/detail-page.tsx

@@ -7,6 +7,7 @@ import { EditorOutlineViewer } from '@affine/core/components/blocksuite/outline-
 import { useAppSettingHelper } from '@affine/core/hooks/affine/use-app-setting-helper';
 import { useDocMetaHelper } from '@affine/core/hooks/use-block-suite-page-meta';
 import { EditorService } from '@affine/core/modules/editor';
+import { EditorSettingService } from '@affine/core/modules/editor-settting';
 import { RecentDocsService } from '@affine/core/modules/quicksearch';
 import { ViewService } from '@affine/core/modules/workbench/services/view';
 import type { PageRootService } from '@blocksuite/blocks';
@@ -71,6 +72,7 @@ const DetailPageImpl = memo(function DetailPageImpl() {
     docService,
     workspaceService,
     globalContextService,
+    editorSettingService,
   } = useServices({
     WorkbenchService,
     ViewService,
@@ -78,6 +80,7 @@ const DetailPageImpl = memo(function DetailPageImpl() {
     DocService,
     WorkspaceService,
     GlobalContextService,
+    EditorSettingService,
   });
   const workbench = workbenchService.workbench;
   const editor = editorService.editor;
@@ -307,9 +310,15 @@ const DetailPageImpl = memo(function DetailPageImpl() {
         </div>
       </ViewBody>
 
-      <ViewSidebarTab tabId="chat" icon={<AiIcon />} unmountOnInactive={false}>
-        <EditorChatPanel editor={editorContainer} ref={chatPanelRef} />
-      </ViewSidebarTab>
+      {editorSettingService.editorSetting.settings$.value.enableAI && (
+        <ViewSidebarTab
+          tabId="chat"
+          icon={<AiIcon />}
+          unmountOnInactive={false}
+        >
+          <EditorChatPanel editor={editorContainer} ref={chatPanelRef} />
+        </ViewSidebarTab>
+      )}
 
       <ViewSidebarTab tabId="journal" icon={<TodayIcon />}>
         <EditorJournalPanel />

packages/frontend/core/src/providers/modal-provider.tsx

@@ -1,4 +1,4 @@
-import { NotificationCenter, notify } from '@affine/component';
+import { ConfirmModal, NotificationCenter, notify } from '@affine/component';
 import { events } from '@affine/electron-api';
 import { WorkspaceFlavour } from '@affine/env/workspace';
 import {
@@ -10,7 +10,7 @@ import {
 } from '@toeverything/infra';
 import { useAtom } from 'jotai';
 import type { ReactElement } from 'react';
-import { useCallback, useEffect } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 
 import type { SettingAtom } from '../atoms';
 import { openSettingModalAtom, openSignOutModalAtom } from '../atoms';
@@ -31,6 +31,7 @@ import { useAsyncCallback } from '../hooks/affine-async-hooks';
 import { useNavigateHelper } from '../hooks/use-navigate-helper';
 import { AuthService } from '../modules/cloud/services/auth';
 import { CreateWorkspaceDialogProvider } from '../modules/create-workspace';
+import { EditorSettingService } from '../modules/editor-settting';
 import { FindInPageModal } from '../modules/find-in-page/view/find-in-page-modal';
 import { ImportTemplateDialogProvider } from '../modules/import-template';
 import { PeekViewManagerModal } from '../modules/peek-view';
@@ -183,6 +184,40 @@ export const SignOutConfirmModal = () => {
   );
 };
 
+export const AIReloadConfirmModal = () => {
+  const editorSettingService = useService(EditorSettingService);
+  const enableAI = useLiveData(
+    editorSettingService.editorSetting.settings$
+  ).enableAI;
+  const [aiState] = useState(enableAI);
+  const [open, setOpen] = useState(false);
+
+  useEffect(() => {
+    setOpen(enableAI !== aiState);
+  }, [aiState, enableAI]);
+
+  const onConfirm = useCallback(() => {
+    window.location.reload();
+  }, []);
+
+  return (
+    <ConfirmModal
+      open={open}
+      onOpenChange={setOpen}
+      onConfirm={onConfirm}
+      confirmButtonOptions={{
+        variant: 'primary',
+      }}
+      title={'You need to reload the page'}
+      description={
+        'AI settings have been updated. Please reload the page to apply the changes.'
+      }
+      cancelText={'Cancel'}
+      confirmText={'Reload'}
+    />
+  );
+};
+
 export const AllWorkspaceModals = (): ReactElement => {
   return (
     <>
@@ -191,6 +226,7 @@ export const AllWorkspaceModals = (): ReactElement => {
       <CreateWorkspaceDialogProvider />
       <AuthModal />
       <SignOutConfirmModal />
+      <AIReloadConfirmModal />
     </>
   );
 };