| 1 | package openai |
| 2 | |
| 3 | import "testing" |
| 4 | |
| 5 | func runSplitter(deltas []string) (reasoning, text string) { |
| 6 | var t thinkSplitter |
| 7 | for _, d := range deltas { |
| 8 | r, txt := t.push(d) |
| 9 | reasoning += r |
| 10 | text += txt |
| 11 | } |
| 12 | r, txt := t.flush() |
| 13 | return reasoning + r, text + txt |
| 14 | } |
| 15 | |
| 16 | func TestThinkSplitter(t *testing.T) { |
| 17 | cases := []struct { |
| 18 | name string |
| 19 | deltas []string |
| 20 | reasoning string |
| 21 | text string |
| 22 | }{ |
| 23 | { |
| 24 | name: "whole block in one delta", |
| 25 | deltas: []string{"<think>reasoning here</think>the answer"}, |
| 26 | reasoning: "reasoning here", |
| 27 | text: "the answer", |
| 28 | }, |
| 29 | { |
| 30 | name: "open tag split across deltas", |
| 31 | deltas: []string{"<th", "ink>chain", " of thought</think>answer"}, |
| 32 | reasoning: "chain of thought", |
| 33 | text: "answer", |
| 34 | }, |
| 35 | { |
| 36 | name: "close tag split across deltas", |
| 37 | deltas: []string{"<think>thinking</thi", "nk>done"}, |
| 38 | reasoning: "thinking", |
| 39 | text: "done", |
| 40 | }, |
| 41 | { |
| 42 | name: "leading whitespace before think is dropped", |
| 43 | deltas: []string{"\n\n <think>r</think>\n\nanswer"}, |
| 44 | reasoning: "r", |
| 45 | text: "answer", |
| 46 | }, |
| 47 | { |
| 48 | name: "no think tag passes through as text", |
| 49 | deltas: []string{"just a normal ", "answer"}, |
| 50 | reasoning: "", |
| 51 | text: "just a normal answer", |
| 52 | }, |
| 53 | { |
| 54 | name: "think mentioned mid-answer is not hijacked", |
| 55 | deltas: []string{"the model emits <think> tags around its reasoning"}, |
| 56 | reasoning: "", |
| 57 | text: "the model emits <think> tags around its reasoning", |
| 58 | }, |
| 59 | { |
| 60 | name: "unterminated think block flushes as reasoning", |
| 61 | deltas: []string{"<think>still thinking when the stream ended"}, |
| 62 | reasoning: "still thinking when the stream ended", |
| 63 | text: "", |
| 64 | }, |
| 65 | { |
| 66 | name: "per-character streaming", |
| 67 | deltas: []string{"<", "t", "h", "i", "n", "k", ">", "a", "</", "think>", "b"}, |
| 68 | reasoning: "a", |
| 69 | text: "b", |
| 70 | }, |
| 71 | } |
| 72 | |
| 73 | for _, tc := range cases { |
| 74 | t.Run(tc.name, func(t *testing.T) { |
| 75 | r, txt := runSplitter(tc.deltas) |
| 76 | if r != tc.reasoning { |
| 77 | t.Errorf("reasoning = %q, want %q", r, tc.reasoning) |
| 78 | } |
| 79 | if txt != tc.text { |
| 80 | t.Errorf("text = %q, want %q", txt, tc.text) |
| 81 | } |
| 82 | }) |
| 83 | } |
| 84 | } |
| 85 |