fix(heuristics): penalize user submit handlers and support history-only generator endpoints

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,13 +1998,22 @@ 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" {
+          isSharedInput := false
+          for _, inID := range dep.Inputs {
+            if inID == outID {
+              isSharedInput = true
+              break
+            }
+          }
+          if !isSharedInput {
             score += 50
           }
         }
       }
     }
+  }
 
   return score
 }
@@ -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,6 +2339,7 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover
     }
 
     if discovery.MessageIndex == -1 {
+      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 {
@@ -2303,6 +2350,7 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover
         }
       }
     }
+  }
 
   // Refine history format or discover tools from bestEndpointInfo
   if bestEndpointInfo != nil {
@@ -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,9 +2452,11 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover
     }
 
     if discovery.MessageIndex == -1 {
+      if discovery.HistoryIndex != 0 && discovery.SystemIndex != 0 && discovery.FunctionsJSONIndex != 0 {
         discovery.MessageIndex = 0
       }
     }
+  }
 
   // 7. Resolve default history format if not set by type inspection
   if discovery.HistoryIndex != -1 {
@@ -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
   }

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)
+  }
+}
+
+