Retriever - Milvus

go get github.com/cloudwego/eino-ext/components/retriever/milvus

package main

import (
        "context"
        "fmt"
        "log"
        "os"
        
        "github.com/cloudwego/eino-ext/components/embedding/ark"
        "github.com/milvus-io/milvus-sdk-go/v2/client"
        
        "github.com/cloudwego/eino-ext/components/retriever/milvus"
)

func main() {
        // Get the environment variables
        addr := os.Getenv("MILVUS_ADDR")
        username := os.Getenv("MILVUS_USERNAME")
        password := os.Getenv("MILVUS_PASSWORD")
        arkApiKey := os.Getenv("ARK_API_KEY")
        arkModel := os.Getenv("ARK_MODEL")
        
        // Create a client
        ctx := context.Background()
        cli, err := client.NewClient(ctx, client.Config{
                Address:  addr,
                Username: username,
                Password: password,
        })
        if err != nil {
                log.Fatalf("Failed to create client: %v", err)
                return
        }
        defer cli.Close()
        
        // Create an embedding model
        emb, err := ark.NewEmbedder(ctx, &ark.EmbeddingConfig{
                APIKey: arkApiKey,
                Model:  arkModel,
        })
        
        // Create a retriever
        retriever, err := milvus.NewRetriever(ctx, &milvus.RetrieverConfig{
                Client:      cli,
                Collection:  "",
                Partition:   nil,
                VectorField: "",
                OutputFields: []string{
                        "id",
                        "content",
                        "metadata",
                },
                DocumentConverter: nil,
                MetricType:        "",
                TopK:              0,
                ScoreThreshold:    5,
                Sp:                nil,
                Embedding:         emb,
        })
        if err != nil {
                log.Fatalf("Failed to create retriever: %v", err)
                return
        }
        
        // Retrieve documents
        documents, err := retriever.Retrieve(ctx, "milvus")
        if err != nil {
                log.Fatalf("Failed to retrieve: %v", err)
                return
        }
        
        // Print the documents
        for i, doc := range documents {
                fmt.Printf("Document %d:\n", i)
                fmt.Printf("title: %s\n", doc.ID)
                fmt.Printf("content: %s\n", doc.Content)
                fmt.Printf("metadata: %v\n", doc.MetaData)
        }
}

type RetrieverConfig struct {
    // Client 是要调用的 milvus 客户端
    // 必需
    Client client.Client

    // Retriever 通用配置
    // Collection 是 milvus 数据库中的集合名称
    // 可选，默认值为 "eino_collection"
    Collection string
    // Partition 是集合的分区名称
    // 可选，默认值为空
    Partition []string
    // VectorField 是集合中的向量字段���称
    // 可选，默认值为 "vector"
    VectorField string
    // OutputFields 是要返回的字段
    // 可选，默认值为空
    OutputFields []string
    // DocumentConverter 是将搜索结果转换为 s.Document 的函数
    // 可选，默认值为 defaultDocumentConverter
    DocumentConverter func(ctx context.Context, doc client.SearchResult) ([]*s.Document, error)
    // MetricType 是向量的度量类型
    // 可选，默认值为 "HAMMING"
    MetricType entity.MetricType
    // TopK 是要返回的前 k 个结果
    // 可选，默认值为 5
    TopK int
    // ScoreThreshold 是搜索结果的阈值
    // 可选，默认值为 0
    ScoreThreshold float64
    // SearchParams
    // 可选，默认值为 entity.IndexAUTOINDEXSearchParam，级别为 1
    Sp entity.SearchParam

    // Embedding 是从 s.Document 的内容中嵌入需要嵌入的值的方法
    // 必需的
    Embedding embedding.Embedder
}