feat(core): add ai model switch ui (#12266)

Close [AI-86](https://linear.app/affine-design/issue/AI-86)

![截屏2025-05-14 11.32.41.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/b92d5c32-fa5a-4afd-93e6-3699347575be.png)

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->

## Summary by CodeRabbit

- **New Features**
  - Introduced AI model switching in chat, allowing users to select from multiple AI models during conversations.
  - Added a floating menu for easy AI model selection within the chat interface.
  - Enabled visibility of the AI model switcher through a new experimental feature flag, configurable in workspace settings (canary builds only).

- **Enhancements**
  - Improved session management in the chat panel for smoother model switching and state handling.
  - Updated localization to support the new AI model switch feature in settings.

- **Bug Fixes**
  - None.

- **Chores**
  - Registered new components and services to support AI model switching functionality.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts

@@ -1,7 +1,7 @@
 import './chat-panel-messages';
 
 import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
-import type { ContextEmbedStatus } from '@affine/graphql';
+import type { ContextEmbedStatus, CopilotSessionType } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
@@ -18,6 +18,7 @@ import type {
   SearchMenuConfig,
 } from '../components/ai-chat-chips';
 import type {
+  AIModelSwitchConfig,
   AINetworkSearchConfig,
   AIReasoningConfig,
 } from '../components/ai-chat-input';
@@ -155,8 +156,8 @@ export class ChatPanel extends SignalWatcher(
   };
 
   private readonly _getSessionId = async () => {
-    if (this._sessionId) {
-      return this._sessionId;
+    if (this.session) {
+      return this.session.id;
     }
     const sessions = (
       (await AIProvider.session?.getSessions(
@@ -164,22 +165,33 @@ export class ChatPanel extends SignalWatcher(
         this.doc.id,
         { action: false }
       )) || []
-    ).filter(session => !session.parentSessionId);
-    const sessionId = sessions.at(-1)?.id;
-    this._sessionId = sessionId;
-    return this._sessionId;
+    ).filter(session => {
+      if (this.parentSessionId) {
+        return session.parentSessionId === this.parentSessionId;
+      } else {
+        return !session.parentSessionId;
+      }
+    });
+    this.session = sessions.at(-1);
+    return this.session?.id;
   };
 
   private readonly _createSessionId = async () => {
-    if (this._sessionId) {
-      return this._sessionId;
+    if (this.session) {
+      return this.session.id;
     }
-    this._sessionId = await AIProvider.session?.createSession({
+    const sessionId = await AIProvider.session?.createSession({
       docId: this.doc.id,
       workspaceId: this.doc.workspace.id,
       promptName: 'Chat With AFFiNE AI',
     });
-    return this._sessionId;
+    if (sessionId) {
+      this.session = await AIProvider.session?.getSession(
+        this.doc.workspace.id,
+        sessionId
+      );
+    }
+    return sessionId;
   };
 
   @property({ attribute: false })
@@ -194,6 +206,9 @@ export class ChatPanel extends SignalWatcher(
   @property({ attribute: false })
   accessor reasoningConfig!: AIReasoningConfig;
 
+  @property({ attribute: false })
+  accessor modelSwitchConfig!: AIModelSwitchConfig;
+
   @property({ attribute: false })
   accessor appSidebarConfig!: AppSidebarConfig;
 
@@ -209,6 +224,9 @@ export class ChatPanel extends SignalWatcher(
   @property({ attribute: false })
   accessor affineFeatureFlagService!: FeatureFlagService;
 
+  @property({ attribute: false })
+  accessor parentSessionId: string | undefined = undefined;
+
   @state()
   accessor isLoading = false;
 
@@ -218,10 +236,10 @@ export class ChatPanel extends SignalWatcher(
   @state()
   accessor embeddingProgress: [number, number] = [0, 0];
 
-  private _isInitialized = false;
+  @state()
+  accessor session: CopilotSessionType | undefined = undefined;
 
-  // always use getSessionId to get the sessionId
-  private _sessionId: string | undefined = undefined;
+  private _isInitialized = false;
 
   private _isSidebarOpen: Signal<boolean | undefined> = signal(false);
 
@@ -252,7 +270,7 @@ export class ChatPanel extends SignalWatcher(
   };
 
   private readonly _resetPanel = () => {
-    this._sessionId = undefined;
+    this.session = undefined;
     this.chatContextValue = DEFAULT_CHAT_CONTEXT_VALUE;
     this.isLoading = false;
     this._isInitialized = false;
@@ -408,6 +426,7 @@ export class ChatPanel extends SignalWatcher(
       <ai-chat-composer
         .host=${this.host}
         .doc=${this.doc}
+        .session=${this.session}
         .getSessionId=${this._getSessionId}
         .createSessionId=${this._createSessionId}
         .chatContextValue=${this.chatContextValue}
@@ -416,6 +435,7 @@ export class ChatPanel extends SignalWatcher(
         .isVisible=${this._isSidebarOpen}
         .networkSearchConfig=${this.networkSearchConfig}
         .reasoningConfig=${this.reasoningConfig}
+        .modelSwitchConfig=${this.modelSwitchConfig}
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
         .trackOptions=${{

packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts

@@ -3,6 +3,7 @@ import type {
   CopilotContextDoc,
   CopilotContextFile,
   CopilotDocType,
+  CopilotSessionType,
 } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
@@ -26,6 +27,7 @@ import type {
 import { isCollectionChip, isDocChip, isTagChip } from '../ai-chat-chips';
 import type {
   AIChatInputContext,
+  AIModelSwitchConfig,
   AINetworkSearchConfig,
   AIReasoningConfig,
 } from '../ai-chat-input';
@@ -53,6 +55,9 @@ export class AIChatComposer extends SignalWatcher(
   @property({ attribute: false })
   accessor doc!: Store;
 
+  @property({ attribute: false })
+  accessor session!: CopilotSessionType | undefined;
+
   @property({ attribute: false })
   accessor getSessionId!: () => Promise<string | undefined>;
 
@@ -85,6 +90,9 @@ export class AIChatComposer extends SignalWatcher(
   @property({ attribute: false })
   accessor searchMenuConfig!: SearchMenuConfig;
 
+  @property({ attribute: false })
+  accessor modelSwitchConfig!: AIModelSwitchConfig;
+
   @property({ attribute: false })
   accessor onChatSuccess: (() => void) | undefined;
 
@@ -124,6 +132,7 @@ export class AIChatComposer extends SignalWatcher(
       <ai-chat-input
         .host=${this.host}
         .chips=${this.chips}
+        .session=${this.session}
         .getSessionId=${this.getSessionId}
         .createSessionId=${this.createSessionId}
         .getContextId=${this._getContextId}
@@ -131,6 +140,7 @@ export class AIChatComposer extends SignalWatcher(
         .updateContext=${this.updateContext}
         .networkSearchConfig=${this.networkSearchConfig}
         .reasoningConfig=${this.reasoningConfig}
+        .modelSwitchConfig=${this.modelSwitchConfig}
         .docDisplayConfig=${this.docDisplayConfig}
         .onChatSuccess=${this.onChatSuccess}
         .trackOptions=${this.trackOptions}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts

@@ -1,4 +1,5 @@
 import { stopPropagation } from '@affine/core/utils';
+import type { CopilotSessionType } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
 import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
@@ -34,6 +35,7 @@ import type { ChatMessage } from '../ai-chat-messages';
 import { MAX_IMAGE_COUNT } from './const';
 import type {
   AIChatInputContext,
+  AIModelSwitchConfig,
   AINetworkSearchConfig,
   AIReasoningConfig,
 } from './type';
@@ -238,6 +240,9 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
   @property({ attribute: false })
   accessor host!: EditorHost;
 
+  @property({ attribute: false })
+  accessor session!: CopilotSessionType | undefined;
+
   @query('image-preview-grid')
   accessor imagePreviewGrid: HTMLDivElement | null = null;
 
@@ -250,6 +255,9 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
   @state()
   accessor focused = false;
 
+  @state()
+  accessor modelId: string | undefined = undefined;
+
   @property({ attribute: false })
   accessor chatContextValue!: AIChatInputContext;
 
@@ -274,6 +282,9 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
   @property({ attribute: false })
   accessor reasoningConfig!: AIReasoningConfig;
 
+  @property({ attribute: false })
+  accessor modelSwitchConfig!: AIModelSwitchConfig;
+
   @property({ attribute: false })
   accessor docDisplayConfig!: DocDisplayConfig;
 
@@ -392,6 +403,16 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
           ${ImageIcon()}
           <affine-tooltip>Upload</affine-tooltip>
         </div>
+        ${this.modelSwitchConfig.visible.value
+          ? html`
+              <ai-chat-models
+                class="chat-input-icon"
+                .modelId=${this.modelId}
+                .session=${this.session}
+                .onModelChange=${this._handleModelChange}
+              ></ai-chat-models>
+            `
+          : nothing}
         ${this.networkSearchConfig.visible.value
           ? html`
               <div
@@ -536,6 +557,10 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
     await this.send(value);
   };
 
+  private readonly _handleModelChange = (modelId: string) => {
+    this.modelId = modelId;
+  };
+
   send = async (text: string) => {
     try {
       const { status, markdown, images } = this.chatContextValue;
@@ -582,6 +607,7 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
         control: this.trackOptions.control,
         webSearch: this._isNetworkActive,
         reasoning: this._isReasoningActive,
+        modelId: this.modelId,
       });
 
       for await (const text of stream) {

packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/type.ts

@@ -14,6 +14,10 @@ export interface AIReasoningConfig {
   setEnabled: (state: boolean) => void;
 }
 
+export interface AIModelSwitchConfig {
+  visible: Signal<boolean | undefined>;
+}
+
 // TODO: remove this type
 export type AIChatInputContext = {
   messages: HistoryMessage[];

packages/frontend/core/src/blocksuite/ai/components/ai-chat-models/ai-chat-models.ts

@@ -0,0 +1,112 @@
+import type { CopilotSessionType } from '@affine/graphql';
+import { createLitPortal } from '@blocksuite/affine/components/portal';
+import { WithDisposable } from '@blocksuite/affine/global/lit';
+import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
+import { ShadowlessElement } from '@blocksuite/affine/std';
+import { flip, offset } from '@floating-ui/dom';
+import { css, html, nothing } from 'lit';
+import { property, query } from 'lit/decorators.js';
+import { repeat } from 'lit/directives/repeat.js';
+
+export class AIChatModels extends WithDisposable(ShadowlessElement) {
+  @property()
+  accessor modelId: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor onModelChange: ((modelId: string) => void) | undefined;
+
+  @property({ attribute: false })
+  accessor session!: CopilotSessionType | undefined;
+
+  @query('.ai-chat-models')
+  accessor modelsButton!: HTMLDivElement;
+
+  private _abortController: AbortController | null = null;
+
+  static override styles = css`
+    ai-chat-models {
+      font-size: 13px;
+      cursor: pointer;
+    }
+  `;
+
+  private readonly _onItemClick = (modelId: string) => {
+    this.onModelChange?.(modelId);
+    this._abortController?.abort();
+    this._abortController = null;
+  };
+
+  private readonly _toggleSwitchModelMenu = () => {
+    if (this._abortController) {
+      this._abortController.abort();
+      return;
+    }
+
+    this._abortController = new AbortController();
+    this._abortController.signal.addEventListener('abort', () => {
+      this._abortController = null;
+    });
+
+    createLitPortal({
+      template: html` <style>
+          .ai-model-list {
+            border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
+            border-radius: 4px;
+            background: ${unsafeCSSVarV2('layer/background/overlayPanel')};
+            box-shadow: ${unsafeCSSVar('overlayPanelShadow')};
+            padding: 8px;
+          }
+          .ai-model-item {
+            font-size: 13px;
+            padding: 4px;
+            cursor: pointer;
+          }
+          .ai-model-item:hover {
+            background: var(--affine-hover-color);
+          }
+        </style>
+        <div class="ai-model-list">
+          ${repeat(
+            this.session?.optionalModels ?? [],
+            modelId => modelId,
+            modelId => {
+              return html`<div
+                class="ai-model-item"
+                @click=${() => this._onItemClick(modelId)}
+              >
+                ${modelId}
+              </div>`;
+            }
+          )}
+        </div>`,
+      portalStyles: {
+        zIndex: 'var(--affine-z-index-popover)',
+      },
+      container: document.body,
+      computePosition: {
+        referenceElement: this.modelsButton,
+        placement: 'top-start',
+        middleware: [offset({ crossAxis: -30, mainAxis: 8 }), flip()],
+        autoUpdate: { animationFrame: true },
+      },
+      abortController: this._abortController,
+      closeOnClickAway: true,
+    });
+  };
+
+  override render() {
+    if (!this.session) {
+      return nothing;
+    }
+
+    return html`
+      <div
+        class="ai-chat-models"
+        @click=${this._toggleSwitchModelMenu}
+        data-testid="ai-chat-models"
+      >
+        ${this.modelId || this.session.model}
+      </div>
+    `;
+  }
+}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-models/index.ts

@@ -0,0 +1 @@
+export * from './ai-chat-models';

packages/frontend/core/src/blocksuite/ai/effects.ts

@@ -38,6 +38,7 @@ import { ChatPanelFileChip } from './components/ai-chat-chips/file-chip';
 import { ChatPanelTagChip } from './components/ai-chat-chips/tag-chip';
 import { AIChatComposer } from './components/ai-chat-composer';
 import { AIChatInput } from './components/ai-chat-input';
+import { AIChatModels } from './components/ai-chat-models/ai-chat-models';
 import { AIHistoryClear } from './components/ai-history-clear';
 import { effects as componentAiItemEffects } from './components/ai-item';
 import { AIScrollableTextRenderer } from './components/ai-scrollable-text-renderer';
@@ -110,6 +111,7 @@ export function registerAIEffects() {
   customElements.define('chat-panel-tag-chip', ChatPanelTagChip);
   customElements.define('chat-panel-collection-chip', ChatPanelCollectionChip);
   customElements.define('chat-panel-chip', ChatPanelChip);
+  customElements.define('ai-chat-models', AIChatModels);
   customElements.define('ai-error-wrapper', AIErrorWrapper);
   customElements.define('ai-slides-renderer', AISlidesRenderer);
   customElements.define('ai-answer-wrapper', AIAnswerWrapper);

packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts

@@ -1,4 +1,5 @@
 // packages/frontend/core/src/blocksuite/ai/hooks/useChatPanelConfig.ts
+import { AIModelSwitchService } from '@affine/core/modules/ai-button/services/model-switch';
 import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import { AIReasoningService } from '@affine/core/modules/ai-button/services/reasoning';
 import { CollectionService } from '@affine/core/modules/collection';
@@ -21,6 +22,7 @@ export function useAIChatConfig() {
 
   const searchService = framework.get(AINetworkSearchService);
   const reasoningService = framework.get(AIReasoningService);
+  const modelSwitchService = framework.get(AIModelSwitchService);
   const docDisplayMetaService = framework.get(DocDisplayMetaService);
   const workspaceService = framework.get(WorkspaceService);
   const searchMenuService = framework.get(SearchMenuService);
@@ -40,6 +42,10 @@ export function useAIChatConfig() {
     setEnabled: reasoningService.setEnabled,
   };
 
+  const modelSwitchConfig = {
+    visible: modelSwitchService.visible,
+  };
+
   const docDisplayConfig = {
     getIcon: (docId: string) => {
       return docDisplayMetaService.icon$(docId, { type: 'lit' }).value;
@@ -124,5 +130,6 @@ export function useAIChatConfig() {
     reasoningConfig,
     docDisplayConfig,
     searchMenuConfig,
+    modelSwitchConfig,
   };
 }

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

@@ -47,6 +47,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
     searchMenuConfig,
     networkSearchConfig,
     reasoningConfig,
+    modelSwitchConfig,
   } = useAIChatConfig();
 
   useEffect(() => {
@@ -73,6 +74,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
       chatPanelRef.current.searchMenuConfig = searchMenuConfig;
       chatPanelRef.current.networkSearchConfig = networkSearchConfig;
       chatPanelRef.current.reasoningConfig = reasoningConfig;
+      chatPanelRef.current.modelSwitchConfig = modelSwitchConfig;
       chatPanelRef.current.extensions = editor.host.std
         .get(ViewExtensionManagerIdentifier)
         .get('preview-page');
@@ -107,6 +109,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
     networkSearchConfig,
     searchMenuConfig,
     reasoningConfig,
+    modelSwitchConfig,
   ]);
 
   return <div className={styles.root} ref={containerRef} />;

packages/frontend/core/src/modules/ai-button/index.ts

@@ -7,6 +7,7 @@ import { FeatureFlagService } from '../feature-flag';
 import { GlobalStateService } from '../storage';
 import { AIButtonProvider } from './provider/ai-button';
 import { AIButtonService } from './services/ai-button';
+import { AIModelSwitchService } from './services/model-switch';
 import { AINetworkSearchService } from './services/network-search';
 import { AIReasoningService } from './services/reasoning';
 
@@ -26,3 +27,7 @@ export function configureAINetworkSearchModule(framework: Framework) {
 export function configureAIReasoningModule(framework: Framework) {
   framework.service(AIReasoningService, [GlobalStateService]);
 }
+
+export function configureAIModelSwitchModule(framework: Framework) {
+  framework.service(AIModelSwitchService, [FeatureFlagService]);
+}

packages/frontend/core/src/modules/ai-button/services/model-switch.ts

@@ -0,0 +1,26 @@
+import {
+  createSignalFromObservable,
+  type Signal,
+} from '@blocksuite/affine/shared/utils';
+import { Service } from '@toeverything/infra';
+
+import type { FeatureFlagService } from '../../feature-flag';
+
+export class AIModelSwitchService extends Service {
+  constructor(private readonly featureFlagService: FeatureFlagService) {
+    super();
+
+    const { signal: visible, cleanup: visibleCleanup } =
+      createSignalFromObservable<boolean | undefined>(
+        this._visible$,
+        undefined
+      );
+    this.visible = visible;
+    this.disposables.push(visibleCleanup);
+  }
+
+  visible: Signal<boolean | undefined>;
+
+  private readonly _visible$ =
+    this.featureFlagService.flags.enable_ai_model_switch.$;
+}

packages/frontend/core/src/modules/feature-flag/constant.ts

@@ -26,6 +26,15 @@ export const AFFINE_FLAGS = {
     configurable: false,
     defaultState: true,
   },
+  enable_ai_model_switch: {
+    category: 'affine',
+    displayName:
+      'com.affine.settings.workspace.experimental-features.enable-ai-model-switch.name',
+    description:
+      'com.affine.settings.workspace.experimental-features.enable-ai-model-switch.description',
+    configurable: isCanaryBuild,
+    defaultState: isCanaryBuild,
+  },
   enable_edgeless_text: {
     category: 'blocksuite',
     bsFlag: 'enable_edgeless_text',

packages/frontend/core/src/modules/index.ts

@@ -3,6 +3,7 @@ import { type Framework } from '@toeverything/infra';
 
 import {
   configureAIButtonModule,
+  configureAIModelSwitchModule,
   configureAINetworkSearchModule,
   configureAIReasoningModule,
 } from './ai-button';
@@ -105,6 +106,7 @@ export function configureCommonModules(framework: Framework) {
   configureCommonGlobalStorageImpls(framework);
   configureAINetworkSearchModule(framework);
   configureAIReasoningModule(framework);
+  configureAIModelSwitchModule(framework);
   configureAIButtonModule(framework);
   configureTemplateDocModule(framework);
   configureBlobManagementModule(framework);