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);