140 lines
3.9 KiB
Go
140 lines
3.9 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// sseServer 起一个假的 SSE 服务端，按脚本逐条吐数据。
|
|
func sseServer(t *testing.T, chunks []string, gap time.Duration) *httptest.Server {
|
|
t.Helper()
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
flusher, _ := w.(http.Flusher)
|
|
for _, c := range chunks {
|
|
fmt.Fprint(w, c)
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
time.Sleep(gap)
|
|
}
|
|
}))
|
|
}
|
|
|
|
func openAIChunk(content, reasoning string) string {
|
|
delta := fmt.Sprintf(`{"choices":[{"delta":{"content":%q,"reasoning_content":%q}}]}`, content, reasoning)
|
|
return "data: " + delta + "\n\n"
|
|
}
|
|
|
|
func TestOpenAIStreamsIncrementally(t *testing.T) {
|
|
srv := sseServer(t, []string{
|
|
openAIChunk("第一段", ""),
|
|
openAIChunk("第二段", ""),
|
|
openAIChunk("第三段", ""),
|
|
}, 60*time.Millisecond)
|
|
defer srv.Close()
|
|
|
|
prov, err := New(Options{Format: FormatOpenAI, BaseURL: srv.URL, Model: "m", APIKey: "k"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
type ev struct {
|
|
text string
|
|
at time.Duration
|
|
}
|
|
var evs []ev
|
|
start := time.Now()
|
|
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
|
|
evs = append(evs, ev{text: d.Text, at: time.Since(start)})
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(evs) != 3 {
|
|
t.Fatalf("应收到 3 段增量，实际 %d", len(evs))
|
|
}
|
|
// 第一段必须在服务端还没吐完时就到，否则就是"一次性返回"。
|
|
if evs[0].at > 50*time.Millisecond {
|
|
t.Fatalf("第一段增量到达太晚(%v)，说明没有真正流式", evs[0].at)
|
|
}
|
|
if got := evs[0].text + evs[1].text + evs[2].text; got != "第一段第二段第三段" {
|
|
t.Fatalf("增量内容不对: %q", got)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIParsesReasoning(t *testing.T) {
|
|
srv := sseServer(t, []string{
|
|
openAIChunk("", "先想一下…"),
|
|
openAIChunk("", "再想一下…"),
|
|
openAIChunk("正式回答", ""),
|
|
}, 10*time.Millisecond)
|
|
defer srv.Close()
|
|
|
|
prov, err := New(Options{Format: FormatOpenAI, BaseURL: srv.URL, Model: "m"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
var think, text strings.Builder
|
|
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
|
|
think.WriteString(d.Reasoning)
|
|
text.WriteString(d.Text)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if think.String() != "先想一下…再想一下…" {
|
|
t.Fatalf("思考内容解析不对: %q", think.String())
|
|
}
|
|
if text.String() != "正式回答" {
|
|
t.Fatalf("正文解析不对: %q", text.String())
|
|
}
|
|
}
|
|
|
|
func TestAnthropicStreamsThinkingAndText(t *testing.T) {
|
|
srv := sseServer(t, []string{
|
|
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"我在想…\"}}\n\n",
|
|
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"答案\"}}\n\n",
|
|
"event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":12}}\n\n",
|
|
}, 60*time.Millisecond)
|
|
defer srv.Close()
|
|
|
|
prov, err := New(Options{Format: FormatAnthropic, BaseURL: srv.URL, Model: "m"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
var think, text strings.Builder
|
|
start := time.Now()
|
|
var firstAt time.Duration
|
|
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
|
|
if firstAt == 0 {
|
|
firstAt = time.Since(start)
|
|
}
|
|
think.WriteString(d.Reasoning)
|
|
text.WriteString(d.Text)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if think.String() != "我在想…" {
|
|
t.Fatalf("思考内容解析不对: %q", think.String())
|
|
}
|
|
if text.String() != "答案" {
|
|
t.Fatalf("正文解析不对: %q", text.String())
|
|
}
|
|
if firstAt > 50*time.Millisecond {
|
|
t.Fatalf("第一段增量到达太晚(%v)，说明没有真正流式", firstAt)
|
|
}
|
|
}
|