优化代码

internal/llm/stream_test.go

@@ -0,0 +1,139 @@
+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)
+	}
+}