3370 lines
118 KiB
Go
3370 lines
118 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)
|
|
}
|
|
// 6. Hy3 with null content and reasoning
|
|
hy3NullContent := `[[null, "Thinking process...", null]]`
|
|
frame6 := ParseGradioStreamOutput(hy3NullContent)
|
|
if !frame6.OK || frame6.Content != "" || frame6.Reasoning != "Thinking process..." {
|
|
t.Errorf("unexpected frame6: %+v", frame6)
|
|
}
|
|
|
|
// 7. Hy3 with 3 elements (answer, reasoning, null tool calls)
|
|
hy3ThreeElem := `[["Mocked answer", "Mocked reasoning", null]]`
|
|
frame7 := ParseGradioStreamOutput(hy3ThreeElem)
|
|
if !frame7.OK || frame7.Content != "Mocked answer" || frame7.Reasoning != "Mocked reasoning" || len(frame7.ToolCalls) != 0 {
|
|
t.Errorf("unexpected frame7: %+v", frame7)
|
|
}
|
|
|
|
// 8. Gradio 6 Chatbot multi-output format with TextMessage list
|
|
g6Raw := `["", [{"role": "user", "content": [{"text": "weather in Tokyo?", "type": "text"}]}, {"role": "assistant", "content": [{"text": "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>", "type": "text"}]}]]`
|
|
frame8 := ParseGradioStreamOutput(g6Raw)
|
|
if !frame8.OK || len(frame8.ToolCalls) != 1 || frame8.ToolCalls[0].Function.Name != "get_weather" {
|
|
t.Errorf("unexpected frame8: %+v", frame8)
|
|
}
|
|
}
|
|
|
|
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 tool result prompt matching hygate behavior
|
|
if msg, ok := data[0].(string); !ok || msg != "Tool result for c1: 4" {
|
|
t.Errorf("expected 'Tool result for c1: 4', 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 prior turns (user, assistant with tool calls)
|
|
hist, ok := data[2].([]map[string]interface{})
|
|
if !ok {
|
|
t.Fatalf("expected history slice of maps, got %T", data[2])
|
|
}
|
|
if len(hist) != 2 {
|
|
t.Fatalf("expected 2 history items (user, assistant), got %d", len(hist))
|
|
}
|
|
if hist[0]["role"] != "user" || hist[0]["content"] != "2+2" {
|
|
t.Errorf("unexpected user history entry: %+v", hist[0])
|
|
}
|
|
if hist[1]["role"] != "assistant" || len(hist[1]["tool_calls"].([]ToolCall)) != 1 {
|
|
t.Errorf("unexpected assistant history entry: %+v", hist[1])
|
|
}
|
|
|
|
// 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])
|
|
}
|
|
|
|
// Test multi-tool turn: 2 tool messages at the end
|
|
reqMulti := req
|
|
reqMulti.Messages = append(reqMulti.Messages, ChatMessage{Role: "tool", Name: "fetch", Content: "done"})
|
|
dataMulti, err := gw.BuildGradioPayload(disc, reqMulti)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed on multi-tool: %v", err)
|
|
}
|
|
// The last tool message is data[0]
|
|
if msg, ok := dataMulti[0].(string); !ok || msg != "Tool result for fetch: done" {
|
|
t.Errorf("expected 'Tool result for fetch: done', got %v", dataMulti[0])
|
|
}
|
|
// The first tool message is in history
|
|
histMulti, _ := dataMulti[2].([]map[string]interface{})
|
|
if len(histMulti) != 3 {
|
|
t.Fatalf("expected 3 history items (user, assistant, tool 1), got %d", len(histMulti))
|
|
}
|
|
if histMulti[2]["role"] != "tool" || histMulti[2]["content"] != "4" {
|
|
t.Errorf("unexpected tool 1 in history: %+v", histMulti[2])
|
|
}
|
|
}
|
|
|
|
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(), "API router") {
|
|
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)
|
|
}
|
|
|
|
// 5. XML toolCall with child tags
|
|
xmlChildTags := "```xml\n<toolCall>\n <name>get_weather</name>\n <arguments>\n <loc>Paris</loc>\n </arguments>\n</toolCall>\n```"
|
|
calls5, rem5, ok5 := DetectToolCalls(xmlChildTags)
|
|
if !ok5 || len(calls5) != 1 {
|
|
t.Fatalf("expected 1 call from xmlChildTags, got %d", len(calls5))
|
|
}
|
|
if calls5[0].Function.Name != "get_weather" {
|
|
t.Errorf("expected get_weather, got %q", calls5[0].Function.Name)
|
|
}
|
|
if !strings.Contains(calls5[0].Function.Arguments, `"Paris"`) {
|
|
t.Errorf("expected Paris in arguments, got %s", calls5[0].Function.Arguments)
|
|
}
|
|
if rem5 != "" {
|
|
t.Errorf("expected empty remaining, got %q", rem5)
|
|
}
|
|
}
|
|
|
|
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, "API router") {
|
|
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], "API router") || !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, "API router") || !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, "API router") || !strings.Contains(singleMsg, "How to configure SSL?") {
|
|
t.Errorf("expected single message to contain augmented instructions and user query, got: %s", singleMsg)
|
|
}
|
|
}
|
|
|
|
func TestToolChoiceHandling(t *testing.T) {
|
|
tools := []Tool{
|
|
{
|
|
Type: "function",
|
|
Function: map[string]interface{}{
|
|
"name": "calculator",
|
|
"description": "Evaluate math expression",
|
|
},
|
|
},
|
|
}
|
|
|
|
// 1. ToolChoice: "none" -> no tool instruction generated
|
|
instrNone := BuildToolInstruction(tools, "none")
|
|
if instrNone != "" {
|
|
t.Errorf("expected empty instruction for tool_choice: none, got: %s", instrNone)
|
|
}
|
|
|
|
// 2. ToolChoice: "required" -> mandatory directive
|
|
instrReq := BuildToolInstruction(tools, "required")
|
|
if !strings.Contains(instrReq, "You MUST call one of the available tools") {
|
|
t.Errorf("expected required directive in instruction, got: %s", instrReq)
|
|
}
|
|
|
|
// 3. ToolChoice: specific function
|
|
instrFn := BuildToolInstruction(tools, map[string]interface{}{
|
|
"type": "function",
|
|
"function": map[string]interface{}{
|
|
"name": "calculator",
|
|
},
|
|
})
|
|
if !strings.Contains(instrFn, "You MUST call the calculator tool") {
|
|
t.Errorf("expected specific function directive in instruction, got: %s", instrFn)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredModelName(t *testing.T) {
|
|
disc := &SpaceDiscovery{
|
|
PrimaryModel: "hy3",
|
|
Models: []string{"hy3", "hunyuan3"},
|
|
}
|
|
|
|
// Without override
|
|
ConfiguredModelName = ""
|
|
list1 := disc.GetModelList()
|
|
if len(list1) < 2 || list1[0].ID != "hy3" {
|
|
t.Errorf("expected default first model hy3, got %+v", list1)
|
|
}
|
|
|
|
// With override
|
|
ConfiguredModelName = "my-custom-hy3"
|
|
defer func() { ConfiguredModelName = "" }()
|
|
|
|
list2 := disc.GetModelList()
|
|
if len(list2) < 3 || list2[0].ID != "my-custom-hy3" {
|
|
t.Errorf("expected first model to be my-custom-hy3, got %+v", list2)
|
|
}
|
|
}
|
|
|
|
func TestMultimodalStateSpaceMockServerCompletion(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/config" {
|
|
cfg := GradioConfigResponse{
|
|
Version: "5.29.0",
|
|
Dependencies: []GradioDependency{
|
|
{
|
|
ID: 6,
|
|
APIName: "chat",
|
|
Inputs: []int{12, 16, 20, 21, 22, 23, 24, 25, 26},
|
|
Outputs: []int{14, 16},
|
|
},
|
|
},
|
|
Components: []GradioComponent{
|
|
{ID: 12, Type: "multimodaltextbox", Props: map[string]interface{}{"label": "Message"}},
|
|
{ID: 16, Type: "state"},
|
|
{ID: 20, Type: "radio", Props: map[string]interface{}{"label": "Model Type", "value": "Chat"}},
|
|
{ID: 21, Type: "checkbox", Props: map[string]interface{}{"label": "Use Internet", "value": false}},
|
|
{ID: 22, Type: "slider", Props: map[string]interface{}{"label": "Max Tokens", "value": 32768}},
|
|
{ID: 23, Type: "slider", Props: map[string]interface{}{"label": "Temperature", "value": 0.8}},
|
|
{ID: 24, Type: "slider", Props: map[string]interface{}{"label": "Top P", "value": 0.95}},
|
|
{ID: 25, Type: "checkbox", Props: map[string]interface{}{"label": "stream", "value": true}},
|
|
{ID: 26, Type: "textbox", Props: map[string]interface{}{"label": "user", "value": "null"}},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(cfg)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
info := GradioAPIInfoResponse{
|
|
NamedEndpoints: map[string]GradioEndpointInfo{
|
|
"/chat": {
|
|
Parameters: []GradioParamInfo{
|
|
{ParameterName: "param_0", Label: "Message", Component: "Multimodaltextbox"},
|
|
{ParameterName: "param_2", Label: "Model Type", Component: "Radio", ParameterDefault: "Chat"},
|
|
{ParameterName: "param_3", Label: "Use Internet", Component: "Checkbox", ParameterDefault: false},
|
|
{ParameterName: "param_4", Label: "Max Tokens", Component: "Slider", ParameterDefault: 32768},
|
|
{ParameterName: "param_5", Label: "Temperature", Component: "Slider", ParameterDefault: 0.8},
|
|
{ParameterName: "param_6", Label: "Top P", Component: "Slider", ParameterDefault: 0.95},
|
|
{ParameterName: "param_7", Label: "stream", Component: "Checkbox", ParameterDefault: true},
|
|
{ParameterName: "param_8", Label: "user", Component: "Textbox", ParameterDefault: "null"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(info)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat" {
|
|
var body struct {
|
|
Data []interface{} `json:"data"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if len(body.Data) != 9 {
|
|
http.Error(w, fmt.Sprintf("expected 9 inputs, got %d", len(body.Data)), http.StatusBadRequest)
|
|
return
|
|
}
|
|
// Check multimodal dict at index 0
|
|
mmDict, ok := body.Data[0].(map[string]interface{})
|
|
if !ok {
|
|
http.Error(w, "input 0 must be multimodal dict", http.StatusBadRequest)
|
|
return
|
|
}
|
|
txt, _ := mmDict["text"].(string)
|
|
evt := "evt_normal"
|
|
if strings.Contains(txt, "get_weather") {
|
|
evt = "evt_tool"
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(GradioJoinResponse{EventID: evt})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat/evt_normal" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"Hello from multimodal space!\", null]\n\n")
|
|
flusher.Flush()
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat/evt_tool" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"```json\\n[{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"location\\\": \\\"Tokyo\\\"}}]\\n```\", null]\n\n")
|
|
flusher.Flush()
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Verify Space Inspection
|
|
disc, err := InspectSpace(gw.client, ts.URL, DefaultUserAgent)
|
|
if err != nil {
|
|
t.Fatalf("InspectSpace failed: %v", err)
|
|
}
|
|
if disc.TotalInputs != 9 {
|
|
t.Errorf("expected TotalInputs 9, got %d", disc.TotalInputs)
|
|
}
|
|
if disc.MessageIndex != 0 {
|
|
t.Errorf("expected MessageIndex 0, got %d", disc.MessageIndex)
|
|
}
|
|
if !disc.MessageIsMultimodal {
|
|
t.Errorf("expected MessageIsMultimodal true")
|
|
}
|
|
if disc.TempIndex != 5 {
|
|
t.Errorf("expected TempIndex 5, got %d", disc.TempIndex)
|
|
}
|
|
if disc.MaxTokensIndex != 4 {
|
|
t.Errorf("expected MaxTokensIndex 4, got %d", disc.MaxTokensIndex)
|
|
}
|
|
if disc.TopPIndex != 6 {
|
|
t.Errorf("expected TopPIndex 6, got %d", disc.TopPIndex)
|
|
}
|
|
if disc.StreamIndex != 7 {
|
|
t.Errorf("expected StreamIndex 7, got %d", disc.StreamIndex)
|
|
}
|
|
|
|
// 2. Normal Chat Completion
|
|
req1 := ChatCompletionRequest{
|
|
Model: "askcyph",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello world"},
|
|
},
|
|
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("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
var resp1 ChatCompletionResponse
|
|
if err := json.NewDecoder(rec1.Body).Decode(&resp1); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if resp1.Choices[0].Message.Content != "Hello from multimodal space!" {
|
|
t.Errorf("unexpected content: %v", resp1.Choices[0].Message.Content)
|
|
}
|
|
|
|
// 3. Tool Calling Chat Completion
|
|
req2 := ChatCompletionRequest{
|
|
Model: "askcyph",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Weather in Tokyo?"},
|
|
},
|
|
Tools: []Tool{
|
|
{
|
|
Type: "function",
|
|
Function: map[string]interface{}{
|
|
"name": "get_weather",
|
|
"description": "Get current weather",
|
|
"parameters": map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"location": map[string]interface{}{"type": "string"},
|
|
},
|
|
"required": []string{"location"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
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("ExecuteChatCompletion tool calling failed: %v", err)
|
|
}
|
|
var resp2 ChatCompletionResponse
|
|
if err := json.NewDecoder(rec2.Body).Decode(&resp2); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if resp2.Choices[0].FinishReason != "tool_calls" {
|
|
t.Fatalf("expected finish_reason 'tool_calls', got %q", resp2.Choices[0].FinishReason)
|
|
}
|
|
if len(resp2.Choices[0].Message.ToolCalls) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d", len(resp2.Choices[0].Message.ToolCalls))
|
|
}
|
|
if resp2.Choices[0].Message.ToolCalls[0].Function.Name != "get_weather" {
|
|
t.Errorf("expected get_weather, got %s", resp2.Choices[0].Message.ToolCalls[0].Function.Name)
|
|
}
|
|
}
|
|
|
|
// TestGradio3DirectPredictProtocol verifies discovery and chat completion against a Gradio 3 space
|
|
// where /gradio_api/info returns 404 and the protocol resolves to /run/predict with tuple pairs history.
|
|
func TestGradio3DirectPredictProtocol(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" || r.URL.Path == "/info" {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
cfg := GradioConfigResponse{
|
|
Version: "3.41.2",
|
|
Mode: "chat_interface",
|
|
Title: "Legacy Gradio 3 Chat",
|
|
Components: []GradioComponent{
|
|
{ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Input"}},
|
|
{ID: 2, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot"}},
|
|
},
|
|
Dependencies: []GradioDependency{
|
|
{
|
|
ID: 0,
|
|
Inputs: []int{1, 2},
|
|
Outputs: []int{2},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(cfg)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/run/predict" {
|
|
var body struct {
|
|
Data []interface{} `json:"data"`
|
|
FnIndex int `json:"fn_index"`
|
|
SessionHash string `json:"session_hash"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if body.FnIndex != 0 {
|
|
http.Error(w, fmt.Sprintf("expected fn_index 0, got %d", body.FnIndex), http.StatusBadRequest)
|
|
return
|
|
}
|
|
msg, _ := body.Data[0].(string)
|
|
reply := "Echo from Gradio 3: " + msg
|
|
respData := map[string]interface{}{
|
|
"data": []interface{}{
|
|
[][]string{
|
|
{msg, reply},
|
|
},
|
|
},
|
|
"is_generating": false,
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(respData)
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
disc := gw.GetDiscovery(ts.URL, DefaultUserAgent)
|
|
|
|
if disc.GradioVersion != "3.41.2" {
|
|
t.Errorf("expected GradioVersion 3.41.2, got %s", disc.GradioVersion)
|
|
}
|
|
if disc.Protocol != "predict" {
|
|
t.Errorf("expected Protocol predict, got %s", disc.Protocol)
|
|
}
|
|
if disc.Flavor != "ChatInterface" {
|
|
t.Errorf("expected Flavor ChatInterface, got %s", disc.Flavor)
|
|
}
|
|
if disc.HistoryFormat != "pairs" {
|
|
t.Errorf("expected HistoryFormat pairs, got %s", disc.HistoryFormat)
|
|
}
|
|
if disc.FnIndex != 0 {
|
|
t.Errorf("expected FnIndex 0, got %d", disc.FnIndex)
|
|
}
|
|
|
|
// Non-streaming completion
|
|
req := ChatCompletionRequest{
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello Gradio 3!"},
|
|
},
|
|
Stream: false,
|
|
}
|
|
b, _ := json.Marshal(req)
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
|
|
rec := httptest.NewRecorder()
|
|
|
|
if err := gw.ExecuteChatCompletion(rec, httpReq, req); err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed on Gradio 3: %v", err)
|
|
}
|
|
|
|
var res ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&res); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if res.Choices[0].Message.Content != "Echo from Gradio 3: Hello Gradio 3!" {
|
|
t.Errorf("unexpected completion content: %v", res.Choices[0].Message.Content)
|
|
}
|
|
}
|
|
|
|
// TestEndpointScoringDisambiguation verifies that chat generation endpoints are selected over
|
|
// UI utility, reset, voting, and feedback endpoints.
|
|
func TestEndpointScoringDisambiguation(t *testing.T) {
|
|
compMap := map[int]GradioComponent{
|
|
1: {ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Message"}},
|
|
2: {ID: 2, Type: "chatbot", Props: map[string]interface{}{"label": "Chat"}},
|
|
3: {ID: 3, Type: "state", Props: map[string]interface{}{"label": "State"}},
|
|
}
|
|
|
|
chatDep := GradioDependency{
|
|
ID: 0,
|
|
Inputs: []int{1, 2, 3},
|
|
Outputs: []int{2},
|
|
Types: GradioDependencyTypes{Generator: true},
|
|
}
|
|
clearDep := GradioDependency{
|
|
ID: 1,
|
|
Inputs: []int{2},
|
|
Outputs: []int{2},
|
|
}
|
|
voteDep := GradioDependency{
|
|
ID: 2,
|
|
Inputs: []int{2},
|
|
Outputs: []int{},
|
|
}
|
|
|
|
chatScore := ScoreCandidateEndpoint("chat", nil, &chatDep, compMap)
|
|
clearScore := ScoreCandidateEndpoint("clear", nil, &clearDep, compMap)
|
|
voteScore := ScoreCandidateEndpoint("vote", nil, &voteDep, compMap)
|
|
resetScore := ScoreCandidateEndpoint("reset_all", nil, &clearDep, compMap)
|
|
|
|
if chatScore <= 0 {
|
|
t.Errorf("expected positive chat score, got %d", chatScore)
|
|
}
|
|
if clearScore >= chatScore {
|
|
t.Errorf("expected chat score > clear score, got chat=%d clear=%d", chatScore, clearScore)
|
|
}
|
|
if voteScore >= chatScore {
|
|
t.Errorf("expected chat score > vote score, got chat=%d vote=%d", chatScore, voteScore)
|
|
}
|
|
if resetScore >= chatScore {
|
|
t.Errorf("expected chat score > reset score, got chat=%d reset=%d", chatScore, resetScore)
|
|
}
|
|
}
|
|
|
|
// TestToolCallingPictureResolution verifies that tool calling mode is accurately resolved
|
|
// based on component topology.
|
|
func TestToolCallingPictureResolution(t *testing.T) {
|
|
// Case A: Native slot
|
|
discNative := NewDefaultSpaceDiscovery("https://tencent-hy3.hf.space")
|
|
discNative.FunctionsJSONIndex = 8
|
|
discNative.SystemIndex = 2
|
|
discNative.HistoryIndex = 1
|
|
discNative.ToolCallMode = "native_slot"
|
|
if discNative.ToolCallMode != "native_slot" {
|
|
t.Errorf("expected native_slot, got %s", discNative.ToolCallMode)
|
|
}
|
|
|
|
// Case B: Dedicated system prompt slot
|
|
discSys := NewDefaultSpaceDiscovery("https://custom-chat.hf.space")
|
|
discSys.FunctionsJSONIndex = -1
|
|
discSys.SystemIndex = 2
|
|
discSys.HistoryIndex = 1
|
|
discSys.ToolCallMode = "prompt_augmented_system"
|
|
if discSys.ToolCallMode != "prompt_augmented_system" {
|
|
t.Errorf("expected prompt_augmented_system, got %s", discSys.ToolCallMode)
|
|
}
|
|
|
|
// Case C: Conversation history (first turn)
|
|
discFirst := NewDefaultSpaceDiscovery("https://chat-only.hf.space")
|
|
discFirst.FunctionsJSONIndex = -1
|
|
discFirst.SystemIndex = -1
|
|
discFirst.HistoryIndex = 1
|
|
discFirst.ToolCallMode = "prompt_augmented_first_turn"
|
|
if discFirst.ToolCallMode != "prompt_augmented_first_turn" {
|
|
t.Errorf("expected prompt_augmented_first_turn, got %s", discFirst.ToolCallMode)
|
|
}
|
|
|
|
// Case D: Single prompt input
|
|
discSingle := NewDefaultSpaceDiscovery("https://single-prompt.hf.space")
|
|
discSingle.FunctionsJSONIndex = -1
|
|
discSingle.SystemIndex = -1
|
|
discSingle.HistoryIndex = -1
|
|
discSingle.ToolCallMode = "prompt_augmented_single_prompt"
|
|
if discSingle.ToolCallMode != "prompt_augmented_single_prompt" {
|
|
t.Errorf("expected prompt_augmented_single_prompt, got %s", discSingle.ToolCallMode)
|
|
}
|
|
}
|
|
|
|
// TestCallToPredictProtocolFallback verifies that if a space reports /call support in /info
|
|
// but /call returns 404 at runtime, gr2gw gracefully falls back to /run/predict.
|
|
func TestCallToPredictProtocolFallback(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
info := GradioAPIInfoResponse{
|
|
NamedEndpoints: map[string]GradioEndpointInfo{
|
|
"/chat": {
|
|
Parameters: []GradioParamInfo{
|
|
{ParameterName: "prompt", Label: "Prompt", Component: "Textbox"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(info)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
cfg := GradioConfigResponse{
|
|
Version: "4.20.0",
|
|
Mode: "interface",
|
|
Components: []GradioComponent{
|
|
{ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Prompt"}},
|
|
},
|
|
Dependencies: []GradioDependency{
|
|
{ID: 0, APIName: "/chat", Inputs: []int{1}, Outputs: []int{1}},
|
|
},
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(cfg)
|
|
return
|
|
}
|
|
|
|
// /call/chat returns 404 (endpoint disabled or unsupported)
|
|
if strings.HasPrefix(r.URL.Path, "/gradio_api/call/") || strings.HasPrefix(r.URL.Path, "/call/") {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
|
|
// Fallback /run/predict works
|
|
if r.URL.Path == "/run/predict" || r.URL.Path == "/gradio_api/run/predict" {
|
|
var body struct {
|
|
Data []interface{} `json:"data"`
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&body)
|
|
msg, _ := body.Data[0].(string)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"data": []interface{}{"Fallback response for: " + msg},
|
|
"is_generating": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
req := ChatCompletionRequest{
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Testing fallback"},
|
|
},
|
|
Stream: false,
|
|
}
|
|
b, _ := json.Marshal(req)
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
|
|
rec := httptest.NewRecorder()
|
|
|
|
if err := gw.ExecuteChatCompletion(rec, httpReq, req); err != nil {
|
|
t.Fatalf("ExecuteChatCompletion fallback failed: %v", err)
|
|
}
|
|
|
|
var res ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&res); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if res.Choices[0].Message.Content != "Fallback response for: Testing fallback" {
|
|
t.Errorf("unexpected content: %v", res.Choices[0].Message.Content)
|
|
}
|
|
}
|
|
|
|
// TestEndpointScoringUserHandlerDepPenalization verifies that user submission helper
|
|
// endpoints with shared input/output textboxes are penalized while streaming generators
|
|
// and bot endpoints like Chat_Message are favored.
|
|
func TestEndpointScoringUserHandlerDepPenalization(t *testing.T) {
|
|
compMap := map[int]GradioComponent{
|
|
5: {ID: 5, Type: "state", Props: map[string]interface{}{"label": "State"}},
|
|
6: {ID: 6, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot"}},
|
|
8: {ID: 8, Type: "textbox", Props: map[string]interface{}{"label": "Message"}},
|
|
27: {ID: 27, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot 2"}},
|
|
29: {ID: 29, Type: "textbox", Props: map[string]interface{}{"label": "Link"}},
|
|
30: {ID: 30, Type: "textbox", Props: map[string]interface{}{"label": "User Message"}},
|
|
}
|
|
|
|
// user: inputs [8, 6], outputs [8, 6] (clears textbox 8)
|
|
userDep := GradioDependency{
|
|
ID: 2,
|
|
Inputs: []int{8, 6},
|
|
Outputs: []int{8, 6},
|
|
Types: GradioDependencyTypes{Generator: false},
|
|
}
|
|
|
|
// user2: inputs [30, 27, 29], outputs [30, 27, 29] (clears textboxes 30 and 29)
|
|
user2Dep := GradioDependency{
|
|
ID: 23,
|
|
Inputs: []int{30, 27, 29},
|
|
Outputs: []int{30, 27, 29},
|
|
Types: GradioDependencyTypes{Generator: false},
|
|
}
|
|
|
|
// Chat_Message: inputs [6, 5], outputs [6, 5], generator: true
|
|
chatMsgDep := GradioDependency{
|
|
ID: 3,
|
|
Inputs: []int{6, 5},
|
|
Outputs: []int{6, 5},
|
|
Types: GradioDependencyTypes{Generator: true},
|
|
}
|
|
|
|
scoreUser := ScoreCandidateEndpoint("/user", nil, &userDep, compMap)
|
|
scoreUser2 := ScoreCandidateEndpoint("/user2", nil, &user2Dep, compMap)
|
|
scoreChatMsg := ScoreCandidateEndpoint("/Chat_Message", nil, &chatMsgDep, compMap)
|
|
|
|
if scoreChatMsg <= 0 {
|
|
t.Errorf("expected positive score for Chat_Message, got %d", scoreChatMsg)
|
|
}
|
|
if scoreChatMsg <= scoreUser {
|
|
t.Errorf("expected Chat_Message score > user score, got chatMsg=%d user=%d", scoreChatMsg, scoreUser)
|
|
}
|
|
if scoreChatMsg <= scoreUser2 {
|
|
t.Errorf("expected Chat_Message score > user2 score, got chatMsg=%d user2=%d", scoreChatMsg, scoreUser2)
|
|
}
|
|
}
|
|
|
|
// TestHistoryOnlyEndpointPayloadBuilding verifies that endpoints with no separate
|
|
// message textbox (MessageIndex == -1, HistoryIndex == 0) correctly append the user prompt
|
|
// into the history array without overwriting other slots.
|
|
func TestHistoryOnlyEndpointPayloadBuilding(t *testing.T) {
|
|
gw := NewGradioGateway("https://test-space.hf.space", "", 10*time.Second)
|
|
|
|
// 1. Gradio 6 messages format
|
|
discG6 := NewDefaultSpaceDiscovery("https://test-space.hf.space")
|
|
discG6.TotalInputs = 2
|
|
discG6.MessageIndex = -1
|
|
discG6.HistoryIndex = 0
|
|
discG6.HistoryFormat = "gradio_messages"
|
|
|
|
req := ChatCompletionRequest{
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello world"},
|
|
},
|
|
}
|
|
|
|
data, err := gw.BuildGradioPayload(discG6, req)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
if len(data) != 2 {
|
|
t.Fatalf("expected 2 inputs, got %d", len(data))
|
|
}
|
|
|
|
gMsgs, ok := data[0].([]map[string]interface{})
|
|
if !ok {
|
|
t.Fatalf("expected []map[string]interface{} for gradio_messages, got %T", data[0])
|
|
}
|
|
if len(gMsgs) != 1 {
|
|
t.Fatalf("expected 1 message in history, got %d", len(gMsgs))
|
|
}
|
|
if gMsgs[0]["role"] != "user" {
|
|
t.Errorf("expected user role, got %v", gMsgs[0]["role"])
|
|
}
|
|
contents, ok := gMsgs[0]["content"].([]map[string]string)
|
|
if !ok || len(contents) != 1 || contents[0]["text"] != "Hello world" {
|
|
t.Errorf("unexpected content structure: %v", gMsgs[0]["content"])
|
|
}
|
|
if data[1] != nil {
|
|
t.Errorf("expected slot 1 (state) to remain nil, got %v", data[1])
|
|
}
|
|
|
|
// 2. Pairs format with MessageIndex == -1
|
|
discPairs := NewDefaultSpaceDiscovery("https://test-space.hf.space")
|
|
discPairs.TotalInputs = 1
|
|
discPairs.MessageIndex = -1
|
|
discPairs.HistoryIndex = 0
|
|
discPairs.HistoryFormat = "pairs"
|
|
|
|
dataPairs, err := gw.BuildGradioPayload(discPairs, req)
|
|
if err != nil {
|
|
t.Fatalf("BuildGradioPayload failed: %v", err)
|
|
}
|
|
pairs, ok := dataPairs[0].([][]string)
|
|
if !ok || len(pairs) != 1 {
|
|
t.Fatalf("expected 1 pair in history, got %T (%v)", dataPairs[0], dataPairs[0])
|
|
}
|
|
if pairs[0][0] != "Hello world" {
|
|
t.Errorf("expected user prompt in pair[0], got %v", pairs[0][0])
|
|
}
|
|
}
|
|
|
|
// TestInspectSpaceBlocksChatbotStateResolution verifies InspectSpace selects /Chat_Message
|
|
// when presented with a Blocks layout containing user handlers and generator functions.
|
|
func TestInspectSpaceBlocksChatbotStateResolution(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"named_endpoints": map[string]interface{}{
|
|
"/user": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{"parameter_name": "user_message", "component": "Textbox", "label": "msg"},
|
|
{"parameter_name": "history", "component": "Chatbot", "label": "chat"},
|
|
},
|
|
},
|
|
"/Chat_Message": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{
|
|
"parameter_name": "history",
|
|
"component": "Chatbot",
|
|
"label": "chat",
|
|
"type": map[string]interface{}{"title": "ChatbotDataMessages"},
|
|
"python_type": map[string]interface{}{"type": "dict(text: str)"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"version": "6.20.0",
|
|
"components": []map[string]interface{}{
|
|
{"id": 5, "type": "state", "props": map[string]interface{}{"label": "State"}},
|
|
{"id": 6, "type": "chatbot", "props": map[string]interface{}{"label": "Chatbot"}},
|
|
{"id": 8, "type": "textbox", "props": map[string]interface{}{"label": "Message"}},
|
|
},
|
|
"dependencies": []map[string]interface{}{
|
|
{
|
|
"id": 2,
|
|
"api_name": "user",
|
|
"inputs": []int{8, 6},
|
|
"outputs": []int{8, 6},
|
|
"types": map[string]interface{}{"generator": false},
|
|
},
|
|
{
|
|
"id": 3,
|
|
"api_name": "Chat_Message",
|
|
"inputs": []int{6, 5},
|
|
"outputs": []int{6, 5},
|
|
"types": map[string]interface{}{"generator": true},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
client := &http.Client{Timeout: 5 * time.Second}
|
|
disc, err := InspectSpace(client, ts.URL, "test-agent")
|
|
if err != nil {
|
|
t.Fatalf("InspectSpace failed: %v", err)
|
|
}
|
|
|
|
if disc.Endpoint != "/Chat_Message" {
|
|
t.Errorf("expected resolved endpoint /Chat_Message, got %s", disc.Endpoint)
|
|
}
|
|
if disc.FnIndex != 3 {
|
|
t.Errorf("expected FnIndex 3, got %d", disc.FnIndex)
|
|
}
|
|
if disc.MessageIndex != -1 {
|
|
t.Errorf("expected MessageIndex -1, got %d", disc.MessageIndex)
|
|
}
|
|
if disc.HistoryIndex != 0 {
|
|
t.Errorf("expected HistoryIndex 0, got %d", disc.HistoryIndex)
|
|
}
|
|
if disc.HistoryFormat != "gradio_messages" {
|
|
t.Errorf("expected HistoryFormat gradio_messages, got %s", disc.HistoryFormat)
|
|
}
|
|
// TotalInputs should be trimmed to 1 to match the exposed API parameters and avoid padding unexposed state
|
|
if disc.TotalInputs != 1 {
|
|
t.Errorf("expected TotalInputs 1 (unexposed trailing state trimmed), got %d", disc.TotalInputs)
|
|
}
|
|
}
|
|
|
|
// TestExecuteCallCompletionTrailingStateTrim verifies that call requests
|
|
// omit session_hash (to avoid breaking SSE streams) and trim unexposed trailing state inputs from the payload.
|
|
func TestExecuteCallCompletionTrailingStateTrim(t *testing.T) {
|
|
var receivedSessionHash string
|
|
var receivedDataLen int
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"named_endpoints": map[string]interface{}{
|
|
"/Chat_Message": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{
|
|
"parameter_name": "history",
|
|
"component": "Chatbot",
|
|
"type": map[string]interface{}{"title": "ChatbotDataMessages"},
|
|
"python_type": map[string]interface{}{"type": "dict(text: str)"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"version": "6.20.0",
|
|
"components": []map[string]interface{}{
|
|
{"id": 5, "type": "state"},
|
|
{"id": 6, "type": "chatbot"},
|
|
},
|
|
"dependencies": []map[string]interface{}{
|
|
{
|
|
"id": 3,
|
|
"api_name": "Chat_Message",
|
|
"inputs": []int{6, 5},
|
|
"outputs": []int{6, 5},
|
|
"types": map[string]interface{}{"generator": true},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/Chat_Message" {
|
|
var body struct {
|
|
Data []interface{} `json:"data"`
|
|
SessionHash string `json:"session_hash"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
receivedSessionHash = body.SessionHash
|
|
receivedDataLen = len(body.Data)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"event_id": "evt-12345",
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/Chat_Message/evt-12345" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.WriteHeader(http.StatusOK)
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: generating\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "event: complete\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
req := ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "hello"},
|
|
},
|
|
}
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
rec := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec, httpReq, req)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if receivedSessionHash != "" {
|
|
t.Errorf("expected empty session_hash in call payload, got %s", receivedSessionHash)
|
|
}
|
|
if receivedDataLen != 1 {
|
|
t.Errorf("expected call data length 1 (unexposed state trimmed), got %d", receivedDataLen)
|
|
}
|
|
}
|
|
|
|
// TestGradio6CallV2Protocol verifies that Gradio 6 spaces resolve protocol call_v2,
|
|
// submit named JSON parameters to /call/v2/{endpoint}, and stream from /call/{endpoint}/{event_id}.
|
|
func TestGradio6CallV2Protocol(t *testing.T) {
|
|
var receivedBody map[string]interface{}
|
|
var callEndpointReached string
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"named_endpoints": map[string]interface{}{
|
|
"/answer": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{
|
|
"parameter_name": "prompt",
|
|
"component": "Textbox",
|
|
"label": "Prompt",
|
|
},
|
|
},
|
|
"code_snippets": map[string]interface{}{
|
|
"bash": "curl -X POST http://localhost:7860/gradio_api/call/v2/answer -s -H \"Content-Type: application/json\" -d '{\"prompt\": \"Hello!!\"}' | awk -F'\"' '{ print $4}' | read EVENT_ID; curl -N http://localhost:7860/gradio_api/call/answer/$EVENT_ID",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"version": "6.13.0",
|
|
"components": []map[string]interface{}{
|
|
{"id": 1, "type": "textbox", "props": map[string]interface{}{"label": "Prompt"}},
|
|
},
|
|
"dependencies": []map[string]interface{}{
|
|
{
|
|
"id": 0,
|
|
"api_name": "answer",
|
|
"inputs": []int{1},
|
|
"outputs": []int{1},
|
|
"types": map[string]interface{}{"generator": true},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/v2/answer" {
|
|
callEndpointReached = r.URL.Path
|
|
if err := json.NewDecoder(r.Body).Decode(&receivedBody); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"event_id": "evt-v2-777",
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/answer/evt-v2-777" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.WriteHeader(http.StatusOK)
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"Hello from Gradio 6 v2!\"]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Verify inspection resolved call_v2 protocol
|
|
disc, err := InspectSpace(gw.client, ts.URL, DefaultUserAgent)
|
|
if err != nil {
|
|
t.Fatalf("InspectSpace failed: %v", err)
|
|
}
|
|
if disc.Protocol != "call_v2" {
|
|
t.Errorf("expected Protocol call_v2, got %s", disc.Protocol)
|
|
}
|
|
if disc.Endpoint != "/answer" {
|
|
t.Errorf("expected Endpoint /answer, got %s", disc.Endpoint)
|
|
}
|
|
|
|
// 2. Execute chat completion
|
|
req := ChatCompletionRequest{
|
|
Model: "default",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello test"},
|
|
},
|
|
}
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
rec := httptest.NewRecorder()
|
|
|
|
err = gw.ExecuteChatCompletion(rec, httpReq, req)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if callEndpointReached != "/gradio_api/call/v2/answer" {
|
|
t.Errorf("expected call to /gradio_api/call/v2/answer, got %s", callEndpointReached)
|
|
}
|
|
if promptVal, ok := receivedBody["prompt"].(string); !ok || promptVal != "Hello test" {
|
|
t.Errorf("expected named param 'prompt' with value 'Hello test', got %v", receivedBody)
|
|
}
|
|
|
|
var resp ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if len(resp.Choices) == 0 || resp.Choices[0].Message.Content != "Hello from Gradio 6 v2!" {
|
|
t.Errorf("unexpected content: %v", resp.Choices[0].Message.Content)
|
|
}
|
|
}
|
|
|
|
// TestGradio6WithoutSnippetResolvesToCall verifies that Gradio 6 spaces without code snippets
|
|
// resolve protocol call (not call_v2), submit standard data array, and omit session_hash.
|
|
func TestGradio6WithoutSnippetResolvesToCall(t *testing.T) {
|
|
var callEndpointReached string
|
|
var receivedPayload map[string]interface{}
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"named_endpoints": map[string]interface{}{
|
|
"/predict": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{
|
|
"parameter_name": "message",
|
|
"component": "Textbox",
|
|
"label": "",
|
|
},
|
|
},
|
|
// No code_snippets present
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"version": "6.5.1",
|
|
"components": []map[string]interface{}{
|
|
{"id": 11, "type": "textbox", "props": map[string]interface{}{"label": ""}},
|
|
{"id": 15, "type": "state"},
|
|
},
|
|
"dependencies": []map[string]interface{}{
|
|
{
|
|
"id": 6,
|
|
"api_name": "predict",
|
|
"inputs": []int{11, 15},
|
|
"outputs": []int{11, 15},
|
|
"types": map[string]interface{}{"generator": true},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/predict" {
|
|
callEndpointReached = r.URL.Path
|
|
if err := json.NewDecoder(r.Body).Decode(&receivedPayload); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"event_id": "evt-call-651",
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/predict/evt-call-651" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.WriteHeader(http.StatusOK)
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: complete\ndata: [\"Hello from standard Gradio 6 call!\"]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Verify inspection resolved protocol "call", not "call_v2"
|
|
disc, err := InspectSpace(gw.client, ts.URL, DefaultUserAgent)
|
|
if err != nil {
|
|
t.Fatalf("InspectSpace failed: %v", err)
|
|
}
|
|
if disc.Protocol != "call" {
|
|
t.Errorf("expected Protocol call, got %s", disc.Protocol)
|
|
}
|
|
if disc.Endpoint != "/predict" {
|
|
t.Errorf("expected Endpoint /predict, got %s", disc.Endpoint)
|
|
}
|
|
if disc.TotalInputs != 1 {
|
|
t.Errorf("expected TotalInputs 1 (trimmed trailing state), got %d", disc.TotalInputs)
|
|
}
|
|
if disc.RawTotalInputs != 2 {
|
|
t.Errorf("expected RawTotalInputs 2, got %d", disc.RawTotalInputs)
|
|
}
|
|
|
|
// 2. Execute chat completion
|
|
req := ChatCompletionRequest{
|
|
Model: "default",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello test"},
|
|
},
|
|
}
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
rec := httptest.NewRecorder()
|
|
|
|
err = gw.ExecuteChatCompletion(rec, httpReq, req)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String())
|
|
}
|
|
if callEndpointReached != "/gradio_api/call/predict" {
|
|
t.Errorf("expected call to /gradio_api/call/predict, got %s", callEndpointReached)
|
|
}
|
|
if _, hasHash := receivedPayload["session_hash"]; hasHash {
|
|
t.Errorf("expected no session_hash in call payload")
|
|
}
|
|
dataSlice, ok := receivedPayload["data"].([]interface{})
|
|
if !ok || len(dataSlice) != 1 || dataSlice[0] != "Hello test" {
|
|
t.Errorf("expected data payload [\"Hello test\"], got %v", receivedPayload["data"])
|
|
}
|
|
}
|
|
|
|
// TestGradioQueueProtocol verifies direct queue join and queue data streaming.
|
|
func TestGradioQueueProtocol(t *testing.T) {
|
|
var joinReached bool
|
|
var receivedFnIndex int
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/queue/join" {
|
|
joinReached = true
|
|
var body struct {
|
|
Data []interface{} `json:"data"`
|
|
FnIndex int `json:"fn_index"`
|
|
SessionHash string `json:"session_hash"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
receivedFnIndex = body.FnIndex
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "evt-q-1"})
|
|
return
|
|
}
|
|
|
|
if strings.HasPrefix(r.URL.Path, "/gradio_api/queue/data") {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "data: {\"msg\":\"process_generating\",\"output\":{\"data\":[\"Hello\",null]},\"success\":true}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "data: {\"msg\":\"process_generating\",\"output\":{\"data\":[[[\"append\",[],\" world!\"]]]},\"success\":true}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "data: {\"msg\":\"process_completed\",\"output\":{\"data\":[\"Hello world!\",null]},\"success\":true}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "data: {\"msg\":\"close_stream\"}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
disc := &SpaceDiscovery{
|
|
SpaceURL: ts.URL,
|
|
APIPrefix: "/gradio_api",
|
|
Protocol: "queue",
|
|
FnIndex: 7,
|
|
TotalInputs: 1,
|
|
RawTotalInputs: 2,
|
|
DefaultInputs: []interface{}{nil, nil},
|
|
ParamMappings: []SpaceParamMapping{
|
|
{InputIndex: 0, ParamType: "message"},
|
|
{InputIndex: 1, ParamType: "state"},
|
|
},
|
|
}
|
|
gw.discoveries[ts.URL] = disc
|
|
disc.LastDiscovered = time.Now()
|
|
|
|
req := ChatCompletionRequest{
|
|
Model: "default",
|
|
Messages: []ChatMessage{{Role: "user", Content: "Hi"}},
|
|
}
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
rec := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec, httpReq, req)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
if !joinReached {
|
|
t.Errorf("expected /gradio_api/queue/join to be reached")
|
|
}
|
|
if receivedFnIndex != 7 {
|
|
t.Errorf("expected fn_index 7, got %d", receivedFnIndex)
|
|
}
|
|
|
|
var resp ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if len(resp.Choices) == 0 || resp.Choices[0].Message.Content != "Hello world!" {
|
|
t.Errorf("expected content 'Hello world!', got %v", resp.Choices)
|
|
}
|
|
}
|
|
|
|
// TestGradioCallFallbackToQueue verifies that when /call returns a protocol or input binding error,
|
|
// the gateway automatically falls back to /queue/join and completes successfully.
|
|
func TestGradioCallFallbackToQueue(t *testing.T) {
|
|
var callReached, queueReached bool
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/call/v2/lisa_stream" {
|
|
callReached = true
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "call-err-1"})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/lisa_stream/call-err-1" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: error\ndata: {\"error\": null}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/queue/join" {
|
|
queueReached = true
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "queue-succ-1"})
|
|
return
|
|
}
|
|
|
|
if strings.HasPrefix(r.URL.Path, "/gradio_api/queue/data") {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "data: {\"msg\":\"process_completed\",\"output\":{\"data\":[\"Recovered via queue!\",null]},\"success\":true}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "data: {\"msg\":\"close_stream\"}\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
disc := &SpaceDiscovery{
|
|
SpaceURL: ts.URL,
|
|
APIPrefix: "/gradio_api",
|
|
Endpoint: "/lisa_stream",
|
|
CleanEndpoint: "lisa_stream",
|
|
Protocol: "call_v2",
|
|
FnIndex: 7,
|
|
TotalInputs: 1,
|
|
RawTotalInputs: 2,
|
|
DefaultInputs: []interface{}{nil, nil},
|
|
ParamMappings: []SpaceParamMapping{
|
|
{InputIndex: 0, ParamName: "message", ParamType: "message"},
|
|
},
|
|
}
|
|
gw.discoveries[ts.URL] = disc
|
|
disc.LastDiscovered = time.Now()
|
|
|
|
req := ChatCompletionRequest{
|
|
Model: "default",
|
|
Messages: []ChatMessage{{Role: "user", Content: "Test fallback"}},
|
|
}
|
|
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
rec := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(rec, httpReq, req)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion failed: %v", err)
|
|
}
|
|
if !callReached {
|
|
t.Errorf("expected /call to be attempted first")
|
|
}
|
|
if !queueReached {
|
|
t.Errorf("expected fallback to /queue/join")
|
|
}
|
|
|
|
var resp ChatCompletionResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
|
t.Fatalf("failed to decode response: %v", err)
|
|
}
|
|
if len(resp.Choices) == 0 || resp.Choices[0].Message.Content != "Recovered via queue!" {
|
|
t.Errorf("expected content 'Recovered via queue!', got %v", resp.Choices)
|
|
}
|
|
if disc.Protocol != "queue" {
|
|
t.Errorf("expected disc.Protocol to be switched to 'queue', got %s", disc.Protocol)
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestParseGradioStreamOutputHunyuanTuple(t *testing.T) {
|
|
// Test Hunyuan 3 4-element tuple during reasoning phase: [content, reasoning, tools, history]
|
|
raw1 := `[["", "Thinking through the query...", [], [{"role": "user", "content": "hi"}]]]`
|
|
f1 := ParseGradioStreamOutput(raw1)
|
|
if !f1.OK {
|
|
t.Fatalf("expected f1.OK to be true")
|
|
}
|
|
if f1.Content != "" {
|
|
t.Errorf("expected empty content during reasoning phase, got %q", f1.Content)
|
|
}
|
|
if f1.Reasoning != "Thinking through the query..." {
|
|
t.Errorf("expected reasoning 'Thinking through the query...', got %q", f1.Reasoning)
|
|
}
|
|
|
|
// Test Hunyuan 3 completion with both content and reasoning
|
|
raw2 := `[["Hello! How can I help?", "Thinking through the query...", [], [{"role": "user", "content": "hi"}]]]`
|
|
f2 := ParseGradioStreamOutput(raw2)
|
|
if !f2.OK {
|
|
t.Fatalf("expected f2.OK to be true")
|
|
}
|
|
if f2.Content != "Hello! How can I help?" {
|
|
t.Errorf("expected content 'Hello! How can I help?', got %q", f2.Content)
|
|
}
|
|
if f2.Reasoning != "Thinking through the query..." {
|
|
t.Errorf("expected reasoning 'Thinking through the query...', got %q", f2.Reasoning)
|
|
}
|
|
|
|
// Test 2-element tuple [content, reasoning]
|
|
raw3 := `[["", "Still thinking..."]]`
|
|
f3 := ParseGradioStreamOutput(raw3)
|
|
if !f3.OK {
|
|
t.Fatalf("expected f3.OK to be true")
|
|
}
|
|
if f3.Content != "" {
|
|
t.Errorf("expected empty content, got %q", f3.Content)
|
|
}
|
|
if f3.Reasoning != "Still thinking..." {
|
|
t.Errorf("expected reasoning 'Still thinking...', got %q", f3.Reasoning)
|
|
}
|
|
}
|
|
|
|
func TestGradioDiffDeltaReasoningContentSeparation(t *testing.T) {
|
|
// Diff appending to reasoning at path [1]
|
|
raw1 := `[[["append", [1], "reasoning delta "], ["append", [3, 1, "reasoning_content"], "reasoning delta "]]]`
|
|
f1 := ParseGradioStreamOutput(raw1)
|
|
if !f1.OK || !f1.IsDelta {
|
|
t.Fatalf("expected f1 to be valid delta frame")
|
|
}
|
|
if f1.Content != "" {
|
|
t.Errorf("expected empty content, got %q", f1.Content)
|
|
}
|
|
if f1.Reasoning != "reasoning delta " {
|
|
t.Errorf("expected reasoning delta 'reasoning delta ', got %q", f1.Reasoning)
|
|
}
|
|
|
|
// Diff appending to content at path [0]
|
|
raw2 := `[[["append", [0], "content delta "], ["append", [3, 1, "content"], "content delta "]]]`
|
|
f2 := ParseGradioStreamOutput(raw2)
|
|
if !f2.OK || !f2.IsDelta {
|
|
t.Fatalf("expected f2 to be valid delta frame")
|
|
}
|
|
if f2.Content != "content delta " {
|
|
t.Errorf("expected content delta 'content delta ', got %q", f2.Content)
|
|
}
|
|
if f2.Reasoning != "" {
|
|
t.Errorf("expected empty reasoning, got %q", f2.Reasoning)
|
|
}
|
|
}
|
|
|
|
func TestHunyuan3CallAndStreamingNoInterleaving(t *testing.T) {
|
|
var callPayloadReceived map[string]interface{}
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/info" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"named_endpoints": map[string]interface{}{
|
|
"/chat": map[string]interface{}{
|
|
"parameters": []map[string]interface{}{
|
|
{"parameter_name": "message", "component": "Api"},
|
|
{"parameter_name": "system_prompt", "component": "Api"},
|
|
{"parameter_name": "history", "component": "Api"},
|
|
{"parameter_name": "think_level", "component": "Api"},
|
|
{"parameter_name": "temperature", "component": "Api"},
|
|
{"parameter_name": "max_tokens", "component": "Api"},
|
|
{"parameter_name": "top_p", "component": "Api"},
|
|
{"parameter_name": "preserved_thinking", "component": "Api"},
|
|
{"parameter_name": "functions_json_str", "component": "Api"},
|
|
},
|
|
"code_snippets": map[string]interface{}{
|
|
"bash": "curl -X POST http://localhost:7860/gradio_api/call/chat -s -H \"Content-Type: application/json\" -d '{\"data\": [\"...\", \"\", null, \"high\", null, 0, 0, null, \"\"]}' | awk -F'\"' '{ print $4}' | read EVENT_ID; curl -N http://localhost:7860/gradio_api/call/chat/$EVENT_ID",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/config" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"version": "6.12.0",
|
|
"components": []map[string]interface{}{
|
|
{"id": 30, "type": "api", "props": map[string]interface{}{"label": "1st"}},
|
|
{"id": 31, "type": "api", "props": map[string]interface{}{"label": "2nd"}},
|
|
{"id": 32, "type": "api", "props": map[string]interface{}{"label": "3rd"}},
|
|
{"id": 33, "type": "api", "props": map[string]interface{}{"label": "4th"}},
|
|
{"id": 34, "type": "api", "props": map[string]interface{}{"label": "5th"}},
|
|
{"id": 35, "type": "api", "props": map[string]interface{}{"label": "6th"}},
|
|
{"id": 36, "type": "api", "props": map[string]interface{}{"label": "7th"}},
|
|
{"id": 37, "type": "api", "props": map[string]interface{}{"label": "8th"}},
|
|
{"id": 38, "type": "api", "props": map[string]interface{}{"label": "9th"}},
|
|
{"id": 39, "type": "api", "props": map[string]interface{}{"label": "out"}},
|
|
},
|
|
"dependencies": []map[string]interface{}{
|
|
{
|
|
"id": 8,
|
|
"api_name": "chat",
|
|
"inputs": []int{30, 31, 32, 33, 34, 35, 36, 37, 38},
|
|
"outputs": []int{39},
|
|
"types": map[string]interface{}{"generator": true},
|
|
},
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat" {
|
|
if err := json.NewDecoder(r.Body).Decode(&callPayloadReceived); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"event_id": "hy3-event-999",
|
|
})
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/gradio_api/call/chat/hy3-event-999" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.WriteHeader(http.StatusOK)
|
|
flusher, _ := w.(http.Flusher)
|
|
|
|
// Event 1: Thinking chunk 1
|
|
fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Thinking about the answer.\", [], []]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
|
|
// Event 2: Thinking chunk 2
|
|
fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Thinking about the answer. Planning response.\", [], []]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
|
|
// Event 3: Content start
|
|
fmt.Fprintf(w, "event: generating\ndata: [[\"Hello \", \"Thinking about the answer. Planning response.\", [], []]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
|
|
// Event 4: Content complete
|
|
fmt.Fprintf(w, "event: complete\ndata: [[\"Hello world!\", \"Thinking about the answer. Planning response.\", [], []]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
|
|
|
// 1. Verify Space Discovery for Hunyuan 3 space
|
|
disc := gw.GetDiscovery(ts.URL, "test-ua")
|
|
if !disc.IsHunyuan3 {
|
|
t.Errorf("expected IsHunyuan3 to be true")
|
|
}
|
|
if disc.Protocol != "call" {
|
|
t.Errorf("expected Protocol 'call', got %q", disc.Protocol)
|
|
}
|
|
if disc.ThinkLevelIndex != 3 {
|
|
t.Errorf("expected ThinkLevelIndex 3, got %d", disc.ThinkLevelIndex)
|
|
}
|
|
if disc.PreservedThinkingIndex != 7 {
|
|
t.Errorf("expected PreservedThinkingIndex 7, got %d", disc.PreservedThinkingIndex)
|
|
}
|
|
if disc.FunctionsJSONIndex != 8 {
|
|
t.Errorf("expected FunctionsJSONIndex 8, got %d", disc.FunctionsJSONIndex)
|
|
}
|
|
|
|
// 2. Verify Streaming Request: NO interleaving between reasoning and content
|
|
reqStream := ChatCompletionRequest{
|
|
Model: "hy3",
|
|
Stream: true,
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
}
|
|
httpReqStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
recStream := httptest.NewRecorder()
|
|
|
|
err := gw.ExecuteChatCompletion(recStream, httpReqStream, reqStream)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion streaming failed: %v", err)
|
|
}
|
|
|
|
// Verify call payload has no session_hash for Hunyuan 3
|
|
if _, hasHash := callPayloadReceived["session_hash"]; hasHash {
|
|
t.Errorf("session_hash should NOT be present in call payload for Hunyuan 3")
|
|
}
|
|
|
|
var streamedReasoning strings.Builder
|
|
var streamedContent strings.Builder
|
|
|
|
lines := strings.Split(recStream.Body.String(), "\n")
|
|
for _, line := range lines {
|
|
line = strings.TrimSpace(line)
|
|
if strings.HasPrefix(line, "data: ") && line != "data: [DONE]" {
|
|
var chunk StreamResponse
|
|
if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &chunk); err == nil {
|
|
if len(chunk.Choices) > 0 {
|
|
delta := chunk.Choices[0].Delta
|
|
if delta.ReasoningContent != "" {
|
|
if delta.Content != "" {
|
|
t.Errorf("interleaving detected: chunk has both reasoning and content: %+v", delta)
|
|
}
|
|
streamedReasoning.WriteString(delta.ReasoningContent)
|
|
}
|
|
if delta.Content != "" {
|
|
if strings.Contains(delta.Content, "Thinking") {
|
|
t.Errorf("leakage detected: content chunk contains reasoning: %q", delta.Content)
|
|
}
|
|
streamedContent.WriteString(delta.Content)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if streamedReasoning.String() != "Thinking about the answer. Planning response." {
|
|
t.Errorf("expected full reasoning 'Thinking about the answer. Planning response.', got %q", streamedReasoning.String())
|
|
}
|
|
if streamedContent.String() != "Hello world!" {
|
|
t.Errorf("expected full content 'Hello world!', got %q", streamedContent.String())
|
|
}
|
|
|
|
// 3. Verify Non-Streaming Request: separate content and reasoning_content
|
|
reqNonStream := ChatCompletionRequest{
|
|
Model: "hy3",
|
|
Stream: false,
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
}
|
|
httpReqNonStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
recNonStream := httptest.NewRecorder()
|
|
|
|
err = gw.ExecuteChatCompletion(recNonStream, httpReqNonStream, reqNonStream)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteChatCompletion non-streaming failed: %v", err)
|
|
}
|
|
|
|
var nonStreamResp ChatCompletionResponse
|
|
if err := json.NewDecoder(recNonStream.Body).Decode(&nonStreamResp); err != nil {
|
|
t.Fatalf("failed to decode non-stream response: %v", err)
|
|
}
|
|
if len(nonStreamResp.Choices) == 0 {
|
|
t.Fatalf("expected at least 1 choice")
|
|
}
|
|
choice := nonStreamResp.Choices[0]
|
|
if choice.Message.ReasoningContent != "Thinking about the answer. Planning response." {
|
|
t.Errorf("expected reasoning_content 'Thinking about the answer. Planning response.', got %q", choice.Message.ReasoningContent)
|
|
}
|
|
if choice.Message.GetContentString() != "Hello world!" {
|
|
t.Errorf("expected content 'Hello world!', got %q", choice.Message.GetContentString())
|
|
}
|
|
}
|
|
|