batch_image_provider_gemini.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/base64"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "mime/multipart"
  11. "net/http"
  12. "net/textproto"
  13. "net/url"
  14. "strings"
  15. "time"
  16. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
  19. )
  20. const defaultGeminiBatchRequeueAfter = 30 * time.Second
  21. type GeminiBatchClient interface {
  22. UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error)
  23. CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error)
  24. GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error)
  25. CancelBatch(ctx context.Context, apiKey string, batchName string) error
  26. DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error)
  27. DeleteFile(ctx context.Context, apiKey string, fileName string) error
  28. }
  29. type GeminiUploadedFile struct {
  30. Name string `json:"name"`
  31. DisplayName string `json:"displayName"`
  32. URI string `json:"uri"`
  33. MimeType string `json:"mimeType"`
  34. }
  35. type GeminiBatchJob struct {
  36. Name string `json:"name"`
  37. State string `json:"state"`
  38. Dest *GeminiBatchDest `json:"dest"`
  39. Response *GeminiBatchResponse `json:"response"`
  40. Error *GeminiBatchError `json:"error"`
  41. Raw map[string]any `json:"-"`
  42. }
  43. type GeminiBatchDest struct {
  44. FileName string `json:"fileName"`
  45. FileNameSnake string `json:"file_name"`
  46. }
  47. type GeminiBatchResponse struct {
  48. ResponsesFile string `json:"responsesFile"`
  49. ResponsesFileSnake string `json:"responses_file"`
  50. InlinedResponses []any `json:"inlinedResponses"`
  51. InlinedResponsesAlt []any `json:"inlined_responses"`
  52. }
  53. type GeminiBatchError struct {
  54. Code string `json:"code"`
  55. Message string `json:"message"`
  56. Status string `json:"status"`
  57. }
  58. type GeminiAPIBatchImageProvider struct {
  59. client GeminiBatchClient
  60. }
  61. func NewGeminiAPIBatchImageProvider(client GeminiBatchClient) *GeminiAPIBatchImageProvider {
  62. if client == nil {
  63. client = NewGeminiBatchHTTPClient("", nil)
  64. }
  65. return &GeminiAPIBatchImageProvider{client: client}
  66. }
  67. func (p *GeminiAPIBatchImageProvider) Name() string {
  68. return BatchImageProviderGeminiAPI
  69. }
  70. func (p *GeminiAPIBatchImageProvider) SupportsAccount(account *Account) bool {
  71. return account != nil &&
  72. account.Platform == PlatformGemini &&
  73. account.Type == AccountTypeAPIKey &&
  74. batchImageProviderAPIKey(account) != ""
  75. }
  76. func (p *GeminiAPIBatchImageProvider) Submit(ctx context.Context, job *BatchImageJob, account *Account, input BatchImageInput) (*BatchProviderJob, error) {
  77. if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
  78. return nil, ErrBatchImageProviderUnsupportedAccount
  79. }
  80. apiKey := batchImageProviderAPIKey(account)
  81. if apiKey == "" {
  82. return nil, ErrBatchImageProviderMissingAPIKey
  83. }
  84. if input.BatchID == "" && job != nil {
  85. input.BatchID = job.BatchID
  86. }
  87. if input.Model == "" && job != nil {
  88. input.Model = job.Model
  89. }
  90. jsonl, err := BuildGeminiBatchJSONL(input)
  91. if err != nil {
  92. return nil, err
  93. }
  94. displayName := strings.TrimSpace(input.DisplayName)
  95. if displayName == "" {
  96. displayName = strings.TrimSpace(input.BatchID)
  97. }
  98. uploaded, err := p.client.UploadJSONL(ctx, apiKey, displayName, bytes.NewReader(jsonl))
  99. if err != nil {
  100. return nil, mapGeminiClientError(err)
  101. }
  102. if uploaded == nil || strings.TrimSpace(uploaded.Name) == "" {
  103. return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini upload response is missing file name", nil)
  104. }
  105. batch, err := p.client.CreateBatch(ctx, apiKey, input.Model, uploaded.Name, displayName)
  106. if err != nil {
  107. return nil, mapGeminiClientError(err)
  108. }
  109. if batch == nil || strings.TrimSpace(batch.Name) == "" {
  110. return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is missing job name", nil)
  111. }
  112. return &BatchProviderJob{
  113. ProviderJobName: batch.Name,
  114. ProviderInputRef: uploaded.Name,
  115. RawState: batch.State,
  116. }, nil
  117. }
  118. func (p *GeminiAPIBatchImageProvider) Get(ctx context.Context, job *BatchImageJob, account *Account) (*BatchProviderStatus, error) {
  119. if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
  120. return nil, ErrBatchImageProviderUnsupportedAccount
  121. }
  122. apiKey := batchImageProviderAPIKey(account)
  123. if apiKey == "" {
  124. return nil, ErrBatchImageProviderMissingAPIKey
  125. }
  126. jobName := batchImageProviderJobName(job)
  127. if jobName == "" {
  128. return nil, ErrBatchImageProviderMissingJobName
  129. }
  130. batch, err := p.client.GetBatch(ctx, apiKey, jobName)
  131. if err != nil {
  132. return nil, mapGeminiClientError(err)
  133. }
  134. if batch == nil {
  135. return nil, geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini batch response is empty", nil)
  136. }
  137. status := mapGeminiBatchState(batch)
  138. if status.InternalState == BatchProviderStateSucceeded {
  139. if geminiBatchHasInlineResults(batch) {
  140. return nil, ErrBatchImageProviderInlineResultUnsupported
  141. }
  142. outputRef := geminiBatchOutputRef(batch)
  143. if outputRef == "" {
  144. status.InternalState = BatchProviderStateFailed
  145. status.Done = true
  146. status.ErrorCode = "GEMINI_RESULT_FILE_MISSING"
  147. status.ErrorMessage = "Gemini batch succeeded without a result file reference"
  148. }
  149. status.ProviderOutputRef = outputRef
  150. }
  151. return status, nil
  152. }
  153. func (p *GeminiAPIBatchImageProvider) Cancel(ctx context.Context, job *BatchImageJob, account *Account) error {
  154. if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
  155. return ErrBatchImageProviderUnsupportedAccount
  156. }
  157. apiKey := batchImageProviderAPIKey(account)
  158. if apiKey == "" {
  159. return ErrBatchImageProviderMissingAPIKey
  160. }
  161. jobName := batchImageProviderJobName(job)
  162. if jobName == "" {
  163. return ErrBatchImageProviderMissingJobName
  164. }
  165. return mapGeminiClientError(p.client.CancelBatch(ctx, apiKey, jobName))
  166. }
  167. func (p *GeminiAPIBatchImageProvider) OpenResult(ctx context.Context, job *BatchImageJob, account *Account) (io.ReadCloser, string, error) {
  168. if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
  169. return nil, "", ErrBatchImageProviderUnsupportedAccount
  170. }
  171. apiKey := batchImageProviderAPIKey(account)
  172. if apiKey == "" {
  173. return nil, "", ErrBatchImageProviderMissingAPIKey
  174. }
  175. outputRef := batchImageProviderOutputRef(job)
  176. if outputRef == "" {
  177. return nil, "", ErrBatchImageProviderMissingResultRef
  178. }
  179. r, contentType, err := p.client.DownloadFile(ctx, apiKey, outputRef)
  180. return r, contentType, mapGeminiClientError(err)
  181. }
  182. func (p *GeminiAPIBatchImageProvider) Cleanup(ctx context.Context, job *BatchImageJob, account *Account, target CleanupTarget) error {
  183. if account == nil || account.Platform != PlatformGemini || account.Type != AccountTypeAPIKey {
  184. return ErrBatchImageProviderUnsupportedAccount
  185. }
  186. apiKey := batchImageProviderAPIKey(account)
  187. if apiKey == "" {
  188. return ErrBatchImageProviderMissingAPIKey
  189. }
  190. switch target {
  191. case CleanupTargetInput:
  192. return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job))
  193. case CleanupTargetOutput:
  194. return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job))
  195. case CleanupTargetAll:
  196. if err := p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderInputRef(job)); err != nil {
  197. return err
  198. }
  199. return p.deleteGeminiFileIfPresent(ctx, apiKey, batchImageProviderOutputRef(job))
  200. default:
  201. return ErrUnsupportedCleanupTarget
  202. }
  203. }
  204. func (p *GeminiAPIBatchImageProvider) deleteGeminiFileIfPresent(ctx context.Context, apiKey, fileName string) error {
  205. if strings.TrimSpace(fileName) == "" {
  206. return nil
  207. }
  208. return mapGeminiClientError(p.client.DeleteFile(ctx, apiKey, fileName))
  209. }
  210. type geminiJSONLLine struct {
  211. Key string `json:"key"`
  212. Request geminiGenerateRequest `json:"request"`
  213. }
  214. type geminiGenerateRequest struct {
  215. Contents []geminiContent `json:"contents"`
  216. GenerationConfig geminiGenerationConfig `json:"generationConfig"`
  217. }
  218. type geminiContent struct {
  219. Parts []geminiPart `json:"parts"`
  220. }
  221. type geminiPart struct {
  222. Text string `json:"text,omitempty"`
  223. InlineData *geminiInlineData `json:"inlineData,omitempty"`
  224. FileData *geminiFileData `json:"fileData,omitempty"`
  225. }
  226. type geminiInlineData struct {
  227. MimeType string `json:"mimeType"`
  228. Data string `json:"data"`
  229. }
  230. type geminiFileData struct {
  231. MimeType string `json:"mimeType"`
  232. FileURI string `json:"fileUri"`
  233. }
  234. type geminiGenerationConfig struct {
  235. ResponseModalities []string `json:"responseModalities"`
  236. }
  237. func BuildGeminiBatchJSONL(input BatchImageInput) ([]byte, error) {
  238. if strings.TrimSpace(input.Model) == "" {
  239. return nil, batchImageProviderInputError("model is required")
  240. }
  241. if len(input.Items) == 0 {
  242. return nil, batchImageProviderInputError("at least one item is required")
  243. }
  244. seen := make(map[string]struct{}, len(input.Items))
  245. var buf bytes.Buffer
  246. enc := json.NewEncoder(&buf)
  247. for _, item := range input.Items {
  248. customID := strings.TrimSpace(item.CustomID)
  249. if customID == "" {
  250. return nil, batchImageProviderInputError("custom_id is required")
  251. }
  252. if _, ok := seen[customID]; ok {
  253. return nil, batchImageProviderInputError("duplicate custom_id %q", customID)
  254. }
  255. seen[customID] = struct{}{}
  256. prompt := strings.TrimSpace(item.Prompt)
  257. if prompt == "" {
  258. return nil, batchImageProviderInputError("prompt is required for custom_id %q", customID)
  259. }
  260. parts, err := batchImageGeminiParts(prompt, item.ReferenceImages)
  261. if err != nil {
  262. return nil, err
  263. }
  264. // TODO(batch-image): add response_mime_type/aspect_ratio/image_size once the
  265. // Gemini batch image REST shape is stabilized for those options.
  266. line := geminiJSONLLine{
  267. Key: customID,
  268. Request: geminiGenerateRequest{
  269. Contents: []geminiContent{{
  270. Parts: parts,
  271. }},
  272. GenerationConfig: geminiGenerationConfig{
  273. ResponseModalities: []string{"TEXT", "IMAGE"},
  274. },
  275. },
  276. }
  277. if err := enc.Encode(line); err != nil {
  278. return nil, err
  279. }
  280. }
  281. return buf.Bytes(), nil
  282. }
  283. func batchImageGeminiParts(prompt string, refs []BatchImageReference) ([]geminiPart, error) {
  284. parts := []geminiPart{{Text: prompt}}
  285. for _, ref := range refs {
  286. mimeType := normalizeBatchImageReferenceMimeType(ref.MimeType)
  287. if mimeType == "" {
  288. return nil, batchImageProviderInputError("reference image mime_type is required")
  289. }
  290. fileURI := strings.TrimSpace(ref.FileURI)
  291. switch {
  292. case len(ref.Data) > 0 && fileURI == "":
  293. parts = append(parts, geminiPart{InlineData: &geminiInlineData{
  294. MimeType: mimeType,
  295. Data: base64.StdEncoding.EncodeToString(ref.Data),
  296. }})
  297. case len(ref.Data) == 0 && fileURI != "":
  298. parts = append(parts, geminiPart{FileData: &geminiFileData{
  299. MimeType: mimeType,
  300. FileURI: fileURI,
  301. }})
  302. default:
  303. return nil, batchImageProviderInputError("reference image must contain exactly one of data or file_uri")
  304. }
  305. }
  306. return parts, nil
  307. }
  308. func mapGeminiBatchState(batch *GeminiBatchJob) *BatchProviderStatus {
  309. state := strings.TrimSpace(batch.State)
  310. normalized := strings.ToUpper(state)
  311. status := &BatchProviderStatus{
  312. RawState: state,
  313. InternalState: BatchProviderStateRunning,
  314. SuggestedRequeueAfter: defaultGeminiBatchRequeueAfter,
  315. }
  316. switch normalized {
  317. case "JOB_STATE_PENDING", "JOB_STATE_QUEUED":
  318. status.InternalState = BatchProviderStateQueued
  319. case "JOB_STATE_RUNNING":
  320. status.InternalState = BatchProviderStateRunning
  321. case "JOB_STATE_SUCCEEDED":
  322. status.InternalState = BatchProviderStateSucceeded
  323. status.Done = true
  324. case "JOB_STATE_FAILED":
  325. status.InternalState = BatchProviderStateFailed
  326. status.Done = true
  327. status.ErrorCode = "GEMINI_BATCH_FAILED"
  328. case "JOB_STATE_CANCELLED":
  329. status.InternalState = BatchProviderStateCancelled
  330. status.Done = true
  331. status.ErrorCode = "GEMINI_BATCH_CANCELLED"
  332. case "JOB_STATE_EXPIRED":
  333. status.InternalState = BatchProviderStateExpired
  334. status.Done = true
  335. status.ErrorCode = "GEMINI_BATCH_EXPIRED"
  336. default:
  337. if batch.Error != nil && (strings.TrimSpace(batch.Error.Message) != "" || strings.TrimSpace(batch.Error.Code) != "") {
  338. status.InternalState = BatchProviderStateFailed
  339. status.Done = true
  340. status.ErrorCode = "GEMINI_BATCH_FAILED"
  341. }
  342. }
  343. if batch.Error != nil {
  344. if code := strings.TrimSpace(batch.Error.Code); code != "" {
  345. status.ErrorCode = code
  346. } else if status.ErrorCode == "" && strings.TrimSpace(batch.Error.Status) != "" {
  347. status.ErrorCode = strings.TrimSpace(batch.Error.Status)
  348. }
  349. status.ErrorMessage = strings.TrimSpace(batch.Error.Message)
  350. }
  351. return status
  352. }
  353. func geminiBatchOutputRef(batch *GeminiBatchJob) string {
  354. if batch == nil {
  355. return ""
  356. }
  357. if batch.Dest != nil {
  358. if v := strings.TrimSpace(batch.Dest.FileName); v != "" {
  359. return v
  360. }
  361. if v := strings.TrimSpace(batch.Dest.FileNameSnake); v != "" {
  362. return v
  363. }
  364. }
  365. if batch.Response != nil {
  366. if v := strings.TrimSpace(batch.Response.ResponsesFile); v != "" {
  367. return v
  368. }
  369. if v := strings.TrimSpace(batch.Response.ResponsesFileSnake); v != "" {
  370. return v
  371. }
  372. }
  373. return ""
  374. }
  375. func geminiBatchHasInlineResults(batch *GeminiBatchJob) bool {
  376. return batch != nil &&
  377. batch.Response != nil &&
  378. (len(batch.Response.InlinedResponses) > 0 || len(batch.Response.InlinedResponsesAlt) > 0)
  379. }
  380. func geminiProviderError(reason, message string, cause error) error {
  381. err := infraerrors.New(http.StatusBadGateway, reason, message)
  382. if cause != nil {
  383. return err.WithCause(cause)
  384. }
  385. return err
  386. }
  387. func mapGeminiClientError(err error) error {
  388. if err == nil {
  389. return nil
  390. }
  391. var apiErr *GeminiAPIError
  392. if errors.As(err, &apiErr) {
  393. switch apiErr.StatusCode {
  394. case http.StatusUnauthorized, http.StatusForbidden:
  395. return geminiProviderError("GEMINI_AUTH_FAILED", "Gemini authentication failed", nil)
  396. case http.StatusTooManyRequests:
  397. return geminiProviderError("GEMINI_RATE_LIMITED", "Gemini rate limit exceeded", nil)
  398. case http.StatusNotFound:
  399. return geminiProviderError("GEMINI_BATCH_NOT_FOUND", "Gemini batch resource was not found", nil)
  400. default:
  401. return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil)
  402. }
  403. }
  404. return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini API request failed", nil)
  405. }
  406. type GeminiBatchHTTPClient struct {
  407. baseURL string
  408. client *http.Client
  409. }
  410. func NewGeminiBatchHTTPClient(baseURL string, client *http.Client) *GeminiBatchHTTPClient {
  411. baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
  412. if baseURL == "" {
  413. baseURL = geminicli.AIStudioBaseURL
  414. }
  415. if client == nil {
  416. client = batchImageDefaultHTTPClient()
  417. }
  418. return &GeminiBatchHTTPClient{baseURL: baseURL, client: client}
  419. }
  420. // batchImageDefaultHTTPClient 返回带连接/握手/响应头超时的共享客户端。
  421. // 不设整体 Timeout:大文件上传与结果流式下载耗时不可预估,
  422. // 但拨号、TLS、等待响应头必须有界,否则挂死的连接会无限占用提交路径。
  423. func batchImageDefaultHTTPClient() *http.Client {
  424. client, err := httpclient.GetClient(httpclient.Options{
  425. ResponseHeaderTimeout: 60 * time.Second,
  426. })
  427. if err != nil {
  428. return http.DefaultClient
  429. }
  430. return client
  431. }
  432. func (c *GeminiBatchHTTPClient) UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error) {
  433. var body bytes.Buffer
  434. writer := multipart.NewWriter(&body)
  435. metadataHeader := textproto.MIMEHeader{}
  436. metadataHeader.Set("Content-Disposition", `form-data; name="metadata"`)
  437. metadataHeader.Set("Content-Type", "application/json; charset=utf-8")
  438. metadataPart, err := writer.CreatePart(metadataHeader)
  439. if err != nil {
  440. return nil, err
  441. }
  442. metadata := map[string]any{"file": map[string]any{"displayName": displayName, "mimeType": "application/jsonl"}}
  443. if err := json.NewEncoder(metadataPart).Encode(metadata); err != nil {
  444. return nil, err
  445. }
  446. fileHeader := textproto.MIMEHeader{}
  447. fileHeader.Set("Content-Disposition", `form-data; name="file"; filename="batch.jsonl"`)
  448. fileHeader.Set("Content-Type", "application/jsonl")
  449. filePart, err := writer.CreatePart(fileHeader)
  450. if err != nil {
  451. return nil, err
  452. }
  453. if _, err := io.Copy(filePart, r); err != nil {
  454. return nil, err
  455. }
  456. if err := writer.Close(); err != nil {
  457. return nil, err
  458. }
  459. req, err := c.newRequest(ctx, http.MethodPost, "/upload/v1beta/files?uploadType=multipart", apiKey, &body)
  460. if err != nil {
  461. return nil, err
  462. }
  463. req.Header.Set("Content-Type", writer.FormDataContentType())
  464. var resp struct {
  465. File *GeminiUploadedFile `json:"file"`
  466. *GeminiUploadedFile
  467. }
  468. if err := c.doJSON(req, &resp); err != nil {
  469. return nil, err
  470. }
  471. if resp.File != nil {
  472. return resp.File, nil
  473. }
  474. return resp.GeminiUploadedFile, nil
  475. }
  476. func (c *GeminiBatchHTTPClient) CreateBatch(ctx context.Context, apiKey string, model string, fileName string, displayName string) (*GeminiBatchJob, error) {
  477. body := map[string]any{
  478. "batch": map[string]any{
  479. "displayName": displayName,
  480. "inputConfig": map[string]any{
  481. "fileName": fileName,
  482. },
  483. },
  484. }
  485. payload, _ := json.Marshal(body)
  486. path := fmt.Sprintf("/v1beta/models/%s:batchGenerateContent", url.PathEscape(strings.TrimSpace(model)))
  487. req, err := c.newRequest(ctx, http.MethodPost, path, apiKey, bytes.NewReader(payload))
  488. if err != nil {
  489. return nil, err
  490. }
  491. req.Header.Set("Content-Type", "application/json")
  492. return c.doBatchJob(req)
  493. }
  494. func (c *GeminiBatchHTTPClient) GetBatch(ctx context.Context, apiKey string, batchName string) (*GeminiBatchJob, error) {
  495. req, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(batchName, "/"), apiKey, nil)
  496. if err != nil {
  497. return nil, err
  498. }
  499. return c.doBatchJob(req)
  500. }
  501. func (c *GeminiBatchHTTPClient) CancelBatch(ctx context.Context, apiKey string, batchName string) error {
  502. req, err := c.newRequest(ctx, http.MethodPost, "/v1beta/"+strings.TrimLeft(batchName, "/")+":cancel", apiKey, nil)
  503. if err != nil {
  504. return err
  505. }
  506. return c.doNoBody(req)
  507. }
  508. func (c *GeminiBatchHTTPClient) DownloadFile(ctx context.Context, apiKey string, fileName string) (io.ReadCloser, string, error) {
  509. metaReq, err := c.newRequest(ctx, http.MethodGet, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil)
  510. if err != nil {
  511. return nil, "", err
  512. }
  513. var metadata struct {
  514. DownloadURI string `json:"downloadUri"`
  515. DownloadURL string `json:"download_url"`
  516. MimeType string `json:"mimeType"`
  517. }
  518. if err := c.doJSON(metaReq, &metadata); err != nil {
  519. return nil, "", err
  520. }
  521. downloadURL := strings.TrimSpace(metadata.DownloadURI)
  522. if downloadURL == "" {
  523. downloadURL = strings.TrimSpace(metadata.DownloadURL)
  524. }
  525. if downloadURL == "" {
  526. downloadURL = c.baseURL + "/v1beta/" + strings.TrimLeft(fileName, "/") + ":download"
  527. }
  528. // 纵深加固:downloadUri 来自上游响应,跟随前校验目标 host,
  529. // 防止异常/被劫持的响应把带 api key 的请求带到任意主机。
  530. if err := validateGeminiDownloadHost(downloadURL, c.baseURL); err != nil {
  531. return nil, "", err
  532. }
  533. req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
  534. if err != nil {
  535. return nil, "", err
  536. }
  537. req.Header.Set("x-goog-api-key", apiKey)
  538. resp, err := c.client.Do(req)
  539. if err != nil {
  540. return nil, "", err
  541. }
  542. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  543. defer func() { _ = resp.Body.Close() }()
  544. return nil, "", readGeminiAPIError(resp)
  545. }
  546. contentType := resp.Header.Get("Content-Type")
  547. if contentType == "" {
  548. contentType = metadata.MimeType
  549. }
  550. if contentType == "" {
  551. contentType = "application/octet-stream"
  552. }
  553. return resp.Body, contentType, nil
  554. }
  555. func (c *GeminiBatchHTTPClient) DeleteFile(ctx context.Context, apiKey string, fileName string) error {
  556. req, err := c.newRequest(ctx, http.MethodDelete, "/v1beta/"+strings.TrimLeft(fileName, "/"), apiKey, nil)
  557. if err != nil {
  558. return err
  559. }
  560. return c.doNoBody(req)
  561. }
  562. func (c *GeminiBatchHTTPClient) doBatchJob(req *http.Request) (*GeminiBatchJob, error) {
  563. var job GeminiBatchJob
  564. if err := c.doJSON(req, &job); err != nil {
  565. return nil, err
  566. }
  567. job.Raw = map[string]any{}
  568. return &job, nil
  569. }
  570. func (c *GeminiBatchHTTPClient) doNoBody(req *http.Request) error {
  571. resp, err := c.client.Do(req)
  572. if err != nil {
  573. return err
  574. }
  575. defer func() { _ = resp.Body.Close() }()
  576. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  577. return readGeminiAPIError(resp)
  578. }
  579. return nil
  580. }
  581. func (c *GeminiBatchHTTPClient) doJSON(req *http.Request, out any) error {
  582. resp, err := c.client.Do(req)
  583. if err != nil {
  584. return err
  585. }
  586. defer func() { _ = resp.Body.Close() }()
  587. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  588. return readGeminiAPIError(resp)
  589. }
  590. return json.NewDecoder(resp.Body).Decode(out)
  591. }
  592. func (c *GeminiBatchHTTPClient) newRequest(ctx context.Context, method, path, apiKey string, body io.Reader) (*http.Request, error) {
  593. if strings.TrimSpace(apiKey) == "" {
  594. return nil, ErrBatchImageProviderMissingAPIKey
  595. }
  596. req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, body)
  597. if err != nil {
  598. return nil, err
  599. }
  600. req.Header.Set("x-goog-api-key", apiKey)
  601. return req, nil
  602. }
  603. // validateGeminiDownloadHost 只允许跟随到 googleapis.com(含子域)
  604. // 或与配置的 baseURL 同 host 的下载地址。
  605. func validateGeminiDownloadHost(downloadURL, baseURL string) error {
  606. parsed, err := url.Parse(downloadURL)
  607. if err != nil {
  608. return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri is invalid", err)
  609. }
  610. if parsed.Scheme != "https" {
  611. return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri must use https", nil)
  612. }
  613. host := strings.ToLower(parsed.Hostname())
  614. if host == "googleapis.com" || strings.HasSuffix(host, ".googleapis.com") {
  615. return nil
  616. }
  617. if base, err := url.Parse(baseURL); err == nil && strings.EqualFold(base.Hostname(), host) {
  618. return nil
  619. }
  620. return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri host is not allowed", nil)
  621. }
  622. type GeminiAPIError struct {
  623. StatusCode int
  624. Code string
  625. Message string
  626. }
  627. func (e *GeminiAPIError) Error() string {
  628. if e == nil {
  629. return "<nil>"
  630. }
  631. if e.Code != "" {
  632. return fmt.Sprintf("gemini api error: status=%d code=%s message=%s", e.StatusCode, e.Code, e.Message)
  633. }
  634. return fmt.Sprintf("gemini api error: status=%d message=%s", e.StatusCode, e.Message)
  635. }
  636. func readGeminiAPIError(resp *http.Response) error {
  637. body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192))
  638. message := string(body)
  639. var parsed struct {
  640. Error struct {
  641. Code any `json:"code"`
  642. Message string `json:"message"`
  643. Status string `json:"status"`
  644. } `json:"error"`
  645. }
  646. if err := json.Unmarshal(body, &parsed); err == nil && parsed.Error.Message != "" {
  647. message = parsed.Error.Message
  648. return &GeminiAPIError{StatusCode: resp.StatusCode, Code: parsed.Error.Status, Message: message}
  649. }
  650. return &GeminiAPIError{StatusCode: resp.StatusCode, Message: message}
  651. }
  652. var _ BatchImageProvider = (*GeminiAPIBatchImageProvider)(nil)
  653. var _ GeminiBatchClient = (*GeminiBatchHTTPClient)(nil)