Skip to content

Commit 5dcd343

Browse files
committed
chore(refactor): re-use iterative parser
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 4786cb6 commit 5dcd343

4 files changed

Lines changed: 78 additions & 57 deletions

File tree

pkg/functions/iterative_parser.go

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -621,25 +621,39 @@ func (p *ChatMsgParser) TryConsumeXMLToolCalls(format *XMLToolCallFormat) (bool,
621621
// Handle Functionary format (JSON parameters inside XML tags) - use regex parser
622622
if format.KeyStart == "" && format.ToolStart == "<function=" {
623623
// Fall back to regex-based parser for Functionary format
624-
results, err := parseFunctionaryFormat(p.input[p.pos:], format)
624+
sub := p.input[p.pos:]
625+
results, err := parseFunctionaryFormat(sub, format)
625626
if err != nil || len(results) == 0 {
626627
return false, nil
627628
}
628629
for _, result := range results {
629630
p.AddToolCall(result.Name, "", result.Arguments)
630631
}
632+
// Advance position past the last tool call end tag
633+
if last := strings.LastIndex(sub, format.ToolEnd); last >= 0 {
634+
p.pos += last + len(format.ToolEnd)
635+
}
631636
return true, nil
632637
}
633638

634639
// Handle JSON-like formats (Apriel-1.5, Xiaomi-MiMo) - use regex parser
635640
if format.ToolStart != "" && strings.Contains(format.ToolStart, "{\"name\"") {
636-
results, err := parseJSONLikeXMLFormat(p.input[p.pos:], format)
641+
sub := p.input[p.pos:]
642+
results, err := parseJSONLikeXMLFormat(sub, format)
637643
if err != nil || len(results) == 0 {
638644
return false, nil
639645
}
640646
for _, result := range results {
641647
p.AddToolCall(result.Name, "", result.Arguments)
642648
}
649+
// Advance position past the last scope/tool end tag
650+
endTag := format.ScopeEnd
651+
if endTag == "" {
652+
endTag = format.ToolEnd
653+
}
654+
if last := strings.LastIndex(sub, endTag); last >= 0 {
655+
p.pos += last + len(endTag)
656+
}
643657
return true, nil
644658
}
645659

pkg/functions/parse.go

Lines changed: 43 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -554,35 +554,64 @@ func getScopeOrToolStart(format *XMLToolCallFormat) string {
554554
return format.ToolStart
555555
}
556556

557+
// ParseResult holds tool calls and any non-tool-call content extracted by the parser.
558+
type ParseResult struct {
559+
ToolCalls []FuncCallResults
560+
Content string
561+
}
562+
557563
// tryParseXMLFromScopeStart finds the first occurrence of scopeStart (or format.ToolStart),
558-
// splits the input there, and parses only the suffix as XML tool calls. Returns (toolCalls, true)
559-
// if any tool calls were parsed, else (nil, false). This mimics llama.cpp's PEG order so that
564+
// splits the input there, and parses only the suffix as XML tool calls. Returns (result, true)
565+
// if any tool calls were parsed, else (empty, false). This mimics llama.cpp's PEG order so that
560566
// reasoning or content before the tool block does not cause "whitespace only before scope" to fail.
561-
func tryParseXMLFromScopeStart(s string, format *XMLToolCallFormat, isPartial bool) ([]FuncCallResults, bool) {
567+
func tryParseXMLFromScopeStart(s string, format *XMLToolCallFormat, isPartial bool) (ParseResult, bool) {
562568
if format == nil {
563-
return nil, false
569+
return ParseResult{}, false
564570
}
565571
scopeStart := getScopeOrToolStart(format)
566572
if scopeStart == "" {
567-
return nil, false
573+
return ParseResult{}, false
568574
}
569575
idx := strings.Index(s, scopeStart)
570576
if idx < 0 {
571-
return nil, false
577+
return ParseResult{}, false
572578
}
573579
toolCallsPart := s[idx:]
574580
parser := NewChatMsgParser(toolCallsPart, isPartial)
575581
success, err := parser.TryConsumeXMLToolCalls(format)
576582
if err != nil {
577583
if _, ok := err.(*ChatMsgPartialException); ok && isPartial {
578-
return parser.ToolCalls(), len(parser.ToolCalls()) > 0
584+
tc := parser.ToolCalls()
585+
if len(tc) > 0 {
586+
return ParseResult{ToolCalls: tc, Content: buildContent(s[:idx], parser)}, true
587+
}
579588
}
580-
return nil, false
589+
return ParseResult{}, false
581590
}
582591
if success && len(parser.ToolCalls()) > 0 {
583-
return parser.ToolCalls(), true
592+
return ParseResult{
593+
ToolCalls: parser.ToolCalls(),
594+
Content: buildContent(s[:idx], parser),
595+
}, true
596+
}
597+
return ParseResult{}, false
598+
}
599+
600+
// buildContent assembles the non-tool-call content from the text before the tool
601+
// block, any content tracked by the parser, and any unconsumed trailing text.
602+
func buildContent(before string, parser *ChatMsgParser) string {
603+
var parts []string
604+
if b := strings.TrimSpace(before); b != "" {
605+
parts = append(parts, b)
606+
}
607+
if pc := strings.TrimSpace(parser.Content()); pc != "" {
608+
parts = append(parts, pc)
609+
}
610+
remaining := parser.Input()[parser.Pos():]
611+
if t := strings.TrimSpace(remaining); t != "" {
612+
parts = append(parts, t)
584613
}
585-
return nil, false
614+
return strings.Join(parts, " ")
586615
}
587616

588617
// ParseXMLIterative parses XML tool calls using the iterative parser
@@ -592,15 +621,15 @@ func tryParseXMLFromScopeStart(s string, format *XMLToolCallFormat, isPartial bo
592621
func ParseXMLIterative(s string, format *XMLToolCallFormat, isPartial bool) ([]FuncCallResults, error) {
593622
// Try split-on-scope first so reasoning/content before tool block is skipped
594623
if format != nil {
595-
if results, ok := tryParseXMLFromScopeStart(s, format, isPartial); ok {
596-
return results, nil
624+
if pr, ok := tryParseXMLFromScopeStart(s, format, isPartial); ok {
625+
return pr.ToolCalls, nil
597626
}
598627
} else {
599628
formats := getAllXMLFormats()
600629
for _, fmtPreset := range formats {
601630
if fmtPreset.format != nil {
602-
if results, ok := tryParseXMLFromScopeStart(s, fmtPreset.format, isPartial); ok {
603-
return results, nil
631+
if pr, ok := tryParseXMLFromScopeStart(s, fmtPreset.format, isPartial); ok {
632+
return pr.ToolCalls, nil
604633
}
605634
}
606635
}

pkg/functions/parse_test.go

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2569,26 +2569,20 @@ def hello():
25692569
})
25702570

25712571
Context("StripToolCallMarkup", func() {
2572-
It("removes closed tool_call blocks", func() {
2573-
input := "Here is my answer <tool_call>{\"name\":\"search\"}</tool_call> and more text"
2572+
It("removes functionary-style function blocks and keeps surrounding text", func() {
2573+
input := `Text before <function=search>{"q":"cats"}</function> text after`
25742574
result := StripToolCallMarkup(input)
2575-
Expect(result).To(Equal("Here is my answer and more text"))
2575+
Expect(result).To(Equal("Text before text after"))
25762576
})
25772577

2578-
It("removes closed function blocks", func() {
2579-
input := "Text before <function=search>{\"q\":\"cats\"}</function> text after"
2578+
It("removes qwen3-coder-style tool_call blocks and keeps preceding text", func() {
2579+
input := `Here is my answer <tool_call><function=search><parameter=q>cats</parameter></function></tool_call>`
25802580
result := StripToolCallMarkup(input)
2581-
Expect(result).To(Equal("Text before text after"))
2582-
})
2583-
2584-
It("removes open-ended tool_call at end of string", func() {
2585-
input := "Some text <tool_call>{\"name\":\"search\"}"
2586-
result := StripToolCallMarkup(input)
2587-
Expect(result).To(Equal("Some text"))
2581+
Expect(result).To(Equal("Here is my answer"))
25882582
})
25892583

25902584
It("returns empty string when content is only tool calls", func() {
2591-
input := "<tool_call>{\"name\":\"search\",\"arguments\":{}}</tool_call>"
2585+
input := `<function=search>{"q":"cats"}</function>`
25922586
result := StripToolCallMarkup(input)
25932587
Expect(result).To(Equal(""))
25942588
})

pkg/functions/strip.go

Lines changed: 12 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,18 @@
11
package functions
22

3-
import (
4-
"regexp"
5-
"strings"
6-
)
3+
import "strings"
74

8-
// toolCallBlockRe matches closed <tool_call>...</tool_call> blocks
9-
var toolCallBlockRe = regexp.MustCompile(`(?s)<tool_call>.*?</tool_call>`)
10-
11-
// functionBlockRe matches closed <function=...>...</function> blocks
12-
var functionBlockRe = regexp.MustCompile(`(?s)<function=[^>]*>.*?</function>`)
13-
14-
// openToolCallRe matches an open-ended <tool_call> at the end of string (no closing tag)
15-
var openToolCallRe = regexp.MustCompile(`(?s)<tool_call>[^<]*$`)
16-
17-
// openFunctionRe matches an open-ended <function=...> at the end of string (no closing tag)
18-
var openFunctionRe = regexp.MustCompile(`(?s)<function=[^>]*>[^<]*$`)
19-
20-
// StripToolCallMarkup removes tool call markup blocks from content.
21-
// It removes closed <tool_call>...</tool_call> and <function=...>...</function> blocks,
22-
// plus open-ended ones at the end of the string.
23-
// Returns the remaining text, trimmed of whitespace.
5+
// StripToolCallMarkup extracts the non-tool-call content from a string
6+
// by reusing the iterative XML parser which already separates content
7+
// from tool calls. Returns the remaining text, trimmed.
248
func StripToolCallMarkup(content string) string {
25-
// Remove closed blocks
26-
content = toolCallBlockRe.ReplaceAllString(content, "")
27-
content = functionBlockRe.ReplaceAllString(content, "")
28-
29-
// Remove open-ended blocks at end of string
30-
content = openToolCallRe.ReplaceAllString(content, "")
31-
content = openFunctionRe.ReplaceAllString(content, "")
32-
9+
for _, fmtPreset := range getAllXMLFormats() {
10+
if fmtPreset.format == nil {
11+
continue
12+
}
13+
if pr, ok := tryParseXMLFromScopeStart(content, fmtPreset.format, false); ok && len(pr.ToolCalls) > 0 {
14+
return strings.TrimSpace(pr.Content)
15+
}
16+
}
3317
return strings.TrimSpace(content)
3418
}

0 commit comments

Comments
 (0)