fix: optimize Gemini stream parsing to support real-time streaming of JSON arrays
This commit is contained in:
@@ -523,9 +523,12 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio
|
||||
first := string(peek[:n])
|
||||
|
||||
if first[0] == '[' {
|
||||
// JSON array format
|
||||
rest, _ := io.ReadAll(ctx)
|
||||
streamGeminiJSONArray(append([]byte(first), rest...), ch, model)
|
||||
// JSON array format — stream parse it in real-time
|
||||
combined := io.MultiReader(
|
||||
strings.NewReader(string(peek[:n])),
|
||||
ctx,
|
||||
)
|
||||
streamGeminiJSONArrayStream(combined, ch, model)
|
||||
return
|
||||
} else if strings.HasPrefix(first, "data:") || strings.HasPrefix(first, "data: ") {
|
||||
// SSE format — pre-pend the peeked bytes then run SSE scanner
|
||||
@@ -548,35 +551,48 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// readAll reads remaining bytes from a reader (keeps the function signature simple
|
||||
// for the JSON array fallback path).
|
||||
func readAll(r io.Reader) []byte {
|
||||
b, _ := io.ReadAll(r)
|
||||
return b
|
||||
}
|
||||
func streamGeminiJSONArrayStream(r io.Reader, ch chan<- *models.ChatCompletionStreamResponse, model string) {
|
||||
dec := json.NewDecoder(r)
|
||||
|
||||
func streamGeminiJSONArray(data []byte, ch chan<- *models.ChatCompletionStreamResponse, model string) {
|
||||
var chunks []geminiStreamChunk
|
||||
if err := json.Unmarshal(data, &chunks); err != nil {
|
||||
fmt.Printf("[Gemini-Stream] JSON array parse error: %v\n", err)
|
||||
// Read open bracket '['
|
||||
t, err := dec.Token()
|
||||
if err != nil {
|
||||
fmt.Printf("[Gemini-Stream] JSON array token error: %v\n", err)
|
||||
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
|
||||
for i := range chunks {
|
||||
if chunks[i].UsageMetadata.TotalTokenCount > 0 {
|
||||
lastUsage = &chunks[i]
|
||||
|
||||
// Read array elements
|
||||
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
|
||||
if len(lastUsage.Candidates) == 0 && lastUsage.UsageMetadata.TotalTokenCount > 0 {
|
||||
emitGeminiChunk(ch, lastUsage, model)
|
||||
}
|
||||
}
|
||||
// Also emit each content-bearing chunk
|
||||
for i := range chunks {
|
||||
emitGeminiChunk(ch, &chunks[i], model)
|
||||
|
||||
// Read close bracket ']'
|
||||
_, _ = 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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user