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

This commit is contained in:
2026-05-04 09:46:47 +08:00
parent be751e76f6
commit a08a7249f0
4 changed files with 237 additions and 0 deletions
+13
View File
@@ -15,6 +15,9 @@ require (
) )
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 dario.cat/mergo v1.0.2 // indirect
github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect
github.com/Microsoft/go-winio v0.6.2 // 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-logr/stdr v1.2.2 // indirect
github.com/go-ole/go-ole v1.2.6 // indirect github.com/go-ole/go-ole v1.2.6 // indirect
github.com/go-viper/mapstructure/v2 v2.4.0 // 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/invopop/jsonschema v0.13.0 // indirect
github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 // indirect github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect
@@ -90,8 +98,13 @@ require (
go.opentelemetry.io/otel/trace v1.41.0 // indirect go.opentelemetry.io/otel/trace v1.41.0 // indirect
go.yaml.in/yaml/v3 v3.0.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect
golang.org/x/crypto v0.48.0 // 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/sync v0.19.0 // indirect
golang.org/x/sys v0.42.0 // indirect golang.org/x/sys v0.42.0 // indirect
golang.org/x/text v0.34.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 gopkg.in/yaml.v3 v3.0.1 // indirect
) )
+25
View File
@@ -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 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8=
dario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA= 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= 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.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 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= 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 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 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 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 h1:D/V0gu4zQ3cL2WKeVNVM4r2gLxGGf6McLwgXzRTo2RQ= 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.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.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.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-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.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.1.0/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/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-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/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 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 h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
+144
View File
@@ -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
}
+55
View File
@@ -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())
}