225 lines
6.5 KiB
Go
225 lines
6.5 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestParseSSETextAndUsage(t *testing.T) {
|
|
body := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"content":"你"},"index":0}]}`,
|
|
`data: {"choices":[{"delta":{"content":"好"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`,
|
|
`data: [DONE]`,
|
|
"",
|
|
}, "\n")
|
|
var deltas []string
|
|
var end Event
|
|
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
|
|
switch ev.Kind {
|
|
case KindTextDelta:
|
|
deltas = append(deltas, ev.Text)
|
|
case KindMessageEnd:
|
|
end = ev
|
|
default:
|
|
t.Fatalf("unexpected event %#v", ev)
|
|
}
|
|
}
|
|
if strings.Join(deltas, "") != "你好" {
|
|
t.Fatalf("deltas %v", deltas)
|
|
}
|
|
if end.Message.StopReason != StopEnd || end.Message.Content != "你好" {
|
|
t.Fatalf("end %#v", end.Message)
|
|
}
|
|
if end.Usage.TotalTokens != 5 {
|
|
t.Fatalf("usage %#v", end.Usage)
|
|
}
|
|
}
|
|
|
|
func TestParseSSEToolCallIncrements(t *testing.T) {
|
|
body := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.md\"}"}}]},"finish_reason":"tool_calls"}]}`,
|
|
`data: [DONE]`,
|
|
"",
|
|
}, "\n")
|
|
var end Event
|
|
deltas := 0
|
|
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
|
|
if ev.Kind == KindToolCallDelta {
|
|
deltas++
|
|
}
|
|
if ev.Kind == KindMessageEnd {
|
|
end = ev
|
|
}
|
|
}
|
|
if deltas != 3 {
|
|
t.Fatalf("tool deltas %d", deltas)
|
|
}
|
|
if end.Message.StopReason != StopToolCalls || len(end.Message.ToolCalls) != 1 {
|
|
t.Fatalf("end %#v", end.Message)
|
|
}
|
|
tc := end.Message.ToolCalls[0]
|
|
if tc.ID != "call_1" || tc.Name != "read_file" || tc.Arguments != `{"path":"a.md"}` {
|
|
t.Fatalf("tool call %#v", tc)
|
|
}
|
|
}
|
|
|
|
func TestParseSSEMalformedIsError(t *testing.T) {
|
|
body := "data: {not-json}\n\n"
|
|
var got Event
|
|
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
|
|
got = ev
|
|
}
|
|
if got.Kind != KindError || got.Message.StopReason != StopError || got.Retryable {
|
|
t.Fatalf("event %#v", got)
|
|
}
|
|
}
|
|
|
|
func TestParseSSECancel(t *testing.T) {
|
|
pr, pw := io.Pipe()
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ch := ParseSSE(ctx, pr)
|
|
go func() {
|
|
_, _ = io.WriteString(pw, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n")
|
|
cancel()
|
|
time.Sleep(20 * time.Millisecond)
|
|
_ = pw.Close()
|
|
}()
|
|
sawAbort := false
|
|
for ev := range ch {
|
|
if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
|
|
sawAbort = true
|
|
}
|
|
}
|
|
if !sawAbort {
|
|
t.Fatal("expected aborted event")
|
|
}
|
|
}
|
|
|
|
func TestClientHTTPErrorIsRetryable(t *testing.T) {
|
|
srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Header.Get("Authorization") != "Bearer secret" {
|
|
t.Errorf("auth header %q", r.Header.Get("Authorization"))
|
|
}
|
|
http.Error(w, "busy", http.StatusTooManyRequests)
|
|
})
|
|
c := &Client{BaseURL: srv, APIKey: "secret", HTTPClient: srvClient()}
|
|
var got Event
|
|
for ev := range c.Stream(context.Background(), "MiniMax-M2.7-highspeed", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
|
|
got = ev
|
|
}
|
|
if got.Kind != KindError || !got.Retryable || got.Message.StopReason != StopError {
|
|
t.Fatalf("event %#v", got)
|
|
}
|
|
}
|
|
|
|
func TestClientStreamsCompletion(t *testing.T) {
|
|
srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n")
|
|
})
|
|
c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: srvClient()}
|
|
var end Event
|
|
for ev := range c.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, []ToolDecl{{Name: "read_file", Parameters: []byte(`{"type":"object"}`)}}, StreamOptions{}) {
|
|
if ev.Kind == KindMessageEnd {
|
|
end = ev
|
|
}
|
|
}
|
|
if end.Message.Content != "ok" || end.Message.StopReason != StopEnd {
|
|
t.Fatalf("end %#v", end.Message)
|
|
}
|
|
}
|
|
|
|
func TestClientCancel(t *testing.T) {
|
|
started := make(chan struct{})
|
|
srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
|
|
close(started)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
fl, _ := w.(http.Flusher)
|
|
_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
|
|
if fl != nil {
|
|
fl.Flush()
|
|
}
|
|
<-r.Context().Done()
|
|
})
|
|
c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: &http.Client{}}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ch := c.Stream(ctx, "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{})
|
|
<-started
|
|
cancel()
|
|
saw := false
|
|
deadline := time.After(2 * time.Second)
|
|
for {
|
|
select {
|
|
case ev, ok := <-ch:
|
|
if !ok {
|
|
if !saw {
|
|
t.Fatal("channel closed without abort")
|
|
}
|
|
return
|
|
}
|
|
if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
|
|
saw = true
|
|
}
|
|
case <-deadline:
|
|
t.Fatal("timed out waiting for cancel")
|
|
}
|
|
}
|
|
}
|
|
|
|
func httpServer(t *testing.T, h http.HandlerFunc) string {
|
|
t.Helper()
|
|
srv := httptest.NewServer(h)
|
|
t.Cleanup(srv.Close)
|
|
return srv.URL
|
|
}
|
|
|
|
func srvClient() *http.Client {
|
|
return &http.Client{Timeout: 2 * time.Second}
|
|
}
|
|
|
|
func TestStreamErrorString(t *testing.T) {
|
|
if (&StreamError{Message: "boom"}).Error() != "boom" {
|
|
t.Fatal("error string")
|
|
}
|
|
}
|
|
|
|
func TestParseSSEReasoningDelta(t *testing.T) {
|
|
body := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"reasoning_content":"先分析需求"},"index":0}]}`,
|
|
`data: {"choices":[{"delta":{"reasoning_content":"再检索标准"},"index":0}]}`,
|
|
`data: {"choices":[{"delta":{"content":"结论"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`,
|
|
`data: [DONE]`,
|
|
"",
|
|
}, "\n")
|
|
var reasoning []string
|
|
var text []string
|
|
var end Event
|
|
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
|
|
switch ev.Kind {
|
|
case KindReasoningDelta:
|
|
reasoning = append(reasoning, ev.ReasoningDelta)
|
|
case KindTextDelta:
|
|
text = append(text, ev.Text)
|
|
case KindMessageEnd:
|
|
end = ev
|
|
}
|
|
}
|
|
if strings.Join(reasoning, "") != "先分析需求再检索标准" {
|
|
t.Fatalf("reasoning deltas %v", reasoning)
|
|
}
|
|
// 正文与思考链分流：Content 只含 text，不含 reasoning
|
|
if strings.Join(text, "") != "结论" {
|
|
t.Fatalf("text deltas %v", text)
|
|
}
|
|
if end.Message.Content != "结论" {
|
|
t.Fatalf("end content %q must not include reasoning", end.Message.Content)
|
|
}
|
|
}
|