diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index d994758f61..da5823a9cf 100644 --- a/packages/backend/server/schema.prisma +++ b/packages/backend/server/schema.prisma @@ -114,13 +114,13 @@ model Workspace { name String? @db.VarChar avatarKey String? @map("avatar_key") @db.VarChar - features WorkspaceFeature[] - docs WorkspaceDoc[] - permissions WorkspaceUserRole[] - docPermissions WorkspaceDocUserRole[] - blobs Blob[] - AiWorkspaceIgnoredDocs AiWorkspaceIgnoredDocs[] - AiWorkspaceFiles AiWorkspaceFiles[] + features WorkspaceFeature[] + docs WorkspaceDoc[] + permissions WorkspaceUserRole[] + docPermissions WorkspaceDocUserRole[] + blobs Blob[] + ignoredDocs AiWorkspaceIgnoredDocs[] + embedFiles AiWorkspaceFiles[] @@map("workspaces") } diff --git a/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts index ab349c816d..a1f1df70a0 100644 --- a/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts @@ -104,22 +104,18 @@ test('should manage copilot workspace ignored docs', async t => { test('should insert and search embedding', async t => { { - await t.context.copilotWorkspace.addWorkspaceFile( - workspace.id, + const { fileId } = await t.context.copilotWorkspace.addFile(workspace.id, { + fileName: 'file1', + mimeType: 'text/plain', + size: 1, + }); + await t.context.copilotWorkspace.addFileEmbeddings(workspace.id, fileId, [ { - fileName: 'file1', - mimeType: 'text/plain', - - size: 1, + index: 0, + content: 'content', + embedding: Array.from({ length: 1024 }, () => 1), }, - [ - { - index: 0, - content: 'content', - embedding: Array.from({ length: 1024 }, () => 1), - }, - ] - ); + ]); { const ret = await t.context.copilotWorkspace.matchWorkspaceFileEmbedding( diff --git a/packages/backend/server/src/base/error/def.ts b/packages/backend/server/src/base/error/def.ts index da253ad483..73b3f54960 100644 --- a/packages/backend/server/src/base/error/def.ts +++ b/packages/backend/server/src/base/error/def.ts @@ -711,15 +711,21 @@ export const USER_FRIENDLY_ERRORS = { }, copilot_transcription_job_exists: { type: 'bad_request', - message: () => 'Transcription job already exists', + message: 'Transcription job already exists', }, copilot_transcription_job_not_found: { type: 'bad_request', - message: () => `Transcription job not found.`, + message: `Transcription job not found.`, }, copilot_transcription_audio_not_provided: { type: 'bad_request', - message: () => `Audio not provided.`, + message: `Audio not provided.`, + }, + copilot_failed_to_add_workspace_file_embedding: { + type: 'internal_server_error', + args: { message: 'string' }, + message: ({ message }) => + `Failed to add workspace file embedding: ${message}`, }, // Quota & Limit errors diff --git a/packages/backend/server/src/base/error/errors.gen.ts b/packages/backend/server/src/base/error/errors.gen.ts index 4c4458bb14..65b6b9f55e 100644 --- a/packages/backend/server/src/base/error/errors.gen.ts +++ b/packages/backend/server/src/base/error/errors.gen.ts @@ -788,6 +788,16 @@ export class CopilotTranscriptionAudioNotProvided extends UserFriendlyError { super('bad_request', 'copilot_transcription_audio_not_provided', message); } } +@ObjectType() +class CopilotFailedToAddWorkspaceFileEmbeddingDataType { + @Field() message!: string +} + +export class CopilotFailedToAddWorkspaceFileEmbedding extends UserFriendlyError { + constructor(args: CopilotFailedToAddWorkspaceFileEmbeddingDataType, message?: string | ((args: CopilotFailedToAddWorkspaceFileEmbeddingDataType) => string)) { + super('internal_server_error', 'copilot_failed_to_add_workspace_file_embedding', message, args); + } +} export class BlobQuotaExceeded extends UserFriendlyError { constructor(message?: string) { @@ -1048,6 +1058,7 @@ export enum ErrorNames { COPILOT_TRANSCRIPTION_JOB_EXISTS, COPILOT_TRANSCRIPTION_JOB_NOT_FOUND, COPILOT_TRANSCRIPTION_AUDIO_NOT_PROVIDED, + COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING, BLOB_QUOTA_EXCEEDED, STORAGE_QUOTA_EXCEEDED, MEMBER_QUOTA_EXCEEDED, @@ -1078,5 +1089,5 @@ registerEnumType(ErrorNames, { export const ErrorDataUnionType = createUnionType({ name: 'ErrorDataUnion', types: () => - [GraphqlBadRequestDataType, HttpRequestErrorDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, CopilotInvalidContextDataType, CopilotContextFileNotSupportedDataType, CopilotFailedToModifyContextDataType, CopilotFailedToMatchContextDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseUpdateParamsDataType, WorkspaceMembersExceedLimitToDowngradeDataType, UnsupportedClientVersionDataType, MentionUserDocAccessDeniedDataType] as const, + [GraphqlBadRequestDataType, HttpRequestErrorDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, CopilotInvalidContextDataType, CopilotContextFileNotSupportedDataType, CopilotFailedToModifyContextDataType, CopilotFailedToMatchContextDataType, CopilotFailedToAddWorkspaceFileEmbeddingDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseUpdateParamsDataType, WorkspaceMembersExceedLimitToDowngradeDataType, UnsupportedClientVersionDataType, MentionUserDocAccessDeniedDataType] as const, }); diff --git a/packages/backend/server/src/models/common/copilot.ts b/packages/backend/server/src/models/common/copilot.ts index 9fe2a56e67..f7a3f64347 100644 --- a/packages/backend/server/src/models/common/copilot.ts +++ b/packages/backend/server/src/models/common/copilot.ts @@ -112,9 +112,10 @@ export const CopilotWorkspaceFileSchema = z.object({ size: z.number(), }); -export type CopilotWorkspaceFile = z.infer< +export type CopilotWorkspaceFileMetadata = z.infer< typeof CopilotWorkspaceFileSchema -> & { +>; +export type CopilotWorkspaceFile = CopilotWorkspaceFileMetadata & { workspaceId: string; fileId: string; createdAt: Date; diff --git a/packages/backend/server/src/models/copilot-workspace.ts b/packages/backend/server/src/models/copilot-workspace.ts index 8cfd0af254..8ab39c269e 100644 --- a/packages/backend/server/src/models/copilot-workspace.ts +++ b/packages/backend/server/src/models/copilot-workspace.ts @@ -5,9 +5,10 @@ import { Transactional } from '@nestjs-cls/transactional'; import { Prisma } from '@prisma/client'; import { BaseModel } from './base'; -import { - type CopilotWorkspaceFile, - type Embedding, +import type { + CopilotWorkspaceFile, + CopilotWorkspaceFileMetadata, + Embedding, FileChunkSimilarity, } from './common'; @@ -95,24 +96,40 @@ export class CopilotWorkspaceConfigModel extends BaseModel { return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`)); } - @Transactional() - async addWorkspaceFile( + async addFile( workspaceId: string, - file: Pick, - embeddings: Embedding[] - ): Promise { + file: CopilotWorkspaceFileMetadata + ): Promise { const fileId = randomUUID(); - await this.db.aiWorkspaceFiles.create({ + const row = await this.db.aiWorkspaceFiles.create({ data: { ...file, workspaceId, fileId }, }); + return row; + } + + async getFile(workspaceId: string, fileId: string) { + const file = await this.db.aiWorkspaceFiles.findFirst({ + where: { + workspaceId, + fileId, + }, + }); + return file; + } + + @Transactional() + async addFileEmbeddings( + workspaceId: string, + fileId: string, + embeddings: Embedding[] + ) { const values = this.processEmbeddings(workspaceId, fileId, embeddings); await this.db.$executeRaw` - INSERT INTO "ai_workspace_file_embeddings" - ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values} - ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING; - `; - return fileId; + INSERT INTO "ai_workspace_file_embeddings" + ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values} + ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING; + `; } async listWorkspaceFiles( @@ -152,5 +169,6 @@ export class CopilotWorkspaceConfigModel extends BaseModel { fileId, }, }); + return true; } } diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/context/job.ts index 4d962c725e..b086a9eb23 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/context/job.ts @@ -13,9 +13,9 @@ import { import { DocReader } from '../../../core/doc'; import { Models } from '../../../models'; import { CopilotStorage } from '../storage'; +import { readStream } from '../utils'; import { OpenAIEmbeddingClient } from './embedding'; import { EmbeddingClient } from './types'; -import { readStream } from './utils'; @Injectable() export class CopilotContextDocJob { diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index 6d20b8896e..c14dc0b891 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -47,10 +47,10 @@ import { import { COPILOT_LOCKER, CopilotType } from '../resolver'; import { ChatSessionService } from '../session'; import { CopilotStorage } from '../storage'; +import { MAX_EMBEDDABLE_SIZE } from '../types'; +import { readStream } from '../utils'; import { CopilotContextDocJob } from './job'; import { CopilotContextService } from './service'; -import { MAX_EMBEDDABLE_SIZE } from './types'; -import { readStream } from './utils'; @InputType() class AddContextCategoryInput { diff --git a/packages/backend/server/src/plugins/copilot/context/types.ts b/packages/backend/server/src/plugins/copilot/context/types.ts index a40cf5874d..3bd919dfd5 100644 --- a/packages/backend/server/src/plugins/copilot/context/types.ts +++ b/packages/backend/server/src/plugins/copilot/context/types.ts @@ -1,6 +1,6 @@ import { File } from 'node:buffer'; -import { CopilotContextFileNotSupported, OneMB } from '../../../base'; +import { CopilotContextFileNotSupported } from '../../../base'; import { Embedding } from '../../../models'; import { parseDoc } from '../../../native'; @@ -46,8 +46,6 @@ declare global { } } -export const MAX_EMBEDDABLE_SIZE = 50 * OneMB; - export type Chunk = { index: number; content: string; diff --git a/packages/backend/server/src/plugins/copilot/context/utils.ts b/packages/backend/server/src/plugins/copilot/context/utils.ts index 604a1772d6..451280d68d 100644 --- a/packages/backend/server/src/plugins/copilot/context/utils.ts +++ b/packages/backend/server/src/plugins/copilot/context/utils.ts @@ -1,8 +1,3 @@ -import { Readable } from 'node:stream'; - -import { readBufferWithLimit } from '../../../base'; -import { MAX_EMBEDDABLE_SIZE } from './types'; - export class GqlSignal implements AsyncDisposable { readonly abortController = new AbortController(); @@ -14,10 +9,3 @@ export class GqlSignal implements AsyncDisposable { this.abortController.abort(); } } - -export function readStream( - readable: Readable, - maxSize = MAX_EMBEDDABLE_SIZE -): Promise { - return readBufferWithLimit(readable, maxSize); -} diff --git a/packages/backend/server/src/plugins/copilot/index.ts b/packages/backend/server/src/plugins/copilot/index.ts index e3f557f32b..faf0973868 100644 --- a/packages/backend/server/src/plugins/copilot/index.ts +++ b/packages/backend/server/src/plugins/copilot/index.ts @@ -7,6 +7,7 @@ import { DocStorageModule } from '../../core/doc'; import { FeatureModule } from '../../core/features'; import { PermissionModule } from '../../core/permission'; import { QuotaModule } from '../../core/quota'; +import { WorkspaceModule } from '../../core/workspaces'; import { CopilotContextDocJob, CopilotContextResolver, @@ -29,6 +30,11 @@ import { CopilotTranscriptionService, } from './transcript'; import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow'; +import { + CopilotWorkspaceEmbeddingConfigResolver, + CopilotWorkspaceEmbeddingResolver, + CopilotWorkspaceService, +} from './workspace'; @Module({ imports: [ @@ -37,6 +43,7 @@ import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow'; QuotaModule, PermissionModule, ServerConfigModule, + WorkspaceModule, ], providers: [ // providers @@ -58,6 +65,10 @@ import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow'; // transcription CopilotTranscriptionService, CopilotTranscriptionResolver, + // workspace embeddings + CopilotWorkspaceService, + CopilotWorkspaceEmbeddingResolver, + CopilotWorkspaceEmbeddingConfigResolver, // gql resolvers UserCopilotResolver, PromptsManagementResolver, diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index 2db4d0a42e..61f7f8fb00 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -1,6 +1,7 @@ import { type Tokenizer } from '@affine/server-native'; import { z } from 'zod'; +import { OneMB } from '../../base'; import { fromModelName } from '../../native'; import type { ChatPrompt } from './prompt'; import { PromptMessageSchema, PureMessageSchema } from './providers'; @@ -116,3 +117,5 @@ export type CopilotContextFile = { // embedding status status: 'in_progress' | 'completed' | 'failed'; }; + +export const MAX_EMBEDDABLE_SIZE = 50 * OneMB; diff --git a/packages/backend/server/src/plugins/copilot/utils.ts b/packages/backend/server/src/plugins/copilot/utils.ts new file mode 100644 index 0000000000..8ba8f138b7 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/utils.ts @@ -0,0 +1,11 @@ +import { Readable } from 'node:stream'; + +import { readBufferWithLimit } from '../../base'; +import { MAX_EMBEDDABLE_SIZE } from './types'; + +export function readStream( + readable: Readable, + maxSize = MAX_EMBEDDABLE_SIZE +): Promise { + return readBufferWithLimit(readable, maxSize); +} diff --git a/packages/backend/server/src/plugins/copilot/workspace/index.ts b/packages/backend/server/src/plugins/copilot/workspace/index.ts new file mode 100644 index 0000000000..8cf7d5b25c --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/workspace/index.ts @@ -0,0 +1,5 @@ +export { + CopilotWorkspaceEmbeddingConfigResolver, + CopilotWorkspaceEmbeddingResolver, +} from './resolver'; +export { CopilotWorkspaceService } from './service'; diff --git a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts new file mode 100644 index 0000000000..d77e76135c --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts @@ -0,0 +1,218 @@ +import { + Args, + Context, + Field, + Mutation, + ObjectType, + Parent, + ResolveField, + Resolver, +} from '@nestjs/graphql'; +import type { Request } from 'express'; +import { SafeIntResolver } from 'graphql-scalars'; +import GraphQLUpload, { + type FileUpload, +} from 'graphql-upload/GraphQLUpload.mjs'; + +import { + BlobQuotaExceeded, + CopilotEmbeddingUnavailable, + CopilotFailedToAddWorkspaceFileEmbedding, + Mutex, + TooManyRequest, + UserFriendlyError, +} from '../../../base'; +import { CurrentUser } from '../../../core/auth'; +import { AccessController } from '../../../core/permission'; +import { WorkspaceType } from '../../../core/workspaces'; +import { CopilotWorkspaceFile, Models } from '../../../models'; +import { COPILOT_LOCKER } from '../resolver'; +import { MAX_EMBEDDABLE_SIZE } from '../types'; +import { CopilotWorkspaceService } from './service'; + +@ObjectType('CopilotWorkspaceConfig') +export class CopilotWorkspaceConfigType { + @Field(() => String) + workspaceId!: string; +} + +@ObjectType('CopilotWorkspaceFile') +export class CopilotWorkspaceFileType implements CopilotWorkspaceFile { + @Field(() => String) + workspaceId!: string; + + @Field(() => String) + fileId!: string; + + @Field(() => String) + fileName!: string; + + @Field(() => String) + mimeType!: string; + + @Field(() => SafeIntResolver) + size!: number; + + @Field(() => Date) + createdAt!: Date; +} + +/** + * Workspace embedding config resolver + * Public apis rate limit: 10 req/m + * Other rate limit: 120 req/m + */ +@Resolver(() => WorkspaceType) +export class CopilotWorkspaceEmbeddingResolver { + constructor(private readonly ac: AccessController) {} + + @ResolveField(() => CopilotWorkspaceConfigType, { + complexity: 2, + }) + async embedding( + @CurrentUser() user: CurrentUser, + @Parent() workspace: WorkspaceType + ): Promise { + await this.ac + .user(user.id) + .workspace(workspace.id) + .assert('Workspace.Read'); + + return { workspaceId: workspace.id }; + } +} + +@Resolver(() => CopilotWorkspaceConfigType) +export class CopilotWorkspaceEmbeddingConfigResolver { + constructor( + private readonly ac: AccessController, + private readonly models: Models, + private readonly mutex: Mutex, + private readonly copilotWorkspace: CopilotWorkspaceService + ) {} + + @ResolveField(() => [String], { + complexity: 2, + }) + async ignoredDocs( + @Parent() config: CopilotWorkspaceConfigType + ): Promise { + return this.models.copilotWorkspace.listIgnoredDocs(config.workspaceId); + } + + @Mutation(() => Number, { + name: 'updateWorkspaceEmbeddingIgnoredDocs', + complexity: 2, + description: 'Update ignored docs', + }) + async updateIgnoredDocs( + @CurrentUser() user: CurrentUser, + @Args('workspaceId', { type: () => String }) + workspaceId: string, + @Args('add', { type: () => [String], nullable: true }) + add?: string[], + @Args('remove', { type: () => [String], nullable: true }) + remove?: string[] + ): Promise { + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Settings.Update'); + return await this.models.copilotWorkspace.updateIgnoredDocs( + workspaceId, + add, + remove + ); + } + + @ResolveField(() => [CopilotWorkspaceFileType], { + complexity: 2, + }) + async files( + @Parent() config: CopilotWorkspaceConfigType + ): Promise { + return this.models.copilotWorkspace.listWorkspaceFiles(config.workspaceId); + } + + @Mutation(() => CopilotWorkspaceFileType, { + name: 'addWorkspaceEmbeddingFiles', + complexity: 2, + description: 'Update workspace embedding files', + }) + async addFiles( + @Context() ctx: { req: Request }, + @CurrentUser() user: CurrentUser, + @Args('workspaceId', { type: () => String }) + workspaceId: string, + @Args({ name: 'blob', type: () => GraphQLUpload }) + content: FileUpload + ): Promise { + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Settings.Update'); + + if (!this.copilotWorkspace.canEmbedding) { + throw new CopilotEmbeddingUnavailable(); + } + + const lockFlag = `${COPILOT_LOCKER}:workspace:${workspaceId}`; + await using lock = await this.mutex.acquire(lockFlag); + if (!lock) { + throw new TooManyRequest('Server is busy'); + } + + const length = Number(ctx.req.headers['content-length']); + if (length && length >= MAX_EMBEDDABLE_SIZE) { + throw new BlobQuotaExceeded(); + } + + try { + const { blobId, file } = await this.copilotWorkspace.addWorkspaceFile( + user.id, + workspaceId, + content + ); + await this.copilotWorkspace.addWorkspaceFileEmbeddingQueue({ + userId: user.id, + workspaceId, + blobId, + fileId: file.fileId, + fileName: file.fileName, + }); + + return file; + } catch (e: any) { + // passthrough user friendly error + if (e instanceof UserFriendlyError) { + throw e; + } + throw new CopilotFailedToAddWorkspaceFileEmbedding({ + message: e.message, + }); + } + } + + @Mutation(() => Boolean, { + name: 'removeWorkspaceEmbeddingFiles', + complexity: 2, + description: 'Remove workspace embedding files', + }) + async removeFiles( + @CurrentUser() user: CurrentUser, + @Args('workspaceId', { type: () => String }) + workspaceId: string, + @Args('fileId', { type: () => String }) + fileId: string + ): Promise { + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Settings.Update'); + + return await this.models.copilotWorkspace.removeWorkspaceFile( + workspaceId, + fileId + ); + } +} diff --git a/packages/backend/server/src/plugins/copilot/workspace/service.ts b/packages/backend/server/src/plugins/copilot/workspace/service.ts new file mode 100644 index 0000000000..60b9f92e5a --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/workspace/service.ts @@ -0,0 +1,87 @@ +import { createHash } from 'node:crypto'; + +import { Injectable, OnApplicationBootstrap } from '@nestjs/common'; + +import { FileUpload, JobQueue } from '../../../base'; +import { Models } from '../../../models'; +import { CopilotStorage } from '../storage'; +import { readStream } from '../utils'; + +declare global { + interface Events { + 'workspace.file.embedding.finished': { + jobId: string; + }; + 'workspace.file.embedding.failed': { + jobId: string; + }; + } + interface Jobs { + 'copilot.workspace.embedding.files': { + userId: string; + workspaceId: string; + blobId: string; + fileId: string; + fileName: string; + }; + } +} + +@Injectable() +export class CopilotWorkspaceService implements OnApplicationBootstrap { + private supportEmbedding = false; + + constructor( + private readonly models: Models, + private readonly queue: JobQueue, + private readonly storage: CopilotStorage + ) {} + + async onApplicationBootstrap() { + const supportEmbedding = + await this.models.copilotContext.checkEmbeddingAvailable(); + if (supportEmbedding) { + this.supportEmbedding = true; + } + } + + get canEmbedding() { + return this.supportEmbedding; + } + + async addWorkspaceFile( + userId: string, + workspaceId: string, + content: FileUpload + ) { + const fileName = content.filename; + const buffer = await readStream(content.createReadStream()); + const blobId = createHash('sha256').update(buffer).digest('base64url'); + await this.storage.put(userId, workspaceId, blobId, buffer); + const file = await this.models.copilotWorkspace.addFile(workspaceId, { + fileName, + mimeType: content.mimetype, + size: buffer.length, + }); + return { blobId, file }; + } + + async getWorkspaceFile(workspaceId: string, fileId: string) { + return await this.models.copilotWorkspace.getFile(workspaceId, fileId); + } + + async addWorkspaceFileEmbeddingQueue( + file: Jobs['copilot.workspace.embedding.files'] + ) { + if (!this.supportEmbedding) return; + + const { userId, workspaceId, blobId, fileId, fileName } = file; + await this.queue.add('copilot.workspace.embedding.files', { + userId, + workspaceId, + blobId, + fileId, + fileName, + }); + } +} diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index f8efd8f56d..9a0080af4e 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -159,6 +159,10 @@ type CopilotDocType { status: ContextEmbedStatus } +type CopilotFailedToAddWorkspaceFileEmbeddingDataType { + message: String! +} + type CopilotFailedToMatchContextDataType { content: String! contextId: String! @@ -268,6 +272,21 @@ type CopilotSessionType { promptName: String! } +type CopilotWorkspaceConfig { + files: [CopilotWorkspaceFile!]! + ignoredDocs: [String!]! + workspaceId: String! +} + +type CopilotWorkspaceFile { + createdAt: DateTime! + fileId: String! + fileName: String! + mimeType: String! + size: SafeInt! + workspaceId: String! +} + input CreateChatMessageInput { attachments: [String!] blobs: [Upload!] @@ -404,7 +423,7 @@ type EditorType { name: String! } -union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CopilotContextFileNotSupportedDataType | CopilotDocNotFoundDataType | CopilotFailedToMatchContextDataType | CopilotFailedToModifyContextDataType | CopilotInvalidContextDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocActionDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | DocUpdateBlockedDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | GraphqlBadRequestDataType | HttpRequestErrorDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidLicenseUpdateParamsDataType | InvalidOauthCallbackCodeDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MemberNotFoundInSpaceDataType | MentionUserDocAccessDeniedDataType | MissingOauthQueryParameterDataType | NotInSpaceDataType | QueryTooLongDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedClientVersionDataType | UnsupportedSubscriptionPlanDataType | ValidationErrorDataType | VersionRejectedDataType | WorkspaceMembersExceedLimitToDowngradeDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType +union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CopilotContextFileNotSupportedDataType | CopilotDocNotFoundDataType | CopilotFailedToAddWorkspaceFileEmbeddingDataType | CopilotFailedToMatchContextDataType | CopilotFailedToModifyContextDataType | CopilotInvalidContextDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocActionDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | DocUpdateBlockedDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | GraphqlBadRequestDataType | HttpRequestErrorDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidLicenseUpdateParamsDataType | InvalidOauthCallbackCodeDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MemberNotFoundInSpaceDataType | MentionUserDocAccessDeniedDataType | MissingOauthQueryParameterDataType | NotInSpaceDataType | QueryTooLongDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedClientVersionDataType | UnsupportedSubscriptionPlanDataType | ValidationErrorDataType | VersionRejectedDataType | WorkspaceMembersExceedLimitToDowngradeDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType enum ErrorNames { ACCESS_DENIED @@ -427,6 +446,7 @@ enum ErrorNames { COPILOT_DOC_NOT_FOUND COPILOT_EMBEDDING_DISABLED COPILOT_EMBEDDING_UNAVAILABLE + COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING COPILOT_FAILED_TO_CREATE_MESSAGE COPILOT_FAILED_TO_GENERATE_TEXT COPILOT_FAILED_TO_MATCH_CONTEXT @@ -907,6 +927,9 @@ type Mutation { """add a file to context""" addContextFile(content: Upload!, options: AddContextFileInput!): CopilotContextFile! + + """Update workspace embedding files""" + addWorkspaceEmbeddingFiles(blob: Upload!, workspaceId: String!): CopilotWorkspaceFile! addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! approveMember(userId: String!, workspaceId: String!): Boolean! @@ -995,6 +1018,9 @@ type Mutation { """remove a file from context""" removeContextFile(options: RemoveContextFileInput!): Boolean! + + """Remove workspace embedding files""" + removeWorkspaceEmbeddingFiles(fileId: String!, workspaceId: String!): Boolean! removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! resumeSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType! retryAudioTranscription(jobId: String!, workspaceId: String!): TranscriptionResultType @@ -1037,6 +1063,9 @@ type Mutation { """Update workspace""" updateWorkspace(input: UpdateWorkspaceInput!): WorkspaceType! + """Update ignored docs""" + updateWorkspaceEmbeddingIgnoredDocs(add: [String!], remove: [String!], workspaceId: String!): Int! + """Upload user avatar""" uploadAvatar(avatar: Upload!): UserType! verifyEmail(token: String!): Boolean! @@ -1703,6 +1732,7 @@ type WorkspaceType { """Get get with given id""" doc(docId: String!): DocType! + embedding: CopilotWorkspaceConfig! """Enable AI""" enableAi: Boolean! diff --git a/packages/common/graphql/src/graphql/copilot-workspace-file-add.gql b/packages/common/graphql/src/graphql/copilot-workspace-file-add.gql new file mode 100644 index 0000000000..d6cee2aceb --- /dev/null +++ b/packages/common/graphql/src/graphql/copilot-workspace-file-add.gql @@ -0,0 +1,9 @@ +mutation addWorkspaceEmbeddingFiles($workspaceId: String!, $blob: Upload!) { + addWorkspaceEmbeddingFiles(workspaceId: $workspaceId, blob: $blob) { + fileId + fileName + mimeType + size + createdAt + } +} diff --git a/packages/common/graphql/src/graphql/copilot-workspace-file-remove.gql b/packages/common/graphql/src/graphql/copilot-workspace-file-remove.gql new file mode 100644 index 0000000000..19652cee1a --- /dev/null +++ b/packages/common/graphql/src/graphql/copilot-workspace-file-remove.gql @@ -0,0 +1,3 @@ +mutation removeWorkspaceEmbeddingFiles($workspaceId: String!, $fileId: String!) { + removeWorkspaceEmbeddingFiles(workspaceId: $workspaceId, fileId: $fileId) +} diff --git a/packages/common/graphql/src/graphql/copilot-workspace-get.gql b/packages/common/graphql/src/graphql/copilot-workspace-get.gql new file mode 100644 index 0000000000..e783a4c445 --- /dev/null +++ b/packages/common/graphql/src/graphql/copilot-workspace-get.gql @@ -0,0 +1,14 @@ +query getWorkspaceEmbeddingConfig($workspaceId: String!) { + workspace(id: $workspaceId) { + embedding { + files { + fileId + fileName + mimeType + size + createdAt + } + ignoredDocs + } + } +} diff --git a/packages/common/graphql/src/graphql/copilot-workspace-ignored-docs.gql b/packages/common/graphql/src/graphql/copilot-workspace-ignored-docs.gql new file mode 100644 index 0000000000..7ecabdb5d6 --- /dev/null +++ b/packages/common/graphql/src/graphql/copilot-workspace-ignored-docs.gql @@ -0,0 +1,3 @@ +mutation updateWorkspaceEmbeddingIgnoredDocs($workspaceId: String!, $add: [String!], $remove: [String!]) { + updateWorkspaceEmbeddingIgnoredDocs(workspaceId: $workspaceId, add: $add, remove: $remove) +} diff --git a/packages/common/graphql/src/graphql/index.ts b/packages/common/graphql/src/graphql/index.ts index 61348b9b4b..375cdf7037 100644 --- a/packages/common/graphql/src/graphql/index.ts +++ b/packages/common/graphql/src/graphql/index.ts @@ -741,6 +741,60 @@ export const getCopilotSessionsQuery = { }`, }; +export const addWorkspaceEmbeddingFilesMutation = { + id: 'addWorkspaceEmbeddingFilesMutation' as const, + op: 'addWorkspaceEmbeddingFiles', + query: `mutation addWorkspaceEmbeddingFiles($workspaceId: String!, $blob: Upload!) { + addWorkspaceEmbeddingFiles(workspaceId: $workspaceId, blob: $blob) { + fileId + fileName + mimeType + size + createdAt + } +}`, + file: true, +}; + +export const removeWorkspaceEmbeddingFilesMutation = { + id: 'removeWorkspaceEmbeddingFilesMutation' as const, + op: 'removeWorkspaceEmbeddingFiles', + query: `mutation removeWorkspaceEmbeddingFiles($workspaceId: String!, $fileId: String!) { + removeWorkspaceEmbeddingFiles(workspaceId: $workspaceId, fileId: $fileId) +}`, +}; + +export const getWorkspaceEmbeddingConfigQuery = { + id: 'getWorkspaceEmbeddingConfigQuery' as const, + op: 'getWorkspaceEmbeddingConfig', + query: `query getWorkspaceEmbeddingConfig($workspaceId: String!) { + workspace(id: $workspaceId) { + embedding { + files { + fileId + fileName + mimeType + size + createdAt + } + ignoredDocs + } + } +}`, +}; + +export const updateWorkspaceEmbeddingIgnoredDocsMutation = { + id: 'updateWorkspaceEmbeddingIgnoredDocsMutation' as const, + op: 'updateWorkspaceEmbeddingIgnoredDocs', + query: `mutation updateWorkspaceEmbeddingIgnoredDocs($workspaceId: String!, $add: [String!], $remove: [String!]) { + updateWorkspaceEmbeddingIgnoredDocs( + workspaceId: $workspaceId + add: $add + remove: $remove + ) +}`, +}; + export const createCheckoutSessionMutation = { id: 'createCheckoutSessionMutation' as const, op: 'createCheckoutSession', diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts index 424785551c..71eaea5088 100644 --- a/packages/common/graphql/src/schema.ts +++ b/packages/common/graphql/src/schema.ts @@ -239,6 +239,11 @@ export interface CopilotDocType { status: Maybe; } +export interface CopilotFailedToAddWorkspaceFileEmbeddingDataType { + __typename?: 'CopilotFailedToAddWorkspaceFileEmbeddingDataType'; + message: Scalars['String']['output']; +} + export interface CopilotFailedToMatchContextDataType { __typename?: 'CopilotFailedToMatchContextDataType'; content: Scalars['String']['output']; @@ -359,6 +364,23 @@ export interface CopilotSessionType { promptName: Scalars['String']['output']; } +export interface CopilotWorkspaceConfig { + __typename?: 'CopilotWorkspaceConfig'; + files: Array; + ignoredDocs: Array; + workspaceId: Scalars['String']['output']; +} + +export interface CopilotWorkspaceFile { + __typename?: 'CopilotWorkspaceFile'; + createdAt: Scalars['DateTime']['output']; + fileId: Scalars['String']['output']; + fileName: Scalars['String']['output']; + mimeType: Scalars['String']['output']; + size: Scalars['SafeInt']['output']; + workspaceId: Scalars['String']['output']; +} + export interface CreateChatMessageInput { attachments?: InputMaybe>; blobs?: InputMaybe>; @@ -507,6 +529,7 @@ export type ErrorDataUnion = | BlobNotFoundDataType | CopilotContextFileNotSupportedDataType | CopilotDocNotFoundDataType + | CopilotFailedToAddWorkspaceFileEmbeddingDataType | CopilotFailedToMatchContextDataType | CopilotFailedToModifyContextDataType | CopilotInvalidContextDataType @@ -572,6 +595,7 @@ export enum ErrorNames { COPILOT_DOC_NOT_FOUND = 'COPILOT_DOC_NOT_FOUND', COPILOT_EMBEDDING_DISABLED = 'COPILOT_EMBEDDING_DISABLED', COPILOT_EMBEDDING_UNAVAILABLE = 'COPILOT_EMBEDDING_UNAVAILABLE', + COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING = 'COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING', COPILOT_FAILED_TO_CREATE_MESSAGE = 'COPILOT_FAILED_TO_CREATE_MESSAGE', COPILOT_FAILED_TO_GENERATE_TEXT = 'COPILOT_FAILED_TO_GENERATE_TEXT', COPILOT_FAILED_TO_MATCH_CONTEXT = 'COPILOT_FAILED_TO_MATCH_CONTEXT', @@ -1040,6 +1064,8 @@ export interface Mutation { addContextDoc: CopilotContextDoc; /** add a file to context */ addContextFile: CopilotContextFile; + /** Update workspace embedding files */ + addWorkspaceEmbeddingFiles: CopilotWorkspaceFile; addWorkspaceFeature: Scalars['Boolean']['output']; approveMember: Scalars['Boolean']['output']; /** Ban an user */ @@ -1107,6 +1133,8 @@ export interface Mutation { removeContextDoc: Scalars['Boolean']['output']; /** remove a file from context */ removeContextFile: Scalars['Boolean']['output']; + /** Remove workspace embedding files */ + removeWorkspaceEmbeddingFiles: Scalars['Boolean']['output']; removeWorkspaceFeature: Scalars['Boolean']['output']; resumeSubscription: SubscriptionType; retryAudioTranscription: Maybe; @@ -1142,6 +1170,8 @@ export interface Mutation { updateUserFeatures: Array; /** Update workspace */ updateWorkspace: WorkspaceType; + /** Update ignored docs */ + updateWorkspaceEmbeddingIgnoredDocs: Scalars['Int']['output']; /** Upload user avatar */ uploadAvatar: UserType; verifyEmail: Scalars['Boolean']['output']; @@ -1171,6 +1201,11 @@ export interface MutationAddContextFileArgs { options: AddContextFileInput; } +export interface MutationAddWorkspaceEmbeddingFilesArgs { + blob: Scalars['Upload']['input']; + workspaceId: Scalars['String']['input']; +} + export interface MutationAddWorkspaceFeatureArgs { feature: FeatureType; workspaceId: Scalars['String']['input']; @@ -1364,6 +1399,11 @@ export interface MutationRemoveContextFileArgs { options: RemoveContextFileInput; } +export interface MutationRemoveWorkspaceEmbeddingFilesArgs { + fileId: Scalars['String']['input']; + workspaceId: Scalars['String']['input']; +} + export interface MutationRemoveWorkspaceFeatureArgs { feature: FeatureType; workspaceId: Scalars['String']['input']; @@ -1494,6 +1534,12 @@ export interface MutationUpdateWorkspaceArgs { input: UpdateWorkspaceInput; } +export interface MutationUpdateWorkspaceEmbeddingIgnoredDocsArgs { + add?: InputMaybe>; + remove?: InputMaybe>; + workspaceId: Scalars['String']['input']; +} + export interface MutationUploadAvatarArgs { avatar: Scalars['Upload']['input']; } @@ -2232,6 +2278,7 @@ export interface WorkspaceType { createdAt: Scalars['DateTime']['output']; /** Get get with given id */ doc: DocType; + embedding: CopilotWorkspaceConfig; /** Enable AI */ enableAi: Scalars['Boolean']['output']; /** Enable doc embedding */ @@ -3169,6 +3216,71 @@ export type GetCopilotSessionsQuery = { } | null; }; +export type AddWorkspaceEmbeddingFilesMutationVariables = Exact<{ + workspaceId: Scalars['String']['input']; + blob: Scalars['Upload']['input']; +}>; + +export type AddWorkspaceEmbeddingFilesMutation = { + __typename?: 'Mutation'; + addWorkspaceEmbeddingFiles: { + __typename?: 'CopilotWorkspaceFile'; + fileId: string; + fileName: string; + mimeType: string; + size: number; + createdAt: string; + }; +}; + +export type RemoveWorkspaceEmbeddingFilesMutationVariables = Exact<{ + workspaceId: Scalars['String']['input']; + fileId: Scalars['String']['input']; +}>; + +export type RemoveWorkspaceEmbeddingFilesMutation = { + __typename?: 'Mutation'; + removeWorkspaceEmbeddingFiles: boolean; +}; + +export type GetWorkspaceEmbeddingConfigQueryVariables = Exact<{ + workspaceId: Scalars['String']['input']; +}>; + +export type GetWorkspaceEmbeddingConfigQuery = { + __typename?: 'Query'; + workspace: { + __typename?: 'WorkspaceType'; + embedding: { + __typename?: 'CopilotWorkspaceConfig'; + ignoredDocs: Array; + files: Array<{ + __typename?: 'CopilotWorkspaceFile'; + fileId: string; + fileName: string; + mimeType: string; + size: number; + createdAt: string; + }>; + }; + }; +}; + +export type UpdateWorkspaceEmbeddingIgnoredDocsMutationVariables = Exact<{ + workspaceId: Scalars['String']['input']; + add?: InputMaybe< + Array | Scalars['String']['input'] + >; + remove?: InputMaybe< + Array | Scalars['String']['input'] + >; +}>; + +export type UpdateWorkspaceEmbeddingIgnoredDocsMutation = { + __typename?: 'Mutation'; + updateWorkspaceEmbeddingIgnoredDocs: number; +}; + export type CreateCheckoutSessionMutationVariables = Exact<{ input: CreateCheckoutSessionInput; }>; @@ -4427,6 +4539,11 @@ export type Queries = variables: GetCopilotSessionsQueryVariables; response: GetCopilotSessionsQuery; } + | { + name: 'getWorkspaceEmbeddingConfigQuery'; + variables: GetWorkspaceEmbeddingConfigQueryVariables; + response: GetWorkspaceEmbeddingConfigQuery; + } | { name: 'getDocRolePermissionsQuery'; variables: GetDocRolePermissionsQueryVariables; @@ -4784,6 +4901,21 @@ export type Mutations = variables: UpdateCopilotSessionMutationVariables; response: UpdateCopilotSessionMutation; } + | { + name: 'addWorkspaceEmbeddingFilesMutation'; + variables: AddWorkspaceEmbeddingFilesMutationVariables; + response: AddWorkspaceEmbeddingFilesMutation; + } + | { + name: 'removeWorkspaceEmbeddingFilesMutation'; + variables: RemoveWorkspaceEmbeddingFilesMutationVariables; + response: RemoveWorkspaceEmbeddingFilesMutation; + } + | { + name: 'updateWorkspaceEmbeddingIgnoredDocsMutation'; + variables: UpdateWorkspaceEmbeddingIgnoredDocsMutationVariables; + response: UpdateWorkspaceEmbeddingIgnoredDocsMutation; + } | { name: 'createCheckoutSessionMutation'; variables: CreateCheckoutSessionMutationVariables; diff --git a/packages/frontend/i18n/src/i18n.gen.ts b/packages/frontend/i18n/src/i18n.gen.ts index 4ee70257fe..4285b90097 100644 --- a/packages/frontend/i18n/src/i18n.gen.ts +++ b/packages/frontend/i18n/src/i18n.gen.ts @@ -8193,6 +8193,12 @@ export function useAFFiNEI18N(): { * `Audio not provided.` */ ["error.COPILOT_TRANSCRIPTION_AUDIO_NOT_PROVIDED"](): string; + /** + * `Failed to add workspace file embedding: {{message}}` + */ + ["error.COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING"](options: { + readonly message: string; + }): string; /** * `You have exceeded your blob size quota.` */ diff --git a/packages/frontend/i18n/src/resources/en.json b/packages/frontend/i18n/src/resources/en.json index fd767d9db5..480e5b3ed2 100644 --- a/packages/frontend/i18n/src/resources/en.json +++ b/packages/frontend/i18n/src/resources/en.json @@ -2027,6 +2027,7 @@ "error.COPILOT_TRANSCRIPTION_JOB_EXISTS": "Transcription job already exists", "error.COPILOT_TRANSCRIPTION_JOB_NOT_FOUND": "Transcription job not found.", "error.COPILOT_TRANSCRIPTION_AUDIO_NOT_PROVIDED": "Audio not provided.", + "error.COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING": "Failed to add workspace file embedding: {{message}}", "error.BLOB_QUOTA_EXCEEDED": "You have exceeded your blob size quota.", "error.STORAGE_QUOTA_EXCEEDED": "You have exceeded your storage quota.", "error.MEMBER_QUOTA_EXCEEDED": "You have exceeded your workspace member quota.",