github.com/alibabacloud-go/tea@v1.3.10/dara/stream.go (about)

     1  package dara
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"encoding/json"
     7  	"io"
     8  	"io/ioutil"
     9  	"strings"
    10  	"fmt"
    11  )
    12  
    13  // 定义 Event 结构体
    14  type SSEEvent struct {
    15  	ID    *string
    16  	Event *string
    17  	Data  *string
    18  	Retry *int
    19  }
    20  
    21  // 解析单个事件
    22  func parseEvent(lines []string) *SSEEvent {
    23  	event := &SSEEvent{}
    24  	for _, line := range lines {
    25  		if strings.HasPrefix(line, "data: ") {
    26  			data := strings.TrimPrefix(line, "data: ") + "\n"
    27  			if event.Data == nil {
    28  				event.Data = new(string)
    29  			}
    30  			*event.Data += data
    31  		} else if strings.HasPrefix(line, "event: ") {
    32  			eventName := strings.TrimPrefix(line, "event: ")
    33  			event.Event = &eventName
    34  		} else if strings.HasPrefix(line, "id: ") {
    35  			id := strings.TrimPrefix(line, "id: ")
    36  			event.ID = &id
    37  		} else if strings.HasPrefix(line, "retry: ") {
    38  			var retry int
    39  			fmt.Sscanf(strings.TrimPrefix(line, "retry: "), "%d", &retry)
    40  			event.Retry = &retry
    41  		}
    42  	}
    43  	// Remove last newline from data
    44  	if event.Data != nil {
    45  		data := strings.TrimRight(*event.Data, "\n")
    46  		event.Data = &data
    47  	}
    48  	return event
    49  }
    50  
    51  func ReadAsBytes(body io.Reader) ([]byte, error) {
    52  	byt, err := ioutil.ReadAll(body)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  	r, ok := body.(io.ReadCloser)
    57  	if ok {
    58  		r.Close()
    59  	}
    60  	return byt, nil
    61  }
    62  
    63  func ReadAsJSON(body io.Reader) (result interface{}, err error) {
    64  	byt, err := ioutil.ReadAll(body)
    65  	if err != nil {
    66  		return
    67  	}
    68  	if string(byt) == "" {
    69  		return
    70  	}
    71  	r, ok := body.(io.ReadCloser)
    72  	if ok {
    73  		r.Close()
    74  	}
    75  	d := json.NewDecoder(bytes.NewReader(byt))
    76  	d.UseNumber()
    77  	err = d.Decode(&result)
    78  	return
    79  }
    80  
    81  func ReadAsString(body io.Reader) (string, error) {
    82  	byt, err := ioutil.ReadAll(body)
    83  	if err != nil {
    84  		return "", err
    85  	}
    86  	r, ok := body.(io.ReadCloser)
    87  	if ok {
    88  		r.Close()
    89  	}
    90  	return string(byt), nil
    91  }
    92  
    93  func ReadAsSSE(body io.ReadCloser, eventChannel chan *SSEEvent, errorChannel chan error) {
    94  
    95  	go func() {
    96  		defer func() {
    97  			body.Close()
    98  			close(eventChannel)
    99  		}()
   100  
   101  		reader := bufio.NewReader(body)
   102  		var eventLines []string
   103  
   104  		for {
   105  			line, err := reader.ReadString('\n')
   106  			if err != nil {
   107  				if err == io.EOF {
   108  					// Handle the end of the stream and possibly pending event
   109  					if len(eventLines) > 0 {
   110  						event := parseEvent(eventLines)
   111  						eventChannel <- event
   112  					}
   113  					errorChannel <- nil
   114  					return
   115  				}
   116  				errorChannel <- err
   117  				return
   118  			}
   119  
   120  			line = strings.TrimRight(line, "\n")
   121  
   122  			if line == "" {
   123  				// End of an SSE event
   124  				if len(eventLines) > 0 {
   125  					event := parseEvent(eventLines)
   126  					eventChannel <- event
   127  					eventLines = []string{} // Reset for the next event
   128  				}
   129  				continue
   130  			}
   131  
   132  			eventLines = append(eventLines, line)
   133  		}
   134  	}()
   135  
   136  	return
   137  }