highperfocused/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2025-11-11 16:57:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

export ThinkingParser

Browse Source
This commit is contained in:
Devon Rifkin
2025-06-05 10:22:32 -07:00
parent 0943001193
commit 0683efa637
3 changed files with 35 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
server/routes.go
View File

@@ -282,12 +282,12 @@ func (s *Server) GenerateHandler(c *gin.Context) {
prompt = b.String() prompt = b.String()
} }
var thinkingState *thinkingParser var thinkingState *ThinkingParser
openingTag, closingTag := inferThinkingTags(m.Template.Template) openingTag, closingTag := inferThinkingTags(m.Template.Template)
if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" { if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" {
thinkingState = &thinkingParser{ thinkingState = &ThinkingParser{
openingTag: openingTag, OpeningTag: openingTag,
closingTag: closingTag, ClosingTag: closingTag,
} }
} }
@@ -316,7 +316,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
} }
if thinkingState != nil { if thinkingState != nil {
thinking, content := thinkingState.addContent(cr.Content) thinking, content := thinkingState.AddContent(cr.Content)
res.Thinking = thinking res.Thinking = thinking
res.Response = content res.Response = content
} }
@@ -1522,12 +1522,12 @@ func (s *Server) ChatHandler(c *gin.Context) {
return return
} }
var thinkingState *thinkingParser var thinkingState *ThinkingParser
openingTag, closingTag := inferThinkingTags(m.Template.Template) openingTag, closingTag := inferThinkingTags(m.Template.Template)
if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" { if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" {
thinkingState = &thinkingParser{ thinkingState = &ThinkingParser{
openingTag: openingTag, OpeningTag: openingTag,
closingTag: closingTag, ClosingTag: closingTag,
} }
} }
@@ -1565,7 +1565,7 @@ func (s *Server) ChatHandler(c *gin.Context) {
} }
if thinkingState != nil { if thinkingState != nil {
thinkingContent, remainingContent := thinkingState.addContent(res.Message.Content) thinkingContent, remainingContent := thinkingState.AddContent(res.Message.Content)
if thinkingContent == "" && remainingContent == "" && !r.Done { if thinkingContent == "" && remainingContent == "" && !r.Done {
// need to accumulate more to decide what to send // need to accumulate more to decide what to send
return return
@@ -1676,11 +1676,11 @@ func filterThinkTags(msgs []api.Message, m *Model) []api.Message {
// change the user output), we should probably perform this filtering // change the user output), we should probably perform this filtering
// for all thinking models (not just qwen3 & deepseek-r1) since it tends // for all thinking models (not just qwen3 & deepseek-r1) since it tends
// to save tokens and improve quality. // to save tokens and improve quality.
thinkingState := &thinkingParser{ thinkingState := &ThinkingParser{
openingTag: "<think>", OpeningTag: "<think>",
closingTag: "</think>", ClosingTag: "</think>",
} }
_, content := thinkingState.addContent(msg.Content) _, content := thinkingState.AddContent(msg.Content)
msgs[i].Content = content msgs[i].Content = content
} }
} }

26
server/thinking.go
View File

@@ -46,17 +46,17 @@ func (s thinkingState) String() string {
} }
} }
type thinkingParser struct { type ThinkingParser struct {
state thinkingState state thinkingState
openingTag string OpeningTag string
closingTag string ClosingTag string
acc strings.Builder acc strings.Builder
} }
// addContent returns the thinking content and the non-thinking content that // AddContent returns the thinking content and the non-thinking content that
// should be immediately sent to the user. It will internally buffer if it needs // should be immediately sent to the user. It will internally buffer if it needs
// to see more raw content to disambiguate // to see more raw content to disambiguate
func (s *thinkingParser) addContent(content string) (string, string) { func (s *ThinkingParser) AddContent(content string) (string, string) {
s.acc.WriteString(content) s.acc.WriteString(content)
var thinkingSb, remainingSb strings.Builder var thinkingSb, remainingSb strings.Builder
@@ -76,12 +76,12 @@ func (s *thinkingParser) addContent(content string) (string, string) {
} }
// the additional bool return is true iff we should continue eating // the additional bool return is true iff we should continue eating
func eat(s *thinkingParser) (string, string, bool) { func eat(s *ThinkingParser) (string, string, bool) {
switch s.state { switch s.state {
case thinkingState_LookingForOpening: case thinkingState_LookingForOpening:
trimmed := strings.TrimLeftFunc(s.acc.String(), unicode.IsSpace) trimmed := strings.TrimLeftFunc(s.acc.String(), unicode.IsSpace)
if strings.HasPrefix(trimmed, s.openingTag) { if strings.HasPrefix(trimmed, s.OpeningTag) {
after := strings.Join(strings.Split(trimmed, s.openingTag)[1:], s.openingTag) after := strings.Join(strings.Split(trimmed, s.OpeningTag)[1:], s.OpeningTag)
after = strings.TrimLeftFunc(after, unicode.IsSpace) after = strings.TrimLeftFunc(after, unicode.IsSpace)
// after might contain more than just thinking tokens, so we continue // after might contain more than just thinking tokens, so we continue
// parsing instead of returning it as thinking tokens here // parsing instead of returning it as thinking tokens here
@@ -93,7 +93,7 @@ func eat(s *thinkingParser) (string, string, bool) {
s.state = thinkingState_Thinking s.state = thinkingState_Thinking
} }
return "", "", true return "", "", true
} else if strings.HasPrefix(s.openingTag, trimmed) { } else if strings.HasPrefix(s.OpeningTag, trimmed) {
// partial opening seen, so let's keep accumulating // partial opening seen, so let's keep accumulating
return "", "", false return "", "", false
} else if trimmed == "" { } else if trimmed == "" {
@@ -119,10 +119,10 @@ func eat(s *thinkingParser) (string, string, bool) {
} }
case thinkingState_Thinking: case thinkingState_Thinking:
acc := s.acc.String() acc := s.acc.String()
if strings.Contains(acc, s.closingTag) { if strings.Contains(acc, s.ClosingTag) {
split := strings.Split(acc, s.closingTag) split := strings.Split(acc, s.ClosingTag)
thinking := split[0] thinking := split[0]
remaining := strings.Join(split[1:], s.closingTag) remaining := strings.Join(split[1:], s.ClosingTag)
remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace) remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
s.acc.Reset() s.acc.Reset()
if remaining == "" { if remaining == "" {
@@ -131,7 +131,7 @@ func eat(s *thinkingParser) (string, string, bool) {
s.state = thinkingState_ThinkingDone s.state = thinkingState_ThinkingDone
} }
return thinking, remaining, false return thinking, remaining, false
} else if overlapLen := overlap(acc, s.closingTag); overlapLen > 0 { } else if overlapLen := overlap(acc, s.ClosingTag); overlapLen > 0 {
thinking := acc[:len(acc)-overlapLen] thinking := acc[:len(acc)-overlapLen]
remaining := acc[len(acc)-overlapLen:] remaining := acc[len(acc)-overlapLen:]
s.acc.Reset() s.acc.Reset()

16
server/thinking_test.go
View File

@@ -26,11 +26,11 @@ func TestExtractThinking(t *testing.T) {
}, },
} }
for i, tt := range tests { for i, tt := range tests {
parser := thinkingParser{ parser := ThinkingParser{
openingTag: "<think>", OpeningTag: "<think>",
closingTag: "</think>", ClosingTag: "</think>",
} }
gotThinking, gotContent := parser.addContent(tt.in) gotThinking, gotContent := parser.AddContent(tt.in)
if gotContent != tt.wantContent || gotThinking != tt.wantThink { if gotContent != tt.wantContent || gotThinking != tt.wantThink {
t.Errorf("case %d: got (%q,%q), want (%q,%q)", i, gotThinking, gotContent, tt.wantThink, tt.wantContent) t.Errorf("case %d: got (%q,%q), want (%q,%q)", i, gotThinking, gotContent, tt.wantThink, tt.wantContent)
} }
@@ -259,15 +259,15 @@ func TestThinkingStreaming(t *testing.T) {
} }
for _, c := range cases { for _, c := range cases {
parser := thinkingParser{ parser := ThinkingParser{
openingTag: "<think>", OpeningTag: "<think>",
closingTag: "</think>", ClosingTag: "</think>",
} }
if c.skip { if c.skip {
continue continue
} }
for i, step := range c.steps { for i, step := range c.steps {
thinking, content := parser.addContent(step.input) thinking, content := parser.AddContent(step.input)
if content != step.wantContent || thinking != step.wantThinking { if content != step.wantContent || thinking != step.wantThinking {
t.Errorf("case %q (step %d): got (%q,%q), want (%q,%q)", c.desc, i, content, thinking, step.wantContent, step.wantThinking) t.Errorf("case %q (step %d): got (%q,%q), want (%q,%q)", c.desc, i, content, thinking, step.wantContent, step.wantThinking)
} }