74 lines
1.5 KiB
Go
74 lines
1.5 KiB
Go
package stream
|
|
|
|
import "strings"
|
|
|
|
type Parser struct {
|
|
inThink bool
|
|
buffer string
|
|
}
|
|
|
|
const (
|
|
thinkOpenTag = "<think>"
|
|
thinkCloseTag = "</think>"
|
|
)
|
|
|
|
func (p *Parser) Accept(delta string) (visible string, reasoning string) {
|
|
p.buffer += delta
|
|
for p.buffer != "" {
|
|
if p.inThink {
|
|
idx := strings.Index(p.buffer, thinkCloseTag)
|
|
if idx >= 0 {
|
|
reasoning += p.buffer[:idx]
|
|
p.buffer = p.buffer[idx+len(thinkCloseTag):]
|
|
p.inThink = false
|
|
continue
|
|
}
|
|
keep := tagPrefixSuffixLen(p.buffer, thinkCloseTag)
|
|
if len(p.buffer) > keep {
|
|
reasoning += p.buffer[:len(p.buffer)-keep]
|
|
p.buffer = p.buffer[len(p.buffer)-keep:]
|
|
}
|
|
return visible, reasoning
|
|
}
|
|
|
|
idx := strings.Index(p.buffer, thinkOpenTag)
|
|
if idx >= 0 {
|
|
visible += p.buffer[:idx]
|
|
p.buffer = p.buffer[idx+len(thinkOpenTag):]
|
|
p.inThink = true
|
|
continue
|
|
}
|
|
keep := tagPrefixSuffixLen(p.buffer, thinkOpenTag)
|
|
if len(p.buffer) > keep {
|
|
visible += p.buffer[:len(p.buffer)-keep]
|
|
p.buffer = p.buffer[len(p.buffer)-keep:]
|
|
}
|
|
return visible, reasoning
|
|
}
|
|
return visible, reasoning
|
|
}
|
|
|
|
func (p *Parser) Flush() (visible string, reasoning string) {
|
|
if p.inThink {
|
|
reasoning = p.buffer
|
|
} else {
|
|
visible = p.buffer
|
|
}
|
|
p.buffer = ""
|
|
p.inThink = false
|
|
return visible, reasoning
|
|
}
|
|
|
|
func tagPrefixSuffixLen(text, tag string) int {
|
|
limit := len(tag) - 1
|
|
if len(text) < limit {
|
|
limit = len(text)
|
|
}
|
|
for i := limit; i > 0; i-- {
|
|
if strings.HasPrefix(tag, text[len(text)-i:]) {
|
|
return i
|
|
}
|
|
}
|
|
return 0
|
|
}
|