From 84667b34406b0cb78b1f21d45faf7bf4214c5366 Mon Sep 17 00:00:00 2001
From: darkskygit <darksky.graphite@toeverything.info>
Date: Wed, 21 May 2025 10:51:35 +0000
Subject: [PATCH] feat(server): workspace embedding status count with files
 (#12420)

fix AI-32
fix AI-132
---
 .../server/src/models/copilot-workspace.ts    | 31 +++++++++++++++++++
 .../src/plugins/copilot/context/resolver.ts   | 13 ++++----
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/packages/backend/server/src/models/copilot-workspace.ts b/packages/backend/server/src/models/copilot-workspace.ts
index 9dcb0703aa..45f23b6900 100644
--- a/packages/backend/server/src/models/copilot-workspace.ts
+++ b/packages/backend/server/src/models/copilot-workspace.ts
@@ -144,6 +144,37 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
     return docIds.filter(id => ignored.has(id));
   }
 
+  @Transactional()
+  async getWorkspaceEmbeddingStatus(workspaceId: string) {
+    const ignoredDocIds = (await this.listIgnoredDocIds(workspaceId)).map(
+      d => d.docId
+    );
+    const snapshotCondition = {
+      workspaceId,
+      AND: [
+        { id: { notIn: ignoredDocIds } },
+        { id: { not: workspaceId } },
+        { id: { not: { contains: '$' } } },
+      ],
+    };
+
+    const [docTotal, docEmbedded, fileTotal, fileEmbedded] = await Promise.all([
+      this.db.snapshot.count({ where: snapshotCondition }),
+      this.db.snapshot.count({
+        where: { ...snapshotCondition, embedding: { some: {} } },
+      }),
+      this.db.aiWorkspaceFiles.count({ where: { workspaceId } }),
+      this.db.aiWorkspaceFiles.count({
+        where: { workspaceId, embeddings: { some: {} } },
+      }),
+    ]);
+
+    return {
+      total: docTotal + fileTotal,
+      embedded: docEmbedded + fileEmbedded,
+    };
+  }
+
   // ================ embeddings ================
 
   async checkEmbeddingAvailable(): Promise<boolean> {
diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts
index b36748d9aa..cb7ad71ed6 100644
--- a/packages/backend/server/src/plugins/copilot/context/resolver.ts
+++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts
@@ -13,7 +13,6 @@ import {
   ResolveField,
   Resolver,
 } from '@nestjs/graphql';
-import { PrismaClient } from '@prisma/client';
 import type { Request } from 'express';
 import { SafeIntResolver } from 'graphql-scalars';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
@@ -241,12 +240,12 @@ class ContextMatchedDocChunk implements DocChunkSimilarity {
 @Resolver(() => CopilotType)
 export class CopilotContextRootResolver {
   constructor(
-    private readonly db: PrismaClient,
     private readonly ac: AccessController,
     private readonly event: EventBus,
     private readonly mutex: RequestMutex,
     private readonly chatSession: ChatSessionService,
-    private readonly context: CopilotContextService
+    private readonly context: CopilotContextService,
+    private readonly models: Models
   ) {}
 
   private async checkChatSession(
@@ -369,10 +368,10 @@ export class CopilotContextRootResolver {
       .assert('Workspace.Copilot');
 
     if (this.context.canEmbedding) {
-      const total = await this.db.snapshot.count({ where: { workspaceId } });
-      const embedded = await this.db.snapshot.count({
-        where: { workspaceId, embedding: { some: {} } },
-      });
+      const { total, embedded } =
+        await this.models.copilotWorkspace.getWorkspaceEmbeddingStatus(
+          workspaceId
+        );
       return { total, embedded };
     }