mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-29 15:31:05 +08:00
Go: implement Encode (embeddings) in Google Gemini driver (#14682)
### What problem does this PR solve? - Implements the `Encode` method in the Google Gemini driver, which was previously a stub returning `not implemented` - Uses the `google.golang.org/genai` SDK's `EmbedContent` API, which routes to the `batchEmbedContents` endpoint internally — all texts are sent in a single request - Adds `text-embedding-004` (max 2048 tokens) to `conf/models/google.json` - Response values are `[]float32` from the SDK and are cast to `[]float64` to satisfy the `ModelDriver` interface ## Files changed - `internal/entity/models/google.go` — full `Encode` implementation - `conf/models/google.json` — adds `text-embedding-004` embedding model ### Type of change - [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
@@ -18,6 +18,13 @@
|
||||
"default_value": true,
|
||||
"clear_thinking": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "text-embedding-004",
|
||||
"max_tokens": 2048,
|
||||
"model_types": [
|
||||
"embedding"
|
||||
]
|
||||
}
|
||||
],
|
||||
"features": {
|
||||
|
||||
@@ -212,9 +212,60 @@ func (z *GoogleModel) ChatStreamlyWithSender(modelName string, messages []Messag
|
||||
return err
|
||||
}
|
||||
|
||||
// Encode encodes a list of texts into embeddings
|
||||
// Encode generates embeddings for a batch of texts using the Gemini embeddings API.
|
||||
// The SDK routes to batchEmbedContents internally, so all texts are sent in one request.
|
||||
func (z *GoogleModel) Encode(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
|
||||
return nil, fmt.Errorf("not implemented")
|
||||
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
|
||||
return nil, fmt.Errorf("api key is required")
|
||||
}
|
||||
if modelName == nil || *modelName == "" {
|
||||
return nil, fmt.Errorf("model name is required")
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
return nil, fmt.Errorf("texts is empty")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout)
|
||||
defer cancel()
|
||||
|
||||
client, err := genai.NewClient(ctx, &genai.ClientConfig{
|
||||
APIKey: *apiConfig.ApiKey,
|
||||
Backend: genai.BackendGeminiAPI,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create client: %w", err)
|
||||
}
|
||||
|
||||
contents := make([]*genai.Content, len(texts))
|
||||
for i, text := range texts {
|
||||
contents[i] = genai.NewContentFromText(text, genai.RoleUser)
|
||||
}
|
||||
|
||||
var cfg *genai.EmbedContentConfig
|
||||
if embeddingConfig != nil && embeddingConfig.Dimension > 0 {
|
||||
dim := int32(embeddingConfig.Dimension)
|
||||
cfg = &genai.EmbedContentConfig{OutputDimensionality: &dim}
|
||||
}
|
||||
|
||||
resp, err := client.Models.EmbedContent(ctx, *modelName, contents, cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to embed content: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Embeddings) != len(texts) {
|
||||
return nil, fmt.Errorf("expected %d embeddings, got %d", len(texts), len(resp.Embeddings))
|
||||
}
|
||||
|
||||
result := make([][]float64, len(resp.Embeddings))
|
||||
for i, emb := range resp.Embeddings {
|
||||
vec := make([]float64, len(emb.Values))
|
||||
for j, v := range emb.Values {
|
||||
vec[j] = float64(v)
|
||||
}
|
||||
result[i] = vec
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (z *GoogleModel) ListModels(apiConfig *APIConfig) ([]string, error) {
|
||||
@@ -245,7 +296,8 @@ func (z *GoogleModel) Balance(apiConfig *APIConfig) (map[string]interface{}, err
|
||||
}
|
||||
|
||||
func (z *GoogleModel) CheckConnection(apiConfig *APIConfig) error {
|
||||
return fmt.Errorf("no such method")
|
||||
_, err := z.ListModels(apiConfig)
|
||||
return err
|
||||
}
|
||||
|
||||
// Rerank calculates similarity scores between query and documents
|
||||
|
||||
Reference in New Issue
Block a user