feat(server): context awareness for copilot (#9611)

fix PD-2167
fix PD-2169
fix PD-2190

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

@@ -1,30 +1,70 @@
-import { Injectable } from '@nestjs/common';
+import { Injectable, OnModuleInit } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
+import OpenAI from 'openai';
 
 import {
   Cache,
+  Config,
   CopilotInvalidContext,
   CopilotSessionNotFound,
+  NoCopilotProviderAvailable,
+  OnEvent,
+  PrismaTransaction,
 } from '../../../base';
+import { OpenAIEmbeddingClient } from './embedding';
 import { ContextSession } from './session';
-import { ContextConfig, ContextConfigSchema } from './types';
+import {
+  ContextConfig,
+  ContextConfigSchema,
+  ContextFile,
+  ContextFileStatus,
+  EmbeddingClient,
+} from './types';
+import { checkEmbeddingAvailable } from './utils';
 
 const CONTEXT_SESSION_KEY = 'context-session';
 
 @Injectable()
-export class CopilotContextService {
+export class CopilotContextService implements OnModuleInit {
+  private supportEmbedding = false;
+  private readonly client: EmbeddingClient | undefined;
+
   constructor(
+    config: Config,
     private readonly cache: Cache,
     private readonly db: PrismaClient
-  ) {}
+  ) {
+    const configure = config.plugins.copilot.openai;
+    if (configure) {
+      this.client = new OpenAIEmbeddingClient(new OpenAI(configure));
+    }
+  }
+
+  async onModuleInit() {
+    const supportEmbedding = await checkEmbeddingAvailable(this.db);
+    if (supportEmbedding) {
+      this.supportEmbedding = true;
+    }
+  }
+
+  get canEmbedding() {
+    return this.supportEmbedding;
+  }
+
+  // public this client to allow overriding in tests
+  get embeddingClient() {
+    return this.client as EmbeddingClient;
+  }
 
   private async saveConfig(
     contextId: string,
     config: ContextConfig,
+    tx?: PrismaTransaction,
     refreshCache = false
   ): Promise<void> {
     if (!refreshCache) {
-      await this.db.aiContext.update({
+      const executor = tx || this.db;
+      await executor.aiContext.update({
         where: { id: contextId },
         data: { config },
       });
@@ -42,8 +82,10 @@ export class CopilotContextService {
       const config = ContextConfigSchema.safeParse(cachedSession);
       if (config.success) {
         return new ContextSession(
+          this.embeddingClient,
           contextId,
           config.data,
+          this.db,
           this.saveConfig.bind(this, contextId)
         );
       }
@@ -60,8 +102,14 @@ export class CopilotContextService {
     config: ContextConfig
   ): Promise<ContextSession> {
     const dispatcher = this.saveConfig.bind(this, contextId);
-    await dispatcher(config, true);
-    return new ContextSession(contextId, config, dispatcher);
+    await dispatcher(config, undefined, true);
+    return new ContextSession(
+      this.embeddingClient,
+      contextId,
+      config,
+      this.db,
+      dispatcher
+    );
   }
 
   async create(sessionId: string): Promise<ContextSession> {
@@ -89,6 +137,10 @@ export class CopilotContextService {
   }
 
   async get(id: string): Promise<ContextSession> {
+    if (!this.embeddingClient) {
+      throw new NoCopilotProviderAvailable('embedding client not configured');
+    }
+
     const context = await this.getCachedSession(id);
     if (context) return context;
     const ret = await this.db.aiContext.findUnique({
@@ -110,4 +162,32 @@ export class CopilotContextService {
     if (existsContext) return this.get(existsContext.id);
     return null;
   }
+
+  @OnEvent('workspace.file.embed.finished')
+  async onFileEmbedFinish({
+    contextId,
+    fileId,
+    chunkSize,
+  }: Events['workspace.file.embed.finished']) {
+    const context = await this.get(contextId);
+    await context.saveFileRecord(fileId, file => ({
+      ...(file as ContextFile),
+      chunkSize,
+      status: ContextFileStatus.finished,
+    }));
+  }
+
+  @OnEvent('workspace.file.embed.failed')
+  async onFileEmbedFailed({
+    contextId,
+    fileId,
+    error,
+  }: Events['workspace.file.embed.failed']) {
+    const context = await this.get(contextId);
+    await context.saveFileRecord(fileId, file => ({
+      ...(file as ContextFile),
+      error,
+      status: ContextFileStatus.failed,
+    }));
+  }
 }