feat: context service

Author: darkskygit

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

@@ -0,0 +1,91 @@
+import { Injectable } from '@nestjs/common';
+import { PrismaClient } from '@prisma/client';
+import { z } from 'zod';
+
+// import type { FileLike } from './types';
+
+const ContextConfigSchema = z.object({
+  files: z
+    .object({
+      id: z.string(),
+      chunk_size: z.number(),
+      name: z.string(),
+    })
+    .array(),
+});
+
+type ContextConfig = z.infer<typeof ContextConfigSchema>;
+
+@Injectable()
+export class CopilotContextService {
+  private readonly sessionCache = new Map<string, ContextSession>();
+
+  constructor(private readonly db: PrismaClient) {}
+
+  private saveContext(
+    workspaceId: string,
+    id: string,
+    config: ContextConfig
+  ): ContextSession {
+    const context = new ContextSession(workspaceId, id, config, this.db);
+    this.sessionCache.set(context.id, context);
+    return context;
+  }
+
+  async getOrCreate(workspaceId: string, id?: string): Promise<ContextSession> {
+    if (id) {
+      const context = this.sessionCache.get(id);
+      if (context) return context;
+      const ret = await this.db.aiContext.findUnique({
+        where: { workspaceId, id },
+        select: { config: true },
+      });
+      if (ret) {
+        const config = ContextConfigSchema.safeParse(ret.config);
+        if (config.success)
+          return this.saveContext(workspaceId, id, config.data);
+        throw new Error('Invalid context config');
+      }
+    }
+
+    const context = await this.db.aiContext.create({
+      data: { workspaceId, config: { files: [] } },
+    });
+    const config = ContextConfigSchema.parse(context.config);
+    return this.saveContext(workspaceId, context.id, config);
+  }
+}
+
+export class ContextSession implements AsyncDisposable {
+  constructor(
+    private readonly wsId: string,
+    private readonly contextId: string,
+    private readonly config: ContextConfig,
+    private readonly db: PrismaClient
+  ) {}
+
+  get workspaceId() {
+    return this.wsId;
+  }
+
+  get id() {
+    return this.contextId;
+  }
+
+  async list() {}
+
+  // async add(content: FileLike, signal?: AbortSignal) {}
+
+  // async remove(fileId: string) {}
+
+  async save() {
+    await this.db.aiContext.update({
+      where: { workspaceId: this.wsId, id: this.contextId },
+      data: { config: this.config },
+    });
+  }
+
+  async [Symbol.asyncDispose]() {
+    await this.save?.();
+  }
+}

packages/backend/server/src/plugins/copilot/index.ts

@@ -5,6 +5,7 @@ import { FeatureModule } from '../../core/features';
 import { PermissionModule } from '../../core/permission';
 import { QuotaModule } from '../../core/quota';
 import { Plugin } from '../registry';
+import { CopilotContextService } from './context';
 import { CopilotController } from './controller';
 import { ChatMessageCache } from './message';
 import { PromptService } from './prompt';
@@ -43,6 +44,7 @@ registerCopilotProvider(PerplexityProvider);
     PromptsManagementResolver,
     CopilotWorkflowService,
     ...CopilotWorkflowExecutors,
+    CopilotContextService,
   ],
   controllers: [CopilotController],
   contributesTo: ServerFeature.Copilot,

packages/backend/server/src/plugins/copilot/providers/openai/index.ts renamed to packages/backend/server/src/plugins/copilot/providers/openai.ts

@@ -6,7 +6,7 @@ import {
   CopilotProviderSideError,
   metrics,
   UserFriendlyError,
-} from '../../../../base';
+} from '../../../base';
 import {
   ChatMessageRole,
   CopilotCapability,
@@ -18,10 +18,8 @@ import {
   CopilotTextToEmbeddingProvider,
   CopilotTextToImageProvider,
   CopilotTextToTextProvider,
-  CopilotTextToTextWithContextProvider,
   PromptMessage,
-} from '../../types';
-import { Context, ContextService } from './context';
+} from '../types';
 
 export const DEFAULT_DIMENSIONS = 256;
 
@@ -30,15 +28,13 @@ const SIMPLE_IMAGE_URL_REGEX = /^(https?:\/\/|data:image\/)/;
 export class OpenAIProvider
   implements
     CopilotTextToTextProvider,
