feat(server): adapt context model (#11028)

expose more field in listContextObject

packages/backend/server/migrations/20250210090228_ai_context_embedding/migration.sql

@@ -37,7 +37,7 @@ BEGIN -- check if pgvector extension is installed
         "file_id" VARCHAR NOT NULL,
         "chunk" INTEGER NOT NULL,
         "content" VARCHAR NOT NULL,
-        "embedding" vector(512) NOT NULL,
+        "embedding" vector(1024) NOT NULL,
         "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
         "updated_at" TIMESTAMPTZ(3) NOT NULL,
 
@@ -50,7 +50,7 @@ BEGIN -- check if pgvector extension is installed
         "doc_id" VARCHAR NOT NULL,
         "chunk" INTEGER NOT NULL,
         "content" VARCHAR NOT NULL,
-        "embedding" vector(512) NOT NULL,
+        "embedding" vector(1024) NOT NULL,
         "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
         "updated_at" TIMESTAMPTZ(3) NOT NULL,
 

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -1,6 +1,7 @@
 import { randomUUID } from 'node:crypto';
 
 import { ProjectRoot } from '@affine-tools/utils/path';
+import { PrismaClient } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
@@ -9,6 +10,7 @@ import { EventBus } from '../base';
 import { ConfigModule } from '../base/config';
 import { AuthService } from '../core/auth';
 import { QuotaModule } from '../core/quota';
+import { ContextCategories } from '../models';
 import { CopilotModule } from '../plugins/copilot';
 import {
   CopilotContextDocJob,
@@ -54,6 +56,7 @@ import { MockCopilotTestProvider, WorkflowTestCases } from './utils/copilot';
 const test = ava as TestFn<{
   auth: AuthService;
   module: TestingModule;
+  db: PrismaClient;
   event: EventBus;
   context: CopilotContextService;
   prompt: PromptService;
@@ -95,6 +98,7 @@ test.before(async t => {
   });
 
   const auth = module.get(AuthService);
+  const db = module.get(PrismaClient);
   const event = module.get(EventBus);
   const context = module.get(CopilotContextService);
   const prompt = module.get(PromptService);
@@ -106,6 +110,7 @@ test.before(async t => {
 
   t.context.module = module;
   t.context.auth = auth;
+  t.context.db = db;
   t.context.event = event;
   t.context.context = context;
   t.context.prompt = prompt;
@@ -1338,47 +1343,112 @@ test('should be able to manage context', async t => {
   {
     const session = await context.create(chatSession);
 
-    await storage.put(userId, session.workspaceId, 'blob', buffer);
+    // file record
+    {
+      await storage.put(userId, session.workspaceId, 'blob', buffer);
+      const file = await session.addFile('blob', 'sample.pdf');
 
-    const file = await session.addFile('blob', 'sample.pdf');
+      const handler = Sinon.spy(event, 'emit');
 
-    const handler = Sinon.spy(event, 'emit');
+      await jobs.embedPendingFile({
+        userId,
+        workspaceId: session.workspaceId,
+        contextId: session.id,
+        blobId: file.blobId,
+        fileId: file.id,
+        fileName: file.name,
+      });
 
-    await jobs.embedPendingFile({
+      t.deepEqual(handler.lastCall.args, [
-      userId,
+        'workspace.file.embed.finished',
-      workspaceId: session.workspaceId,
+        {
-      contextId: session.id,
+          contextId: session.id,
-      blobId: file.blobId,
+          fileId: file.id,
-      fileId: file.id,
+          chunkSize: 1,
-      fileName: file.name,
+        },
+      ]);
+
+      const list = session.files;
+      t.deepEqual(
+        list.map(f => f.id),
+        [file.id],
+        'should list file id'
+      );
+
+      const result = await session.matchFileChunks('test', 1, undefined, 1);
+      t.is(result.length, 1, 'should match context');
+      t.is(result[0].fileId, file.id, 'should match file id');
+    }
+
+    // doc record
+    const docId = randomUUID();
+    await t.context.db.snapshot.create({
+      data: {
+        workspaceId: session.workspaceId,
+        id: docId,
+        blob: Buffer.from([1, 1]),
+        state: Buffer.from([1, 1]),
+        updatedAt: new Date(),
+        createdAt: new Date(),
+      },
     });
 
-    t.deepEqual(handler.lastCall.args, [
+    {
-      'workspace.file.embed.finished',
+      await session.addDocRecord(docId);
-      {
+      const docs = session.docs.map(d => d.id);
-        contextId: session.id,
+      t.deepEqual(docs, [docId], 'should list doc id');
-        fileId: file.id,
-        chunkSize: 1,
-      },
-    ]);
 
-    const list = session.listFiles();
+      await session.removeDocRecord(docId);
-    t.deepEqual(
+      t.deepEqual(session.docs, [], 'should remove doc id');
-      list.map(f => f.id),
+    }
-      [file.id],
-      'should list file id'
-    );
 
-    const docId = randomUUID();
+    // tag record
-    await session.addDocRecord(docId);
+    {
-    const docs = session.listDocs().map(d => d.id);
+      const tagId = randomUUID();
-    t.deepEqual(docs, [docId], 'should list doc id');
+      await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId]);
+      const tags = session.tags.map(t => t.id);
+      t.deepEqual(tags, [tagId], 'should list tag id');
 
-    await session.removeDocRecord(docId);
+      await session.removeCategoryRecord(ContextCategories.Tag, tagId);
-    t.deepEqual(session.listDocs(), [], 'should remove doc id');
+      t.deepEqual(session.tags, [], 'should remove tag id');
 
-    const result = await session.matchFileChunks('test', 1, undefined, 1);
+      await t.throwsAsync(
-    t.is(result.length, 1, 'should match context');
+        session.addCategoryRecord(ContextCategories.Tag, tagId, [
-    t.is(result[0].fileId, file.id, 'should match file id');
+          'not-exists-doc',
+        ]),
+        {
+          instanceOf: Error,
+        },
+        'should throw error if doc id not exists'
+      );
+    }
+
+    // collection record
+    {
+      const collectionId = randomUUID();
+      await session.addCategoryRecord(
+        ContextCategories.Collection,
+        collectionId,
+        [docId]
+      );
+      const collection = session.collections.map(l => l.id);
+      t.deepEqual(collection, [collectionId], 'should list collection id');
+
+      await session.removeCategoryRecord(
+        ContextCategories.Collection,
+        collectionId
+      );
+      t.deepEqual(session.collections, [], 'should remove collection id');
+
+      await t.throwsAsync(
+        session.addCategoryRecord(ContextCategories.Collection, collectionId, [
+          'not-exists-doc',
+        ]),
+        {
+          instanceOf: Error,
+        },
+        'should throw error if doc id not exists'
+      );
+    }
   }
 });

packages/backend/server/src/__tests__/models/copilot-context.spec.ts

@@ -84,6 +84,7 @@ test('should update context', async t => {
   const doc = {
     id: docId,
     createdAt: Date.now(),
+    status: null,
   };
   config?.docs.push(doc);
   await t.context.copilotContext.update(contextId, { config });
@@ -96,16 +97,20 @@ test('should insert embedding by doc id', async t => {
   const { id: contextId } = await t.context.copilotContext.create(session.id);
 
   {
-    await t.context.copilotContext.insertEmbedding(contextId, 'file-id', [
+    await t.context.copilotContext.insertContentEmbedding(
-      {
+      contextId,
-        index: 0,
+      'file-id',
-        content: 'content',
+      [
-        embedding: Array.from({ length: 512 }, () => 1),
+        {
-      },
+          index: 0,
-    ]);
+          content: 'content',
+          embedding: Array.from({ length: 512 }, () => 1),
+        },
+      ]
+    );
 
     {
-      const ret = await t.context.copilotContext.matchEmbedding(
+      const ret = await t.context.copilotContext.matchContentEmbedding(
         Array.from({ length: 512 }, () => 0.9),
         contextId,
         1,
@@ -117,7 +122,7 @@ test('should insert embedding by doc id', async t => {
 
     {
       await t.context.copilotContext.deleteEmbedding(contextId, 'file-id');
-      const ret = await t.context.copilotContext.matchEmbedding(
+      const ret = await t.context.copilotContext.matchContentEmbedding(
         Array.from({ length: 512 }, () => 0.9),
         contextId,
         1,

packages/backend/server/src/base/error/def.ts

@@ -650,6 +650,10 @@ export const USER_FRIENDLY_ERRORS = {
     args: { docId: 'string' },
     message: ({ docId }) => `Doc ${docId} not found.`,
   },
+  copilot_docs_not_found: {
+    type: 'resource_not_found',
+    message: () => `Some docs not found.`,
+  },
   copilot_message_not_found: {
     type: 'resource_not_found',
     args: { messageId: 'string' },

packages/backend/server/src/base/error/errors.gen.ts

@@ -664,6 +664,12 @@ export class CopilotDocNotFound extends UserFriendlyError {
     super('resource_not_found', 'copilot_doc_not_found', message, args);
   }
 }
+
+export class CopilotDocsNotFound extends UserFriendlyError {
+  constructor(message?: string) {
+    super('resource_not_found', 'copilot_docs_not_found', message);
+  }
+}
 @ObjectType()
 class CopilotMessageNotFoundDataType {
   @Field() messageId!: string
@@ -997,6 +1003,7 @@ export enum ErrorNames {
   UNSPLASH_IS_NOT_CONFIGURED,
   COPILOT_ACTION_TAKEN,
   COPILOT_DOC_NOT_FOUND,
+  COPILOT_DOCS_NOT_FOUND,
   COPILOT_MESSAGE_NOT_FOUND,
   COPILOT_PROMPT_NOT_FOUND,
   COPILOT_PROMPT_INVALID,

packages/backend/server/src/models/common/copilot.ts

@@ -34,6 +34,13 @@ export enum ContextCategories {
 export const ContextDocSchema = z.object({
   id: z.string(),
   createdAt: z.number(),
+  status: z
+    .enum([
+      ContextEmbedStatus.processing,
+      ContextEmbedStatus.finished,
+      ContextEmbedStatus.failed,
+    ])
+    .nullable(),
 });
 
 export const ContextFileSchema = z.object({

packages/backend/server/src/models/copilot-context.ts

@@ -6,7 +6,6 @@ import { Prisma } from '@prisma/client';
 import { CopilotSessionNotFound } from '../base';
 import { BaseModel } from './base';
 import {
-  ChunkSimilarity,
   ContextConfigSchema,
   ContextDoc,
   ContextEmbedStatus,
@@ -24,7 +23,7 @@ type UpdateCopilotContextInput = Pick<CopilotContext, 'config'>;
  */
 @Injectable()
 export class CopilotContextModel extends BaseModel {
-  // contexts
+  // ================ contexts ================
 
   async create(sessionId: string) {
     const session = await this.db.aiSession.findFirst({
@@ -113,7 +112,7 @@ export class CopilotContextModel extends BaseModel {
     return ret.count > 0;
   }
 
-  // embeddings
+  // ================ embeddings ================
 
   async checkEmbeddingAvailable(): Promise<boolean> {
     const [{ count }] = await this.db.$queryRaw<
@@ -157,7 +156,7 @@ export class CopilotContextModel extends BaseModel {
     return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
   }
 
-  async insertEmbedding(
+  async insertContentEmbedding(
     contextId: string,
     fileId: string,
     embeddings: Embedding[]
@@ -172,12 +171,12 @@ export class CopilotContextModel extends BaseModel {
   `;
   }
 
-  async matchEmbedding(
+  async matchContentEmbedding(
     embedding: number[],
     contextId: string,
     topK: number,
     threshold: number
-  ): Promise<ChunkSimilarity[]> {
+  ): Promise<FileChunkSimilarity[]> {
     const similarityChunks = await this.db.$queryRaw<
       Array<FileChunkSimilarity>
     >`
@@ -214,7 +213,7 @@ export class CopilotContextModel extends BaseModel {
     workspaceId: string,
     topK: number,
     threshold: number
-  ): Promise<ChunkSimilarity[]> {
+  ): Promise<DocChunkSimilarity[]> {
     const similarityChunks = await this.db.$queryRaw<Array<DocChunkSimilarity>>`
        SELECT "doc_id" as "docId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
        FROM "ai_workspace_embeddings"

packages/backend/server/src/models/doc.ts

@@ -185,6 +185,23 @@ export class DocModel extends BaseModel {
     });
   }
 
+  /**
+   * Check if all doc exists in the workspace.
+   * Ignore pending updates.
+   */
+  async existsAll(workspaceId: string, docIds: string[]) {
+    const count = await this.db.snapshot.count({
+      where: {
+        workspaceId,
+        id: { in: docIds },
+      },
+    });
+    if (count === docIds.length) {
+      return true;
+    }
+    return false;
+  }
+
   /**
    * Detect a doc exists or not, including updates
    */

packages/backend/server/src/plugins/copilot/context/embedding.ts

@@ -1,6 +1,7 @@
 import OpenAI from 'openai';
 
-import { Embedding, EmbeddingClient } from './types';
+import { Embedding } from '../../../models';
+import { EmbeddingClient } from './types';
 
 export class OpenAIEmbeddingClient extends EmbeddingClient {
   constructor(private readonly client: OpenAI) {
@@ -15,6 +16,7 @@ export class OpenAIEmbeddingClient extends EmbeddingClient {
       {
         input,
         model: 'text-embedding-3-large',
+        dimensions: 1024,
         encoding_format: 'float',
       },
       { signal }

packages/backend/server/src/plugins/copilot/context/index.ts

@@ -1,7 +1,3 @@
 export { CopilotContextDocJob } from './job';
 export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
 export { CopilotContextService } from './service';
-export {
-  type ContextFile,
-  ContextEmbedStatus as ContextFileStatus,
-} from './types';

packages/backend/server/src/plugins/copilot/context/job.ts

@@ -1,7 +1,4 @@
-import { randomUUID } from 'node:crypto';
-
 import { Injectable, OnModuleInit } from '@nestjs/common';
-import { Prisma, PrismaClient } from '@prisma/client';
 import OpenAI from 'openai';
 
 import {
@@ -15,10 +12,11 @@ import {
   OnJob,
 } from '../../../base';
 import { DocReader } from '../../../core/doc';
+import { Models } from '../../../models';
 import { CopilotStorage } from '../storage';
 import { OpenAIEmbeddingClient } from './embedding';
-import { Embedding, EmbeddingClient } from './types';
+import { EmbeddingClient } from './types';
-import { checkEmbeddingAvailable, readStream } from './utils';
+import { readStream } from './utils';
 
 declare global {
   interface Jobs {
@@ -45,10 +43,10 @@ export class CopilotContextDocJob implements OnModuleInit {
 
   constructor(
     config: Config,
-    private readonly db: PrismaClient,
     private readonly doc: DocReader,
     private readonly event: EventBus,
     private readonly logger: AFFiNELogger,
+    private readonly models: Models,
     private readonly queue: JobQueue,
     private readonly storage: CopilotStorage
   ) {
@@ -60,7 +58,8 @@ export class CopilotContextDocJob implements OnModuleInit {
   }
 
   async onModuleInit() {
-    this.supportEmbedding = await checkEmbeddingAvailable(this.db);
+    this.supportEmbedding =
+      await this.models.copilotContext.checkEmbeddingAvailable();
   }
 
   // public this client to allow overriding in tests
@@ -91,23 +90,6 @@ export class CopilotContextDocJob implements OnModuleInit {
     }
   }
 
-  private processEmbeddings(
-    contextOrWorkspaceId: string,
-    fileOrDocId: string,
-    embeddings: Embedding[]
-  ) {
-    const groups = embeddings.map(e => [
-      randomUUID(),
-      contextOrWorkspaceId,
-      fileOrDocId,
-      e.index,
-      e.content,
-      Prisma.raw(`'[${e.embedding.join(',')}]'`),
-      new Date(),
-    ]);
-    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
-  }
-
   async readCopilotBlob(
     userId: string,
     workspaceId: string,
@@ -145,14 +127,11 @@ export class CopilotContextDocJob implements OnModuleInit {
 
       for (const chunk of chunks) {
         const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
-        const values = this.processEmbeddings(contextId, fileId, embeddings);
+        await this.models.copilotContext.insertContentEmbedding(
-
+          contextId,
-        await this.db.$executeRaw`
+          fileId,
-          INSERT INTO "ai_context_embeddings"
+          embeddings
-          ("id", "context_id", "file_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
+        );
-          ON CONFLICT (context_id, file_id, chunk) DO UPDATE SET
-          content = EXCLUDED.content, embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
-        `;
       }
 
       this.event.emit('workspace.file.embed.finished', {
@@ -188,13 +167,11 @@ export class CopilotContextDocJob implements OnModuleInit {
         );
 
         for (const chunks of embeddings) {
-          const values = this.processEmbeddings(workspaceId, docId, chunks);
+          await this.models.copilotContext.insertWorkspaceEmbedding(
-          await this.db.$executeRaw`
+            workspaceId,
-              INSERT INTO "ai_workspace_embeddings"
+            docId,
-              ("workspace_id", "doc_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
+            chunks
-              ON CONFLICT (context_id, file_id, chunk) DO UPDATE SET
+          );
-              embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
-            `;
         }
       }
     } catch (e: any) {

packages/backend/server/src/plugins/copilot/context/resolver.ts

@@ -34,25 +34,41 @@ import {
 } from '../../../base';
 import { CurrentUser } from '../../../core/auth';
 import { AccessController } from '../../../core/permission';
-import { COPILOT_LOCKER, CopilotType } from '../resolver';
-import { ChatSessionService } from '../session';
-import { CopilotStorage } from '../storage';
-import { CopilotContextDocJob } from './job';
-import { CopilotContextService } from './service';
 import {
   ContextCategories,
   ContextCategory,
   ContextDoc,
   ContextEmbedStatus,
-  type ContextFile,
+  ContextFile,
   DocChunkSimilarity,
   FileChunkSimilarity,
-  MAX_EMBEDDABLE_SIZE,
+  Models,
-} from './types';
+} from '../../../models';
+import { COPILOT_LOCKER, CopilotType } from '../resolver';
+import { ChatSessionService } from '../session';
+import { CopilotStorage } from '../storage';
+import { CopilotContextDocJob } from './job';
+import { CopilotContextService } from './service';
+import { MAX_EMBEDDABLE_SIZE } from './types';
 import { readStream } from './utils';
 
 @InputType()
-class AddRemoveContextCategoryInput {
+class AddContextCategoryInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => ContextCategories)
+  type!: ContextCategories;
+
+  @Field(() => String)
+  categoryId!: string;
+
+  @Field(() => [String], { nullable: true })
+  docs!: string[] | null;
+}
+
+@InputType()
+class RemoveContextCategoryInput {
   @Field(() => String)
   contextId!: string;
 
@@ -111,21 +127,7 @@ export class CopilotContextType {
 registerEnumType(ContextCategories, { name: 'ContextCategories' });
 
 @ObjectType()
-class CopilotContextCategory implements ContextCategory {
+class CopilotDocType implements ContextDoc {
-  @Field(() => ID)
-  id!: string;
-
-  @Field(() => ContextCategories)
-  type!: ContextCategories;
-
-  @Field(() => SafeIntResolver)
-  createdAt!: number;
-}
-
-registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' });
-
-@ObjectType()
-class CopilotContextDoc implements ContextDoc {
   @Field(() => ID)
   id!: string;
 
@@ -136,6 +138,29 @@ class CopilotContextDoc implements ContextDoc {
   createdAt!: number;
 }
 
+@ObjectType()
+class CopilotContextCategory implements Omit<ContextCategory, 'docs'> {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => ContextCategories)
+  type!: ContextCategories;
+
+  @Field(() => [CopilotDocType])
+  docs!: CopilotDocType[];
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' });
+
+@ObjectType()
+class CopilotContextDoc extends CopilotDocType {
+  @Field(() => String, { nullable: true })
+  error!: string | null;
+}
+
 @ObjectType()
 class CopilotContextFile implements ContextFile {
   @Field(() => ID)
@@ -338,6 +363,7 @@ export class CopilotContextRootResolver {
 export class CopilotContextResolver {
   constructor(
     private readonly ac: AccessController,
+    private readonly models: Models,
     private readonly mutex: RequestMutex,
     private readonly context: CopilotContextService,
     private readonly jobs: CopilotContextDocJob,
@@ -354,13 +380,61 @@ export class CopilotContextResolver {
     return controller.signal;
   }
 
+  @ResolveField(() => [CopilotContextCategory], {
+    description: 'list collections in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async collections(
+    @Parent() context: CopilotContextType
+  ): Promise<ContextCategory[]> {
+    const session = await this.context.get(context.id);
+    const collections = session.collections;
+    await this.models.copilotContext.mergeDocStatus(
+      session.workspaceId,
+      collections.flatMap(c => c.docs)
+    );
+
+    return collections;
+  }
+
+  @ResolveField(() => [CopilotContextCategory], {
+    description: 'list tags in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async tags(
+    @Parent() context: CopilotContextType
+  ): Promise<ContextCategory[]> {
+    const session = await this.context.get(context.id);
+    const tags = session.tags;
+    await this.models.copilotContext.mergeDocStatus(
+      session.workspaceId,
+      tags.flatMap(c => c.docs)
+    );
+
+    return tags;
+  }
+
   @ResolveField(() => [CopilotContextDoc], {
     description: 'list files in context',
   })
   @CallMetric('ai', 'context_file_list')
   async docs(@Parent() context: CopilotContextType): Promise<ContextDoc[]> {
     const session = await this.context.get(context.id);
-    return session.listDocs();
+    const docs = session.docs;
+    await this.models.copilotContext.mergeDocStatus(session.workspaceId, docs);
+
+    return docs;
+  }
+
+  @ResolveField(() => [CopilotContextFile], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async files(
+    @Parent() context: CopilotContextType
+  ): Promise<CopilotContextFile[]> {
+    const session = await this.context.get(context.id);
+    return session.files;
   }
 
   @Mutation(() => CopilotContextCategory, {
@@ -368,18 +442,33 @@ export class CopilotContextResolver {
   })
   @CallMetric('ai', 'context_category_add')
   async addContextCategory(
-    @Args({ name: 'options', type: () => AddRemoveContextCategoryInput })
+    @Args({ name: 'options', type: () => AddContextCategoryInput })
-    options: AddRemoveContextCategoryInput
+    options: AddContextCategoryInput
-  ) {
+  ): Promise<CopilotContextCategory> {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
     const session = await this.context.get(options.contextId);
 
     try {
-      return await session.addCategoryRecord(options.type, options.categoryId);
+      const records = await session.addCategoryRecord(
+        options.type,
+        options.categoryId,
+        options.docs || []
+      );
+
+      if (options.docs) {
+        await this.jobs.addDocEmbeddingQueue(
+          options.docs.map(docId => ({
+            workspaceId: session.workspaceId,
+            docId,
+          }))
+        );
+      }
+
+      return records;
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,
@@ -393,8 +482,8 @@ export class CopilotContextResolver {
   })
   @CallMetric('ai', 'context_category_remove')
   async removeContextCategory(
-    @Args({ name: 'options', type: () => AddRemoveContextCategoryInput })
+    @Args({ name: 'options', type: () => RemoveContextCategoryInput })
-    options: AddRemoveContextCategoryInput
+    options: RemoveContextCategoryInput
   ) {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
@@ -432,7 +521,16 @@ export class CopilotContextResolver {
     const session = await this.context.get(options.contextId);
 
     try {
-      return await session.addDocRecord(options.docId);
+      const record = await session.addDocRecord(options.docId);
+
+      await this.jobs.addDocEmbeddingQueue([
+        {
+          workspaceId: session.workspaceId,
+          docId: options.docId,
+        },
+      ]);
+
+      return record;
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,
@@ -466,17 +564,6 @@ export class CopilotContextResolver {
     }
   }
 
-  @ResolveField(() => [CopilotContextFile], {
-    description: 'list files in context',
-  })
-  @CallMetric('ai', 'context_file_list')
-  async files(
-    @Parent() context: CopilotContextType
-  ): Promise<CopilotContextFile[]> {
-    const session = await this.context.get(context.id);
-    return session.listFiles();
-  }
-
   @Mutation(() => CopilotContextFile, {
     description: 'add a file to context',
   })

packages/backend/server/src/plugins/copilot/context/service.ts

@@ -1,27 +1,23 @@
 import { Injectable, OnModuleInit } from '@nestjs/common';
-import { PrismaClient } from '@prisma/client';
 import OpenAI from 'openai';
 
 import {
   Cache,
   Config,
   CopilotInvalidContext,
-  CopilotSessionNotFound,
   NoCopilotProviderAvailable,
   OnEvent,
-  PrismaTransaction,
 } from '../../../base';
-import { OpenAIEmbeddingClient } from './embedding';
-import { ContextSession } from './session';
 import {
   ContextConfig,
   ContextConfigSchema,
   ContextEmbedStatus,
   ContextFile,
-  EmbeddingClient,
+  Models,
-  MinimalContextConfigSchema,
+} from '../../../models';
-} from './types';
+import { OpenAIEmbeddingClient } from './embedding';
-import { checkEmbeddingAvailable } from './utils';
+import { ContextSession } from './session';
+import { EmbeddingClient } from './types';
 
 const CONTEXT_SESSION_KEY = 'context-session';
 
@@ -33,7 +29,7 @@ export class CopilotContextService implements OnModuleInit {
   constructor(
     config: Config,
     private readonly cache: Cache,
-    private readonly db: PrismaClient
+    private readonly models: Models
   ) {
     const configure = config.plugins.copilot.openai;
     if (configure) {
@@ -42,7 +38,8 @@ export class CopilotContextService implements OnModuleInit {
   }
 
   async onModuleInit() {
-    const supportEmbedding = await checkEmbeddingAvailable(this.db);
+    const supportEmbedding =
+      await this.models.copilotContext.checkEmbeddingAvailable();
     if (supportEmbedding) {
       this.supportEmbedding = true;
     }
@@ -60,15 +57,10 @@ export class CopilotContextService implements OnModuleInit {
   private async saveConfig(
     contextId: string,
     config: ContextConfig,
-    tx?: PrismaTransaction,
     refreshCache = false
   ): Promise<void> {
     if (!refreshCache) {
-      const executor = tx || this.db;
+      await this.models.copilotContext.update(contextId, { config });
-      await executor.aiContext.update({
-        where: { id: contextId },
-        data: { config },
-      });
     }
     await this.cache.set(`${CONTEXT_SESSION_KEY}:${contextId}`, config);
   }
@@ -86,7 +78,7 @@ export class CopilotContextService implements OnModuleInit {
           this.embeddingClient,
           contextId,
           config.data,
-          this.db,
+          this.models,
           this.saveConfig.bind(this, contextId)
         );
       }
@@ -103,41 +95,22 @@ export class CopilotContextService implements OnModuleInit {
     config: ContextConfig
   ): Promise<ContextSession> {
     const dispatcher = this.saveConfig.bind(this, contextId);
-    await dispatcher(config, undefined, true);
+    await dispatcher(config, true);
     return new ContextSession(
       this.embeddingClient,
       contextId,
       config,
-      this.db,
+      this.models,
       dispatcher
     );
   }
 
   async create(sessionId: string): Promise<ContextSession> {
-    const session = await this.db.aiSession.findFirst({
-      where: { id: sessionId },
-      select: { workspaceId: true },
-    });
-    if (!session) {
-      throw new CopilotSessionNotFound();
-    }
-
     // keep the context unique per session
     const existsContext = await this.getBySessionId(sessionId);
     if (existsContext) return existsContext;
 
-    const context = await this.db.aiContext.create({
+    const context = await this.models.copilotContext.create(sessionId);
-      data: {
-        sessionId,
-        config: {
-          workspaceId: session.workspaceId,
-          docs: [],
-          files: [],
-          categories: [],
-        },
-      },
-    });
-
     const config = ContextConfigSchema.parse(context.config);
     return await this.cacheSession(context.id, config);
   }
@@ -149,34 +122,16 @@ export class CopilotContextService implements OnModuleInit {
 
     const context = await this.getCachedSession(id);
     if (context) return context;
-    const ret = await this.db.aiContext.findUnique({
+    const config = await this.models.copilotContext.getConfig(id);
-      where: { id },
+    if (config) {
-      select: { config: true },
+      return this.cacheSession(id, config);
-    });
-    if (ret) {
-      const config = ContextConfigSchema.safeParse(ret.config);
-      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 });
   }
 
   async getBySessionId(sessionId: string): Promise<ContextSession | null> {
-    const existsContext = await this.db.aiContext.findFirst({
+    const existsContext =
-      where: { sessionId },
+      await this.models.copilotContext.getBySessionId(sessionId);
-      select: { id: true },
-    });
     if (existsContext) return this.get(existsContext.id);
     return null;
   }

packages/backend/server/src/plugins/copilot/context/session.ts

@@ -1,30 +1,25 @@
-import { PrismaClient } from '@prisma/client';
 import { nanoid } from 'nanoid';
 
-import { PrismaTransaction } from '../../../base';
+import { CopilotDocsNotFound } from '../../../base';
 import {
-  ChunkSimilarity,
   ContextCategories,
+  ContextCategory,
   ContextConfig,
   ContextDoc,
   ContextEmbedStatus,
   ContextFile,
   ContextList,
-  DocChunkSimilarity,
+  Models,
-  EmbeddingClient,
+} from '../../../models';
-  FileChunkSimilarity,
+import { EmbeddingClient } from './types';
-} from './types';
 
 export class ContextSession implements AsyncDisposable {
   constructor(
     private readonly client: EmbeddingClient,
     private readonly contextId: string,
     private readonly config: ContextConfig,
-    private readonly db: PrismaClient,
+    private readonly models: Models,
-    private readonly dispatcher?: (
+    private readonly dispatcher?: (config: ContextConfig) => Promise<void>
-      config: ContextConfig,
-      tx?: PrismaTransaction
-    ) => Promise<void>
   ) {}
 
   get id() {
@@ -35,11 +30,28 @@ export class ContextSession implements AsyncDisposable {
     return this.config.workspaceId;
   }
 
-  listDocs(): ContextDoc[] {
+  get categories(): ContextCategory[] {
-    return [...this.config.docs];
+    return this.config.categories.map(c => ({
+      ...c,
+      docs: c.docs.map(d => ({ ...d })),
+    }));
   }
 
-  listFiles() {
+  get tags() {
+    const categories = this.config.categories;
+    return categories.filter(c => c.type === ContextCategories.Tag);
+  }
+
+  get collections() {
+    const categories = this.config.categories;
+    return categories.filter(c => c.type === ContextCategories.Collection);
+  }
+
+  get docs(): ContextDoc[] {
+    return this.config.docs.map(d => ({ ...d }));
+  }
+
+  get files() {
     return this.config.files.map(f => ({ ...f }));
   }
 
@@ -50,14 +62,25 @@ export class ContextSession implements AsyncDisposable {
     ) as ContextList;
   }
 
-  async addCategoryRecord(type: ContextCategories, id: string) {
+  async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
+    const existDocs = await this.models.doc.existsAll(this.workspaceId, docs);
+    if (!existDocs) {
+      throw new CopilotDocsNotFound();
+    }
+
     const category = this.config.categories.find(
       c => c.type === type && c.id === id
     );
     if (category) {
       return category;
     }
-    const record = { id, type, createdAt: Date.now() };
+    const createdAt = Date.now();
+    const record = {
+      id,
+      type,
+      docs: docs.map(id => ({ id, createdAt, status: null })),
+      createdAt,
+    };
     this.config.categories.push(record);
     await this.save();
     return record;
@@ -122,14 +145,10 @@ export class ContextSession implements AsyncDisposable {
   }
 
   async removeFile(fileId: string): Promise<boolean> {
-    return await this.db.$transaction(async tx => {
+    await this.models.copilotContext.deleteEmbedding(this.contextId, fileId);
-      await tx.aiContextEmbedding.deleteMany({
+    this.config.files = this.config.files.filter(f => f.id !== fileId);
-        where: { contextId: this.contextId, fileId },
+    await this.save();
-      });
+    return true;
-      this.config.files = this.config.files.filter(f => f.id !== fileId);
-      await this.save(tx);
-      return true;
-    });
   }
 
   /**
@@ -145,21 +164,18 @@ export class ContextSession implements AsyncDisposable {
     topK: number = 5,
     signal?: AbortSignal,
     threshold: number = 0.7
-  ): Promise<FileChunkSimilarity[]> {
+  ) {
     const embedding = await this.client
       .getEmbeddings([content], signal)
       .then(r => r?.[0]?.embedding);
     if (!embedding) return [];
-    const similarityChunks = await this.db.$queryRaw<
+
-      Array<FileChunkSimilarity>
+    return this.models.copilotContext.matchContentEmbedding(
-    >`
+      embedding,
-      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      this.id,
-      FROM "ai_context_embeddings"
+      topK,
-      WHERE context_id = ${this.id}
+      threshold
-      ORDER BY "distance" ASC
+    );
-      LIMIT ${topK};
-    `;
-    return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
   /**
@@ -175,19 +191,18 @@ export class ContextSession implements AsyncDisposable {
     topK: number = 5,
     signal?: AbortSignal,
     threshold: number = 0.7
-  ): Promise<ChunkSimilarity[]> {
+  ) {
     const embedding = await this.client
       .getEmbeddings([content], signal)
       .then(r => r?.[0]?.embedding);
     if (!embedding) return [];
-    const similarityChunks = await this.db.$queryRaw<Array<DocChunkSimilarity>>`
+
-      SELECT "doc_id" as "docId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
+    return this.models.copilotContext.matchWorkspaceEmbedding(
-      FROM "ai_workspace_embeddings"
+      embedding,
-      WHERE "workspace_id" = ${this.workspaceId}
+      this.id,
-      ORDER BY "distance" ASC
+      topK,
-      LIMIT ${topK};
+      threshold
-    `;
+    );
-    return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
   async saveFileRecord(
@@ -195,8 +210,7 @@ export class ContextSession implements AsyncDisposable {
     cb: (
       record: Pick<ContextFile, 'id' | 'status'> &
         Partial<Omit<ContextFile, 'id' | 'status'>>
-    ) => ContextFile,
+    ) => ContextFile
-    tx?: PrismaTransaction
   ) {
     const files = this.config.files;
     const file = files.find(f => f.id === fileId);
@@ -206,11 +220,11 @@ export class ContextSession implements AsyncDisposable {
       const file = { id: fileId, status: ContextEmbedStatus.processing };
       files.push(cb(file));
     }
-    await this.save(tx);
+    await this.save();
   }
 
-  async save(tx?: PrismaTransaction) {
+  async save() {
-    await this.dispatcher?.(this.config, tx);
+    await this.dispatcher?.(this.config);
   }
 
   async [Symbol.asyncDispose]() {

packages/backend/server/src/plugins/copilot/context/types.ts

@@ -1,8 +1,7 @@
 import { File } from 'node:buffer';
 
-import { z } from 'zod';
-
 import { CopilotContextFileNotSupported, OneMB } from '../../../base';
+import { Embedding } from '../../../models';
 import { parseDoc } from '../../../native';
 
 declare global {
@@ -26,99 +25,11 @@ declare global {
 
 export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
 
-export enum ContextEmbedStatus {
-  processing = 'processing',
-  finished = 'finished',
-  failed = 'failed',
-}
-
-export enum ContextCategories {
-  Tag = 'tag',
-  Collection = 'collection',
-}
-
-export const ContextConfigSchema = z.object({
-  workspaceId: z.string(),
-  files: z
-    .object({
-      id: z.string(),
-      chunkSize: z.number(),
-      name: z.string(),
-      status: z.enum([
-        ContextEmbedStatus.processing,
-        ContextEmbedStatus.finished,
-        ContextEmbedStatus.failed,
-      ]),
-      error: z.string().nullable(),
-      blobId: z.string(),
-      createdAt: z.number(),
-    })
-    .array(),
-  docs: z
-    .object({
-      id: z.string(),
-      // status for workspace doc embedding progress
-      // only exists when the client submits the doc embedding task
-      status: z
-        .enum([
-          ContextEmbedStatus.processing,
-          ContextEmbedStatus.finished,
-          ContextEmbedStatus.failed,
-        ])
-        .nullable(),
-      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;
-export type ContextList = ContextListItem[];
-
 export type Chunk = {
   index: number;
   content: string;
 };
 
-export type ChunkSimilarity = {
-  chunk: number;
-  content: string;
-  distance: number | null;
-};
-
-export type FileChunkSimilarity = ChunkSimilarity & {
-  fileId: string;
-};
-
-export type DocChunkSimilarity = ChunkSimilarity & {
-  docId: string;
-};
-
-export type Embedding = {
-  /**
-   * The index of the embedding in the list of embeddings.
-   */
-  index: number;
-  content: string;
-  embedding: Array<number>;
-};
-
 export abstract class EmbeddingClient {
   async getFileEmbeddings(
     file: File,

packages/backend/server/src/plugins/copilot/context/utils.ts

@@ -1,7 +1,5 @@
 import { Readable } from 'node:stream';
 
-import { PrismaClient } from '@prisma/client';
-
 import { readBufferWithLimit } from '../../../base';
 import { MAX_EMBEDDABLE_SIZE } from './types';
 
@@ -17,17 +15,6 @@ export class GqlSignal implements AsyncDisposable {
   }
 }
 
-export async function checkEmbeddingAvailable(
-  db: PrismaClient
-): Promise<boolean> {
-  const [{ count }] = await db.$queryRaw<
-    {
-      count: number;
-    }[]
-  >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_context_embeddings', 'ai_workspace_embeddings')`;
-  return Number(count) === 2;
-}
-
 export function readStream(
   readable: Readable,
   maxSize = MAX_EMBEDDABLE_SIZE

packages/backend/server/src/schema.gql

@@ -2,6 +2,13 @@
 # THIS FILE WAS AUTOMATICALLY GENERATED (DO NOT MODIFY)
 # ------------------------------------------------------
 
+input AddContextCategoryInput {
+  categoryId: String!
+  contextId: String!
+  docs: [String!]
+  type: ContextCategories!
+}
+
 input AddContextDocInput {
   contextId: String!
   docId: String!
@@ -12,12 +19,6 @@ input AddContextFileInput {
   contextId: String!
 }
 
-input AddRemoveContextCategoryInput {
-  categoryId: String!
-  contextId: String!
-  type: ContextCategories!
-}
-
 enum AiJobStatus {
   claimed
   failed
@@ -98,6 +99,9 @@ type Copilot {
 }
 
 type CopilotContext {
+  """list collections in context"""
+  collections: [CopilotContextCategory!]!
+
   """list files in context"""
   docs: [CopilotContextDoc!]!
 
@@ -110,17 +114,22 @@ type CopilotContext {
 
   """match workspace doc content"""
   matchWorkspaceContext(content: String!, limit: SafeInt): ContextMatchedDocChunk!
+
+  """list tags in context"""
+  tags: [CopilotContextCategory!]!
   workspaceId: String!
 }
 
 type CopilotContextCategory {
   createdAt: SafeInt!
+  docs: [CopilotDocType!]!
   id: ID!
   type: ContextCategories!
 }
 
 type CopilotContextDoc {
   createdAt: SafeInt!
+  error: String
   id: ID!
   status: ContextEmbedStatus
 }
@@ -144,6 +153,12 @@ type CopilotDocNotFoundDataType {
   docId: String!
 }
 
+type CopilotDocType {
+  createdAt: SafeInt!
+  id: ID!
+  status: ContextEmbedStatus
+}
+
 type CopilotFailedToMatchContextDataType {
   content: String!
   contextId: String!
@@ -405,6 +420,7 @@ enum ErrorNames {
   CAPTCHA_VERIFICATION_FAILED
   COPILOT_ACTION_TAKEN
   COPILOT_CONTEXT_FILE_NOT_SUPPORTED
+  COPILOT_DOCS_NOT_FOUND
   COPILOT_DOC_NOT_FOUND
   COPILOT_EMBEDDING_UNAVAILABLE
   COPILOT_FAILED_TO_CREATE_MESSAGE
@@ -838,7 +854,7 @@ type Mutation {
   activateLicense(license: String!, workspaceId: String!): License!
 
   """add a category to context"""
-  addContextCategory(options: AddRemoveContextCategoryInput!): CopilotContextCategory!
+  addContextCategory(options: AddContextCategoryInput!): CopilotContextCategory!
 
   """add a doc to context"""
   addContextDoc(options: AddContextDocInput!): CopilotContextDoc!
@@ -926,7 +942,7 @@ type Mutation {
   removeAvatar: RemoveAvatar!
 
   """remove a category from context"""
-  removeContextCategory(options: AddRemoveContextCategoryInput!): Boolean!
+  removeContextCategory(options: RemoveContextCategoryInput!): Boolean!
 
   """remove a doc from context"""
   removeContextDoc(options: RemoveContextDocInput!): Boolean!
@@ -1192,6 +1208,12 @@ type RemoveAvatar {
   success: Boolean!
 }
 
+input RemoveContextCategoryInput {
+  categoryId: String!
+  contextId: String!
+  type: ContextCategories!
+}
+
 input RemoveContextDocInput {
   contextId: String!
   docId: String!

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

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

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

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

packages/common/graphql/src/graphql/copilot-context-doc-add.gql

@@ -3,5 +3,6 @@ mutation addContextDoc($options: AddContextDocInput!) {
     id
     createdAt
     status
+    error
   }
 }

packages/common/graphql/src/graphql/copilot-context-list-object.gql

@@ -9,6 +9,7 @@ query listContextObject(
         docs {
           id
           status
+          error
           createdAt
         }
         files {
@@ -20,6 +21,22 @@ query listContextObject(
           status
           createdAt
         }
+        tags {
+          id
+          docs {
+            id
+            status
+          }
+          createdAt
+        }
+        collections {
+          id
+          docs {
+            id
+            status
+          }
+          createdAt
+        }
       }
     }
   }

packages/common/graphql/src/graphql/index.ts

@@ -137,11 +137,16 @@ export const changePasswordMutation = {
 export const addContextCategoryMutation = {
   id: 'addContextCategoryMutation' as const,
   op: 'addContextCategory',
-  query: `mutation addContextCategory($options: AddRemoveContextCategoryInput!) {
+  query: `mutation addContextCategory($options: AddContextCategoryInput!) {
   addContextCategory(options: $options) {
     id
     createdAt
     type
+    docs {
+      id
+      createdAt
+      status
+    }
   }
 }`,
 };
@@ -149,7 +154,7 @@ export const addContextCategoryMutation = {
 export const removeContextCategoryMutation = {
   id: 'removeContextCategoryMutation' as const,
   op: 'removeContextCategory',
-  query: `mutation removeContextCategory($options: AddRemoveContextCategoryInput!) {
+  query: `mutation removeContextCategory($options: RemoveContextCategoryInput!) {
   removeContextCategory(options: $options)
 }`,
 };
@@ -170,6 +175,7 @@ export const addContextDocMutation = {
     id
     createdAt
     status
+    error
   }
 }`,
 };
@@ -236,6 +242,7 @@ export const listContextObjectQuery = {
         docs {
           id
           status
+          error
           createdAt
         }
         files {
@@ -247,6 +254,22 @@ export const listContextObjectQuery = {
           status
           createdAt
         }
+        tags {
+          id
+          docs {
+            id
+            status
+          }
+          createdAt
+        }
+        collections {
+          id
+          docs {
+            id
+            status
+          }
+          createdAt
+        }
       }
     }
   }

packages/common/graphql/src/schema.ts

@@ -37,6 +37,13 @@ export interface Scalars {
   Upload: { input: File; output: File };
 }
 
+export interface AddContextCategoryInput {
+  categoryId: Scalars['String']['input'];
+  contextId: Scalars['String']['input'];
+  docs?: InputMaybe<Array<Scalars['String']['input']>>;
+  type: ContextCategories;
+}
+
 export interface AddContextDocInput {
   contextId: Scalars['String']['input'];
   docId: Scalars['String']['input'];
@@ -47,12 +54,6 @@ export interface AddContextFileInput {
   contextId: Scalars['String']['input'];
 }
 
-export interface AddRemoveContextCategoryInput {
-  categoryId: Scalars['String']['input'];
-  contextId: Scalars['String']['input'];
-  type: ContextCategories;
-}
-
 export enum AiJobStatus {
   claimed = 'claimed',
   failed = 'failed',
@@ -164,6 +165,8 @@ export interface CopilotSessionsArgs {
 
 export interface CopilotContext {
   __typename?: 'CopilotContext';
+  /** list collections in context */
+  collections: Array<CopilotContextCategory>;
   /** list files in context */
   docs: Array<CopilotContextDoc>;
   /** list files in context */
@@ -173,6 +176,8 @@ export interface CopilotContext {
   matchContext: Array<ContextMatchedFileChunk>;
   /** match workspace doc content */
   matchWorkspaceContext: ContextMatchedDocChunk;
+  /** list tags in context */
+  tags: Array<CopilotContextCategory>;
   workspaceId: Scalars['String']['output'];
 }
 
@@ -190,6 +195,7 @@ export interface CopilotContextMatchWorkspaceContextArgs {
 export interface CopilotContextCategory {
   __typename?: 'CopilotContextCategory';
   createdAt: Scalars['SafeInt']['output'];
+  docs: Array<CopilotDocType>;
   id: Scalars['ID']['output'];
   type: ContextCategories;
 }
@@ -197,6 +203,7 @@ export interface CopilotContextCategory {
 export interface CopilotContextDoc {
   __typename?: 'CopilotContextDoc';
   createdAt: Scalars['SafeInt']['output'];
+  error: Maybe<Scalars['String']['output']>;
   id: Scalars['ID']['output'];
   status: Maybe<ContextEmbedStatus>;
 }
@@ -223,6 +230,13 @@ export interface CopilotDocNotFoundDataType {
   docId: Scalars['String']['output'];
 }
 
+export interface CopilotDocType {
+  __typename?: 'CopilotDocType';
+  createdAt: Scalars['SafeInt']['output'];
+  id: Scalars['ID']['output'];
+  status: Maybe<ContextEmbedStatus>;
+}
+
 export interface CopilotFailedToMatchContextDataType {
   __typename?: 'CopilotFailedToMatchContextDataType';
   content: Scalars['String']['output'];
@@ -549,6 +563,7 @@ export enum ErrorNames {
   CAPTCHA_VERIFICATION_FAILED = 'CAPTCHA_VERIFICATION_FAILED',
   COPILOT_ACTION_TAKEN = 'COPILOT_ACTION_TAKEN',
   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_UNAVAILABLE = 'COPILOT_EMBEDDING_UNAVAILABLE',
   COPILOT_FAILED_TO_CREATE_MESSAGE = 'COPILOT_FAILED_TO_CREATE_MESSAGE',
@@ -1103,7 +1118,7 @@ export interface MutationActivateLicenseArgs {
 }
 
 export interface MutationAddContextCategoryArgs {
-  options: AddRemoveContextCategoryInput;
+  options: AddContextCategoryInput;
 }
 
 export interface MutationAddContextDocArgs {
@@ -1297,7 +1312,7 @@ export interface MutationReleaseDeletedBlobsArgs {
 }
 
 export interface MutationRemoveContextCategoryArgs {
-  options: AddRemoveContextCategoryInput;
+  options: RemoveContextCategoryInput;
 }
 
 export interface MutationRemoveContextDocArgs {
@@ -1695,6 +1710,12 @@ export interface RemoveAvatar {
   success: Scalars['Boolean']['output'];
 }
 
+export interface RemoveContextCategoryInput {
+  categoryId: Scalars['String']['input'];
+  contextId: Scalars['String']['input'];
+  type: ContextCategories;
+}
+
 export interface RemoveContextDocInput {
   contextId: Scalars['String']['input'];
   docId: Scalars['String']['input'];
@@ -2423,7 +2444,7 @@ export type ChangePasswordMutation = {
 };
 
 export type AddContextCategoryMutationVariables = Exact<{
-  options: AddRemoveContextCategoryInput;
+  options: AddContextCategoryInput;
 }>;
 
 export type AddContextCategoryMutation = {
@@ -2433,11 +2454,17 @@ export type AddContextCategoryMutation = {
     id: string;
     createdAt: number;
     type: ContextCategories;
+    docs: Array<{
+      __typename?: 'CopilotDocType';
+      id: string;
+      createdAt: number;
+      status: ContextEmbedStatus | null;
+    }>;
   };
 };
 
 export type RemoveContextCategoryMutationVariables = Exact<{
-  options: AddRemoveContextCategoryInput;
+  options: RemoveContextCategoryInput;
 }>;
 
 export type RemoveContextCategoryMutation = {
@@ -2466,6 +2493,7 @@ export type AddContextDocMutation = {
     id: string;
     createdAt: number;
     status: ContextEmbedStatus | null;
+    error: string | null;
   };
 };
 
@@ -2550,6 +2578,7 @@ export type ListContextObjectQuery = {
           __typename?: 'CopilotContextDoc';
           id: string;
           status: ContextEmbedStatus | null;
+          error: string | null;
           createdAt: number;
         }>;
         files: Array<{
@@ -2562,6 +2591,26 @@ export type ListContextObjectQuery = {
           status: ContextEmbedStatus;
           createdAt: number;
         }>;
+        tags: Array<{
+          __typename?: 'CopilotContextCategory';
+          id: string;
+          createdAt: number;
+          docs: Array<{
+            __typename?: 'CopilotDocType';
+            id: string;
+            status: ContextEmbedStatus | null;
+          }>;
+        }>;
+        collections: Array<{
+          __typename?: 'CopilotContextCategory';
+          id: string;
+          createdAt: number;
+          docs: Array<{
+            __typename?: 'CopilotDocType';
+            id: string;
+            status: ContextEmbedStatus | null;
+          }>;
+        }>;
       }>;
     };
   } | null;

packages/frontend/i18n/src/i18n.gen.ts

@@ -7816,6 +7816,10 @@ export function useAFFiNEI18N(): {
     ["error.COPILOT_DOC_NOT_FOUND"](options: {
         readonly docId: string;
     }): string;
+    /**
+      * `Some docs not found.`
+      */
+    ["error.COPILOT_DOCS_NOT_FOUND"](): string;
     /**
       * `Copilot message {{messageId}} not found.`
       */

packages/frontend/i18n/src/resources/en.json

@@ -1928,6 +1928,7 @@
   "error.UNSPLASH_IS_NOT_CONFIGURED": "Unsplash is not configured.",
   "error.COPILOT_ACTION_TAKEN": "Action has been taken, no more messages allowed.",
   "error.COPILOT_DOC_NOT_FOUND": "Doc {{docId}} not found.",
+  "error.COPILOT_DOCS_NOT_FOUND": "Some docs not found.",
   "error.COPILOT_MESSAGE_NOT_FOUND": "Copilot message {{messageId}} not found.",
   "error.COPILOT_PROMPT_NOT_FOUND": "Copilot prompt {{name}} not found.",
   "error.COPILOT_PROMPT_INVALID": "Copilot prompt is invalid.",