Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-08-11 05:58:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(server): improve context management (#15448)

Browse Source 
#### PR Dependency Tree


* **PR #15448** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

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

* **New Features**
* Added workspace artifact upload, browsing, removal, deduplication, and
library ownership support.
* Copilot now supports scoped document and artifact search, canvas
reading, live editor context, and frontend tools.
* Added scope and focus selectors with source-resolution receipts in
chat.
* Added embedding health, progress, synchronization, and retrieval
capabilities.
* Added BYOK policy visibility, provider restrictions, endpoint dialect
selection, and validation.
* Added delegated editor interactions and userdata document
authorization.

* **Bug Fixes**
* Improved attachment handling, cancellation, access control, retrieval
fallbacks, workspace synchronization, and configuration validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-08-10 09:27:58 +08:00
committed by GitHub
parent 42322d13fe
commit ee899a267b
311 changed files with 20468 additions and 14806 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/frontend/core/src/blocksuite/ai/actions/types.ts
+10 -98
View File
@@ -1,14 +1,6 @@
import type { AIToolsConfig } from '@affine/core/modules/ai-button';
import type {
AddContextFileInput,
ContextMatchedDocChunk,
ContextMatchedFileChunk,
ContextWorkspaceEmbeddingStatus,
CopilotChatHistoryFragment,
CopilotContextBlob,
CopilotContextCategory,
CopilotContextDoc,
CopilotContextFile,
CopilotHistories,
getCopilotHistoriesQuery,
QueryChatHistoriesInput,
@@ -126,7 +118,6 @@ declare global {
interface AIDocContextOption {
docId: string;
docTitle: string;
docContent: string;
tags: string;
createDate: string;
updatedDate: string;
@@ -152,6 +143,16 @@ declare global {
selectedMarkdown?: string;
html?: string;
};
scopeSelectors?: Array<{
kind: 'document' | 'tag' | 'collection' | 'favorite';
id: string;
name?: string;
}>;
focusSelectors?: Array<{
kind: 'document' | 'tag' | 'collection' | 'favorite';
id: string;
name?: string;
}>;
}
interface TranslateOptions extends AITextActionOptions {
@@ -275,95 +276,6 @@ declare global {
): Promise<AIActionTextResponse<T>>;
}
type AIDocsAndFilesContext = {
docs: CopilotContextDoc[];
files: CopilotContextFile[];
tags: CopilotContextCategory[];
collections: CopilotContextCategory[];
blobs: CopilotContextBlob[];
};
interface AIContextService {
createContext: (
workspaceId: string,
sessionId: string
) => Promise<string>;
getContextId: (
workspaceId: string,
sessionId: string
) => Promise<string | undefined>;
addContextDoc: (options: {
contextId: string;
docId: string;
}) => Promise<CopilotContextDoc>;
removeContextDoc: (options: {
contextId: string;
docId: string;
}) => Promise<boolean>;
addContextFile: (
file: File,
options: AddContextFileInput
) => Promise<CopilotContextFile>;
removeContextFile: (options: {
contextId: string;
fileId: string;
}) => Promise<boolean>;
addContextTag: (options: {
contextId: string;
tagId: string;
docIds: string[];
}) => Promise<CopilotContextCategory>;
removeContextTag: (options: {
contextId: string;
tagId: string;
}) => Promise<boolean>;
addContextCollection: (options: {
contextId: string;
collectionId: string;
docIds: string[];
}) => Promise<CopilotContextCategory>;
removeContextCollection: (options: {
contextId: string;
collectionId: string;
}) => Promise<boolean>;
getContextDocsAndFiles: (
workspaceId: string,
sessionId: string,
contextId: string
) => Promise<AIDocsAndFilesContext | undefined>;
pollContextDocsAndFiles: (
workspaceId: string,
sessionId: string,
contextId: string,
onPoll: (result: AIDocsAndFilesContext | undefined) => void,
abortSignal: AbortSignal
) => Promise<void>;
pollEmbeddingStatus: (
workspaceId: string,
onPoll: (result: ContextWorkspaceEmbeddingStatus) => void,
abortSignal: AbortSignal
) => Promise<void>;
matchContext: (
content: string,
contextId?: string,
workspaceId?: string,
limit?: number,
scopedThreshold?: number,
threshold?: number
) => Promise<{
files?: ContextMatchedFileChunk[];
docs?: ContextMatchedDocChunk[];
}>;
addContextBlob: (options: {
blobId: string;
contextId: string;
}) => Promise<CopilotContextBlob>;
removeContextBlob: (options: {
blobId: string;
contextId: string;
}) => Promise<boolean>;
}
// TODO(@Peng): should be refactored to get rid of implement details (like messages, action, role, etc.)
interface AIHistory {
sessionId: string;
packages/frontend/core/src/blocksuite/ai/chat-panel/message/assistant.ts
+4 -2
View File
@@ -107,7 +107,7 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
'content' in this.item &&
this.item.content &&
this.item.content.includes('[^') &&
/\[\^\d+\]:{"type":"doc","docId":"[^"]+"}/.test(this.item.content);
/\[\^[^\]]+\]:{"type":"doc","docId":"[^"]+"}/.test(this.item.content);
return html`<div class="user-info">
<chat-assistant-avatar .status=${this.status}></chat-assistant-avatar>
@@ -127,7 +127,9 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
${streamObjects?.length
? this.renderStreamObjects(streamObjects)
: this.renderRichText(content)}
${shouldRenderError ? AIChatErrorRenderer(error, host) : nothing}
${shouldRenderError
? AIChatErrorRenderer(error, host, () => this.retry())
: nothing}
${this.renderEditorActions()}
`;
}
packages/frontend/core/src/blocksuite/ai/chat-panel/message/user.ts
+27
View File
@@ -1,3 +1,4 @@
import { I18n } from '@affine/i18n';
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { ShadowlessElement } from '@blocksuite/affine/std';
import { css, html, nothing } from 'lit';
@@ -31,6 +32,14 @@ export class ChatMessageUser extends WithDisposable(ShadowlessElement) {
.text-content-wrapper {
align-self: flex-end;
}
.scope-receipt {
align-self: flex-end;
margin-top: 6px;
color: var(--affine-text-secondary-color);
font-size: 11px;
text-align: right;
}
`;
@property({ attribute: false })
@@ -41,6 +50,15 @@ export class ChatMessageUser extends WithDisposable(ShadowlessElement) {
renderContent() {
const { item } = this;
const receipt = item.scopeSnapshot;
const showReceipt = receipt && receipt.selectors.length > 0;
const resolvedCount = receipt
? receipt.requiredDocIds.length + receipt.requiredArtifactIds.length
: 0;
const selectorNames = receipt?.selectors.map(selector => {
if (selector.name) return selector.name;
return I18n[`com.affine.ai.chat-panel.scope.${selector.kind}`]();
});
return html`
${item.attachments
@@ -56,6 +74,15 @@ export class ChatMessageUser extends WithDisposable(ShadowlessElement) {
>
<chat-content-pure-text .text=${item.content}></chat-content-pure-text>
</div>
${showReceipt
? html`<div class="scope-receipt" data-testid="chat-scope-receipt">
${selectorNames?.join(', ')} ·
${I18n['com.affine.ai.chat-panel.scope.sources']({
count: String(resolvedCount),
})}
· ${new Date(receipt.resolvedAt).toLocaleString()}
</div>`
: nothing}
`;
}
packages/frontend/core/src/blocksuite/ai/components/ai-chat-add-context/ai-chat-add-context.ts
+18
View File
@@ -1,3 +1,4 @@
import { I18n } from '@affine/i18n';
import { createLitPortal } from '@blocksuite/affine/components/portal';
import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
import { ShadowlessElement } from '@blocksuite/affine/std';
@@ -20,6 +21,10 @@ export class AIChatAddContext extends SignalWatcher(
align-items: center;
justify-content: center;
cursor: pointer;
&[aria-disabled='true'] {
cursor: not-allowed;
}
}
`;
@@ -50,18 +55,31 @@ export class AIChatAddContext extends SignalWatcher(
private abortController: AbortController | null = null;
override render() {
const disabled = !this.searchMenuConfig.addContextAvailable;
return html`
<div
class="ai-chat-add-context"
aria-disabled=${disabled}
data-testid="chat-panel-with-button"
@click=${this.toggleAddDocMenu}
>
${PlusIcon()}
${disabled
? html`<affine-tooltip>
${I18n[
'com.affine.ai.chat-panel.local-workspace-context-unavailable'
]()}
</affine-tooltip>`
: null}
</div>
`;
}
private readonly toggleAddDocMenu = () => {
if (!this.searchMenuConfig.addContextAvailable) {
return;
}
if (this.abortController) {
this.abortController.abort();
return;
packages/frontend/core/src/blocksuite/ai/components/ai-chat-add-context/type.ts
+1
View File
@@ -6,6 +6,7 @@ import type {
import type { LinkedMenuGroup } from '@blocksuite/affine/widgets/linked-doc';
export interface SearchMenuConfig {
addContextAvailable: boolean;
getDocMenuGroup: (
query: string,
action: SearchDocMenuAction,
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/add-popover.ts
+3 -3
View File
@@ -487,7 +487,7 @@ export class ChatPanelAddPopover extends SignalWatcher(
this.abortController.abort();
await this.addChip({
docId: meta.id,
state: 'processing',
state: 'finished',
});
const mode = this.docDisplayConfig.getDocPrimaryMode(meta.id);
const method = meta.id === this.docId ? 'cur-doc' : 'doc';
@@ -498,7 +498,7 @@ export class ChatPanelAddPopover extends SignalWatcher(
this.abortController.abort();
await this.addChip({
tagId: tag.id,
state: 'processing',
state: 'finished',
});
this._track('tags');
};
@@ -507,7 +507,7 @@ export class ChatPanelAddPopover extends SignalWatcher(
this.abortController.abort();
await this.addChip({
collectionId: collection.id,
state: 'processing',
state: 'finished',
});
this._track('collections');
};
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/attachment-utils.ts
+1 -1
View File
@@ -29,7 +29,7 @@ export async function addFilesToChat(
}
await addChip({
file,
state: 'processing',
state: 'finished',
});
})
);
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/candidates-popover.ts
+1 -1
View File
@@ -105,7 +105,7 @@ export class ChatPanelCandidatesPopover extends SignalWatcher(
private readonly _addDocChip = (docId: string) => {
this.addChip({
docId,
state: 'processing',
state: 'finished',
});
};
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
-39
View File
@@ -13,7 +13,6 @@ import { isEqual } from 'lodash-es';
import type { ChatChip, DocChip, DocDisplayConfig, FileChip } from './type';
import {
estimateTokenCount,
getChipKey,
isAttachmentChip,
isCollectionChip,
@@ -23,9 +22,6 @@ import {
isTagChip,
} from './utils';
// 100k tokens limit for the docs context
const MAX_TOKEN_COUNT = 100000;
const MAX_CANDIDATES = 3;
export class ChatPanelChips extends SignalWatcher(
@@ -149,9 +145,7 @@ export class ChatPanelChips extends SignalWatcher(
.chip=${chip}
.independentMode=${this.independentMode}
.addChip=${this.addChip}
.updateChip=${this.updateChip}
.removeChip=${this.removeChip}
.checkTokenLimit=${this._checkTokenLimit}
.docDisplayConfig=${this.docDisplayConfig}
></chat-panel-doc-chip>`;
}
@@ -277,39 +271,6 @@ export class ChatPanelChips extends SignalWatcher(
});
};
private readonly _checkTokenLimit = (
newChip: DocChip,
newTokenCount: number
) => {
const estimatedTokens = this.chips.reduce((acc, chip) => {
if (isFileChip(chip) || isTagChip(chip) || isCollectionChip(chip)) {
return acc;
}
if (isDocChip(chip) && chip.docId === newChip.docId) {
return acc + newTokenCount;
}
if (
isDocChip(chip) &&
chip.markdown?.value &&
chip.state === 'finished'
) {
const tokenCount =
chip.tokenCount ?? estimateTokenCount(chip.markdown.value);
return acc + tokenCount;
}
if (isSelectedContextChip(chip)) {
const tokenCount =
estimateTokenCount(chip.combinedElementsMarkdown ?? '') +
estimateTokenCount(chip.snapshot ?? '') +
estimateTokenCount(chip.html ?? '');
return acc + tokenCount;
}
return acc;
}, 0);
return estimatedTokens <= MAX_TOKEN_COUNT;
};
private readonly _updateReferenceDocs = () => {
const docIds = this.chips
.filter(isDocChip)
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/doc-chip.ts
+3 -75
View File
@@ -2,15 +2,11 @@ import track from '@affine/track';
import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
import { ShadowlessElement } from '@blocksuite/affine/std';
import { Signal } from '@preact/signals-core';
import { html, type PropertyValues } from 'lit';
import { html } from 'lit';
import { property } from 'lit/decorators.js';
import throttle from 'lodash-es/throttle';
import { extractMarkdownFromDoc } from '../../utils/extract';
import type { DocChip, DocDisplayConfig } from './type';
import { estimateTokenCount, getChipIcon, getChipTooltip } from './utils';
const EXTRACT_DOC_THROTTLE = 1000;
import { getChipIcon, getChipTooltip } from './utils';
export class ChatPanelDocChip extends SignalWatcher(
WithDisposable(ShadowlessElement)
@@ -24,18 +20,9 @@ export class ChatPanelDocChip extends SignalWatcher(
@property({ attribute: false })
accessor addChip!: (chip: DocChip) => void;
@property({ attribute: false })
accessor updateChip!: (chip: DocChip, options: Partial<DocChip>) => void;
@property({ attribute: false })
accessor removeChip!: (chip: DocChip) => void;
@property({ attribute: false })
accessor checkTokenLimit!: (
newChip: DocChip,
newTokenCount: number
) => boolean;
@property({ attribute: false })
accessor docDisplayConfig!: DocDisplayConfig;
@@ -49,27 +36,6 @@ export class ChatPanelDocChip extends SignalWatcher(
);
this.chipName = signal;
this.disposables.add(cleanup);
const doc = this.docDisplayConfig.getDoc(this.chip.docId);
if (doc) {
this.disposables.add(
doc.slots.blockUpdated.subscribe(
throttle(this.autoUpdateChip, EXTRACT_DOC_THROTTLE)
)
);
this.autoUpdateChip();
}
}
override updated(changedProperties: PropertyValues): void {
super.updated(changedProperties);
if (
changedProperties.has('chip') &&
this.chip.state === 'processing' &&
!this.chip.markdown
) {
this.processDocChip().catch(console.error);
}
}
override disconnectedCallback() {
@@ -81,7 +47,7 @@ export class ChatPanelDocChip extends SignalWatcher(
if (this.chip.state === 'candidate') {
this.addChip({
...this.chip,
state: 'processing',
state: 'finished',
});
const mode = this.docDisplayConfig.getDocPrimaryMode(this.chip.docId);
const page = this.independentMode
@@ -99,44 +65,6 @@ export class ChatPanelDocChip extends SignalWatcher(
this.removeChip(this.chip);
};
private readonly autoUpdateChip = () => {
if (this.chip.state !== 'candidate') {
this.processDocChip().catch(console.error);
}
};
private readonly processDocChip = async () => {
try {
const doc = this.docDisplayConfig.getDoc(this.chip.docId);
if (!doc) {
throw new Error('Document not found');
}
if (!doc.ready) {
doc.load();
}
const value = await extractMarkdownFromDoc(doc);
const tokenCount = estimateTokenCount(value);
if (this.checkTokenLimit(this.chip, tokenCount)) {
const markdown = this.chip.markdown ?? new Signal<string>('');
markdown.value = value;
this.updateChip(this.chip, {
markdown,
tokenCount,
});
} else {
this.updateChip(this.chip, {
state: 'failed',
tooltip: 'Content exceeds token limit',
});
}
} catch (e) {
this.updateChip(this.chip, {
state: 'failed',
tooltip: e instanceof Error ? e.message : 'Failed to extract markdown',
});
}
};
override render() {
const { state, docId } = this.chip;
const isLoading = state === 'processing';
packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
+1 -2
View File
@@ -18,8 +18,6 @@ export interface BaseChip {
export interface DocChip extends BaseChip {
docId: string;
markdown?: Signal<string> | null;
tokenCount?: number | null;
}
export interface FileChip extends BaseChip {
@@ -85,5 +83,6 @@ export interface DocDisplayConfig {
signal: Signal<{ id: string; name: string }[]>;
cleanup: () => void;
};
getCollectionTitle: (collectionId: string) => string;
getCollectionPageIds: (collectionId: string) => string[];
}
packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts
+41 -55
View File
@@ -146,9 +146,6 @@ export class AIChatComposer extends SignalWatcher(
@state()
accessor isChipsCollapsed = false;
@state()
accessor embeddingCompleted = false;
override render() {
return html`
<chat-panel-chips
@@ -190,7 +187,6 @@ export class AIChatComposer extends SignalWatcher(
.portalContainer=${this.portalContainer}
.onChatSuccess=${this.onChatSuccess}
.trackOptions=${this.trackOptions}
.isContextProcessing=${this.isContextProcessing}
></ai-chat-input>
<div class="chat-panel-footer">
<ai-chat-composer-tip
@@ -206,17 +202,16 @@ export class AIChatComposer extends SignalWatcher(
override connectedCallback() {
super.connectedCallback();
this._disposables.add(
AIAppEvents.requestOpenWithChat.subscribe(this.beforeChatContextSend)
AIAppEvents.requestOpenWithChat.subscribe(this.beforeChatSourceSend)
);
this._disposables.add(
AIAppEvents.requestSendWithChat.subscribe(this.beforeChatContextSend)
AIAppEvents.requestSendWithChat.subscribe(this.beforeChatSourceSend)
);
this.initComposer().catch(console.error);
}
override disconnectedCallback() {
super.disconnectedCallback();
this.runtime?.dispatch({ type: 'stopContextPolling' }).catch(console.error);
}
protected override willUpdate(changedProperties: PropertyValues): void {
@@ -240,7 +235,7 @@ export class AIChatComposer extends SignalWatcher(
}
}
private readonly beforeChatContextSend = (
private readonly beforeChatSourceSend = (
params: AISendParams | AIChatParams | null
) => {
if (!params) return;
@@ -265,19 +260,15 @@ export class AIChatComposer extends SignalWatcher(
}
};
private get isContextProcessing() {
return this.chips.some(chip => chip.state === 'processing');
}
private readonly toChipState = (state?: string): ChipState => {
if (state === 'finished' || state === 'processing' || state === 'failed') {
return state;
}
return 'processing';
return 'finished';
};
private readonly runtimeItemToChip = (
item: AIChatSnapshot['composer']['context']['items'][number]
item: AIChatSnapshot['composer']['scopeSelection']['items'][number]
): ChatChip => {
switch (item.kind) {
case 'doc':
@@ -318,15 +309,20 @@ export class AIChatComposer extends SignalWatcher(
createdAt: item.createdAt,
tooltip: item.tooltip,
};
case 'favorite':
return {
sourceId: item.favoriteId,
name: item.name ?? 'Favorites',
state: 'finished',
};
}
};
private readonly syncChipsFromRuntimeSnapshot = (
snapshot = this.runtimeSnapshot
) => {
const context = snapshot?.composer.context;
const context = snapshot?.composer.scopeSelection;
if (!context) return;
this.embeddingCompleted = context.embeddingCompleted;
const selectedChips = this.chips.filter(isSelectedContextChip);
this.updateChips([
...context.items.map(this.runtimeItemToChip),
@@ -340,9 +336,14 @@ export class AIChatComposer extends SignalWatcher(
private readonly chipToContextItem = (
chip: ChatChip
): AIChatSnapshot['composer']['context']['items'][number] | null => {
): AIChatSnapshot['composer']['scopeSelection']['items'][number] | null => {
if (isDocChip(chip)) {
return { kind: 'doc', docId: chip.docId, state: chip.state };
return {
kind: 'doc',
docId: chip.docId,
name: this.docDisplayConfig.getTitle(chip.docId),
state: chip.state,
};
}
if (isFileChip(chip)) {
return {
@@ -357,6 +358,7 @@ export class AIChatComposer extends SignalWatcher(
return {
kind: 'tag',
tagId: chip.tagId,
name: this.docDisplayConfig.getTagTitle(chip.tagId),
docIds: this.docDisplayConfig.getTagPageIds(chip.tagId),
state: chip.state,
};
@@ -365,6 +367,7 @@ export class AIChatComposer extends SignalWatcher(
return {
kind: 'collection',
collectionId: chip.collectionId,
name: this.docDisplayConfig.getCollectionTitle(chip.collectionId),
docIds: this.docDisplayConfig.getCollectionPageIds(chip.collectionId),
state: chip.state,
};
@@ -376,7 +379,6 @@ export class AIChatComposer extends SignalWatcher(
};
private readonly initChips = async () => {
await this.runtime?.dispatch({ type: 'loadContext' });
this.syncChipsFromRuntime();
};
@@ -417,14 +419,13 @@ export class AIChatComposer extends SignalWatcher(
return;
}
this.updateChips([...this.chips, chip]);
await this.addToContext(chip);
await this.pollContextDocsAndFiles();
await this.addSource(chip);
};
private readonly removeChip = async (chip: ChatChip) => {
const chips = omitChip(this.chips, chip);
this.updateChips(chips);
await this.removeFromContext(chip);
await this.removeSource(chip);
};
private readonly addSelectedContextChip = async () => {
@@ -444,7 +445,7 @@ export class AIChatComposer extends SignalWatcher(
this.addChip(
{
docId,
state: 'processing',
state: 'finished',
},
true
)
@@ -454,7 +455,7 @@ export class AIChatComposer extends SignalWatcher(
{
sourceId: attachment.sourceId,
name: attachment.name,
state: 'processing',
state: 'finished',
},
true
)
@@ -469,7 +470,7 @@ export class AIChatComposer extends SignalWatcher(
}
};
private readonly addToContext = async (chip: ChatChip) => {
private readonly addSource = async (chip: ChatChip) => {
if (isDocChip(chip)) {
return await this.addDocToContext(chip);
}
@@ -491,8 +492,13 @@ export class AIChatComposer extends SignalWatcher(
private readonly addDocToContext = async (chip: DocChip) => {
try {
await this.runtime?.dispatch({
type: 'addContextItem',
item: { kind: 'doc', docId: chip.docId, state: chip.state },
type: 'addScopeSelector',
item: {
kind: 'doc',
docId: chip.docId,
name: this.docDisplayConfig.getTitle(chip.docId),
state: chip.state,
},
});
this.syncChipsFromRuntime();
} catch (e) {
@@ -506,7 +512,7 @@ export class AIChatComposer extends SignalWatcher(
private readonly addFileToContext = async (chip: FileChip) => {
try {
await this.runtime?.dispatch({
type: 'addContextItem',
type: 'addScopeSelector',
item: {
kind: 'file',
file: chip.file,
@@ -527,10 +533,11 @@ export class AIChatComposer extends SignalWatcher(
private readonly addTagToContext = async (chip: TagChip) => {
try {
await this.runtime?.dispatch({
type: 'addContextItem',
type: 'addScopeSelector',
item: {
kind: 'tag',
tagId: chip.tagId,
name: this.docDisplayConfig.getTagTitle(chip.tagId),
docIds: this.docDisplayConfig.getTagPageIds(chip.tagId),
state: chip.state,
},
@@ -547,10 +554,11 @@ export class AIChatComposer extends SignalWatcher(
private readonly addCollectionToContext = async (chip: CollectionChip) => {
try {
await this.runtime?.dispatch({
type: 'addContextItem',
type: 'addScopeSelector',
item: {
kind: 'collection',
collectionId: chip.collectionId,
name: this.docDisplayConfig.getCollectionTitle(chip.collectionId),
docIds: this.docDisplayConfig.getCollectionPageIds(chip.collectionId),
state: chip.state,
},
@@ -570,7 +578,7 @@ export class AIChatComposer extends SignalWatcher(
) => {
try {
await this.runtime?.dispatch({
type: 'addContextItem',
type: 'addScopeSelector',
item: { kind: 'blob', blobId: chip.sourceId, state: chip.state },
});
this.syncChipsFromRuntime();
@@ -583,9 +591,7 @@ export class AIChatComposer extends SignalWatcher(
}
};
private readonly removeFromContext = async (
chip: ChatChip
): Promise<boolean> => {
private readonly removeSource = async (chip: ChatChip): Promise<boolean> => {
if (isSelectedContextChip(chip)) {
this.updateContext({
...this.chatContextValue,
@@ -597,7 +603,7 @@ export class AIChatComposer extends SignalWatcher(
const item = this.chipToContextItem(chip);
if (!item) return true;
try {
await this.runtime?.dispatch({ type: 'removeContextItem', item });
await this.runtime?.dispatch({ type: 'removeScopeSelector', item });
this.syncChipsFromRuntime();
return true;
} catch {
@@ -621,30 +627,10 @@ export class AIChatComposer extends SignalWatcher(
});
};
private readonly pollContextDocsAndFiles = async () => {
if (!this.runtime) return;
await this.runtime.dispatch({ type: 'startContextPolling' });
this.syncChipsFromRuntime();
};
private readonly pollEmbeddingStatus = async () => {
if (!this.runtime) return;
await this.runtime.dispatch({ type: 'pollEmbeddingStatus' });
this.syncChipsFromRuntime();
};
private readonly initComposer = async () => {
const userId = AIAppEvents.userInfo.value?.id;
if (!userId || !this.session) return;
await this.initChips();
const needPoll = this.chips.some(
chip =>
chip.state === 'processing' || isTagChip(chip) || isCollectionChip(chip)
);
if (needPoll) {
await this.pollContextDocsAndFiles();
}
await this.pollEmbeddingStatus();
};
}
packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
+34 -60
View File
@@ -199,6 +199,10 @@ export class AIChatInput extends SignalWatcher(
color: ${unsafeCSSVarV2('icon/secondary')} !important;
}
}
.chat-input-icon[aria-disabled='true']:hover {
background-color: transparent;
}
}
.chat-panel-input {
@@ -341,9 +345,6 @@ export class AIChatInput extends SignalWatcher(
@property({ attribute: false })
accessor runtimeSnapshot: AIChatSnapshot | null | undefined;
@property({ attribute: false })
accessor isContextProcessing!: boolean | undefined;
@query('image-preview-grid')
accessor imagePreviewGrid: HTMLDivElement | null = null;
@@ -544,7 +545,7 @@ export class AIChatInput extends SignalWatcher(
if (entity?.type === 'doc' && entity.id) {
this.addChip({
docId: entity.id,
state: 'processing',
state: 'finished',
}).catch(console.error);
this._trackDragDrop('doc');
}
@@ -619,7 +620,10 @@ export class AIChatInput extends SignalWatcher(
data-testid="chat-panel-input"
></textarea>
<div class="chat-panel-input-actions">
<div class="chat-input-icon">
<div
class="chat-input-icon"
aria-disabled=${!this.searchMenuConfig.addContextAvailable}
>
<ai-chat-add-context
.docId=${this.docId}
.independentMode=${this.independentMode}
@@ -671,10 +675,6 @@ export class AIChatInput extends SignalWatcher(
return true;
}
if (this.isContextProcessing) {
return true;
}
return false;
}
@@ -827,16 +827,11 @@ export class AIChatInput extends SignalWatcher(
send = async (text: string) => {
if (!this.runtime) return;
const { markdown, images, snapshot, combinedElementsMarkdown, html } =
this.chatContextValue;
const userInput = (markdown ? `${markdown}\n` : '') + text;
const { images } = this.chatContextValue;
const imageAttachments = await Promise.all(
images?.map(image => readBlobAsURL(image))
);
const contexts = await this._getMatchedContexts();
const enableSendDetailedObject =
this.affineFeatureFlagService.flags.enable_send_detailed_object_to_ai
.value;
const userInfo = AIAppEvents.userInfo.value;
this.updateContext({
@@ -846,17 +841,8 @@ export class AIChatInput extends SignalWatcher(
});
await this.runtime.dispatch({
type: 'send',
input: userInput,
contexts: {
...contexts,
selectedSnapshot:
snapshot && enableSendDetailedObject ? snapshot : undefined,
selectedMarkdown:
combinedElementsMarkdown && enableSendDetailedObject
? combinedElementsMarkdown
: undefined,
html: html || undefined,
},
input: text,
contexts,
attachments: images,
attachmentPreviews: imageAttachments,
isRootSession: this.isRootSession,
@@ -875,43 +861,31 @@ export class AIChatInput extends SignalWatcher(
};
private async _getMatchedContexts() {
const docContexts = new Map<
string,
{ docId: string; docContent: string }
>();
const docIds = new Set(
this.chips
.filter(isDocChip)
.filter(chip => chip.state !== 'candidate')
.map(chip => chip.docId)
);
this.chips.forEach(chip => {
if (isDocChip(chip) && !!chip.markdown?.value) {
docContexts.set(chip.docId, {
docId: chip.docId,
docContent: chip.markdown.value,
});
}
});
const docs: BlockSuitePresets.AIDocContextOption[] = Array.from(
docContexts.values()
).map(doc => {
const docMeta = this.docDisplayConfig.getDocMeta(doc.docId);
const docTitle = this.docDisplayConfig.getTitle(doc.docId);
const tags = docMeta?.tags
? docMeta.tags
.map(tagId => this.docDisplayConfig.getTagTitle(tagId))
.join(',')
: '';
return {
docId: doc.docId,
docContent: doc.docContent,
docTitle,
tags,
createDate: docMeta?.createDate
const docs: BlockSuitePresets.AIDocContextOption[] = Array.from(docIds).map(
docId => {
const docMeta = this.docDisplayConfig.getDocMeta(docId);
const docTitle = this.docDisplayConfig.getTitle(docId);
const tags = docMeta?.tags
? docMeta.tags
.map(tagId => this.docDisplayConfig.getTagTitle(tagId))
.join(',')
: '';
const createDate = docMeta?.createDate
? new Date(docMeta.createDate).toISOString()
: '',
updatedDate: docMeta?.updatedDate
: '';
const updatedDate = docMeta?.updatedDate
? new Date(docMeta.updatedDate).toISOString()
: '',
};
});
: '';
return { docId, docTitle, tags, createDate, updatedDate };
}
);
return { docs, files: [] };
}
packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/ai-chat-messages.spec.ts
+213 -1
View File
@@ -1,8 +1,18 @@
/**
* @vitest-environment happy-dom
*/
import { render } from 'lit';
import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
import { ChatMessageUser } from '../../chat-panel/message/user';
import { AIChatErrorRenderer, AIErrorWrapper } from '../../messages/error';
import {
SelectedSourcesLimitExceededError,
SelectedSourcesProcessingError,
} from '../../provider/error';
import { ChatContentStreamObjects } from '../ai-message-content/stream-objects';
import { ToolFailedCard } from '../ai-tools/tool-failed-card';
import { ToolResultCard } from '../ai-tools/tool-result-card';
import { AIChatMessages } from './ai-chat-messages';
describe('AIChatMessages scrolling', () => {
@@ -94,7 +104,7 @@ describe('AIChatMessages scrolling', () => {
expect(scrollToEnd).toHaveBeenCalled();
});
test('message keys are scoped by active tab', () => {
test('message keys, scope receipts, and live reads render for the active chat', async () => {
const element = {} as AIChatMessages;
const message = {
id: 'message-1',
@@ -123,5 +133,207 @@ describe('AIChatMessages scrolling', () => {
expect(firstKey).toBe('session-1:message-1');
expect(secondKey).toBe('session-2:message-1');
if (!customElements.get('chat-message-user')) {
customElements.define('chat-message-user', ChatMessageUser);
}
const userMessage = document.createElement('chat-message-user');
userMessage.item = {
id: 'user-message-1',
role: 'user',
content: 'question',
createdAt: new Date().toISOString(),
scopeSnapshot: {
resolvedAt: '2026-08-08T10:00:00.000Z',
selectors: [
{ kind: 'document', id: 'doc-1', name: 'Product notes' },
{ kind: 'artifact', id: 'artifact-1', name: 'brief.pdf' },
],
requiredDocIds: ['doc-1'],
requiredArtifactIds: ['artifact-1'],
},
};
document.body.append(userMessage);
await userMessage.updateComplete;
expect(
userMessage
.querySelector('[data-testid="chat-scope-receipt"]')
?.textContent?.replace(/\s+/g, ' ')
).toContain('Product notes, brief.pdf · 2 sources');
userMessage.remove();
const userMessageWithoutReceipt =
document.createElement('chat-message-user');
userMessageWithoutReceipt.item = {
id: 'user-message-2',
role: 'user',
content: 'new question',
createdAt: new Date().toISOString(),
scopeSnapshot: {
resolvedAt: '2026-08-08T10:00:00.000Z',
selectors: [],
requiredDocIds: [],
requiredArtifactIds: [],
},
};
document.body.append(userMessageWithoutReceipt);
await userMessageWithoutReceipt.updateComplete;
expect(
userMessageWithoutReceipt.querySelector('.scope-receipt')
).toBeNull();
userMessageWithoutReceipt.remove();
if (!customElements.get('tool-result-card')) {
customElements.define('tool-result-card', ToolResultCard);
}
if (!customElements.get('tool-call-failed')) {
customElements.define('tool-call-failed', ToolFailedCard);
}
if (!customElements.get('chat-content-stream-objects')) {
customElements.define(
'chat-content-stream-objects',
ChatContentStreamObjects
);
}
const liveRead = document.createElement(
'chat-content-stream-objects'
) as ChatContentStreamObjects;
liveRead.host = {
std: { store: { meta: { title: 'Getting Started' } } },
} as never;
liveRead.answer = [
{
type: 'tool-result',
toolCallId: 'call_provider_1',
toolName: 'frontend_snapshot_document',
args: { view: 'outline' },
result: {
editor_state_id: 'state-1',
mode: 'page',
outline: [
{
id: 'block-1',
flavour: 'affine:paragraph',
text: { content: 'Welcome to AFFiNE', truncated: false },
},
],
truncated: false,
},
},
{
type: 'tool-result',
toolCallId: 'call_provider_2',
toolName: 'frontend_snapshot_document',
args: { view: 'outline' },
result: {
error: {
code: 'VIEW_NOT_AVAILABLE',
message: 'The requested live editor view is not available.',
retryable: false,
},
},
},
{
type: 'tool-result',
toolCallId: 'call_provider_3',
toolName: 'doc_canvas_read',
args: { doc_id: 'doc-1', target: { kind: 'overview' } },
result: {
doc_id: 'doc-1',
source: {
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
title: 'Getting Started',
visibility: 'edgeless',
},
counts: { blocks: 12, elements: 8 },
blocks: [],
elements: [],
},
},
{
type: 'tool-result',
toolCallId: 'call_provider_2',
toolName: 'frontend_snapshot_document',
args: { view: 'outline' },
result: {
error: {
code: 'VIEW_NOT_AVAILABLE',
message: 'The requested live editor view is not available.',
retryable: false,
},
},
},
];
document.body.append(liveRead);
await liveRead.updateComplete;
const resultCard =
liveRead.querySelector<ToolResultCard>('tool-result-card');
expect(resultCard?.name).toBe('Read outline of "Getting Started"');
expect(resultCard?.results).toEqual([
expect.objectContaining({
title: 'Getting Started',
content: 'Welcome to AFFiNE',
}),
]);
expect(
liveRead.querySelector<ToolFailedCard>('tool-call-failed')?.name
).toBe('This view is not available in the current editor mode');
expect(
[...liveRead.querySelectorAll<ToolResultCard>('tool-result-card')].map(
card => card.name
)
).toContain('Read canvas of "Getting Started"');
expect(
liveRead.querySelector<HTMLDetailsElement>('.tool-group')?.open
).toBe(true);
expect(
liveRead.querySelector('.tool-group-summary')?.textContent
).toContain('3 actions · 1 failed');
liveRead.remove();
if (!customElements.get('ai-error-wrapper')) {
customElements.define('ai-error-wrapper', AIErrorWrapper);
}
const retry = vi.fn();
const errorContainer = document.createElement('div');
document.body.append(errorContainer);
render(
AIChatErrorRenderer(
new SelectedSourcesProcessingError('processing'),
null,
retry
),
errorContainer
);
const error =
errorContainer.querySelector<AIErrorWrapper>('ai-error-wrapper');
await error?.updateComplete;
expect(error?.text).toContain('still processing');
expect(error?.actionTooltip).toBe('');
error?.shadowRoot
?.querySelector<HTMLElement>('[data-testid="ai-error-action-button"]')
?.click();
expect(retry).toHaveBeenCalledOnce();
render(
AIChatErrorRenderer(new SelectedSourcesProcessingError('processing')),
errorContainer
);
const processingWithoutRetry =
errorContainer.querySelector<AIErrorWrapper>('ai-error-wrapper');
await processingWithoutRetry?.updateComplete;
expect(processingWithoutRetry?.showAction).toBe(false);
render(
AIChatErrorRenderer(new SelectedSourcesLimitExceededError('limit')),
errorContainer
);
const limitError =
errorContainer.querySelector<AIErrorWrapper>('ai-error-wrapper');
await limitError?.updateComplete;
expect(limitError?.showAction).toBe(false);
errorContainer.remove();
});
});
packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/type.ts
+22
View File
@@ -36,6 +36,28 @@ const ChatMessageSchema = z.object({
userId: z.string().optional(),
userName: z.string().optional(),
avatarUrl: z.string().optional(),
scopeSnapshot: z
.object({
resolvedAt: z.string(),
selectors: z.array(
z.object({
kind: z.enum([
'document',
'tag',
'collection',
'favorite',
'artifact',
]),
id: z.string(),
name: z.string().optional(),
})
),
requiredDocIds: z.array(z.string()),
requiredArtifactIds: z.array(z.string()),
})
.passthrough()
.nullable()
.optional(),
});
export const ChatMessagesSchema = z.array(ChatMessageSchema);
packages/frontend/core/src/blocksuite/ai/components/ai-message-content/stream-objects.ts
+433 -19
View File
@@ -1,5 +1,6 @@
import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
import type { PeekViewService } from '@affine/core/modules/peek-view';
import { I18n } from '@affine/i18n';
import { WithDisposable } from '@blocksuite/affine/global/lit';
import type { ColorScheme } from '@blocksuite/affine/model';
import {
@@ -9,14 +10,148 @@ import {
} from '@blocksuite/affine/std';
import type { ExtensionType } from '@blocksuite/affine/store';
import type { NotificationService } from '@blocksuite/affine-shared/services';
import {
EdgelessIcon,
PageIcon,
ToggleDownIcon,
ToolIcon,
ViewIcon,
} from '@blocksuite/icons/lit';
import type { Signal } from '@preact/signals-core';
import { css, html, nothing } from 'lit';
import { property } from 'lit/decorators.js';
import { property, state } from 'lit/decorators.js';
import { classMap } from 'lit/directives/class-map.js';
import type { AffineAIPanelState } from '../../widgets/ai-panel/type';
import type { DocDisplayConfig } from '../ai-chat-chips';
import type { StreamObject } from '../ai-chat-messages';
const frontendReadTools = new Set([
'frontend_get_editor_state',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
]);
function object(value: unknown): Record<string, unknown> | undefined {
return value && typeof value === 'object'
? (value as Record<string, unknown>)
: undefined;
}
function projectedText(value: unknown): string | undefined {
const item = object(value);
const projected = object(item?.value) ?? item;
const text = projected?.text;
if (typeof text === 'string') return text;
const content = object(text)?.content;
if (typeof content === 'string') return content;
const title = projected?.title;
if (typeof title === 'string') return title;
const titleContent = object(title)?.content;
return typeof titleContent === 'string' ? titleContent : undefined;
}
function frontendReadPreview(result: unknown) {
const value = object(result);
if (!value) return undefined;
const viewport = object(value.viewport);
const items = [
value.outline,
value.selection_neighborhood,
value.items,
value.nodes,
value.blocks,
value.elements,
viewport?.elements,
].find((item): item is unknown[] => Array.isArray(item) && item.length > 0);
const preview = items
?.map(projectedText)
.filter((text): text is string => !!text)
.slice(0, 4)
.join('\n');
if (preview) return preview;
return typeof value.mode === 'string'
? I18n['com.affine.ai.chat-panel.tool.live.mode']({ mode: value.mode })
: undefined;
}
function frontendReadError(result: Record<string, unknown>) {
const error = object(result.error);
if (error?.code === 'VIEW_NOT_AVAILABLE') {
return I18n['com.affine.ai.chat-panel.tool.live.view-unavailable']();
}
if (typeof error?.message === 'string') return error.message;
if (typeof result.message === 'string') return result.message;
return I18n['com.affine.ai.chat-panel.tool.live.failed']();
}
function canvasReadPreview(result: Record<string, unknown>) {
const preview = frontendReadPreview(result);
if (preview) return preview;
const counts = object(result.counts);
const parts: string[] = [];
if (typeof counts?.blocks === 'number') {
parts.push(
I18n['com.affine.ai.chat-panel.tool.live.blocks']({
count: String(counts.blocks),
})
);
}
if (typeof counts?.elements === 'number') {
parts.push(
I18n['com.affine.ai.chat-panel.tool.live.elements']({
count: String(counts.elements),
})
);
}
return parts.join(' · ') || undefined;
}
type ToolStreamObject = Extract<
StreamObject,
{ type: 'tool-call' | 'tool-result' }
>;
function isToolObject(value: StreamObject): value is ToolStreamObject {
return value.type === 'tool-call' || value.type === 'tool-result';
}
function toolFailed(value: StreamObject) {
if (value.type !== 'tool-result') return false;
const result = object(value.result);
return !result || !!result.error || result.type === 'error';
}
type StreamGroup =
| { type: 'item'; item: StreamObject }
| { type: 'tools'; items: ToolStreamObject[]; key: string };
function groupStreamObjects(answer: StreamObject[]): StreamGroup[] {
const groups: StreamGroup[] = [];
for (let index = 0; index < answer.length; ) {
if (!isToolObject(answer[index])) {
groups.push({ type: 'item', item: answer[index] });
index += 1;
continue;
}
const items: ToolStreamObject[] = [];
while (index < answer.length) {
const item = answer[index];
if (!isToolObject(item)) break;
items.push(item);
index += 1;
}
const callIds = [...new Set(items.map(item => item.toolCallId))];
if (callIds.length < 3) {
groups.push(...items.map(item => ({ type: 'item' as const, item })));
continue;
}
groups.push({ type: 'tools', items, key: callIds.join(':') });
}
return groups;
}
export class ChatContentStreamObjects extends WithDisposable(
ShadowlessElement
) {
@@ -27,6 +162,97 @@ export class ChatContentStreamObjects extends WithDisposable(
border-radius: 8px;
background-color: rgba(0, 0, 0, 0.05);
}
.tool-group {
margin: 8px 0;
border: 0.5px solid var(--affine-border-color);
border-radius: 8px;
overflow: hidden;
}
.tool-group-summary {
display: flex;
align-items: center;
gap: 8px;
padding: 12px;
color: var(--affine-text-secondary-color);
cursor: pointer;
list-style: none;
user-select: none;
}
.tool-group-summary::-webkit-details-marker {
display: none;
}
.tool-group-summary:focus-visible {
outline: 2px solid currentColor;
outline-offset: -2px;
}
.tool-group-icon,
.tool-group-toggle {
display: flex;
width: 24px;
height: 24px;
align-items: center;
justify-content: center;
}
.tool-group-icon svg,
.tool-group-toggle svg {
width: 24px;
height: 24px;
}
.tool-group-title {
flex: 1;
min-width: 0;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
font-size: 14px;
font-weight: 500;
line-height: 24px;
}
.tool-group-toggle {
transition: transform 0.23s ease;
transform: rotate(-90deg);
}
.tool-group[open] .tool-group-toggle {
transform: rotate(0deg);
}
.tool-group-content {
border-top: 0.5px solid var(--affine-border-color);
}
.tool-group-content .ai-tool-call-wrapper,
.tool-group-content .ai-tool-result-wrapper,
.tool-group-content .ai-tool-failed-wrapper {
margin: 0;
border: 0;
border-bottom: 0.5px solid var(--affine-border-color);
border-radius: 0;
}
.tool-group-content > :last-child .ai-tool-call-wrapper,
.tool-group-content > :last-child .ai-tool-result-wrapper,
.tool-group-content > :last-child .ai-tool-failed-wrapper {
border-bottom: 0;
}
.tool-group.failed > .tool-group-summary {
color: var(--affine-error-color);
}
@media (prefers-reduced-motion: reduce) {
.tool-group-toggle {
transition: none;
}
}
`;
@property({ attribute: false })
@@ -68,11 +294,144 @@ export class ChatContentStreamObjects extends WithDisposable(
@property({ attribute: false })
accessor onOpenDoc!: (docId: string, sessionId?: string) => void;
@state()
private accessor toolGroupOverrides = new Map<string, boolean>();
private renderFrontendRead(streamObject: StreamObject) {
if (
(streamObject.type !== 'tool-call' &&
streamObject.type !== 'tool-result') ||
!frontendReadTools.has(streamObject.toolName)
) {
return nothing;
}
const title =
this.host?.std.store.meta?.title ||
I18n['com.affine.ai.chat-panel.tool.live.current-document']();
const result =
streamObject.type === 'tool-result'
? object(streamObject.result)
: undefined;
const isCanvasSnapshot =
streamObject.toolName === 'frontend_snapshot_document' &&
(!!object(result?.viewport) ||
object(streamObject.args)?.view === 'viewport');
const labels = {
frontend_get_editor_state: [
I18n['com.affine.ai.chat-panel.tool.live.state-checking'](),
I18n['com.affine.ai.chat-panel.tool.live.state-checked'](),
],
frontend_read_selection: [
I18n['com.affine.ai.chat-panel.tool.live.selection-reading'](),
I18n['com.affine.ai.chat-panel.tool.live.selection-read'](),
],
frontend_read_nodes: [
I18n['com.affine.ai.chat-panel.tool.live.content-reading'](),
I18n['com.affine.ai.chat-panel.tool.live.content-read'](),
],
frontend_snapshot_document: [
isCanvasSnapshot
? I18n['com.affine.ai.chat-panel.tool.live.canvas-reading']({ title })
: I18n['com.affine.ai.chat-panel.tool.live.outline-reading']({
title,
}),
isCanvasSnapshot
? I18n['com.affine.ai.chat-panel.tool.live.canvas-read']({ title })
: I18n['com.affine.ai.chat-panel.tool.live.outline-read']({ title }),
],
} as const;
const [callLabel, resultLabel] =
labels[streamObject.toolName as keyof typeof labels];
if (streamObject.type === 'tool-call') {
return html`<tool-call-card
.name=${callLabel}
.icon=${ViewIcon()}
.width=${this.width}
></tool-call-card>`;
}
if (!result || result.error || result.type === 'error') {
return html`<tool-call-failed
.name=${result
? frontendReadError(result)
: I18n['com.affine.ai.chat-panel.tool.live.failed']()}
.icon=${ViewIcon()}
></tool-call-failed>`;
}
return html`<tool-result-card
.name=${resultLabel}
.icon=${ViewIcon()}
.width=${this.width}
.results=${[
{
title:
streamObject.toolName === 'frontend_snapshot_document'
? title
: resultLabel,
icon: PageIcon(),
content: frontendReadPreview(result),
},
]}
></tool-result-card>`;
}
private renderCanvasRead(streamObject: StreamObject) {
if (
(streamObject.type !== 'tool-call' &&
streamObject.type !== 'tool-result') ||
streamObject.toolName !== 'doc_canvas_read'
) {
return nothing;
}
if (streamObject.type === 'tool-call') {
return html`<tool-call-card
.name=${I18n['com.affine.ai.chat-panel.tool.canvas.reading']()}
.icon=${ViewIcon()}
.width=${this.width}
></tool-call-card>`;
}
const result = object(streamObject.result);
if (!result || result.error || result.type === 'error') {
return html`<tool-call-failed
.name=${result
? frontendReadError(result)
: I18n['com.affine.ai.chat-panel.tool.canvas.failed']()}
.icon=${ViewIcon()}
></tool-call-failed>`;
}
const title = object(result.source)?.title;
const name =
typeof title === 'string'
? I18n['com.affine.ai.chat-panel.tool.canvas.read']({ title })
: I18n['com.affine.ai.chat-panel.tool.canvas.read-untitled']();
return html`<tool-result-card
.name=${name}
.icon=${ViewIcon()}
.width=${this.width}
.results=${[
{
title:
typeof title === 'string'
? title
: I18n['com.affine.ai.chat-panel.tool.canvas.content'](),
icon: EdgelessIcon(),
content: canvasReadPreview(result),
},
]}
></tool-result-card>`;
}
private renderToolCall(streamObject: StreamObject) {
if (streamObject.type !== 'tool-call') {
return nothing;
}
if (frontendReadTools.has(streamObject.toolName)) {
return this.renderFrontendRead(streamObject);
}
if (streamObject.toolName === 'doc_canvas_read') {
return this.renderCanvasRead(streamObject);
}
switch (streamObject.toolName) {
case 'web_crawl_exa':
return html`
@@ -122,6 +481,7 @@ export class ChatContentStreamObjects extends WithDisposable(
.peekViewService=${this.peekViewService}
></doc-semantic-search-result>`;
case 'doc_keyword_search':
case 'doc_search':
return html`<doc-keyword-search-result
.data=${streamObject}
.width=${this.width}
@@ -167,6 +527,13 @@ export class ChatContentStreamObjects extends WithDisposable(
return nothing;
}
if (frontendReadTools.has(streamObject.toolName)) {
return this.renderFrontendRead(streamObject);
}
if (streamObject.toolName === 'doc_canvas_read') {
return this.renderCanvasRead(streamObject);
}
switch (streamObject.toolName) {
case 'web_crawl_exa':
return html`
@@ -220,6 +587,7 @@ export class ChatContentStreamObjects extends WithDisposable(
.onOpenDoc=${this.onOpenDoc}
></doc-semantic-search-result>`;
case 'doc_keyword_search':
case 'doc_search':
return html`<doc-keyword-search-result
.data=${streamObject}
.width=${this.width}
@@ -277,26 +645,72 @@ export class ChatContentStreamObjects extends WithDisposable(
></chat-content-rich-text>`;
}
private renderStreamObject(data: StreamObject) {
switch (data.type) {
case 'text-delta':
return this.renderRichText(data.textDelta);
case 'reasoning':
return html`
<div class="reasoning-wrapper">
${this.renderRichText(data.textDelta)}
</div>
`;
case 'tool-call':
return this.renderToolCall(data);
case 'tool-result':
return this.renderToolResult(data);
}
}
private renderToolGroup(group: Extract<StreamGroup, { type: 'tools' }>) {
const callCount = new Set(group.items.map(item => item.toolCallId)).size;
const failedCount = new Set(
group.items.filter(toolFailed).map(item => item.toolCallId)
).size;
const defaultOpen = this.state !== 'finished' || failedCount > 0;
const open = this.toolGroupOverrides.get(group.key) ?? defaultOpen;
const title = failedCount
? I18n['com.affine.ai.chat-panel.tool-group.failed']({
count: String(callCount),
failed: String(failedCount),
})
: this.state === 'finished'
? I18n['com.affine.ai.chat-panel.tool-group.completed']({
count: String(callCount),
})
: I18n['com.affine.ai.chat-panel.tool-group.running']({
count: String(callCount),
});
return html`<details
class=${classMap({ 'tool-group': true, failed: failedCount > 0 })}
.open=${open}
@toggle=${(event: Event) => {
const details = event.currentTarget as HTMLDetailsElement;
if (details.open === open) return;
this.toolGroupOverrides = new Map(this.toolGroupOverrides).set(
group.key,
details.open
);
}}
>
<summary class="tool-group-summary">
<span class="tool-group-icon">${ToolIcon()}</span>
<span class="tool-group-title">${title}</span>
<span class="tool-group-toggle">${ToggleDownIcon()}</span>
</summary>
<div class="tool-group-content">
${group.items.map(item => this.renderStreamObject(item))}
</div>
</details>`;
}
protected override render() {
return html`<div>
${this.answer.map(data => {
switch (data.type) {
case 'text-delta':
return this.renderRichText(data.textDelta);
case 'reasoning':
return html`
<div class="reasoning-wrapper">
${this.renderRichText(data.textDelta)}
</div>
`;
case 'tool-call':
return this.renderToolCall(data);
case 'tool-result':
return this.renderToolResult(data);
default:
return nothing;
}
})}
${groupStreamObjects(this.answer).map(group =>
group.type === 'tools'
? this.renderToolGroup(group)
: this.renderStreamObject(group.item)
)}
</div>`;
}
}
packages/frontend/core/src/blocksuite/ai/components/ai-tools/doc-keyword-search-result.ts
+10 -3
View File
@@ -22,7 +22,11 @@ interface DocKeywordSearchToolResult {
toolCallId: string;
toolName: string;
args: { query: string };
result: Array<{ title: string; docId: string }> | ToolError | null;
result:
| Array<{ title: string; docId: string }>
| { hits: Array<{ title: string; doc_id: string; excerpt?: string }> }
| ToolError
| null;
}
export class DocKeywordSearchResult extends WithDisposable(ShadowlessElement) {
@@ -64,9 +68,12 @@ export class DocKeywordSearchResult extends WithDisposable(ShadowlessElement) {
.icon=${SearchIcon()}
></tool-call-failed>`;
}
const items = Array.isArray(result)
? result
: result.hits.map(item => ({ title: item.title, docId: item.doc_id }));
let results: ToolResult[] = [];
try {
results = result.map(item => ({
results = items.map(item => ({
title: item.title,
icon: PageIcon(),
onClick: () => {
@@ -82,7 +89,7 @@ export class DocKeywordSearchResult extends WithDisposable(ShadowlessElement) {
console.error('Failed to parse result', err);
}
return html`<tool-result-card
.name=${`Found ${result.length} pages for "${this.data.args.query}"`}
.name=${`Found ${items.length} pages for "${this.data.args.query}"`}
.icon=${SearchIcon()}
.width=${this.width}
.results=${results}
packages/frontend/core/src/blocksuite/ai/components/playground/chat.ts
+4 -8
View File
@@ -325,19 +325,15 @@ export class PlaygroundChat extends SignalWatcher(
}
override render() {
const embeddingCount =
this.runtimeSnapshot?.composer.context.embeddingCount;
const done = embeddingCount?.finished ?? 0;
const total =
done + (embeddingCount?.processing ?? 0) + (embeddingCount?.failed ?? 0);
const isEmbedding = total > 0 && done < total;
const isSynchronizing =
this.runtimeSnapshot?.composer.scopeSelection.syncing ?? false;
return html`<div class="chat-panel-container">
<div class="chat-panel-title">
<div class="chat-panel-title-text">
${isEmbedding
${isSynchronizing
? html`<span data-testid="chat-panel-embedding-progress"
>Embedding ${done}/${total}</span
>Synchronizing sources</span
>`
: 'AFFiNE AI'}
</div>
packages/frontend/core/src/blocksuite/ai/index.ts
+1
View File
@@ -6,5 +6,6 @@ export * from './messages';
export { AIChatBlockPeekViewTemplate } from './peek-view/chat-block-peek-view';
export * from './provider';
export * from './runtime/chat';
export * from './runtime/frontend';
export * from './runtime/request';
export * from './utils/edgeless';
packages/frontend/core/src/blocksuite/ai/messages/error.ts
+63 -17
View File
@@ -1,3 +1,4 @@
import { I18n } from '@affine/i18n';
import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
import { scrollbarStyle } from '@blocksuite/affine/shared/styles';
import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
@@ -12,6 +13,10 @@ import {
AIAppEvents,
type AIError,
PaymentRequiredError,
SelectedSourcesFailedError,
SelectedSourcesLimitExceededError,
SelectedSourcesProcessingError,
SelectedSourcesUnavailableError,
UnauthorizedError,
} from '../provider';
@@ -147,20 +152,22 @@ export class AIErrorWrapper extends SignalWatcher(WithDisposable(LitElement)) {
: nothing}
</div>
</div>
<div class="action">
<span
class="action-button"
@click=${this.onClick}
data-testid="ai-error-action-button"
>
${this.actionText}
${this.actionTooltip
? html`<affine-tooltip tip-position="top">
${this.actionTooltip}
</affine-tooltip>`
: nothing}
</span>
</div>
${this.showAction
? html`<div class="action">
<span
class="action-button"
@click=${this.onClick}
data-testid="ai-error-action-button"
>
${this.actionText}
${this.actionTooltip
? html`<affine-tooltip tip-position="top">
${this.actionTooltip}
</affine-tooltip>`
: nothing}
</span>
</div>`
: nothing}
</div>`;
}
@@ -182,6 +189,9 @@ export class AIErrorWrapper extends SignalWatcher(WithDisposable(LitElement)) {
@property({ attribute: false })
accessor showDetailPanel: boolean = false;
@property({ attribute: false })
accessor showAction: boolean = true;
@property({ attribute: 'data-testid', reflect: true })
accessor testId = 'ai-error';
}
@@ -207,13 +217,15 @@ type ErrorProps = {
errorMessage?: string;
actionText?: string;
actionTooltip?: string;
onClick?: () => void;
showAction?: boolean;
};
const generalErrorText =
'An error occurred, If this issue persists please let us know.';
const GeneralErrorRenderer = (props: ErrorProps = {}) => {
const onClick = () => {
const contactSupport = () => {
window.open('mailto:support@toeverything.info', '_blank');
};
@@ -223,15 +235,49 @@ const GeneralErrorRenderer = (props: ErrorProps = {}) => {
.showDetailPanel=${!!props.errorMessage}
.actionText=${props.actionText ?? 'Contact us'}
.actionTooltip=${props.actionTooltip ?? 'support@toeverything.info'}
.onClick=${onClick}
.onClick=${props.onClick ?? contactSupport}
.showAction=${props.showAction ?? true}
></ai-error-wrapper>`;
};
export function AIChatErrorRenderer(error: AIError, host?: EditorHost | null) {
export function AIChatErrorRenderer(
error: AIError,
host?: EditorHost | null,
retry?: () => void
) {
if (error instanceof PaymentRequiredError) {
return PaymentRequiredErrorRenderer(host);
} else if (error instanceof UnauthorizedError) {
return LoginRequiredErrorRenderer(host);
} else if (error instanceof SelectedSourcesProcessingError) {
return GeneralErrorRenderer({
text: I18n['com.affine.ai.error.selectedSourcesProcessing'](),
actionText: I18n['com.affine.ai.error.retry'](),
actionTooltip: '',
onClick: retry,
showAction: !!retry,
});
} else if (error instanceof SelectedSourcesFailedError) {
return GeneralErrorRenderer({
text: I18n['com.affine.ai.error.selectedSourcesFailed'](),
actionText: I18n['com.affine.ai.error.retry'](),
actionTooltip: '',
onClick: retry,
showAction: !!retry,
});
} else if (error instanceof SelectedSourcesUnavailableError) {
return GeneralErrorRenderer({
text: I18n['com.affine.ai.error.selectedSourcesUnavailable'](),
actionText: I18n['com.affine.ai.error.retry'](),
actionTooltip: '',
onClick: retry,
showAction: !!retry,
});
} else if (error instanceof SelectedSourcesLimitExceededError) {
return GeneralErrorRenderer({
text: I18n['com.affine.ai.error.selectedSourcesLimitExceeded'](),
showAction: false,
});
} else {
return GeneralErrorRenderer({
errorMessage: error.message,
packages/frontend/core/src/blocksuite/ai/provider/error.ts
+41 -1
View File
@@ -7,6 +7,10 @@ export enum AIErrorType {
PaymentRequired = 'PaymentRequired',
Unauthorized = 'Unauthorized',
RequestTimeout = 'RequestTimeout',
SelectedSourcesProcessing = 'SelectedSourcesProcessing',
SelectedSourcesFailed = 'SelectedSourcesFailed',
SelectedSourcesUnavailable = 'SelectedSourcesUnavailable',
SelectedSourcesLimitExceeded = 'SelectedSourcesLimitExceeded',
}
export class UnauthorizedError extends BaseAIError {
@@ -44,8 +48,44 @@ export class RequestTimeoutError extends BaseAIError {
}
}
export class SelectedSourcesProcessingError extends BaseAIError {
readonly type = AIErrorType.SelectedSourcesProcessing;
constructor(message: string) {
super(message);
}
}
export class SelectedSourcesFailedError extends BaseAIError {
readonly type = AIErrorType.SelectedSourcesFailed;
constructor(message: string) {
super(message);
}
}
export class SelectedSourcesUnavailableError extends BaseAIError {
readonly type = AIErrorType.SelectedSourcesUnavailable;
constructor(message: string) {
super(message);
}
}
export class SelectedSourcesLimitExceededError extends BaseAIError {
readonly type = AIErrorType.SelectedSourcesLimitExceeded;
constructor(message: string) {
super(message);
}
}
export type AIError =
| UnauthorizedError
| PaymentRequiredError
| GeneralNetworkError
| RequestTimeoutError;
| RequestTimeoutError
| SelectedSourcesProcessingError
| SelectedSourcesFailedError
| SelectedSourcesUnavailableError
| SelectedSourcesLimitExceededError;
packages/frontend/core/src/blocksuite/ai/runtime/chat/actions.ts
+5 -11
View File
@@ -1,15 +1,12 @@
import type { CopilotChatHistoryFragment } from '@affine/graphql';
import type { AIChatContextItem, AIChatScope } from './state';
import type { AIChatScope, AIChatScopeSelector } from './state';
export type AIChatSendOptions = {
input?: string;
contexts?: {
docs?: unknown;
files?: unknown;
selectedSnapshot?: unknown;
selectedMarkdown?: unknown;
html?: unknown;
};
attachments?: (string | Blob | File)[];
attachmentPreviews?: string[];
@@ -48,13 +45,10 @@ export type AIChatAction =
| { type: 'setRouteTarget'; routeTargetId?: string }
| { type: 'addAttachment'; attachment: string | Blob | File }
| { type: 'removeAttachment'; index: number }
| { type: 'addContextItem'; item: AIChatContextItem }
| { type: 'removeContextItem'; item: AIChatContextItem }
| { type: 'loadContext' }
| { type: 'pollContext' }
| { type: 'startContextPolling' }
| { type: 'stopContextPolling' }
| { type: 'pollEmbeddingStatus' }
| { type: 'addScopeSelector'; item: AIChatScopeSelector }
| { type: 'removeScopeSelector'; item: AIChatScopeSelector }
| { type: 'addFocusSelector'; item: AIChatScopeSelector }
| { type: 'removeFocusSelector'; item: AIChatScopeSelector }
| ({ type: 'send' } & AIChatSendOptions)
| { type: 'retry'; messageId: string }
| { type: 'stop' };
packages/frontend/core/src/blocksuite/ai/runtime/chat/runtime.spec.ts
+120 -166
View File
@@ -4,6 +4,7 @@
import type { CopilotChatHistoryFragment } from '@affine/graphql';
import { describe, expect, test, vi } from 'vitest';
import { SelectedSourcesProcessingError } from '../../provider/error';
import type { AIRequestService } from '../request';
import { AIChatRuntime } from './runtime';
import {
@@ -66,32 +67,12 @@ function createRequest(
createSessionWithHistory: vi.fn().mockResolvedValue(session()),
updateSession: vi.fn().mockResolvedValue(undefined),
cleanupSessions: vi.fn().mockResolvedValue(undefined),
getActiveEditorContext: vi.fn().mockReturnValue(undefined),
executeAction: vi.fn().mockResolvedValue(stream(['hello'])),
waitForSelectedSources: vi.fn().mockResolvedValue(undefined),
histories: {
ids: vi.fn().mockResolvedValue([]),
},
context: {
createContext: vi.fn().mockResolvedValue('context-1'),
getContextId: vi.fn().mockResolvedValue(undefined),
addContextDoc: vi.fn().mockResolvedValue(undefined),
removeContextDoc: vi.fn().mockResolvedValue(undefined),
addContextFile: vi
.fn()
.mockResolvedValue({ id: 'file-1', status: 'processing' }),
removeContextFile: vi.fn().mockResolvedValue(undefined),
addContextTag: vi.fn().mockResolvedValue(undefined),
removeContextTag: vi.fn().mockResolvedValue(undefined),
addContextCollection: vi.fn().mockResolvedValue(undefined),
removeContextCollection: vi.fn().mockResolvedValue(undefined),
getContextDocsAndFiles: vi.fn().mockResolvedValue(undefined),
matchContext: vi.fn().mockResolvedValue({ files: [], docs: [] }),
addContextBlob: vi
.fn()
.mockResolvedValue({ id: 'blob-1', status: 'processing' }),
removeContextBlob: vi.fn().mockResolvedValue(undefined),
pollContextDocsAndFiles: vi.fn(),
pollEmbeddingStatus: vi.fn(),
},
...overrides,
} as unknown as AIRequestService;
}
@@ -128,6 +109,7 @@ describe('AIChatRuntime', () => {
id: 'message-1',
role: 'user',
content: 'previous chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -149,7 +131,7 @@ describe('AIChatRuntime', () => {
expect(runtime.getSnapshot().messages).toEqual(fullSession.messages);
});
test('send creates a session once and ignores duplicate sends while transmitting', async () => {
test('send tolerates optional editor activation and ignores duplicate sends while transmitting', async () => {
let release!: () => void;
const blockedStream = {
async *[Symbol.asyncIterator]() {
@@ -160,6 +142,7 @@ describe('AIChatRuntime', () => {
},
};
const request = createRequest({
getActiveEditorContext: vi.fn().mockReturnValue('{"mode":"page"}'),
executeAction: vi.fn().mockResolvedValue(blockedStream),
});
const runtime = createRuntime(request);
@@ -175,8 +158,37 @@ describe('AIChatRuntime', () => {
expect(request.createSessionWithHistory).toHaveBeenCalledTimes(1);
expect(request.executeAction).toHaveBeenCalledTimes(1);
expect(request.executeAction).toHaveBeenCalledWith(
'chat',
expect.objectContaining({ liveEditorContext: '{"mode":"page"}' })
);
expect(runtime.getSnapshot().messages.at(-1)?.content).toBe('done');
expect(runtime.getSnapshot().uiPolicy.canCreateNewSession).toBe(true);
const activationError = new Error('editor unavailable');
const consoleError = vi
.spyOn(console, 'error')
.mockImplementation(() => undefined);
const fallbackRequest = createRequest({
activateEditor: vi.fn().mockRejectedValue(activationError),
getActiveEditorContext: vi.fn().mockReturnValue('stale context'),
executeAction: vi.fn().mockResolvedValue(stream(['fallback'])),
});
const fallbackRuntime = createRuntime(fallbackRequest);
await fallbackRuntime.dispatch({ type: 'initialize' });
await fallbackRuntime.dispatch({ type: 'send', input: 'hello' });
expect(consoleError).toHaveBeenCalledWith(activationError);
expect(fallbackRequest.executeAction).toHaveBeenCalledWith(
'chat',
expect.objectContaining({ liveEditorContext: undefined })
);
expect(fallbackRuntime.getSnapshot().messages.at(-1)?.content).toBe(
'fallback'
);
expect(fallbackRuntime.getSnapshot().status).toBe('success');
consoleError.mockRestore();
});
test('send binds an unbound session to the active doc after success', async () => {
@@ -218,6 +230,7 @@ describe('AIChatRuntime', () => {
id: 'message-1',
role: 'user',
content: 'existing chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -281,6 +294,7 @@ describe('AIChatRuntime', () => {
id: 'message-1',
role: 'user',
content: 'first chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -297,6 +311,7 @@ describe('AIChatRuntime', () => {
id: 'message-2',
role: 'user',
content: 'second chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -313,6 +328,7 @@ describe('AIChatRuntime', () => {
id: 'message-1',
role: 'user',
content: 'first chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -360,6 +376,7 @@ describe('AIChatRuntime', () => {
id: 'message-1',
role: 'user',
content: 'old chat',
scopeSnapshot: null,
attachments: [],
streamObjects: [],
createdAt: new Date().toISOString(),
@@ -567,7 +584,7 @@ describe('AIChatRuntime', () => {
});
test('retry failure commits error status and keeps the retried assistant placeholder', async () => {
const error = new Error('retry failed');
const error = new SelectedSourcesProcessingError('retry failed');
const request = createRequest({
executeAction: vi.fn().mockRejectedValue(error),
});
@@ -581,6 +598,7 @@ describe('AIChatRuntime', () => {
role: 'user',
content: 'hello',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -589,6 +607,7 @@ describe('AIChatRuntime', () => {
role: 'assistant',
content: 'old',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -649,6 +668,7 @@ describe('AIChatRuntime', () => {
role: 'user',
content: 'hello',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -657,6 +677,7 @@ describe('AIChatRuntime', () => {
role: 'assistant',
content: 'old',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -695,6 +716,7 @@ describe('AIChatRuntime', () => {
role: 'user',
content: 'hello',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -703,6 +725,7 @@ describe('AIChatRuntime', () => {
role: 'assistant',
content: 'old',
createdAt: new Date().toISOString(),
scopeSnapshot: null,
attachments: null,
streamObjects: null,
},
@@ -778,156 +801,87 @@ describe('AIChatRuntime', () => {
expect(runtime.getSnapshot().history.recent[0].sessionId).toBe('recent');
});
test('context add remove and poll preserve operation order', async () => {
test('selected sources stay local, synchronize before send, and clear after success', async () => {
const request = createRequest();
const runtime = createRuntime(request);
const file = new File(['attachment'], 'attachment.txt', {
type: 'text/plain',
});
const image = new File(['image'], 'image.png', { type: 'image/png' });
await runtime.dispatch({ type: 'initialize' });
await runtime.dispatch({
type: 'addScopeSelector',
item: { kind: 'doc', docId: 'doc-1' },
});
await runtime.dispatch({
type: 'addScopeSelector',
item: { kind: 'doc', docId: 'doc-2' },
});
await runtime.dispatch({
type: 'removeScopeSelector',
item: { kind: 'doc', docId: 'doc-1' },
});
await runtime.dispatch({
type: 'addFocusSelector',
item: { kind: 'tag', tagId: 'tag-1', docIds: [] },
});
await runtime.dispatch({
type: 'addFocusSelector',
item: { kind: 'tag', tagId: 'tag-2', docIds: ['doc-3'] },
});
await runtime.dispatch({
type: 'removeFocusSelector',
item: { kind: 'tag', tagId: 'tag-1', docIds: [] },
});
await runtime.dispatch({
type: 'addScopeSelector',
item: { kind: 'file', file },
});
await runtime.dispatch({
type: 'send',
input: 'question',
attachments: [image],
});
expect(request.waitForSelectedSources).toHaveBeenCalledTimes(1);
expect(request.waitForSelectedSources).toHaveBeenCalledWith([
'doc-2',
'doc-3',
]);
expect(request.executeAction).toHaveBeenCalledWith(
'chat',
expect.objectContaining({
scopeSelectors: [{ kind: 'document', id: 'doc-2' }],
focusSelectors: [{ kind: 'tag', id: 'tag-2' }],
attachments: [file, image],
})
);
expect(runtime.getSnapshot().composer.scopeSelection.items).toEqual([]);
expect(runtime.getSnapshot().composer.focus.items).toEqual([
{ kind: 'tag', tagId: 'tag-2', docIds: ['doc-3'] },
]);
});
test('selected source synchronization failure does not submit a message', async () => {
const request = createRequest({
waitForSelectedSources: vi.fn().mockRejectedValue(new Error('offline')),
});
const runtime = createRuntime(request);
await runtime.dispatch({ type: 'initialize' });
await runtime.dispatch({
type: 'addContextItem',
item: { kind: 'doc', docId: 'doc-2' },
});
await runtime.dispatch({
type: 'addContextItem',
item: { kind: 'blob', blobId: 'blob-1' },
});
await runtime.dispatch({
type: 'removeContextItem',
item: { kind: 'doc', docId: 'doc-2' },
});
(
request.context.getContextDocsAndFiles as ReturnType<typeof vi.fn>
).mockResolvedValue({
blobs: [{ blobId: 'blob-1', status: 'finished' }],
});
await runtime.dispatch({ type: 'pollContext' });
expect(request.context.createContext).toHaveBeenCalledTimes(1);
expect(request.context.addContextDoc).toHaveBeenCalledWith({
contextId: 'context-1',
docId: 'doc-2',
});
expect(request.context.removeContextDoc).toHaveBeenCalledWith({
contextId: 'context-1',
docId: 'doc-2',
});
expect(runtime.getSnapshot().composer.context.items).toEqual([
{ kind: 'blob', blobId: 'blob-1', state: 'finished' },
]);
});
test('loadContext restores existing session context without creating a new context', async () => {
const request = createRequest();
(
request.context.getContextId as ReturnType<typeof vi.fn>
).mockResolvedValue('context-1');
(
request.context.getContextDocsAndFiles as ReturnType<typeof vi.fn>
).mockResolvedValue({
docs: [{ id: 'doc-2', status: 'finished', createdAt: 2 }],
files: [
{
id: 'file-1',
blobId: 'blob-file-1',
name: 'note.pdf',
status: 'processing',
createdAt: 1,
},
],
tags: [
{
id: 'tag-1',
docs: [{ id: 'tag-doc', status: 'failed' }],
createdAt: 3,
},
],
collections: [],
blobs: [],
});
const runtime = createRuntime(request);
await runtime.dispatch({
type: 'openSessionObject',
session: session(),
type: 'addScopeSelector',
item: { kind: 'collection', collectionId: 'collection-1', docIds: [] },
});
await runtime.dispatch({ type: 'loadContext' });
await runtime.dispatch({ type: 'send', input: 'question' });
expect(request.context.createContext).not.toHaveBeenCalled();
expect(runtime.getSnapshot().composer.context.contextId).toBe('context-1');
expect(runtime.getSnapshot().composer.context.items).toEqual([
expect.objectContaining({
kind: 'file',
fileId: 'file-1',
blobId: 'blob-file-1',
state: 'processing',
}),
{ kind: 'doc', docId: 'doc-2', state: 'finished', createdAt: 2 },
{
kind: 'tag',
tagId: 'tag-1',
docIds: ['tag-doc'],
state: 'finished',
createdAt: 3,
tooltip: undefined,
},
]);
expect(runtime.getSnapshot().composer.context.embeddingCount).toEqual({
finished: 1,
processing: 1,
failed: 1,
});
});
test('pollEmbeddingStatus updates composer embedding completion state', async () => {
const request = createRequest();
(request.context.pollEmbeddingStatus as ReturnType<typeof vi.fn>)
.mockImplementationOnce(async (_workspaceId, onPoll) => {
onPoll({ embedded: 1, total: 2 });
})
.mockImplementationOnce(async (_workspaceId, onPoll) => {
onPoll({ embedded: 2, total: 2 });
});
const runtime = createRuntime(request);
await runtime.dispatch({ type: 'pollEmbeddingStatus' });
expect(runtime.getSnapshot().composer.context.embeddingCompleted).toBe(
false
expect(request.executeAction).not.toHaveBeenCalled();
expect(runtime.getSnapshot().status).toBe('error');
expect(runtime.getSnapshot().composer.scopeSelection.items).toHaveLength(1);
expect(runtime.getSnapshot().composer.scopeSelection.error?.message).toBe(
'offline'
);
await runtime.dispatch({ type: 'pollEmbeddingStatus' });
expect(runtime.getSnapshot().composer.context.embeddingCompleted).toBe(
true
);
});
test('startContextPolling owns context polling lifecycle', async () => {
const request = createRequest();
(
request.context.getContextDocsAndFiles as ReturnType<typeof vi.fn>
).mockResolvedValue({
docs: [{ docId: 'doc-2', status: 'finished' }],
});
(
request.context.getContextId as ReturnType<typeof vi.fn>
).mockResolvedValue('context-1');
const runtime = createRuntime(request);
await runtime.dispatch({
type: 'openSessionObject',
session: session(),
});
await runtime.dispatch({ type: 'loadContext' });
await runtime.dispatch({ type: 'startContextPolling' });
await waitUntil(() => {
expect(request.context.getContextDocsAndFiles).toHaveBeenCalledTimes(2);
});
expect(runtime.getSnapshot().composer.context.polling).toBe(false);
expect(runtime.getSnapshot().composer.context.embeddingCount).toEqual({
finished: 1,
processing: 0,
failed: 0,
});
});
test('fork strategy creates child session from parent without doc tab restrictions', async () => {
packages/frontend/core/src/blocksuite/ai/runtime/chat/runtime.ts
+145 -481
View File
@@ -4,9 +4,9 @@ import type { AIRequestService } from '../request';
import type { AIChatAction, AIChatSendOptions } from './actions';
import type { AIChatSessionStrategy } from './session-strategy';
import {
type AIChatContextItem,
type AIChatMessage,
type AIChatScope,
type AIChatScopeSelector,
type AIChatSnapshot,
type AIChatStatus,
type AIChatTab,
@@ -21,42 +21,11 @@ type RuntimeOptions = {
strategy: AIChatSessionStrategy;
};
type ContextStatus = 'finished' | 'processing' | 'failed';
type ContextObject = {
id?: string;
docId?: string;
blobId?: string;
name?: string;
status?: ContextStatus;
error?: string | null;
createdAt?: number | null;
docs?: ContextObject[];
};
type ContextData = {
docs?: ContextObject[];
files?: ContextObject[];
tags?: ContextObject[];
collections?: ContextObject[];
blobs?: ContextObject[];
};
type EmbeddingStatus = {
embedded: number;
total: number;
};
const CONTEXT_POLLING_INTERVAL = 10000;
export class AIChatRuntime {
private readonly listeners = new Set<() => void>();
private requestSeq = 0;
private historyRequestSeq = 0;
private contextRequestSeq = 0;
private streamAbortController: AbortController | null = null;
private contextPollingAbortController: AbortController | null = null;
private embeddingStatusAbortController: AbortController | null = null;
private createSessionPromiseKey: string | null = null;
private createSessionPromise: Promise<
CopilotChatHistoryFragment | null | undefined
@@ -100,8 +69,6 @@ export class AIChatRuntime {
this.createSessionPromise = null;
this.createSessionPromiseKey = null;
this.streamAbortController?.abort();
this.contextPollingAbortController?.abort();
this.embeddingStatusAbortController?.abort();
this.listeners.clear();
}
@@ -170,26 +137,32 @@ export class AIChatRuntime {
),
});
return;
case 'addContextItem':
await this.addContextItem(action.item);
case 'addScopeSelector':
this.addScopeSelector(action.item);
return;
case 'removeContextItem':
await this.removeContextItem(action.item);
case 'removeScopeSelector':
this.removeScopeSelector(action.item);
return;
case 'loadContext':
await this.loadContext();
case 'addFocusSelector':
this.updateComposer({
focus: {
items: this.mergeSelectors(
this.snapshot.composer.focus.items,
action.item
),
},
});
return;
case 'startContextPolling':
this.startContextPolling();
return;
case 'stopContextPolling':
this.stopContextPolling();
return;
case 'pollContext':
await this.pollContext();
return;
case 'pollEmbeddingStatus':
this.pollEmbeddingStatus();
case 'removeFocusSelector':
this.updateComposer({
focus: {
items: this.snapshot.composer.focus.items.filter(
item =>
this.getScopeSelectorKey(item) !==
this.getScopeSelectorKey(action.item)
),
},
});
return;
}
}
@@ -344,6 +317,16 @@ export class AIChatRuntime {
return this.snapshot.messages.findLast(message => message.role === 'user');
}
private async activateEditorContext(sessionId: string) {
try {
await this.options.request.activateEditor?.(sessionId);
return this.options.request.getActiveEditorContext();
} catch (error) {
console.error(error);
return undefined;
}
}
private async send(options: AIChatSendOptions, retryExisting = false) {
const content = options.input || this.snapshot.composer.text;
if (!content.trim() || !this.snapshot.uiPolicy.canSend) return;
@@ -374,6 +357,37 @@ export class AIChatRuntime {
if (!this.snapshot.activeSessionId) {
this.openSessionObject(session, true);
}
const liveEditorContext = await this.activateEditorContext(
session.sessionId
);
const scopeSelectors = this.snapshot.composer.scopeSelection.items
.map(item => this.selectorInput(item))
.filter(selector => selector !== null);
const focusSelectors = this.snapshot.composer.focus.items
.map(item => this.selectorInput(item))
.filter(selector => selector !== null);
if (scopeSelectors.length || focusSelectors.length) {
this.updateScopeSelection({ syncing: true, error: null });
const selectedDocIds = [
...this.snapshot.composer.scopeSelection.items,
...this.snapshot.composer.focus.items,
].flatMap(item => {
switch (item.kind) {
case 'doc':
return [item.docId];
case 'tag':
case 'collection':
return item.docIds;
default:
return [];
}
});
await this.options.request.waitForSelectedSources([
...new Set(selectedDocIds),
]);
if (seq !== this.requestSeq) return;
this.updateScopeSelection({ syncing: false });
}
const stream = (await this.options.request.executeAction('chat', {
workspaceId: this.snapshot.scope.workspaceId,
@@ -384,8 +398,13 @@ export class AIChatRuntime {
sessionId: session.sessionId,
input: content,
contexts: options.contexts,
attachments: options.attachments ?? this.snapshot.composer.attachments,
contextId: this.snapshot.composer.context.contextId,
scopeSelectors,
focusSelectors,
liveEditorContext,
attachments: [
...this.snapshot.composer.attachments,
...(options.attachments ?? []),
],
reasoning: options.reasoning ?? this.snapshot.composer.reasoning,
toolsConfig: options.toolsConfig ?? this.snapshot.composer.toolsConfig,
routeTargetId:
@@ -410,13 +429,21 @@ export class AIChatRuntime {
...this.snapshot.composer,
text: '',
attachments: [],
scopeSelection: {
...this.snapshot.composer.scopeSelection,
items: [],
syncing: false,
error: null,
},
},
});
await this.refreshLastMessageId(session.sessionId).catch(console.error);
await this.bindActiveSessionToDoc().catch(console.error);
} catch (error) {
if (seq !== this.requestSeq) return;
this.commit({ status: 'error', error: this.toError(error) });
const resolved = this.toError(error);
this.updateScopeSelection({ syncing: false, error: resolved });
this.commit({ status: 'error', error: resolved });
}
}
@@ -440,9 +467,13 @@ export class AIChatRuntime {
messages: this.resetLastAssistantMessage(this.snapshot.messages),
});
try {
const liveEditorContext = await this.activateEditorContext(
this.snapshot.activeSessionId
);
const stream = (await this.options.request.executeAction('chat', {
workspaceId: this.snapshot.scope.workspaceId,
sessionId: this.snapshot.activeSessionId,
liveEditorContext,
retry: true,
stream: true,
signal: this.streamAbortController.signal,
@@ -525,466 +556,99 @@ export class AIChatRuntime {
});
}
private updateContextState(
patch: Partial<AIChatSnapshot['composer']['context']>
private updateScopeSelection(
patch: Partial<AIChatSnapshot['composer']['scopeSelection']>
) {
this.updateComposer({
context: {
...this.snapshot.composer.context,
scopeSelection: {
...this.snapshot.composer.scopeSelection,
...patch,
},
});
}
private async getContextId() {
const createdSession = this.snapshot.activeSessionId
? null
: await this.ensureSession();
if (createdSession) {
this.openSessionObject(createdSession, true);
private addScopeSelector(item: AIChatScopeSelector) {
if (item.kind === 'file') {
this.updateComposer({
attachments: [...this.snapshot.composer.attachments, item.file],
});
} else if (item.kind === 'blob') {
this.updateComposer({
attachments: [...this.snapshot.composer.attachments, item.blobId],
});
}
const sessionId =
this.snapshot.activeSessionId ?? createdSession?.sessionId ?? null;
if (!sessionId) return null;
this.updateScopeSelection({
error: null,
items: this.mergeSelectors(
this.snapshot.composer.scopeSelection.items,
item
),
});
}
const cached = this.snapshot.composer.context.contextId;
if (cached) return cached;
const { workspaceId } = this.snapshot.scope;
const existing = await this.options.request.context.getContextId(
workspaceId,
sessionId
private removeScopeSelector(item: AIChatScopeSelector) {
const key = this.getScopeSelectorKey(item);
const attachments = this.snapshot.composer.attachments.filter(
attachment => {
if (item.kind === 'file') return attachment !== item.file;
if (item.kind === 'blob') return attachment !== item.blobId;
return true;
}
);
const contextId =
existing ??
(await this.options.request.context.createContext(
workspaceId,
sessionId
));
this.updateContextState({ contextId });
return contextId;
}
private async addContextItem(item: AIChatContextItem) {
const seq = ++this.contextRequestSeq;
this.updateContextState({ loading: true, error: null });
try {
const contextId = await this.getContextId();
if (!contextId) throw new Error('Context not found');
const nextItem = await this.persistContextItem(contextId, item);
if (seq !== this.contextRequestSeq) return;
this.updateContextState({
loading: false,
items: [...this.snapshot.composer.context.items, nextItem],
});
} catch (error) {
if (seq !== this.contextRequestSeq) return;
this.updateContextState({ loading: false, error: this.toError(error) });
}
}
private async removeContextItem(item: AIChatContextItem) {
const seq = ++this.contextRequestSeq;
this.updateContextState({ loading: true, error: null });
try {
const contextId = this.snapshot.composer.context.contextId;
if (contextId) {
await this.deleteContextItem(contextId, item);
}
if (seq !== this.contextRequestSeq) return;
this.updateContextState({
loading: false,
items: this.snapshot.composer.context.items.filter(
existing =>
this.getContextItemKey(existing) !== this.getContextItemKey(item)
this.updateComposer({
attachments,
scopeSelection: {
...this.snapshot.composer.scopeSelection,
items: this.snapshot.composer.scopeSelection.items.filter(
existing => this.getScopeSelectorKey(existing) !== key
),
});
} catch (error) {
if (seq !== this.contextRequestSeq) return;
this.updateContextState({ loading: false, error: this.toError(error) });
}
}
private async pollContext() {
const seq = ++this.contextRequestSeq;
const sessionId = this.snapshot.activeSessionId;
const contextId = this.snapshot.composer.context.contextId;
if (!sessionId || !contextId) return;
this.updateContextState({ polling: true, error: null });
try {
const context = await this.options.request.context.getContextDocsAndFiles(
this.snapshot.scope.workspaceId,
sessionId,
contextId
);
if (seq !== this.contextRequestSeq) return;
this.updateContextState({
polling: false,
items: this.mergePolledContextItems(context),
embeddingCount: this.getContextEmbeddingCount(context),
});
} catch (error) {
if (seq !== this.contextRequestSeq) return;
this.updateContextState({ polling: false, error: this.toError(error) });
}
}
private startContextPolling() {
this.stopContextPolling();
this.contextPollingAbortController = new AbortController();
const signal = this.contextPollingAbortController.signal;
void this.pollContextUntilIdle(signal).catch(error => {
if (signal.aborted) return;
this.updateContextState({ polling: false, error: this.toError(error) });
},
});
}
private stopContextPolling() {
this.contextPollingAbortController?.abort();
this.contextPollingAbortController = null;
private mergeSelectors(
items: AIChatScopeSelector[],
item: AIChatScopeSelector
) {
const key = this.getScopeSelectorKey(item);
return [
...items.filter(existing => this.getScopeSelectorKey(existing) !== key),
item,
];
}
private async pollContextUntilIdle(signal: AbortSignal) {
while (!signal.aborted) {
await this.pollContext();
if (signal.aborted) return;
if (this.snapshot.composer.context.embeddingCount.processing === 0) {
this.stopContextPolling();
return;
}
await this.waitForContextPollingInterval(signal);
}
}
private waitForContextPollingInterval(signal: AbortSignal) {
return new Promise<void>(resolve => {
const timeout = setTimeout(resolve, CONTEXT_POLLING_INTERVAL);
signal.addEventListener(
'abort',
() => {
clearTimeout(timeout);
resolve();
},
{ once: true }
);
});
}
private async loadContext() {
const seq = ++this.contextRequestSeq;
const sessionId = this.snapshot.activeSessionId;
if (!sessionId) return;
this.updateContextState({ loading: true, error: null });
try {
const { workspaceId } = this.snapshot.scope;
const contextId = await this.options.request.context.getContextId(
workspaceId,
sessionId
);
if (!contextId) {
if (seq !== this.contextRequestSeq) return;
this.updateContextState({
contextId: null,
items: [],
loading: false,
embeddingCount: { finished: 0, processing: 0, failed: 0 },
});
return;
}
const context = await this.options.request.context.getContextDocsAndFiles(
workspaceId,
sessionId,
contextId
);
if (seq !== this.contextRequestSeq) return;
this.updateContextState({
contextId,
loading: false,
items: this.contextDataToItems(context),
embeddingCount: this.getContextEmbeddingCount(context),
});
} catch (error) {
if (seq !== this.contextRequestSeq) return;
this.updateContextState({ loading: false, error: this.toError(error) });
}
}
private pollEmbeddingStatus() {
this.embeddingStatusAbortController?.abort();
this.embeddingStatusAbortController = new AbortController();
const signal = this.embeddingStatusAbortController.signal;
void this.options.request.context
.pollEmbeddingStatus(
this.snapshot.scope.workspaceId,
status => {
if (signal.aborted) return;
this.updateContextState({
embeddingCompleted: this.isEmbeddingCompleted(status),
});
},
signal
)
.catch(error => {
if (signal.aborted) return;
this.updateContextState({
embeddingCompleted: false,
error: this.toError(error),
});
});
}
private async persistContextItem(
contextId: string,
item: AIChatContextItem
): Promise<AIChatContextItem> {
private selectorInput(item: AIChatScopeSelector) {
switch (item.kind) {
case 'doc':
await this.options.request.context.addContextDoc({
contextId,
docId: item.docId,
});
return item;
case 'file': {
const file = await this.options.request.context.addContextFile(
item.file,
{ contextId }
);
return {
...item,
fileId: file.id,
blobId: file.blobId ?? item.blobId,
state: file.status,
createdAt: file.createdAt,
tooltip: file.error ?? undefined,
};
}
return { kind: 'document', id: item.docId, name: item.name };
case 'tag':
await this.options.request.context.addContextTag({
contextId,
tagId: item.tagId,
docIds: item.docIds,
});
return item;
return { kind: 'tag', id: item.tagId, name: item.name };
case 'collection':
await this.options.request.context.addContextCollection({
contextId,
collectionId: item.collectionId,
docIds: item.docIds,
});
return item;
case 'blob': {
const blob = await this.options.request.context.addContextBlob({
contextId,
blobId: item.blobId,
});
return {
...item,
state: blob.status || item.state,
createdAt: blob.createdAt,
};
}
}
}
private deleteContextItem(contextId: string, item: AIChatContextItem) {
switch (item.kind) {
case 'doc':
return this.options.request.context.removeContextDoc({
contextId,
docId: item.docId,
});
return { kind: 'collection', id: item.collectionId, name: item.name };
case 'favorite':
return { kind: 'favorite', id: item.favoriteId, name: item.name };
case 'file':
if (!item.fileId) return Promise.resolve();
return this.options.request.context.removeContextFile({
contextId,
fileId: item.fileId,
});
case 'tag':
return this.options.request.context.removeContextTag({
contextId,
tagId: item.tagId,
});
case 'collection':
return this.options.request.context.removeContextCollection({
contextId,
collectionId: item.collectionId,
});
case 'blob':
return this.options.request.context.removeContextBlob({
contextId,
blobId: item.blobId,
});
return null;
}
}
private mergePolledContextItems(context: unknown) {
if (!context || typeof context !== 'object') {
return this.snapshot.composer.context.items;
}
const data = context as ContextData;
const docs = [
...(data.docs ?? []),
...(data.tags ?? []).flatMap(tag => tag.docs ?? []),
...(data.collections ?? []).flatMap(collection => collection.docs ?? []),
];
return this.snapshot.composer.context.items.map(item => {
if (item.kind === 'doc') {
const doc = docs.find(
candidate =>
candidate.docId === item.docId || candidate.id === item.docId
);
return doc?.status
? { ...item, state: doc.status, tooltip: doc.error ?? undefined }
: item;
}
if (item.kind === 'file') {
const file = data.files?.find(
candidate =>
candidate.id === item.fileId ||
candidate.blobId === item.blobId ||
candidate.blobId === item.fileId
);
return file?.status
? { ...item, state: file.status, tooltip: file.error ?? undefined }
: item;
}
if (item.kind === 'blob') {
const blob = data.blobs?.find(
candidate =>
candidate.blobId === item.blobId || candidate.id === item.blobId
);
return blob?.status
? { ...item, state: blob.status, tooltip: blob.error ?? undefined }
: item;
}
return item;
});
}
private contextDataToItems(context: unknown): AIChatContextItem[] {
if (!context || typeof context !== 'object') return [];
const data = context as ContextData;
const items: AIChatContextItem[] = [
...(data.docs ?? []).flatMap(doc =>
doc.id
? [
{
kind: 'doc' as const,
docId: doc.id,
state: doc.status,
createdAt: doc.createdAt ?? undefined,
tooltip: doc.error ?? undefined,
},
]
: []
),
...(data.files ?? []).flatMap(file =>
file.id && file.name
? [
{
kind: 'file' as const,
file: new File([], file.name),
fileId: file.id,
blobId: file.blobId,
state: file.status,
createdAt: file.createdAt ?? undefined,
tooltip: file.error ?? undefined,
},
]
: []
),
...(data.tags ?? []).flatMap(tag =>
tag.id
? [
{
kind: 'tag' as const,
tagId: tag.id,
docIds: (tag.docs ?? []).flatMap(doc =>
doc.id ? [doc.id] : []
),
state: 'finished',
createdAt: tag.createdAt ?? undefined,
tooltip: tag.error ?? undefined,
},
]
: []
),
...(data.collections ?? []).flatMap(collection =>
collection.id
? [
{
kind: 'collection' as const,
collectionId: collection.id,
docIds: (collection.docs ?? []).flatMap(doc =>
doc.id ? [doc.id] : []
),
state: 'finished',
createdAt: collection.createdAt ?? undefined,
tooltip: collection.error ?? undefined,
},
]
: []
),
...(data.blobs ?? []).flatMap(blob =>
(blob.blobId ?? blob.id)
? [
{
kind: 'blob' as const,
blobId: blob.blobId ?? blob.id ?? '',
state: blob.status,
createdAt: blob.createdAt ?? undefined,
tooltip: blob.error ?? undefined,
},
]
: []
),
];
return items.sort((a, b) => (a.createdAt ?? 0) - (b.createdAt ?? 0));
}
private getContextEmbeddingCount(
context: unknown
): AIChatSnapshot['composer']['context']['embeddingCount'] {
const count = { finished: 0, processing: 0, failed: 0 };
if (!context || typeof context !== 'object') return count;
const data = context as ContextData;
const docs = [
...(data.docs ?? []),
...(data.tags ?? []).flatMap(tag => tag.docs ?? []),
...(data.collections ?? []).flatMap(collection => collection.docs ?? []),
];
for (const item of [
...docs,
...(data.files ?? []),
...(data.blobs ?? []),
]) {
if (item.status) count[item.status]++;
}
return count;
}
private isEmbeddingCompleted(status: unknown) {
if (!status || typeof status !== 'object') return false;
const { embedded, total } = status as EmbeddingStatus;
return embedded === total;
}
private getContextItemKey(item: AIChatContextItem) {
private getScopeSelectorKey(item: AIChatScopeSelector) {
switch (item.kind) {
case 'doc':
return `doc:${item.docId}`;
case 'file':
return `file:${item.fileId ?? item.file.name}`;
return `document:${item.docId}`;
case 'file': {
const id = item.fileId ?? item.blobId;
return id ? `file:${id}` : item.file;
}
case 'tag':
return `tag:${item.tagId}`;
case 'collection':
return `collection:${item.collectionId}`;
case 'blob':
return `blob:${item.blobId}`;
case 'favorite':
return `favorite:${item.favoriteId}`;
}
}
packages/frontend/core/src/blocksuite/ai/runtime/chat/state.ts
+33 -19
View File
@@ -45,6 +45,18 @@ export type AIChatMessage = {
userId?: string;
userName?: string;
avatarUrl?: string;
scopeSnapshot?: AIChatScopeReceipt | null;
};
export type AIChatScopeReceipt = {
resolvedAt: string;
selectors: Array<{
kind: 'document' | 'tag' | 'collection' | 'favorite' | 'artifact';
id: string;
name?: string;
}>;
requiredDocIds: string[];
requiredArtifactIds: string[];
};
export type AIChatTab =
@@ -79,10 +91,11 @@ export type AIChatHistoryGroups = {
error: Error | null;
};
export type AIChatContextItem =
export type AIChatScopeSelector =
| {
kind: 'doc';
docId: string;
name?: string;
state?: string;
createdAt?: number;
tooltip?: string;
@@ -99,6 +112,7 @@ export type AIChatContextItem =
| {
kind: 'tag';
tagId: string;
name?: string;
docIds: string[];
state?: string;
createdAt?: number;
@@ -107,6 +121,7 @@ export type AIChatContextItem =
| {
kind: 'collection';
collectionId: string;
name?: string;
docIds: string[];
state?: string;
createdAt?: number;
@@ -118,22 +133,28 @@ export type AIChatContextItem =
state?: string;
createdAt?: number;
tooltip?: string;
}
| {
kind: 'favorite';
favoriteId: string;
name?: string;
};
export type AIChatContextState = {
contextId: string | null;
items: AIChatContextItem[];
loading: boolean;
polling: boolean;
export type AIChatScopeSelection = {
items: AIChatScopeSelector[];
syncing: boolean;
error: Error | null;
embeddingCompleted: boolean;
embeddingCount: Record<'finished' | 'processing' | 'failed', number>;
};
export type AIChatFocusState = {
items: AIChatScopeSelector[];
};
export type AIChatComposerState = {
text: string;
attachments: (string | Blob | File)[];
context: AIChatContextState;
scopeSelection: AIChatScopeSelection;
focus: AIChatFocusState;
reasoning: boolean;
toolsConfig?: AIToolsConfig;
routeTargetId?: string;
@@ -172,19 +193,12 @@ export function createInitialComposerState(): AIChatComposerState {
return {
text: '',
attachments: [],
context: {
contextId: null,
scopeSelection: {
items: [],
loading: false,
polling: false,
syncing: false,
error: null,
embeddingCompleted: false,
embeddingCount: {
finished: 0,
processing: 0,
failed: 0,
},
},
focus: { items: [] },
reasoning: false,
};
}
packages/frontend/core/src/blocksuite/ai/runtime/frontend/delegated-editor-host.ts
+388
View File
@@ -0,0 +1,388 @@
import type { NbstoreService } from '@affine/core/modules/storage';
import type {
DelegatedToolCancel,
DelegatedToolName,
DelegatedToolRequest,
} from '@affine/realtime';
import type { EditorHost } from '@blocksuite/affine/std';
import { GfxControllerIdentifier } from '@blocksuite/affine/std/gfx';
import type { Subscription } from 'rxjs';
import {
getLiveEditorMode,
getLiveSelectionIds,
lightEditorContext,
readEditorState,
readNodes,
readSelection,
snapshotDocument,
} from './live-projection';
type Realtime = Pick<NbstoreService['realtime'], 'request' | 'subscribe'>;
type HostOptions = {
realtime: Realtime;
host: EditorHost;
sessionId: string;
workspaceId: string;
docId: string;
};
const capabilities: DelegatedToolName[] = [
'frontend_get_editor_state',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
];
const STATE_UPSERT_DELAY_MS = 150;
const LEASE_RENEWAL_MS = 10_000;
export class DelegatedEditorHost {
readonly clientId = crypto.randomUUID();
private editorStateId = crypto.randomUUID();
private subscription?: Subscription;
private heartbeat?: ReturnType<typeof setTimeout>;
private stateUpsert?: ReturnType<typeof setTimeout>;
private blockSubscription?: { unsubscribe(): void };
private selectionSubscription?: { unsubscribe(): void };
private viewportSubscriptions: Array<{ unsubscribe(): void }> = [];
private upsertInFlight?: Promise<void>;
private upsertRequested = false;
private disposed = false;
private selectionSignature = '';
private metadataSignature = '';
private publishedEditorStateId?: string;
private readonly inFlight = new Map<
string,
{
identity: ReturnType<DelegatedEditorHost['identity']>;
abort: AbortController;
}
>();
private readonly focusChanged = () => {
void this.requestUpsert().catch(console.error);
};
constructor(private readonly options: HostOptions) {}
context() {
return JSON.stringify({
workspace_id: this.options.workspaceId,
doc_id: this.options.docId,
session_id: this.options.sessionId,
...lightEditorContext(this.options.host, this.editorStateId),
});
}
async start() {
this.disposed = false;
this.subscription = this.options.realtime
.subscribe('copilot.delegated.tool.requested', {
clientId: this.clientId,
})
.subscribe(event => {
if (event.type === 'request') {
void this.respond(event).catch(console.error);
} else if (event.type === 'cancel') {
this.cancel(event);
}
});
const changed = () => {
this.editorStateId = crypto.randomUUID();
this.scheduleStateUpsert();
};
this.blockSubscription =
this.options.host.store.slots.blockUpdated.subscribe(changed);
this.selectionSignature = this.getSelectionSignature();
this.metadataSignature = this.getMetadataSignature();
this.selectionSubscription =
this.options.host.selection.slots.changed.subscribe(() => {
const signature = this.getSelectionSignature();
if (signature === this.selectionSignature) return;
this.selectionSignature = signature;
changed();
});
const viewport = this.options.host.std.get(
GfxControllerIdentifier
).viewport;
const viewportChanged = () => {
if (getLiveEditorMode(this.options.host) === 'edgeless') changed();
};
this.viewportSubscriptions = [
viewport.viewportUpdated.subscribe(viewportChanged),
viewport.sizeUpdated.subscribe(viewportChanged),
];
window.addEventListener('focus', this.focusChanged);
window.addEventListener('blur', this.focusChanged);
document.addEventListener('visibilitychange', this.focusChanged);
await this.sync();
}
dispose() {
this.disposed = true;
this.subscription?.unsubscribe();
this.blockSubscription?.unsubscribe();
this.selectionSubscription?.unsubscribe();
this.viewportSubscriptions.forEach(subscription =>
subscription.unsubscribe()
);
this.viewportSubscriptions = [];
window.removeEventListener('focus', this.focusChanged);
window.removeEventListener('blur', this.focusChanged);
document.removeEventListener('visibilitychange', this.focusChanged);
for (const request of this.inFlight.values()) request.abort.abort();
this.inFlight.clear();
if (this.heartbeat) clearTimeout(this.heartbeat);
if (this.stateUpsert) clearTimeout(this.stateUpsert);
const release = () =>
this.options.realtime.request('copilot.delegated.editor.release', {
clientId: this.clientId,
editorStateId: this.publishedEditorStateId ?? this.editorStateId,
});
void (this.upsertInFlight ?? Promise.resolve())
.catch(() => {})
.then(release)
.catch(console.error);
}
async sync() {
while (!this.disposed) {
this.refreshMetadataState();
if (this.stateUpsert) {
clearTimeout(this.stateUpsert);
this.stateUpsert = undefined;
}
if (this.upsertInFlight) {
await this.upsertInFlight;
continue;
}
if (this.publishedEditorStateId === this.editorStateId) return;
await this.requestUpsert();
}
}
private getSelectionSignature() {
return JSON.stringify([
getLiveEditorMode(this.options.host),
getLiveSelectionIds(this.options.host),
]);
}
private getMetadataSignature() {
return JSON.stringify([
getLiveEditorMode(this.options.host),
this.options.host.store.readonly$.value,
]);
}
private refreshMetadataState(scheduleUpsert = false) {
const signature = this.getMetadataSignature();
const changed =
!!this.metadataSignature && signature !== this.metadataSignature;
if (changed) {
this.editorStateId = crypto.randomUUID();
if (scheduleUpsert) this.scheduleStateUpsert();
}
this.metadataSignature = signature;
return changed;
}
private scheduleStateUpsert() {
if (this.stateUpsert || this.disposed) return;
this.stateUpsert = setTimeout(() => {
this.stateUpsert = undefined;
void this.requestUpsert().catch(console.error);
}, STATE_UPSERT_DELAY_MS);
}
private scheduleHeartbeat() {
if (this.heartbeat) clearTimeout(this.heartbeat);
if (this.disposed) return;
this.heartbeat = setTimeout(() => {
this.heartbeat = undefined;
void this.requestUpsert().catch(console.error);
}, LEASE_RENEWAL_MS);
}
private requestUpsert() {
if (this.disposed) return Promise.resolve();
this.refreshMetadataState();
this.upsertRequested = true;
if (this.upsertInFlight) return this.upsertInFlight;
this.upsertInFlight = (async () => {
while (this.upsertRequested && !this.disposed) {
this.upsertRequested = false;
if (this.stateUpsert) {
clearTimeout(this.stateUpsert);
this.stateUpsert = undefined;
}
const editorStateId = this.editorStateId;
try {
await this.options.realtime.request(
'copilot.delegated.editor.upsert',
{
clientId: this.clientId,
sessionId: this.options.sessionId,
workspaceId: this.options.workspaceId,
docId: this.options.docId,
editorStateId,
mode: getLiveEditorMode(this.options.host),
readonly: this.options.host.store.readonly$.value,
focused:
document.visibilityState === 'visible' && document.hasFocus(),
capabilities,
}
);
this.publishedEditorStateId = editorStateId;
} finally {
this.scheduleHeartbeat();
}
}
})().finally(() => {
this.upsertInFlight = undefined;
if (this.upsertRequested && !this.disposed && !this.stateUpsert) {
this.stateUpsert = setTimeout(() => {
this.stateUpsert = undefined;
void this.requestUpsert().catch(console.error);
}, STATE_UPSERT_DELAY_MS);
}
});
return this.upsertInFlight;
}
private async respond(request: DelegatedToolRequest) {
if (
request.sessionId !== this.options.sessionId ||
request.workspaceId !== this.options.workspaceId ||
request.docId !== this.options.docId ||
request.clientId !== this.clientId
) {
await this.sendError(request, 'EDITOR_CONTEXT_CHANGED');
return;
}
this.refreshMetadataState(true);
if (request.deadlineAt <= Date.now()) {
await this.sendError(request, 'FRONTEND_TIMEOUT');
return;
}
if (request.editorStateId !== this.editorStateId) {
await this.sendError(request, 'EDITOR_STATE_CHANGED');
return;
}
const identity = this.identity(request);
const abort = new AbortController();
this.inFlight.set(request.requestId, { identity, abort });
try {
const result = await Promise.resolve(this.execute(request));
if (abort.signal.aborted) return;
this.refreshMetadataState(true);
if (request.editorStateId !== this.editorStateId) {
await this.sendError(request, 'EDITOR_STATE_CHANGED');
return;
}
const resultError = this.resultError(result);
if (resultError) {
await this.options.realtime.request('copilot.delegated.tool.respond', {
...this.identity(request),
error: resultError,
});
return;
}
await this.options.realtime.request('copilot.delegated.tool.respond', {
...this.identity(request),
result,
});
} catch {
if (abort.signal.aborted) return;
await this.sendError(request, 'FRONTEND_READ_FAILED');
} finally {
this.inFlight.delete(request.requestId);
}
}
private cancel(event: DelegatedToolCancel) {
const request = this.inFlight.get(event.requestId);
if (request && this.sameIdentity(request.identity, event)) {
request.abort.abort();
this.inFlight.delete(event.requestId);
}
}
private execute(request: DelegatedToolRequest) {
switch (request.tool) {
case 'frontend_get_editor_state':
return readEditorState(this.options.host, this.editorStateId);
case 'frontend_read_selection':
return readSelection(
this.options.host,
this.editorStateId,
request.args
);
case 'frontend_read_nodes':
return readNodes(this.options.host, this.editorStateId, request.args);
case 'frontend_snapshot_document':
return snapshotDocument(
this.options.host,
this.editorStateId,
request.args
);
}
}
private identity(request: DelegatedToolRequest) {
const {
requestId,
runId,
toolCallId,
sessionId,
workspaceId,
docId,
clientId,
editorStateId,
} = request;
return {
requestId,
runId,
toolCallId,
sessionId,
workspaceId,
docId,
clientId,
editorStateId,
};
}
private sameIdentity(
expected: ReturnType<DelegatedEditorHost['identity']>,
actual: ReturnType<DelegatedEditorHost['identity']>
) {
return Object.entries(expected).every(
([key, value]) => actual[key as keyof typeof actual] === value
);
}
private sendError(request: DelegatedToolRequest, code: string) {
return this.options.realtime.request('copilot.delegated.tool.respond', {
...this.identity(request),
error: {
code,
message: 'The focused editor changed before the read completed.',
retryable: true,
},
});
}
private resultError(result: unknown) {
if (!result || typeof result !== 'object' || !('error' in result)) {
return null;
}
const error = result.error;
if (!error || typeof error !== 'object' || !('code' in error)) return null;
return {
code: String(error.code),
message: 'The requested live editor view is not available.',
retryable: false,
};
}
}
packages/frontend/core/src/blocksuite/ai/runtime/frontend/index.ts
+1
View File
@@ -0,0 +1 @@
export { DelegatedEditorHost } from './delegated-editor-host';
packages/frontend/core/src/blocksuite/ai/runtime/frontend/live-projection.ts
+365
View File
@@ -0,0 +1,365 @@
import { DocModeProvider } from '@blocksuite/affine/shared/services';
import type { EditorHost } from '@blocksuite/affine/std';
import {
GfxControllerIdentifier,
type GfxModel,
isPrimitiveModel,
} from '@blocksuite/affine/std/gfx';
import type { BlockModel } from '@blocksuite/affine/store';
import { Bound } from '@blocksuite/global/gfx';
import {
getSelectedModels,
getSelectedTextContent,
} from '../../utils/selection-utils';
const textLimit = (value: string, limit: number) => ({
content: value.slice(0, limit),
truncated: value.length > limit,
});
const MAX_LOCATOR_IDS = 50;
const MAX_RELATION_IDS = 200;
function boundedInteger(value: unknown, fallback: number, maximum: number) {
return typeof value === 'number' && Number.isInteger(value) && value >= 0
? Math.min(value, maximum)
: fallback;
}
function boundsOf(element: GfxModel) {
const bounds = Bound.deserialize(element.xywh);
return {
x: bounds.x,
y: bounds.y,
width: bounds.w,
height: bounds.h,
};
}
function elementType(element: GfxModel) {
const type = isPrimitiveModel(element) ? element.type : element.flavour;
return type.startsWith('affine:') ? type.slice('affine:'.length) : type;
}
function record(value: unknown): Record<string, unknown> | undefined {
return value && typeof value === 'object'
? (value as Record<string, unknown>)
: undefined;
}
function elementProperties(element: GfxModel) {
const serialized =
record(isPrimitiveModel(element) ? element.serialize() : element.props) ??
{};
return {
...serialized,
...record(serialized.props),
};
}
function textValue(value: unknown) {
if (typeof value === 'string') return value;
const object = record(value);
if (!object || typeof object.toString !== 'function') return undefined;
const text = object.toString();
return text === '[object Object]' ? undefined : text;
}
function relations(element: GfxModel, props: Record<string, unknown>) {
const pickIds = (key: string) => {
const value = props[key];
if (Array.isArray(value)) {
return value.filter((item): item is string => typeof item === 'string');
}
const entries = record(value);
return entries
? Object.entries(entries)
.filter(([, included]) => included !== false && included != null)
.map(([id]) => id)
: undefined;
};
const source = record(props.source);
const target = record(props.target);
const group = element.group;
const frame = element.groups.find(
candidate => 'flavour' in candidate && candidate.flavour === 'affine:frame'
);
const serializedGroup = group
? record(isPrimitiveModel(group) ? group.serialize() : group.props)
: undefined;
const groupProps = serializedGroup
? {
...serializedGroup,
...record(serializedGroup.props),
}
: undefined;
const groupDetail = record(record(groupProps?.children)?.[element.id]);
const childIds = (
pickIds('childElementIds') ??
pickIds('children') ??
[]
).sort();
return {
frame_id: frame?.id,
child_ids: childIds.slice(0, MAX_RELATION_IDS),
child_ids_truncated: childIds.length > MAX_RELATION_IDS,
source_id: typeof source?.id === 'string' ? source.id : undefined,
target_id: typeof target?.id === 'string' ? target.id : undefined,
parent_id:
typeof groupDetail?.parent === 'string'
? groupDetail.parent
: group && group !== frame
? group.id
: undefined,
index:
typeof groupDetail?.index === 'string' ? groupDetail.index : undefined,
};
}
function projectElement(element: GfxModel, limit: number) {
const props = elementProperties(element);
const text = textValue(props.text) ?? textValue(props.label);
const title = textValue(props.title);
const pointCount = Array.isArray(props.points)
? props.points.length
: props.pointCount;
return {
id: element.id,
type: elementType(element),
bounds: boundsOf(element),
text: text ? textLimit(text, limit) : undefined,
title: title ? textLimit(title, limit) : undefined,
point_count:
typeof pointCount === 'number' && Number.isInteger(pointCount)
? pointCount
: undefined,
...relations(element, props),
};
}
function projectBlock(model: BlockModel, limit: number) {
const text = model.text?.toString() ?? '';
return {
id: model.id,
flavour: model.flavour,
text: textLimit(text, limit),
child_ids: model.children.slice(0, MAX_RELATION_IDS).map(child => child.id),
child_ids_truncated: model.children.length > MAX_RELATION_IDS,
};
}
function neighborhoodOf(models: BlockModel[], distance: number) {
const neighborhood = new Map<string, BlockModel>();
for (const model of models) {
const siblings = model.parent?.children ?? [];
const index = siblings.indexOf(model);
for (
let cursor = Math.max(0, index - distance);
cursor <= Math.min(siblings.length - 1, index + distance);
cursor++
) {
const sibling = siblings[cursor];
if (sibling) neighborhood.set(sibling.id, sibling);
}
}
return [...neighborhood.values()];
}
export function getLiveEditorMode(host: EditorHost) {
return host.std.get(DocModeProvider).getEditorMode() || 'page';
}
export function getLiveSelectionIds(host: EditorHost) {
if (getLiveEditorMode(host) === 'edgeless') {
return host.std
.get(GfxControllerIdentifier)
.selection.selectedElements.map(element => element.id);
}
return (getSelectedModels(host) ?? []).map(model => model.id);
}
export function readEditorState(host: EditorHost, editorStateId: string) {
const mode = getLiveEditorMode(host);
const selectionIds = getLiveSelectionIds(host);
return {
editor_state_id: editorStateId,
mode,
readonly: host.store.readonly$.value,
selection: {
kind: mode === 'edgeless' ? 'elements' : 'blocks',
ids: selectionIds.slice(0, MAX_LOCATOR_IDS),
truncated: selectionIds.length > MAX_LOCATOR_IDS,
},
capabilities: [
'frontend_get_editor_state',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
],
};
}
export function lightEditorContext(host: EditorHost, editorStateId: string) {
const state = readEditorState(host, editorStateId);
if (state.mode === 'edgeless') {
const preview = host.std
.get(GfxControllerIdentifier)
.selection.selectedElements.slice(0, 3)
.map(element => {
const projected = projectElement(element, 200);
return {
id: projected.id,
type: projected.type,
text: projected.text,
title: projected.title,
};
});
return { ...state, preview };
}
const preview = (getSelectedModels(host) ?? []).slice(0, 3).map(model => ({
id: model.id,
type: model.flavour,
text: textLimit(model.text?.toString() ?? '', 200),
}));
return { ...state, preview };
}
export async function readSelection(
host: EditorHost,
editorStateId: string,
args: Record<string, unknown>
) {
const limit = boundedInteger(args.limit, 10_000, 50_000);
const mode = getLiveEditorMode(host);
if (mode === 'edgeless') {
const selectedElements = host.std.get(GfxControllerIdentifier).selection
.selectedElements;
const elements = selectedElements
.slice(0, MAX_LOCATOR_IDS)
.map(element => projectElement(element, limit));
return {
editor_state_id: editorStateId,
mode,
elements,
truncated: selectedElements.length > MAX_LOCATOR_IDS,
};
}
const models = getSelectedModels(host) ?? [];
const neighborhood = boundedInteger(args.neighborhood, 0, 20);
if (args.format === 'structure') {
const blocks = models.slice(0, 50);
const nearby = neighborhoodOf(models, neighborhood).slice(0, 100);
return {
editor_state_id: editorStateId,
mode,
blocks: blocks.map(model => projectBlock(model, limit)),
neighborhood: nearby.map(model =>
projectBlock(model, Math.min(limit, 500))
),
truncated: blocks.length < models.length,
};
}
const format = args.format === 'text' ? 'plain-text' : 'markdown';
const content = await getSelectedTextContent(host, format);
return {
editor_state_id: editorStateId,
mode,
...textLimit(content, limit),
block_ids: models.slice(0, MAX_LOCATOR_IDS).map(model => model.id),
block_ids_truncated: models.length > MAX_LOCATOR_IDS,
neighborhood: neighborhoodOf(models, neighborhood)
.slice(0, 100)
.map(model => projectBlock(model, Math.min(limit, 500))),
};
}
export function readNodes(
host: EditorHost,
editorStateId: string,
args: Record<string, unknown>
) {
const limit = boundedInteger(args.limit, 10_000, 50_000);
const blockIds = Array.isArray(args.block_ids)
? args.block_ids
.filter((id): id is string => typeof id === 'string')
.slice(0, 50)
: [];
const elementIds = Array.isArray(args.element_ids)
? args.element_ids
.filter((id): id is string => typeof id === 'string')
.slice(0, 50)
: [];
const gfx = host.std.get(GfxControllerIdentifier);
return {
editor_state_id: editorStateId,
items: [
...blockIds.map(id => {
const block = host.store.getBlock(id)?.model;
return block
? { id, kind: 'block', value: projectBlock(block, limit) }
: { id, kind: 'block', error: { code: 'NODE_NOT_FOUND' } };
}),
...elementIds.map(id => {
const element = gfx.getElementById<GfxModel>(id);
return element
? { id, kind: 'element', value: projectElement(element, limit) }
: { id, kind: 'element', error: { code: 'NODE_NOT_FOUND' } };
}),
],
};
}
export function snapshotDocument(
host: EditorHost,
editorStateId: string,
args: Record<string, unknown>
) {
const limit = boundedInteger(args.limit, 50, 200);
const mode = getLiveEditorMode(host);
if (mode === 'edgeless') {
const gfx = host.std.get(GfxControllerIdentifier);
const bounds = gfx.viewport.viewportBounds;
const visible = gfx.gfxElements.filter(element =>
bounds.isIntersectWithBound(element.elementBound)
);
const elements = visible
.slice(0, limit)
.map(element => projectElement(element, 500));
return {
editor_state_id: editorStateId,
mode,
viewport: {
bounds: {
x: bounds.x,
y: bounds.y,
width: bounds.w,
height: bounds.h,
},
elements,
},
truncated: visible.length > limit,
};
}
if (args.view === 'selection_neighborhood') {
const selected = getSelectedModels(host) ?? [];
const neighborhood = neighborhoodOf(selected, 2);
const blocks = neighborhood.slice(0, limit);
return {
editor_state_id: editorStateId,
mode,
selection_neighborhood: blocks.map(model => projectBlock(model, 500)),
truncated: neighborhood.length > limit,
};
}
const outline = host.store
.getBlocksByFlavour('affine:note')
.flatMap(note => note.model.children);
const blocks = outline.slice(0, limit).map(model => projectBlock(model, 500));
return {
editor_state_id: editorStateId,
mode,
outline: blocks,
truncated: outline.length > limit,
};
}
packages/frontend/core/src/blocksuite/ai/runtime/request/__fixtures__/live-projection-contract.json
+120
View File
@@ -0,0 +1,120 @@
{
"schemaVersion": 1,
"fixtureId": "fixture-canvas-contract-v1",
"documents": [
{
"docId": "doc-01",
"title": "Project Aurora",
"revision": "fixture-revision-01",
"pageBlocks": [
{
"id": "block-01",
"type": "paragraph",
"visibility": "page",
"text": "Page fact alpha"
},
{
"id": "note-01",
"type": "note",
"visibility": "both",
"text": "Shared fact beta",
"bounds": { "x": 20, "y": 20, "width": 180, "height": 100 }
},
{
"id": "frame-01",
"type": "frame",
"visibility": "edgeless",
"title": "Delivery lane",
"bounds": { "x": 0, "y": 0, "width": 640, "height": 360 },
"childIds": ["shape-01", "text-01", "connector-01"]
},
{
"id": "text-block-01",
"type": "edgeless-text",
"visibility": "edgeless",
"text": "Canvas fact gamma",
"bounds": { "x": 260, "y": 70, "width": 180, "height": 40 }
}
],
"surfaceElements": [
{
"id": "shape-01",
"type": "shape",
"text": "Start node",
"bounds": { "x": 60, "y": 180, "width": 120, "height": 60 }
},
{
"id": "text-01",
"type": "text",
"text": "Finish node",
"bounds": { "x": 420, "y": 180, "width": 120, "height": 60 }
},
{
"id": "connector-01",
"type": "connector",
"label": "handoff",
"sourceId": "shape-01",
"targetId": "text-01",
"bounds": { "x": 180, "y": 205, "width": 240, "height": 10 }
},
{
"id": "group-01",
"type": "group",
"title": "Milestones",
"children": ["shape-01", "text-01"],
"bounds": { "x": 40, "y": 160, "width": 520, "height": 100 }
},
{
"id": "brush-01",
"type": "brush",
"pointCount": 4,
"bounds": { "x": 30, "y": 300, "width": 200, "height": 20 }
},
{
"id": "mindmap-01",
"type": "mindmap",
"text": "Root idea",
"children": ["mindmap-02"],
"bounds": { "x": 700, "y": 100, "width": 140, "height": 60 }
},
{
"id": "mindmap-02",
"type": "mindmap",
"text": "Leaf idea",
"parentId": "mindmap-01",
"index": "0",
"bounds": { "x": 900, "y": 100, "width": 140, "height": 60 }
},
{
"id": "unknown-01",
"type": "fixture-unknown",
"bounds": { "x": 700, "y": 260, "width": 80, "height": 80 }
}
]
}
],
"expectations": [
{
"docId": "doc-01",
"counts": {
"shape": 1,
"text": 1,
"connector": 1,
"group": 1,
"brush": 1,
"mindmap": 2,
"unknown": 1
},
"bothUnitIds": ["block:note-01"],
"connector": {
"id": "connector-01",
"sourceId": "shape-01",
"targetId": "text-01"
},
"frame": {
"id": "frame-01",
"childIds": ["connector-01", "shape-01", "text-01"]
}
}
]
}
packages/frontend/core/src/blocksuite/ai/runtime/request/action-definitions.ts
+6 -6
View File
@@ -59,17 +59,17 @@ export const actionDefinitions = {
| {
docs?: unknown;
files?: unknown;
selectedSnapshot?: unknown;
selectedMarkdown?: unknown;
html?: unknown;
}
| undefined;
return {
docs: contexts?.docs,
files: contexts?.files,
selectedSnapshot: contexts?.selectedSnapshot,
selectedMarkdown: contexts?.selectedMarkdown,
html: contexts?.html,
liveEditorContext:
typeof options.liveEditorContext === 'string'
? options.liveEditorContext
: undefined,
scopeSelectors: options.scopeSelectors,
focusSelectors: options.focusSelectors,
...(options.docId ? { currentDocId: options.docId } : {}),
};
},
packages/frontend/core/src/blocksuite/ai/runtime/request/byok-local-lease.ts
+29 -1
View File
@@ -1,6 +1,13 @@
import { apis, type ClientHandler } from '@affine/electron-api';
import { UserFriendlyError } from '@affine/error';
import {
type ByokAttachmentKind,
type ByokAttachmentSource,
type ByokEndpointKind,
type ByokModelFeature,
type ByokModelInput,
type ByokModelOutput,
type ByokOpenAiDialect,
ByokProvider,
createWorkspaceByokLocalLeaseMutation,
} from '@affine/graphql';
@@ -68,7 +75,28 @@ export async function createWorkspaceByokLocalLease(
name: provider.name,
description: provider.description ?? null,
credential: provider.credential,
definition: provider.definition,
definition: {
endpoint: {
kind: provider.definition.endpoint.kind as ByokEndpointKind,
url: provider.definition.endpoint.url,
dialect: provider.definition.endpoint.dialect as
| ByokOpenAiDialect
| null
| undefined,
},
models: provider.definition.models.map(model => ({
...model,
capabilities: model.capabilities.map(capability => ({
input: capability.input as ByokModelInput[],
output: capability.output as ByokModelOutput[],
features: capability.features as ByokModelFeature[],
attachmentKinds:
capability.attachmentKinds as ByokAttachmentKind[],
attachmentSources:
capability.attachmentSources as ByokAttachmentSource[],
})),
})),
},
enabled: provider.enabled ?? true,
},
]
packages/frontend/core/src/blocksuite/ai/runtime/request/copilot-client.spec.ts
+31 -1
View File
@@ -3,7 +3,13 @@
*/
import { describe, expect, test, vi } from 'vitest';
import { CopilotClient, Endpoint } from './copilot-client';
import {
SelectedSourcesFailedError,
SelectedSourcesLimitExceededError,
SelectedSourcesProcessingError,
SelectedSourcesUnavailableError,
} from '../../provider/error';
import { CopilotClient, Endpoint, resolveError } from './copilot-client';
describe('CopilotClient action streams', () => {
test('routes action endpoint outside the deprecated workflow path', () => {
@@ -48,5 +54,29 @@ describe('CopilotClient action streams', () => {
{ workspaceId: 'workspace-1' },
{ timeoutMs: 10000 }
);
expect(
resolveError({
name: 'COPILOT_SELECTED_SOURCES_PROCESSING',
type: 'BAD_REQUEST',
message: 'processing',
})
).toBeInstanceOf(SelectedSourcesProcessingError);
for (const [name, expected] of [
['COPILOT_SELECTED_SOURCES_FAILED', SelectedSourcesFailedError],
['COPILOT_SELECTED_SOURCES_UNAVAILABLE', SelectedSourcesUnavailableError],
[
'COPILOT_SELECTED_SOURCES_LIMIT_EXCEEDED',
SelectedSourcesLimitExceededError,
],
] as const) {
expect(
resolveError({
name,
type: 'INTERNAL_SERVER_ERROR',
message: 'unavailable',
})
).toBeInstanceOf(expected);
}
});
});
packages/frontend/core/src/blocksuite/ai/runtime/request/copilot-client.ts
+16 -165
View File
@@ -3,12 +3,7 @@ import type { AIToolsConfig } from '@affine/core/modules/ai-button';
import type { NbstoreService } from '@affine/core/modules/storage';
import { UserFriendlyError } from '@affine/error';
import {
addContextBlobMutation,
addContextCategoryMutation,
addContextDocMutation,
addContextFileMutation,
cleanupCopilotSessionMutation,
createCopilotContextMutation,
createCopilotMessageMutation,
createCopilotSessionMutation,
createCopilotSessionWithHistoryMutation,
@@ -19,16 +14,9 @@ import {
getCopilotSessionQuery,
getCopilotSessionsQuery,
type GraphQLQuery,
listContextObjectQuery,
listContextQuery,
matchContextQuery,
type PaginationInput,
type QueryOptions,
type QueryResponse,
removeContextBlobMutation,
removeContextCategoryMutation,
removeContextDocMutation,
removeContextFileMutation,
type RequestOptions,
updateCopilotSessionMutation,
} from '@affine/graphql';
@@ -37,6 +25,10 @@ import { getCurrentStore } from '@toeverything/infra';
import {
GeneralNetworkError,
PaymentRequiredError,
SelectedSourcesFailedError,
SelectedSourcesLimitExceededError,
SelectedSourcesProcessingError,
SelectedSourcesUnavailableError,
UnauthorizedError,
} from '../../provider/error';
@@ -64,6 +56,18 @@ function isAbortError(error: UserFriendlyError) {
}
function codeToError(error: UserFriendlyError) {
if (error.name === 'COPILOT_SELECTED_SOURCES_PROCESSING') {
return new SelectedSourcesProcessingError(error.message);
}
if (error.name === 'COPILOT_SELECTED_SOURCES_FAILED') {
return new SelectedSourcesFailedError(error.message);
}
if (error.name === 'COPILOT_SELECTED_SOURCES_UNAVAILABLE') {
return new SelectedSourcesUnavailableError(error.message);
}
if (error.name === 'COPILOT_SELECTED_SOURCES_LIMIT_EXCEEDED') {
return new SelectedSourcesLimitExceededError(error.message);
}
switch (error.status) {
case 401:
return new UnauthorizedError();
@@ -326,141 +330,6 @@ export class CopilotClient {
}
}
async createContext(workspaceId: string, sessionId: string) {
const res = await this.gql({
query: createCopilotContextMutation,
variables: {
workspaceId,
sessionId,
},
});
return res.createCopilotContext;
}
async getContextId(workspaceId: string, sessionId: string) {
const res = await this.gql({
query: listContextQuery,
variables: {
workspaceId,
sessionId,
},
});
return res.currentUser?.copilot?.contexts?.[0]?.id || undefined;
}
async addContextDoc(options: OptionsField<typeof addContextDocMutation>) {
const res = await this.gql({
query: addContextDocMutation,
variables: {
options,
},
});
return res.addContextDoc;
}
async removeContextDoc(
options: OptionsField<typeof removeContextDocMutation>
) {
const res = await this.gql({
query: removeContextDocMutation,
variables: {
options,
},
});
return res.removeContextDoc;
}
async addContextFile(
content: File,
options: OptionsField<typeof addContextFileMutation>
) {
const res = await this.gql({
query: addContextFileMutation,
variables: {
content,
options,
},
timeout: 60000,
});
return res.addContextFile;
}
async removeContextFile(
options: OptionsField<typeof removeContextFileMutation>
) {
const res = await this.gql({
query: removeContextFileMutation,
variables: {
options,
},
});
return res.removeContextFile;
}
async addContextCategory(
options: OptionsField<typeof addContextCategoryMutation>
) {
const res = await this.gql({
query: addContextCategoryMutation,
variables: {
options,
},
});
return res.addContextCategory;
}
async removeContextCategory(
options: OptionsField<typeof removeContextCategoryMutation>
) {
const res = await this.gql({
query: removeContextCategoryMutation,
variables: {
options,
},
});
return res.removeContextCategory;
}
async getContextDocsAndFiles(
workspaceId: string,
sessionId: string,
contextId: string
) {
const res = await this.gql({
query: listContextObjectQuery,
variables: {
workspaceId,
sessionId,
contextId,
},
});
return res.currentUser?.copilot?.contexts?.[0];
}
async matchContext(
content: string,
contextId?: string,
workspaceId?: string,
limit?: number,
scopedThreshold?: number,
threshold?: number
) {
const res = await this.gql({
query: matchContextQuery,
variables: {
content,
contextId,
workspaceId,
limit,
scopedThreshold,
threshold,
},
});
const { matchFiles: files, matchWorkspaceDocs: docs } =
res.currentUser?.copilot?.contexts?.[0] || {};
return { files, docs };
}
// Text or image to text
chatTextStream(
{
@@ -563,22 +432,4 @@ export class CopilotClient {
{ timeoutMs: 10000 }
);
}
addContextBlob(options: OptionsField<typeof addContextBlobMutation>) {
return this.gql({
query: addContextBlobMutation,
variables: {
options,
},
}).then(res => res.addContextBlob);
}
removeContextBlob(options: OptionsField<typeof removeContextBlobMutation>) {
return this.gql({
query: removeContextBlobMutation,
variables: {
options,
},
}).then(res => res.removeContextBlob);
}
}
packages/frontend/core/src/blocksuite/ai/runtime/request/message-transport.ts
+19 -5
View File
@@ -65,6 +65,11 @@ async function resizeImage(blob: Blob | File): Promise<Blob | null> {
return null;
}
function jpegFileName(name: string) {
const stem = name.replace(/\.[^./\\]+$/, '');
return `${stem || 'image'}.jpg`;
}
interface CreateMessageOptions {
client: CopilotClient;
sessionId: string;
@@ -99,12 +104,21 @@ async function createMessage({
options.attachments = stringAttachments;
options.blobs = (
await Promise.all(
blobs.map(resizeImage).map(async blob => {
const file = await blob;
blobs.map(async blob => {
const file = blob.type.startsWith('image/')
? await resizeImage(blob)
: blob;
if (!file) return null;
return new File([file], sessionId, {
type: file.type,
});
const resized = file !== blob;
return new File(
[file],
resized
? jpegFileName(blob instanceof File ? blob.name : sessionId)
: blob instanceof File
? blob.name
: sessionId,
{ type: file.type }
);
})
)
).filter(Boolean) as File[];
packages/frontend/core/src/blocksuite/ai/runtime/request/service.spec.ts
+255 -1
View File
@@ -1,9 +1,17 @@
/**
* @vitest-environment happy-dom
*/
import { readFileSync } from 'node:fs';
import { join } from 'node:path';
import { UserFriendlyError } from '@affine/error';
import type { EditorHost } from '@blocksuite/affine/std';
import type { GfxModel } from '@blocksuite/affine/std/gfx';
import { BehaviorSubject, Subject } from 'rxjs';
import { beforeEach, describe, expect, test, vi } from 'vitest';
import { DelegatedEditorHost } from '../frontend/delegated-editor-host';
import { readNodes } from '../frontend/live-projection';
import { type CopilotClient, Endpoint } from './copilot-client';
import { textToText, toImage } from './message-transport';
import { AIRequestService } from './service';
@@ -123,7 +131,6 @@ describe('runtime request transport BYOK local lease handling', () => {
name: 'OpenAI',
credential: 'sk-local',
definition: {
version: 1,
endpoint: { kind: 'provider_default' },
models: [{ modelId: 'model-1', capabilities: [] }],
},
@@ -204,6 +211,253 @@ describe('AIRequestService action definitions', () => {
electronApis.byokStorage = undefined;
});
test('manages the active delegated editor and its live projection contract', async () => {
const service = new AIRequestService(createClient());
const started: string[] = [];
const synced: string[] = [];
const disposed: string[] = [];
service.setActiveEditorFactory(sessionId => ({
start: async () => {
started.push(sessionId);
},
sync: async () => {
synced.push(sessionId);
},
dispose: () => {
disposed.push(sessionId);
},
context: () => JSON.stringify({ session_id: sessionId }),
}));
await service.activateEditor('session-1');
await service.activateEditor('session-1');
await service.activateEditor('session-2');
expect(service.getActiveEditorContext()).toBe(
JSON.stringify({ session_id: 'session-2' })
);
service.setActiveEditorFactory(undefined);
expect(started).toEqual(['session-1', 'session-2']);
expect(synced).toEqual(['session-1']);
expect(disposed).toEqual(['session-1', 'session-2']);
vi.useFakeTimers();
try {
const blockUpdated$ = new Subject<void>();
const selectionChanged$ = new Subject<void>();
const viewportUpdated$ = new Subject<void>();
const viewportSizeUpdated$ = new Subject<void>();
const toolRequests$ = new Subject<never>();
const selectedElements: Array<{ id: string }> = [];
const requests: Array<{ op: string; input: Record<string, unknown> }> =
[];
let holdNextUpsert = false;
let resolveHeldUpsert: (() => void) | undefined;
const realtime = {
subscribe: vi.fn(() => toolRequests$),
request: vi.fn(async (op: string, input: Record<string, unknown>) => {
requests.push({ op, input });
if (holdNextUpsert && op.endsWith('.upsert')) {
await new Promise<void>(resolve => {
resolveHeldUpsert = resolve;
});
}
return { ok: true };
}),
};
const delegatedHost = new DelegatedEditorHost({
realtime: realtime as never,
host: {
store: {
readonly$: new BehaviorSubject(false),
slots: Object.fromEntries([['blockUpdated', blockUpdated$]]),
},
selection: {
slots: Object.fromEntries([['changed', selectionChanged$]]),
},
std: {
get: () => ({
getEditorMode: () => 'edgeless',
selection: { selectedElements },
viewport: Object.fromEntries([
['viewportUpdated', viewportUpdated$],
['sizeUpdated', viewportSizeUpdated$],
]),
}),
},
} as unknown as EditorHost,
sessionId: 'session-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
});
await delegatedHost.start();
selectionChanged$.next();
await vi.advanceTimersByTimeAsync(150);
expect(requests.filter(item => item.op.endsWith('.upsert'))).toHaveLength(
1
);
selectedElements.push({ id: 'element-1' });
selectionChanged$.next();
selectionChanged$.next();
await vi.advanceTimersByTimeAsync(150);
const selectionUpserts = requests.filter(item =>
item.op.endsWith('.upsert')
);
expect(selectionUpserts).toHaveLength(2);
expect(selectionUpserts[1].input.editorStateId).not.toBe(
selectionUpserts[0].input.editorStateId
);
blockUpdated$.next();
blockUpdated$.next();
blockUpdated$.next();
await vi.advanceTimersByTimeAsync(150);
expect(requests.filter(item => item.op.endsWith('.upsert'))).toHaveLength(
3
);
holdNextUpsert = true;
blockUpdated$.next();
await vi.advanceTimersByTimeAsync(150);
blockUpdated$.next();
const sync = delegatedHost.sync();
holdNextUpsert = false;
resolveHeldUpsert?.();
await sync;
const synchronizedUpserts = requests.filter(item =>
item.op.endsWith('.upsert')
);
expect(synchronizedUpserts).toHaveLength(5);
expect(synchronizedUpserts[4].input.editorStateId).not.toBe(
synchronizedUpserts[3].input.editorStateId
);
delegatedHost.dispose();
} finally {
vi.useRealTimers();
}
const fixture = JSON.parse(
readFileSync(
join(
process.cwd(),
'packages/frontend/core/src/blocksuite/ai/runtime/request/__fixtures__/live-projection-contract.json'
),
'utf8'
)
);
const document = fixture.documents[0];
const expectation = fixture.expectations[0];
const frameSource = document.pageBlocks.find(
(block: { id: string }) => block.id === expectation.frame.id
);
const frame = {
id: frameSource.id,
flavour: `affine:${frameSource.type}`,
xywh: JSON.stringify([
frameSource.bounds.x,
frameSource.bounds.y,
frameSource.bounds.width,
frameSource.bounds.height,
]),
props: {
title: frameSource.title,
childElementIds: frameSource.childIds,
},
group: null,
groups: [],
} as unknown as GfxModel;
const elements = document.surfaceElements.map(
(source: Record<string, unknown>) =>
({
id: source.id,
flavour: source.type,
xywh: JSON.stringify([
(source.bounds as { x: number }).x,
(source.bounds as { y: number }).y,
(source.bounds as { width: number }).width,
(source.bounds as { height: number }).height,
]),
props: {
...source,
source: source.sourceId ? { id: source.sourceId } : undefined,
target: source.targetId ? { id: source.targetId } : undefined,
},
group: source.id === expectation.connector.id ? frame : null,
groups: source.id === expectation.connector.id ? [frame] : [],
}) as unknown as GfxModel
);
const mindmapRoot = elements.find(
(element: GfxModel) => element.id === 'mindmap-01'
);
const mindmapLeaf = elements.find(
(element: GfxModel) => element.id === 'mindmap-02'
);
if (!mindmapRoot || !mindmapLeaf) {
throw new Error('Anonymous contract fixture is missing mindmap nodes');
}
(
mindmapRoot as unknown as { props: Record<string, unknown> }
).props.children = {
[mindmapLeaf.id]: {
parent: mindmapRoot.id,
index: document.surfaceElements.find(
(element: { id: string }) => element.id === mindmapLeaf.id
).index,
},
};
(mindmapLeaf as unknown as { group: GfxModel }).group = mindmapRoot;
(mindmapLeaf as unknown as { groups: GfxModel[] }).groups = [mindmapRoot];
const models = new Map(
[frame, ...elements].map(model => [model.id, model])
);
const projection = readNodes(
{
store: { getBlock: () => undefined },
std: {
get: () => ({ getElementById: (id: string) => models.get(id) }),
},
} as unknown as EditorHost,
'editor-state-1',
{ element_ids: [...models.keys()] }
);
const values = projection.items.map(item =>
'value' in item && item.value && 'type' in item.value
? item.value
: undefined
);
const frameValue = values.find(value => value?.id === expectation.frame.id);
const connectorValue = values.find(
value => value?.id === expectation.connector.id
);
const mindmapRootValue = values.find(value => value?.id === 'mindmap-01');
const mindmapLeafValue = values.find(value => value?.id === 'mindmap-02');
const brushValue = values.find(value => value?.id === 'brush-01');
const unknownValue = values.find(value => value?.id === 'unknown-01');
expect(frameValue).toMatchObject({
type: 'frame',
child_ids: expectation.frame.childIds,
});
expect(connectorValue).toMatchObject({
type: 'connector',
frame_id: expectation.frame.id,
source_id: expectation.connector.sourceId,
target_id: expectation.connector.targetId,
});
expect(mindmapRootValue).toMatchObject({ child_ids: ['mindmap-02'] });
expect(mindmapLeafValue).toMatchObject({
parent_id: 'mindmap-01',
index: '0',
});
expect(brushValue).toMatchObject({ type: 'brush', point_count: 4 });
expect(unknownValue).toMatchObject({ type: 'fixture-unknown' });
expect(values.slice(1).map(value => value?.type)).toEqual(
document.surfaceElements.map((element: { type: string }) => element.type)
);
});
test('routes action-stream requests through action endpoint', async () => {
const client = createClient();
const service = new AIRequestService(client);
packages/frontend/core/src/blocksuite/ai/runtime/request/service.ts
+103 -130
View File
@@ -1,6 +1,5 @@
import type { NbstoreService } from '@affine/core/modules/storage';
import {
ContextCategories,
type CopilotChatHistoryFragment,
type getCopilotHistoriesQuery,
type GraphQLQuery,
@@ -34,6 +33,20 @@ export type AIRequestActionEvent = {
};
export class AIRequestService {
private activeEditorFactory?: (sessionId: string) => {
start(): Promise<void>;
sync(): Promise<void>;
dispose(): void;
context(): string;
};
private activeEditor?: {
sessionId: string;
sync(): Promise<void>;
dispose(): void;
context(): string;
};
private activeEditorGeneration = 0;
private activeEditorActivation = Promise.resolve();
private lastActionSessionId = '';
private readonly actionHistory: {
action: AIActionId;
@@ -41,7 +54,89 @@ export class AIRequestService {
}[] = [];
readonly actionEvents$ = new Subject<AIRequestActionEvent>();
constructor(readonly client: CopilotClientType) {}
constructor(
readonly client: CopilotClientType,
private readonly syncSelectedSources?: (docIds: string[]) => Promise<void>
) {}
async waitForSelectedSources(docIds: string[]) {
if (!this.syncSelectedSources) {
throw new Error('Selected sources cannot be synchronized');
}
let timeout: ReturnType<typeof setTimeout> | undefined;
try {
await Promise.race([
this.syncSelectedSources(docIds),
new Promise<never>((_, reject) => {
timeout = setTimeout(
() =>
reject(new Error('Selected source synchronization timed out')),
15000
);
}),
]);
} finally {
if (timeout) clearTimeout(timeout);
}
}
setActiveEditorFactory(
factory:
| ((sessionId: string) => {
start(): Promise<void>;
sync(): Promise<void>;
dispose(): void;
context(): string;
})
| undefined
) {
this.activeEditorGeneration++;
this.activeEditor?.dispose();
this.activeEditor = undefined;
this.activeEditorFactory = factory;
}
activateEditor(sessionId: string) {
const activation = this.activeEditorActivation.then(() =>
this.activateEditorNow(sessionId)
);
this.activeEditorActivation = activation.catch(() => {});
return activation;
}
getActiveEditorContext() {
return this.activeEditor?.context();
}
private async activateEditorNow(sessionId: string) {
if (this.activeEditor?.sessionId === sessionId) {
await this.activeEditor.sync();
return;
}
this.activeEditor?.dispose();
this.activeEditor = undefined;
const generation = this.activeEditorGeneration;
const host = this.activeEditorFactory?.(sessionId);
if (!host) {
return;
}
try {
await host.start();
} catch (error) {
host.dispose();
throw error;
}
if (generation !== this.activeEditorGeneration) {
host.dispose();
return;
}
this.activeEditor = {
sessionId,
sync: () => host.sync(),
dispose: () => host.dispose(),
context: () => host.context(),
};
}
isReady() {
return true;
@@ -180,132 +275,6 @@ export class AIRequestService {
},
};
context = {
createContext: (workspaceId: string, sessionId: string) =>
this.client.createContext(workspaceId, sessionId),
getContextId: (workspaceId: string, sessionId: string) =>
this.client.getContextId(workspaceId, sessionId),
addContextDoc: (options: { contextId: string; docId: string }) =>
this.client.addContextDoc(options),
removeContextDoc: (options: { contextId: string; docId: string }) =>
this.client.removeContextDoc(options),
addContextFile: (
file: File,
options: Parameters<CopilotClient['addContextFile']>[1]
) => this.client.addContextFile(file, options),
removeContextFile: (options: { contextId: string; fileId: string }) =>
this.client.removeContextFile(options),
addContextTag: (options: {
contextId: string;
tagId: string;
docIds: string[];
}) =>
this.client.addContextCategory({
contextId: options.contextId,
type: ContextCategories.Tag,
categoryId: options.tagId,
docs: options.docIds,
}),
removeContextTag: (options: { contextId: string; tagId: string }) =>
this.client.removeContextCategory({
contextId: options.contextId,
type: ContextCategories.Tag,
categoryId: options.tagId,
}),
addContextCollection: (options: {
contextId: string;
collectionId: string;
docIds: string[];
}) =>
this.client.addContextCategory({
contextId: options.contextId,
type: ContextCategories.Collection,
categoryId: options.collectionId,
docs: options.docIds,
}),
removeContextCollection: (options: {
contextId: string;
collectionId: string;
}) =>
this.client.removeContextCategory({
contextId: options.contextId,
type: ContextCategories.Collection,
categoryId: options.collectionId,
}),
getContextDocsAndFiles: (
workspaceId: string,
sessionId: string,
contextId: string
) => this.client.getContextDocsAndFiles(workspaceId, sessionId, contextId),
matchContext: (
content: string,
contextId?: string,
workspaceId?: string,
limit?: number,
scopedThreshold?: number,
threshold?: number
) =>
this.client.matchContext(
content,
contextId,
workspaceId,
limit,
scopedThreshold,
threshold
),
addContextBlob: (options: { blobId: string; contextId: string }) =>
this.client.addContextBlob({
contextId: options.contextId,
blobId: options.blobId,
}),
removeContextBlob: (options: { blobId: string; contextId: string }) =>
this.client.removeContextBlob({
contextId: options.contextId,
blobId: options.blobId,
}),
pollContextDocsAndFiles: async (
workspaceId: string,
sessionId: string,
contextId: string,
onPoll: (
result: BlockSuitePresets.AIDocsAndFilesContext | undefined
) => void,
abortSignal: AbortSignal
) => {
let attempts = 0;
const minInterval = 1000;
const maxInterval = 30 * 1000;
while (!abortSignal.aborted) {
const result = await this.client.getContextDocsAndFiles(
workspaceId,
sessionId,
contextId
);
onPoll(result);
const interval = Math.min(
minInterval * Math.pow(1.5, attempts),
maxInterval
);
attempts++;
await new Promise(resolve => setTimeout(resolve, interval));
}
},
pollEmbeddingStatus: async (
workspaceId: string,
onPoll: (
result: Awaited<ReturnType<CopilotClientType['getEmbeddingStatus']>>
) => void,
abortSignal: AbortSignal
) => {
const interval = 10 * 1000;
while (!abortSignal.aborted) {
onPoll(await this.client.getEmbeddingStatus(workspaceId));
await new Promise(resolve => setTimeout(resolve, interval));
}
},
};
forkChat(options: BlockSuitePresets.AIForkChatSessionOptions) {
return this.client.forkSession(options);
}
@@ -392,7 +361,11 @@ export function createAIRequestService(
url: string,
eventSourceInitDict?: EventSourceInit
) => EventSource,
realtime: Pick<NbstoreService['realtime'], 'request'>
realtime: Pick<NbstoreService['realtime'], 'request'>,
syncSelectedSources?: (docIds: string[]) => Promise<void>
) {
return new AIRequestService(new CopilotClient(gql, eventSource, realtime));
return new AIRequestService(
new CopilotClient(gql, eventSource, realtime),
syncSelectedSources
);
}
packages/frontend/core/src/components/hooks/affine/use-ai-chat-config.ts
+8
View File
@@ -82,6 +82,13 @@ export function useAIChatConfig() {
const collectionMetas$ = collectionService.collectionMetas$;
return createSignalFromObservable(collectionMetas$, []);
},
getCollectionTitle: (collectionId: string) => {
return (
collectionService.collectionMetas$.value.find(
collection => collection.id === collectionId
)?.name ?? ''
);
},
getCollectionPageIds: (collectionId: string) => {
const collection$ = collectionService.collection$(collectionId);
// TODO: lack of documents that meet the collection rules
@@ -90,6 +97,7 @@ export function useAIChatConfig() {
};
const searchMenuConfig = {
addContextAvailable: workspaceService.workspace.flavour !== 'local',
getDocMenuGroup: (
query: string,
action: SearchDocMenuAction,
packages/frontend/core/src/components/providers/workspace-side-effects.tsx
+10 -2
View File
@@ -164,7 +164,14 @@ export const WorkspaceSideEffects = () => {
createAIRequestService(
graphqlService.gql,
eventSourceService.eventSource,
nbstoreService.realtime
nbstoreService.realtime,
async docIds => {
await Promise.all(
[currentWorkspace.id, 'db$docProperties', ...docIds].map(docId =>
currentWorkspace.engine.doc.waitForSynced(docId)
)
);
}
),
globalDialogService,
authService
@@ -173,9 +180,10 @@ export const WorkspaceSideEffects = () => {
dispose();
};
}, [
currentWorkspace.engine.doc,
currentWorkspace.id,
eventSourceService,
nbstoreService,
workspaceDialogService,
graphqlService,
globalDialogService,
authService,
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/add-key-modal.tsx
+84 -38
View File
@@ -1,5 +1,8 @@
import { Button, Input, Modal, notify } from '@affine/component';
import {
ByokCustomEndpointMode,
ByokEndpointKind,
ByokOpenAiDialect,
ByokProvider,
createWorkspaceByokProfileMutation,
probeWorkspaceByokDraftMutation,
@@ -15,7 +18,6 @@ import {
byokT,
endpointHintKey,
providerLabels,
shouldShowEndpoint,
storageLabel,
} from './metadata';
import { ModelSelector } from './model-selector';
@@ -41,7 +43,6 @@ export const AddKeyModal = ({
localStorageSupported,
canAddServerKey,
canAddLocalKey,
isSelfHosted,
gql,
}: {
workspaceId: string;
@@ -55,7 +56,6 @@ export const AddKeyModal = ({
localStorageSupported: boolean;
canAddServerKey: boolean;
canAddLocalKey: boolean;
isSelfHosted: boolean;
gql?: GqlFn;
}) => {
const t = useI18n();
@@ -67,6 +67,7 @@ export const AddKeyModal = ({
const [apiKey, setApiKey] = useState('');
const [customEndpoint, setCustomEndpoint] = useState(false);
const [endpoint, setEndpoint] = useState('');
const [dialect, setDialect] = useState<ByokOpenAiDialect | null>(null);
const [models, setModels] = useState<ModelDeclaration[]>([]);
const [testStatus, setTestStatus] = useState<'passed' | 'failed' | null>(
null
@@ -76,14 +77,16 @@ export const AddKeyModal = ({
const busyRef = useRef(false);
const localStorageUnavailable = !localStorageSupported || !canAddLocalKey;
const localStorageDisabled = !!editingKey || localStorageUnavailable;
const showCustomEndpoint = shouldShowEndpoint(
isSelfHosted,
settings.customEndpointSupported
);
const customEndpointMode = settings.policy.customEndpointMode;
const showCustomEndpoint =
provider === ByokProvider.openai &&
customEndpointMode !== ByokCustomEndpointMode.unavailable;
const customEndpointEnabled =
customEndpointMode === ByokCustomEndpointMode.enabled;
const endpointHint = endpointHintKey(
settings.customEndpointSupported,
settings.privateEndpointSupported
customEndpointMode,
settings.policy.privateEndpointSupported
);
const providerCatalog = useMemo(
() => catalogModels(settings, provider),
@@ -103,7 +106,11 @@ export const AddKeyModal = ({
);
setApiKey('');
setEndpoint(editingKey?.definition.endpoint.url ?? '');
setCustomEndpoint(editingKey?.definition.endpoint.kind === 'custom');
setDialect(editingKey?.definition.endpoint.dialect ?? null);
setCustomEndpoint(
editingKey?.definition.endpoint.kind ===
ByokEndpointKind.openai_compatible
);
setModels(
editingKey?.definition.models ?? defaultModels(settings, nextProvider)
);
@@ -113,13 +120,20 @@ export const AddKeyModal = ({
const definition = useMemo<ByokDefinition>(
() => ({
version: editingKey?.definition.version ?? 1,
endpoint: customEndpoint
? { kind: 'custom', url: endpoint }
: { kind: 'provider_default', url: null },
? {
kind: ByokEndpointKind.openai_compatible,
url: endpoint,
dialect,
}
: {
kind: ByokEndpointKind.provider_default,
url: null,
dialect: null,
},
models,
}),
[customEndpoint, editingKey?.definition.version, endpoint, models]
[customEndpoint, dialect, endpoint, models]
);
const invalidateTest = () => setTestStatus(null);
@@ -297,7 +311,7 @@ export const AddKeyModal = ({
models.length > 0 &&
models.every(model => model.modelId.trim() && model.capabilities.length) &&
new Set(models.map(model => model.modelId.trim())).size === models.length &&
(!customEndpoint || !!endpoint.trim());
(!customEndpoint || (!!endpoint.trim() && dialect !== null));
return (
<Modal
@@ -325,11 +339,14 @@ export const AddKeyModal = ({
const next = event.target.value as ByokProvider;
setProvider(next);
setName(providerLabels[next]);
setCustomEndpoint(false);
setEndpoint('');
setDialect(null);
setModels(defaultModels(settings, next));
invalidateTest();
}}
>
{settings.allowedProviders.map(item => (
{settings.policy.allowedProviders.map(item => (
<option key={item} value={item}>
{providerLabels[item]}
</option>
@@ -408,42 +425,71 @@ export const AddKeyModal = ({
<input
type="checkbox"
checked={customEndpoint}
disabled={!settings.customEndpointSupported}
disabled={!customEndpointEnabled}
onChange={event => {
setCustomEndpoint(event.target.checked);
if (event.target.checked && !editingKey) setModels([]);
if (!event.target.checked)
setEndpoint('');
setDialect(null);
if (event.target.checked && !editingKey) {
setModels([]);
} else if (!event.target.checked) {
setModels(defaultModels(settings, provider));
}
invalidateTest();
}}
/>
{byokT(t, 'endpoint.use-custom')}
</label>
{!settings.customEndpointSupported && endpointHint ? (
{!customEndpointEnabled && endpointHint ? (
<span className={styles.fieldHint}>
{byokT(t, endpointHint)}
</span>
) : null}
{customEndpoint ? (
<label className={styles.endpointField}>
<span className={styles.label}>
{byokT(t, 'field.endpoint')}
</span>
<Input
size="large"
value={endpoint}
onChange={value => {
setEndpoint(value);
invalidateTest();
}}
placeholder="https://api.example.com/v1"
/>
{endpointHint ? (
<span className={styles.fieldHint}>
{byokT(t, endpointHint)}
<>
<label className={styles.endpointField}>
<span className={styles.label}>
{byokT(t, 'field.endpoint')}
</span>
) : null}
</label>
<Input
size="large"
value={endpoint}
onChange={value => {
setEndpoint(value);
invalidateTest();
}}
placeholder="https://api.example.com/v1"
/>
{endpointHint ? (
<span className={styles.fieldHint}>
{byokT(t, endpointHint)}
</span>
) : null}
</label>
<label className={styles.field}>
<span className={styles.label}>
{byokT(t, 'field.dialect')}
</span>
<select
className={styles.input}
value={dialect ?? ''}
onChange={event => {
setDialect(event.target.value as ByokOpenAiDialect);
invalidateTest();
}}
>
<option value="" disabled>
{byokT(t, 'placeholder.dialect')}
</option>
<option value={ByokOpenAiDialect.responses}>
{byokT(t, 'dialect.responses')}
</option>
<option value={ByokOpenAiDialect.chat_completions}>
{byokT(t, 'dialect.chat-completions')}
</option>
</select>
</label>
</>
) : null}
</>
) : null}
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/index.spec.tsx
+106 -23
View File
@@ -24,6 +24,57 @@ const ByokProvider = vi.hoisted(() => ({
gemini: 'gemini',
fal: 'fal',
}));
const ByokEnums = vi.hoisted(() => ({
ByokCustomEndpointMode: {
unavailable: 'unavailable',
disabled: 'disabled',
enabled: 'enabled',
},
ByokEndpointKind: {
provider_default: 'provider_default',
openai_compatible: 'openai_compatible',
},
ByokOpenAiDialect: {
responses: 'responses',
chat_completions: 'chat_completions',
},
ByokModelInput: {
text: 'text',
image: 'image',
audio: 'audio',
file: 'file',
},
ByokModelOutput: {
text: 'text',
object: 'object',
structured: 'structured',
embedding: 'embedding',
rerank: 'rerank',
image: 'image',
},
ByokModelFeature: {
tool_calling: 'tool_calling',
reasoning: 'reasoning',
web_search: 'web_search',
},
ByokAttachmentKind: { image: 'image', audio: 'audio', file: 'file' },
ByokAttachmentSource: {
url: 'url',
data: 'data',
bytes: 'bytes',
file_handle: 'file_handle',
},
ByokProbeOperation: {
chat: 'chat',
structured: 'structured',
tool_calling: 'tool_calling',
vision: 'vision',
embedding: 'embedding',
rerank: 'rerank',
image: 'image',
transcript: 'transcript',
},
}));
const createMutation = vi.hoisted(() => Symbol('create'));
const probeMutation = vi.hoisted(() => Symbol('probe'));
const replaceMutation = vi.hoisted(() => Symbol('replace'));
@@ -118,6 +169,7 @@ vi.mock('@affine/component', () => ({
vi.mock('@affine/graphql', () => ({
ByokProvider,
...ByokEnums,
createWorkspaceByokProfileMutation: createMutation,
probeWorkspaceByokDraftMutation: probeMutation,
replaceWorkspaceByokProfileMutation: replaceMutation,
@@ -138,15 +190,20 @@ const textCapability = {
attachmentSources: [],
};
function settings(customEndpointSupported = true) {
function settings(
customEndpointMode = ByokEnums.ByokCustomEndpointMode.enabled
) {
return {
workspaceId: 'workspace-1',
entitled: true,
serverEntitled: true,
localEntitled: false,
allowedProviders: Object.values(ByokProvider),
customEndpointSupported,
privateEndpointSupported: false,
policy: {
enabled: true,
allowedProviders: Object.values(ByokProvider),
customEndpointMode,
privateEndpointSupported: false,
},
localStorageSupported: false,
keys: [],
catalog: {
@@ -181,22 +238,30 @@ describe('BYOK settings behavior', () => {
});
test.each([
[false, false, 'endpoint.custom-disabled'],
[true, false, 'endpoint.private-disabled'],
[true, true, null],
[
ByokEnums.ByokCustomEndpointMode.disabled,
false,
'endpoint.custom-disabled',
],
[
ByokEnums.ByokCustomEndpointMode.enabled,
false,
'endpoint.private-disabled',
],
[ByokEnums.ByokCustomEndpointMode.enabled, true, null],
] as const)(
'maps endpoint policy custom=%s private=%s',
(customEndpointSupported, privateEndpointSupported, expected) => {
expect(
endpointHintKey(customEndpointSupported, privateEndpointSupported)
).toBe(expected);
'maps endpoint policy mode=%s private=%s',
(mode, privateEndpointSupported, expected) => {
expect(endpointHintKey(mode as never, privateEndpointSupported)).toBe(
expected
);
}
);
test('shows a disabled custom endpoint control with its self-hosted policy hint', () => {
const props = {
workspaceId: 'workspace-1',
settings: settings(false) as never,
settings: settings(ByokEnums.ByokCustomEndpointMode.disabled) as never,
editingKey: null,
open: true,
onOpenChange: vi.fn(),
@@ -208,7 +273,7 @@ describe('BYOK settings behavior', () => {
canAddLocalKey: false,
gql: vi.fn() as never,
};
const { rerender } = render(<AddKeyModal {...props} isSelfHosted />);
const { rerender } = render(<AddKeyModal {...props} />);
expect(
(
@@ -219,7 +284,14 @@ describe('BYOK settings behavior', () => {
).toBe(true);
expect(screen.getByText('custom-disabled')).toBeTruthy();
rerender(<AddKeyModal {...props} isSelfHosted={false} />);
rerender(
<AddKeyModal
{...props}
settings={
settings(ByokEnums.ByokCustomEndpointMode.unavailable) as never
}
/>
);
expect(screen.queryByRole('checkbox', { name: 'use-custom' })).toBeNull();
expect(screen.queryByText('custom-disabled')).toBeNull();
});
@@ -238,7 +310,6 @@ describe('BYOK settings behavior', () => {
localStorageSupported={false}
canAddServerKey
canAddLocalKey={false}
isSelfHosted={false}
gql={vi.fn() as never}
/>
);
@@ -335,7 +406,6 @@ describe('BYOK settings behavior', () => {
localStorageSupported={false}
canAddServerKey
canAddLocalKey={false}
isSelfHosted={false}
gql={gql as never}
/>
);
@@ -399,7 +469,6 @@ describe('BYOK settings behavior', () => {
localStorageSupported={false}
canAddServerKey
canAddLocalKey={false}
isSelfHosted={false}
gql={gql as never}
/>
);
@@ -421,7 +490,7 @@ describe('BYOK settings behavior', () => {
).toBe(false);
});
test('preserves definition version and server revision while editing', async () => {
test('preserves server revision and writes the breaking definition shape', async () => {
type MockOperation = {
query: symbol;
variables?: { input?: Record<string, unknown> };
@@ -462,8 +531,11 @@ describe('BYOK settings behavior', () => {
sortOrder: 0,
revision: 7,
definition: {
version: 3,
endpoint: { kind: 'provider_default', url: null },
endpoint: {
kind: ByokEnums.ByokEndpointKind.provider_default,
url: null,
dialect: null,
},
models: [
{
modelId: 'model-a',
@@ -483,7 +555,6 @@ describe('BYOK settings behavior', () => {
localStorageSupported={false}
canAddServerKey
canAddLocalKey={false}
isSelfHosted={false}
gql={gql as never}
/>
);
@@ -495,7 +566,19 @@ describe('BYOK settings behavior', () => {
);
expect(replaceCall?.[0].variables?.input).toMatchObject({
expectedRevision: 7,
definition: { version: 3 },
definition: {
endpoint: {
kind: ByokEnums.ByokEndpointKind.provider_default,
url: null,
dialect: null,
},
},
});
expect(
Object.hasOwn(
replaceCall?.[0].variables?.input?.definition as object,
'version'
)
).toBe(false);
});
});
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/index.tsx
+7 -7
View File
@@ -10,7 +10,6 @@ import {
type GraphQLQuery,
probeWorkspaceByokProfileMutation,
reorderWorkspaceByokProfilesMutation,
ServerDeploymentType,
workspaceByokSettingsQuery as byokSettingsQuery,
} from '@affine/graphql';
import { useI18n } from '@affine/i18n';
@@ -115,14 +114,16 @@ export const WorkspaceByokSetting = () => {
return a.sortOrder - b.sortOrder;
});
}, [localKeys, settings?.keys]);
const canAddServerKey = settings?.serverEntitled ?? false;
const policyAllowsCreation =
(settings?.policy.enabled ?? false) &&
(settings?.policy.allowedProviders.length ?? 0) > 0;
const canAddServerKey =
(settings?.serverEntitled ?? false) && policyAllowsCreation;
const canAddLocalKey =
(settings?.localEntitled ?? false) &&
(settings?.localStorageSupported ?? false);
(settings?.localStorageSupported ?? false) &&
policyAllowsCreation;
const canManageKeys = canAddServerKey || canAddLocalKey;
const isSelfHosted =
workspaceServer.server?.config$.value.type ===
ServerDeploymentType.Selfhosted;
const clearAll = useCallback(async () => {
if (!settings) {
@@ -397,7 +398,6 @@ export const WorkspaceByokSetting = () => {
localStorageSupported={settings.localStorageSupported}
canAddServerKey={canAddServerKey}
canAddLocalKey={canAddLocalKey}
isSelfHosted={isSelfHosted}
gql={workspaceServer.server?.gql as GqlFn | undefined}
/>
</>
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/metadata.ts
+20 -16
View File
@@ -1,4 +1,10 @@
import { ByokProvider } from '@affine/graphql';
import {
ByokCustomEndpointMode,
ByokModelFeature,
ByokModelInput,
ByokModelOutput,
ByokProvider,
} from '@affine/graphql';
import type { I18nInstance } from '@affine/i18n';
import { type ByokKey, ByokStorage } from './types';
@@ -25,10 +31,10 @@ export function storageLabel(t: I18nInstance, storage: ByokStorage) {
}
export function endpointHintKey(
customEndpointSupported: boolean,
mode: ByokCustomEndpointMode,
privateEndpointSupported: boolean
) {
if (!customEndpointSupported) {
if (mode === ByokCustomEndpointMode.disabled) {
return 'endpoint.custom-disabled';
}
if (!privateEndpointSupported) {
@@ -37,25 +43,23 @@ export function endpointHintKey(
return null;
}
export function shouldShowEndpoint(
isSelfHosted: boolean,
customEndpointSupported: boolean
) {
return isSelfHosted || customEndpointSupported;
}
export function capabilitiesFor(key: Pick<ByokKey, 'definition'>) {
const capabilities = key.definition.models.flatMap(model =>
model.enabled ? model.capabilities : []
);
const labels = new Set<string>();
for (const capability of capabilities) {
if (capability.output.includes('text')) labels.add('Text');
if (capability.input.includes('image')) labels.add('Image input');
if (capability.output.includes('image')) labels.add('Image generate');
if (capability.features.includes('tools')) labels.add('Actions');
if (capability.input.includes('audio')) labels.add('Transcript');
if (capability.output.includes('embedding')) labels.add('Indexing');
if (capability.output.includes(ByokModelOutput.text)) labels.add('Text');
if (capability.input.includes(ByokModelInput.image))
labels.add('Image input');
if (capability.output.includes(ByokModelOutput.image))
labels.add('Image generate');
if (capability.features.includes(ByokModelFeature.tool_calling))
labels.add('Actions');
if (capability.input.includes(ByokModelInput.audio))
labels.add('Transcript');
if (capability.output.includes(ByokModelOutput.embedding))
labels.add('Indexing');
}
return [...labels];
}
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.spec.ts
+27 -10
View File
@@ -1,3 +1,10 @@
import {
ByokAttachmentKind,
ByokAttachmentSource,
ByokModelFeature,
ByokModelInput,
ByokModelOutput,
} from '@affine/graphql';
import { describe, expect, test } from 'vitest';
import {
@@ -13,11 +20,16 @@ describe('BYOK model capabilities', () => {
enabled: true,
capabilities: [
{
input: ['text', 'image'],
output: ['text'],
features: ['tools'],
attachmentKinds: ['image'],
attachmentSources: ['url', 'data', 'bytes', 'file_handle'],
input: [ByokModelInput.text, ByokModelInput.image],
output: [ByokModelOutput.text],
features: [ByokModelFeature.tool_calling],
attachmentKinds: [ByokAttachmentKind.image],
attachmentSources: [
ByokAttachmentSource.url,
ByokAttachmentSource.data,
ByokAttachmentSource.bytes,
ByokAttachmentSource.file_handle,
],
},
],
};
@@ -27,11 +39,16 @@ describe('BYOK model capabilities', () => {
test('preserves a rich capability when its represented uses stay selected', () => {
const capability = {
input: ['text', 'image'],
output: ['text'],
features: ['tools'],
attachmentKinds: ['image'],
attachmentSources: ['url', 'data', 'bytes', 'file_handle'],
input: [ByokModelInput.text, ByokModelInput.image],
output: [ByokModelOutput.text],
features: [ByokModelFeature.tool_calling],
attachmentKinds: [ByokAttachmentKind.image],
attachmentSources: [
ByokAttachmentSource.url,
ByokAttachmentSource.data,
ByokAttachmentSource.bytes,
ByokAttachmentSource.file_handle,
],
};
const model: ModelDeclaration = {
modelId: 'multimodal-tools',
packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/byok/model-utils.ts
+53 -22
View File
@@ -1,4 +1,12 @@
import type { ByokProvider } from '@affine/graphql';
import {
ByokAttachmentKind,
ByokAttachmentSource,
ByokModelFeature,
ByokModelInput,
ByokModelOutput,
ByokProbeOperation,
type ByokProvider,
} from '@affine/graphql';
import type { ByokDefinition, ByokSettings } from './types';
@@ -28,42 +36,62 @@ export const useCases: { id: UseCase; labelKey: string }[] = [
export function capabilityForUseCase(useCase: UseCase): Capability {
switch (useCase) {
case 'actions':
return modelCapability(['text'], ['text'], ['tools']);
return modelCapability(
[ByokModelInput.text],
[ByokModelOutput.text],
[ByokModelFeature.tool_calling]
);
case 'structured':
return modelCapability(['text'], ['structured']);
return modelCapability(
[ByokModelInput.text],
[ByokModelOutput.structured]
);
case 'vision':
return modelCapability(
['text', 'image'],
['text'],
[ByokModelInput.text, ByokModelInput.image],
[ByokModelOutput.text],
[],
['image'],
['url', 'data', 'bytes', 'file_handle']
[ByokAttachmentKind.image],
[
ByokAttachmentSource.url,
ByokAttachmentSource.data,
ByokAttachmentSource.bytes,
ByokAttachmentSource.file_handle,
]
);
case 'image':
return modelCapability(['text'], ['image']);
return modelCapability([ByokModelInput.text], [ByokModelOutput.image]);
case 'transcript':
return modelCapability(
['audio'],
['structured'],
[ByokModelInput.audio],
[ByokModelOutput.structured],
[],
['audio'],
['url', 'data', 'bytes', 'file_handle']
[ByokAttachmentKind.audio],
[
ByokAttachmentSource.url,
ByokAttachmentSource.data,
ByokAttachmentSource.bytes,
ByokAttachmentSource.file_handle,
]
);
case 'embedding':
return modelCapability(['text'], ['embedding']);
return modelCapability(
[ByokModelInput.text],
[ByokModelOutput.embedding]
);
case 'rerank':
return modelCapability(['text'], ['rerank']);
return modelCapability([ByokModelInput.text], [ByokModelOutput.rerank]);
default:
return modelCapability(['text'], ['text']);
return modelCapability([ByokModelInput.text], [ByokModelOutput.text]);
}
}
function modelCapability(
input: string[],
output: string[],
features: string[] = [],
attachmentKinds: string[] = [],
attachmentSources: string[] = []
input: ByokModelInput[],
output: ByokModelOutput[],
features: ByokModelFeature[] = [],
attachmentKinds: ByokAttachmentKind[] = [],
attachmentSources: ByokAttachmentSource[] = []
): Capability {
return { input, output, features, attachmentKinds, attachmentSources };
}
@@ -78,7 +106,7 @@ function matchesCapability(value: Capability, useCase: UseCase) {
'attachmentSources',
] as const;
return fields.every(field =>
expected[field].every(item => value[field].includes(item))
expected[field].every(item => new Set<string>(value[field]).has(item))
);
}
@@ -127,7 +155,10 @@ export function probeChecks(models: ModelDeclaration[], includeImage: boolean) {
)
.map(useCase => ({
modelId: model.modelId,
operation: useCase === 'actions' ? 'tools' : useCase,
operation:
useCase === 'actions'
? ByokProbeOperation.tool_calling
: ByokProbeOperation[useCase],
}))
);
}
packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx
+10 -2
View File
@@ -57,15 +57,23 @@ function useAIRequestService() {
const graphqlService = useService(GraphQLService);
const eventSourceService = useService(EventSourceService);
const nbstoreService = useService(NbstoreService);
const workspace = useService(WorkspaceService).workspace;
return useMemo(
() =>
createAIRequestService(
graphqlService.gql,
eventSourceService.eventSource,
nbstoreService.realtime
nbstoreService.realtime,
async docIds => {
await Promise.all(
[workspace.id, 'db$docProperties', ...docIds].map(docId =>
workspace.engine.doc.waitForSynced(docId)
)
);
}
),
[graphqlService, eventSourceService, nbstoreService]
[graphqlService, eventSourceService, nbstoreService, workspace]
);
}
packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx
+46 -12
View File
@@ -3,6 +3,7 @@ import {
AIAppEvents,
AIChatRuntime,
createAIRequestService,
DelegatedEditorHost,
DocAIChatSessionStrategy,
useAIChatElement,
useAIChatRuntime,
@@ -39,6 +40,7 @@ import { PeekViewService } from '@affine/core/modules/peek-view';
import { NbstoreService } from '@affine/core/modules/storage';
import { AppThemeService } from '@affine/core/modules/theme';
import { WorkbenchService } from '@affine/core/modules/workbench';
import { WorkspaceService } from '@affine/core/modules/workspace';
import { useI18n } from '@affine/i18n';
import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference';
import { DocModeProvider } from '@blocksuite/affine/shared/services';
@@ -78,6 +80,7 @@ export const EditorChatPanel = ({
const eventSourceService = useService(EventSourceService);
const nbstoreService = useService(NbstoreService);
const workbench = useService(WorkbenchService).workbench;
const workspace = useService(WorkspaceService).workspace;
const t = useI18n();
const { closeConfirmModal, openConfirmModal } = useConfirmModal();
@@ -111,9 +114,21 @@ export const EditorChatPanel = ({
createAIRequestService(
graphqlService.gql,
eventSourceService.eventSource,
nbstoreService.realtime
nbstoreService.realtime,
async docIds => {
await Promise.all(
[workspace.id, 'db$docProperties', ...docIds].map(docId =>
workspace.engine.doc.waitForSynced(docId)
)
);
}
),
[eventSourceService.eventSource, graphqlService.gql, nbstoreService]
[
eventSourceService.eventSource,
graphqlService.gql,
nbstoreService,
workspace,
]
);
const [pendingSessionId] = useState(() => {
@@ -148,6 +163,32 @@ export const EditorChatPanel = ({
snapshot?.sessions.find(
item => item.sessionId === snapshot.activeSessionId
) ?? null;
useEffect(() => {
if (!host || !workspaceId) return;
requestService.setActiveEditorFactory(
sessionId =>
new DelegatedEditorHost({
realtime: nbstoreService.realtime,
host,
sessionId,
workspaceId,
docId: doc.id,
})
);
if (session?.sessionId) {
void requestService
.activateEditor(session.sessionId)
.catch(console.error);
}
return () => requestService.setActiveEditorFactory(undefined);
}, [
doc.id,
host,
nbstoreService.realtime,
requestService,
session?.sessionId,
workspaceId,
]);
const appSidebarConfig = useMemo<AppSidebarConfig>(() => {
return {
getWidth: () =>
@@ -438,11 +479,7 @@ export const EditorChatPanel = ({
chatTabsContainerRef.current = node;
}, []);
const embeddingCount = snapshot?.composer.context.embeddingCount;
const done = embeddingCount?.finished ?? 0;
const total =
done + (embeddingCount?.processing ?? 0) + (embeddingCount?.failed ?? 0);
const isEmbedding = total > 0 && done < total;
const isSynchronizing = snapshot?.composer.scopeSelection.syncing ?? false;
const hasRuntimeSnapshot = !!snapshot;
return (
@@ -460,12 +497,9 @@ export const EditorChatPanel = ({
<div className={styles.container}>
<div className={styles.header}>
<div className={styles.title}>
{isEmbedding ? (
{isSynchronizing ? (
<span data-testid="chat-panel-embedding-progress">
{t.t('com.affine.ai.chat-panel.embedding-progress', {
done,
total,
})}
Synchronizing sources
</span>
) : (
t['com.affine.ai.chat-panel.title']()
packages/frontend/core/src/modules/workspace-indexer-embedding/entities/additional-attachments.ts
+33 -7
View File
@@ -30,6 +30,7 @@ interface Attachments {
}
export class AdditionalAttachments extends Entity {
private refreshTimer?: ReturnType<typeof setTimeout>;
error$ = new LiveData<any>(null);
attachments$ = new LiveData<Attachments>({
edges: [],
@@ -69,15 +70,39 @@ export class AdditionalAttachments extends Entity {
mergeMap(value => {
const patched = {
...value,
edges: value.edges.map(edge => ({
...edge,
node: {
...edge.node,
status: 'uploaded' as const,
},
})),
edges: value.edges.map(edge => {
const status = edge.node.embeddingStatus;
if (
status !== 'processing' &&
status !== 'ready' &&
status !== 'failed'
) {
throw new Error(`Unknown attachment status: ${status}`);
}
const normalizedStatus: PersistedAttachmentFile['status'] =
status;
return {
...edge,
node: {
...edge.node,
status: normalizedStatus,
},
};
}),
};
this.attachments$.next(patched);
clearTimeout(this.refreshTimer);
if (
patched.edges.some(
(edge: { node: PersistedAttachmentFile }) =>
edge.node.status === 'processing'
)
) {
this.refreshTimer = setTimeout(
() => this.getAttachments({ first: COUNT_PER_PAGE, after: null }),
1000
);
}
return EMPTY;
}),
catchErrorInto(this.error$, error => {
@@ -148,6 +173,7 @@ export class AdditionalAttachments extends Entity {
};
override dispose(): void {
clearTimeout(this.refreshTimer);
this.getAttachments.unsubscribe();
}
}
packages/frontend/core/src/modules/workspace-indexer-embedding/stores/embedding.ts
+12 -12
View File
@@ -1,12 +1,12 @@
import type { WorkspaceServerService } from '@affine/core/modules/cloud';
import type { NbstoreService } from '@affine/core/modules/storage';
import {
addWorkspaceEmbeddingFilesMutation,
addWorkspaceArtifactMutation,
addWorkspaceEmbeddingIgnoredDocsMutation,
getAllWorkspaceEmbeddingIgnoredDocsQuery,
getWorkspaceEmbeddingFilesQuery,
getWorkspaceArtifactsQuery,
type PaginationInput,
removeWorkspaceEmbeddingFilesMutation,
removeWorkspaceArtifactMutation,
removeWorkspaceEmbeddingIgnoredDocsMutation,
setEnableDocEmbeddingMutation,
} from '@affine/graphql';
@@ -104,7 +104,7 @@ export class EmbeddingStore extends Store {
}
await this.workspaceServerService.server.gql({
query: addWorkspaceEmbeddingFilesMutation,
query: addWorkspaceArtifactMutation,
variables: {
workspaceId,
blob,
@@ -125,7 +125,7 @@ export class EmbeddingStore extends Store {
async removeEmbeddingFile(
workspaceId: string,
fileId: string,
artifactId: string,
signal?: AbortSignal
) {
if (!this.workspaceServerService.server) {
@@ -133,10 +133,10 @@ export class EmbeddingStore extends Store {
}
await this.workspaceServerService.server.gql({
query: removeWorkspaceEmbeddingFilesMutation,
query: removeWorkspaceArtifactMutation,
variables: {
workspaceId,
fileId,
artifactId,
},
context: { signal },
});
@@ -144,11 +144,11 @@ export class EmbeddingStore extends Store {
async removeEmbeddingFiles(
workspaceId: string,
fileIds: string[],
artifactIds: string[],
signal?: AbortSignal
) {
for (const fileId of fileIds) {
await this.removeEmbeddingFile(workspaceId, fileId, signal);
for (const artifactId of artifactIds) {
await this.removeEmbeddingFile(workspaceId, artifactId, signal);
}
}
@@ -162,14 +162,14 @@ export class EmbeddingStore extends Store {
}
const data = await this.workspaceServerService.server.gql({
query: getWorkspaceEmbeddingFilesQuery,
query: getWorkspaceArtifactsQuery,
variables: {
workspaceId,
pagination,
},
context: { signal },
});
return data.workspace.embedding.files;
return data.workspace.embedding.artifacts;
}
async getEmbeddingProgress(workspaceId: string, signal?: AbortSignal) {
packages/frontend/core/src/modules/workspace-indexer-embedding/types.ts
+3 -3
View File
@@ -1,10 +1,10 @@
export interface PersistedAttachmentFile {
fileId: string;
artifactId: string;
fileName: string;
mimeType: string;
mediaType: string;
size: number;
createdAt: string;
status: 'uploaded';
status: 'processing' | 'ready' | 'failed';
}
export interface LocalAttachmentFile {
packages/frontend/core/src/modules/workspace-indexer-embedding/utils.ts
+2 -2
View File
@@ -11,7 +11,7 @@ import type {
export function isPersistedAttachment(
attachment: AttachmentFile
): attachment is PersistedAttachmentFile {
return 'fileId' in attachment;
return 'artifactId' in attachment;
}
export function isErrorAttachment(
@@ -34,7 +34,7 @@ export function isLocalAttachment(
export function getAttachmentId(attachment: AttachmentFile): string {
if (isPersistedAttachment(attachment)) {
return attachment.fileId;
return attachment.artifactId;
}
return attachment.localId;
}
packages/frontend/core/src/modules/workspace-indexer-embedding/view/attachments.tsx
+9 -3
View File
@@ -75,13 +75,19 @@ const ErrorItem: React.FC<{ attachment: ErrorAttachmentFile }> = ({
const PersistedItem: React.FC<{ attachment: PersistedAttachmentFile }> = ({
attachment,
}) => {
const Icon = getAttachmentFileIconRC(attachment.mimeType);
const Icon = getAttachmentFileIconRC(attachment.mediaType);
return (
<div
className={attachmentTitle}
data-testid="workspace-embedding-setting-attachment-persisted-item"
data-testid={`workspace-embedding-setting-attachment-${attachment.status}-item`}
>
<Icon style={{ marginRight: 4 }} />
{attachment.status === 'processing' ? (
<Loading />
) : attachment.status === 'failed' ? (
<WarningIcon />
) : (
<Icon style={{ marginRight: 4 }} />
)}
<span className="attachment-title-text">{attachment.fileName}</span>
</div>
);