-
-
Notifications
You must be signed in to change notification settings - Fork 24.7k
feat(milvus): make distance metric type configurable #6579
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,58 @@ | ||
| import { MetricType } from '@zilliz/milvus2-sdk-node' | ||
|
|
||
| // getBaseClasses(Milvus) is called in the node constructor; stub the utils module | ||
| // so the test does not pull the full src/utils dependency graph. | ||
| jest.mock('../../../src/utils', () => ({ | ||
| getBaseClasses: jest.fn(() => ['Milvus', 'VectorStore']), | ||
| getCredentialData: jest.fn(), | ||
| getCredentialParam: jest.fn(), | ||
| FLOWISE_CHATID: 'chatId' | ||
| })) | ||
|
|
||
| const { nodeClass: Milvus_VectorStores, resolveMilvusMetricType } = require('./Milvus') | ||
|
|
||
| describe('Milvus resolveMilvusMetricType', () => { | ||
| it('maps L2 (any case) to MetricType.L2', () => { | ||
| expect(resolveMilvusMetricType('L2')).toBe(MetricType.L2) | ||
| expect(resolveMilvusMetricType('l2')).toBe(MetricType.L2) | ||
| }) | ||
|
|
||
| it('maps COSINE (any case) to MetricType.COSINE', () => { | ||
| expect(resolveMilvusMetricType('COSINE')).toBe(MetricType.COSINE) | ||
| expect(resolveMilvusMetricType('cosine')).toBe(MetricType.COSINE) | ||
| expect(resolveMilvusMetricType(' Cosine ')).toBe(MetricType.COSINE) | ||
| }) | ||
|
|
||
| it('maps IP (any case) to MetricType.IP', () => { | ||
| expect(resolveMilvusMetricType('IP')).toBe(MetricType.IP) | ||
| expect(resolveMilvusMetricType('ip')).toBe(MetricType.IP) | ||
| }) | ||
|
|
||
| it('falls back to L2 for empty, undefined, or unknown input (backward compatible)', () => { | ||
| expect(resolveMilvusMetricType(undefined)).toBe(MetricType.L2) | ||
| expect(resolveMilvusMetricType('')).toBe(MetricType.L2) | ||
| expect(resolveMilvusMetricType('HAMMING')).toBe(MetricType.L2) | ||
| }) | ||
| }) | ||
|
|
||
| describe('Milvus node Metric Type input', () => { | ||
| let node: any | ||
|
|
||
| beforeEach(() => { | ||
| node = new Milvus_VectorStores() | ||
| }) | ||
|
|
||
| it('exposes a metricType options input with L2/COSINE/IP and an L2 default', () => { | ||
| const input = node.inputs.find((i: any) => i.name === 'metricType') | ||
| expect(input).toBeDefined() | ||
| expect(input.type).toBe('options') | ||
| expect(input.default).toBe('L2') | ||
| const optionNames = input.options.map((o: any) => o.name) | ||
| expect(optionNames).toEqual(['L2', 'COSINE', 'IP']) | ||
| }) | ||
|
|
||
| it('keeps the metricType input optional so existing flows are unaffected', () => { | ||
| const input = node.inputs.find((i: any) => i.name === 'metricType') | ||
| expect(input.optional).toBe(true) | ||
| }) | ||
| }) |
| Original file line number | Diff line number | Diff line change | ||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -11,6 +11,25 @@ interface InsertRow { | |||||||||||||||
| [x: string]: string | number[] | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| /** | ||||||||||||||||
| * Resolve a user-provided metric type string to a Milvus metric type. | ||||||||||||||||
| * Accepts the values exposed by the node's "Metric Type" input (L2, COSINE, IP), | ||||||||||||||||
| * case-insensitively, and falls back to L2 for empty or unrecognized input so | ||||||||||||||||
| * existing flows keep their previous behavior. Returns the string-literal union | ||||||||||||||||
| * used by @langchain/community's Milvus `indexCreateParams.metric_type`. | ||||||||||||||||
| */ | ||||||||||||||||
| const resolveMilvusMetricType = (value?: string): 'L2' | 'COSINE' | 'IP' => { | ||||||||||||||||
| switch ((value ?? '').trim().toUpperCase()) { | ||||||||||||||||
| case 'COSINE': | ||||||||||||||||
| return 'COSINE' | ||||||||||||||||
| case 'IP': | ||||||||||||||||
| return 'IP' | ||||||||||||||||
| case 'L2': | ||||||||||||||||
| default: | ||||||||||||||||
| return 'L2' | ||||||||||||||||
| } | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| class Milvus_VectorStores implements INode { | ||||||||||||||||
| label: string | ||||||||||||||||
| name: string | ||||||||||||||||
|
|
@@ -28,7 +47,7 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
| constructor() { | ||||||||||||||||
| this.label = 'Milvus' | ||||||||||||||||
| this.name = 'milvus' | ||||||||||||||||
| this.version = 2.1 | ||||||||||||||||
| this.version = 2.2 | ||||||||||||||||
| this.type = 'Milvus' | ||||||||||||||||
| this.icon = 'milvus.svg' | ||||||||||||||||
| this.category = 'Vector Stores' | ||||||||||||||||
|
|
@@ -112,6 +131,21 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
| additionalParams: true, | ||||||||||||||||
| optional: true | ||||||||||||||||
| }, | ||||||||||||||||
| { | ||||||||||||||||
| label: 'Metric Type', | ||||||||||||||||
| name: 'metricType', | ||||||||||||||||
| description: | ||||||||||||||||
| 'Distance metric used when creating the index on upsert and when querying. Must match the metric the Milvus collection was created with. Defaults to L2.', | ||||||||||||||||
| type: 'options', | ||||||||||||||||
| options: [ | ||||||||||||||||
| { label: 'L2 (Euclidean)', name: 'L2' }, | ||||||||||||||||
| { label: 'Cosine Similarity', name: 'COSINE' }, | ||||||||||||||||
| { label: 'Inner Product (IP)', name: 'IP' } | ||||||||||||||||
| ], | ||||||||||||||||
| default: 'L2', | ||||||||||||||||
| additionalParams: true, | ||||||||||||||||
| optional: true | ||||||||||||||||
| }, | ||||||||||||||||
| { | ||||||||||||||||
| label: 'Secure', | ||||||||||||||||
| name: 'secure', | ||||||||||||||||
|
|
@@ -194,6 +228,9 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
| // partition | ||||||||||||||||
| const partitionName = nodeData.inputs?.milvusPartition ?? '_default' | ||||||||||||||||
|
|
||||||||||||||||
| // metric type | ||||||||||||||||
| const metricType = resolveMilvusMetricType(nodeData.inputs?.metricType as string) | ||||||||||||||||
|
|
||||||||||||||||
| // init MilvusLibArgs | ||||||||||||||||
| const milVusArgs: MilvusLibArgs = { | ||||||||||||||||
| url: address, | ||||||||||||||||
|
|
@@ -229,7 +266,13 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| try { | ||||||||||||||||
| const vectorStore = await MilvusUpsert.fromDocuments(finalDocs, embeddings, milVusArgs) | ||||||||||||||||
| // Mirror MilvusUpsert.fromDocuments, but inject the configured metric type | ||||||||||||||||
| // into indexCreateParams before the index is created on first upsert. | ||||||||||||||||
| const vectorStore = new MilvusUpsert(embeddings, milVusArgs) | ||||||||||||||||
| if (vectorStore.indexCreateParams) { | ||||||||||||||||
| vectorStore.indexCreateParams.metric_type = metricType | ||||||||||||||||
| } | ||||||||||||||||
| await vectorStore.addDocuments(finalDocs) | ||||||||||||||||
|
|
||||||||||||||||
| // Avoid Illegal Invocation | ||||||||||||||||
| vectorStore.similaritySearchVectorWithScore = async (query: number[], k: number, filter?: string) => { | ||||||||||||||||
|
|
@@ -276,6 +319,9 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
| // partition | ||||||||||||||||
| const partitionName = nodeData.inputs?.milvusPartition ?? '_default' | ||||||||||||||||
|
|
||||||||||||||||
| // metric type | ||||||||||||||||
| const metricType = resolveMilvusMetricType(nodeData.inputs?.metricType as string) | ||||||||||||||||
|
|
||||||||||||||||
| // init MilvusLibArgs | ||||||||||||||||
| const milVusArgs: MilvusLibArgs = { | ||||||||||||||||
| url: address, | ||||||||||||||||
|
|
@@ -308,6 +354,12 @@ class Milvus_VectorStores implements INode { | |||||||||||||||
|
|
||||||||||||||||
| const vectorStore = await Milvus.fromExistingCollection(embeddings, milVusArgs) | ||||||||||||||||
|
|
||||||||||||||||
| // Ensure the search path uses the user-selected metric (and its matching score | ||||||||||||||||
| // normalization) instead of the LangChain default of L2. | ||||||||||||||||
| if (vectorStore.indexCreateParams) { | ||||||||||||||||
| vectorStore.indexCreateParams.metric_type = metricType | ||||||||||||||||
| } | ||||||||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Similarly, if
Suggested change
References
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Addressed in 5e153b0: same spread fix applied on the search path so the metric is never silently dropped. |
||||||||||||||||
|
|
||||||||||||||||
| // Avoid Illegal Invocation | ||||||||||||||||
| vectorStore.similaritySearchVectorWithScore = async (query: number[], k: number, filter?: string) => { | ||||||||||||||||
| return await similaritySearchVectorWithScore(query, k, vectorStore, milvusFilter, filter) | ||||||||||||||||
|
|
@@ -471,7 +523,7 @@ class MilvusUpsert extends Milvus { | |||||||||||||||
| field_name: this.vectorField, | ||||||||||||||||
| index_name: `myindex_${Date.now().toString()}`, | ||||||||||||||||
| index_type: IndexType.AUTOINDEX, | ||||||||||||||||
| metric_type: MetricType.L2 | ||||||||||||||||
| metric_type: this.indexCreateParams?.metric_type ?? MetricType.L2 | ||||||||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In TypeScript,
Suggested change
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Addressed in 5e153b0: cast the value to MetricType at the createIndex call site for strict type-checking. Thanks for the review. |
||||||||||||||||
| }) | ||||||||||||||||
| if (resp.error_code !== ErrorCode.SUCCESS) { | ||||||||||||||||
| throw new Error(`Error creating index`) | ||||||||||||||||
|
|
@@ -491,4 +543,4 @@ class MilvusUpsert extends Milvus { | |||||||||||||||
| } | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| module.exports = { nodeClass: Milvus_VectorStores } | ||||||||||||||||
| module.exports = { nodeClass: Milvus_VectorStores, resolveMilvusMetricType } | ||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If
vectorStore.indexCreateParamsis undefined, the checkif (vectorStore.indexCreateParams)will prevent setting the metric type, causing it to silently fall back to the default (L2) and ignoring the user's configuration. It is safer to initializeindexCreateParamsif it is not already defined. Spreadingundefinedis safe and does not require a fallback.References
nullorundefinedwithin an object literal is safe and does not require a?? {}fallback, as it evaluates to an empty object.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Addressed in 5e153b0: indexCreateParams is now built with object spread so the selected metric is always applied.