ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/stream.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+)
+
+var streamError = []byte("data: {\"error\":{\"message\":\"Batch stream ended without a completed response\",\"type\":\"batch_error\"}}\n\n")
+
+type events struct {
+	pending  []byte
+	complete bool
+	done     bool
+}
+
+func (e *events) feed(body []byte, end bool) []byte {
+	e.pending = append(e.pending, body...)
+	var output []byte
+	for !e.done {
+		cut, size := bytes.Index(e.pending, []byte("\n\n")), 2
+		if cr := bytes.Index(e.pending, []byte("\r\n\r\n")); cr >= 0 && (cut < 0 || cr < cut) {
+			cut, size = cr, 4
+		}
+		if cut < 0 {
+			break
+		}
+		event := e.pending[:cut+size]
+		var data []byte
+		for line := range bytes.SplitSeq(e.pending[:cut], []byte("\n")) {
+			if bytes.HasPrefix(line, []byte("data:")) {
+				if len(data) > 0 {
+					data = append(data, '\n')
+				}
+				data = append(data, bytes.TrimSpace(line[5:])...)
+			}
+		}
+		if bytes.Equal(data, []byte("[DONE]")) {
+			e.done = true
+			if !e.complete {
+				output = append(output, streamError...)
+			} else {
+				output = append(output, event...)
+			}
+		} else {
+			var chunk struct {
+				Choices []struct {
+					Finish *string `json:"finish_reason"`
+				} `json:"choices"`
+			}
+			if json.Unmarshal(data, &chunk) == nil {
+				for _, choice := range chunk.Choices {
+					if choice.Finish != nil && *choice.Finish != "" {
+						e.complete = true
+					}
+				}
+			}
+			output = append(output, event...)
+		}
+		e.pending = e.pending[cut+size:]
+	}
+	if end && !e.done {
+		output = append(output, streamError...)
+		e.done = true
+	}
+	return output
+}