1301 lines
45 KiB
Go
1301 lines
45 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestParseSOCKS5URL(t *testing.T) {
|
|
tests := []struct {
|
|
input string
|
|
expected *SOCKS5Config
|
|
}{
|
|
{"", nil},
|
|
{"127.0.0.1:1080", &SOCKS5Config{Address: "127.0.0.1:1080"}},
|
|
{"socks5://127.0.0.1:9050", &SOCKS5Config{Address: "127.0.0.1:9050"}},
|
|
{"socks5h://user:pass@10.0.0.1:1080", &SOCKS5Config{Address: "10.0.0.1:1080", Username: "user", Password: "pass"}},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
cfg, err := ParseSOCKS5URL(tc.input)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error for %q: %v", tc.input, err)
|
|
}
|
|
if tc.expected == nil {
|
|
if cfg != nil {
|
|
t.Errorf("expected nil config, got %+v", cfg)
|
|
}
|
|
continue
|
|
}
|
|
if cfg.Address != tc.expected.Address || cfg.Username != tc.expected.Username || cfg.Password != tc.expected.Password {
|
|
t.Errorf("for %q, expected %+v, got %+v", tc.input, tc.expected, cfg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestChatMessageGetContentString(t *testing.T) {
|
|
m1 := ChatMessage{Role: "user", Content: "hello world"}
|
|
if m1.GetContentString() != "hello world" {
|
|
t.Errorf("expected 'hello world', got %q", m1.GetContentString())
|
|
}
|
|
|
|
m2 := ChatMessage{
|
|
Role: "user",
|
|
Content: []interface{}{
|
|
map[string]interface{}{"type": "text", "text": "part 1 "},
|
|
map[string]interface{}{"type": "text", "text": "part 2"},
|
|
},
|
|
}
|
|
if m2.GetContentString() != "part 1 part 2" {
|
|
t.Errorf("expected 'part 1 part 2', got %q", m2.GetContentString())
|
|
}
|
|
}
|
|
|
|
func TestExtractThinking(t *testing.T) {
|
|
content := "<think>Let me calculate 2+2.</think>The answer is 4."
|
|
clean, reasoning := ExtractThinking(content)
|
|
if reasoning != "Let me calculate 2+2." {
|
|
t.Errorf("expected reasoning 'Let me calculate 2+2.', got %q", reasoning)
|
|
}
|
|
if clean != "The answer is 4." {
|
|
t.Errorf("expected clean 'The answer is 4.', got %q", clean)
|
|
}
|
|
}
|
|
|
|
func TestDetectToolCalls(t *testing.T) {
|
|
xmlContent := `<tool_call>
|
|
{"name": "get_weather", "arguments": {"city": "Paris"}}
|
|
</tool_call>`
|
|
calls, rem, ok := DetectToolCalls(xmlContent)
|
|
if !ok || len(calls) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d (ok: %v)", len(calls), ok)
|
|
}
|
|
if calls[0].Function.Name != "get_weather" {
|
|
t.Errorf("expected function name get_weather, got %q", calls[0].Function.Name)
|
|
}
|
|
if rem != "" {
|
|
t.Errorf("expected empty remaining content, got %q", rem)
|
|
}
|
|
|
|
jsonContent := `{"name": "calculator", "arguments": {"expr": "1+1"}}`
|
|
calls2, rem2, ok2 := DetectToolCalls(jsonContent)
|
|
if !ok2 || len(calls2) != 1 {
|
|
t.Fatalf("expected 1 tool call from JSON, got %d", len(calls2))
|
|
}
|
|
if calls2[0].Function.Name != "calculator" {
|
|
t.Errorf("expected function calculator, got %q", calls2[0].Function.Name)
|
|
}
|
|
if rem2 != "" {
|
|
t.Errorf("expected empty remaining, got %q", rem2)
|
|
}
|
|
}
|
|
|
|
func TestStreamThinkingFilter(t *testing.T) {
|
|
filter := NewStreamThinkingFilter()
|
|
var contentParts []string
|
|
var reasoningParts []string
|
|
|
|
onContent := func(s string) { contentParts = append(contentParts, s) }
|
|
onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }
|
|
|
|
chunks := []string{"<thi", "nk>Thinking de", "eply</th", "ink>Here is your answer."}
|
|
for _, c := range chunks {
|
|
filter.Feed(c, onContent, onReasoning)
|
|
}
|
|
filter.Flush(onContent, onReasoning)
|
|
|
|
fullReasoning := strings.Join(reasoningParts, "")
|
|
fullContent := strings.Join(contentParts, "")
|
|
|
|
if fullReasoning != "Thinking deeply" {
|
|
t.Errorf("expected reasoning 'Thinking deeply', got %q", fullReasoning)
|
|
}
|
|
if fullContent != "Here is your answer." {
|
|
t.Errorf("expected content 'Here is your answer.', got %q", fullContent)
|
|
}
|
|
}
|
|
|
|
func TestStreamToolCallFilter(t *testing.T) {
|
|
filter := NewStreamToolCallFilter()
|
|
var contentParts []string
|
|
var toolCalls []ToolCall
|
|
|
|
onContent := func(s string) { contentParts = append(contentParts, s) }
|
|
onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }
|
|
|
|
chunks := []string{
|
|
"Searching now: ",
|
|
"<tool_c",
|
|
"all>\n{\"name\": \"search_web\", \"arguments\": {\"query\": \"golang\"}}\n</tool_",
|
|
"call>",
|
|
" Done.",
|
|
}
|
|
|
|
for _, c := range chunks {
|
|
filter.Feed(c, onContent, onToolCall)
|
|
}
|
|
filter.Flush(onContent, onToolCall)
|
|
|
|
if len(toolCalls) != 1 {
|
|
t.Fatalf("expected 1 emitted tool call, got %d", len(toolCalls))
|
|
}
|
|
if toolCalls[0].Function.Name != "search_web" {
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestMockGradioServerCompletion(t *testing.T) {
|
|
// Setup a mock Gradio server
|
|
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"},
|
|
},
|
|
Returns: []GradioParamInfo{
|
|
{ParameterName: "response", Component: "Json"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
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_123"})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat_fn/evt_123" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatal("expected flusher")
|
|
}
|
|
fmt.Fprintf(w, "event: generating\ndata: [\"Hello \", null]\n\n")
|
|
flusher.Flush()
|
|
fmt.Fprintf(w, "event: generating\ndata: [\"Hello world!\", null]\n\n")
|
|
flusher.Flush()
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"Hello world!\", null]\n\n")
|
|
flusher.Flush()
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Test Non-streaming request
|
|
reqBody := ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hi"},
|
|
},
|
|
Stream: false,
|
|
}
|
|
b, _ := json.Marshal(reqBody)
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec, httpReq, reqBody)
|
|
if err != nil {
|
|
t.Fatalf("unexpected completion error: %v", err)
|
|
}
|
|
|
|
var resp ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
|
t.Fatalf("failed to decode completion response: %v", err)
|
|
}
|
|
if len(resp.Choices) != 1 {
|
|
t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
|
|
}
|
|
if resp.Choices[0].Message.GetContentString() != "Hello world!" {
|
|
t.Errorf("expected 'Hello world!', got %q", resp.Choices[0].Message.GetContentString())
|
|
}
|
|
|
|
// 2. Test Streaming request
|
|
reqBodyStream := reqBody
|
|
reqBodyStream.Stream = true
|
|
recStream := httptest.NewRecorder()
|
|
err = gw.ExecuteChatCompletion(recStream, httpReq, reqBodyStream)
|
|
if err != nil {
|
|
t.Fatalf("unexpected streaming error: %v", err)
|
|
}
|
|
streamOutput := recStream.Body.String()
|
|
if !strings.Contains(streamOutput, "data: [DONE]") {
|
|
t.Errorf("expected stream to contain [DONE], got:\n%s", streamOutput)
|
|
}
|
|
if !strings.Contains(streamOutput, "Hello world!") && !strings.Contains(streamOutput, "world!") {
|
|
t.Errorf("expected stream output to contain delta tokens, got:\n%s", streamOutput)
|
|
}
|
|
}
|
|
|
|
func TestParseGradioStreamOutput(t *testing.T) {
|
|
// 1. Standard 1D Gradio array
|
|
frame1 := ParseGradioStreamOutput(`["Hello from 1D", null]`)
|
|
if !frame1.OK || frame1.Content != "Hello from 1D" || frame1.Reasoning != "" || len(frame1.ToolCalls) != 0 {
|
|
t.Errorf("unexpected frame1: %+v", frame1)
|
|
}
|
|
|
|
// 2. Hy3 2D array frame with reasoning
|
|
hy3Raw := `[["Hello answer", "Let me think deeply...", [], [{"role": "user", "content": "hi"}]]]`
|
|
frame2 := ParseGradioStreamOutput(hy3Raw)
|
|
if !frame2.OK || frame2.Content != "Hello answer" || frame2.Reasoning != "Let me think deeply..." || len(frame2.ToolCalls) != 0 {
|
|
t.Errorf("unexpected frame2: %+v", frame2)
|
|
}
|
|
|
|
// 3. Hy3 2D array frame with tool calls
|
|
hy3ToolRaw := `[["", "Calling weather tool", [{"id": "call_abc", "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\": \"Tokyo\"}"}}], []]]`
|
|
frame3 := ParseGradioStreamOutput(hy3ToolRaw)
|
|
if !frame3.OK || frame3.Content != "" || frame3.Reasoning != "Calling weather tool" || len(frame3.ToolCalls) != 1 {
|
|
t.Fatalf("unexpected frame3: %+v", frame3)
|
|
}
|
|
if frame3.ToolCalls[0].ID != "call_abc" || frame3.ToolCalls[0].Function.Name != "get_weather" {
|
|
t.Errorf("unexpected tool call in frame3: %+v", frame3.ToolCalls[0])
|
|
}
|
|
|
|
// 4. Chat pairs
|
|
pairRaw := `[[["user prompt", "assistant answer"]]]`
|
|
frame4 := ParseGradioStreamOutput(pairRaw)
|
|
if !frame4.OK || frame4.Content != "assistant answer" {
|
|
t.Errorf("unexpected frame4: %+v", frame4)
|
|
}
|
|
|
|
// 5. Messages array
|
|
msgRaw := `[[{"role": "assistant", "content": "msg answer", "reasoning_content": "msg think"}]]`
|
|
frame5 := ParseGradioStreamOutput(msgRaw)
|
|
if !frame5.OK || frame5.Content != "msg answer" || frame5.Reasoning != "msg think" {
|
|
t.Errorf("unexpected frame5: %+v", frame5)
|
|
}
|
|
}
|
|
|
|
func TestHunyuan3BuildPayload(t *testing.T) {
|
|
gw := &GradioGateway{}
|
|
disc := &SpaceDiscovery{
|
|
TotalInputs: 9,
|
|
MessageIndex: 0,
|
|
SystemIndex: 1,
|
|
HistoryIndex: 2,
|
|
ThinkLevelIndex: 3,
|
|
TempIndex: 4,
|
|
MaxTokensIndex: 5,
|
|
TopPIndex: 6,
|
|
FunctionsJSONIndex: 8,
|
|
IsHunyuan3: true,
|
|
}
|
|
|
|
temp := 0.2
|
|
req := ChatCompletionRequest{
|
|
Model: "hy3",
|
|
ReasoningEffort: "low",
|
|
Temperature: &temp,
|
|
Tools: []Tool{
|
|
{
|
|
Type: "function",
|
|
Function: map[string]interface{}{
|
|
"name": "calc",
|
|
},
|
|
},
|
|
},
|
|
Messages: []ChatMessage{
|
|
{Role: "system", Content: "Be helpful"},
|
|
{Role: "user", Content: "2+2"},
|
|
{
|
|
Role: "assistant",
|
|
ReasoningContent: "Thinking...",
|
|
ToolCalls: []ToolCall{
|
|
{ID: "c1", Type: "function", Function: ToolCallFunction{Name: "calc", Arguments: `{"expr":"2+2"}`}},
|
|
},
|
|
},
|
|
{Role: "tool", ToolCallID: "c1", Content: "4"},
|
|
},
|
|
}
|
|
|
|
data, err := gw.BuildGradioPayload(disc, req)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
|
|
if len(data) != 9 {
|
|
t.Fatalf("expected 9 payload items, got %d", len(data))
|
|
}
|
|
|
|
// Message parameter (0): should be prompt continuation since last was tool
|
|
if msg, ok := data[0].(string); !ok || msg != "Please proceed based on the tool results." {
|
|
t.Errorf("expected continuation prompt, got %v", data[0])
|
|
}
|
|
|
|
// System parameter (1)
|
|
if sys, ok := data[1].(string); !ok || sys != "Be helpful" {
|
|
t.Errorf("expected 'Be helpful', got %v", data[1])
|
|
}
|
|
|
|
// History parameter (2): should contain all messages including the tool turn
|
|
hist, ok := data[2].([]map[string]interface{})
|
|
if !ok {
|
|
t.Fatalf("expected history slice of maps, got %T", data[2])
|
|
}
|
|
if len(hist) != 3 {
|
|
t.Fatalf("expected 3 history items (user, assistant, tool), got %d", len(hist))
|
|
}
|
|
if hist[2]["role"] != "tool" || hist[2]["content"] != "4" || hist[2]["tool_call_id"] != "c1" {
|
|
t.Errorf("unexpected tool history entry: %+v", hist[2])
|
|
}
|
|
|
|
// ThinkLevel parameter (3)
|
|
if data[3] != "low" {
|
|
t.Errorf("expected think_level 'low', got %v", data[3])
|
|
}
|
|
|
|
// Temp parameter (4)
|
|
if data[4] != 0.2 {
|
|
t.Errorf("expected temp 0.2, got %v", data[4])
|
|
}
|
|
|
|
// FunctionsJSON parameter (8)
|
|
fnStr, ok := data[8].(string)
|
|
if !ok || !strings.Contains(fnStr, "calc") {
|
|
t.Errorf("expected functions_json_str to contain 'calc', got %v", data[8])
|
|
}
|
|
}
|
|
|
|
func TestHunyuan3MockServerCompletion(t *testing.T) {
|
|
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": {
|
|
Parameters: []GradioParamInfo{
|
|
{ParameterName: "message"},
|
|
{ParameterName: "system_prompt"},
|
|
{ParameterName: "history"},
|
|
{ParameterName: "think_level"},
|
|
{ParameterName: "temperature"},
|
|
{ParameterName: "max_tokens"},
|
|
{ParameterName: "top_p"},
|
|
{ParameterName: "preserved_thinking"},
|
|
{ParameterName: "functions_json_str"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(resp)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_hy3"})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat/evt_hy3" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatal("expected flusher")
|
|
}
|
|
// Frame 1: Reasoning delta
|
|
fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Reasoning part 1 \", [], []]]\n\n")
|
|
flusher.Flush()
|
|
// Frame 2: Tool call initiated
|
|
fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Reasoning part 1 and 2\", [{\"id\": \"call_hy3\", \"type\": \"function\", \"function\": {\"name\": \"search\", \"arguments\": \"{\\\"q\\\": \\\"tencent\\\"}\"}}], []]]\n\n")
|
|
flusher.Flush()
|
|
// Frame 3: Completion
|
|
fmt.Fprintf(w, "event: complete\ndata: [[\"\", \"Reasoning part 1 and 2\", [{\"id\": \"call_hy3\", \"type\": \"function\", \"function\": {\"name\": \"search\", \"arguments\": \"{\\\"q\\\": \\\"tencent\\\"}\"}}], []]]\n\n")
|
|
flusher.Flush()
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Non-streaming tool call test
|
|
reqBody := ChatCompletionRequest{
|
|
Model: "hy3",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "search for tencent"},
|
|
},
|
|
Stream: false,
|
|
}
|
|
b, _ := json.Marshal(reqBody)
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec, httpReq, reqBody)
|
|
if err != nil {
|
|
t.Fatalf("unexpected completion error: %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.Errorf("expected finish_reason 'tool_calls', got %q", resp.Choices[0].FinishReason)
|
|
}
|
|
if resp.Choices[0].Message.ReasoningContent != "Reasoning part 1 and 2" {
|
|
t.Errorf("expected native reasoning, got %q", resp.Choices[0].Message.ReasoningContent)
|
|
}
|
|
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 != "search" {
|
|
t.Errorf("expected function 'search', got %q", resp.Choices[0].Message.ToolCalls[0].Function.Name)
|
|
}
|
|
|
|
// 2. Streaming tool call test
|
|
reqBodyStream := reqBody
|
|
reqBodyStream.Stream = true
|
|
recStream := httptest.NewRecorder()
|
|
err = gw.ExecuteChatCompletion(recStream, httpReq, reqBodyStream)
|
|
if err != nil {
|
|
t.Fatalf("unexpected streaming error: %v", err)
|
|
}
|
|
|
|
streamOut := recStream.Body.String()
|
|
if !strings.Contains(streamOut, "reasoning_content") {
|
|
t.Errorf("expected stream to contain reasoning_content, got:\n%s", streamOut)
|
|
}
|
|
if !strings.Contains(streamOut, "tool_calls") {
|
|
t.Errorf("expected stream to contain tool_calls, got:\n%s", streamOut)
|
|
}
|
|
if !strings.Contains(streamOut, "call_hy3") {
|
|
t.Errorf("expected stream to contain tool call ID call_hy3, got:\n%s", streamOut)
|
|
}
|
|
if !strings.Contains(streamOut, "\"finish_reason\":\"tool_calls\"") {
|
|
t.Errorf("expected stream finish_reason tool_calls, got:\n%s", streamOut)
|
|
}
|
|
}
|
|
|
|
func TestUniversalToolCallingTransformMessages(t *testing.T) {
|
|
req := ChatCompletionRequest{
|
|
Tools: []Tool{
|
|
{
|
|
Type: "function",
|
|
Function: map[string]interface{}{
|
|
"name": "get_weather",
|
|
"description": "Get current weather",
|
|
},
|
|
},
|
|
},
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "What is the weather in Tokyo and Paris?"},
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []ToolCall{
|
|
{ID: "call_tokyo", Type: "function", Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Tokyo"}`}},
|
|
{ID: "call_paris", Type: "function", Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Paris"}`}},
|
|
},
|
|
},
|
|
{Role: "tool", ToolCallID: "call_tokyo", Content: `{"temp": 20}`},
|
|
{Role: "tool", ToolCallID: "call_paris", Content: `{"temp": 15}`},
|
|
},
|
|
}
|
|
|
|
processed, toolInstruction, hasSystem := TransformMessages(req)
|
|
if !hasSystem {
|
|
t.Errorf("expected hasSystem to be true after injecting tool instructions")
|
|
}
|
|
if toolInstruction == "" {
|
|
t.Errorf("expected non-empty toolInstruction")
|
|
}
|
|
|
|
// Expect:
|
|
// [0] System message with tool instructions
|
|
// [1] User message: "What is the weather in Tokyo and Paris?"
|
|
// [2] Assistant message with <tool_call> blocks
|
|
// [3] User message with coalesced <tool_response> blocks
|
|
if len(processed) != 4 {
|
|
t.Fatalf("expected 4 processed messages, got %d", len(processed))
|
|
}
|
|
|
|
if processed[0].Role != "system" || !strings.Contains(processed[0].GetContentString(), "Tool Calling Instructions") {
|
|
t.Errorf("unexpected message 0: %+v", processed[0])
|
|
}
|
|
|
|
if processed[1].Role != "user" || processed[1].GetContentString() != "What is the weather in Tokyo and Paris?" {
|
|
t.Errorf("unexpected message 1: %+v", processed[1])
|
|
}
|
|
|
|
if processed[2].Role != "assistant" || !strings.Contains(processed[2].GetContentString(), "get_weather") {
|
|
t.Errorf("unexpected message 2: %+v", processed[2])
|
|
}
|
|
|
|
respContent := processed[3].GetContentString()
|
|
if processed[3].Role != "user" {
|
|
t.Errorf("expected coalesced message 3 to have role user, got %q", processed[3].Role)
|
|
}
|
|
if !strings.Contains(respContent, `{"name": "get_weather", "content": {"temp": 20}}`) {
|
|
t.Errorf("expected resolved function name get_weather for tokyo, got:\n%s", respContent)
|
|
}
|
|
if !strings.Contains(respContent, `{"name": "get_weather", "content": {"temp": 15}}`) {
|
|
t.Errorf("expected resolved function name get_weather for paris, got:\n%s", respContent)
|
|
}
|
|
if !strings.Contains(respContent, "Please answer the user's request based on the tool results.") {
|
|
t.Errorf("expected continuation prompt in coalesced message, got:\n%s", respContent)
|
|
}
|
|
}
|
|
|
|
func TestUniversalToolCallDetectionVariants(t *testing.T) {
|
|
// 1. Array of tool calls inside <tool_calls> tag
|
|
multiXML := `<tool_calls>
|
|
[
|
|
{"name": "get_weather", "arguments": {"city": "Tokyo"}},
|
|
{"name": "get_weather", "arguments": {"city": "Paris"}}
|
|
]
|
|
</tool_calls>`
|
|
calls1, rem1, ok1 := DetectToolCalls(multiXML)
|
|
if !ok1 || len(calls1) != 2 {
|
|
t.Fatalf("expected 2 tool calls from <tool_calls>, got %d", len(calls1))
|
|
}
|
|
if calls1[0].Function.Name != "get_weather" || calls1[1].Function.Name != "get_weather" {
|
|
t.Errorf("unexpected function names: %+v", calls1)
|
|
}
|
|
if rem1 != "" {
|
|
t.Errorf("expected empty remaining, got %q", rem1)
|
|
}
|
|
|
|
// 2. <function_call> tag
|
|
fnCallXML := `Some preamble before call.
|
|
<function_call>
|
|
{"name": "search", "arguments": {"q": "golang"}}
|
|
</function_call>
|
|
Some postamble.`
|
|
calls2, rem2, ok2 := DetectToolCalls(fnCallXML)
|
|
if !ok2 || len(calls2) != 1 {
|
|
t.Fatalf("expected 1 tool call from <function_call>, got %d", len(calls2))
|
|
}
|
|
if calls2[0].Function.Name != "search" {
|
|
t.Errorf("expected function search, got %q", calls2[0].Function.Name)
|
|
}
|
|
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)
|
|
}
|
|
|
|
// 3. [TOOL_CALLS] bracket syntax
|
|
bracketXML := `[TOOL_CALLS]
|
|
{"name": "calculate", "arguments": {"x": 42}}
|
|
[/TOOL_CALLS]`
|
|
calls3, rem3, ok3 := DetectToolCalls(bracketXML)
|
|
if !ok3 || len(calls3) != 1 {
|
|
t.Fatalf("expected 1 tool call from [TOOL_CALLS], got %d", len(calls3))
|
|
}
|
|
if calls3[0].Function.Name != "calculate" {
|
|
t.Errorf("expected function calculate, got %q", calls3[0].Function.Name)
|
|
}
|
|
if rem3 != "" {
|
|
t.Errorf("expected empty remaining, got %q", rem3)
|
|
}
|
|
|
|
// 4. Raw JSON array without tags
|
|
rawArray := `[{"name": "f1", "arguments": {}}, {"name": "f2", "arguments": {}}]`
|
|
calls4, rem4, ok4 := DetectToolCalls(rawArray)
|
|
if !ok4 || len(calls4) != 2 {
|
|
t.Fatalf("expected 2 calls from raw array, got %d", len(calls4))
|
|
}
|
|
if calls4[0].Function.Name != "f1" || calls4[1].Function.Name != "f2" {
|
|
t.Errorf("unexpected names from raw array: %+v", calls4)
|
|
}
|
|
if rem4 != "" {
|
|
t.Errorf("expected empty remaining, got %q", rem4)
|
|
}
|
|
}
|
|
|
|
func TestUniversalStreamToolCallFilterVariants(t *testing.T) {
|
|
filter := NewStreamToolCallFilter()
|
|
var contentParts []string
|
|
var toolCalls []ToolCall
|
|
|
|
onContent := func(s string) { contentParts = append(contentParts, s) }
|
|
onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }
|
|
|
|
// Stream using [TOOL_CALLS] across multiple chunk boundaries
|
|
chunks := []string{
|
|
"Preamble text: ",
|
|
"[TOOL_",
|
|
"CALLS]\n{\"name\": \"browse\", \"arguments\": {\"url\": \"example.com\"}}\n[/TOOL_",
|
|
"CALLS]",
|
|
" Completed.",
|
|
}
|
|
|
|
for _, c := range chunks {
|
|
filter.Feed(c, onContent, onToolCall)
|
|
}
|
|
filter.Flush(onContent, onToolCall)
|
|
|
|
if len(toolCalls) != 1 {
|
|
t.Fatalf("expected 1 tool call from stream filter, got %d", len(toolCalls))
|
|
}
|
|
if toolCalls[0].Function.Name != "browse" {
|
|
t.Errorf("expected function browse, got %q", toolCalls[0].Function.Name)
|
|
}
|
|
if !filter.emittedCall {
|
|
t.Errorf("expected emittedCall to be true")
|
|
}
|
|
fullContent := strings.Join(contentParts, "")
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestBuildGradioPayloadGenericSpaces(t *testing.T) {
|
|
gw := &GradioGateway{}
|
|
|
|
req := ChatCompletionRequest{
|
|
Tools: []Tool{
|
|
{Type: "function", Function: map[string]interface{}{"name": "lookup"}},
|
|
},
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "What is 10+10?"},
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []ToolCall{
|
|
{ID: "c1", Type: "function", Function: ToolCallFunction{Name: "lookup", Arguments: `{"q":"10+10"}`}},
|
|
},
|
|
},
|
|
{Role: "tool", ToolCallID: "c1", Content: `{"result": 20}`},
|
|
},
|
|
}
|
|
|
|
// 1. Space with native system prompt input (SystemIndex: 0, MessageIndex: 1, HistoryIndex: 2)
|
|
discWithSystem := NewDefaultSpaceDiscovery("https://space-1.hf.space")
|
|
discWithSystem.TotalInputs = 3
|
|
discWithSystem.SystemIndex = 0
|
|
discWithSystem.MessageIndex = 1
|
|
discWithSystem.HistoryIndex = 2
|
|
discWithSystem.HistoryFormat = "pairs"
|
|
|
|
data1, err := gw.BuildGradioPayload(discWithSystem, req)
|
|
if err != nil {
|
|
t.Fatalf("failed to build payload 1: %v", err)
|
|
}
|
|
sysStr, ok := data1[0].(string)
|
|
if !ok || !strings.Contains(sysStr, "Tool Calling Instructions") {
|
|
t.Errorf("expected system prompt at index 0, got %v", data1[0])
|
|
}
|
|
msgStr, ok := data1[1].(string)
|
|
if !ok || !strings.Contains(msgStr, "Please answer the user's request based on the tool result.") {
|
|
t.Errorf("expected coalesced tool prompt at index 1, got %v", data1[1])
|
|
}
|
|
pairs1, ok := data1[2].([][]string)
|
|
if !ok || len(pairs1) != 1 {
|
|
t.Fatalf("expected 1 history pair at index 2, got %T (%v)", data1[2], data1[2])
|
|
}
|
|
if pairs1[0][0] != "What is 10+10?" || !strings.Contains(pairs1[0][1], "lookup") {
|
|
t.Errorf("unexpected history pair: %+v", pairs1[0])
|
|
}
|
|
|
|
// 2. Space without system prompt (SystemIndex: -1, MessageIndex: 0, HistoryIndex: 1)
|
|
discNoSystem := NewDefaultSpaceDiscovery("https://space-2.hf.space")
|
|
discNoSystem.TotalInputs = 2
|
|
discNoSystem.SystemIndex = -1
|
|
discNoSystem.MessageIndex = 0
|
|
discNoSystem.HistoryIndex = 1
|
|
discNoSystem.HistoryFormat = "pairs"
|
|
|
|
data2, err := gw.BuildGradioPayload(discNoSystem, req)
|
|
if err != nil {
|
|
t.Fatalf("failed to build payload 2: %v", err)
|
|
}
|
|
pairs2, ok := data2[1].([][]string)
|
|
if !ok || len(pairs2) != 1 {
|
|
t.Fatalf("expected 1 history pair at index 1, got %T (%v)", data2[1], data2[1])
|
|
}
|
|
// Instructions prepended to the first user turn:
|
|
if !strings.Contains(pairs2[0][0], "Tool Calling Instructions") || !strings.Contains(pairs2[0][0], "What is 10+10?") {
|
|
t.Errorf("expected system instructions prepended to first pair user message, got: %q", pairs2[0][0])
|
|
}
|
|
if !strings.Contains(pairs2[0][1], "lookup") {
|
|
t.Errorf("expected assistant tool call in pair bot turn, got: %q", pairs2[0][1])
|
|
}
|
|
|
|
// 3. Single-textbox space (SystemIndex: -1, MessageIndex: 0, HistoryIndex: -1)
|
|
discSingleInput := NewDefaultSpaceDiscovery("https://space-3.hf.space")
|
|
discSingleInput.TotalInputs = 1
|
|
discSingleInput.SystemIndex = -1
|
|
discSingleInput.MessageIndex = 0
|
|
discSingleInput.HistoryIndex = -1
|
|
|
|
data3, err := gw.BuildGradioPayload(discSingleInput, req)
|
|
if err != nil {
|
|
t.Fatalf("failed to build payload 3: %v", err)
|
|
}
|
|
transcript, ok := data3[0].(string)
|
|
if !ok {
|
|
t.Fatalf("expected string transcript, got %T", data3[0])
|
|
}
|
|
if !strings.Contains(transcript, "# Instructions") || !strings.Contains(transcript, "User: What is 10+10?") || !strings.Contains(transcript, "Assistant: <tool_call>") || !strings.Contains(transcript, "# Current Request") {
|
|
t.Errorf("unexpected single-input transcript: %s", transcript)
|
|
}
|
|
}
|
|
|
|
func TestGenericSpaceMockServerToolCalling(t *testing.T) {
|
|
var lastReceivedData []interface{}
|
|
|
|
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: "system_prompt"},
|
|
{ParameterName: "message"},
|
|
{ParameterName: "history"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(resp)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat_fn" {
|
|
var body map[string]interface{}
|
|
json.NewDecoder(r.Body).Decode(&body)
|
|
if dataSlice, ok := body["data"].([]interface{}); ok {
|
|
lastReceivedData = dataSlice
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_generic"})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat_fn/evt_generic" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatal("expected flusher")
|
|
}
|
|
msgStr := ""
|
|
if len(lastReceivedData) > 1 {
|
|
msgStr, _ = lastReceivedData[1].(string)
|
|
}
|
|
|
|
if strings.Contains(msgStr, "<tool_response>") {
|
|
// Turn 2: answer
|
|
fmt.Fprintf(w, "event: generating\ndata: [\"The weather in Tokyo is 20 C.\", null]\n\n")
|
|
flusher.Flush()
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"The weather in Tokyo is 20 C.\", null]\n\n")
|
|
flusher.Flush()
|
|
} else {
|
|
// Turn 1: tool call
|
|
fmt.Fprintf(w, "event: generating\ndata: [\"<tool_call>\\n{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"city\\\": \\\"Tokyo\\\"}}\\n</tool_call>\", null]\n\n")
|
|
flusher.Flush()
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"<tool_call>\\n{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"city\\\": \\\"Tokyo\\\"}}\\n</tool_call>\", null]\n\n")
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// Turn 1: User question with tools
|
|
req1 := ChatCompletionRequest{
|
|
Model: "generic-bot",
|
|
Tools: []Tool{
|
|
{Type: "function", Function: map[string]interface{}{"name": "get_weather"}},
|
|
},
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Weather in Tokyo?"},
|
|
},
|
|
Stream: false,
|
|
}
|
|
|
|
b1, _ := json.Marshal(req1)
|
|
httpReq1 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b1))
|
|
rec1 := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec1, httpReq1, req1)
|
|
if err != nil {
|
|
t.Fatalf("Turn 1 execution failed: %v", err)
|
|
}
|
|
|
|
var resp1 ChatCompletionResponse
|
|
if err := json.NewDecoder(rec1.Body).Decode(&resp1); err != nil {
|
|
t.Fatalf("Turn 1 decode failed: %v", err)
|
|
}
|
|
if resp1.Choices[0].FinishReason != "tool_calls" {
|
|
t.Fatalf("expected finish_reason 'tool_calls', got %q", resp1.Choices[0].FinishReason)
|
|
}
|
|
if len(resp1.Choices[0].Message.ToolCalls) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d", len(resp1.Choices[0].Message.ToolCalls))
|
|
}
|
|
tc := resp1.Choices[0].Message.ToolCalls[0]
|
|
if tc.Function.Name != "get_weather" {
|
|
t.Fatalf("expected function name get_weather, got %q", tc.Function.Name)
|
|
}
|
|
|
|
// Turn 2: Send tool response
|
|
req2 := ChatCompletionRequest{
|
|
Model: "generic-bot",
|
|
Tools: []Tool{
|
|
{Type: "function", Function: map[string]interface{}{"name": "get_weather"}},
|
|
},
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Weather in Tokyo?"},
|
|
resp1.Choices[0].Message,
|
|
{Role: "tool", ToolCallID: tc.ID, Content: `{"temp": 20}`},
|
|
},
|
|
Stream: false,
|
|
}
|
|
|
|
b2, _ := json.Marshal(req2)
|
|
httpReq2 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b2))
|
|
rec2 := httptest.NewRecorder()
|
|
|
|
err = gw.ExecuteChatCompletion(rec2, httpReq2, req2)
|
|
if err != nil {
|
|
t.Fatalf("Turn 2 execution failed: %v", err)
|
|
}
|
|
|
|
var resp2 ChatCompletionResponse
|
|
if err := json.NewDecoder(rec2.Body).Decode(&resp2); err != nil {
|
|
t.Fatalf("Turn 2 decode failed: %v", err)
|
|
}
|
|
if resp2.Choices[0].FinishReason != "stop" {
|
|
t.Errorf("expected finish_reason 'stop', got %q", resp2.Choices[0].FinishReason)
|
|
}
|
|
if resp2.Choices[0].Message.GetContentString() != "The weather in Tokyo is 20 C." {
|
|
t.Errorf("expected final answer, got %q", resp2.Choices[0].Message.GetContentString())
|
|
}
|
|
}
|
|
|
|
func TestExtractGradioErrorMessage(t *testing.T) {
|
|
cases := []struct {
|
|
input string
|
|
expected string
|
|
}{
|
|
{
|
|
input: `{"error": "Client error '402 Payment Required'"}`,
|
|
expected: "Client error '402 Payment Required'",
|
|
},
|
|
{
|
|
input: `{"message": "Rate limit exceeded"}`,
|
|
expected: "Rate limit exceeded",
|
|
},
|
|
{
|
|
input: `{"error": null, "title": "Validation Error"}`,
|
|
expected: "Validation Error",
|
|
},
|
|
{
|
|
input: `{"error": null}`,
|
|
expected: "internal space error (check Gradio inputs/types)",
|
|
},
|
|
{
|
|
input: `raw server failure`,
|
|
expected: "raw server failure",
|
|
},
|
|
}
|
|
|
|
for _, c := range cases {
|
|
got := extractGradioErrorMessage(c.input)
|
|
if got != c.expected {
|
|
t.Errorf("extractGradioErrorMessage(%q) = %q, expected %q", c.input, got, c.expected)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestStreamerLazyStartAndErrors(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
s := NewStreamer(rec, nil, "cmpl-1", 12345, "test-model")
|
|
if s.started {
|
|
t.Errorf("expected streamer to start as not started")
|
|
}
|
|
|
|
// Ensure Role/headers only sent on first write
|
|
s.Content("Hello")
|
|
if !s.started {
|
|
t.Errorf("expected streamer to be started after Content")
|
|
}
|
|
s.Finish("stop")
|
|
s.Done()
|
|
|
|
body := rec.Body.String()
|
|
if !strings.Contains(body, `"role":"assistant"`) {
|
|
t.Errorf("expected role chunk in body: %s", body)
|
|
}
|
|
if !strings.Contains(body, `"content":"Hello"`) {
|
|
t.Errorf("expected content chunk in body: %s", body)
|
|
}
|
|
if !strings.Contains(body, "[DONE]") {
|
|
t.Errorf("expected [DONE] in body: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestUpstreamGradioErrorPropagation(t *testing.T) {
|
|
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: "err-event"})
|
|
return
|
|
}
|
|
if r.URL.Path == "/gradio_api/call/chat_fn/err-event" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: error\ndata: {\"error\": \"Quota exceeded: 402 Payment Required\"}\n\n")
|
|
flusher.Flush()
|
|
return
|
|
}
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Non-streaming error should return error with upstream message
|
|
recNonStream := httptest.NewRecorder()
|
|
httpReqNonStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
err := gw.ExecuteChatCompletion(recNonStream, httpReqNonStream, ChatCompletionRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
|
|
Stream: false,
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("expected error from non-streaming upstream failure, got nil")
|
|
}
|
|
if !strings.Contains(err.Error(), "402 Payment Required") {
|
|
t.Errorf("expected error to mention 402 Payment Required, got %v", err)
|
|
}
|
|
|
|
// 2. Streaming error before start should return error with upstream message
|
|
recStream := httptest.NewRecorder()
|
|
httpReqStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
err = gw.ExecuteChatCompletion(recStream, httpReqStream, ChatCompletionRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
|
|
Stream: true,
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("expected error from streaming upstream failure, got nil")
|
|
}
|
|
if !strings.Contains(err.Error(), "402 Payment Required") {
|
|
t.Errorf("expected error to mention 402 Payment Required, got %v", err)
|
|
}
|
|
// And nothing should have been written to body
|
|
if strings.Contains(recStream.Body.String(), "[DONE]") {
|
|
t.Errorf("did not expect [DONE] on upstream error: %s", recStream.Body.String())
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestSystemPromptAugmentationWithoutNativeToolCalling(t *testing.T) {
|
|
gw := NewGradioGateway("https://generic-space.hf.space", "", 10*time.Second)
|
|
|
|
tools := []Tool{
|
|
{
|
|
Type: "function",
|
|
Function: map[string]interface{}{
|
|
"name": "search_docs",
|
|
"description": "Search local documentation",
|
|
"parameters": map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"query": map[string]interface{}{"type": "string"},
|
|
},
|
|
"required": []string{"query"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// Case 1: Space with SystemIndex and space default prompt, client sends no system message
|
|
discWithSys := NewDefaultSpaceDiscovery("https://generic-space.hf.space")
|
|
discWithSys.TotalInputs = 2
|
|
discWithSys.MessageIndex = 0
|
|
discWithSys.SystemIndex = 1
|
|
discWithSys.DefaultSystemPrompt = "You are a specialized documentation bot."
|
|
|
|
req1 := ChatCompletionRequest{
|
|
Tools: tools,
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "How to configure SSL?"},
|
|
},
|
|
}
|
|
|
|
payload1, err := gw.BuildGradioPayload(discWithSys, req1)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
sysStr, ok := payload1[1].(string)
|
|
if !ok {
|
|
t.Fatalf("expected string at SystemIndex 1, got %T", payload1[1])
|
|
}
|
|
if !strings.Contains(sysStr, "You are a specialized documentation bot.") {
|
|
t.Errorf("expected default system prompt to be retained, got: %s", sysStr)
|
|
}
|
|
if !strings.Contains(sysStr, "Tool Calling Instructions") || !strings.Contains(sysStr, "search_docs") {
|
|
t.Errorf("expected tool calling instructions and function name in system prompt, got: %s", sysStr)
|
|
}
|
|
|
|
// Case 2: Space with SystemIndex, client sends their own system message
|
|
req2 := ChatCompletionRequest{
|
|
Tools: tools,
|
|
Messages: []ChatMessage{
|
|
{Role: "system", Content: "You are an expert developer assistant."},
|
|
{Role: "user", Content: "How to configure SSL?"},
|
|
},
|
|
}
|
|
payload2, err := gw.BuildGradioPayload(discWithSys, req2)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
sysStr2, ok := payload2[1].(string)
|
|
if !ok {
|
|
t.Fatalf("expected string at SystemIndex 1, got %T", payload2[1])
|
|
}
|
|
if !strings.Contains(sysStr2, "You are an expert developer assistant.") {
|
|
t.Errorf("expected client system message, got: %s", sysStr2)
|
|
}
|
|
if !strings.Contains(sysStr2, "search_docs") {
|
|
t.Errorf("expected search_docs tool instruction, got: %s", sysStr2)
|
|
}
|
|
|
|
// Case 3: Space without SystemIndex, single input textbox
|
|
discSingle := NewDefaultSpaceDiscovery("https://single-input.hf.space")
|
|
discSingle.TotalInputs = 1
|
|
discSingle.MessageIndex = 0
|
|
discSingle.SystemIndex = -1
|
|
discSingle.HistoryIndex = -1
|
|
|
|
payload3, err := gw.BuildGradioPayload(discSingle, req1)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
singleMsg, ok := payload3[0].(string)
|
|
if !ok {
|
|
t.Fatalf("expected string at index 0, got %T", payload3[0])
|
|
}
|
|
if !strings.Contains(singleMsg, "Tool Calling Instructions") || !strings.Contains(singleMsg, "How to configure SSL?") {
|
|
t.Errorf("expected single message to contain augmented instructions and user query, got: %s", singleMsg)
|
|
}
|
|
}
|
|
|
|
|
|
|