alihan/claude-context-mcp
1
0
Fork 0
mirror of https://github.com/zilliztech/claude-context.git synced 2025-10-06 01:10:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: ChengZi <chen.zhang@zilliz.com>
This commit is contained in:
Cheney Zhang
2025-08-11 14:13:29 +08:00
committed by GitHub
parent a7fd7fa2a9
commit d385f6fc28
2 changed files with 58 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
packages/core/src/vectordb/milvus-restful-vectordb.ts
View File

@@ -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;

30
packages/core/src/vectordb/milvus-vectordb.ts
View File

@@ -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)