fix(gateway): aggregate all text chunks in non-streaming Gemini responses
Previously, collectGeminiSSE() only returned the last chunk received
from the upstream streaming response when converting to non-streaming.
This caused incomplete responses where only the final text fragment
was returned to clients.
For example, a request asking to "count from 1 to 10" would only
return "\n" (the last chunk) instead of "1\n2\n3\n...\n10\n".
This was especially problematic for JSON structured output where
the opening brace "{" from the first chunk was lost, resulting
in invalid JSON like: colors": ["red", "blue"]}
The fix:
- Collect all text parts from each SSE chunk into a slice
- Merge all collected text parts into the final response
- Reuse the same pattern as handleGeminiStreamToNonStreaming
in antigravity_gateway_service.go
Fixes: non-streaming responses returning incomplete text
Fixes: structured output (JSON schema) returning invalid JSON
This commit is contained in:
@@ -1972,6 +1972,7 @@ func collectGeminiSSE(body io.Reader, isOAuth bool) (map[string]any, *ClaudeUsag
|
|||||||
|
|
||||||
var last map[string]any
|
var last map[string]any
|
||||||
var lastWithParts map[string]any
|
var lastWithParts map[string]any
|
||||||
|
var collectedTextParts []string // Collect all text parts for aggregation
|
||||||
usage := &ClaudeUsage{}
|
usage := &ClaudeUsage{}
|
||||||
|
|
||||||
for {
|
for {
|
||||||
@@ -1983,7 +1984,7 @@ func collectGeminiSSE(body io.Reader, isOAuth bool) (map[string]any, *ClaudeUsag
|
|||||||
switch payload {
|
switch payload {
|
||||||
case "", "[DONE]":
|
case "", "[DONE]":
|
||||||
if payload == "[DONE]" {
|
if payload == "[DONE]" {
|
||||||
return pickGeminiCollectResult(last, lastWithParts), usage, nil
|
return mergeCollectedTextParts(pickGeminiCollectResult(last, lastWithParts), collectedTextParts), usage, nil
|
||||||
}
|
}
|
||||||
default:
|
default:
|
||||||
var parsed map[string]any
|
var parsed map[string]any
|
||||||
@@ -2002,6 +2003,12 @@ func collectGeminiSSE(body io.Reader, isOAuth bool) (map[string]any, *ClaudeUsag
|
|||||||
}
|
}
|
||||||
if parts := extractGeminiParts(parsed); len(parts) > 0 {
|
if parts := extractGeminiParts(parsed); len(parts) > 0 {
|
||||||
lastWithParts = parsed
|
lastWithParts = parsed
|
||||||
|
// Collect text from each part for aggregation
|
||||||
|
for _, part := range parts {
|
||||||
|
if text, ok := part["text"].(string); ok && text != "" {
|
||||||
|
collectedTextParts = append(collectedTextParts, text)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2016,7 +2023,7 @@ func collectGeminiSSE(body io.Reader, isOAuth bool) (map[string]any, *ClaudeUsag
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return pickGeminiCollectResult(last, lastWithParts), usage, nil
|
return mergeCollectedTextParts(pickGeminiCollectResult(last, lastWithParts), collectedTextParts), usage, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func pickGeminiCollectResult(last map[string]any, lastWithParts map[string]any) map[string]any {
|
func pickGeminiCollectResult(last map[string]any, lastWithParts map[string]any) map[string]any {
|
||||||
@@ -2029,6 +2036,83 @@ func pickGeminiCollectResult(last map[string]any, lastWithParts map[string]any)
|
|||||||
return map[string]any{}
|
return map[string]any{}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// mergeCollectedTextParts merges all collected text chunks into the final response.
|
||||||
|
// This fixes the issue where non-streaming responses only returned the last chunk
|
||||||
|
// instead of the complete aggregated text.
|
||||||
|
func mergeCollectedTextParts(response map[string]any, textParts []string) map[string]any {
|
||||||
|
if len(textParts) == 0 {
|
||||||
|
return response
|
||||||
|
}
|
||||||
|
|
||||||
|
// Join all text parts
|
||||||
|
mergedText := strings.Join(textParts, "")
|
||||||
|
|
||||||
|
// Deep copy response
|
||||||
|
result := make(map[string]any)
|
||||||
|
for k, v := range response {
|
||||||
|
result[k] = v
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get or create candidates
|
||||||
|
candidates, ok := result["candidates"].([]any)
|
||||||
|
if !ok || len(candidates) == 0 {
|
||||||
|
candidates = []any{map[string]any{}}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get first candidate
|
||||||
|
candidate, ok := candidates[0].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
candidate = make(map[string]any)
|
||||||
|
candidates[0] = candidate
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get or create content
|
||||||
|
content, ok := candidate["content"].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
content = map[string]any{"role": "model"}
|
||||||
|
candidate["content"] = content
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get existing parts
|
||||||
|
existingParts, ok := content["parts"].([]any)
|
||||||
|
if !ok {
|
||||||
|
existingParts = []any{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find and update first text part, or create new one
|
||||||
|
newParts := make([]any, 0, len(existingParts)+1)
|
||||||
|
textUpdated := false
|
||||||
|
|
||||||
|
for _, p := range existingParts {
|
||||||
|
pm, ok := p.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
newParts = append(newParts, p)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, hasText := pm["text"]; hasText && !textUpdated {
|
||||||
|
// Replace with merged text
|
||||||
|
newPart := make(map[string]any)
|
||||||
|
for k, v := range pm {
|
||||||
|
newPart[k] = v
|
||||||
|
}
|
||||||
|
newPart["text"] = mergedText
|
||||||
|
newParts = append(newParts, newPart)
|
||||||
|
textUpdated = true
|
||||||
|
} else {
|
||||||
|
newParts = append(newParts, pm)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !textUpdated {
|
||||||
|
newParts = append([]any{map[string]any{"text": mergedText}}, newParts...)
|
||||||
|
}
|
||||||
|
|
||||||
|
content["parts"] = newParts
|
||||||
|
result["candidates"] = candidates
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
type geminiNativeStreamResult struct {
|
type geminiNativeStreamResult struct {
|
||||||
usage *ClaudeUsage
|
usage *ClaudeUsage
|
||||||
firstTokenMs *int
|
firstTokenMs *int
|
||||||
|
|||||||
Reference in New Issue
Block a user