diff --git a/gr2gw.go b/gr2gw.go index 96f2b60..bb9085d 100644 --- a/gr2gw.go +++ b/gr2gw.go @@ -1865,9 +1865,33 @@ func ScoreCandidateEndpoint(apiName string, parameters []GradioParamInfo, dep *G score -= 200 } - if lowerName == "chat" || lowerName == "chat_fn" { + // Penalize user-submission helper endpoints (common in Gradio Blocks multi-step chat interfaces + // where a "user" function simply appends input to history and clears the textbox) + isUserHandler := lowerName == "user" || strings.HasPrefix(lowerName, "user_") || + strings.Contains(lowerName, "add_text") || strings.Contains(lowerName, "add_msg") || + strings.Contains(lowerName, "add_message") || strings.Contains(lowerName, "append_to_history") + if isUserHandler { + score -= 600 + } else if strings.HasPrefix(lowerName, "user") { + // Check if name is user followed by digits/underscores (e.g. user2, user3, user_1) + suffix := strings.TrimPrefix(lowerName, "user") + isDigitsOrUnder := true + for _, r := range suffix { + if (r < '0' || r > '9') && r != '_' { + isDigitsOrUnder = false + break + } + } + if isDigitsOrUnder && len(suffix) > 0 { + score -= 600 + } + } + + if lowerName == "chat" || lowerName == "chat_fn" || lowerName == "bot" || lowerName == "bot_fn" || lowerName == "chat_message" { score += 150 - } else if strings.Contains(lowerName, "chat") || strings.Contains(lowerName, "conversation") || strings.Contains(lowerName, "dialogue") || strings.Contains(lowerName, "chatbot") { + } else if strings.Contains(lowerName, "chat") || strings.Contains(lowerName, "conversation") || + strings.Contains(lowerName, "dialogue") || strings.Contains(lowerName, "chatbot") || + strings.Contains(lowerName, "bot") { score += 120 } if lowerName == "predict" || lowerName == "generate" { @@ -1921,7 +1945,20 @@ func ScoreCandidateEndpoint(apiName string, parameters []GradioParamInfo, dep *G if dep != nil { if dep.Types.Generator { - score += 40 + score += 150 + } + // Check if any textbox component in dep.Inputs is also in dep.Outputs (textbox-clearing UI handler) + for _, inID := range dep.Inputs { + if comp, exists := compMap[inID]; exists { + if strings.ToLower(comp.Type) == "textbox" { + for _, outID := range dep.Outputs { + if outID == inID { + score -= 600 + break + } + } + } + } } for _, inID := range dep.Inputs { if comp, exists := compMap[inID]; exists { @@ -1961,9 +1998,18 @@ func ScoreCandidateEndpoint(apiName string, parameters []GradioParamInfo, dep *G if comp, exists := compMap[outID]; exists { cType := strings.ToLower(comp.Type) if cType == "chatbot" { - score += 90 + score += 100 } else if cType == "textbox" || cType == "markdown" { - score += 50 + isSharedInput := false + for _, inID := range dep.Inputs { + if inID == outID { + isSharedInput = true + break + } + } + if !isSharedInput { + score += 50 + } } } } @@ -2269,7 +2315,7 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover mapping.ParamType = "message" discovery.MessageIndex = idx discovery.MessageIsMultimodal = true - } else if strings.Contains(pName, "message") || strings.Contains(pLabel, "message") || strings.Contains(cLabel, "message") || strings.Contains(cLabel, "prompt") || strings.Contains(cLabel, "query") || (discovery.MessageIndex == -1 && idx == 0) { + } else if strings.Contains(pName, "message") || strings.Contains(pLabel, "message") || strings.Contains(cLabel, "message") || strings.Contains(cLabel, "prompt") || strings.Contains(cLabel, "query") || (discovery.MessageIndex == -1 && idx == 0 && (cType == "textbox" || pComp == "textbox")) { mapping.ParamType = "message" discovery.MessageIndex = idx discovery.MessageIsMultimodal = false @@ -2293,11 +2339,13 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover } if discovery.MessageIndex == -1 { - discovery.MessageIndex = 0 - if len(bestMatchingDep.Inputs) > 0 { - if comp, exists := compMap[bestMatchingDep.Inputs[0]]; exists { - if strings.ToLower(comp.Type) == "multimodaltextbox" { - discovery.MessageIsMultimodal = true + if discovery.HistoryIndex != 0 && discovery.SystemIndex != 0 && discovery.FunctionsJSONIndex != 0 { + discovery.MessageIndex = 0 + if len(bestMatchingDep.Inputs) > 0 { + if comp, exists := compMap[bestMatchingDep.Inputs[0]]; exists { + if strings.ToLower(comp.Type) == "multimodaltextbox" { + discovery.MessageIsMultimodal = true + } } } } @@ -2374,7 +2422,7 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover } else if strings.Contains(pPyType, "textmessage") || strings.Contains(pPyType, "dict(text: str") || strings.Contains(bTypeStr, "textmessage") || strings.Contains(bTypeStr, "chatbotdatamessages") { discovery.HistoryFormat = "gradio_messages" } - } else if strings.Contains(pLabel, "message") || strings.Contains(pName, "message") || (strings.Contains(pLabel, "prompt") && !strings.Contains(pLabel, "system")) || (strings.Contains(pName, "prompt") && !strings.Contains(pName, "system")) || strings.Contains(pLabel, "query") || strings.Contains(pName, "query") || strings.Contains(pLabel, "question") || strings.Contains(pName, "question") || (discovery.MessageIndex == -1 && (pComp == "textbox" || idx == 0)) { + } else if strings.Contains(pLabel, "message") || strings.Contains(pName, "message") || (strings.Contains(pLabel, "prompt") && !strings.Contains(pLabel, "system")) || (strings.Contains(pName, "prompt") && !strings.Contains(pName, "system")) || strings.Contains(pLabel, "query") || strings.Contains(pName, "query") || strings.Contains(pLabel, "question") || strings.Contains(pName, "question") || (discovery.MessageIndex == -1 && (pComp == "textbox" || (idx == 0 && pComp != "chatbot"))) { discovery.MessageIndex = idx mapping.ParamType = "message" discovery.MessageIsMultimodal = false @@ -2404,7 +2452,9 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover } if discovery.MessageIndex == -1 { - discovery.MessageIndex = 0 + if discovery.HistoryIndex != 0 && discovery.SystemIndex != 0 && discovery.FunctionsJSONIndex != 0 { + discovery.MessageIndex = 0 + } } } @@ -2817,6 +2867,9 @@ func (g *GradioGateway) BuildGradioPayload(disc *SpaceDiscovery, req ChatComplet } pairs = append(pairs, []string{u, a}) } + if disc.MessageIndex == -1 && promptMessageText != "" { + pairs = append(pairs, []string{promptMessageText, ""}) + } data[disc.HistoryIndex] = pairs } else if disc.HistoryFormat == "gradio_messages" { var gMsgs []map[string]interface{} @@ -2834,15 +2887,34 @@ func (g *GradioGateway) BuildGradioPayload(disc *SpaceDiscovery, req ChatComplet } gMsgs = append(gMsgs, gMsg) } + if disc.MessageIndex == -1 && promptMessageText != "" { + gMsgs = append(gMsgs, map[string]interface{}{ + "role": "user", + "content": []map[string]string{ + {"text": promptMessageText, "type": "text"}, + }, + }) + } if gMsgs == nil { gMsgs = []map[string]interface{}{} } data[disc.HistoryIndex] = gMsgs } else { + if disc.MessageIndex == -1 && promptMessageText != "" { + historyArray = append(historyArray, map[string]interface{}{ + "role": "user", + "content": promptMessageText, + }) + } data[disc.HistoryIndex] = historyArray } } + // If neither message nor history slot was mapped, populate slot 0 with the prompt + if disc.MessageIndex == -1 && disc.HistoryIndex == -1 && len(data) > 0 { + data[0] = promptMessageText + } + if disc.SystemIndex >= 0 && disc.SystemIndex < len(data) { data[disc.SystemIndex] = systemPromptStr } diff --git a/gr2gw_test.go b/gr2gw_test.go index 684de14..08588f1 100644 --- a/gr2gw_test.go +++ b/gr2gw_test.go @@ -1878,4 +1878,211 @@ func TestCallToPredictProtocolFallback(t *testing.T) { } } +// 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) + } + if disc.TotalInputs != 2 { + t.Errorf("expected TotalInputs 2, got %d", disc.TotalInputs) + } +} + +