66 lines
1.5 KiB
Go
66 lines
1.5 KiB
Go
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
|
|
}
|