| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719 |
- package service
- import (
- "bytes"
- "context"
- "encoding/base64"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "mime/multipart"
- "net/http"
- "net/textproto"
- "net/url"
- "strings"
- "time"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
- "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
- )
- const defaultGeminiBatchRequeueAfter = 30 * time.Second
- type GeminiBatchClient interface {
- UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error)
- CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error)
- GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error)
- CancelBatch(ctx context.Context, apiKey string, batchName string) error
- DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error)
- DeleteFile(ctx context.Context, apiKey string, fileName string) error
- }
- type GeminiUploadedFile struct {
- Name string `json:"name"`
- DisplayName string `json:"displayName"`
- URI string `json:"uri"`
- MimeType string `json:"mimeType"`
- }
- type GeminiBatchJob struct {
- Name string `json:"name"`
- State string `json:"state"`
- Dest *GeminiBatchDest `json:"dest"`
- Response *GeminiBatchResponse `json:"response"`
- Error *GeminiBatchError `json:"error"`
- Raw map[string]any `json:"-"`
- }
- type GeminiBatchDest struct {
- FileName string `json:"fileName"`
- FileNameSnake string `json:"file_name"`
- }
- type GeminiBatchResponse struct {
- ResponsesFile string `json:"responsesFile"`
- ResponsesFileSnake string `json:"responses_file"`
- InlinedResponses []any `json:"inlinedResponses"`
- InlinedResponsesAlt []any `json:"inlined_responses"`
- }
- type GeminiBatchError struct {
- Code string `json:"code"`
- Message string `json:"message"`
- Status string `json:"status"`
- }
- type GeminiAPIBatchImageProvider struct {
- client GeminiBatchClient
- }
- func NewGeminiAPIBatchImageProvider(client GeminiBatchClient) *GeminiAPIBatchImageProvider {
- if client == nil {
- client = NewGeminiBatchHTTPClient("", nil)
- }
- return &GeminiAPIBatchImageProvider{client: client}
- }
- func (p *GeminiAPIBatchImageProvider) Name() string {
- return BatchImageProviderGeminiAPI
- }
- func (p *GeminiAPIBatchImageProvider) SupportsAccount(account *Account) bool {
- return account != nil &&
- account.Platform == PlatformGemini &&
- account.Type == AccountTypeAPIKey &&
- batchImageProviderAPIKey(account) != ""
- }
- func (p *GeminiAPIBatchImageProvider) Submit(ctx context.Context, job *BatchImageJob, account *Account, input BatchImageInput) (*BatchProviderJob, error) {
- if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
- return nil, ErrBatchImageProviderUnsupportedAccount
- }
- apiKey := batchImageProviderAPIKey(account)
- if apiKey == "" {
- return nil, ErrBatchImageProviderMissingAPIKey
- }
- if input.BatchID == "" && job != nil {
- input.BatchID = job.BatchID
- }
- if input.Model == "" && job != nil {
- input.Model = job.Model
- }
- jsonl, err := BuildGeminiBatchJSONL(input)
- if err != nil {
- return nil, err
- }
- displayName := strings.TrimSpace(input.DisplayName)
- if displayName == "" {
- displayName = strings.TrimSpace(input.BatchID)
- }
- uploaded, err := p.client.UploadJSONL(ctx, apiKey, displayName, bytes.NewReader(jsonl))
- if err != nil {
- return nil, mapGeminiClientError(err)
- }
- if uploaded == nil || strings.TrimSpace(uploaded.Name) == "" {
- return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini upload response is missing file name", nil)
- }
- batch, err := p.client.CreateBatch(ctx, apiKey, input.Model, uploaded.Name, displayName)
- if err != nil {
- return nil, mapGeminiClientError(err)
- }
- if batch == nil || strings.TrimSpace(batch.Name) == "" {
- return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is missing job name", nil)
- }
- return &BatchProviderJob{
- ProviderJobName: batch.Name,
- ProviderInputRef: uploaded.Name,
- RawState: batch.State,
- }, nil
- }
- func (p *GeminiAPIBatchImageProvider) Get(ctx context.Context, job *BatchImageJob, account *Account) (*BatchProviderStatus, error) {
- if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
- return nil, ErrBatchImageProviderUnsupportedAccount
- }
- apiKey := batchImageProviderAPIKey(account)
- if apiKey == "" {
- return nil, ErrBatchImageProviderMissingAPIKey
- }
- jobName := batchImageProviderJobName(job)
- if jobName == "" {
- return nil, ErrBatchImageProviderMissingJobName
- }
- batch, err := p.client.GetBatch(ctx, apiKey, jobName)
- if err != nil {
- return nil, mapGeminiClientError(err)
- }
- if batch == nil {
- return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is empty", nil)
- }
- status := mapGeminiBatchState(batch)
- if status.InternalState == BatchProviderStateSucceeded {
- if geminiBatchHasInlineResults(batch) {
- return nil, ErrBatchImageProviderInlineResultUnsupported
- }
- outputRef := geminiBatchOutputRef(batch)
- if outputRef == "" {
- status.InternalState = BatchProviderStateFailed
- status.Done = true
- status.ErrorCode = "GEMINI_RESULT_FILE_MISSING"
- status.ErrorMessage = "Gemini batch succeeded without a result file reference"
- }
- status.ProviderOutputRef = outputRef
- }
- return status, nil
- }
- func (p *GeminiAPIBatchImageProvider) Cancel(ctx context.Context, job *BatchImageJob, account *Account) error {
- if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
- return ErrBatchImageProviderUnsupportedAccount
- }
- apiKey := batchImageProviderAPIKey(account)
- if apiKey == "" {
- return ErrBatchImageProviderMissingAPIKey
- }
- jobName := batchImageProviderJobName(job)
- if jobName == "" {
- return ErrBatchImageProviderMissingJobName
- }
- return mapGeminiClientError(p.client.CancelBatch(ctx, apiKey, jobName))
- }
- func (p *GeminiAPIBatchImageProvider) OpenResult(ctx context.Context, job *BatchImageJob, account *Account) (io.ReadCloser, string, error) {
- if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
- return nil, "", ErrBatchImageProviderUnsupportedAccount
- }
- apiKey := batchImageProviderAPIKey(account)
- if apiKey == "" {
- return nil, "", ErrBatchImageProviderMissingAPIKey
- }
- outputRef := batchImageProviderOutputRef(job)
- if outputRef == "" {
- return nil, "", ErrBatchImageProviderMissingResultRef
- }
- r, contentType, err := p.client.DownloadFile(ctx, apiKey, outputRef)
- return r, contentType, mapGeminiClientError(err)
- }
- func (p *GeminiAPIBatchImageProvider) Cleanup(ctx context.Context, job *BatchImageJob, account *Account, target CleanupTarget) error {
- if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
- return ErrBatchImageProviderUnsupportedAccount
- }
- apiKey := batchImageProviderAPIKey(account)
- if apiKey == "" {
- return ErrBatchImageProviderMissingAPIKey
- }
- switch target {
- case CleanupTargetInput:
- return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job))
- case CleanupTargetOutput:
- return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job))
- case CleanupTargetAll:
- if err := p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job)); err != nil {
- return err
- }
- return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job))
- default:
- return ErrUnsupportedCleanupTarget
- }
- }
- func (p *GeminiAPIBatchImageProvider) deleteGeminiFileIfPresent(ctx context.Context, apiKey, fileName string) error {
- if strings.TrimSpace(fileName) == "" {
- return nil
- }
- return mapGeminiClientError(p.client.DeleteFile(ctx, apiKey, fileName))
- }
- type geminiJSONLLine struct {
- Key string `json:"key"`
- Request geminiGenerateRequest `json:"request"`
- }
- type geminiGenerateRequest struct {
- Contents []geminiContent `json:"contents"`
- GenerationConfig geminiGenerationConfig `json:"generationConfig"`
- }
- type geminiContent struct {
- Parts []geminiPart `json:"parts"`
- }
- type geminiPart struct {
- Text string `json:"text,omitempty"`
- InlineData *geminiInlineData `json:"inlineData,omitempty"`
- FileData *geminiFileData `json:"fileData,omitempty"`
- }
- type geminiInlineData struct {
- MimeType string `json:"mimeType"`
- Data string `json:"data"`
- }
- type geminiFileData struct {
- MimeType string `json:"mimeType"`
- FileURI string `json:"fileUri"`
- }
- type geminiGenerationConfig struct {
- ResponseModalities []string `json:"responseModalities"`
- }
- func BuildGeminiBatchJSONL(input BatchImageInput) ([]byte, error) {
- if strings.TrimSpace(input.Model) == "" {
- return nil, batchImageProviderInputError("model is required")
- }
- if len(input.Items) == 0 {
- return nil, batchImageProviderInputError("at least one item is required")
- }
- seen := make(map[string]struct{}, len(input.Items))
- var buf bytes.Buffer
- enc := json.NewEncoder(&buf)
- for _, item := range input.Items {
- customID := strings.TrimSpace(item.CustomID)
- if customID == "" {
- return nil, batchImageProviderInputError("custom_id is required")
- }
- if _, ok := seen[customID]; ok {
- return nil, batchImageProviderInputError("duplicate custom_id %q", customID)
- }
- seen[customID] = struct{}{}
- prompt := strings.TrimSpace(item.Prompt)
- if prompt == "" {
- return nil, batchImageProviderInputError("prompt is required for custom_id %q", customID)
- }
- parts, err := batchImageGeminiParts(prompt, item.ReferenceImages)
- if err != nil {
- return nil, err
- }
- // TODO(batch-image): add response_mime_type/aspect_ratio/image_size once the
- // Gemini batch image REST shape is stabilized for those options.
- line := geminiJSONLLine{
- Key: customID,
- Request: geminiGenerateRequest{
- Contents: []geminiContent{{
- Parts: parts,
- }},
- GenerationConfig: geminiGenerationConfig{
- ResponseModalities: []string{"TEXT", "IMAGE"},
- },
- },
- }
- if err := enc.Encode(line); err != nil {
- return nil, err
- }
- }
- return buf.Bytes(), nil
- }
- func batchImageGeminiParts(prompt string, refs []BatchImageReference) ([]geminiPart, error) {
- parts := []geminiPart{{Text: prompt}}
- for _, ref := range refs {
- mimeType := normalizeBatchImageReferenceMimeType(ref.MimeType)
- if mimeType == "" {
- return nil, batchImageProviderInputError("reference image mime_type is required")
- }
- fileURI := strings.TrimSpace(ref.FileURI)
- switch {
- case len(ref.Data) > 0 && fileURI == "":
- parts = append(parts, geminiPart{InlineData: &geminiInlineData{
- MimeType: mimeType,
- Data: base64.StdEncoding.EncodeToString(ref.Data),
- }})
- case len(ref.Data) == 0 && fileURI != "":
- parts = append(parts, geminiPart{FileData: &geminiFileData{
- MimeType: mimeType,
- FileURI: fileURI,
- }})
- default:
- return nil, batchImageProviderInputError("reference image must contain exactly one of data or file_uri")
- }
- }
- return parts, nil
- }
- func mapGeminiBatchState(batch *GeminiBatchJob) *BatchProviderStatus {
- state := strings.TrimSpace(batch.State)
- normalized := strings.ToUpper(state)
- status := &BatchProviderStatus{
- RawState: state,
- InternalState: BatchProviderStateRunning,
- SuggestedRequeueAfter: defaultGeminiBatchRequeueAfter,
- }
- switch normalized {
- case "JOB_STATE_PENDING", "JOB_STATE_QUEUED":
- status.InternalState = BatchProviderStateQueued
- case "JOB_STATE_RUNNING":
- status.InternalState = BatchProviderStateRunning
- case "JOB_STATE_SUCCEEDED":
- status.InternalState = BatchProviderStateSucceeded
- status.Done = true
- case "JOB_STATE_FAILED":
- status.InternalState = BatchProviderStateFailed
- status.Done = true
- status.ErrorCode = "GEMINI_BATCH_FAILED"
- case "JOB_STATE_CANCELLED":
- status.InternalState = BatchProviderStateCancelled
- status.Done = true
- status.ErrorCode = "GEMINI_BATCH_CANCELLED"
- case "JOB_STATE_EXPIRED":
- status.InternalState = BatchProviderStateExpired
- status.Done = true
- status.ErrorCode = "GEMINI_BATCH_EXPIRED"
- default:
- if batch.Error != nil && (strings.TrimSpace(batch.Error.Message) != "" || strings.TrimSpace(batch.Error.Code) != "") {
- status.InternalState = BatchProviderStateFailed
- status.Done = true
- status.ErrorCode = "GEMINI_BATCH_FAILED"
- }
- }
- if batch.Error != nil {
- if code := strings.TrimSpace(batch.Error.Code); code != "" {
- status.ErrorCode = code
- } else if status.ErrorCode == "" && strings.TrimSpace(batch.Error.Status) != "" {
- status.ErrorCode = strings.TrimSpace(batch.Error.Status)
- }
- status.ErrorMessage = strings.TrimSpace(batch.Error.Message)
- }
- return status
- }
- func geminiBatchOutputRef(batch *GeminiBatchJob) string {
- if batch == nil {
- return ""
- }
- if batch.Dest != nil {
- if v := strings.TrimSpace(batch.Dest.FileName); v != "" {
- return v
- }
- if v := strings.TrimSpace(batch.Dest.FileNameSnake); v != "" {
- return v
- }
- }
- if batch.Response != nil {
- if v := strings.TrimSpace(batch.Response.ResponsesFile); v != "" {
- return v
- }
- if v := strings.TrimSpace(batch.Response.ResponsesFileSnake); v != "" {
- return v
- }
- }
- return ""
- }
- func geminiBatchHasInlineResults(batch *GeminiBatchJob) bool {
- return batch != nil &&
- batch.Response != nil &&
- (len(batch.Response.InlinedResponses) > 0 || len(batch.Response.InlinedResponsesAlt) > 0)
- }
- func geminiProviderError(reason, message string, cause error) error {
- err := infraerrors.New(http.StatusBadGateway, reason, message)
- if cause != nil {
- return err.WithCause(cause)
- }
- return err
- }
- func mapGeminiClientError(err error) error {
- if err == nil {
- return nil
- }
- var apiErr *GeminiAPIError
- if errors.As(err, &apiErr) {
- switch apiErr.StatusCode {
- case http.StatusUnauthorized, http.StatusForbidden:
- return geminiProviderError("GEMINI_AUTH_FAILED", "Gemini authentication failed", nil)
- case http.StatusTooManyRequests:
- return geminiProviderError("GEMINI_RATE_LIMITED", "Gemini rate limit exceeded", nil)
- case http.StatusNotFound:
- return geminiProviderError("GEMINI_BATCH_NOT_FOUND", "Gemini batch resource was not found", nil)
- default:
- return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil)
- }
- }
- return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil)
- }
- type GeminiBatchHTTPClient struct {
- baseURL string
- client *http.Client
- }
- func NewGeminiBatchHTTPClient(baseURL string, client *http.Client) *GeminiBatchHTTPClient {
- baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
- if baseURL == "" {
- baseURL = geminicli.AIStudioBaseURL
- }
- if client == nil {
- client = batchImageDefaultHTTPClient()
- }
- return &GeminiBatchHTTPClient{baseURL: baseURL, client: client}
- }
- // batchImageDefaultHTTPClient 返回带连接/握手/响应头超时的共享客户端。
- // 不设整体 Timeout:大文件上传与结果流式下载耗时不可预估,
- // 但拨号、TLS、等待响应头必须有界,否则挂死的连接会无限占用提交路径。
- func batchImageDefaultHTTPClient() *http.Client {
- client, err := httpclient.GetClient(httpclient.Options{
- ResponseHeaderTimeout: 60 * time.Second,
- })
- if err != nil {
- return http.DefaultClient
- }
- return client
- }
- func (c *GeminiBatchHTTPClient) UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error) {
- var body bytes.Buffer
- writer := multipart.NewWriter(&body)
- metadataHeader := textproto.MIMEHeader{}
- metadataHeader.Set("Content-Disposition", `form-data; name="metadata"`)
- metadataHeader.Set("Content-Type", "application/json; charset=utf-8")
- metadataPart, err := writer.CreatePart(metadataHeader)
- if err != nil {
- return nil, err
- }
- metadata := map[string]any{"file": map[string]any{"displayName": displayName, "mimeType": "application/jsonl"}}
- if err := json.NewEncoder(metadataPart).Encode(metadata); err != nil {
- return nil, err
- }
- fileHeader := textproto.MIMEHeader{}
- fileHeader.Set("Content-Disposition", `form-data; name="file"; filename="batch.jsonl"`)
- fileHeader.Set("Content-Type", "application/jsonl")
- filePart, err := writer.CreatePart(fileHeader)
- if err != nil {
- return nil, err
- }
- if _, err := io.Copy(filePart, r); err != nil {
- return nil, err
- }
- if err := writer.Close(); err != nil {
- return nil, err
- }
- req, err := c.newRequest(ctx, http.MethodPost, "/upload/v1beta/files?uploadType=multipart", apiKey, &body)
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", writer.FormDataContentType())
- var resp struct {
- File *GeminiUploadedFile `json:"file"`
- *GeminiUploadedFile
- }
- if err := c.doJSON(req, &resp); err != nil {
- return nil, err
- }
- if resp.File != nil {
- return resp.File, nil
- }
- return resp.GeminiUploadedFile, nil
- }
- func (c *GeminiBatchHTTPClient) CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error) {
- body := map[string]any{
- "batch": map[string]any{
- "displayName": displayName,
- "inputConfig": map[string]any{
- "fileName": fileName,
- },
- },
- }
- payload, _ := json.Marshal(body)
- path := fmt.Sprintf("/v1beta/models/%s:batchGenerateContent", url.PathEscape(strings.TrimSpace(model)))
- req, err := c.newRequest(ctx, http.MethodPost, path, apiKey, bytes.NewReader(payload))
- if err != nil {
- return nil, err
- }
- req.Header.Set("Content-Type", "application/json")
- return c.doBatchJob(req)
- }
- func (c *GeminiBatchHTTPClient) GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error) {
- req, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(batchName, "/"), apiKey, nil)
- if err != nil {
- return nil, err
- }
- return c.doBatchJob(req)
- }
- func (c *GeminiBatchHTTPClient) CancelBatch(ctx context.Context, apiKey string, batchName string) error {
- req, err := c.newRequest(ctx, http.MethodPost, "/v1beta/"+strings.TrimLeft(batchName, "/")+":cancel", apiKey, nil)
- if err != nil {
- return err
- }
- return c.doNoBody(req)
- }
- func (c *GeminiBatchHTTPClient) DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error) {
- metaReq, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil)
- if err != nil {
- return nil, "", err
- }
- var metadata struct {
- DownloadURI string `json:"downloadUri"`
- DownloadURL string `json:"download_url"`
- MimeType string `json:"mimeType"`
- }
- if err := c.doJSON(metaReq, &metadata); err != nil {
- return nil, "", err
- }
- downloadURL := strings.TrimSpace(metadata.DownloadURI)
- if downloadURL == "" {
- downloadURL = strings.TrimSpace(metadata.DownloadURL)
- }
- if downloadURL == "" {
- downloadURL = c.baseURL + "/v1beta/" + strings.TrimLeft(fileName, "/") + ":download"
- }
- // 纵深加固:downloadUri 来自上游响应,跟随前校验目标 host,
- // 防止异常/被劫持的响应把带 api key 的请求带到任意主机。
- if err := validateGeminiDownloadHost(downloadURL, c.baseURL); err != nil {
- return nil, "", err
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
- if err != nil {
- return nil, "", err
- }
- req.Header.Set("x-goog-api-key", apiKey)
- resp, err := c.client.Do(req)
- if err != nil {
- return nil, "", err
- }
- if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- defer func() { _ = resp.Body.Close() }()
- return nil, "", readGeminiAPIError(resp)
- }
- contentType := resp.Header.Get("Content-Type")
- if contentType == "" {
- contentType = metadata.MimeType
- }
- if contentType == "" {
- contentType = "application/octet-stream"
- }
- return resp.Body, contentType, nil
- }
- func (c *GeminiBatchHTTPClient) DeleteFile(ctx context.Context, apiKey string, fileName string) error {
- req, err := c.newRequest(ctx, http.MethodDelete, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil)
- if err != nil {
- return err
- }
- return c.doNoBody(req)
- }
- func (c *GeminiBatchHTTPClient) doBatchJob(req *http.Request) (*GeminiBatchJob, error) {
- var job GeminiBatchJob
- if err := c.doJSON(req, &job); err != nil {
- return nil, err
- }
- job.Raw = map[string]any{}
- return &job, nil
- }
- func (c *GeminiBatchHTTPClient) doNoBody(req *http.Request) error {
- resp, err := c.client.Do(req)
- if err != nil {
- return err
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- return readGeminiAPIError(resp)
- }
- return nil
- }
- func (c *GeminiBatchHTTPClient) doJSON(req *http.Request, out any) error {
- resp, err := c.client.Do(req)
- if err != nil {
- return err
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- return readGeminiAPIError(resp)
- }
- return json.NewDecoder(resp.Body).Decode(out)
- }
- func (c *GeminiBatchHTTPClient) newRequest(ctx context.Context, method, path, apiKey string, body io.Reader) (*http.Request, error) {
- if strings.TrimSpace(apiKey) == "" {
- return nil, ErrBatchImageProviderMissingAPIKey
- }
- req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, body)
- if err != nil {
- return nil, err
- }
- req.Header.Set("x-goog-api-key", apiKey)
- return req, nil
- }
- // validateGeminiDownloadHost 只允许跟随到 googleapis.com(含子域)
- // 或与配置的 baseURL 同 host 的下载地址。
- func validateGeminiDownloadHost(downloadURL, baseURL string) error {
- parsed, err := url.Parse(downloadURL)
- if err != nil {
- return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri is invalid", err)
- }
- if parsed.Scheme != "https" {
- return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri must use https", nil)
- }
- host := strings.ToLower(parsed.Hostname())
- if host == "googleapis.com" || strings.HasSuffix(host, ".googleapis.com") {
- return nil
- }
- if base, err := url.Parse(baseURL); err == nil && strings.EqualFold(base.Hostname(), host) {
- return nil
- }
- return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri host is not allowed", nil)
- }
- type GeminiAPIError struct {
- StatusCode int
- Code string
- Message string
- }
- func (e *GeminiAPIError) Error() string {
- if e == nil {
- return "<nil>"
- }
- if e.Code != "" {
- return fmt.Sprintf("gemini api error: status=%d code=%s message=%s", e.StatusCode, e.Code, e.Message)
- }
- return fmt.Sprintf("gemini api error: status=%d message=%s", e.StatusCode, e.Message)
- }
- func readGeminiAPIError(resp *http.Response) error {
- body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192))
- message := string(body)
- var parsed struct {
- Error struct {
- Code any `json:"code"`
- Message string `json:"message"`
- Status string `json:"status"`
- } `json:"error"`
- }
- if err := json.Unmarshal(body, &parsed); err == nil && parsed.Error.Message != "" {
- message = parsed.Error.Message
- return &GeminiAPIError{StatusCode: resp.StatusCode, Code: parsed.Error.Status, Message: message}
- }
- return &GeminiAPIError{StatusCode: resp.StatusCode, Message: message}
- }
- var _ BatchImageProvider = (*GeminiAPIBatchImageProvider)(nil)
- var _ GeminiBatchClient = (*GeminiBatchHTTPClient)(nil)
|