diff --git a/packages/common/graphql/src/graphql/copilot-context-list-object.gql b/packages/common/graphql/src/graphql/copilot-context-list-object.gql index b1167a9727..3f102ce0e8 100644 --- a/packages/common/graphql/src/graphql/copilot-context-list-object.gql +++ b/packages/common/graphql/src/graphql/copilot-context-list-object.gql @@ -22,18 +22,22 @@ query listContextObject( createdAt } tags { + type id docs { id status + createdAt } createdAt } collections { + type id docs { id status + createdAt } createdAt } diff --git a/packages/common/graphql/src/graphql/index.ts b/packages/common/graphql/src/graphql/index.ts index bd6ce13cb5..a7448b193b 100644 --- a/packages/common/graphql/src/graphql/index.ts +++ b/packages/common/graphql/src/graphql/index.ts @@ -255,18 +255,22 @@ export const listContextObjectQuery = { createdAt } tags { + type id docs { id status + createdAt } createdAt } collections { + type id docs { id status + createdAt } createdAt } diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts index 6bed265434..f6345629eb 100644 --- a/packages/common/graphql/src/schema.ts +++ b/packages/common/graphql/src/schema.ts @@ -2593,22 +2593,26 @@ export type ListContextObjectQuery = { }>; tags: Array<{ __typename?: 'CopilotContextCategory'; + type: ContextCategories; id: string; createdAt: number; docs: Array<{ __typename?: 'CopilotDocType'; id: string; status: ContextEmbedStatus | null; + createdAt: number; }>; }>; collections: Array<{ __typename?: 'CopilotContextCategory'; + type: ContextCategories; id: string; createdAt: number; docs: Array<{ __typename?: 'CopilotDocType'; id: string; status: ContextEmbedStatus | null; + createdAt: number; }>; }>; }>; diff --git a/packages/frontend/core/src/blocksuite/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/ai/actions/types.ts index 35a9df92c4..4f38f6c028 100644 --- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts +++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts @@ -1,6 +1,7 @@ import type { ChatHistoryOrder, ContextMatchedFileChunk, + CopilotContextCategory, CopilotContextDoc, CopilotContextFile, CopilotSessionType, @@ -245,6 +246,8 @@ declare global { type AIDocsAndFilesContext = { docs: CopilotContextDoc[]; files: CopilotContextFile[]; + tags: CopilotContextCategory[]; + collections: CopilotContextCategory[]; }; interface AIContextService { @@ -275,6 +278,24 @@ declare global { contextId: string; fileId: string; }) => Promise; + addContextTag: (options: { + contextId: string; + tagId: string; + docIds: string[]; + }) => Promise; + removeContextTag: (options: { + contextId: string; + tagId: string; + }) => Promise; + addContextCollection: (options: { + contextId: string; + collectionId: string; + docIds: string[]; + }) => Promise; + removeContextCollection: (options: { + contextId: string; + collectionId: string; + }) => Promise; getContextDocsAndFiles: ( workspaceId: string, sessionId: string, diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts index 3b8cbb7398..95bf18d5d3 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts @@ -1,8 +1,10 @@ +import type { TagMeta } from '@affine/core/components/page-list'; import type { SearchCollectionMenuAction, SearchDocMenuAction, SearchTagMenuAction, } from '@affine/core/modules/search-menu/services'; +import type { Collection } from '@affine/env/filter'; import type { LinkedMenuGroup } from '@blocksuite/affine/blocks/root'; import type { Store } from '@blocksuite/affine/store'; import type { Signal } from '@preact/signals-core'; @@ -40,6 +42,15 @@ export interface DocDisplayConfig { >; cleanup: () => void; }; + getTags: () => { + signal: Signal; + cleanup: () => void; + }; + getTagPageIds: (tagId: string) => string[]; + getCollections: () => { + signal: Signal; + cleanup: () => void; + }; } export interface SearchMenuConfig { diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts index c208826dad..9f9d7bee58 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts @@ -83,6 +83,7 @@ export interface BaseChip { */ state: ChipState; tooltip?: string | null; + createdAt?: number | null; } export interface DocChip extends BaseChip { @@ -99,13 +100,10 @@ export interface FileChip extends BaseChip { export interface TagChip extends BaseChip { tagId: string; - tagName: string; - tagColor: string; } export interface CollectionChip extends BaseChip { collectionId: string; - collectionName: string; } -export type ChatChip = DocChip | FileChip; +export type ChatChip = DocChip | FileChip | TagChip | CollectionChip; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts index ff6cd59916..580c42a516 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts @@ -1,3 +1,5 @@ +import type { TagMeta } from '@affine/core/components/page-list'; +import type { Collection } from '@affine/env/filter'; import { type EditorHost, ShadowlessElement, @@ -17,14 +19,18 @@ import type { DocDisplayConfig, SearchMenuConfig } from './chat-config'; import type { ChatChip, ChatContextValue, + CollectionChip, DocChip, FileChip, + TagChip, } from './chat-context'; import { estimateTokenCount, getChipKey, + isCollectionChip, isDocChip, isFileChip, + isTagChip, } from './components/utils'; // 100k tokens limit for the docs context @@ -95,6 +101,10 @@ export class ChatPanelChips extends SignalWatcher( }> > = signal([]); + private _tags: Signal = signal([]); + + private _collections: Signal = signal([]); + private _cleanup: (() => void) | null = null; private _docIds: string[] = []; @@ -133,6 +143,30 @@ export class ChatPanelChips extends SignalWatcher( .removeChip=${this._removeChip} >`; } + if (isTagChip(chip)) { + const tag = this._tags.value.find(tag => tag.id === chip.tagId); + if (!tag) { + return null; + } + return html``; + } + if (isCollectionChip(chip)) { + const collection = this._collections.value.find( + collection => collection.id === chip.collectionId + ); + if (!collection) { + return null; + } + return html``; + } return null; } )} @@ -144,6 +178,17 @@ export class ChatPanelChips extends SignalWatcher( `; } + override connectedCallback(): void { + super.connectedCallback(); + const tags = this.docDisplayConfig.getTags(); + this._tags = tags.signal; + this._disposables.add(tags.cleanup); + + const collections = this.docDisplayConfig.getCollections(); + this._collections = collections.signal; + this._disposables.add(collections.cleanup); + } + protected override updated(_changedProperties: PropertyValues): void { if ( _changedProperties.has('chatContextValue') && @@ -204,15 +249,8 @@ export class ChatPanelChips extends SignalWatcher( private readonly _addChip = async (chip: ChatChip) => { this.isCollapsed = false; - // remove the chip if it already exists - const chips = this.chatContextValue.chips.filter(item => { - if (isDocChip(chip)) { - return !isDocChip(item) || item.docId !== chip.docId; - } else { - return !isFileChip(item) || item.file !== chip.file; - } - }); + const chips = this._omitChip(this.chatContextValue.chips, chip); this.updateContext({ chips: [...chips, chip], }); @@ -227,13 +265,10 @@ export class ChatPanelChips extends SignalWatcher( chip: ChatChip, options: Partial ) => { - const index = this.chatContextValue.chips.findIndex(item => { - if (isDocChip(chip)) { - return isDocChip(item) && item.docId === chip.docId; - } else { - return isFileChip(item) && item.file === chip.file; - } - }); + const index = this._findChipIndex(this.chatContextValue.chips, chip); + if (index === -1) { + return; + } const nextChip: ChatChip = { ...chip, ...options, @@ -248,21 +283,13 @@ export class ChatPanelChips extends SignalWatcher( }; private readonly _removeChip = async (chip: ChatChip) => { - if (isDocChip(chip)) { - this.updateContext({ - chips: this.chatContextValue.chips.filter(item => { - return !isDocChip(item) || item.docId !== chip.docId; - }), - }); + const chips = this._omitChip(this.chatContextValue.chips, chip); + this.updateContext({ + chips, + }); + if (chips.length < this.chatContextValue.chips.length) { + await this._removeFromContext(chip); } - if (isFileChip(chip)) { - this.updateContext({ - chips: this.chatContextValue.chips.filter(item => { - return !isFileChip(item) || item.file !== chip.file; - }), - }); - } - await this._removeFromContext(chip); }; private readonly _addToContext = async (chip: ChatChip) => { @@ -271,29 +298,111 @@ export class ChatPanelChips extends SignalWatcher( return; } if (isDocChip(chip)) { + return await this._addDocToContext(chip); + } + if (isFileChip(chip)) { + return await this._addFileToContext(chip); + } + if (isTagChip(chip)) { + return await this._addTagToContext(chip); + } + if (isCollectionChip(chip)) { + return await this._addCollectionToContext(chip); + } + return null; + }; + + private readonly _addDocToContext = async (chip: DocChip) => { + const contextId = await this.getContextId(); + if (!contextId || !AIProvider.context) { + return; + } + try { await AIProvider.context.addContextDoc({ contextId, docId: chip.docId, }); + } catch (e) { + this._updateChip(chip, { + state: 'failed', + tooltip: e instanceof Error ? e.message : 'Add context doc error', + }); } - if (isFileChip(chip)) { - try { - const blobId = await this.host.doc.blobSync.set(chip.file); - const contextFile = await AIProvider.context.addContextFile(chip.file, { - contextId, - blobId, - }); - this._updateChip(chip, { - state: contextFile.status, - blobId: contextFile.blobId, - fileId: contextFile.id, - }); - } catch (e) { - this._updateChip(chip, { - state: 'failed', - tooltip: e instanceof Error ? e.message : 'Add context file error', - }); - } + }; + + private readonly _addFileToContext = async (chip: FileChip) => { + const contextId = await this.getContextId(); + if (!contextId || !AIProvider.context) { + return; + } + try { + const blobId = await this.host.doc.blobSync.set(chip.file); + const contextFile = await AIProvider.context.addContextFile(chip.file, { + contextId, + blobId, + }); + this._updateChip(chip, { + state: contextFile.status, + blobId: contextFile.blobId, + fileId: contextFile.id, + }); + } catch (e) { + this._updateChip(chip, { + state: 'failed', + tooltip: e instanceof Error ? e.message : 'Add context file error', + }); + } + }; + + private readonly _addTagToContext = async (chip: TagChip) => { + const contextId = await this.getContextId(); + if (!contextId || !AIProvider.context) { + return; + } + try { + // TODO: server side docIds calculation + const docIds = this.docDisplayConfig.getTagPageIds(chip.tagId); + await AIProvider.context.addContextTag({ + contextId, + tagId: chip.tagId, + docIds, + }); + this._updateChip(chip, { + state: 'finished', + }); + } catch (e) { + this._updateChip(chip, { + state: 'failed', + tooltip: e instanceof Error ? e.message : 'Add context tag error', + }); + } + }; + + private readonly _addCollectionToContext = async (chip: CollectionChip) => { + const contextId = await this.getContextId(); + if (!contextId || !AIProvider.context) { + return; + } + try { + const collection = this._collections.value.find( + collection => collection.id === chip.collectionId + ); + // TODO: server side docIds calculation + const docIds = collection?.allowList ?? []; + await AIProvider.context.addContextCollection({ + contextId, + collectionId: chip.collectionId, + docIds, + }); + this._updateChip(chip, { + state: 'finished', + }); + } catch (e) { + this._updateChip(chip, { + state: 'failed', + tooltip: + e instanceof Error ? e.message : 'Add context collection error', + }); } }; @@ -316,6 +425,18 @@ export class ChatPanelChips extends SignalWatcher( fileId: chip.fileId, }); } + if (isTagChip(chip)) { + return await AIProvider.context.removeContextTag({ + contextId, + tagId: chip.tagId, + }); + } + if (isCollectionChip(chip)) { + return await AIProvider.context.removeContextCollection({ + contextId, + collectionId: chip.collectionId, + }); + } return true; }; @@ -324,7 +445,7 @@ export class ChatPanelChips extends SignalWatcher( newTokenCount: number ) => { const estimatedTokens = this.chatContextValue.chips.reduce((acc, chip) => { - if (isFileChip(chip)) { + if (isFileChip(chip) || isTagChip(chip) || isCollectionChip(chip)) { return acc; } if (chip.docId === newChip.docId) { @@ -355,4 +476,44 @@ export class ChatPanelChips extends SignalWatcher( this.referenceDocs = signal; this._cleanup = cleanup; }; + + private readonly _omitChip = (chips: ChatChip[], chip: ChatChip) => { + return chips.filter(item => { + if (isDocChip(chip)) { + return !isDocChip(item) || item.docId !== chip.docId; + } + if (isFileChip(chip)) { + return !isFileChip(item) || item.file !== chip.file; + } + if (isTagChip(chip)) { + return !isTagChip(item) || item.tagId !== chip.tagId; + } + if (isCollectionChip(chip)) { + return ( + !isCollectionChip(item) || item.collectionId !== chip.collectionId + ); + } + return true; + }); + }; + + private readonly _findChipIndex = (chips: ChatChip[], chip: ChatChip) => { + return chips.findIndex(item => { + if (isDocChip(chip)) { + return isDocChip(item) && item.docId === chip.docId; + } + if (isFileChip(chip)) { + return isFileChip(item) && item.file === chip.file; + } + if (isTagChip(chip)) { + return isTagChip(item) && item.tagId === chip.tagId; + } + if (isCollectionChip(chip)) { + return ( + isCollectionChip(item) && item.collectionId === chip.collectionId + ); + } + return -1; + }); + }; } diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/add-popover.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/add-popover.ts index b2b8ff6cfb..84df4b6181 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/add-popover.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/add-popover.ts @@ -453,11 +453,19 @@ export class ChatPanelAddPopover extends SignalWatcher( this.abortController.abort(); }; - private readonly _addTagChip = (_tag: TagMeta) => { + private readonly _addTagChip = (tag: TagMeta) => { + this.addChip({ + tagId: tag.id, + state: 'processing', + }); this.abortController.abort(); }; - private readonly _addCollectionChip = (_collection: CollectionMeta) => { + private readonly _addCollectionChip = (collection: CollectionMeta) => { + this.addChip({ + collectionId: collection.id, + state: 'processing', + }); this.abortController.abort(); }; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/collection-chip.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/collection-chip.ts index 31379e7990..3567daa080 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/collection-chip.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/collection-chip.ts @@ -1,3 +1,4 @@ +import type { Collection } from '@affine/env/filter'; import { ShadowlessElement } from '@blocksuite/affine/block-std'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; import { CollectionsIcon } from '@blocksuite/icons/lit'; @@ -13,19 +14,31 @@ export class ChatPanelCollectionChip extends SignalWatcher( @property({ attribute: false }) accessor chip!: CollectionChip; + @property({ attribute: false }) + accessor removeChip!: (chip: CollectionChip) => void; + + @property({ attribute: false }) + accessor collection!: Collection; + override render() { - const { state, collectionName } = this.chip; + const { state } = this.chip; + const { name } = this.collection; const isLoading = state === 'processing'; - const tooltip = getChipTooltip(state, collectionName, this.chip.tooltip); + const tooltip = getChipTooltip(state, name, this.chip.tooltip); const collectionIcon = CollectionsIcon(); const icon = getChipIcon(state, collectionIcon); return html``; } + + private readonly onChipDelete = () => { + this.removeChip(this.chip); + }; } diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts index cf4684646b..0163ed5756 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts @@ -10,7 +10,7 @@ import throttle from 'lodash-es/throttle'; import { extractMarkdownFromDoc } from '../../utils/extract'; import type { DocDisplayConfig } from '../chat-config'; -import type { ChatChip, DocChip } from '../chat-context'; +import type { DocChip } from '../chat-context'; import { estimateTokenCount, getChipIcon, getChipTooltip } from './utils'; const EXTRACT_DOC_THROTTLE = 1000; @@ -22,13 +22,13 @@ export class ChatPanelDocChip extends SignalWatcher( accessor chip!: DocChip; @property({ attribute: false }) - accessor addChip!: (chip: ChatChip) => void; + accessor addChip!: (chip: DocChip) => void; @property({ attribute: false }) - accessor updateChip!: (chip: ChatChip, options: Partial) => void; + accessor updateChip!: (chip: DocChip, options: Partial) => void; @property({ attribute: false }) - accessor removeChip!: (chip: ChatChip) => void; + accessor removeChip!: (chip: DocChip) => void; @property({ attribute: false }) accessor checkTokenLimit!: ( diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/file-chip.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/file-chip.ts index 7a8feef65b..9ec37deaf9 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/file-chip.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/file-chip.ts @@ -4,7 +4,7 @@ import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; import { html } from 'lit'; import { property } from 'lit/decorators.js'; -import type { ChatChip, FileChip } from '../chat-context'; +import type { FileChip } from '../chat-context'; import { getChipIcon, getChipTooltip } from './utils'; export class ChatPanelFileChip extends SignalWatcher( @@ -14,7 +14,7 @@ export class ChatPanelFileChip extends SignalWatcher( accessor chip!: FileChip; @property({ attribute: false }) - accessor removeChip!: (chip: ChatChip) => void; + accessor removeChip!: (chip: FileChip) => void; override render() { const { state, file } = this.chip; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/tag-chip.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/tag-chip.ts index 2fea7e9058..3acc9e9d2d 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/tag-chip.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/tag-chip.ts @@ -1,3 +1,4 @@ +import type { TagMeta } from '@affine/core/components/page-list'; import { ShadowlessElement } from '@blocksuite/affine/block-std'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; import { css, html } from 'lit'; @@ -27,23 +28,35 @@ export class ChatPanelTagChip extends SignalWatcher( @property({ attribute: false }) accessor chip!: TagChip; + @property({ attribute: false }) + accessor removeChip!: (chip: TagChip) => void; + + @property({ attribute: false }) + accessor tag!: TagMeta; + override render() { - const { state, tagName, tagColor } = this.chip; + const { state } = this.chip; + const { title, color } = this.tag; const isLoading = state === 'processing'; - const tooltip = getChipTooltip(state, tagName, this.chip.tooltip); + const tooltip = getChipTooltip(state, title, this.chip.tooltip); const tagIcon = html`
-
+
`; const icon = getChipIcon(state, tagIcon); return html``; } + + private readonly onChipDelete = () => { + this.removeChip(this.chip); + }; } diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts index 5df05e501f..a88a49af4c 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts @@ -1,9 +1,15 @@ -import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql'; import { LoadingIcon } from '@blocksuite/affine/blocks/image'; import { WarningIcon } from '@blocksuite/icons/lit'; import { type TemplateResult } from 'lit'; -import type { ChatChip, ChipState, DocChip, FileChip } from '../chat-context'; +import type { + ChatChip, + ChipState, + CollectionChip, + DocChip, + FileChip, + TagChip, +} from '../chat-context'; export function getChipTooltip( state: ChipState, @@ -48,16 +54,12 @@ export function isFileChip(chip: ChatChip): chip is FileChip { return 'file' in chip && chip.file instanceof File; } -export function isDocContext( - context: CopilotContextDoc | CopilotContextFile -): context is CopilotContextDoc { - return !('blobId' in context); +export function isTagChip(chip: ChatChip): chip is TagChip { + return 'tagId' in chip; } -export function isFileContext( - context: CopilotContextDoc | CopilotContextFile -): context is CopilotContextFile { - return 'blobId' in context; +export function isCollectionChip(chip: ChatChip): chip is CollectionChip { + return 'collectionId' in chip; } export function getChipKey(chip: ChatChip) { @@ -65,7 +67,13 @@ export function getChipKey(chip: ChatChip) { return chip.docId; } if (isFileChip(chip)) { - return chip.fileId; + return chip.file.name; + } + if (isTagChip(chip)) { + return chip.tagId; + } + if (isCollectionChip(chip)) { + return chip.collectionId; } return null; } diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts index a8c3eab42e..94a3853e01 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts @@ -1,7 +1,12 @@ import './chat-panel-input'; import './chat-panel-messages'; -import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql'; +import type { + ContextEmbedStatus, + CopilotContextDoc, + CopilotContextFile, + CopilotDocType, +} from '@affine/graphql'; import type { EditorHost } from '@blocksuite/affine/block-std'; import { ShadowlessElement } from '@blocksuite/affine/block-std'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; @@ -32,11 +37,13 @@ import type { ChatChip, ChatContextValue, ChatItem, + CollectionChip, DocChip, FileChip, + TagChip, } from './chat-context'; import type { ChatPanelMessages } from './chat-panel-messages'; -import { isDocChip, isDocContext } from './components/utils'; +import { isCollectionChip, isDocChip, isTagChip } from './components/utils'; const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = { quote: '', @@ -183,44 +190,61 @@ export class ChatPanel extends SignalWatcher( } // context initialized, show the chips - const { docs = [], files = [] } = - (await AIProvider.context?.getContextDocsAndFiles( - this.doc.workspace.id, - this._chatSessionId, - this._chatContextId - )) || {}; - const list = [...docs, ...files].sort( - (a, b) => - new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime() - ); - const chips: ChatChip[] = await Promise.all( - list.map(async item => { - if (isDocContext(item)) { - return { - docId: item.id, - state: item.status || 'processing', - } as DocChip; - } - const file = await this.host.doc.blobSync.get(item.blobId); - if (!file) { - return { - blobId: item.id, - file: new File([], item.name), - state: 'failed', - tooltip: 'File not found in blob storage', - } as FileChip; - } else { - return { - file: new File([file], item.name), - blobId: item.blobId, - fileId: item.id, - state: item.status, - tooltip: item.error, - } as FileChip; - } + const { + docs = [], + files = [], + tags = [], + collections = [], + } = (await AIProvider.context?.getContextDocsAndFiles( + this.doc.workspace.id, + this._chatSessionId, + this._chatContextId + )) || {}; + + const docChips: DocChip[] = docs.map(doc => ({ + docId: doc.id, + state: doc.status || 'processing', + tooltip: doc.error, + createdAt: doc.createdAt, + })); + + const fileChips: FileChip[] = await Promise.all( + files.map(async file => { + const blob = await this.host.doc.blobSync.get(file.blobId); + return { + file: new File(blob ? [blob] : [], file.name), + blobId: file.blobId, + fileId: file.id, + state: blob ? file.status : 'failed', + tooltip: blob ? file.error : 'File not found in blob storage', + createdAt: file.createdAt, + }; }) ); + const tagChips: TagChip[] = tags.map(tag => ({ + tagId: tag.id, + state: 'finished', + createdAt: tag.createdAt, + })); + + const collectionChips: CollectionChip[] = collections.map(collection => ({ + collectionId: collection.id, + state: 'finished', + createdAt: collection.createdAt, + })); + + const chips: ChatChip[] = [ + ...docChips, + ...fileChips, + ...tagChips, + ...collectionChips, + ].sort((a, b) => { + const aTime = a.createdAt ?? Date.now(); + const bTime = b.createdAt ?? Date.now(); + return aTime - bTime; + }); + this.chatContextValue = { ...this.chatContextValue, chips, @@ -385,38 +409,55 @@ export class ChatPanel extends SignalWatcher( this._abortPoll(); return; } - const { docs = [], files = [] } = result; - const hashMap = new Map(); - const totalCount = docs.length + files.length; - let processingCount = 0; + const { + docs: sDocs = [], + files = [], + tags = [], + collections = [], + } = result; + const docs = [ + ...sDocs, + ...tags.flatMap(tag => tag.docs), + ...collections.flatMap(collection => collection.docs), + ]; + const hashMap = new Map< + string, + CopilotContextDoc | CopilotDocType | CopilotContextFile + >(); + const count: Record = { + finished: 0, + processing: 0, + failed: 0, + }; docs.forEach(doc => { hashMap.set(doc.id, doc); - if (doc.status === 'processing') { - processingCount++; - } + doc.status && count[doc.status]++; }); files.forEach(file => { hashMap.set(file.id, file); - if (file.status === 'processing') { - processingCount++; - } + file.status && count[file.status]++; }); const nextChips = this.chatContextValue.chips.map(chip => { + if (isTagChip(chip) || isCollectionChip(chip)) { + return chip; + } const id = isDocChip(chip) ? chip.docId : chip.fileId; const item = id && hashMap.get(id); if (item && item.status) { return { ...chip, state: item.status, + tooltip: 'error' in item ? item.error : undefined, }; } return chip; }); + const total = count.finished + count.processing + count.failed; this.updateContext({ chips: nextChips, - embeddingProgress: [totalCount - processingCount, totalCount], + embeddingProgress: [count.finished, total], }); - if (processingCount === 0) { + if (count.processing === 0) { this._abortPoll(); } }; diff --git a/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts b/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts index 0bba0a0181..4916de9ed1 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts +++ b/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts @@ -6,6 +6,7 @@ import { import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-login-required'; import type { UserFriendlyError } from '@affine/error'; import { + addContextCategoryMutation, addContextDocMutation, addContextFileMutation, cleanupCopilotSessionMutation, @@ -22,6 +23,7 @@ import { matchContextQuery, type QueryOptions, type QueryResponse, + removeContextCategoryMutation, removeContextDocMutation, removeContextFileMutation, type RequestOptions, @@ -290,6 +292,30 @@ export class CopilotClient { return res.removeContextFile; } + async addContextCategory( + options: OptionsField + ) { + const res = await this.gql({ + query: addContextCategoryMutation, + variables: { + options, + }, + }); + return res.addContextCategory; + } + + async removeContextCategory( + options: OptionsField + ) { + const res = await this.gql({ + query: removeContextCategoryMutation, + variables: { + options, + }, + }); + return res.removeContextCategory; + } + async getContextDocsAndFiles( workspaceId: string, sessionId: string, diff --git a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx index 88179a6011..652cf4ce18 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx +++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx @@ -2,6 +2,7 @@ import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onb import type { GlobalDialogService } from '@affine/core/modules/dialogs'; import { type ChatHistoryOrder, + ContextCategories, type getCopilotHistoriesQuery, type RequestOptions, } from '@affine/graphql'; @@ -444,6 +445,47 @@ Could you make a new website based on these notes and send back just the html fi }) => { return client.removeContextFile(options); }, + addContextTag: async (options: { + contextId: string; + tagId: string; + docIds: string[]; + }) => { + return client.addContextCategory({ + contextId: options.contextId, + type: ContextCategories.Tag, + categoryId: options.tagId, + docs: options.docIds, + }); + }, + removeContextTag: async (options: { contextId: string; tagId: string }) => { + return client.removeContextCategory({ + contextId: options.contextId, + type: ContextCategories.Tag, + categoryId: options.tagId, + }); + }, + addContextCollection: async (options: { + contextId: string; + collectionId: string; + docIds: string[]; + }) => { + return client.addContextCategory({ + contextId: options.contextId, + type: ContextCategories.Collection, + categoryId: options.collectionId, + docs: options.docIds, + }); + }, + removeContextCollection: async (options: { + contextId: string; + collectionId: string; + }) => { + return client.removeContextCategory({ + contextId: options.contextId, + type: ContextCategories.Collection, + categoryId: options.collectionId, + }); + }, getContextDocsAndFiles: async ( workspaceId: string, sessionId: string, diff --git a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx index d49619b312..ee32fcab99 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx +++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx @@ -2,9 +2,11 @@ import { ChatPanel } from '@affine/core/blocksuite/ai'; import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-editor'; import { enableFootnoteConfigExtension } from '@affine/core/blocksuite/extensions'; import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search'; +import { CollectionService } from '@affine/core/modules/collection'; import { DocDisplayMetaService } from '@affine/core/modules/doc-display-meta'; import { DocsSearchService } from '@affine/core/modules/docs-search'; import { SearchMenuService } from '@affine/core/modules/search-menu/services'; +import { TagService } from '@affine/core/modules/tag'; import { WorkbenchService } from '@affine/core/modules/workbench'; import { WorkspaceService } from '@affine/core/modules/workspace'; import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference'; @@ -62,7 +64,8 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel( const searchMenuService = framework.get(SearchMenuService); const workbench = framework.get(WorkbenchService).workbench; const docsSearchService = framework.get(DocsSearchService); - + const tagService = framework.get(TagService); + const collectionService = framework.get(CollectionService); chatPanelRef.current.appSidebarConfig = { getWidth: () => { const width$ = workbench.sidebarWidth$; @@ -96,6 +99,19 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel( const docs$ = docsSearchService.watchRefsFrom(docIds); return createSignalFromObservable(docs$, []); }, + getTags: () => { + const tagMetas$ = tagService.tagList.tagMetas$; + return createSignalFromObservable(tagMetas$, []); + }, + getTagPageIds: (tagId: string) => { + const tag$ = tagService.tagList.tagByTagId$(tagId); + if (!tag$) return []; + return tag$.value?.pageIds$.value ?? []; + }, + getCollections: () => { + const collections$ = collectionService.collections$; + return createSignalFromObservable(collections$, []); + }, }; chatPanelRef.current.searchMenuConfig = {