mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-03-18 05:10:00 +08:00
### What problem does this PR solve? Implement Search() in Infinity in GO. The function can handle the following request. "search '曹操' on datasets 'infinity'" "search '常胜将军' on datasets 'infinity'" "search '卓越儒雅' on datasets 'infinity'" "search '辅佐刘禅北伐中原' on datasets 'infinity'" The output is exactly the same as request to python Search() ### Type of change - [ ] New Feature (non-breaking change which adds functionality)
991 lines
27 KiB
Go
991 lines
27 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package cli
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"strings"
|
|
"syscall"
|
|
"unsafe"
|
|
)
|
|
|
|
// PasswordPromptFunc is a function type for password input
|
|
type PasswordPromptFunc func(prompt string) (string, error)
|
|
|
|
// RAGFlowClient handles API interactions with the RAGFlow server
|
|
type RAGFlowClient struct {
|
|
HTTPClient *HTTPClient
|
|
ServerType string // "admin" or "user"
|
|
PasswordPrompt PasswordPromptFunc // Function for password input
|
|
}
|
|
|
|
// NewRAGFlowClient creates a new RAGFlow client
|
|
func NewRAGFlowClient(serverType string) *RAGFlowClient {
|
|
httpClient := NewHTTPClient()
|
|
// Set port from configuration file based on server type
|
|
if serverType == "admin" {
|
|
httpClient.Port = 9381
|
|
} else {
|
|
httpClient.Port = 9380
|
|
}
|
|
|
|
return &RAGFlowClient{
|
|
HTTPClient: httpClient,
|
|
ServerType: serverType,
|
|
}
|
|
}
|
|
|
|
// LoginUserInteractive performs interactive login with username and password
|
|
func (c *RAGFlowClient) LoginUserInteractive(username, password string) error {
|
|
// First, ping the server to check if it's available
|
|
// For admin mode, use /admin/ping with useAPIBase=true
|
|
// For user mode, use /system/ping with useAPIBase=false
|
|
var pingPath string
|
|
var useAPIBase bool
|
|
if c.ServerType == "admin" {
|
|
pingPath = "/admin/ping"
|
|
useAPIBase = true
|
|
} else {
|
|
pingPath = "/system/ping"
|
|
useAPIBase = false
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("GET", pingPath, useAPIBase, "web", nil, nil)
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
fmt.Println("Can't access server for login (connection failed)")
|
|
return err
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
|
|
// Check response - admin returns JSON with message "PONG", user returns plain "pong"
|
|
resJSON, err := resp.JSON()
|
|
if err == nil {
|
|
// Admin mode returns {"code":0,"message":"PONG"}
|
|
if msg, ok := resJSON["message"].(string); !ok || msg != "PONG" {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
} else {
|
|
// User mode returns plain "pong"
|
|
if string(resp.Body) != "pong" {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
}
|
|
|
|
// If password is not provided, prompt for it
|
|
if password == "" {
|
|
fmt.Printf("password for %s: ", username)
|
|
var err error
|
|
password, err = readPassword()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to read password: %w", err)
|
|
}
|
|
password = strings.TrimSpace(password)
|
|
}
|
|
|
|
// Login
|
|
token, err := c.loginUser(username, password)
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
fmt.Println("Can't access server for login (connection failed)")
|
|
return err
|
|
}
|
|
|
|
c.HTTPClient.LoginToken = token
|
|
fmt.Printf("Login user %s successfully\n", username)
|
|
return nil
|
|
}
|
|
|
|
// LoginUser performs user login
|
|
func (c *RAGFlowClient) LoginUser(cmd *Command) error {
|
|
// First, ping the server to check if it's available
|
|
// For admin mode, use /admin/ping with useAPIBase=true
|
|
// For user mode, use /system/ping with useAPIBase=false
|
|
var pingPath string
|
|
var useAPIBase bool
|
|
if c.ServerType == "admin" {
|
|
pingPath = "/admin/ping"
|
|
useAPIBase = true
|
|
} else {
|
|
pingPath = "/system/ping"
|
|
useAPIBase = false
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("GET", pingPath, useAPIBase, "web", nil, nil)
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
fmt.Println("Can't access server for login (connection failed)")
|
|
return err
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
|
|
// Check response - admin returns JSON with message "PONG", user returns plain "pong"
|
|
resJSON, err := resp.JSON()
|
|
if err == nil {
|
|
// Admin mode returns {"code":0,"message":"PONG"}
|
|
if msg, ok := resJSON["message"].(string); !ok || msg != "PONG" {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
} else {
|
|
// User mode returns plain "pong"
|
|
if string(resp.Body) != "pong" {
|
|
fmt.Println("Server is down")
|
|
return fmt.Errorf("server is down")
|
|
}
|
|
}
|
|
|
|
email, ok := cmd.Params["email"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("email not provided")
|
|
}
|
|
|
|
// Get password from user input (hidden)
|
|
var password string
|
|
if c.PasswordPrompt != nil {
|
|
pwd, err := c.PasswordPrompt(fmt.Sprintf("password for %s: ", email))
|
|
if err != nil {
|
|
return fmt.Errorf("failed to read password: %w", err)
|
|
}
|
|
password = pwd
|
|
} else {
|
|
fmt.Printf("password for %s: ", email)
|
|
pwd, err := readPassword()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to read password: %w", err)
|
|
}
|
|
password = pwd
|
|
}
|
|
password = strings.TrimSpace(password)
|
|
|
|
// Login
|
|
token, err := c.loginUser(email, password)
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
fmt.Println("Can't access server for login (connection failed)")
|
|
return err
|
|
}
|
|
|
|
c.HTTPClient.LoginToken = token
|
|
fmt.Printf("Login user %s successfully\n", email)
|
|
return nil
|
|
}
|
|
|
|
// loginUser performs the actual login request
|
|
func (c *RAGFlowClient) loginUser(email, password string) (string, error) {
|
|
// Encrypt password using scrypt (same as Python implementation)
|
|
encryptedPassword, err := EncryptPassword(password)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to encrypt password: %w", err)
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"email": email,
|
|
"password": encryptedPassword,
|
|
}
|
|
|
|
var path string
|
|
if c.ServerType == "admin" {
|
|
path = "/admin/login"
|
|
} else {
|
|
path = "/user/login"
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("POST", path, c.ServerType == "admin", "", nil, payload)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
resJSON, err := resp.JSON()
|
|
if err != nil {
|
|
return "", fmt.Errorf("login failed: invalid JSON response (%w)", err)
|
|
}
|
|
|
|
code, ok := resJSON["code"].(float64)
|
|
if !ok || code != 0 {
|
|
msg, _ := resJSON["message"].(string)
|
|
return "", fmt.Errorf("login failed: %s", msg)
|
|
}
|
|
|
|
token := resp.Headers.Get("Authorization")
|
|
if token == "" {
|
|
return "", fmt.Errorf("login failed: missing Authorization header")
|
|
}
|
|
|
|
return token, nil
|
|
}
|
|
|
|
// 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) (map[string]interface{}, 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
|
|
result, err := c.HTTPClient.RequestWithIterations("GET", "/system/ping", false, "web", nil, nil, iterations)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// Single ping 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 && string(resp.Body) == "pong" {
|
|
fmt.Println("Server is alive")
|
|
} else {
|
|
fmt.Printf("Error: %d\n", resp.StatusCode)
|
|
}
|
|
return nil, 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) (map[string]interface{}, 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("POST", "/kb/list", false, "web", nil, nil, iterations)
|
|
}
|
|
|
|
// Normal mode
|
|
resp, err := c.HTTPClient.Request("POST", "/kb/list", false, "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))
|
|
}
|
|
|
|
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 list datasets: %s", msg)
|
|
}
|
|
|
|
data, ok := resJSON["data"].(map[string]interface{})
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid response format")
|
|
}
|
|
|
|
kbs, ok := data["kbs"].([]interface{})
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid response format: kbs not found")
|
|
}
|
|
|
|
// Convert to slice of maps
|
|
tableData := make([]map[string]interface{}, 0, len(kbs))
|
|
for _, kb := range kbs {
|
|
if kbMap, ok := kb.(map[string]interface{}); ok {
|
|
// Remove avatar field
|
|
delete(kbMap, "avatar")
|
|
tableData = append(tableData, kbMap)
|
|
}
|
|
}
|
|
|
|
PrintTableSimple(tableData)
|
|
return nil, nil
|
|
}
|
|
|
|
// ListDatasets lists datasets for a specific user (admin mode)
|
|
// Returns (result_map, error) - result_map is non-nil for benchmark mode
|
|
func (c *RAGFlowClient) ListDatasets(cmd *Command) (map[string]interface{}, error) {
|
|
if c.ServerType != "admin" {
|
|
return nil, fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return nil, fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
// 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", fmt.Sprintf("/admin/users/%s/datasets", userName), true, "admin", nil, nil, iterations)
|
|
}
|
|
|
|
fmt.Printf("Listing all datasets of user: %s\n", userName)
|
|
|
|
resp, err := c.HTTPClient.Request("GET", fmt.Sprintf("/admin/users/%s/datasets", userName), true, "admin", 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))
|
|
}
|
|
|
|
resJSON, err := resp.JSON()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("invalid JSON response: %w", err)
|
|
}
|
|
|
|
data, ok := resJSON["data"].([]interface{})
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid response format")
|
|
}
|
|
|
|
// Convert to slice of maps and remove avatar
|
|
tableData := make([]map[string]interface{}, 0, len(data))
|
|
for _, item := range data {
|
|
if itemMap, ok := item.(map[string]interface{}); ok {
|
|
delete(itemMap, "avatar")
|
|
tableData = append(tableData, itemMap)
|
|
}
|
|
}
|
|
|
|
PrintTableSimple(tableData)
|
|
return nil, nil
|
|
}
|
|
|
|
// readPassword reads password from terminal without echoing
|
|
func readPassword() (string, error) {
|
|
// Check if stdin is a terminal by trying to get terminal size
|
|
if isTerminal() {
|
|
// Use stty to disable echo
|
|
cmd := exec.Command("stty", "-echo")
|
|
cmd.Stdin = os.Stdin
|
|
if err := cmd.Run(); err != nil {
|
|
// Fallback: read normally
|
|
return readPasswordFallback()
|
|
}
|
|
defer func() {
|
|
// Re-enable echo
|
|
cmd := exec.Command("stty", "echo")
|
|
cmd.Stdin = os.Stdin
|
|
cmd.Run()
|
|
}()
|
|
|
|
reader := bufio.NewReader(os.Stdin)
|
|
password, err := reader.ReadString('\n')
|
|
fmt.Println() // New line after password input
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return strings.TrimSpace(password), nil
|
|
}
|
|
|
|
// Fallback for non-terminal input (e.g., piped input)
|
|
return readPasswordFallback()
|
|
}
|
|
|
|
// isTerminal checks if stdin is a terminal
|
|
func isTerminal() bool {
|
|
var termios syscall.Termios
|
|
_, _, err := syscall.Syscall6(syscall.SYS_IOCTL, os.Stdin.Fd(), syscall.TCGETS, uintptr(unsafe.Pointer(&termios)), 0, 0, 0)
|
|
return err == 0
|
|
}
|
|
|
|
// readPasswordFallback reads password as plain text (fallback mode)
|
|
func readPasswordFallback() (string, error) {
|
|
reader := bufio.NewReader(os.Stdin)
|
|
password, err := reader.ReadString('\n')
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return strings.TrimSpace(password), 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) (map[string]interface{}, 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
|
|
}
|
|
|
|
// ExecuteCommand executes a parsed command
|
|
// Returns benchmark result map for commands that support it (e.g., ping_server with iterations > 1)
|
|
func (c *RAGFlowClient) ExecuteCommand(cmd *Command) (map[string]interface{}, error) {
|
|
switch cmd.Type {
|
|
case "login_user":
|
|
return nil, c.LoginUser(cmd)
|
|
case "ping_server":
|
|
return c.PingServer(cmd)
|
|
case "benchmark":
|
|
return nil, c.RunBenchmark(cmd)
|
|
case "list_user_datasets":
|
|
return c.ListUserDatasets(cmd)
|
|
case "list_datasets":
|
|
return c.ListDatasets(cmd)
|
|
case "search_on_datasets":
|
|
return c.SearchOnDatasets(cmd)
|
|
case "list_users":
|
|
return c.ListUsers(cmd)
|
|
case "grant_admin":
|
|
return nil, c.GrantAdmin(cmd)
|
|
case "revoke_admin":
|
|
return nil, c.RevokeAdmin(cmd)
|
|
case "show_current_user":
|
|
return c.ShowCurrentUser(cmd)
|
|
case "create_user":
|
|
return nil, c.CreateUser(cmd)
|
|
case "activate_user":
|
|
return nil, c.ActivateUser(cmd)
|
|
case "alter_user":
|
|
return nil, c.AlterUserPassword(cmd)
|
|
case "drop_user":
|
|
return nil, c.DropUser(cmd)
|
|
// TODO: Implement other commands
|
|
default:
|
|
return nil, fmt.Errorf("command '%s' would be executed with API", cmd.Type)
|
|
}
|
|
}
|
|
|
|
// ListUsers lists all users (admin mode only)
|
|
// Returns (result_map, error) - result_map is non-nil for benchmark mode
|
|
func (c *RAGFlowClient) ListUsers(cmd *Command) (map[string]interface{}, error) {
|
|
if c.ServerType != "admin" {
|
|
return nil, fmt.Errorf("this command is only allowed in ADMIN 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", "/admin/users", true, "admin", nil, nil, iterations)
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("GET", "/admin/users", true, "admin", nil, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to list users: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return nil, fmt.Errorf("failed to list users: 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 list users: %s", msg)
|
|
}
|
|
|
|
data, ok := resJSON["data"].([]interface{})
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid response format")
|
|
}
|
|
|
|
// Convert to slice of maps and remove sensitive fields
|
|
tableData := make([]map[string]interface{}, 0, len(data))
|
|
for _, item := range data {
|
|
if itemMap, ok := item.(map[string]interface{}); ok {
|
|
// Remove sensitive fields
|
|
delete(itemMap, "password")
|
|
delete(itemMap, "access_token")
|
|
tableData = append(tableData, itemMap)
|
|
}
|
|
}
|
|
|
|
PrintTableSimple(tableData)
|
|
return nil, nil
|
|
}
|
|
|
|
// GrantAdmin grants admin privileges to a user (admin mode only)
|
|
func (c *RAGFlowClient) GrantAdmin(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("PUT", fmt.Sprintf("/admin/users/%s/admin", userName), true, "admin", nil, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to grant admin: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to grant admin: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 grant admin: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("Admin role granted to user: %s\n", userName)
|
|
return nil
|
|
}
|
|
|
|
// RevokeAdmin revokes admin privileges from a user (admin mode only)
|
|
func (c *RAGFlowClient) RevokeAdmin(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("DELETE", fmt.Sprintf("/admin/users/%s/admin", userName), true, "admin", nil, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to revoke admin: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to revoke admin: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 revoke admin: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("Admin role revoked from user: %s\n", userName)
|
|
return nil
|
|
}
|
|
|
|
// ShowCurrentUser shows the current logged-in user information
|
|
// TODO: Implement showing current user information when API is available
|
|
func (c *RAGFlowClient) ShowCurrentUser(cmd *Command) (map[string]interface{}, error) {
|
|
// TODO: Call the appropriate API to get current user information
|
|
// Currently there is no /admin/user/info or /user/info API available
|
|
// The /admin/auth API only verifies authorization, does not return user info
|
|
return nil, fmt.Errorf("command 'SHOW CURRENT USER' is not yet implemented")
|
|
}
|
|
|
|
// CreateUser creates a new user (admin mode only)
|
|
func (c *RAGFlowClient) CreateUser(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
password, ok := cmd.Params["password"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("password not provided")
|
|
}
|
|
|
|
// Encrypt password using RSA
|
|
encryptedPassword, err := EncryptPassword(password)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to encrypt password: %w", err)
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"username": userName,
|
|
"password": encryptedPassword,
|
|
"role": "user",
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("POST", "/admin/users", true, "admin", nil, payload)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create user: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to create user: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 create user: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("User created successfully: %s\n", userName)
|
|
return nil
|
|
}
|
|
|
|
// ActivateUser activates or deactivates a user (admin mode only)
|
|
func (c *RAGFlowClient) ActivateUser(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
activateStatus, ok := cmd.Params["activate_status"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("activate_status not provided")
|
|
}
|
|
|
|
// Validate activate_status
|
|
if activateStatus != "on" && activateStatus != "off" {
|
|
return fmt.Errorf("activate_status must be 'on' or 'off'")
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"activate_status": activateStatus,
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("PUT", fmt.Sprintf("/admin/users/%s/activate", userName), true, "admin", nil, payload)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to update user activate status: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to update user activate status: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 update user activate status: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("User '%s' activate status set to '%s'\n", userName, activateStatus)
|
|
return nil
|
|
}
|
|
|
|
// AlterUserPassword changes a user's password (admin mode only)
|
|
func (c *RAGFlowClient) AlterUserPassword(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
password, ok := cmd.Params["password"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("password not provided")
|
|
}
|
|
|
|
// Encrypt password using RSA
|
|
encryptedPassword, err := EncryptPassword(password)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to encrypt password: %w", err)
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"new_password": encryptedPassword,
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("PUT", fmt.Sprintf("/admin/users/%s/password", userName), true, "admin", nil, payload)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to change user password: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to change user password: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 change user password: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("Password changed for user: %s\n", userName)
|
|
return nil
|
|
}
|
|
|
|
// DropUser deletes a user (admin mode only)
|
|
func (c *RAGFlowClient) DropUser(cmd *Command) error {
|
|
if c.ServerType != "admin" {
|
|
return fmt.Errorf("this command is only allowed in ADMIN mode")
|
|
}
|
|
|
|
userName, ok := cmd.Params["user_name"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("user_name not provided")
|
|
}
|
|
|
|
resp, err := c.HTTPClient.Request("DELETE", fmt.Sprintf("/admin/users/%s", userName), true, "admin", nil, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to delete user: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("failed to delete user: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
|
}
|
|
|
|
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 delete user: %s", msg)
|
|
}
|
|
|
|
fmt.Printf("User deleted: %s\n", userName)
|
|
return nil
|
|
}
|