yao/agent/output/message/utils.go
Max ef49827faf Refactor message handling for improved streaming and event management
- Updated message structures to replace 'StreamStartData' and 'StreamEndData' with 'EventStreamStartData' and 'EventStreamEndData' for better clarity and consistency.
- Introduced 'EventMessageStartData' and 'EventMessageEndData' to represent individual message lifecycle events, enhancing the granularity of message tracking.
- Refactored the 'streamState' and 'groupTracker' to utilize the new message structures, improving the organization and handling of streaming events.
- Enhanced the context management by integrating an ID generator for unique message identifiers, facilitating better tracking of message sequences.
- Updated documentation and tests to reflect the new message structures and ensure proper functionality across the system.
2025-11-26 18:30:43 +08:00

90 lines
2.7 KiB
Go

package message
import (
"fmt"
"sync/atomic"
gonanoid "github.com/matoous/go-nanoid/v2"
)
// IDGenerator generates unique IDs within a context (e.g., one conversation stream)
// Each Context should have its own IDGenerator to ensure IDs are unique within that context
type IDGenerator struct {
chunkCounter uint64
messageCounter uint64
blockCounter uint64
threadCounter uint64
}
// NewIDGenerator creates a new ID generator for a context
func NewIDGenerator() *IDGenerator {
return &IDGenerator{}
}
// GenerateChunkID generates a unique chunk ID with prefix C
// Format: C1, C2, C3...
func (g *IDGenerator) GenerateChunkID() string {
id := atomic.AddUint64(&g.chunkCounter, 1)
return fmt.Sprintf("C%d", id)
}
// GenerateMessageID generates a unique message ID with prefix M
// Format: M1, M2, M3...
func (g *IDGenerator) GenerateMessageID() string {
id := atomic.AddUint64(&g.messageCounter, 1)
return fmt.Sprintf("M%d", id)
}
// GenerateBlockID generates a unique block ID with prefix B
// Format: B1, B2, B3...
func (g *IDGenerator) GenerateBlockID() string {
id := atomic.AddUint64(&g.blockCounter, 1)
return fmt.Sprintf("B%d", id)
}
// GenerateThreadID generates a unique thread ID with prefix T
// Format: T1, T2, T3...
func (g *IDGenerator) GenerateThreadID() string {
id := atomic.AddUint64(&g.threadCounter, 1)
return fmt.Sprintf("T%d", id)
}
// Reset resets all counters (useful for testing)
func (g *IDGenerator) Reset() {
atomic.StoreUint64(&g.chunkCounter, 0)
atomic.StoreUint64(&g.messageCounter, 0)
atomic.StoreUint64(&g.blockCounter, 0)
atomic.StoreUint64(&g.threadCounter, 0)
}
// GetCounters returns current counter values (for debugging/testing)
func (g *IDGenerator) GetCounters() (chunk, message, block, thread uint64) {
return atomic.LoadUint64(&g.chunkCounter),
atomic.LoadUint64(&g.messageCounter),
atomic.LoadUint64(&g.blockCounter),
atomic.LoadUint64(&g.threadCounter)
}
// GenerateNanoID generates a unique ID using nanoid
// Returns a 21-character URL-safe string
// This is a static function that doesn't depend on the generator's counter
func GenerateNanoID() string {
id, err := gonanoid.New()
if err != nil {
// Fallback to timestamp-based ID if nanoid fails
return fmt.Sprintf("id_%d", atomic.AddUint64(new(uint64), 1))
}
return id
}
// GenerateCustomID generates a custom ID with prefix and nanoid
// Format: prefix_nanoid (e.g., "msg_V1StGXR8_Z5jdHi6B-myT")
// This is a static function that doesn't depend on the generator's counter
func GenerateCustomID(prefix string) string {
id, err := gonanoid.New()
if err != nil {
// Fallback to timestamp-based ID
return fmt.Sprintf("%s_%d", prefix, atomic.AddUint64(new(uint64), 1))
}
return fmt.Sprintf("%s_%s", prefix, id)
}