fix(server): query workspace embed files (#11982)

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

- **New Features**
	- Expanded file chunk matching to include both context and workspace file embeddings, providing broader and more relevant search results.
- **Improvements**
	- Enhanced result ranking by introducing a re-ranking step for combined embedding matches, improving the relevance of returned file chunks.
	- Adjusted file count reporting to reflect the total number of workspace files instead of ignored documents for more accurate workspace file statistics.
	- Renamed and streamlined workspace file management methods for clearer and more consistent API usage.
- **Bug Fixes**
	- Prevented embedding similarity queries when embedding is disabled for a workspace, improving system behavior consistency.
- **Tests**
	- Added comprehensive tests to verify workspace embedding management, including enabling, matching, and disabling embedding functionality.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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

@@ -124,24 +124,38 @@ export class CopilotContextDocJob {
 
       for (const chunk of chunks) {
         const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
-        await this.models.copilotContext.insertContentEmbedding(
-          contextId,
-          fileId,
-          embeddings
-        );
+        if (contextId) {
+          // for context files
+          await this.models.copilotContext.insertFileEmbedding(
+            contextId,
+            fileId,
+            embeddings
+          );
+        } else {
+          // for workspace files
+          await this.models.copilotWorkspace.insertFileEmbeddings(
+            workspaceId,
+            fileId,
+            embeddings
+          );
+        }
       }
 
-      this.event.emit('workspace.file.embed.finished', {
-        contextId,
-        fileId,
-        chunkSize: total,
-      });
+      if (contextId) {
+        this.event.emit('workspace.file.embed.finished', {
+          contextId,
+          fileId,
+          chunkSize: total,
+        });
+      }
     } catch (error: any) {
-      this.event.emit('workspace.file.embed.failed', {
-        contextId,
-        fileId,
-        error: mapAnyError(error).message,
-      });
+      if (contextId) {
+        this.event.emit('workspace.file.embed.failed', {
+          contextId,
+          fileId,
+          error: mapAnyError(error).message,
+        });
+      }
 
       // passthrough error to job queue
       throw error;

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

@@ -8,6 +8,7 @@ import {
   ContextEmbedStatus,
   ContextFile,
   ContextList,
+  FileChunkSimilarity,
   Models,
 } from '../../../models';
 import { EmbeddingClient } from './types';
@@ -176,18 +177,28 @@ export class ContextSession implements AsyncDisposable {
     topK: number = 5,
     signal?: AbortSignal,
     threshold: number = 0.7
-  ) {
+  ): Promise<FileChunkSimilarity[]> {
     const embedding = await this.client
       .getEmbeddings([content], signal)
       .then(r => r?.[0]?.embedding);
     if (!embedding) return [];
 
-    return this.models.copilotContext.matchContentEmbedding(
-      embedding,
-      this.id,
-      topK,
-      threshold
-    );
+    const [context, workspace] = await Promise.all([
+      this.models.copilotContext.matchFileEmbedding(
+        embedding,
+        this.id,
+        topK,
+        threshold
+      ),
+      this.models.copilotWorkspace.matchFileEmbedding(
+        this.workspaceId,
+        embedding,
+        topK,
+        threshold
+      ),
+    ]);
+
+    return this.client.reRank([...context, ...workspace]);
   }
 
   /**

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

@@ -1,7 +1,7 @@
 import { File } from 'node:buffer';
 
 import { CopilotContextFileNotSupported } from '../../../base';
-import { Embedding } from '../../../models';
+import { ChunkSimilarity, Embedding } from '../../../models';
 import { parseDoc } from '../../../native';
 
 declare global {
@@ -36,7 +36,7 @@ declare global {
     };
 
     'copilot.embedding.files': {
-      contextId: string;
+      contextId?: string;
       userId: string;
       workspaceId: string;
       blobId: string;
@@ -114,6 +114,15 @@ export abstract class EmbeddingClient {
     return embeddings.map(e => ({ ...e, index: chunks[e.index].index }));
   }
 
+  async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
+    embeddings: Chunk[]
+  ): Promise<Chunk[]> {
+    // sort by distance with ascending order
+    return embeddings.sort(
+      (a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
+    );
+  }
+
   abstract getEmbeddings(
     input: string[],
     signal?: AbortSignal

packages/backend/server/src/plugins/copilot/workspace/resolver.ts

@@ -135,7 +135,7 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
     @Parent() config: CopilotWorkspaceConfigType,
     @Args('pagination', PaginationInput.decode) pagination: PaginationInput
   ): Promise<PaginatedCopilotWorkspaceFileType> {
-    const [files, totalCount] = await this.copilotWorkspace.listWorkspaceFiles(
+    const [files, totalCount] = await this.copilotWorkspace.listFiles(
       config.workspaceId,
       pagination
     );
@@ -177,12 +177,12 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
     }
 
     try {
-      const { blobId, file } = await this.copilotWorkspace.addWorkspaceFile(
+      const { blobId, file } = await this.copilotWorkspace.addFile(
         user.id,
         workspaceId,
         content
       );
-      await this.copilotWorkspace.addWorkspaceFileEmbeddingQueue({
+      await this.copilotWorkspace.queueFileEmbedding({
         userId: user.id,
         workspaceId,
         blobId,
@@ -219,6 +219,6 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
       .workspace(workspaceId)
       .assert('Workspace.Settings.Update');
 
-    return await this.copilotWorkspace.removeWorkspaceFile(workspaceId, fileId);
+    return await this.copilotWorkspace.removeFile(workspaceId, fileId);
   }
 }

packages/backend/server/src/plugins/copilot/workspace/service.ts

@@ -53,11 +53,7 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
     ]);
   }
 
-  async addWorkspaceFile(
-    userId: string,
-    workspaceId: string,
-    content: FileUpload
-  ) {
+  async addFile(userId: string, workspaceId: string, content: FileUpload) {
     const fileName = content.filename;
     const buffer = await readStream(content.createReadStream());
     const blobId = createHash('sha256').update(buffer).digest('base64url');
@@ -70,29 +66,25 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
     return { blobId, file };
   }
 
-  async getWorkspaceFile(workspaceId: string, fileId: string) {
+  async getFile(workspaceId: string, fileId: string) {
     return await this.models.copilotWorkspace.getFile(workspaceId, fileId);
   }
 
-  async listWorkspaceFiles(
+  async listFiles(
     workspaceId: string,
     pagination?: {
       includeRead?: boolean;
     } & PaginationInput
   ) {
     return await Promise.all([
-      this.models.copilotWorkspace.listWorkspaceFiles(workspaceId, pagination),
-      this.models.copilotWorkspace.countIgnoredDocs(workspaceId),
+      this.models.copilotWorkspace.listFiles(workspaceId, pagination),
+      this.models.copilotWorkspace.countFiles(workspaceId),
     ]);
   }
 
-  async addWorkspaceFileEmbeddingQueue(
-    file: Jobs['copilot.workspace.embedding.files']
-  ) {
-    if (!this.supportEmbedding) return;
-
+  async queueFileEmbedding(file: Jobs['copilot.embedding.files']) {
     const { userId, workspaceId, blobId, fileId, fileName } = file;
-    await this.queue.add('copilot.workspace.embedding.files', {
+    await this.queue.add('copilot.embedding.files', {
       userId,
       workspaceId,
       blobId,
@@ -101,10 +93,7 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
     });
   }
 
-  async removeWorkspaceFile(workspaceId: string, fileId: string) {
-    return await this.models.copilotWorkspace.removeWorkspaceFile(
-      workspaceId,
-      fileId
-    );
+  async removeFile(workspaceId: string, fileId: string) {
+    return await this.models.copilotWorkspace.removeFile(workspaceId, fileId);
   }
 }

packages/backend/server/src/plugins/copilot/workspace/types.ts

@@ -13,15 +13,6 @@ declare global {
       jobId: string;
     };
   }
-  interface Jobs {
-    'copilot.workspace.embedding.files': {
-      userId: string;
-      workspaceId: string;
-      blobId: string;
-      fileId: string;
-      fileName: string;
-    };
-  }
 }
 
 @ObjectType('CopilotWorkspaceIgnoredDoc')