feat(server): trigger workspace embedding (#12328)

fix AI-127

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

- **New Features**
  - Added automated event handling for workspace updates and document embedding, streamlining document embedding workflows.
  - Introduced detection and queuing of documents needing embedding, excluding ignored documents.
- **Improvements**
  - Enhanced performance of embedding-related searches by filtering results at the database level.
  - Increased concurrency for embedding job processing to improve throughput.
- **Bug Fixes**
  - Improved error handling and fallback for missing document titles during embedding.
  - Added safeguards to skip invalid embedding jobs based on document identifiers.
- **Tests**
  - Expanded test coverage for document embedding and ignored document filtering.
  - Updated end-to-end tests to use dynamic content for improved reliability.
  - Added synchronization waits in document creation utilities to improve test stability.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
darkskygit
2025-05-20 05:16:45 +00:00
parent 3c982d2b91
commit 6f9361caee
13 changed files with 213 additions and 35 deletions
@@ -73,3 +73,19 @@ Generated by [AVA](https://avajs.dev).
name: 'file1',
},
]
> should find docs to embed
1
> should not find docs to embed
0
> should find docs to embed
1
> should not find docs to embed
0
@@ -1,8 +1,12 @@
import { randomUUID } from 'node:crypto';
import { PrismaClient, User, Workspace } from '@prisma/client';
import ava, { TestFn } from 'ava';
import { Config } from '../../base';
import { CopilotContextModel } from '../../models/copilot-context';
import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace';
import { DocModel } from '../../models/doc';
import { UserModel } from '../../models/user';
import { WorkspaceModel } from '../../models/workspace';
import { createTestingModule, type TestingModule } from '../utils';
@@ -12,8 +16,10 @@ interface Context {
config: Config;
module: TestingModule;
db: PrismaClient;
doc: DocModel;
user: UserModel;
workspace: WorkspaceModel;
copilotContext: CopilotContextModel;
copilotWorkspace: CopilotWorkspaceConfigModel;
}
@@ -23,8 +29,10 @@ test.before(async t => {
const module = await createTestingModule();
t.context.user = module.get(UserModel);
t.context.workspace = module.get(WorkspaceModel);
t.context.copilotContext = module.get(CopilotContextModel);
t.context.copilotWorkspace = module.get(CopilotWorkspaceConfigModel);
t.context.db = module.get(PrismaClient);
t.context.doc = module.get(DocModel);
t.context.config = module.get(Config);
t.context.module = module;
});
@@ -136,6 +144,61 @@ test('should insert and search embedding', async t => {
);
}
}
{
const docId = randomUUID();
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
});
const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed');
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
const afterInsertEmbedding =
await t.context.copilotWorkspace.findDocsToEmbed(workspace.id);
t.snapshot(afterInsertEmbedding.length, 'should not find docs to embed');
}
{
const docId = randomUUID();
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
});
const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed');
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, [docId]);
const afterAddIgnoreDocs = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(afterAddIgnoreDocs.length, 'should not find docs to embed');
}
});
test('should check embedding table', async t => {