feat(server): tag and collection record for context (#10926)

fix CLOUD-174

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',
   })

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 });
   }

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<ContextDoc> {
     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<ContextFile> {

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<typeof ContextConfigSchema>;
+export type ContextCategory = z.infer<
+  typeof ContextConfigSchema
+>['categories'][number];
 export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
 export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
 export type ContextListItem = ContextDoc | ContextFile;

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!
 

packages/frontend/graphql/src/graphql/copilot-context-category-add.gql

@@ -0,0 +1,7 @@
+mutation addContextCategory($options: AddRemoveContextCategoryInput!) {
+  addContextCategory(options: $options) {
+    id
+    createdAt
+    type
+  }
+}

packages/frontend/graphql/src/graphql/copilot-context-category-remove.gql

@@ -0,0 +1,3 @@
+mutation removeContextCategory($options: AddRemoveContextCategoryInput!) {
+  removeContextCategory(options: $options)
+}

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

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<Scalars['SafeInt']['input']>;
 }
 
+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;