Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 325bcd1

Browse files
zc277584121zc277584121
committed
enhance milvus vector db
- add index_name when creating index - optimize collection limit check with better decoupling design - optimize waitForIndexReady to replace client.getIndexState(cause it's about to be deprecated) Signed-off-by: ChengZi <[email protected]>
1 parent 84253a3 commit 325bcd1

5 files changed

Lines changed: 139 additions & 89 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,4 @@ evaluation/repos
6666
repos
6767

6868
evaluation/retrieval_results*
69+
.venv/

‎packages/core/src/vectordb/milvus-restful-vectordb.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export interface MilvusRestfulConfig {
1919
}
2020

2121
/**
22+
* TODO: Change this usage to checkCollectionLimit()
2223
* Wrapper function to handle collection creation with limit detection
2324
* This is the single point where collection limit errors are detected and handled
2425
*/
@@ -774,4 +775,16 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase {
774775
throw error;
775776
}
776777
}
778+
779+
/**
780+
* Check collection limit
781+
* Returns true if collection can be created, false if limit exceeded
782+
* TODO: Implement proper collection limit checking for REST API
783+
*/
784+
async checkCollectionLimit(): Promise<boolean> {
785+
// TODO: Implement REST API version of collection limit checking
786+
// For now, always return true to maintain compatibility
787+
console.warn('⚠️ checkCollectionLimit not implemented for REST API - returning true');
788+
return true;
789+
}
777790
}

‎packages/core/src/vectordb/milvus-vectordb.ts‎

