Support Gradio queue protocol with diff streaming and call fallback
This commit is contained in:
@@ -1205,6 +1205,18 @@ func extractGradioErrorMessage(dataStr string) string {
|
|||||||
return clean
|
return clean
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func isProtocolOrInputError(errMsg string) bool {
|
||||||
|
lower := strings.ToLower(errMsg)
|
||||||
|
return strings.Contains(lower, "check gradio inputs/types") ||
|
||||||
|
strings.Contains(lower, "null error") ||
|
||||||
|
strings.Contains(lower, "typeerror") ||
|
||||||
|
strings.Contains(lower, "missing") ||
|
||||||
|
strings.Contains(lower, "required positional argument") ||
|
||||||
|
strings.Contains(lower, "internal server error") ||
|
||||||
|
strings.Contains(lower, "not found") ||
|
||||||
|
strings.Contains(lower, "unknown upstream gradio error")
|
||||||
|
}
|
||||||
|
|
||||||
type Streamer struct {
|
type Streamer struct {
|
||||||
w http.ResponseWriter
|
w http.ResponseWriter
|
||||||
flusher http.Flusher
|
flusher http.Flusher
|
||||||
@@ -1622,6 +1634,7 @@ type SpaceDiscovery struct {
|
|||||||
Models []string `json:"models"`
|
Models []string `json:"models"`
|
||||||
PrimaryModel string `json:"primary_model"`
|
PrimaryModel string `json:"primary_model"`
|
||||||
TotalInputs int `json:"total_inputs"`
|
TotalInputs int `json:"total_inputs"`
|
||||||
|
RawTotalInputs int `json:"raw_total_inputs,omitempty"`
|
||||||
ParamMappings []SpaceParamMapping `json:"param_mappings"`
|
ParamMappings []SpaceParamMapping `json:"param_mappings"`
|
||||||
DefaultInputs []interface{} `json:"default_inputs"`
|
DefaultInputs []interface{} `json:"default_inputs"`
|
||||||
MessageIsMultimodal bool `json:"message_is_multimodal"`
|
MessageIsMultimodal bool `json:"message_is_multimodal"`
|
||||||
@@ -2270,6 +2283,7 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover
|
|||||||
// 6. Correlate with config.dependencies to determine exact input count & state padding
|
// 6. Correlate with config.dependencies to determine exact input count & state padding
|
||||||
if bestMatchingDep != nil {
|
if bestMatchingDep != nil {
|
||||||
discovery.TotalInputs = len(bestMatchingDep.Inputs)
|
discovery.TotalInputs = len(bestMatchingDep.Inputs)
|
||||||
|
discovery.RawTotalInputs = len(bestMatchingDep.Inputs)
|
||||||
discovery.DefaultInputs = make([]interface{}, len(bestMatchingDep.Inputs))
|
discovery.DefaultInputs = make([]interface{}, len(bestMatchingDep.Inputs))
|
||||||
discovery.ParamMappings = nil
|
discovery.ParamMappings = nil
|
||||||
discovery.MessageIndex = -1
|
discovery.MessageIndex = -1
|
||||||
@@ -3062,9 +3076,45 @@ type GradioOutputFrame struct {
|
|||||||
Content string
|
Content string
|
||||||
Reasoning string
|
Reasoning string
|
||||||
ToolCalls []ToolCall
|
ToolCalls []ToolCall
|
||||||
|
IsDelta bool
|
||||||
OK bool
|
OK bool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// extractGradioDiffDelta extracts delta string from Gradio 6 streaming diff operations
|
||||||
|
func extractGradioDiffDelta(v []interface{}) (string, bool) {
|
||||||
|
if len(v) == 0 {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
if len(v) >= 3 {
|
||||||
|
if op, ok := v[0].(string); ok && (op == "append" || op == "add") {
|
||||||
|
if delta, ok := v[2].(string); ok {
|
||||||
|
return delta, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, item := range v {
|
||||||
|
if innerList, ok := item.([]interface{}); ok && len(innerList) > 0 {
|
||||||
|
for _, opItem := range innerList {
|
||||||
|
if opArr, ok := opItem.([]interface{}); ok && len(opArr) >= 3 {
|
||||||
|
if op, ok := opArr[0].(string); ok && (op == "append" || op == "add") {
|
||||||
|
if delta, ok := opArr[2].(string); ok {
|
||||||
|
return delta, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(innerList) >= 3 {
|
||||||
|
if op, ok := innerList[0].(string); ok && (op == "append" || op == "add") {
|
||||||
|
if delta, ok := innerList[2].(string); ok {
|
||||||
|
return delta, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
|
||||||
// ParseGradioStreamOutput extracts structured content, reasoning, and tool calls from Gradio output
|
// ParseGradioStreamOutput extracts structured content, reasoning, and tool calls from Gradio output
|
||||||
func ParseGradioStreamOutput(rawJSON string) (frame GradioOutputFrame) {
|
func ParseGradioStreamOutput(rawJSON string) (frame GradioOutputFrame) {
|
||||||
defer func() {
|
defer func() {
|
||||||
@@ -3093,6 +3143,13 @@ func ParseGradioStreamOutput(rawJSON string) (frame GradioOutputFrame) {
|
|||||||
return frame
|
return frame
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if delta, ok := extractGradioDiffDelta(v); ok {
|
||||||
|
frame.Content = delta
|
||||||
|
frame.IsDelta = true
|
||||||
|
frame.OK = true
|
||||||
|
return frame
|
||||||
|
}
|
||||||
|
|
||||||
// 1. Check if v[0] is an inner slice with len >= 3 (e.g. Hy3: [[content, reasoning, tool_calls, history]])
|
// 1. Check if v[0] is an inner slice with len >= 3 (e.g. Hy3: [[content, reasoning, tool_calls, history]])
|
||||||
if inner, ok := v[0].([]interface{}); ok && len(inner) >= 3 {
|
if inner, ok := v[0].([]interface{}); ok && len(inner) >= 3 {
|
||||||
s0, _ := inner[0].(string)
|
s0, _ := inner[0].(string)
|
||||||
@@ -3344,8 +3401,20 @@ func (g *GradioGateway) executePredictCompletion(w http.ResponseWriter, r *http.
|
|||||||
fnIndex = 0
|
fnIndex = 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
predictData := make([]interface{}, len(gradioData))
|
||||||
|
copy(predictData, gradioData)
|
||||||
|
if disc.RawTotalInputs > len(predictData) {
|
||||||
|
for i := len(predictData); i < disc.RawTotalInputs; i++ {
|
||||||
|
var defVal interface{}
|
||||||
|
if i < len(disc.DefaultInputs) {
|
||||||
|
defVal = disc.DefaultInputs[i]
|
||||||
|
}
|
||||||
|
predictData = append(predictData, defVal)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
payloadMap := map[string]interface{}{
|
payloadMap := map[string]interface{}{
|
||||||
"data": gradioData,
|
"data": predictData,
|
||||||
"fn_index": fnIndex,
|
"fn_index": fnIndex,
|
||||||
"session_hash": GenerateUUID(),
|
"session_hash": GenerateUUID(),
|
||||||
}
|
}
|
||||||
@@ -3453,6 +3522,417 @@ func (g *GradioGateway) executePredictCompletion(w http.ResponseWriter, r *http.
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// executeQueueCompletion handles completions using Gradio queue SSE protocol (/queue/join + /queue/data).
|
||||||
|
func (g *GradioGateway) executeQueueCompletion(w http.ResponseWriter, r *http.Request, disc *SpaceDiscovery, gradioData []interface{}, req ChatCompletionRequest, completionID string, createdTime int64, modelName, effUA string) error {
|
||||||
|
fnIndex := disc.FnIndex
|
||||||
|
if fnIndex < 0 {
|
||||||
|
fnIndex = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
queueData := make([]interface{}, len(gradioData))
|
||||||
|
copy(queueData, gradioData)
|
||||||
|
if disc.RawTotalInputs > len(queueData) {
|
||||||
|
for i := len(queueData); i < disc.RawTotalInputs; i++ {
|
||||||
|
var defVal interface{}
|
||||||
|
if i < len(disc.DefaultInputs) {
|
||||||
|
defVal = disc.DefaultInputs[i]
|
||||||
|
}
|
||||||
|
queueData = append(queueData, defVal)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
sessionHash := GenerateUUID()
|
||||||
|
payloadMap := map[string]interface{}{
|
||||||
|
"data": queueData,
|
||||||
|
"fn_index": fnIndex,
|
||||||
|
"session_hash": sessionHash,
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonPayload, err := json.Marshal(payloadMap)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to marshal queue payload: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var joinURLs []string
|
||||||
|
if disc.APIPrefix != "" {
|
||||||
|
joinURLs = append(joinURLs, fmt.Sprintf("%s%s/queue/join", disc.SpaceURL, disc.APIPrefix))
|
||||||
|
}
|
||||||
|
joinURLs = append(joinURLs, fmt.Sprintf("%s/queue/join", disc.SpaceURL))
|
||||||
|
|
||||||
|
var resp *http.Response
|
||||||
|
var lastErr error
|
||||||
|
|
||||||
|
for _, targetURL := range joinURLs {
|
||||||
|
curURL := targetURL
|
||||||
|
makeReq := func() (*http.Request, error) {
|
||||||
|
req, err := http.NewRequest("POST", curURL, bytes.NewBuffer(jsonPayload))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("User-Agent", effUA)
|
||||||
|
return req, nil
|
||||||
|
}
|
||||||
|
resp, lastErr = DoWithFibonacciRetry(g.client, makeReq, 3)
|
||||||
|
if lastErr == nil && resp != nil && resp.StatusCode == http.StatusOK {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if resp != nil {
|
||||||
|
resp.Body.Close()
|
||||||
|
resp = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp == nil {
|
||||||
|
log.Printf("Gradio /queue/join unavailable (%v), falling back to /run/predict protocol...", lastErr)
|
||||||
|
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
|
||||||
|
// 2. Connect to Gradio SSE queue data stream
|
||||||
|
var dataURLs []string
|
||||||
|
if disc.APIPrefix != "" {
|
||||||
|
dataURLs = append(dataURLs, fmt.Sprintf("%s%s/queue/data?session_hash=%s", disc.SpaceURL, disc.APIPrefix, sessionHash))
|
||||||
|
}
|
||||||
|
dataURLs = append(dataURLs, fmt.Sprintf("%s/queue/data?session_hash=%s", disc.SpaceURL, sessionHash))
|
||||||
|
|
||||||
|
var streamResp *http.Response
|
||||||
|
for _, targetURL := range dataURLs {
|
||||||
|
curURL := targetURL
|
||||||
|
makeStreamReq := func() (*http.Request, error) {
|
||||||
|
req, err := http.NewRequest("GET", curURL, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
req.Header.Set("Accept", "text/event-stream")
|
||||||
|
req.Header.Set("User-Agent", effUA)
|
||||||
|
return req, nil
|
||||||
|
}
|
||||||
|
streamResp, lastErr = DoWithFibonacciRetry(g.client, makeStreamReq, 3)
|
||||||
|
if lastErr == nil && streamResp != nil && streamResp.StatusCode == http.StatusOK {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if streamResp != nil {
|
||||||
|
streamResp.Body.Close()
|
||||||
|
streamResp = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if streamResp == nil {
|
||||||
|
log.Printf("Gradio /queue/data unavailable (%v), falling back to /run/predict protocol...", lastErr)
|
||||||
|
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
|
}
|
||||||
|
defer streamResp.Body.Close()
|
||||||
|
|
||||||
|
type queueMsg struct {
|
||||||
|
Msg string `json:"msg"`
|
||||||
|
EventID string `json:"event_id,omitempty"`
|
||||||
|
Success *bool `json:"success,omitempty"`
|
||||||
|
Output map[string]interface{} `json:"output,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Handle Non-Streaming vs Streaming
|
||||||
|
if !req.Stream {
|
||||||
|
reader := bufio.NewReader(streamResp.Body)
|
||||||
|
var latestFrame GradioOutputFrame
|
||||||
|
|
||||||
|
for {
|
||||||
|
line, err := reader.ReadString('\n')
|
||||||
|
if err != nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
line = strings.TrimRight(line, "\r\n")
|
||||||
|
|
||||||
|
if strings.HasPrefix(line, "data: ") {
|
||||||
|
dataStr := strings.TrimPrefix(line, "data: ")
|
||||||
|
var qMsg queueMsg
|
||||||
|
if err := json.Unmarshal([]byte(dataStr), &qMsg); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Msg == "close_stream" {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Msg == "process_generating" || qMsg.Msg == "process_completed" {
|
||||||
|
if qMsg.Success != nil && !*qMsg.Success {
|
||||||
|
errBytes, _ := json.Marshal(qMsg.Output)
|
||||||
|
errMsg := extractGradioErrorMessage(string(errBytes))
|
||||||
|
log.Printf("Upstream Gradio error: %s", errMsg)
|
||||||
|
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Output != nil && qMsg.Output["data"] != nil {
|
||||||
|
dataBytes, _ := json.Marshal(qMsg.Output["data"])
|
||||||
|
frame := ParseGradioStreamOutput(string(dataBytes))
|
||||||
|
if frame.OK {
|
||||||
|
if frame.IsDelta {
|
||||||
|
latestFrame.Content += frame.Content
|
||||||
|
} else {
|
||||||
|
if frame.Content != "" || latestFrame.Content == "" {
|
||||||
|
latestFrame.Content = frame.Content
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if frame.Reasoning != "" || latestFrame.Reasoning == "" {
|
||||||
|
latestFrame.Reasoning = frame.Reasoning
|
||||||
|
}
|
||||||
|
if len(frame.ToolCalls) > 0 {
|
||||||
|
latestFrame.ToolCalls = frame.ToolCalls
|
||||||
|
}
|
||||||
|
latestFrame.OK = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if qMsg.Msg == "process_completed" {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !latestFrame.OK {
|
||||||
|
return fmt.Errorf("upstream Gradio space returned empty or unparseable response")
|
||||||
|
}
|
||||||
|
|
||||||
|
cleanText := latestFrame.Content
|
||||||
|
reasoning := latestFrame.Reasoning
|
||||||
|
toolCalls := latestFrame.ToolCalls
|
||||||
|
hasTools := len(toolCalls) > 0
|
||||||
|
|
||||||
|
if reasoning == "" {
|
||||||
|
cleanText, reasoning = ExtractThinking(cleanText)
|
||||||
|
}
|
||||||
|
if !hasTools {
|
||||||
|
toolCalls, cleanText, hasTools = DetectToolCalls(cleanText)
|
||||||
|
}
|
||||||
|
|
||||||
|
finishReason := "stop"
|
||||||
|
var finalContent interface{} = cleanText
|
||||||
|
if hasTools && len(toolCalls) > 0 {
|
||||||
|
finishReason = "tool_calls"
|
||||||
|
if strings.TrimSpace(cleanText) == "" {
|
||||||
|
finalContent = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{
|
||||||
|
Content: finalContent,
|
||||||
|
ReasoningContent: reasoning,
|
||||||
|
ToolCalls: toolCalls,
|
||||||
|
FinishReason: finishReason,
|
||||||
|
})
|
||||||
|
disc.Protocol = "queue"
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Streaming Mode
|
||||||
|
flusher, _ := w.(http.Flusher)
|
||||||
|
streamer := NewStreamer(w, flusher, completionID, createdTime, modelName)
|
||||||
|
|
||||||
|
thinkFilter := NewStreamThinkingFilter()
|
||||||
|
toolFilter := NewStreamToolCallFilter()
|
||||||
|
|
||||||
|
reader := bufio.NewReader(streamResp.Body)
|
||||||
|
var prevContent string
|
||||||
|
var prevReasoning string
|
||||||
|
prevToolArgs := make(map[int]string)
|
||||||
|
nativeReasoningSeen := false
|
||||||
|
nativeToolCallsSeen := false
|
||||||
|
var streamErr error
|
||||||
|
|
||||||
|
for {
|
||||||
|
line, err := reader.ReadString('\n')
|
||||||
|
if err != nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
line = strings.TrimRight(line, "\r\n")
|
||||||
|
|
||||||
|
if strings.HasPrefix(line, "data: ") {
|
||||||
|
dataStr := strings.TrimPrefix(line, "data: ")
|
||||||
|
var qMsg queueMsg
|
||||||
|
if err := json.Unmarshal([]byte(dataStr), &qMsg); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Msg == "close_stream" {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Msg == "process_generating" || qMsg.Msg == "process_completed" {
|
||||||
|
if qMsg.Success != nil && !*qMsg.Success {
|
||||||
|
errBytes, _ := json.Marshal(qMsg.Output)
|
||||||
|
errMsg := extractGradioErrorMessage(string(errBytes))
|
||||||
|
log.Printf("Upstream Gradio error: %s", errMsg)
|
||||||
|
if !streamer.started {
|
||||||
|
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
|
}
|
||||||
|
streamer.Error(errMsg)
|
||||||
|
streamErr = fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
if qMsg.Output != nil && qMsg.Output["data"] != nil {
|
||||||
|
dataBytes, _ := json.Marshal(qMsg.Output["data"])
|
||||||
|
frame := ParseGradioStreamOutput(string(dataBytes))
|
||||||
|
if frame.OK {
|
||||||
|
// 1. Native reasoning handling
|
||||||
|
if frame.Reasoning != "" || nativeReasoningSeen {
|
||||||
|
nativeReasoningSeen = true
|
||||||
|
if len(frame.Reasoning) > len(prevReasoning) && strings.HasPrefix(frame.Reasoning, prevReasoning) {
|
||||||
|
delta := frame.Reasoning[len(prevReasoning):]
|
||||||
|
streamer.Reasoning(delta)
|
||||||
|
prevReasoning = frame.Reasoning
|
||||||
|
} else if prevReasoning == "" && frame.Reasoning != "" {
|
||||||
|
streamer.Reasoning(frame.Reasoning)
|
||||||
|
prevReasoning = frame.Reasoning
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Native tool call handling
|
||||||
|
if len(frame.ToolCalls) > 0 || nativeToolCallsSeen {
|
||||||
|
nativeToolCallsSeen = true
|
||||||
|
for i, tc := range frame.ToolCalls {
|
||||||
|
prevArgs := prevToolArgs[i]
|
||||||
|
fullArgs := tc.Function.Arguments
|
||||||
|
if len(fullArgs) > len(prevArgs) && strings.HasPrefix(fullArgs, prevArgs) {
|
||||||
|
deltaArgs := fullArgs[len(prevArgs):]
|
||||||
|
iCopy := i
|
||||||
|
streamer.ToolCallDelta(ToolCall{
|
||||||
|
Index: &iCopy,
|
||||||
|
ID: tc.ID,
|
||||||
|
Type: tc.Type,
|
||||||
|
Function: ToolCallFunction{
|
||||||
|
Name: tc.Function.Name,
|
||||||
|
Arguments: deltaArgs,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
prevToolArgs[i] = fullArgs
|
||||||
|
} else if prevArgs == "" {
|
||||||
|
iCopy := i
|
||||||
|
streamer.ToolCallDelta(ToolCall{
|
||||||
|
Index: &iCopy,
|
||||||
|
ID: tc.ID,
|
||||||
|
Type: tc.Type,
|
||||||
|
Function: ToolCallFunction{
|
||||||
|
Name: tc.Function.Name,
|
||||||
|
Arguments: fullArgs,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
prevToolArgs[i] = fullArgs
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Content handling
|
||||||
|
currentText := frame.Content
|
||||||
|
var delta string
|
||||||
|
if frame.IsDelta {
|
||||||
|
delta = frame.Content
|
||||||
|
prevContent += delta
|
||||||
|
} else {
|
||||||
|
if strings.HasPrefix(currentText, prevContent) {
|
||||||
|
delta = currentText[len(prevContent):]
|
||||||
|
} else if prevContent == "" {
|
||||||
|
delta = currentText
|
||||||
|
} else {
|
||||||
|
delta = currentText
|
||||||
|
}
|
||||||
|
prevContent = currentText
|
||||||
|
}
|
||||||
|
|
||||||
|
if delta != "" {
|
||||||
|
if nativeReasoningSeen || nativeToolCallsSeen {
|
||||||
|
if nativeReasoningSeen && nativeToolCallsSeen {
|
||||||
|
streamer.Content(delta)
|
||||||
|
} else if nativeReasoningSeen {
|
||||||
|
toolFilter.Feed(delta, func(cleanChunk string) {
|
||||||
|
if cleanChunk != "" {
|
||||||
|
streamer.Content(cleanChunk)
|
||||||
|
}
|
||||||
|
}, func(tc ToolCall) {
|
||||||
|
streamer.ToolCallDelta(tc)
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
thinkFilter.Feed(delta, func(contentChunk string) {
|
||||||
|
if contentChunk != "" {
|
||||||
|
streamer.Content(contentChunk)
|
||||||
|
}
|
||||||
|
}, func(reasoningChunk string) {
|
||||||
|
if reasoningChunk != "" {
|
||||||
|
streamer.Reasoning(reasoningChunk)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
thinkFilter.Feed(delta, func(contentChunk string) {
|
||||||
|
toolFilter.Feed(contentChunk, func(cleanChunk string) {
|
||||||
|
if cleanChunk != "" {
|
||||||
|
streamer.Content(cleanChunk)
|
||||||
|
}
|
||||||
|
}, func(tc ToolCall) {
|
||||||
|
streamer.ToolCallDelta(tc)
|
||||||
|
})
|
||||||
|
}, func(reasoningChunk string) {
|
||||||
|
if reasoningChunk != "" {
|
||||||
|
streamer.Reasoning(reasoningChunk)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if qMsg.Msg == "process_completed" {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if streamErr != nil {
|
||||||
|
return streamErr
|
||||||
|
}
|
||||||
|
|
||||||
|
// Flush remaining tokens in filters if used
|
||||||
|
if !nativeReasoningSeen {
|
||||||
|
thinkFilter.Flush(func(contentChunk string) {
|
||||||
|
if !nativeToolCallsSeen {
|
||||||
|
toolFilter.Feed(contentChunk, func(cleanChunk string) {
|
||||||
|
if cleanChunk != "" {
|
||||||
|
streamer.Content(cleanChunk)
|
||||||
|
}
|
||||||
|
}, func(tc ToolCall) {
|
||||||
|
streamer.ToolCallDelta(tc)
|
||||||
|
})
|
||||||
|
} else if contentChunk != "" {
|
||||||
|
streamer.Content(contentChunk)
|
||||||
|
}
|
||||||
|
}, func(reasoningChunk string) {
|
||||||
|
if reasoningChunk != "" {
|
||||||
|
streamer.Reasoning(reasoningChunk)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
if !nativeToolCallsSeen {
|
||||||
|
toolFilter.Flush(func(cleanChunk string) {
|
||||||
|
if cleanChunk != "" {
|
||||||
|
streamer.Content(cleanChunk)
|
||||||
|
}
|
||||||
|
}, func(tc ToolCall) {
|
||||||
|
streamer.ToolCallDelta(tc)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
finishReason := "stop"
|
||||||
|
if nativeToolCallsSeen || toolFilter.emittedCall {
|
||||||
|
finishReason = "tool_calls"
|
||||||
|
}
|
||||||
|
|
||||||
|
streamer.Finish(finishReason)
|
||||||
|
streamer.Done()
|
||||||
|
disc.Protocol = "queue"
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// ExecuteChatCompletion handles both streaming and non-streaming requests.
|
// ExecuteChatCompletion handles both streaming and non-streaming requests.
|
||||||
func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Request, req ChatCompletionRequest) error {
|
func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Request, req ChatCompletionRequest) error {
|
||||||
effUA := EffectiveUserAgent(r)
|
effUA := EffectiveUserAgent(r)
|
||||||
@@ -3488,12 +3968,18 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
if disc.Protocol == "predict" {
|
if disc.Protocol == "predict" {
|
||||||
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
}
|
}
|
||||||
|
if disc.Protocol == "queue" {
|
||||||
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
|
}
|
||||||
|
|
||||||
isV2 := disc.Protocol == "call_v2"
|
isV2 := disc.Protocol == "call_v2"
|
||||||
buildPayload := func(v2 bool) ([]byte, error) {
|
buildPayload := func(v2 bool) ([]byte, error) {
|
||||||
if v2 {
|
if v2 {
|
||||||
v2Payload := make(map[string]interface{})
|
v2Payload := make(map[string]interface{})
|
||||||
for idx, val := range gradioData {
|
for idx, val := range gradioData {
|
||||||
|
if idx < len(disc.ParamMappings) && disc.ParamMappings[idx].ParamType == "state" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
pName := ""
|
pName := ""
|
||||||
if idx < len(disc.ParamMappings) {
|
if idx < len(disc.ParamMappings) {
|
||||||
if disc.ParamMappings[idx].ParamName != "" {
|
if disc.ParamMappings[idx].ParamName != "" {
|
||||||
@@ -3602,9 +4088,9 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "405") {
|
if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "405") || strings.Contains(err.Error(), "500") {
|
||||||
log.Printf("Gradio /call endpoint unavailable (%v), falling back to /run/predict protocol...", err)
|
log.Printf("Gradio /call endpoint unavailable (%v), falling back to /queue/join protocol...", err)
|
||||||
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
}
|
}
|
||||||
return fmt.Errorf("upstream Gradio call error: %w", err)
|
return fmt.Errorf("upstream Gradio call error: %w", err)
|
||||||
}
|
}
|
||||||
@@ -3613,8 +4099,8 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
|
|
||||||
var joinRes GradioJoinResponse
|
var joinRes GradioJoinResponse
|
||||||
if err := json.NewDecoder(resp.Body).Decode(&joinRes); err != nil || joinRes.EventID == "" {
|
if err := json.NewDecoder(resp.Body).Decode(&joinRes); err != nil || joinRes.EventID == "" {
|
||||||
log.Printf("Gradio /call returned non-SSE response, falling back to /run/predict protocol...")
|
log.Printf("Gradio /call returned non-SSE response, falling back to /queue/join protocol...")
|
||||||
return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 2. Connect to Gradio SSE EventStream
|
// 2. Connect to Gradio SSE EventStream
|
||||||
@@ -3652,7 +4138,8 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return fmt.Errorf("upstream Gradio stream error: %w", err)
|
log.Printf("Gradio /call stream connection failed (%v), falling back to /queue/join protocol...", err)
|
||||||
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
}
|
}
|
||||||
defer streamResp.Body.Close()
|
defer streamResp.Body.Close()
|
||||||
|
|
||||||
@@ -3687,6 +4174,10 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
errMsg := extractGradioErrorMessage(dataStr)
|
errMsg := extractGradioErrorMessage(dataStr)
|
||||||
|
if isProtocolOrInputError(errMsg) {
|
||||||
|
log.Printf("Upstream Gradio /call stream error (%s), falling back to /queue/join protocol...", errMsg)
|
||||||
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
|
}
|
||||||
log.Printf("Upstream Gradio error: %s", errMsg)
|
log.Printf("Upstream Gradio error: %s", errMsg)
|
||||||
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
}
|
}
|
||||||
@@ -3785,10 +4276,15 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
errMsg := extractGradioErrorMessage(dataStr)
|
errMsg := extractGradioErrorMessage(dataStr)
|
||||||
log.Printf("Upstream Gradio error: %s", errMsg)
|
|
||||||
if !streamer.started {
|
if !streamer.started {
|
||||||
|
if isProtocolOrInputError(errMsg) {
|
||||||
|
log.Printf("Upstream Gradio /call stream error (%s), falling back to /queue/join protocol...", errMsg)
|
||||||
|
return g.executeQueueCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA)
|
||||||
|
}
|
||||||
|
log.Printf("Upstream Gradio error: %s", errMsg)
|
||||||
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
return fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
}
|
}
|
||||||
|
log.Printf("Upstream Gradio error: %s", errMsg)
|
||||||
streamer.Error(errMsg)
|
streamer.Error(errMsg)
|
||||||
streamErr = fmt.Errorf("upstream Gradio error: %s", errMsg)
|
streamErr = fmt.Errorf("upstream Gradio error: %s", errMsg)
|
||||||
break
|
break
|
||||||
|
|||||||
+190
@@ -2321,6 +2321,196 @@ func TestGradio6CallV2Protocol(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user