diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md new file mode 100644 index 0000000000..b4d85cc168 --- /dev/null +++ b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md @@ -0,0 +1,11 @@ +# Snapshot report for `src/__tests__/models/copilot-context.spec.ts` + +The actual snapshot is saved in `copilot-context.spec.ts.snap`. + +Generated by [AVA](https://avajs.dev). + +## should insert embedding by doc id + +> should return empty array when embedding deleted + + [] diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap new file mode 100644 index 0000000000..d76cf43ae7 Binary files /dev/null and b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts index 7a57708556..d6a014f892 100644 --- a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts @@ -116,7 +116,7 @@ test('should insert embedding by doc id', async t => { } { - await t.context.copilotContext.deleteEmbedding(contextId, 'file-id'); + await t.context.copilotContext.deleteFileEmbedding(contextId, 'file-id'); const ret = await t.context.copilotContext.matchFileEmbedding( Array.from({ length: 1024 }, () => 0.9), contextId, @@ -169,6 +169,20 @@ test('should insert embedding by doc id', async t => { t.is(ret.length, 1); t.is(ret[0].content, 'content'); } + + { + await t.context.copilotContext.deleteWorkspaceEmbedding( + workspace.id, + docId + ); + const ret = await t.context.copilotContext.matchWorkspaceEmbedding( + Array.from({ length: 1024 }, () => 0.9), + workspace.id, + 1, + 1 + ); + t.snapshot(ret, 'should return empty array when embedding deleted'); + } } }); diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts index 2a392815a4..ab34530666 100644 --- a/packages/backend/server/src/models/copilot-context.ts +++ b/packages/backend/server/src/models/copilot-context.ts @@ -168,6 +168,12 @@ export class CopilotContextModel extends BaseModel { `; } + async deleteFileEmbedding(contextId: string, fileId: string) { + await this.db.aiContextEmbedding.deleteMany({ + where: { contextId, fileId }, + }); + } + async matchFileEmbedding( embedding: number[], contextId: string, @@ -205,6 +211,12 @@ export class CopilotContextModel extends BaseModel { `; } + async deleteWorkspaceEmbedding(workspaceId: string, docId: string) { + await this.db.aiWorkspaceEmbedding.deleteMany({ + where: { workspaceId, docId }, + }); + } + async matchWorkspaceEmbedding( embedding: number[], workspaceId: string, @@ -222,10 +234,4 @@ export class CopilotContextModel extends BaseModel { `; return similarityChunks.filter(c => Number(c.distance) <= threshold); } - - async deleteEmbedding(contextId: string, fileId: string) { - await this.db.aiContextEmbedding.deleteMany({ - where: { contextId, fileId }, - }); - } } diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/context/job.ts index 70c7412765..c272d2c261 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/context/job.ts @@ -87,6 +87,30 @@ export class CopilotContextDocJob { } } + // @OnEvent('doc.indexer.updated') + async addDocEmbeddingQueueFromEvent( + // TODO(@darkskygit): replace this with real event type + doc: { workspaceId: string; docId: string } //Events['doc.indexer.updated'], + ) { + if (!this.supportEmbedding) return; + + await this.queue.add('copilot.embedding.docs', { + workspaceId: doc.workspaceId, + docId: doc.workspaceId, + }); + } + + // @OnEvent('doc.indexer.deleted') + async deleteDocEmbeddingQueueFromEvent( + // TODO(@darkskygit): replace this with real event type + doc: { workspaceId: string; docId: string } //Events['doc.indexer.deleted'], + ) { + await this.models.copilotContext.deleteWorkspaceEmbedding( + doc.workspaceId, + doc.docId + ); + } + async readCopilotBlob( userId: string, workspaceId: string, diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index c4ea10bc4d..05dd9ff6af 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -164,7 +164,10 @@ export class ContextSession implements AsyncDisposable { } async removeFile(fileId: string): Promise { - await this.models.copilotContext.deleteEmbedding(this.contextId, fileId); + await this.models.copilotContext.deleteFileEmbedding( + this.contextId, + fileId + ); this.config.files = this.config.files.filter(f => f.id !== fileId); await this.save(); return true;