package pregel import ( "context" "fmt" "sync" "ragflow/internal/harness/graph/channels" ) // ChannelRead represents a read operation from channels. // It encapsulates the logic for reading state from multiple channels. type ChannelRead struct { registry *channels.Registry selector ChannelSelector transformer ChannelTransformer mu sync.RWMutex } // ChannelSelector selects which channels to read. type ChannelSelector interface { Select(registry *channels.Registry) ([]string, error) } // ChannelTransformer transforms the raw channel values. type ChannelTransformer interface { Transform(values map[string]any) (map[string]any, error) } // NewChannelRead creates a new channel read operation. func NewChannelRead(registry *channels.Registry, opts ...ChannelReadOption) *ChannelRead { cr := &ChannelRead{ registry: registry, selector: &AllChannelsSelector{}, transformer: &IdentityTransformer{}, } for _, opt := range opts { opt(cr) } return cr } // ChannelReadOption configures a ChannelRead. type ChannelReadOption func(*ChannelRead) // WithSelector sets the channel selector. func WithSelector(selector ChannelSelector) ChannelReadOption { return func(cr *ChannelRead) { cr.selector = selector } } // WithTransformer sets the channel transformer. func WithTransformer(transformer ChannelTransformer) ChannelReadOption { return func(cr *ChannelRead) { cr.transformer = transformer } } // Read performs the read operation. func (cr *ChannelRead) Read(ctx context.Context) (map[string]any, error) { cr.mu.RLock() defer cr.mu.RUnlock() // Select channels to read channelNames, err := cr.selector.Select(cr.registry) if err != nil { return nil, fmt.Errorf("channel selection failed: %w", err) } // Read values from channels values := make(map[string]any) for _, name := range channelNames { if ch, ok := cr.registry.Get(name); ok { val, err := ch.Get() if err == nil { values[name] = val } } } // Transform values if cr.transformer != nil { transformed, err := cr.transformer.Transform(values) if err != nil { return nil, fmt.Errorf("channel transformation failed: %w", err) } values = transformed } return values, nil } // ReadChannel reads a single channel by name. func (cr *ChannelRead) ReadChannel(ctx context.Context, name string) (any, error) { cr.mu.RLock() defer cr.mu.RUnlock() if ch, ok := cr.registry.Get(name); ok { return ch.Get() } return nil, fmt.Errorf("channel not found: %s", name) } // HasChannel checks if a channel exists. func (cr *ChannelRead) HasChannel(name string) bool { cr.mu.RLock() defer cr.mu.RUnlock() _, ok := cr.registry.Get(name) return ok } // ListChannels returns all available channel names. func (cr *ChannelRead) ListChannels() []string { cr.mu.RLock() defer cr.mu.RUnlock() return cr.registry.List() } // ==================== Channel Selectors ==================== // AllChannelsSelector selects all available channels. type AllChannelsSelector struct{} func (s *AllChannelsSelector) Select(registry *channels.Registry) ([]string, error) { return registry.List(), nil } // SpecificChannelsSelector selects specific channels. type SpecificChannelsSelector struct { channels []string } // NewSpecificChannelsSelector creates a selector for specific channels. func NewSpecificChannelsSelector(channels ...string) *SpecificChannelsSelector { return &SpecificChannelsSelector{channels: channels} } func (s *SpecificChannelsSelector) Select(registry *channels.Registry) ([]string, error) { result := make([]string, 0, len(s.channels)) for _, name := range s.channels { if _, ok := registry.Get(name); ok { result = append(result, name) } } return result, nil } // PrefixChannelsSelector selects channels with a specific prefix. type PrefixChannelsSelector struct { prefix string } // NewPrefixChannelsSelector creates a selector for channels with a prefix. func NewPrefixChannelsSelector(prefix string) *PrefixChannelsSelector { return &PrefixChannelsSelector{prefix: prefix} } func (s *PrefixChannelsSelector) Select(registry *channels.Registry) ([]string, error) { all := registry.List() result := make([]string, 0) for _, name := range all { if len(name) >= len(s.prefix) && name[:len(s.prefix)] == s.prefix { result = append(result, name) } } return result, nil } // AvailableChannelsSelector selects only available (non-empty) channels. type AvailableChannelsSelector struct{} func (s *AvailableChannelsSelector) Select(registry *channels.Registry) ([]string, error) { all := registry.List() result := make([]string, 0) for _, name := range all { if ch, ok := registry.Get(name); ok && ch.IsAvailable() { result = append(result, name) } } return result, nil } // ==================== Channel Transformers ==================== // IdentityTransformer returns values as-is. type IdentityTransformer struct{} func (t *IdentityTransformer) Transform(values map[string]any) (map[string]any, error) { return values, nil } // MappingTransformer renames channels. type MappingTransformer struct { mappings map[string]string } // NewMappingTransformer creates a transformer with channel name mappings. func NewMappingTransformer(mappings map[string]string) *MappingTransformer { return &MappingTransformer{mappings: mappings} } func (t *MappingTransformer) Transform(values map[string]any) (map[string]any, error) { result := make(map[string]any) for oldName, value := range values { newName := oldName if mapped, ok := t.mappings[oldName]; ok { newName = mapped } result[newName] = value } return result, nil } // FilterTransformer filters channels. type FilterTransformer struct { filter map[string]bool } // NewFilterTransformer creates a transformer that filters channels. func NewFilterTransformer(keep ...string) *FilterTransformer { filter := make(map[string]bool) for _, name := range keep { filter[name] = true } return &FilterTransformer{filter: filter} } func (t *FilterTransformer) Transform(values map[string]any) (map[string]any, error) { result := make(map[string]any) for name, value := range values { if t.filter[name] { result[name] = value } } return result, nil } // DefaultTransformer provides default values for missing channels. type DefaultTransformer struct { defaults map[string]any } // NewDefaultTransformer creates a transformer with default values. func NewDefaultTransformer(defaults map[string]any) *DefaultTransformer { return &DefaultTransformer{defaults: defaults} } func (t *DefaultTransformer) Transform(values map[string]any) (map[string]any, error) { result := make(map[string]any) for name, defValue := range t.defaults { if value, ok := values[name]; ok { result[name] = value } else { result[name] = defValue } } // Add any extra values that don't have defaults for name, value := range values { if _, ok := result[name]; !ok { result[name] = value } } return result, nil } // MergingTransformer merges channels into a single value. type MergingTransformer struct { target string merger func([]any) (any, error) } // NewMergingTransformer creates a transformer that merges channels. func NewMergingTransformer(target string, merger func([]any) (any, error)) *MergingTransformer { return &MergingTransformer{ target: target, merger: merger, } } func (t *MergingTransformer) Transform(values map[string]any) (map[string]any, error) { // Collect all values items := make([]any, 0, len(values)) for _, value := range values { items = append(items, value) } // Merge merged, err := t.merger(items) if err != nil { return nil, err } // Return merged value under target key return map[string]any{t.target: merged}, nil } // ==================== Triggers ==================== // Trigger determines when a channel read should execute. type Trigger interface { ShouldTrigger(registry *channels.Registry) bool } // AlwaysTrigger always triggers. type AlwaysTrigger struct{} func (t *AlwaysTrigger) ShouldTrigger(registry *channels.Registry) bool { return true } // AnyAvailableTrigger triggers when any selected channel is available. type AnyAvailableTrigger struct { channels []string } // NewAnyAvailableTrigger creates a trigger that fires when any channel is available. func NewAnyAvailableTrigger(channels ...string) *AnyAvailableTrigger { return &AnyAvailableTrigger{channels: channels} } func (t *AnyAvailableTrigger) ShouldTrigger(registry *channels.Registry) bool { for _, name := range t.channels { if ch, ok := registry.Get(name); ok && ch.IsAvailable() { return true } } return false } // AllAvailableTrigger triggers when all selected channels are available. type AllAvailableTrigger struct { channels []string } // NewAllAvailableTrigger creates a trigger that fires when all channels are available. func NewAllAvailableTrigger(channels ...string) *AllAvailableTrigger { return &AllAvailableTrigger{channels: channels} } func (t *AllAvailableTrigger) ShouldTrigger(registry *channels.Registry) bool { for _, name := range t.channels { if ch, ok := registry.Get(name); !ok || !ch.IsAvailable() { return false } } return true } // ChannelChangedTrigger triggers when a specific channel's value changes. type ChannelChangedTrigger struct { channel string lastVersion int64 } // NewChannelChangedTrigger creates a trigger that fires when a channel changes. func NewChannelChangedTrigger(channel string) *ChannelChangedTrigger { return &ChannelChangedTrigger{ channel: channel, lastVersion: -1, } } func (t *ChannelChangedTrigger) ShouldTrigger(registry *channels.Registry) bool { if ch, ok := registry.Get(t.channel); ok { // Track the channel version to detect actual changes (not just availability). // GetVersion returns -1 if the channel does not support versioning, in which // case we fall back to IsAvailable() for backward compatibility. version := int64(ch.GetVersion()) if version >= 0 && version != t.lastVersion { t.lastVersion = version return true } // Fallback for channels without version tracking. return version < 0 && ch.IsAvailable() } return false } // ==================== Utility Functions ==================== // ReadContext represents the context of a channel read operation. type ReadContext struct { Node string Step int Triggers []Trigger Readers map[string]*ChannelRead } // NewReadContext creates a new read context. func NewReadContext(node string, step int) *ReadContext { return &ReadContext{ Node: node, Step: step, Triggers: make([]Trigger, 0), Readers: make(map[string]*ChannelRead), } } // AddReader adds a channel reader. func (rc *ReadContext) AddReader(name string, reader *ChannelRead) { rc.Readers[name] = reader } // GetReader gets a channel reader by name. func (rc *ReadContext) GetReader(name string) *ChannelRead { return rc.Readers[name] } // ShouldExecute checks if any trigger fires. func (rc *ReadContext) ShouldExecute(registry *channels.Registry) bool { if len(rc.Triggers) == 0 { return true } for _, trigger := range rc.Triggers { if trigger.ShouldTrigger(registry) { return true } } return false } // ReadAll executes all readers and combines their results. func (rc *ReadContext) ReadAll(ctx context.Context) (map[string]any, error) { combined := make(map[string]any) for name, reader := range rc.Readers { values, err := reader.Read(ctx) if err != nil { return nil, fmt.Errorf("reader %s failed: %w", name, err) } for k, v := range values { combined[name+"."+k] = v } } return combined, nil }