Merge pull request #1361 from trheyi/main
Implement thread-safe message and block metadata storage
This commit is contained in:
Max
GitHub
commit
a5776ef6da
5 changed files with 115 additions and 98 deletions
|
|
@ -200,24 +200,78 @@ func (s *streamState) handleToolCall(data []byte) int {
|
|||
// Track current message type
|
||||
s.currentType = message.TypeToolCall
|
||||
|
||||
// Append to buffer
|
||||
// Append to buffer for message_end event
|
||||
s.buffer = append(s.buffer, data...)
|
||||
s.chunkCount++
|
||||
s.messageSeq++
|
||||
|
||||
// Parse the tool call delta data (JSON array from OpenAI)
|
||||
var toolCallArray []map[string]interface{}
|
||||
if err := jsoniter.Unmarshal(data, &toolCallArray); err != nil {
|
||||
// If parse fails, log and skip this chunk
|
||||
return 0
|
||||
}
|
||||
|
||||
// Extract tool call fields from delta
|
||||
// OpenAI delta typically has one element, but we handle arrays safely
|
||||
var props map[string]interface{}
|
||||
var deltaAction string
|
||||
var deltaPath string
|
||||
|
||||
if len(toolCallArray) == 1 {
|
||||
// Single tool call - flatten to props root level
|
||||
tc := toolCallArray[0]
|
||||
props = map[string]interface{}{}
|
||||
|
||||
// Static fields (only in first chunk): use merge
|
||||
if id, ok := tc["id"].(string); ok {
|
||||
props["id"] = id
|
||||
}
|
||||
if typ, ok := tc["type"].(string); ok {
|
||||
props["type"] = typ
|
||||
}
|
||||
if index, ok := tc["index"].(float64); ok {
|
||||
props["index"] = int(index)
|
||||
}
|
||||
if fn, ok := tc["function"].(map[string]interface{}); ok {
|
||||
if name, ok := fn["name"].(string); ok {
|
||||
props["name"] = name
|
||||
}
|
||||
// Arguments field: use append
|
||||
if args, ok := fn["arguments"].(string); ok {
|
||||
props["arguments"] = args
|
||||
// If this chunk has arguments, use append action for arguments field
|
||||
deltaAction = "append"
|
||||
deltaPath = "arguments"
|
||||
}
|
||||
}
|
||||
|
||||
// If no arguments in this chunk, use merge for other fields
|
||||
if deltaAction == "" {
|
||||
deltaAction = "merge"
|
||||
}
|
||||
} else {
|
||||
// Multiple tool calls in delta (rare) - keep as array
|
||||
props = map[string]interface{}{
|
||||
"calls": toolCallArray,
|
||||
}
|
||||
deltaAction = "merge"
|
||||
}
|
||||
|
||||
// Send delta message
|
||||
// - ChunkID: Unique chunk ID (C1, C2, C3...) for this fragment
|
||||
// - MessageID: Same for all chunks of this logical message (frontend merges by message_id)
|
||||
// - DeltaAction: "replace" for tool call raw data (each chunk contains complete state, not incremental)
|
||||
// - DeltaAction: "append" for arguments chunks, "merge" for id/type/name chunks
|
||||
// - DeltaPath: "arguments" when appending arguments field
|
||||
// OpenAI sends: first chunk has id/type/name, subsequent chunks only have arguments fragments
|
||||
msg := &message.Message{
|
||||
ChunkID: s.ctx.IDGenerator.GenerateChunkID(), // Unique chunk ID
|
||||
MessageID: s.currentGroupID, // Message ID for merging (all chunks share this)
|
||||
Type: message.TypeToolCall,
|
||||
Delta: true,
|
||||
DeltaAction: "replace", // Replace entire raw field with latest state
|
||||
Props: map[string]interface{}{
|
||||
"raw": string(data), // Raw tool call JSON data
|
||||
},
|
||||
DeltaAction: deltaAction, // "append" for arguments, "merge" for static fields
|
||||
DeltaPath: deltaPath, // "arguments" when appending
|
||||
Props: props, // Flattened tool call fields
|
||||
}
|
||||
|
||||
if err := s.ctx.Send(msg); err != nil {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
||||
|
|
|
|||
|
|
@ -587,6 +587,8 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
}
|
||||
|
||||
// Notify handler of tool call progress
|
||||
// Send the raw delta from OpenAI (as JSON bytes)
|
||||
// Handler will convert to object for frontend merge
|
||||
if handler != nil {
|
||||
toolCallData, _ := jsoniter.Marshal(delta.ToolCalls)
|
||||
handler(message.ChunkToolCall, toolCallData)
|
||||
|
|
|
|||
|
|
@ -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