|
| 1 | +package providers |
| 2 | + |
| 3 | +import ( |
| 4 | + "bufio" |
| 5 | + "bytes" |
| 6 | + "encoding/json" |
| 7 | + "io" |
| 8 | +) |
| 9 | + |
| 10 | +// chatDonePayload terminates a chat completions SSE stream. |
| 11 | +var chatDonePayload = []byte("data: [DONE]\n\n") |
| 12 | + |
| 13 | +// peekForNonSSE inspects up to this many leading bytes to classify the upstream |
| 14 | +// response. SSE payloads begin with a field name (data:, event:, id:, retry:) or |
| 15 | +// a ':' comment; a buffered JSON completion begins with '{'. 512 bytes comfortably |
| 16 | +// clears any leading whitespace or comment lines without buffering real streams. |
| 17 | +const peekForNonSSE = 512 |
| 18 | + |
| 19 | +// EnsureChatCompletionSSE normalizes a chat completions stream so the client |
| 20 | +// always receives well-formed Server-Sent Events terminated by data: [DONE]. |
| 21 | +// |
| 22 | +// Some OpenAI-compatible upstreams ignore stream:true and reply with a single |
| 23 | +// buffered application/json completion (no data: framing, no [DONE]). Forwarding |
| 24 | +// that verbatim under a text/event-stream content type leaves SSE clients waiting |
| 25 | +// forever for an end-of-stream marker that never arrives. When the upstream body |
| 26 | +// is detected as a buffered JSON object it is re-emitted as one SSE chunk plus a |
| 27 | +// terminal [DONE]; genuine SSE streams pass through untouched with no buffering. |
| 28 | +func EnsureChatCompletionSSE(stream io.ReadCloser) io.ReadCloser { |
| 29 | + if stream == nil { |
| 30 | + return nil |
| 31 | + } |
| 32 | + |
| 33 | + reader := bufio.NewReaderSize(stream, peekForNonSSE) |
| 34 | + if firstNonSpaceByte(reader, peekForNonSSE) != '{' { |
| 35 | + // Genuine SSE (or empty): stream through unchanged, no buffering. |
| 36 | + return &bufferedReadCloser{Reader: reader, closer: stream} |
| 37 | + } |
| 38 | + |
| 39 | + // The '{' that classified this body is already buffered, so io.ReadAll |
| 40 | + // always returns at least that byte; a mid-read failure still yields the |
| 41 | + // partial bytes. Either way bufferedCompletionToSSE forwards what arrived |
| 42 | + // (raw when the JSON is truncated) and appends [DONE], so generated content |
| 43 | + // is never dropped and the client always receives a terminator. |
| 44 | + body, _ := io.ReadAll(reader) |
| 45 | + _ = stream.Close() //nolint:errcheck |
| 46 | + return io.NopCloser(bytes.NewReader(bufferedCompletionToSSE(body))) |
| 47 | +} |
| 48 | + |
| 49 | +// firstNonSpaceByte reports the first non-whitespace byte buffered by reader, |
| 50 | +// peeking one byte further at a time so a genuine SSE stream is classified from |
| 51 | +// its first token without blocking until a full buffer fills. It never consumes |
| 52 | +// input, so a passed-through stream is forwarded byte-for-byte. Returns 0 when |
| 53 | +// the stream ends, errors, or yields only whitespace within max bytes. |
| 54 | +func firstNonSpaceByte(r *bufio.Reader, max int) byte { |
| 55 | + for i := 1; i <= max; i++ { |
| 56 | + prefix, err := r.Peek(i) |
| 57 | + if len(prefix) < i { |
| 58 | + _ = err // EOF or error before any non-space byte was found |
| 59 | + return 0 |
| 60 | + } |
| 61 | + switch b := prefix[i-1]; b { |
| 62 | + case ' ', '\t', '\r', '\n': |
| 63 | + continue |
| 64 | + default: |
| 65 | + return b |
| 66 | + } |
| 67 | + } |
| 68 | + return 0 |
| 69 | +} |
| 70 | + |
| 71 | +// bufferedCompletionToSSE wraps a buffered chat completion JSON object as a |
| 72 | +// single SSE chunk followed by the terminal [DONE] marker. The object field is |
| 73 | +// rewritten to chat.completion.chunk and each choice's message is moved to delta |
| 74 | +// so OpenAI SSE clients parse it as a streaming chunk. If the body does not parse |
| 75 | +// as a JSON object it is forwarded as-is so no data is lost, still followed by |
| 76 | +// [DONE] so the client stops waiting. |
| 77 | +func bufferedCompletionToSSE(body []byte) []byte { |
| 78 | + payload := body |
| 79 | + var obj map[string]any |
| 80 | + if err := json.Unmarshal(body, &obj); err == nil { |
| 81 | + obj["object"] = "chat.completion.chunk" |
| 82 | + if choices, ok := obj["choices"].([]any); ok { |
| 83 | + for _, c := range choices { |
| 84 | + choice, ok := c.(map[string]any) |
| 85 | + if !ok { |
| 86 | + continue |
| 87 | + } |
| 88 | + if msg, ok := choice["message"]; ok { |
| 89 | + choice["delta"] = msg |
| 90 | + delete(choice, "message") |
| 91 | + } |
| 92 | + } |
| 93 | + } |
| 94 | + if encoded, err := json.Marshal(obj); err == nil { |
| 95 | + payload = encoded |
| 96 | + } |
| 97 | + } |
| 98 | + |
| 99 | + var out bytes.Buffer |
| 100 | + out.WriteString("data: ") |
| 101 | + out.Write(payload) |
| 102 | + out.WriteString("\n\n") |
| 103 | + out.Write(chatDonePayload) |
| 104 | + return out.Bytes() |
| 105 | +} |
| 106 | + |
| 107 | +// bufferedReadCloser pairs a buffered reader with the underlying stream's Close. |
| 108 | +type bufferedReadCloser struct { |
| 109 | + *bufio.Reader |
| 110 | + closer io.Closer |
| 111 | +} |
| 112 | + |
| 113 | +func (b *bufferedReadCloser) Close() error { return b.closer.Close() } |
0 commit comments