sundynix-agentix/sundynix-mcp-go/internal/memory/store.go

// Package memory 是偏好记忆的存储后端（第 5 层 I/O 型工具持有）。
// 常驻画像存 Postgres，按 sundynix_ 前缀约定 + AutoMigrate 自动迁移。
package memory

import (
	"context"
	"fmt"
	"log"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/bwmarrin/snowflake"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/logger"
)

// sfNode 是本服务的雪花 ID 生成器。node=2 与网关(node=1)区分，避免共享 PG 中 ID 冲突。
var sfNode *snowflake.Node

func init() { sfNode, _ = snowflake.NewNode(2) }

// NewID 生成字符串型雪花 ID（项目级 DB 规约：主键统一用它）。
func NewID() string { return sfNode.Generate().String() }

// BaseModel 是所有 DB 映射结构体的基础字段（项目级规约）：
// 字符串雪花 ID 主键 + 创建/更新时间 + GORM 软删（带索引）。与网关 store.BaseModel 同规约。
type BaseModel struct {
	ID        string         `gorm:"primaryKey;size:24" json:"id"`
	CreatedAt time.Time      `json:"created_at"`
	UpdatedAt time.Time      `json:"updated_at"`
	DeletedAt gorm.DeletedAt `gorm:"index" json:"-"`
}

// BeforeCreate 在插入前补雪花 ID。
func (b *BaseModel) BeforeCreate(*gorm.DB) error {
	if b.ID == "" {
		b.ID = NewID()
	}
	return nil
}

// Profile 是用户常驻画像的一条 key/value 偏好（always-on memory）。
// 套用 BaseModel 规约：雪花 ID 主键；(user_id, key) 唯一索引 —— 同一用户同一键 upsert 覆盖。
type Profile struct {
	BaseModel
	UserID     string    `gorm:"column:user_id;size:64;uniqueIndex:idx_profile_uk"`
	Key        string    `gorm:"size:64;uniqueIndex:idx_profile_uk"`
	Value      string    `gorm:"type:text"`
	Importance float64   `gorm:"column:importance"`   // 1~10，consolidate 时 LLM 打分（poignancy）→ 读路径权重
	LastSeenAt time.Time `gorm:"column:last_seen_at"` // 最近被印证时间 → Recency 衰减依据
}

// TableName 固定表名，遵守 sundynix_ 前缀约定。
func (Profile) TableName() string { return "sundynix_user_profile" }

// Store 封装画像读写。db 为 nil 表示降级（无 Postgres 时记忆功能空转，不阻断工具服务）。
type Store struct{ db *gorm.DB }

// Open 连接 Postgres 并自动迁移 sundynix_user_profile。连接失败不 fatal：返回降级实例。
func Open(dsn string) *Store {
	db, err := gorm.Open(postgres.New(postgres.Config{DSN: dsn}), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		log.Printf("[memory] postgres 不可用，记忆降级（召回为空）: %v", err)
		return &Store{}
	}
	// 连接池上限，防高并发打爆 PG（max_connections 默认 100）；默认 25，可经 DB_MAX_OPEN_CONNS 调。
	if sqlDB, derr := db.DB(); derr == nil {
		maxOpen := 25
		if v := os.Getenv("DB_MAX_OPEN_CONNS"); v != "" {
			if n, perr := strconv.Atoi(v); perr == nil && n > 0 {
				maxOpen = n
			}
		}
		sqlDB.SetMaxOpenConns(maxOpen)
		sqlDB.SetMaxIdleConns(5)
		sqlDB.SetConnMaxLifetime(time.Hour)
	}
	// 一次性迁移：旧表用复合主键 (user_id,key) 无 id/时间戳，与雪花规约不兼容。
	migrateLegacyProfile(db)
	if err := db.AutoMigrate(&Profile{}); err != nil {
		log.Printf("[memory] AutoMigrate 失败，记忆降级: %v", err)
		return &Store{}
	}
	log.Println("[memory] postgres connected, migrated sundynix_user_profile (雪花 id + 软删 规约)")
	return &Store{db: db}
}

// migrateLegacyProfile 检测旧 Profile 表（无 id 列）则备份偏好 → 重建为雪花规约 → 回灌。
func migrateLegacyProfile(db *gorm.DB) {
	m := db.Migrator()
	if !m.HasTable("sundynix_user_profile") {
		return // 全新库，AutoMigrate 直接建新表
	}
	if m.HasColumn(&Profile{}, "id") {
		return // 已是新规约
	}
	log.Println("[memory] 检测到旧 Profile 表（复合主键无 id），迁移到雪花规约（保留偏好）")
	var saved []struct {
		UserID string `gorm:"column:user_id"`
		Key    string `gorm:"column:key"`
		Value  string `gorm:"column:value"`
	}
	db.Raw(`SELECT user_id, "key", value FROM sundynix_user_profile`).Scan(&saved)
	if err := m.DropTable("sundynix_user_profile"); err != nil {
		log.Printf("[memory] 迁移 drop 旧表失败: %v", err)
		return
	}
	if err := db.AutoMigrate(&Profile{}); err != nil {
		log.Printf("[memory] 迁移建新表失败: %v", err)
		return
	}
	for _, r := range saved {
		_ = db.Create(&Profile{UserID: r.UserID, Key: r.Key, Value: r.Value}).Error
	}
	log.Printf("[memory] 已回灌 %d 条偏好（新雪花 id）", len(saved))
}

