feat(server): integrate blob to context (#13491)

Author: darkskygit

packages/backend/server/src/models/blob.ts

@@ -67,12 +67,17 @@ export class BlobModel extends BaseModel {
     });
   }
 
-  async list(workspaceId: string) {
+  async list(
+    workspaceId: string,
+    options?: { where: Prisma.BlobWhereInput; select?: Prisma.BlobSelect }
+  ) {
     return await this.db.blob.findMany({
       where: {
+        ...options?.where,
         workspaceId,
         deletedAt: null,
       },
+      select: options?.select,
     });
   }
 

packages/backend/server/src/models/copilot-workspace.ts

@@ -6,13 +6,14 @@ import { Prisma, PrismaClient } from '@prisma/client';
 
 import { PaginationInput } from '../base';
 import { BaseModel } from './base';
-import type {
-  BlobChunkSimilarity,
-  CopilotWorkspaceFile,
-  CopilotWorkspaceFileMetadata,
-  Embedding,
-  FileChunkSimilarity,
-  IgnoredDoc,
+import {
+  type BlobChunkSimilarity,
+  clearEmbeddingContent,
+  type CopilotWorkspaceFile,
+  type CopilotWorkspaceFileMetadata,
+  type Embedding,
+  type FileChunkSimilarity,
+  type IgnoredDoc,
 } from './common';
 
 @Injectable()
@@ -413,6 +414,33 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
     return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
+  async getBlobContent(
+    workspaceId: string,
+    blobId: string,
+    chunk?: number
+  ): Promise<string | undefined> {
+    const blob = await this.db.aiWorkspaceBlobEmbedding.findMany({
+      where: { workspaceId, blobId, chunk },
+      select: { content: true },
+      orderBy: { chunk: 'asc' },
+    });
+    return blob?.map(f => clearEmbeddingContent(f.content)).join('\n');
+  }
+
+  async getBlobChunkSizes(workspaceId: string, blobIds: string[]) {
+    const sizes = await this.db.aiWorkspaceBlobEmbedding.groupBy({
+      by: ['blobId'],
+      _count: { chunk: true },
+      where: { workspaceId, blobId: { in: blobIds } },
+    });
+    return sizes.reduce((acc, cur) => {
+      if (cur._count.chunk) {
+        acc.set(cur.blobId, cur._count.chunk);
+      }
+      return acc;
+    }, new Map<string, number>());
+  }
+
   @Transactional()
   async insertBlobEmbeddings(
     workspaceId: string,

packages/backend/server/src/plugins/copilot/context/session.ts

@@ -55,7 +55,7 @@ export class ContextSession implements AsyncDisposable {
     return this.config.docs.map(d => ({ ...d }));
   }
 
-  get files() {
+  get files(): Required<ContextFile>[] {
     return this.config.files.map(f => this.fulfillFile(f));
   }
 
@@ -135,6 +135,36 @@ export class ContextSession implements AsyncDisposable {
     return record;
   }
 
+  async getBlobMetadata() {
+    const blobIds = this.blobs.map(b => b.id);
+    const blobs = await this.models.blob.list(this.config.workspaceId, {
+      where: { key: { in: blobIds } },
+      select: { key: true, mime: true },
+    });
+    const blobChunkSizes = await this.models.copilotWorkspace.getBlobChunkSizes(
+      this.config.workspaceId,
+      blobIds
+    );
+    return blobs
+      .filter(b => !!blobChunkSizes.get(b.key))
+      .map(b => ({
+        id: b.key,
+        mimeType: b.mime,
+        chunkSize: blobChunkSizes.get(b.key),
+      }));
+  }
+
+  async getBlobContent(
+    blobId: string,
+    chunk?: number
+  ): Promise<string | undefined> {
+    return this.models.copilotWorkspace.getBlobContent(
+      this.config.workspaceId,
+      blobId,
+      chunk
+    );
+  }
+
   async removeBlobRecord(blobId: string): Promise<boolean> {
     const index = this.config.blobs.findIndex(b => b.id === blobId);
     if (index >= 0) {

packages/backend/server/src/plugins/copilot/controller.ts

@@ -208,8 +208,14 @@ export class CopilotController implements BeforeApplicationShutdown {
 
     const context = await this.context.getBySessionId(sessionId);
     const contextParams =
-      Array.isArray(context?.files) && context.files.length > 0
-        ? { contextFiles: context.files }
+      (Array.isArray(context?.files) && context.files.length > 0) ||
+      (Array.isArray(context?.blobs) && context.blobs.length > 0)
+        ? {
+            contextFiles: [
+              ...context.files,
+              ...(await context.getBlobMetadata()),
+            ],
+          }
         : {};
     const lastParams = latestMessage
       ? {

packages/backend/server/src/plugins/copilot/tools/blob-read.ts

@@ -33,7 +33,11 @@ export const buildBlobContentGetter = (
       return;
     }
 
-    const content = await context?.getFileContent(blobId, chunk);
+    const [file, blob] = await Promise.all([
+      context?.getFileContent(blobId, chunk),
+      context?.getBlobContent(blobId, chunk),
+    ]);
+    const content = file?.trim() || blob?.trim();
     if (!content) {
       return;
     }