From e98f035f97d50560ad0154d56b231927d961ee3e Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Fri, 13 Jun 2025 12:37:05 +0800 Subject: [PATCH] feat(server): split embedding client (#12809) --- packages/backend/server/src/__tests__/copilot.e2e.ts | 12 ++++++------ .../backend/server/src/__tests__/copilot.spec.ts | 12 ++++++------ .../server/src/plugins/copilot/context/index.ts | 1 - .../server/src/plugins/copilot/context/resolver.ts | 4 ++-- .../server/src/plugins/copilot/context/service.ts | 3 +-- .../server/src/plugins/copilot/context/session.ts | 6 ++++-- .../server/src/plugins/copilot/context/utils.ts | 11 ----------- .../{context/embedding.ts => embedding/client.ts} | 8 ++++---- .../server/src/plugins/copilot/embedding/index.ts | 4 ++++ .../plugins/copilot/{context => embedding}/job.ts | 6 +++--- .../plugins/copilot/{context => embedding}/types.ts | 0 packages/backend/server/src/plugins/copilot/index.ts | 4 ++-- 12 files changed, 32 insertions(+), 39 deletions(-) delete mode 100644 packages/backend/server/src/plugins/copilot/context/utils.ts rename packages/backend/server/src/plugins/copilot/{context/embedding.ts => embedding/client.ts} (96%) create mode 100644 packages/backend/server/src/plugins/copilot/embedding/index.ts rename packages/backend/server/src/plugins/copilot/{context => embedding}/job.ts (98%) rename packages/backend/server/src/plugins/copilot/{context => embedding}/types.ts (100%) diff --git a/packages/backend/server/src/__tests__/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot.e2e.ts index 3e5537a0d..bd65c5b0f 100644 --- a/packages/backend/server/src/__tests__/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot.e2e.ts @@ -11,11 +11,11 @@ import { JobQueue } from '../base'; import { ConfigModule } from '../base/config'; import { AuthService } from '../core/auth'; import { DocReader } from '../core/doc'; +import { CopilotContextService } from '../plugins/copilot/context'; import { - CopilotContextDocJob, - CopilotContextService, -} from '../plugins/copilot/context'; -import { MockEmbeddingClient } from '../plugins/copilot/context/embedding'; + CopilotEmbeddingJob, + MockEmbeddingClient, +} from '../plugins/copilot/embedding'; import { prompts, PromptService } from '../plugins/copilot/prompt'; import { CopilotProviderFactory, @@ -65,7 +65,7 @@ const test = ava as TestFn<{ app: TestingApp; db: PrismaClient; context: CopilotContextService; - jobs: CopilotContextDocJob; + jobs: CopilotEmbeddingJob; prompt: PromptService; factory: CopilotProviderFactory; storage: CopilotStorage; @@ -115,7 +115,7 @@ test.before(async t => { const context = app.get(CopilotContextService); const prompt = app.get(PromptService); const storage = app.get(CopilotStorage); - const jobs = app.get(CopilotContextDocJob); + const jobs = app.get(CopilotEmbeddingJob); t.context.app = app; t.context.db = db; diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts index 547878928..44724e8c9 100644 --- a/packages/backend/server/src/__tests__/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot.spec.ts @@ -13,11 +13,11 @@ import { AuthService } from '../core/auth'; import { QuotaModule } from '../core/quota'; import { ContextCategories, WorkspaceModel } from '../models'; import { CopilotModule } from '../plugins/copilot'; +import { CopilotContextService } from '../plugins/copilot/context'; import { - CopilotContextDocJob, - CopilotContextService, -} from '../plugins/copilot/context'; -import { MockEmbeddingClient } from '../plugins/copilot/context/embedding'; + CopilotEmbeddingJob, + MockEmbeddingClient, +} from '../plugins/copilot/embedding'; import { prompts, PromptService } from '../plugins/copilot/prompt'; import { CopilotProviderFactory, @@ -69,7 +69,7 @@ const test = ava as TestFn<{ workspaceEmbedding: CopilotWorkspaceService; factory: CopilotProviderFactory; session: ChatSessionService; - jobs: CopilotContextDocJob; + jobs: CopilotEmbeddingJob; storage: CopilotStorage; workflow: CopilotWorkflowService; executors: { @@ -127,7 +127,7 @@ test.before(async t => { const storage = module.get(CopilotStorage); const context = module.get(CopilotContextService); - const jobs = module.get(CopilotContextDocJob); + const jobs = module.get(CopilotEmbeddingJob); const transcript = module.get(CopilotTranscriptionService); const workspaceEmbedding = module.get(CopilotWorkspaceService); diff --git a/packages/backend/server/src/plugins/copilot/context/index.ts b/packages/backend/server/src/plugins/copilot/context/index.ts index e0562903e..d3afb3caa 100644 --- a/packages/backend/server/src/plugins/copilot/context/index.ts +++ b/packages/backend/server/src/plugins/copilot/context/index.ts @@ -1,3 +1,2 @@ -export { CopilotContextDocJob } from './job'; export { CopilotContextResolver, CopilotContextRootResolver } from './resolver'; export { CopilotContextService } from './service'; diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index ed2a7e402..fa15f8e06 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -44,12 +44,12 @@ import { FileChunkSimilarity, Models, } from '../../../models'; +import { CopilotEmbeddingJob } from '../embedding'; import { COPILOT_LOCKER, CopilotType } from '../resolver'; import { ChatSessionService } from '../session'; import { CopilotStorage } from '../storage'; import { MAX_EMBEDDABLE_SIZE } from '../types'; import { readStream } from '../utils'; -import { CopilotContextDocJob } from './job'; import { CopilotContextService } from './service'; @InputType() @@ -387,7 +387,7 @@ export class CopilotContextResolver { private readonly models: Models, private readonly mutex: RequestMutex, private readonly context: CopilotContextService, - private readonly jobs: CopilotContextDocJob, + private readonly jobs: CopilotEmbeddingJob, private readonly storage: CopilotStorage ) {} diff --git a/packages/backend/server/src/plugins/copilot/context/service.ts b/packages/backend/server/src/plugins/copilot/context/service.ts index ab727712e..47c618e81 100644 --- a/packages/backend/server/src/plugins/copilot/context/service.ts +++ b/packages/backend/server/src/plugins/copilot/context/service.ts @@ -14,11 +14,10 @@ import { ContextFile, Models, } from '../../../models'; +import { type EmbeddingClient, getEmbeddingClient } from '../embedding'; import { PromptService } from '../prompt'; import { CopilotProviderFactory } from '../providers'; -import { getEmbeddingClient } from './embedding'; import { ContextSession } from './session'; -import type { EmbeddingClient } from './types'; const CONTEXT_SESSION_KEY = 'context-session'; diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index a84ac4872..855aa04a1 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -11,11 +11,11 @@ import { FileChunkSimilarity, Models, } from '../../../models'; -import { EmbeddingClient } from './types'; +import { EmbeddingClient } from '../embedding'; export class ContextSession implements AsyncDisposable { constructor( - private readonly client: EmbeddingClient, + private readonly client: EmbeddingClient | undefined, private readonly contextId: string, private readonly config: ContextConfig, private readonly models: Models, @@ -204,6 +204,7 @@ export class ContextSession implements AsyncDisposable { scopedThreshold: number = 0.85, threshold: number = 0.5 ): Promise { + if (!this.client) return []; const embedding = await this.client.getEmbedding(content, signal); if (!embedding) return []; @@ -256,6 +257,7 @@ export class ContextSession implements AsyncDisposable { scopedThreshold: number = 0.85, threshold: number = 0.5 ) { + if (!this.client) return []; const embedding = await this.client.getEmbedding(content, signal); if (!embedding) return []; diff --git a/packages/backend/server/src/plugins/copilot/context/utils.ts b/packages/backend/server/src/plugins/copilot/context/utils.ts deleted file mode 100644 index 451280d68..000000000 --- a/packages/backend/server/src/plugins/copilot/context/utils.ts +++ /dev/null @@ -1,11 +0,0 @@ -export class GqlSignal implements AsyncDisposable { - readonly abortController = new AbortController(); - - get signal() { - return this.abortController.signal; - } - - async [Symbol.asyncDispose]() { - this.abortController.abort(); - } -} diff --git a/packages/backend/server/src/plugins/copilot/context/embedding.ts b/packages/backend/server/src/plugins/copilot/embedding/client.ts similarity index 96% rename from packages/backend/server/src/plugins/copilot/context/embedding.ts rename to packages/backend/server/src/plugins/copilot/embedding/client.ts index 240d0c26c..f2d4c7b6e 100644 --- a/packages/backend/server/src/plugins/copilot/context/embedding.ts +++ b/packages/backend/server/src/plugins/copilot/embedding/client.ts @@ -4,11 +4,11 @@ import { CopilotPromptNotFound, CopilotProviderNotSupported, } from '../../../base'; -import type { ChunkSimilarity, Embedding } from '../../../models'; -import type { PromptService } from '../prompt'; +import { ChunkSimilarity, Embedding } from '../../../models'; +import { PromptService } from '../prompt'; import { type CopilotProvider, - type CopilotProviderFactory, + CopilotProviderFactory, type ModelFullConditions, ModelInputType, ModelOutputType, @@ -22,7 +22,7 @@ import { const RERANK_PROMPT = 'Rerank results'; -export class ProductionEmbeddingClient extends EmbeddingClient { +class ProductionEmbeddingClient extends EmbeddingClient { private readonly logger = new Logger(ProductionEmbeddingClient.name); constructor( diff --git a/packages/backend/server/src/plugins/copilot/embedding/index.ts b/packages/backend/server/src/plugins/copilot/embedding/index.ts new file mode 100644 index 000000000..0091adf30 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/embedding/index.ts @@ -0,0 +1,4 @@ +export { getEmbeddingClient, MockEmbeddingClient } from './client'; +export { CopilotEmbeddingJob } from './job'; +export type { Chunk, DocFragment } from './types'; +export { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types'; diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/embedding/job.ts similarity index 98% rename from packages/backend/server/src/plugins/copilot/context/job.ts rename to packages/backend/server/src/plugins/copilot/embedding/job.ts index 9562b786e..c7784ea07 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/embedding/job.ts @@ -18,12 +18,12 @@ import { PromptService } from '../prompt'; import { CopilotProviderFactory } from '../providers'; import { CopilotStorage } from '../storage'; import { readStream } from '../utils'; -import { getEmbeddingClient } from './embedding'; +import { getEmbeddingClient } from './client'; import type { Chunk, DocFragment } from './types'; import { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types'; @Injectable() -export class CopilotContextDocJob { +export class CopilotEmbeddingJob { private readonly workspaceJobAbortController: Map = new Map(); @@ -40,7 +40,7 @@ export class CopilotContextDocJob { private readonly queue: JobQueue, private readonly storage: CopilotStorage ) { - this.logger.setContext(CopilotContextDocJob.name); + this.logger.setContext(CopilotEmbeddingJob.name); } @OnEvent('config.init') diff --git a/packages/backend/server/src/plugins/copilot/context/types.ts b/packages/backend/server/src/plugins/copilot/embedding/types.ts similarity index 100% rename from packages/backend/server/src/plugins/copilot/context/types.ts rename to packages/backend/server/src/plugins/copilot/embedding/types.ts diff --git a/packages/backend/server/src/plugins/copilot/index.ts b/packages/backend/server/src/plugins/copilot/index.ts index faf097386..31a3b3947 100644 --- a/packages/backend/server/src/plugins/copilot/index.ts +++ b/packages/backend/server/src/plugins/copilot/index.ts @@ -9,12 +9,12 @@ import { PermissionModule } from '../../core/permission'; import { QuotaModule } from '../../core/quota'; import { WorkspaceModule } from '../../core/workspaces'; import { - CopilotContextDocJob, CopilotContextResolver, CopilotContextRootResolver, CopilotContextService, } from './context'; import { CopilotController } from './controller'; +import { CopilotEmbeddingJob } from './embedding'; import { ChatMessageCache } from './message'; import { PromptService } from './prompt'; import { CopilotProviderFactory, CopilotProviders } from './providers'; @@ -61,7 +61,7 @@ import { // context CopilotContextResolver, CopilotContextService, - CopilotContextDocJob, + CopilotEmbeddingJob, // transcription CopilotTranscriptionService, CopilotTranscriptionResolver,