fix: optimize Gemini stream parsing to support real-time streaming of JSON arrays
This commit is contained in:
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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,37 +551,50 @@ 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) {
|
||||||
scanner := bufio.NewScanner(r)
|
scanner := bufio.NewScanner(r)
|
||||||
|
|||||||
Reference in New Issue
Block a user