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!