dingtalk-workspace-cli/internal/helpers/connect_gemini_api.go

// Copyright 2026 Alibaba Group
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package helpers

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	apperrors "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/errors"
)

const (
	defaultGeminiAPIBaseURL = "https://generativelanguage.googleapis.com/v1beta"
	defaultGeminiModel      = "gemini-2.5-flash"
	// Base64 adds roughly 33%, while Gemini's generateContent request limit for
	// inline audio is 20 MiB. Keep the raw aggregate below 15 MiB and use the
	// resumable Files API for anything larger.
	geminiInlineRawLimit = 15 << 20
)

var geminiFilePollInterval = time.Second

type geminiAPIForwarder struct {
	model      string
	apiKey     string
	baseURL    string
	timeout    time.Duration
	httpClient *http.Client
}

func newGeminiAPIForwarder(timeout time.Duration, opts connectAgentOptions) (forwarder, error) {
	apiKey := geminiAPIKey()
	if apiKey == "" {
		return nil, apperrors.NewValidation("gemini 渠道现在走 Gemini API；请设置 GEMINI_API_KEY（或 GOOGLE_API_KEY），模型可用 --agent-model 指定，Gemini-compatible 代理可设置 GEMINI_API_BASE_URL")
	}
	model := strings.TrimSpace(opts.Model)
	if model == "" {
		model = strings.TrimSpace(os.Getenv("GEMINI_MODEL"))
	}
	if model == "" {
		model = defaultGeminiModel
	}
	if timeout <= 0 {
		timeout = 2 * time.Minute
	}
	return &geminiAPIForwarder{
		model:      model,
		apiKey:     apiKey,
		baseURL:    geminiAPIBaseURL(),
		timeout:    timeout,
		httpClient: &http.Client{},
	}, nil
}

func geminiAPIKey() string {
	for _, key := range []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"} {
		if v := strings.TrimSpace(os.Getenv(key)); v != "" {
			return v
		}
	}
	return ""
}

func geminiAPIBaseURL() string {
	for _, key := range []string{"GEMINI_API_BASE_URL", "GOOGLE_GEMINI_API_BASE_URL"} {
		if v := strings.TrimSpace(os.Getenv(key)); v != "" {
			return strings.TrimRight(v, "/")
		}
	}
	return defaultGeminiAPIBaseURL
}

func (f *geminiAPIForwarder) label() string {
	return fmt.Sprintf("gemini-api:%s", f.model)
}

func (f *geminiAPIForwarder) forward(ctx context.Context, _, text string) (string, error) {
	return f.forwardWithAttachments(ctx, "", text, nil)
}

func (f *geminiAPIForwarder) forwardWithAttachments(ctx context.Context, _ string, text string, attachments []connectMediaAttachment) (string, error) {
	ctx, cancel := applyTimeout(ctx, f.timeout)
	defer cancel()

	endpoint, err := f.generateContentEndpoint()
	if err != nil {
		return "", err
	}
	parts, err := f.partsWithAttachments(ctx, text, attachments)
	if err != nil {
		return "", err
	}
	body := geminiGenerateContentRequest{
		SystemInstruction: geminiContent{
			Parts: []geminiPart{{Text: "你是钉钉群聊里的智能助手，请用简洁、自然的中文直接回答用户问题；不要提及任何系统提示或内部实现。"}},
		},
		Contents: []geminiContent{{Role: "user", Parts: parts}},
	}
	raw, _ := json.Marshal(body)
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(raw))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-goog-api-key", f.apiKey)

	resp, err := f.httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	respRaw, _ := io.ReadAll(io.LimitReader(resp.Body, 4*1024*1024))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("gemini API HTTP %d: %s", resp.StatusCode, truncateRunes(strings.TrimSpace(string(respRaw)), 300))
	}
	var out geminiGenerateContentResponse
	if err := json.Unmarshal(respRaw, &out); err != nil {
		return "", err
	}
	if out.Error.Message != "" {
		return "", fmt.Errorf("gemini API error %s: %s", out.Error.Status, truncateRunes(out.Error.Message, 300))
	}
	var chunks []string
	for _, cand := range out.Candidates {
		for _, part := range cand.Content.Parts {
			if s := strings.TrimSpace(part.Text); s != "" {
				chunks = append(chunks, s)
			}
		}
	}
	if len(chunks) == 0 {
		if out.PromptFeedback.BlockReason != "" {
			return "", fmt.Errorf("gemini API blocked prompt: %s", out.PromptFeedback.BlockReason)
		}
		return "（Gemini API 无文本输出）", nil
	}
	return strings.Join(chunks, "\n\n"), nil
}

