From 570dc79e3d9846f753767c887061f843ab01960e Mon Sep 17 00:00:00 2001 From: darkskygit Date: Thu, 17 Apr 2025 09:57:33 +0000 Subject: [PATCH] feat(server): stop embedding in doc embedding disabled workspace (#11761) fix AI-33 --- packages/backend/server/src/base/error/def.ts | 4 +++ .../server/src/base/error/errors.gen.ts | 7 +++++ .../server/src/plugins/copilot/context/job.ts | 6 +++++ .../src/plugins/copilot/context/resolver.ts | 27 ++++++++++++++++++- packages/backend/server/src/schema.gql | 1 + packages/common/graphql/src/schema.ts | 1 + packages/frontend/i18n/src/i18n.gen.ts | 4 +++ packages/frontend/i18n/src/resources/en.json | 1 + 8 files changed, 50 insertions(+), 1 deletion(-) diff --git a/packages/backend/server/src/base/error/def.ts b/packages/backend/server/src/base/error/def.ts index 7f520ccb9..da253ad48 100644 --- a/packages/backend/server/src/base/error/def.ts +++ b/packages/backend/server/src/base/error/def.ts @@ -701,6 +701,10 @@ export const USER_FRIENDLY_ERRORS = { message: ({ contextId, content, message }) => `Failed to match context ${contextId} with "${escape(content)}": ${message}`, }, + copilot_embedding_disabled: { + type: 'action_forbidden', + message: `Embedding feature is disabled, please contact the administrator to enable it in the workspace settings.`, + }, copilot_embedding_unavailable: { type: 'action_forbidden', message: `Embedding feature not available, you may need to install pgvector extension to your database`, diff --git a/packages/backend/server/src/base/error/errors.gen.ts b/packages/backend/server/src/base/error/errors.gen.ts index 0a6a4bdbb..4c4458bb1 100644 --- a/packages/backend/server/src/base/error/errors.gen.ts +++ b/packages/backend/server/src/base/error/errors.gen.ts @@ -759,6 +759,12 @@ export class CopilotFailedToMatchContext extends UserFriendlyError { } } +export class CopilotEmbeddingDisabled extends UserFriendlyError { + constructor(message?: string) { + super('action_forbidden', 'copilot_embedding_disabled', message); + } +} + export class CopilotEmbeddingUnavailable extends UserFriendlyError { constructor(message?: string) { super('action_forbidden', 'copilot_embedding_unavailable', message); @@ -1037,6 +1043,7 @@ export enum ErrorNames { COPILOT_CONTEXT_FILE_NOT_SUPPORTED, COPILOT_FAILED_TO_MODIFY_CONTEXT, COPILOT_FAILED_TO_MATCH_CONTEXT, + COPILOT_EMBEDDING_DISABLED, COPILOT_EMBEDDING_UNAVAILABLE, COPILOT_TRANSCRIPTION_JOB_EXISTS, COPILOT_TRANSCRIPTION_JOB_NOT_FOUND, diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/context/job.ts index 4d962c725..56470d629 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/context/job.ts @@ -77,6 +77,12 @@ export class CopilotContextDocJob { contextId?: string ) { if (!this.supportEmbedding) return; + const allowEmbedding = await this.models.workspace.allowEmbedding( + docs[0]?.workspaceId + ); + if (!allowEmbedding) { + return; + } for (const { workspaceId, docId } of docs) { await this.queue.add('copilot.embedding.docs', { diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index ccb6fc33b..7bf8e53a8 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -21,6 +21,7 @@ import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs'; import { BlobQuotaExceeded, CallMetric, + CopilotEmbeddingDisabled, CopilotEmbeddingUnavailable, CopilotFailedToMatchContext, CopilotFailedToModifyContext, @@ -231,6 +232,7 @@ export class CopilotContextRootResolver { private readonly db: PrismaClient, private readonly ac: AccessController, private readonly event: EventBus, + private readonly models: Models, private readonly mutex: RequestMutex, private readonly chatSession: ChatSessionService, private readonly context: CopilotContextService @@ -346,7 +348,10 @@ export class CopilotContextRootResolver { .allowLocal() .assert('Workspace.Copilot'); - if (this.context.canEmbedding) { + if ( + this.context.canEmbedding && + (await this.models.workspace.allowEmbedding(workspaceId)) + ) { const total = await this.db.snapshot.count({ where: { workspaceId } }); const embedded = await this.db.snapshot.count({ where: { workspaceId, embedding: { isNot: null } }, @@ -452,6 +457,13 @@ export class CopilotContextResolver { } const session = await this.context.get(options.contextId); + const allowEmbedding = await this.models.workspace.allowEmbedding( + session.workspaceId + ); + if (!allowEmbedding) { + throw new CopilotEmbeddingDisabled(); + } + try { const records = await session.addCategoryRecord( options.type, @@ -521,6 +533,13 @@ export class CopilotContextResolver { } const session = await this.context.get(options.contextId); + const allowEmbedding = await this.models.workspace.allowEmbedding( + session.workspaceId + ); + if (!allowEmbedding) { + throw new CopilotEmbeddingDisabled(); + } + try { const record = await session.addDocRecord(options.docId); @@ -714,6 +733,12 @@ export class CopilotContextResolver { .workspace(session.workspaceId) .allowLocal() .assert('Workspace.Copilot'); + const allowEmbedding = await this.models.workspace.allowEmbedding( + session.workspaceId + ); + if (!allowEmbedding) { + return []; + } try { return await session.matchWorkspaceChunks( diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index f0947de82..7f429a5c2 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -425,6 +425,7 @@ enum ErrorNames { COPILOT_CONTEXT_FILE_NOT_SUPPORTED COPILOT_DOCS_NOT_FOUND COPILOT_DOC_NOT_FOUND + COPILOT_EMBEDDING_DISABLED COPILOT_EMBEDDING_UNAVAILABLE COPILOT_FAILED_TO_CREATE_MESSAGE COPILOT_FAILED_TO_GENERATE_TEXT diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts index 02fa89b71..528e28a8d 100644 --- a/packages/common/graphql/src/schema.ts +++ b/packages/common/graphql/src/schema.ts @@ -570,6 +570,7 @@ export enum ErrorNames { COPILOT_CONTEXT_FILE_NOT_SUPPORTED = 'COPILOT_CONTEXT_FILE_NOT_SUPPORTED', COPILOT_DOCS_NOT_FOUND = 'COPILOT_DOCS_NOT_FOUND', COPILOT_DOC_NOT_FOUND = 'COPILOT_DOC_NOT_FOUND', + COPILOT_EMBEDDING_DISABLED = 'COPILOT_EMBEDDING_DISABLED', COPILOT_EMBEDDING_UNAVAILABLE = 'COPILOT_EMBEDDING_UNAVAILABLE', COPILOT_FAILED_TO_CREATE_MESSAGE = 'COPILOT_FAILED_TO_CREATE_MESSAGE', COPILOT_FAILED_TO_GENERATE_TEXT = 'COPILOT_FAILED_TO_GENERATE_TEXT', diff --git a/packages/frontend/i18n/src/i18n.gen.ts b/packages/frontend/i18n/src/i18n.gen.ts index 9ed5ef846..2dd951069 100644 --- a/packages/frontend/i18n/src/i18n.gen.ts +++ b/packages/frontend/i18n/src/i18n.gen.ts @@ -8137,6 +8137,10 @@ export function useAFFiNEI18N(): { contextId: string; message: string; }>): string; + /** + * `Embedding feature is disabled, please contact the administrator to enable it in the workspace settings.` + */ + ["error.COPILOT_EMBEDDING_DISABLED"](): string; /** * `Embedding feature not available, you may need to install pgvector extension to your database` */ diff --git a/packages/frontend/i18n/src/resources/en.json b/packages/frontend/i18n/src/resources/en.json index dd5e40117..2659f7de5 100644 --- a/packages/frontend/i18n/src/resources/en.json +++ b/packages/frontend/i18n/src/resources/en.json @@ -2012,6 +2012,7 @@ "error.COPILOT_CONTEXT_FILE_NOT_SUPPORTED": "File {{fileName}} is not supported to use as context: {{message}}", "error.COPILOT_FAILED_TO_MODIFY_CONTEXT": "Failed to modify context {{contextId}}: {{message}}", "error.COPILOT_FAILED_TO_MATCH_CONTEXT": "Failed to match context {{contextId}} with \"%7B%7Bcontent%7D%7D\": {{message}}", + "error.COPILOT_EMBEDDING_DISABLED": "Embedding feature is disabled, please contact the administrator to enable it in the workspace settings.", "error.COPILOT_EMBEDDING_UNAVAILABLE": "Embedding feature not available, you may need to install pgvector extension to your database", "error.COPILOT_TRANSCRIPTION_JOB_EXISTS": "Transcription job already exists", "error.COPILOT_TRANSCRIPTION_JOB_NOT_FOUND": "Transcription job not found.",