mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-18 21:57:27 +08:00
fix(go-agent): filter SSE done sentinel from model stream (#17034)
## Summary - Filter the `[DONE]` transport sentinel from Eino model streams - Add regression coverage for stream messages and callbacks ## Testing - `bash build.sh --test ./internal/entity/models` - `bash build.sh --test ./internal/agent/component` - `bash build.sh --test -run 'Test(AgentChatCompletions_Stream|RunAgent_Stream)' ./internal/handler`
This commit is contained in:
@@ -293,6 +293,12 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts
|
||||
if content == nil && reasoning == nil {
|
||||
return nil
|
||||
}
|
||||
// Provider drivers use the OpenAI-compatible [DONE] sentinel to
|
||||
// signal the end of their transport stream. It is not assistant
|
||||
// content and must not reach Eino's message stream or callback.
|
||||
if content != nil && *content == "[DONE]" {
|
||||
return nil
|
||||
}
|
||||
msg := &schema.Message{Role: schema.Assistant}
|
||||
if content != nil {
|
||||
msg.Content = *content
|
||||
|
||||
@@ -2,11 +2,53 @@ package models
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestEinoChatModelStreamFiltersDoneSentinel(t *testing.T) {
|
||||
modelName := "chat"
|
||||
driver := &streamSentinelDriver{captureToolDriver: &captureToolDriver{}}
|
||||
var callbacks []string
|
||||
base := NewChatModel(driver, &modelName, &APIConfig{})
|
||||
model := NewEinoChatModel(base, &ChatConfig{
|
||||
StreamCallback: func(content, _ string) {
|
||||
if content != "" {
|
||||
callbacks = append(callbacks, content)
|
||||
}
|
||||
},
|
||||
})
|
||||
|
||||
stream, err := model.Stream(context.Background(), []*schema.Message{schema.UserMessage("hello")})
|
||||
if err != nil {
|
||||
t.Fatalf("Stream: %v", err)
|
||||
}
|
||||
|
||||
var messages []string
|
||||
for {
|
||||
msg, recvErr := stream.Recv()
|
||||
if errors.Is(recvErr, io.EOF) {
|
||||
break
|
||||
}
|
||||
if recvErr != nil {
|
||||
t.Fatalf("stream.Recv: %v", recvErr)
|
||||
}
|
||||
if msg != nil {
|
||||
messages = append(messages, msg.Content)
|
||||
}
|
||||
}
|
||||
|
||||
if len(messages) != 2 || messages[0] != "answer" || messages[1] != "DONE!" {
|
||||
t.Fatalf("stream messages = %#v, want [answer DONE!]", messages)
|
||||
}
|
||||
if len(callbacks) != 2 || callbacks[0] != "answer" || callbacks[1] != "DONE!" {
|
||||
t.Fatalf("stream callbacks = %#v, want [answer DONE!]", callbacks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoChatModelGenerateSendsBoundTools(t *testing.T) {
|
||||
apiKey := "key"
|
||||
modelName := "chat"
|
||||
@@ -148,6 +190,23 @@ type captureToolDriver struct {
|
||||
lastConfig *ChatConfig
|
||||
}
|
||||
|
||||
type streamSentinelDriver struct {
|
||||
*captureToolDriver
|
||||
}
|
||||
|
||||
func (d *streamSentinelDriver) ChatStreamlyWithSender(_ string, _ []Message, _ *APIConfig, _ *ChatConfig, sender func(*string, *string) error) error {
|
||||
answer := "answer"
|
||||
if err := sender(&answer, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
visibleDone := "DONE!"
|
||||
if err := sender(&visibleDone, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
done := "[DONE]"
|
||||
return sender(&done, nil)
|
||||
}
|
||||
|
||||
func (d *captureToolDriver) NewInstance(baseURL map[string]string) ModelDriver { return d }
|
||||
func (d *captureToolDriver) Name() string { return "capture" }
|
||||
func (d *captureToolDriver) ChatWithMessages(_ string, _ []Message, _ *APIConfig, cfg *ChatConfig) (*ChatResponse, error) {
|
||||
|
||||
Reference in New Issue
Block a user