113 lines
2.5 KiB
Go
113 lines
2.5 KiB
Go
package proxy
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"strings"
|
|
)
|
|
|
|
type sseEventTracker struct {
|
|
buf []byte
|
|
limit int
|
|
overflow bool
|
|
recognized bool
|
|
terminal bool
|
|
afterTerminal bool
|
|
}
|
|
|
|
func newSSEEventTracker(limit int64) sseEventTracker {
|
|
if limit > int64(^uint(0)>>1) {
|
|
limit = int64(^uint(0) >> 1)
|
|
}
|
|
return sseEventTracker{limit: int(limit)}
|
|
}
|
|
|
|
func (t *sseEventTracker) Write(p []byte) {
|
|
if t.overflow {
|
|
return
|
|
}
|
|
if t.limit > 0 && len(p) > t.limit-len(t.buf) {
|
|
t.buf = nil
|
|
t.overflow = true
|
|
return
|
|
}
|
|
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 t.terminal {
|
|
t.afterTerminal = true
|
|
}
|
|
recognized, terminal := classifySSEEvent(event)
|
|
t.recognized = t.recognized || recognized
|
|
if terminal {
|
|
t.terminal = true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (t *sseEventTracker) Complete() bool {
|
|
return !t.overflow && len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
|
|
}
|
|
|
|
func completeSSEEvent(buf []byte) (int, int) {
|
|
lf := bytes.Index(buf, []byte("\n\n"))
|
|
crlf := bytes.Index(buf, []byte("\r\n\r\n"))
|
|
cr := bytes.Index(buf, []byte("\r\r"))
|
|
end, separator := lf, 2
|
|
if crlf >= 0 && (end < 0 || crlf < end) {
|
|
end, separator = crlf, 4
|
|
}
|
|
if cr >= 0 && (end < 0 || cr < end) {
|
|
end, separator = cr, 2
|
|
}
|
|
return end, separator
|
|
}
|
|
|
|
func classifySSEEvent(event []byte) (recognized, terminal bool) {
|
|
var data strings.Builder
|
|
normalized := strings.ReplaceAll(string(event), "\r\n", "\n")
|
|
for _, line := range strings.Split(strings.ReplaceAll(normalized, "\r", "\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, true
|
|
}
|
|
var envelope struct {
|
|
Type string `json:"type"`
|
|
Choices json.RawMessage `json:"choices"`
|
|
Candidates []struct {
|
|
FinishReason string `json:"finishReason"`
|
|
} `json:"candidates"`
|
|
}
|
|
if json.Unmarshal([]byte(payload), &envelope) != nil {
|
|
return false, false
|
|
}
|
|
if envelope.Type != "" {
|
|
return true, envelope.Type == "message_stop" || envelope.Type == "response.completed"
|
|
}
|
|
if envelope.Choices != nil {
|
|
return true, false
|
|
}
|
|
if envelope.Candidates != nil {
|
|
for _, candidate := range envelope.Candidates {
|
|
if candidate.FinishReason != "" {
|
|
return true, true
|
|
}
|
|
}
|
|
return true, false
|
|
}
|
|
return false, false
|
|
}
|