youngkingdom/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-05-03 00:37:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update go cli (#13717)

Browse Source 
### 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>
This commit is contained in:
Jin Hai
2026-03-24 20:08:36 +08:00
committed by GitHub
parent d84b688b91
commit b308cd3a02
28 changed files with 3239 additions and 955 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1101
internal/cli/admin_command.go Normal file
View File

File diff suppressed because it is too large Load Diff

156
internal/cli/benchmark.go
View File

@ -34,7 +34,7 @@ type BenchmarkResult struct {
}
// RunBenchmark runs a benchmark with the given concurrency and iterations
func (c *RAGFlowClient) RunBenchmark(cmd *Command) error {
func (c *RAGFlowClient) RunBenchmark(cmd *Command) (ResponseIf, error) {
concurrency, ok := cmd.Params["concurrency"].(int)
if !ok {
concurrency = 1
@ -47,29 +47,26 @@ func (c *RAGFlowClient) RunBenchmark(cmd *Command) error {
nestedCmd, ok := cmd.Params["command"].(*Command)
if !ok {
return fmt.Errorf("benchmark command not found")
return nil, fmt.Errorf("benchmark command not found")
}
if concurrency < 1 {
return fmt.Errorf("concurrency must be greater than 0")
return nil, fmt.Errorf("concurrency must be greater than 0")
}
// Add iterations to the nested command
nestedCmd.Params["iterations"] = iterations
if concurrency == 1 {
return c.runBenchmarkSingle(concurrency, iterations, nestedCmd)
return c.runBenchmarkSingle(iterations, nestedCmd)
}
return c.runBenchmarkConcurrent(concurrency, iterations, nestedCmd)
}
// runBenchmarkSingle runs benchmark with single concurrency (sequential execution)
func (c *RAGFlowClient) runBenchmarkSingle(concurrency, iterations int, nestedCmd *Command) error {
func (c *RAGFlowClient) runBenchmarkSingle(iterations int, nestedCmd *Command) (*BenchmarkResponse, error) {
commandType := nestedCmd.Type
startTime := time.Now()
responseList := make([]*Response, 0, iterations)
// For search_on_datasets, convert dataset names to IDs first
if commandType == "search_on_datasets" && iterations > 1 {
datasets, _ := nestedCmd.Params["datasets"].(string)
@ -79,7 +76,7 @@ func (c *RAGFlowClient) runBenchmarkSingle(concurrency, iterations int, nestedCm
name = strings.TrimSpace(name)
id, err := c.getDatasetID(name)
if err != nil {
return err
return nil, err
}
datasetIDs = append(datasetIDs, id)
}
@ -87,86 +84,67 @@ func (c *RAGFlowClient) runBenchmarkSingle(concurrency, iterations int, nestedCm
}
// Check if command supports native benchmark (iterations > 1)
supportsNative := false
if iterations > 1 {
result, err := c.ExecuteCommand(nestedCmd)
if err == nil && result != nil {
// Command supports benchmark natively
supportsNative = true
duration, _ := result["duration"].(float64)
respList, _ := result["response_list"].([]*Response)
responseList = respList
// convert result to BenchmarkResponse
benchmarkResponse := result.(*BenchmarkResponse)
benchmarkResponse.Concurrency = 1
return benchmarkResponse, err
}
// Calculate and print results
successCount := 0
for _, resp := range responseList {
if isSuccess(resp, commandType) {
successCount++
}
}
result, err := c.ExecuteCommand(nestedCmd)
if err != nil {
fmt.Printf("fail to execute: %s", commandType)
return nil, err
}
qps := float64(0)
if duration > 0 {
qps = float64(iterations) / duration
}
fmt.Printf("command: %s, Concurrency: %d, iterations: %d\n", commandType, concurrency, iterations)
fmt.Printf("total duration: %.4fs, QPS: %.2f, COMMAND_COUNT: %d, SUCCESS: %d, FAILURE: %d\n",
duration, qps, iterations, successCount, iterations-successCount)
return nil
var benchmarkResponse BenchmarkResponse
switch result.Type() {
case "common":
commonResponse := result.(*CommonResponse)
benchmarkResponse.Code = commonResponse.Code
benchmarkResponse.Duration = commonResponse.Duration
if commonResponse.Code == 0 {
benchmarkResponse.SuccessCount = 1
} else {
benchmarkResponse.FailureCount = 1
}
}
// Manual execution: run iterations times
if !supportsNative {
// Remove iterations param to avoid native benchmark
delete(nestedCmd.Params, "iterations")
for i := 0; i < iterations; i++ {
singleResult, err := c.ExecuteCommand(nestedCmd)
if err != nil {
// Command failed, add a failed response
responseList = append(responseList, &Response{StatusCode: 0})
continue
}
// For commands that return a single response (like ping with iterations=1)
if singleResult != nil {
if respList, ok := singleResult["response_list"].([]*Response); ok {
responseList = append(responseList, respList...)
}
} else {
// Command executed successfully but returned no data
// Mark as success for now
responseList = append(responseList, &Response{StatusCode: 200, Body: []byte("pong")})
}
case "simple":
simpleResponse := result.(*SimpleResponse)
benchmarkResponse.Code = simpleResponse.Code
benchmarkResponse.Duration = simpleResponse.Duration
if simpleResponse.Code == 0 {
benchmarkResponse.SuccessCount = 1
} else {
benchmarkResponse.FailureCount = 1
}
}
duration := time.Since(startTime).Seconds()
successCount := 0
for _, resp := range responseList {
if isSuccess(resp, commandType) {
successCount++
case "show":
dataResponse := result.(*CommonDataResponse)
benchmarkResponse.Code = dataResponse.Code
benchmarkResponse.Duration = dataResponse.Duration
if dataResponse.Code == 0 {
benchmarkResponse.SuccessCount = 1
} else {
benchmarkResponse.FailureCount = 1
}
case "data":
kvResponse := result.(*KeyValueResponse)
benchmarkResponse.Code = kvResponse.Code
benchmarkResponse.Duration = kvResponse.Duration
if kvResponse.Code == 0 {
benchmarkResponse.SuccessCount = 1
} else {
benchmarkResponse.FailureCount = 1
}
default:
return nil, fmt.Errorf("unsupported command type: %s", result.Type())
}
qps := float64(0)
if duration > 0 {
qps = float64(iterations) / duration
}
// Print results
fmt.Printf("command: %s, Concurrency: %d, iterations: %d\n", commandType, concurrency, iterations)
fmt.Printf("total duration: %.4fs, QPS: %.2f, COMMAND_COUNT: %d, SUCCESS: %d, FAILURE: %d\n",
duration, qps, iterations, successCount, iterations-successCount)
return nil
benchmarkResponse.Concurrency = 1
return &benchmarkResponse, nil
}
// runBenchmarkConcurrent runs benchmark with multiple concurrent workers
func (c *RAGFlowClient) runBenchmarkConcurrent(concurrency, iterations int, nestedCmd *Command) error {
func (c *RAGFlowClient) runBenchmarkConcurrent(concurrency, iterations int, nestedCmd *Command) (*BenchmarkResponse, error) {
results := make([]map[string]interface{}, concurrency)
var wg sync.WaitGroup
@ -179,7 +157,7 @@ func (c *RAGFlowClient) runBenchmarkConcurrent(concurrency, iterations int, nest
name = strings.TrimSpace(name)
id, err := c.getDatasetID(name)
if err != nil {
return err
return nil, err
}
datasetIDs = append(datasetIDs, id)
}
@ -228,17 +206,15 @@ func (c *RAGFlowClient) runBenchmarkConcurrent(concurrency, iterations int, nest
}
totalCommands := iterations * concurrency
qps := float64(0)
if totalDuration > 0 {
qps = float64(totalCommands) / totalDuration
}
// Print results
fmt.Printf("command: %s, Concurrency: %d, iterations: %d\n", commandType, concurrency, iterations)
fmt.Printf("total duration: %.4fs, QPS: %.2f, COMMAND_COUNT: %d, SUCCESS: %d, FAILURE: %d\n",
totalDuration, qps, totalCommands, successCount, totalCommands-successCount)
var benchmarkResponse BenchmarkResponse
benchmarkResponse.Duration = totalDuration
benchmarkResponse.Code = 0
benchmarkResponse.SuccessCount = successCount
benchmarkResponse.FailureCount = totalCommands - successCount
benchmarkResponse.Concurrency = concurrency
return nil
return &benchmarkResponse, nil
}
// executeBenchmarkSilent executes a command for benchmark without printing output
@ -250,7 +226,7 @@ func (c *RAGFlowClient) executeBenchmarkSilent(cmd *Command, iterations int) []*
var err error
switch cmd.Type {
case "ping_server":
case "ping":
resp, err = c.HTTPClient.Request("GET", "/system/ping", false, "web", nil, nil)
case "list_user_datasets":
resp, err = c.HTTPClient.Request("POST", "/kb/list", false, "web", nil, nil)
@ -290,7 +266,7 @@ func isSuccess(resp *Response, commandType string) bool {
}
switch commandType {
case "ping_server":
case "ping":
return resp.StatusCode == 200 && string(resp.Body) == "pong"
case "list_user_datasets", "list_datasets", "search_on_datasets":
// Check status code and JSON response code for dataset commands

401
internal/cli/cli.go
View File

@ -17,15 +17,285 @@
package cli
import (
"errors"
"flag"
"fmt"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"github.com/peterh/liner"
"golang.org/x/term"
"gopkg.in/yaml.v3"
)
// ConfigFile represents the rf.yml configuration file structure
type ConfigFile struct {
Host string `yaml:"host"`
User string `yaml:"user"`
Password string `yaml:"password"`
APIToken string `yaml:"api_token"`
}
// ConnectionArgs holds the parsed command line arguments
type ConnectionArgs struct {
Host string
Port int
Password string
Key string
Type string
Username string
Command string
ShowHelp bool
}
// LoadDefaultConfigFile reads the rf.yml file from current directory if it exists
func LoadDefaultConfigFile() (*ConfigFile, error) {
// Try to read rf.yml from current directory
data, err := os.ReadFile("rf.yml")
if err != nil {
// File doesn't exist, return nil without error
if os.IsNotExist(err) {
return nil, nil
}
return nil, err
}
var config ConfigFile
if err = yaml.Unmarshal(data, &config); err != nil {
return nil, fmt.Errorf("failed to parse rf.yml: %v", err)
}
return &config, nil
}
// LoadConfigFileFromPath reads a config file from the specified path
func LoadConfigFileFromPath(path string) (*ConfigFile, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read config file %s: %v", path, err)
}
var config ConfigFile
if err = yaml.Unmarshal(data, &config); err != nil {
return nil, fmt.Errorf("failed to parse config file %s: %v", path, err)
}
return &config, nil
}
// parseHostPort parses a host:port string and returns host and port
func parseHostPort(hostPort string) (string, int, error) {
if hostPort == "" {
return "", -1, nil
}
// Split host and port
parts := strings.Split(hostPort, ":")
if len(parts) != 2 {
return "", -1, fmt.Errorf("invalid host format, expected host:port, got: %s", hostPort)
}
host := parts[0]
port, err := strconv.Atoi(parts[1])
if err != nil {
return "", -1, fmt.Errorf("invalid port number: %s", parts[1])
}
return host, port, nil
}
// ParseConnectionArgs parses command line arguments similar to Python's parse_connection_args
func ParseConnectionArgs(args []string) (*ConnectionArgs, error) {
// First, scan args to check for help and config file
var configFilePath string
var hasOtherFlags bool
for i := 0; i < len(args); i++ {
arg := args[i]
if arg == "--help" {
return &ConnectionArgs{ShowHelp: true}, nil
} else if arg == "-f" && i+1 < len(args) {
configFilePath = args[i+1]
i++
} else if strings.HasPrefix(arg, "-") && arg != "-f" {
// Check if it's a flag (not a command)
if arg == "-h" || arg == "-p" || arg == "-w" || arg == "-k" ||
arg == "-u" || arg == "-admin" || arg == "-user" {
hasOtherFlags = true
}
}
}
// Load config file with priority: -f > rf.yml > none
var config *ConfigFile
var err error
if configFilePath != "" {
// User specified config file via -f
config, err = LoadConfigFileFromPath(configFilePath)
if err != nil {
return nil, err
}
} else {
// Try default rf.yml
config, err = LoadDefaultConfigFile()
if err != nil {
return nil, err
}
}
if config != nil {
if hasOtherFlags {
return nil, fmt.Errorf("cannot use command line flags (-h, -p, -w, -k, -u, -admin, -user) when using config file. Please use config file or command line flags, not both")
}
return buildArgsFromConfig(config, args)
}
// Create a new flag set
fs := flag.NewFlagSet("ragflow_cli", flag.ContinueOnError)
// Define flags
host := fs.String("h", "127.0.0.1", "Admin or RAGFlow service host")
port := fs.Int("p", -1, "Admin or RAGFlow service port (default: 9381 for admin, 9380 for user)")
password := fs.String("w", "", "Superuser password")
key := fs.String("k", "", "API key for authentication")
_ = fs.String("f", "", "Path to config file (YAML format)") // Already parsed above
_ = fs.Bool("admin", false, "Run in admin mode (default)")
userMode := fs.Bool("user", false, "Run in user mode")
username := fs.String("u", "", "Username (email). In admin mode defaults to admin@ragflow.io, in user mode required")
// Parse the arguments
if err = fs.Parse(args); err != nil {
return nil, fmt.Errorf("failed to parse arguments: %v", err)
}
// Otherwise, use command line flags
return buildArgsFromFlags(host, port, password, key, userMode, username, fs.Args())
}
// buildArgsFromConfig builds ConnectionArgs from config file
func buildArgsFromConfig(config *ConfigFile, remainingArgs []string) (*ConnectionArgs, error) {
result := &ConnectionArgs{}
// Parse host:port from config file
if config.Host != "" {
host, port, err := parseHostPort(config.Host)
if err != nil {
return nil, fmt.Errorf("invalid host in config file: %v", err)
}
result.Host = host
result.Port = port
} else {
result.Host = "127.0.0.1"
}
// Apply auth info from config
result.Username = config.User
result.Password = config.Password
result.Key = config.APIToken
// Determine mode: if config has auth info, use user mode
if config.User != "" || config.APIToken != "" {
result.Type = "user"
} else {
result.Type = "admin"
result.Username = "admin@ragflow.io"
}
// Set default port if not specified in config
if result.Port == -1 {
if result.Type == "admin" {
result.Port = 9381
} else {
result.Port = 9380
}
}
// Get command from remaining args (no need for quotes or semicolon)
if len(remainingArgs) > 0 {
result.Command = strings.Join(remainingArgs, " ") + ";"
}
return result, nil
}
// buildArgsFromFlags builds ConnectionArgs from command line flags
func buildArgsFromFlags(host *string, port *int, password *string, key *string, userMode *bool, username *string, remainingArgs []string) (*ConnectionArgs, error) {
result := &ConnectionArgs{
Host: *host,
Port: *port,
Password: *password,
Key: *key,
Username: *username,
}
// Determine mode
if *userMode {
result.Type = "user"
} else {
result.Type = "admin"
}
// Set default port based on type if not specified
if result.Port == -1 {
if result.Type == "admin" {
result.Port = 9383
} else {
result.Port = 9384
}
}
// Determine username based on mode
if result.Type == "admin" && result.Username == "" {
result.Username = "admin@ragflow.io"
}
// Get command from remaining args (no need for quotes or semicolon)
if len(remainingArgs) > 0 {
result.Command = strings.Join(remainingArgs, " ") + ";"
}
return result, nil
}
// PrintUsage prints the CLI usage information
func PrintUsage() {
fmt.Println(`RAGFlow CLI Client
Usage: ragflow_cli [options] [command]
Options:
-h string Admin or RAGFlow service host (default "127.0.0.1")
-p int Admin or RAGFlow service port (default 9381 for admin, 9380 for user)
-w string Superuser password
-k string API key for authentication
-f string Path to config file (YAML format)
-admin Run in admin mode (default)
-user Run in user mode
-u string Username (email). In admin mode defaults to admin@ragflow.io
--help Show this help message
Configuration File:
The CLI will automatically read rf.yml from the current directory if it exists.
Use -f to specify a custom config file path.
Config file format:
host: 127.0.0.1:9380
user: your-email@example.com
password: your-password
api_token: your-api-token
When using a config file, you cannot use other command line flags except -help.
The command line is only for the SQL command.
Commands:
SQL commands like: LOGIN USER 'email'; LIST USERS; etc.
`)
}
// HistoryFile returns the path to the history file
func HistoryFile() string {
return os.Getenv("HOME") + "/" + historyFileName
@ -39,26 +309,95 @@ type CLI struct {
prompt string
running bool
line *liner.State
args *ConnectionArgs
}
// NewCLI creates a new CLI instance
func NewCLI() (*CLI, error) {
return NewCLIWithArgs(nil)
}
// NewCLIWithArgs creates a new CLI instance with connection arguments
func NewCLIWithArgs(args *ConnectionArgs) (*CLI, error) {
// Create liner first
line := liner.NewLiner()
// Determine server type
serverType := "user"
if args != nil && args.Type != "" {
serverType = args.Type
}
// Create client with password prompt using liner
client := NewRAGFlowClient("user") // Default to user mode
client := NewRAGFlowClient(serverType)
client.PasswordPrompt = line.PasswordPrompt
// Apply connection arguments if provided
client.HTTPClient.Host = args.Host
client.HTTPClient.Port = args.Port
// Set prompt based on server type
prompt := "RAGFlow> "
if serverType == "admin" {
prompt = "RAGFlow(admin)> "
}
return &CLI{
prompt: "RAGFlow> ",
prompt: prompt,
client: client,
line: line,
args: args,
}, nil
}
// Run starts the interactive CLI
func (c *CLI) Run() error {
if c.args.Type == "admin" {
// Allow 3 attempts for password verification
maxAttempts := 3
for attempt := 1; attempt <= maxAttempts; attempt++ {
var input string
var err error
// Check if terminal supports password masking
if term.IsTerminal(int(os.Stdin.Fd())) {
input, err = c.line.PasswordPrompt("Please input your password: ")
} else {
// Terminal doesn't support password masking, use regular prompt
fmt.Println("Warning: This terminal does not support secure password input")
input, err = c.line.Prompt("Please input your password (will be visible): ")
}
if err != nil {
fmt.Printf("Error reading input: %v\n", err)
return err
}
input = strings.TrimSpace(input)
if input == "" {
if attempt < maxAttempts {
fmt.Println("Password cannot be empty, please try again")
continue
}
return errors.New("no password provided after 3 attempts")
}
// Set the password for verification
c.args.Password = input
if err = c.VerifyAuth(); err != nil {
if attempt < maxAttempts {
fmt.Printf("Authentication failed: %v (%d/%d attempts)\n", err, attempt, maxAttempts)
continue
}
return fmt.Errorf("authentication failed after %d attempts: %v", maxAttempts, err)
}
// Authentication successful
break
}
}
c.running = true
// Load history from file
@ -99,8 +438,8 @@ func (c *CLI) Run() error {
c.line.AppendHistory(input)
}
if err := c.execute(input); err != nil {
fmt.Printf("Error: %v\n", err)
if err = c.execute(input); err != nil {
fmt.Printf("CLI error: %v\n", err)
}
}
@ -124,7 +463,11 @@ func (c *CLI) execute(input string) error {
}
// Execute the command using the client
_, err = c.client.ExecuteCommand(cmd)
var result ResponseIf
result, err = c.client.ExecuteCommand(cmd)
if result != nil {
result.PrintOut()
}
return err
}
@ -143,14 +486,12 @@ func (c *CLI) handleMetaCommand(cmd *Command) error {
fmt.Print("\033[H\033[2J")
case "admin":
c.client.ServerType = "admin"
c.client.HTTPClient.Port = 9381
c.prompt = "RAGFlow(admin)> "
fmt.Println("Switched to ADMIN mode (port 9381)")
fmt.Println("Switched to ADMIN mode")
case "user":
c.client.ServerType = "user"
c.client.HTTPClient.Port = 9380
c.prompt = "RAGFlow> "
fmt.Println("Switched to USER mode (port 9380)")
fmt.Println("Switched to USER mode")
case "host":
if len(args) == 0 {
fmt.Printf("Current host: %s\n", c.client.HTTPClient.Host)
@ -195,7 +536,7 @@ Meta Commands:
\q or \quit - Exit CLI
\c or \clear - Clear screen
SQL Commands (User Mode):
Commands (User Mode):
LOGIN USER 'email'; - Login as user
REGISTER USER 'name' AS 'nickname' PASSWORD 'pwd'; - Register new user
SHOW VERSION; - Show version info
@ -205,9 +546,14 @@ SQL Commands (User Mode):
LIST CHATS; - List user chats
LIST MODEL PROVIDERS; - List model providers
LIST DEFAULT MODELS; - List default models
LIST TOKENS; - List API tokens
CREATE TOKEN; - Create new API token
DROP TOKEN 'token_value'; - Delete an API token
SET TOKEN 'token_value'; - Set and validate API token
SHOW TOKEN; - Show current API token
UNSET TOKEN; - Remove current API token
SQL Commands (Admin Mode):
LOGIN USER 'email'; - Login as admin
Commands (Admin Mode):
LIST USERS; - List all users
SHOW USER 'email'; - Show user details
CREATE USER 'email' 'password'; - Create new user
@ -254,3 +600,34 @@ func RunInteractive() error {
return cli.Run()
}
// RunSingleCommand executes a single command and exits
func (c *CLI) RunSingleCommand(command string) error {
// Execute the command
if err := c.execute(command); err != nil {
return err
}
return nil
}
// VerifyAuth verifies authentication if needed
func (c *CLI) VerifyAuth() error {
if c.args == nil {
return nil
}
if c.args.Username == "" {
return fmt.Errorf("username is required")
}
if c.args.Password == "" {
return fmt.Errorf("password is required")
}
// Create login command with username and password
cmd := NewCommand("login_user")
cmd.Params["email"] = c.args.Username
cmd.Params["password"] = c.args.Password
_, err := c.client.ExecuteCommand(cmd)
return err
}

922
internal/cli/client.go
View File

File diff suppressed because it is too large Load Diff

110
internal/cli/http_client.go
View File

@ -31,12 +31,13 @@ type HTTPClient struct {
Host string
Port int
APIVersion string
APIKey string
APIToken string
LoginToken string
ConnectTimeout time.Duration
ReadTimeout time.Duration
VerifySSL bool
client *http.Client
useAPIToken bool
}
// NewHTTPClient creates a new HTTP client
@ -85,8 +86,8 @@ func (c *HTTPClient) Headers(authKind string, extra map[string]string) map[strin
headers := make(map[string]string)
switch authKind {
case "api":
if c.APIKey != "" {
headers["Authorization"] = fmt.Sprintf("Bearer %s", c.APIKey)
if c.APIToken != "" {
headers["Authorization"] = fmt.Sprintf("Bearer %s", c.APIToken)
}
case "web", "admin":
if c.LoginToken != "" {
@ -104,6 +105,7 @@ type Response struct {
StatusCode int
Body []byte
Headers http.Header
Duration float64
}
// JSON parses the response body as JSON
@ -142,11 +144,14 @@ func (c *HTTPClient) Request(method, path string, useAPIBase bool, authKind stri
req.Header.Set(k, v)
}
resp, err := c.client.Do(req)
var resp *http.Response
startTime := time.Now()
resp, err = c.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
duration := time.Since(startTime).Seconds()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
@ -157,21 +162,93 @@ func (c *HTTPClient) Request(method, path string, useAPIBase bool, authKind stri
StatusCode: resp.StatusCode,
Body: respBody,
Headers: resp.Header.Clone(),
Duration: duration,
}, nil
}
// Request makes an HTTP request
func (c *HTTPClient) RequestWith2URL(method, webPath string, apiPath string, headers map[string]string, jsonBody map[string]interface{}) (*Response, error) {
var path string
var useAPIBase bool
var authKind string
if c.useAPIToken {
path = apiPath
useAPIBase = true
authKind = "api"
} else {
path = webPath
useAPIBase = false
authKind = "web"
}
url := c.BuildURL(path, useAPIBase)
mergedHeaders := c.Headers(authKind, headers)
var body io.Reader
if jsonBody != nil {
jsonData, err := json.Marshal(jsonBody)
if err != nil {
return nil, err
}
body = bytes.NewReader(jsonData)
if mergedHeaders == nil {
mergedHeaders = make(map[string]string)
}
mergedHeaders["Content-Type"] = "application/json"
}
req, err := http.NewRequest(method, url, body)
if err != nil {
return nil, err
}
for k, v := range mergedHeaders {
req.Header.Set(k, v)
}
var resp *http.Response
startTime := time.Now()
resp, err = c.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
duration := time.Since(startTime).Seconds()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
return &Response{
StatusCode: resp.StatusCode,
Body: respBody,
Headers: resp.Header.Clone(),
Duration: duration,
}, nil
}
// RequestWithIterations makes multiple HTTP requests for benchmarking
// Returns a map with "duration" (total time in seconds) and "response_list"
func (c *HTTPClient) RequestWithIterations(method, path string, useAPIBase bool, authKind string, headers map[string]string, jsonBody map[string]interface{}, iterations int) (map[string]interface{}, error) {
func (c *HTTPClient) RequestWithIterations(method, path string, useAPIBase bool, authKind string, headers map[string]string, jsonBody map[string]interface{}, iterations int) (*BenchmarkResponse, error) {
response := new(BenchmarkResponse)
if iterations <= 1 {
start := time.Now()
resp, err := c.Request(method, path, useAPIBase, authKind, headers, jsonBody)
totalDuration := time.Since(start).Seconds()
if err != nil {
return nil, err
}
return map[string]interface{}{
"duration": 0.0,
"response_list": []*Response{resp},
}, nil
response.Code = resp.StatusCode
response.Duration = totalDuration
if response.Code == 0 {
response.SuccessCount = 1
} else {
response.FailureCount = 1
}
return response, nil
}
url := c.BuildURL(path, useAPIBase)
@ -232,10 +309,17 @@ func (c *HTTPClient) RequestWithIterations(method, path string, useAPIBase bool,
totalDuration += time.Since(start).Seconds()
}
return map[string]interface{}{
"duration": totalDuration,
"response_list": responseList,
}, nil
response.Code = 0
response.Duration = totalDuration
for _, resp := range responseList {
if resp.StatusCode == 200 {
response.SuccessCount++
} else {
response.FailureCount++
}
}
return response, nil
}
// RequestJSON makes an HTTP request and returns JSON response

6
internal/cli/lexer.go
View File

@ -215,6 +215,8 @@ func (l *Lexer) lookupIdent(ident string) Token {
return Token{Type: TokenOn, Value: ident}
case "SET":
return Token{Type: TokenSet, Value: ident}
case "UNSET":
return Token{Type: TokenUnset, Value: ident}
case "RESET":
return Token{Type: TokenReset, Value: ident}
case "VERSION":
@ -287,6 +289,10 @@ func (l *Lexer) lookupIdent(ident string) Token {
return Token{Type: TokenBenchmark, Value: ident}
case "PING":
return Token{Type: TokenPing, Value: ident}
case "TOKEN":
return Token{Type: TokenToken, Value: ident}
case "TOKENS":
return Token{Type: TokenTokens, Value: ident}
default:
return Token{Type: TokenIdentifier, Value: ident}
}

369
internal/cli/parser.go
View File

@ -102,6 +102,8 @@ func (p *Parser) parseSQLCommand() (*Command, error) {
return p.parseRevokeCommand()
case TokenSet:
return p.parseSetCommand()
case TokenUnset:
return p.parseUnsetCommand()
case TokenReset:
return p.parseResetCommand()
case TokenGenerate:
@ -189,18 +191,20 @@ func (p *Parser) parseLoginUser() (*Command, error) {
cmd.Params["email"] = email
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parsePingServer() (*Command, error) {
cmd := NewCommand("ping_server")
cmd := NewCommand("ping")
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -208,7 +212,6 @@ func (p *Parser) parsePingServer() (*Command, error) {
func (p *Parser) parseRegisterCommand() (*Command, error) {
cmd := NewCommand("register_user")
p.nextToken() // consume REGISTER
if err := p.expectPeek(TokenUser); err != nil {
return nil, err
}
@ -245,8 +248,9 @@ func (p *Parser) parseRegisterCommand() (*Command, error) {
cmd.Params["password"] = password
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
@ -258,38 +262,51 @@ func (p *Parser) parseListCommand() (*Command, error) {
switch p.curToken.Type {
case TokenServices:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_services"), nil
case TokenUsers:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_users"), nil
case TokenRoles:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_roles"), nil
case TokenVars:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_variables"), nil
case TokenConfigs:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_configs"), nil
case TokenTokens:
p.nextToken()
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_tokens"), nil
case TokenEnvs:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_environments"), nil
case TokenDatasets:
@ -304,8 +321,9 @@ func (p *Parser) parseListCommand() (*Command, error) {
return p.parseListDefaultModels()
case TokenChats:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_user_chats"), nil
case TokenFiles:
@ -334,8 +352,9 @@ func (p *Parser) parseListDatasets() (*Command, error) {
p.nextToken()
}
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -361,8 +380,9 @@ func (p *Parser) parseListAgents() (*Command, error) {
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -379,12 +399,13 @@ func (p *Parser) parseListKeys() (*Command, error) {
return nil, err
}
cmd := NewCommand("list_keys")
cmd := NewCommand("list_tokens")
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -395,8 +416,9 @@ func (p *Parser) parseListModelProviders() (*Command, error) {
return nil, fmt.Errorf("expected PROVIDERS")
}
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_user_model_providers"), nil
}
@ -407,8 +429,9 @@ func (p *Parser) parseListDefaultModels() (*Command, error) {
return nil, fmt.Errorf("expected MODELS")
}
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("list_user_default_models"), nil
}
@ -433,8 +456,9 @@ func (p *Parser) parseListFiles() (*Command, error) {
cmd.Params["dataset_name"] = datasetName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -445,18 +469,27 @@ func (p *Parser) parseShowCommand() (*Command, error) {
switch p.curToken.Type {
case TokenVersion:
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("show_version"), nil
case TokenToken:
p.nextToken()
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("show_token"), nil
case TokenCurrent:
p.nextToken()
if p.curToken.Type != TokenUser {
return nil, fmt.Errorf("expected USER after CURRENT")
}
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("show_current_user"), nil
case TokenUser:
@ -485,8 +518,9 @@ func (p *Parser) parseShowUser() (*Command, error) {
cmd := NewCommand("show_user_permission")
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -500,8 +534,9 @@ func (p *Parser) parseShowUser() (*Command, error) {
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -517,8 +552,9 @@ func (p *Parser) parseShowRole() (*Command, error) {
cmd.Params["role_name"] = roleName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -534,8 +570,9 @@ func (p *Parser) parseShowVariable() (*Command, error) {
cmd.Params["var_name"] = varName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -551,8 +588,9 @@ func (p *Parser) parseShowService() (*Command, error) {
cmd.Params["number"] = serviceNum
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -571,11 +609,24 @@ func (p *Parser) parseCreateCommand() (*Command, error) {
return p.parseCreateDataset()
case TokenChat:
return p.parseCreateChat()
case TokenToken:
return p.parseCreateToken()
default:
return nil, fmt.Errorf("unknown CREATE target: %s", p.curToken.Value)
}
}
func (p *Parser) parseCreateToken() (*Command, error) {
p.nextToken() // consume TOKEN
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("create_token"), nil
}
func (p *Parser) parseCreateUser() (*Command, error) {
p.nextToken() // consume USER
userName, err := p.parseQuotedString()
@ -595,8 +646,9 @@ func (p *Parser) parseCreateUser() (*Command, error) {
cmd.Params["role"] = "user"
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -622,8 +674,9 @@ func (p *Parser) parseCreateRole() (*Command, error) {
p.nextToken()
}
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -651,8 +704,9 @@ func (p *Parser) parseCreateModelProvider() (*Command, error) {
cmd.Params["provider_key"] = providerKey
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -704,8 +758,9 @@ func (p *Parser) parseCreateDataset() (*Command, error) {
return nil, fmt.Errorf("expected PARSER or PIPELINE")
}
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -721,8 +776,9 @@ func (p *Parser) parseCreateChat() (*Command, error) {
cmd.Params["chat_name"] = chatName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -743,11 +799,32 @@ func (p *Parser) parseDropCommand() (*Command, error) {
return p.parseDropChat()
case TokenKey:
return p.parseDropKey()
case TokenToken:
return p.parseDropToken()
default:
return nil, fmt.Errorf("unknown DROP target: %s", p.curToken.Value)
}
}
func (p *Parser) parseDropToken() (*Command, error) {
p.nextToken() // consume TOKEN
tokenValue, err := p.parseQuotedString()
if err != nil {
return nil, err
}
cmd := NewCommand("drop_token")
cmd.Params["token"] = tokenValue
p.nextToken()
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseDropUser() (*Command, error) {
p.nextToken() // consume USER
userName, err := p.parseQuotedString()
@ -759,8 +836,9 @@ func (p *Parser) parseDropUser() (*Command, error) {
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -776,8 +854,9 @@ func (p *Parser) parseDropRole() (*Command, error) {
cmd.Params["role_name"] = roleName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -798,8 +877,9 @@ func (p *Parser) parseDropModelProvider() (*Command, error) {
cmd.Params["provider_name"] = providerName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -815,8 +895,9 @@ func (p *Parser) parseDropDataset() (*Command, error) {
cmd.Params["dataset_name"] = datasetName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -832,15 +913,16 @@ func (p *Parser) parseDropChat() (*Command, error) {
cmd.Params["chat_name"] = chatName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseDropKey() (*Command, error) {
p.nextToken() // consume KEY
key, err := p.parseQuotedString()
token, err := p.parseQuotedString()
if err != nil {
return nil, err
}
@ -856,13 +938,14 @@ func (p *Parser) parseDropKey() (*Command, error) {
return nil, err
}
cmd := NewCommand("drop_key")
cmd.Params["key"] = key
cmd := NewCommand("drop_token")
cmd.Params["token"] = token
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -905,8 +988,9 @@ func (p *Parser) parseAlterUser() (*Command, error) {
cmd.Params["password"] = password
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for SHOW TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -936,8 +1020,9 @@ func (p *Parser) parseAlterUser() (*Command, error) {
cmd.Params["role_name"] = roleName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -954,16 +1039,17 @@ func (p *Parser) parseActivateUser() (*Command, error) {
status := p.curToken.Value
if status != "on" && status != "off" {
return nil, fmt.Errorf("expected 'on' or 'off', got %s", p.curToken.Value)
}
}
cmd := NewCommand("activate_user")
cmd.Params["user_name"] = userName
cmd.Params["activate_status"] = status
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
}
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -994,8 +1080,9 @@ func (p *Parser) parseAlterRole() (*Command, error) {
cmd.Params["description"] = description
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1021,8 +1108,9 @@ func (p *Parser) parseGrantAdmin() (*Command, error) {
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1064,8 +1152,9 @@ func (p *Parser) parseGrantPermission() (*Command, error) {
cmd.Params["role_name"] = roleName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1091,8 +1180,9 @@ func (p *Parser) parseRevokeAdmin() (*Command, error) {
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1134,8 +1224,9 @@ func (p *Parser) parseRevokePermission() (*Command, error) {
cmd.Params["role_name"] = roleName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1172,6 +1263,9 @@ func (p *Parser) parseSetCommand() (*Command, error) {
if p.curToken.Type == TokenDefault {
return p.parseSetDefault()
}
if p.curToken.Type == TokenToken {
return p.parseSetToken()
}
return nil, fmt.Errorf("unknown SET target: %s", p.curToken.Value)
}
@ -1194,8 +1288,9 @@ func (p *Parser) parseSetVariable() (*Command, error) {
cmd.Params["var_value"] = varValue
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1234,9 +1329,29 @@ func (p *Parser) parseSetDefault() (*Command, error) {
cmd.Params["model_id"] = modelID
p.nextToken()
if err := p.expectSemicolon(); err != nil {
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseSetToken() (*Command, error) {
p.nextToken() // consume TOKEN
tokenValue, err := p.parseQuotedString()
if err != nil {
return nil, err
}
cmd := NewCommand("set_token")
cmd.Params["token"] = tokenValue
p.nextToken()
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1270,8 +1385,9 @@ func (p *Parser) parseResetCommand() (*Command, error) {
cmd.Params["model_type"] = modelType
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1296,12 +1412,13 @@ func (p *Parser) parseGenerateCommand() (*Command, error) {
return nil, err
}
cmd := NewCommand("generate_key")
cmd := NewCommand("generate_token")
cmd.Params["user_name"] = userName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1333,8 +1450,9 @@ func (p *Parser) parseImportCommand() (*Command, error) {
cmd.Params["dataset_name"] = datasetName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1366,8 +1484,9 @@ func (p *Parser) parseSearchCommand() (*Command, error) {
cmd.Params["datasets"] = datasets
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1404,8 +1523,9 @@ func (p *Parser) parseParseDataset() (*Command, error) {
cmd.Params["method"] = method
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1436,8 +1556,9 @@ func (p *Parser) parseParseDocs() (*Command, error) {
cmd.Params["dataset_name"] = datasetName
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1482,6 +1603,8 @@ func (p *Parser) parseUserStatement() (*Command, error) {
return p.parseDropCommand()
case TokenSet:
return p.parseSetCommand()
case TokenUnset:
return p.parseUnsetCommand()
case TokenReset:
return p.parseResetCommand()
case TokenList:
@ -1513,8 +1636,9 @@ func (p *Parser) parseStartupCommand() (*Command, error) {
cmd.Params["number"] = serviceNum
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1535,8 +1659,9 @@ func (p *Parser) parseShutdownCommand() (*Command, error) {
cmd.Params["number"] = serviceNum
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
@ -1557,12 +1682,28 @@ func (p *Parser) parseRestartCommand() (*Command, error) {
cmd.Params["number"] = serviceNum
p.nextToken()
if err := p.expectSemicolon(); err != nil {
return nil, err
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return cmd, nil
}
func (p *Parser) parseUnsetCommand() (*Command, error) {
p.nextToken() // consume UNSET
if p.curToken.Type != TokenToken {
return nil, fmt.Errorf("expected TOKEN after UNSET")
}
p.nextToken()
// Semicolon is optional for UNSET TOKEN
if p.curToken.Type == TokenSemicolon {
p.nextToken()
}
return NewCommand("unset_token"), nil
}
func tokenTypeToString(t int) string {
// Simplified for error messages
return fmt.Sprintf("token(%d)", t)

3
internal/cli/types.go
View File

@ -95,6 +95,9 @@ const (
TokenSync
TokenBenchmark
TokenPing
TokenToken
TokenTokens
TokenUnset
// Literals
TokenIdentifier

548
internal/cli/user_command.go Normal file
View File

@ -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
}