Agent-System/internal/ctxmgr/recategorize.go

package ctxmgr

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

// RecategorizeResult 存储重分类批处理结果统计
type RecategorizeResult struct {
	TotalProcessed int `json:"total_processed"`
	TotalUpdated   int `json:"total_updated"`
}

type recatMsg struct {
	ID        int64
	SessionID string
	Role      string
	Intent    string
	Branch    string
	Metadata  map[string]any
	Category  string
	CreatedAt time.Time
}

// RecategorizeMessages 存量消息重分类回填引擎 (Spec §3.4 A4)
//
// 1. assistant 消息：
//   - metadata 含有 turn_intent -> 直接按 IntentCategoryMap 映射；
//   - metadata 无 turn_intent -> 回溯同轮用户消息（时间间隔 < 60s 且 role=user 的前序消息）推导其意图；
//   - 兜底按现有 DeriveMessageCategory 判定；
//
// 2. user 消息：
//   - 按最新意图词表重新判定，纠偏遗漏映射（如 CONFIRM_ACTION、EXECUTE_STAGE_ADVANCE 等）；
//
// 3. 幂等更新：仅当计算出的 newCategory != 当前 category 时才执行持久化。
func (m *ContextManager) RecategorizeMessages(ctx context.Context, sessionID string) (*RecategorizeResult, error) {
	result := &RecategorizeResult{}

	if m.pool != nil {
		query := `
			SELECT id, session_id, role, intent, branch, metadata, category, created_at
			FROM chat_messages
			WHERE ($1 = '' OR session_id = $1) AND deleted_at IS NULL
			ORDER BY session_id, created_at ASC, id ASC
		`
		rows, err := m.pool.Query(ctx, query, sessionID)
		if err != nil {
			return nil, fmt.Errorf("query messages for recategorize failed: %w", err)
		}
		defer rows.Close()

		var msgs []recatMsg
		for rows.Next() {
			var r recatMsg
			var metaJSON []byte
			if err := rows.Scan(&r.ID, &r.SessionID, &r.Role, &r.Intent, &r.Branch, &metaJSON, &r.Category, &r.CreatedAt); err != nil {
				return nil, fmt.Errorf("scan recategorize row failed: %w", err)
			}
			r.Metadata = make(map[string]any)
			if len(metaJSON) > 0 {
				_ = json.Unmarshal(metaJSON, &r.Metadata)
			}
			msgs = append(msgs, r)
		}
		if err := rows.Err(); err != nil {
			return nil, fmt.Errorf("scan messages for recategorize failed: %w", err)
		}

		var lastUserMsg *recatMsg
		for _, msg := range msgs {
			result.TotalProcessed++

			var targetIntent = msg.Intent
			var targetMeta = msg.Metadata
			if targetMeta == nil {
				targetMeta = make(map[string]any)
			}

			if msg.Role == "user" {
				copyMsg := msg
				lastUserMsg = &copyMsg
			} else if msg.Role == "assistant" {
				// 若元数据中无 turn_intent，尝试回溯 60s 内的同轮用户消息
				if _, ok := targetMeta["turn_intent"]; !ok && lastUserMsg != nil && lastUserMsg.SessionID == msg.SessionID {
					if msg.CreatedAt.Sub(lastUserMsg.CreatedAt) <= 60*time.Second {
						if lastUserMsg.Intent != "" {
							targetMeta["turn_intent"] = lastUserMsg.Intent
						}
					}
				}
			}

			newCategory := DeriveMessageCategory(msg.Role, targetIntent, msg.Branch, targetMeta)

			if newCategory != msg.Category {
				_, err := m.pool.Exec(ctx, `
					UPDATE chat_messages
					SET category = $1, metadata = $2
					WHERE id = $3
				`, newCategory, targetMeta, msg.ID)
				if err != nil {
					return nil, fmt.Errorf("update message %d category failed: %w", msg.ID, err)
				}
				result.TotalUpdated++
			}
		}

		return result, nil
	}

	// 内存模式 fallback
	m.mu.Lock()
	defer m.mu.Unlock()

	for sid, list := range m.messages {
		if sessionID != "" && sid != sessionID {
			continue
		}

		var lastUserMsg *ChatMessageRecord
		for i := range list {
			result.TotalProcessed++
			msg := &list[i]
			if msg.DeletedAt != nil {
				continue
			}

			meta := msg.Metadata
			if meta == nil {
				meta = make(map[string]any)
			}

			if msg.Role == "user" {
				lastUserMsg = msg
			} else if msg.Role == "assistant" {
				if _, ok := meta["turn_intent"]; !ok && lastUserMsg != nil && lastUserMsg.SessionID == msg.SessionID {
					if msg.CreatedAt.Sub(lastUserMsg.CreatedAt) <= 60*time.Second {
						if lastUserMsg.Intent != "" {
							meta["turn_intent"] = lastUserMsg.Intent
						}
					}
				}
			}

			newCategory := DeriveMessageCategory(msg.Role, msg.Intent, msg.Branch, meta)
			if newCategory != msg.Category {
				msg.Category = newCategory
				msg.Metadata = meta
				result.TotalUpdated++
			}
		}
	}

	return result, nil
}