// 读路径打分参数（Generative Agents 公式的 Recency + Importance 两项；Relevance 待接 Milvus）。
const (
	memTopN            = 30   // 注入上限（截断，控 context）
	wRecency           = 0.4  // 最近性权重
	wImportance        = 0.6  // 重要度权重
	recencyDecayPerDay = 0.98 // 每天衰减因子（指数）
	defaultImportance  = 5.0  // 旧/未评分行的兜底重要度（避免被不公平遗忘）
)

// Get 返回某用户画像，渲染为可注入 prompt 的多行文本。
// 按 Score = wRecency·Recency + wImportance·Importance 降序，截断 top-N（控 context + 自然遗忘）。
func (s *Store) Get(ctx context.Context, userID string) (string, error) {
	if s.db == nil || userID == "" {
		return "", nil
	}
	var rows []Profile
	if err := s.db.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
		return "", err
	}
	ranked := rankProfiles(rows, time.Now(), memTopN)
	var b strings.Builder
	for _, r := range ranked {
		fmt.Fprintf(&b, "- %s：%s\n", r.Key, r.Value)
	}
	return strings.TrimRight(b.String(), "\n"), nil
}

// List 返回某用户全部 active 偏好（结构化，供管理面板查看/编辑），按 Score 降序、不截断。
func (s *Store) List(ctx context.Context, userID string) ([]Profile, error) {
	if s.db == nil || userID == "" {
		return nil, nil
	}
	var rows []Profile
	if err := s.db.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
		return nil, err
	}
	return rankProfiles(rows, time.Now(), 0), nil
}

// rankProfiles 纯函数：按 Score 降序排序（同分 key 升序稳定），截断 top-N（topN<=0 不截断）。
func rankProfiles(rows []Profile, now time.Time, topN int) []Profile {
	out := make([]Profile, len(rows))
	copy(out, rows)
	sort.SliceStable(out, func(i, j int) bool {
		si, sj := profileScore(out[i], now), profileScore(out[j], now)
		if si != sj {
			return si > sj
		}
		return out[i].Key < out[j].Key
	})
	if topN > 0 && len(out) > topN {
		out = out[:topN]
	}
	return out
}

// profileScore 计算一条偏好的 Recency+Importance 综合分（各归一到 [0,1]）。
func profileScore(p Profile, now time.Time) float64 {
	imp := p.Importance
	if imp <= 0 {
		imp = defaultImportance
	}
	return wRecency*recencyScore(now, p.LastSeenAt) + wImportance*(imp/10)
}

// recencyScore 指数衰减的最近性分：last_seen 越久越低；未记时间视为新鲜(1)。
func recencyScore(now, last time.Time) float64 {
	if last.IsZero() {
		return 1.0
	}
	days := now.Sub(last).Hours() / 24
	if days < 0 {
		days = 0
	}
	return math.Pow(recencyDecayPerDay, days)
}

// Upsert 写入/更新一条画像偏好（(user_id,key) 冲突即覆盖 value/importance，保留原 id；
// 置 last_seen=now 作"印证"）。importance<=0 时不覆盖旧值（NOOP 印证场景只 bump 时间）。
func (s *Store) Upsert(ctx context.Context, userID, key, value string, importance float64) error {
	if s.db == nil {
		return fmt.Errorf("memory store disabled")
	}
	now := time.Now()
	updates := map[string]any{"value": value, "updated_at": now, "last_seen_at": now}
	if importance > 0 {
		updates["importance"] = importance
	}
	return s.db.WithContext(ctx).Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "user_id"}, {Name: "key"}},
		DoUpdates: clause.Assignments(updates),
	}).Create(&Profile{UserID: userID, Key: key, Value: value, Importance: importance, LastSeenAt: now}).Error
}

// Touch 仅刷新某条偏好的 last_seen（NOOP 印证：被再次提及但内容不变，强化 Recency）。
func (s *Store) Touch(ctx context.Context, userID, key string) error {
	if s.db == nil {
		return nil
	}
	return s.db.WithContext(ctx).Model(&Profile{}).
		Where("user_id = ? AND key = ?", userID, key).
		Update("last_seen_at", time.Now()).Error
}

// Delete 软删一条偏好（(user_id,key)）—— 置 deleted_at，行保留可审计/恢复，正常查询自动过滤。
func (s *Store) Delete(ctx context.Context, userID, key string) error {
	if s.db == nil {
		return nil
	}
	return s.db.WithContext(ctx).Where("user_id = ? AND key = ?", userID, key).Delete(&Profile{}).Error
}

// Close 释放连接。
func (s *Store) Close() {
	if s.db == nil {
		return
	}
	if sqlDB, err := s.db.DB(); err == nil {
		_ = sqlDB.Close()
	}
}