You've already forked agentic-coding-workflow
feat(infra/llm): Gemini client with stream + count_tokens via google.golang.org/genai
This commit is contained in:
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user