Retriever - Milvus v1 (Legacy)

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 is the milvus client to call
    // Required
    Client client.Client

    // Retriever common configuration
    // Collection is the collection name in milvus database
    // Optional, default value is "eino_collection"
    Collection string
    // Partition is the partition name of the collection
    // Optional, default value is empty
    Partition []string
    // VectorField is the vector field name in the collection
    // Optional, default value is "vector"
    VectorField string
    // OutputFields are the fields to return
    // Optional, default value is empty
    OutputFields []string
    // DocumentConverter is the function to convert search result to s.Document
    // Optional, default value is defaultDocumentConverter
    DocumentConverter func(ctx context.Context, doc client.SearchResult) ([]*s.Document, error)
    // MetricType is the metric type of the vector
    // Optional, default value is "HAMMING"
    MetricType entity.MetricType
    // TopK is the top k results to return
    // Optional, default value is 5
    TopK int
    // ScoreThreshold is the threshold of the search result
    // Optional, default value is 0
    ScoreThreshold float64
    // SearchParams
    // Optional, default value is entity.IndexAUTOINDEXSearchParam, level 1
    Sp entity.SearchParam

    // Embedding is the method to embed values from s.Document content that need to be embedded
    // Required
    Embedding embedding.Embedder
}