feat(core): extract selected docs (#13426)

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

## Summary by CodeRabbit

* **New Features**
* Added support for handling and extracting embedded document references
within selected content in AI chat features.
* Documents associated with selected context chips are now properly
managed alongside attachments, improving context handling in AI chat
interactions.

* **Bug Fixes**
* Ensured that the state of context chips accurately reflects the
presence of attachments and documents.

* **Documentation**
* Updated type definitions to include support for document references in
relevant AI chat contexts.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

> CLOSE AF-2783

packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts

@@ -39,6 +39,7 @@ export interface CollectionChip extends BaseChip {
 export interface SelectedContextChip extends BaseChip {
   uuid: string;
   attachments: { sourceId: string; name: string }[];
+  docs: string[];
   snapshot: string | null;
   combinedElementsMarkdown: string | null;
 }

packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts

@@ -412,15 +412,16 @@ export class AIChatComposer extends SignalWatcher(
   };
 
   private readonly addSelectedContextChip = async () => {
-    const { attachments, snapshot, combinedElementsMarkdown } =
+    const { attachments, snapshot, combinedElementsMarkdown, docs } =
       this.chatContextValue;
     await this.removeSelectedContextChip();
     const chip: SelectedContextChip = {
       uuid: uuidv4(),
       attachments,
+      docs,
       snapshot,
       combinedElementsMarkdown,
-      state: 'processing',
+      state: attachments.length > 0 ? 'processing' : 'finished',
     };
     await this.addChip(chip, true);
   };
@@ -545,7 +546,7 @@ export class AIChatComposer extends SignalWatcher(
   private readonly addSelectedContextChipToContext = async (
     chip: SelectedContextChip
   ) => {
-    const { attachments } = chip;
+    const { attachments, docs } = chip;
     const contextId = await this.createContextId();
     if (!contextId || !AIProvider.context) {
       throw new Error('Context not found');
@@ -554,6 +555,14 @@ export class AIChatComposer extends SignalWatcher(
       blobIds: attachments.map(attachment => attachment.sourceId),
       contextId,
     });
+    await Promise.all(
+      docs.map(docId =>
+        AIProvider.context?.addContextDoc({
+          contextId,
+          docId,
+        })
+      )
+    );
   };
 
   private readonly removeFromContext = async (
@@ -589,11 +598,19 @@ export class AIChatComposer extends SignalWatcher(
         });
       }
       if (isSelectedContextChip(chip)) {
-        const { attachments } = chip;
-        return await AIProvider.context.removeContextBlobs({
+        const { attachments, docs } = chip;
+        await AIProvider.context.removeContextBlobs({
           contextId,
           blobIds: attachments.map(attachment => attachment.sourceId),
         });
+        await Promise.all(
+          docs.map(docId =>
+            AIProvider.context?.removeContextDoc({
+              contextId,
+              docId,
+            })
+          )
+        );
       }
       return true;
     } catch {

packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts

@@ -52,6 +52,7 @@ const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   snapshot: null,
   attachments: [],
   combinedElementsMarkdown: null,
+  docs: [],
 };
 
 export class AIChatContent extends SignalWatcher(

packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts

@@ -18,5 +18,7 @@ export type ChatContextValue = {
   attachments: { sourceId: string; name: string }[];
   // combined markdown of the selected elements
   combinedElementsMarkdown: string | null;
+  // docs of the selected content
+  docs: string[];
   abortController: AbortController | null;
 };

packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/type.ts

@@ -30,5 +30,5 @@ export type AIChatInputContext = {
   abortController: AbortController | null;
 } & Pick<
   ChatContextValue,
-  'snapshot' | 'combinedElementsMarkdown' | 'attachments'
+  'snapshot' | 'combinedElementsMarkdown' | 'attachments' | 'docs'
 >;

packages/frontend/core/src/blocksuite/ai/components/playground/chat.ts

@@ -48,6 +48,7 @@ const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   snapshot: null,
   attachments: [],
   combinedElementsMarkdown: null,
+  docs: [],
 };
 
 export class PlaygroundChat extends SignalWatcher(

packages/frontend/core/src/blocksuite/ai/utils/extract.ts

@@ -2,6 +2,8 @@ import { WorkspaceImpl } from '@affine/core/modules/workspace/impls/workspace';
 import { getSurfaceBlock } from '@blocksuite/affine/blocks/surface';
 import {
   DatabaseBlockModel,
+  EmbedLinkedDocModel,
+  EmbedSyncedDocModel,
   ImageBlockModel,
   NoteBlockModel,
   NoteDisplayMode,
@@ -65,6 +67,7 @@ async function extractEdgelessSelected(
   let markdown = '';
   const attachments: ChatContextValue['attachments'] = [];
   const images: File[] = [];
+  const docs: ChatContextValue['docs'] = [];
 
   if (selectedElements.length) {
     const transformer = host.store.getTransformer();
@@ -117,6 +120,12 @@ async function extractEdgelessSelected(
           needSnapshot = true;
           const props = getElementProps(element, new Map());
           surface.addElement(props);
+        } else if (
+          element instanceof EmbedSyncedDocModel ||
+          element instanceof EmbedLinkedDocModel
+        ) {
+          const docId = element.props.pageId;
+          docs.push(docId);
         }
       }
 
@@ -144,6 +153,7 @@ async function extractEdgelessSelected(
     snapshot: snapshot ? JSON.stringify(snapshot) : null,
     combinedElementsMarkdown: markdown.length ? markdown : null,
     attachments,
+    docs,
   };
 }