feat(server): adapt context model (#11028)

expose more field in listContextObject

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

@@ -1,6 +1,7 @@
 import OpenAI from 'openai';
 
-import { Embedding, EmbeddingClient } from './types';
+import { Embedding } from '../../../models';
+import { EmbeddingClient } from './types';
 
 export class OpenAIEmbeddingClient extends EmbeddingClient {
   constructor(private readonly client: OpenAI) {
@@ -15,6 +16,7 @@ export class OpenAIEmbeddingClient extends EmbeddingClient {
       {
         input,
         model: 'text-embedding-3-large',
+        dimensions: 1024,
         encoding_format: 'float',
       },
       { signal }

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

@@ -1,7 +1,3 @@
 export { CopilotContextDocJob } from './job';
 export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
 export { CopilotContextService } from './service';
-export {
-  type ContextFile,
-  ContextEmbedStatus as ContextFileStatus,
-} from './types';

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

@@ -1,7 +1,4 @@
-import { randomUUID } from 'node:crypto';
-
 import { Injectable, OnModuleInit } from '@nestjs/common';
-import { Prisma, PrismaClient } from '@prisma/client';
 import OpenAI from 'openai';
 
 import {
@@ -15,10 +12,11 @@ import {
   OnJob,
 } from '../../../base';
 import { DocReader } from '../../../core/doc';
+import { Models } from '../../../models';
 import { CopilotStorage } from '../storage';
 import { OpenAIEmbeddingClient } from './embedding';
-import { Embedding, EmbeddingClient } from './types';
-import { checkEmbeddingAvailable, readStream } from './utils';
+import { EmbeddingClient } from './types';
+import { readStream } from './utils';
 
 declare global {
   interface Jobs {
@@ -45,10 +43,10 @@ export class CopilotContextDocJob implements OnModuleInit {
 
   constructor(
     config: Config,
-    private readonly db: PrismaClient,
     private readonly doc: DocReader,
     private readonly event: EventBus,
     private readonly logger: AFFiNELogger,
+    private readonly models: Models,
     private readonly queue: JobQueue,
     private readonly storage: CopilotStorage
   ) {
@@ -60,7 +58,8 @@ export class CopilotContextDocJob implements OnModuleInit {
   }
 
   async onModuleInit() {
-    this.supportEmbedding = await checkEmbeddingAvailable(this.db);
+    this.supportEmbedding =
+      await this.models.copilotContext.checkEmbeddingAvailable();
   }
 
   // public this client to allow overriding in tests
@@ -91,23 +90,6 @@ export class CopilotContextDocJob implements OnModuleInit {
     }
   }
 
-  private processEmbeddings(
-    contextOrWorkspaceId: string,
-    fileOrDocId: string,
-    embeddings: Embedding[]
-  ) {
-    const groups = embeddings.map(e => [
-      randomUUID(),
-      contextOrWorkspaceId,
-      fileOrDocId,
-      e.index,
-      e.content,
-      Prisma.raw(`'[${e.embedding.join(',')}]'`),
-      new Date(),
-    ]);
-    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
-  }
-
   async readCopilotBlob(
     userId: string,
     workspaceId: string,
@@ -145,14 +127,11 @@ export class CopilotContextDocJob implements OnModuleInit {
 
       for (const chunk of chunks) {
         const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
-        const values = this.processEmbeddings(contextId, fileId, embeddings);
-
-        await this.db.$executeRaw`
-          INSERT INTO "ai_context_embeddings"
-          ("id", "context_id", "file_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
-          ON CONFLICT (context_id, file_id, chunk) DO UPDATE SET
-          content = EXCLUDED.content, embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
-        `;
+        await this.models.copilotContext.insertContentEmbedding(
+          contextId,
+          fileId,
+          embeddings
+        );
       }
 
       this.event.emit('workspace.file.embed.finished', {
@@ -188,13 +167,11 @@ export class CopilotContextDocJob implements OnModuleInit {
         );
 
         for (const chunks of embeddings) {
-          const values = this.processEmbeddings(workspaceId, docId, chunks);
-          await this.db.$executeRaw`
-              INSERT INTO "ai_workspace_embeddings"
-              ("workspace_id", "doc_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
-              ON CONFLICT (context_id, file_id, chunk) DO UPDATE SET
-              embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
-            `;
+          await this.models.copilotContext.insertWorkspaceEmbedding(
+            workspaceId,
+            docId,
+            chunks
+          );
         }
       }
     } catch (e: any) {

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

@@ -34,25 +34,41 @@ import {
 } from '../../../base';
 import { CurrentUser } from '../../../core/auth';
 import { AccessController } from '../../../core/permission';
-import { COPILOT_LOCKER, CopilotType } from '../resolver';
-import { ChatSessionService } from '../session';
-import { CopilotStorage } from '../storage';
-import { CopilotContextDocJob } from './job';
-import { CopilotContextService } from './service';
 import {
   ContextCategories,
   ContextCategory,
   ContextDoc,
   ContextEmbedStatus,
-  type ContextFile,
+  ContextFile,
   DocChunkSimilarity,
   FileChunkSimilarity,
-  MAX_EMBEDDABLE_SIZE,
-} from './types';
+  Models,
+} from '../../../models';
+import { COPILOT_LOCKER, CopilotType } from '../resolver';
+import { ChatSessionService } from '../session';
+import { CopilotStorage } from '../storage';
+import { CopilotContextDocJob } from './job';
+import { CopilotContextService } from './service';
+import { MAX_EMBEDDABLE_SIZE } from './types';
 import { readStream } from './utils';
 
 @InputType()
-class AddRemoveContextCategoryInput {
+class AddContextCategoryInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => ContextCategories)
+  type!: ContextCategories;
+
+  @Field(() => String)
+  categoryId!: string;
+
+  @Field(() => [String], { nullable: true })
+  docs!: string[] | null;
+}
+
+@InputType()
+class RemoveContextCategoryInput {
   @Field(() => String)
   contextId!: string;
 
@@ -111,21 +127,7 @@ export class CopilotContextType {
 registerEnumType(ContextCategories, { name: 'ContextCategories' });
 
 @ObjectType()
-class CopilotContextCategory implements ContextCategory {
-  @Field(() => ID)
-  id!: string;
-
-  @Field(() => ContextCategories)
-  type!: ContextCategories;
-
-  @Field(() => SafeIntResolver)
-  createdAt!: number;
-}
-
-registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' });
-
-@ObjectType()
-class CopilotContextDoc implements ContextDoc {
+class CopilotDocType implements ContextDoc {
   @Field(() => ID)
   id!: string;
 
@@ -136,6 +138,29 @@ class CopilotContextDoc implements ContextDoc {
   createdAt!: number;
 }
 
+@ObjectType()
+class CopilotContextCategory implements Omit<ContextCategory, 'docs'> {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => ContextCategories)
+  type!: ContextCategories;
+
+  @Field(() => [CopilotDocType])
+  docs!: CopilotDocType[];
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' });
+
+@ObjectType()
+class CopilotContextDoc extends CopilotDocType {
+  @Field(() => String, { nullable: true })
+  error!: string | null;
+}
+
 @ObjectType()
 class CopilotContextFile implements ContextFile {
   @Field(() => ID)
@@ -338,6 +363,7 @@ export class CopilotContextRootResolver {
 export class CopilotContextResolver {
   constructor(
     private readonly ac: AccessController,
+    private readonly models: Models,
     private readonly mutex: RequestMutex,
     private readonly context: CopilotContextService,
     private readonly jobs: CopilotContextDocJob,
@@ -354,13 +380,61 @@ export class CopilotContextResolver {
     return controller.signal;
   }
 
+  @ResolveField(() => [CopilotContextCategory], {
+    description: 'list collections in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async collections(
+    @Parent() context: CopilotContextType
+  ): Promise<ContextCategory[]> {
+    const session = await this.context.get(context.id);
+    const collections = session.collections;
+    await this.models.copilotContext.mergeDocStatus(
+      session.workspaceId,
+      collections.flatMap(c => c.docs)
+    );
+
+    return collections;
+  }
+
+  @ResolveField(() => [CopilotContextCategory], {
+    description: 'list tags in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async tags(
+    @Parent() context: CopilotContextType
+  ): Promise<ContextCategory[]> {
+    const session = await this.context.get(context.id);
+    const tags = session.tags;
+    await this.models.copilotContext.mergeDocStatus(
+      session.workspaceId,
+      tags.flatMap(c => c.docs)
+    );
+
+    return tags;
+  }
+
   @ResolveField(() => [CopilotContextDoc], {
     description: 'list files in context',
   })
   @CallMetric('ai', 'context_file_list')
   async docs(@Parent() context: CopilotContextType): Promise<ContextDoc[]> {
     const session = await this.context.get(context.id);
-    return session.listDocs();
+    const docs = session.docs;
+    await this.models.copilotContext.mergeDocStatus(session.workspaceId, docs);
+
+    return docs;
+  }
+
+  @ResolveField(() => [CopilotContextFile], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async files(
+    @Parent() context: CopilotContextType
+  ): Promise<CopilotContextFile[]> {
+    const session = await this.context.get(context.id);
+    return session.files;
   }
 
   @Mutation(() => CopilotContextCategory, {
@@ -368,18 +442,33 @@ export class CopilotContextResolver {
   })
   @CallMetric('ai', 'context_category_add')
   async addContextCategory(
-    @Args({ name: 'options', type: () => AddRemoveContextCategoryInput })
-    options: AddRemoveContextCategoryInput
-  ) {
+    @Args({ name: 'options', type: () => AddContextCategoryInput })
+    options: AddContextCategoryInput
+  ): Promise<CopilotContextCategory> {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
     const session = await this.context.get(options.contextId);
 
     try {
-      return await session.addCategoryRecord(options.type, options.categoryId);
+      const records = await session.addCategoryRecord(
+        options.type,
+        options.categoryId,
+        options.docs || []
+      );
+
+      if (options.docs) {
+        await this.jobs.addDocEmbeddingQueue(
+          options.docs.map(docId => ({
+            workspaceId: session.workspaceId,
+            docId,
+          }))
+        );
+      }
+
+      return records;
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,
@@ -393,8 +482,8 @@ export class CopilotContextResolver {
   })
   @CallMetric('ai', 'context_category_remove')
   async removeContextCategory(
-    @Args({ name: 'options', type: () => AddRemoveContextCategoryInput })
-    options: AddRemoveContextCategoryInput
+    @Args({ name: 'options', type: () => RemoveContextCategoryInput })
+    options: RemoveContextCategoryInput
   ) {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
@@ -432,7 +521,16 @@ export class CopilotContextResolver {
     const session = await this.context.get(options.contextId);
 
     try {
-      return await session.addDocRecord(options.docId);
+      const record = await session.addDocRecord(options.docId);
+
+      await this.jobs.addDocEmbeddingQueue([
+        {
+          workspaceId: session.workspaceId,
+          docId: options.docId,
+        },
+      ]);
+
+      return record;
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,
@@ -466,17 +564,6 @@ export class CopilotContextResolver {
     }
   }
 
-  @ResolveField(() => [CopilotContextFile], {
-    description: 'list files in context',
-  })
-  @CallMetric('ai', 'context_file_list')
-  async files(
-    @Parent() context: CopilotContextType
-  ): Promise<CopilotContextFile[]> {
-    const session = await this.context.get(context.id);
-    return session.listFiles();
-  }
-
   @Mutation(() => CopilotContextFile, {
     description: 'add a file to context',
   })

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

@@ -1,27 +1,23 @@
 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,
   ContextEmbedStatus,
   ContextFile,
-  EmbeddingClient,
-  MinimalContextConfigSchema,
-} from './types';
-import { checkEmbeddingAvailable } from './utils';
+  Models,
+} from '../../../models';
+import { OpenAIEmbeddingClient } from './embedding';
+import { ContextSession } from './session';
+import { EmbeddingClient } from './types';
 
 const CONTEXT_SESSION_KEY = 'context-session';
 
@@ -33,7 +29,7 @@ export class CopilotContextService implements OnModuleInit {
   constructor(
     config: Config,
     private readonly cache: Cache,
-    private readonly db: PrismaClient
+    private readonly models: Models
   ) {
     const configure = config.plugins.copilot.openai;
     if (configure) {
@@ -42,7 +38,8 @@ export class CopilotContextService implements OnModuleInit {
   }
 
   async onModuleInit() {
-    const supportEmbedding = await checkEmbeddingAvailable(this.db);
+    const supportEmbedding =
+      await this.models.copilotContext.checkEmbeddingAvailable();
     if (supportEmbedding) {
       this.supportEmbedding = true;
     }
@@ -60,15 +57,10 @@ export class CopilotContextService implements OnModuleInit {
   private async saveConfig(
     contextId: string,
     config: ContextConfig,
-    tx?: PrismaTransaction,
     refreshCache = false
   ): Promise<void> {
     if (!refreshCache) {
-      const executor = tx || this.db;
-      await executor.aiContext.update({
-        where: { id: contextId },
-        data: { config },
-      });
+      await this.models.copilotContext.update(contextId, { config });
     }
     await this.cache.set(`${CONTEXT_SESSION_KEY}:${contextId}`, config);
   }
@@ -86,7 +78,7 @@ export class CopilotContextService implements OnModuleInit {
           this.embeddingClient,
           contextId,
           config.data,
-          this.db,
+          this.models,
           this.saveConfig.bind(this, contextId)
         );
       }
@@ -103,41 +95,22 @@ export class CopilotContextService implements OnModuleInit {
     config: ContextConfig
   ): Promise<ContextSession> {
     const dispatcher = this.saveConfig.bind(this, contextId);
-    await dispatcher(config, undefined, true);
+    await dispatcher(config, true);
     return new ContextSession(
       this.embeddingClient,
       contextId,
       config,
-      this.db,
+      this.models,
       dispatcher
     );
   }
 
   async create(sessionId: string): Promise<ContextSession> {
-    const session = await this.db.aiSession.findFirst({
-      where: { id: sessionId },
-      select: { workspaceId: true },
-    });
-    if (!session) {
-      throw new CopilotSessionNotFound();
-    }
-
     // keep the context unique per session
     const existsContext = await this.getBySessionId(sessionId);
     if (existsContext) return existsContext;
 
-    const context = await this.db.aiContext.create({
-      data: {
-        sessionId,
-        config: {
-          workspaceId: session.workspaceId,
-          docs: [],
-          files: [],
-          categories: [],
-        },
-      },
-    });
-
+    const context = await this.models.copilotContext.create(sessionId);
     const config = ContextConfigSchema.parse(context.config);
     return await this.cacheSession(context.id, config);
   }
@@ -149,34 +122,16 @@ export class CopilotContextService implements OnModuleInit {
 
     const context = await this.getCachedSession(id);
     if (context) return context;
-    const ret = await this.db.aiContext.findUnique({
-      where: { id },
-      select: { config: true },
-    });
-    if (ret) {
-      const config = ContextConfigSchema.safeParse(ret.config);
-      if (config.success) {
-        return this.cacheSession(id, config.data);
-      }
-      const minimalConfig = MinimalContextConfigSchema.safeParse(ret.config);
-      if (minimalConfig.success) {
-        // fulfill the missing fields
-        return this.cacheSession(id, {
-          ...minimalConfig.data,
-          docs: [],
-          files: [],
-          categories: [],
-        });
-      }
+    const config = await this.models.copilotContext.getConfig(id);
+    if (config) {
+      return this.cacheSession(id, config);
     }
     throw new CopilotInvalidContext({ contextId: id });
   }
 
   async getBySessionId(sessionId: string): Promise<ContextSession | null> {
-    const existsContext = await this.db.aiContext.findFirst({
-      where: { sessionId },
-      select: { id: true },
-    });
+    const existsContext =
+      await this.models.copilotContext.getBySessionId(sessionId);
     if (existsContext) return this.get(existsContext.id);
     return null;
   }

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

@@ -1,30 +1,25 @@
-import { PrismaClient } from '@prisma/client';
 import { nanoid } from 'nanoid';
 
-import { PrismaTransaction } from '../../../base';
+import { CopilotDocsNotFound } from '../../../base';
 import {
-  ChunkSimilarity,
   ContextCategories,
+  ContextCategory,
   ContextConfig,
   ContextDoc,
   ContextEmbedStatus,
   ContextFile,
   ContextList,
-  DocChunkSimilarity,
-  EmbeddingClient,
-  FileChunkSimilarity,
-} from './types';
+  Models,
+} from '../../../models';
+import { EmbeddingClient } from './types';
 
 export class ContextSession implements AsyncDisposable {
   constructor(
     private readonly client: EmbeddingClient,
     private readonly contextId: string,
     private readonly config: ContextConfig,
-    private readonly db: PrismaClient,
-    private readonly dispatcher?: (
-      config: ContextConfig,
-      tx?: PrismaTransaction
-    ) => Promise<void>
+    private readonly models: Models,
+    private readonly dispatcher?: (config: ContextConfig) => Promise<void>
   ) {}
 
   get id() {
@@ -35,11 +30,28 @@ export class ContextSession implements AsyncDisposable {
     return this.config.workspaceId;
   }
 
-  listDocs(): ContextDoc[] {
-    return [...this.config.docs];
+  get categories(): ContextCategory[] {
+    return this.config.categories.map(c => ({
+      ...c,
+      docs: c.docs.map(d => ({ ...d })),
+    }));
   }
 
-  listFiles() {
+  get tags() {
+    const categories = this.config.categories;
+    return categories.filter(c => c.type === ContextCategories.Tag);
+  }
+
+  get collections() {
+    const categories = this.config.categories;
+    return categories.filter(c => c.type === ContextCategories.Collection);
+  }
+
+  get docs(): ContextDoc[] {
+    return this.config.docs.map(d => ({ ...d }));
+  }
+
+  get files() {
     return this.config.files.map(f => ({ ...f }));
   }
 
@@ -50,14 +62,25 @@ export class ContextSession implements AsyncDisposable {
     ) as ContextList;
   }
 
-  async addCategoryRecord(type: ContextCategories, id: string) {
+  async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
+    const existDocs = await this.models.doc.existsAll(this.workspaceId, docs);
+    if (!existDocs) {
+      throw new CopilotDocsNotFound();
+    }
+
     const category = this.config.categories.find(
       c => c.type === type && c.id === id
     );
     if (category) {
       return category;
     }
-    const record = { id, type, createdAt: Date.now() };
+    const createdAt = Date.now();
+    const record = {
+      id,
+      type,
+      docs: docs.map(id => ({ id, createdAt, status: null })),
+      createdAt,
+    };
     this.config.categories.push(record);
     await this.save();
     return record;
@@ -122,14 +145,10 @@ export class ContextSession implements AsyncDisposable {
   }
 
   async removeFile(fileId: string): Promise<boolean> {
-    return await this.db.$transaction(async tx => {
-      await tx.aiContextEmbedding.deleteMany({
-        where: { contextId: this.contextId, fileId },
-      });
-      this.config.files = this.config.files.filter(f => f.id !== fileId);
-      await this.save(tx);
-      return true;
-    });
+    await this.models.copilotContext.deleteEmbedding(this.contextId, fileId);
+    this.config.files = this.config.files.filter(f => f.id !== fileId);
+    await this.save();
+    return true;
   }
 
   /**
@@ -145,21 +164,18 @@ export class ContextSession implements AsyncDisposable {
     topK: number = 5,
     signal?: AbortSignal,
     threshold: number = 0.7
-  ): Promise<FileChunkSimilarity[]> {
+  ) {
     const embedding = await this.client
       .getEmbeddings([content], signal)
       .then(r => r?.[0]?.embedding);
     if (!embedding) return [];
-    const similarityChunks = await this.db.$queryRaw<
-      Array<FileChunkSimilarity>
-    >`
-      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
-      FROM "ai_context_embeddings"
-      WHERE context_id = ${this.id}
-      ORDER BY "distance" ASC
-      LIMIT ${topK};
-    `;
-    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+
+    return this.models.copilotContext.matchContentEmbedding(
+      embedding,
+      this.id,
+      topK,
+      threshold
+    );
   }
 
   /**
@@ -175,19 +191,18 @@ export class ContextSession implements AsyncDisposable {
     topK: number = 5,
     signal?: AbortSignal,
     threshold: number = 0.7
-  ): Promise<ChunkSimilarity[]> {
+  ) {
     const embedding = await this.client
       .getEmbeddings([content], signal)
       .then(r => r?.[0]?.embedding);
     if (!embedding) return [];
-    const similarityChunks = await this.db.$queryRaw<Array<DocChunkSimilarity>>`
-      SELECT "doc_id" as "docId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
-      FROM "ai_workspace_embeddings"
-      WHERE "workspace_id" = ${this.workspaceId}
-      ORDER BY "distance" ASC
-      LIMIT ${topK};
-    `;
-    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+
+    return this.models.copilotContext.matchWorkspaceEmbedding(
+      embedding,
+      this.id,
+      topK,
+      threshold
+    );
   }
 
   async saveFileRecord(
@@ -195,8 +210,7 @@ export class ContextSession implements AsyncDisposable {
     cb: (
       record: Pick<ContextFile, 'id' | 'status'> &
         Partial<Omit<ContextFile, 'id' | 'status'>>
-    ) => ContextFile,
-    tx?: PrismaTransaction
+    ) => ContextFile
   ) {
     const files = this.config.files;
     const file = files.find(f => f.id === fileId);
@@ -206,11 +220,11 @@ export class ContextSession implements AsyncDisposable {
       const file = { id: fileId, status: ContextEmbedStatus.processing };
       files.push(cb(file));
     }
-    await this.save(tx);
+    await this.save();
   }
 
-  async save(tx?: PrismaTransaction) {
-    await this.dispatcher?.(this.config, tx);
+  async save() {
+    await this.dispatcher?.(this.config);
   }
 
   async [Symbol.asyncDispose]() {

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

@@ -1,8 +1,7 @@
 import { File } from 'node:buffer';
 
-import { z } from 'zod';
-
 import { CopilotContextFileNotSupported, OneMB } from '../../../base';
+import { Embedding } from '../../../models';
 import { parseDoc } from '../../../native';
 
 declare global {
@@ -26,99 +25,11 @@ declare global {
 
 export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
 
-export enum ContextEmbedStatus {
-  processing = 'processing',
-  finished = 'finished',
-  failed = 'failed',
-}
-
-export enum ContextCategories {
-  Tag = 'tag',
-  Collection = 'collection',
-}
-
-export const ContextConfigSchema = z.object({
-  workspaceId: z.string(),
-  files: z
-    .object({
-      id: z.string(),
-      chunkSize: z.number(),
-      name: z.string(),
-      status: z.enum([
-        ContextEmbedStatus.processing,
-        ContextEmbedStatus.finished,
-        ContextEmbedStatus.failed,
-      ]),
-      error: z.string().nullable(),
-      blobId: z.string(),
-      createdAt: z.number(),
-    })
-    .array(),
-  docs: z
-    .object({
-      id: z.string(),
-      // status for workspace doc embedding progress
-      // only exists when the client submits the doc embedding task
-      status: z
-        .enum([
-          ContextEmbedStatus.processing,
-          ContextEmbedStatus.finished,
-          ContextEmbedStatus.failed,
-        ])
-        .nullable(),
-      createdAt: z.number(),
-    })
-    .array(),
-  categories: z
-    .object({
-      id: z.string(),
-      type: z.enum([ContextCategories.Tag, ContextCategories.Collection]),
-      createdAt: z.number(),
-    })
-    .array(),
-});
-
-export const MinimalContextConfigSchema = ContextConfigSchema.pick({
-  workspaceId: true,
-});
-
-export type ContextConfig = z.infer<typeof ContextConfigSchema>;
-export type ContextCategory = z.infer<
-  typeof ContextConfigSchema
->['categories'][number];
-export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
-export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
-export type ContextListItem = ContextDoc | ContextFile;
-export type ContextList = ContextListItem[];
-
 export type Chunk = {
   index: number;
   content: string;
 };
 
-export type ChunkSimilarity = {
-  chunk: number;
-  content: string;
-  distance: number | null;
-};
-
-export type FileChunkSimilarity = ChunkSimilarity & {
-  fileId: string;
-};
-
-export type DocChunkSimilarity = ChunkSimilarity & {
-  docId: string;
-};
-
-export type Embedding = {
-  /**
-   * The index of the embedding in the list of embeddings.
-   */
-  index: number;
-  content: string;
-  embedding: Array<number>;
-};
-
 export abstract class EmbeddingClient {
   async getFileEmbeddings(
     file: File,

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

@@ -1,7 +1,5 @@
 import { Readable } from 'node:stream';
 
-import { PrismaClient } from '@prisma/client';
-
 import { readBufferWithLimit } from '../../../base';
 import { MAX_EMBEDDABLE_SIZE } from './types';
 
@@ -17,17 +15,6 @@ export class GqlSignal implements AsyncDisposable {
   }
 }
 
-export async function checkEmbeddingAvailable(
-  db: PrismaClient
-): Promise<boolean> {
-  const [{ count }] = await db.$queryRaw<
-    {
-      count: number;
-    }[]
-  >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_context_embeddings', 'ai_workspace_embeddings')`;
-  return Number(count) === 2;
-}
-
 export function readStream(
   readable: Readable,
   maxSize = MAX_EMBEDDABLE_SIZE