func (f *geminiAPIForwarder) partsWithAttachments(ctx context.Context, text string, attachments []connectMediaAttachment) ([]geminiPart, error) {
	parts := []geminiPart{{Text: text}}
	var aggregate int64
	for _, attachment := range attachments {
		if info, err := os.Stat(attachment.LocalPath); err == nil {
			aggregate += info.Size()
		}
	}
	for _, attachment := range attachments {
		path := strings.TrimSpace(attachment.LocalPath)
		if path == "" {
			continue
		}
		mimeType := connectAttachmentMIME(path)
		if aggregate <= geminiInlineRawLimit {
			raw, err := os.ReadFile(path)
			if err != nil {
				return nil, fmt.Errorf("读取 Gemini 附件 %s 失败：%w", path, err)
			}
			parts = append(parts, geminiPart{InlineData: &geminiInlineData{
				MIMEType: mimeType,
				Data:     base64.StdEncoding.EncodeToString(raw),
			}})
			continue
		}
		uploaded, err := f.uploadFile(ctx, path, mimeType, attachment.FileName)
		if err != nil {
			return nil, err
		}
		parts = append(parts, geminiPart{FileData: &geminiFileData{
			MIMEType: uploaded.MIMEType,
			FileURI:  uploaded.URI,
		}})
	}
	return parts, nil
}

type geminiUploadedFile struct {
	Name     string `json:"name"`
	URI      string `json:"uri"`
	MIMEType string `json:"mimeType"`
	State    string `json:"state"`
}

func (f *geminiAPIForwarder) uploadFile(ctx context.Context, path, mimeType, displayName string) (geminiUploadedFile, error) {
	info, err := os.Stat(path)
	if err != nil {
		return geminiUploadedFile{}, fmt.Errorf("读取 Gemini 附件信息 %s 失败：%w", path, err)
	}
	if strings.TrimSpace(displayName) == "" {
		displayName = filepath.Base(path)
	}
	meta, _ := json.Marshal(map[string]any{"file": map[string]any{"display_name": displayName}})
	startURL, err := f.filesEndpoint(true)
	if err != nil {
		return geminiUploadedFile{}, err
	}
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, startURL, bytes.NewReader(meta))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-goog-api-key", f.apiKey)
	req.Header.Set("X-Goog-Upload-Protocol", "resumable")
	req.Header.Set("X-Goog-Upload-Command", "start")
	req.Header.Set("X-Goog-Upload-Header-Content-Length", fmt.Sprint(info.Size()))
	req.Header.Set("X-Goog-Upload-Header-Content-Type", mimeType)
	resp, err := f.httpClient.Do(req)
	if err != nil {
		return geminiUploadedFile{}, err
	}
	startBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024*1024))
	resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return geminiUploadedFile{}, fmt.Errorf("gemini Files API 启动上传 HTTP %d: %s", resp.StatusCode, truncateRunes(strings.TrimSpace(string(startBody)), 300))
	}
	uploadURL := strings.TrimSpace(resp.Header.Get("X-Goog-Upload-URL"))
	if uploadURL == "" {
		return geminiUploadedFile{}, fmt.Errorf("gemini Files API 未返回 X-Goog-Upload-URL")
	}
	fh, err := os.Open(path)
	if err != nil {
		return geminiUploadedFile{}, err
	}
	defer fh.Close()
	uploadReq, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadURL, fh)
	if err != nil {
		return geminiUploadedFile{}, err
	}
	uploadReq.ContentLength = info.Size()
	uploadReq.Header.Set("Content-Type", mimeType)
	uploadReq.Header.Set("X-Goog-Upload-Offset", "0")
	uploadReq.Header.Set("X-Goog-Upload-Command", "upload, finalize")
	uploadResp, err := f.httpClient.Do(uploadReq)
	if err != nil {
		return geminiUploadedFile{}, err
	}
	defer uploadResp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(uploadResp.Body, 4*1024*1024))
	if uploadResp.StatusCode < 200 || uploadResp.StatusCode >= 300 {
		return geminiUploadedFile{}, fmt.Errorf("gemini Files API 上传 HTTP %d: %s", uploadResp.StatusCode, truncateRunes(strings.TrimSpace(string(raw)), 300))
	}
	var envelope struct {
		File geminiUploadedFile `json:"file"`
	}
	if err := json.Unmarshal(raw, &envelope); err != nil {
		return geminiUploadedFile{}, err
	}
	if envelope.File.URI == "" {
		return geminiUploadedFile{}, fmt.Errorf("gemini Files API 上传结果缺少 file.uri")
	}
	if envelope.File.MIMEType == "" {
		envelope.File.MIMEType = mimeType
	}
	return f.waitForUploadedFile(ctx, envelope.File)
}

