/
event_reader.go
68 lines (56 loc) · 1.91 KB
/
event_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
//go:build go1.18
// +build go1.18
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See License.txt in the project root for license information.
package azopenai
import (
"bufio"
"encoding/json"
"errors"
"io"
"strings"
)
// EventReader streams events dynamically from an OpenAI endpoint.
type EventReader[T any] struct {
reader io.ReadCloser // Required for Closing
scanner *bufio.Scanner
}
func newEventReader[T any](r io.ReadCloser) *EventReader[T] {
return &EventReader[T]{reader: r, scanner: bufio.NewScanner(r)}
}
// Read reads the next event from the stream.
// Returns io.EOF when there are no further events.
func (er *EventReader[T]) Read() (T, error) {
// https://html.spec.whatwg.org/multipage/server-sent-events.html
for er.scanner.Scan() { // Scan while no error
line := er.scanner.Text() // Get the line & interpret the event stream:
if line == "" || line[0] == ':' { // If the line is blank or is a comment, skip it
continue
}
if strings.Contains(line, ":") { // If the line contains a U+003A COLON character (:), process the field
tokens := strings.SplitN(line, ":", 2)
tokens[0], tokens[1] = strings.TrimSpace(tokens[0]), strings.TrimSpace(tokens[1])
var data T
switch tokens[0] {
case "data": // return the deserialized JSON object
if tokens[1] == "[DONE]" { // If data is [DONE], end of stream was reached
return data, io.EOF
}
err := json.Unmarshal([]byte(tokens[1]), &data)
return data, err
default: // Any other event type is an unexpected
return data, errors.New("unexpected event type: " + tokens[0])
}
// Unreachable
}
}
scannerErr := er.scanner.Err()
if scannerErr == nil {
return *new(T), errors.New("incomplete stream")
}
return *new(T), scannerErr
}
// Close closes the EventReader and any applicable inner stream state.
func (er *EventReader[T]) Close() error {
return er.reader.Close()
}