From c9a0127bc58fc7a02b0228671c0b114e56a40072 Mon Sep 17 00:00:00 2001 From: wangbinteng Date: Thu, 9 Jul 2026 15:54:08 +0800 Subject: [PATCH] feat: add matryoshka embedding truncation --- .../DocumentStoreVS/DocStoreVector.ts | 13 +++++-- packages/components/src/index.ts | 1 + .../src/matryoshkaEmbeddings.test.ts | 37 +++++++++++++++++++ .../components/src/matryoshkaEmbeddings.ts | 35 ++++++++++++++++++ .../src/services/documentstore/index.ts | 31 +++++++++++++--- .../views/docstore/VectorStoreConfigure.jsx | 17 ++++++++- 6 files changed, 124 insertions(+), 10 deletions(-) create mode 100644 packages/components/src/matryoshkaEmbeddings.test.ts create mode 100644 packages/components/src/matryoshkaEmbeddings.ts diff --git a/packages/components/nodes/vectorstores/DocumentStoreVS/DocStoreVector.ts b/packages/components/nodes/vectorstores/DocumentStoreVS/DocStoreVector.ts index 0f228d1fb79..bdee7e6c31f 100644 --- a/packages/components/nodes/vectorstores/DocumentStoreVS/DocStoreVector.ts +++ b/packages/components/nodes/vectorstores/DocumentStoreVS/DocStoreVector.ts @@ -1,4 +1,5 @@ import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, INodeOutputsValue, INodeParams } from '../../../src/Interface' +import { applyMatryoshkaTruncation, MATRYOSHKA_TRUNCATE_DIMENSIONS } from '../../../src/matryoshkaEmbeddings' import { DataSource } from 'typeorm' class DocStore_VectorStores implements INode { @@ -118,8 +119,11 @@ class DocStore_VectorStores implements INode { const _createEmbeddingsObject = async (componentNodes: ICommonObject, data: ICommonObject, options: ICommonObject): Promise => { // prepare embedding node data const embeddingComponent = componentNodes[data.embeddingName] + const embeddingConfig = { ...data.embeddingConfig } + const matryoshkaTruncateDimensions = embeddingConfig[MATRYOSHKA_TRUNCATE_DIMENSIONS] + delete embeddingConfig[MATRYOSHKA_TRUNCATE_DIMENSIONS] const embeddingNodeData: any = { - inputs: { ...data.embeddingConfig }, + inputs: embeddingConfig, outputs: { output: 'document' }, id: `${embeddingComponent.name}_0`, label: embeddingComponent.label, @@ -127,15 +131,16 @@ const _createEmbeddingsObject = async (componentNodes: ICommonObject, data: ICom category: embeddingComponent.category, inputParams: embeddingComponent.inputs || [] } - if (data.embeddingConfig.credential) { - embeddingNodeData.credential = data.embeddingConfig.credential + if (embeddingConfig.credential) { + embeddingNodeData.credential = embeddingConfig.credential } // init embedding object const embeddingNodeInstanceFilePath = embeddingComponent.filePath as string const embeddingNodeModule = await import(embeddingNodeInstanceFilePath) const embeddingNodeInstance = new embeddingNodeModule.nodeClass() - return await embeddingNodeInstance.init(embeddingNodeData, '', options) + const embeddingObj = await embeddingNodeInstance.init(embeddingNodeData, '', options) + return applyMatryoshkaTruncation(embeddingObj, matryoshkaTruncateDimensions) } const _createVectorStoreNodeData = (componentNodes: ICommonObject, data: ICommonObject, embeddingObj: any) => { diff --git a/packages/components/src/index.ts b/packages/components/src/index.ts index 8ebf5ffba3a..09e75d496ed 100644 --- a/packages/components/src/index.ts +++ b/packages/components/src/index.ts @@ -12,6 +12,7 @@ export * from './handler' export * from './headerValidation' export * from './httpSecurity' export * from './Interface' +export * from './matryoshkaEmbeddings' export * from './recordManagerSecurity' export * from './sanitizeDataSourceOptions' export * from './speechToText' diff --git a/packages/components/src/matryoshkaEmbeddings.test.ts b/packages/components/src/matryoshkaEmbeddings.test.ts new file mode 100644 index 00000000000..f7c836db73d --- /dev/null +++ b/packages/components/src/matryoshkaEmbeddings.test.ts @@ -0,0 +1,37 @@ +import { EmbeddingsInterface } from '@langchain/core/embeddings' +import { applyMatryoshkaTruncation } from './matryoshkaEmbeddings' + +const makeEmbeddings = (): EmbeddingsInterface => ({ + embedDocuments: jest.fn(async () => [ + [1, 2, 3, 4], + [5, 6, 7, 8] + ]), + embedQuery: jest.fn(async () => [9, 10, 11, 12]) +}) + +describe('applyMatryoshkaTruncation', () => { + it('returns the original embeddings when no truncate dimension is configured', async () => { + const embeddings = makeEmbeddings() + const result = applyMatryoshkaTruncation(embeddings, undefined) + + expect(result).toBe(embeddings) + await expect(result.embedQuery('hello')).resolves.toEqual([9, 10, 11, 12]) + }) + + it('truncates document and query embeddings to the configured dimension', async () => { + const embeddings = makeEmbeddings() + const result = applyMatryoshkaTruncation(embeddings, '2') + + await expect(result.embedDocuments(['a', 'b'])).resolves.toEqual([ + [1, 2], + [5, 6] + ]) + await expect(result.embedQuery('hello')).resolves.toEqual([9, 10]) + }) + + it('rejects invalid truncate dimensions', () => { + expect(() => applyMatryoshkaTruncation(makeEmbeddings(), '0')).toThrow('Matryoshka truncate dimension must be a positive integer') + expect(() => applyMatryoshkaTruncation(makeEmbeddings(), '2.5')).toThrow('Matryoshka truncate dimension must be a positive integer') + expect(() => applyMatryoshkaTruncation(makeEmbeddings(), 'abc')).toThrow('Matryoshka truncate dimension must be a positive integer') + }) +}) diff --git a/packages/components/src/matryoshkaEmbeddings.ts b/packages/components/src/matryoshkaEmbeddings.ts new file mode 100644 index 00000000000..43e3fe4e3c7 --- /dev/null +++ b/packages/components/src/matryoshkaEmbeddings.ts @@ -0,0 +1,35 @@ +import { EmbeddingsInterface } from '@langchain/core/embeddings' + +export const MATRYOSHKA_TRUNCATE_DIMENSIONS = 'matryoshkaTruncateDimensions' + +const parseTruncateDimension = (value: unknown): number | undefined => { + if (value === undefined || value === null || value === '') return undefined + + const dimension = typeof value === 'number' ? value : Number(value) + if (!Number.isInteger(dimension) || dimension <= 0) { + throw new Error('Matryoshka truncate dimension must be a positive integer') + } + + return dimension +} + +const truncateVector = (vector: number[], dimension: number): number[] => vector.slice(0, dimension) + +export const applyMatryoshkaTruncation = >(embeddings: T, truncateDimensions: unknown): T => { + const dimension = parseTruncateDimension(truncateDimensions) + if (!dimension) return embeddings + + const wrappedEmbeddings = Object.create(embeddings) as T + + wrappedEmbeddings.embedDocuments = async (documents: string[]): Promise => { + const vectors = await embeddings.embedDocuments(documents) + return vectors.map((vector) => truncateVector(vector, dimension)) + } + + wrappedEmbeddings.embedQuery = async (document: string): Promise => { + const vector = await embeddings.embedQuery(document) + return truncateVector(vector, dimension) + } + + return wrappedEmbeddings +} diff --git a/packages/server/src/services/documentstore/index.ts b/packages/server/src/services/documentstore/index.ts index a7dc444f1d6..b9f46f63014 100644 --- a/packages/server/src/services/documentstore/index.ts +++ b/packages/server/src/services/documentstore/index.ts @@ -7,6 +7,8 @@ import { getFileFromUpload, ICommonObject, IDocument, + applyMatryoshkaTruncation, + MATRYOSHKA_TRUNCATE_DIMENSIONS, mapExtToInputField, mapMimeTypeToInputField, removeFilesFromStorage, @@ -1583,8 +1585,11 @@ const _createEmbeddingsObject = async ( ): Promise => { // prepare embedding node data const embeddingComponent = componentNodes[data.embeddingName] + const embeddingConfig = { ...data.embeddingConfig } + const matryoshkaTruncateDimensions = embeddingConfig[MATRYOSHKA_TRUNCATE_DIMENSIONS] + delete embeddingConfig[MATRYOSHKA_TRUNCATE_DIMENSIONS] const embeddingNodeData: any = { - inputs: { ...data.embeddingConfig }, + inputs: embeddingConfig, outputs: { output: 'document' }, id: `${embeddingComponent.name}_0`, label: embeddingComponent.label, @@ -1592,12 +1597,28 @@ const _createEmbeddingsObject = async ( category: embeddingComponent.category, inputParams: embeddingComponent.inputs || [] } - if (data.embeddingConfig.credential) { - embeddingNodeData.credential = data.embeddingConfig.credential + if (embeddingConfig.credential) { + embeddingNodeData.credential = embeddingConfig.credential } // save to upsert history - if (upsertHistory) upsertHistory['flowData'] = saveUpsertFlowData(embeddingNodeData, upsertHistory) + if (upsertHistory) { + const embeddingHistoryNodeData = matryoshkaTruncateDimensions + ? { + ...embeddingNodeData, + inputs: { ...embeddingConfig, [MATRYOSHKA_TRUNCATE_DIMENSIONS]: matryoshkaTruncateDimensions }, + inputParams: [ + ...embeddingNodeData.inputParams, + { + label: 'Matryoshka Truncate Dimensions', + name: MATRYOSHKA_TRUNCATE_DIMENSIONS, + type: 'number' + } + ] + } + : embeddingNodeData + upsertHistory['flowData'] = saveUpsertFlowData(embeddingHistoryNodeData, upsertHistory) + } // init embedding object const embeddingNodeInstanceFilePath = embeddingComponent.filePath as string @@ -1607,7 +1628,7 @@ const _createEmbeddingsObject = async ( if (!embeddingObj) { throw new InternalFlowiseError(StatusCodes.INTERNAL_SERVER_ERROR, `Failed to create EmbeddingObj`) } - return embeddingObj + return applyMatryoshkaTruncation(embeddingObj, matryoshkaTruncateDimensions) } const _createRecordManagerObject = async ( diff --git a/packages/ui/src/views/docstore/VectorStoreConfigure.jsx b/packages/ui/src/views/docstore/VectorStoreConfigure.jsx index 13636b77f79..7901ae94565 100644 --- a/packages/ui/src/views/docstore/VectorStoreConfigure.jsx +++ b/packages/ui/src/views/docstore/VectorStoreConfigure.jsx @@ -45,6 +45,21 @@ import useNotifier from '@/utils/useNotifier' // const const steps = ['Embeddings', 'Vector Store', 'Record Manager'] +const matryoshkaTruncateDimensionsInput = { + label: 'Matryoshka Truncate Dimensions', + name: 'matryoshkaTruncateDimensions', + type: 'number', + optional: true, + additionalParams: true, + description: 'Optionally keep only the first N dimensions from document and query embeddings.' +} + +const addMatryoshkaInput = (nodeData) => { + if (!nodeData.inputParams?.some((inputParam) => inputParam.name === matryoshkaTruncateDimensionsInput.name)) { + nodeData.inputParams = [...(nodeData.inputParams || []), matryoshkaTruncateDimensionsInput] + } + return nodeData +} const VectorStoreConfigure = () => { const navigate = useNavigate() @@ -123,7 +138,7 @@ const VectorStoreConfigure = () => { nodeData.inputs = documentStore.embeddingConfig.config nodeData.credential = documentStore.embeddingConfig.config.credential } - setSelectedEmbeddingsProvider(nodeData) + setSelectedEmbeddingsProvider(addMatryoshkaInput(nodeData)) setShowEmbeddingsListDialog(false) }