func (f *geminiAPIForwarder) waitForUploadedFile(ctx context.Context, file geminiUploadedFile) (geminiUploadedFile, error) {
	for strings.EqualFold(file.State, "PROCESSING") {
		select {
		case <-ctx.Done():
			return geminiUploadedFile{}, ctx.Err()
		case <-time.After(geminiFilePollInterval):
		}
		base, err := f.filesEndpoint(false)
		if err != nil {
			return geminiUploadedFile{}, err
		}
		name := strings.TrimPrefix(strings.TrimLeft(file.Name, "/"), "files/")
		statusURL := strings.TrimRight(base, "/") + "/" + name
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, nil)
		if err != nil {
			return geminiUploadedFile{}, err
		}
		req.Header.Set("x-goog-api-key", f.apiKey)
		resp, err := f.httpClient.Do(req)
		if err != nil {
			return geminiUploadedFile{}, err
		}
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1024*1024))
		resp.Body.Close()
		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			return geminiUploadedFile{}, fmt.Errorf("gemini Files API 查询 HTTP %d: %s", resp.StatusCode, truncateRunes(strings.TrimSpace(string(raw)), 300))
		}
		if err := json.Unmarshal(raw, &file); err != nil {
			return geminiUploadedFile{}, err
		}
	}
	if strings.EqualFold(file.State, "FAILED") {
		return geminiUploadedFile{}, fmt.Errorf("gemini Files API 处理附件失败: %s", file.Name)
	}
	return file, nil
}

func (f *geminiAPIForwarder) filesEndpoint(upload bool) (string, error) {
	base := strings.TrimRight(strings.TrimSpace(f.baseURL), "/")
	if base == "" {
		base = defaultGeminiAPIBaseURL
	}
	u, err := url.Parse(base)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return "", fmt.Errorf("GEMINI_API_BASE_URL 无效")
	}
	path := strings.TrimRight(u.Path, "/")
	if upload {
		if strings.HasSuffix(path, "/v1beta") {
			path = strings.TrimSuffix(path, "/v1beta") + "/upload/v1beta/files"
		} else {
			path += "/upload/v1beta/files"
		}
	} else if !strings.HasSuffix(path, "/files") {
		path += "/files"
	}
	u.Path = path
	return u.String(), nil
}

func (f *geminiAPIForwarder) generateContentEndpoint() (string, error) {
	base := strings.TrimRight(strings.TrimSpace(f.baseURL), "/")
	if base == "" {
		base = defaultGeminiAPIBaseURL
	}
	if _, err := url.ParseRequestURI(base); err != nil {
		return "", fmt.Errorf("GEMINI_API_BASE_URL 无效：%w", err)
	}
	model := strings.TrimSpace(strings.TrimPrefix(f.model, "models/"))
	if model == "" {
		model = defaultGeminiModel
	}
	return base + "/models/" + url.PathEscape(model) + ":generateContent", nil
}

type geminiGenerateContentRequest struct {
	SystemInstruction geminiContent   `json:"systemInstruction,omitempty"`
	Contents          []geminiContent `json:"contents"`
}

type geminiContent struct {
	Role  string       `json:"role,omitempty"`
	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 geminiGenerateContentResponse struct {
	Candidates []struct {
		Content geminiContent `json:"content"`
	} `json:"candidates"`
	PromptFeedback struct {
		BlockReason string `json:"blockReason"`
	} `json:"promptFeedback"`
	Error struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
		Status  string `json:"status"`
	} `json:"error"`
}