From 85e6926a14059f227f6c94be8957bff17e55fd02 Mon Sep 17 00:00:00 2001 From: mhowto Date: Mon, 18 Aug 2025 16:44:19 +0800 Subject: [PATCH] feat: intergrate gemini embedding (#783) --- .../application/base/appinfra/app_infra.go | 56 +++++++++++++++++++ backend/go.mod | 3 +- backend/go.sum | 4 ++ backend/infra/impl/embedding/wrap/gemini.go | 33 +++++++++++ docker/.env.debug.example | 11 +++- docker/.env.example | 11 +++- 6 files changed, 115 insertions(+), 3 deletions(-) create mode 100644 backend/infra/impl/embedding/wrap/gemini.go diff --git a/backend/application/base/appinfra/app_infra.go b/backend/application/base/appinfra/app_infra.go index d516d4d0..e08b356b 100644 --- a/backend/application/base/appinfra/app_infra.go +++ b/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") diff --git a/backend/go.mod b/backend/go.mod index f06a9352..118e8f48 100755 --- a/backend/go.mod +++ b/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 diff --git a/backend/go.sum b/backend/go.sum index 39e3cc96..31515acd 100644 --- a/backend/go.sum +++ b/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= diff --git a/backend/infra/impl/embedding/wrap/gemini.go b/backend/infra/impl/embedding/wrap/gemini.go new file mode 100644 index 00000000..44111288 --- /dev/null +++ b/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 +} diff --git a/docker/.env.debug.example b/docker/.env.debug.example index 131461ed..b4a7dd6e 100644 --- a/docker/.env.debug.example +++ b/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 diff --git a/docker/.env.example b/docker/.env.example index 3f739cf3..d0cdd3ec 100644 --- a/docker/.env.example +++ b/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