Add files via upload

internal/multiagent/eino_assistant_stream_event_handler.go

@@ -0,0 +1,167 @@
+package multiagent
+
+import (
+	"context"
+	"errors"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/schema"
+	"go.uber.org/zap"
+)
+
+type einoAssistantStreamEventHandlerConfig struct {
+	Context                   context.Context
+	ConversationID            string
+	OrchMode                  string
+	Progress                  func(eventType, message string, data interface{})
+	Logger                    *zap.Logger
+	SnapshotMCPIDs            func() []string
+	StreamsMainAssistant      func(agent string) bool
+	EinoRoleTag               func(agent string) string
+	RunProgress               *einoRunProgressTracker
+	StdoutSuppressor          *einoExecuteStdoutSuppressor
+	AssistantOutput           *einoAssistantOutputAccumulator
+	RunMessages               *einoRunMessageAccumulator
+	Usage                     *einoRunUsageAccumulator
+	ToolCallCompletion        *einoStreamToolCallCompletionHandler
+	NextMainStreamID          func() string
+	NextReasoningStreamID     func() string
+	NextSubAgentReplyStreamID func() string
+}
+
+type einoAssistantStreamEventHandler struct {
+	ctx                       context.Context
+	conversationID            string
+	orchMode                  string
+	progress                  func(eventType, message string, data interface{})
+	logger                    *zap.Logger
+	snapshotMCPIDs            func() []string
+	streamsMainAssistant      func(agent string) bool
+	einoRoleTag               func(agent string) string
+	runProgress               *einoRunProgressTracker
+	stdoutSuppressor          *einoExecuteStdoutSuppressor
+	assistantOutput           *einoAssistantOutputAccumulator
+	runMessages               *einoRunMessageAccumulator
+	usage                     *einoRunUsageAccumulator
+	toolCallCompletion        *einoStreamToolCallCompletionHandler
+	nextMainStreamID          func() string
+	nextReasoningStreamID     func() string
+	nextSubAgentReplyStreamID func() string
+}
+
+func newEinoAssistantStreamEventHandler(cfg einoAssistantStreamEventHandlerConfig) *einoAssistantStreamEventHandler {
+	if cfg.Context == nil {
+		cfg.Context = context.Background()
+	}
+	if cfg.SnapshotMCPIDs == nil {
+		cfg.SnapshotMCPIDs = func() []string { return nil }
+	}
+	if cfg.StreamsMainAssistant == nil {
+		cfg.StreamsMainAssistant = func(string) bool { return true }
+	}
+	if cfg.EinoRoleTag == nil {
+		cfg.EinoRoleTag = func(string) string { return "" }
+	}
+	if cfg.NextMainStreamID == nil {
+		cfg.NextMainStreamID = func() string { return "eino-main" }
+	}
+	if cfg.NextReasoningStreamID == nil {
+		cfg.NextReasoningStreamID = func() string { return "eino-reasoning" }
+	}
+	if cfg.NextSubAgentReplyStreamID == nil {
+		cfg.NextSubAgentReplyStreamID = func() string { return "eino-sub-reply" }
+	}
+	return &einoAssistantStreamEventHandler{
+		ctx:                       cfg.Context,
+		conversationID:            cfg.ConversationID,
+		orchMode:                  cfg.OrchMode,
+		progress:                  cfg.Progress,
+		logger:                    cfg.Logger,
+		snapshotMCPIDs:            cfg.SnapshotMCPIDs,
+		streamsMainAssistant:      cfg.StreamsMainAssistant,
+		einoRoleTag:               cfg.EinoRoleTag,
+		runProgress:               cfg.RunProgress,
+		stdoutSuppressor:          cfg.StdoutSuppressor,
+		assistantOutput:           cfg.AssistantOutput,
+		runMessages:               cfg.RunMessages,
+		usage:                     cfg.Usage,
+		toolCallCompletion:        cfg.ToolCallCompletion,
+		nextMainStreamID:          cfg.NextMainStreamID,
+		nextReasoningStreamID:     cfg.NextReasoningStreamID,
+		nextSubAgentReplyStreamID: cfg.NextSubAgentReplyStreamID,
+	}
+}
+
+func (h *einoAssistantStreamEventHandler) Handle(mv *adk.MessageVariant, agentName string) (handled bool, recvErr error) {
+	if h == nil || mv == nil || !mv.IsStreaming || mv.MessageStream == nil || mv.Role == schema.Tool {
+		return false, nil
+	}
+	mainStreamID := h.nextMainStreamID()
+	mainEmitter := newEinoMainResponseStreamEmitter(
+		h.conversationID, h.orchMode, agentName, mainStreamID, h.mainIteration(agentName), h.progress, h.snapshotMCPIDs,
+	)
+	reasoningEmitter := newEinoReasoningStreamEmitter(
+		h.conversationID,
+		h.orchMode,
+		agentName,
+		h.einoRoleTag(agentName),
+		h.progress,
+		h.nextReasoningStreamID,
+	)
+	var toolStreamFragments []schema.ToolCall
+	var streamUsage *schema.TokenUsage
+	subReplyEmitter := newEinoSubAgentReplyEmitter(
+		h.conversationID,
+		agentName,
+		h.progress,
+		h.nextSubAgentReplyStreamID,
+	)
+	mainAssistantStream := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{
+		AgentName:        agentName,
+		Emitter:          mainEmitter,
+		StdoutSuppressor: h.stdoutSuppressor,
+		AssistantOutput:  h.assistantOutput,
+		RunMessages:      h.runMessages,
+	})
+	recvErr = recvEinoSchemaMessageStreamWithContext(h.ctx, mv.MessageStream, 8, func(chunk *schema.Message) {
+		reasoningEmitter.EmitDelta(chunk.ReasoningContent)
+		if chunk.Content != "" {
+			if h.streamsMainAssistant(agentName) {
+				mainAssistantStream.EmitDelta(chunk.Content)
+			} else if !h.streamsMainAssistant(agentName) {
+				subReplyEmitter.EmitDelta(chunk.Content)
+			}
+		}
+		if len(chunk.ToolCalls) > 0 {
+			toolStreamFragments = append(toolStreamFragments, chunk.ToolCalls...)
+		}
+		if chunk.ResponseMeta != nil && chunk.ResponseMeta.Usage != nil {
+			streamUsage = maxEinoTokenUsage(streamUsage, chunk.ResponseMeta.Usage)
+		}
+	})
+	if recvErr != nil && !errors.Is(recvErr, context.Canceled) && h.logger != nil {
+		h.logger.Warn("eino stream recv error, flushing incomplete stream",
+			zap.Error(recvErr),
+			zap.String("agent", agentName),
+			zap.Int("toolFragments", len(toolStreamFragments)))
+	}
+	reasoningEmitter.Finish()
+	if h.streamsMainAssistant(agentName) {
+		mainAssistantStream.Finish()
+	}
+	subReplyEmitter.Finish()
+	if h.toolCallCompletion != nil {
+		h.toolCallCompletion.Complete(toolStreamFragments, agentName)
+	}
+	if h.usage != nil {
+		h.usage.AddUsage(streamUsage)
+	}
+	return true, recvErr
+}
+
+func (h *einoAssistantStreamEventHandler) mainIteration(agentName string) int {
+	if h == nil || h.runProgress == nil {
+		return 0
+	}
+	return h.runProgress.MainIteration(agentName)
+}