feat: intergrate gemini embedding (#783)
This commit is contained in:
mhowto
GitHub
parent
a9b87c188b
commit
85e6926a14
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue