milvus index receiver

go.mod

@@ -11,6 +11,8 @@ require (
 	cloud.google.com/go v0.118.0 // indirect
 	cloud.google.com/go/auth v0.16.0 // indirect
 	cloud.google.com/go/compute/metadata v0.6.0 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/blues/jsonata-go v1.5.4 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
@@ -40,13 +42,24 @@ require (
 	github.com/lib/pq v1.10.9 // indirect
 	github.com/mailru/easyjson v0.9.0 // indirect
 	github.com/mbleigh/raymond v0.0.0-20250414171441-6b3a58ab9e0a // indirect
+	github.com/microsoft/kiota-abstractions-go v1.9.2 // indirect
+	github.com/microsoft/kiota-authentication-azure-go v1.3.0 // indirect
+	github.com/microsoft/kiota-http-go v1.5.2 // indirect
+	github.com/microsoft/kiota-serialization-form-go v1.1.2 // indirect
+	github.com/microsoft/kiota-serialization-json-go v1.1.2 // indirect
+	github.com/microsoft/kiota-serialization-multipart-go v1.1.2 // indirect
+	github.com/microsoft/kiota-serialization-text-go v1.1.2 // indirect
+	github.com/microsoftgraph/msgraph-sdk-go v1.72.0 // indirect
+	github.com/microsoftgraph/msgraph-sdk-go-core v1.3.2 // indirect
 	github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a // indirect
 	github.com/milvus-io/milvus-sdk-go/v2 v2.4.2 // indirect
 	github.com/ollama/ollama v0.6.5 // indirect
 	github.com/pgvector/pgvector-go v0.3.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
-	github.com/rogpeppe/go-internal v1.13.1 // indirect
+	github.com/rogpeppe/go-internal v1.14.1 // indirect
+	github.com/std-uritemplate/std-uritemplate/go/v2 v2.0.3 // indirect
+	github.com/stretchr/testify v1.10.0 // indirect
 	github.com/test-go/testify v1.1.4 // indirect
 	github.com/tidwall/gjson v1.18.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect

go.sum

@@ -6,6 +6,10 @@ cloud.google.com/go/auth v0.16.0/go.mod h1:1howDHJ5IETh/LwYs3ZxvlkXF48aSqqJUM+5o
 cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I=
 cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 h1:Gt0j3wceWMwPmiazCa8MzMA0MfhmPIz0Qp0FJ6qcM0U=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0/go.mod h1:Ot/6aikWnKWi4l9QB7qVSwa8iMphQNqkWALMoNT3rzM=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.0 h1:Bg8m3nq/X1DeePkAbCfb6ml6F3F0IunEhE8TMh+lY48=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.0/go.mod h1:j2chePtV91HrC22tGoRX3sGY42uF13WzmmV80/OdVAA=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/CloudyKit/fastprinter v0.0.0-20200109182630-33d98a066a53/go.mod h1:+3IMCy2vIlbG1XG/0ggNQv0SvxCAIpPM5b1nCz56Xno=
 github.com/CloudyKit/jet/v3 v3.0.0/go.mod h1:HKQPgSJmdK8hdoAbKUUWajkHyHo4RaU5rMdUywE7VMo=
@@ -192,6 +196,24 @@ github.com/mbleigh/raymond v0.0.0-20250414171441-6b3a58ab9e0a h1:v2cBA3xWKv2cIOV
 github.com/mbleigh/raymond v0.0.0-20250414171441-6b3a58ab9e0a/go.mod h1:Y6ghKH+ZijXn5d9E7qGGZBmjitx7iitZdQiIW97EpTU=
 github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
 github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
+github.com/microsoft/kiota-abstractions-go v1.9.2 h1:3U5VgN2YGe3lsu1pyuS0t5jxv1llxX2ophwX8ewE6wQ=
+github.com/microsoft/kiota-abstractions-go v1.9.2/go.mod h1:f06pl3qSyvUHEfVNkiRpXPkafx7khZqQEb71hN/pmuU=
+github.com/microsoft/kiota-authentication-azure-go v1.3.0 h1:PWH6PgtzhJjnmvR6N1CFjriwX09Kv7S5K3vL6VbPVrg=
+github.com/microsoft/kiota-authentication-azure-go v1.3.0/go.mod h1:l/MPGUVvD7xfQ+MYSdZaFPv0CsLDqgSOp8mXwVgArIs=
+github.com/microsoft/kiota-http-go v1.5.2 h1:xqvo4ssWwSvCJw2yuRocKFTxm3Y1iN+a4rrhuTYtBWg=
+github.com/microsoft/kiota-http-go v1.5.2/go.mod h1:L+5Ri+SzwELnUcNA0cpbFKp/pBbvypLh3Cd1PR6sjx0=
+github.com/microsoft/kiota-serialization-form-go v1.1.2 h1:SD6MATqNw+Dc5beILlsb/D87C36HKC/Zw7l+N9+HY2A=
+github.com/microsoft/kiota-serialization-form-go v1.1.2/go.mod h1:m4tY2JT42jAZmgbqFwPy3zGDF+NPJACuyzmjNXeuHio=
+github.com/microsoft/kiota-serialization-json-go v1.1.2 h1:eJrPWeQ665nbjO0gsHWJ0Bw6V/ZHHU1OfFPaYfRG39k=
+github.com/microsoft/kiota-serialization-json-go v1.1.2/go.mod h1:deaGt7fjZarywyp7TOTiRsjfYiyWxwJJPQZytXwYQn8=
+github.com/microsoft/kiota-serialization-multipart-go v1.1.2 h1:1pUyA1QgIeKslQwbk7/ox1TehjlCUUT3r1f8cNlkvn4=
+github.com/microsoft/kiota-serialization-multipart-go v1.1.2/go.mod h1:j2K7ZyYErloDu7Kuuk993DsvfoP7LPWvAo7rfDpdPio=
+github.com/microsoft/kiota-serialization-text-go v1.1.2 h1:7OfKFlzdjpPygca/+OtqafkEqCWR7+94efUFGC28cLw=
+github.com/microsoft/kiota-serialization-text-go v1.1.2/go.mod h1:QNTcswkBPFY3QVBFmzfk00UMNViKQtV0AQKCrRw5ibM=
+github.com/microsoftgraph/msgraph-sdk-go v1.72.0 h1:yKYJ46dJgX6XdI79zI2kotiBPIYktMQs1ad2DVz81rk=
+github.com/microsoftgraph/msgraph-sdk-go v1.72.0/go.mod h1:5ncg4aauxM5XKHo/xvAq7Cjl6+Dqu6lOtoihSGKtDt4=
+github.com/microsoftgraph/msgraph-sdk-go-core v1.3.2 h1:5jCUSosTKaINzPPQXsz7wsHWwknyBmJSu8+ZWxx3kdQ=
+github.com/microsoftgraph/msgraph-sdk-go-core v1.3.2/go.mod h1:iD75MK3LX8EuwjDYCmh0hkojKXK6VKME33u4daCo3cE=
 github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a h1:0B/8Fo66D8Aa23Il0yrQvg1KKz92tE/BJ5BvkUxxAAk=
 github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
 github.com/milvus-io/milvus-sdk-go/v2 v2.4.2 h1:Xqf+S7iicElwYoS2Zly8Nf/zKHuZsNy1xQajfdtygVY=
@@ -230,6 +252,8 @@ github.com/rogpeppe/go-internal v1.8.1/go.mod h1:JeRgkft04UBgHMgCIwADu4Pn6Mtm5d4
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
 github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
+github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
+github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
 github.com/ryanuber/columnize v2.1.0+incompatible/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts=
 github.com/schollz/closestmatch v2.1.0+incompatible/go.mod h1:RtP1ddjLong6gTkbtmuhtR2uUrrJOpYzYRvbcPAid+g=
@@ -244,6 +268,8 @@ github.com/spf13/cobra v0.0.5/go.mod h1:3K3wKZymM7VvHMDS9+Akkh4K60UwM26emMESw8tL
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
 github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
 github.com/spf13/viper v1.3.2/go.mod h1:ZiWeW+zYFKm7srdB9IoDzzZXaJaI5eL9QjNiN/DMA2s=
+github.com/std-uritemplate/std-uritemplate/go/v2 v2.0.3 h1:7hth9376EoQEd1hH4lAp3vnaLP2UMyxuMMghLKzDHyU=
+github.com/std-uritemplate/std-uritemplate/go/v2 v2.0.3/go.mod h1:Z5KcoM0YLC7INlNhEezeIZ0TZNYf7WSNO0Lvah4DSeQ=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=

plugins/milvus/milvus_test.go

@@ -103,11 +103,11 @@ func TestMilvusIntegration(t *testing.T) {
     }
 
     // Clean up collection after test
-    defer func() {
-        if err := dropCollection(ctx, m.client, collectionName); err != nil {
-            t.Errorf("Cleanup failed: %v", err)
-        }
-    }()
+    // 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
@@ -186,200 +186,3 @@ func TestMilvusIntegration(t *testing.T) {
     })
 }
 
-
-
-
-
-
-
-
-
-
-
-// // 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