1package provider
2
3// Google Vertex AI client.
4//
5// Wire format is identical to the public Generative Language API
6// (generateContent / streamGenerateContent), so we reuse the existing
7// geminiClient and rewrite the outgoing URL + auth via a RoundTripper.
8//
9// URL shape:
10//
11// POST https://{location}-aiplatform.googleapis.com/v1/projects/{project}/
12// locations/{location}/publishers/google/models/{model}:streamGenerateContent?alt=sse
13//
14// Auth: two paths supported.
15//
16// - **API key** (GOOGLE_CLOUD_API_KEY): `x-goog-api-key` header. Simplest
17// setup; works for users who have created a Vertex AI API key in the
18// GCP console. We use this directly, no token exchange.
19// - **Service-account JSON** pointed to by GOOGLE_APPLICATION_CREDENTIALS:
20// we read the JSON, build a JWT signed with RS256 using the private
21// key, exchange it at `https://oauth2.googleapis.com/token` for a
22// short-lived (1h) access token, cache it in memory, and refresh on
23// demand. Sent as `Authorization: Bearer <token>`.
24// - **gcloud user OAuth** (~/.config/gcloud/application_default_credentials.json
25// with `type: "authorized_user"`): we read the JSON, pull the
26// client_id/client_secret/refresh_token, exchange the refresh token
27// at `https://oauth2.googleapis.com/token` for an access token,
28// same cache + refresh logic as service-account.
29//
30// Config env vars (read at construction time):
31// GOOGLE_CLOUD_PROJECT required
32// GOOGLE_CLOUD_LOCATION required (e.g. us-central1)
33// GOOGLE_CLOUD_API_KEY optional; preferred over service-account
34// GOOGLE_APPLICATION_CREDENTIALS optional; service-account JSON path
35
36import (
37 "context"
38 "crypto"
39 "crypto/rand"
40 "crypto/rsa"
41 "crypto/sha256"
42 "crypto/x509"
43 "encoding/base64"
44 "encoding/json"
45 "encoding/pem"
46 "fmt"
47 "io"
48 "net/http"
49 "net/url"
50 "os"
51 "strings"
52 "sync"
53 "time"
54)
55
56type vertexConfig struct {
57 project string
58 location string
59 apiKey string // empty if using OAuth-style auth
60
61 // Auth mode. Exactly one of (saEmail+saPrivateKey) or
62 // (userClientID+userClientSecret+userRefreshToken) is populated when
63 // apiKey is empty.
64 saEmail string
65 saPrivateKey *rsa.PrivateKey
66 saTokenURI string
67
68 userClientID string
69 userClientSecret string
70 userRefreshToken string
71 userTokenURI string
72}
73
74// cacheKey identifies the credential set for the token cache. We key on
75// the auth principal (SA email or user client_id) so multiple configs
76// in the same process don't trample each other.
77func (c *vertexConfig) cacheKey() string {
78 if c.saEmail != "" {
79 return "sa:" + c.saEmail
80 }
81 return "user:" + c.userClientID
82}
83
84func loadVertexConfig() (*vertexConfig, error) {
85 cfg := &vertexConfig{
86 project: os.Getenv("GOOGLE_CLOUD_PROJECT"),
87 location: os.Getenv("GOOGLE_CLOUD_LOCATION"),
88 apiKey: os.Getenv("GOOGLE_CLOUD_API_KEY"),
89 }
90 if cfg.location == "" {
91 cfg.location = "us-central1"
92 }
93 if cfg.project == "" {
94 return nil, fmt.Errorf("vertex: GOOGLE_CLOUD_PROJECT not set")
95 }
96 if cfg.apiKey != "" {
97 return cfg, nil
98 }
99 // Try service-account JSON.
100 credPath := os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
101 if credPath == "" {
102 // ADC default path. Mirrors `gcloud auth application-default login`.
103 if home, err := os.UserHomeDir(); err == nil {
104 candidate := home + "/.config/gcloud/application_default_credentials.json"
105 if _, err := os.Stat(candidate); err == nil {
106 credPath = candidate
107 }
108 }
109 }
110 if credPath == "" {
111 return nil, fmt.Errorf("vertex: no auth — set GOOGLE_CLOUD_API_KEY or GOOGLE_APPLICATION_CREDENTIALS")
112 }
113 b, err := os.ReadFile(credPath)
114 if err != nil {
115 return nil, fmt.Errorf("vertex: read credentials %q: %w", credPath, err)
116 }
117 var raw struct {
118 Type string `json:"type"`
119 ClientEmail string `json:"client_email"`
120 PrivateKey string `json:"private_key"`
121 TokenURI string `json:"token_uri"`
122 ClientID string `json:"client_id"`
123 ClientSecret string `json:"client_secret"`
124 RefreshToken string `json:"refresh_token"`
125 }
126 if err := json.Unmarshal(b, &raw); err != nil {
127 return nil, fmt.Errorf("vertex: parse credentials: %w", err)
128 }
129 switch raw.Type {
130 case "service_account":
131 if raw.ClientEmail == "" || raw.PrivateKey == "" {
132 return nil, fmt.Errorf("vertex: service_account JSON missing client_email or private_key")
133 }
134 block, _ := pem.Decode([]byte(raw.PrivateKey))
135 if block == nil {
136 return nil, fmt.Errorf("vertex: private_key is not PEM-encoded")
137 }
138 key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
139 if err != nil {
140 // Try PKCS#1 as a fallback (older keys).
141 k1, err2 := x509.ParsePKCS1PrivateKey(block.Bytes)
142 if err2 != nil {
143 return nil, fmt.Errorf("vertex: parse private_key: %w", err)
144 }
145 cfg.saPrivateKey = k1
146 } else {
147 rsaKey, ok := key.(*rsa.PrivateKey)
148 if !ok {
149 return nil, fmt.Errorf("vertex: private_key is not RSA")
150 }
151 cfg.saPrivateKey = rsaKey
152 }
153 cfg.saEmail = raw.ClientEmail
154 cfg.saTokenURI = raw.TokenURI
155 if cfg.saTokenURI == "" {
156 cfg.saTokenURI = "https://oauth2.googleapis.com/token"
157 }
158 case "authorized_user":
159 // gcloud auth application-default login produces this shape.
160 if raw.ClientID == "" || raw.ClientSecret == "" || raw.RefreshToken == "" {
161 return nil, fmt.Errorf("vertex: authorized_user JSON missing client_id/client_secret/refresh_token")
162 }
163 cfg.userClientID = raw.ClientID
164 cfg.userClientSecret = raw.ClientSecret
165 cfg.userRefreshToken = raw.RefreshToken
166 cfg.userTokenURI = raw.TokenURI
167 if cfg.userTokenURI == "" {
168 cfg.userTokenURI = "https://oauth2.googleapis.com/token"
169 }
170 default:
171 return nil, fmt.Errorf("vertex: unsupported credentials type %q (expected service_account or authorized_user)", raw.Type)
172 }
173 return cfg, nil
174}
175
176// vertexTokenCache caches one access token per service-account email.
177type vertexTokenCache struct {
178 mu sync.Mutex
179 tokens map[string]struct {
180 value string
181 expiresAt time.Time
182 }
183 http *http.Client
184}
185
186var vertexCache = &vertexTokenCache{
187 tokens: map[string]struct {
188 value string
189 expiresAt time.Time
190 }{},
191 http: &http.Client{Timeout: 30 * time.Second},
192}
193
194func b64url(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }
195
196// signServiceAccountJWT builds and signs a JWT for the
197// `urn:ietf:params:oauth:grant-type:jwt-bearer` token exchange.
198func signServiceAccountJWT(cfg *vertexConfig) (string, error) {
199 header := map[string]string{"alg": "RS256", "typ": "JWT"}
200 now := time.Now().Unix()
201 claims := map[string]any{
202 "iss": cfg.saEmail,
203 "scope": "https://www.googleapis.com/auth/cloud-platform",
204 "aud": cfg.saTokenURI,
205 "exp": now + 3600,
206 "iat": now,
207 }
208 hb, _ := json.Marshal(header)
209 cb, _ := json.Marshal(claims)
210 signingInput := b64url(hb) + "." + b64url(cb)
211 hashed := sha256.Sum256([]byte(signingInput))
212 sig, err := rsa.SignPKCS1v15(rand.Reader, cfg.saPrivateKey, crypto.SHA256, hashed[:])
213 if err != nil {
214 return "", err
215 }
216 return signingInput + "." + b64url(sig), nil
217}
218
219func (c *vertexTokenCache) get(ctx context.Context, cfg *vertexConfig) (string, error) {
220 key := cfg.cacheKey()
221 c.mu.Lock()
222 tok, ok := c.tokens[key]
223 c.mu.Unlock()
224 if ok && time.Now().Before(tok.expiresAt) {
225 return tok.value, nil
226 }
227
228 var form url.Values
229 var endpoint string
230 switch {
231 case cfg.saPrivateKey != nil:
232 jwt, err := signServiceAccountJWT(cfg)
233 if err != nil {
234 return "", fmt.Errorf("vertex: sign jwt: %w", err)
235 }
236 form = url.Values{
237 "grant_type": {"urn:ietf:params:oauth:grant-type:jwt-bearer"},
238 "assertion": {jwt},
239 }
240 endpoint = cfg.saTokenURI
241 case cfg.userRefreshToken != "":
242 form = url.Values{
243 "grant_type": {"refresh_token"},
244 "client_id": {cfg.userClientID},
245 "client_secret": {cfg.userClientSecret},
246 "refresh_token": {cfg.userRefreshToken},
247 }
248 endpoint = cfg.userTokenURI
249 default:
250 return "", fmt.Errorf("vertex: no usable auth (need service_account or authorized_user creds)")
251 }
252
253 req, err := http.NewRequestWithContext(ctx, "POST", endpoint, strings.NewReader(form.Encode()))
254 if err != nil {
255 return "", err
256 }
257 req.Header.Set("content-type", "application/x-www-form-urlencoded")
258 resp, err := c.http.Do(req)
259 if err != nil {
260 return "", fmt.Errorf("vertex: token exchange: %w", err)
261 }
262 defer resp.Body.Close()
263 body, _ := io.ReadAll(resp.Body)
264 if resp.StatusCode != http.StatusOK {
265 return "", fmt.Errorf("vertex: token exchange http %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
266 }
267 var out struct {
268 AccessToken string `json:"access_token"`
269 ExpiresIn int `json:"expires_in"`
270 }
271 if err := json.Unmarshal(body, &out); err != nil {
272 return "", fmt.Errorf("vertex: parse token: %w", err)
273 }
274 if out.AccessToken == "" {
275 return "", fmt.Errorf("vertex: empty access token")
276 }
277 exp := time.Now().Add(time.Duration(out.ExpiresIn-60) * time.Second)
278 c.mu.Lock()
279 c.tokens[key] = struct {
280 value string
281 expiresAt time.Time
282 }{out.AccessToken, exp}
283 c.mu.Unlock()
284 return out.AccessToken, nil
285}
286
287// vertexTransport rewrites the outgoing Gemini URL into the Vertex shape
288// and replaces the x-goog-api-key header with either another
289// x-goog-api-key (for the API-key flow) or Authorization: Bearer (for
290// the service-account flow).
291type vertexTransport struct {
292 inner http.RoundTripper
293 cfg *vertexConfig
294}
295
296func (t *vertexTransport) RoundTrip(req *http.Request) (*http.Response, error) {
297 // Path looks like /v1beta/models/{model}:streamGenerateContent. Pull
298 // out the model id and rewrite.
299 const prefix = "/v1beta/models/"
300 if !strings.HasPrefix(req.URL.Path, prefix) {
301 return nil, fmt.Errorf("vertex: unexpected request path %q", req.URL.Path)
302 }
303 rest := strings.TrimPrefix(req.URL.Path, prefix)
304 // rest = "<model>:streamGenerateContent"
305 sep := strings.IndexByte(rest, ':')
306 if sep < 0 {
307 return nil, fmt.Errorf("vertex: malformed request path %q", req.URL.Path)
308 }
309 modelID := rest[:sep]
310 verb := rest[sep:] // ":streamGenerateContent"
311 newPath := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/google/models/%s%s",
312 t.cfg.project, t.cfg.location, modelID, verb)
313 clone := req.Clone(req.Context())
314 clone.URL.Host = t.cfg.location + "-aiplatform.googleapis.com"
315 clone.URL.Path = newPath
316 // Drop existing auth header and pick the right one.
317 clone.Header.Del("x-goog-api-key")
318 if t.cfg.apiKey != "" {
319 clone.Header.Set("x-goog-api-key", t.cfg.apiKey)
320 } else {
321 tok, err := vertexCache.get(req.Context(), t.cfg)
322 if err != nil {
323 return nil, err
324 }
325 clone.Header.Set("Authorization", "Bearer "+tok)
326 }
327 return t.inner.RoundTrip(clone)
328}
329
330// NewVertex returns a Vertex AI client. The apiKey argument is ignored
331// in favor of env-based config (GOOGLE_CLOUD_API_KEY or
332// GOOGLE_APPLICATION_CREDENTIALS), since Vertex's auth model doesn't fit
333// the "just paste a key" interface other providers use.
334func NewVertex(_ string, _ string) Client {
335 cfg, err := loadVertexConfig()
336 if err != nil {
337 return &unimplementedClient{name: "google-vertex", hint: err.Error()}
338 }
339 inner := &geminiClient{
340 apiKey: "vertex-placeholder", // overwritten by transport
341 baseURL: "https://" + cfg.location + "-aiplatform.googleapis.com",
342 http: &http.Client{
343 Transport: &vertexTransport{inner: http.DefaultTransport, cfg: cfg},
344 Timeout: 0,
345 },
346 }
347 // Wrap so Name() reports "google-vertex" instead of "google".
348 return &renamedClient{inner: inner, name: "google-vertex"}
349}
350
351// renamedClient wraps a Client and overrides its Name(). Used for
352// provider-id remapping (Vertex reuses gemini's wire client but should
353// report as google-vertex for cost / cache / log purposes).
354type renamedClient struct {
355 inner Client
356 name string
357}
358
359func (r *renamedClient) Name() string { return r.name }
360func (r *renamedClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
361 return r.inner.Stream(ctx, req)
362}