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 { 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',