fix: optimize Gemini stream parsing to support real-time streaming of JSON arrays

This commit is contained in:
2026-07-17 18:52:54 +00:00
parent 293cf057b9
commit 9980123f97
2 changed files with 97 additions and 26 deletions
+55
View File
@@ -3,6 +3,7 @@ package providers
import ( import (
"encoding/json" "encoding/json"
"reflect" "reflect"
"strings"
"testing" "testing"
"gophergate/internal/models" "gophergate/internal/models"
@@ -209,3 +210,57 @@ func TestEmitGeminiChunk_ToolCalls(t *testing.T) {
t.Fatalf("expected response on channel") t.Fatalf("expected response on channel")
} }
} }
func TestStreamGeminiJSONArrayStream(t *testing.T) {
input := `[
{
"candidates": [
{
"content": {
"parts": [
{"text": "Hello"}
]
}
}
]
},
{
"candidates": [
{
"content": {
"parts": [
{"text": " world"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 2,
"totalTokenCount": 7
}
}
]`
ch := make(chan *models.ChatCompletionStreamResponse, 5)
defer close(ch)
streamGeminiJSONArrayStream(strings.NewReader(input), ch, "gemini-2.5-flash")
var texts []string
for i := 0; i < 2; i++ {
select {
case resp := <-ch:
if len(resp.Choices) > 0 && resp.Choices[0].Delta.Content != nil {
texts = append(texts, *resp.Choices[0].Delta.Content)
}
default:
t.Fatalf("expected chunk %d", i+1)
}
}
joined := strings.Join(texts, "")
if joined != "Hello world" {
t.Errorf("expected 'Hello world', got %q", joined)
}
}
+41 -25
View File
@@ -523,9 +523,12 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio
first := string(peek[:n]) first := string(peek[:n])
if first[0] == '[' { if first[0] == '[' {
// JSON array format // JSON array format — stream parse it in real-time
rest, _ := io.ReadAll(ctx) combined := io.MultiReader(
streamGeminiJSONArray(append([]byte(first), rest...), ch, model) strings.NewReader(string(peek[:n])),
ctx,
)
streamGeminiJSONArrayStream(combined, ch, model)
return return
} else if strings.HasPrefix(first, "data:") || strings.HasPrefix(first, "data: ") { } else if strings.HasPrefix(first, "data:") || strings.HasPrefix(first, "data: ") {
// SSE format — pre-pend the peeked bytes then run SSE scanner // SSE format — pre-pend the peeked bytes then run SSE scanner
@@ -548,36 +551,49 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio
return ch, nil return ch, nil
} }
// readAll reads remaining bytes from a reader (keeps the function signature simple func streamGeminiJSONArrayStream(r io.Reader, ch chan<- *models.ChatCompletionStreamResponse, model string) {
// for the JSON array fallback path). dec := json.NewDecoder(r)
func readAll(r io.Reader) []byte {
b, _ := io.ReadAll(r)
return b
}
func streamGeminiJSONArray(data []byte, ch chan<- *models.ChatCompletionStreamResponse, model string) { // Read open bracket '['
var chunks []geminiStreamChunk t, err := dec.Token()
if err := json.Unmarshal(data, &chunks); err != nil { if err != nil {
fmt.Printf("[Gemini-Stream] JSON array parse error: %v\n", err) fmt.Printf("[Gemini-Stream] JSON array token error: %v\n", err)
return return
} }
// Track the last chunk with usage for the final emission delim, ok := t.(json.Delim)
if !ok || delim != '[' {
fmt.Printf("[Gemini-Stream] JSON array expected '[', got %v\n", t)
return
}
var lastUsage *geminiStreamChunk var lastUsage *geminiStreamChunk
for i := range chunks {
if chunks[i].UsageMetadata.TotalTokenCount > 0 { // Read array elements
lastUsage = &chunks[i] for dec.More() {
var chunk geminiStreamChunk
if err := dec.Decode(&chunk); err != nil {
fmt.Printf("[Gemini-Stream] JSON array decode error: %v\n", err)
return
}
if chunk.UsageMetadata.TotalTokenCount > 0 {
temp := chunk
lastUsage = &temp
}
// Emit content-bearing chunks immediately
if len(chunk.Candidates) > 0 {
emitGeminiChunk(ch, &chunk, model)
} }
} }
if lastUsage != nil {
// Emit a synthetic final chunk with usage data // Read close bracket ']'
if len(lastUsage.Candidates) == 0 && lastUsage.UsageMetadata.TotalTokenCount > 0 { _, _ = dec.Token()
// Emit synthetic final chunk with usage if we collected it and it was not yet emitted
if lastUsage != nil && len(lastUsage.Candidates) == 0 && lastUsage.UsageMetadata.TotalTokenCount > 0 {
emitGeminiChunk(ch, lastUsage, model) emitGeminiChunk(ch, lastUsage, model)
} }
}
// Also emit each content-bearing chunk
for i := range chunks {
emitGeminiChunk(ch, &chunks[i], model)
}
} }
func streamGeminiSSE(r io.Reader, ch chan<- *models.ChatCompletionStreamResponse, model string) { func streamGeminiSSE(r io.Reader, ch chan<- *models.ChatCompletionStreamResponse, model string) {