milvus write ok

plugins/milvus/milvus_test.go

@@ -66,16 +66,22 @@ func dropCollection(ctx context.Context, client client.Client, collectionName st
 
 func TestMilvusIntegration(t *testing.T) {
     ctx := context.Background()
-    // Initialize Genkit
-    g := &genkit.Genkit{}
 
     // Initialize Milvus plugin
-    m := &Milvus{
+    ms := Milvus{
         Addr: "54.92.111.204:19530", // Milvus gRPC endpoint
     }
-    err := m.Init(ctx, g)
+
+    // Initialize Genkit with Milvus plugin
+    g, err := genkit.Init(ctx, genkit.WithPlugins(&ms))
     if err != nil {
-        t.Fatalf("Milvus.Init failed: %v", err)
+        t.Fatalf("genkit.Init failed: %v", err)
+    }
+
+    // Get the Milvus client for cleanup
+    m, ok := genkit.LookupPlugin(g, provider).(*Milvus)
+    if !ok {
+        t.Fatalf("Failed to lookup Milvus plugin")
     }
     defer m.client.Close()
 
@@ -87,7 +93,7 @@ func TestMilvusIntegration(t *testing.T) {
         Collection:      collectionName,
         Dimension:       768, // Match mock embedder dimension
         Embedder:        &MockEmbedder{},
-        EmbedderOptions: map[string]interface{}{},
+        EmbedderOptions: map[string]interface{}{}, // Explicitly set as map
     }
 
     // Define indexer and retriever
@@ -116,7 +122,7 @@ func TestMilvusIntegration(t *testing.T) {
             },
         }
         req := &ai.IndexerRequest{Documents: documents}
-        err = indexer.Index(ctx, req)
+        err := indexer.Index(ctx, req)
         if err != nil {
             t.Fatalf("Index failed: %v", err)
         }
@@ -152,7 +158,7 @@ func TestMilvusIntegration(t *testing.T) {
 
     t.Run("Empty Index", func(t *testing.T) {
         req := &ai.IndexerRequest{Documents: []*ai.Document{}}
-        err = indexer.Index(ctx, req)
+        err := indexer.Index(ctx, req)
         assert.NoError(t, err, "Indexing empty documents should succeed")
     })
 
@@ -161,9 +167,219 @@ func TestMilvusIntegration(t *testing.T) {
             Query:   &ai.Document{Content: []*ai.Part{ai.NewTextPart("Hello world")}},
             Options: &RetrieverOptions{MetricType: "INVALID"},
         }
-        _, err = retriever.Retrieve(ctx, queryReq)
+        _, err := retriever.Retrieve(ctx, queryReq)
         assert.Error(t, err, "Should fail with invalid metric type")
         assert.Contains(t, err.Error(), "unsupported metric type")
     })
+
+    t.Run("Invalid Embedder Options", func(t *testing.T) {
+        // Test with invalid EmbedderOptions type
+        invalidCfg := CollectionConfig{
+            Collection:      collectionName + "_invalid",
+            Dimension:       768,
+            Embedder:        &MockEmbedder{},
+            EmbedderOptions: "not-a-map", // Invalid type
+        }
+        _, _, err := DefineIndexerAndRetriever(ctx, g, invalidCfg)
+        assert.Error(t, err, "Should fail with invalid EmbedderOptions type")
+        assert.Contains(t, err.Error(), "EmbedderOptions must be a map[string]interface{}")
+    })
 }
 
