Implement thread-safe message and block metadata storage
- Introduced a new messageMetadataStore type for managing message and block metadata with thread-safe operations. - Added methods for setting, getting, and updating message and block metadata, enhancing concurrency handling. - Removed redundant metadata store implementation from types.go, centralizing functionality in message.go for improved organization.
This commit is contained in:
Max
parent
a868b4d741
commit
b45165b0ae
3 changed files with 53 additions and 92 deletions
|
|
@ -3,8 +3,61 @@ package context
|
|||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// messageMetadataStore provides thread-safe storage for message and block metadata
|
||||
type messageMetadataStore struct {
|
||||
messages map[string]*MessageMetadata // Message metadata by MessageID
|
||||
blocks map[string]*BlockMetadata // Block metadata by BlockID
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// newMessageMetadataStore creates a new message metadata store
|
||||
func newMessageMetadataStore() *messageMetadataStore {
|
||||
return &messageMetadataStore{
|
||||
messages: make(map[string]*MessageMetadata),
|
||||
blocks: make(map[string]*BlockMetadata),
|
||||
}
|
||||
}
|
||||
|
||||
// setMessage stores metadata for a message (thread-safe)
|
||||
func (s *messageMetadataStore) setMessage(messageID string, metadata *MessageMetadata) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.messages[messageID] = metadata
|
||||
}
|
||||
|
||||
// getMessage retrieves metadata for a message (thread-safe)
|
||||
func (s *messageMetadataStore) getMessage(messageID string) *MessageMetadata {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.messages[messageID]
|
||||
}
|
||||
|
||||
// setBlock stores metadata for a block (thread-safe)
|
||||
func (s *messageMetadataStore) setBlock(blockID string, metadata *BlockMetadata) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.blocks[blockID] = metadata
|
||||
}
|
||||
|
||||
// getBlock retrieves metadata for a block (thread-safe)
|
||||
func (s *messageMetadataStore) getBlock(blockID string) *BlockMetadata {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.blocks[blockID]
|
||||
}
|
||||
|
||||
// updateBlock updates block metadata (thread-safe)
|
||||
func (s *messageMetadataStore) updateBlock(blockID string, update func(*BlockMetadata)) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if block, exists := s.blocks[blockID]; exists {
|
||||
update(block)
|
||||
}
|
||||
}
|
||||
|
||||
// UnmarshalJSON custom unmarshaler for Message to handle Content field
|
||||
func (m *Message) UnmarshalJSON(data []byte) error {
|
||||
// Define a temporary struct to avoid infinite recursion
|
||||
|
|
|
|||
|
|
@ -214,58 +214,6 @@ type BlockMetadata struct {
|
|||
MessageCount int // Number of messages in this block
|
||||
}
|
||||
|
||||
// messageMetadataStore provides thread-safe storage for message and block metadata
|
||||
type messageMetadataStore struct {
|
||||
messages map[string]*MessageMetadata // Message metadata by MessageID
|
||||
blocks map[string]*BlockMetadata // Block metadata by BlockID
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// newMessageMetadataStore creates a new message metadata store
|
||||
func newMessageMetadataStore() *messageMetadataStore {
|
||||
return &messageMetadataStore{
|
||||
messages: make(map[string]*MessageMetadata),
|
||||
blocks: make(map[string]*BlockMetadata),
|
||||
}
|
||||
}
|
||||
|
||||
// setMessage stores metadata for a message (thread-safe)
|
||||
func (s *messageMetadataStore) setMessage(messageID string, metadata *MessageMetadata) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.messages[messageID] = metadata
|
||||
}
|
||||
|
||||
// getMessage retrieves metadata for a message (thread-safe)
|
||||
func (s *messageMetadataStore) getMessage(messageID string) *MessageMetadata {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.messages[messageID]
|
||||
}
|
||||
|
||||
// setBlock stores metadata for a block (thread-safe)
|
||||
func (s *messageMetadataStore) setBlock(blockID string, metadata *BlockMetadata) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.blocks[blockID] = metadata
|
||||
}
|
||||
|
||||
// getBlock retrieves metadata for a block (thread-safe)
|
||||
func (s *messageMetadataStore) getBlock(blockID string) *BlockMetadata {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.blocks[blockID]
|
||||
}
|
||||
|
||||
// updateBlock updates block metadata (thread-safe)
|
||||
func (s *messageMetadataStore) updateBlock(blockID string, update func(*BlockMetadata)) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if block, exists := s.blocks[blockID]; exists {
|
||||
update(block)
|
||||
}
|
||||
}
|
||||
|
||||
// Context the context
|
||||
type Context struct {
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
package types
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
store "github.com/yaoapp/yao/agent/store/types"
|
||||
)
|
||||
|
|
@ -15,26 +14,15 @@ type DSL struct {
|
|||
StoreSetting store.Setting `json:"store" yaml:"store"` // The store setting of the assistant
|
||||
Cache string `json:"cache" yaml:"cache"` // The cache store of the assistant, if not set, default is "__yao.agent.cache"
|
||||
|
||||
// AuthSetting *Auth `json:"auth,omitempty" yaml:"auth,omitempty"` // Authenticate Settings
|
||||
// UploadSetting *Upload `json:"upload,omitempty" yaml:"upload,omitempty"` // Upload Settings
|
||||
// KnowledgeSetting *Knowledge `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base Settings
|
||||
|
||||
// Global External Settings - model capabilities, tools, etc.
|
||||
// ===============================
|
||||
Models map[string]assistant.ModelCapabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
|
||||
|
||||
// Agent API Settings
|
||||
// ===============================s
|
||||
// Guard string `json:"guard,omitempty" yaml:"guard,omitempty"` // The guard of the assistant
|
||||
// Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` // The allowed domains of the assistant
|
||||
|
||||
// Internal
|
||||
// ===============================
|
||||
// ID string `json:"-" yaml:"-"` // The id of the instance
|
||||
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
|
||||
Store store.Store `json:"-" yaml:"-"` // The store of the assistant
|
||||
// Vision *vision.Vision `json:"-" yaml:"-"`
|
||||
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
|
||||
}
|
||||
|
||||
// Uses the default assistant settings
|
||||
|
|
@ -49,34 +37,6 @@ type Uses struct {
|
|||
Fetch string `json:"fetch,omitempty" yaml:"fetch,omitempty"` // The assistant for fetching the http/https/ftp/sftp/etc. file, and return the file's content. if not set, use the http process to fetch the file.
|
||||
}
|
||||
|
||||
// Auth Authenticate Settings
|
||||
// ===============================
|
||||
type Auth struct {
|
||||
Models *AuthModels `json:"models,omitempty" yaml:"models,omitempty"` // The models of the user, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
|
||||
Fields *AuthFields `json:"fields,omitempty" yaml:"fields,omitempty"` // The fields of the user model, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
|
||||
SessionFields *AuthSessionFields `json:"session_fields,omitempty" yaml:"session_fields,omitempty"` // The session fields of the user, it is used to handle the user, and the user is a user in the database.
|
||||
}
|
||||
|
||||
// AuthModels the auth model
|
||||
type AuthModels struct {
|
||||
User string `json:"user,omitempty" yaml:"user,omitempty"` // default is admin.user, The user model is a special model, it is used to handle the user, and the user is a user in the database.
|
||||
Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // The guest model is a special model, it is used to handle the guest user, and the guest user is not a user in the database.
|
||||
}
|
||||
|
||||
// AuthSessionFields the auth session field
|
||||
type AuthSessionFields struct {
|
||||
ID string `json:"id,omitempty" yaml:"id,omitempty"` // the field name of the user id, default is user_id
|
||||
Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is user_roles. the value must be an JSON array string.
|
||||
Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // the field name of the guest user, default is guest_id
|
||||
}
|
||||
|
||||
// AuthFields the auth field
|
||||
type AuthFields struct {
|
||||
ID string `json:"id,omitempty" yaml:"id,omitempty"` // the field name of the user id, default is id
|
||||
Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is roles, it must be an JSON field.
|
||||
Permission string `json:"permission,omitempty" yaml:"permission,omitempty"` // the field name of the user permission, default is permission
|
||||
}
|
||||
|
||||
// Mention Structure
|
||||
// ===============================
|
||||
type Mention struct {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue