wizardchen
a1a07e6818
Introduce KnowledgePostProcess task to decouple image multimodal processing from downstream operations like summary, question generation, and knowledge graph extraction. This allows unifying the processing logic across documents and images and ensures all text content, including from images, is properly extracted before post-processing.
127 lines
3.8 KiB
Go
127 lines
3.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
|
|
DataSourceService interfaces.DataSourceService
|
|
ChunkExtractor interfaces.TaskHandler `name:"chunkExtractor"`
|
|
DataTableSummary interfaces.TaskHandler `name:"dataTableSummary"`
|
|
ImageMultimodal interfaces.TaskHandler `name:"imageMultimodal"`
|
|
KnowledgePostProcess interfaces.TaskHandler `name:"knowledgePostProcess"`
|
|
}
|
|
|
|
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"),
|
|
Username: os.Getenv("REDIS_USERNAME"),
|
|
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.ChunkExtractor.Handle)
|
|
mux.HandleFunc(types.TypeDataTableSummary, params.DataTableSummary.Handle)
|
|
|
|
// Register document processing handler
|
|
mux.HandleFunc(types.TypeDocumentProcess, params.KnowledgeService.ProcessDocument)
|
|
|
|
// Register manual knowledge processing handler (cleanup + re-indexing)
|
|
mux.HandleFunc(types.TypeManualProcess, params.KnowledgeService.ProcessManualUpdate)
|
|
|
|
// Register FAQ import handler (includes dry run mode)
|
|
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 knowledge move handler
|
|
mux.HandleFunc(types.TypeKnowledgeMove, params.KnowledgeService.ProcessKnowledgeMove)
|
|
|
|
// Register knowledge list delete handler
|
|
mux.HandleFunc(types.TypeKnowledgeListDelete, params.KnowledgeService.ProcessKnowledgeListDelete)
|
|
|
|
// Register index delete handler
|
|
mux.HandleFunc(types.TypeIndexDelete, params.TagService.ProcessIndexDelete)
|
|
|
|
// Register KB delete handler
|
|
mux.HandleFunc(types.TypeKBDelete, params.KnowledgeBaseService.ProcessKBDelete)
|
|
|
|
// Register image multimodal handler
|
|
mux.HandleFunc(types.TypeImageMultimodal, params.ImageMultimodal.Handle)
|
|
|
|
// Register knowledge post process handler
|
|
mux.HandleFunc(types.TypeKnowledgePostProcess, params.KnowledgePostProcess.Handle)
|
|
|
|
// Register data source sync handler
|
|
mux.HandleFunc(types.TypeDataSourceSync, params.DataSourceService.ProcessSync)
|
|
|
|
go func() {
|
|
// Start the server
|
|
if err := params.Server.Run(mux); err != nil {
|
|
log.Fatalf("could not run server: %v", err)
|
|
}
|
|
}()
|
|
return mux
|
|
}
|