feat: intergrate gemini embedding (#783)

backend/application/base/appinfra/app_infra.go

@@ -25,8 +25,10 @@ import (
 	"strings"
 	"time"
 
+	"google.golang.org/genai"
 	"gorm.io/gorm"
 
+	"github.com/cloudwego/eino-ext/components/embedding/gemini"
 	"github.com/cloudwego/eino-ext/components/embedding/ollama"
 	"github.com/cloudwego/eino-ext/components/embedding/openai"
 	"github.com/milvus-io/milvus/client/v2/milvusclient"
@@ -499,7 +501,61 @@ func getEmbedding(ctx context.Context) (embedding.Embedder, error) {
 		if err != nil {
 			return nil, fmt.Errorf("init ollama embedding failed, err=%w", err)
 		}
+	case "gemini":
+		var (
+			geminiEmbeddingBaseURL  = os.Getenv("GEMINI_EMBEDDING_BASE_URL")
+			geminiEmbeddingModel    = os.Getenv("GEMINI_EMBEDDING_MODEL")
+			geminiEmbeddingApiKey   = os.Getenv("GEMINI_EMBEDDING_API_KEY")
+			geminiEmbeddingDims     = os.Getenv("GEMINI_EMBEDDING_DIMS")
+			geminiEmbeddingBackend  = os.Getenv("GEMINI_EMBEDDING_BACKEND") // "1" for BackendGeminiAPI / "2" for BackendVertexAI
+			geminiEmbeddingProject  = os.Getenv("GEMINI_EMBEDDING_PROJECT")
+			geminiEmbeddingLocation = os.Getenv("GEMINI_EMBEDDING_LOCATION")
+		)
 
+		if len(geminiEmbeddingModel) == 0 {
+			return nil, fmt.Errorf("GEMINI_EMBEDDING_MODEL environment variable is required")
+		}
+		if len(geminiEmbeddingApiKey) == 0 {
+			return nil, fmt.Errorf("GEMINI_EMBEDDING_API_KEY environment variable is required")
+		}
+		if len(geminiEmbeddingDims) == 0 {
+			return nil, fmt.Errorf("GEMINI_EMBEDDING_DIMS environment variable is required")
+		}
+		if len(geminiEmbeddingBackend) == 0 {
+			return nil, fmt.Errorf("GEMINI_EMBEDDING_BACKEND environment variable is required")
+		}
+
+		dims, convErr := strconv.ParseInt(geminiEmbeddingDims, 10, 64)
+		if convErr != nil {
+			return nil, fmt.Errorf("invalid GEMINI_EMBEDDING_DIMS value: %s, err=%w", geminiEmbeddingDims, convErr)
+		}
+
+		backend, convErr := strconv.ParseInt(geminiEmbeddingBackend, 10, 64)
+		if convErr != nil {
+			return nil, fmt.Errorf("invalid GEMINI_EMBEDDING_BACKEND value: %s, err=%w", geminiEmbeddingBackend, convErr)
+		}
+
+		geminiCli, err := genai.NewClient(ctx, &genai.ClientConfig{
+			APIKey:   geminiEmbeddingApiKey,
+			Backend:  genai.Backend(backend),
+			Project:  geminiEmbeddingProject,
+			Location: geminiEmbeddingLocation,
+			HTTPOptions: genai.HTTPOptions{
+				BaseURL: geminiEmbeddingBaseURL,
+			},
+		})
+		if err != nil {
+			return nil, fmt.Errorf("init gemini client failed, err=%w", err)
+		}
+
+		emb, err = wrap.NewGeminiEmbedder(ctx, &gemini.EmbeddingConfig{
+			Client:               geminiCli,
+			Model:                geminiEmbeddingModel,
+			OutputDimensionality: ptr.Of(int32(dims)),
+		}, dims, batchSize)
+		if err != nil {
+			return nil, fmt.Errorf("init gemini embedding failed, err=%w", err)
+		}
 	case "http":
 		var (
 			httpEmbeddingBaseURL = os.Getenv("HTTP_EMBEDDING_ADDR")

backend/go.mod

@@ -74,7 +74,7 @@ require (
 	github.com/yuin/goldmark v1.4.13
 	golang.org/x/image v0.22.0
 	golang.org/x/oauth2 v0.23.0
-	google.golang.org/genai v1.13.0
+	google.golang.org/genai v1.18.0
 )
 
 require (
@@ -85,6 +85,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.37 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.7.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.18 // indirect
+	github.com/cloudwego/eino-ext/components/embedding/gemini v0.0.0-20250814083140-54b99ff82f8e // indirect
 	github.com/cloudwego/gopkg v0.1.4 // indirect
 	github.com/evanphx/json-patch v4.12.0+incompatible // indirect
 	github.com/extrame/ole2 v0.0.0-20160812065207-d69429661ad7 // indirect

backend/go.sum

@@ -928,6 +928,8 @@ github.com/cloudwego/eino v0.3.55 h1:lMZrGtEh0k3qykQTLNXSXuAa98OtF2tS43GMHyvN7nA
 github.com/cloudwego/eino v0.3.55/go.mod h1:wUjz990apdsaOraOXdh6CdhVXq8DJsOvLsVlxNTcNfY=
 github.com/cloudwego/eino-ext/components/embedding/ark v0.1.0 h1:AuJsMdaTXc+dGUDQp82MifLYK8oiJf4gLQPUETmKISM=
 github.com/cloudwego/eino-ext/components/embedding/ark v0.1.0/go.mod h1:0FZG/KRBl3hGWkNsm55UaXyVa6PDVIy5u+QvboAB+cY=
+github.com/cloudwego/eino-ext/components/embedding/gemini v0.0.0-20250814083140-54b99ff82f8e h1:46D2fFDbUysA7kUD5x/wK3huneMEvTQfuWcHqI3M6iQ=
+github.com/cloudwego/eino-ext/components/embedding/gemini v0.0.0-20250814083140-54b99ff82f8e/go.mod h1:mz3PGQenODaRelcH+lmX012PAHT8vnuHsiL6EgFw3FA=
 github.com/cloudwego/eino-ext/components/embedding/ollama v0.0.0-20250728060543-79ec300857b8 h1:uJrs6SmfYnca8A+k9+3qJ4MYwYHMncUlGac1mYQT+Ak=
 github.com/cloudwego/eino-ext/components/embedding/ollama v0.0.0-20250728060543-79ec300857b8/go.mod h1:nav79aUcd+UR24dLA+7l7RcHCMlg26zbDAKvjONdrw0=
 github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20250522060253-ddb617598b09 h1:C8RjF193iguUuevkuv0q4SC+XGlM/DlJEgic7l8OUAI=
@@ -2614,6 +2616,8 @@ google.golang.org/appengine v1.6.6/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCID
 google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
 google.golang.org/genai v1.13.0 h1:LRhwx5PU+bXhfnXyPEHu2kt9yc+MpvuYbajxSorOJjg=
 google.golang.org/genai v1.13.0/go.mod h1:QPj5NGJw+3wEOHg+PrsWwJKvG6UC84ex5FR7qAYsN/M=
+google.golang.org/genai v1.18.0 h1:fTmK7y30CO0CL8xRyyFSjTkd1MNbYUeFUehvDyU/2gQ=
+google.golang.org/genai v1.18.0/go.mod h1:QPj5NGJw+3wEOHg+PrsWwJKvG6UC84ex5FR7qAYsN/M=
 google.golang.org/genproto v0.0.0-20180518175338-11a468237815/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20190307195333-5fe7a883aa19/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=

backend/infra/impl/embedding/wrap/gemini.go

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * 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.
+ */
+
+package wrap
+
+import (
+	"context"
+
+	"github.com/cloudwego/eino-ext/components/embedding/gemini"
+
+	contract "github.com/coze-dev/coze-studio/backend/infra/contract/embedding"
+)
+
+func NewGeminiEmbedder(ctx context.Context, config *gemini.EmbeddingConfig, dimensions int64, batchSize int) (contract.Embedder, error) {
+	emb, err := gemini.NewEmbedder(ctx, config)
+	if err != nil {
+		return nil, err
+	}
+	return &denseOnlyWrap{dims: dimensions, batchSize: batchSize, Embedder: emb}, nil
+}

docker/.env.debug.example

@@ -101,7 +101,7 @@ export VIKING_DB_MODEL_NAME="" # if vikingdb model name is not set, you need to
 # The Embedding model relied on by knowledge base vectorization does not need to be configured
 # if the vector database comes with built-in Embedding functionality (such as VikingDB). Currently,
 # Coze Studio supports four access methods: openai, ark, ollama, and custom http. Users can simply choose one of them when using
-# embedding type: openai / ark / ollama / http
+# embedding type: openai / ark / ollama / http / gemini
 export EMBEDDING_TYPE="ark"
 export EMBEDDING_MAX_BATCH_SIZE=100
 
@@ -126,6 +126,15 @@ export OLLAMA_EMBEDDING_BASE_URL="" # (string, required) Ollama embedding base_u
 export OLLAMA_EMBEDDING_MODEL=""    # (string, required) Ollama embedding model
 export OLLAMA_EMBEDDING_DIMS=""     # (int,    required) Ollama embedding dimensions
 
+# gemini embedding
+export GEMINI_EMBEDDING_BASE_URL=""                  # (string, required) Gemini embedding base_url
+export GEMINI_EMBEDDING_MODEL="gemini-embedding-001" # (string, required) Gemini embedding model.
+export GEMINI_EMBEDDING_API_KEY=""                   # (string, required) Gemini embedding api_key
+export GEMINI_EMBEDDING_DIMS=2048                    # (int,    required) Gemini embedding dimensions
+export GEMINI_EMBEDDING_BACKEND="1"                  # (string, required) Gemini embedding backend, should be "1" for BackendGeminiAPI / "2" for BackendVertexAI.
+export GEMINI_EMBEDDING_PROJECT=""                   # (string, optional) Gemini embedding project
+export GEMINI_EMBEDDING_LOCATION=""                  # (string, optional) Gemini embedding location
+
 # http embedding
 export HTTP_EMBEDDING_ADDR=""   # (string, required) http embedding address
 export HTTP_EMBEDDING_DIMS=1024 # (string, required) http embedding dimensions

docker/.env.example

@@ -99,7 +99,7 @@ export VIKING_DB_MODEL_NAME="" # if vikingdb model name is not set, you need to
 # The Embedding model relied on by knowledge base vectorization does not need to be configured
 # if the vector database comes with built-in Embedding functionality (such as VikingDB). Currently,
 # Coze Studio supports four access methods: openai, ark, ollama, and custom http. Users can simply choose one of them when using
-# embedding type: ark / openai / ollama / http
+# embedding type: ark / openai / ollama / gemini / http
 export EMBEDDING_TYPE="ark"
 export EMBEDDING_MAX_BATCH_SIZE=100
 
@@ -124,6 +124,15 @@ export OLLAMA_EMBEDDING_BASE_URL="" # (string, required) Ollama embedding base_u
 export OLLAMA_EMBEDDING_MODEL=""    # (string, required) Ollama embedding model
 export OLLAMA_EMBEDDING_DIMS=""     # (int,    required) Ollama embedding dimensions
 
+# gemini embedding
+export GEMINI_EMBEDDING_BASE_URL=""                  # (string, required) Gemini embedding base_url
+export GEMINI_EMBEDDING_MODEL="gemini-embedding-001" # (string, required) Gemini embedding model.
+export GEMINI_EMBEDDING_API_KEY=""                   # (string, required) Gemini embedding api_key
+export GEMINI_EMBEDDING_DIMS=2048                    # (int,    required) Gemini embedding dimensions
+export GEMINI_EMBEDDING_BACKEND="1"                  # (string, required) Gemini embedding backend, should be "1" for BackendGeminiAPI / "2" for BackendVertexAI.
+export GEMINI_EMBEDDING_PROJECT=""                   # (string, optional) Gemini embedding project
+export GEMINI_EMBEDDING_LOCATION=""                  # (string, optional) Gemini embedding location
+
 # http embedding
 export HTTP_EMBEDDING_ADDR=""   # (string, required) http embedding address
 export HTTP_EMBEDDING_DIMS=1024 # (string, required) http embedding dimensions