Repository: toeverything/AFFiNE
Pull request #11804: feat(server): workspace file embedding & ignored docs model impl
Author: login: darkskygit
Labels: name: app:server
name: test
description: Related to test cases
Created: 2025-04-18T07:52:06Z | Merged: 2025-04-21T05:53:33Z

fix AI-30
fix AI-31



Diff:
diff --git a/packages/backend/server/migrations/20250417062046_workspace_level_embed_file_doc/migration.sql b/packages/backend/server/migrations/20250417062046_workspace_level_embed_file_doc/migration.sql
new file mode 100644
index 0000000000000..85e311ab8b9c9
--- /dev/null
+++ b/packages/backend/server/migrations/20250417062046_workspace_level_embed_file_doc/migration.sql
@@ -0,0 +1,81 @@
+-- CreateTable
+CREATE TABLE "ai_workspace_ignored_docs" (
+    "workspace_id" VARCHAR NOT NULL,
+    "doc_id" VARCHAR NOT NULL,
+    "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "ai_workspace_ignored_docs_pkey" PRIMARY KEY ("workspace_id","doc_id")
+);
+
+-- CreateTable
+CREATE TABLE "ai_workspace_files" (
+    "workspace_id" VARCHAR NOT NULL,
+    "file_id" VARCHAR NOT NULL,
+    "file_name" VARCHAR NOT NULL,
+    "mime_type" VARCHAR NOT NULL,
+    "size" INTEGER NOT NULL,
+    "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "ai_workspace_files_pkey" PRIMARY KEY ("workspace_id","file_id")
+);
+
+-- AddForeignKey
+ALTER TABLE "ai_workspace_ignored_docs" ADD CONSTRAINT "ai_workspace_ignored_docs_workspace_id_fkey" FOREIGN KEY ("workspace_id") REFERENCES "workspaces"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "ai_workspace_files" ADD CONSTRAINT "ai_workspace_files_workspace_id_fkey" FOREIGN KEY ("workspace_id") REFERENCES "workspaces"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+DO $$
+DECLARE error_message TEXT;
+BEGIN -- check if pgvector extension is installed
+  IF NOT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN
+    BEGIN
+      -- CreateExtension
+      CREATE EXTENSION IF NOT EXISTS "vector";
+    EXCEPTION
+      WHEN OTHERS THEN
+        -- if not found and cannot create extension, raise the exception
+        error_message := 'pgvector extension not found.' || E'\n' ||
+        '****************************************************************************' || E'\n' ||
+        '*                                                                          *' || E'\n' ||
+        '*   NOTICE: From AFFiNE 0.20 onwards, the copilot module will depend       *' || E'\n' ||
+        '*           on pgvector.                                                   *' || E'\n' ||
+        '*                                                                          *' || E'\n' ||
+        '*   1. If you are using the official PostgreSQL Docker container,          *' || E'\n' ||
+        '*      please switch to the pgvector/pgvector:pg${VERSION} container,      *' || E'\n' ||
+        '*      where ${VERSION} is the major version of your PostgreSQL container. *' || E'\n' ||
+        '*                                                                          *' || E'\n' ||
+        '*   2. If you are using a self-installed PostgreSQL, please follow the     *' || E'\n' ||
+        '*      the official pgvector installation guide to install it into your    *' || E'\n' ||
+        '*      database: https://github.com/pgvector/pgvector?tab=readme-ov-       *' || E'\n' ||
+        '*      file#installation-notes---linux-and-mac                             *' || E'\n' ||
+        '*                                                                          *' || E'\n' ||
+        '****************************************************************************';
+
+        RAISE WARNING '%', error_message;
+    END;
+  END IF;
+  -- check again, initialize the tables if the extension is installed
+  IF EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN
+    -- CreateTable
+    CREATE TABLE "ai_workspace_file_embeddings" (
+        "workspace_id" VARCHAR NOT NULL,
+        "file_id" VARCHAR NOT NULL,
+        "chunk" INTEGER NOT NULL,
+        "content" VARCHAR NOT NULL,
+        "embedding" vector(1024) NOT NULL,
+        "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+        CONSTRAINT "ai_workspace_file_embeddings_pkey" PRIMARY KEY ("workspace_id","file_id")
+    );
+
+    -- CreateIndex
+    CREATE INDEX "ai_workspace_file_embeddings_idx" ON "ai_workspace_file_embeddings" USING hnsw (embedding vector_cosine_ops);
+
+    -- CreateIndex
+    CREATE UNIQUE INDEX "ai_workspace_file_embeddings_workspace_id_file_id_chunk_key" ON "ai_workspace_file_embeddings"("workspace_id", "file_id", "chunk");
+
+    -- AddForeignKey
+    ALTER TABLE "ai_workspace_file_embeddings" ADD CONSTRAINT "ai_workspace_file_embeddings_workspace_id_file_id_fkey" FOREIGN KEY ("workspace_id", "file_id") REFERENCES "ai_workspace_files"("workspace_id", "file_id") ON DELETE CASCADE ON UPDATE CASCADE;
+  END IF;
+END $$;
diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma
index 274a20fe3d995..d994758f61a0a 100644
--- a/packages/backend/server/schema.prisma
+++ b/packages/backend/server/schema.prisma
@@ -114,11 +114,13 @@ model Workspace {
   name               String?  @db.VarChar
   avatarKey          String?  @map("avatar_key") @db.VarChar
 
-  features       WorkspaceFeature[]
-  docs           WorkspaceDoc[]
-  permissions    WorkspaceUserRole[]
-  docPermissions WorkspaceDocUserRole[]
-  blobs          Blob[]
+  features               WorkspaceFeature[]
+  docs                   WorkspaceDoc[]
+  permissions            WorkspaceUserRole[]
+  docPermissions         WorkspaceDocUserRole[]
+  blobs                  Blob[]
+  AiWorkspaceIgnoredDocs AiWorkspaceIgnoredDocs[]
+  AiWorkspaceFiles       AiWorkspaceFiles[]
 
   @@map("workspaces")
 }
