agentic-coding-workflow/internal/infra/llm/gemini.go

package llm

import (
	"context"
	"fmt"

	"google.golang.org/genai"
)

// GeminiClient 通过 google.golang.org/genai SDK 调用 Generative AI API。
// baseURL 可改写以走 Vertex AI 或内部网关。
type GeminiClient struct {
	sdk *genai.Client
}

// Compile-time assertion: GeminiClient must satisfy LLMClient + Embedder.
var (
	_ LLMClient = (*GeminiClient)(nil)
	_ Embedder  = (*GeminiClient)(nil)
)

// NewGemini 构造 GeminiClient。apiKey 不可空；baseURL 可空（用 SDK 默认端点）。
func NewGemini(apiKey, baseURL string) (*GeminiClient, error) {
	cfg := &genai.ClientConfig{
		APIKey:  apiKey,
		Backend: genai.BackendGeminiAPI,
	}
	if baseURL != "" {
		// SDK v1.55: HTTPOptions.BaseURL 改写 API 端点，用于 httptest 或自定义网关。
		cfg.HTTPOptions = genai.HTTPOptions{BaseURL: baseURL}
	}
	c, err := genai.NewClient(context.Background(), cfg)
	if err != nil {
		return nil, fmt.Errorf("gemini client: %w", err)
	}
	return &GeminiClient{sdk: c}, nil
}

func (c *GeminiClient) Provider() string { return "gemini" }

// toGeminiContents 把中立 Message 切片转换为 SDK *Content 切片。
// image/document 的 Data 字段为原始字节，Gemini SDK 直接接受（InlineData 不需要 base64 编码）。
func toGeminiContents(req Request) []*genai.Content {
	out := make([]*genai.Content, 0, len(req.Messages))
	for _, m := range req.Messages {
		role := "user"
		if m.Role == RoleAssistant {
			role = "model"
		}
		parts := make([]*genai.Part, 0, len(m.Blocks))
		for _, b := range m.Blocks {
			switch b.Type {
			case "text":
				parts = append(parts, genai.NewPartFromText(b.Text))
			case "image", "document":
				// SDK v1.55: Blob.MIMEType 大写，Data 为原始字节（SDK 内部处理编码）。
				parts = append(parts, &genai.Part{InlineData: &genai.Blob{
					MIMEType: b.MimeType,
					Data:     b.Data,
				}})
			}
		}
		out = append(out, &genai.Content{Role: role, Parts: parts})
	}
	return out
}

func (c *GeminiClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
	cfg := &genai.GenerateContentConfig{
		MaxOutputTokens: int32(req.MaxOutputTokens),
	}
	if req.SystemPrompt != "" {
		cfg.SystemInstruction = &genai.Content{
			Parts: []*genai.Part{genai.NewPartFromText(req.SystemPrompt)},
		}
	}
	if req.Reasoning {
		// SDK v1.55: ThinkingConfig.IncludeThoughts 是 bool，不是 *bool；无需 genai.Ptr()。
		cfg.ThinkingConfig = &genai.ThinkingConfig{IncludeThoughts: true}
	}

	// SDK v1.55: GenerateContentStream 返回 iter.Seq2[*GenerateContentResponse, error]，
	// 可用 Go 1.23+ range-over-func 语法直接遍历。
	iter := c.sdk.Models.GenerateContentStream(ctx, modelID, toGeminiContents(req), cfg)
	out := make(chan StreamEvent, 16)

	// send 在 ctx 取消时立即放弃发送，防止下游 reader 已退出导致 goroutine 永久阻塞。
	send := func(ev StreamEvent) bool {
		select {
		case <-ctx.Done():
			return false
		case out <- ev:
			return true
		}
	}

	go func() {
		defer close(out)
		var lastReason string
		for resp, err := range iter {
			if err != nil {
				send(StreamEvent{Type: EventError, Err: fmt.Errorf("gemini stream: %w", err)})
				return
			}
			if len(resp.Candidates) == 0 {
				continue
			}
			cand := resp.Candidates[0]
			if cand.Content != nil {
				for _, p := range cand.Content.Parts {
					if p.Thought {
						if !send(StreamEvent{Type: EventThinkingDelta, Text: p.Text}) {
							return
						}
					} else if p.Text != "" {
						if !send(StreamEvent{Type: EventTextDelta, Text: p.Text}) {
							return
						}
					}
				}
			}
			if resp.UsageMetadata != nil {
				if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
					PromptTokens:     int(resp.UsageMetadata.PromptTokenCount),
					CompletionTokens: int(resp.UsageMetadata.CandidatesTokenCount),
					ThinkingTokens:   int(resp.UsageMetadata.ThoughtsTokenCount),
				}}) {
					return
				}
			}
			if cand.FinishReason != "" {
				lastReason = string(cand.FinishReason)
			}
		}
		send(StreamEvent{Type: EventDone, Reason: lastReason})
	}()
	return out, nil
}

// CountTokens 调用 Gemini 远程 CountTokens 端点计算 token 数。
func (c *GeminiClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) {
	res, err := c.sdk.Models.CountTokens(ctx, modelID, toGeminiContents(req), nil)
	if err != nil {
		return 0, fmt.Errorf("gemini count_tokens: %w", err)
	}
	return int(res.TotalTokens), nil
}

// EmbedDims returns the platform-standard embedding dimension. gemini-embedding
// honors OutputDimensionality so Embed pins the output to PlatformEmbeddingDims.
func (c *GeminiClient) EmbedDims(string) int { return PlatformEmbeddingDims }

// Embed calls EmbedContent with all inputs batched as separate Contents, pinning
// OutputDimensionality to the fixed pgvector column width. Returns one vector per
// input in order.
func (c *GeminiClient) Embed(ctx context.Context, modelID string, inputs []string) ([][]float32, error) {
	if len(inputs) == 0 {
		return nil, nil
	}
	contents := make([]*genai.Content, len(inputs))
	for i, in := range inputs {
		contents[i] = &genai.Content{Role: "user", Parts: []*genai.Part{genai.NewPartFromText(in)}}
	}
	dims := int32(PlatformEmbeddingDims)
	resp, err := c.sdk.Models.EmbedContent(ctx, modelID, contents, &genai.EmbedContentConfig{
		OutputDimensionality: &dims,
	})
	if err != nil {
		return nil, fmt.Errorf("gemini embeddings: %w", err)
	}
	if len(resp.Embeddings) != len(inputs) {
		return nil, fmt.Errorf("gemini embeddings: got %d vectors for %d inputs", len(resp.Embeddings), len(inputs))
	}
	out := make([][]float32, len(resp.Embeddings))
	for i, e := range resp.Embeddings {
		v := make([]float32, len(e.Values))
		copy(v, e.Values)
		out[i] = v
	}
	return out, nil
}