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

expose more field in listContextObject

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({
-      userId,
-      workspaceId: session.workspaceId,
-      contextId: session.id,
-      blobId: file.blobId,
-      fileId: file.id,
-      fileName: file.name,
+      t.deepEqual(handler.lastCall.args, [
+        'workspace.file.embed.finished',
+        {
+          contextId: session.id,
+          fileId: file.id,
+          chunkSize: 1,
+        },
+      ]);
+
+      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',
-      {
-        contextId: session.id,
-        fileId: file.id,
-        chunkSize: 1,
-      },
-    ]);
+    {
+      await session.addDocRecord(docId);
+      const docs = session.docs.map(d => d.id);
+      t.deepEqual(docs, [docId], 'should list doc id');
 
-    const list = session.listFiles();
-    t.deepEqual(
-      list.map(f => f.id),
-      [file.id],
-      'should list file id'
-    );
+      await session.removeDocRecord(docId);
+      t.deepEqual(session.docs, [], 'should remove doc id');
+    }
 
-    const docId = randomUUID();
-    await session.addDocRecord(docId);
-    const docs = session.listDocs().map(d => d.id);
-    t.deepEqual(docs, [docId], 'should list doc id');
+    // tag record
+    {
+      const tagId = randomUUID();
+      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);
-    t.deepEqual(session.listDocs(), [], 'should remove doc id');
+      await session.removeCategoryRecord(ContextCategories.Tag, tagId);
+      t.deepEqual(session.tags, [], 'should remove tag 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');
+      await t.throwsAsync(
+        session.addCategoryRecord(ContextCategories.Tag, tagId, [
+          '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', [
-      {
-        index: 0,
-        content: 'content',
-        embedding: Array.from({ length: 512 }, () => 1),
-      },
-    ]);
+    await t.context.copilotContext.insertContentEmbedding(
+      contextId,
+      'file-id',
+      [
+        {
+          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,