@@ -481,6 +483,53 @@ model AiWorkspaceEmbedding {
   @@map("ai_workspace_embeddings")
 }
 
+model AiWorkspaceIgnoredDocs {
+  workspaceId String @map("workspace_id") @db.VarChar
+  docId       String @map("doc_id") @db.VarChar
+
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
+
+  workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
+
+  @@id([workspaceId, docId])
+  @@map("ai_workspace_ignored_docs")
+}
+
+model AiWorkspaceFiles {
+  workspaceId String @map("workspace_id") @db.VarChar
+  fileId      String @map("file_id") @db.VarChar
+  fileName    String @map("file_name") @db.VarChar
+  mimeType    String @map("mime_type") @db.VarChar
+  size        Int    @db.Integer
+
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
+
+  workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
+
+  embeddings AiWorkspaceFileEmbedding[]
+
+  @@id([workspaceId, fileId])
+  @@map("ai_workspace_files")
+}
+
+model AiWorkspaceFileEmbedding {
+  workspaceId String                      @map("workspace_id") @db.VarChar
+  fileId      String                      @map("file_id") @db.VarChar
+  // a file can be divided into multiple chunks and embedded separately.
+  chunk       Int                         @db.Integer
+  content     String                      @db.VarChar
+  embedding   Unsupported("vector(1024)")
+
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
+
+  file AiWorkspaceFiles @relation(fields: [workspaceId, fileId], references: [workspaceId, fileId], onDelete: Cascade)
+
+  @@id([workspaceId, fileId])
+  @@unique([workspaceId, fileId, chunk])
+  @@index([embedding], map: "ai_workspace_file_embeddings_idx")
+  @@map("ai_workspace_file_embeddings")
+}
+
 enum AiJobStatus {
   pending
   running
diff --git a/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts
new file mode 100644
index 0000000000000..ab349c816d291
--- /dev/null
+++ b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts
@@ -0,0 +1,149 @@
+import { PrismaClient, User, Workspace } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { Config } from '../../base';
+import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace';
+import { UserModel } from '../../models/user';
+import { WorkspaceModel } from '../../models/workspace';
+import { createTestingModule, type TestingModule } from '../utils';
+
+interface Context {
+  config: Config;
+  module: TestingModule;
+  db: PrismaClient;
+  user: UserModel;
+  workspace: WorkspaceModel;
+  copilotWorkspace: CopilotWorkspaceConfigModel;
+}
+
+const test = ava as TestFn<Context>;
+
+test.before(async t => {
+  const module = await createTestingModule();
+  t.context.user = module.get(UserModel);
+  t.context.workspace = module.get(WorkspaceModel);
+  t.context.copilotWorkspace = module.get(CopilotWorkspaceConfigModel);
+  t.context.db = module.get(PrismaClient);
+  t.context.config = module.get(Config);
+  t.context.module = module;
+});
+
+let user: User;
+let workspace: Workspace;
+
+let docId = 'doc1';
+
+test.beforeEach(async t => {
+  await t.context.module.initTestingDB();
+  user = await t.context.user.create({
+    email: 'test@affine.pro',
+  });
+  workspace = await t.context.workspace.create(user.id);
+});
+
+test.after(async t => {
+  await t.context.module.close();
+});
+
+test('should manage copilot workspace ignored docs', async t => {
+  const ignoredDocs = await t.context.copilotWorkspace.listIgnoredDocs(
+    workspace.id
+  );
+  t.deepEqual(ignoredDocs, []);
+
+  {
+    const count = await t.context.copilotWorkspace.updateIgnoredDocs(
+      workspace.id,
+      [docId]
+    );
+    t.is(count, 1, 'should add ignored doc');
+
+    const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
+    t.deepEqual(ret, [docId], 'should return added doc');
+
+    const check = await t.context.copilotWorkspace.checkIgnoredDocs(
+      workspace.id,
+      [docId]
+    );
+    t.deepEqual(check, [docId], 'should return ignored docs in workspace');
+  }
+
+  {
+    const count = await t.context.copilotWorkspace.updateIgnoredDocs(
+      workspace.id,
+      [docId]
+    );
+    t.is(count, 1, 'should not add ignored doc again');
+
+    const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
+    t.deepEqual(ret, [docId], 'should not add ignored doc again');
+  }
+
+  {
+    const count = await t.context.copilotWorkspace.updateIgnoredDocs(
+      workspace.id,
+      ['new_doc']
+    );
+    t.is(count, 2, 'should add new ignored doc');
+
+    const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
+    t.deepEqual(ret, [docId, 'new_doc'], 'should add ignored doc');
+  }
+
+  {
+    await t.context.copilotWorkspace.updateIgnoredDocs(
+      workspace.id,
+      undefined,
+      [docId]
+    );
+
+    const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
+    t.deepEqual(ret, ['new_doc'], 'should remove ignored doc');
+  }
+});
+
+test('should insert and search embedding', async t => {
+  {
+    await t.context.copilotWorkspace.addWorkspaceFile(
+      workspace.id,
+      {
+        fileName: 'file1',
+        mimeType: 'text/plain',
+
+        size: 1,
+      },
+      [
+        {
+          index: 0,
+          content: 'content',
+          embedding: Array.from({ length: 1024 }, () => 1),
+        },
+      ]
+    );
+
+    {
+      const ret = await t.context.copilotWorkspace.matchWorkspaceFileEmbedding(
+        workspace.id,
+        Array.from({ length: 1024 }, () => 0.9),
+        1,
+        1
+      );
+      t.is(ret.length, 1);
+      t.is(ret[0].content, 'content');
+    }
+  }
+});
+
+test('should check embedding table', async t => {
+  {
+    const ret = await t.context.copilotWorkspace.checkEmbeddingAvailable();
+    t.true(ret, 'should return true when embedding table is available');
+  }
+
+  // {
+  //   await t.context.db
+  //     .$executeRaw`DROP TABLE IF EXISTS "ai_workspace_file_embeddings"`;
+  //   const ret = await t.context.copilotWorkspace.checkEmbeddingAvailable();
+  //   t.false(ret, 'should return false when embedding table is not available');
+  // }
+});
diff --git a/packages/backend/server/src/models/common/copilot.ts b/packages/backend/server/src/models/common/copilot.ts
index 2ba40e16ff64b..9fe2a56e6757d 100644
--- a/packages/backend/server/src/models/common/copilot.ts
+++ b/packages/backend/server/src/models/common/copilot.ts
@@ -105,3 +105,17 @@ export type FileChunkSimilarity = ChunkSimilarity & {
 export type DocChunkSimilarity = ChunkSimilarity & {
   docId: string;
 };
+
+export const CopilotWorkspaceFileSchema = z.object({
+  fileName: z.string(),
+  mimeType: z.string(),
+  size: z.number(),
+});
+
+export type CopilotWorkspaceFile = z.infer<
+  typeof CopilotWorkspaceFileSchema
+> & {
+  workspaceId: string;
+  fileId: string;
+  createdAt: Date;
+};
diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts
index 712527671d341..670ac2874207f 100644
--- a/packages/backend/server/src/models/copilot-context.ts
+++ b/packages/backend/server/src/models/copilot-context.ts
@@ -177,12 +177,12 @@ export class CopilotContextModel extends BaseModel {
     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 = ${contextId}
-    ORDER BY "distance" ASC
-    LIMIT ${topK};
-  `;
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_context_embeddings"
+      WHERE context_id = ${contextId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
     return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
@@ -198,11 +198,11 @@ export class CopilotContextModel extends BaseModel {
       false
     );
     await this.db.$executeRaw`
-        INSERT INTO "ai_workspace_embeddings"
-        ("workspace_id", "doc_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
-        ON CONFLICT (workspace_id, doc_id, chunk) DO UPDATE SET
-        embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
-      `;
+      INSERT INTO "ai_workspace_embeddings"
+      ("workspace_id", "doc_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
+      ON CONFLICT (workspace_id, doc_id, chunk) DO UPDATE SET
+      embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
+    `;
   }
 
   async matchWorkspaceEmbedding(
@@ -212,12 +212,12 @@ export class CopilotContextModel extends BaseModel {
     threshold: number
   ): Promise<DocChunkSimilarity[]> {
     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" = ${workspaceId}
-       ORDER BY "distance" ASC
-       LIMIT ${topK};
-     `;
+      SELECT "doc_id" as "docId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
+      FROM "ai_workspace_embeddings"
+      WHERE "workspace_id" = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
     return similarityChunks.filter(c => Number(c.distance) <= threshold);
   }
 
diff --git a/packages/backend/server/src/models/copilot-workspace.ts b/packages/backend/server/src/models/copilot-workspace.ts
new file mode 100644
index 0000000000000..8cfd0af254212
--- /dev/null
+++ b/packages/backend/server/src/models/copilot-workspace.ts
@@ -0,0 +1,156 @@
+import { randomUUID } from 'node:crypto';
+
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  @Transactional()
+  async updateIgnoredDocs(
+    workspaceId: string,
+    add: string[] = [],
+    remove: string[] = []
+  ) {
+    const removed = new Set(remove);
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r.filter(id => !removed.has(id)))
+    );
+    const added = add.filter(id => !ignored.has(id));
+
+    if (added.length) {
+      await this.db.aiWorkspaceIgnoredDocs.createMany({
+        data: added.map(docId => ({
+          workspaceId,
+          docId,
+        })),
+      });
+    }
+
+    if (removed.size) {
+      await this.db.aiWorkspaceIgnoredDocs.deleteMany({
+        where: {
+          workspaceId,
+          docId: {
+            in: Array.from(removed),
+          },
+        },
+      });
+    }
+
+    return added.length + ignored.size;
+  }
+
+  async listIgnoredDocs(workspaceId: string): Promise<string[]> {
+    const row = await this.db.aiWorkspaceIgnoredDocs.findMany({
+      where: {
+        workspaceId,
+      },
+      select: {
+        docId: true,
+      },
+    });
+    return row.map(r => r.docId);
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: Pick<CopilotWorkspaceFile, 'fileName' | 'mimeType' | 'size'>,
+    embeddings: Embedding[]
+  ): Promise<string> {
+    const fileId = randomUUID();
+    await this.db.aiWorkspaceFiles.create({
+      data: { ...file, workspaceId, fileId },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+    return fileId;
+  }
+
+  async listWorkspaceFiles(
+    workspaceId: string
+  ): Promise<CopilotWorkspaceFile[]> {
+    const files = await this.db.aiWorkspaceFiles.findMany({
+      where: {
+        workspaceId,
+      },
+    });
+    return files;
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    // embeddings will be removed by foreign key constraint
+    await this.db.aiWorkspaceFiles.deleteMany({
+      where: {
+        workspaceId,
+        fileId,
+      },
+    });
+  }
+}
diff --git a/packages/backend/server/src/models/index.ts b/packages/backend/server/src/models/index.ts
index e0cbb0754b51a..aea37ad0c05c7 100644
--- a/packages/backend/server/src/models/index.ts
+++ b/packages/backend/server/src/models/index.ts
@@ -11,6 +11,7 @@ import { AppConfigModel } from './config';
 import { CopilotContextModel } from './copilot-context';
 import { CopilotJobModel } from './copilot-job';
 import { CopilotSessionModel } from './copilot-session';
+import { CopilotWorkspaceConfigModel } from './copilot-workspace';
 import { DocModel } from './doc';
 import { DocUserModel } from './doc-user';
 import { FeatureModel } from './feature';
@@ -44,6 +45,7 @@ const MODELS = {
   userSettings: UserSettingsModel,
   copilotSession: CopilotSessionModel,
   copilotContext: CopilotContextModel,
+  copilotWorkspace: CopilotWorkspaceConfigModel,
   copilotJob: CopilotJobModel,
   appConfig: AppConfigModel,
 };


Reviews:
user: login: graphite-app[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T08:34:13Z
user: login: fengmk2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-18T14:57:42Z
user: login: fengmk2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-18T14:59:01Z
user: login: fengmk2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-18T15:01:14Z
user: login: darkskygit
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-19T17:58:10Z
user: login: darkskygit
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-19T17:58:24Z
user: login: darkskygit
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-19T18:17:32Z
user: login: yoyoyohamapi
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-21T03:44:03Z
user: login: graphite-app[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-21T04:34:59Z

Review comments:
diff_hunk: @@ -0,0 +1,215 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  CopilotWorkspaceConfigPayloadSchema,
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  private async init(workspaceId: string) {
+    const count = await this.db.aiWorkspaceConfig.count({
+      where: { workspaceId },
+    });
+    if (!count) {
+      await this.db.aiWorkspaceConfig.create({
+        data: {
+          workspaceId,
+          payload: {},
+        },
+      });
+    }
+  }
+
+  @Transactional()
+  async addIgnoredDocs(workspaceId: string, docIds: string[]) {
+    await this.init(workspaceId);
+
+    const ignored = await this.getIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+    const added = docIds.filter(id => !ignored.has(id));
+
+    if (added.length) {
+      await this.db.aiWorkspaceConfig.update({
+        where: {
+          workspaceId,
+        },
+        data: {
+          ignoreDocs: {
+            push: added,
+          },
+          updatedAt: new Date(),
+        },
+      });
+    }
+    return added.length;
+  }
+
+  async getIgnoredDocs(workspaceId: string) {
+    const row = await this.db.aiWorkspaceConfig.findFirst({
+      where: {
+        workspaceId,
+      },
+      select: {
+        ignoreDocs: true,
+      },
+    });
+    return row?.ignoreDocs ?? [];
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.getIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  async removeIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const removed = new Set(docIds);
+
+    const ignored = await this.getIgnoredDocs(workspaceId).then(r =>
+      r.filter(id => !removed.has(id))
+    );
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        ignoreDocs: ignored,
+        updatedAt: new Date(),
+      },
+    });
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  private async getPayload(workspaceId: string) {
+    const payload = await this.db.aiWorkspaceConfig
+      .findFirst({
+        where: {
+          workspaceId,
+        },
+        select: {
+          payload: true,
+        },
+      })
+      .then(r => r?.payload ?? { files: {} });
+
+    return CopilotWorkspaceConfigPayloadSchema.parse(payload);
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: CopilotWorkspaceFile,
+    embeddings: Embedding[]
+  ) {
+    await this.init(workspaceId);
+
+    const { fileId, ...fileData } = file;
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (payload.files[fileId]) {
+      return;
+    }
+    payload.files[fileId] = fileData;
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (!payload.files[fileId]) {
+      return;
+    }
+    delete payload.files[fileId];
+    const files = Object.keys(payload.files);
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        files,
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    await this.db.$executeRaw`
+        DELETE FROM "ai_workspace_file_embeddings"
+        WHERE "workspaceId" = ${workspaceId} AND "file_id" = ${fileId};
+    `;
path: packages/backend/server/src/models/copilot-workspace.ts
user: login: graphite-app[bot]
body: There's a column name mismatch in the DELETE query. The query uses `"workspaceId"` but the actual column name in the database schema is `"workspace_id"` (as defined in the migration and model). This will cause the query to fail.

The correct query should be:

```sql
DELETE FROM "ai_workspace_file_embeddings"
WHERE "workspace_id" = ${workspaceId} AND "file_id" = ${fileId};
```
```suggestion
    await this.db.$executeRaw`
        DELETE FROM "ai_workspace_file_embeddings"
        WHERE "workspace_id" = ${workspaceId} AND "file_id" = ${fileId};
    `;
```
  

*Spotted by [Diamond](https://app.graphite.dev/diamond/?org=toeverything&ref=ai-review-comment)*<i class='graphite__hidden'><br /><br />Is this helpful? React 👍 or 👎 to let us know.</i>
created_at: 2025-04-18T08:34:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,209 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  CopilotWorkspaceConfigPayloadSchema,
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  private async init(workspaceId: string) {
+    const count = await this.db.aiWorkspaceConfig.count({
+      where: { workspaceId },
+    });
+    if (!count) {
+      await this.db.aiWorkspaceConfig.create({
+        data: {
+          workspaceId,
+          payload: {},
+        },
+      });
+    }
+  }
+
+  @Transactional()
+  async updateIgnoredDocs(
+    workspaceId: string,
+    add: string[] = [],
+    remove: string[] = []
+  ) {
+    await this.init(workspaceId);
+
+    const removed = new Set(remove);
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r.filter(id => !removed.has(id)))
+    );
+    const added = add.filter(id => !ignored.has(id));
+    const finalIgnored = [...ignored, ...added];
+
+    if (added.length || removed.size) {
+      await this.db.aiWorkspaceConfig.update({
+        where: {
+          workspaceId,
+        },
+        data: {
+          ignoreDocs: finalIgnored,
+          updatedAt: new Date(),
+        },
+      });
+    }
+    return finalIgnored.length;
+  }
+
+  async listIgnoredDocs(workspaceId: string) {
+    const row = await this.db.aiWorkspaceConfig.findFirst({
+      where: {
+        workspaceId,
+      },
+      select: {
+        ignoreDocs: true,
+      },
+    });
+    return row?.ignoreDocs ?? [];
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  private async getPayload(workspaceId: string) {
+    const payload = await this.db.aiWorkspaceConfig
+      .findFirst({
+        where: {
+          workspaceId,
+        },
+        select: {
+          payload: true,
+        },
+      })
+      .then(r => r?.payload ?? { files: {} });
+
+    return CopilotWorkspaceConfigPayloadSchema.parse(payload);
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: CopilotWorkspaceFile,
+    embeddings: Embedding[]
+  ) {
+    await this.init(workspaceId);
+
+    const { fileId, ...fileData } = file;
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (payload.files[fileId]) {
+      return;
+    }
+    payload.files[fileId] = fileData;
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+  }
+
+  async listWorkspaceFiles(
+    workspaceId: string
+  ): Promise<CopilotWorkspaceFile[]> {
+    const row = await this.getPayload(workspaceId);
+    return Object.entries(row.files || {}).map(([fileId, file]) => ({
+      fileId,
+      ...file,
+    }));
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (!payload.files[fileId]) {
+      return;
+    }
+    delete payload.files[fileId];
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    await this.db.$executeRaw`
+        DELETE FROM "ai_workspace_file_embeddings"
+        WHERE "workspaceId" = ${workspaceId} AND "file_id" = ${fileId};
path: packages/backend/server/src/models/copilot-workspace.ts
user: login: fengmk2
body: should be `workspace_id`?
created_at: 2025-04-18T14:57:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,209 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  CopilotWorkspaceConfigPayloadSchema,
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  private async init(workspaceId: string) {
+    const count = await this.db.aiWorkspaceConfig.count({
+      where: { workspaceId },
+    });
+    if (!count) {
+      await this.db.aiWorkspaceConfig.create({
+        data: {
+          workspaceId,
+          payload: {},
+        },
+      });
+    }
+  }
+
+  @Transactional()
+  async updateIgnoredDocs(
+    workspaceId: string,
+    add: string[] = [],
+    remove: string[] = []
+  ) {
+    await this.init(workspaceId);
+
+    const removed = new Set(remove);
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r.filter(id => !removed.has(id)))
+    );
+    const added = add.filter(id => !ignored.has(id));
+    const finalIgnored = [...ignored, ...added];
+
+    if (added.length || removed.size) {
+      await this.db.aiWorkspaceConfig.update({
+        where: {
+          workspaceId,
+        },
+        data: {
+          ignoreDocs: finalIgnored,
+          updatedAt: new Date(),
+        },
+      });
+    }
+    return finalIgnored.length;
+  }
+
+  async listIgnoredDocs(workspaceId: string) {
+    const row = await this.db.aiWorkspaceConfig.findFirst({
+      where: {
+        workspaceId,
+      },
+      select: {
+        ignoreDocs: true,
+      },
+    });
+    return row?.ignoreDocs ?? [];
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  private async getPayload(workspaceId: string) {
+    const payload = await this.db.aiWorkspaceConfig
+      .findFirst({
+        where: {
+          workspaceId,
+        },
+        select: {
+          payload: true,
+        },
+      })
+      .then(r => r?.payload ?? { files: {} });
+
+    return CopilotWorkspaceConfigPayloadSchema.parse(payload);
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: CopilotWorkspaceFile,
+    embeddings: Embedding[]
+  ) {
+    await this.init(workspaceId);
+
+    const { fileId, ...fileData } = file;
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (payload.files[fileId]) {
+      return;
+    }
+    payload.files[fileId] = fileData;
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+  }
+
+  async listWorkspaceFiles(
+    workspaceId: string
+  ): Promise<CopilotWorkspaceFile[]> {
+    const row = await this.getPayload(workspaceId);
+    return Object.entries(row.files || {}).map(([fileId, file]) => ({
+      fileId,
+      ...file,
+    }));
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (!payload.files[fileId]) {
+      return;
+    }
+    delete payload.files[fileId];
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    await this.db.$executeRaw`
path: packages/backend/server/src/models/copilot-workspace.ts
user: login: fengmk2
body: Do all table queries involving `vector` type field require handwritten SQL?
created_at: 2025-04-18T14:59:00Z
author_association: COLLABORATOR
diff_hunk: @@ -441,6 +441,22 @@ model AiContext {
   @@map("ai_contexts")
 }
 
+model AiWorkspaceConfig {
+  workspaceId String   @map("workspace_id") @db.VarChar
+  // ignored embed docs in the workspace
+  ignoreDocs  String[] @unique @map("ignore_docs") @db.VarChar
path: packages/backend/server/schema.prisma
user: login: fengmk2
body: How large will the data volume be for ignoreDocs? Why not create a new association table to do it?
created_at: 2025-04-18T15:01:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,209 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  CopilotWorkspaceConfigPayloadSchema,
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  private async init(workspaceId: string) {
+    const count = await this.db.aiWorkspaceConfig.count({
+      where: { workspaceId },
+    });
+    if (!count) {
+      await this.db.aiWorkspaceConfig.create({
+        data: {
+          workspaceId,
+          payload: {},
+        },
+      });
+    }
+  }
+
+  @Transactional()
+  async updateIgnoredDocs(
+    workspaceId: string,
+    add: string[] = [],
+    remove: string[] = []
+  ) {
+    await this.init(workspaceId);
+
+    const removed = new Set(remove);
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r.filter(id => !removed.has(id)))
+    );
+    const added = add.filter(id => !ignored.has(id));
+    const finalIgnored = [...ignored, ...added];
+
+    if (added.length || removed.size) {
+      await this.db.aiWorkspaceConfig.update({
+        where: {
+          workspaceId,
+        },
+        data: {
+          ignoreDocs: finalIgnored,
+          updatedAt: new Date(),
+        },
+      });
+    }
+    return finalIgnored.length;
+  }
+
+  async listIgnoredDocs(workspaceId: string) {
+    const row = await this.db.aiWorkspaceConfig.findFirst({
+      where: {
+        workspaceId,
+      },
+      select: {
+        ignoreDocs: true,
+      },
+    });
+    return row?.ignoreDocs ?? [];
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  private async getPayload(workspaceId: string) {
+    const payload = await this.db.aiWorkspaceConfig
+      .findFirst({
+        where: {
+          workspaceId,
+        },
+        select: {
+          payload: true,
+        },
+      })
+      .then(r => r?.payload ?? { files: {} });
+
+    return CopilotWorkspaceConfigPayloadSchema.parse(payload);
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: CopilotWorkspaceFile,
+    embeddings: Embedding[]
+  ) {
+    await this.init(workspaceId);
+
+    const { fileId, ...fileData } = file;
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (payload.files[fileId]) {
+      return;
+    }
+    payload.files[fileId] = fileData;
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+  }
+
+  async listWorkspaceFiles(
+    workspaceId: string
+  ): Promise<CopilotWorkspaceFile[]> {
+    const row = await this.getPayload(workspaceId);
+    return Object.entries(row.files || {}).map(([fileId, file]) => ({
+      fileId,
+      ...file,
+    }));
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (!payload.files[fileId]) {
+      return;
+    }
+    delete payload.files[fileId];
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    await this.db.$executeRaw`
+        DELETE FROM "ai_workspace_file_embeddings"
+        WHERE "workspaceId" = ${workspaceId} AND "file_id" = ${fileId};
path: packages/backend/server/src/models/copilot-workspace.ts
user: login: darkskygit
body: fixed in next pr, let me cp it later
created_at: 2025-04-19T17:58:10Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,209 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+import {
+  CopilotWorkspaceConfigPayloadSchema,
+  type CopilotWorkspaceFile,
+  type Embedding,
+  FileChunkSimilarity,
+} from './common';
+
+@Injectable()
+export class CopilotWorkspaceConfigModel extends BaseModel {
+  private async init(workspaceId: string) {
+    const count = await this.db.aiWorkspaceConfig.count({
+      where: { workspaceId },
+    });
+    if (!count) {
+      await this.db.aiWorkspaceConfig.create({
+        data: {
+          workspaceId,
+          payload: {},
+        },
+      });
+    }
+  }
+
+  @Transactional()
+  async updateIgnoredDocs(
+    workspaceId: string,
+    add: string[] = [],
+    remove: string[] = []
+  ) {
+    await this.init(workspaceId);
+
+    const removed = new Set(remove);
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r.filter(id => !removed.has(id)))
+    );
+    const added = add.filter(id => !ignored.has(id));
+    const finalIgnored = [...ignored, ...added];
+
+    if (added.length || removed.size) {
+      await this.db.aiWorkspaceConfig.update({
+        where: {
+          workspaceId,
+        },
+        data: {
+          ignoreDocs: finalIgnored,
+          updatedAt: new Date(),
+        },
+      });
+    }
+    return finalIgnored.length;
+  }
+
+  async listIgnoredDocs(workspaceId: string) {
+    const row = await this.db.aiWorkspaceConfig.findFirst({
+      where: {
+        workspaceId,
+      },
+      select: {
+        ignoreDocs: true,
+      },
+    });
+    return row?.ignoreDocs ?? [];
+  }
+
+  @Transactional()
+  async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
+    const ignored = await this.listIgnoredDocs(workspaceId).then(
+      r => new Set(r)
+    );
+
+    return docIds.filter(id => ignored.has(id));
+  }
+
+  // ================ embeddings ================
+
+  async checkEmbeddingAvailable(): Promise<boolean> {
+    const [{ count }] = await this.db.$queryRaw<
+      { count: number }[]
+    >`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_file_embeddings')`;
+    return Number(count) === 1;
+  }
+
+  private processEmbeddings(
+    workspaceId: string,
+    fileId: string,
+    embeddings: Embedding[]
+  ) {
+    const groups = embeddings.map(e =>
+      [
+        workspaceId,
+        fileId,
+        e.index,
+        e.content,
+        Prisma.raw(`'[${e.embedding.join(',')}]'`),
+      ].filter(v => v !== undefined)
+    );
+    return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
+  }
+
+  private async getPayload(workspaceId: string) {
+    const payload = await this.db.aiWorkspaceConfig
+      .findFirst({
+        where: {
+          workspaceId,
+        },
+        select: {
+          payload: true,
+        },
+      })
+      .then(r => r?.payload ?? { files: {} });
+
+    return CopilotWorkspaceConfigPayloadSchema.parse(payload);
+  }
+
+  @Transactional()
+  async addWorkspaceFile(
+    workspaceId: string,
+    file: CopilotWorkspaceFile,
+    embeddings: Embedding[]
+  ) {
+    await this.init(workspaceId);
+
+    const { fileId, ...fileData } = file;
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (payload.files[fileId]) {
+      return;
+    }
+    payload.files[fileId] = fileData;
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    const values = this.processEmbeddings(workspaceId, fileId, embeddings);
+    await this.db.$executeRaw`
+        INSERT INTO "ai_workspace_file_embeddings"
+        ("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
+        ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
+    `;
+  }
+
+  async listWorkspaceFiles(
+    workspaceId: string
+  ): Promise<CopilotWorkspaceFile[]> {
+    const row = await this.getPayload(workspaceId);
+    return Object.entries(row.files || {}).map(([fileId, file]) => ({
+      fileId,
+      ...file,
+    }));
+  }
+
+  async matchWorkspaceFileEmbedding(
+    workspaceId: string,
+    embedding: number[],
+    topK: number,
+    threshold: number
+  ): Promise<FileChunkSimilarity[]> {
+    const similarityChunks = await this.db.$queryRaw<
+      Array<FileChunkSimilarity>
+    >`
+      SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance" 
+      FROM "ai_workspace_file_embeddings"
+      WHERE workspace_id = ${workspaceId}
+      ORDER BY "distance" ASC
+      LIMIT ${topK};
+    `;
+    return similarityChunks.filter(c => Number(c.distance) <= threshold);
+  }
+
+  async removeWorkspaceFile(workspaceId: string, fileId: string) {
+    const payload = await this.getPayload(workspaceId);
+    if (!payload.files) {
+      payload.files = {};
+    }
+    if (!payload.files[fileId]) {
+      return;
+    }
+    delete payload.files[fileId];
+
+    await this.db.aiWorkspaceConfig.update({
+      where: {
+        workspaceId,
+      },
+      data: {
+        payload,
+        updatedAt: new Date(),
+      },
+    });
+
+    await this.db.$executeRaw`
path: packages/backend/server/src/models/copilot-workspace.ts
user: login: darkskygit
body: yes, prisma not have vector support yet
created_at: 2025-04-19T17:58:24Z
author_association: MEMBER
diff_hunk: @@ -441,6 +441,22 @@ model AiContext {
   @@map("ai_contexts")
 }
 
+model AiWorkspaceConfig {
+  workspaceId String   @map("workspace_id") @db.VarChar
+  // ignored embed docs in the workspace
+  ignoreDocs  String[] @unique @map("ignore_docs") @db.VarChar
path: packages/backend/server/schema.prisma
user: login: darkskygit
body: since it can only restrict specific docs (server cannot know the associations like collection/tag), it is foreseeable that the data volume will be small and the write frequency will be low
if you are still concerned about performance, I will split file and doc into separate tables later
created_at: 2025-04-19T18:17:32Z
author_association: MEMBER
diff_hunk: @@ -105,3 +105,17 @@ export type FileChunkSimilarity = ChunkSimilarity & {
 export type DocChunkSimilarity = ChunkSimilarity & {
   docId: string;
 };
+
+export const CopilotWorkspaceFileSchema = z.object({
+  fileName: z.string(),
+  mimeType: z.string(),
+  size: z.number(),
+});
+
+export type CopilotWorkspaceFile = z.infer<
+  typeof CopilotWorkspaceFileSchema
+> & {
path: packages/backend/server/src/models/common/copilot.ts
user: login: graphite-app[bot]
body: It appears that the `z` object is being used to define a schema without an import statement. Please add the appropriate import for the zod library:

```typescript
import { z } from 'zod';
```

This will ensure proper type checking and prevent potential runtime errors. The schema definition itself looks good, but explicit imports help maintain code clarity and prevent "magic globals" that can cause confusion for other developers.

*Spotted by [Diamond](https://app.graphite.dev/diamond/?org=toeverything&ref=ai-review-comment) (based on [custom rules](https://app.graphite.dev/diamond/?org=toeverything&ref=ai-review-comment&view=settings))*<i class='graphite__hidden'><br /><br />Is this helpful? React 👍 or 👎 to let us know.</i>
created_at: 2025-04-21T04:34:58Z
author_association: CONTRIBUTOR

PR comments:
user: login: darkskygit
created_at: 2025-04-18T07:52:24Z
body: * **#11859** <a href="https://app.graphite.dev/github/pr/toeverything/AFFiNE/11859?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a>
* **#11809** <a href="https://app.graphite.dev/github/pr/toeverything/AFFiNE/11809?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a>
* **#11804** <a href="https://app.graphite.dev/github/pr/toeverything/AFFiNE/11804?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a> 👈 <a href="https://app.graphite.dev/github/pr/toeverything/AFFiNE/11804?utm_source=stack-comment-view-in-graphite" target="_blank">(View in Graphite)</a>
* `canary`
<hr />

## How to use the Graphite Merge Queue

Add either label to this PR to merge it via the merge queue:
- <kbd>merge</kbd> - adds this PR to the back of the merge queue
- <kbd>hotfix</kbd> - for urgent hot fixes, skip the queue and merge this PR next

You must have a Graphite account in order to use the merge queue. Sign up using [this link](https://app.graphite.dev/invite/github/toeverything?ref=merge-queue-instructions-in-stack-comment&prId=11804).

<h2></h2>
<sub>

An organization admin has enabled the [Graphite Merge Queue](https://graphite.dev/docs/graphite-merge-queue) in this repository.

</sub>

<sub>

Please do not merge from GitHub as this will restart CI on PRs being processed by the merge queue.

</sub>

<h2></h2>

This stack of pull requests is managed by <a href="https://graphite.dev?utm-source=stack-comment"><b>Graphite</b></a>. Learn more about <a href="https://stacking.dev/?utm_source=stack-comment">stacking</a>.
<!-- Current dependencies on/for this PR: -->
author_association: MEMBER
user: login: codecov[bot]
created_at: 2025-04-18T07:57:34Z
body: ## [Codecov](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) Report
Attention: Patch coverage is `89.94709%` with `19 lines` in your changes missing coverage. Please review.
> Project coverage is 55.83%. Comparing base [(`93b7c28`)](https://app.codecov.io/gh/toeverything/AFFiNE/commit/93b7c288cbbc481ecfc9400ac06cb3b1d05df06e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) to head [(`f2adb9f`)](https://app.codecov.io/gh/toeverything/AFFiNE/commit/f2adb9f72c6df2e7a0ab48b4b371c928f198a43f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything).
> Report is 1 commits behind head on canary.

| [Files with missing lines](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) | Patch % | Lines |
|---|---|---|
| [...ges/backend/server/src/models/copilot-workspace.ts](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804?src=pr&el=tree&filepath=packages%2Fbackend%2Fserver%2Fsrc%2Fmodels%2Fcopilot-workspace.ts&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything#diff-cGFja2FnZXMvYmFja2VuZC9zZXJ2ZXIvc3JjL21vZGVscy9jb3BpbG90LXdvcmtzcGFjZS50cw==) | 87.82% | [19 Missing :warning: ](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           canary   #11804      +/-   ##
==========================================
+ Coverage   55.78%   55.83%   +0.04%     
==========================================
  Files        2532     2533       +1     
  Lines      113441   113613     +172     
  Branches    18407    18421      +14     
==========================================
+ Hits        63287    63432     +145     
- Misses      48361    48381      +20     
- Partials     1793     1800       +7     
```

| [Flag](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) | Coverage Δ | |
|---|---|---|
| [server-test](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) | `79.21% <89.94%> (+0.04%)` | :arrow_up: |
| [unittest](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything) | `31.12% <ø> (-0.01%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/toeverything/AFFiNE/pull/11804?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=toeverything).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
author_association: NONE
user: login: graphite-app[bot]
created_at: 2025-04-21T05:32:51Z
body: ### Merge activity

* **Apr 21, 1:32 AM EDT**: `darkskygit` added this pull request to the [Graphite merge queue](https://app.graphite.dev/merges?org=toeverything&repo=AFFiNE).
* **Apr 21, 5:47 AM UTC**: The [Graphite merge queue](https://app.graphite.dev/merges?org=toeverything&repo=AFFiNE) couldn't merge this PR because **it was not satisfying all requirements** (Failed CI: '3, 2, 1 Launch', 'Native Unit Test').
* **Apr 21, 1:53 AM EDT**: `darkskygit` added this pull request to the [Graphite merge queue](https://app.graphite.dev/merges?org=toeverything&repo=AFFiNE).
* **Apr 21, 1:53 AM EDT**: Merged by the [Graphite merge queue](https://app.graphite.dev/merges?org=toeverything&repo=AFFiNE).

author_association: CONTRIBUTOR