ollama source for Momentry Core verification

model/parsers/nemotron3nano.go

@@ -0,0 +1,188 @@
+package parsers
+
+import (
+	"strings"
+	"unicode"
+
+	"github.com/ollama/ollama/api"
+)
+
+type Nemotron3NanoParserState int
+
+const (
+	Nemotron3NanoCollectingThinking Nemotron3NanoParserState = iota
+	Nemotron3NanoSkipWhitespaceAfterThinking
+	Nemotron3NanoCollectingContent
+)
+
+const (
+	nemotronThinkOpen    = "<think>"
+	nemotronThinkClose   = "</think>"
+	nemotronToolCallOpen = "<tool_call>"
+)
+
+type Nemotron3NanoParser struct {
+	state                  Nemotron3NanoParserState
+	buffer                 strings.Builder
+	toolParser             *Qwen3CoderParser
+	maybeThinkingOpenAtBOL bool
+	skipThinkingLeadingWS  bool
+}
+
+func (p *Nemotron3NanoParser) HasToolSupport() bool     { return true }
+func (p *Nemotron3NanoParser) HasThinkingSupport() bool { return true }
+
+func (p *Nemotron3NanoParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
+	p.toolParser = &Qwen3CoderParser{}
+	p.toolParser.Init(tools, nil, nil)
+	p.buffer.Reset()
+	p.maybeThinkingOpenAtBOL = false
+	p.skipThinkingLeadingWS = false
+
+	thinkingEnabled := thinkValue == nil || thinkValue.Bool()
+	prefill := lastMessage != nil && lastMessage.Role == "assistant"
+
+	if !thinkingEnabled || (prefill && lastMessage.Content != "") {
+		p.state = Nemotron3NanoCollectingContent
+	} else {
+		p.state = Nemotron3NanoCollectingThinking
+		p.maybeThinkingOpenAtBOL = true
+	}
+
+	return tools
+}
+
+func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
+	if p.state == Nemotron3NanoCollectingContent {
+		return p.toolParser.Add(s, done)
+	}
+
+	if p.state == Nemotron3NanoSkipWhitespaceAfterThinking {
+		s = strings.TrimLeftFunc(s, unicode.IsSpace)
+		if s == "" {
+			return "", "", nil, nil
+		}
+		p.state = Nemotron3NanoCollectingContent
+		return p.toolParser.Add(s, done)
+	}
+
+	// Nemotron3NanoCollectingThinking - buffer and look for end markers
+	p.buffer.WriteString(s)
+	if p.skipThinkingLeadingWS {
+		trimmed := strings.TrimLeftFunc(p.buffer.String(), unicode.IsSpace)
+		p.buffer.Reset()
+		p.buffer.WriteString(trimmed)
+		if trimmed == "" {
+			return "", "", nil, nil
+		}
+		p.skipThinkingLeadingWS = false
+	}
+	if p.stripOpeningThinkTag() {
+		return p.Add("", done)
+	}
+	if p.maybeThinkingOpenAtBOL {
+		bufStr := p.buffer.String()
+		trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
+		if trimmed == "" || overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
+			if len(trimmed) != len(bufStr) {
+				p.buffer.Reset()
+				p.buffer.WriteString(trimmed)
+			}
+			return "", "", nil, nil
+		}
+	}
+	bufStr := p.buffer.String()
+
+	// Look for end of thinking: </think> or <tool_call> (model may skip </think>)
+	thinkIdx := strings.Index(bufStr, nemotronThinkClose)
+	toolIdx := strings.Index(bufStr, nemotronToolCallOpen)
+
+	var endIdx int = -1
+	var remainder string
+
+	if thinkIdx != -1 && (toolIdx == -1 || thinkIdx < toolIdx) {
+		endIdx = thinkIdx
+		remainder = strings.TrimLeftFunc(bufStr[thinkIdx+len(nemotronThinkClose):], unicode.IsSpace)
+	} else if toolIdx != -1 {
+		endIdx = toolIdx
+		remainder = bufStr[toolIdx:] // Include <tool_call> tag
+	}
+
+	if endIdx != -1 {
+		thinking = strings.TrimRightFunc(bufStr[:endIdx], unicode.IsSpace)
+		p.buffer.Reset()
+
+		if remainder == "" {
+			p.state = Nemotron3NanoSkipWhitespaceAfterThinking
+		} else {
+			p.state = Nemotron3NanoCollectingContent
+			content, _, calls, err = p.toolParser.Add(remainder, done)
+		}
+		return content, thinking, calls, err
+	}
+
+	// No end marker - emit unambiguous thinking
+	thinking = p.emitThinking(bufStr)
+	return "", thinking, nil, nil
+}
+
+// emitThinking returns unambiguous thinking content, keeping potential partial tags in buffer
+func (p *Nemotron3NanoParser) emitThinking(bufStr string) string {
+	// Check for partial </think> or <tool_call> at end
+	thinkOverlap := overlap(bufStr, nemotronThinkClose)
+	toolOverlap := overlap(bufStr, nemotronToolCallOpen)
+	maxOverlap := max(thinkOverlap, toolOverlap)
+
+	if maxOverlap > 0 {
+		unambiguous := bufStr[:len(bufStr)-maxOverlap]
+		unambiguous = strings.TrimRightFunc(unambiguous, unicode.IsSpace)
+		p.buffer.Reset()
+		p.buffer.WriteString(bufStr[len(bufStr)-maxOverlap:])
+		return unambiguous
+	}
+
+	// No partial tags - emit all but trailing whitespace
+	wsLen := trailingWhitespaceLen(bufStr)
+	if wsLen > 0 {
+		unambiguous := bufStr[:len(bufStr)-wsLen]
+		p.buffer.Reset()
+		p.buffer.WriteString(bufStr[len(bufStr)-wsLen:])
+		return unambiguous
+	}
+
+	// Nothing to hold back
+	p.buffer.Reset()
+	return bufStr
+}
+
+func (p *Nemotron3NanoParser) stripOpeningThinkTag() bool {
+	if !p.maybeThinkingOpenAtBOL {
+		return false
+	}
+
+	bufStr := p.buffer.String()
+	trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
+	if trimmed == "" {
+		p.buffer.Reset()
+		return false
+	}
+
+	if strings.HasPrefix(trimmed, nemotronThinkOpen) {
+		p.buffer.Reset()
+		p.buffer.WriteString(strings.TrimLeftFunc(trimmed[len(nemotronThinkOpen):], unicode.IsSpace))
+		p.maybeThinkingOpenAtBOL = false
+		p.skipThinkingLeadingWS = true
+		return true
+	}
+
+	if overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
+		if len(trimmed) != len(bufStr) {
+			p.buffer.Reset()
+			p.buffer.WriteString(trimmed)
+		}
+		return false
+	}
+
+	p.maybeThinkingOpenAtBOL = false
+	return false
+}