mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-19 14:11:04 +08:00
Go: add cli command, list dataset documents (#14948)
### What problem does this PR solve? ``` +---------------------+----------------------------------+-------------+-----------------+---------+--------+------+ | created_at | id | meta_fields | name | size | status | type | +---------------------+----------------------------------+-------------+-----------------+---------+--------+------+ | 2026-05-08 19:35:08 | f6aa38bb4ad111f1ba6338a74640adcc | map[] | abc.pdf | 3387987 | 1 | pdf | +---------------------+----------------------------------+-------------+-----------------+---------+--------+------+ ``` ### 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:
@@ -203,6 +203,8 @@ func (c *RAGFlowClient) ExecuteUserCommand(cmd *Command) (ResponseIf, error) {
|
||||
return c.RunBenchmark(cmd)
|
||||
case "list_datasets":
|
||||
return c.ListDatasets(cmd)
|
||||
case "list_dataset_documents":
|
||||
return c.ListDatasetDocumentUserCommand(cmd)
|
||||
case "search_on_datasets":
|
||||
return c.SearchOnDatasets(cmd)
|
||||
case "create_token":
|
||||
|
||||
@@ -431,6 +431,8 @@ func (l *Lexer) lookupIdent(ident string) Token {
|
||||
return Token{Type: TokenChunks, Value: ident}
|
||||
case "DOCUMENT":
|
||||
return Token{Type: TokenDocument, Value: ident}
|
||||
case "DOCUMENTS":
|
||||
return Token{Type: TokenDocuments, Value: ident}
|
||||
case "TAGS":
|
||||
return Token{Type: TokenTag, Value: ident}
|
||||
case "REGION":
|
||||
|
||||
@@ -85,6 +85,42 @@ func (r *CommonDataResponse) PrintOut() {
|
||||
}
|
||||
}
|
||||
|
||||
type ListDocumentsResponse struct {
|
||||
Code int `json:"code"`
|
||||
Data map[string]interface{} `json:"data"`
|
||||
Message string `json:"message"`
|
||||
Duration float64
|
||||
OutputFormat OutputFormat
|
||||
}
|
||||
|
||||
func (r *ListDocumentsResponse) Type() string {
|
||||
return "list_documents"
|
||||
}
|
||||
|
||||
func (r *ListDocumentsResponse) TimeCost() float64 {
|
||||
return r.Duration
|
||||
}
|
||||
|
||||
func (r *ListDocumentsResponse) SetOutputFormat(format OutputFormat) {
|
||||
r.OutputFormat = format
|
||||
}
|
||||
|
||||
func (r *ListDocumentsResponse) PrintOut() {
|
||||
if r.Code == 0 {
|
||||
total := r.Data["total"].(float64)
|
||||
fmt.Printf("Total: %0.0f\n", total)
|
||||
docs := r.Data["docs"].([]interface{})
|
||||
table := make([]map[string]interface{}, 0)
|
||||
for _, doc := range docs {
|
||||
table = append(table, doc.(map[string]interface{}))
|
||||
}
|
||||
PrintTableSimpleByFormat(table, r.OutputFormat)
|
||||
} else {
|
||||
fmt.Println("ERROR")
|
||||
fmt.Printf("%d, %s\n", r.Code, r.Message)
|
||||
}
|
||||
}
|
||||
|
||||
type SimpleResponse struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
|
||||
@@ -149,6 +149,7 @@ const (
|
||||
TokenChunk
|
||||
TokenChunks
|
||||
TokenDocument
|
||||
TokenDocuments
|
||||
TokenTag
|
||||
TokenRegion
|
||||
TokenURL
|
||||
|
||||
@@ -391,7 +391,63 @@ func (c *RAGFlowClient) ListDatasets(cmd *Command) (ResponseIf, error) {
|
||||
|
||||
var result CommonResponse
|
||||
if err = json.Unmarshal(resp.Body, &result); err != nil {
|
||||
return nil, fmt.Errorf("list users failed: invalid JSON (%w)", err)
|
||||
return nil, fmt.Errorf("list datasets failed: invalid JSON (%w)", err)
|
||||
}
|
||||
|
||||
if result.Code != 0 {
|
||||
return nil, fmt.Errorf("%s", result.Message)
|
||||
}
|
||||
result.Duration = resp.Duration
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// ListDatasetDocumentUserCommand lists dataset documents
|
||||
func (c *RAGFlowClient) ListDatasetDocumentUserCommand(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
|
||||
}
|
||||
|
||||
// Determine auth kind based on whether API token is being used
|
||||
if c.HTTPClient.LoginToken == "" && !c.HTTPClient.useAPIToken {
|
||||
return nil, fmt.Errorf("no authorization")
|
||||
}
|
||||
|
||||
datasetID, ok := cmd.Params["dataset_id"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no dataset id")
|
||||
}
|
||||
|
||||
page := 1
|
||||
pageSize := 10
|
||||
keywords := ""
|
||||
returnEmptyMetadata := "true"
|
||||
url := fmt.Sprintf("/datasets/%s/documents?page=%d&page_size=%d&keywords=%s&return_empty_metadata=%s", datasetID, page, pageSize, keywords, returnEmptyMetadata)
|
||||
|
||||
if iterations > 1 {
|
||||
// Benchmark mode - return raw result for benchmark stats
|
||||
return c.HTTPClient.RequestWithIterations("GET", url, "web", nil, nil, iterations)
|
||||
}
|
||||
|
||||
// Normal mode
|
||||
resp, err := c.HTTPClient.Request("GET", url, "web", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list documents: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("failed to list documents: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
|
||||
}
|
||||
|
||||
var result ListDocumentsResponse
|
||||
if err = json.Unmarshal(resp.Body, &result); err != nil {
|
||||
return nil, fmt.Errorf("list documents failed: invalid JSON (%w)", err)
|
||||
}
|
||||
|
||||
if result.Code != 0 {
|
||||
|
||||
@@ -136,6 +136,8 @@ func (p *Parser) parseListCommand() (*Command, error) {
|
||||
return NewCommand("list_environments"), nil
|
||||
case TokenDatasets:
|
||||
return p.parseListDatasets()
|
||||
case TokenDocuments:
|
||||
return p.parseListDatasetDocuments()
|
||||
case TokenAgents:
|
||||
return p.parseListAgents()
|
||||
case TokenTokens:
|
||||
@@ -181,6 +183,31 @@ func (p *Parser) parseListDatasets() (*Command, error) {
|
||||
return cmd, nil
|
||||
}
|
||||
|
||||
func (p *Parser) parseListDatasetDocuments() (*Command, error) {
|
||||
p.nextToken() // consume DOCUMENTS
|
||||
|
||||
if p.curToken.Type != TokenFrom {
|
||||
return nil, fmt.Errorf("expected FROM")
|
||||
}
|
||||
p.nextToken()
|
||||
|
||||
datasetID, err := p.parseQuotedString()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.nextToken()
|
||||
|
||||
cmd := NewCommand("list_dataset_documents")
|
||||
cmd.Params["dataset_id"] = datasetID
|
||||
|
||||
// Semicolon is optional for UNSET TOKEN
|
||||
if p.curToken.Type == TokenSemicolon {
|
||||
p.nextToken()
|
||||
}
|
||||
|
||||
return cmd, nil
|
||||
}
|
||||
|
||||
func (p *Parser) parseListAgents() (*Command, error) {
|
||||
p.nextToken() // consume AGENTS
|
||||
|
||||
|
||||
Reference in New Issue
Block a user