@@ -25,8 +25,10 @@ import (
2525 "strings"
2626 "time"
2727
28+ "google.golang.org/genai"
2829 "gorm.io/gorm"
2930
31+ "github.com/cloudwego/eino-ext/components/embedding/gemini"
3032 "github.com/cloudwego/eino-ext/components/embedding/ollama"
3133 "github.com/cloudwego/eino-ext/components/embedding/openai"
3234 "github.com/milvus-io/milvus/client/v2/milvusclient"
@@ -499,7 +501,61 @@ func getEmbedding(ctx context.Context) (embedding.Embedder, error) {
499501 if err != nil {
500502 return nil , fmt .Errorf ("init ollama embedding failed, err=%w" , err )
501503 }
504+ case "gemini" :
505+ var (
506+ geminiEmbeddingBaseURL = os .Getenv ("GEMINI_EMBEDDING_BASE_URL" )
507+ geminiEmbeddingModel = os .Getenv ("GEMINI_EMBEDDING_MODEL" )
508+ geminiEmbeddingApiKey = os .Getenv ("GEMINI_EMBEDDING_API_KEY" )
509+ geminiEmbeddingDims = os .Getenv ("GEMINI_EMBEDDING_DIMS" )
510+ geminiEmbeddingBackend = os .Getenv ("GEMINI_EMBEDDING_BACKEND" ) // "1" for BackendGeminiAPI / "2" for BackendVertexAI
511+ geminiEmbeddingProject = os .Getenv ("GEMINI_EMBEDDING_PROJECT" )
512+ geminiEmbeddingLocation = os .Getenv ("GEMINI_EMBEDDING_LOCATION" )
513+ )
514+
515+ if len (geminiEmbeddingModel ) == 0 {
516+ return nil , fmt .Errorf ("GEMINI_EMBEDDING_MODEL environment variable is required" )
517+ }
518+ if len (geminiEmbeddingApiKey ) == 0 {
519+ return nil , fmt .Errorf ("GEMINI_EMBEDDING_API_KEY environment variable is required" )
520+ }
521+ if len (geminiEmbeddingDims ) == 0 {
522+ return nil , fmt .Errorf ("GEMINI_EMBEDDING_DIMS environment variable is required" )
523+ }
524+ if len (geminiEmbeddingBackend ) == 0 {
525+ return nil , fmt .Errorf ("GEMINI_EMBEDDING_BACKEND environment variable is required" )
526+ }
527+
528+ dims , convErr := strconv .ParseInt (geminiEmbeddingDims , 10 , 64 )
529+ if convErr != nil {
530+ return nil , fmt .Errorf ("invalid GEMINI_EMBEDDING_DIMS value: %s, err=%w" , geminiEmbeddingDims , convErr )
531+ }
502532
533+ backend , convErr := strconv .ParseInt (geminiEmbeddingBackend , 10 , 64 )
534+ if convErr != nil {
535+ return nil , fmt .Errorf ("invalid GEMINI_EMBEDDING_BACKEND value: %s, err=%w" , geminiEmbeddingBackend , convErr )
536+ }
537+
538+ geminiCli , err := genai .NewClient (ctx , & genai.ClientConfig {
539+ APIKey : geminiEmbeddingApiKey ,
540+ Backend : genai .Backend (backend ),
541+ Project : geminiEmbeddingProject ,
542+ Location : geminiEmbeddingLocation ,
543+ HTTPOptions : genai.HTTPOptions {
544+ BaseURL : geminiEmbeddingBaseURL ,
545+ },
546+ })
547+ if err != nil {
548+ return nil , fmt .Errorf ("init gemini client failed, err=%w" , err )
549+ }
550+
551+ emb , err = wrap .NewGeminiEmbedder (ctx , & gemini.EmbeddingConfig {
552+ Client : geminiCli ,
553+ Model : geminiEmbeddingModel ,
554+ OutputDimensionality : ptr .Of (int32 (dims )),
555+ }, dims , batchSize )
556+ if err != nil {
557+ return nil , fmt .Errorf ("init gemini embedding failed, err=%w" , err )
558+ }
503559 case "http" :
504560 var (
505561 httpEmbeddingBaseURL = os .Getenv ("HTTP_EMBEDDING_ADDR" )
0 commit comments