66 lines
1.6 KiB
Go
66 lines
1.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"os/signal"
|
|
"rag/consts/public"
|
|
"rag/controller"
|
|
"rag/service"
|
|
"strings"
|
|
"syscall"
|
|
|
|
"gitea.com/red-future/common/http"
|
|
"gitea.com/red-future/common/jaeger"
|
|
"gitea.com/red-future/common/utils"
|
|
gmq "github.com/bjang03/gmq/core/gmq"
|
|
"github.com/bjang03/gmq/mq"
|
|
"github.com/bjang03/gmq/types"
|
|
_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
)
|
|
|
|
func main() {
|
|
ctx := context.Background()
|
|
defer jaeger.ShutDown(ctx)
|
|
|
|
// 注册路由
|
|
http.RouteRegister([]interface{}{
|
|
controller.Dataset,
|
|
controller.Document,
|
|
controller.DocumentVector,
|
|
controller.Keyword,
|
|
})
|
|
|
|
if err := utils.InitGseTool(ctx); err != nil {
|
|
g.Log().Error(ctx, "gse 分词工具初始化失败:", err)
|
|
}
|
|
|
|
redisAddress := g.Cfg().MustGet(ctx, "redis.default.address").String()
|
|
redisAddressList := strings.Split(redisAddress, ":")
|
|
gmq.GmqRegister(public.GmqMsgPluginsName, &mq.RedisConn{
|
|
RedisConfig: mq.RedisConfig{
|
|
Addr: redisAddressList[0],
|
|
Port: redisAddressList[1],
|
|
},
|
|
})
|
|
if err := gmq.GetGmq(public.GmqMsgPluginsName).GmqSubscribe(ctx, &mq.RedisSubMessage{
|
|
SubMessage: types.SubMessage{
|
|
Topic: public.KnowledgeDocumentVectorTopic,
|
|
ConsumerName: public.KnowledgeDocumentVectorConsumer,
|
|
AutoAck: public.KnowledgeDocumentVectorAutoAck,
|
|
FetchCount: public.KnowledgeDocumentVectorCount,
|
|
HandleFunc: service.DocumentVector.DocsChunkMsg,
|
|
},
|
|
}); err != nil {
|
|
return
|
|
}
|
|
|
|
// 等待退出信号
|
|
quit := make(chan os.Signal, 1)
|
|
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
|
<-quit
|
|
|
|
g.Log().Info(ctx, "服务正在关闭...")
|
|
}
|