重构了一下 rag的方法, 使用 goframe的框架, 还有redis连接部分

This commit is contained in:
Cold
2025-12-06 18:04:29 +08:00
committed by 张斌
parent f7cb007491
commit 4b2b5e6177
16 changed files with 398 additions and 260 deletions

View File

@@ -2,11 +2,12 @@ package rabbitmq
import (
"context"
"fmt"
"sync"
"time"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/util/gconv"
amqp "github.com/rabbitmq/amqp091-go"
)
@@ -33,13 +34,7 @@ func Init(ctx context.Context, cfg *Config) error {
var err error
once.Do(func() {
// 构建连接字符串
url := fmt.Sprintf("amqp://%s:%s@%s:%d/%s",
cfg.Username,
cfg.Password,
cfg.Host,
cfg.Port,
cfg.VHost,
)
url := "amqp://" + cfg.Username + ":" + cfg.Password + "@" + cfg.Host + ":" + gconv.String(cfg.Port) + "/" + cfg.VHost
// 创建连接
conn, err = amqp.Dial(url)
@@ -89,7 +84,7 @@ func GetChannel() (*amqp.Channel, error) {
defer mu.RUnlock()
if channel == nil || channel.IsClosed() {
return nil, fmt.Errorf("RabbitMQ Channel 未初始化或已关闭")
return nil, gerror.New("RabbitMQ Channel 未初始化或已关闭")
}
return channel, nil
@@ -101,7 +96,7 @@ func GetConnection() (*amqp.Connection, error) {
defer mu.RUnlock()
if conn == nil || conn.IsClosed() {
return nil, fmt.Errorf("RabbitMQ 连接未初始化或已关闭")
return nil, gerror.New("RabbitMQ 连接未初始化或已关闭")
}
return conn, nil
@@ -160,13 +155,7 @@ func reconnect(ctx context.Context) {
VHost: g.Cfg().MustGet(ctx, "rabbitmq.vhost", "/").String(),
}
url := fmt.Sprintf("amqp://%s:%s@%s:%d/%s",
cfg.Username,
cfg.Password,
cfg.Host,
cfg.Port,
cfg.VHost,
)
url := "amqp://" + cfg.Username + ":" + cfg.Password + "@" + cfg.Host + ":" + gconv.String(cfg.Port) + "/" + cfg.VHost
var err error
conn, err = amqp.Dial(url)
@@ -190,7 +179,7 @@ func reconnect(ctx context.Context) {
}
// Close 关闭连接
func Close(ctx context.Context) error {
func Close(ctx context.Context) (err error) {
mu.Lock()
defer mu.Unlock()
@@ -201,21 +190,21 @@ func Close(ctx context.Context) error {
}
if channel != nil {
if err := channel.Close(); err != nil {
if err = channel.Close(); err != nil {
g.Log().Errorf(ctx, "关闭 RabbitMQ Channel 失败: %v", err)
}
channel = nil
}
if conn != nil {
if err := conn.Close(); err != nil {
if err = conn.Close(); err != nil {
g.Log().Errorf(ctx, "关闭 RabbitMQ 连接失败: %v", err)
return err
return
}
conn = nil
}
watcherStarted = false
g.Log().Info(ctx, "RabbitMQ 连接已关闭")
return nil
return
}

View File

@@ -2,9 +2,9 @@ package rabbitmq
import (
"context"
"encoding/json"
"fmt"
"github.com/gogf/gf/v2/encoding/gjson"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
amqp "github.com/rabbitmq/amqp091-go"
)
@@ -74,7 +74,7 @@ func NewConsumer(queue string, handler MessageHandler, opts ...ConsumerOption) *
}
// Start 启动消费者
func (c *Consumer) Start(ctx context.Context) error {
func (c *Consumer) Start(ctx context.Context) (err error) {
// 创建可取消的 context
workerCtx, cancel := context.WithCancel(ctx)
c.cancel = cancel
@@ -90,7 +90,7 @@ func (c *Consumer) Start(ctx context.Context) error {
false, // global: false 表示仅应用于当前 channel
)
if err != nil {
return fmt.Errorf("设置 QoS 失败: %v", err)
return gerror.Newf("设置 QoS 失败: %v", err)
}
// 开始消费
@@ -104,7 +104,7 @@ func (c *Consumer) Start(ctx context.Context) error {
nil, // args
)
if err != nil {
return fmt.Errorf("开始消费失败: %v", err)
return gerror.Newf("开始消费失败: %v", err)
}
g.Log().Infof(ctx, "消费者已启动: queue=%s, prefetch=%d, workers=%d",
@@ -115,7 +115,7 @@ func (c *Consumer) Start(ctx context.Context) error {
go c.worker(workerCtx, i, msgs)
}
return nil
return
}
// worker 工作协程
@@ -168,8 +168,8 @@ func StartTypedConsumer[T any](
// 包装处理函数
wrappedHandler := func(ctx context.Context, body []byte) error {
var msg T
if err := json.Unmarshal(body, &msg); err != nil {
return fmt.Errorf("反序列化消息失败: %v", err)
if err := gjson.DecodeTo(body, &msg); err != nil {
return gerror.Newf("反序列化消息失败: %v", err)
}
return handler(ctx, &msg)

View File

@@ -2,9 +2,9 @@ package rabbitmq
import (
"context"
"encoding/json"
"fmt"
"github.com/gogf/gf/v2/encoding/gjson"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
amqp "github.com/rabbitmq/amqp091-go"
)
@@ -24,16 +24,16 @@ func NewPublisher(exchange, routingKey string) *Publisher {
}
// Publish 发布消息
func (p *Publisher) Publish(ctx context.Context, message interface{}) error {
func (p *Publisher) Publish(ctx context.Context, message interface{}) (err error) {
ch, err := GetChannel()
if err != nil {
return err
}
// 序列化消息
body, err := json.Marshal(message)
body, err := gjson.Encode(message)
if err != nil {
return fmt.Errorf("消息序列化失败: %v", err)
return gerror.Newf("消息序列化失败: %v", err)
}
// 发布消息
@@ -59,21 +59,21 @@ func (p *Publisher) Publish(ctx context.Context, message interface{}) error {
g.Log().Debugf(ctx, "消息发布成功: exchange=%s, routingKey=%s",
p.exchange, p.routingKey)
return nil
return
}
// PublishDelayed 发布延时消息
// delaySeconds: 延时秒数
func (p *Publisher) PublishDelayed(ctx context.Context, message interface{}, delaySeconds int) error {
func (p *Publisher) PublishDelayed(ctx context.Context, message interface{}, delaySeconds int) (err error) {
ch, err := GetChannel()
if err != nil {
return err
}
// 序列化消息
body, err := json.Marshal(message)
body, err := gjson.Encode(message)
if err != nil {
return fmt.Errorf("消息序列化失败: %v", err)
return gerror.Newf("消息序列化失败: %v", err)
}
// 发布延时消息(需要 rabbitmq_delayed_message_exchange 插件)
@@ -102,13 +102,13 @@ func (p *Publisher) PublishDelayed(ctx context.Context, message interface{}, del
g.Log().Debugf(ctx, "延时消息发布成功: exchange=%s, routingKey=%s, delay=%ds",
p.exchange, p.routingKey, delaySeconds)
return nil
return
}
// PublishBatch 批量发布消息
func (p *Publisher) PublishBatch(ctx context.Context, messages []interface{}) error {
func (p *Publisher) PublishBatch(ctx context.Context, messages []interface{}) (err error) {
if len(messages) == 0 {
return nil
return
}
ch, err := GetChannel()
@@ -117,7 +117,7 @@ func (p *Publisher) PublishBatch(ctx context.Context, messages []interface{}) er
}
for i, message := range messages {
body, err := json.Marshal(message)
body, err := gjson.Encode(message)
if err != nil {
g.Log().Errorf(ctx, "消息 %d 序列化失败: %v", i, err)
continue
@@ -143,5 +143,5 @@ func (p *Publisher) PublishBatch(ctx context.Context, messages []interface{}) er
}
g.Log().Infof(ctx, "批量发布完成: 共 %d 条消息", len(messages))
return nil
return
}

View File

@@ -2,8 +2,8 @@ package rabbitmq
import (
"context"
"fmt"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
amqp "github.com/rabbitmq/amqp091-go"
)
@@ -35,7 +35,7 @@ type BindingConfig struct {
}
// DeclareQueue 声明队列
func DeclareQueue(ctx context.Context, cfg *QueueConfig) error {
func DeclareQueue(ctx context.Context, cfg *QueueConfig) (err error) {
ch, err := GetChannel()
if err != nil {
return err
@@ -56,11 +56,11 @@ func DeclareQueue(ctx context.Context, cfg *QueueConfig) error {
}
g.Log().Infof(ctx, "队列声明成功: %s", cfg.Name)
return nil
return
}
// DeclareExchange 声明 Exchange
func DeclareExchange(ctx context.Context, cfg *ExchangeConfig) error {
func DeclareExchange(ctx context.Context, cfg *ExchangeConfig) (err error) {
ch, err := GetChannel()
if err != nil {
return err
@@ -82,11 +82,11 @@ func DeclareExchange(ctx context.Context, cfg *ExchangeConfig) error {
}
g.Log().Infof(ctx, "Exchange 声明成功: %s (type=%s)", cfg.Name, cfg.Type)
return nil
return
}
// BindQueue 绑定队列到 Exchange
func BindQueue(ctx context.Context, cfg *BindingConfig) error {
func BindQueue(ctx context.Context, cfg *BindingConfig) (err error) {
ch, err := GetChannel()
if err != nil {
return err
@@ -108,7 +108,7 @@ func BindQueue(ctx context.Context, cfg *BindingConfig) error {
g.Log().Infof(ctx, "队列绑定成功: queue=%s → exchange=%s (routingKey=%s)",
cfg.Queue, cfg.Exchange, cfg.RoutingKey)
return nil
return
}
// SetupDelayExchange 设置延时 Exchange需要 rabbitmq_delayed_message_exchange 插件)
@@ -165,9 +165,9 @@ func SetupQueueWithDLX(ctx context.Context, queueName, dlxExchange, dlxRoutingKe
}
// SetupBasicTopology 设置基础拓扑(适用于小红书客服场景)
func SetupBasicTopology(ctx context.Context) error {
func SetupBasicTopology(ctx context.Context) (err error) {
// 1. 声明普通 Exchange
err := DeclareExchange(ctx, &ExchangeConfig{
err = DeclareExchange(ctx, &ExchangeConfig{
Name: "ragflow_exchange",
Type: "direct",
Durable: true,
@@ -179,7 +179,7 @@ func SetupBasicTopology(ctx context.Context) error {
// 2. 声明延时 Exchange
err = SetupDelayExchange(ctx, "delay_exchange")
if err != nil {
return fmt.Errorf("延时 Exchange 声明失败(可能未安装插件): %v", err)
return gerror.Newf("延时 Exchange 声明失败(可能未安装插件): %v", err)
}
// 3. 声明死信队列
@@ -227,5 +227,5 @@ func SetupBasicTopology(ctx context.Context) error {
}
g.Log().Info(ctx, "RabbitMQ 拓扑结构设置完成")
return nil
return
}