common/redis/redis.go

package redis

import (
	"context"
	"strings"
	"time"

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

// GRedisClient GoFrame gredis 客户端，统一使用
var GRedisClient *gredis.Redis

// RedisClient GRedisClient 的别名，保持向后兼容
var RedisClient *gredis.Redis

func init() {
	// 初始化 GoFrame gredis 客户端
	GRedisClient = g.Redis()
	RedisClient = GRedisClient // 别名指向同一个客户端
}

// Stream 和消费者组常量
const (
	// RAGFlow 请求 Stream Key
	RAGFlowRequestStreamKey = "ragflow:request:stream"
	// RAGFlow 消费者组名称
	RAGFlowConsumerGroup = "ragflow:consumer:group"
	// 会话最后活跃时间 Key 前缀
	SessionLastActiveKeyPrefix = "ragflow:session:"
)

// StreamMessage Redis Stream 消息结构
type StreamMessage struct {
	ID     string                 // 消息ID（自动生成）
	Values map[string]interface{} // 消息内容
}

// InitStreamGroup 初始化 Stream 和消费者组
// 使用 gredis Do() 方法执行 XGROUP CREATE 命令
func InitStreamGroup(ctx context.Context, streamKey, groupName string) error {
	// XGROUP CREATE streamKey groupName 0 MKSTREAM
	_, err := GRedisClient.Do(ctx, "XGROUP", "CREATE", streamKey, groupName, "0", "MKSTREAM")
	if err != nil {
		// 如果组已存在，忽略错误
		errStr := err.Error()
		if strings.Contains(errStr, "BUSYGROUP") || strings.Contains(errStr, "already exists") {
			return nil
		}
		return err
	}
	return nil
}

// AddToStream 将消息添加到 Stream
// 使用 gredis Do() 方法执行 XADD 命令
func AddToStream(ctx context.Context, streamKey string, values map[string]interface{}) (string, error) {
	// XADD streamKey * field1 value1 field2 value2 ...
	args := []interface{}{streamKey, "*"} // "*" 自动生成ID
	for key, val := range values {
		args = append(args, key, val)
	}

	result, err := GRedisClient.Do(ctx, "XADD", args...)
	if err != nil {
		return "", err
	}

	// 返回消息ID
	messageID := result.String()
	return messageID, nil
}

// ReadFromStream 从 Stream 读取消息（消费者组模式）
// 使用 gredis Do() 方法执行 XREADGROUP 命令
func ReadFromStream(ctx context.Context, streamKey, groupName, consumerName string, count int64, blockMs int64) ([]StreamMessage, error) {
	// XREADGROUP GROUP groupName consumerName COUNT count BLOCK blockMs STREAMS streamKey >
	result, err := GRedisClient.Do(ctx,
		"XREADGROUP", "GROUP", groupName, consumerName,
		"COUNT", count,
		"BLOCK", blockMs,
		"STREAMS", streamKey, ">",
	)

	if err != nil {
		return nil, err
	}

	// 解析返回值
	// 格式: [[streamKey, [[msgID, [field1, value1, field2, value2, ...]], ...]]]
	messages := []StreamMessage{}

	if result == nil {
		// 超时或没有数据
		return messages, nil
	}

	// 类型断言：result.Val() 返回 interface{}
	streamsArray, ok := result.Val().([]interface{})
	if !ok || len(streamsArray) == 0 {
		return messages, nil
	}

	// 遍历每个 stream
	for _, streamData := range streamsArray {
		streamArray, ok := streamData.([]interface{})
		if !ok || len(streamArray) < 2 {
			continue
		}

		// streamArray[0] 是 streamKey, streamArray[1] 是消息数组
		messagesArray, ok := streamArray[1].([]interface{})
		if !ok {
			continue
		}

		// 解析每条消息
		for _, msgData := range messagesArray {
			msgArray, ok := msgData.([]interface{})
			if !ok || len(msgArray) < 2 {
				continue
			}

			// msgArray[0] 是 ID, msgArray[1] 是字段数组
			msgID := gconv.String(msgArray[0])
			fieldsArray, ok := msgArray[1].([]interface{})
			if !ok {
				continue
			}

			// 解析字段为 map
			values := make(map[string]interface{})
			for i := 0; i < len(fieldsArray); i += 2 {
				if i+1 < len(fieldsArray) {
					key := gconv.String(fieldsArray[i])
					val := fieldsArray[i+1]
					values[key] = val
				}
			}

			messages = append(messages, StreamMessage{
				ID:     msgID,
				Values: values,
			})
		}
	}

	return messages, nil
}

// AckMessage 确认消息已处理
// 使用 gredis Do() 方法执行 XACK 命令
func AckMessage(ctx context.Context, streamKey, groupName string, messageIDs ...string) error {
	// XACK streamKey groupName messageID1 messageID2 ...
	args := []interface{}{streamKey, groupName}
	for _, id := range messageIDs {
		args = append(args, id)
	}

	_, err := GRedisClient.Do(ctx, "XACK", args...)
	return err
}

// GetStreamLength 获取 Stream 当前长度
// 使用 gredis Do() 方法执行 XLEN 命令
func GetStreamLength(ctx context.Context, streamKey string) (int64, error) {
	// XLEN streamKey
	result, err := GRedisClient.Do(ctx, "XLEN", streamKey)
	if err != nil {
		return 0, err
	}

	length := gconv.Int64(result)
	return length, nil
}

// PendingMessage Pending 消息结构
type PendingMessage struct {
	ID         string // 消息ID
	Consumer   string // 消费者名称
	Idle       int64  // 空闲时间(毫秒)
	RetryCount int64  // 重试次数
}

// GetPendingMessages 获取待处理消息
// 使用 gredis Do() 方法执行 XPENDING 命令
func GetPendingMessages(ctx context.Context, streamKey, groupName string, start, end string, count int64) ([]PendingMessage, error) {
	// XPENDING streamKey groupName start end count
	result, err := GRedisClient.Do(ctx, "XPENDING", streamKey, groupName, start, end, count)
	if err != nil {
		return nil, err
	}

	if result == nil {
		return []PendingMessage{}, nil
	}

	// 解析返回值：[[ID, consumer, idle, retryCount], ...]
	pendingArray, ok := result.Val().([]interface{})
	if !ok {
		return []PendingMessage{}, nil
	}

	var messages []PendingMessage
	for _, item := range pendingArray {
		itemArray, ok := item.([]interface{})
		if !ok || len(itemArray) < 4 {
			continue
		}

		messages = append(messages, PendingMessage{
			ID:         gconv.String(itemArray[0]),
			Consumer:   gconv.String(itemArray[1]),
			Idle:       gconv.Int64(itemArray[2]),
			RetryCount: gconv.Int64(itemArray[3]),
		})
	}

	return messages, nil
}

// ClaimPendingMessage 认领超时的 Pending 消息
// 使用 gredis Do() 方法执行 XCLAIM 命令
func ClaimPendingMessage(ctx context.Context, streamKey, groupName, consumerName string, minIdleTime int64, messageIDs ...string) ([]StreamMessage, error) {
	// XCLAIM streamKey groupName consumerName minIdleTime messageID1 messageID2 ...
	args := []interface{}{streamKey, groupName, consumerName, minIdleTime}
	for _, id := range messageIDs {
		args = append(args, id)
	}

	result, err := GRedisClient.Do(ctx, "XCLAIM", args...)
	if err != nil {
		return nil, err
	}

	if result == nil {
		return []StreamMessage{}, nil
	}

	// 解析返回值：类似 XREADGROUP
	messagesArray, ok := result.Val().([]interface{})
	if !ok {
		return []StreamMessage{}, nil
	}

	var messages []StreamMessage
	for _, msgData := range messagesArray {
		msgArray, ok := msgData.([]interface{})
		if !ok || len(msgArray) < 2 {
			continue
		}

		msgID := gconv.String(msgArray[0])
		fieldsArray, ok := msgArray[1].([]interface{})
		if !ok {
			continue
		}

		values := make(map[string]interface{})
		for i := 0; i < len(fieldsArray); i += 2 {
			if i+1 < len(fieldsArray) {
				key := gconv.String(fieldsArray[i])
				val := fieldsArray[i+1]
				values[key] = val
			}
		}

		messages = append(messages, StreamMessage{
			ID:     msgID,
			Values: values,
		})
	}

	return messages, nil
}

// SetSessionLastActive 设置用户最后活跃时间
// 使用 gredis SetEX 方法
func SetSessionLastActive(ctx context.Context, userId string) error {
	key := SessionLastActiveKeyPrefix + userId + ":last_active"
	timestamp := time.Now().Unix()

	// SETEX key 7200 value (7200秒 = 2小时)
	_, err := GRedisClient.Do(ctx, "SETEX", key, 7200, timestamp)
	return err
}

// GetSessionLastActive 获取用户最后活跃时间
// 使用 gredis Get 方法
func GetSessionLastActive(ctx context.Context, userId string) (int64, error) {
	key := SessionLastActiveKeyPrefix + userId + ":last_active"
	result, err := GRedisClient.Get(ctx, key)
	if err != nil {
		return 0, err
	}

	if result.IsEmpty() {
		return 0, nil
	}

	timestamp := gconv.Int64(result.Val())
	return timestamp, nil
}

// IsUserActive 检查用户是否在指定时间范围内活跃过
// 用于追问逻辑：如果用户最近活跃过，则不发送追问消息
// 参数：
//   - userId: 用户ID
//   - seconds: 时间范围（秒），例如传入300表示检查5分钟内是否活跃
//
// 返回：
//   - bool: true表示用户在指定时间内活跃过
//   - error: 操作失败时返回错误
func IsUserActive(ctx context.Context, userId string, seconds int64) (bool, error) {
	lastActive, err := GetSessionLastActive(ctx, userId)
	if err != nil {
		return false, err
	}

	if lastActive == 0 {
		return false, nil // 未找到记录，视为不活跃
	}

	now := time.Now().Unix()
	return (now - lastActive) < seconds, nil
}

// SetSessionCache 缓存用户的 RAGFlow Session ID
// 使用 gredis SetEX 方法
func SetSessionCache(ctx context.Context, userId, sessionId string) error {
	key := SessionLastActiveKeyPrefix + userId + ":session_id"

	// SETEX key 604800 value (604800秒 = 7天)
	_, err := GRedisClient.Do(ctx, "SETEX", key, 604800, sessionId)
	return err
}

// GetSessionCache 获取缓存的 RAGFlow Session ID
// 使用 gredis Get 方法
func GetSessionCache(ctx context.Context, userId string) (string, error) {
	key := SessionLastActiveKeyPrefix + userId + ":session_id"
	result, err := GRedisClient.Get(ctx, key)
	if err != nil {
		return "", err
	}

	if result.IsEmpty() {
		return "", nil
	}

	return result.String(), nil
}