zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-05 07:10:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: use NATS as the message queue (#15327)

Browse Source 
### What problem does this PR solve?

```
RAGFlow(admin)> mq publish 'msg2';
SUCCESS
RAGFlow(admin)> mq publish 'msg3';
SUCCESS
RAGFlow(admin)> mq list;
+---------+---------------+
| message | subject       |
+---------+---------------+
| msg1    | tasks.RAGFLOW |
| msg2    | tasks.RAGFLOW |
| msg3    | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq pull 2;
+---------+---------------+
| message | subject       |
+---------+---------------+
| msg1    | tasks.RAGFLOW |
| msg2    | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq pull noack;
+---------+---------------+
| message | subject       |
+---------+---------------+
| abc     | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq show
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+
| ack_pending_count | consumer_count | memory | message_count | pending_count | redelivered_count | waiting_count |
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+
| 2                 | 1              | 0      | 2             | 0             | 1                 | 0             |
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+

RAGFlow(admin)> list ingestors;
+--------------+-------------------------------------------+--------+
| host         | name                                      | status |
+--------------+-------------------------------------------+--------+
| 192.168.1.38 | ingestor-8f0e4bd5650a4ac58b0151969fbf6935 | alive  |
+--------------+-------------------------------------------+--------+

RAGFlow(admin)> list ingestion tasks;
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+
| document_id                      | id                               | status    | step | user        | user_id                          |
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+
| ffe64fae423411f1a2d938a74640adcc | 90d3d0f6528941c1ac8eb0360effccc4 | COMPLETED | 5    | aaa@aaa.com | 2ba4881420fa11f19e9c38a74640adcc |
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+

RAGFlow(admin)> remove ingestion tasks '90d3d0f6528941c1ac8eb0360effccc4';
+---------+----------------------------------+
| delete  | task_id                          |
+---------+----------------------------------+
| success | 90d3d0f6528941c1ac8eb0360effccc4 |
+---------+----------------------------------+

RAGFlow(admin)> stop ingestion tasks 'e89e20d9a25848a1b79bd9345ddbfe1d';
+----------+----------------------------------+
| status   | task_id                          |
+----------+----------------------------------+
| STOPPING | e89e20d9a25848a1b79bd9345ddbfe1d |
+----------+----------------------------------+

# Publish a message
RAGFlow(admin)> mq publish 'cdd';
SUCCESS

# List current tasks in the message queue
RAGFlow(admin)> mq list
+----------------------------------+---------------+
| message                          | subject       |
+----------------------------------+---------------+
| 7ce392a3c1624cd2be4b5276e8825059 | tasks.RAGFLOW |
+----------------------------------+---------------+

# Consume a task from the message queue
RAGFlow(admin)> mq pull
+------+-----+----------------+
| ack  | id  | type           |
+------+-----+----------------+
| true | cdd | ingestion_test |
+------+-----+----------------+

# User mode
# List ingestion tasks, followed by dataset id
RAGFlow(user)> list ingestion tasks from '0abe79f9423311f1ad8d38a74640adcc';
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-05-30T20:21:06+08:00 | 1780143666289 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | 8d758cd14a8b4ba8ab505003fb52017d |        | COMPLETED | 2026-05-30T20:21:26+08:00 | 1780143686431 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

RAGFlow(user)> list ingestion tasks;
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-06-02T19:02:31+08:00 | 1780398151417 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | e89e20d9a25848a1b79bd9345ddbfe1d |        | COMPLETED | 2026-06-02T19:02:52+08:00 | 1780398172208 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

# Create an ingestion task
# First argument is document id, second argument is dataset id
RAGFlow(user)> start ingestion 'ffe64fae423411f1a2d938a74640adcc' from '0abe79f9423311f1ad8d38a74640adcc';
+----------------------------------+-------------------------------------------+
| document_id                      | result                                    |
+----------------------------------+-------------------------------------------+
| ffe64fae423411f1a2d938a74640adcc | task_id: 8d758cd14a8b4ba8ab505003fb52017d |
+----------------------------------+-------------------------------------------+

# Pause an ingestion task, first argument is ingestion id
RAGFlow(user)> stop ingestion '8d758cd14a8b4ba8ab505003fb52017d';
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-05-30T20:21:06+08:00 | 1780143666289 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | 8d758cd14a8b4ba8ab505003fb52017d |        | COMPLETED | 2026-05-30T20:21:26+08:00 | 1780143686431 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

# Delete an ingestion task
RAGFlow(api/default)> remove ingestion tasks 'f366450a27d54677aec1c7090add30f0';
+---------+----------------------------------+
| remove  | task_id                          |
+---------+----------------------------------+
| success | f366450a27d54677aec1c7090add30f0 |
+---------+----------------------------------+

```

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-06-12 14:56:44 +08:00
committed by GitHub
parent 30724140d2
commit e96bc37d06
45 changed files with 2911 additions and 2337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
internal/admin/handler.go
View File

@@ -17,12 +17,14 @@
package admin
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"ragflow/internal/cache"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/engine"
"ragflow/internal/server"
"ragflow/internal/service"
"ragflow/internal/utility"
@@ -203,15 +205,6 @@ func (h *Handler) AuthCheck(c *gin.Context) {
successNoData(c, "Admin is authorized")
}
// ListTasks handle list tasks
func (h *Handler) ListTasks(c *gin.Context) {
tasks, err := h.service.ListTasks()
if err != nil {
errorResponse(c, err.Error(), 500)
}
success(c, tasks, "Get all tasks")
}
// ListUsers handle list users
func (h *Handler) ListUsers(c *gin.Context) {
users, err := h.service.ListUsers()
@@ -261,7 +254,7 @@ func (h *Handler) GetUser(c *gin.Context) {
userDetails, err := h.service.GetUserDetails(username)
if err != nil {
if errors.Is(err, ErrUserNotFound) {
if errors.Is(err, common.ErrUserNotFound) {
errorResponse(c, "User not found", 404)
return
}
@@ -1256,57 +1249,206 @@ func (h *Handler) SetLogLevel(c *gin.Context) {
success(c, gin.H{"level": req.Level}, "Log level updated successfully")
}
type StartIngestionTaskRequest struct {
FileURI string `json:"uri" binding:"required"`
From string `json:"from" binding:"required"`
func (h *Handler) ListMessagesFromQueue(c *gin.Context) {
msgQueueEngine := engine.GetMessageQueueEngine()
messages, err := msgQueueEngine.ListMessages("ingestion", false)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
var result []map[string]string
for _, message := range messages {
var taskMessage common.TaskMessage
err = json.Unmarshal([]byte(message["message"]), &taskMessage)
if err != nil {
return
}
result = append(result, map[string]string{
"subject": message["subject"],
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
})
}
success(c, result, "List messages from queue successfully")
}
func (h *Handler) StartIngestionTask(c *gin.Context) {
var req StartIngestionTaskRequest
type PublishMessageToQueueRequest struct {
Message string `json:"message" binding:"required"`
}
func (h *Handler) PublishMessageToQueue(c *gin.Context) {
var req PublishMessageToQueueRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "file uri and from is required", 400)
errorResponse(c, "message is required", 400)
return
}
taskID := common.GenerateUUID()
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: taskID,
TaskType: "start_ingestion_task",
Config: req.FileURI,
ComeFrom: req.From,
})
taskMessage := common.TaskMessage{
TaskID: req.Message,
TaskType: common.TaskTypeIngestionTest,
}
success(c, gin.H{"task_id": taskID}, "Send task for ingestion successfully")
// convert task
taskMessageStr, err := json.Marshal(taskMessage)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
msgQueueEngine := engine.GetMessageQueueEngine()
err = msgQueueEngine.PublishTask("tasks.RAGFLOW", taskMessageStr)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, nil, "Publish message successfully")
}
type PullMessageFromQueueRequest struct {
MessageCount int `json:"message_count" binding:"required"`
AckPolicy string `json:"ack_policy" binding:"required"`
}
func (h *Handler) PullMessageFromQueue(c *gin.Context) {
var req PullMessageFromQueueRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, fmt.Sprintf("message count and ack_policy are required, error: %s", err.Error()), 400)
return
}
msgQueueEngine := engine.GetMessageQueueEngine()
err := msgQueueEngine.InitConsumer("tasks.RAGFLOW")
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
messages, err := msgQueueEngine.GetMessages(req.MessageCount)
var result []map[string]string
if req.AckPolicy == "ACK" {
for _, message := range messages {
taskMessage := message.GetMessage()
resultMessage := map[string]string{
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
}
err = message.Ack()
if err == nil {
resultMessage["ack"] = "true"
} else {
resultMessage["ack"] = "false"
}
result = append(result, resultMessage)
}
} else {
for _, message := range messages {
taskMessage := message.GetMessage()
resultMessage := map[string]string{
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
}
if err == nil {
resultMessage["nack"] = "true"
} else {
resultMessage["nack"] = "false"
}
result = append(result, resultMessage)
}
}
success(c, result, "Pull messages from queue successfully")
}
func (h *Handler) ShowMessageQueue(c *gin.Context) {
msgQueueEngine := engine.GetMessageQueueEngine()
result, err := msgQueueEngine.ShowMessageQueue()
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, result, "show message queue successfully")
}
type RemoveIngestionTaskRequest struct {
Tasks []string `json:"tasks" binding:"required"`
}
func (h *Handler) RemoveIngestionTasks(c *gin.Context) {
var req RemoveIngestionTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "task id is required", 400)
return
}
tasks, err := h.service.RemoveIngestionTasks(req.Tasks)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, tasks, "Remove tasks successfully")
}
type StopIngestionTaskRequest struct {
TaskID string `json:"task_id" binding:"required"`
From string `json:"from" binding:"required"`
Tasks []string `json:"tasks" binding:"required"`
}
func (h *Handler) StopIngestionTask(c *gin.Context) {
func (h *Handler) StopIngestionTasks(c *gin.Context) {
var req StopIngestionTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "task id and from is required", 400)
return
}
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: req.TaskID,
TaskType: "cancel_ingestion_task",
ComeFrom: req.From,
})
tasks, err := h.service.StopIngestionTasks(req.Tasks)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, gin.H{"task_id": req.TaskID}, "Cancel task successfully")
var result []map[string]string
for _, task := range tasks {
result = append(result, map[string]string{
"task_id": task.ID,
"status": task.Status,
})
}
success(c, result, "Stop tasks successfully")
}
func (h *Handler) ListIngestors(c *gin.Context) {
ingestionMgr := GetIngestionManager()
ingestors, err := ingestionMgr.ListIngestors()
// ListIngestionTasks
func (h *Handler) ListIngestionTasks(c *gin.Context) {
tasks, err := h.service.ListIngestionTasks()
if err != nil {
errorResponse(c, err.Error(), 500)
}
success(c, ingestors, "Get all tasks")
success(c, tasks, "Get all tasks")
}
func (h *Handler) ListIngestors(c *gin.Context) {
serverList := GlobalServerStore.ListInfos()
var ingestorResults []map[string]string
now := time.Now()
for _, ingestorServer := range serverList {
if ingestorServer.ServerType == common.ServerTypeIngestion {
ingestorResult := map[string]string{}
ingestorResult["name"] = ingestorServer.ServerName
ingestorResult["host"] = ingestorServer.Host
ingestorResult["status"] = ingestorServer.Version
if now.Sub(ingestorServer.Timestamp) < 30*time.Second {
ingestorResult["status"] = "alive"
} else {
ingestorResult["status"] = "timeout"
}
ingestorResults = append(ingestorResults, ingestorResult)
}
}
success(c, ingestorResults, "Get all tasks")
}
type ShutdownIngestorRequest struct {
@@ -1321,11 +1463,11 @@ func (h *Handler) ShutdownIngestor(c *gin.Context) {
}
taskID := common.GenerateUUID()
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: taskID,
TaskType: "shutdown_ingestor",
AssignedTo: req.IngestorID,
})
//ingestionManager.SubmitTask(&common.TaskAssignment{
// TaskId: taskID,
// TaskType: "SHUTDOWN",
// AssignedTo: req.IngestorID,
//})
success(c, gin.H{"task_id": taskID, "ingestor_id": req.IngestorID}, "Shutdown ingestor")
}
@@ -1364,14 +1506,3 @@ func (h *Handler) Reports(c *gin.Context) {
responseWithCode(c, message, http.StatusOK, errCode)
}
// ListIngestionTasks
func (h *Handler) ListIngestionTasks(c *gin.Context) {
tasks, err := h.service.ListIngestionTasks()
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, tasks, "")
}