Jayden/AFFiNE-Mirror
1
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-02-12 20:38:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(server): improve embedding & rerank speed (#12666)

Browse Source 
fix AI-109
This commit is contained in:
darkskygit
2025-06-03 11:12:34 +00:00
parent 2288cbe54d
commit 44e1eb503f
12 changed files with 232 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
packages/backend/server/src/plugins/copilot/context/job.ts
View File

@@ -4,7 +4,6 @@ import {
AFFiNELogger,
BlobNotFound,
CallMetric,
Config,
CopilotContextFileNotSupported,
DocNotFound,
EventBus,
@@ -15,9 +14,11 @@ import {
} from '../../../base';
import { DocReader } from '../../../core/doc';
import { Models } from '../../../models';
import { PromptService } from '../prompt';
import { CopilotProviderFactory } from '../providers';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
import { OpenAIEmbeddingClient } from './embedding';
import { getEmbeddingClient } from './embedding';
import type { Chunk, DocFragment } from './types';
import { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types';
@@ -30,11 +31,12 @@ export class CopilotContextDocJob {
private client: EmbeddingClient | undefined;
constructor(
private readonly config: Config,
private readonly doc: DocReader,
private readonly event: EventBus,
private readonly logger: AFFiNELogger,
private readonly models: Models,
private readonly providerFactory: CopilotProviderFactory,
private readonly prompt: PromptService,
private readonly queue: JobQueue,
private readonly storage: CopilotStorage
) {
@@ -54,10 +56,8 @@ export class CopilotContextDocJob {
private async setup() {
this.supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (this.supportEmbedding && this.config.copilot.providers.openai.apiKey) {
this.client = new OpenAIEmbeddingClient(
this.config.copilot.providers.openai
);
if (this.supportEmbedding) {
this.client = await getEmbeddingClient(this.providerFactory, this.prompt);
}
}
@@ -89,6 +89,14 @@ export class CopilotContextDocJob {
if (!this.supportEmbedding) return;
for (const { workspaceId, docId } of docs) {
const jobId = `workspace:embedding:${workspaceId}:${docId}`;
const job = await this.queue.get(jobId, 'copilot.embedding.docs');
// if the job exists and is older than 5 minute, remove it
if (job && job.timestamp + 5 * 60 * 1000 < Date.now()) {
this.logger.verbose(`Removing old embedding job ${jobId}`);
await this.queue.remove(jobId, 'copilot.embedding.docs');
}
await this.queue.add(
'copilot.embedding.docs',
{
@@ -99,6 +107,7 @@ export class CopilotContextDocJob {
{
jobId: `workspace:embedding:${workspaceId}:${docId}`,
priority: options?.priority ?? 1,
timestamp: Date.now(),
}
);
}
@@ -336,6 +345,9 @@ export class CopilotContextDocJob {
workspaceId,
docId
);
this.logger.verbose(
`Check if doc ${docId} in workspace ${workspaceId} needs embedding: ${needEmbedding}`
);
if (needEmbedding) {
if (signal.aborted) return;
const fragment = await this.getDocFragment(workspaceId, docId);