ensure collection is loaded before search/query (#148)

Signed-off-by: ChengZi <chen.zhang@zilliz.com>

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<void> {
+        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<void> {
         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<VectorSearchResult[]> {
         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<void> {
         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<Record<string, any>[]> {
         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<void> {
         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<HybridSearchResult[]> {
         await this.ensureInitialized();
+        await this.ensureLoaded(collectionName);
 
         try {
             const restfulConfig = this.config as MilvusRestfulConfig;

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<void> {
+        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<void> {
         await this.ensureInitialized();
 
@@ -212,6 +234,7 @@ export class MilvusVectorDatabase implements VectorDatabase {
 
     async insert(collectionName: string, documents: VectorDocument[]): Promise<void> {
         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<VectorSearchResult[]> {
         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<void> {
         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<Record<string, any>[]> {
         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<void> {
         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<HybridSearchResult[]> {
         await this.ensureInitialized();
+        await this.ensureLoaded(collectionName);
 
         try {
             // Generate OpenAI embedding for the first search request (dense)