rebuild index upon restart

Signed-off-by: Huamin Chen <[email protected]>

Author: rootfs

src/semantic-router/pkg/cache/hybrid_cache.go

@@ -98,6 +98,9 @@ type HybridCacheOptions struct {
 
 	// Milvus settings
 	MilvusConfigPath string
+
+	// Startup settings
+	DisableRebuildOnStartup bool // Skip rebuilding HNSW index from Milvus on startup (default: false, meaning rebuild IS enabled)
 }
 
 // NewHybridCache creates a new hybrid cache instance
@@ -153,6 +156,26 @@ func NewHybridCache(options HybridCacheOptions) (*HybridCache, error) {
 	observability.Infof("Hybrid cache initialized: HNSW(M=%d, ef=%d), maxMemory=%d",
 		options.HNSWM, options.HNSWEfConstruction, options.MaxMemoryEntries)
 
+	// Rebuild HNSW index from Milvus on startup (enabled by default)
+	// This ensures the in-memory index is populated after a restart
+	// Set DisableRebuildOnStartup=true to skip this step (not recommended for production)
+	if !options.DisableRebuildOnStartup {
+		observability.Infof("Hybrid cache: rebuilding HNSW index from Milvus...")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+		defer cancel()
+
+		if err := cache.RebuildFromMilvus(ctx); err != nil {
+			observability.Warnf("Hybrid cache: failed to rebuild HNSW index from Milvus: %v", err)
+			observability.Warnf("Hybrid cache: continuing with empty HNSW index")
+			// Don't fail initialization, just log warning and continue with empty index
+		} else {
+			observability.Infof("Hybrid cache: HNSW index rebuild complete")
+		}
+	} else {
+		observability.Warnf("Hybrid cache: skipping HNSW index rebuild (DisableRebuildOnStartup=true)")
+		observability.Warnf("Hybrid cache: index will be empty until entries are added")
+	}
+
 	return cache, nil
 }
 
@@ -161,6 +184,83 @@ func (h *HybridCache) IsEnabled() bool {
 	return h.enabled
 }
 
