package assistant import ( "fmt" "strings" "time" "github.com/google/uuid" agentcontext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" storetypes "github.com/yaoapp/yao/agent/store/types" ) // InitializeConversation prepares conversation context (synchronous) // KB collection is now initialized when user logs in (see openapi/user/login.go) func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error { // Reserved for future conversation initialization logic return nil } // InitializeConversationAsync prepares conversation context asynchronously func (ast *Assistant) InitializeConversationAsync(ctx *agentcontext.Context, options ...*agentcontext.Options) { go ast.InitializeConversation(ctx, options...) } // GetChatKBID returns the KB collection ID for a chat session // Same team + user always returns the same ID (deterministic) // Format: chat_{team}_{user} or chat_user_{user} if no team func GetChatKBID(teamID, userID string) string { // Sanitize IDs: replace invalid chars with underscores cleanTeamID := sanitizeCollectionID(teamID) cleanUserID := sanitizeCollectionID(userID) if cleanTeamID != "" { return fmt.Sprintf("chat_%s_%s", cleanTeamID, cleanUserID) } return fmt.Sprintf("chat_user_%s", cleanUserID) } // sanitizeCollectionID replaces invalid characters with underscores // Collection IDs only allow: a-z, A-Z, 0-9, and underscore func sanitizeCollectionID(id string) string { if id == "" { return "" } // Replace any character that is not alphanumeric or underscore with underscore result := make([]byte, len(id)) for i := 0; i < len(id); i++ { c := id[i] if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' { result[i] = c } else { result[i] = '_' } } return string(result) } // mergeChatMetadata merges default metadata with chat context information func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext.Context) map[string]interface{} { metadata := make(map[string]interface{}) // Copy default metadata for k, v := range defaultMetadata { metadata[k] = v } // Add chat-specific metadata (only for internal tracking, not displayed) metadata["chat_id"] = ctx.ChatID metadata["team_id"] = ctx.Authorized.TeamID metadata["user_id"] = ctx.Authorized.UserID // Get locale from context, default to zh-CN if not set locale := ctx.Locale if locale == "" { locale = "zh-CN" } locale = strings.ToLower(locale) // Use i18n for name and description (fixed, not showing user/team IDs) if _, exists := metadata["name"]; !exists { metadata["name"] = i18n.T(locale, "kb.chat.name") } if _, exists := metadata["description"]; !exists { metadata["description"] = i18n.T(locale, "kb.chat.description") } return metadata } // ============================================================================= // Chat Buffer Integration // ============================================================================= // InitBuffer initializes the chat buffer for the context // Should be called at the start of Stream() for root stack only func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) { // Only initialize for root stack if ctx.Stack == nil || !ctx.Stack.IsRoot() { return } // Skip if buffer already exists if ctx.Buffer != nil { return } // Skip if History is disabled in options if ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History { ctx.Logger.Debug("Buffer skipped: Skip.History is true") return } // Generate request ID if not set requestID := ctx.RequestID() if requestID == "" { requestID = uuid.New().String() } // Get connector and mode from options connector := "" mode := "" if ctx.Stack.Options != nil { connector = ctx.Stack.Options.Connector mode = ctx.Stack.Options.Mode } ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector, mode) ctx.Logger.Debug("Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID) } // BufferUserInput adds user input messages to the buffer // Should be called after InitBuffer func (ast *Assistant) BufferUserInput(ctx *agentcontext.Context, inputMessages []agentcontext.Message) { if ctx.Buffer == nil { return } // Only root stack should buffer user input // Delegated agents share the same buffer but should not duplicate user input if ctx.Stack != nil && !ctx.Stack.IsRoot() { return } // Convert input messages to buffer format for _, msg := range inputMessages { // Extract content from message var content interface{} var name string content = msg.Content if msg.Name != nil { name = *msg.Name } ctx.Buffer.AddUserInput(content, name) } } // UpdateSpaceSnapshot updates the context memory snapshot in the buffer // Only captures Context-level memory (request-scoped temporary data) for recovery func (ast *Assistant) UpdateSpaceSnapshot(ctx *agentcontext.Context) { if ctx.Buffer == nil || ctx.Memory == nil || ctx.Memory.Context == nil { return } snapshot := ctx.Memory.Context.Snapshot() ctx.Buffer.SetSpaceSnapshot(snapshot) } // BeginStep starts tracking an execution step // Returns the step for further updates func (ast *Assistant) BeginStep(ctx *agentcontext.Context, stepType string, input map[string]interface{}) *agentcontext.BufferedStep { if ctx.Buffer == nil { return nil } // Update space snapshot before beginning step ast.UpdateSpaceSnapshot(ctx) return ctx.Buffer.BeginStep(stepType, input, ctx.Stack) } // CompleteStep marks the current step as completed func (ast *Assistant) CompleteStep(ctx *agentcontext.Context, output map[string]interface{}) { if ctx.Buffer == nil { return } ctx.Buffer.CompleteStep(output) } // FlushBuffer saves all buffered data to the database // Should be called in defer block at the end of Stream() func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, err error) { if ctx.Buffer == nil { return } // Only flush for root stack if ctx.Stack == nil || !ctx.Stack.IsRoot() { return } // Get chat store chatStore := GetChatStore() if chatStore == nil { ctx.Logger.Error("Chat store not available, cannot flush buffer") return } // Mark current step as failed/interrupted if needed if finalStatus != agentcontext.StepStatusCompleted && err != nil { ctx.Buffer.FailCurrentStep(finalStatus, err) } // 1. Save all messages (user input + assistant responses) messages := ast.convertBufferedMessages(ctx.Buffer.GetMessages()) if len(messages) > 0 { if saveErr := chatStore.SaveMessages(ctx.ChatID, messages); saveErr != nil { ctx.Logger.Error("Failed to save messages: %v", saveErr) } else { ctx.Logger.Debug("Saved %d messages for chat=%s", len(messages), ctx.ChatID) } } // 2. Update chat last_message_at, last_connector, and last_mode if len(messages) > 0 { now := time.Now() updates := map[string]interface{}{ "last_message_at": now, } // Also update last_connector if available if connector := ctx.Buffer.Connector(); connector != "" { updates["last_connector"] = connector } // Also update last_mode if available if mode := ctx.Buffer.Mode(); mode != "" { updates["last_mode"] = mode } if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil { ctx.Logger.Debug("Failed to update chat: %v", updateErr) } } // 3. Only save resume steps on error/interrupt (not on success) if finalStatus != agentcontext.StepStatusCompleted { steps := ast.convertBufferedSteps(ctx.Buffer.GetStepsForResume(finalStatus)) if len(steps) > 0 { if saveErr := chatStore.SaveResume(steps); saveErr != nil { ctx.Logger.Error("Failed to save resume steps: %v", saveErr) } else { ctx.Logger.Debug("Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus) } } } } // convertBufferedMessages converts BufferedMessage slice to store Message slice func (ast *Assistant) convertBufferedMessages(buffered []*agentcontext.BufferedMessage) []*storetypes.Message { if len(buffered) == 0 { return nil } messages := make([]*storetypes.Message, len(buffered)) for i, msg := range buffered { messages[i] = &storetypes.Message{ MessageID: msg.MessageID, ChatID: msg.ChatID, RequestID: msg.RequestID, Role: msg.Role, Type: msg.Type, Props: msg.Props, BlockID: msg.BlockID, ThreadID: msg.ThreadID, AssistantID: msg.AssistantID, Connector: msg.Connector, Mode: msg.Mode, Sequence: msg.Sequence, Metadata: msg.Metadata, CreatedAt: msg.CreatedAt, UpdatedAt: msg.CreatedAt, } } return messages } // convertBufferedSteps converts BufferedStep slice to store Resume slice func (ast *Assistant) convertBufferedSteps(buffered []*agentcontext.BufferedStep) []*storetypes.Resume { if len(buffered) == 0 { return nil } steps := make([]*storetypes.Resume, len(buffered)) for i, step := range buffered { steps[i] = &storetypes.Resume{ ResumeID: step.ResumeID, ChatID: step.ChatID, RequestID: step.RequestID, AssistantID: step.AssistantID, StackID: step.StackID, StackParentID: step.StackParentID, StackDepth: step.StackDepth, Type: step.Type, Status: step.Status, Input: step.Input, Output: step.Output, SpaceSnapshot: step.SpaceSnapshot, Error: step.Error, Sequence: step.Sequence, Metadata: step.Metadata, CreatedAt: step.CreatedAt, UpdatedAt: step.CreatedAt, } } return steps } // EnsureChat ensures a chat session exists, creates if not func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error { if ctx.ChatID == "" { return nil // No chat ID, skip } // Skip if history is disabled if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History { return nil // Skip.History is true, don't create chat session } chatStore := GetChatStore() if chatStore == nil { return nil // No store, skip } // Check if chat exists _, err := chatStore.GetChat(ctx.ChatID) if err == nil { return nil // Chat exists } // Create new chat with permission fields chat := &storetypes.Chat{ ChatID: ctx.ChatID, AssistantID: ast.ID, Status: "active", Share: "private", Sort: 0, CreatedAt: time.Now(), UpdatedAt: time.Now(), } // Set last_connector from options (user selected connector) if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Connector != "" { chat.LastConnector = ctx.Stack.Options.Connector } // Set permission fields from authorized info if ctx.Authorized != nil { chat.CreatedBy = ctx.Authorized.UserID chat.UpdatedBy = ctx.Authorized.UserID chat.TeamID = ctx.Authorized.TeamID chat.TenantID = ctx.Authorized.TenantID } return chatStore.CreateChat(chat) } // GetChatStore returns the chat store instance // Returns nil if storage is not configured func GetChatStore() storetypes.ChatStore { if storage == nil { return nil } return storage } // GetStore returns the full store instance (implements both ChatStore and AssistantStore) // Returns nil if storage is not configured func GetStore() storetypes.Store { if storage == nil { return nil } return storage } // ============================================================================= // Deprecated methods (kept for compatibility) // ============================================================================= func (ast *Assistant) saveChat(ctx *agentcontext.Context, input []agentcontext.Message, opts *agentcontext.Options) error { _ = ctx _ = input _ = opts return nil }