zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-28 03:38:11 +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

196
internal/cli/admin_command.go
View File

@@ -1307,6 +1307,7 @@ func (c *CLI) ListAdminIngestors(cmd *Command) (ResponseIf, error) {
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) ListAdminIngestionTasks(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
@@ -1334,21 +1335,20 @@ func (c *CLI) ListAdminIngestionTasks(cmd *Command) (ResponseIf, error) {
return &result, nil
}
func (c *CLI) AdminStartIngestionCommand(cmd *Command) (ResponseIf, error) {
func (c *CLI) AdminStopIngestionCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
fileURI, ok := cmd.Params["uri"].(string)
tasks, ok := cmd.Params["tasks"].([]string)
if !ok {
return nil, fmt.Errorf("uri not provided")
}
payload := map[string]interface{}{
"uri": fileURI,
"from": "CLI",
"tasks": tasks,
}
resp, err := c.AdminServerClient.Request("POST", "/admin/ingestion", "admin", nil, payload)
resp, err := c.AdminServerClient.Request("PUT", "/admin/ingestion/tasks", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to ingest file: %w", err)
}
@@ -1357,7 +1357,7 @@ func (c *CLI) AdminStartIngestionCommand(cmd *Command) (ResponseIf, error) {
return nil, fmt.Errorf("failed to ingest file: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonDataResponse
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("ingest file failed: invalid JSON (%w)", err)
}
@@ -1370,21 +1370,20 @@ func (c *CLI) AdminStartIngestionCommand(cmd *Command) (ResponseIf, error) {
return &result, nil
}
func (c *CLI) AdminStopIngestionCommand(cmd *Command) (ResponseIf, error) {
func (c *CLI) AdminRemoveIngestionCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
taskID, ok := cmd.Params["task_id"].(string)
tasks, ok := cmd.Params["tasks"].([]string)
if !ok {
return nil, fmt.Errorf("uri not provided")
}
payload := map[string]interface{}{
"task_id": taskID,
"from": "CLI",
"tasks": tasks,
}
resp, err := c.AdminServerClient.Request("DELETE", "/admin/ingestion", "admin", nil, payload)
resp, err := c.AdminServerClient.Request("DELETE", "/admin/ingestion/tasks", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to ingest file: %w", err)
}
@@ -1393,7 +1392,7 @@ func (c *CLI) AdminStopIngestionCommand(cmd *Command) (ResponseIf, error) {
return nil, fmt.Errorf("failed to ingest file: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonDataResponse
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("ingest file failed: invalid JSON (%w)", err)
}
@@ -1440,3 +1439,176 @@ func (c *CLI) AdminShutdownIngestor(cmd *Command) (ResponseIf, error) {
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserListMessageQueueCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
pending, ok := cmd.Params["pending"].(bool)
if !ok {
return nil, fmt.Errorf("pending not provided")
}
payload := map[string]interface{}{
"pending": pending,
}
resp, err := c.AdminServerClient.Request("GET", "/admin/queue/messages", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to list tasks in message queue: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to list tasks in message queue: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("list tasks in message queue failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserPublishMessageCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
message, ok := cmd.Params["message"].(string)
if !ok {
return nil, fmt.Errorf("message not provided")
}
payload := map[string]interface{}{
"message": message,
}
resp, err := c.AdminServerClient.Request("POST", "/admin/queue/messages", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to publish message: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to publish message: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result SimpleResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("publish message failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserPullMessageCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
messageCount, ok := cmd.Params["message_count"].(int)
if !ok {
return nil, fmt.Errorf("message_count not provided")
}
ackPolicy, ok := cmd.Params["ack_policy"].(string)
if !ok {
return nil, fmt.Errorf("ack_policy not provided")
}
payload := map[string]interface{}{
"message_count": messageCount,
"ack_policy": ackPolicy,
}
resp, err := c.AdminServerClient.Request("PUT", "/admin/queue/messages", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to pull message: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to pull message: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("pull message failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserShowMessageQueueCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
resp, err := c.AdminServerClient.Request("GET", "/admin/queue", "admin", nil, nil)
if err != nil {
return nil, fmt.Errorf("failed to show message queue: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to show message queue: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonDataResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("show message queue failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) AdminRemoveServiceCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil {
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
}
serviceNumber, ok := cmd.Params["service_number"].(int)
if !ok {
return nil, fmt.Errorf("service_number not provided")
}
payload := map[string]interface{}{
"service_number": serviceNumber,
}
resp, err := c.AdminServerClient.Request("DELETE", "/admin/services", "admin", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to remove unavailable service: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to remove unavailable service: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result SimpleResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("remove unavailable service failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}

171
internal/cli/admin_parser.go
View File

@@ -16,7 +16,10 @@
package cli
import "fmt"
import (
"fmt"
"strings"
)
// Command parsers
func (p *Parser) parseAdminLoginUser() (*Command, error) {
@@ -190,8 +193,6 @@ func (p *Parser) parseAdminListCommand() (*Command, error) {
return NewCommand("list_user_chats"), nil
case TokenFiles:
return p.parseAdminListFiles()
case TokenTasks:
return p.parseAdminListTasks()
case TokenIngestors:
return p.parseAdminListIngestors()
case TokenIngestion:
@@ -376,12 +377,6 @@ func (p *Parser) parseAdminListFiles() (*Command, error) {
return cmd, nil
}
func (p *Parser) parseAdminListTasks() (*Command, error) {
p.nextToken() // consume TASKS
cmd := NewCommand("list_admin_tasks")
return cmd, nil
}
func (p *Parser) parseAdminListIngestors() (*Command, error) {
p.nextToken() // consume TASKS
cmd := NewCommand("admin_list_ingestors")
@@ -389,14 +384,80 @@ func (p *Parser) parseAdminListIngestors() (*Command, error) {
return cmd, nil
}
func (p *Parser) parseAdminStopIngestionTasks() (*Command, error) {
p.nextToken() // consume STOP
if p.curToken.Type != TokenIngestion {
return nil, fmt.Errorf("expected INGESTION")
}
p.nextToken()
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expected TASKS")
}
p.nextToken() // consume TASK
taskString, err := p.parseQuotedString()
if err != nil {
return nil, err
}
tasks := strings.Split(taskString, " ")
p.nextToken() // consume TASK
cmd := NewCommand("admin_stop_ingestion_tasks")
cmd.Params["tasks"] = tasks
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseAdminRemoveIngestionTasks() (*Command, error) {
p.nextToken() // consume Ingestion
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expected TASKS")
}
p.nextToken() // consume TASKS
taskString, err := p.parseQuotedString()
if err != nil {
return nil, err
}
tasks := strings.Split(taskString, " ")
p.nextToken() // consume TASKS
cmd := NewCommand("admin_remove_ingestion_tasks")
cmd.Params["tasks"] = tasks
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseAdminListIngestionTasks() (*Command, error) {
p.nextToken() // consume Ingestion
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expected TASKS")
}
p.nextToken() // consume TASKS
cmd := NewCommand("list_admin_ingestion_tasks")
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@@ -1752,7 +1813,6 @@ func (p *Parser) parseAdminIngestCommand() (*Command, error) {
}
return cmd, nil
}
func (p *Parser) parseAdminUnsetCommand() (*Command, error) {
p.nextToken() // consume UNSET
@@ -1767,3 +1827,94 @@ func (p *Parser) parseAdminUnsetCommand() (*Command, error) {
}
return NewCommand("unset_token"), nil
}
func (p *Parser) parseMessageQueueCommand() (*Command, error) {
p.nextToken() // consume MESSAGE_QUEUE
var cmd *Command
switch p.curToken.Type {
case TokenShow:
p.nextToken()
cmd = NewCommand("user_show_message_queue_command")
case TokenList:
p.nextToken() // consume LIST
cmd = NewCommand("user_list_message_queue_command")
if p.curToken.Type == TokenPending {
cmd.Params["pending"] = true
p.nextToken() // consume PENDING
} else {
cmd.Params["pending"] = false
}
case TokenPublish:
p.nextToken() // consume PUBLISH
message, err := p.parseQuotedString()
if err != nil {
return nil, fmt.Errorf("expected message after PUBLISH")
}
p.nextToken() // consume message
cmd = NewCommand("user_publish_message_command")
cmd.Params["message"] = message
case TokenPull:
p.nextToken() // consume PULL
messageCount, err := p.parseNumber()
if err != nil {
messageCount = 1
} else {
p.nextToken() // consume NUMBER
}
if messageCount <= 0 || messageCount > 100 {
return nil, fmt.Errorf("message count cannot be less than 0 or greater than 100")
}
cmd = NewCommand("user_pull_message_command")
cmd.Params["message_count"] = messageCount
if p.curToken.Type == TokenNoACK {
cmd.Params["ack_policy"] = "NOACK"
p.nextToken() // consume NOACK
} else {
cmd.Params["ack_policy"] = "ACK"
}
default:
return nil, fmt.Errorf("expected WITH")
}
// Semicolon is optional
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseAdminRemoveCommand() (*Command, error) {
p.nextToken() // consume MESSAGE_QUEUE
var cmd *Command
switch p.curToken.Type {
case TokenService:
p.nextToken() // consume SERVICE
serviceNum, err := p.parseNumber()
if err != nil {
return nil, fmt.Errorf("expected service number after SERVICE")
}
p.nextToken() // consume service number
cmd = NewCommand("admin_remove_service_command")
cmd.Params["service_number"] = serviceNum
case TokenIngestion:
return p.parseAdminRemoveIngestionTasks()
default:
return nil, fmt.Errorf("expected SERVICE")
}
// Semicolon is optional
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}

30
internal/cli/cli_http.go
View File

@@ -16,7 +16,9 @@
package cli
import "fmt"
import (
"fmt"
)
// ExecuteCommand executes a parsed command
// Returns benchmark result map for commands that support it (e.g., ping_server with iterations > 1)
@@ -103,14 +105,25 @@ func (c *CLI) ExecuteAdminCommand(cmd *Command) (ResponseIf, error) {
return c.ListAdminTasks(cmd)
case "admin_list_ingestors":
return c.ListAdminIngestors(cmd)
case "admin_start_ingestion_command":
return c.AdminStartIngestionCommand(cmd)
case "admin_stop_ingestion_command":
case "admin_stop_ingestion_tasks":
return c.AdminStopIngestionCommand(cmd)
case "admin_remove_ingestion_tasks":
return c.AdminRemoveIngestionCommand(cmd)
case "admin_shutdown_ingestor_command":
return c.AdminShutdownIngestor(cmd)
case "list_admin_ingestion_tasks":
return c.ListAdminIngestionTasks(cmd)
case "user_list_message_queue_command":
return c.UserListMessageQueueCommand(cmd)
case "user_publish_message_command":
return c.UserPublishMessageCommand(cmd)
case "user_pull_message_command":
return c.UserPullMessageCommand(cmd)
case "user_show_message_queue_command":
return c.UserShowMessageQueueCommand(cmd)
case "admin_remove_service_command":
return c.AdminRemoveServiceCommand(cmd)
// TODO: Implement other commands
case "show_admin_server":
return c.ShowAdminServer(cmd)
case "show_api_server":
@@ -285,6 +298,15 @@ func (c *CLI) ExecuteUserCommand(cmd *Command) (ResponseIf, error) {
return c.GetMetadata(cmd)
case "parse_documents_user_command":
return c.ParseDocumentsUserCommand(cmd)
case "user_start_ingestion_command":
return c.UserStartIngestionCommand(cmd)
case "user_stop_ingestion_command":
return c.UserStopIngestionCommand(cmd)
case "user_list_ingestion_tasks":
return c.ListUserIngestionTasks(cmd)
case "user_remove_task_command":
return c.UserRemoveTaskCommand(cmd)
// TODO: Implement other commands
case "user_parse_local_file_command":
return c.UserParseLocalFile(cmd)
case "show_admin_server":

10
internal/cli/lexer.go
View File

@@ -477,6 +477,16 @@ func (l *Lexer) lookupIdent(ident string) Token {
return Token{Type: TokenIngestors, Value: ident}
case "INGESTION":
return Token{Type: TokenIngestion, Value: ident}
case "MQ":
return Token{Type: TokenMQ, Value: ident}
case "PUBLISH":
return Token{Type: TokenPublish, Value: ident}
case "PULL":
return Token{Type: TokenPull, Value: ident}
case "PENDING":
return Token{Type: TokenPending, Value: ident}
case "NOACK":
return Token{Type: TokenNoACK, Value: ident}
case "LOG":
return Token{Type: TokenLog, Value: ident}
case "LEVEL":

13
internal/cli/parser.go
View File

@@ -124,10 +124,12 @@ func (p *Parser) parseAdminCommand() (*Command, error) {
return p.parseAdminShutdownCommand()
case TokenRestart:
return p.parseAdminRestartCommand()
case TokenStart:
return p.parseStartIngestion()
case TokenMQ:
return p.parseMessageQueueCommand()
case TokenRemove:
return p.parseAdminRemoveCommand()
case TokenStop:
return p.parseStopIngestion()
return p.parseAdminStopIngestionTasks()
case TokenAdd:
return p.parseAdminAddCommand()
case TokenDelete:
@@ -216,6 +218,11 @@ func (p *Parser) parseUserCommand() (*Command, error) {
return p.parseOCRCommand()
case TokenCheck:
return p.parseCheckCommand()
case TokenStart:
return p.parseUserStartIngestion()
case TokenStop:
return p.parseUserStopIngestion()
case TokenSave:
return p.parseUserSaveCommand()
case TokenUse:

5
internal/cli/types.go
View File

@@ -170,6 +170,11 @@ const (
TokenStart
TokenStop
TokenIngestion
TokenMQ
TokenPublish
TokenPull
TokenPending
TokenNoACK
TokenLog
TokenLevel
TokenDebug

169
internal/cli/user_command.go
View File

@@ -3273,6 +3273,173 @@ func formatRequestError(action string, err error) error {
}
}
func (c *CLI) ListUserIngestionTasks(cmd *Command) (ResponseIf, error) {
if c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].APIToken == nil && c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].LoginToken == nil {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.Config.CLIMode != APIMode {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
datasetID, ok := cmd.Params["dataset_id"].(*string)
if !ok {
datasetID = nil
}
payload := map[string]interface{}{
"dataset_id": datasetID,
}
resp, err := c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].Request("GET", "/datasets/ingestion/tasks", "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to list ingestion tasks: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to list ingestion tasks:: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("list ingestion tasks: failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserStartIngestionCommand(cmd *Command) (ResponseIf, error) {
if c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].APIToken == nil && c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].LoginToken == nil {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.Config.CLIMode != APIMode {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
documentID, ok := cmd.Params["document_id"].(string)
if !ok {
return nil, fmt.Errorf("document_id not provided")
}
datasetID, ok := cmd.Params["dataset_id"].(string)
if !ok {
return nil, fmt.Errorf("dataset_id not provided")
}
payload := map[string]interface{}{
"documents": []string{documentID},
"dataset_id": datasetID,
}
url := fmt.Sprintf("/datasets/%s/documents/parse", datasetID)
resp, err := c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].Request("POST", url, "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to ingest file: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to ingest file: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("ingest file failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserStopIngestionCommand(cmd *Command) (ResponseIf, error) {
if c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].APIToken == nil && c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].LoginToken == nil {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.Config.CLIMode != APIMode {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
tasks, ok := cmd.Params["tasks"].([]string)
if !ok {
return nil, fmt.Errorf("uri not provided")
}
payload := map[string]interface{}{
"tasks": tasks,
}
resp, err := c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].Request("PUT", "/datasets/ingestion/tasks", "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to ingest file: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to ingest file: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("ingest file failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) UserRemoveTaskCommand(cmd *Command) (ResponseIf, error) {
if c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].APIToken == nil && c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].LoginToken == nil {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.Config.CLIMode != APIMode {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
tasks, ok := cmd.Params["tasks"].([]string)
if !ok {
return nil, fmt.Errorf("tasks not provided")
}
payload := map[string]interface{}{
"tasks": tasks,
}
resp, err := c.APIServerClientMap[c.Config.APIClientConfig.CurrentAPIServer].Request("DELETE", "/datasets/ingestion/tasks", "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to remove tasks: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to remove tasks: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
}
var result CommonResponse
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("remove tasks failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *CLI) ChunkCommand(cmd *Command) (ResponseIf, error) {
if c.Config.CLIMode != APIMode {
return nil, fmt.Errorf("this command is only allowed in USER mode")
@@ -3294,12 +3461,12 @@ func (c *CLI) ChunkCommand(cmd *Command) (ResponseIf, error) {
if explain {
fmt.Printf("Explain chunk file: %s, DSL: %s\n", filename, dsl)
} else {
// TODO: not implemented
fmt.Printf("Chunk file: %s, DSL: %s\n", filename, dsl)
}
var result SimpleResponse
result.Code = 0
result.Message = fmt.Sprintf("Success to chunk %s", filename)
return &result, nil
}

126
internal/cli/user_parser.go
View File

@@ -159,6 +159,8 @@ func (p *Parser) parseListCommand() (*Command, error) {
return p.parseListProviders()
case TokenInstances:
return p.parseListInstances()
case TokenIngestion:
return p.parseUserListIngestionTasks()
case TokenDefault:
return p.parseListDefaultModels()
case TokenAvailable:
@@ -1387,6 +1389,8 @@ func (p *Parser) parseRemoveCommand() (*Command, error) {
return p.parseRemoveTags()
case TokenChunks, TokenAll:
return p.parseRemoveChunk()
case TokenIngestion:
return p.parseUserRemoveTask()
case TokenModel:
return p.parseRemoveInstanceModel()
default:
@@ -3892,7 +3896,7 @@ optionsLoop:
if err != nil {
return nil, err
}
cmd.Params["embed_model"] = embedModel
cmd.Params["embedding_model"] = embedModel
p.nextToken()
case TokenDocParse:
p.nextToken()
@@ -4474,6 +4478,42 @@ func (p *Parser) parseRemoveChunk() (*Command, error) {
return cmd, nil
}
func (p *Parser) parseUserStartIngestion() (*Command, error) {
p.nextToken() // consume Start
if p.curToken.Type != TokenIngestion {
return nil, fmt.Errorf("expect INGESTION")
}
p.nextToken() // consume Ingestion
documentID, err := p.parseQuotedString()
if err != nil {
return nil, err
}
p.nextToken()
if p.curToken.Type != TokenFrom {
return nil, fmt.Errorf("expect FROM")
}
p.nextToken() // consume FROM
datasetID, err := p.parseQuotedString()
if err != nil {
return nil, err
}
cmd := NewCommand("user_start_ingestion_command")
cmd.Params["document_id"] = documentID
cmd.Params["dataset_id"] = datasetID
p.nextToken()
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
// parseShowTask parses SHOW ADMIN SERVER
func (p *Parser) parseUserShowAdmin() (*Command, error) {
p.nextToken() // consume ADMIN
@@ -4494,6 +4534,90 @@ func (p *Parser) parseUserShowAdmin() (*Command, error) {
return cmd, nil
}
func (p *Parser) parseUserStopIngestion() (*Command, error) {
p.nextToken() // consume Stop
if p.curToken.Type != TokenIngestion {
return nil, fmt.Errorf("expect INGESTION")
}
p.nextToken() // consume Ingestion
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expect TASKS")
}
p.nextToken()
taskStr, err := p.parseQuotedString()
if err != nil {
return nil, err
}
tasks := strings.Split(taskStr, " ")
cmd := NewCommand("user_stop_ingestion_command")
cmd.Params["tasks"] = tasks
p.nextToken()
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseUserListIngestionTasks() (*Command, error) {
p.nextToken() // consume Ingestion
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expected TASKS")
}
p.nextToken() // consume TASKS
cmd := NewCommand("user_list_ingestion_tasks")
if p.curToken.Type == TokenFrom {
p.nextToken()
datasetID, err := p.parseQuotedString()
if err != nil {
return nil, err
}
cmd.Params["dataset_id"] = datasetID
}
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseUserRemoveTask() (*Command, error) {
p.nextToken() // consume Ingestion
if p.curToken.Type != TokenTasks {
return nil, fmt.Errorf("expected TASKS")
}
p.nextToken() // consume TASKS
taskStr, err := p.parseQuotedString()
if err != nil {
return nil, err
}
cmd := NewCommand("user_remove_task_command")
tasks := strings.Split(taskStr, " ")
cmd.Params["tasks"] = tasks
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
// parseShowTask parses SHOW API SERVER <server_name>
func (p *Parser) parseUserShowAPI() (*Command, error) {
p.nextToken() // consume API