Tuzki
/
coze_studio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: intergrate gemini embedding (#783)

Browse Source
main
mhowto 2 months ago committed by GitHub
parent a9b87c188b
commit 85e6926a14
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 115 additions and 3 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 56
      backend/application/base/appinfra/app_infra.go
  2. 3
      backend/go.mod
  3. 4
      backend/go.sum
  4. 33
      backend/infra/impl/embedding/wrap/gemini.go
  5. 11
      docker/.env.debug.example
  6. 11
      docker/.env.example

56
backend/application/base/appinfra/app_infra.go
Unescape View File

@ -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")

3
backend/go.mod
Unescape View File

@ -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

4
backend/go.sum
Unescape View File

@ -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=

33
backend/infra/impl/embedding/wrap/gemini.go
Unescape View File

@ -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
}

11
docker/.env.debug.example
Unescape View File

@ -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

11
docker/.env.example
Unescape View File

@ -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

Write Preview
Loading…
Cancel
Save