wizardchen
105 lines
2.8 KiB
Go
105 lines
2.8 KiB
Go
package router
|
|
|
|
import (
|
|
"log"
|
|
"os"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
|
"github.com/hibiken/asynq"
|
|
"go.uber.org/dig"
|
|
)
|
|
|
|
type AsynqTaskParams struct {
|
|
dig.In
|
|
|
|
Server *asynq.Server
|
|
KnowledgeService interfaces.KnowledgeService
|
|
KnowledgeBaseService interfaces.KnowledgeBaseService
|
|
TagService interfaces.KnowledgeTagService
|
|
ChunkExtracter interfaces.TaskHandler `name:"chunkExtracter"`
|
|
DataTableSummary interfaces.TaskHandler `name:"dataTableSummary"`
|
|
}
|
|
|
|
func getAsynqRedisClientOpt() *asynq.RedisClientOpt {
|
|
db := 0
|
|
if dbStr := os.Getenv("REDIS_DB"); dbStr != "" {
|
|
if parsed, err := strconv.Atoi(dbStr); err == nil {
|
|
db = parsed
|
|
}
|
|
}
|
|
opt := &asynq.RedisClientOpt{
|
|
Addr: os.Getenv("REDIS_ADDR"),
|
|
Password: os.Getenv("REDIS_PASSWORD"),
|
|
ReadTimeout: 100 * time.Millisecond,
|
|
WriteTimeout: 200 * time.Millisecond,
|
|
DB: db,
|
|
}
|
|
return opt
|
|
}
|
|
|
|
func NewAsyncqClient() (*asynq.Client, error) {
|
|
opt := getAsynqRedisClientOpt()
|
|
client := asynq.NewClient(opt)
|
|
err := client.Ping()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return client, nil
|
|
}
|
|
|
|
func NewAsynqServer() *asynq.Server {
|
|
opt := getAsynqRedisClientOpt()
|
|
srv := asynq.NewServer(
|
|
opt,
|
|
asynq.Config{
|
|
Queues: map[string]int{
|
|
"critical": 6, // Highest priority queue
|
|
"default": 3, // Default priority queue
|
|
"low": 1, // Lowest priority queue
|
|
},
|
|
},
|
|
)
|
|
return srv
|
|
}
|
|
|
|
func RunAsynqServer(params AsynqTaskParams) *asynq.ServeMux {
|
|
// Create a new mux and register all handlers
|
|
mux := asynq.NewServeMux()
|
|
|
|
// Register extract handlers - router will dispatch to appropriate handler
|
|
mux.HandleFunc(types.TypeChunkExtract, params.ChunkExtracter.Handle)
|
|
mux.HandleFunc(types.TypeDataTableSummary, params.DataTableSummary.Handle)
|
|
|
|
// Register document processing handler
|
|
mux.HandleFunc(types.TypeDocumentProcess, params.KnowledgeService.ProcessDocument)
|
|
|
|
// Register FAQ import handler
|
|
mux.HandleFunc(types.TypeFAQImport, params.KnowledgeService.ProcessFAQImport)
|
|
|
|
// Register question generation handler
|
|
mux.HandleFunc(types.TypeQuestionGeneration, params.KnowledgeService.ProcessQuestionGeneration)
|
|
|
|
// Register summary generation handler
|
|
mux.HandleFunc(types.TypeSummaryGeneration, params.KnowledgeService.ProcessSummaryGeneration)
|
|
|
|
// Register KB clone handler
|
|
mux.HandleFunc(types.TypeKBClone, params.KnowledgeService.ProcessKBClone)
|
|
|
|
// Register index delete handler
|
|
mux.HandleFunc(types.TypeIndexDelete, params.TagService.ProcessIndexDelete)
|
|
|
|
// Register KB delete handler
|
|
mux.HandleFunc(types.TypeKBDelete, params.KnowledgeBaseService.ProcessKBDelete)
|
|
|
|
go func() {
|
|
// Start the server
|
|
if err := params.Server.Run(mux); err != nil {
|
|
log.Fatalf("could not run server: %v", err)
|
|
}
|
|
}()
|
|
return mux
|
|
}
|