From b0aa2c90fda4130c23aeccdc9641b726275d4f09 Mon Sep 17 00:00:00 2001 From: darkskygit Date: Mon, 17 Mar 2025 14:17:01 +0000 Subject: [PATCH] feat(server): tag and collection record for context (#10926) fix CLOUD-174 --- .../src/plugins/copilot/context/resolver.ts | 81 +++++++++++++++++++ .../src/plugins/copilot/context/service.ts | 22 ++++- .../src/plugins/copilot/context/session.ts | 28 ++++++- .../src/plugins/copilot/context/types.ts | 19 +++++ packages/backend/server/src/schema.gql | 23 ++++++ .../graphql/copilot-context-category-add.gql | 7 ++ .../copilot-context-category-remove.gql | 3 + .../frontend/graphql/src/graphql/index.ts | 20 +++++ packages/frontend/graphql/src/schema.ts | 63 +++++++++++++++ 9 files changed, 262 insertions(+), 4 deletions(-) create mode 100644 packages/frontend/graphql/src/graphql/copilot-context-category-add.gql create mode 100644 packages/frontend/graphql/src/graphql/copilot-context-category-remove.gql diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index 41605c51b1..1a22f9cdfc 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -40,6 +40,8 @@ import { CopilotStorage } from '../storage'; import { CopilotContextDocJob } from './job'; import { CopilotContextService } from './service'; import { + ContextCategories, + ContextCategory, ContextDoc, ContextEmbedStatus, type ContextFile, @@ -49,6 +51,18 @@ import { } from './types'; import { readStream } from './utils'; +@InputType() +class AddRemoveContextCategoryInput { + @Field(() => String) + contextId!: string; + + @Field(() => ContextCategories) + type!: ContextCategories; + + @Field(() => String) + categoryId!: string; +} + @InputType() class AddContextDocInput { @Field(() => String) @@ -94,6 +108,20 @@ export class CopilotContextType { workspaceId!: string; } +registerEnumType(ContextCategories, { name: 'ContextCategories' }); + +@ObjectType() +class CopilotContextCategory implements ContextCategory { + @Field(() => ID) + id!: string; + + @Field(() => ContextCategories) + type!: ContextCategories; + + @Field(() => SafeIntResolver) + createdAt!: number; +} + registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' }); @ObjectType() @@ -335,6 +363,59 @@ export class CopilotContextResolver { return session.listDocs(); } + @Mutation(() => CopilotContextCategory, { + description: 'add a category to context', + }) + @CallMetric('ai', 'context_category_add') + async addContextCategory( + @Args({ name: 'options', type: () => AddRemoveContextCategoryInput }) + options: AddRemoveContextCategoryInput + ) { + const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; + await using lock = await this.mutex.acquire(lockFlag); + if (!lock) { + return new TooManyRequest('Server is busy'); + } + const session = await this.context.get(options.contextId); + + try { + return await session.addCategoryRecord(options.type, options.categoryId); + } catch (e: any) { + throw new CopilotFailedToModifyContext({ + contextId: options.contextId, + message: e.message, + }); + } + } + + @Mutation(() => Boolean, { + description: 'remove a category from context', + }) + @CallMetric('ai', 'context_category_remove') + async removeContextCategory( + @Args({ name: 'options', type: () => AddRemoveContextCategoryInput }) + options: AddRemoveContextCategoryInput + ) { + const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; + await using lock = await this.mutex.acquire(lockFlag); + if (!lock) { + return new TooManyRequest('Server is busy'); + } + const session = await this.context.get(options.contextId); + + try { + return await session.removeCategoryRecord( + options.type, + options.categoryId + ); + } catch (e: any) { + throw new CopilotFailedToModifyContext({ + contextId: options.contextId, + message: e.message, + }); + } + } + @Mutation(() => CopilotContextDoc, { description: 'add a doc to context', }) diff --git a/packages/backend/server/src/plugins/copilot/context/service.ts b/packages/backend/server/src/plugins/copilot/context/service.ts index e178b2c2a4..5ca8560490 100644 --- a/packages/backend/server/src/plugins/copilot/context/service.ts +++ b/packages/backend/server/src/plugins/copilot/context/service.ts @@ -19,6 +19,7 @@ import { ContextEmbedStatus, ContextFile, EmbeddingClient, + MinimalContextConfigSchema, } from './types'; import { checkEmbeddingAvailable } from './utils'; @@ -128,7 +129,12 @@ export class CopilotContextService implements OnModuleInit { const context = await this.db.aiContext.create({ data: { sessionId, - config: { workspaceId: session.workspaceId, docs: [], files: [] }, + config: { + workspaceId: session.workspaceId, + docs: [], + files: [], + categories: [], + }, }, }); @@ -149,7 +155,19 @@ export class CopilotContextService implements OnModuleInit { }); if (ret) { const config = ContextConfigSchema.safeParse(ret.config); - if (config.success) return this.cacheSession(id, config.data); + if (config.success) { + return this.cacheSession(id, config.data); + } + const minimalConfig = MinimalContextConfigSchema.safeParse(ret.config); + if (minimalConfig.success) { + // fulfill the missing fields + return this.cacheSession(id, { + ...minimalConfig.data, + docs: [], + files: [], + categories: [], + }); + } } throw new CopilotInvalidContext({ contextId: id }); } diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index c2195fde96..b52139f787 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -4,6 +4,7 @@ import { nanoid } from 'nanoid'; import { PrismaTransaction } from '../../../base'; import { ChunkSimilarity, + ContextCategories, ContextConfig, ContextDoc, ContextEmbedStatus, @@ -49,6 +50,30 @@ export class ContextSession implements AsyncDisposable { ) as ContextList; } + async addCategoryRecord(type: ContextCategories, id: string) { + const category = this.config.categories.find( + c => c.type === type && c.id === id + ); + if (category) { + return category; + } + const record = { id, type, createdAt: Date.now() }; + this.config.categories.push(record); + await this.save(); + return record; + } + + async removeCategoryRecord(type: ContextCategories, id: string) { + const index = this.config.categories.findIndex( + c => c.type === type && c.id === id + ); + if (index >= 0) { + this.config.categories.splice(index, 1); + await this.save(); + } + return true; + } + async addDocRecord(docId: string): Promise { const doc = this.config.docs.find(f => f.id === docId); if (doc) { @@ -65,9 +90,8 @@ export class ContextSession implements AsyncDisposable { if (index >= 0) { this.config.docs.splice(index, 1); await this.save(); - return true; } - return false; + return true; } async addFile(blobId: string, name: string): Promise { diff --git a/packages/backend/server/src/plugins/copilot/context/types.ts b/packages/backend/server/src/plugins/copilot/context/types.ts index 471145649b..c325f8fafc 100644 --- a/packages/backend/server/src/plugins/copilot/context/types.ts +++ b/packages/backend/server/src/plugins/copilot/context/types.ts @@ -32,6 +32,11 @@ export enum ContextEmbedStatus { failed = 'failed', } +export enum ContextCategories { + Tag = 'tag', + Collection = 'collection', +} + export const ContextConfigSchema = z.object({ workspaceId: z.string(), files: z @@ -64,9 +69,23 @@ export const ContextConfigSchema = z.object({ createdAt: z.number(), }) .array(), + categories: z + .object({ + id: z.string(), + type: z.enum([ContextCategories.Tag, ContextCategories.Collection]), + createdAt: z.number(), + }) + .array(), +}); + +export const MinimalContextConfigSchema = ContextConfigSchema.pick({ + workspaceId: true, }); export type ContextConfig = z.infer; +export type ContextCategory = z.infer< + typeof ContextConfigSchema +>['categories'][number]; export type ContextDoc = z.infer['docs'][number]; export type ContextFile = z.infer['files'][number]; export type ContextListItem = ContextDoc | ContextFile; diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index 6609e8cee8..6050039555 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -12,6 +12,12 @@ input AddContextFileInput { contextId: String! } +input AddRemoveContextCategoryInput { + categoryId: String! + contextId: String! + type: ContextCategories! +} + type AlreadyInSpaceDataType { spaceId: String! } @@ -35,6 +41,11 @@ type ChatMessage { role: String! } +enum ContextCategories { + Collection + Tag +} + enum ContextEmbedStatus { failed finished @@ -92,6 +103,12 @@ type CopilotContext { workspaceId: String! } +type CopilotContextCategory { + createdAt: SafeInt! + id: ID! + type: ContextCategories! +} + type CopilotContextDoc { createdAt: SafeInt! id: ID! @@ -803,6 +820,9 @@ type Mutation { acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean! activateLicense(license: String!, workspaceId: String!): License! + """add a category to context""" + addContextCategory(options: AddRemoveContextCategoryInput!): CopilotContextCategory! + """add a doc to context""" addContextDoc(options: AddContextDocInput!): CopilotContextDoc! @@ -887,6 +907,9 @@ type Mutation { """Remove user avatar""" removeAvatar: RemoveAvatar! + """remove a category from context""" + removeContextCategory(options: AddRemoveContextCategoryInput!): Boolean! + """remove a doc from context""" removeContextDoc(options: RemoveContextDocInput!): Boolean! diff --git a/packages/frontend/graphql/src/graphql/copilot-context-category-add.gql b/packages/frontend/graphql/src/graphql/copilot-context-category-add.gql new file mode 100644 index 0000000000..326fb4a38c --- /dev/null +++ b/packages/frontend/graphql/src/graphql/copilot-context-category-add.gql @@ -0,0 +1,7 @@ +mutation addContextCategory($options: AddRemoveContextCategoryInput!) { + addContextCategory(options: $options) { + id + createdAt + type + } +} diff --git a/packages/frontend/graphql/src/graphql/copilot-context-category-remove.gql b/packages/frontend/graphql/src/graphql/copilot-context-category-remove.gql new file mode 100644 index 0000000000..05d57c21fe --- /dev/null +++ b/packages/frontend/graphql/src/graphql/copilot-context-category-remove.gql @@ -0,0 +1,3 @@ +mutation removeContextCategory($options: AddRemoveContextCategoryInput!) { + removeContextCategory(options: $options) +} diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts index c88eb5106d..5a8312bd0f 100644 --- a/packages/frontend/graphql/src/graphql/index.ts +++ b/packages/frontend/graphql/src/graphql/index.ts @@ -134,6 +134,26 @@ export const changePasswordMutation = { }`, }; +export const addContextCategoryMutation = { + id: 'addContextCategoryMutation' as const, + op: 'addContextCategory', + query: `mutation addContextCategory($options: AddRemoveContextCategoryInput!) { + addContextCategory(options: $options) { + id + createdAt + type + } +}`, +}; + +export const removeContextCategoryMutation = { + id: 'removeContextCategoryMutation' as const, + op: 'removeContextCategory', + query: `mutation removeContextCategory($options: AddRemoveContextCategoryInput!) { + removeContextCategory(options: $options) +}`, +}; + export const createCopilotContextMutation = { id: 'createCopilotContextMutation' as const, op: 'createCopilotContext', diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts index 46e78fea25..bd5ec0a545 100644 --- a/packages/frontend/graphql/src/schema.ts +++ b/packages/frontend/graphql/src/schema.ts @@ -47,6 +47,12 @@ export interface AddContextFileInput { contextId: Scalars['String']['input']; } +export interface AddRemoveContextCategoryInput { + categoryId: Scalars['String']['input']; + contextId: Scalars['String']['input']; + type: ContextCategories; +} + export interface AlreadyInSpaceDataType { __typename?: 'AlreadyInSpaceDataType'; spaceId: Scalars['String']['output']; @@ -73,6 +79,11 @@ export interface ChatMessage { role: Scalars['String']['output']; } +export enum ContextCategories { + Collection = 'Collection', + Tag = 'Tag', +} + export enum ContextEmbedStatus { failed = 'failed', finished = 'finished', @@ -163,6 +174,13 @@ export interface CopilotContextMatchWorkspaceContextArgs { limit?: InputMaybe; } +export interface CopilotContextCategory { + __typename?: 'CopilotContextCategory'; + createdAt: Scalars['SafeInt']['output']; + id: Scalars['ID']['output']; + type: ContextCategories; +} + export interface CopilotContextDoc { __typename?: 'CopilotContextDoc'; createdAt: Scalars['SafeInt']['output']; @@ -938,6 +956,8 @@ export interface Mutation { __typename?: 'Mutation'; acceptInviteById: Scalars['Boolean']['output']; activateLicense: License; + /** add a category to context */ + addContextCategory: CopilotContextCategory; /** add a doc to context */ addContextDoc: CopilotContextDoc; /** add a file to context */ @@ -1002,6 +1022,8 @@ export interface Mutation { releaseDeletedBlobs: Scalars['Boolean']['output']; /** Remove user avatar */ removeAvatar: RemoveAvatar; + /** remove a category from context */ + removeContextCategory: Scalars['Boolean']['output']; /** remove a doc from context */ removeContextDoc: Scalars['Boolean']['output']; /** remove a file from context */ @@ -1054,6 +1076,10 @@ export interface MutationActivateLicenseArgs { workspaceId: Scalars['String']['input']; } +export interface MutationAddContextCategoryArgs { + options: AddRemoveContextCategoryInput; +} + export interface MutationAddContextDocArgs { options: AddContextDocInput; } @@ -1240,6 +1266,10 @@ export interface MutationReleaseDeletedBlobsArgs { workspaceId: Scalars['String']['input']; } +export interface MutationRemoveContextCategoryArgs { + options: AddRemoveContextCategoryInput; +} + export interface MutationRemoveContextDocArgs { options: RemoveContextDocInput; } @@ -2319,6 +2349,29 @@ export type ChangePasswordMutation = { changePassword: boolean; }; +export type AddContextCategoryMutationVariables = Exact<{ + options: AddRemoveContextCategoryInput; +}>; + +export type AddContextCategoryMutation = { + __typename?: 'Mutation'; + addContextCategory: { + __typename?: 'CopilotContextCategory'; + id: string; + createdAt: number; + type: ContextCategories; + }; +}; + +export type RemoveContextCategoryMutationVariables = Exact<{ + options: AddRemoveContextCategoryInput; +}>; + +export type RemoveContextCategoryMutation = { + __typename?: 'Mutation'; + removeContextCategory: boolean; +}; + export type CreateCopilotContextMutationVariables = Exact<{ workspaceId: Scalars['String']['input']; sessionId: Scalars['String']['input']; @@ -4308,6 +4361,16 @@ export type Mutations = variables: ChangePasswordMutationVariables; response: ChangePasswordMutation; } + | { + name: 'addContextCategoryMutation'; + variables: AddContextCategoryMutationVariables; + response: AddContextCategoryMutation; + } + | { + name: 'removeContextCategoryMutation'; + variables: RemoveContextCategoryMutationVariables; + response: RemoveContextCategoryMutation; + } | { name: 'createCopilotContextMutation'; variables: CreateCopilotContextMutationVariables;