66 lines
1.3 KiB
Go
66 lines
1.3 KiB
Go
package proxy
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"strings"
|
|
)
|
|
|
|
type sseEventTracker struct {
|
|
buf []byte
|
|
terminal bool
|
|
}
|
|
|
|
func (t *sseEventTracker) Write(p []byte) {
|
|
t.buf = append(t.buf, p...)
|
|
for {
|
|
end, separator := completeSSEEvent(t.buf)
|
|
if end < 0 {
|
|
return
|
|
}
|
|
event := t.buf[:end]
|
|
t.buf = t.buf[end+separator:]
|
|
if terminalSSEEvent(event) {
|
|
t.terminal = true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (t *sseEventTracker) Complete() bool { return len(t.buf) == 0 }
|
|
|
|
func completeSSEEvent(buf []byte) (int, int) {
|
|
lf := bytes.Index(buf, []byte("\n\n"))
|
|
crlf := bytes.Index(buf, []byte("\r\n\r\n"))
|
|
if crlf >= 0 && (lf < 0 || crlf < lf) {
|
|
return crlf, 4
|
|
}
|
|
if lf >= 0 {
|
|
return lf, 2
|
|
}
|
|
return -1, 0
|
|
}
|
|
|
|
func terminalSSEEvent(event []byte) bool {
|
|
var data strings.Builder
|
|
for _, line := range strings.Split(strings.ReplaceAll(string(event), "\r\n", "\n"), "\n") {
|
|
if !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
if data.Len() > 0 {
|
|
data.WriteByte('\n')
|
|
}
|
|
data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
|
}
|
|
payload := data.String()
|
|
if payload == "[DONE]" {
|
|
return true
|
|
}
|
|
var envelope struct {
|
|
Type string `json:"type"`
|
|
}
|
|
if json.Unmarshal([]byte(payload), &envelope) != nil {
|
|
return false
|
|
}
|
|
return envelope.Type == "message_stop" || envelope.Type == "response.completed"
|
|
}
|