From a08a7249f0038031e77c01541e2eb6aa69262efe Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Mon, 4 May 2026 09:46:47 +0800 Subject: [PATCH] feat(infra/llm): Gemini client with stream + count_tokens via google.golang.org/genai --- go.mod | 13 +++ go.sum | 25 ++++++ internal/infra/llm/gemini.go | 144 ++++++++++++++++++++++++++++++ internal/infra/llm/gemini_test.go | 55 ++++++++++++ 4 files changed, 237 insertions(+) create mode 100644 internal/infra/llm/gemini.go create mode 100644 internal/infra/llm/gemini_test.go diff --git a/go.mod b/go.mod index 8669704..f7eca3d 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,9 @@ require ( ) require ( + cloud.google.com/go v0.121.6 // indirect + cloud.google.com/go/auth v0.16.4 // indirect + cloud.google.com/go/compute/metadata v0.8.0 // indirect dario.cat/mergo v1.0.2 // indirect github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect github.com/Microsoft/go-winio v0.6.2 // indirect @@ -40,6 +43,11 @@ require ( github.com/go-logr/stdr v1.2.2 // indirect github.com/go-ole/go-ole v1.2.6 // indirect github.com/go-viper/mapstructure/v2 v2.4.0 // indirect + github.com/google/go-cmp v0.7.0 // indirect + github.com/google/s2a-go v0.1.9 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect + github.com/googleapis/gax-go/v2 v2.15.0 // indirect + github.com/gorilla/websocket v1.5.3 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect @@ -90,8 +98,13 @@ require ( go.opentelemetry.io/otel/trace v1.41.0 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect golang.org/x/crypto v0.48.0 // indirect + golang.org/x/net v0.49.0 // indirect golang.org/x/sync v0.19.0 // indirect golang.org/x/sys v0.42.0 // indirect golang.org/x/text v0.34.0 // indirect + google.golang.org/genai v1.55.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20250818200422-3122310a409c // indirect + google.golang.org/grpc v1.74.2 // indirect + google.golang.org/protobuf v1.36.7 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index dfc6d97..5f207c7 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,9 @@ +cloud.google.com/go v0.121.6 h1:waZiuajrI28iAf40cWgycWNgaXPO06dupuS+sgibK6c= +cloud.google.com/go v0.121.6/go.mod h1:coChdst4Ea5vUpiALcYKXEpR1S9ZgXbhEzzMcMR66vI= +cloud.google.com/go/auth v0.16.4 h1:fXOAIQmkApVvcIn7Pc2+5J8QTMVbUGLscnSVNl11su8= +cloud.google.com/go/auth v0.16.4/go.mod h1:j10ncYwjX/g3cdX7GpEzsdM+d+ZNsXAbb6qXA7p1Y5M= +cloud.google.com/go/compute/metadata v0.8.0 h1:HxMRIbao8w17ZX6wBnjhcDkW6lTFpgcaobyVfZWqRLA= +cloud.google.com/go/compute/metadata v0.8.0/go.mod h1:sYOGTp851OV9bOFJ9CH7elVvyzopvWQFNNghtDQ/Biw= dario.cat/mergo v1.0.2 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8= dario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA= github.com/AdaLogics/go-fuzz-headers v0.0.0-20240806141605-e8a1dd7889d6 h1:He8afgbRMd7mFxO99hRNu+6tazq8nFF9lIwo9JFroBk= @@ -71,8 +77,16 @@ github.com/golang-migrate/migrate/v4 v4.19.1/go.mod h1:CTcgfjxhaUtsLipnLoQRWCrjY github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= +github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= +github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= +github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81vgd/bo= +github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 h1:D/V0gu4zQ3cL2WKeVNVM4r2gLxGGf6McLwgXzRTo2RQ= @@ -220,6 +234,8 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= +golang.org/x/net v0.49.0 h1:eeHFmOGUTtaaPSGNmjBKpbng9MulQsJURQUAfUwY++o= +golang.org/x/net v0.49.0/go.mod h1:/ysNB2EvaqvesRkuLAyjI1ycPZlQHM3q01F02UY/MV8= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -259,6 +275,15 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/genai v1.55.0 h1:iLHGk4Bj/IZ/GNNZb7hYqwSJMRBvqLeu2Hb6YQ+rYGw= +google.golang.org/genai v1.55.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= +google.golang.org/genproto v0.0.0-20250603155806-513f23925822 h1:rHWScKit0gvAPuOnu87KpaYtjK5zBMLcULh7gxkCXu4= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250818200422-3122310a409c h1:qXWI/sQtv5UKboZ/zUk7h+mrf/lXORyI+n9DKDAusdg= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250818200422-3122310a409c/go.mod h1:gw1tLEfykwDz2ET4a12jcXt4couGAm7IwsVaTy0Sflo= +google.golang.org/grpc v1.74.2 h1:WoosgB65DlWVC9FqI82dGsZhWFNBSLjQ84bjROOpMu4= +google.golang.org/grpc v1.74.2/go.mod h1:CtQ+BGjaAIXHs/5YS3i473GqwBBa1zGQNevxdeBEXrM= +google.golang.org/protobuf v1.36.7 h1:IgrO7UwFQGJdRNXH/sQux4R1Dj1WAKcLElzeeRaXV2A= +google.golang.org/protobuf v1.36.7/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= diff --git a/internal/infra/llm/gemini.go b/internal/infra/llm/gemini.go new file mode 100644 index 0000000..1c963f0 --- /dev/null +++ b/internal/infra/llm/gemini.go @@ -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 +} diff --git a/internal/infra/llm/gemini_test.go b/internal/infra/llm/gemini_test.go new file mode 100644 index 0000000..e62f8a2 --- /dev/null +++ b/internal/infra/llm/gemini_test.go @@ -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()) +}