From 4bf9161e57ea57fec16cdc0a20407b67192c2a9a Mon Sep 17 00:00:00 2001 From: darkskygit Date: Tue, 25 Mar 2025 15:10:56 +0000 Subject: [PATCH] chore(server): ignore non-exists doc embedding (#11153) --- .github/actions/server-test-env/action.yml | 2 +- .../server/src/__tests__/copilot.spec.ts | 102 +++++++++++------- .../src/plugins/copilot/context/session.ts | 26 +++-- 3 files changed, 81 insertions(+), 49 deletions(-) diff --git a/.github/actions/server-test-env/action.yml b/.github/actions/server-test-env/action.yml index 5994d68900..260e5209d6 100644 --- a/.github/actions/server-test-env/action.yml +++ b/.github/actions/server-test-env/action.yml @@ -19,5 +19,5 @@ runs: NODE_ENV: test run: | yarn affine @affine/server prisma generate - yarn affine @affine/server prisma migrate reset -f + yarn affine @affine/server prisma migrate deploy yarn affine @affine/server data-migration run diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts index 13ccf77f76..f559ca8239 100644 --- a/packages/backend/server/src/__tests__/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot.spec.ts @@ -1381,19 +1381,24 @@ test('should be able to manage context', async t => { } // 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(), - }, - }); + + const addDoc = async () => { + 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(), + }, + }); + return docId; + }; { + const docId = await addDoc(); await session.addDocRecord(docId); const docs = session.docs.map(d => d.id); t.deepEqual(docs, [docId], 'should list doc id'); @@ -1405,50 +1410,65 @@ test('should be able to manage context', async t => { // 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'); + + const docId1 = await addDoc(); + const docId2 = await addDoc(); + + { + await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId1]); + const tags = session.tags.map(t => t.id); + t.deepEqual(tags, [tagId], 'should list tag id'); + + const docs = session.tags.flatMap(l => l.docs.map(d => d.id)); + t.deepEqual(docs, [docId1], 'should list doc ids'); + } + + { + await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId2]); + + const docs = session.tags.flatMap(l => l.docs.map(d => d.id)); + t.deepEqual(docs, [docId1, docId2], 'should list doc ids'); + } await session.removeCategoryRecord(ContextCategories.Tag, tagId); t.deepEqual(session.tags, [], 'should remove tag 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'); + + const docId1 = await addDoc(); + const docId2 = await addDoc(); + { + await session.addCategoryRecord( + ContextCategories.Collection, + collectionId, + [docId1] + ); + const collection = session.collections.map(l => l.id); + t.deepEqual(collection, [collectionId], 'should list collection id'); + + const docs = session.collections.flatMap(l => l.docs.map(d => d.id)); + t.deepEqual(docs, [docId1], 'should list doc ids'); + } + + { + await session.addCategoryRecord( + ContextCategories.Collection, + collectionId, + [docId2] + ); + + const docs = session.collections.flatMap(l => l.docs.map(d => d.id)); + t.deepEqual(docs, [docId1, docId2], 'should list doc ids'); + } 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' - ); } } }); diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index 99b8e97dea..1d1be07517 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -1,6 +1,5 @@ import { nanoid } from 'nanoid'; -import { CopilotDocsNotFound } from '../../../base'; import { ContextCategories, ContextCategory, @@ -63,22 +62,35 @@ export class ContextSession implements AsyncDisposable { } 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) { + const missingDocs = docs.filter( + docId => !category.docs.some(d => d.id === docId) + ); + if (missingDocs.length) { + category.docs.push( + ...missingDocs.map(id => ({ + id, + createdAt: Date.now(), + status: ContextEmbedStatus.processing, + })) + ); + await this.save(); + } + return category; } const createdAt = Date.now(); const record = { id, type, - docs: docs.map(id => ({ id, createdAt, status: null })), + docs: docs.map(id => ({ + id, + createdAt, + status: ContextEmbedStatus.processing, + })), createdAt, }; this.config.categories.push(record);