mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-17 01:56:27 +08:00
feat(server): split embedding client (#12809)
This commit is contained in:
DarkSky
@@ -0,0 +1,219 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
|
||||
import {
|
||||
CopilotPromptNotFound,
|
||||
CopilotProviderNotSupported,
|
||||
} from '../../../base';
|
||||
import { ChunkSimilarity, Embedding } from '../../../models';
|
||||
import { PromptService } from '../prompt';
|
||||
import {
|
||||
type CopilotProvider,
|
||||
CopilotProviderFactory,
|
||||
type ModelFullConditions,
|
||||
ModelInputType,
|
||||
ModelOutputType,
|
||||
} from '../providers';
|
||||
import {
|
||||
EMBEDDING_DIMENSIONS,
|
||||
EmbeddingClient,
|
||||
getReRankSchema,
|
||||
type ReRankResult,
|
||||
} from './types';
|
||||
|
||||
const RERANK_PROMPT = 'Rerank results';
|
||||
|
||||
class ProductionEmbeddingClient extends EmbeddingClient {
|
||||
private readonly logger = new Logger(ProductionEmbeddingClient.name);
|
||||
|
||||
constructor(
|
||||
private readonly providerFactory: CopilotProviderFactory,
|
||||
private readonly prompt: PromptService
|
||||
) {
|
||||
super();
|
||||
}
|
||||
|
||||
override async configured(): Promise<boolean> {
|
||||
const embedding = await this.providerFactory.getProvider({
|
||||
outputType: ModelOutputType.Embedding,
|
||||
});
|
||||
const result = Boolean(embedding);
|
||||
if (!result) {
|
||||
this.logger.warn(
|
||||
'Copilot embedding client is not configured properly, please check your configuration.'
|
||||
);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private async getProvider(
|
||||
cond: ModelFullConditions
|
||||
): Promise<CopilotProvider> {
|
||||
const provider = await this.providerFactory.getProvider(cond);
|
||||
if (!provider) {
|
||||
throw new CopilotProviderNotSupported({
|
||||
provider: 'embedding',
|
||||
kind: cond.outputType || 'embedding',
|
||||
});
|
||||
}
|
||||
return provider;
|
||||
}
|
||||
|
||||
async getEmbeddings(input: string[]): Promise<Embedding[]> {
|
||||
const provider = await this.getProvider({
|
||||
outputType: ModelOutputType.Embedding,
|
||||
});
|
||||
this.logger.verbose(
|
||||
`Using provider ${provider.type} for embedding: ${input.join(', ')}`
|
||||
);
|
||||
|
||||
const embeddings = await provider.embedding(
|
||||
{ inputTypes: [ModelInputType.Text] },
|
||||
input,
|
||||
{ dimensions: EMBEDDING_DIMENSIONS }
|
||||
);
|
||||
|
||||
return Array.from(embeddings.entries()).map(([index, embedding]) => ({
|
||||
index,
|
||||
embedding,
|
||||
content: input[index],
|
||||
}));
|
||||
}
|
||||
|
||||
private getTargetId<T extends ChunkSimilarity>(embedding: T) {
|
||||
return 'docId' in embedding
|
||||
? embedding.docId
|
||||
: 'fileId' in embedding
|
||||
? embedding.fileId
|
||||
: '';
|
||||
}
|
||||
|
||||
private async getEmbeddingRelevance<
|
||||
Chunk extends ChunkSimilarity = ChunkSimilarity,
|
||||
>(
|
||||
query: string,
|
||||
embeddings: Chunk[],
|
||||
signal?: AbortSignal
|
||||
): Promise<ReRankResult> {
|
||||
if (!embeddings.length) return [];
|
||||
|
||||
const prompt = await this.prompt.get(RERANK_PROMPT);
|
||||
if (!prompt) {
|
||||
throw new CopilotPromptNotFound({ name: RERANK_PROMPT });
|
||||
}
|
||||
const provider = await this.getProvider({ modelId: prompt.model });
|
||||
const schema = getReRankSchema(embeddings.length);
|
||||
|
||||
const ranks = await provider.structure(
|
||||
{ modelId: prompt.model },
|
||||
prompt.finish({
|
||||
query,
|
||||
results: embeddings.map(e => ({
|
||||
targetId: this.getTargetId(e),
|
||||
chunk: e.chunk,
|
||||
content: e.content,
|
||||
})),
|
||||
schema,
|
||||
}),
|
||||
{ maxRetries: 3, signal }
|
||||
);
|
||||
|
||||
try {
|
||||
return schema.parse(JSON.parse(ranks)).ranks;
|
||||
} catch (error) {
|
||||
this.logger.error('Failed to parse rerank results', error);
|
||||
// silent error, will fallback to default sorting in parent method
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
override async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
|
||||
query: string,
|
||||
embeddings: Chunk[],
|
||||
topK: number,
|
||||
signal?: AbortSignal
|
||||
): Promise<Chunk[]> {
|
||||
// search in context and workspace may find same chunks, de-duplicate them
|
||||
const { deduped: dedupedEmbeddings } = embeddings.reduce(
|
||||
(acc, e) => {
|
||||
const key = `${this.getTargetId(e)}:${e.chunk}`;
|
||||
if (!acc.seen.has(key)) {
|
||||
acc.seen.add(key);
|
||||
acc.deduped.push(e);
|
||||
}
|
||||
return acc;
|
||||
},
|
||||
{ deduped: [] as Chunk[], seen: new Set<string>() }
|
||||
);
|
||||
const sortedEmbeddings = dedupedEmbeddings.toSorted(
|
||||
(a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
|
||||
);
|
||||
|
||||
const chunks = sortedEmbeddings.reduce(
|
||||
(acc, e) => {
|
||||
const targetId = 'docId' in e ? e.docId : 'fileId' in e ? e.fileId : '';
|
||||
const key = `${targetId}:${e.chunk}`;
|
||||
acc[key] = e;
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, Chunk>
|
||||
);
|
||||
|
||||
try {
|
||||
// 4.1 mini's context windows large enough to handle all embeddings
|
||||
const ranks = await this.getEmbeddingRelevance(
|
||||
query,
|
||||
sortedEmbeddings,
|
||||
signal
|
||||
);
|
||||
if (sortedEmbeddings.length !== ranks.length) {
|
||||
// llm return wrong result, fallback to default sorting
|
||||
this.logger.warn(
|
||||
`Batch size mismatch: expected ${sortedEmbeddings.length}, got ${ranks.length}`
|
||||
);
|
||||
return await super.reRank(query, dedupedEmbeddings, topK, signal);
|
||||
}
|
||||
|
||||
const highConfidenceChunks = ranks
|
||||
.flat()
|
||||
.toSorted((a, b) => b.scores.score - a.scores.score)
|
||||
.filter(r => r.scores.score > 5)
|
||||
.map(r => chunks[`${r.scores.targetId}:${r.scores.chunk}`])
|
||||
.filter(Boolean);
|
||||
|
||||
this.logger.verbose(
|
||||
`ReRank completed: ${highConfidenceChunks.length} high-confidence results found`
|
||||
);
|
||||
return highConfidenceChunks.slice(0, topK);
|
||||
} catch (error) {
|
||||
this.logger.warn('ReRank failed, falling back to default sorting', error);
|
||||
return await super.reRank(query, dedupedEmbeddings, topK, signal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let EMBEDDING_CLIENT: EmbeddingClient | undefined;
|
||||
export async function getEmbeddingClient(
|
||||
providerFactory: CopilotProviderFactory,
|
||||
prompt: PromptService
|
||||
): Promise<EmbeddingClient | undefined> {
|
||||
if (EMBEDDING_CLIENT) {
|
||||
return EMBEDDING_CLIENT;
|
||||
}
|
||||
const client = new ProductionEmbeddingClient(providerFactory, prompt);
|
||||
if (await client.configured()) {
|
||||
EMBEDDING_CLIENT = client;
|
||||
}
|
||||
return EMBEDDING_CLIENT;
|
||||
}
|
||||
|
||||
export class MockEmbeddingClient extends EmbeddingClient {
|
||||
async getEmbeddings(input: string[]): Promise<Embedding[]> {
|
||||
return input.map((_, i) => ({
|
||||
index: i,
|
||||
content: input[i],
|
||||
embedding: Array.from({ length: EMBEDDING_DIMENSIONS }, () =>
|
||||
Math.random()
|
||||
),
|
||||
}));
|
||||
}
|
||||
}
|
||||
@@ -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';
|
||||
@@ -0,0 +1,403 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
|
||||
import {
|
||||
AFFiNELogger,
|
||||
BlobNotFound,
|
||||
CallMetric,
|
||||
CopilotContextFileNotSupported,
|
||||
DocNotFound,
|
||||
EventBus,
|
||||
JobQueue,
|
||||
mapAnyError,
|
||||
OnEvent,
|
||||
OnJob,
|
||||
} from '../../../base';
|
||||
import { DocReader } from '../../../core/doc';
|
||||
import { Models } from '../../../models';
|
||||
import { PromptService } from '../prompt';
|
||||
import { CopilotProviderFactory } from '../providers';
|
||||
import { CopilotStorage } from '../storage';
|
||||
import { readStream } from '../utils';
|
||||
import { getEmbeddingClient } from './client';
|
||||
import type { Chunk, DocFragment } from './types';
|
||||
import { EMBEDDING_DIMENSIONS, EmbeddingClient } from './types';
|
||||
|
||||
@Injectable()
|
||||
export class CopilotEmbeddingJob {
|
||||
private readonly workspaceJobAbortController: Map<string, AbortController> =
|
||||
new Map();
|
||||
|
||||
private supportEmbedding = false;
|
||||
private client: EmbeddingClient | undefined;
|
||||
|
||||
constructor(
|
||||
private readonly doc: DocReader,
|
||||
private readonly event: EventBus,
|
||||
private readonly logger: AFFiNELogger,
|
||||
private readonly models: Models,
|
||||
private readonly providerFactory: CopilotProviderFactory,
|
||||
private readonly prompt: PromptService,
|
||||
private readonly queue: JobQueue,
|
||||
private readonly storage: CopilotStorage
|
||||
) {
|
||||
this.logger.setContext(CopilotEmbeddingJob.name);
|
||||
}
|
||||
|
||||
@OnEvent('config.init')
|
||||
async onConfigInit() {
|
||||
await this.setup();
|
||||
}
|
||||
|
||||
@OnEvent('config.changed')
|
||||
async onConfigChanged() {
|
||||
await this.setup();
|
||||
}
|
||||
|
||||
private async setup() {
|
||||
this.supportEmbedding =
|
||||
await this.models.copilotContext.checkEmbeddingAvailable();
|
||||
if (this.supportEmbedding) {
|
||||
this.client = await getEmbeddingClient(this.providerFactory, this.prompt);
|
||||
}
|
||||
}
|
||||
|
||||
// public this client to allow overriding in tests
|
||||
get embeddingClient() {
|
||||
return this.client as EmbeddingClient;
|
||||
}
|
||||
|
||||
@CallMetric('ai', 'addFileEmbeddingQueue')
|
||||
async addFileEmbeddingQueue(file: Jobs['copilot.embedding.files']) {
|
||||
if (!this.supportEmbedding) return;
|
||||
|
||||
const { userId, workspaceId, contextId, blobId, fileId, fileName } = file;
|
||||
await this.queue.add('copilot.embedding.files', {
|
||||
userId,
|
||||
workspaceId,
|
||||
contextId,
|
||||
blobId,
|
||||
fileId,
|
||||
fileName,
|
||||
});
|
||||
}
|
||||
|
||||
@OnEvent('workspace.doc.embedding')
|
||||
async addDocEmbeddingQueue(
|
||||
docs: Events['workspace.doc.embedding'],
|
||||
options?: { contextId: string; priority: number }
|
||||
) {
|
||||
if (!this.supportEmbedding) return;
|
||||
|
||||
for (const { workspaceId, docId } of docs) {
|
||||
const jobId = `workspace:embedding:${workspaceId}:${docId}`;
|
||||
const job = await this.queue.get(jobId, 'copilot.embedding.docs');
|
||||
// if the job exists and is older than 5 minute, remove it
|
||||
if (job && job.timestamp + 5 * 60 * 1000 < Date.now()) {
|
||||
this.logger.verbose(`Removing old embedding job ${jobId}`);
|
||||
await this.queue.remove(jobId, 'copilot.embedding.docs');
|
||||
}
|
||||
|
||||
await this.queue.add(
|
||||
'copilot.embedding.docs',
|
||||
{
|
||||
contextId: options?.contextId,
|
||||
workspaceId,
|
||||
docId,
|
||||
},
|
||||
{
|
||||
jobId: `workspace:embedding:${workspaceId}:${docId}`,
|
||||
priority: options?.priority ?? 1,
|
||||
timestamp: Date.now(),
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@OnEvent('workspace.updated')
|
||||
async onWorkspaceConfigUpdate({
|
||||
id,
|
||||
enableDocEmbedding,
|
||||
}: Events['workspace.updated']) {
|
||||
// trigger workspace embedding
|
||||
this.event.emit('workspace.embedding', {
|
||||
workspaceId: id,
|
||||
enableDocEmbedding,
|
||||
});
|
||||
}
|
||||
|
||||
@OnEvent('workspace.embedding')
|
||||
async addWorkspaceEmbeddingQueue({
|
||||
workspaceId,
|
||||
enableDocEmbedding,
|
||||
}: Events['workspace.embedding']) {
|
||||
if (!this.supportEmbedding || !this.embeddingClient) return;
|
||||
|
||||
if (enableDocEmbedding === undefined) {
|
||||
enableDocEmbedding =
|
||||
await this.models.workspace.allowEmbedding(workspaceId);
|
||||
}
|
||||
|
||||
if (enableDocEmbedding) {
|
||||
const toBeEmbedDocIds =
|
||||
await this.models.copilotWorkspace.findDocsToEmbed(workspaceId);
|
||||
this.logger.debug(
|
||||
`Trigger embedding for ${toBeEmbedDocIds.length} docs in workspace ${workspaceId}`
|
||||
);
|
||||
for (const docId of toBeEmbedDocIds) {
|
||||
await this.queue.add(
|
||||
'copilot.embedding.docs',
|
||||
{
|
||||
workspaceId,
|
||||
docId,
|
||||
},
|
||||
{
|
||||
jobId: `workspace:embedding:${workspaceId}:${docId}`,
|
||||
priority: 1,
|
||||
}
|
||||
);
|
||||
}
|
||||
} else {
|
||||
const controller = this.workspaceJobAbortController.get(workspaceId);
|
||||
if (controller) {
|
||||
controller.abort();
|
||||
this.workspaceJobAbortController.delete(workspaceId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@OnEvent('doc.indexer.updated')
|
||||
async addDocEmbeddingQueueFromEvent(doc: Events['doc.indexer.updated']) {
|
||||
if (!this.supportEmbedding || !this.embeddingClient) return;
|
||||
|
||||
await this.queue.add(
|
||||
'copilot.embedding.docs',
|
||||
{
|
||||
workspaceId: doc.workspaceId,
|
||||
docId: doc.docId,
|
||||
},
|
||||
{
|
||||
jobId: `workspace:embedding:${doc.workspaceId}:${doc.docId}`,
|
||||
priority: 2,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@OnEvent('doc.indexer.deleted')
|
||||
async deleteDocEmbeddingQueueFromEvent(doc: Events['doc.indexer.deleted']) {
|
||||
await this.queue.remove(
|
||||
`workspace:embedding:${doc.workspaceId}:${doc.docId}`,
|
||||
'copilot.embedding.docs'
|
||||
);
|
||||
await this.models.copilotContext.deleteWorkspaceEmbedding(
|
||||
doc.workspaceId,
|
||||
doc.docId
|
||||
);
|
||||
}
|
||||
|
||||
private async readCopilotBlob(
|
||||
userId: string,
|
||||
workspaceId: string,
|
||||
blobId: string,
|
||||
fileName: string
|
||||
) {
|
||||
const { body } = await this.storage.get(userId, workspaceId, blobId);
|
||||
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
|
||||
const buffer = await readStream(body);
|
||||
return new File([buffer], fileName);
|
||||
}
|
||||
|
||||
@OnJob('copilot.embedding.files')
|
||||
async embedPendingFile({
|
||||
userId,
|
||||
workspaceId,
|
||||
contextId,
|
||||
blobId,
|
||||
fileId,
|
||||
fileName,
|
||||
}: Jobs['copilot.embedding.files']) {
|
||||
if (!this.supportEmbedding || !this.embeddingClient) return;
|
||||
|
||||
try {
|
||||
const file = await this.readCopilotBlob(
|
||||
userId,
|
||||
workspaceId,
|
||||
blobId,
|
||||
fileName
|
||||
);
|
||||
|
||||
// no need to check if embeddings is empty, will throw internally
|
||||
const chunks = await this.embeddingClient.getFileChunks(file);
|
||||
const total = chunks.reduce((acc, c) => acc + c.length, 0);
|
||||
|
||||
for (const chunk of chunks) {
|
||||
const embeddings = await this.embeddingClient.generateEmbeddings(chunk);
|
||||
if (contextId) {
|
||||
// for context files
|
||||
await this.models.copilotContext.insertFileEmbedding(
|
||||
contextId,
|
||||
fileId,
|
||||
embeddings
|
||||
);
|
||||
} else {
|
||||
// for workspace files
|
||||
await this.models.copilotWorkspace.insertFileEmbeddings(
|
||||
workspaceId,
|
||||
fileId,
|
||||
embeddings
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (contextId) {
|
||||
this.event.emit('workspace.file.embed.finished', {
|
||||
contextId,
|
||||
fileId,
|
||||
chunkSize: total,
|
||||
});
|
||||
}
|
||||
} catch (error: any) {
|
||||
if (contextId) {
|
||||
this.event.emit('workspace.file.embed.failed', {
|
||||
contextId,
|
||||
fileId,
|
||||
error: mapAnyError(error).message,
|
||||
});
|
||||
}
|
||||
|
||||
// passthrough error to job queue
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private async getDocFragment(
|
||||
workspaceId: string,
|
||||
docId: string
|
||||
): Promise<DocFragment | null> {
|
||||
const docContent = await this.doc.getFullDocContent(workspaceId, docId);
|
||||
const authors = await this.models.doc.getAuthors(workspaceId, docId);
|
||||
if (docContent && authors) {
|
||||
const { title = 'Untitled', summary } = docContent;
|
||||
const { createdAt, updatedAt, createdByUser, updatedByUser } = authors;
|
||||
return {
|
||||
title,
|
||||
summary,
|
||||
createdAt: createdAt.toDateString(),
|
||||
updatedAt: updatedAt.toDateString(),
|
||||
createdBy: createdByUser?.name,
|
||||
updatedBy: updatedByUser?.name,
|
||||
};
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private formatDocChunks(chunks: Chunk[], fragment: DocFragment): Chunk[] {
|
||||
return chunks.map(chunk => ({
|
||||
index: chunk.index,
|
||||
content: [
|
||||
`Title: ${fragment.title}`,
|
||||
`Created at: ${fragment.createdAt}`,
|
||||
`Updated at: ${fragment.updatedAt}`,
|
||||
fragment.createdBy ? `Created by: ${fragment.createdBy}` : undefined,
|
||||
fragment.updatedBy ? `Updated by: ${fragment.updatedBy}` : undefined,
|
||||
chunk.content,
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join('\n'),
|
||||
}));
|
||||
}
|
||||
|
||||
private getWorkspaceSignal(workspaceId: string) {
|
||||
let controller = this.workspaceJobAbortController.get(workspaceId);
|
||||
if (!controller) {
|
||||
controller = new AbortController();
|
||||
this.workspaceJobAbortController.set(workspaceId, controller);
|
||||
}
|
||||
return controller.signal;
|
||||
}
|
||||
|
||||
private async fulfillEmptyEmbedding(workspaceId: string, docId: string) {
|
||||
const emptyEmbedding = {
|
||||
index: 0,
|
||||
content: '',
|
||||
embedding: Array.from({ length: EMBEDDING_DIMENSIONS }, () => 0),
|
||||
};
|
||||
await this.models.copilotContext.insertWorkspaceEmbedding(
|
||||
workspaceId,
|
||||
docId,
|
||||
[emptyEmbedding]
|
||||
);
|
||||
}
|
||||
|
||||
@OnJob('copilot.embedding.docs')
|
||||
async embedPendingDocs({
|
||||
contextId,
|
||||
workspaceId,
|
||||
docId,
|
||||
}: Jobs['copilot.embedding.docs']) {
|
||||
if (!this.supportEmbedding || !this.embeddingClient) return;
|
||||
if (workspaceId === docId || docId.includes('$')) return;
|
||||
const signal = this.getWorkspaceSignal(workspaceId);
|
||||
|
||||
try {
|
||||
const needEmbedding =
|
||||
await this.models.copilotWorkspace.checkDocNeedEmbedded(
|
||||
workspaceId,
|
||||
docId
|
||||
);
|
||||
this.logger.verbose(
|
||||
`Check if doc ${docId} in workspace ${workspaceId} needs embedding: ${needEmbedding}`
|
||||
);
|
||||
if (needEmbedding) {
|
||||
if (signal.aborted) return;
|
||||
const fragment = await this.getDocFragment(workspaceId, docId);
|
||||
if (fragment) {
|
||||
// fast fall for empty doc, journal is easily to create a empty doc
|
||||
if (fragment.summary.trim()) {
|
||||
const embeddings = await this.embeddingClient.getFileEmbeddings(
|
||||
new File(
|
||||
[fragment.summary],
|
||||
`${fragment.title || 'Untitled'}.md`
|
||||
),
|
||||
chunks => this.formatDocChunks(chunks, fragment),
|
||||
signal
|
||||
);
|
||||
|
||||
for (const chunks of embeddings) {
|
||||
await this.models.copilotContext.insertWorkspaceEmbedding(
|
||||
workspaceId,
|
||||
docId,
|
||||
chunks
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// for empty doc, insert empty embedding
|
||||
await this.fulfillEmptyEmbedding(workspaceId, docId);
|
||||
}
|
||||
} else if (contextId) {
|
||||
throw new DocNotFound({ spaceId: workspaceId, docId });
|
||||
}
|
||||
}
|
||||
} catch (error: any) {
|
||||
if (contextId) {
|
||||
this.event.emit('workspace.doc.embed.failed', {
|
||||
contextId,
|
||||
docId,
|
||||
});
|
||||
}
|
||||
if (
|
||||
error instanceof CopilotContextFileNotSupported &&
|
||||
error.message.includes('no content found')
|
||||
) {
|
||||
this.logger.warn(
|
||||
`Doc ${docId} in workspace ${workspaceId} has no content, fulfilling empty embedding.`
|
||||
);
|
||||
// if the doc is empty, we still need to fulfill the embedding
|
||||
await this.fulfillEmptyEmbedding(workspaceId, docId);
|
||||
return;
|
||||
}
|
||||
|
||||
// passthrough error to job queue
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
import { File } from 'node:buffer';
|
||||
|
||||
import { z } from 'zod';
|
||||
|
||||
import { CopilotContextFileNotSupported } from '../../../base';
|
||||
import type { PageDocContent } from '../../../core/utils/blocksuite';
|
||||
import { ChunkSimilarity, Embedding } from '../../../models';
|
||||
import { parseDoc } from '../../../native';
|
||||
|
||||
declare global {
|
||||
interface Events {
|
||||
'workspace.embedding': {
|
||||
workspaceId: string;
|
||||
enableDocEmbedding?: boolean;
|
||||
};
|
||||
|
||||
'workspace.doc.embedding': Array<{
|
||||
workspaceId: string;
|
||||
docId: string;
|
||||
}>;
|
||||
|
||||
'workspace.doc.embed.failed': {
|
||||
contextId: string;
|
||||
docId: string;
|
||||
};
|
||||
|
||||
'workspace.file.embed.finished': {
|
||||
contextId: string;
|
||||
fileId: string;
|
||||
chunkSize: number;
|
||||
};
|
||||
|
||||
'workspace.file.embed.failed': {
|
||||
contextId: string;
|
||||
fileId: string;
|
||||
error: string;
|
||||
};
|
||||
}
|
||||
interface Jobs {
|
||||
'copilot.embedding.docs': {
|
||||
contextId?: string;
|
||||
workspaceId: string;
|
||||
docId: string;
|
||||
};
|
||||
|
||||
'copilot.embedding.files': {
|
||||
contextId?: string;
|
||||
userId: string;
|
||||
workspaceId: string;
|
||||
blobId: string;
|
||||
fileId: string;
|
||||
fileName: string;
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
export type DocFragment = PageDocContent & {
|
||||
createdAt: string;
|
||||
createdBy?: string;
|
||||
updatedAt: string;
|
||||
updatedBy?: string;
|
||||
};
|
||||
|
||||
export type Chunk = {
|
||||
index: number;
|
||||
content: string;
|
||||
};
|
||||
|
||||
export const EMBEDDING_DIMENSIONS = 1024;
|
||||
|
||||
export abstract class EmbeddingClient {
|
||||
async configured() {
|
||||
return true;
|
||||
}
|
||||
|
||||
async getFileEmbeddings(
|
||||
file: File,
|
||||
chunkMapper: (chunk: Chunk[]) => Chunk[],
|
||||
signal?: AbortSignal
|
||||
): Promise<Embedding[][]> {
|
||||
const chunks = await this.getFileChunks(file, signal);
|
||||
const chunkedEmbeddings = await Promise.all(
|
||||
chunks.map(chunk => this.generateEmbeddings(chunkMapper(chunk)))
|
||||
);
|
||||
return chunkedEmbeddings;
|
||||
}
|
||||
|
||||
async getFileChunks(file: File, signal?: AbortSignal): Promise<Chunk[][]> {
|
||||
const buffer = Buffer.from(await file.arrayBuffer());
|
||||
let doc;
|
||||
try {
|
||||
doc = await parseDoc(file.name, buffer);
|
||||
} catch (e: any) {
|
||||
throw new CopilotContextFileNotSupported({
|
||||
fileName: file.name,
|
||||
message: e?.message || e?.toString?.() || 'format not supported',
|
||||
});
|
||||
}
|
||||
if (doc && !signal?.aborted) {
|
||||
if (!doc.chunks.length) {
|
||||
throw new CopilotContextFileNotSupported({
|
||||
fileName: file.name,
|
||||
message: 'no content found',
|
||||
});
|
||||
}
|
||||
const input = doc.chunks.toSorted((a, b) => a.index - b.index);
|
||||
// chunk input into 128 every array
|
||||
const chunks: Chunk[][] = [];
|
||||
for (let i = 0; i < input.length; i += 128) {
|
||||
chunks.push(input.slice(i, i + 128));
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
throw new CopilotContextFileNotSupported({
|
||||
fileName: file.name,
|
||||
message: 'failed to parse file',
|
||||
});
|
||||
}
|
||||
|
||||
async generateEmbeddings(
|
||||
chunks: Chunk[],
|
||||
signal?: AbortSignal
|
||||
): Promise<Embedding[]> {
|
||||
const retry = 3;
|
||||
|
||||
let embeddings: Embedding[] = [];
|
||||
let error = null;
|
||||
for (let i = 0; i < retry; i++) {
|
||||
try {
|
||||
embeddings = await this.getEmbeddings(
|
||||
chunks.map(c => c.content),
|
||||
signal
|
||||
);
|
||||
break;
|
||||
} catch (e) {
|
||||
error = e;
|
||||
}
|
||||
}
|
||||
if (error) throw error;
|
||||
|
||||
// fix the index of the embeddings
|
||||
return embeddings.map(e => ({ ...e, index: chunks[e.index].index }));
|
||||
}
|
||||
|
||||
async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
|
||||
_query: string,
|
||||
embeddings: Chunk[],
|
||||
topK: number,
|
||||
_signal?: AbortSignal
|
||||
): Promise<Chunk[]> {
|
||||
// sort by distance with ascending order
|
||||
return embeddings
|
||||
.toSorted((a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity))
|
||||
.slice(0, topK);
|
||||
}
|
||||
|
||||
async getEmbedding(query: string, signal?: AbortSignal) {
|
||||
const embedding = await this.getEmbeddings([query], signal);
|
||||
return embedding?.[0]?.embedding;
|
||||
}
|
||||
|
||||
abstract getEmbeddings(
|
||||
input: string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<Embedding[]>;
|
||||
}
|
||||
|
||||
const ReRankItemSchema = z.object({
|
||||
scores: z.object({
|
||||
reason: z
|
||||
.string()
|
||||
.describe(
|
||||
'Think step by step, describe in 20 words the reason for giving this score.'
|
||||
),
|
||||
chunk: z.string().describe('The chunk index of the search result.'),
|
||||
targetId: z.string().describe('The id of the target.'),
|
||||
score: z
|
||||
.number()
|
||||
.min(0)
|
||||
.max(10)
|
||||
.describe(
|
||||
'The relevance score of the results should be 0-10, with 0 being the least relevant and 10 being the most relevant.'
|
||||
),
|
||||
}),
|
||||
});
|
||||
|
||||
export const getReRankSchema = (size: number) =>
|
||||
z.object({
|
||||
ranks: ReRankItemSchema.array().describe(
|
||||
`A array of scores. Make sure to score all ${size} results.`
|
||||
),
|
||||
});
|
||||
|
||||
export type ReRankResult = z.infer<ReturnType<typeof getReRankSchema>>['ranks'];
|
||||
Reference in New Issue
Block a user