diff --git a/gr2gw.go b/gr2gw.go index 67ed18e..7258b1c 100644 --- a/gr2gw.go +++ b/gr2gw.go @@ -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 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 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 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)>[^<]*`) + s = reTags.ReplaceAllString(s, "") + reArgTags := regexp.MustCompile(`(?s)<(?:arguments|parameters|args|input)>[\s\S]*?`) + s = reArgTags.ReplaceAllString(s, "") + + extraTags := []string{ + "", "", + "", "", + "", "", + "", "", + "", "", + } + 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, "") && strings.Contains(content, "") { start := strings.Index(content, "") @@ -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,17 +1644,33 @@ 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 != "" { - onContent(safe) - } - f.buf = f.buf[len(f.buf)-matchLen:] - break } else { - onContent(f.buf) - f.buf = "" - break + 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)-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 { @@ -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 { - onContent(f.buf) + 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, diff --git a/gr2gw_test.go b/gr2gw_test.go index 0a21dc8..b450001 100644 --- a/gr2gw_test.go +++ b/gr2gw_test.go @@ -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\n```", ""}, + {"", ""}, + {"", ""}, + {"[TOOL_CALLS][/TOOL_CALLS]", ""}, + {"foo", ""}, + {"```\n```", ""}, + {"```xml\n```", ""}, + {"```json\n```", ""}, + {"Some text\n```xml\n\n```", "Some text"}, + {"Some text", "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\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n\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", + "\n{\"name\": \"calc\", \"arguments\": {\"expr\": \"2+2\"}}\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 != "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\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"London\"}}\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 != "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{ + "\n{\"name\": \"search\", \"arguments\": {\"q\": \"rust\"}}\n", + "\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: "\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\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Madrid\"}}\n\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\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n\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: "\n\n{\"name\": \"search\", \"arguments\": {}}\n\n\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) + } +}