model: update lfm2 parser/renderer for optional thinking (#16359)

This commit is contained in:
Jeffrey Morgan
2026-06-14 20:37:08 -07:00
committed by GitHub
parent 7ea692cb2b
commit 993acc7504
4 changed files with 201 additions and 24 deletions

View File

@@ -14,7 +14,12 @@ import (
type LFM2ParserState int
const (
LFM2CollectingThinking LFM2ParserState = iota
// LFM2LookingForThinking is the initial state when thinking is enabled.
// LFM2 models emit an explicit <think> tag only when they reason; a direct
// answer has no tag at all. This state waits to see which one the output
// begins with before committing to thinking or content.
LFM2LookingForThinking LFM2ParserState = iota
LFM2CollectingThinking
LFM2CollectingContent
LFM2CollectingToolCalls
)
@@ -70,8 +75,10 @@ func (p *LFM2Parser) setInitialState(lastMessage *api.Message, thinkValue *api.T
return
}
p.state = LFM2CollectingThinking
p.needsThinkingLeadingTrim = true
// Thinking is enabled, but the model decides per-turn whether to reason. Wait
// for a leading <think> tag before treating output as thinking; otherwise it's
// a direct answer.
p.state = LFM2LookingForThinking
}
func (p *LFM2Parser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
@@ -109,6 +116,18 @@ func (lfm2EventToolCall) isLFM2Event() {}
func (p *LFM2Parser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
p.buffer.WriteString(s)
// On the final chunk a partial "<think>" prefix can never complete, so commit
// the buffered output as a direct answer rather than withholding it.
if done && p.state == LFM2LookingForThinking {
trimmed := strings.TrimLeftFunc(p.buffer.String(), unicode.IsSpace)
if !strings.HasPrefix(trimmed, lfm2ThinkingOpenTag) {
p.buffer.Reset()
p.buffer.WriteString(trimmed)
p.state = LFM2CollectingContent
}
}
events := p.parseEvents()
var toolCalls []api.ToolCall
@@ -180,6 +199,28 @@ func (p *LFM2Parser) eat() ([]lfm2Event, bool) {
}
switch p.state {
case LFM2LookingForThinking:
// Decide whether this turn is a reasoning turn (begins with <think>) or a
// direct answer (no tag). Leading whitespace is ignored either way.
trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
if strings.HasPrefix(trimmed, lfm2ThinkingOpenTag) {
after := trimmed[len(lfm2ThinkingOpenTag):]
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = LFM2CollectingThinking
p.needsThinkingLeadingTrim = true
return events, true
}
if trimmed == "" || strings.HasPrefix(lfm2ThinkingOpenTag, trimmed) {
// Only whitespace so far, or a partial "<think>" prefix; wait for more.
return events, false
}
// Direct answer: no thinking block.
p.buffer.Reset()
p.buffer.WriteString(trimmed)
p.state = LFM2CollectingContent
return events, true
case LFM2CollectingThinking:
// Strip opening <think> tag if present
if strings.HasPrefix(bufStr, lfm2ThinkingOpenTag) {

View File

@@ -25,14 +25,36 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_content",
input: "I need to think about this...</think>The answer is 42.",
input: "<think>I need to think about this...</think>The answer is 42.",
expectedThinking: "I need to think about this...",
expectedContent: "The answer is 42.",
hasThinking: true,
},
{
name: "direct_answer_with_thinking_enabled",
input: "The answer is 42.",
expectedContent: "The answer is 42.",
hasThinking: true,
},
{
name: "direct_answer_with_tool_call",
input: "I'll check the weather.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
expectedContent: "I'll check the weather.",
expectedCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{
"location": "Paris",
}),
},
},
},
hasThinking: true,
},
{
name: "thinking_with_newlines",
input: "Let me think:\n- Point 1\n- Point 2</think>\n\nHere's my answer.",
input: "<think>Let me think:\n- Point 1\n- Point 2</think>\n\nHere's my answer.",
expectedThinking: "Let me think:\n- Point 1\n- Point 2",
expectedContent: "Here's my answer.",
hasThinking: true,
@@ -98,7 +120,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_with_tool_call",
input: "Let me check the weather...</think>I'll get that for you.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
input: "<think>Let me check the weather...</think>I'll get that for you.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
expectedThinking: "Let me check the weather...",
expectedContent: "I'll get that for you.",
expectedCalls: []api.ToolCall{
@@ -121,7 +143,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "only_thinking",
input: "Just thinking content</think>",
input: "<think>Just thinking content</think>",
expectedThinking: "Just thinking content",
expectedContent: "",
hasThinking: true,
@@ -140,7 +162,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_with_unicode",
input: "我在思考这个问题...</think>答案是42。",
input: "<think>我在思考这个问题...</think>答案是42。",
expectedThinking: "我在思考这个问题...",
expectedContent: "答案是42。",
hasThinking: true,
@@ -164,7 +186,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_with_special_chars",
input: "Let me calculate: 2+2=4 & 3*3=9...</think>The results are correct!",
input: "<think>Let me calculate: 2+2=4 & 3*3=9...</think>The results are correct!",
expectedThinking: "Let me calculate: 2+2=4 & 3*3=9...",
expectedContent: "The results are correct!",
hasThinking: true,
@@ -228,7 +250,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_then_python_tool_call",
input: "I should check the status...</think>Let me look that up.<|tool_call_start|>[get_status(id=\"123\")]<|tool_call_end|>",
input: "<think>I should check the status...</think>Let me look that up.<|tool_call_start|>[get_status(id=\"123\")]<|tool_call_end|>",
expectedThinking: "I should check the status...",
expectedContent: "Let me look that up.",
expectedCalls: []api.ToolCall{
@@ -291,7 +313,7 @@ func TestLFM2Parser(t *testing.T) {
},
{
name: "thinking_directly_to_tool_call",
input: "Let me run this command...</think><|tool_call_start|>[bash(command='ls')]<|tool_call_end|>",
input: "<think>Let me run this command...</think><|tool_call_start|>[bash(command='ls')]<|tool_call_end|>",
expectedThinking: "Let me run this command...",
expectedContent: "",
expectedCalls: []api.ToolCall{
@@ -350,11 +372,17 @@ func TestLFM2Parser_Streaming(t *testing.T) {
},
{
name: "streaming_thinking",
chunks: []string{"I need to ", "think about this", "...</think>", "The answer is 42."},
chunks: []string{"<think>", "I need to ", "think about this", "...</think>", "The answer is 42."},
expectedThinking: "I need to think about this...",
expectedContent: "The answer is 42.",
hasThinking: true,
},
{
name: "streaming_direct_answer",
chunks: []string{"The answer ", "is ", "42."},
expectedContent: "The answer is 42.",
hasThinking: true,
},
{
name: "streaming_tool_call",
chunks: []string{"I'll check weather.", "<|tool_call_start|>", "[get_weather(", "location=\"Paris\")]", "<|tool_call_end|>"},
@@ -373,7 +401,7 @@ func TestLFM2Parser_Streaming(t *testing.T) {
},
{
name: "streaming_thinking_with_partial_tag",
chunks: []string{"Thinking about this", "...</", "think>", "Done thinking."},
chunks: []string{"<think>", "Thinking about this", "...</", "think>", "Done thinking."},
expectedThinking: "Thinking about this...",
expectedContent: "Done thinking.",
hasThinking: true,
@@ -401,6 +429,22 @@ func TestLFM2Parser_Streaming(t *testing.T) {
},
hasThinking: false,
},
{
// Opening <think> tag split across chunks must still be recognized.
name: "streaming_thinking_split_open_tag",
chunks: []string{"<th", "ink>", "reasoning", "</think>", "answer"},
expectedThinking: "reasoning",
expectedContent: "answer",
hasThinking: true,
},
{
// A direct answer starting with '<' is ambiguous until enough arrives
// to rule out an opening <think> tag.
name: "streaming_direct_answer_starting_with_angle",
chunks: []string{"<", "html> is a tag"},
expectedContent: "<html> is a tag",
hasThinking: true,
},
{
// Test that leading whitespace after <think> is trimmed even when in separate chunks
name: "streaming_thinking_whitespace_after_tag",
@@ -506,9 +550,9 @@ func TestLFM2Parser_Init(t *testing.T) {
t.Errorf("Init() returned tools mismatch (-want +got):\n%s", diff)
}
// Test initial state is set to thinking when enabled
if parser.state != LFM2CollectingThinking {
t.Errorf("Expected initial state to be LFM2CollectingThinking, got %v", parser.state)
// Test initial state looks for a leading <think> tag when thinking is enabled
if parser.state != LFM2LookingForThinking {
t.Errorf("Expected initial state to be LFM2LookingForThinking, got %v", parser.state)
}
}
@@ -1071,18 +1115,24 @@ func TestLFM2Parser_EdgeCases(t *testing.T) {
}{
{
name: "multiple_think_close_tags",
input: "First thought</think>Second thought</think>Final content",
input: "<think>First thought</think>Second thought</think>Final content",
expectedThinking: "First thought",
expectedContent: "Second thought</think>Final content",
hasThinking: true,
},
{
name: "empty_thinking_content",
input: "</think>Just content",
name: "empty_thinking_block",
input: "<think></think>Just content",
expectedThinking: "",
expectedContent: "Just content",
hasThinking: true,
},
{
name: "direct_answer_with_leading_whitespace",
input: " \n Hello there",
expectedContent: "Hello there",
hasThinking: true,
},
{
name: "thinking_disabled_with_think_tags",
input: "Some content</think>More content",

View File

@@ -17,6 +17,8 @@ type LFM2Renderer struct {
const lfm2BOSToken = "<|startoftext|>"
const (
lfm2ThinkingOpenTag = "<think>"
lfm2ThinkingCloseTag = "</think>"
lfm2ToolListStartTag = "<|tool_list_start|>"
lfm2ToolListEndTag = "<|tool_list_end|>"
lfm2ToolCallStartTag = "<|tool_call_start|>"
@@ -285,11 +287,6 @@ func (r *LFM2Renderer) Render(messages []api.Message, tools []api.Tool, thinkVal
content := r.renderMessageContent(message, imageOffset)
imageOffset += len(message.Images)
if message.Role == "assistant" && !keepPastThinking && i != lastAssistantIndex {
if idx := strings.LastIndex(content, "</think>"); idx >= 0 {
content = strings.TrimSpace(content[idx+len("</think>"):])
}
}
if message.Role == "assistant" && len(message.ToolCalls) > 0 && !strings.Contains(content, lfm2ToolCallStartTag) {
if strings.TrimSpace(content) == "" {
content = lfm2RenderToolCalls(message.ToolCalls) + content
@@ -297,6 +294,23 @@ func (r *LFM2Renderer) Render(messages []api.Message, tools []api.Tool, thinkVal
content = lfm2RenderToolCalls(message.ToolCalls) + "\n" + content
}
}
// Reconstruct the inline <think>...</think> block from the separate
// Thinking field so reasoning turns round-trip in the model's own format:
// thinking precedes any tool calls and content. A direct answer carries no
// Thinking, so nothing is added. Only the thinking variant emits these tags;
// the non-thinking renderer must never send them, including for a trailing
// assistant prefill (which is exempt from the stripping below).
if r.IsThinking && message.Role == "assistant" && message.Thinking != "" && !strings.Contains(content, lfm2ThinkingCloseTag) {
content = lfm2ThinkingOpenTag + message.Thinking + lfm2ThinkingCloseTag + content
}
// Drop reasoning from earlier assistant turns unless thinking is kept; the
// <think>...</think> block is a clean prefix, so everything after the close
// tag (tool calls and content) is preserved.
if message.Role == "assistant" && !keepPastThinking && i != lastAssistantIndex {
if idx := strings.LastIndex(content, lfm2ThinkingCloseTag); idx >= 0 {
content = strings.TrimSpace(content[idx+len(lfm2ThinkingCloseTag):])
}
}
if message.Role == "tool" && !strings.Contains(content, lfm2ToolResponseStartTag) {
content = lfm2ToolResponseStartTag + content + lfm2ToolResponseEndTag
}

View File

@@ -159,6 +159,78 @@ func TestLFM2Renderer_ChatTemplateParity(t *testing.T) {
thinkValue: &api.ThinkValue{Value: true},
expected: "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\n<think>reason1</think>A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
},
{
name: "thinking_field_reconstructed_when_enabled",
renderer: &LFM2Renderer{IsThinking: true},
messages: []api.Message{
{Role: "user", Content: "Q1"},
{Role: "assistant", Content: "A1", Thinking: "reason1"},
{Role: "user", Content: "Q2"},
{Role: "assistant", Content: "A2", Thinking: "reason2"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\n<think>reason1</think>A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
},
{
name: "thinking_field_stripped_for_non_last_when_disabled",
renderer: &LFM2Renderer{IsThinking: true},
messages: []api.Message{
{Role: "user", Content: "Q1"},
{Role: "assistant", Content: "A1", Thinking: "reason1"},
{Role: "user", Content: "Q2"},
{Role: "assistant", Content: "A2", Thinking: "reason2"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\nA1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
},
{
name: "thinking_precedes_tool_calls",
renderer: &LFM2Renderer{IsThinking: true},
messages: []api.Message{
{Role: "user", Content: "Weather?"},
{
Role: "assistant",
Content: "",
Thinking: "Let me check the weather.",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{
"location": "Paris",
}),
},
},
},
},
{Role: "tool", Content: "22C"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|startoftext|><|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think>Let me check the weather.</think><|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|><|im_end|>\n<|im_start|>tool\n<|tool_response_start|>22C<|tool_response_end|><|im_end|>\n<|im_start|>assistant\n",
},
{
name: "direct_answer_history_has_no_think_tags",
renderer: &LFM2Renderer{IsThinking: true},
messages: []api.Message{
{Role: "user", Content: "Q1"},
{Role: "assistant", Content: "A1"},
{Role: "user", Content: "Q2"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\nA1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n",
},
{
// The non-thinking variant must never emit <think> tags, even for a
// trailing assistant prefill (which is exempt from thinking stripping).
name: "non_thinking_renderer_drops_thinking_metadata_on_prefill",
renderer: &LFM2Renderer{IsThinking: false},
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Hello", Thinking: "some reasoning"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: "<|startoftext|><|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello",
},
{
name: "arbitrary_roles_are_rendered_verbatim",
renderer: &LFM2Renderer{IsThinking: false},