From 993acc7504e7527d0764a1309dcb1d9036a0e389 Mon Sep 17 00:00:00 2001 From: Jeffrey Morgan Date: Sun, 14 Jun 2026 20:37:08 -0700 Subject: [PATCH] model: update lfm2 parser/renderer for optional thinking (#16359) --- model/parsers/lfm2.go | 47 +++++++++++++++++++-- model/parsers/lfm2_test.go | 82 +++++++++++++++++++++++++++++------- model/renderers/lfm2.go | 24 ++++++++--- model/renderers/lfm2_test.go | 72 +++++++++++++++++++++++++++++++ 4 files changed, 201 insertions(+), 24 deletions(-) diff --git a/model/parsers/lfm2.go b/model/parsers/lfm2.go index 9dc8ec02e..88b1b9f16 100644 --- a/model/parsers/lfm2.go +++ b/model/parsers/lfm2.go @@ -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 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 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 "" 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 ) 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 "" 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 tag if present if strings.HasPrefix(bufStr, lfm2ThinkingOpenTag) { diff --git a/model/parsers/lfm2_test.go b/model/parsers/lfm2_test.go index 416fcb8ce..9bc725056 100644 --- a/model/parsers/lfm2_test.go +++ b/model/parsers/lfm2_test.go @@ -25,14 +25,36 @@ func TestLFM2Parser(t *testing.T) { }, { name: "thinking_content", - input: "I need to think about this...The answer is 42.", + input: "I need to think about this...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\n\nHere's my answer.", + input: "Let me think:\n- Point 1\n- Point 2\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...I'll get that for you.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>", + input: "Let me check the weather...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", + input: "Just thinking content", expectedThinking: "Just thinking content", expectedContent: "", hasThinking: true, @@ -140,7 +162,7 @@ func TestLFM2Parser(t *testing.T) { }, { name: "thinking_with_unicode", - input: "我在思考这个问题...答案是42。", + input: "我在思考这个问题...答案是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...The results are correct!", + input: "Let me calculate: 2+2=4 & 3*3=9...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...Let me look that up.<|tool_call_start|>[get_status(id=\"123\")]<|tool_call_end|>", + input: "I should check the status...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...<|tool_call_start|>[bash(command='ls')]<|tool_call_end|>", + input: "Let me run this command...<|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", "...", "The answer is 42."}, + chunks: []string{"", "I need to ", "think about this", "...", "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", "...", "Done thinking."}, + chunks: []string{"", "Thinking about this", "...", "Done thinking."}, expectedThinking: "Thinking about this...", expectedContent: "Done thinking.", hasThinking: true, @@ -401,6 +429,22 @@ func TestLFM2Parser_Streaming(t *testing.T) { }, hasThinking: false, }, + { + // Opening tag split across chunks must still be recognized. + name: "streaming_thinking_split_open_tag", + chunks: []string{"", "reasoning", "", "answer"}, + expectedThinking: "reasoning", + expectedContent: "answer", + hasThinking: true, + }, + { + // A direct answer starting with '<' is ambiguous until enough arrives + // to rule out an opening tag. + name: "streaming_direct_answer_starting_with_angle", + chunks: []string{"<", "html> is a tag"}, + expectedContent: " is a tag", + hasThinking: true, + }, { // Test that leading whitespace after 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 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 thoughtSecond thoughtFinal content", + input: "First thoughtSecond thoughtFinal content", expectedThinking: "First thought", expectedContent: "Second thoughtFinal content", hasThinking: true, }, { - name: "empty_thinking_content", - input: "Just content", + name: "empty_thinking_block", + input: "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 contentMore content", diff --git a/model/renderers/lfm2.go b/model/renderers/lfm2.go index 8bfacf23b..1d6b7518c 100644 --- a/model/renderers/lfm2.go +++ b/model/renderers/lfm2.go @@ -17,6 +17,8 @@ type LFM2Renderer struct { const lfm2BOSToken = "<|startoftext|>" const ( + lfm2ThinkingOpenTag = "" + lfm2ThinkingCloseTag = "" 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, ""); idx >= 0 { - content = strings.TrimSpace(content[idx+len(""):]) - } - } 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 ... 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 + // ... 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 } diff --git a/model/renderers/lfm2_test.go b/model/renderers/lfm2_test.go index 6e27af9bf..bd3062256 100644 --- a/model/renderers/lfm2_test.go +++ b/model/renderers/lfm2_test.go @@ -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\nreason1A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\nreason2A2", }, + { + 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\nreason1A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\nreason2A2", + }, + { + 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\nreason2A2", + }, + { + 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\nLet me check the weather.<|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 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},