Files
2026-07-13 13:00:08 +08:00

85 lines
2.1 KiB
Go

package openai
import "testing"
func runSplitter(deltas []string) (reasoning, text string) {
var t thinkSplitter
for _, d := range deltas {
r, txt := t.push(d)
reasoning += r
text += txt
}
r, txt := t.flush()
return reasoning + r, text + txt
}
func TestThinkSplitter(t *testing.T) {
cases := []struct {
name string
deltas []string
reasoning string
text string
}{
{
name: "whole block in one delta",
deltas: []string{"<think>reasoning here</think>the answer"},
reasoning: "reasoning here",
text: "the answer",
},
{
name: "open tag split across deltas",
deltas: []string{"<th", "ink>chain", " of thought</think>answer"},
reasoning: "chain of thought",
text: "answer",
},
{
name: "close tag split across deltas",
deltas: []string{"<think>thinking</thi", "nk>done"},
reasoning: "thinking",
text: "done",
},
{
name: "leading whitespace before think is dropped",
deltas: []string{"\n\n <think>r</think>\n\nanswer"},
reasoning: "r",
text: "answer",
},
{
name: "no think tag passes through as text",
deltas: []string{"just a normal ", "answer"},
reasoning: "",
text: "just a normal answer",
},
{
name: "think mentioned mid-answer is not hijacked",
deltas: []string{"the model emits <think> tags around its reasoning"},
reasoning: "",
text: "the model emits <think> tags around its reasoning",
},
{
name: "unterminated think block flushes as reasoning",
deltas: []string{"<think>still thinking when the stream ended"},
reasoning: "still thinking when the stream ended",
text: "",
},
{
name: "per-character streaming",
deltas: []string{"<", "t", "h", "i", "n", "k", ">", "a", "</", "think>", "b"},
reasoning: "a",
text: "b",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
r, txt := runSplitter(tc.deltas)
if r != tc.reasoning {
t.Errorf("reasoning = %q, want %q", r, tc.reasoning)
}
if txt != tc.text {
t.Errorf("text = %q, want %q", txt, tc.text)
}
})
}
}