61 lines
1.3 KiB
Go
61 lines
1.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"os/signal"
|
|
"rag/consts/public"
|
|
"rag/controller"
|
|
"rag/service"
|
|
"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.DocumentChunk,
|
|
controller.Keyword,
|
|
controller.RAGQuery,
|
|
})
|
|
|
|
err := utils.InitGseTool(ctx)
|
|
if err != nil {
|
|
g.Log().Error(ctx, "gse 分词工具初始化失败:", err)
|
|
}
|
|
|
|
gmq.Init("config.yml")
|
|
|
|
if err := gmq.GetGmq("primary").GmqSubscribe(ctx, &mq.RedisSubMessage{
|
|
SubMessage: types.SubMessage{
|
|
Topic: public.KnowledgeDocumentChunkTopic,
|
|
ConsumerName: public.KnowledgeDocumentChunkConsumer,
|
|
AutoAck: public.KnowledgeDocumentChunkAutoAck,
|
|
FetchCount: public.KnowledgeDocumentChunkBatchSize,
|
|
HandleFunc: service.DocumentChunk.DocsChunkMsg,
|
|
},
|
|
}); err != nil {
|
|
return
|
|
}
|
|
|
|
// 等待退出信号
|
|
quit := make(chan os.Signal, 1)
|
|
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
|
<-quit
|
|
|
|
g.Log().Info(ctx, "服务正在关闭...")
|
|
}
|