diff --git a/packages/core/src/vectordb/milvus-restful-vectordb.ts b/packages/core/src/vectordb/milvus-restful-vectordb.ts index 65fc8c9..ff44caf 100644 --- a/packages/core/src/vectordb/milvus-restful-vectordb.ts +++ b/packages/core/src/vectordb/milvus-restful-vectordb.ts @@ -103,6 +103,29 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { } } + /** + * Ensure collection is loaded before search/query operations + */ + protected async ensureLoaded(collectionName: string): Promise { + try { + const restfulConfig = this.config as MilvusRestfulConfig; + // Check if collection is loaded + const response = await this.makeRequest('/collections/get_load_state', 'POST', { + collectionName, + dbName: restfulConfig.database + }); + + const loadState = response.data?.loadState; + if (loadState !== 'LoadStateLoaded') { + console.log(`🔄 Loading collection '${collectionName}' to memory...`); + await this.loadCollection(collectionName); + } + } catch (error) { + console.error(`❌ Failed to ensure collection '${collectionName}' is loaded:`, error); + throw error; + } + } + /** * Make HTTP request to Milvus REST API */ @@ -327,6 +350,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async insert(collectionName: string, documents: VectorDocument[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const restfulConfig = this.config as MilvusRestfulConfig; @@ -358,6 +382,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async search(collectionName: string, queryVector: number[], options?: SearchOptions): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); const topK = options?.topK || 10; @@ -427,6 +452,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async delete(collectionName: string, ids: string[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const restfulConfig = this.config as MilvusRestfulConfig; @@ -450,6 +476,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async query(collectionName: string, filter: string, outputFields: string[], limit?: number): Promise[]> { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const restfulConfig = this.config as MilvusRestfulConfig; @@ -613,6 +640,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async insertHybrid(collectionName: string, documents: VectorDocument[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const restfulConfig = this.config as MilvusRestfulConfig; @@ -648,6 +676,7 @@ export class MilvusRestfulVectorDatabase implements VectorDatabase { async hybridSearch(collectionName: string, searchRequests: HybridSearchRequest[], options?: HybridSearchOptions): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const restfulConfig = this.config as MilvusRestfulConfig; diff --git a/packages/core/src/vectordb/milvus-vectordb.ts b/packages/core/src/vectordb/milvus-vectordb.ts index 768f684..6a69d5b 100644 --- a/packages/core/src/vectordb/milvus-vectordb.ts +++ b/packages/core/src/vectordb/milvus-vectordb.ts @@ -1,4 +1,4 @@ -import { MilvusClient, DataType, MetricType, FunctionType } from '@zilliz/milvus2-sdk-node'; +import { MilvusClient, DataType, MetricType, FunctionType, LoadState } from '@zilliz/milvus2-sdk-node'; import { VectorDocument, SearchOptions, @@ -99,6 +99,28 @@ export class MilvusVectorDatabase implements VectorDatabase { } } + /** + * Ensure collection is loaded before search/query operations + */ + protected async ensureLoaded(collectionName: string): Promise { + try { + // Check if collection is loaded + const result = await this.client!.getLoadState({ + collection_name: collectionName + }); + + if (result.state !== LoadState.LoadStateLoaded) { + console.log(`🔄 Loading collection '${collectionName}' to memory...`); + await this.client!.loadCollection({ + collection_name: collectionName, + }); + } + } catch (error) { + console.error(`❌ Failed to ensure collection '${collectionName}' is loaded:`, error); + throw error; + } + } + async createCollection(collectionName: string, dimension: number, description?: string): Promise { await this.ensureInitialized(); @@ -212,6 +234,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async insert(collectionName: string, documents: VectorDocument[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); console.log('Inserting documents into collection:', collectionName); const data = documents.map(doc => ({ @@ -233,6 +256,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async search(collectionName: string, queryVector: number[], options?: SearchOptions): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); const searchParams: any = { collection_name: collectionName, @@ -269,6 +293,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async delete(collectionName: string, ids: string[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); await this.client!.delete({ collection_name: collectionName, @@ -278,6 +303,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async query(collectionName: string, filter: string, outputFields: string[], limit?: number): Promise[]> { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { const queryParams: any = { @@ -422,6 +448,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async insertHybrid(collectionName: string, documents: VectorDocument[]): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); const data = documents.map(doc => ({ id: doc.id, @@ -442,6 +469,7 @@ export class MilvusVectorDatabase implements VectorDatabase { async hybridSearch(collectionName: string, searchRequests: HybridSearchRequest[], options?: HybridSearchOptions): Promise { await this.ensureInitialized(); + await this.ensureLoaded(collectionName); try { // Generate OpenAI embedding for the first search request (dense)