347 lines
10 KiB
Go
347 lines
10 KiB
Go
package redis
|
||
|
||
import (
|
||
"context"
|
||
"strconv"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
var RedisClient *redis.Client
|
||
|
||
func init() {
|
||
// 从 GoFrame 配置读取 Redis 配置
|
||
ctx := context.Background()
|
||
|
||
// 读取 Redis 配置
|
||
addr := g.Cfg().MustGet(ctx, "redis.default.address").String()
|
||
|
||
password := g.Cfg().MustGet(ctx, "redis.default.pass", "").String()
|
||
db := g.Cfg().MustGet(ctx, "redis.default.db", 0).Int()
|
||
|
||
// 读取超时配置
|
||
dialTimeout := g.Cfg().MustGet(ctx, "redis.default.dialTimeout", "30s").Duration()
|
||
readTimeout := g.Cfg().MustGet(ctx, "redis.default.readTimeout", "30s").Duration()
|
||
writeTimeout := g.Cfg().MustGet(ctx, "redis.default.writeTimeout", "30s").Duration()
|
||
|
||
// 创建 Redis 客户端
|
||
RedisClient = redis.NewClient(&redis.Options{
|
||
Addr: addr,
|
||
Password: password,
|
||
DB: db,
|
||
DialTimeout: dialTimeout,
|
||
ReadTimeout: readTimeout,
|
||
WriteTimeout: writeTimeout,
|
||
// 不设置 Protocol（让 go-redis 自动协商）
|
||
// Protocol: 2,
|
||
})
|
||
}
|
||
|
||
// 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 和消费者组
|
||
// 在应用启动时调用一次，创建 Stream 和消费者组
|
||
// 使用 GoFrame Do() 方法执行 XGROUP CREATE 命令
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - groupName: 消费者组名称
|
||
//
|
||
// 返回：error 初始化失败时返回错误
|
||
func InitStreamGroup(ctx context.Context, streamKey, groupName string) error {
|
||
// 使用 XGroupCreateMkStream 创建消费者组
|
||
// 如果 Stream 不存在会自动创建 (MKSTREAM)
|
||
// "0": 从 Stream 开头开始消费
|
||
err := RedisClient.XGroupCreateMkStream(ctx, streamKey, groupName, "0").Err()
|
||
if err != nil {
|
||
// 如果组已存在，忽略 BUSYGROUP 错误
|
||
if err.Error() == "BUSYGROUP Consumer Group name already exists" {
|
||
return nil
|
||
}
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// AddToStream 将消息添加到 Stream
|
||
// 用于 Controller 层将 RAGFlow 请求推入 Stream
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - values: 消息内容（键值对）
|
||
//
|
||
// 返回：
|
||
// - string: 消息ID
|
||
// - error: 添加失败时返回错误
|
||
func AddToStream(ctx context.Context, streamKey string, values map[string]interface{}) (string, error) {
|
||
// 使用 XAdd 添加消息到 Stream
|
||
messageID, err := RedisClient.XAdd(ctx, &redis.XAddArgs{
|
||
Stream: streamKey,
|
||
Values: values,
|
||
}).Result()
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
return messageID, nil
|
||
}
|
||
|
||
// ReadFromStream 从 Stream 读取消息（消费者组模式）
|
||
// 后台 Goroutine 使用此方法从 Stream 中取出请求进行处理
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - groupName: 消费者组名称
|
||
// - consumerName: 消费者名称（唯一标识）
|
||
// - count: 每次读取的消息数量
|
||
// - blockMs: 阻塞时间（毫秒），0表示不阻塞
|
||
//
|
||
// 返回：
|
||
// - []StreamMessage: 消息列表
|
||
// - error: 读取失败时返回错误
|
||
func ReadFromStream(ctx context.Context, streamKey, groupName, consumerName string, count int64, blockMs int64) ([]StreamMessage, error) {
|
||
// 使用 XReadGroup 从消费者组读取消息
|
||
// ">" 表示读取未被消费的新消息（只获取新消息）
|
||
// 如果使用 "0" 或其他 ID，则返回 Pending 消息（未确认的消息）
|
||
streams, err := RedisClient.XReadGroup(ctx, &redis.XReadGroupArgs{
|
||
Group: groupName,
|
||
Consumer: consumerName,
|
||
Streams: []string{streamKey, ">"}, // Stream名称 + 起始ID
|
||
Count: count,
|
||
Block: time.Duration(blockMs) * time.Millisecond,
|
||
}).Result()
|
||
|
||
// 处理错误：超时或没有数据时返回 redis.Nil
|
||
if err != nil {
|
||
if err == redis.Nil {
|
||
// 超时或没有数据，返回空数组
|
||
return []StreamMessage{}, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
// 解析返回的消息
|
||
var messages []StreamMessage
|
||
for _, stream := range streams {
|
||
for _, msg := range stream.Messages {
|
||
messages = append(messages, StreamMessage{
|
||
ID: msg.ID,
|
||
Values: msg.Values,
|
||
})
|
||
}
|
||
}
|
||
|
||
return messages, nil
|
||
}
|
||
|
||
// AckMessage 确认消息已处理
|
||
// 处理完消息后必须调用此方法确认，否则消息会保留在 Pending List (PEL)
|
||
// 确认后消息会从 PEL 中移除
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - groupName: 消费者组名称
|
||
// - messageIDs: 要确认的消息ID列表
|
||
//
|
||
// 返回：error 确认失败时返回错误
|
||
func AckMessage(ctx context.Context, streamKey, groupName string, messageIDs ...string) error {
|
||
// 使用 XAck 确认消息
|
||
// 返回值是成功确认的消息数量
|
||
count, err := RedisClient.XAck(ctx, streamKey, groupName, messageIDs...).Result()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
// 可以检查 count 是否等于 len(messageIDs)
|
||
_ = count
|
||
return nil
|
||
}
|
||
|
||
// GetStreamLength 获取 Stream 当前长度
|
||
// 用于监控 Stream 消息积压情况
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
//
|
||
// 返回：
|
||
// - int64: Stream 中消息数量
|
||
// - error: 操作失败时返回错误
|
||
func GetStreamLength(ctx context.Context, streamKey string) (int64, error) {
|
||
// 使用 XLen 获取 Stream 长度
|
||
length, err := RedisClient.XLen(ctx, streamKey).Result()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
return length, nil
|
||
}
|
||
|
||
// GetPendingMessages 获取待处理消息（未确认的消息）
|
||
// 用于监控和重试失败的消息
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - groupName: 消费者组名称
|
||
// - start: 起始ID，"-" 表示最小ID
|
||
// - end: 结束ID，"+" 表示最大ID
|
||
// - count: 返回数量
|
||
//
|
||
// 返回：
|
||
// - []redis.XPendingExt: Pending 消息列表
|
||
// - error: 操作失败时返回错误
|
||
func GetPendingMessages(ctx context.Context, streamKey, groupName string, start, end string, count int64) ([]redis.XPendingExt, error) {
|
||
// 使用 XPendingExt 获取详细的 Pending 消息
|
||
pending, err := RedisClient.XPendingExt(ctx, &redis.XPendingExtArgs{
|
||
Stream: streamKey,
|
||
Group: groupName,
|
||
Start: start,
|
||
End: end,
|
||
Count: count,
|
||
}).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return pending, nil
|
||
}
|
||
|
||
// ClaimPendingMessage 认领超时的 Pending 消息
|
||
// 当某个消费者故障后，其他消费者可以认领其未完成的消息
|
||
// 参数：
|
||
// - streamKey: Stream 键名
|
||
// - groupName: 消费者组名称
|
||
// - consumerName: 新消费者名称
|
||
// - minIdleTime: 消息空闲时间（毫秒），超过此时间才能被认领
|
||
// - messageIDs: 要认领的消息ID列表
|
||
//
|
||
// 返回：
|
||
// - []StreamMessage: 认领的消息列表
|
||
// - error: 操作失败时返回错误
|
||
func ClaimPendingMessage(ctx context.Context, streamKey, groupName, consumerName string, minIdleTime int64, messageIDs ...string) ([]StreamMessage, error) {
|
||
// 使用 XClaim 认领消息
|
||
msgs, err := RedisClient.XClaim(ctx, &redis.XClaimArgs{
|
||
Stream: streamKey,
|
||
Group: groupName,
|
||
Consumer: consumerName,
|
||
MinIdle: time.Duration(minIdleTime) * time.Millisecond,
|
||
Messages: messageIDs,
|
||
}).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 转换为 StreamMessage
|
||
var messages []StreamMessage
|
||
for _, msg := range msgs {
|
||
messages = append(messages, StreamMessage{
|
||
ID: msg.ID,
|
||
Values: msg.Values,
|
||
})
|
||
}
|
||
|
||
return messages, nil
|
||
}
|
||
|
||
// SetSessionLastActive 设置用户最后活跃时间
|
||
// 用于控制是否发送追问：用户回复后更新活跃时间，避免重复追问
|
||
// 过期时间：2小时，超过2小时未活跃的记录会自动删除
|
||
// 参数：
|
||
// - userId: 用户ID
|
||
//
|
||
// 返回：error 设置失败时返回错误
|
||
func SetSessionLastActive(ctx context.Context, userId string) error {
|
||
key := SessionLastActiveKeyPrefix + userId + ":last_active"
|
||
timestamp := time.Now().Unix()
|
||
|
||
// 设置过期时间为 2 小时
|
||
return RedisClient.Set(ctx, key, timestamp, 2*time.Hour).Err()
|
||
}
|
||
|
||
// GetSessionLastActive 获取用户最后活跃时间
|
||
// 参数：
|
||
// - userId: 用户ID
|
||
//
|
||
// 返回：
|
||
// - int64: Unix时间戳，未找到返回0
|
||
// - error: 操作失败时返回错误
|
||
func GetSessionLastActive(ctx context.Context, userId string) (int64, error) {
|
||
key := SessionLastActiveKeyPrefix + userId + ":last_active"
|
||
result, err := RedisClient.Get(ctx, key).Result()
|
||
if err == redis.Nil {
|
||
return 0, nil // 未找到返回 0
|
||
}
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
|
||
// 将字符串转换为 int64
|
||
timestamp, err := strconv.ParseInt(result, 10, 64)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
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
|
||
// 避免每次请求都创建新 Session，提高性能
|
||
// 过期时间：7天，超过7天未使用的Session会自动清理
|
||
// 参数：
|
||
// - userId: 用户ID
|
||
// - sessionId: RAGFlow返回的Session ID
|
||
//
|
||
// 返回：error 设置失败时返回错误
|
||
func SetSessionCache(ctx context.Context, userId, sessionId string) error {
|
||
key := SessionLastActiveKeyPrefix + userId + ":session_id"
|
||
return RedisClient.Set(ctx, key, sessionId, 7*24*time.Hour).Err()
|
||
}
|
||
|
||
// GetSessionCache 获取缓存的 RAGFlow Session ID
|
||
// 如果缓存中存在则直接使用，不存在则需要创建新Session
|
||
// 参数：
|
||
// - userId: 用户ID
|
||
//
|
||
// 返回：
|
||
// - string: Session ID，未找到返回空字符串
|
||
// - error: 操作失败时返回错误
|
||
func GetSessionCache(ctx context.Context, userId string) (string, error) {
|
||
key := SessionLastActiveKeyPrefix + userId + ":session_id"
|
||
result, err := RedisClient.Get(ctx, key).Result()
|
||
if err == redis.Nil {
|
||
return "", nil // 未找到返回空字符串
|
||
}
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return result, nil
|
||
}
|