feat(server): attachment embedding (#13348)

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

* **New Features**
* Added support for managing "blobs" in Copilot context, including
adding and removing blobs via new GraphQL mutations and UI fields.
* Introduced tracking and querying of blob embeddings within workspaces,
enabling search and similarity matching for blob content.
* Extended Copilot context and workspace APIs, schema, and UI to display
and manage blobs alongside existing documents and files.

* **Bug Fixes**
* Updated context and embedding status logic to handle blobs, ensuring
accurate status reporting and embedding management.

* **Tests**
* Added and updated test cases and snapshots to cover blob embedding
insertion, matching, and removal scenarios.

* **Documentation**
* Updated GraphQL schema and TypeScript types to reflect new
blob-related fields and mutations.

* **Chores**
* Refactored and cleaned up code to support new blob entity and
embedding logic, including renaming and updating internal methods and
types.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/backend/server/src/models/common/copilot.ts

@@ -37,6 +37,11 @@ const ContextEmbedStatusSchema = z.enum([
   ContextEmbedStatus.failed,
 ]);
 
+const ContextBlobSchema = z.object({
+  id: z.string(),
+  createdAt: z.number(),
+});
+
 const ContextDocSchema = z.object({
   id: z.string(),
   createdAt: z.number(),
@@ -64,6 +69,9 @@ export const ContextCategorySchema = z.object({
 
 export const ContextConfigSchema = z.object({
   workspaceId: z.string(),
+  blobs: ContextBlobSchema.merge(
+    z.object({ status: ContextEmbedStatusSchema.optional() })
+  ).array(),
   files: ContextFileSchema.array(),
   docs: ContextDocSchema.merge(
     z.object({ status: ContextEmbedStatusSchema.optional() })
@@ -77,10 +85,9 @@ export const MinimalContextConfigSchema = ContextConfigSchema.pick({
 
 export type ContextCategory = z.infer<typeof ContextCategorySchema>;
 export type ContextConfig = z.infer<typeof ContextConfigSchema>;
+export type ContextBlob = z.infer<typeof ContextConfigSchema>['blobs'][number];
 export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
 export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
-export type ContextListItem = ContextDoc | ContextFile;
-export type ContextList = ContextListItem[];
 
 // embeddings
 
@@ -106,6 +113,10 @@ export type FileChunkSimilarity = ChunkSimilarity & {
   mimeType: string;
 };
 
+export type BlobChunkSimilarity = ChunkSimilarity & {
+  blobId: string;
+};
+
 export type DocChunkSimilarity = ChunkSimilarity & {
   docId: string;
 };

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

@@ -6,6 +6,7 @@ import { Prisma } from '@prisma/client';
 import { CopilotSessionNotFound } from '../base';
 import { BaseModel } from './base';
 import {
+  ContextBlob,
   ContextConfigSchema,
   ContextDoc,
   ContextEmbedStatus,
@@ -39,6 +40,7 @@ export class CopilotContextModel extends BaseModel {
         sessionId,
         config: {
           workspaceId: session.workspaceId,
+          blobs: [],
           docs: [],
           files: [],
           categories: [],
@@ -66,10 +68,11 @@ export class CopilotContextModel extends BaseModel {
       if (minimalConfig.success) {
         // fulfill the missing fields
         return {
-          ...minimalConfig.data,
+          blobs: [],
           docs: [],
           files: [],
           categories: [],
+          ...minimalConfig.data,
         };
       }
     }
@@ -83,10 +86,35 @@ export class CopilotContextModel extends BaseModel {
     return row;
   }
 
+  async mergeBlobStatus(
+    workspaceId: string,
+    blobs: ContextBlob[]
+  ): Promise<ContextBlob[]> {
+    const canEmbedding = await this.checkEmbeddingAvailable();
+    const finishedBlobs = canEmbedding
+      ? await this.listWorkspaceBlobEmbedding(
+          workspaceId,
+          Array.from(new Set(blobs.map(blob => blob.id)))
+        )
+      : [];
+    const finishedBlobSet = new Set(finishedBlobs);
+
+    for (const blob of blobs) {
+      const status = finishedBlobSet.has(blob.id)
+        ? ContextEmbedStatus.finished
+        : undefined;
+      // NOTE: when the blob has not been synchronized to the server or is in the embedding queue
+      // the status will be empty, fallback to processing if no status is provided
+      blob.status = status || blob.status || ContextEmbedStatus.processing;
+    }
+
+    return blobs;
+  }
+
   async mergeDocStatus(workspaceId: string, docs: ContextDoc[]) {
     const canEmbedding = await this.checkEmbeddingAvailable();
     const finishedDoc = canEmbedding
-      ? await this.listWorkspaceEmbedding(
+      ? await this.listWorkspaceDocEmbedding(
           workspaceId,
           Array.from(new Set(docs.map(doc => doc.id)))
         )
@@ -126,7 +154,23 @@ export class CopilotContextModel extends BaseModel {
     return Number(count) === 2;
   }
 
-  async listWorkspaceEmbedding(workspaceId: string, docIds?: string[]) {
+  async listWorkspaceBlobEmbedding(
+    workspaceId: string,
+    blobIds?: string[]
+  ): Promise<string[]> {
+    const existsIds = await this.db.aiWorkspaceBlobEmbedding
+      .groupBy({
+        where: {
+          workspaceId,
+          blobId: blobIds ? { in: blobIds } : undefined,
+        },
+        by: ['blobId'],
+      })
+      .then(r => r.map(r => r.blobId));
+    return existsIds;
+  }
+
+  async listWorkspaceDocEmbedding(workspaceId: string, docIds?: string[]) {
     const existsIds = await this.db.aiWorkspaceEmbedding
       .groupBy({
         where: {

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

@@ -7,6 +7,7 @@ import { Prisma, PrismaClient } from '@prisma/client';
 import { PaginationInput } from '../base';
 import { BaseModel } from './base';
 import type {
+  BlobChunkSimilarity,
   CopilotWorkspaceFile,
   CopilotWorkspaceFileMetadata,
   Embedding,
@@ -256,19 +257,19 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
   async checkEmbeddingAvailable(): Promise<boolean> {
     const [{ count }] = await this.db.$queryRaw<
       { count: number }[]
-    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_embeddings', 'ai_workspace_file_embeddings')`;
-    return Number(count) === 2;
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_embeddings', 'ai_workspace_file_embeddings', 'ai_workspace_blob_embeddings')`;
+    return Number(count) === 3;
   }
 
   private processEmbeddings(
     workspaceId: string,
-    fileId: string,
+    fileOrBlobId: string,
     embeddings: Embedding[]
   ) {
     const groups = embeddings.map(e =>
       [
         workspaceId,
-        fileId,
+        fileOrBlobId,
         e.index,
         e.content,
         Prisma.raw(`'[${e.embedding.join(',')}]'`),
@@ -378,6 +379,61 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
     return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
+  @Transactional()
+  async insertBlobEmbeddings(
+    workspaceId: string,
+    blobId: string,
+    embeddings: Embedding[]
+  ) {
+    if (embeddings.length === 0) {
+      this.logger.warn(
+        `No embeddings provided for workspaceId: ${workspaceId}, blobId: ${blobId}. Skipping insertion.`
+      );
+      return;
+    }
+
+    const values = this.processEmbeddings(workspaceId, blobId, embeddings);
+    await this.db.$executeRaw`
+          INSERT INTO "ai_workspace_blob_embeddings"
+          ("workspace_id", "blob_id", "chunk", "content", "embedding") VALUES ${values}
+          ON CONFLICT (workspace_id, blob_id, chunk) DO NOTHING;
+      `;
+  }
+
+  async matchBlobEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<BlobChunkSimilarity[]> {
+    if (!(await this.allowEmbedding(workspaceId))) {
+      return [];
+    }
+
+    const similarityChunks = await this.db.$queryRaw<
+      Array<BlobChunkSimilarity>
+    >`
+      SELECT
+        e."blob_id" as "blobId",
+        e."chunk",
+        e."content",
+        e."embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_blob_embeddings" e
+      WHERE e.workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeBlob(workspaceId: string, blobId: string) {
+    await this.db.$executeRaw`
+      DELETE FROM "ai_workspace_blob_embeddings"
+      WHERE workspace_id = ${workspaceId} AND blob_id = ${blobId};
+    `;
+    return true;
+  }
+
   async removeFile(workspaceId: string, fileId: string) {
     // embeddings will be removed by foreign key constraint
     await this.db.aiWorkspaceFiles.deleteMany({