From 4db525958130bd8eb9b38d2b9332f5c954a821a7 Mon Sep 17 00:00:00 2001
From: miken90 <nguyenminhcanh1990@gmail.com>
Date: Sat, 28 Feb 2026 09:16:00 +0700
Subject: [PATCH] fix(memory): update Gemini embedding to gemini-embedding-001
 with dimension support (#10)

- Replace deprecated text-embedding-004 with gemini-embedding-001
- Add WithDimensions() to OpenAIEmbeddingProvider for dimension truncation
- Pass dimensions=1536 for Gemini to match pgvector column size
- Allow config override of Gemini embedding model via embedding_model
---
 cmd/gateway_agents.go         |  7 ++++++-
 internal/memory/embeddings.go | 18 ++++++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/cmd/gateway_agents.go b/cmd/gateway_agents.go
index a1f85d75..22bec2ac 100644
--- a/cmd/gateway_agents.go
+++ b/cmd/gateway_agents.go
@@ -237,7 +237,12 @@ func createEmbeddingProvider(name string, cfg *config.Config, memCfg *config.Mem
 		if cfg.Providers.Gemini.APIKey == "" {
 			return nil
 		}
-		return memory.NewOpenAIEmbeddingProvider("gemini", cfg.Providers.Gemini.APIKey, "https://generativelanguage.googleapis.com/v1beta/openai", "text-embedding-004")
+		geminiModel := "gemini-embedding-001"
+		if memCfg != nil && memCfg.EmbeddingModel != "" {
+			geminiModel = memCfg.EmbeddingModel
+		}
+		return memory.NewOpenAIEmbeddingProvider("gemini", cfg.Providers.Gemini.APIKey, "https://generativelanguage.googleapis.com/v1beta/openai", geminiModel).
+			WithDimensions(1536)
 	}
 	return nil
 }
diff --git a/internal/memory/embeddings.go b/internal/memory/embeddings.go
index b0959545..20046f22 100644
--- a/internal/memory/embeddings.go
+++ b/internal/memory/embeddings.go
@@ -25,10 +25,11 @@ type EmbeddingProvider interface {
 // OpenAIEmbeddingProvider uses the OpenAI-compatible embedding API.
 // Works with OpenAI, OpenRouter, and any compatible endpoint.
 type OpenAIEmbeddingProvider struct {
-	name   string
-	model  string
-	apiKey string
-	apiURL string
+	name       string
+	model      string
+	apiKey     string
+	apiURL     string
+	dimensions int // optional: truncate output to this many dimensions (0 = use model default)
 }
 
 // NewOpenAIEmbeddingProvider creates a provider for OpenAI-compatible embedding APIs.
@@ -48,6 +49,12 @@ func NewOpenAIEmbeddingProvider(name, apiKey, apiURL, model string) *OpenAIEmbed
 	}
 }
 
+// WithDimensions sets the output dimensions for models that support dimension truncation.
+func (p *OpenAIEmbeddingProvider) WithDimensions(d int) *OpenAIEmbeddingProvider {
+	p.dimensions = d
+	return p
+}
+
 func (p *OpenAIEmbeddingProvider) Name() string  { return p.name }
 func (p *OpenAIEmbeddingProvider) Model() string { return p.model }
 
@@ -56,6 +63,9 @@ func (p *OpenAIEmbeddingProvider) Embed(ctx context.Context, texts []string) ([]
 		"input": texts,
 		"model": p.model,
 	}
+	if p.dimensions > 0 {
+		reqBody["dimensions"] = p.dimensions
+	}
 
 	bodyJSON, err := json.Marshal(reqBody)
 	if err != nil {