Lines changed: 105 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { MilvusClient, DataType, MetricType, FunctionType, LoadState, IndexState } from '@zilliz/milvus2-sdk-node';
1+
import { MilvusClient, DataType, MetricType, FunctionType, LoadState } from '@zilliz/milvus2-sdk-node';
22
import {
33
VectorDocument,
44
SearchOptions,
@@ -7,7 +7,6 @@ import {
77
HybridSearchRequest,
88
HybridSearchOptions,
99
HybridSearchResult,
10-
COLLECTION_LIMIT_MESSAGE
1110
} from './types';
1211
import { ClusterManager } from './zilliz-utils';
1312

@@ -19,27 +18,7 @@ export interface MilvusConfig {
1918
ssl?: boolean;
2019
}
2120

22-
/**
23-
* Wrapper function to handle collection creation with limit detection for gRPC client
24-
* This is the single point where collection limit errors are detected and handled
25-
*/
26-
async function createCollectionWithLimitCheck(
27-
client: MilvusClient,
28-
createCollectionParams: any
29-
): Promise<void> {
30-
try {
31-
await client.createCollection(createCollectionParams);
32-
} catch (error: any) {
33-
// Check if the error message contains the collection limit exceeded pattern
34-
const errorMessage = error.message || error.toString() || '';
35-
if (/exceeded the limit number of collections/i.test(errorMessage)) {
36-
// Throw the exact message string, not an Error object
37-
throw COLLECTION_LIMIT_MESSAGE;
38-
}
39-
// Re-throw other errors as-is
40-
throw error;
41-
}
42-
}
21+
4322

4423
export class MilvusVectorDatabase implements VectorDatabase {
4524
protected config: MilvusConfig;
@@ -127,95 +106,103 @@ export class MilvusVectorDatabase implements VectorDatabase {
127106

128107
/**
129108
* Wait for an index to be ready before proceeding
130-
* Polls index status with exponential backoff up to 30 seconds
109+
* Polls index build progress with exponential backoff up to 60 seconds
131110
*/
132-
protected async waitForIndexReady(collectionName: string, fieldName: string): Promise<void> {
111+
protected async waitForIndexReady(
112+
collectionName: string,
113+
fieldName: string,
114+
maxWaitTime: number = 60000, // 60 seconds
115+
initialInterval: number = 500, // 500ms
116+
maxInterval: number = 5000, // 5 seconds
117+
backoffMultiplier: number = 1.5
118+
): Promise<void> {
133119
if (!this.client) {
134120
throw new Error('MilvusClient is not initialized. Call ensureInitialized() first.');
135121
}
136122

137-
const maxWaitTime = 60000; // 60 seconds
138-
const initialInterval = 500; // 500ms
139-
const maxInterval = 5000; // 5 seconds
140-
const backoffMultiplier = 1.5;
141-
142123
let interval = initialInterval;
143124
const startTime = Date.now();
144-
125+
145126
console.log(`⏳ Waiting for index on field '${fieldName}' in collection '${collectionName}' to be ready...`);
146-
127+
147128
while (Date.now() - startTime < maxWaitTime) {
148129
try {
149-
const indexStateResult = await this.client.getIndexState({
130+
const indexBuildProgress = await this.client.getIndexBuildProgress({
150131
collection_name: collectionName,
151132
field_name: fieldName
152133
});
153-
154-
// Debug logging to understand the state value
155-
console.log(`📊 Index state for '${fieldName}': raw=${indexStateResult.state}, type=${typeof indexStateResult.state}, IndexState.Finished=${IndexState.Finished}`);
156-
console.log(`📊 Full response:`, JSON.stringify(indexStateResult));
157-
158-
// Check both numeric and potential string values
159-
// Cast to any to bypass TypeScript checks while debugging
160-
const stateValue = indexStateResult.state as any;
161-
if (stateValue === IndexState.Finished ||
162-
stateValue === 3 ||
163-
stateValue === 'Finished') {
164-
console.log(`✅ Index on field '${fieldName}' is ready!`);
134+
135+
// Debug logging to understand the progress
136+
console.log(`📊 Index build progress for '${fieldName}': indexed_rows=${indexBuildProgress.indexed_rows}, total_rows=${indexBuildProgress.total_rows}`);
137+
console.log(`📊 Full response:`, JSON.stringify(indexBuildProgress));
138+
139+
// Check if index building is complete
140+
if (indexBuildProgress.indexed_rows === indexBuildProgress.total_rows) {
141+
console.log(`✅ Index on field '${fieldName}' is ready! (${indexBuildProgress.indexed_rows}/${indexBuildProgress.total_rows} rows indexed)`);
165142
return;
166143
}
167-
168-
if (indexStateResult.state === IndexState.Failed) {
169-
throw new Error(`Index creation failed for field '${fieldName}' in collection '${collectionName}'`);
144+
145+
// Check for error status
146+
if (indexBuildProgress.status && indexBuildProgress.status.error_code !== 'Success') {
147+
// Handle known issue with older Milvus versions where sparse vector index progress returns incorrect error
148+
if (indexBuildProgress.status.reason && indexBuildProgress.status.reason.includes('index duplicates[indexName=]')) {
149+
console.log(`⚠️ Index progress check returned known older Milvus issue: ${indexBuildProgress.status.reason}`);
150+
console.log(`⚠️ This is a known issue with older Milvus versions - treating as index ready`);
151+
return; // Treat as ready since this is a false error
152+
}
153+
throw new Error(`Index creation failed for field '${fieldName}' in collection '${collectionName}': ${indexBuildProgress.status.reason}`);
170154
}
171-
155+
156+
console.log(`📊 Index building in progress: ${indexBuildProgress.indexed_rows}/${indexBuildProgress.total_rows} rows indexed`);
157+
172158
// Wait with exponential backoff
173159
await new Promise(resolve => setTimeout(resolve, interval));
174160
interval = Math.min(interval * backoffMultiplier, maxInterval);
175-
161+
176162
} catch (error) {
177-
console.error(`❌ Error checking index state for field '${fieldName}':`, error);
163+
console.error(`❌ Error checking index build progress for field '${fieldName}':`, error);
178164
throw error;
179165
}
180166
}
181-
167+
182168
throw new Error(`Timeout waiting for index on field '${fieldName}' in collection '${collectionName}' to be ready after ${maxWaitTime}ms`);
183169
}
184170

185171
/**
186172
* Load collection with retry logic and exponential backoff
187173
* Retries up to 5 times with exponential backoff
188174
*/
189-
protected async loadCollectionWithRetry(collectionName: string): Promise<void> {
175+
protected async loadCollectionWithRetry(
176+
collectionName: string,
177+
maxRetries: number = 5,
178+
initialInterval: number = 1000, // 1 second
179+
backoffMultiplier: number = 2
180+
): Promise<void> {
190181
if (!this.client) {
191182
throw new Error('MilvusClient is not initialized. Call ensureInitialized() first.');
192183
}
193184

194-
const maxRetries = 5;
195-
const initialInterval = 1000; // 1 second
196-
const backoffMultiplier = 2;
197-
198185
let attempt = 1;
199186
let interval = initialInterval;
200-
187+
201188
while (attempt <= maxRetries) {
202189
try {
203190
console.log(`🔄 Loading collection '${collectionName}' to memory (attempt ${attempt}/${maxRetries})...`);
204-
191+
205192
await this.client.loadCollection({
206193
collection_name: collectionName,
207194
});
208-
195+
209196
console.log(`✅ Collection '${collectionName}' loaded successfully!`);
210197
return;
211-
198+
212199
} catch (error) {
213200
console.error(`❌ Failed to load collection '${collectionName}' on attempt ${attempt}:`, error);
214-
201+
215202
if (attempt === maxRetries) {
216203
throw new Error(`Failed to load collection '${collectionName}' after ${maxRetries} attempts: ${error}`);
217204
}
218-
205+
219206
// Wait with exponential backoff before retry
220207
console.log(`⏳ Retrying collection load in ${interval}ms...`);
221208
await new Promise(resolve => setTimeout(resolve, interval));
@@ -290,12 +277,13 @@ export class MilvusVectorDatabase implements VectorDatabase {
290277
throw new Error('MilvusClient is not initialized. Call ensureInitialized() first.');
291278
}
292279

293-
await createCollectionWithLimitCheck(this.client, createCollectionParams);
280+
await this.client.createCollection(createCollectionParams);
294281

295282
// Create index
296283
const indexParams = {
297284
collection_name: collectionName,
298285
field_name: 'vector',
286+
index_name: 'vector_index',
299287
index_type: 'AUTOINDEX',
300288
metric_type: MetricType.COSINE,
301289
};
@@ -556,13 +544,14 @@ export class MilvusVectorDatabase implements VectorDatabase {
556544
throw new Error('MilvusClient is not initialized. Call ensureInitialized() first.');
557545
}
558546

559-
await createCollectionWithLimitCheck(this.client, createCollectionParams);
547+
await this.client.createCollection(createCollectionParams);
560548

561549
// Create indexes for both vector fields
562550
// Index for dense vector
563551
const denseIndexParams = {
564552
collection_name: collectionName,
565553
field_name: 'vector',
554+
index_name: 'vector_index',
566555
index_type: 'AUTOINDEX',
567556
metric_type: MetricType.COSINE,
568557
};
@@ -576,10 +565,12 @@ export class MilvusVectorDatabase implements VectorDatabase {
576565
const sparseIndexParams = {
577566
collection_name: collectionName,
578567
field_name: 'sparse_vector',
568+
index_name: 'sparse_vector_index',
579569
index_type: 'SPARSE_INVERTED_INDEX',
580570
metric_type: MetricType.BM25,
581571
};
582572
console.log(`🔧 Creating sparse vector index for field 'sparse_vector' in collection '${collectionName}'...`);
573+
583574
await this.client.createIndex(sparseIndexParams);
584575

585576
// Wait for sparse vector index to be ready
@@ -722,4 +713,53 @@ export class MilvusVectorDatabase implements VectorDatabase {
722713
throw error;
723714
}
724715
}
716+
717+
/**
718+
* Wrapper method to handle collection creation with limit detection for gRPC client
719+
* Returns true if collection can be created, false if limit exceeded
720+
*/
721+
async checkCollectionLimit(): Promise<boolean> {
722+
if (!this.client) {
723+
throw new Error('MilvusClient is not initialized. Call ensureInitialized() first.');
724+
}
725+
726+
const collectionName = `dummy_collection_${Date.now()}`;
727+
const createCollectionParams = {
728+
collection_name: collectionName,
729+
description: 'Test collection for limit check',
730+
fields: [
731+
{
732+
name: 'id',
733+
data_type: DataType.VarChar,
734+
max_length: 512,
735+
is_primary_key: true,
736+
},
737+
{
738+
name: 'vector',
739+
data_type: DataType.FloatVector,
740+
dim: 128,
741+
}
742+
]
743+
};
744+
745+
try {
746+
await this.client.createCollection(createCollectionParams);
747+
// Immediately drop the collection after successful creation
748+
if (await this.client.hasCollection({ collection_name: collectionName })) {
749+
await this.client.dropCollection({
750+
collection_name: collectionName,
751+
});
752+
}
753+
return true;
754+
} catch (error: any) {
755+
// Check if the error message contains the collection limit exceeded pattern
756+
const errorMessage = error.message || error.toString() || '';
757+
if (/exceeded the limit number of collections/i.test(errorMessage)) {
758+
// Return false for collection limit exceeded
759+
return false;
760+
}
761+
// Re-throw other errors as-is
762+
throw error;
763+
}
764+
}
725765
}

‎packages/core/src/vectordb/types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,12 @@ export interface VectorDatabase {
125125
* @param limit Maximum number of results
126126
*/
127127
query(collectionName: string, filter: string, outputFields: string[], limit?: number): Promise<Record<string, any>[]>;
128+
129+
/**
130+
* Check collection limit
131+
* Returns true if collection can be created, false if limit exceeded
132+
*/
133+
checkCollectionLimit(): Promise<boolean>;
128134
}
129135

130136
/**

‎packages/mcp/src/handlers.ts‎

Lines changed: 14 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -230,21 +230,9 @@ export class ToolHandlers {
230230
// CRITICAL: Pre-index collection creation validation
231231
try {
232232
console.log(`[INDEX-VALIDATION] 🔍 Validating collection creation capability`);
233-
//dummy collection name
234-
const collectionName = `dummy_collection_${Date.now()}`;
235-
await this.context.getVectorDatabase().createCollection(collectionName, 128);
236-
if (await this.context.getVectorDatabase().hasCollection(collectionName)) {
237-
console.log(`[INDEX-VALIDATION] ℹ️ Dummy collection created successfully`);
238-
await this.context.getVectorDatabase().dropCollection(collectionName);
239-
} else {
240-
console.log(`[INDEX-VALIDATION] ❌ Dummy collection creation failed`);
241-
}
242-
console.log(`[INDEX-VALIDATION] ✅ Collection creation validation completed`);
243-
} catch (validationError: any) {
244-
const errorMessage = typeof validationError === 'string' ? validationError :
245-
(validationError instanceof Error ? validationError.message : String(validationError));
233+
const canCreateCollection = await this.context.getVectorDatabase().checkCollectionLimit();
246234

247-
if (errorMessage === COLLECTION_LIMIT_MESSAGE || errorMessage.includes(COLLECTION_LIMIT_MESSAGE)) {
235+
if (!canCreateCollection) {
248236
console.error(`[INDEX-VALIDATION] ❌ Collection limit validation failed: ${absolutePath}`);
249237

250238
// CRITICAL: Immediately return the COLLECTION_LIMIT_MESSAGE to MCP client
@@ -255,17 +243,19 @@ export class ToolHandlers {
255243
}],
256244
isError: true
257245
};
258-
} else {
259-
// Handle other collection creation errors
260-
console.error(`[INDEX-VALIDATION] ❌ Collection creation validation failed:`, validationError);
261-
return {
262-
content: [{
263-
type: "text",
264-
text: `Error validating collection creation: ${validationError.message || validationError}`
265-
}],
266-
isError: true
267-
};
268246
}
247+
248+
console.log(`[INDEX-VALIDATION] ✅ Collection creation validation completed`);
249+
} catch (validationError: any) {
250+
// Handle other collection creation errors
251+
console.error(`[INDEX-VALIDATION] ❌ Collection creation validation failed:`, validationError);
252+
return {
253+
content: [{
254+
type: "text",
255+
text: `Error validating collection creation: ${validationError.message || validationError}`
256+
}],
257+
isError: true
258+
};
269259
}
270260

271261
// Add custom extensions if provided

0 commit comments

Comments
 (0)