-    CopilotTextToTextWithContextProvider,
     CopilotTextToEmbeddingProvider,
     CopilotTextToImageProvider,
     CopilotImageToTextProvider
 {
   static readonly type = CopilotProviderType.OpenAI;
   static readonly capabilities = [
     CopilotCapability.TextToText,
-    CopilotCapability.TextToTextWithContext,
     CopilotCapability.TextToEmbedding,
     CopilotCapability.TextToImage,
     CopilotCapability.ImageToText,
@@ -63,12 +59,11 @@ export class OpenAIProvider
 
   private readonly logger = new Logger(OpenAIProvider.type);
   private readonly instance: OpenAI;
-  private readonly context: ContextService;
+
   private existsModels: string[] | undefined;
 
   constructor(config: ClientOptions) {
     this.instance = new OpenAI(config);
-    this.context = new ContextService(this.instance);
   }
 
   static assetsConfig(config: ClientOptions) {
@@ -83,10 +78,6 @@ export class OpenAIProvider
     return OpenAIProvider.capabilities;
   }
 
-  getContext(name: string, id?: string): Promise<Context> {
-    return this.context.getOrCreate(name, id);
-  }
-
   async isModelAvailable(model: string): Promise<boolean> {
     const knownModels = this.availableModels.includes(model);
     if (knownModels) return true;

packages/backend/server/src/plugins/copilot/providers/openai/context.ts

@@ -166,7 +166,6 @@ export enum CopilotProviderType {
 
 export enum CopilotCapability {
   TextToText = 'text-to-text',
-  TextToTextWithContext = 'text-to-text-with-context',
   TextToEmbedding = 'text-to-embedding',
   TextToImage = 'text-to-image',
   ImageToImage = 'image-to-image',
@@ -184,14 +183,6 @@ const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
 
 export type CopilotChatOptions = z.infer<typeof CopilotChatOptionsSchema>;
 
-const CopilotChatWithContextOptionsSchema = CopilotProviderOptionsSchema.merge(
-  z.object({ contextId: z.string().optional() }).strict()
-);
-
-export type CopilotChatWithContextOptions = z.infer<
-  typeof CopilotChatWithContextOptionsSchema
->;
-
 const CopilotEmbeddingOptionsSchema = CopilotProviderOptionsSchema.extend({
   dimensions: z.number(),
 }).optional();
@@ -214,7 +205,7 @@ export type CopilotContextFile = {
   id: string; // fileId
   created_at: number;
   // embedding status
-  status: 'in_progress' | 'completed' | 'cancelled' | 'failed';
+  status: 'in_progress' | 'completed' | 'failed';
 };
 
 /**
@@ -241,13 +232,6 @@ export interface FileLike extends BlobLike {
   readonly name: string;
 }
 
-export interface CopilotContext {
-  list(): Promise<CopilotContextFile[]>;
-  add(content: FileLike, signal?: AbortSignal): Promise<string>;
-  addByFileId(context: string, signal?: AbortSignal): Promise<string>;
-  remove(fileId: string): Promise<boolean>;
-}
-
 export interface CopilotProvider {
   readonly type: CopilotProviderType;
   getCapabilities(): CopilotCapability[];
@@ -267,20 +251,6 @@ export interface CopilotTextToTextProvider extends CopilotProvider {
   ): AsyncIterable<string>;
 }
 
-export interface CopilotTextToTextWithContextProvider extends CopilotProvider {
-  generateText(
-    messages: PromptMessage[],
-    model?: string,
-    options?: CopilotChatWithContextOptions
-  ): Promise<string>;
-  generateTextStream(
-    messages: PromptMessage[],
-    model?: string,
-    options?: CopilotChatOptions
-  ): AsyncIterable<string>;
-  getContext(name: string, id?: string): Promise<CopilotContext>;
-}
-
 export interface CopilotTextToEmbeddingProvider extends CopilotProvider {
   generateEmbedding(
     messages: string[] | string,
@@ -330,7 +300,6 @@ export interface CopilotImageToImageProvider extends CopilotProvider {
 
 export type CapabilityToCopilotProvider = {
   [CopilotCapability.TextToText]: CopilotTextToTextProvider;
-  [CopilotCapability.TextToTextWithContext]: CopilotTextToTextWithContextProvider;
   [CopilotCapability.TextToEmbedding]: CopilotTextToEmbeddingProvider;
   [CopilotCapability.TextToImage]: CopilotTextToImageProvider;
   [CopilotCapability.ImageToText]: CopilotImageToTextProvider;