package executor import ( "context" "fmt" "io" "net/http" "strings" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" 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" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) // GeminiExecutorRefactored is a stateless executor for the official Gemini API. type GeminiExecutorRefactored struct { cfg *config.Config base *BaseExecutor } // NewGeminiExecutorRefactored creates a new Gemini executor instance. func NewGeminiExecutorRefactored(cfg *config.Config) *GeminiExecutorRefactored { provider := &GeminiProvider{} return &GeminiExecutorRefactored{ cfg: cfg, base: NewBaseExecutor(cfg, provider), } } // Identifier returns the executor identifier. func (e *GeminiExecutorRefactored) Identifier() string { return "gemini" } // PrepareRequest injects Gemini credentials into the outgoing HTTP request. func (e *GeminiExecutorRefactored) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if req == nil { return nil } apiKey, bearer := geminiCreds(auth) if apiKey != "" { req.Header.Set("x-goog-api-key", apiKey) req.Header.Del("Authorization") } else if bearer != "" { req.Header.Set("Authorization", "Bearer "+bearer) req.Header.Del("x-goog-api-key") } applyGeminiHeadersRefactored(req, auth) return nil } // HttpRequest injects Gemini credentials into the request and executes it. func (e *GeminiExecutorRefactored) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { if req == nil { return nil, fmt.Errorf("gemini 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 Gemini API. func (e *GeminiExecutorRefactored) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { if opts.Alt == "responses/compact" { return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } return e.base.Execute(ctx, auth, req, opts) } // ExecuteStream performs a streaming request to the Gemini API. func (e *GeminiExecutorRefactored) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } return e.base.ExecuteStream(ctx, auth, req, opts) } // CountTokens counts tokens for the given request using the Gemini API. func (e *GeminiExecutorRefactored) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { baseModel := thinking.ParseSuffix(req.Model).ModelName apiKey, bearer := geminiCreds(auth) from := opts.SourceFormat to := sdktranslator.FromString("gemini") translatedReq := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) translatedReq, err := thinking.ApplyThinking(translatedReq, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { return cliproxyexecutor.Response{}, err } translatedReq = fixGeminiImageAspectRatio(baseModel, translatedReq) respCtx := context.WithValue(ctx, "alt", opts.Alt) translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools") translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig") translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings") translatedReq, _ = sjson.SetBytes(translatedReq, "model", baseModel) baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, baseModel, "countTokens") httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(translatedReq))) if err != nil { return cliproxyexecutor.Response{}, err } httpReq.Header.Set("Content-Type", "application/json") if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } else { httpReq.Header.Set("Authorization", "Bearer "+bearer) } applyGeminiHeadersRefactored(httpReq, 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: url, Method: http.MethodPost, Headers: httpReq.Header.Clone(), Body: translatedReq, Provider: e.Identifier(), AuthID: authID, AuthLabel: authLabel, AuthType: authType, AuthValue: authValue, }) httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) resp, err := httpClient.Do(httpReq) if err != nil { recordAPIResponseError(ctx, e.cfg, err) return cliproxyexecutor.Response{}, err } defer resp.Body.Close() recordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone()) data, err := io.ReadAll(resp.Body) if err != nil { recordAPIResponseError(ctx, e.cfg, err) return cliproxyexecutor.Response{}, err } appendAPIResponseChunk(ctx, e.cfg, data) if resp.StatusCode < 200 || resp.StatusCode >= 300 { logWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", resp.StatusCode, summarizeErrorBody(resp.Header.Get("Content-Type"), data)) return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(data)} } count := gjson.GetBytes(data, "totalTokens").Int() translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, data) return cliproxyexecutor.Response{Payload: []byte(translated), Headers: resp.Header.Clone()}, nil } // Refresh refreshes the authentication credentials (no-op for Gemini API key). func (e *GeminiExecutorRefactored) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { return auth, nil } func geminiCreds(a *cliproxyauth.Auth) (apiKey, bearer string) { if a == nil { return "", "" } if a.Attributes != nil { if v := a.Attributes["api_key"]; v != "" { apiKey = v } } if a.Metadata != nil { // GeminiTokenStorage.Token is a map that may contain access_token if v, ok := a.Metadata["access_token"].(string); ok && v != "" { bearer = v } if token, ok := a.Metadata["token"].(map[string]any); ok && token != nil { if v, ok2 := token["access_token"].(string); ok2 && v != "" { bearer = v } } } return } func resolveGeminiBaseURL(auth *cliproxyauth.Auth) string { base := glEndpoint if auth != nil && auth.Attributes != nil { if custom := strings.TrimSpace(auth.Attributes["base_url"]); custom != "" { base = strings.TrimRight(custom, "/") } } if base == "" { return glEndpoint } return base } func applyGeminiHeadersRefactored(req *http.Request, auth *cliproxyauth.Auth) { var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(req, attrs) }