mirror of
https://gitee.com/dromara/mayfly-go
synced 2026-05-20 09:55:25 +08:00
36 lines
607 B
Go
36 lines
607 B
Go
package middleware
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
func singleChunkReader(msg string) *schema.StreamReader[string] {
|
|
r, w := schema.Pipe[string](1)
|
|
_ = w.Send(msg, nil)
|
|
w.Close()
|
|
return r
|
|
}
|
|
|
|
func safeWrapReader(sr *schema.StreamReader[string]) *schema.StreamReader[string] {
|
|
r, w := schema.Pipe[string](64)
|
|
go func() {
|
|
defer w.Close()
|
|
for {
|
|
chunk, err := sr.Recv()
|
|
if errors.Is(err, io.EOF) {
|
|
return
|
|
}
|
|
if err != nil {
|
|
_ = w.Send(fmt.Sprintf("\n[tool error] %v", err), nil)
|
|
return
|
|
}
|
|
_ = w.Send(chunk, nil)
|
|
}
|
|
}()
|
|
return r
|
|
}
|