feat(infra/llm): Gemini client with stream + count_tokens via google.golang.org/genai

2026-05-04 09:46:47 +08:00
parent be751e76f6
commit a08a7249f0
4 changed files with 237 additions and 0 deletions
@@ -0,0 +1,144 @@
+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.
+var _ LLMClient = (*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
+}
@@ -0,0 +1,55 @@
+package llm
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// Gemini SSE 行格式：每行 "data: {json}\n"，body 内含完整 candidates。
+const geminiStreamFixture = `data: {"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"finishReason":null}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1}}
+
+data: {"candidates":[{"content":{"parts":[{"text":" world"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}
+
+`
+
+func TestGemini_Stream_TextDelta(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte(geminiStreamFixture))
+	}))
+	defer srv.Close()
+
+	c, _ := NewGemini("k", srv.URL)
+	ch, err := c.Stream(context.Background(), "gemini-2.0-flash", Request{
+		Messages:        []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}},
+		MaxOutputTokens: 100,
+	})
+	require.NoError(t, err)
+
+	var deltas []string
+	var sawUsage, sawDone bool
+	for ev := range ch {
+		switch ev.Type {
+		case EventTextDelta:
+			deltas = append(deltas, ev.Text)
+		case EventUsage:
+			sawUsage = true
+			require.Equal(t, 10, ev.Usage.PromptTokens)
+		case EventDone:
+			sawDone = true
+		}
+	}
+	require.Equal(t, "Hello world", strings.Join(deltas, ""))
+	require.True(t, sawUsage)
+	require.True(t, sawDone)
+}
+
+func TestGemini_Provider(t *testing.T) {
+	c, _ := NewGemini("k", "http://localhost")
+	require.Equal(t, "gemini", c.Provider())
+}