Make prompt-augmented tool calling robust across non-native spaces

2026-09-07 15:19:23 +03:00
parent cc2fa234ac
commit 21c11c5fdd
2 changed files with 1065 additions and 86 deletions
@@ -2729,3 +2729,231 @@ func TestFinalizeOutputToolCallSanitization(t *testing.T) {
    t.Errorf("unexpected tool calls: %+v", tcs4)
  }
 }
+
+func TestPythonFunctionCallParsing(t *testing.T) {
+  // 1. Direct function call with kwargs
+  pyCall := `get_weather(city="Paris", units="metric")`
+  tcs1, rem1, ok1 := DetectToolCalls(pyCall)
+  if !ok1 || len(tcs1) != 1 {
+    t.Fatalf("expected 1 call from python call, got %d (ok=%v)", len(tcs1), ok1)
+  }
+  if tcs1[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs1[0].Function.Name)
+  }
+  var args1 map[string]interface{}
+  if err := json.Unmarshal([]byte(tcs1[0].Function.Arguments), &args1); err != nil {
+    t.Fatalf("invalid json args: %v", err)
+  }
+  if args1["city"] != "Paris" || args1["units"] != "metric" {
+    t.Errorf("unexpected args: %+v", args1)
+  }
+  if rem1 != "" {
+    t.Errorf("expected empty remaining, got %q", rem1)
+  }
+
+  // 2. Python call with single quotes and booleans
+  pyCall2 := `set_config(key='debug_mode', enabled=True, timeout=30, retries=None)`
+  tcs2, _, ok2 := DetectToolCalls(pyCall2)
+  if !ok2 || len(tcs2) != 1 {
+    t.Fatalf("expected 1 call, got %d", len(tcs2))
+  }
+  var args2 map[string]interface{}
+  if err := json.Unmarshal([]byte(tcs2[0].Function.Arguments), &args2); err != nil {
+    t.Fatalf("invalid json args: %v", err)
+  }
+  if args2["key"] != "debug_mode" || args2["enabled"] != true || args2["timeout"] != float64(30) || args2["retries"] != nil {
+    t.Errorf("unexpected args: %+v", args2)
+  }
+
+  // 3. Python call inside <tool_call> tag
+  tagPy := `<tool_call>
+get_weather(city="Tokyo")
+</tool_call>`
+  tcs3, _, ok3 := DetectToolCalls(tagPy)
+  if !ok3 || len(tcs3) != 1 {
+    t.Fatalf("expected 1 call from tagPy, got %d", len(tcs3))
+  }
+  if tcs3[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs3[0].Function.Name)
+  }
+
+  // 4. Multiple calls in list
+  multiPy := `[calc(expr='2+2'), search(q='golang')]`
+  tcs4, _, ok4 := DetectToolCalls(multiPy)
+  if !ok4 || len(tcs4) != 2 {
+    t.Fatalf("expected 2 calls from multiPy, got %d", len(tcs4))
+  }
+  if tcs4[0].Function.Name != "calc" || tcs4[1].Function.Name != "search" {
+    t.Errorf("unexpected multiPy calls: %+v", tcs4)
+  }
+}
+
+func TestDynamicXMLTagsWithAttributes(t *testing.T) {
+  // 1. Anthropic-style <invoke> with <parameter>
+  anthropic := `<invoke name="get_weather">
+<parameter name="city">Paris</parameter>
+</invoke>`
+  tcs1, _, ok1 := DetectToolCalls(anthropic)
+  if !ok1 || len(tcs1) != 1 {
+    t.Fatalf("expected 1 call from anthropic invoke, got %d", len(tcs1))
+  }
+  if tcs1[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs1[0].Function.Name)
+  }
+  if !strings.Contains(tcs1[0].Function.Arguments, `"Paris"`) {
+    t.Errorf("expected Paris in arguments, got %s", tcs1[0].Function.Arguments)
+  }
+
+  // 2. <tool_call name="get_weather">
+  attrTag := `<tool_call name="get_weather">
+{"city": "Berlin"}
+</tool_call>`
+  tcs2, _, ok2 := DetectToolCalls(attrTag)
+  if !ok2 || len(tcs2) != 1 {
+    t.Fatalf("expected 1 call from attrTag, got %d", len(tcs2))
+  }
+  if tcs2[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs2[0].Function.Name)
+  }
+
+  // 3. <call:get_weather>
+  colonTag := `<call:get_weather>
+{"city": "Madrid"}
+</call:get_weather>`
+  tcs3, _, ok3 := DetectToolCalls(colonTag)
+  if !ok3 || len(tcs3) != 1 {
+    t.Fatalf("expected 1 call from colonTag, got %d", len(tcs3))
+  }
+  if tcs3[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs3[0].Function.Name)
+  }
+
+  // 4. <function=get_weather>
+  eqTag := `<function=get_weather>
+{"city": "Rome"}
+</function>`
+  tcs4, _, ok4 := DetectToolCalls(eqTag)
+  if !ok4 || len(tcs4) != 1 {
+    t.Fatalf("expected 1 call from eqTag, got %d", len(tcs4))
+  }
+  if tcs4[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs4[0].Function.Name)
+  }
+
+  // 5. <command name="get_weather">
+  cmdTag := `<command name="get_weather">
+{"city": "Lisbon"}
+</command>`
+  tcs5, _, ok5 := DetectToolCalls(cmdTag)
+  if !ok5 || len(tcs5) != 1 {
+    t.Fatalf("expected 1 call from cmdTag, got %d", len(tcs5))
+  }
+  if tcs5[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs5[0].Function.Name)
+  }
+}
+
+func TestReActToolCallParsing(t *testing.T) {
+  react := `I will query the weather database.
+Action: get_weather
+Action Input: {"city": "Paris"}`
+  tcs, rem, ok := DetectToolCalls(react)
+  if !ok || len(tcs) != 1 {
+    t.Fatalf("expected 1 call from ReAct, got %d", len(tcs))
+  }
+  if tcs[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs[0].Function.Name)
+  }
+  if !strings.Contains(rem, "I will query the weather database.") {
+    t.Errorf("expected preamble commentary preserved, got %q", rem)
+  }
+  if strings.Contains(rem, "Action:") || strings.Contains(rem, "Action Input:") {
+    t.Errorf("expected ReAct markers scrubbed, got %q", rem)
+  }
+}
+
+func TestPythonSingleQuoteJSON(t *testing.T) {
+  dirtyJSON := `<tool_call>
+{'name': 'get_weather', 'arguments': {'city': 'Paris', 'active': True}}
+</tool_call>`
+  tcs, _, ok := DetectToolCalls(dirtyJSON)
+  if !ok || len(tcs) != 1 {
+    t.Fatalf("expected 1 call from dirtyJSON, got %d", len(tcs))
+  }
+  if tcs[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs[0].Function.Name)
+  }
+  var args map[string]interface{}
+  if err := json.Unmarshal([]byte(tcs[0].Function.Arguments), &args); err != nil {
+    t.Fatalf("failed to unmarshal parsed args: %v", err)
+  }
+  if args["active"] != true || args["city"] != "Paris" {
+    t.Errorf("unexpected args: %+v", args)
+  }
+}
+
+func TestFencedJSONToolCallWithoutTags(t *testing.T) {
+  fenced := "Let me check that for you.\n```json\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n```\n"
+  tcs, rem, ok := DetectToolCalls(fenced)
+  if !ok || len(tcs) != 1 {
+    t.Fatalf("expected 1 call from fenced json, got %d", len(tcs))
+  }
+  if tcs[0].Function.Name != "get_weather" {
+    t.Errorf("expected get_weather, got %s", tcs[0].Function.Name)
+  }
+  if !strings.Contains(rem, "Let me check that for you.") {
+    t.Errorf("expected commentary preserved, got %q", rem)
+  }
+}
+
+func TestMultiTurnToolResponseNotWrappedInQuery(t *testing.T) {
+  gw := &GradioGateway{}
+  disc := &SpaceDiscovery{
+    SpaceURL:            "https://test-space.hf.space",
+    TotalInputs:         2,
+    MessageIndex:        0,
+    HistoryIndex:        1,
+    HistoryFormat:       "gradio_messages",
+    SystemIndex:         -1,
+    FunctionsJSONIndex:  -1,
+    DefaultSystemPrompt: "",
+  }
+
+  req := ChatCompletionRequest{
+    Tools: []Tool{
+      {
+        Type: "function",
+        Function: map[string]interface{}{
+          "name": "get_weather",
+        },
+      },
+    },
+    Messages: []ChatMessage{
+      {Role: "user", Content: "What is the weather in Paris?"},
+      {
+        Role: "assistant",
+        ToolCalls: []ToolCall{
+          {ID: "call_123", Type: "function", Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Paris"}`}},
+        },
+      },
+      {Role: "tool", ToolCallID: "call_123", Name: "get_weather", Content: `{"temp": 22}`},
+    },
+  }
+
+  payload, err := gw.BuildGradioPayload(disc, req)
+  if err != nil {
+    t.Fatalf("BuildGradioPayload failed: %v", err)
+  }
+
+  msgStr, ok := payload[0].(string)
+  if !ok {
+    t.Fatalf("expected string at index 0, got %T", payload[0])
+  }
+
+  if strings.Contains(msgStr, "Query: <tool_response>") {
+    t.Errorf("message input should NOT contain 'Query: <tool_response>', got:\n%s", msgStr)
+  }
+  if !strings.Contains(msgStr, "<tool_response>") {
+    t.Errorf("message input should contain tool response, got:\n%s", msgStr)
+  }
+}