You've already forked agentic-coding-workflow
56 lines
1.6 KiB
Go
56 lines
1.6 KiB
Go
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())
|
|
}
|