// Package executor provides runtime execution capabilities for various AI service providers. // This file implements the Vertex AI Gemini executor that talks to Google Vertex AI // endpoints using service account credentials or API keys. package executor import ( "bufio" "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strings" "time" vertexauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/vertex" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" "golang.org/x/oauth2" "golang.org/x/oauth2/google" ) const ( // vertexAPIVersion aligns with current public Vertex Generative AI API. vertexAPIVersion = "v1" // defaultVertexLocation is the default location for Vertex AI resources. defaultVertexLocation = "us-central1" ) // GeminiVertexExecutor sends requests to Vertex AI Gemini endpoints using service account credentials. type GeminiVertexExecutor struct { cfg *config.Config } // NewGeminiVertexExecutor creates a new Vertex AI Gemini executor instance. func NewGeminiVertexExecutor(cfg *config.Config) *GeminiVertexExecutor { return &GeminiVertexExecutor{cfg: cfg} } // Identifier returns the executor identifier. func (e *GeminiVertexExecutor) Identifier() string { return "vertex" } // vertexContext holds resolved configuration and credentials for a request. type vertexContext struct { IsAPIKey bool APIKey string ProjectID string Location string ServiceAccountJSON []byte BaseURL string CompatConfig *config.VertexCompatKey } // PrepareRequest injects Vertex credentials into the outgoing HTTP request. func (e *GeminiVertexExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if req == nil { return nil } ctx := req.Context() vCtx, err := e.resolveVertexContext(auth) if err != nil { return err } return e.setVertexHeaders(ctx, req, vCtx, auth) } // HttpRequest injects Vertex credentials into the request and executes it. func (e *GeminiVertexExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { if req == nil { return nil, fmt.Errorf("vertex executor: request is nil") } if ctx == nil { ctx = req.Context() } httpReq := req.WithContext(ctx) if err := e.PrepareRequest(httpReq, auth); err != nil { return nil, err } httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) return httpClient.Do(httpReq) } // Execute performs a non-streaming request to the Vertex AI API. func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { vCtx, err := e.resolveVertexContext(auth) if err != nil { return resp, err } baseModel := thinking.ParseSuffix(req.Model).ModelName reporter := newUsageReporter(ctx, e.Identifier(), baseModel, auth) defer reporter.trackFailure(ctx, &err) body, _, err := e.preparePayload(ctx, req, opts, baseModel, false) if err != nil { return resp, err } action := getVertexAction(baseModel, false) if req.Metadata != nil { if a, _ := req.Metadata["action"].(string); a == "countTokens" { action = "countTokens" } } url := e.buildVertexURL(vCtx, baseModel, action, opts.Alt) body, _ = sjson.DeleteBytes(body, "session_id") httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) if errNewReq != nil { return resp, errNewReq } if err := e.setVertexHeaders(ctx, httpReq, vCtx, auth); err != nil { return resp, err } e.recordRequest(ctx, httpReq, body, auth) httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { recordAPIResponseError(ctx, e.cfg, errDo) return resp, errDo } defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) } }() recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { return resp, e.handleErrorResponse(ctx, httpResp) } data, errRead := io.ReadAll(httpResp.Body) if errRead != nil { recordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } appendAPIResponseChunk(ctx, e.cfg, data) reporter.publish(ctx, parseGeminiUsage(data)) // Handle Imagen models conversion if isImagenModel(baseModel) { data = convertImagenToGeminiResponse(data, baseModel) } from := opts.SourceFormat to := sdktranslator.FromString("gemini") var param any out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, ¶m) resp = cliproxyexecutor.Response{Payload: []byte(out)} return resp, nil } // ExecuteStream performs a streaming request to the Vertex AI API. func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) { vCtx, err := e.resolveVertexContext(auth) if err != nil { return nil, err } baseModel := thinking.ParseSuffix(req.Model).ModelName reporter := newUsageReporter(ctx, e.Identifier(), baseModel, auth) defer reporter.trackFailure(ctx, &err) body, _, err := e.preparePayload(ctx, req, opts, baseModel, true) if err != nil { return nil, err } action := getVertexAction(baseModel, true) url := e.buildVertexURL(vCtx, baseModel, action, opts.Alt) // Add alt=sse if not overridden and not imagen (imagen doesn't stream) if !isImagenModel(baseModel) && !strings.Contains(url, "alt=") { if strings.Contains(url, "?") { url += "&alt=sse" } else { url += "?alt=sse" } } body, _ = sjson.DeleteBytes(body, "session_id") httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) if errNewReq != nil { return nil, errNewReq } if err := e.setVertexHeaders(ctx, httpReq, vCtx, auth); err != nil { return nil, err } e.recordRequest(ctx, httpReq, body, auth) httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { recordAPIResponseError(ctx, e.cfg, errDo) return nil, errDo } recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { return nil, e.handleErrorResponse(ctx, httpResp) } out := make(chan cliproxyexecutor.StreamChunk) stream = out go func() { defer close(out) defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) } }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, streamScannerBuffer) from := opts.SourceFormat to := sdktranslator.FromString("gemini") var param any for scanner.Scan() { line := scanner.Bytes() appendAPIResponseChunk(ctx, e.cfg, line) if detail, ok := parseGeminiStreamUsage(line); ok { reporter.publish(ctx, detail) } lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), ¶m) for i := range lines { out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])} } } lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, []byte("[DONE]"), ¶m) for i := range lines { out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])} } if errScan := scanner.Err(); errScan != nil { recordAPIResponseError(ctx, e.cfg, errScan) reporter.publishFailure(ctx) out <- cliproxyexecutor.StreamChunk{Err: errScan} } }() return stream, nil } // CountTokens counts tokens for the given request using the Vertex AI API. func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { vCtx, err := e.resolveVertexContext(auth) if err != nil { return cliproxyexecutor.Response{}, err } baseModel := thinking.ParseSuffix(req.Model).ModelName body, _, err := e.preparePayload(ctx, req, opts, baseModel, false) if err != nil { return cliproxyexecutor.Response{}, err } // CountTokens specific cleanup body, _ = sjson.DeleteBytes(body, "tools") body, _ = sjson.DeleteBytes(body, "generationConfig") body, _ = sjson.DeleteBytes(body, "safetySettings") url := e.buildVertexURL(vCtx, baseModel, "countTokens", opts.Alt) respCtx := context.WithValue(ctx, "alt", opts.Alt) httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(body)) if errNewReq != nil { return cliproxyexecutor.Response{}, errNewReq } if err := e.setVertexHeaders(ctx, httpReq, vCtx, auth); err != nil { return cliproxyexecutor.Response{}, err } e.recordRequest(ctx, httpReq, body, auth) httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { recordAPIResponseError(ctx, e.cfg, errDo) return cliproxyexecutor.Response{}, errDo } defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) } }() recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { return cliproxyexecutor.Response{}, e.handleErrorResponse(ctx, httpResp) } data, errRead := io.ReadAll(httpResp.Body) if errRead != nil { recordAPIResponseError(ctx, e.cfg, errRead) return cliproxyexecutor.Response{}, errRead } appendAPIResponseChunk(ctx, e.cfg, data) from := opts.SourceFormat to := sdktranslator.FromString("gemini") count := gjson.GetBytes(data, "totalTokens").Int() out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data) return cliproxyexecutor.Response{Payload: []byte(out)}, nil } // Refresh refreshes the authentication credentials (no-op for Vertex). func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } // resolveVertexContext resolves the execution context (API Key vs Service Account) and configuration. func (e *GeminiVertexExecutor) resolveVertexContext(auth *cliproxyauth.Auth) (*vertexContext, error) { if auth == nil { return nil, fmt.Errorf("vertex executor: auth is nil") } // 1. Try API Key apiKey, baseURL := vertexAPICreds(auth) if apiKey != "" { compat := e.resolveVertexConfig(auth) if compat != nil && compat.BaseURL != "" { baseURL = compat.BaseURL } if baseURL == "" { baseURL = "https://generativelanguage.googleapis.com" } return &vertexContext{ IsAPIKey: true, APIKey: apiKey, BaseURL: baseURL, CompatConfig: compat, }, nil } // 2. Fallback to Service Account projectID, location, saJSON, err := vertexCreds(auth) if err != nil { return nil, err } return &vertexContext{ IsAPIKey: false, ProjectID: projectID, Location: location, ServiceAccountJSON: saJSON, BaseURL: vertexBaseURL(location), }, nil } // buildVertexURL constructs the API URL based on the context and action. func (e *GeminiVertexExecutor) buildVertexURL(vCtx *vertexContext, model, action, alt string) string { var url string if vCtx.IsAPIKey { // API Key: base/v1/publishers/google/models/{model}:{action} base := strings.TrimRight(vCtx.BaseURL, "/") url = fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", base, vertexAPIVersion, model, action) } else { // Service Account: base/v1/projects/{id}/locations/{loc}/publishers/google/models/{model}:{action} base := strings.TrimRight(vCtx.BaseURL, "/") url = fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", base, vertexAPIVersion, vCtx.ProjectID, vCtx.Location, model, action) } if alt != "" && action != "countTokens" { if strings.Contains(url, "?") { url += fmt.Sprintf("&alt=%s", alt) } else { url += fmt.Sprintf("?$alt=%s", alt) } } return url } // setVertexHeaders applies the necessary headers for authentication and configuration. func (e *GeminiVertexExecutor) setVertexHeaders(ctx context.Context, req *http.Request, vCtx *vertexContext, auth *cliproxyauth.Auth) error { req.Header.Set("Content-Type", "application/json") if vCtx.IsAPIKey { req.Header.Set("x-goog-api-key", vCtx.APIKey) req.Header.Del("Authorization") // Apply compat headers if available if vCtx.CompatConfig != nil && len(vCtx.CompatConfig.Headers) > 0 { for k, v := range vCtx.CompatConfig.Headers { req.Header.Set(k, v) } } } else { // Service Account needs Access Token token, err := vertexAccessToken(ctx, e.cfg, auth, vCtx.ServiceAccountJSON) if err != nil { return err } if strings.TrimSpace(token) == "" { return statusErr{code: http.StatusUnauthorized, msg: "missing access token"} } req.Header.Set("Authorization", "Bearer "+token) req.Header.Del("x-goog-api-key") } // Apply headers from auth attributes (overrides compat headers if conflict) applyGeminiHeaders(req, auth) return nil } // preparePayload handles common payload preparation logic. func (e *GeminiVertexExecutor) preparePayload(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string, isStream bool) (body []byte, originalTranslated []byte, err error) { // Handle Imagen models separately if isImagenModel(baseModel) { body, err = convertToImagenRequest(req.Payload) return body, nil, err } from := opts.SourceFormat to := sdktranslator.FromString("gemini") originalPayload := bytes.Clone(req.Payload) if len(opts.OriginalRequest) > 0 { originalPayload = bytes.Clone(opts.OriginalRequest) } originalTranslated = sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, isStream) body = sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), isStream) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { return nil, nil, err } body = fixGeminiImageAspectRatio(baseModel, body) requestedModel := payloadRequestedModel(opts, req.Model) body = applyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel) body, _ = sjson.SetBytes(body, "model", baseModel) return body, originalTranslated, nil } // recordRequest helper to log the upstream request. func (e *GeminiVertexExecutor) recordRequest(ctx context.Context, req *http.Request, body []byte, auth *cliproxyauth.Auth) { var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID authLabel = auth.Label authType, authValue = auth.AccountInfo() } recordAPIRequest(ctx, e.cfg, upstreamRequestLog{ URL: req.URL.String(), Method: req.Method, Headers: req.Header.Clone(), Body: body, Provider: e.Identifier(), AuthID: authID, AuthLabel: authLabel, AuthType: authType, AuthValue: authValue, }) } // handleErrorResponse handles non-2xx responses. func (e *GeminiVertexExecutor) handleErrorResponse(ctx context.Context, resp *http.Response) error { b, _ := io.ReadAll(resp.Body) appendAPIResponseChunk(ctx, e.cfg, b) logWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", resp.StatusCode, summarizeErrorBody(resp.Header.Get("Content-Type"), b)) return statusErr{code: resp.StatusCode, msg: string(b)} } // --- Helper Functions --- // isImagenModel checks if the model name is an Imagen image generation model. func isImagenModel(model string) bool { lowerModel := strings.ToLower(model) return strings.Contains(lowerModel, "imagen") } // getVertexAction returns the appropriate action for the given model. func getVertexAction(model string, isStream bool) string { if isImagenModel(model) { return "predict" } if isStream { return "streamGenerateContent" } return "generateContent" } // convertImagenToGeminiResponse converts Imagen API response to Gemini format. func convertImagenToGeminiResponse(data []byte, model string) []byte { predictions := gjson.GetBytes(data, "predictions") if !predictions.Exists() || !predictions.IsArray() { return data } parts := make([]map[string]any, 0) for _, pred := range predictions.Array() { imageData := pred.Get("bytesBase64Encoded").String() mimeType := pred.Get("mimeType").String() if mimeType == "" { mimeType = "image/png" } if imageData != "" { parts = append(parts, map[string]any{ "inlineData": map[string]any{ "mimeType": mimeType, "data": imageData, }, }) } } responseId := fmt.Sprintf("imagen-%d", time.Now().UnixNano()) response := map[string]any{ "candidates": []map[string]any{{ "content": map[string]any{ "parts": parts, "role": "model", }, "finishReason": "STOP", }}, "responseId": responseId, "modelVersion": model, "usageMetadata": map[string]any{ "promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0, }, } result, err := json.Marshal(response) if err != nil { return data } return result } // convertToImagenRequest converts a Gemini-style request to Imagen API format. func convertToImagenRequest(payload []byte) ([]byte, error) { prompt := "" contentsText := gjson.GetBytes(payload, "contents.0.parts.0.text") if contentsText.Exists() { prompt = contentsText.String() } if prompt == "" { messagesText := gjson.GetBytes(payload, "messages.#.content") if messagesText.Exists() && messagesText.IsArray() { for _, msg := range messagesText.Array() { if msg.String() != "" { prompt = msg.String() break } } } } if prompt == "" { directPrompt := gjson.GetBytes(payload, "prompt") if directPrompt.Exists() { prompt = directPrompt.String() } } if prompt == "" { return nil, fmt.Errorf("imagen: no prompt found in request") } imagenReq := map[string]any{ "instances": []map[string]any{ {"prompt": prompt}, }, "parameters": map[string]any{ "sampleCount": 1, }, } if aspectRatio := gjson.GetBytes(payload, "aspectRatio"); aspectRatio.Exists() { imagenReq["parameters"].(map[string]any)["aspectRatio"] = aspectRatio.String() } if sampleCount := gjson.GetBytes(payload, "sampleCount"); sampleCount.Exists() { imagenReq["parameters"].(map[string]any)["sampleCount"] = int(sampleCount.Int()) } if negativePrompt := gjson.GetBytes(payload, "negativePrompt"); negativePrompt.Exists() { imagenReq["instances"].([]map[string]any)[0]["negativePrompt"] = negativePrompt.String() } return json.Marshal(imagenReq) } // vertexCreds extracts project, location and raw service account JSON from auth metadata. func vertexCreds(a *cliproxyauth.Auth) (projectID, location string, serviceAccountJSON []byte, err error) { if a == nil || a.Metadata == nil { return "", "", nil, fmt.Errorf("vertex executor: missing auth metadata") } if v, ok := a.Metadata["project_id"].(string); ok { projectID = strings.TrimSpace(v) } if projectID == "" { if v, ok := a.Metadata["project"].(string); ok { projectID = strings.TrimSpace(v) } } if projectID == "" { return "", "", nil, fmt.Errorf("vertex executor: missing project_id in credentials") } if v, ok := a.Metadata["location"].(string); ok && strings.TrimSpace(v) != "" { location = strings.TrimSpace(v) } else { location = defaultVertexLocation } var sa map[string]any if raw, ok := a.Metadata["service_account"].(map[string]any); ok { sa = raw } if sa == nil { return "", "", nil, fmt.Errorf("vertex executor: missing service_account in credentials") } normalized, errNorm := vertexauth.NormalizeServiceAccountMap(sa) if errNorm != nil { return "", "", nil, fmt.Errorf("vertex executor: %w", errNorm) } saJSON, errMarshal := json.Marshal(normalized) if errMarshal != nil { return "", "", nil, fmt.Errorf("vertex executor: marshal service_account failed: %w", errMarshal) } return projectID, location, saJSON, nil } // vertexAPICreds extracts API key and base URL from auth attributes. func vertexAPICreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { if a == nil { return "", "" } if a.Attributes != nil { apiKey = a.Attributes["api_key"] baseURL = a.Attributes["base_url"] } if apiKey == "" && a.Metadata != nil { if v, ok := a.Metadata["access_token"].(string); ok { apiKey = v } } return } func vertexBaseURL(location string) string { loc := strings.TrimSpace(location) if loc == "" { loc = defaultVertexLocation } return fmt.Sprintf("https://%s-aiplatform.googleapis.com", loc) } func vertexAccessToken(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, saJSON []byte) (string, error) { if httpClient := newProxyAwareHTTPClient(ctx, cfg, auth, 0); httpClient != nil { ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient) } creds, errCreds := google.CredentialsFromJSON(ctx, saJSON, "https://www.googleapis.com/auth/cloud-platform") if errCreds != nil { return "", fmt.Errorf("vertex executor: parse service account json failed: %w", errCreds) } tok, errTok := creds.TokenSource.Token() if errTok != nil { return "", fmt.Errorf("vertex executor: get access token failed: %w", errTok) } return tok.AccessToken, nil } // resolveVertexConfig finds the matching vertex-api-key configuration entry for the given auth. func (e *GeminiVertexExecutor) resolveVertexConfig(auth *cliproxyauth.Auth) *config.VertexCompatKey { if auth == nil || e.cfg == nil { return nil } var attrKey, attrBase string if auth.Attributes != nil { attrKey = strings.TrimSpace(auth.Attributes["api_key"]) attrBase = strings.TrimSpace(auth.Attributes["base_url"]) } for i := range e.cfg.VertexCompatAPIKey { entry := &e.cfg.VertexCompatAPIKey[i] cfgKey := strings.TrimSpace(entry.APIKey) cfgBase := strings.TrimSpace(entry.BaseURL) if attrKey != "" && attrBase != "" { if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { return entry } continue } if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { return entry } } if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { return entry } } if attrKey != "" { for i := range e.cfg.VertexCompatAPIKey { entry := &e.cfg.VertexCompatAPIKey[i] if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { return entry } } } return nil }