use cut prefix

This commit is contained in:
Michael Yang 2023-10-16 16:31:56 -07:00
parent 094df37563
commit b36b0b71f8

View file

@ -531,16 +531,14 @@ func (llm *llama) Predict(ctx context.Context, prevContext []int, prompt string,
// This handles the request cancellation // This handles the request cancellation
return ctx.Err() return ctx.Err()
default: default:
line := scanner.Text() line := scanner.Bytes()
if line == "" { if len(line) == 0 {
continue continue
} }
// Read data from the server-side event stream if evt, ok := bytes.CutPrefix(line, []byte("data: ")); ok {
if strings.HasPrefix(line, "data: ") {
evt := line[6:]
var p prediction var p prediction
if err := json.Unmarshal([]byte(evt), &p); err != nil { if err := json.Unmarshal(evt, &p); err != nil {
return fmt.Errorf("error unmarshaling llm prediction response: %v", err) return fmt.Errorf("error unmarshaling llm prediction response: %v", err)
} }