luxferre/gr2gw
1
0
Fork 0
Code Actions Projects Activity

Sanitize tool call markers and enclosing code fences in completions

Browse Source
This commit is contained in:
Luxferre
2026-09-07 15:03:44 +03:00
parent e505f44f4a
commit cc2fa234ac
2 changed files with 404 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gr2gw.go
+199 -80
View File
@@ -476,18 +476,18 @@ func BuildToolInstruction(tools []Tool, toolChoice interface{}) string {
sampleFnName = "function_name"
}
directive := "If a tool is relevant, emit the tool call XML. If no tools are relevant, answer the query directly."
directive := "If a tool is relevant, emit the tool call in <tool_call> tags with JSON content. If no tools are relevant, answer the query directly."
if tcStr, ok := toolChoice.(string); ok && tcStr == "required" {
directive = "You MUST call one of the available tools for this query and emit the tool call XML."
directive = "You MUST call one of the available tools for this query and emit the tool call in <tool_call> tags with JSON content."
} else if tcMap, ok := toolChoice.(map[string]interface{}); ok {
if fnMap, ok := tcMap["function"].(map[string]interface{}); ok {
if fnName, ok := fnMap["name"].(string); ok && fnName != "" {
directive = fmt.Sprintf("You MUST call the %s tool for this query and emit the tool call XML.", fnName)
directive = fmt.Sprintf("You MUST call the %s tool for this query and emit the tool call in <tool_call> tags with JSON content.", fnName)
}
}
}
return fmt.Sprintf(`You are an API router and assistant. Convert the user query into the appropriate tool call XML using the available tools.
return fmt.Sprintf(`You are an API router and assistant. Convert the user query into the appropriate tool call using the available tools.
Available Tools:
%s
@@ -953,6 +953,49 @@ func parseXMLToolCall(block string) ([]ToolCall, bool) {
return nil, false
}
func scrubToolMarkers(text string) string {
s := text
for _, pair := range ToolTagPairs {
s = strings.ReplaceAll(s, pair.Start, "")
s = strings.ReplaceAll(s, pair.End, "")
}
reTags := regexp.MustCompile(`(?s)<(?:name|function|action|call)>[^<]*</(?:name|function|action|call)>`)
s = reTags.ReplaceAllString(s, "")
reArgTags := regexp.MustCompile(`(?s)<(?:arguments|parameters|args|input)>[\s\S]*?</(?:arguments|parameters|args|input)>`)
s = reArgTags.ReplaceAllString(s, "")
extraTags := []string{
"<name>", "</name>",
"<arguments>", "</arguments>",
"<parameters>", "</parameters>",
"<argument>", "</argument>",
"<parameter>", "</parameter>",
}
for _, tag := range extraTags {
s = strings.ReplaceAll(s, tag, "")
}
reEmptyFences := regexp.MustCompile("(?s)```(?:xml|json|ya?ml)?\\s*```")
s = reEmptyFences.ReplaceAllString(s, "")
trimmed := strings.TrimSpace(s)
if trimmed == "```xml" || trimmed == "```json" || trimmed == "```" {
return ""
}
if strings.Trim(trimmed, "`\r\n\t ") == "" {
return ""
}
if strings.Count(s, "```") == 1 {
reLeadingFence := regexp.MustCompile("^\\s*```(?:xml|json|ya?ml)?\\s*\\n?")
reTrailingFence := regexp.MustCompile("\\n?\\s*```(?:xml|json|ya?ml)?\\s*$")
s = reLeadingFence.ReplaceAllString(s, "")
s = reTrailingFence.ReplaceAllString(s, "")
}
return strings.TrimSpace(s)
}
func ExtractToolCallBlocks(content string) (blocks []string, remaining string) {
remaining = content
@@ -981,7 +1024,16 @@ func ExtractToolCallBlocks(content string) (blocks []string, remaining string) {
if eIdx != -1 && (nextSIdx == -1 || eIdx < nextSIdx) {
blockEndPos := eIdx + len(pair.End)
blockText = remaining[sIdx:blockEndPos]
remaining = strings.TrimSpace(remaining[:sIdx] + remaining[blockEndPos:])
before := remaining[:sIdx]
after := remaining[blockEndPos:]
reFenceOpen := regexp.MustCompile("(?s)\\n?\\s*```(?:xml|json|ya?ml)?\\s*$")
reFenceClose := regexp.MustCompile("^\\s*```\\s*\\n?")
if reFenceOpen.MatchString(before) && reFenceClose.MatchString(after) {
before = reFenceOpen.ReplaceAllString(before, "")
after = reFenceClose.ReplaceAllString(after, "")
}
remaining = strings.TrimSpace(before + after)
} else if nextSIdx != -1 {
blockEndPos := nextSIdx
blockText = remaining[sIdx:blockEndPos]
@@ -1009,12 +1061,7 @@ func DetectToolCalls(content string) ([]ToolCall, string, bool) {
}
if len(calls) > 0 {
remaining = strings.TrimSpace(remaining)
if remaining == "```xml" || remaining == "```" || remaining == "```json" {
remaining = ""
}
reEmptyFence := regexp.MustCompile("(?m)^```(?:xml|json)?\\s*\\n?\\s*```\\s*$")
remaining = strings.TrimSpace(reEmptyFence.ReplaceAllString(remaining, ""))
remaining = scrubToolMarkers(remaining)
return calls, remaining, true
}
@@ -1025,6 +1072,53 @@ func DetectToolCalls(content string) ([]ToolCall, string, bool) {
return nil, content, false
}
func finalizeOutput(frame GradioOutputFrame) (finalContent interface{}, reasoning string, toolCalls []ToolCall, finishReason string) {
cleanText := frame.Content
reasoning = frame.Reasoning
toolCalls = frame.ToolCalls
hasTools := len(toolCalls) > 0
if reasoning == "" {
cleanText, reasoning = ExtractThinking(cleanText)
}
if !hasTools {
toolCalls, cleanText, hasTools = DetectToolCalls(cleanText)
} else {
if extraCalls, extraClean, extraHas := DetectToolCalls(cleanText); extraHas && len(extraCalls) > 0 {
for _, ec := range extraCalls {
duplicate := false
for _, tc := range toolCalls {
if tc.Function.Name == ec.Function.Name && tc.Function.Arguments == ec.Function.Arguments {
duplicate = true
break
}
}
if !duplicate {
toolCalls = append(toolCalls, ec)
}
}
cleanText = extraClean
} else {
cleanText = scrubToolMarkers(cleanText)
}
}
finishReason = "stop"
finalContent = cleanText
if hasTools && len(toolCalls) > 0 {
finishReason = "tool_calls"
cleanText = scrubToolMarkers(cleanText)
if strings.TrimSpace(cleanText) == "" {
finalContent = nil
} else {
finalContent = strings.TrimSpace(cleanText)
}
}
return finalContent, reasoning, toolCalls, finishReason
}
func ExtractThinking(content string) (string, string) {
if strings.Contains(content, "<think>") && strings.Contains(content, "</think>") {
start := strings.Index(content, "<think>")
@@ -1462,6 +1556,13 @@ func (f *StreamThinkingFilter) Flush(onContent func(string), onReasoning func(st
// Stateful tool call tag filter for streaming
// ---------------------------------------------------------------------------
var potentialFencePrefixes = []string{
"```",
"```x", "```xm", "```xml", "```xml\r", "```xml\n", "```xml\r\n",
"```j", "```js", "```jso", "```json", "```json\r", "```json\n", "```json\r\n",
"```\r", "```\n", "```\r\n",
}
type StreamToolCallFilter struct {
inToolCall bool
buf string
@@ -1476,11 +1577,43 @@ func NewStreamToolCallFilter() *StreamToolCallFilter {
return &StreamToolCallFilter{}
}
func (f *StreamToolCallFilter) cleanTrailingAfterToolCall() {
for {
stripped := false
for _, pair := range ToolTagPairs {
trimmed := strings.TrimLeft(f.buf, " \t\r\n")
if strings.HasPrefix(trimmed, pair.End) {
idx := strings.Index(f.buf, pair.End)
f.buf = f.buf[idx+len(pair.End):]
stripped = true
break
}
}
if !stripped {
break
}
}
reCloseFence := regexp.MustCompile("^\\s*```\\s*\\n?")
f.buf = reCloseFence.ReplaceAllString(f.buf, "")
if f.buf == "\n" || f.buf == "\r\n" || f.buf == "\n\n" {
f.buf = ""
}
}
func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onToolCall func(ToolCall)) {
f.buf += chunk
for len(f.buf) > 0 {
if !f.inToolCall {
if f.emittedCall {
f.cleanTrailingAfterToolCall()
if len(f.buf) == 0 {
break
}
}
earliestIdx := -1
var matchedPair ToolTagPair
@@ -1495,6 +1628,15 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool
if earliestIdx != -1 {
before := f.buf[:earliestIdx]
reTrailingFence := regexp.MustCompile("(?s)\\n?\\s*```(?:xml|json|ya?ml)?\\s*$")
if reTrailingFence.MatchString(before) {
before = reTrailingFence.ReplaceAllString(before, "")
}
if strings.TrimSpace(before) == "" {
before = ""
} else {
before = strings.TrimRight(before, "\r\n")
}
if before != "" {
onContent(before)
}
@@ -1502,18 +1644,34 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool
f.activePair = matchedPair
f.activeEndTag = matchedPair.End
f.buf = f.buf[earliestIdx+len(matchedPair.Start):]
} else if matchLen := hasPrefixOf(f.buf, toolStartPrefixes); matchLen > 0 {
safe := f.buf[:len(f.buf)-matchLen]
if safe != "" {
} else {
matchLen := hasPrefixOf(f.buf, toolStartPrefixes)
fenceLen := hasPrefixOf(f.buf, potentialFencePrefixes)
holdLen := matchLen
if fenceLen > holdLen {
holdLen = fenceLen
}
if holdLen > 0 {
safe := f.buf[:len(f.buf)-holdLen]
if matchLen > 0 {
reTrailingFence := regexp.MustCompile("(?s)\\n?\\s*```(?:xml|json|ya?ml)?\\s*$")
safe = reTrailingFence.ReplaceAllString(safe, "")
safe = strings.TrimRight(safe, "\r\n")
}
if strings.TrimSpace(safe) != "" {
onContent(safe)
}
f.buf = f.buf[len(f.buf)-matchLen:]
f.buf = f.buf[len(f.buf)-holdLen:]
break
} else if len(f.buf) < 16 && strings.TrimSpace(f.buf) == "" {
break
} else {
onContent(f.buf)
f.buf = ""
break
}
}
} else {
if idx := strings.Index(f.buf, f.activeEndTag); idx != -1 {
f.toolCallBuf += f.buf[:idx]
@@ -1540,6 +1698,10 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool
onContent(f.activePair.Start + f.toolCallBuf + f.activePair.End)
}
f.toolCallBuf = ""
if f.emittedCall {
f.cleanTrailingAfterToolCall()
}
} else if matchLen := hasSuffixPrefixOf(f.buf, f.activeEndTag); matchLen > 0 {
safe := f.buf[:len(f.buf)-matchLen]
f.toolCallBuf += safe
@@ -1578,7 +1740,14 @@ func (f *StreamToolCallFilter) Flush(onContent func(string), onToolCall func(Too
f.toolCallBuf = ""
}
if len(f.buf) > 0 {
if f.emittedCall {
clean := scrubToolMarkers(f.buf)
if strings.TrimSpace(clean) != "" {
onContent(strings.TrimSpace(clean))
}
} else {
onContent(f.buf)
}
f.buf = ""
}
}
@@ -3518,26 +3687,7 @@ func (g *GradioGateway) executePredictCompletion(w http.ResponseWriter, r *http.
return fmt.Errorf("upstream Gradio space returned empty or unparseable response")
}
cleanText := frame.Content
reasoning := frame.Reasoning
toolCalls := frame.ToolCalls
hasTools := len(toolCalls) > 0
if reasoning == "" {
cleanText, reasoning = ExtractThinking(cleanText)
}
if !hasTools {
toolCalls, cleanText, hasTools = DetectToolCalls(cleanText)
}
finishReason := "stop"
var finalContent interface{} = cleanText
if hasTools && len(toolCalls) > 0 {
finishReason = "tool_calls"
if strings.TrimSpace(cleanText) == "" {
finalContent = nil
}
}
finalContent, reasoning, toolCalls, finishReason := finalizeOutput(frame)
if !req.Stream {
WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{
@@ -3554,14 +3704,17 @@ func (g *GradioGateway) executePredictCompletion(w http.ResponseWriter, r *http.
if reasoning != "" {
streamer.Reasoning(reasoning)
}
if hasTools && len(toolCalls) > 0 {
if len(toolCalls) > 0 {
for i, tc := range toolCalls {
iCopy := i
tc.Index = &iCopy
streamer.ToolCallDelta(tc)
}
} else if cleanText != "" {
streamer.Content(cleanText)
}
if finalContent != nil {
if s, ok := finalContent.(string); ok && s != "" {
streamer.Content(s)
}
}
streamer.Finish(finishReason)
streamer.Done()
@@ -3715,7 +3868,9 @@ func (g *GradioGateway) executeQueueCompletion(w http.ResponseWriter, r *http.Re
if frame.IsDelta {
latestFrame.Content += frame.Content
} else {
if frame.Content != "" || latestFrame.Content == "" {
if len(frame.ToolCalls) > 0 {
latestFrame.Content = frame.Content
} else if frame.Content != "" || latestFrame.Content == "" {
latestFrame.Content = frame.Content
}
}
@@ -3739,26 +3894,7 @@ func (g *GradioGateway) executeQueueCompletion(w http.ResponseWriter, r *http.Re
return fmt.Errorf("upstream Gradio space returned empty or unparseable response")
}
cleanText := latestFrame.Content
reasoning := latestFrame.Reasoning
toolCalls := latestFrame.ToolCalls
hasTools := len(toolCalls) > 0
if reasoning == "" {
cleanText, reasoning = ExtractThinking(cleanText)
}
if !hasTools {
toolCalls, cleanText, hasTools = DetectToolCalls(cleanText)
}
finishReason := "stop"
var finalContent interface{} = cleanText
if hasTools && len(toolCalls) > 0 {
finishReason = "tool_calls"
if strings.TrimSpace(cleanText) == "" {
finalContent = nil
}
}
finalContent, reasoning, toolCalls, finishReason := finalizeOutput(latestFrame)
WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{
Content: finalContent,
@@ -4228,7 +4364,9 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
return fmt.Errorf("upstream Gradio error: %s", errMsg)
}
if frame := ParseGradioStreamOutput(dataStr); frame.OK {
if frame.Content != "" || latestFrame.Content == "" {
if len(frame.ToolCalls) > 0 {
latestFrame.Content = frame.Content
} else if frame.Content != "" || latestFrame.Content == "" {
latestFrame.Content = frame.Content
}
if frame.Reasoning != "" || latestFrame.Reasoning == "" {
@@ -4250,26 +4388,7 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
}
cleanText := latestFrame.Content
reasoning := latestFrame.Reasoning
toolCalls := latestFrame.ToolCalls
hasTools := len(toolCalls) > 0
if reasoning == "" {
cleanText, reasoning = ExtractThinking(cleanText)
}
if !hasTools {
toolCalls, cleanText, hasTools = DetectToolCalls(cleanText)
}
finishReason := "stop"
var finalContent interface{} = cleanText
if hasTools && len(toolCalls) > 0 {
finishReason = "tool_calls"
if strings.TrimSpace(cleanText) == "" {
finalContent = nil
}
}
finalContent, reasoning, toolCalls, finishReason := finalizeOutput(latestFrame)
WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{
Content: finalContent,
gr2gw_test.go
+198
View File
@@ -2526,8 +2526,206 @@ func TestGradioCallFallbackToQueue(t *testing.T) {
}
}
func TestScrubToolMarkers(t *testing.T) {
cases := []struct {
input string
expected string
}{
{"```xml\n\n```", ""},
{"```xml\n<tool_call></tool_call>\n```", ""},
{"</tool_call>", ""},
{"<tool_call>", ""},
{"[TOOL_CALLS][/TOOL_CALLS]", ""},
{"<name>foo</name><arguments></arguments>", ""},
{"```\n```", ""},
{"```xml\n```", ""},
{"```json\n```", ""},
{"Some text\n```xml\n\n```", "Some text"},
{"Some text</tool_call>", "Some text"},
}
for _, c := range cases {
got := scrubToolMarkers(c.input)
if got != c.expected {
t.Errorf("scrubToolMarkers(%q) = %q; expected %q", c.input, got, c.expected)
}
}
}
func TestStreamToolCallFilterNoLeakFencesOrMarkers(t *testing.T) {
// 1. Tool call fully enclosed in code fences in a single chunk
{
filter := NewStreamToolCallFilter()
var contentChunks []string
var calls []ToolCall
filter.Feed("```xml\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n</tool_call>\n```",
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
filter.Flush(
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
t.Fatalf("expected 1 tool call 'get_weather', got: %+v", calls)
}
if len(contentChunks) > 0 {
t.Fatalf("expected zero content chunks leaked, got: %v", contentChunks)
}
}
// 2. Chunks split across fence and tags
{
filter := NewStreamToolCallFilter()
var contentChunks []string
var calls []ToolCall
chunks := []string{
"```",
"xml\n",
"<tool_",
"call>\n{\"name\": \"calc\", \"arguments\": {\"expr\": \"2+2\"}}\n</tool_",
"call>",
"\n```",
}
for _, c := range chunks {
filter.Feed(c,
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
}
filter.Flush(
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
if len(calls) != 1 || calls[0].Function.Name != "calc" {
t.Fatalf("expected 1 tool call 'calc', got: %+v", calls)
}
if len(contentChunks) > 0 {
t.Fatalf("expected zero content chunks leaked, got: %v", contentChunks)
}
}
// 3. Commentary before fenced tool call
{
filter := NewStreamToolCallFilter()
var contentChunks []string
var calls []ToolCall
chunks := []string{
"I'll look up the weather for you.",
"\n```xml\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"London\"}}\n</tool_call>\n```",
}
for _, c := range chunks {
filter.Feed(c,
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
}
filter.Flush(
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
t.Fatalf("expected 1 tool call 'get_weather', got: %+v", calls)
}
fullContent := strings.Join(contentChunks, "")
if fullContent != "I'll look up the weather for you." {
t.Fatalf("expected only preamble commentary, got: %q", fullContent)
}
}
// 4. Duplicate close tags and closing fences
{
filter := NewStreamToolCallFilter()
var contentChunks []string
var calls []ToolCall
chunks := []string{
"<tool_call>\n{\"name\": \"search\", \"arguments\": {\"q\": \"rust\"}}\n</tool_call>",
"</tool_call>\n```\n",
}
for _, c := range chunks {
filter.Feed(c,
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
}
filter.Flush(
func(s string) { contentChunks = append(contentChunks, s) },
func(tc ToolCall) { calls = append(calls, tc) },
)
if len(calls) != 1 || calls[0].Function.Name != "search" {
t.Fatalf("expected 1 tool call 'search', got: %+v", calls)
}
if len(contentChunks) > 0 {
t.Fatalf("expected zero content chunks leaked, got: %v", contentChunks)
}
}
}
func TestFinalizeOutputToolCallSanitization(t *testing.T) {
// 1. Frame with tool call in Content and tool_calls populated
frame1 := GradioOutputFrame{
Content: "<tool_call>\n",
ToolCalls: []ToolCall{
{Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Berlin"}`}},
},
OK: true,
}
content1, _, tcs1, finish1 := finalizeOutput(frame1)
if finish1 != "tool_calls" {
t.Errorf("expected finish_reason 'tool_calls', got %s", finish1)
}
if content1 != nil {
t.Errorf("expected content to be nil, got: %v", content1)
}
if len(tcs1) != 1 || tcs1[0].Function.Name != "get_weather" {
t.Errorf("unexpected tool calls: %+v", tcs1)
}
// 2. Frame with tool call inside code fence in Content
frame2 := GradioOutputFrame{
Content: "```xml\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Madrid\"}}\n</tool_call>\n```",
OK: true,
}
content2, _, tcs2, finish2 := finalizeOutput(frame2)
if finish2 != "tool_calls" {
t.Errorf("expected finish_reason 'tool_calls', got %s", finish2)
}
if content2 != nil {
t.Errorf("expected content to be nil, got: %v", content2)
}
if len(tcs2) != 1 || tcs2[0].Function.Name != "get_weather" {
t.Errorf("unexpected tool calls: %+v", tcs2)
}
// 3. Frame with commentary and tool call
frame3 := GradioOutputFrame{
Content: "Checking the weather for Tokyo.\n```xml\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>\n```",
OK: true,
}
content3, _, tcs3, finish3 := finalizeOutput(frame3)
if finish3 != "tool_calls" {
t.Errorf("expected finish_reason 'tool_calls', got %s", finish3)
}
if content3 != "Checking the weather for Tokyo." {
t.Errorf("expected commentary preserved without fences, got: %v", content3)
}
if len(tcs3) != 1 || tcs3[0].Function.Name != "get_weather" {
t.Errorf("unexpected tool calls: %+v", tcs3)
}
// 4. Frame with duplicate end tags and unclosed fences
frame4 := GradioOutputFrame{
Content: "<tool_call>\n<tool_call>\n{\"name\": \"search\", \"arguments\": {}}\n</tool_call>\n</tool_call>\n```",
OK: true,
}
content4, _, tcs4, finish4 := finalizeOutput(frame4)
if finish4 != "tool_calls" {
t.Errorf("expected finish_reason 'tool_calls', got %s", finish4)
}
if content4 != nil {
t.Errorf("expected content to be nil, got: %v", content4)
}
if len(tcs4) != 1 || tcs4[0].Function.Name != "search" {
t.Errorf("unexpected tool calls: %+v", tcs4)
}
}