Sanitize tool call markers and enclosing code fences in completions
This commit is contained in:
Luxferre
+198
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user