feat(server): stop embedding in doc embedding disabled workspace (#11761)

fix AI-33
This commit is contained in:
darkskygit
2025-04-17 09:57:33 +00:00
parent e577bb7aa9
commit 570dc79e3d
8 changed files with 50 additions and 1 deletions

View File

@@ -701,6 +701,10 @@ export const USER_FRIENDLY_ERRORS = {
message: ({ contextId, content, message }) => message: ({ contextId, content, message }) =>
`Failed to match context ${contextId} with "${escape(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: { copilot_embedding_unavailable: {
type: 'action_forbidden', type: 'action_forbidden',
message: `Embedding feature not available, you may need to install pgvector extension to your database`, message: `Embedding feature not available, you may need to install pgvector extension to your database`,

View File

@@ -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 { export class CopilotEmbeddingUnavailable extends UserFriendlyError {
constructor(message?: string) { constructor(message?: string) {
super('action_forbidden', 'copilot_embedding_unavailable', message); super('action_forbidden', 'copilot_embedding_unavailable', message);
@@ -1037,6 +1043,7 @@ export enum ErrorNames {
COPILOT_CONTEXT_FILE_NOT_SUPPORTED, COPILOT_CONTEXT_FILE_NOT_SUPPORTED,
COPILOT_FAILED_TO_MODIFY_CONTEXT, COPILOT_FAILED_TO_MODIFY_CONTEXT,
COPILOT_FAILED_TO_MATCH_CONTEXT, COPILOT_FAILED_TO_MATCH_CONTEXT,
COPILOT_EMBEDDING_DISABLED,
COPILOT_EMBEDDING_UNAVAILABLE, COPILOT_EMBEDDING_UNAVAILABLE,
COPILOT_TRANSCRIPTION_JOB_EXISTS, COPILOT_TRANSCRIPTION_JOB_EXISTS,
COPILOT_TRANSCRIPTION_JOB_NOT_FOUND, COPILOT_TRANSCRIPTION_JOB_NOT_FOUND,

View File

@@ -77,6 +77,12 @@ export class CopilotContextDocJob {
contextId?: string contextId?: string
) { ) {
if (!this.supportEmbedding) return; if (!this.supportEmbedding) return;
const allowEmbedding = await this.models.workspace.allowEmbedding(
docs[0]?.workspaceId
);
if (!allowEmbedding) {
return;
}
for (const { workspaceId, docId } of docs) { for (const { workspaceId, docId } of docs) {
await this.queue.add('copilot.embedding.docs', { await this.queue.add('copilot.embedding.docs', {

View File

@@ -21,6 +21,7 @@ import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
import { import {
BlobQuotaExceeded, BlobQuotaExceeded,
CallMetric, CallMetric,
CopilotEmbeddingDisabled,
CopilotEmbeddingUnavailable, CopilotEmbeddingUnavailable,
CopilotFailedToMatchContext, CopilotFailedToMatchContext,
CopilotFailedToModifyContext, CopilotFailedToModifyContext,
@@ -231,6 +232,7 @@ export class CopilotContextRootResolver {
private readonly db: PrismaClient, private readonly db: PrismaClient,
private readonly ac: AccessController, private readonly ac: AccessController,
private readonly event: EventBus, private readonly event: EventBus,
private readonly models: Models,
private readonly mutex: RequestMutex, private readonly mutex: RequestMutex,
private readonly chatSession: ChatSessionService, private readonly chatSession: ChatSessionService,
private readonly context: CopilotContextService private readonly context: CopilotContextService
@@ -346,7 +348,10 @@ export class CopilotContextRootResolver {
.allowLocal() .allowLocal()
.assert('Workspace.Copilot'); .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 total = await this.db.snapshot.count({ where: { workspaceId } });
const embedded = await this.db.snapshot.count({ const embedded = await this.db.snapshot.count({
where: { workspaceId, embedding: { isNot: null } }, where: { workspaceId, embedding: { isNot: null } },
@@ -452,6 +457,13 @@ export class CopilotContextResolver {
} }
const session = await this.context.get(options.contextId); const session = await this.context.get(options.contextId);
const allowEmbedding = await this.models.workspace.allowEmbedding(
session.workspaceId
);
if (!allowEmbedding) {
throw new CopilotEmbeddingDisabled();
}
try { try {
const records = await session.addCategoryRecord( const records = await session.addCategoryRecord(
options.type, options.type,
@@ -521,6 +533,13 @@ export class CopilotContextResolver {
} }
const session = await this.context.get(options.contextId); const session = await this.context.get(options.contextId);
const allowEmbedding = await this.models.workspace.allowEmbedding(
session.workspaceId
);
if (!allowEmbedding) {
throw new CopilotEmbeddingDisabled();
}
try { try {
const record = await session.addDocRecord(options.docId); const record = await session.addDocRecord(options.docId);
@@ -714,6 +733,12 @@ export class CopilotContextResolver {
.workspace(session.workspaceId) .workspace(session.workspaceId)
.allowLocal() .allowLocal()
.assert('Workspace.Copilot'); .assert('Workspace.Copilot');
const allowEmbedding = await this.models.workspace.allowEmbedding(
session.workspaceId
);
if (!allowEmbedding) {
return [];
}
try { try {
return await session.matchWorkspaceChunks( return await session.matchWorkspaceChunks(

View File

@@ -425,6 +425,7 @@ 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_UNAVAILABLE COPILOT_EMBEDDING_UNAVAILABLE
COPILOT_FAILED_TO_CREATE_MESSAGE COPILOT_FAILED_TO_CREATE_MESSAGE
COPILOT_FAILED_TO_GENERATE_TEXT COPILOT_FAILED_TO_GENERATE_TEXT

View File

@@ -570,6 +570,7 @@ export enum ErrorNames {
COPILOT_CONTEXT_FILE_NOT_SUPPORTED = 'COPILOT_CONTEXT_FILE_NOT_SUPPORTED', COPILOT_CONTEXT_FILE_NOT_SUPPORTED = 'COPILOT_CONTEXT_FILE_NOT_SUPPORTED',
COPILOT_DOCS_NOT_FOUND = 'COPILOT_DOCS_NOT_FOUND', COPILOT_DOCS_NOT_FOUND = 'COPILOT_DOCS_NOT_FOUND',
COPILOT_DOC_NOT_FOUND = 'COPILOT_DOC_NOT_FOUND', COPILOT_DOC_NOT_FOUND = 'COPILOT_DOC_NOT_FOUND',
COPILOT_EMBEDDING_DISABLED = 'COPILOT_EMBEDDING_DISABLED',
COPILOT_EMBEDDING_UNAVAILABLE = 'COPILOT_EMBEDDING_UNAVAILABLE', COPILOT_EMBEDDING_UNAVAILABLE = 'COPILOT_EMBEDDING_UNAVAILABLE',
COPILOT_FAILED_TO_CREATE_MESSAGE = 'COPILOT_FAILED_TO_CREATE_MESSAGE', COPILOT_FAILED_TO_CREATE_MESSAGE = 'COPILOT_FAILED_TO_CREATE_MESSAGE',
COPILOT_FAILED_TO_GENERATE_TEXT = 'COPILOT_FAILED_TO_GENERATE_TEXT', COPILOT_FAILED_TO_GENERATE_TEXT = 'COPILOT_FAILED_TO_GENERATE_TEXT',

View File

@@ -8137,6 +8137,10 @@ export function useAFFiNEI18N(): {
contextId: string; contextId: string;
message: string; message: string;
}>): 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` * `Embedding feature not available, you may need to install pgvector extension to your database`
*/ */

View File

@@ -2012,6 +2012,7 @@
"error.COPILOT_CONTEXT_FILE_NOT_SUPPORTED": "File {{fileName}} is not supported to use as context: {{message}}", "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_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_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_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_EXISTS": "Transcription job already exists",
"error.COPILOT_TRANSCRIPTION_JOB_NOT_FOUND": "Transcription job not found.", "error.COPILOT_TRANSCRIPTION_JOB_NOT_FOUND": "Transcription job not found.",