fix: workflow onEndWithStream handles stream asynchronously (#518)
This commit is contained in:
parent
08d1f6bf67
commit
3fe4031531
|
|
@ -496,48 +496,49 @@ func (w *WorkflowHandler) OnEndWithStreamOutput(ctx context.Context, info *callb
|
||||||
return ctx
|
return ctx
|
||||||
}
|
}
|
||||||
|
|
||||||
// consumes the stream synchronously because the Exit node has already processed this stream synchronously.
|
safego.Go(ctx, func() {
|
||||||
defer output.Close()
|
defer output.Close()
|
||||||
fullOutput := make(map[string]any)
|
fullOutput := make(map[string]any)
|
||||||
for {
|
for {
|
||||||
chunk, e := output.Recv()
|
chunk, e := output.Recv()
|
||||||
if e != nil {
|
if e != nil {
|
||||||
if e == io.EOF {
|
if e == io.EOF {
|
||||||
break
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, ok := schema.GetSourceName(e); ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
logs.Errorf("workflow OnEndWithStreamOutput failed to receive stream output: %v", e)
|
||||||
|
_ = w.OnError(ctx, info, e)
|
||||||
|
return
|
||||||
}
|
}
|
||||||
|
fullOutput, e = nodes.ConcatTwoMaps(fullOutput, chunk.(map[string]any))
|
||||||
if _, ok := schema.GetSourceName(e); ok {
|
if e != nil {
|
||||||
continue
|
logs.Errorf("failed to concat two maps: %v", e)
|
||||||
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
logs.Errorf("workflow OnEndWithStreamOutput failed to receive stream output: %v", e)
|
|
||||||
_ = w.OnError(ctx, info, e)
|
|
||||||
return ctx
|
|
||||||
}
|
}
|
||||||
fullOutput, e = nodes.ConcatTwoMaps(fullOutput, chunk.(map[string]any))
|
|
||||||
if e != nil {
|
|
||||||
logs.Errorf("failed to concat two maps: %v", e)
|
|
||||||
return ctx
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
c := GetExeCtx(ctx)
|
c := GetExeCtx(ctx)
|
||||||
e := &Event{
|
e := &Event{
|
||||||
Type: WorkflowSuccess,
|
Type: WorkflowSuccess,
|
||||||
Context: c,
|
Context: c,
|
||||||
Duration: time.Since(time.UnixMilli(c.StartTime)),
|
Duration: time.Since(time.UnixMilli(c.StartTime)),
|
||||||
Output: fullOutput,
|
Output: fullOutput,
|
||||||
}
|
|
||||||
|
|
||||||
if c.TokenCollector != nil {
|
|
||||||
usage := c.TokenCollector.wait()
|
|
||||||
e.Token = &TokenInfo{
|
|
||||||
InputToken: int64(usage.PromptTokens),
|
|
||||||
OutputToken: int64(usage.CompletionTokens),
|
|
||||||
TotalToken: int64(usage.TotalTokens),
|
|
||||||
}
|
}
|
||||||
}
|
|
||||||
w.ch <- e
|
if c.TokenCollector != nil {
|
||||||
|
usage := c.TokenCollector.wait()
|
||||||
|
e.Token = &TokenInfo{
|
||||||
|
InputToken: int64(usage.PromptTokens),
|
||||||
|
OutputToken: int64(usage.CompletionTokens),
|
||||||
|
TotalToken: int64(usage.TotalTokens),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
w.ch <- e
|
||||||
|
})
|
||||||
|
|
||||||
return ctx
|
return ctx
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue