|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "log" |
| 7 | + "time" |
| 8 | + |
| 9 | + candle_binding "github.com/vllm-project/semantic-router/candle-binding" |
| 10 | + "github.com/vllm-project/semantic-router/src/semantic-router/pkg/vectorstore" |
| 11 | +) |
| 12 | + |
| 13 | +type doc struct { |
| 14 | + id, fileID, filename, content string |
| 15 | +} |
| 16 | + |
| 17 | +var sampleDocs = []doc{ |
| 18 | + // Europe |
| 19 | + {"c1", "f1", "france.txt", "The capital of France is Paris. It is known for the Eiffel Tower."}, |
| 20 | + {"c2", "f1", "france.txt", "France is a country in Western Europe with a rich cultural heritage."}, |
| 21 | + {"c3", "f2", "germany.txt", "Berlin is the capital of Germany. It is famous for the Brandenburg Gate."}, |
| 22 | + {"c4", "f2", "germany.txt", "Germany is the largest economy in Europe and a leader in engineering."}, |
| 23 | + // Asia |
| 24 | + {"c5", "f3", "japan.txt", "Tokyo is the capital of Japan. It is one of the most populous cities in the world."}, |
| 25 | + {"c6", "f3", "japan.txt", "Japan is an island nation in East Asia known for its technology and cuisine."}, |
| 26 | + {"c7", "f4", "india.txt", "New Delhi is the capital of India. Mumbai is the most populated city in India."}, |
| 27 | + {"c8", "f4", "india.txt", "India is the most populous country in the world with over 1.4 billion people."}, |
| 28 | + {"c9", "f5", "china.txt", "Beijing is the capital of China. Shanghai is the largest city by population."}, |
| 29 | + {"c10", "f5", "china.txt", "China has the second largest economy in the world and a long history of innovation."}, |
| 30 | +} |
| 31 | + |
| 32 | +func main() { |
| 33 | + fmt.Println("Valkey Vector Store Backend Example") |
| 34 | + fmt.Println("====================================") |
| 35 | + |
| 36 | + ctx := context.Background() |
| 37 | + backend := initBackend() |
| 38 | + defer backend.Close() |
| 39 | + |
| 40 | + storeID := fmt.Sprintf("demo_%d", time.Now().UnixNano()) |
| 41 | + createCollection(ctx, backend, storeID) |
| 42 | + defer cleanupCollection(ctx, backend, storeID) |
| 43 | + |
| 44 | + embedAndInsert(ctx, backend, storeID) |
| 45 | + time.Sleep(500 * time.Millisecond) |
| 46 | + runSearches(ctx, backend, storeID) |
| 47 | + runFilteredSearch(ctx, backend, storeID) |
| 48 | + |
| 49 | + fmt.Println("\n✓ Example completed successfully!") |
| 50 | +} |
| 51 | + |
| 52 | +func initBackend() *vectorstore.ValkeyBackend { |
| 53 | + fmt.Println("\n1. Initializing embedding model...") |
| 54 | + if err := candle_binding.InitModel("sentence-transformers/all-MiniLM-L6-v2", true); err != nil { |
| 55 | + log.Fatalf("Failed to initialize embedding model: %v", err) |
| 56 | + } |
| 57 | + fmt.Println("✓ Embedding model initialized") |
| 58 | + |
| 59 | + fmt.Println("\n2. Connecting to Valkey...") |
| 60 | + backend, err := vectorstore.NewValkeyBackend(vectorstore.ValkeyBackendConfig{ |
| 61 | + Host: "localhost", |
| 62 | + Port: 6379, |
| 63 | + CollectionPrefix: "example_vs_", |
| 64 | + MetricType: "COSINE", |
| 65 | + ConnectTimeout: 5, |
| 66 | + }) |
| 67 | + if err != nil { |
| 68 | + log.Fatalf("Failed to connect to Valkey: %v", err) |
| 69 | + } |
| 70 | + fmt.Println("✓ Connected to Valkey") |
| 71 | + return backend |
| 72 | +} |
| 73 | + |
| 74 | +func createCollection(ctx context.Context, backend *vectorstore.ValkeyBackend, storeID string) { |
| 75 | + dimension := 384 |
| 76 | + fmt.Printf("\n3. Creating collection %q (dimension=%d)...\n", storeID, dimension) |
| 77 | + if err := backend.CreateCollection(ctx, storeID, dimension); err != nil { |
| 78 | + log.Fatalf("Failed to create collection: %v", err) |
| 79 | + } |
| 80 | + fmt.Println("✓ Collection created") |
| 81 | +} |
| 82 | + |
| 83 | +func cleanupCollection(ctx context.Context, backend *vectorstore.ValkeyBackend, storeID string) { |
| 84 | + fmt.Printf("\n7. Cleaning up collection %q...\n", storeID) |
| 85 | + if err := backend.DeleteCollection(ctx, storeID); err != nil { |
| 86 | + log.Printf("Warning: cleanup failed: %v", err) |
| 87 | + } else { |
| 88 | + fmt.Println("✓ Collection deleted") |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +func embedAndInsert(ctx context.Context, backend *vectorstore.ValkeyBackend, storeID string) { |
| 93 | + fmt.Println("\n4. Embedding and inserting documents...") |
| 94 | + chunks := make([]vectorstore.EmbeddedChunk, 0, len(sampleDocs)) |
| 95 | + for i, d := range sampleDocs { |
| 96 | + embedding, err := candle_binding.GetEmbedding(d.content, 0) |
| 97 | + if err != nil { |
| 98 | + log.Fatalf("Failed to embed document %d: %v", i, err) |
| 99 | + } |
| 100 | + chunks = append(chunks, vectorstore.EmbeddedChunk{ |
| 101 | + ID: d.id, FileID: d.fileID, Filename: d.filename, |
| 102 | + Content: d.content, Embedding: embedding, |
| 103 | + ChunkIndex: i, VectorStoreID: storeID, |
| 104 | + }) |
| 105 | + } |
| 106 | + if err := backend.InsertChunks(ctx, storeID, chunks); err != nil { |
| 107 | + log.Fatalf("Failed to insert chunks: %v", err) |
| 108 | + } |
| 109 | + fmt.Printf("✓ Inserted %d chunks\n", len(chunks)) |
| 110 | +} |
| 111 | + |
| 112 | +func runSearches(ctx context.Context, backend *vectorstore.ValkeyBackend, storeID string) { |
| 113 | + fmt.Println("\n5. Searching for similar documents (threshold=0.80)...") |
| 114 | + for _, query := range []string{ |
| 115 | + "What is the capital of France?", |
| 116 | + "Tell me about German engineering", |
| 117 | + "Most populated city in Asia", |
| 118 | + } { |
| 119 | + fmt.Printf("\n Query: %q\n", query) |
| 120 | + qEmb, err := candle_binding.GetEmbedding(query, 0) |
| 121 | + if err != nil { |
| 122 | + log.Fatalf("Failed to embed query: %v", err) |
| 123 | + } |
| 124 | + results, err := backend.Search(ctx, storeID, qEmb, 3, 0.80, nil) |
| 125 | + if err != nil { |
| 126 | + log.Fatalf("Search failed: %v", err) |
| 127 | + } |
| 128 | + if len(results) == 0 { |
| 129 | + fmt.Println(" (no results above threshold)") |
| 130 | + } |
| 131 | + for rank, r := range results { |
| 132 | + fmt.Printf(" #%d [%.4f] %s: %s\n", rank+1, r.Score, r.Filename, truncate(r.Content, 70)) |
| 133 | + } |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +func runFilteredSearch(ctx context.Context, backend *vectorstore.ValkeyBackend, storeID string) { |
| 138 | + fmt.Println("\n6. Searching with file_id filter (only germany.txt)...") |
| 139 | + qEmb, err := candle_binding.GetEmbedding("capital city", 0) |
| 140 | + if err != nil { |
| 141 | + log.Fatalf("Failed to embed query: %v", err) |
| 142 | + } |
| 143 | + results, err := backend.Search(ctx, storeID, qEmb, 5, 0.0, map[string]interface{}{"file_id": "f2"}) |
| 144 | + if err != nil { |
| 145 | + log.Fatalf("Filtered search failed: %v", err) |
| 146 | + } |
| 147 | + for rank, r := range results { |
| 148 | + fmt.Printf(" #%d [%.4f] %s: %s\n", rank+1, r.Score, r.Filename, truncate(r.Content, 70)) |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +func truncate(s string, maxLen int) string { |
| 153 | + if len(s) <= maxLen { |
| 154 | + return s |
| 155 | + } |
| 156 | + return s[:maxLen-3] + "..." |
| 157 | +} |
0 commit comments