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

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

import (
	"context"
	"encoding/binary"
	"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 衰减依据
	Embedding  []byte    `gorm:"column:embedding"`    // value 的嵌入向量（float32 小端打包）；召回时对 query 算余弦 → Relevance
}

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

// Embedder 是把文本转向量的最小接口（memory 包不硬依赖 rag 内部；由 gateway 注入 rag.Engine）。
// 返回每条文本一个向量；未配置/失败时 memory 优雅降级为 Recency+Importance 两项打分。
type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
}

// Store 封装画像读写。db 为 nil 表示降级（无 Postgres 时记忆功能空转，不阻断工具服务）。
// emb 为 nil 或未配置时，Relevance 打分静默跳过（回落两项打分，行为与升级前一致）。
type Store struct {
	db  *gorm.DB
	emb Embedder
}

// SetEmbedder 注入嵌入能力（gateway 装配时传 rag.Engine）。nil 安全。
func (s *Store) SetEmbedder(e Embedder) { s.emb = e }

// 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）。
const (
	memTopN            = 30   // 注入上限（截断，控 context）
	recencyDecayPerDay = 0.98 // 每天衰减因子（指数）
	defaultImportance  = 5.0  // 旧/未评分行的兜底重要度（避免被不公平遗忘）

	// 无 query / 无 embedder（回落两项）—— 保持升级前的行为与权重。
	wRecency    = 0.4
	wImportance = 0.6

	// relevance 模式（召回带 query 且 embedder 就绪）：三项加权，语义相关性主导但不淹没另两项。
	wRecencyR    = 0.25
	wImportanceR = 0.35
	wRelevance   = 0.4
)

// Get 返回某用户画像，渲染为可注入 prompt 的多行文本。
// query 非空且 embedder 就绪时，融入 Relevance（对当前任务的语义相关性）三项打分；
// 否则回落 Recency+Importance 两项（升级前行为）。截断 top-N（控 context + 自然遗忘 + 相关性优先）。
func (s *Store) Get(ctx context.Context, userID, query 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
	}
	rel := s.relevance(ctx, rows, query)
	ranked := rankProfiles(rows, time.Now(), memTopN, rel)
	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
}

// relevance 计算每条偏好对 query 的语义相关性（profile.ID → 余弦[0,1]）。
// 无 embedder / query 空 / query 嵌入失败 → 返回 nil（打分回落两项，绝不因此报错）。
// 行无缓存向量（存量未回填 / 嵌入曾失败）→ 该行不进 map，等同 relevance 0（不加分不减分）。
func (s *Store) relevance(ctx context.Context, rows []Profile, query string) map[string]float64 {
	if s.emb == nil || strings.TrimSpace(query) == "" {
		return nil
	}
	qv, err := s.emb.Embed(ctx, []string{query})
	if err != nil || len(qv) == 0 || len(qv[0]) == 0 {
		return nil // 嵌入不可用：优雅回落，不影响召回
	}
	q := qv[0]
	rel := make(map[string]float64, len(rows))
	for _, r := range rows {
		v := decodeVec(r.Embedding)
		if len(v) == len(q) && len(v) > 0 {
			rel[r.ID] = cosine01(q, v)
		}
	}
	return rel
}

// 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), nil // 管理面板：无 query，两项打分
}

// rankProfiles 纯函数：按 Score 降序排序（同分 key 升序稳定），截断 top-N（topN<=0 不截断）。
// rel 非 nil = relevance 模式（三项打分，缺失 ID 视为 relevance 0）；nil = 两项打分（升级前行为）。
func rankProfiles(rows []Profile, now time.Time, topN int, rel map[string]float64) []Profile {
	out := make([]Profile, len(rows))
	copy(out, rows)
	sort.SliceStable(out, func(i, j int) bool {
		si, sj := profileScore(out[i], now, rel), profileScore(out[j], now, rel)
		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 综合分（各项归一到 [0,1]）。rel==nil → 两项（Recency+Importance）；否则三项加入 Relevance。
func profileScore(p Profile, now time.Time, rel map[string]float64) float64 {
	imp := p.Importance
	if imp <= 0 {
		imp = defaultImportance
	}
	recency := recencyScore(now, p.LastSeenAt)
	impN := imp / 10
	if rel == nil {
		return wRecency*recency + wImportance*impN
	}
	return wRecencyR*recency + wImportanceR*impN + wRelevance*rel[p.ID] // 缺失 ID → 0
}

// ---- 向量编解码 + 余弦（float32 小端打包存 PG bytea；召回时内存算相似度）----

func encodeVec(v []float32) []byte {
	b := make([]byte, 4*len(v))
	for i, f := range v {
		binary.LittleEndian.PutUint32(b[i*4:], math.Float32bits(f))
	}
	return b
}

func decodeVec(b []byte) []float32 {
	if len(b) == 0 || len(b)%4 != 0 {
		return nil
	}
	v := make([]float32, len(b)/4)
	for i := range v {
		v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:]))
	}
	return v
}

// cosine01 余弦相似度截到 [0,1]（负相关记 0，不奖励相反语义）。
func cosine01(a, b []float32) float64 {
	var dot, na, nb float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		na += float64(a[i]) * float64(a[i])
		nb += float64(b[i]) * float64(b[i])
	}
	if na == 0 || nb == 0 {
		return 0
	}
	c := dot / (math.Sqrt(na) * math.Sqrt(nb))
	if c < 0 {
		return 0
	}
	if c > 1 {
		return 1
	}
	return c
}

// 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
	}
	p := &Profile{UserID: userID, Key: key, Value: value, Importance: importance, LastSeenAt: now}
	// 写入即向量化 value（供召回算 Relevance）。embedder 未配置/失败 → 留空向量，召回自动回落，不阻断写入。
	if s.emb != nil && strings.TrimSpace(value) != "" {
		if vecs, err := s.emb.Embed(ctx, []string{value}); err == nil && len(vecs) > 0 && len(vecs[0]) > 0 {
			enc := encodeVec(vecs[0])
			p.Embedding = enc
			updates["embedding"] = enc
		}
	}
	return s.db.WithContext(ctx).Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "user_id"}, {Name: "key"}},
		DoUpdates: clause.Assignments(updates),
	}).Create(p).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()
	}
}