feat(core): add ai file context api (#10842)

Close [BS-2349](https://linear.app/affine-design/issue/BS-2349). ### What Changed? - Add file context graphql apis - Pass matched file chunks to LLM [录屏2025-02-19 23.27.47.mov <span class="graphite__hidden">(uploaded via Graphite)</span> <img class="graphite__hidden" src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/8e8a98ca-6959-4bb6-9759-b51d97cede49.mov" />](https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/8e8a98ca-6959-4bb6-9759-b51d97cede49.mov)
2026-02-13 21:05:19 +00:00 · 2025-03-14 04:29:54 +00:00
parent 8880cef20b
commit daccb2c865
18 changed files with 251 additions and 99 deletions
--- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts
+++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts
@@ -1,5 +1,6 @@
 import type {
  ChatHistoryOrder,
+  ContextMatchedFileChunk,
  CopilotContextDoc,
  CopilotContextFile,
  CopilotSessionType,
@@ -10,7 +11,7 @@ import type { EditorHost } from '@blocksuite/affine/block-std';
 import type { GfxModel } from '@blocksuite/affine/block-std/gfx';
 import type { BlockModel } from '@blocksuite/affine/store';

-import type { DocContext } from '../chat-panel/chat-context';
+import type { DocContext, FileContext } from '../chat-panel/chat-context';

 export const translateLangs = [
  'English',
@@ -114,7 +115,10 @@ declare global {
    interface ChatOptions extends AITextActionOptions {
      sessionId?: string;
      isRootSession?: boolean;
-      docs?: DocContext[];
+      contexts?: {
+        docs: DocContext[];
+        files: FileContext[];
+      };
    }

    interface TranslateOptions extends AITextActionOptions {
@@ -250,19 +254,22 @@ declare global {
      addContextDoc: (options: {
        contextId: string;
        docId: string;
-      }) => Promise<{ id: string; createdAt: number }>;
+      }) => Promise<CopilotContextDoc>;
      removeContextDoc: (options: {
        contextId: string;
        docId: string;
      }) => Promise<boolean>;
-      addContextFile: (options: {
-        contextId: string;
-        fileId: string;
-      }) => Promise<void>;
+      addContextFile: (
+        file: File,
+        options: {
+          contextId: string;
+          blobId: string;
+        }
+      ) => Promise<CopilotContextFile>;
      removeContextFile: (options: {
        contextId: string;
        fileId: string;
-      }) => Promise<void>;
+      }) => Promise<boolean>;
      getContextDocsAndFiles: (
        workspaceId: string,
        sessionId: string,
@@ -274,6 +281,11 @@ declare global {
          }
        | undefined
      >;
+      matchContext: (
+        contextId: string,
+        content: string,
+        limit?: number
+      ) => Promise<ContextMatchedFileChunk[] | undefined>;
    }

    // TODO(@Peng): should be refactored to get rid of implement details (like messages, action, role, etc.)
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts
@@ -36,11 +36,18 @@ export type ChatStatus =

 export interface DocContext {
  docId: string;
-  plaintext?: string;
-  markdown?: string;
-  images?: File[];
+  refIndex: number;
+  markdown: string;
 }

+export type FileContext = {
+  blobId: string;
+  refIndex: number;
+  fileName: string;
+  fileType: string;
+  chunks: string;
+};
+
 export type ChatContextValue = {
  // history messages of the chat
  items: ChatItem[];
@@ -73,19 +80,19 @@ export interface BaseChip {
   * failed: the chip is failed to process
   */
  state: ChipState;
-  tooltip?: string;
+  tooltip?: string | null;
 }

 export interface DocChip extends BaseChip {
  docId: string;
-  markdown?: Signal<string>;
-  tokenCount?: number;
+  markdown?: Signal<string> | null;
+  tokenCount?: number | null;
 }

 export interface FileChip extends BaseChip {
-  fileName: string;
-  fileId: string;
-  fileType: string;
+  file: File;
+  fileId?: string | null;
+  blobId?: string | null;
 }

 export interface TagChip extends BaseChip {
--- 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
@@ -112,6 +112,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
          if (isFileChip(chip)) {
            return html`<chat-panel-file-chip
              .chip=${chip}
+              .removeChip=${this._removeChip}
            ></chat-panel-file-chip>`;
          }
          return null;
@@ -174,14 +175,15 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
  };

  private readonly _addChip = async (chip: ChatChip) => {
+    this.isCollapsed = false;
    if (
      this.chatContextValue.chips.length === 1 &&
      this.chatContextValue.chips[0].state === 'candidate'
    ) {
-      await this._addToContext(chip);
      this.updateContext({
        chips: [chip],
      });
+      await this._addToContext(chip);
      return;
    }
    // remove the chip if it already exists
@@ -189,16 +191,16 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
      if (isDocChip(chip)) {
        return !isDocChip(item) || item.docId !== chip.docId;
      } else {
-        return !isFileChip(item) || item.fileId !== chip.fileId;
+        return !isFileChip(item) || item.file !== chip.file;
      }
    });
+    this.updateContext({
+      chips: [...chips, chip],
+    });
    if (chips.length < this.chatContextValue.chips.length) {
      await this._removeFromContext(chip);
    }
    await this._addToContext(chip);
-    this.updateContext({
-      chips: [...chips, chip],
-    });
  };

  private readonly _updateChip = (
@@ -209,7 +211,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
      if (isDocChip(chip)) {
        return isDocChip(item) && item.docId === chip.docId;
      } else {
-        return isFileChip(item) && item.fileId === chip.fileId;
+        return isFileChip(item) && item.file === chip.file;
      }
    });
    const nextChip: ChatChip = {
@@ -237,7 +239,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
      await this._removeFromContext(chip);
      this.updateContext({
        chips: this.chatContextValue.chips.filter(item => {
-          return !isFileChip(item) || item.fileId !== chip.fileId;
+          return !isFileChip(item) || item.file !== chip.file;
        }),
      });
    }
@@ -254,10 +256,23 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
        docId: chip.docId,
      });
    } else {
-      await AIProvider.context.addContextFile({
-        contextId,
-        fileId: chip.fileId,
-      });
+      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: 'success',
+          blobId: contextFile.blobId,
+          fileId: contextFile.id,
+        });
+      } catch (e) {
+        this._updateChip(chip, {
+          state: 'failed',
+          tooltip: e instanceof Error ? e.message : 'Add context file error',
+        });
+      }
    }
  };

