Merge pull request #10987 from ollama/drifkin/export-thinking-parser

export ThinkingParser

server/routes.go

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

server/thinking.go

@@ -46,17 +46,17 @@ func (s thinkingState) String() string {
 	}
 }
 
-type thinkingParser struct {
+type ThinkingParser struct {
 	state      thinkingState
-	openingTag string
-	closingTag string
+	OpeningTag string
+	ClosingTag string
 	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
 // 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)
 
 	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
-func eat(s *thinkingParser) (string, string, bool) {
+func eat(s *ThinkingParser) (string, string, bool) {
 	switch s.state {
 	case thinkingState_LookingForOpening:
 		trimmed := strings.TrimLeftFunc(s.acc.String(), unicode.IsSpace)
-		if strings.HasPrefix(trimmed, s.openingTag) {
-			after := strings.Join(strings.Split(trimmed, s.openingTag)[1:], s.openingTag)
+		if strings.HasPrefix(trimmed, s.OpeningTag) {
+			after := strings.Join(strings.Split(trimmed, s.OpeningTag)[1:], s.OpeningTag)
 			after = strings.TrimLeftFunc(after, unicode.IsSpace)
 			// after might contain more than just thinking tokens, so we continue
 			// parsing instead of returning it as thinking tokens here
@@ -93,7 +93,7 @@ func eat(s *thinkingParser) (string, string, bool) {
 				s.state = thinkingState_Thinking
 			}
 			return "", "", true
-		} else if strings.HasPrefix(s.openingTag, trimmed) {
+		} else if strings.HasPrefix(s.OpeningTag, trimmed) {
 			// partial opening seen, so let's keep accumulating
 			return "", "", false
 		} else if trimmed == "" {
@@ -119,10 +119,10 @@ func eat(s *thinkingParser) (string, string, bool) {
 		}
 	case thinkingState_Thinking:
 		acc := s.acc.String()
-		if strings.Contains(acc, s.closingTag) {
-			split := strings.Split(acc, s.closingTag)
+		if strings.Contains(acc, s.ClosingTag) {
+			split := strings.Split(acc, s.ClosingTag)
 			thinking := split[0]
-			remaining := strings.Join(split[1:], s.closingTag)
+			remaining := strings.Join(split[1:], s.ClosingTag)
 			remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
 			s.acc.Reset()
 			if remaining == "" {
@@ -131,7 +131,7 @@ func eat(s *thinkingParser) (string, string, bool) {
 				s.state = thinkingState_ThinkingDone
 			}
 			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]
 			remaining := acc[len(acc)-overlapLen:]
 			s.acc.Reset()

server/thinking_test.go

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