tools fix

gr2gw_test.go

@@ -149,8 +149,8 @@ func TestStreamToolCallFilter(t *testing.T) {
     t.Errorf("expected tool name 'search_web', got %q", toolCalls[0].Function.Name)
   }
   fullContent := strings.Join(contentParts, "")
-  if fullContent != "Searching now:  Done." {
-    t.Errorf("expected 'Searching now:  Done.', got %q", fullContent)
+  if fullContent != "Searching now: " {
+    t.Errorf("expected 'Searching now: ' (post-call text suppressed), got %q", fullContent)
   }
 }
 
@@ -639,8 +639,11 @@ Some postamble.`
   if strings.Contains(rem2, "function_call") {
     t.Errorf("expected tag stripped from remaining, got %q", rem2)
   }
-  if !strings.Contains(rem2, "Some preamble") || !strings.Contains(rem2, "Some postamble") {
-    t.Errorf("expected surrounding text preserved in remaining, got %q", rem2)
+  if !strings.Contains(rem2, "Some preamble") {
+    t.Errorf("expected preamble text preserved in remaining, got %q", rem2)
+  }
+  if strings.Contains(rem2, "Some postamble") {
+    t.Errorf("expected postamble text discarded from remaining on tool call, got %q", rem2)
   }
 
   // 3. [TOOL_CALLS] bracket syntax
@@ -723,8 +726,8 @@ func TestUniversalStreamToolCallFilterVariants(t *testing.T) {
   if strings.Contains(fullContent, "TOOL_CALLS") {
     t.Errorf("tag leaked into stream content: %q", fullContent)
   }
-  if fullContent != "Preamble text:  Completed." {
-    t.Errorf("unexpected streamed content: %q", fullContent)
+  if fullContent != "Preamble text: " {
+    t.Errorf("unexpected streamed content (expected post-call text suppressed): %q", fullContent)
   }
 }
 
@@ -3367,3 +3370,237 @@ func TestHunyuan3CallAndStreamingNoInterleaving(t *testing.T) {
   }
 }
 
+func TestDetectToolCallsDiscardsPostambleAndCitationDisclaimer(t *testing.T) {
+  raw := `<tool_call>
+{"name": "get_current_weather", "arguments": {"location": "Tokyo"}}
+</tool_call>
+
+*Web evidence was retrieved, but the response did not include a valid source citation.*`
+
+  tcs, rem, ok := DetectToolCalls(raw)
+  if !ok || len(tcs) != 1 {
+    t.Fatalf("expected 1 tool call, got %d (ok: %v)", len(tcs), ok)
+  }
+  if tcs[0].Function.Name != "get_current_weather" {
+    t.Errorf("expected function get_current_weather, got %q", tcs[0].Function.Name)
+  }
+  if rem != "" {
+    t.Errorf("expected empty remaining content, got %q", rem)
+  }
+
+  frame := GradioOutputFrame{Content: raw}
+  finalContent, reasoning, finalCalls, finishReason := finalizeOutput(frame)
+  if finishReason != "tool_calls" {
+    t.Errorf("expected finish_reason 'tool_calls', got %q", finishReason)
+  }
+  if finalContent != nil {
+    t.Errorf("expected finalContent to be nil, got %v", finalContent)
+  }
+  if reasoning != "" {
+    t.Errorf("expected empty reasoning, got %q", reasoning)
+  }
+  if len(finalCalls) != 1 {
+    t.Fatalf("expected 1 final tool call, got %d", len(finalCalls))
+  }
+}
+
+func TestDetectToolCallsPreservesPreambleBeforeToolCall(t *testing.T) {
+  raw := "I will check the weather in Tokyo for you.\n" +
+    "```xml\n" +
+    "<tool_call>\n" +
+    "{\"name\": \"get_current_weather\", \"arguments\": {\"location\": \"Tokyo\"}}\n" +
+    "</tool_call>\n" +
+    "```\n" +
+    "*Web evidence was retrieved, but the response did not include a valid source citation.*"
+
+  tcs, rem, ok := DetectToolCalls(raw)
+  if !ok || len(tcs) != 1 {
+    t.Fatalf("expected 1 tool call, got %d (ok: %v)", len(tcs), ok)
+  }
+  if tcs[0].Function.Name != "get_current_weather" {
+    t.Errorf("expected function get_current_weather, got %q", tcs[0].Function.Name)
+  }
+  if rem != "I will check the weather in Tokyo for you." {
+    t.Errorf("expected preamble preserved without postamble, got %q", rem)
+  }
+
+  frame := GradioOutputFrame{Content: raw}
+  finalContent, _, finalCalls, finishReason := finalizeOutput(frame)
+  if finishReason != "tool_calls" {
+    t.Errorf("expected finish_reason 'tool_calls', got %q", finishReason)
+  }
+  if finalContentStr, ok := finalContent.(string); !ok || finalContentStr != "I will check the weather in Tokyo for you." {
+    t.Errorf("expected finalContent 'I will check the weather in Tokyo for you.', got %v", finalContent)
+  }
+  if len(finalCalls) != 1 {
+    t.Fatalf("expected 1 tool call, got %d", len(finalCalls))
+  }
+}
+
+func TestStreamToolCallFilterPostCallLeakSuppression(t *testing.T) {
+  // Test 1: Tool call followed by citation disclaimer
+  filter1 := NewStreamToolCallFilter()
+  var contentParts1 []string
+  var toolCalls1 []ToolCall
+
+  onContent1 := func(s string) { contentParts1 = append(contentParts1, s) }
+  onToolCall1 := func(tc ToolCall) { toolCalls1 = append(toolCalls1, tc) }
+
+  chunks1 := []string{
+    "<tool_call>\n",
+    "{\"name\": \"get_current_weather\", \"arguments\": {\"location\": \"Tokyo\"}}\n",
+    "</tool_call>\n",
+    "\n*Web evidence was retrieved, but the response did not include a valid source citation.*",
+  }
+
+  for _, c := range chunks1 {
+    filter1.Feed(c, onContent1, onToolCall1)
+  }
+  filter1.Flush(onContent1, onToolCall1)
+
+  if len(toolCalls1) != 1 {
+    t.Fatalf("expected 1 tool call, got %d", len(toolCalls1))
+  }
+  if toolCalls1[0].Function.Name != "get_current_weather" {
+    t.Errorf("expected get_current_weather, got %q", toolCalls1[0].Function.Name)
+  }
+  if len(contentParts1) > 0 {
+    t.Errorf("expected 0 content parts leaked after tool call, got %v", contentParts1)
+  }
+
+  // Test 2: Parallel tool calls followed by trailing residue
+  filter2 := NewStreamToolCallFilter()
+  var contentParts2 []string
+  var toolCalls2 []ToolCall
+
+  onContent2 := func(s string) { contentParts2 = append(contentParts2, s) }
+  onToolCall2 := func(tc ToolCall) { toolCalls2 = append(toolCalls2, tc) }
+
+  chunks2 := []string{
+    "<tool_call>{\"name\": \"call_a\", \"arguments\": {}}</tool_call>\n",
+    "<tool_call>{\"name\": \"call_b\", \"arguments\": {}}</tool_call>\n",
+    "Residual text after parallel calls that should be suppressed.",
+  }
+
+  for _, c := range chunks2 {
+    filter2.Feed(c, onContent2, onToolCall2)
+  }
+  filter2.Flush(onContent2, onToolCall2)
+
+  if len(toolCalls2) != 2 {
+    t.Fatalf("expected 2 tool calls, got %d", len(toolCalls2))
+  }
+  if toolCalls2[0].Function.Name != "call_a" || toolCalls2[1].Function.Name != "call_b" {
+    t.Errorf("unexpected tool calls: %+v", toolCalls2)
+  }
+  if len(contentParts2) > 0 {
+    t.Errorf("expected 0 content parts leaked, got %v", contentParts2)
+  }
+
+  // Test 3: Preamble + Tool call + Postamble
+  filter3 := NewStreamToolCallFilter()
+  var contentParts3 []string
+  var toolCalls3 []ToolCall
+
+  onContent3 := func(s string) { contentParts3 = append(contentParts3, s) }
+  onToolCall3 := func(tc ToolCall) { toolCalls3 = append(toolCalls3, tc) }
+
+  chunks3 := []string{
+    "Thinking about your query: ",
+    "<tool_call>{\"name\": \"search\", \"arguments\": {\"q\": \"go\"}}</tool_call>",
+    " Trailing hallucinated answer that must not be emitted.",
+  }
+
+  for _, c := range chunks3 {
+    filter3.Feed(c, onContent3, onToolCall3)
+  }
+  filter3.Flush(onContent3, onToolCall3)
+
+  if len(toolCalls3) != 1 {
+    t.Fatalf("expected 1 tool call, got %d", len(toolCalls3))
+  }
+  fullContent3 := strings.Join(contentParts3, "")
+  if fullContent3 != "Thinking about your query: " {
+    t.Errorf("expected preamble 'Thinking about your query: ', got %q", fullContent3)
+  }
+}
+
+func TestWebSearchParameterMappingAndSuppression(t *testing.T) {
+  gw := &GradioGateway{}
+
+  // 1. Radio search_mode with choices ["Auto search", "Always search", "Direct"]
+  discRadio := NewDefaultSpaceDiscovery("https://test-search-space.hf.space")
+  discRadio.TotalInputs = 3
+  discRadio.MessageIndex = 0
+  discRadio.HistoryIndex = 1
+  discRadio.WebSearchIndex = 2
+  discRadio.ParamMappings = []SpaceParamMapping{
+    {InputIndex: 0, ComponentType: "textbox", Label: "message", ParamType: "message"},
+    {InputIndex: 1, ComponentType: "chatbot", Label: "history", ParamType: "history"},
+    {
+      InputIndex:    2,
+      ComponentType: "radio",
+      Label:         "search_mode",
+      ParamName:     "search_mode",
+      ParamType:     "web_search",
+      DefaultValue:  "Auto search",
+      Choices:       []string{"Auto search", "Always search", "Direct"},
+    },
+  }
+  discRadio.DefaultInputs = []interface{}{"", nil, "Auto search"}
+
+  // Normal request without tools retains space default ("Auto search")
+  reqNoTools := ChatCompletionRequest{
+    Model: "test-model",
+    Messages: []ChatMessage{
+      {Role: "user", Content: "Hello world"},
+    },
+  }
+  dataNoTools, err := gw.BuildGradioPayload(discRadio, reqNoTools)
+  if err != nil {
+    t.Fatalf("failed to build payload without tools: %v", err)
+  }
+  if dataNoTools[2] != "Auto search" {
+    t.Errorf("expected 'Auto search' default retained, got %v", dataNoTools[2])
+  }
+
+  // Request with explicit tools disables web search ("Direct")
+  reqWithTools := ChatCompletionRequest{
+    Model: "test-model",
+    Tools: []Tool{
+      {Type: "function", Function: map[string]interface{}{"name": "get_weather"}},
+    },
+    Messages: []ChatMessage{
+      {Role: "user", Content: "Weather in Tokyo?"},
+    },
+  }
+  dataWithTools, err := gw.BuildGradioPayload(discRadio, reqWithTools)
+  if err != nil {
+    t.Fatalf("failed to build payload with tools: %v", err)
+  }
+  if dataWithTools[2] != "Direct" {
+    t.Errorf("expected search_mode disabled to 'Direct', got %v", dataWithTools[2])
+  }
+
+  // 2. Checkbox web_search with bool default
+  discCheck := NewDefaultSpaceDiscovery("https://test-check-space.hf.space")
+  discCheck.TotalInputs = 2
+  discCheck.MessageIndex = 0
+  discCheck.WebSearchIndex = 1
+  discCheck.ParamMappings = []SpaceParamMapping{
+    {InputIndex: 0, ComponentType: "textbox", Label: "message", ParamType: "message"},
+    {InputIndex: 1, ComponentType: "checkbox", Label: "enable_web_search", ParamType: "web_search", DefaultValue: true},
+  }
+  discCheck.DefaultInputs = []interface{}{"", true}
+
+  dataCheckNoTools, _ := gw.BuildGradioPayload(discCheck, reqNoTools)
+  if dataCheckNoTools[1] != true {
+    t.Errorf("expected bool true retained without tools, got %v", dataCheckNoTools[1])
+  }
+
+  dataCheckWithTools, _ := gw.BuildGradioPayload(discCheck, reqWithTools)
+  if dataCheckWithTools[1] != false {
+    t.Errorf("expected bool false set with tools, got %v", dataCheckWithTools[1])
+  }
+}
+