fix(tools): recover tool calls from tool_use_failed errors and update default space to digital-twin
This commit is contained in:
+180
@@ -1024,4 +1024,184 @@ func TestUpstreamGradioErrorPropagation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractFailedGeneration(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "json wrapped",
|
||||
input: `{"error": {"code": 400, "failed_generation": "{\"name\": \"test_fn\", \"arguments\": {\"a\": 1}}"}}`,
|
||||
expected: `{"name": "test_fn", "arguments": {"a": 1}}`,
|
||||
},
|
||||
{
|
||||
name: "python repr with single quotes and escaped quotes and newlines",
|
||||
input: `upstream Gradio error: {'message': 'Tool choice is none, but model called a tool', 'type': 'invalid_request_error', 'code': 'tool_use_failed', 'failed_generation': '{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\\n}', 'status_code': 400}`,
|
||||
expected: "{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\n}",
|
||||
},
|
||||
{
|
||||
name: "python repr with standard single quotes",
|
||||
input: `{'code': 'tool_use_failed', 'failed_generation': '{"name": "calc", "arguments": {"x": 5}}'}`,
|
||||
expected: `{"name": "calc", "arguments": {"x": 5}}`,
|
||||
},
|
||||
{
|
||||
name: "raw object failed_generation",
|
||||
input: `{"failed_generation": {"name": "calc", "arguments": {"x": 5}}}`,
|
||||
expected: `{"name": "calc", "arguments": {"x": 5}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
got, ok := extractFailedGeneration(c.input)
|
||||
if !ok {
|
||||
t.Fatalf("extractFailedGeneration failed to find failed_generation in %q", c.input)
|
||||
}
|
||||
tcs, _, has := DetectToolCalls(got)
|
||||
if !has || len(tcs) == 0 {
|
||||
t.Fatalf("DetectToolCalls failed on extracted generation %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolUseFailedRecovery(t *testing.T) {
|
||||
errPayload := `upstream Gradio error: {'message': 'Tool choice is none, but model called a tool', 'type': 'invalid_request_error', 'code': 'tool_use_failed', 'failed_generation': '{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\\n}', 'status_code': 400}`
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/gradio_api/info" {
|
||||
resp := GradioAPIInfoResponse{
|
||||
NamedEndpoints: map[string]GradioEndpointInfo{
|
||||
"/chat_fn": {
|
||||
Parameters: []GradioParamInfo{
|
||||
{ParameterName: "message", Component: "Textbox"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/gradio_api/call/chat_fn" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_fail"})
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/gradio_api/call/chat_fn/evt_fail" {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
flusher := w.(http.Flusher)
|
||||
fmt.Fprintf(w, "event: error\ndata: %s\n\n", errPayload)
|
||||
flusher.Flush()
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
||||
|
||||
// 1. Non-streaming test
|
||||
recNonStream := httptest.NewRecorder()
|
||||
httpReqNonStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
err := gw.ExecuteChatCompletion(recNonStream, httpReqNonStream, ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []ChatMessage{{Role: "user", Content: "List files"}},
|
||||
Stream: false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected non-streaming to succeed by recovering tool call, got error: %v", err)
|
||||
}
|
||||
var nonStreamResp ChatCompletionResponse
|
||||
if err := json.NewDecoder(recNonStream.Body).Decode(&nonStreamResp); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if nonStreamResp.Choices[0].FinishReason != "tool_calls" {
|
||||
t.Fatalf("expected finish_reason 'tool_calls', got %q", nonStreamResp.Choices[0].FinishReason)
|
||||
}
|
||||
if len(nonStreamResp.Choices[0].Message.ToolCalls) != 1 {
|
||||
t.Fatalf("expected 1 tool call, got %d", len(nonStreamResp.Choices[0].Message.ToolCalls))
|
||||
}
|
||||
tc := nonStreamResp.Choices[0].Message.ToolCalls[0]
|
||||
if tc.Function.Name != "repo_browser.print_tree" {
|
||||
t.Errorf("expected function repo_browser.print_tree, got %q", tc.Function.Name)
|
||||
}
|
||||
|
||||
// 2. Streaming test
|
||||
recStream := httptest.NewRecorder()
|
||||
httpReqStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
err = gw.ExecuteChatCompletion(recStream, httpReqStream, ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []ChatMessage{{Role: "user", Content: "List files"}},
|
||||
Stream: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected streaming to succeed by recovering tool call, got error: %v", err)
|
||||
}
|
||||
bodyStream := recStream.Body.String()
|
||||
if !strings.Contains(bodyStream, "repo_browser.print_tree") {
|
||||
t.Errorf("expected streaming body to contain repo_browser.print_tree, got: %s", bodyStream)
|
||||
}
|
||||
if !strings.Contains(bodyStream, `"finish_reason":"tool_calls"`) {
|
||||
t.Errorf("expected streaming body to contain finish_reason tool_calls, got: %s", bodyStream)
|
||||
}
|
||||
if !strings.Contains(bodyStream, "[DONE]") {
|
||||
t.Errorf("expected streaming body to contain [DONE], got: %s", bodyStream)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolUseFailedImmediateCallRecovery(t *testing.T) {
|
||||
errPayload := `{"error": {"code": 400, "message": "Tool choice is none, but model called a tool", "failed_generation": "{\"name\": \"calculator\", \"arguments\": {\"expr\": \"2+2\"}}"}}`
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/gradio_api/info" {
|
||||
resp := GradioAPIInfoResponse{
|
||||
NamedEndpoints: map[string]GradioEndpointInfo{
|
||||
"/chat_fn": {
|
||||
Parameters: []GradioParamInfo{
|
||||
{ParameterName: "message", Component: "Textbox"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/gradio_api/call/chat_fn") || strings.HasPrefix(r.URL.Path, "/call/chat_fn") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
w.Write([]byte(errPayload))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
err := gw.ExecuteChatCompletion(rec, httpReq, ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []ChatMessage{{Role: "user", Content: "Calculate"}},
|
||||
Stream: false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected call-level error recovery to succeed, got: %v", err)
|
||||
}
|
||||
var resp ChatCompletionResponse
|
||||
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if resp.Choices[0].FinishReason != "tool_calls" {
|
||||
t.Fatalf("expected finish_reason 'tool_calls', got %q", resp.Choices[0].FinishReason)
|
||||
}
|
||||
if len(resp.Choices[0].Message.ToolCalls) != 1 {
|
||||
t.Fatalf("expected 1 tool call, got %d", len(resp.Choices[0].Message.ToolCalls))
|
||||
}
|
||||
if resp.Choices[0].Message.ToolCalls[0].Function.Name != "calculator" {
|
||||
t.Errorf("expected calculator function, got %q", resp.Choices[0].Message.ToolCalls[0].Function.Name)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user