feat(server): cleanup trashed doc's embedding (#13201)

fix AI-359

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

* **New Features**
* Added automated cleanup of embeddings for documents deleted or trashed
from workspaces.
* Introduced a new job to schedule and perform this cleanup per
workspace daily and on demand.
  * Added new GraphQL mutation to manually trigger the cleanup process.
* Added the ability to list workspaces with flexible filtering and
selection options.

* **Improvements**
* Enhanced document status handling to more accurately reflect embedding
presence.
* Refined internal methods for managing and checking document
embeddings.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2025-07-15 20:00:33 +08:00
committed by GitHub
parent a4b535a42a
commit cc1d5b497a
12 changed files with 148 additions and 22 deletions
@@ -12,6 +12,7 @@ import {
OnJob,
} from '../../../base';
import { DocReader } from '../../../core/doc';
import { readAllDocIdsFromWorkspaceSnapshot } from '../../../core/utils/blocksuite';
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
@@ -134,10 +135,30 @@ export class CopilotEmbeddingJob {
if (enableDocEmbedding) {
const toBeEmbedDocIds =
await this.models.copilotWorkspace.findDocsToEmbed(workspaceId);
if (!toBeEmbedDocIds.length) {
return;
}
// filter out trashed docs
const rootSnapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!rootSnapshot) {
this.logger.warn(
`Root snapshot for workspace ${workspaceId} not found, skipping embedding.`
);
return;
}
const allDocIds = new Set(
readAllDocIdsFromWorkspaceSnapshot(rootSnapshot.blob)
);
this.logger.log(
`Trigger embedding for ${toBeEmbedDocIds.length} docs in workspace ${workspaceId}`
);
for (const docId of toBeEmbedDocIds) {
const finalToBeEmbedDocIds = toBeEmbedDocIds.filter(docId =>
allDocIds.has(docId)
);
for (const docId of finalToBeEmbedDocIds) {
await this.queue.add(
'copilot.embedding.docs',
{
@@ -422,4 +443,39 @@ export class CopilotEmbeddingJob {
);
}
}
@OnJob('copilot.embedding.cleanupTrashedDocEmbeddings')
async cleanupTrashedDocEmbeddings({
workspaceId,
}: Jobs['copilot.embedding.cleanupTrashedDocEmbeddings']) {
const workspace = await this.models.workspace.get(workspaceId);
if (!workspace) {
this.logger.warn(`workspace ${workspaceId} not found`);
return;
}
const snapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!snapshot) {
this.logger.warn(`workspace snapshot ${workspaceId} not found`);
return;
}
const docIdsInWorkspace = readAllDocIdsFromWorkspaceSnapshot(snapshot.blob);
const docIdsInEmbedding =
await this.models.copilotContext.listWorkspaceEmbedding(workspaceId);
const docIdsInWorkspaceSet = new Set(docIdsInWorkspace);
const deletedDocIds = docIdsInEmbedding.filter(
docId => !docIdsInWorkspaceSet.has(docId)
);
for (const docId of deletedDocIds) {
await this.models.copilotContext.deleteWorkspaceEmbedding(
workspaceId,
docId
);
}
}
}
@@ -61,6 +61,10 @@ declare global {
fileId: string;
fileName: string;
};
'copilot.embedding.cleanupTrashedDocEmbeddings': {
workspaceId: string;
};
}
}