youngkingdom/coze-studio
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: intergrate gemini embedding (#783)

Browse Source
This commit is contained in:
mhowto
2025-08-18 16:44:19 +08:00
committed by GitHub
parent a9b87c188b
commit 85e6926a14
6 changed files with 115 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
backend/application/base/appinfra/app_infra.go
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
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
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 Normal file
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
}