fix(go-models): apply custom Google base URLs (#15385)

## Summary - Add custom `base_url` support to the Google Go model driver. - Preserve Google URL suffix configuration when creating custom base URL driver instances. - Validate Google chat/stream request inputs before constructing the SDK client. - Cover Google model listing, connection checks, base URL resolution, and request validation with focused tests. ## What changed - `GoogleModel.NewInstance` now returns a Google driver configured with the supplied base URL map. - Google SDK client creation now resolves configured base URLs through `genai.HTTPOptions.BaseURL`. - Base URL lookup supports configured regions, empty-region keys, and `default` fallback. - Google chat, streaming chat, embeddings, and model listing now reject blank API keys before creating SDK clients. - Google chat and streaming chat now reject blank model names locally, and streaming chat rejects a nil sender. - Existing message handling, embeddings, pagination, and provider errors are preserved. ## Why Google custom model instances could not use configured base URLs because `NewInstance` returned `nil` and the SDK client path ignored the driver base URL map. The request validation keeps invalid Google calls from reaching SDK client construction with blank credentials or incomplete chat inputs.
2026-07-12 22:55:45 +08:00 · 2026-06-01 04:24:29 -07:00
parent fb3bd3de02
commit f0e4f2d5d8
2 changed files with 188 additions and 33 deletions
--- 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) {
--- 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"},