omniforge support added
This commit is contained in:
Luxferre
+676
@@ -4058,4 +4058,680 @@ func TestQueueToolCallingStreamingParity(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerFlavorMultimodalDictMessageAndHistory(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/config" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"version": "6.20.0",
|
||||
"api_prefix": "/gradio_api",
|
||||
"mode": "server",
|
||||
"components": []map[string]interface{}{
|
||||
{
|
||||
"id": 1,
|
||||
"type": "api",
|
||||
"props": map[string]interface{}{
|
||||
"label": "1st",
|
||||
"_api_info": map[string]interface{}{"type": "object"},
|
||||
},
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "api",
|
||||
"props": map[string]interface{}{
|
||||
"label": "2nd",
|
||||
"_api_info": map[string]interface{}{"type": "array"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"dependencies": []map[string]interface{}{
|
||||
{
|
||||
"id": 0,
|
||||
"api_name": "chat",
|
||||
"inputs": []int{1, 2},
|
||||
"outputs": []int{1},
|
||||
"types": map[string]interface{}{"generator": true},
|
||||
},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if r.URL.Path == "/info" || 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",
|
||||
"type": map[string]interface{}{"type": "object"},
|
||||
"python_type": map[string]interface{}{"type": "dict(str, Any)"},
|
||||
},
|
||||
{
|
||||
"parameter_name": "history",
|
||||
"component": "Api",
|
||||
"type": map[string]interface{}{"type": "array"},
|
||||
"python_type": map[string]interface{}{"type": "list[Any]"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
disc, err := InspectSpace(http.DefaultClient, ts.URL, "test-agent")
|
||||
if err != nil {
|
||||
t.Fatalf("InspectSpace failed: %v", err)
|
||||
}
|
||||
|
||||
if disc.Flavor != "server" {
|
||||
t.Errorf("expected Flavor 'server', got %q", disc.Flavor)
|
||||
}
|
||||
if !disc.MessageIsMultimodal {
|
||||
t.Errorf("expected MessageIsMultimodal=true for dict message, got false")
|
||||
}
|
||||
if disc.HistoryFormat != "messages" {
|
||||
t.Errorf("expected HistoryFormat='messages' for server flavor, got %q", disc.HistoryFormat)
|
||||
}
|
||||
|
||||
gw := &GradioGateway{}
|
||||
req := ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Hello world"},
|
||||
},
|
||||
}
|
||||
|
||||
payload, err := gw.BuildGradioPayload(disc, req)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildGradioPayload failed: %v", err)
|
||||
}
|
||||
|
||||
if len(payload) != 2 {
|
||||
t.Fatalf("expected 2 payload slots, got %d", len(payload))
|
||||
}
|
||||
|
||||
// Slot 0: Multimodal dictionary with "text" and "files"
|
||||
msgDict, ok := payload[0].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected slot 0 to be map[string]interface{}, got %T", payload[0])
|
||||
}
|
||||
if msgDict["text"] != "Hello world" {
|
||||
t.Errorf("expected text='Hello world', got %v", msgDict["text"])
|
||||
}
|
||||
if files, ok := msgDict["files"].([]interface{}); !ok || len(files) != 0 {
|
||||
t.Errorf("expected files to be empty slice, got %v", msgDict["files"])
|
||||
}
|
||||
|
||||
// Slot 1: History must be empty slice, never nil
|
||||
histArr, ok := payload[1].([]map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected slot 1 to be []map[string]interface{}, got %T", payload[1])
|
||||
}
|
||||
if histArr == nil {
|
||||
t.Errorf("expected history to be non-nil slice, got nil")
|
||||
}
|
||||
if len(histArr) != 0 {
|
||||
t.Errorf("expected 0 history items for single turn, got %d", len(histArr))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGradio6ChatInterfaceInspectionAndPayload(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/config" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"version": "6.28.0",
|
||||
"api_prefix": "/gradio_api",
|
||||
"mode": "blocks",
|
||||
"components": []map[string]interface{}{
|
||||
{"id": 12, "type": "textbox", "props": map[string]interface{}{"label": "message"}},
|
||||
{"id": 24, "type": "state", "props": map[string]interface{}{"label": "state"}},
|
||||
{"id": 4, "type": "image", "props": map[string]interface{}{"label": "Image (optional)"}},
|
||||
{"id": 5, "type": "checkbox", "props": map[string]interface{}{"label": "Web search", "value": true}},
|
||||
{"id": 6, "type": "checkbox", "props": map[string]interface{}{"label": "Thinking mode", "value": false}},
|
||||
{"id": 7, "type": "slider", "props": map[string]interface{}{"label": "Temperature", "value": 0.7}},
|
||||
{"id": 8, "type": "slider", "props": map[string]interface{}{"label": "Top-p", "value": 0.9}},
|
||||
{"id": 9, "type": "slider", "props": map[string]interface{}{"label": "Max tokens", "value": 8192}},
|
||||
},
|
||||
"dependencies": []map[string]interface{}{
|
||||
{
|
||||
"id": 0,
|
||||
"api_name": "chat",
|
||||
"inputs": []int{12, 24, 4, 5, 6, 7, 8, 9},
|
||||
"outputs": []int{12},
|
||||
"types": map[string]interface{}{"generator": true},
|
||||
},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if r.URL.Path == "/info" || 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": "Textbox", "type": map[string]interface{}{"type": "string"}},
|
||||
{"parameter_name": "image", "component": "Image", "type": map[string]interface{}{"type": "string"}},
|
||||
{"parameter_name": "web_search", "component": "Checkbox", "type": map[string]interface{}{"type": "boolean"}, "default": true},
|
||||
{"parameter_name": "enable_thinking", "component": "Checkbox", "type": map[string]interface{}{"type": "boolean"}, "default": false},
|
||||
{"parameter_name": "temperature", "component": "Slider", "type": map[string]interface{}{"type": "number"}, "default": 0.7},
|
||||
{"parameter_name": "top_p", "component": "Slider", "type": map[string]interface{}{"type": "number"}, "default": 0.9},
|
||||
{"parameter_name": "max_tokens", "component": "Slider", "type": map[string]interface{}{"type": "integer"}, "default": 8192},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
disc, err := InspectSpace(http.DefaultClient, ts.URL, "test-agent")
|
||||
if err != nil {
|
||||
t.Fatalf("InspectSpace failed: %v", err)
|
||||
}
|
||||
|
||||
if disc.MessageIndex != 0 {
|
||||
t.Errorf("expected MessageIndex=0, got %d", disc.MessageIndex)
|
||||
}
|
||||
if disc.ImageIndex != 2 {
|
||||
t.Errorf("expected ImageIndex=2, got %d", disc.ImageIndex)
|
||||
}
|
||||
if disc.WebSearchIndex != 3 {
|
||||
t.Errorf("expected WebSearchIndex=3, got %d", disc.WebSearchIndex)
|
||||
}
|
||||
if disc.ThinkLevelIndex != 4 {
|
||||
t.Errorf("expected ThinkLevelIndex=4, got %d", disc.ThinkLevelIndex)
|
||||
}
|
||||
if disc.TempIndex != 5 {
|
||||
t.Errorf("expected TempIndex=5, got %d", disc.TempIndex)
|
||||
}
|
||||
if disc.TopPIndex != 6 {
|
||||
t.Errorf("expected TopPIndex=6, got %d", disc.TopPIndex)
|
||||
}
|
||||
if disc.MaxTokensIndex != 7 {
|
||||
t.Errorf("expected MaxTokensIndex=7, got %d", disc.MaxTokensIndex)
|
||||
}
|
||||
if disc.IsHunyuan3 {
|
||||
t.Errorf("expected IsHunyuan3=false for generic Gradio space, got true")
|
||||
}
|
||||
|
||||
gw := &GradioGateway{}
|
||||
webSearchFalse := false
|
||||
tempVal := 0.5
|
||||
topPVal := 0.8
|
||||
req := ChatCompletionRequest{
|
||||
Model: "qwen",
|
||||
Messages: []ChatMessage{
|
||||
{
|
||||
Role: "user",
|
||||
Content: []interface{}{
|
||||
map[string]interface{}{"type": "text", "text": "Describe this photo"},
|
||||
map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": "https://example.com/cat.png"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
ReasoningEffort: "high",
|
||||
WebSearch: &webSearchFalse,
|
||||
Temperature: &tempVal,
|
||||
TopP: &topPVal,
|
||||
MaxTokens: 2048,
|
||||
}
|
||||
|
||||
payload, err := gw.BuildGradioPayload(disc, req)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildGradioPayload failed: %v", err)
|
||||
}
|
||||
|
||||
if len(payload) != 8 {
|
||||
t.Fatalf("expected 8 payload items, got %d", len(payload))
|
||||
}
|
||||
|
||||
// Slot 0: prompt text
|
||||
if msg, ok := payload[0].(string); !ok || !strings.Contains(msg, "Describe this photo") {
|
||||
t.Errorf("expected slot 0 to contain 'Describe this photo', got %v", payload[0])
|
||||
}
|
||||
|
||||
// Slot 2: primary image FileData
|
||||
imgMap, ok := payload[2].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected slot 2 to be FileData map, got %T", payload[2])
|
||||
}
|
||||
if imgMap["url"] != "https://example.com/cat.png" {
|
||||
t.Errorf("expected image url 'https://example.com/cat.png', got %v", imgMap["url"])
|
||||
}
|
||||
|
||||
// Slot 3: web search disabled
|
||||
if ws, ok := payload[3].(bool); !ok || ws != false {
|
||||
t.Errorf("expected slot 3 (web_search) to be false, got %v", payload[3])
|
||||
}
|
||||
|
||||
// Slot 4: thinking mode checkbox enabled for "high" effort
|
||||
if think, ok := payload[4].(bool); !ok || think != true {
|
||||
t.Errorf("expected slot 4 (enable_thinking) to be true, got %v", payload[4])
|
||||
}
|
||||
|
||||
// Slot 5: temperature
|
||||
if temp, ok := payload[5].(float64); !ok || temp != 0.5 {
|
||||
t.Errorf("expected slot 5 (temperature) to be 0.5, got %v", payload[5])
|
||||
}
|
||||
|
||||
// Slot 6: top_p
|
||||
if topP, ok := payload[6].(float64); !ok || topP != 0.8 {
|
||||
t.Errorf("expected slot 6 (top_p) to be 0.8, got %v", payload[6])
|
||||
}
|
||||
|
||||
// Slot 7: max_tokens
|
||||
if maxTokens, ok := payload[7].(int); !ok || maxTokens != 2048 {
|
||||
t.Errorf("expected slot 7 (max_tokens) to be 2048, got %v", payload[7])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractThinkingAndDetailsHtmlFormat(t *testing.T) {
|
||||
// 1. Closed details tag with thinking
|
||||
raw1 := `<details open><summary>Thinking</summary>
|
||||
|
||||
I need to analyze this prompt step by step.
|
||||
</details>
|
||||
|
||||
Here is the answer.`
|
||||
|
||||
rem1, think1 := ExtractThinking(raw1)
|
||||
if think1 != "I need to analyze this prompt step by step." {
|
||||
t.Errorf("unexpected reasoning: %q", think1)
|
||||
}
|
||||
if rem1 != "Here is the answer." {
|
||||
t.Errorf("unexpected clean content: %q", rem1)
|
||||
}
|
||||
|
||||
// 2. Preceding web search results details tag preserved
|
||||
raw2 := `<details><summary>Web results</summary>
|
||||
|
||||
1. https://example.com
|
||||
Snippet
|
||||
</details>
|
||||
|
||||
<details open><summary>Thinking</summary>
|
||||
|
||||
Evaluating results...
|
||||
</details>
|
||||
|
||||
Answer based on web results.`
|
||||
|
||||
rem2, think2 := ExtractThinking(raw2)
|
||||
if think2 != "Evaluating results..." {
|
||||
t.Errorf("unexpected reasoning: %q", think2)
|
||||
}
|
||||
if !strings.Contains(rem2, "<details><summary>Web results</summary>") {
|
||||
t.Errorf("expected web results to remain in rem, got: %q", rem2)
|
||||
}
|
||||
if !strings.Contains(rem2, "Answer based on web results.") {
|
||||
t.Errorf("expected final answer to remain in rem, got: %q", rem2)
|
||||
}
|
||||
if strings.Contains(rem2, "Thinking") {
|
||||
t.Errorf("thinking tag should have been stripped from rem, got: %q", rem2)
|
||||
}
|
||||
|
||||
// 3. Open details tag during streaming
|
||||
raw3 := `<details open><summary>Thinking</summary>
|
||||
|
||||
Streaming thought step 1...`
|
||||
rem3, think3 := ExtractThinking(raw3)
|
||||
if think3 != "Streaming thought step 1..." {
|
||||
t.Errorf("unexpected reasoning: %q", think3)
|
||||
}
|
||||
if rem3 != "" {
|
||||
t.Errorf("expected empty content while thinking is open, got %q", rem3)
|
||||
}
|
||||
|
||||
// 4. ParseGradioStreamOutput separation test
|
||||
streamJSON := `["<details open><summary>Thinking</summary>\n\nReasoning...\n</details>\n\nClean answer."]`
|
||||
frame := ParseGradioStreamOutput(streamJSON)
|
||||
if !frame.OK {
|
||||
t.Fatalf("expected frame.OK to be true")
|
||||
}
|
||||
if frame.Reasoning != "Reasoning..." {
|
||||
t.Errorf("expected frame.Reasoning='Reasoning...', got %q", frame.Reasoning)
|
||||
}
|
||||
if frame.Content != "Clean answer." {
|
||||
t.Errorf("expected frame.Content='Clean answer.', got %q", frame.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPausedSpaceDetectionAndImmediateExit(t *testing.T) {
|
||||
callCount := 0
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
w.Write([]byte("The space is paused"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
makeReq := func() (*http.Request, error) {
|
||||
return http.NewRequest("POST", ts.URL, bytes.NewBufferString("{}"))
|
||||
}
|
||||
|
||||
_, err := DoWithFibonacciRetry(http.DefaultClient, makeReq, 5)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error from paused space, got nil")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(err.Error()), "space is paused") {
|
||||
t.Errorf("expected error to mention 'space is paused', got %v", err)
|
||||
}
|
||||
if callCount != 1 {
|
||||
t.Errorf("expected exactly 1 attempt on paused space (no retries), got %d attempts", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOmniForgeSpaceInspectionAndPayload(t *testing.T) {
|
||||
cfgJSON := `{
|
||||
"version": "6.26.0",
|
||||
"mode": "blocks",
|
||||
"title": "OmniForge AI",
|
||||
"api_prefix": "/gradio_api",
|
||||
"components": [
|
||||
{
|
||||
"id": 7,
|
||||
"type": "textbox",
|
||||
"props": {"label": "Ask OmniForge", "value": ""}
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "chatbot",
|
||||
"props": {"label": "OmniForge Conversation", "type": "messages"}
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "dropdown",
|
||||
"props": {
|
||||
"label": "Specialty mode",
|
||||
"value": "General Engineering",
|
||||
"choices": ["General Engineering", "Diagnostics and Repair"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "dropdown",
|
||||
"props": {
|
||||
"label": "Hugging Face model",
|
||||
"value": "deepseek-ai/DeepSeek-V4.1-Flash",
|
||||
"choices": [
|
||||
"deepseek-ai/DeepSeek-V4.1-Flash",
|
||||
"openai/gpt-oss-120b",
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
"meta-llama/Llama-3.1-8B-Instruct"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "state",
|
||||
"props": {"name": "state"},
|
||||
"skip_api": true
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "browserstate",
|
||||
"props": {"default_value": [], "name": "browserstate"},
|
||||
"skip_api": false
|
||||
}
|
||||
],
|
||||
"dependencies": [
|
||||
{
|
||||
"id": 0,
|
||||
"targets": [[10, "click"]],
|
||||
"inputs": [7, 6, 13, 14, 1, 2],
|
||||
"outputs": [6, 7],
|
||||
"api_name": "chat",
|
||||
"connection": "sse",
|
||||
"protocol": "call_v2"
|
||||
}
|
||||
]
|
||||
}`
|
||||
|
||||
infoJSON := `{
|
||||
"named_endpoints": {
|
||||
"/chat": {
|
||||
"parameters": [
|
||||
{"label": "message", "parameter_name": "message", "component": "Textbox"},
|
||||
{"label": "history", "parameter_name": "history", "component": "Chatbot"},
|
||||
{"label": "specialty_value", "parameter_name": "specialty_value", "component": "Dropdown"},
|
||||
{"label": "model_value", "parameter_name": "model_value", "component": "Dropdown"},
|
||||
{"label": "memory_value", "parameter_name": "memory_value", "component": "BrowserState", "parameter_default": []}
|
||||
]
|
||||
}
|
||||
}
|
||||
}`
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.HasSuffix(r.URL.Path, "/config") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(cfgJSON))
|
||||
return
|
||||
}
|
||||
if strings.HasSuffix(r.URL.Path, "/info") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(infoJSON))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
disc, err := InspectSpace(http.DefaultClient, ts.URL, "gr2gw-test")
|
||||
if err != nil {
|
||||
t.Fatalf("InspectSpace failed: %v", err)
|
||||
}
|
||||
|
||||
if !disc.IsOmniForge {
|
||||
t.Errorf("expected disc.IsOmniForge=true, got false")
|
||||
}
|
||||
if disc.TotalInputs != 6 {
|
||||
t.Errorf("expected TotalInputs=6, got %d", disc.TotalInputs)
|
||||
}
|
||||
if disc.ModelIndex != 3 {
|
||||
t.Errorf("expected ModelIndex=3, got %d", disc.ModelIndex)
|
||||
}
|
||||
if disc.HistoryIndex != 1 {
|
||||
t.Errorf("expected HistoryIndex=1, got %d", disc.HistoryIndex)
|
||||
}
|
||||
if disc.MessageIndex != 0 {
|
||||
t.Errorf("expected MessageIndex=0, got %d", disc.MessageIndex)
|
||||
}
|
||||
if disc.PrimaryModel != "deepseek-ai/DeepSeek-V4.1-Flash" {
|
||||
t.Errorf("expected PrimaryModel='deepseek-ai/DeepSeek-V4.1-Flash', got %q", disc.PrimaryModel)
|
||||
}
|
||||
if !disc.MatchesModel("gpt-oss-120b") {
|
||||
t.Errorf("expected MatchesModel('gpt-oss-120b') to be true")
|
||||
}
|
||||
if !disc.MatchesModel("openai/gpt-oss-120b") {
|
||||
t.Errorf("expected MatchesModel('openai/gpt-oss-120b') to be true")
|
||||
}
|
||||
|
||||
gw := &GradioGateway{}
|
||||
req := ChatCompletionRequest{
|
||||
Model: "gpt-oss-120b",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Hello"},
|
||||
{Role: "assistant", Content: "Hi there!"},
|
||||
{Role: "user", Content: "Tell me a joke"},
|
||||
},
|
||||
}
|
||||
|
||||
payload, err := gw.BuildGradioPayload(disc, req)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildGradioPayload failed: %v", err)
|
||||
}
|
||||
|
||||
if len(payload) != 6 {
|
||||
t.Fatalf("expected payload len 6, got %d", len(payload))
|
||||
}
|
||||
|
||||
// Model selection mapped to slot 3
|
||||
if payload[3] != "openai/gpt-oss-120b" {
|
||||
t.Errorf("expected payload[3]='openai/gpt-oss-120b', got %v", payload[3])
|
||||
}
|
||||
|
||||
// History slot mapped to slot 1: should be empty slice for OmniForge
|
||||
if hist, ok := payload[1].([]map[string]interface{}); !ok || len(hist) != 0 {
|
||||
t.Errorf("expected payload[1] to be empty slice, got %v", payload[1])
|
||||
}
|
||||
|
||||
// Multi-turn history composed into promptMessageText in slot 0
|
||||
promptStr, ok := payload[0].(string)
|
||||
if !ok {
|
||||
t.Fatalf("expected payload[0] to be string, got %T", payload[0])
|
||||
}
|
||||
if !strings.Contains(promptStr, "# Conversation History") ||
|
||||
!strings.Contains(promptStr, "User: Hello") ||
|
||||
!strings.Contains(promptStr, "Assistant: Hi there!") ||
|
||||
!strings.Contains(promptStr, "# Current Request\nUser: Tell me a joke") {
|
||||
t.Errorf("promptStr missing multi-turn history: %s", promptStr)
|
||||
}
|
||||
|
||||
// Slot 5 is browserstate (memory_value)
|
||||
if mem, ok := payload[5].([]interface{}); !ok && payload[5] == nil {
|
||||
t.Errorf("expected payload[5] to be non-nil slice, got %v", payload[5])
|
||||
} else if ok && len(mem) != 0 {
|
||||
t.Errorf("expected payload[5] to be empty slice, got %v", mem)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOmniForgeMockServerCompletion(t *testing.T) {
|
||||
var receivedPayload map[string]interface{}
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/gradio_api/call/v2/chat" && r.Method == "POST" {
|
||||
json.NewDecoder(r.Body).Decode(&receivedPayload)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"event_id": "test-event-omni"}`))
|
||||
return
|
||||
}
|
||||
if (r.URL.Path == "/gradio_api/call/chat/test-event-omni" || r.URL.Path == "/gradio_api/call/v2/chat/test-event-omni") && r.Method == "GET" {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
flusher, _ := w.(http.Flusher)
|
||||
|
||||
w.Write([]byte("event: generating\ndata: [[{\"role\": \"user\", \"metadata\": null, \"content\": [{\"text\": \"hello\", \"type\": \"text\"}], \"options\": null}, {\"role\": \"assistant\", \"metadata\": null, \"content\": [{\"text\": \"<think>Thinking about joke</think>Why did the chicken cross the road?\", \"type\": \"text\"}], \"options\": null}], \"\"]\n\n"))
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
w.Write([]byte("event: complete\ndata: null\n\n"))
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
disc := &SpaceDiscovery{
|
||||
SpaceURL: ts.URL,
|
||||
APIPrefix: "/gradio_api",
|
||||
CleanEndpoint: "chat",
|
||||
Protocol: "call_v2",
|
||||
IsOmniForge: true,
|
||||
TotalInputs: 6,
|
||||
MessageIndex: 0,
|
||||
HistoryIndex: 1,
|
||||
ModelIndex: 3,
|
||||
PrimaryModel: "openai/gpt-oss-120b",
|
||||
LastDiscovered: time.Now(),
|
||||
ParamMappings: []SpaceParamMapping{
|
||||
{InputIndex: 0, ParamName: "message", ComponentType: "textbox"},
|
||||
{InputIndex: 1, ParamName: "history", ComponentType: "chatbot"},
|
||||
{InputIndex: 2, ParamName: "specialty_value", ComponentType: "dropdown", DefaultValue: "General Engineering"},
|
||||
{InputIndex: 3, ParamName: "model_value", ComponentType: "dropdown", Choices: []string{"openai/gpt-oss-120b"}, DefaultValue: "openai/gpt-oss-120b"},
|
||||
{InputIndex: 4, ParamName: "state", ParamType: "state", ComponentType: "state"},
|
||||
{InputIndex: 5, ParamName: "memory_value", ComponentType: "browserstate", DefaultValue: []interface{}{}},
|
||||
},
|
||||
}
|
||||
|
||||
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
|
||||
gw.discoveries[ts.URL] = disc
|
||||
|
||||
// 1. Non-streaming test
|
||||
req := ChatCompletionRequest{
|
||||
Model: "openai/gpt-oss-120b",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Tell me a joke"},
|
||||
},
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
// Check received payload by mock server
|
||||
if receivedPayload["model_value"] != "openai/gpt-oss-120b" {
|
||||
t.Errorf("expected model_value='openai/gpt-oss-120b', got %v", receivedPayload["model_value"])
|
||||
}
|
||||
if receivedPayload["memory_value"] == nil {
|
||||
t.Errorf("expected memory_value to be non-nil in v2 payload")
|
||||
}
|
||||
if _, hasState := receivedPayload["state"]; hasState {
|
||||
t.Errorf("expected unexposed state component to NOT be in v2 payload")
|
||||
}
|
||||
|
||||
var resp ChatCompletionResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
t.Fatalf("expected choices in response")
|
||||
}
|
||||
|
||||
choice := resp.Choices[0]
|
||||
if choice.Message.ReasoningContent != "Thinking about joke" {
|
||||
t.Errorf("expected reasoning='Thinking about joke', got %q", choice.Message.ReasoningContent)
|
||||
}
|
||||
if choice.Message.Content != "Why did the chicken cross the road?" {
|
||||
t.Errorf("expected content='Why did the chicken cross the road?', got %q", choice.Message.Content)
|
||||
}
|
||||
|
||||
// 2. Streaming test
|
||||
reqStream := ChatCompletionRequest{
|
||||
Model: "openai/gpt-oss-120b",
|
||||
Stream: true,
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Tell me a joke"},
|
||||
},
|
||||
}
|
||||
|
||||
recStream := httptest.NewRecorder()
|
||||
errStream := gw.ExecuteChatCompletion(recStream, httpReq, reqStream)
|
||||
if errStream != nil {
|
||||
t.Fatalf("streaming completion failed: %v", errStream)
|
||||
}
|
||||
|
||||
bodyStr := recStream.Body.String()
|
||||
if !strings.Contains(bodyStr, "Thinking about joke") {
|
||||
t.Errorf("streaming response missing reasoning token: %s", bodyStr)
|
||||
}
|
||||
if !strings.Contains(bodyStr, "Why did the chicken cross the road?") {
|
||||
t.Errorf("streaming response missing content token: %s", bodyStr)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user