Fix Milvus Vector Database Tests (#602)

Signed-off-by: Yehudit Kerido <ykerido@ykerido-thinkpadp1gen7.raanaii.csb>
Co-authored-by: Yehudit Kerido <ykerido@ykerido-thinkpadp1gen7.raanaii.csb>

Author: yehudit1987

.github/workflows/test-and-build.yml

@@ -83,13 +83,48 @@ jobs:
           HF_HUB_DISABLE_TELEMETRY: 1
         run: make download-models
 
+      - name: Start Milvus service
+        run: |
+          echo "Starting Milvus vector database..."
+          docker run -d \
+            --name milvus-semantic-cache \
+            --security-opt seccomp:unconfined \
+            -e ETCD_USE_EMBED=true \
+            -e ETCD_DATA_DIR=/var/lib/milvus/etcd \
+            -e ETCD_CONFIG_PATH=/milvus/configs/advanced/etcd.yaml \
+            -e COMMON_STORAGETYPE=local \
+            -e CLUSTER_ENABLED=false \
+            -p 19530:19530 \
+            -p 9091:9091 \
+            milvusdb/milvus:v2.3.3 \
+            milvus run standalone
+          
+          echo "Waiting for Milvus to be ready..."
+          sleep 20
+          
+          # Verify Milvus is responsive
+          timeout 30 bash -c 'until docker logs milvus-semantic-cache 2>&1 | grep -q "Proxy successfully started"; do sleep 2; done' || true
+          
+          echo "Milvus is ready at localhost:19530"
+          docker ps --filter "name=milvus-semantic-cache"
+
       - name: Run semantic router tests
         run: make test
         env:
           CI: true
           CI_MINIMAL_MODELS: ${{ github.event_name == 'pull_request' }}
           CGO_ENABLED: 1
           LD_LIBRARY_PATH: ${{ github.workspace }}/candle-binding/target/release
+          MILVUS_URI: localhost:19530
+          SKIP_MILVUS_TESTS: false
+
+      - name: Stop Milvus service
+        if: always()
+        run: |
+          echo "Stopping Milvus container..."
+          docker stop milvus-semantic-cache || true
+          docker rm milvus-semantic-cache || true
+          echo "Milvus container cleaned up"
 
       - name: Upload test artifacts on failure
         if: failure()

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

@@ -1203,28 +1203,14 @@ func TestHybridCacheDisabled(t *testing.T) {
 
 // TestHybridCacheBasicOperations tests basic cache operations
 func TestHybridCacheBasicOperations(t *testing.T) {
-	// Skip if Milvus is not configured
-	if os.Getenv("MILVUS_URI") == "" {
-		t.Skip("Skipping: MILVUS_URI not set")
-	}
+	t.Log("Starting TestHybridCacheBasicOperations - this may take 30-60 seconds...")
 
 	// Create a test Milvus config
-	milvusConfig := "/tmp/test_milvus_config.yaml"
-	err := os.WriteFile(milvusConfig, []byte(`
-milvus:
-  address: "localhost:19530"
-  collection_name: "test_hybrid_cache"
-  dimension: 384
-  index_type: "HNSW"
-  metric_type: "IP"
-  params:
-    M: 16
-    efConstruction: 200
-`), 0o644)
+	milvusConfig, cleanup, err := createTestMilvusConfig("test_hybrid_cache", 200, true)
 	if err != nil {
 		t.Fatalf("Failed to create test config: %v", err)
 	}
-	defer os.Remove(milvusConfig)
+	defer cleanup()
 
 	cache, err := NewHybridCache(HybridCacheOptions{
 		Enabled:             true,
@@ -1260,7 +1246,8 @@ milvus:
 	}
 
 	// Test FindSimilar with exact same query (should hit)
-	time.Sleep(100 * time.Millisecond) // Allow indexing to complete
+	// Wait for Milvus to index the entry
+	time.Sleep(2 * time.Second)
 
 	response, found, err := cache.FindSimilar("gpt-4", testQuery)
 	if err != nil {
@@ -1303,25 +1290,13 @@ milvus:
 
 // TestHybridCachePendingRequest tests pending request flow
 func TestHybridCachePendingRequest(t *testing.T) {
-	// Skip if Milvus is not configured
-	if os.Getenv("MILVUS_URI") == "" {
-		t.Skip("Skipping: MILVUS_URI not set")
-	}
+	t.Log("Starting TestHybridCachePendingRequest - this may take 30-60 seconds...")
 
-	milvusConfig := "/tmp/test_milvus_pending_config.yaml"
-	err := os.WriteFile(milvusConfig, []byte(`
-milvus:
-  address: "localhost:19530"
-  collection_name: "test_hybrid_pending"
-  dimension: 384
-  index_type: "HNSW"
-  metric_type: "IP"
-`),
-		0o644)
+	milvusConfig, cleanup, err := createTestMilvusConfig("test_hybrid_pending", 64, true)
 	if err != nil {
 		t.Fatalf("Failed to create test config: %v", err)
 	}
-	defer os.Remove(milvusConfig)
+	defer cleanup()
 
 	cache, err := NewHybridCache(HybridCacheOptions{
 		Enabled:             true,
@@ -1367,25 +1342,13 @@ milvus:
 
 // TestHybridCacheEviction tests memory eviction behavior
 func TestHybridCacheEviction(t *testing.T) {
-	// Skip if Milvus is not configured
-	if os.Getenv("MILVUS_URI") == "" {
-		t.Skip("Skipping: MILVUS_URI not set")
-	}
+	t.Log("Starting TestHybridCacheEviction - this may take 30-60 seconds...")
 
-	milvusConfig := "/tmp/test_milvus_eviction_config.yaml"
-	err := os.WriteFile(milvusConfig, []byte(`
-milvus:
-  address: "localhost:19530"
-  collection_name: "test_hybrid_eviction"
-  dimension: 384
-  index_type: "HNSW"
-  metric_type: "IP"
-`),
-		0o644)
+	milvusConfig, cleanup, err := createTestMilvusConfig("test_hybrid_eviction", 64, true)
 	if err != nil {
 		t.Fatalf("Failed to create test config: %v", err)
 	}
-	defer os.Remove(milvusConfig)
+	defer cleanup()
 
 	// Create cache with very small memory limit
 	cache, err := NewHybridCache(HybridCacheOptions{
@@ -1418,7 +1381,8 @@ milvus:
 
 	// All entries should still be in Milvus
 	// Try to find a recent entry (should be in memory)
-	time.Sleep(100 * time.Millisecond)
+	// Wait for Milvus to index all entries
+	time.Sleep(2 * time.Second)
 	_, found, err := cache.FindSimilar("gpt-4", "Query number 9")
 	if err != nil {
 		t.Fatalf("FindSimilar failed: %v", err)
@@ -1438,25 +1402,13 @@ milvus:
 
 // TestHybridCacheLocalCacheHit tests local cache hot path
 func TestHybridCacheLocalCacheHit(t *testing.T) {
-	// Skip if Milvus is not configured
-	if os.Getenv("MILVUS_URI") == "" {
-		t.Skip("Skipping: MILVUS_URI not set")
-	}
+	t.Log("Starting TestHybridCacheLocalCacheHit - this may take 30-60 seconds...")
 
-	milvusConfig := "/tmp/test_milvus_local_config.yaml"
-	err := os.WriteFile(milvusConfig, []byte(`
-milvus:
-  address: "localhost:19530"
-  collection_name: "test_hybrid_local"
-  dimension: 384
-  index_type: "HNSW"
-  metric_type: "IP"
-`),
-		0o644)
+	milvusConfig, cleanup, err := createTestMilvusConfig("test_hybrid_local", 64, true)
 	if err != nil {
 		t.Fatalf("Failed to create test config: %v", err)
 	}
-	defer os.Remove(milvusConfig)
+	defer cleanup()
 
 	cache, err := NewHybridCache(HybridCacheOptions{
 		Enabled:             true,
@@ -1478,16 +1430,18 @@ milvus:
 		t.Fatalf("Failed to add entry: %v", err)
 	}
 
-	time.Sleep(100 * time.Millisecond)
+	// Wait longer for Milvus to index the entry
+	time.Sleep(2 * time.Second)
 
 	// First search - should populate local cache
-	_, found, err := cache.FindSimilar("gpt-4", testQuery)
+	response1, found, err := cache.FindSimilar("gpt-4", testQuery)
 	if err != nil {
 		t.Fatalf("FindSimilar failed: %v", err)
 	}
 	if !found {
 		t.Fatal("Expected to find entry")
 	}
+	t.Logf("First search returned: %s", string(response1))
 
 	// Second search - should hit local cache (much faster)
 	startTime := time.Now()
@@ -1499,6 +1453,7 @@ milvus:
 	if !found {
 		t.Fatal("Expected to find entry in local cache")
 	}
+	t.Logf("Second search returned: %s", string(response))
 	if string(response) != string(testResponse) {
 		t.Errorf("Response mismatch: got %s, want %s", string(response), string(testResponse))
 	}
@@ -1776,6 +1731,47 @@ func (dcc *DatabaseCallCounter) Reset() {
 	atomic.StoreInt64(&dcc.calls, 0)
 }
 
+// createTestMilvusConfig creates a temporary Milvus config file for testing
+// Returns the path to the config file and cleanup function
+func createTestMilvusConfig(collectionName string, efConstruction int, dropOnStartup bool) (string, func(), error) {
+	configPath := fmt.Sprintf("/tmp/test_milvus_%s_config.yaml", collectionName)
+
+	configYAML := fmt.Sprintf(`connection:
+  host: "localhost"
+  port: 19530
+  timeout: 30
+collection:
+  name: "%s"
+  vector_field:
+    name: "embedding"
+    dimension: 384
+    metric_type: "IP"
+  index:
+    type: "HNSW"
+    params:
+      M: 16
+      efConstruction: %d
+search:
+  params:
+    ef: 64
+  topk: 10
+development:
+  auto_create_collection: true
+  drop_collection_on_startup: %t
+`, collectionName, efConstruction, dropOnStartup)
+
+	err := os.WriteFile(configPath, []byte(configYAML), 0o644)
+	if err != nil {
+		return "", nil, fmt.Errorf("failed to create test config: %w", err)
+	}
+
+	cleanup := func() {
+		os.Remove(configPath)
+	}
+
+	return configPath, cleanup, nil
+}
+
 // getMilvusConfigPath returns the path to milvus.yaml config file
 func getMilvusConfigPath() string {
 	// Check for environment variable first
@@ -2475,7 +2471,14 @@ func writeBenchmarkResultToCSV(file *os.File, result BenchmarkResult) {
 
 // TestHybridVsMilvusSmoke is a quick smoke test to verify both caches work
 func TestHybridVsMilvusSmoke(t *testing.T) {
-	t.Skip("Skipping smoke test in short mode")
+	t.Log("Starting TestHybridVsMilvusSmoke - this may take 2-3 minutes...")
+
+	// Create test Milvus config
+	milvusConfig, cleanup, err := createTestMilvusConfig("test_smoke_cache", 64, true)
+	if err != nil {
+		t.Fatalf("Failed to create test config: %v", err)
+	}
+	defer cleanup()
 
 	// Initialize BERT model
 	useCPU := os.Getenv("USE_CPU") != "false"
@@ -2490,7 +2493,7 @@ func TestHybridVsMilvusSmoke(t *testing.T) {
 			Enabled:             true,
 			SimilarityThreshold: 0.85,
 			TTLSeconds:          3600,
-			ConfigPath:          getMilvusConfigPath(),
+			ConfigPath:          milvusConfig,
 		})
 		if err != nil {
 			t.Fatalf("Failed to create Milvus cache: %v", err)
@@ -2531,7 +2534,7 @@ func TestHybridVsMilvusSmoke(t *testing.T) {
 			MaxMemoryEntries:    1000,
 			HNSWM:               16,
 			HNSWEfConstruction:  200,
-			MilvusConfigPath:    getMilvusConfigPath(),
+			MilvusConfigPath:    milvusConfig,
 		})
 		if err != nil {
 			t.Fatalf("Failed to create Hybrid cache: %v", err)

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

@@ -761,28 +761,37 @@ func (h *HybridCache) evictOneUnsafe() {
 		return
 	}
 
-	// Simple FIFO eviction: remove oldest entry
+	// Simple FIFO eviction: remove oldest entry (index 0)
 	victimIdx := 0
 
-	// Could use LRU/LFU here by tracking access times/counts
-	// For now, just evict the first entry
-
 	// Get milvusID before removing from map (for logging)
 	milvusID := h.idMap[victimIdx]
 
-	// Remove from structures
-	delete(h.idMap, victimIdx)
+	// Remove the embedding from the slice
+	h.embeddings = h.embeddings[1:]
+
+	// Rebuild idMap with adjusted indices (all indices shift down by 1)
+	newIDMap := make(map[int]string, len(h.idMap)-1)
+	for idx, id := range h.idMap {
+		if idx > victimIdx {
+			newIDMap[idx-1] = id // Shift index down
+		}
+		// Skip victimIdx (it's being evicted)
+	}
+	h.idMap = newIDMap
 
-	// Note: We don't remove from Milvus (data persists there)
-	// We also don't rebuild HNSW (mark as stale)
+	// Mark HNSW index as stale (needs rebuild with new indices)
 	h.hnswIndex.markStale()
 
 	atomic.AddInt64(&h.evictCount, 1)
 
+	logging.Debugf("HybridCache.evictOne: evicted entry at index %d (milvus_id=%s), new size=%d",
+		victimIdx, milvusID, len(h.embeddings))
 	logging.LogEvent("hybrid_cache_evicted", map[string]interface{}{
 		"backend":     "hybrid",
 		"milvus_id":   milvusID,
 		"hnsw_index":  victimIdx,
+		"new_size":    len(h.embeddings),
 		"max_entries": h.maxMemoryEntries,
 	})
 }

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

@@ -374,7 +374,7 @@ func (c *MilvusCache) createCollection() error {
 	}
 
 	// Create index with updated API
-	index, err := entity.NewIndexHNSW(entity.MetricType(c.config.Collection.VectorField.MetricType), c.config.Collection.Index.Params.EfConstruction, c.config.Collection.Index.Params.M)
+	index, err := entity.NewIndexHNSW(entity.MetricType(c.config.Collection.VectorField.MetricType), c.config.Collection.Index.Params.M, c.config.Collection.Index.Params.EfConstruction)
 	if err != nil {
 		return fmt.Errorf("failed to create HNSW index: %w", err)
 	}
@@ -843,11 +843,12 @@ func (c *MilvusCache) GetByID(ctx context.Context, requestID string) ([]byte, er
 	logging.Debugf("MilvusCache.GetByID: fetching requestID='%s'", requestID)
 
 	// Query Milvus by request_id (primary key)
+	// Filter for non-empty responses to avoid race condition with pending entries
 	queryResult, err := c.client.Query(
 		ctx,
 		c.collectionName,
 		[]string{}, // Empty partitions means search all
-		fmt.Sprintf("request_id == \"%s\"", requestID),
+		fmt.Sprintf("request_id == \"%s\" && response_body != \"\"", requestID),
 		[]string{"response_body"}, // Only fetch document, not embedding!
 	)
 	if err != nil {

tools/make/build-run-test.mk

@@ -37,6 +37,7 @@ test-semantic-router: build-router
 	@$(LOG_TARGET)
 	@export LD_LIBRARY_PATH=${PWD}/candle-binding/target/release && \
 	export SKIP_MILVUS_TESTS=$${SKIP_MILVUS_TESTS:-true} && \
+	export SR_TEST_MODE=true && \
 		cd src/semantic-router && CGO_ENABLED=1 go test -v ./...
 
 # Test the Rust library and the Go binding