diff --git a/internal/entity/models/google.go b/internal/entity/models/google.go index 2702a04384..7c221bac4c 100644 --- a/internal/entity/models/google.go +++ b/internal/entity/models/google.go @@ -48,11 +48,8 @@ func collectGoogleModelNames(ctx context.Context, listPage func(context.Context, } } -var googleListModels = func(ctx context.Context, apiKey string) ([]string, error) { - client, err := genai.NewClient(ctx, &genai.ClientConfig{ - APIKey: apiKey, - Backend: genai.BackendGeminiAPI, - }) +var googleListModels = func(ctx context.Context, config *genai.ClientConfig) ([]string, error) { + client, err := genai.NewClient(ctx, config) if err != nil { return nil, err } @@ -86,27 +83,40 @@ func NewGoogleModel(baseURL map[string]string, urlSuffix URLSuffix) *GoogleModel } func (g *GoogleModel) NewInstance(baseURL map[string]string) ModelDriver { - return nil + return NewGoogleModel(baseURL, g.URLSuffix) } func (g *GoogleModel) Name() string { return "google" } +func (g *GoogleModel) clientConfig(apiKey string, apiConfig *APIConfig) *genai.ClientConfig { + return &genai.ClientConfig{APIKey: apiKey, Backend: genai.BackendGeminiAPI, HTTPOptions: genai.HTTPOptions{BaseURL: g.baseURL(apiConfig)}} +} + +func (g *GoogleModel) baseURL(apiConfig *APIConfig) string { + if apiConfig != nil && apiConfig.Region != nil { + if baseURL := strings.TrimSpace(g.BaseURL[*apiConfig.Region]); baseURL != "" { + return baseURL + } + } + return strings.TrimSpace(g.BaseURL["default"]) +} + func (g *GoogleModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) { - if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { return nil, fmt.Errorf("api key is nil or empty") } + if strings.TrimSpace(modelName) == "" { + return nil, fmt.Errorf("model name is empty") + } if len(messages) == 0 { return nil, fmt.Errorf("messages is empty") } ctx := context.Background() - client, err := genai.NewClient(ctx, &genai.ClientConfig{ - APIKey: *apiConfig.ApiKey, - Backend: genai.BackendGeminiAPI, - }) + client, err := genai.NewClient(ctx, g.clientConfig(strings.TrimSpace(*apiConfig.ApiKey), apiConfig)) if err != nil { return nil, err } @@ -171,12 +181,18 @@ func (g *GoogleModel) ChatStreamlyWithSender(modelName string, messages []Messag if len(messages) == 0 { return fmt.Errorf("messages is empty") } + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { + return fmt.Errorf("api key is nil or empty") + } + if strings.TrimSpace(modelName) == "" { + return fmt.Errorf("model name is empty") + } + if sender == nil { + return fmt.Errorf("sender is nil") + } ctx := context.Background() - client, err := genai.NewClient(ctx, &genai.ClientConfig{ - APIKey: *apiConfig.ApiKey, - Backend: genai.BackendGeminiAPI, - }) + client, err := genai.NewClient(ctx, g.clientConfig(strings.TrimSpace(*apiConfig.ApiKey), apiConfig)) if err != nil { return err } @@ -262,7 +278,7 @@ func (g *GoogleModel) ChatStreamlyWithSender(modelName string, messages []Messag // Embed 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 (g *GoogleModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) { - if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { return nil, fmt.Errorf("api key is required") } if modelName == nil || *modelName == "" { @@ -275,10 +291,7 @@ func (g *GoogleModel) Embed(modelName *string, texts []string, apiConfig *APICon ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout) defer cancel() - client, err := genai.NewClient(ctx, &genai.ClientConfig{ - APIKey: *apiConfig.ApiKey, - Backend: genai.BackendGeminiAPI, - }) + client, err := genai.NewClient(ctx, g.clientConfig(strings.TrimSpace(*apiConfig.ApiKey), apiConfig)) if err != nil { return nil, fmt.Errorf("failed to create client: %w", err) } @@ -323,7 +336,7 @@ func (g *GoogleModel) ListModels(apiConfig *APIConfig) ([]string, error) { return nil, fmt.Errorf("api key is required") } - return googleListModels(context.Background(), *apiConfig.ApiKey) + return googleListModels(context.Background(), g.clientConfig(strings.TrimSpace(*apiConfig.ApiKey), apiConfig)) } func (g *GoogleModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) { diff --git a/internal/entity/models/google_test.go b/internal/entity/models/google_test.go index 5b09c7a168..c9fe8bac9b 100644 --- a/internal/entity/models/google_test.go +++ b/internal/entity/models/google_test.go @@ -7,11 +7,13 @@ import ( "strings" "sync" "testing" + + "google.golang.org/genai" ) var googleListModelsMu sync.Mutex -func withGoogleListModelsStub(t *testing.T, fn func(context.Context, string) ([]string, error)) { +func withGoogleListModelsStub(t *testing.T, fn func(context.Context, *genai.ClientConfig) ([]string, error)) { t.Helper() googleListModelsMu.Lock() @@ -52,7 +54,7 @@ func TestGoogleModelListModelsRequiresAPIKey(t *testing.T) { } calls := 0 - withGoogleListModelsStub(t, func(context.Context, string) ([]string, error) { + withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]string, error) { calls++ return nil, nil }) @@ -80,16 +82,17 @@ func TestGoogleModelListModelsRequiresAPIKey(t *testing.T) { func TestGoogleModelListModelsReturnsModelNames(t *testing.T) { model := &GoogleModel{} apiKey := "test-api-key" + configuredAPIKey := " " + apiKey + " " expected := []string{"models/gemini-2.5-flash", "models/gemini-2.5-pro"} - withGoogleListModelsStub(t, func(_ context.Context, gotAPIKey string) ([]string, error) { - if gotAPIKey != apiKey { - t.Fatalf("expected API key %q, got %q", apiKey, gotAPIKey) + withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]string, error) { + if config.APIKey != apiKey { + t.Fatalf("expected API key %q, got %q", apiKey, config.APIKey) } return expected, nil }) - models, err := model.ListModels(&APIConfig{ApiKey: &apiKey}) + models, err := model.ListModels(&APIConfig{ApiKey: &configuredAPIKey}) if err != nil { t.Fatalf("expected no error, got %v", err) } @@ -99,14 +102,18 @@ func TestGoogleModelListModelsReturnsModelNames(t *testing.T) { } func TestGoogleModelCheckConnectionUsesListModels(t *testing.T) { - model := &GoogleModel{} + customBaseURL := "https://check-connection.example.test/google" + model := NewGoogleModel(map[string]string{"default": customBaseURL}, URLSuffix{}) apiKey := "test-api-key" calls := 0 - withGoogleListModelsStub(t, func(_ context.Context, gotAPIKey string) ([]string, error) { + withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]string, error) { calls++ - if gotAPIKey != apiKey { - t.Fatalf("expected API key %q, got %q", apiKey, gotAPIKey) + if config.APIKey != apiKey { + t.Fatalf("expected API key %q, got %q", apiKey, config.APIKey) + } + if config.HTTPOptions.BaseURL != customBaseURL { + t.Fatalf("expected base URL %q, got %q", customBaseURL, config.HTTPOptions.BaseURL) } return []string{"models/gemini-2.5-flash"}, nil }) @@ -123,7 +130,7 @@ func TestGoogleModelCheckConnectionRequiresAPIKey(t *testing.T) { model := &GoogleModel{} calls := 0 - withGoogleListModelsStub(t, func(context.Context, string) ([]string, error) { + withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]string, error) { calls++ return nil, nil }) @@ -175,7 +182,7 @@ func TestGoogleModelCheckConnectionReturnsListModelsError(t *testing.T) { apiKey := "test-api-key" listErr := errors.New("list models failed") - withGoogleListModelsStub(t, func(context.Context, string) ([]string, error) { + withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]string, error) { return nil, listErr }) @@ -185,6 +192,141 @@ func TestGoogleModelCheckConnectionReturnsListModelsError(t *testing.T) { } } +func TestGoogleModelChatStreamlyRequiresAPIKey(t *testing.T) { + model := &GoogleModel{} + messages := []Message{{Role: "user", Content: "hello"}} + cases := []struct { + name string + apiConfig *APIConfig + }{ + {name: "nil config"}, + {name: "nil api key", apiConfig: &APIConfig{}}, + {name: "empty api key", apiConfig: &APIConfig{ApiKey: stringPtr("")}}, + {name: "blank api key", apiConfig: &APIConfig{ApiKey: stringPtr(" \t\n ")}}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + err := model.ChatStreamlyWithSender("gemini-2.5-flash", messages, tc.apiConfig, nil, func(*string, *string) error { + t.Errorf("sender should not be called without an API key") + return nil + }) + if err == nil { + t.Fatal("expected an API key error") + } + if !strings.Contains(err.Error(), "api key is nil or empty") { + t.Fatalf("expected API key error, got %v", err) + } + }) + } +} + +func TestGoogleModelChatRequiresModelName(t *testing.T) { + model := &GoogleModel{} + apiKey := "test-api-key" + messages := []Message{{Role: "user", Content: "hello"}} + + response, err := model.ChatWithMessages("", messages, &APIConfig{ApiKey: &apiKey}, nil) + if err == nil { + t.Fatal("expected a model name error") + } + if !strings.Contains(err.Error(), "model name is empty") { + t.Fatalf("expected model name error, got %v", err) + } + if response != nil { + t.Fatalf("expected no response, got %v", response) + } + + err = model.ChatStreamlyWithSender("", messages, &APIConfig{ApiKey: &apiKey}, nil, func(*string, *string) error { + t.Errorf("sender should not be called without a model name") + return nil + }) + if err == nil { + t.Fatal("expected a model name error") + } + if !strings.Contains(err.Error(), "model name is empty") { + t.Fatalf("expected model name error, got %v", err) + } + + err = model.ChatStreamlyWithSender("gemini-2.5-flash", messages, &APIConfig{ApiKey: &apiKey}, nil, nil) + if err == nil { + t.Fatal("expected a sender error") + } + if !strings.Contains(err.Error(), "sender is nil") { + t.Fatalf("expected sender error, got %v", err) + } +} + +func TestGoogleModelNewInstancePreservesCustomBaseURL(t *testing.T) { + model := NewGoogleModel(map[string]string{"default": "https://generativelanguage.googleapis.com"}, URLSuffix{Models: "v1beta/models"}) + customBaseURL := map[string]string{"default": "https://example.test/google"} + + instance := model.NewInstance(customBaseURL) + google, ok := instance.(*GoogleModel) + if !ok { + t.Fatalf("expected *GoogleModel, got %T", instance) + } + if google.BaseURL["default"] != customBaseURL["default"] { + t.Fatalf("expected base URL %q, got %q", customBaseURL["default"], google.BaseURL["default"]) + } + if google.URLSuffix != model.URLSuffix { + t.Fatalf("expected URL suffix %v, got %v", model.URLSuffix, google.URLSuffix) + } +} + +func TestGoogleModelListModelsPassesBaseURL(t *testing.T) { + apiKey := "test-api-key" + cases := []struct { + name string + baseURL map[string]string + region *string + expectedBaseURL string + }{ + { + name: "default custom base URL", + baseURL: map[string]string{"default": "https://example.test/google"}, + expectedBaseURL: "https://example.test/google", + }, + { + name: "regional custom base URL", + baseURL: map[string]string{"east": "https://east.example.test/google", "default": "https://default.example.test/google"}, + region: stringPtr("east"), + expectedBaseURL: "https://east.example.test/google", + }, + { + name: "empty region custom base URL", + baseURL: map[string]string{"": "https://empty-region.example.test/google"}, + region: stringPtr(""), + expectedBaseURL: "https://empty-region.example.test/google", + }, + { + name: "missing region falls back to default base URL", + baseURL: map[string]string{"default": "https://default.example.test/google"}, + region: stringPtr("missing"), + expectedBaseURL: "https://default.example.test/google", + }, + { + name: "SDK default base URL", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + model := NewGoogleModel(tc.baseURL, URLSuffix{}) + withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]string, error) { + if config.HTTPOptions.BaseURL != tc.expectedBaseURL { + t.Fatalf("expected base URL %q, got %q", tc.expectedBaseURL, config.HTTPOptions.BaseURL) + } + return []string{"models/gemini-2.5-flash"}, nil + }) + + if _, err := model.ListModels(&APIConfig{ApiKey: &apiKey, Region: tc.region}); err != nil { + t.Fatalf("expected no error, got %v", err) + } + }) + } +} + func TestCollectGoogleModelNamesPaginates(t *testing.T) { pages := []googleModelPage{ {items: []string{"models/gemini-2.5-flash"}, nextPageToken: "page-2"},