From f4e7595f4b15119664661bacfaf9a4f71d9e1454 Mon Sep 17 00:00:00 2001
From: darkskygit <darksky.graphite@toeverything.info>
Date: Wed, 28 May 2025 04:36:37 +0000
Subject: [PATCH] feat(server): add copilot embedding feature (#12590)

fix AI-154
---
 packages/backend/server/src/core/config/types.ts             | 1 +
 packages/backend/server/src/models/copilot-workspace.ts      | 4 ++--
 .../backend/server/src/plugins/copilot/workspace/service.ts  | 5 ++++-
 packages/backend/server/src/schema.gql                       | 1 +
 packages/common/graphql/src/schema.ts                        | 1 +
 5 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/packages/backend/server/src/core/config/types.ts b/packages/backend/server/src/core/config/types.ts
index b579d00be7..c0f4aea578 100644
--- a/packages/backend/server/src/core/config/types.ts
+++ b/packages/backend/server/src/core/config/types.ts
@@ -5,6 +5,7 @@ import { DeploymentType } from '../../env';
 export enum ServerFeature {
   Captcha = 'captcha',
   Copilot = 'copilot',
+  CopilotEmbedding = 'copilot_embedding',
   Payment = 'payment',
   OAuth = 'oauth',
   Indexer = 'indexer',
diff --git a/packages/backend/server/src/models/copilot-workspace.ts b/packages/backend/server/src/models/copilot-workspace.ts
index 7a1d9c9354..5b1b9d9128 100644
--- a/packages/backend/server/src/models/copilot-workspace.ts
+++ b/packages/backend/server/src/models/copilot-workspace.ts
@@ -232,8 +232,8 @@ 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_file_embeddings')`;
-    return Number(count) === 1;
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_embeddings', 'ai_workspace_file_embeddings')`;
+    return Number(count) === 2;
   }
 
   private processEmbeddings(
diff --git a/packages/backend/server/src/plugins/copilot/workspace/service.ts b/packages/backend/server/src/plugins/copilot/workspace/service.ts
index 34307a8bd6..0091ad976a 100644
--- a/packages/backend/server/src/plugins/copilot/workspace/service.ts
+++ b/packages/backend/server/src/plugins/copilot/workspace/service.ts
@@ -3,6 +3,7 @@ import { createHash } from 'node:crypto';
 import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
 
 import { FileUpload, JobQueue, PaginationInput } from '../../../base';
+import { ServerFeature, ServerService } from '../../../core';
 import { Models } from '../../../models';
 import { CopilotStorage } from '../storage';
 import { readStream } from '../utils';
@@ -12,6 +13,7 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
   private supportEmbedding = false;
 
   constructor(
+    private readonly server: ServerService,
     private readonly models: Models,
     private readonly queue: JobQueue,
     private readonly storage: CopilotStorage
@@ -19,8 +21,9 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
 
   async onApplicationBootstrap() {
     const supportEmbedding =
-      await this.models.copilotContext.checkEmbeddingAvailable();
+      await this.models.copilotWorkspace.checkEmbeddingAvailable();
     if (supportEmbedding) {
+      this.server.enableFeature(ServerFeature.CopilotEmbedding);
       this.supportEmbedding = true;
     }
   }
diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql
index 2e5c14d30c..ff30e3190e 100644
--- a/packages/backend/server/src/schema.gql
+++ b/packages/backend/server/src/schema.gql
@@ -1606,6 +1606,7 @@ enum ServerDeploymentType {
 enum ServerFeature {
   Captcha
   Copilot
+  CopilotEmbedding
   Indexer
   OAuth
   Payment
diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts
index cb90699500..1725f57530 100644
--- a/packages/common/graphql/src/schema.ts
+++ b/packages/common/graphql/src/schema.ts
@@ -2169,6 +2169,7 @@ export enum ServerDeploymentType {
 export enum ServerFeature {
   Captcha = 'Captcha',
   Copilot = 'Copilot',
+  CopilotEmbedding = 'CopilotEmbedding',
   Indexer = 'Indexer',
   OAuth = 'OAuth',
   Payment = 'Payment',