129 lines
3.1 KiB
Go
129 lines
3.1 KiB
Go
package llm
|
|
|
|
import "strings"
|
|
|
|
// thinkFilter removes inline reasoning blocks ("<think>…</think>") that some
|
|
// OpenAI-compatible backends (e.g. MiniMax-M2.x on vLLM, seen on Dahl) put into
|
|
// the normal content stream instead of a separate reasoning field. Tags may be
|
|
// split across chunks. Text inside a block is never streamed to the user.
|
|
type thinkFilter struct {
|
|
pending string // undecided tail that may be the start of a tag
|
|
inThink bool
|
|
think strings.Builder // reasoning text of the current (unclosed) block
|
|
out strings.Builder // visible content
|
|
sawOpen bool
|
|
}
|
|
|
|
const (
|
|
openTag = "<think>"
|
|
closeTag = "</think>"
|
|
gap = "\n\n\n"
|
|
)
|
|
|
|
// push consumes a chunk and returns the text that may be shown now.
|
|
func (f *thinkFilter) push(s string) string {
|
|
buf := f.pending + s
|
|
f.pending = ""
|
|
var vis strings.Builder
|
|
for buf != "" {
|
|
tag := openTag
|
|
if f.inThink {
|
|
tag = closeTag
|
|
}
|
|
i, n := strings.Index(buf, tag), len(tag)
|
|
// Dahl/vLLM sometimes swallows "</think>" and leaves only a "\n\n\n" gap;
|
|
// treat that gap as the end of the block so the answer still streams live.
|
|
if f.inThink {
|
|
if j := strings.Index(buf, gap); j >= 0 && (i < 0 || j < i) {
|
|
i, n = j, len(gap)
|
|
}
|
|
}
|
|
if i >= 0 {
|
|
f.emit(&vis, buf[:i])
|
|
buf = buf[i+n:]
|
|
if f.inThink {
|
|
f.think.Reset()
|
|
} else {
|
|
f.sawOpen = true
|
|
}
|
|
f.inThink = !f.inThink
|
|
continue
|
|
}
|
|
// a stray closing tag outside a block is dropped
|
|
if !f.inThink {
|
|
if i := strings.Index(buf, closeTag); i >= 0 {
|
|
f.emit(&vis, buf[:i])
|
|
buf = buf[i+len(closeTag):]
|
|
continue
|
|
}
|
|
}
|
|
keep := partialSuffix(buf, tag)
|
|
if f.inThink {
|
|
if k := partialSuffix(buf, gap); k > keep {
|
|
keep = k
|
|
}
|
|
} else {
|
|
if k := partialSuffix(buf, closeTag); k > keep {
|
|
keep = k
|
|
}
|
|
}
|
|
f.emit(&vis, buf[:len(buf)-keep])
|
|
f.pending = buf[len(buf)-keep:]
|
|
buf = ""
|
|
}
|
|
return vis.String()
|
|
}
|
|
|
|
func (f *thinkFilter) emit(vis *strings.Builder, s string) {
|
|
if s == "" {
|
|
return
|
|
}
|
|
if f.inThink {
|
|
f.think.WriteString(s)
|
|
return
|
|
}
|
|
// drop whitespace directly after a closed block at the start of the answer
|
|
if f.out.Len() == 0 {
|
|
s = strings.TrimLeft(s, " \n\r\t")
|
|
if s == "" {
|
|
return
|
|
}
|
|
}
|
|
f.out.WriteString(s)
|
|
vis.WriteString(s)
|
|
}
|
|
|
|
// finish flushes the remainder. If a block was never closed (observed when the
|
|
// backend swallows "</think>"), the answer is recovered from the text after the
|
|
// first blank-line gap so the user does not end up with nothing.
|
|
func (f *thinkFilter) finish() string {
|
|
var vis strings.Builder
|
|
if f.pending != "" {
|
|
p := f.pending
|
|
f.pending = ""
|
|
f.emit(&vis, p)
|
|
}
|
|
if f.inThink {
|
|
t := f.think.String()
|
|
f.inThink = false
|
|
f.think.Reset()
|
|
ans := ""
|
|
if i := strings.Index(t, "\n\n\n"); i >= 0 {
|
|
ans = t[i:]
|
|
} else if f.out.Len() == 0 {
|
|
ans = t // better a visible reasoning than an empty reply
|
|
}
|
|
f.emit(&vis, strings.TrimSpace(ans))
|
|
}
|
|
return vis.String()
|
|
}
|
|
|
|
// partialSuffix returns the length of the longest suffix of s that is a proper prefix of tag.
|
|
func partialSuffix(s, tag string) int {
|
|
for n := len(tag) - 1; n > 0; n-- {
|
|
if len(s) >= n && strings.HasSuffix(s, tag[:n]) {
|
|
return n
|
|
}
|
|
}
|
|
return 0
|
|
}
|