Update go cli (#13717)

### What problem does this PR solve? Go cli ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
2026-07-14 00:48:26 +08:00 · 2026-03-24 20:08:36 +08:00
parent d84b688b91
commit b308cd3a02
28 changed files with 3239 additions and 955 deletions
--- a/internal/cli/user_command.go
+++ b/internal/cli/user_command.go
@@ -0,0 +1,548 @@
+package cli
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// PingServer pings the server to check if it's alive
+// Returns benchmark result map if iterations > 1, otherwise prints status
+func (c *RAGFlowClient) PingServer(cmd *Command) (ResponseIf, error) {
+	// Get iterations from command params (for benchmark)
+	iterations := 1
+	if val, ok := cmd.Params["iterations"].(int); ok && val > 1 {
+		iterations = val
+	}
+
+	if iterations > 1 {
+		// Benchmark mode: multiple iterations
+		return c.HTTPClient.RequestWithIterations("GET", "/system/ping", false, "web", nil, nil, iterations)
+	}
+
+	// Single mode
+	resp, err := c.HTTPClient.Request("GET", "/system/ping", false, "web", nil, nil)
+	if err != nil {
+		fmt.Printf("Error: %v\n", err)
+		fmt.Println("Server is down")
+		return nil, err
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to ping: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	var result SimpleResponse
+	result.Message = string(resp.Body)
+	result.Code = 0
+	return &result, nil
+}
+
+// Show server version to show RAGFlow server version
+// Returns benchmark result map if iterations > 1, otherwise prints status
+func (c *RAGFlowClient) ShowServerVersion(cmd *Command) (ResponseIf, error) {
+	// Get iterations from command params (for benchmark)
+	iterations := 1
+	if val, ok := cmd.Params["iterations"].(int); ok && val > 1 {
+		iterations = val
+	}
+
+	if iterations > 1 {
+		// Benchmark mode: multiple iterations
+		return c.HTTPClient.RequestWithIterations("GET", "/system/version", false, "web", nil, nil, iterations)
+	}
+
+	// Single mode
+	resp, err := c.HTTPClient.Request("GET", "/system/version", false, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to show version: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to show version: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	var result KeyValueResponse
+	if err = json.Unmarshal(resp.Body, &result); err != nil {
+		return nil, fmt.Errorf("show version failed: invalid JSON (%w)", err)
+	}
+	result.Key = "version"
+	result.Duration = resp.Duration
+
+	return &result, nil
+}
+
+func (c *RAGFlowClient) RegisterUser(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in ADMIN mode")
+	}
+
+	// Check for benchmark iterations
+	var ok bool
+	_, ok = cmd.Params["iterations"].(int)
+	if ok {
+		return nil, fmt.Errorf("failed to register user in benchmark statement")
+	}
+
+	var email string
+	email, ok = cmd.Params["user_name"].(string)
+	if !ok {
+		return nil, fmt.Errorf("no email")
+	}
+
+	var password string
+	password, ok = cmd.Params["password"].(string)
+	if !ok {
+		return nil, fmt.Errorf("no password")
+	}
+
+	var nickname string
+	nickname, ok = cmd.Params["nickname"].(string)
+	if !ok {
+		return nil, fmt.Errorf("no nickname")
+	}
+
+	payload := map[string]interface{}{
+		"email":    email,
+		"password": password,
+		"nickname": nickname,
+	}
+
+	resp, err := c.HTTPClient.Request("POST", "/user/register", false, "admin", nil, payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to register user: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to register user: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	var result RegisterResponse
+	if err = json.Unmarshal(resp.Body, &result); err != nil {
+		return nil, fmt.Errorf("register user failed: invalid JSON (%w)", err)
+	}
+
+	if result.Code != 0 {
+		return nil, fmt.Errorf("%s", result.Message)
+	}
+	result.Duration = resp.Duration
+
+	return &result, nil
+}
+
+// ListUserDatasets lists datasets for current user (user mode)
+// Returns (result_map, error) - result_map is non-nil for benchmark mode
+func (c *RAGFlowClient) ListUserDatasets(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	// Check for benchmark iterations
+	iterations := 1
+	if val, ok := cmd.Params["iterations"].(int); ok && val > 1 {
+		iterations = val
+	}
+
+	if iterations > 1 {
+		// Benchmark mode - return raw result for benchmark stats
+		return c.HTTPClient.RequestWithIterations("GET", "/datasets", true, "web", nil, nil, iterations)
+	}
+
+	// Normal mode
+	resp, err := c.HTTPClient.Request("GET", "/datasets", true, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list datasets: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to list datasets: 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 users failed: invalid JSON (%w)", err)
+	}
+
+	if result.Code != 0 {
+		return nil, fmt.Errorf("%s", result.Message)
+	}
+	result.Duration = resp.Duration
+
+	return &result, nil
+}
+
+// getDatasetID gets dataset ID by name
+func (c *RAGFlowClient) getDatasetID(datasetName string) (string, error) {
+	resp, err := c.HTTPClient.Request("POST", "/kb/list", false, "web", nil, nil)
+	if err != nil {
+		return "", fmt.Errorf("failed to list datasets: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return "", fmt.Errorf("failed to list datasets: HTTP %d", resp.StatusCode)
+	}
+
+	resJSON, err := resp.JSON()
+	if err != nil {
+		return "", fmt.Errorf("invalid JSON response: %w", err)
+	}
+
+	code, ok := resJSON["code"].(float64)
+	if !ok || code != 0 {
+		msg, _ := resJSON["message"].(string)
+		return "", fmt.Errorf("failed to list datasets: %s", msg)
+	}
+
+	data, ok := resJSON["data"].(map[string]interface{})
+	if !ok {
+		return "", fmt.Errorf("invalid response format")
+	}
+
+	kbs, ok := data["kbs"].([]interface{})
+	if !ok {
+		return "", fmt.Errorf("invalid response format: kbs not found")
+	}
+
+	for _, kb := range kbs {
+		if kbMap, ok := kb.(map[string]interface{}); ok {
+			if name, _ := kbMap["name"].(string); name == datasetName {
+				if id, _ := kbMap["id"].(string); id != "" {
+					return id, nil
+				}
+			}
+		}
+	}
+
+	return "", fmt.Errorf("dataset '%s' not found", datasetName)
+}
+
+// formatEmptyArray converts empty arrays to "[]" string
+func formatEmptyArray(v interface{}) string {
+	if v == nil {
+		return "[]"
+	}
+	switch val := v.(type) {
+	case []interface{}:
+		if len(val) == 0 {
+			return "[]"
+		}
+	case []string:
+		if len(val) == 0 {
+			return "[]"
+		}
+	case []int:
+		if len(val) == 0 {
+			return "[]"
+		}
+	}
+	return fmt.Sprintf("%v", v)
+}
+
+// SearchOnDatasets searches for chunks in specified datasets
+// Returns (result_map, error) - result_map is non-nil for benchmark mode
+func (c *RAGFlowClient) SearchOnDatasets(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	question, ok := cmd.Params["question"].(string)
+	if !ok {
+		return nil, fmt.Errorf("question not provided")
+	}
+
+	datasets, ok := cmd.Params["datasets"].(string)
+	if !ok {
+		return nil, fmt.Errorf("datasets not provided")
+	}
+
+	// Parse dataset names (comma-separated) and convert to IDs
+	datasetNames := strings.Split(datasets, ",")
+	datasetIDs := make([]string, 0, len(datasetNames))
+	for _, name := range datasetNames {
+		name = strings.TrimSpace(name)
+		id, err := c.getDatasetID(name)
+		if err != nil {
+			return nil, err
+		}
+		datasetIDs = append(datasetIDs, id)
+	}
+
+	// Check for benchmark iterations
+	iterations := 1
+	if val, ok := cmd.Params["iterations"].(int); ok && val > 1 {
+		iterations = val
+	}
+
+	payload := map[string]interface{}{
+		"kb_id":                    datasetIDs,
+		"question":                 question,
+		"similarity_threshold":     0.2,
+		"vector_similarity_weight": 0.3,
+	}
+
+	if iterations > 1 {
+		// Benchmark mode - return raw result for benchmark stats
+		return c.HTTPClient.RequestWithIterations("POST", "/chunk/retrieval_test", false, "web", nil, payload, iterations)
+	}
+
+	// Normal mode
+	resp, err := c.HTTPClient.Request("POST", "/chunk/retrieval_test", false, "web", nil, payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to search on datasets: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to search on datasets: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	resJSON, err := resp.JSON()
+	if err != nil {
+		return nil, fmt.Errorf("invalid JSON response: %w", err)
+	}
+
+	code, ok := resJSON["code"].(float64)
+	if !ok || code != 0 {
+		msg, _ := resJSON["message"].(string)
+		return nil, fmt.Errorf("failed to search on datasets: %s", msg)
+	}
+
+	data, ok := resJSON["data"].(map[string]interface{})
+	if !ok {
+		return nil, fmt.Errorf("invalid response format")
+	}
+
+	chunks, ok := data["chunks"].([]interface{})
+	if !ok {
+		return nil, fmt.Errorf("invalid response format: chunks not found")
+	}
+
+	// Convert to slice of maps for printing
+	tableData := make([]map[string]interface{}, 0, len(chunks))
+	for _, chunk := range chunks {
+		if chunkMap, ok := chunk.(map[string]interface{}); ok {
+			row := map[string]interface{}{
+				"id":                chunkMap["chunk_id"],
+				"content":           chunkMap["content_with_weight"],
+				"document_id":       chunkMap["doc_id"],
+				"dataset_id":        chunkMap["kb_id"],
+				"docnm_kwd":         chunkMap["docnm_kwd"],
+				"image_id":          chunkMap["image_id"],
+				"similarity":        chunkMap["similarity"],
+				"term_similarity":   chunkMap["term_similarity"],
+				"vector_similarity": chunkMap["vector_similarity"],
+			}
+			// Add optional fields that may be empty arrays
+			if v, ok := chunkMap["doc_type_kwd"]; ok {
+				row["doc_type_kwd"] = formatEmptyArray(v)
+			}
+			if v, ok := chunkMap["important_kwd"]; ok {
+				row["important_kwd"] = formatEmptyArray(v)
+			}
+			if v, ok := chunkMap["mom_id"]; ok {
+				row["mom_id"] = formatEmptyArray(v)
+			}
+			if v, ok := chunkMap["positions"]; ok {
+				row["positions"] = formatEmptyArray(v)
+			}
+			if v, ok := chunkMap["content_ltks"]; ok {
+				row["content_ltks"] = v
+			}
+			tableData = append(tableData, row)
+		}
+	}
+
+	PrintTableSimple(tableData)
+	return nil, nil
+}
+
+// CreateToken creates a new API token
+func (c *RAGFlowClient) CreateToken(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	resp, err := c.HTTPClient.Request("POST", "/tokens", true, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create token: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to create token: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	var createResult CommonDataResponse
+	if err = json.Unmarshal(resp.Body, &createResult); err != nil {
+		return nil, fmt.Errorf("create token failed: invalid JSON (%w)", err)
+	}
+
+	if createResult.Code != 0 {
+		return nil, fmt.Errorf("%s", createResult.Message)
+	}
+
+	var result SimpleResponse
+	result.Code = 0
+	result.Message = "Token created successfully"
+	result.Duration = resp.Duration
+	return &result, nil
+}
+
+// ListTokens lists all API tokens for the current user
+func (c *RAGFlowClient) ListTokens(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	resp, err := c.HTTPClient.Request("GET", "/tokens", true, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list tokens: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to list tokens: 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 tokens failed: invalid JSON (%w)", err)
+	}
+
+	if result.Code != 0 {
+		return nil, fmt.Errorf("%s", result.Message)
+	}
+	result.Duration = resp.Duration
+	return &result, nil
+}
+
+// DropToken deletes an API token
+func (c *RAGFlowClient) DropToken(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	token, ok := cmd.Params["token"].(string)
+	if !ok {
+		return nil, fmt.Errorf("token not provided")
+	}
+
+	resp, err := c.HTTPClient.Request("DELETE", fmt.Sprintf("/tokens/%s", token), true, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to drop token: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to drop token: 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("drop token failed: invalid JSON (%w)", err)
+	}
+
+	if result.Code != 0 {
+		return nil, fmt.Errorf("%s", result.Message)
+	}
+	result.Duration = resp.Duration
+	return &result, nil
+}
+
+// SetToken sets the API token after validating it
+func (c *RAGFlowClient) SetToken(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	token, ok := cmd.Params["token"].(string)
+	if !ok {
+		return nil, fmt.Errorf("token not provided")
+	}
+
+	// Save current token to restore if validation fails
+	savedToken := c.HTTPClient.APIToken
+	savedUseAPIToken := c.HTTPClient.useAPIToken
+
+	// Set the new token temporarily for validation
+	c.HTTPClient.APIToken = token
+	c.HTTPClient.useAPIToken = true
+
+	// Validate token by calling list tokens API
+	resp, err := c.HTTPClient.Request("GET", "/tokens", true, "api", nil, nil)
+	if err != nil {
+		// Restore original token on error
+		c.HTTPClient.APIToken = savedToken
+		c.HTTPClient.useAPIToken = savedUseAPIToken
+		return nil, fmt.Errorf("failed to validate token: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		// Restore original token on error
+		c.HTTPClient.APIToken = savedToken
+		c.HTTPClient.useAPIToken = savedUseAPIToken
+		return nil, fmt.Errorf("token validation failed: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
+	}
+
+	var result CommonResponse
+	if err = json.Unmarshal(resp.Body, &result); err != nil {
+		// Restore original token on error
+		c.HTTPClient.APIToken = savedToken
+		c.HTTPClient.useAPIToken = savedUseAPIToken
+		return nil, fmt.Errorf("token validation failed: invalid JSON (%w)", err)
+	}
+
+	if result.Code != 0 {
+		// Restore original token on error
+		c.HTTPClient.APIToken = savedToken
+		c.HTTPClient.useAPIToken = savedUseAPIToken
+		return nil, fmt.Errorf("token validation failed: %s", result.Message)
+	}
+
+	// Token is valid, keep it set
+	var successResult SimpleResponse
+	successResult.Code = 0
+	successResult.Message = "API token set successfully"
+	successResult.Duration = resp.Duration
+	return &successResult, nil
+}
+
+// ShowToken displays the current API token
+func (c *RAGFlowClient) ShowToken(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	if c.HTTPClient.APIToken == "" {
+		return nil, fmt.Errorf("no API token is currently set")
+	}
+
+	//fmt.Printf("Token: %s\n", c.HTTPClient.APIToken)
+
+	var result CommonResponse
+	result.Code = 0
+	result.Message = ""
+	result.Data = []map[string]interface{}{
+		{
+			"token": c.HTTPClient.APIToken,
+		},
+	}
+	result.Duration = 0
+	return &result, nil
+}
+
+// UnsetToken removes the current API token
+func (c *RAGFlowClient) UnsetToken(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	if c.HTTPClient.APIToken == "" {
+		return nil, fmt.Errorf("no API token is currently set")
+	}
+
+	c.HTTPClient.APIToken = ""
+	c.HTTPClient.useAPIToken = false
+
+	var result SimpleResponse
+	result.Code = 0
+	result.Message = "API token unset successfully"
+	result.Duration = 0
+	return &result, nil
+}