From caa7deea76db46ed0b0505eda2e14c021d11584d Mon Sep 17 00:00:00 2001
From: TheShigure7 <2947458856@qq.com>
Date: Sat, 9 May 2026 17:21:37 +0800
Subject: [PATCH 1/2] fix: support gemini-embedding-001 provider config

---
 pkg/embedding/gemini.go       | 29 ++++++++++++++++++++++++++++-
 pkg/embedding/manager_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/pkg/embedding/gemini.go b/pkg/embedding/gemini.go
index f09fc39b..52ccb635 100644
--- a/pkg/embedding/gemini.go
+++ b/pkg/embedding/gemini.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"time"
 )
 
@@ -20,7 +21,10 @@ type GeminiProvider struct {
 type GeminiOption func(*GeminiProvider)
 
 func WithGeminiModel(model string) GeminiOption {
-	return func(p *GeminiProvider) { p.model = model }
+	return func(p *GeminiProvider) {
+		p.model = normalizeGeminiModel(model)
+		p.dimension = geminiModelDimension(p.model)
+	}
 }
 
 func NewGeminiProvider(apiKey string, opts ...GeminiOption) (*GeminiProvider, error) {
@@ -71,6 +75,7 @@ func (p *GeminiProvider) EmbedBatch(ctx context.Context, texts []string) ([][]fl
 
 func (p *GeminiProvider) embedSingle(ctx context.Context, text string) ([]float32, error) {
 	payload := map[string]any{
+		"model": p.model,
 		"content": map[string]any{
 			"parts": []any{
 				map[string]string{"text": text},
@@ -108,3 +113,25 @@ func (p *GeminiProvider) embedSingle(ctx context.Context, text string) ([]float3
 
 	return result.Embedding.Values, nil
 }
+
+func normalizeGeminiModel(model string) string {
+	model = strings.TrimSpace(model)
+	if model == "" {
+		return "models/text-embedding-004"
+	}
+	if !strings.Contains(model, "/") {
+		return "models/" + model
+	}
+	return model
+}
+
+func geminiModelDimension(model string) int {
+	switch normalizeGeminiModel(model) {
+	case "models/gemini-embedding-001":
+		return 3072
+	case "models/text-embedding-004":
+		return 768
+	default:
+		return 0
+	}
+}
diff --git a/pkg/embedding/manager_test.go b/pkg/embedding/manager_test.go
index a726082e..19ea9712 100644
--- a/pkg/embedding/manager_test.go
+++ b/pkg/embedding/manager_test.go
@@ -224,6 +224,34 @@ func TestGeminiProvider(t *testing.T) {
 	}
 }
 
+func TestGeminiProviderGeminiEmbedding001(t *testing.T) {
+	p, err := NewGeminiProvider("test-key", WithGeminiModel("gemini-embedding-001"))
+	if err != nil {
+		t.Fatalf("failed to create provider: %v", err)
+	}
+
+	if p.model != "models/gemini-embedding-001" {
+		t.Fatalf("expected normalized model path, got %q", p.model)
+	}
+	if p.Dimension() != 3072 {
+		t.Fatalf("expected dimension 3072, got %d", p.Dimension())
+	}
+}
+
+func TestGeminiProviderNormalizesModelPath(t *testing.T) {
+	p, err := NewGeminiProvider("test-key", WithGeminiModel("text-embedding-004"))
+	if err != nil {
+		t.Fatalf("failed to create provider: %v", err)
+	}
+
+	if p.model != "models/text-embedding-004" {
+		t.Fatalf("expected normalized model path, got %q", p.model)
+	}
+	if p.Dimension() != 768 {
+		t.Fatalf("expected dimension 768, got %d", p.Dimension())
+	}
+}
+
 func TestOllamaProvider(t *testing.T) {
 	p := NewOllamaProvider(
 		WithOllamaBaseURL("http://localhost:11434"),

From a53e90c93b4122a8402cab5ab7df5c2181e2ce52 Mon Sep 17 00:00:00 2001
From: TheShigure7 <2947458856@qq.com>
Date: Sat, 9 May 2026 23:21:58 +0800
Subject: [PATCH 2/2] fix(embedding): keep default dimension for unknown gemini
 models

---
 pkg/embedding/gemini.go       |  4 +++-
 pkg/embedding/manager_test.go | 14 ++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/pkg/embedding/gemini.go b/pkg/embedding/gemini.go
index 52ccb635..fd8e6cb7 100644
--- a/pkg/embedding/gemini.go
+++ b/pkg/embedding/gemini.go
@@ -23,7 +23,9 @@ type GeminiOption func(*GeminiProvider)
 func WithGeminiModel(model string) GeminiOption {
 	return func(p *GeminiProvider) {
 		p.model = normalizeGeminiModel(model)
-		p.dimension = geminiModelDimension(p.model)
+		if dim := geminiModelDimension(p.model); dim > 0 {
+			p.dimension = dim
+		}
 	}
 }
 
diff --git a/pkg/embedding/manager_test.go b/pkg/embedding/manager_test.go
index 19ea9712..c3bffd6d 100644
--- a/pkg/embedding/manager_test.go
+++ b/pkg/embedding/manager_test.go
@@ -252,6 +252,20 @@ func TestGeminiProviderNormalizesModelPath(t *testing.T) {
 	}
 }
 
+func TestGeminiProviderUnknownModelKeepsDefaultDimension(t *testing.T) {
+	p, err := NewGeminiProvider("test-key", WithGeminiModel("custom-gemini-embed"))
+	if err != nil {
+		t.Fatalf("failed to create provider: %v", err)
+	}
+
+	if p.model != "models/custom-gemini-embed" {
+		t.Fatalf("expected normalized model path, got %q", p.model)
+	}
+	if p.Dimension() != 768 {
+		t.Fatalf("expected default dimension 768 for unknown Gemini models, got %d", p.Dimension())
+	}
+}
+
 func TestOllamaProvider(t *testing.T) {
 	p := NewOllamaProvider(
 		WithOllamaBaseURL("http://localhost:11434"),