luulam/AFFiNE
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat(server): split embedding client (#12809)

Browse Source
This commit is contained in:
DarkSky
2025-06-13 12:37:05 +08:00
committed by GitHub
parent 1d4bc81e90
commit e98f035f97
12 changed files with 32 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
packages/backend/server/src/__tests__/copilot.e2e.ts
View File

@@ -11,11 +11,11 @@ import { JobQueue } from '../base';
import { ConfigModule } from '../base/config'; import { ConfigModule } from '../base/config';
import { AuthService } from '../core/auth'; import { AuthService } from '../core/auth';
import { DocReader } from '../core/doc'; import { DocReader } from '../core/doc';
import { CopilotContextService } from '../plugins/copilot/context';
import { import {
CopilotContextDocJob, CopilotEmbeddingJob,
CopilotContextService, MockEmbeddingClient,
} from '../plugins/copilot/context'; } from '../plugins/copilot/embedding';
import { MockEmbeddingClient } from '../plugins/copilot/context/embedding';
import { prompts, PromptService } from '../plugins/copilot/prompt'; import { prompts, PromptService } from '../plugins/copilot/prompt';
import { import {
CopilotProviderFactory, CopilotProviderFactory,
@@ -65,7 +65,7 @@ const test = ava as TestFn<{
app: TestingApp; app: TestingApp;
db: PrismaClient; db: PrismaClient;
context: CopilotContextService; context: CopilotContextService;
jobs: CopilotContextDocJob; jobs: CopilotEmbeddingJob;
prompt: PromptService; prompt: PromptService;
factory: CopilotProviderFactory; factory: CopilotProviderFactory;
storage: CopilotStorage; storage: CopilotStorage;
@@ -115,7 +115,7 @@ test.before(async t => {
const context = app.get(CopilotContextService); const context = app.get(CopilotContextService);
const prompt = app.get(PromptService); const prompt = app.get(PromptService);
const storage = app.get(CopilotStorage); const storage = app.get(CopilotStorage);
const jobs = app.get(CopilotContextDocJob); const jobs = app.get(CopilotEmbeddingJob);
t.context.app = app; t.context.app = app;
t.context.db = db; t.context.db = db;

12
packages/backend/server/src/__tests__/copilot.spec.ts
View File

@@ -13,11 +13,11 @@ import { AuthService } from '../core/auth';
import { QuotaModule } from '../core/quota'; import { QuotaModule } from '../core/quota';
import { ContextCategories, WorkspaceModel } from '../models'; import { ContextCategories, WorkspaceModel } from '../models';
import { CopilotModule } from '../plugins/copilot'; import { CopilotModule } from '../plugins/copilot';
import { CopilotContextService } from '../plugins/copilot/context';
import { import {
CopilotContextDocJob, CopilotEmbeddingJob,
CopilotContextService, MockEmbeddingClient,
} from '../plugins/copilot/context'; } from '../plugins/copilot/embedding';
import { MockEmbeddingClient } from '../plugins/copilot/context/embedding';
import { prompts, PromptService } from '../plugins/copilot/prompt'; import { prompts, PromptService } from '../plugins/copilot/prompt';
import { import {
CopilotProviderFactory, CopilotProviderFactory,
@@ -69,7 +69,7 @@ const test = ava as TestFn<{
workspaceEmbedding: CopilotWorkspaceService; workspaceEmbedding: CopilotWorkspaceService;
factory: CopilotProviderFactory; factory: CopilotProviderFactory;
session: ChatSessionService; session: ChatSessionService;
jobs: CopilotContextDocJob; jobs: CopilotEmbeddingJob;
storage: CopilotStorage; storage: CopilotStorage;
workflow: CopilotWorkflowService; workflow: CopilotWorkflowService;
executors: { executors: {
@@ -127,7 +127,7 @@ test.before(async t => {
const storage = module.get(CopilotStorage); const storage = module.get(CopilotStorage);
const context = module.get(CopilotContextService); const context = module.get(CopilotContextService);
const jobs = module.get(CopilotContextDocJob); const jobs = module.get(CopilotEmbeddingJob);
const transcript = module.get(CopilotTranscriptionService); const transcript = module.get(CopilotTranscriptionService);
const workspaceEmbedding = module.get(CopilotWorkspaceService); const workspaceEmbedding = module.get(CopilotWorkspaceService);

1
packages/backend/server/src/plugins/copilot/context/index.ts
View File

@@ -1,3 +1,2 @@
export { CopilotContextDocJob } from './job';
export { CopilotContextResolver, CopilotContextRootResolver } from './resolver'; export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
export { CopilotContextService } from './service'; export { CopilotContextService } from './service';

4
packages/backend/server/src/plugins/copilot/context/resolver.ts
View File

@@ -44,12 +44,12 @@ import {
FileChunkSimilarity, FileChunkSimilarity,
Models, Models,
} from '../../../models'; } from '../../../models';
import { CopilotEmbeddingJob } from '../embedding';
import { COPILOT_LOCKER, CopilotType } from '../resolver'; import { COPILOT_LOCKER, CopilotType } from '../resolver';
import { ChatSessionService } from '../session'; import { ChatSessionService } from '../session';
import { CopilotStorage } from '../storage'; import { CopilotStorage } from '../storage';
import { MAX_EMBEDDABLE_SIZE } from '../types'; import { MAX_EMBEDDABLE_SIZE } from '../types';
import { readStream } from '../utils'; import { readStream } from '../utils';
import { CopilotContextDocJob } from './job';
import { CopilotContextService } from './service'; import { CopilotContextService } from './service';
@InputType() @InputType()
@@ -387,7 +387,7 @@ export class CopilotContextResolver {
private readonly models: Models, private readonly models: Models,
private readonly mutex: RequestMutex, private readonly mutex: RequestMutex,
private readonly context: CopilotContextService, private readonly context: CopilotContextService,
private readonly jobs: CopilotContextDocJob, private readonly jobs: CopilotEmbeddingJob,
private readonly storage: CopilotStorage private readonly storage: CopilotStorage
) {} ) {}

3
packages/backend/server/src/plugins/copilot/context/service.ts
View File

@@ -14,11 +14,10 @@ import {
ContextFile, ContextFile,
Models, Models,
} from '../../../models'; } from '../../../models';
import { type EmbeddingClient, getEmbeddingClient } from '../embedding';
import { PromptService } from '../prompt'; import { PromptService } from '../prompt';
import { CopilotProviderFactory } from '../providers'; import { CopilotProviderFactory } from '../providers';
import { getEmbeddingClient } from './embedding';
import { ContextSession } from './session'; import { ContextSession } from './session';
import type { EmbeddingClient } from './types';
const CONTEXT_SESSION_KEY = 'context-session'; const CONTEXT_SESSION_KEY = 'context-session';

6
packages/backend/server/src/plugins/copilot/context/session.ts
View File

@@ -11,11 +11,11 @@ import {
FileChunkSimilarity, FileChunkSimilarity,
Models, Models,
} from '../../../models'; } from '../../../models';
import { EmbeddingClient } from './types'; import { EmbeddingClient } from '../embedding';
export class ContextSession implements AsyncDisposable { export class ContextSession implements AsyncDisposable {
constructor( constructor(
private readonly client: EmbeddingClient, private readonly client: EmbeddingClient | undefined,
private readonly contextId: string, private readonly contextId: string,
private readonly config: ContextConfig, private readonly config: ContextConfig,
private readonly models: Models, private readonly models: Models,
@@ -204,6 +204,7 @@ export class ContextSession implements AsyncDisposable {
scopedThreshold: number = 0.85, scopedThreshold: number = 0.85,
threshold: number = 0.5 threshold: number = 0.5
): Promise<FileChunkSimilarity[]> { ): Promise<FileChunkSimilarity[]> {
if (!this.client) return [];
const embedding = await this.client.getEmbedding(content, signal); const embedding = await this.client.getEmbedding(content, signal);
if (!embedding) return []; if (!embedding) return [];
@@ -256,6 +257,7 @@ export class ContextSession implements AsyncDisposable {
scopedThreshold: number = 0.85, scopedThreshold: number = 0.85,
threshold: number = 0.5 threshold: number = 0.5
) { ) {
if (!this.client) return [];
const embedding = await this.client.getEmbedding(content, signal); const embedding = await this.client.getEmbedding(content, signal);
if (!embedding) return []; if (!embedding) return [];

11
packages/backend/server/src/plugins/copilot/context/utils.ts
View File

@@ -1,11 +0,0 @@
export class GqlSignal implements AsyncDisposable {
readonly abortController = new AbortController();
get signal() {
return this.abortController.signal;
}
async [Symbol.asyncDispose]() {
this.abortController.abort();
}
}

8
packages/backend/server/src/plugins/copilot/context/embedding.ts → packages/backend/server/src/plugins/copilot/embedding/client.ts
View File

@@ -4,11 +4,11 @@ import {
CopilotPromptNotFound, CopilotPromptNotFound,
CopilotProviderNotSupported, CopilotProviderNotSupported,
} from '../../../base'; } from '../../../base';
import type { ChunkSimilarity, Embedding } from '../../../models'; import { ChunkSimilarity, Embedding } from '../../../models';
import type { PromptService } from '../prompt'; import { PromptService } from '../prompt';
import { import {
type CopilotProvider, type CopilotProvider,
type CopilotProviderFactory, CopilotProviderFactory,
type ModelFullConditions, type ModelFullConditions,
ModelInputType, ModelInputType,
ModelOutputType, ModelOutputType,
@@ -22,7 +22,7 @@ import {
const RERANK_PROMPT = 'Rerank results'; const RERANK_PROMPT = 'Rerank results';
export class ProductionEmbeddingClient extends EmbeddingClient { class ProductionEmbeddingClient extends EmbeddingClient {
private readonly logger = new Logger(ProductionEmbeddingClient.name); private readonly logger = new Logger(ProductionEmbeddingClient.name);
constructor( constructor(

4
packages/backend/server/src/plugins/copilot/embedding/index.ts Normal file
View File

@@ -0,0 +1,4 @@
export { getEmbeddingClient, MockEmbeddingClient } from './client';
export { CopilotEmbeddingJob } from './job';
export type { Chunk, DocFragment } from './types';
export { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types';

6
packages/backend/server/src/plugins/copilot/context/job.ts → packages/backend/server/src/plugins/copilot/embedding/job.ts
View File

@@ -18,12 +18,12 @@ import { PromptService } from '../prompt';
import { CopilotProviderFactory } from '../providers'; import { CopilotProviderFactory } from '../providers';
import { CopilotStorage } from '../storage'; import { CopilotStorage } from '../storage';
import { readStream } from '../utils'; import { readStream } from '../utils';
import { getEmbeddingClient } from './embedding'; import { getEmbeddingClient } from './client';
import type { Chunk, DocFragment } from './types'; import type { Chunk, DocFragment } from './types';
import { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types'; import { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types';
@Injectable() @Injectable()
export class CopilotContextDocJob { export class CopilotEmbeddingJob {
private readonly workspaceJobAbortController: Map<string, AbortController> = private readonly workspaceJobAbortController: Map<string, AbortController> =
new Map(); new Map();
@@ -40,7 +40,7 @@ export class CopilotContextDocJob {
private readonly queue: JobQueue, private readonly queue: JobQueue,
private readonly storage: CopilotStorage private readonly storage: CopilotStorage
) { ) {
this.logger.setContext(CopilotContextDocJob.name); this.logger.setContext(CopilotEmbeddingJob.name);
} }
@OnEvent('config.init') @OnEvent('config.init')

0
packages/backend/server/src/plugins/copilot/context/types.ts → packages/backend/server/src/plugins/copilot/embedding/types.ts
View File

4
packages/backend/server/src/plugins/copilot/index.ts
View File

@@ -9,12 +9,12 @@ import { PermissionModule } from '../../core/permission';
import { QuotaModule } from '../../core/quota'; import { QuotaModule } from '../../core/quota';
import { WorkspaceModule } from '../../core/workspaces'; import { WorkspaceModule } from '../../core/workspaces';
import { import {
CopilotContextDocJob,
CopilotContextResolver, CopilotContextResolver,
CopilotContextRootResolver, CopilotContextRootResolver,
CopilotContextService, CopilotContextService,
} from './context'; } from './context';
import { CopilotController } from './controller'; import { CopilotController } from './controller';
import { CopilotEmbeddingJob } from './embedding';
import { ChatMessageCache } from './message'; import { ChatMessageCache } from './message';
import { PromptService } from './prompt'; import { PromptService } from './prompt';
import { CopilotProviderFactory, CopilotProviders } from './providers'; import { CopilotProviderFactory, CopilotProviders } from './providers';
@@ -61,7 +61,7 @@ import {
// context // context
CopilotContextResolver, CopilotContextResolver,
CopilotContextService, CopilotContextService,
CopilotContextDocJob, CopilotEmbeddingJob,
// transcription // transcription
CopilotTranscriptionService, CopilotTranscriptionService,
CopilotTranscriptionResolver, CopilotTranscriptionResolver,