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 {