feat(server): split embedding client (#12809)

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;

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);
 

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';

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
   ) {}
 

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';
 

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<FileChunkSimilarity[]> {
+    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 [];
 

packages/backend/server/src/plugins/copilot/context/utils.ts

@@ -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();
-  }
-}

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(

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';

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<string, AbortController> =
     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')

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,