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

JimmFly · JimmFly · commit f688c057eba2 · 2024-09-04T08:44:05.000Z
close PD-1658 https://github.com/user-attachments/assets/2f3d1d26-1879-4d95-b80c-7c0965cefbd0
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/ai-spec.ts b/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,23 +85,23 @@ 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) {
         setupEdgelessCopilot(view.component);
       }
 
       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);
       }
     });
   }
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/code-toolbar/setup-code-toolbar.ts b/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',
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/edgeless/index.ts b/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({
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/format-bar/setup-format-bar.ts b/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(
     [
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/image-toolbar/setup-image-toolbar.ts b/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(
     [
       {
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/slash-menu/setup-slash-menu.ts b/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']) => {
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/entries/space/setup-space.ts b/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');
diff --git a/packages/frontend/core/src/components/affine/ai-onboarding/index.tsx b/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>
   );
 };
diff --git a/packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx b/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 />
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx b/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(() => {
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts b/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();
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts b/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(
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
diff --git a/packages/frontend/core/src/modules/editor-settting/schema.ts b/packages/frontend/core/src/modules/editor-settting/schema.ts
diff --git a/packages/frontend/core/src/pages/workspace/detail-page/detail-page.tsx b/packages/frontend/core/src/pages/workspace/detail-page/detail-page.tsx
diff --git a/packages/frontend/core/src/providers/modal-provider.tsx b/packages/frontend/core/src/providers/modal-provider.tsx
diff --git a/packages/frontend/i18n/src/resources/en.json b/packages/frontend/i18n/src/resources/en.json

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ const buttonOptions: AskAIButtonOptions = {`
`12`	`12`	`import type { AskAIButtonOptions } from '../../_common/components/ask-ai-button';`
`13`	`13`	`import { buildAICodeItemGroups } from '../../_common/config';`
`14`	`14`
`15`		`-export function setupCodeToolbarEntry(codeToolbar: AffineCodeToolbarWidget) {`
	`15`	`+export function setupCodeToolbarAIEntry(codeToolbar: AffineCodeToolbarWidget) {`
`16`	`16`	`codeToolbar.addPrimaryItems([`
`17`	`17`	`{`
`18`	`18`	`type: 'ask-ai',`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ export function setupEdgelessCopilot(widget: EdgelessCopilotWidget) {`
`17`	`17`	`widget.groups = edgelessActionGroups;`
`18`	`18`	`}`
`19`	`19`
`20`		`-export function setupEdgelessElementToolbarEntry(`
	`20`	`+export function setupEdgelessElementToolbarAIEntry(`
`21`	`21`	`widget: EdgelessElementToolbarWidget`
`22`	`22`	`) {`
`23`	`23`	`widget.registerEntry({`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ import { html, type TemplateResult } from 'lit';`
`8`	`8`
`9`	`9`	`import { AIItemGroups } from '../../_common/config';`
`10`	`10`
`11`		`-export function setupFormatBarEntry(formatBar: AffineFormatBarWidget) {`
	`11`	`+export function setupFormatBarAIEntry(formatBar: AffineFormatBarWidget) {`
`12`	`12`	`toolbarDefaultConfig(formatBar);`
`13`	`13`	`formatBar.addRawConfigItems(`
`14`	`14`	`[`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,9 @@ const buttonOptions: AskAIButtonOptions = {`
`13`	`13`	`panelWidth: 300,`
`14`	`14`	`};`
`15`	`15`
`16`		`-export function setupImageToolbarEntry(imageToolbar: AffineImageToolbarWidget) {`
	`16`	`+export function setupImageToolbarAIEntry(`
	`17`	`+ imageToolbar: AffineImageToolbarWidget`
	`18`	`+) {`
`17`	`19`	`imageToolbar.addPrimaryItems(`
`18`	`20`	`[`
`19`	`21`	`{`