prompts-core/service/session/prompt_session_redis_service.go

package session

import (
	"context"
	"encoding/json"
	"fmt"
	"prompts-core/common/util"
	"prompts-core/model/dto"
	"time"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

const (
	// RedisKeySessionHistory 会话历史缓存 key: session:history:{tenantId}:{sessionId}:{nodeId}
	RedisKeySessionHistory = "session:history:%d:%s:%s"
)

// formatRedisKey 格式化 Redis key
func formatRedisKey(tenantID uint64, sessionID, nodeID string) string {
	return fmt.Sprintf(RedisKeySessionHistory, tenantID, sessionID, nodeID)
}

// ============================================
// 写操作
// ============================================

// SaveToRedis 保存一轮对话到 Redis ZSET
func SaveToRedis(ctx context.Context, tenantID uint64, sessionID, nodeID string, round *dto.HistoryRound) error {
	key := formatRedisKey(tenantID, sessionID, nodeID)
	maxRounds := util.GetMaxRounds(ctx)
	expireSeconds := int64(util.GetExpireMinutes(ctx) * 60)

	b, err := json.Marshal(round)
	if err != nil {
		return fmt.Errorf("序列化会话数据失败: %w", err)
	}

	score := float64(time.Now().UnixMilli())

	if _, err = g.Redis().Do(ctx, "ZADD", key, score, string(b)); err != nil {
		return fmt.Errorf("ZADD失败: %w", err)
	}
	if _, err = g.Redis().Do(ctx, "ZREMRANGEBYRANK", key, 0, -(maxRounds + 1)); err != nil {
		return fmt.Errorf("裁剪失败: %w", err)
	}
	if _, err = g.Redis().Do(ctx, "EXPIRE", key, expireSeconds); err != nil {
		return fmt.Errorf("设置过期失败: %w", err)
	}

	return nil
}

// DeleteSessionHistory 删除整个 session 下所有 node 的缓存
func DeleteSessionHistory(ctx context.Context, tenantID uint64, sessionID string) error {
	pattern := fmt.Sprintf(RedisKeySessionHistory, tenantID, sessionID, "*")
	keys, err := g.Redis().Do(ctx, "KEYS", pattern)
	if err != nil {
		return err
	}
	for _, key := range keys.Strings() {
		_, _ = g.Redis().Do(ctx, "DEL", key)
	}
	return nil
}

// DeleteRedisMessages 批量删除指定 node 下的消息
func DeleteRedisMessages(ctx context.Context, tenantID uint64, sessionID, nodeID string, msgIDs []int64) error {
	key := formatRedisKey(tenantID, sessionID, nodeID)
	for _, msgID := range msgIDs {
		cursor := "0"
		for {
			result, err := g.Redis().Do(ctx, "ZSCAN", key, cursor, "MATCH", fmt.Sprintf("*\"id\":%d*", msgID), "COUNT", 10)
			if err != nil {
				g.Log().Warningf(ctx, "[会话Redis] ZSCAN失败 msgID=%d err=%v", msgID, err)
				break
			}
			parts := result.Strings()
			if len(parts) < 2 {
				break
			}
			cursor = parts[0]
			for _, member := range parts[1:] {
				_, _ = g.Redis().Do(ctx, "ZREM", key, member)
			}
			if cursor == "0" {
				break
			}
		}
	}
	return nil
}

// ============================================
// 读操作
// ============================================

// GetFromRedis 从 Redis ZSET 获取会话历史
func GetFromRedis(ctx context.Context, tenantID uint64, sessionID, nodeID string) ([]dto.HistoryRound, error) {
	key := formatRedisKey(tenantID, sessionID, nodeID)
	maxRounds := util.GetMaxRounds(ctx)

	result, err := g.Redis().Do(ctx, "ZREVRANGE", key, 0, maxRounds-1)
	if err != nil {
		return nil, fmt.Errorf("ZREVRANGE失败: %w", err)
	}

	if result == nil || result.IsNil() {
		return []dto.HistoryRound{}, nil
	}

	return parseRounds(result.Strings()), nil
}

// ============================================
// 解析
// ============================================

func parseRounds(members []string) []dto.HistoryRound {
	rounds := make([]dto.HistoryRound, 0, len(members))
	for _, member := range members {
		var round dto.HistoryRound
		if err := json.Unmarshal([]byte(member), &round); err != nil {
			continue
		}
		if round.User != nil || round.Assistant != nil {
			rounds = append(rounds, round)
		}
	}
	return rounds
}

func flattenRounds(rounds []dto.HistoryRound) []dto.FlatMessage {
	var messages []dto.FlatMessage
	for i := len(rounds) - 1; i >= 0; i-- {
		if rounds[i].User != nil && gconv.String(rounds[i].User["content"]) != "" {
			messages = append(messages, dto.FlatMessage{
				Role:    gconv.String(rounds[i].User["role"]),
				Content: gconv.String(rounds[i].User["content"]),
			})
		}
		if rounds[i].Assistant != nil && gconv.String(rounds[i].Assistant["content"]) != "" {
			messages = append(messages, dto.FlatMessage{
				Role:    gconv.String(rounds[i].Assistant["role"]),
				Content: gconv.String(rounds[i].Assistant["content"]),
			})
		}
	}
	return messages
}