From f7a094053e0e684c7c6a384529f38255449864eb Mon Sep 17 00:00:00 2001 From: Wu Yue Date: Wed, 30 Jul 2025 10:10:39 +0800 Subject: [PATCH] feat(core): add ai workspace all docs switch (#13345) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close [AI-397](https://linear.app/affine-design/issue/AI-397) 截屏2025-07-29 11 54 20 ## Summary by CodeRabbit * **New Features** * Introduced an AI tools configuration service, allowing users to customize AI tool usage (e.g., workspace search, reading docs) in chat and AI features. * Added a toggle in chat preferences for enabling or disabling workspace-wide document search. * AI chat components now respect user-configured tool settings across chat, retry, and playground scenarios. * **Improvements** * Enhanced chat and AI interfaces to propagate and honor user tool configuration throughout the frontend and backend. * Made draft and tool configuration services optional and safely handled their absence in chat components. --- .../server/src/plugins/copilot/controller.ts | 14 ++++-- .../src/plugins/copilot/providers/types.ts | 44 ++++++++-------- .../server/src/plugins/copilot/types.ts | 29 ++++++++++- .../server/src/plugins/copilot/utils.ts | 33 +++++++++++- .../core/src/blocksuite/ai/actions/types.ts | 2 + .../src/blocksuite/ai/chat-panel/ai-title.ts | 5 ++ .../src/blocksuite/ai/chat-panel/index.ts | 10 +++- .../ai-chat-composer/ai-chat-composer.ts | 11 +++- .../ai-chat-content/ai-chat-content.ts | 35 +++++++++---- .../components/ai-chat-input/ai-chat-input.ts | 44 ++++++++++------ .../ai-chat-input/preference-popup.ts | 18 +++++++ .../ai-chat-messages/ai-chat-messages.ts | 5 ++ .../ai/components/playground/chat.ts | 6 +++ .../ai/components/playground/content.ts | 5 ++ .../ai/peek-view/chat-block-peek-view.ts | 18 ++++++- .../blocksuite/ai/provider/copilot-client.ts | 13 ++++- .../src/blocksuite/ai/provider/request.ts | 4 ++ .../desktop/pages/workspace/chat/index.tsx | 6 ++- .../pages/workspace/detail-page/tabs/chat.tsx | 7 ++- .../core/src/modules/ai-button/index.ts | 9 ++++ .../ai-button/services/tools-config.ts | 50 +++++++++++++++++++ packages/frontend/core/src/modules/index.ts | 2 + .../view/ai-chat-block-peek-view/index.tsx | 12 ++++- 23 files changed, 321 insertions(+), 61 deletions(-) create mode 100644 packages/frontend/core/src/modules/ai-button/services/tools-config.ts diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts index 8bd66878f8..e35a0724d4 100644 --- a/packages/backend/server/src/plugins/copilot/controller.ts +++ b/packages/backend/server/src/plugins/copilot/controller.ts @@ -56,7 +56,7 @@ import { StreamObjectParser } from './providers/utils'; import { ChatSession, ChatSessionService } from './session'; import { CopilotStorage } from './storage'; import { ChatMessage, ChatQuerySchema } from './types'; -import { getSignal } from './utils'; +import { getSignal, getTools } from './utils'; import { CopilotWorkflowService, GraphExecutorState } from './workflow'; export interface ChatEvent { @@ -244,7 +244,8 @@ export class CopilotController implements BeforeApplicationShutdown { info.finalMessage = finalMessage.filter(m => m.role !== 'system'); metrics.ai.counter('chat_calls').add(1, { model }); - const { reasoning, webSearch } = ChatQuerySchema.parse(query); + const { reasoning, webSearch, toolsConfig } = + ChatQuerySchema.parse(query); const content = await provider.text({ modelId: model }, finalMessage, { ...session.config.promptConfig, signal: getSignal(req).signal, @@ -253,6 +254,7 @@ export class CopilotController implements BeforeApplicationShutdown { workspace: session.config.workspaceId, reasoning, webSearch, + tools: getTools(session.config.promptConfig?.tools, toolsConfig), }); session.push({ @@ -306,7 +308,8 @@ export class CopilotController implements BeforeApplicationShutdown { } }); - const { messageId, reasoning, webSearch } = ChatQuerySchema.parse(query); + const { messageId, reasoning, webSearch, toolsConfig } = + ChatQuerySchema.parse(query); const source$ = from( provider.streamText({ modelId: model }, finalMessage, { @@ -317,6 +320,7 @@ export class CopilotController implements BeforeApplicationShutdown { workspace: session.config.workspaceId, reasoning, webSearch, + tools: getTools(session.config.promptConfig?.tools, toolsConfig), }) ).pipe( connect(shared$ => @@ -398,7 +402,8 @@ export class CopilotController implements BeforeApplicationShutdown { } }); - const { messageId, reasoning, webSearch } = ChatQuerySchema.parse(query); + const { messageId, reasoning, webSearch, toolsConfig } = + ChatQuerySchema.parse(query); const source$ = from( provider.streamObject({ modelId: model }, finalMessage, { @@ -409,6 +414,7 @@ export class CopilotController implements BeforeApplicationShutdown { workspace: session.config.workspaceId, reasoning, webSearch, + tools: getTools(session.config.promptConfig?.tools, toolsConfig), }) ).pipe( connect(shared$ => diff --git a/packages/backend/server/src/plugins/copilot/providers/types.ts b/packages/backend/server/src/plugins/copilot/providers/types.ts index b480c1b30f..34a5bd9fba 100644 --- a/packages/backend/server/src/plugins/copilot/providers/types.ts +++ b/packages/backend/server/src/plugins/copilot/providers/types.ts @@ -57,28 +57,28 @@ export const VertexSchema: JSONSchema = { // ========== prompt ========== +export const PromptToolsSchema = z + .enum([ + 'codeArtifact', + 'conversationSummary', + // work with morph + 'docEdit', + // work with indexer + 'docRead', + 'docKeywordSearch', + // work with embeddings + 'docSemanticSearch', + // work with exa/model internal tools + 'webSearch', + // artifact tools + 'docCompose', + // section editing + 'sectionEdit', + ]) + .array(); + export const PromptConfigStrictSchema = z.object({ - tools: z - .enum([ - 'codeArtifact', - 'conversationSummary', - // work with morph - 'docEdit', - // work with indexer - 'docRead', - 'docKeywordSearch', - // work with embeddings - 'docSemanticSearch', - // work with exa/model internal tools - 'webSearch', - // artifact tools - 'docCompose', - // section editing - 'sectionEdit', - ]) - .array() - .nullable() - .optional(), + tools: PromptToolsSchema.nullable().optional(), // params requirements requireContent: z.boolean().nullable().optional(), requireAttachment: z.boolean().nullable().optional(), @@ -107,6 +107,8 @@ export const PromptConfigSchema = export type PromptConfig = z.infer; +export type PromptTools = z.infer; + // ========== message ========== export const EmbeddingMessage = z.array(z.string().trim().min(1)).min(1); diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index dc8db5e18b..af6da6c7e4 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -16,6 +16,23 @@ const zMaybeString = z.preprocess(val => { return s === '' || s == null ? undefined : s; }, z.string().min(1).optional()); +const ToolsConfigSchema = z.preprocess( + val => { + // if val is a string, try to parse it as JSON + if (typeof val === 'string') { + try { + return JSON.parse(val); + } catch { + return {}; + } + } + return val || {}; + }, + z.record(z.enum(['searchWorkspace', 'readingDocs']), z.boolean()).default({}) +); + +export type ToolsConfig = z.infer; + export const ChatQuerySchema = z .object({ messageId: zMaybeString, @@ -23,15 +40,25 @@ export const ChatQuerySchema = z retry: zBool, reasoning: zBool, webSearch: zBool, + toolsConfig: ToolsConfigSchema, }) .catchall(z.string()) .transform( - ({ messageId, modelId, retry, reasoning, webSearch, ...params }) => ({ + ({ messageId, modelId, retry, reasoning, webSearch, + toolsConfig, + ...params + }) => ({ + messageId, + modelId, + retry, + reasoning, + webSearch, + toolsConfig, params, }) ); diff --git a/packages/backend/server/src/plugins/copilot/utils.ts b/packages/backend/server/src/plugins/copilot/utils.ts index caa1c4ead1..e22f4b7584 100644 --- a/packages/backend/server/src/plugins/copilot/utils.ts +++ b/packages/backend/server/src/plugins/copilot/utils.ts @@ -3,7 +3,8 @@ import { Readable } from 'node:stream'; import type { Request } from 'express'; import { readBufferWithLimit } from '../../base'; -import { MAX_EMBEDDABLE_SIZE } from './types'; +import { PromptTools } from './providers'; +import { MAX_EMBEDDABLE_SIZE, ToolsConfig } from './types'; export function readStream( readable: Readable, @@ -49,3 +50,33 @@ export function getSignal(req: Request): SignalReturnType { onConnectionClosed: cb => (callback = cb), }; } + +export function getTools( + tools?: PromptTools | null, + toolsConfig?: ToolsConfig +) { + if (!tools || !toolsConfig) { + return tools; + } + let result: PromptTools = tools; + (Object.keys(toolsConfig) as Array).forEach(key => { + const value = toolsConfig[key]; + switch (key) { + case 'searchWorkspace': + if (value === false) { + result = result.filter(tool => { + return tool !== 'docKeywordSearch' && tool !== 'docSemanticSearch'; + }); + } + break; + case 'readingDocs': + if (value === false) { + result = result.filter(tool => { + return tool !== 'docRead'; + }); + } + break; + } + }); + return result; +} diff --git a/packages/frontend/core/src/blocksuite/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/ai/actions/types.ts index 9571060d27..e3df4e69f6 100644 --- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts +++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts @@ -1,3 +1,4 @@ +import type { AIToolsConfig } from '@affine/core/modules/ai-button'; import type { AddContextFileInput, ContextMatchedDocChunk, @@ -142,6 +143,7 @@ declare global { webSearch?: boolean; reasoning?: boolean; modelId?: string; + toolsConfig?: AIToolsConfig | undefined; contexts?: { docs: AIDocContextOption[]; files: AIFileContextOption[]; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/ai-title.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/ai-title.ts index e01dc80312..0b9c6c289f 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/ai-title.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/ai-title.ts @@ -1,3 +1,4 @@ +import type { AIToolsConfigService } from '@affine/core/modules/ai-button'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { FeatureFlagService } from '@affine/core/modules/feature-flag'; import type { AppThemeService } from '@affine/core/modules/theme'; @@ -105,6 +106,9 @@ export class AIChatPanelTitle extends SignalWatcher( @property({ attribute: false }) accessor notificationService!: NotificationService; + @property({ attribute: false }) + accessor aiToolsConfigService!: AIToolsConfigService; + @property({ attribute: false }) accessor session!: CopilotChatHistoryFragment | null | undefined; @@ -142,6 +146,7 @@ export class AIChatPanelTitle extends SignalWatcher( .affineThemeService=${this.affineThemeService} .notificationService=${this.notificationService} .affineWorkspaceDialogService=${this.affineWorkspaceDialogService} + .aiToolsConfigService=${this.aiToolsConfigService} > `; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts index c9376de683..79fb0de9dd 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts @@ -1,4 +1,7 @@ -import type { AIDraftService } from '@affine/core/modules/ai-button'; +import type { + AIDraftService, + AIToolsConfigService, +} from '@affine/core/modules/ai-button'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { FeatureFlagService } from '@affine/core/modules/feature-flag'; import type { AppThemeService } from '@affine/core/modules/theme'; @@ -119,6 +122,9 @@ export class ChatPanel extends SignalWatcher( @property({ attribute: false }) accessor aiDraftService!: AIDraftService; + @property({ attribute: false }) + accessor aiToolsConfigService!: AIToolsConfigService; + @state() accessor session: CopilotChatHistoryFragment | null | undefined; @@ -387,6 +393,7 @@ export class ChatPanel extends SignalWatcher( .affineWorkspaceDialogService=${this.affineWorkspaceDialogService} .affineThemeService=${this.affineThemeService} .notificationService=${this.notificationService} + .aiToolsConfigService=${this.aiToolsConfigService} .session=${this.session} .status=${this.status} .embeddingProgress=${this.embeddingProgress} @@ -413,6 +420,7 @@ export class ChatPanel extends SignalWatcher( .affineThemeService=${this.affineThemeService} .notificationService=${this.notificationService} .aiDraftService=${this.aiDraftService} + .aiToolsConfigService=${this.aiToolsConfigService} .onEmbeddingProgressChange=${this.onEmbeddingProgressChange} .onContextChange=${this.onContextChange} .width=${this.sidebarWidth} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts index d2d521957a..820cf4cac3 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts @@ -1,6 +1,9 @@ import './ai-chat-composer-tip'; -import type { AIDraftService } from '@affine/core/modules/ai-button'; +import type { + AIDraftService, + AIToolsConfigService, +} from '@affine/core/modules/ai-button'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { ContextEmbedStatus, @@ -118,7 +121,10 @@ export class AIChatComposer extends SignalWatcher( accessor notificationService!: NotificationService; @property({ attribute: false }) - accessor aiDraftService!: AIDraftService; + accessor aiDraftService: AIDraftService | undefined; + + @property({ attribute: false }) + accessor aiToolsConfigService!: AIToolsConfigService; @state() accessor chips: ChatChip[] = []; @@ -166,6 +172,7 @@ export class AIChatComposer extends SignalWatcher( .docDisplayConfig=${this.docDisplayConfig} .searchMenuConfig=${this.searchMenuConfig} .aiDraftService=${this.aiDraftService} + .aiToolsConfigService=${this.aiToolsConfigService} .portalContainer=${this.portalContainer} .onChatSuccess=${this.onChatSuccess} .trackOptions=${this.trackOptions} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts index 79c110e97e..edb86fb46c 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts @@ -1,4 +1,7 @@ -import type { AIDraftService } from '@affine/core/modules/ai-button'; +import type { + AIDraftService, + AIToolsConfigService, +} from '@affine/core/modules/ai-button'; import type { AIDraftState } from '@affine/core/modules/ai-button/services/ai-draft'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { FeatureFlagService } from '@affine/core/modules/feature-flag'; @@ -153,7 +156,10 @@ export class AIChatContent extends SignalWatcher( accessor notificationService!: NotificationService; @property({ attribute: false }) - accessor aiDraftService!: AIDraftService; + accessor aiDraftService: AIDraftService | undefined; + + @property({ attribute: false }) + accessor aiToolsConfigService!: AIToolsConfigService; @property({ attribute: false }) accessor onEmbeddingProgressChange: @@ -273,6 +279,9 @@ export class AIChatContent extends SignalWatcher( }; private readonly updateDraft = async (context: Partial) => { + if (!this.aiDraftService) { + return; + } const draft: Partial = pick(context, [ 'quote', 'images', @@ -344,15 +353,17 @@ export class AIChatContent extends SignalWatcher( this.initChatContent().catch(console.error); - this.aiDraftService - .getDraft() - .then(draft => { - this.chatContextValue = { - ...this.chatContextValue, - ...draft, - }; - }) - .catch(console.error); + if (this.aiDraftService) { + this.aiDraftService + .getDraft() + .then(draft => { + this.chatContextValue = { + ...this.chatContextValue, + ...draft, + }; + }) + .catch(console.error); + } this._disposables.add( AIProvider.slots.actions.subscribe(({ event }) => { @@ -405,6 +416,7 @@ export class AIChatContent extends SignalWatcher( .affineFeatureFlagService=${this.affineFeatureFlagService} .affineThemeService=${this.affineThemeService} .notificationService=${this.notificationService} + .aiToolsConfigService=${this.aiToolsConfigService} .networkSearchConfig=${this.networkSearchConfig} .reasoningConfig=${this.reasoningConfig} .width=${this.width} @@ -434,6 +446,7 @@ export class AIChatContent extends SignalWatcher( .affineWorkspaceDialogService=${this.affineWorkspaceDialogService} .notificationService=${this.notificationService} .aiDraftService=${this.aiDraftService} + .aiToolsConfigService=${this.aiToolsConfigService} .trackOptions=${{ where: 'chat-panel', control: 'chat-send', diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts index fa9c9b3f2c..2093c0eecb 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts @@ -1,4 +1,7 @@ -import type { AIDraftService } from '@affine/core/modules/ai-button'; +import type { + AIDraftService, + AIToolsConfigService, +} from '@affine/core/modules/ai-button'; import type { CopilotChatHistoryFragment } from '@affine/graphql'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme'; @@ -353,7 +356,10 @@ export class AIChatInput extends SignalWatcher( accessor searchMenuConfig!: SearchMenuConfig; @property({ attribute: false }) - accessor aiDraftService!: AIDraftService; + accessor aiDraftService: AIDraftService | undefined; + + @property({ attribute: false }) + accessor aiToolsConfigService!: AIToolsConfigService; @property({ attribute: false }) accessor isRootSession: boolean = true; @@ -406,13 +412,15 @@ export class AIChatInput extends SignalWatcher( protected override firstUpdated(changedProperties: PropertyValues): void { super.firstUpdated(changedProperties); - this.aiDraftService - .getDraft() - .then(draft => { - this.textarea.value = draft.input; - this.isInputEmpty = !this.textarea.value.trim(); - }) - .catch(console.error); + if (this.aiDraftService) { + this.aiDraftService + .getDraft() + .then(draft => { + this.textarea.value = draft.input; + this.isInputEmpty = !this.textarea.value.trim(); + }) + .catch(console.error); + } } protected override render() { @@ -493,6 +501,7 @@ export class AIChatInput extends SignalWatcher( .networkSearchVisible=${!!this.networkSearchConfig.visible.value} .isNetworkActive=${this._isNetworkActive} .onNetworkActiveChange=${this._toggleNetworkSearch} + .toolsConfigService=${this.aiToolsConfigService} > ${status === 'transmitting' || status === 'loading' ? html`