+
+
+
+
+
+
+
+
+
+
+// // Copyright 2025 Google LLC
+// //
+// // Licensed under the Apache License, Version 2.0 (the "License");
+// // you may not use this file except in compliance with the License.
+// // You may obtain a copy of the License at
+// //
+// //     http://www.apache.org/licenses/LICENSE-2.0
+// //
+// // Unless required by applicable law or agreed to in writing, software
+// // distributed under the License is distributed on an "AS IS" BASIS,
+// // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// // See the License for the specific language governing permissions and
+// // limitations under the License.
+// //
+// // SPDX-License-Identifier: Apache-2.0
+
+// package milvus
+
+// import (
+// 	"context"
+// 	"fmt"
+// 	"strings"
+// 	"testing"
+// 	"time"
+
+// 	"github.com/firebase/genkit/go/ai"
+// 	"github.com/firebase/genkit/go/genkit"
+// 	"github.com/milvus-io/milvus-sdk-go/v2/client"
+// 	"github.com/test-go/testify/assert"
+// )
+
+// // MockEmbedder is a mock implementation of ai.Embedder.
+// type MockEmbedder struct{}
+
+// func (m *MockEmbedder) Name() string {
+//     return "mock-embedder"
+// }
+
+// func (m *MockEmbedder) Embed(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) {
+//     resp := &ai.EmbedResponse{}
+//     for range req.Input {
+//         // Generate a simple embedding (768-dimensional vector of ones)
+//         embedding := make([]float32, 768)
+//         for i := range embedding {
+//             embedding[i] = 1.0
+//         }
+//         resp.Embeddings = append(resp.Embeddings, &ai.Embedding{Embedding: embedding})
+//     }
+//     return resp, nil
+// }
+
+// // dropCollection cleans up a test collection.
+// func dropCollection(ctx context.Context, client client.Client, collectionName string) error {
+//     exists, err := client.HasCollection(ctx, collectionName)
+//     if err != nil {
+//         return fmt.Errorf("check collection: %w", err)
+//     }
+//     if exists {
+//         err = client.DropCollection(ctx, collectionName)
+//         if err != nil {
+//             return fmt.Errorf("drop collection: %w", err)
+//         }
+//     }
+//     return nil
+// }
+
+// func TestMilvusIntegration(t *testing.T) {
+//     ctx := context.Background()
+
+//     // Initialize Milvus plugin
+//     ms := Milvus{
+//         Addr: "54.92.111.204:19530", // Milvus gRPC endpoint
+//     }
+
+//     // Initialize Genkit with Milvus plugin
+//     g, err := genkit.Init(ctx, genkit.WithPlugins(&ms))
+//     if err != nil {
+//         t.Fatalf("genkit.Init failed: %v", err)
+//     }
+
+//     // Get the Milvus client for cleanup
+//     m, ok := genkit.LookupPlugin(g, provider).(*Milvus)
+//     if !ok {
+//         t.Fatalf("Failed to lookup Milvus plugin")
+//     }
+//     defer m.client.Close()
+
+//     // Generate unique collection name
+//     collectionName := fmt.Sprintf("test_collection_%d", time.Now().UnixNano())
+
+//     // Configure collection
+//     cfg := CollectionConfig{
+//         Collection:      collectionName,
+//         Dimension:       768, // Match mock embedder dimension
+//         Embedder:        &MockEmbedder{},
+//         EmbedderOptions: map[string]interface{}{}, // Explicitly set as map
+//     }
+
+//     // Define indexer and retriever
+//     indexer, retriever, err := DefineIndexerAndRetriever(ctx, g, cfg)
+//     if err != nil {
+//         t.Fatalf("DefineIndexerAndRetriever failed: %v", err)
+//     }
+
+//     // Clean up collection after test
+//     defer func() {
+//         if err := dropCollection(ctx, m.client, collectionName); err != nil {
+//             t.Errorf("Cleanup failed: %v", err)
+//         }
+//     }()
+
+//     t.Run("Index and Retrieve", func(t *testing.T) {
+//         // Index documents
+//         documents := []*ai.Document{
+//             {
+//                 Content:  []*ai.Part{ai.NewTextPart("Hello world")},
+//                 Metadata: map[string]interface{}{"id": int64(1), "category": "greeting"},
+//             },
+//             {
+//                 Content:  []*ai.Part{ai.NewTextPart("AI is amazing")},
+//                 Metadata: map[string]interface{}{"id": int64(2), "category": "example"},
+//             },
+//         }
+//         req := &ai.IndexerRequest{Documents: documents}
+//         err := indexer.Index(ctx, req)
+//         if err != nil {
+//             t.Fatalf("Index failed: %v", err)
+//         }
+
+//         // Wait briefly to ensure Milvus processes the index
+//         time.Sleep(1 * time.Second)
+
+//         // Retrieve documents
+//         queryReq := &ai.RetrieverRequest{
+//             Query: &ai.Document{Content: []*ai.Part{ai.NewTextPart("Hello world")}},
+//             Options: &RetrieverOptions{
+//                 Count:      2,
+//                 MetricType: "L2",
+//             },
+//         }
+//         resp, err := retriever.Retrieve(ctx, queryReq)
+//         if err != nil {
+//             t.Fatalf("Retrieve failed: %v", err)
+//         }
+
+//         // Verify results
+//         assert.NotNil(t, resp, "Response should not be nil")
+//         assert.NotEmpty(t, resp.Documents, "Should return at least one document")
+//         for _, doc := range resp.Documents {
+//             assert.NotEmpty(t, doc.Content[0].Text, "Document text should not be empty")
+//             // Note: Mock embedder returns identical vectors, so results may not be exact
+//             if strings.Contains(doc.Content[0].Text, "Hello world") || strings.Contains(doc.Content[0].Text, "AI is amazing") {
+//                 continue
+//             }
+//             t.Errorf("Unexpected document text: %s", doc.Content[0].Text)
+//         }
+//     })
+
+//     t.Run("Empty Index", func(t *testing.T) {
+//         req := &ai.IndexerRequest{Documents: []*ai.Document{}}
+//         err := indexer.Index(ctx, req)
+//         assert.NoError(t, err, "Indexing empty documents should succeed")
+//     })
+
+//     t.Run("Invalid Retrieve Options", func(t *testing.T) {
+//         queryReq := &ai.RetrieverRequest{
+//             Query:   &ai.Document{Content: []*ai.Part{ai.NewTextPart("Hello world")}},
+//             Options: &RetrieverOptions{MetricType: "INVALID"},
+//         }
+//         _, err := retriever.Retrieve(ctx, queryReq)
+//         assert.Error(t, err, "Should fail with invalid metric type")
+//         assert.Contains(t, err.Error(), "unsupported metric type")
+//     })
+
+//     t.Run("Invalid Embedder Options", func(t *testing.T) {
+//         // Test with invalid EmbedderOptions type
+//         invalidCfg := CollectionConfig{
+//             Collection:      collectionName + "_invalid",
+//             Dimension:       768,
+//             Embedder:        &MockEmbedder{},
+//             EmbedderOptions: "not-a-map", // Invalid type
+//         }
+//         _, _, err := DefineIndexerAndRetriever(ctx, g, invalidCfg)
+//         assert.Error(t, err, "Should fail with invalid EmbedderOptions type")
+//         assert.Contains(t, err.Error(), "EmbedderOptions must be a map[string]interface{}")
+//     })
+// }
\ No newline at end of file