+// RebuildFromMilvus rebuilds the in-memory HNSW index from persistent Milvus storage
+// This is called on startup to recover the index after a restart
+func (h *HybridCache) RebuildFromMilvus(ctx context.Context) error {
+	if !h.enabled {
+		return nil
+	}
+
+	start := time.Now()
+	observability.Infof("HybridCache.RebuildFromMilvus: starting HNSW index rebuild from Milvus")
+
+	// Query all entries from Milvus
+	requestIDs, embeddings, err := h.milvusCache.GetAllEntries(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get entries from Milvus: %w", err)
+	}
+
+	if len(requestIDs) == 0 {
+		observability.Infof("HybridCache.RebuildFromMilvus: no entries to rebuild, starting with empty index")
+		return nil
+	}
+
+	observability.Infof("HybridCache.RebuildFromMilvus: rebuilding HNSW index with %d entries", len(requestIDs))
+
+	// Lock for the entire rebuild process
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	// Clear existing index
+	h.embeddings = make([][]float32, 0, len(embeddings))
+	h.idMap = make(map[int]string)
+	h.hnswIndex = newHNSWIndex(h.hnswIndex.M, h.hnswIndex.efConstruction)
+
+	// Rebuild HNSW index with progress logging
+	batchSize := 1000
+	for i, embedding := range embeddings {
+		// Check memory limits
+		if len(h.embeddings) >= h.maxMemoryEntries {
+			observability.Warnf("HybridCache.RebuildFromMilvus: reached max memory entries (%d), stopping rebuild at %d/%d",
+				h.maxMemoryEntries, i, len(embeddings))
+			break
+		}
+
+		// Add to HNSW
+		entryIndex := len(h.embeddings)
+		h.embeddings = append(h.embeddings, embedding)
+		h.idMap[entryIndex] = requestIDs[i]
+		h.addNodeHybrid(entryIndex, embedding)
+
+		// Progress logging for large datasets
+		if (i+1)%batchSize == 0 {
+			elapsed := time.Since(start)
+			rate := float64(i+1) / elapsed.Seconds()
+			remaining := len(embeddings) - (i + 1)
+			eta := time.Duration(float64(remaining)/rate) * time.Second
+			observability.Infof("HybridCache.RebuildFromMilvus: progress %d/%d (%.1f%%, %.0f entries/sec, ETA: %v)",
+				i+1, len(embeddings), float64(i+1)/float64(len(embeddings))*100, rate, eta)
+		}
+	}
+
+	elapsed := time.Since(start)
+	rate := float64(len(h.embeddings)) / elapsed.Seconds()
+	observability.Infof("HybridCache.RebuildFromMilvus: rebuild complete - %d entries in %v (%.0f entries/sec)",
+		len(h.embeddings), elapsed, rate)
+
+	observability.LogEvent("hybrid_cache_rebuilt", map[string]interface{}{
+		"backend":           "hybrid",
+		"entries_loaded":    len(h.embeddings),
+		"entries_in_milvus": len(embeddings),
+		"duration_seconds":  elapsed.Seconds(),
+		"entries_per_sec":   rate,
+	})
+
+	metrics.UpdateCacheEntries("hybrid", len(h.embeddings))
+
+	return nil
+}
+
 // AddPendingRequest stores a request awaiting its response
 func (h *HybridCache) AddPendingRequest(requestID string, model string, query string, requestBody []byte) error {
 	start := time.Now()

src/semantic-router/pkg/cache/milvus_cache.go

@@ -753,6 +753,80 @@ func (c *MilvusCache) FindSimilarWithThreshold(model string, query string, thres
 	return responseBody, true, nil
 }
 
+// GetAllEntries retrieves all entries from Milvus for HNSW index rebuilding
+// Returns slices of request_ids and embeddings for efficient bulk loading
+func (c *MilvusCache) GetAllEntries(ctx context.Context) ([]string, [][]float32, error) {
+	start := time.Now()
+
+	if !c.enabled {
+		return nil, nil, fmt.Errorf("milvus cache is not enabled")
+	}
+
+	observability.Infof("MilvusCache.GetAllEntries: querying all entries for HNSW rebuild")
+
+	// Query all entries with embeddings and request_ids
+	// Filter to only get entries with complete responses (not pending)
+	queryResult, err := c.client.Query(
+		ctx,
+		c.collectionName,
+		[]string{},              // Empty partitions means search all
+		"response_body != \"\"", // Only get complete entries
+		[]string{"request_id", c.config.Collection.VectorField.Name}, // Get IDs and embeddings
+	)
+
+	if err != nil {
+		observability.Warnf("MilvusCache.GetAllEntries: query failed: %v", err)
+		return nil, nil, fmt.Errorf("milvus query all failed: %w", err)
+	}
+
+	if len(queryResult) < 2 {
+		observability.Infof("MilvusCache.GetAllEntries: no entries found or incomplete result")
+		return []string{}, [][]float32{}, nil
+	}
+
+	// Extract request IDs (first column)
+	requestIDColumn, ok := queryResult[0].(*entity.ColumnVarChar)
+	if !ok {
+		return nil, nil, fmt.Errorf("unexpected request_id column type: %T", queryResult[0])
+	}
+
+	// Extract embeddings (second column)
+	embeddingColumn, ok := queryResult[1].(*entity.ColumnFloatVector)
+	if !ok {
+		return nil, nil, fmt.Errorf("unexpected embedding column type: %T", queryResult[1])
+	}
+
+	if requestIDColumn.Len() != embeddingColumn.Len() {
+		return nil, nil, fmt.Errorf("column length mismatch: request_ids=%d, embeddings=%d",
+			requestIDColumn.Len(), embeddingColumn.Len())
+	}
+
+	entryCount := requestIDColumn.Len()
+	requestIDs := make([]string, entryCount)
+
+	// Extract request IDs from column
+	for i := 0; i < entryCount; i++ {
+		requestID, err := requestIDColumn.ValueByIdx(i)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to get request_id at index %d: %w", i, err)
+		}
+		requestIDs[i] = requestID
+	}
+
+	// Extract embeddings directly from column data
+	embeddings := embeddingColumn.Data()
+	if len(embeddings) != entryCount {
+		return nil, nil, fmt.Errorf("embedding data length mismatch: got %d, expected %d",
+			len(embeddings), entryCount)
+	}
+
+	elapsed := time.Since(start)
+	observability.Infof("MilvusCache.GetAllEntries: loaded %d entries in %v (%.0f entries/sec)",
+		entryCount, elapsed, float64(entryCount)/elapsed.Seconds())
+
+	return requestIDs, embeddings, nil
+}
+
 // GetByID retrieves a document from Milvus by its request ID
 // This is much more efficient than FindSimilar when you already know the ID
 // Used by hybrid cache to fetch documents after local HNSW search