@@ -271,7 +286,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
        contextId,
        docId: chip.docId,
      });
-    } else {
+    } else if (isFileChip(chip) && chip.fileId) {
      await AIProvider.context.removeContextFile({
        contextId,
        fileId: chip.fileId,
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts
@@ -19,8 +19,13 @@ import { AIProvider } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
 import { readBlobAsURL } from '../utils/image';
 import type { AINetworkSearchConfig } from './chat-config';
-import type { ChatContextValue, ChatMessage, DocContext } from './chat-context';
-import { isDocChip } from './components/utils';
+import type {
+  ChatContextValue,
+  ChatMessage,
+  DocContext,
+  FileContext,
+} from './chat-context';
+import { isDocChip, isFileChip } from './components/utils';
 import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';

 const MaximumImageCount = 32;
@@ -199,6 +204,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
  @property({ attribute: false })
  accessor getSessionId!: () => Promise<string | undefined>;

+  @property({ attribute: false })
+  accessor getContextId!: () => Promise<string | undefined>;
+
  @property({ attribute: false })
  accessor updateContext!: (context: Partial<ChatContextValue>) => void;

@@ -218,7 +226,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
  private get _isNetworkDisabled() {
    return (
      !!this.chatContextValue.images.length ||
-      !!this.chatContextValue.chips.filter(chip => chip.state !== 'candidate')
+      !!this.chatContextValue.chips.filter(chip => chip.state === 'success')
        .length
    );
  }
@@ -452,7 +460,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
  };

  send = async (text: string) => {
-    const { status, markdown, chips, images } = this.chatContextValue;
+    const { status, markdown, images } = this.chatContextValue;
    if (status === 'loading' || status === 'transmitting') return;
    if (!text) return;

@@ -498,17 +506,12 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {

      const abortController = new AbortController();
      const sessionId = await this.getSessionId();
-      const docs: DocContext[] = chips
-        .filter(isDocChip)
-        .filter(chip => !!chip.markdown?.value && chip.state === 'success')
-        .map(chip => ({
-          docId: chip.docId,
-          markdown: chip.markdown?.value || '',
-        }));
+
+      const contexts = await this._getMatchedContexts(userInput);
      const stream = AIProvider.actions.chat?.({
        sessionId,
        input: userInput,
-        docs: docs,
+        contexts,
        docId: doc.id,
        attachments: images,
        workspaceId: doc.workspace.id,
@@ -550,6 +553,44 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
      this.updateContext({ abortController: null });
    }
  };
+
+  private async _getMatchedContexts(userInput: string) {
+    const contextId = await this.getContextId();
+    const matched = contextId
+      ? (await AIProvider.context?.matchContext(contextId, userInput)) || []
+      : [];
+    const contexts = this.chatContextValue.chips.reduce(
+      (acc, chip, index) => {
+        if (chip.state !== 'success') {
+          return acc;
+        }
+        if (isDocChip(chip) && !!chip.markdown?.value) {
+          acc.docs.push({
+            docId: chip.docId,
+            refIndex: index + 1,
+            markdown: chip.markdown.value,
+          });
+        }
+        if (isFileChip(chip) && chip.blobId) {
+          const matchedChunks = matched
+            .filter(chunk => chunk.fileId === chip.fileId)
+            .map(chunk => chunk.content);
+          if (matchedChunks.length > 0) {
+            acc.files.push({
+              blobId: chip.blobId,
+              refIndex: index + 1,
+              fileName: chip.file.name,
+              fileType: chip.file.type,
+              chunks: matchedChunks.join('\n'),
+            });
+          }
+        }
+        return acc;
+      },
+      { docs: [], files: [] } as { docs: DocContext[]; files: FileContext[] }
+    );
+    return contexts;
+  }
 }

 declare global {
--- 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
@@ -1,8 +1,10 @@
+import { toast } from '@affine/component';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
-import type { LinkedMenuGroup } from '@blocksuite/affine/blocks/root';
+import { type LinkedMenuGroup } from '@blocksuite/affine/blocks/root';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { scrollbarStyle } from '@blocksuite/affine/shared/styles';
-import { SearchIcon } from '@blocksuite/icons/lit';
+import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
+import { SearchIcon, UploadIcon } from '@blocksuite/icons/lit';
 import type { DocMeta } from '@blocksuite/store';
 import { css, html } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
@@ -138,9 +140,35 @@ export class ChatPanelAddPopover extends SignalWatcher(
            })
          : html`<div class="no-result">No Result</div>`}
      </div>
+      <div class="divider"></div>
+      <div class="upload-wrapper">
+        <icon-button
+          width="280px"
+          height="30px"
+          data-id="upload"
+          .text=${'Upload files (pdf, txt, csv)'}
+          @click=${this._addFileChip}
+        >
+          ${UploadIcon()}
+        </icon-button>
+      </div>
    </div>`;
  }

+  private readonly _addFileChip = async () => {
+    const file = await openFileOrFiles();
+    if (!file) return;
+    if (file.size > 50 * 1024 * 1024) {
+      toast('You can only upload files less than 50MB');
+      return;
+    }
+    this.addChip({
+      file,
+      state: 'processing',
+    });
+    this.abortController.abort();
+  };
+
  private _onInput(event: Event) {
    this._query = (event.target as HTMLInputElement).value;
    this._updateDocGroup();
--- 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
@@ -125,7 +125,7 @@ export class ChatPanelDocChip extends SignalWatcher(
    } catch (e) {
      this.updateChip(this.chip, {
        state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Failed to process document',
+        tooltip: e instanceof Error ? e.message : 'Failed to extract markdown',
      });
    }
  };
--- 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 { FileChip } from '../chat-context';
+import type { ChatChip, FileChip } from '../chat-context';
 import { getChipIcon, getChipTooltip } from './utils';

 export class ChatPanelFileChip extends SignalWatcher(
@@ -13,19 +13,28 @@ export class ChatPanelFileChip extends SignalWatcher(
  @property({ attribute: false })
  accessor chip!: FileChip;

+  @property({ attribute: false })
+  accessor removeChip!: (chip: ChatChip) => void;
+
  override render() {
-    const { state, fileName, fileType } = this.chip;
+    const { state, file } = this.chip;
    const isLoading = state === 'processing';
-    const tooltip = getChipTooltip(state, fileName, this.chip.tooltip);
+    const tooltip = getChipTooltip(state, file.name, this.chip.tooltip);
+    const fileType = file.name.split('.').pop() ?? '';
    const fileIcon = getAttachmentFileIcon(fileType);
    const icon = getChipIcon(state, fileIcon);

    return html`<chat-panel-chip
      .state=${state}
-      .name=${fileName}
+      .name=${file.name}
      .tooltip=${tooltip}
      .icon=${icon}
      .closeable=${!isLoading}
+      .onChipDelete=${this.onChipDelete}
    ></chat-panel-chip>`;
  }
+
+  private readonly onChipDelete = () => {
+    this.removeChip(this.chip);
+  };
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts
@@ -8,7 +8,7 @@ import type { ChatChip, ChipState, DocChip, FileChip } from '../chat-context';
 export function getChipTooltip(
  state: ChipState,
  name: string,
-  tooltip?: string
+  tooltip?: string | null
 ) {
  if (tooltip) {
    return tooltip;
@@ -20,7 +20,7 @@ export function getChipTooltip(
    return 'Processing...';
  }
  if (state === 'failed') {
-    return 'Failed to process';
+    return 'Failed to add to context';
  }
  return name;
 }
@@ -45,7 +45,7 @@ export function isDocChip(chip: ChatChip): chip is DocChip {
 }

 export function isFileChip(chip: ChatChip): chip is FileChip {
-  return 'fileId' in chip;
+  return 'file' in chip && chip.file instanceof File;
 }

 export function isDocContext(
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
@@ -28,6 +28,7 @@ import type {
  DocSearchMenuConfig,
 } from './chat-config';
 import type {
+  ChatChip,
  ChatContextValue,
  ChatItem,
  DocChip,
@@ -180,7 +181,6 @@ export class ChatPanel extends SignalWatcher(
    }

    // context initialized, show the chips
-    let chips: (DocChip | FileChip)[] = [];
    const { docs = [], files = [] } =
      (await AIProvider.context?.getContextDocsAndFiles(
        this.doc.workspace.id,
@@ -191,23 +191,34 @@ export class ChatPanel extends SignalWatcher(
      (a, b) =>
        new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
    );
-    chips = list.map(item => {
-      let chip: DocChip | FileChip;
-      if (isDocContext(item)) {
-        chip = {
-          docId: item.id,
-          state: 'processing',
-        };
-      } else {
-        chip = {
-          fileId: item.id,
-          state: item.status === 'finished' ? 'success' : item.status,
-          fileName: item.name,
-          fileType: '',
-        };
-      }
-      return chip;
-    });
+    const chips: ChatChip[] = await Promise.all(
+      list.map(async item => {
+        if (isDocContext(item)) {
+          return {
+            docId: item.id,
+            state: '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 === 'finished' ? 'success' : item.status,
+            tooltip: item.error,
+          } as FileChip;
+        }
+      })
+    );
+
    this.chatContextValue = {
      ...this.chatContextValue,
      chips,
@@ -489,6 +500,7 @@ export class ChatPanel extends SignalWatcher(
      <chat-panel-input
        .chatContextValue=${this.chatContextValue}
        .getSessionId=${this._getSessionId}
+        .getContextId=${this._getContextId}
        .networkSearchConfig=${this.networkSearchConfig}
        .updateContext=${this.updateContext}
        .host=${this.host}
--- a/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts
+++ b/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts
@@ -7,6 +7,7 @@ import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-
 import type { UserFriendlyError } from '@affine/error';
 import {
  addContextDocMutation,
+  addContextFileMutation,
  cleanupCopilotSessionMutation,
  createCopilotContextMutation,
  createCopilotMessageMutation,
@@ -22,6 +23,7 @@ import {
  type QueryOptions,
  type QueryResponse,
  removeContextDocMutation,
+  removeContextFileMutation,
  type RequestOptions,
  updateCopilotSessionMutation,
 } from '@affine/graphql';
@@ -261,12 +263,30 @@ export class CopilotClient {
    return res.removeContextDoc;
  }

-  async addContextFile() {
-    return;
+  async addContextFile(
+    content: File,
+    options: OptionsField<typeof addContextFileMutation>
+  ) {
+    const res = await this.gql({
+      query: addContextFileMutation,
+      variables: {
+        content,
+        options,
+      },
+    });
+    return res.addContextFile;
  }

-  async removeContextFile() {
-    return;
+  async removeContextFile(
+    options: OptionsField<typeof removeContextFileMutation>
+  ) {
+    const res = await this.gql({
+      query: removeContextFileMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.removeContextFile;
  }

  async getContextDocsAndFiles(
--- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
+++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
@@ -36,21 +36,12 @@ export function setupAIProvider(
 ) {
  //#region actions
  AIProvider.provide('chat', options => {
-    const { input, docs, ...rest } = options;
-    const params = docs?.length
-      ? {
-          docs: docs.map((doc, i) => ({
-            docId: doc.docId,
-            markdown: doc.markdown,
-            index: i + 1,
-          })),
-        }
-      : undefined;
+    const { input, contexts, ...rest } = options;
    return textToText({
      ...rest,
      client,
      content: input,
-      params,
+      params: contexts,
    });
  });

@@ -441,11 +432,17 @@ Could you make a new website based on these notes and send back just the html fi
    removeContextDoc: async (options: { contextId: string; docId: string }) => {
      return client.removeContextDoc(options);
    },
-    addContextFile: async () => {
-      return client.addContextFile();
+    addContextFile: async (
+      file: File,
+      options: { contextId: string; blobId: string }
+    ) => {
+      return client.addContextFile(file, options);
    },
-    removeContextFile: async () => {
-      return client.removeContextFile();
+    removeContextFile: async (options: {
+      contextId: string;
+      fileId: string;
+    }) => {
+      return client.removeContextFile(options);
    },
    getContextDocsAndFiles: async (
      workspaceId: string,
@@ -454,6 +451,13 @@ Could you make a new website based on these notes and send back just the html fi
    ) => {
      return client.getContextDocsAndFiles(workspaceId, sessionId, contextId);
    },
+    matchContext: async (
+      contextId: string,
+      content: string,
+      limit?: number
+    ) => {
+      return client.matchContext(contextId, content, limit);
+    },
  });

  AIProvider.provide('histories', {