diff --git a/client.go b/client.go index 3e3db755..bbf85a0e 100644 --- a/client.go +++ b/client.go @@ -293,6 +293,8 @@ func NewInternalAPIClient(ctx context.Context, cc *ClientConfig) (*InternalAPICl if cc.HTTPOptions.BaseURL == "" { if cc.Location == "global" || cc.APIKey != "" { cc.HTTPOptions.BaseURL = "https://aiplatform.googleapis.com/" + } else if cc.Location == "us" { + cc.HTTPOptions.BaseURL = fmt.Sprintf("https://aiplatform.%s.rep.googleapis.com/", cc.Location) } else { cc.HTTPOptions.BaseURL = fmt.Sprintf("https://%s-aiplatform.googleapis.com/", cc.Location) } diff --git a/client_test.go b/client_test.go index 03f639bb..d29cfb2a 100644 --- a/client_test.go +++ b/client_test.go @@ -963,6 +963,32 @@ func TestClientConfigHTTPOptions(t *testing.T) { expectedBaseURL: "https://aiplatform.googleapis.com/", expectedAPIVersion: "v1beta1", }, + { + name: "Vertex AI Backend with us location", + clientConfig: ClientConfig{ + Backend: BackendVertexAI, + Project: "test-project", + Location: "us", + HTTPOptions: HTTPOptions{}, + Credentials: &auth.Credentials{}, + }, + expectedBaseURL: "https://aiplatform.us.rep.googleapis.com/", + expectedAPIVersion: "v1beta1", + }, + { + name: "Vertex AI Backend with us location and base URL override", + clientConfig: ClientConfig{ + Backend: BackendVertexAI, + Project: "test-project", + Location: "us", + HTTPOptions: HTTPOptions{ + BaseURL: "https://my-custom-url.com/", + }, + Credentials: &auth.Credentials{}, + }, + expectedBaseURL: "https://my-custom-url.com/", + expectedAPIVersion: "v1beta1", + }, { name: "Google AI Backend with HTTP Client Timeout and no HTTPOptions", clientConfig: ClientConfig{