Implement UpdateDataset and UpdateMetadata in GO (#13928)

### What problem does this PR solve? Implement UpdateDataset and UpdateMetadata in GO Add cli: UPDATE CHUNK <chunk_id> OF DATASET <dataset_name> SET <update_fields> REMOVE TAGS 'tag1', 'tag2' from DATASET 'dataset_name'; SET METADATA OF DOCUMENT <doc_id> TO <meta> ### Type of change - [ ] Refactoring
2026-08-24 17:10:12 +08:00 · 2026-04-07 09:44:51 +08:00
parent 60ec5880e5
commit 49386bc1b5
27 changed files with 1620 additions and 171 deletions
--- a/internal/cli/admin_parser.go
+++ b/internal/cli/admin_parser.go
@@ -35,6 +35,17 @@ func (p *Parser) parseAdminLoginUser() (*Command, error) {
 	cmd.Params["email"] = email

 	p.nextToken()
+	// Optional: PASSWORD 'password'
+	if p.curToken.Type == TokenPassword {
+		p.nextToken()
+		password, err := p.parseQuotedString()
+		if err != nil {
+			return nil, err
+		}
+		cmd.Params["password"] = password
+		p.nextToken()
+	}
+
 	// Semicolon is optional for UNSET TOKEN
 	if p.curToken.Type == TokenSemicolon {
 		p.nextToken()
--- a/internal/cli/cli.go
+++ b/internal/cli/cli.go
@@ -333,6 +333,7 @@ func looksLikeSQL(s string) bool {
 		"LOGIN ", "REGISTER ", "PING", "GRANT ", "REVOKE ",
 		"SET ", "UNSET ", "UPDATE ", "DELETE ", "INSERT ",
 		"SELECT ", "DESCRIBE ", "EXPLAIN ", "ADD ", "ENABLE ", "DISABLE ", "CHAT ", "USE", "THINK",
+		"REMOVE ",
 	}
 	for _, prefix := range sqlPrefixes {
 		if strings.HasPrefix(s, prefix) {
@@ -988,6 +989,7 @@ Meta Commands:

 Commands (User Mode):
  LOGIN USER 'email';                                    - Login as user
+  LOGIN USER 'email' PASSWORD 'pwd';                     - Login as user with password
  REGISTER USER 'name' AS 'nickname' PASSWORD 'pwd';     - Register new user
  SHOW VERSION;                                          - Show version info
  PING;                                                  - Ping server
--- a/internal/cli/client.go
+++ b/internal/cli/client.go
@@ -262,6 +262,12 @@ func (c *RAGFlowClient) ExecuteUserCommand(cmd *Command) (ResponseIf, error) {
 		return c.InsertDatasetFromFile(cmd)
 	case "insert_metadata_from_file":
 		return c.InsertMetadataFromFile(cmd)
+	case "update_chunk":
+		return c.UpdateChunk(cmd)
+	case "set_meta":
+		return c.SetMeta(cmd)
+	case "rm_tags":
+		return c.RmTags(cmd)
 	// TODO: Implement other commands
 	default:
 		return nil, fmt.Errorf("command '%s' would be executed with API", cmd.Type)
--- a/internal/cli/http_client.go
+++ b/internal/cli/http_client.go
@@ -84,10 +84,19 @@ func (c *HTTPClient) BuildURL(path string, useAPIBase bool) string {
 // Headers builds the request headers
 func (c *HTTPClient) Headers(authKind string, extra map[string]string) map[string]string {
 	headers := make(map[string]string)
-	if c.APIToken != "" {
-		headers["Authorization"] = fmt.Sprintf("Bearer %s", c.APIToken)
-	} else if c.LoginToken != "" {
-		headers["Authorization"] = c.LoginToken
+
+	switch authKind {
+	case "api":
+		if c.APIToken != "" {
+			headers["Authorization"] = fmt.Sprintf("Bearer %s", c.APIToken)
+		} else if c.LoginToken != "" {
+			// Fallback to login token for API requests (user mode)
+			headers["Authorization"] = fmt.Sprintf("Bearer %s", c.LoginToken)
+		}
+	case "web", "admin":
+		if c.LoginToken != "" {
+			headers["Authorization"] = c.LoginToken
+		}
 	}

 	for k, v := range extra {
--- a/internal/cli/lexer.go
+++ b/internal/cli/lexer.go
@@ -327,6 +327,16 @@ func (l *Lexer) lookupIdent(ident string) Token {
 		return Token{Type: TokenMetadata, Value: ident}
 	case "USE":
 		return Token{Type: TokenUse, Value: ident}
+	case "UPDATE":
+		return Token{Type: TokenUpdate, Value: ident}
+	case "REMOVE":
+		return Token{Type: TokenRemove, Value: ident}
+	case "CHUNK":
+		return Token{Type: TokenChunk, Value: ident}
+	case "DOCUMENT":
+		return Token{Type: TokenDocument, Value: ident}
+	case "TAGS":
+		return Token{Type: TokenTag, Value: ident}
 	default:
 		return Token{Type: TokenIdentifier, Value: ident}
 	}
--- a/internal/cli/parser.go
+++ b/internal/cli/parser.go
@@ -196,6 +196,10 @@ func (p *Parser) parseUserCommand() (*Command, error) {
 		return p.parseThinkCommand()
 	case TokenUse:
 		return p.parseUseCommand()
+	case TokenUpdate:
+		return p.parseUpdateCommand()
+	case TokenRemove:
+		return p.parseRemoveCommand()
 	default:
 		return nil, fmt.Errorf("unknown command: %s", p.curToken.Value)
 	}
@@ -233,7 +237,7 @@ func (p *Parser) expectSemicolon() error {
 }

 func isKeyword(tokenType int) bool {
-	return tokenType >= TokenLogin && tokenType <= TokenMetadata
+	return tokenType >= TokenLogin && tokenType <= TokenTag
 }

 // isCECommand checks if the given string is a ContextEngine command
--- a/internal/cli/types.go
+++ b/internal/cli/types.go
@@ -115,6 +115,11 @@ const (
 	TokenInsert
 	TokenFile
 	TokenMetadata
+	TokenUpdate
+	TokenRemove
+	TokenChunk
+	TokenDocument
+	TokenTag

 	// Literals
 	TokenIdentifier
--- a/internal/cli/user_command.go
+++ b/internal/cli/user_command.go
@@ -199,13 +199,13 @@ func (c *RAGFlowClient) ListUserDatasets(cmd *Command) (ResponseIf, error) {

 // 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)
+	resp, err := c.HTTPClient.Request("GET", "/datasets", true, "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)
+		return "", fmt.Errorf("failed to list datasets: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
 	}

 	resJSON, err := resp.JSON()
@@ -219,17 +219,12 @@ func (c *RAGFlowClient) getDatasetID(datasetName string) (string, error) {
 		return "", fmt.Errorf("failed to list datasets: %s", msg)
 	}

-	data, ok := resJSON["data"].(map[string]interface{})
+	data, ok := resJSON["data"].([]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 {
+	for _, kb := range data {
 		if kbMap, ok := kb.(map[string]interface{}); ok {
 			if name, _ := kbMap["name"].(string); name == datasetName {
 				if id, _ := kbMap["id"].(string); id != "" {
@@ -1487,3 +1482,195 @@ func (c *RAGFlowClient) InsertMetadataFromFile(cmd *Command) (ResponseIf, error)
 	result.Duration = 0
 	return &result, nil
 }
+
+// UpdateChunk updates a chunk in a dataset
+func (c *RAGFlowClient) UpdateChunk(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	chunkID, ok := cmd.Params["chunk_id"].(string)
+	if !ok {
+		return nil, fmt.Errorf("chunk_id not provided")
+	}
+
+	datasetName, ok := cmd.Params["dataset_name"].(string)
+	if !ok {
+		return nil, fmt.Errorf("dataset_name not provided")
+	}
+
+	jsonBody, ok := cmd.Params["json_body"].(string)
+	if !ok {
+		return nil, fmt.Errorf("json_body not provided")
+	}
+
+	// Look up dataset_id from dataset_name
+	datasetID, err := c.getDatasetID(datasetName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dataset ID: %w", err)
+	}
+
+	// Try to get doc_id from the chunk retrieval endpoint
+	getResp, err := c.HTTPClient.Request("GET", "/chunk/get?chunk_id="+chunkID, false, "web", nil, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get chunk info: %w", err)
+	}
+
+	var docID string
+	if getResp.StatusCode == 200 {
+		getJSON, err := getResp.JSON()
+		if err == nil {
+			if data, ok := getJSON["data"].(map[string]interface{}); ok {
+				if d, ok := data["doc_id"].(string); ok {
+					docID = d
+				}
+			}
+		}
+	}
+
+	if docID == "" {
+		return nil, fmt.Errorf("could not find document_id for chunk %s. Please provide document_id explicitly", chunkID)
+	}
+
+	// Parse the JSON body
+	var payload map[string]interface{}
+	if err := json.Unmarshal([]byte(jsonBody), &payload); err != nil {
+		return nil, fmt.Errorf("invalid JSON body: %w", err)
+	}
+
+	path := fmt.Sprintf("/datasets/%s/documents/%s/chunks/%s", datasetID, docID, chunkID)
+	resp, err := c.HTTPClient.Request("PUT", path, true, "api", nil, payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to update chunk: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to update chunk: 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 {
+		return nil, fmt.Errorf("invalid response format: code is not a number")
+	}
+
+	var result SimpleResponse
+	result.Code = int(code)
+	if result.Code == 0 {
+		result.Message = fmt.Sprintf("Success to update chunk: %s", chunkID)
+	} else {
+		result.Message = fmt.Sprintf("Failed to update chunk: %v", resJSON)
+	}
+	result.Duration = 0
+	return &result, nil
+}
+
+// SetMeta sets metadata for a document
+func (c *RAGFlowClient) SetMeta(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	docID, ok := cmd.Params["doc_id"].(string)
+	if !ok {
+		return nil, fmt.Errorf("doc_id not provided")
+	}
+
+	metaJSON, ok := cmd.Params["meta"].(string)
+	if !ok {
+		return nil, fmt.Errorf("meta not provided")
+	}
+
+	payload := map[string]interface{}{
+		"doc_id": docID,
+		"meta":   metaJSON,
+	}
+
+	resp, err := c.HTTPClient.Request("POST", "/document/set_meta", false, "web", nil, payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to set metadata: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to set metadata: 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 {
+		return nil, fmt.Errorf("invalid response format: code is not a number")
+	}
+
+	var result SimpleResponse
+	result.Code = int(code)
+	if result.Code == 0 {
+		result.Message = fmt.Sprintf("Success to set metadata for document: %s", docID)
+	} else {
+		result.Message = fmt.Sprintf("Failed to set metadata: %v", resJSON)
+	}
+	result.Duration = 0
+	return &result, nil
+}
+
+// RmTags removes tags from chunks in a dataset
+func (c *RAGFlowClient) RmTags(cmd *Command) (ResponseIf, error) {
+	if c.ServerType != "user" {
+		return nil, fmt.Errorf("this command is only allowed in USER mode")
+	}
+
+	datasetName, ok := cmd.Params["dataset_name"].(string)
+	if !ok {
+		return nil, fmt.Errorf("dataset_name not provided")
+	}
+
+	kbID, err := c.getDatasetID(datasetName)
+	if err != nil {
+		return nil, err
+	}
+
+	tags, ok := cmd.Params["tags"].([]string)
+	if !ok {
+		return nil, fmt.Errorf("tags not provided")
+	}
+
+	payload := map[string]interface{}{
+		"tags": tags,
+	}
+
+	resp, err := c.HTTPClient.Request("POST", "/kb/"+kbID+"/rm_tags", false, "web", nil, payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to remove tags: %w", err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("failed to remove tags: 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 {
+		return nil, fmt.Errorf("invalid response format: code is not a number")
+	}
+
+	var result SimpleResponse
+	result.Code = int(code)
+	if result.Code == 0 {
+		result.Message = fmt.Sprintf("Success to remove tags from dataset: %s", kbID)
+	} else {
+		result.Message = fmt.Sprintf("Failed to remove tags: %v", resJSON)
+	}
+	result.Duration = 0
+	return &result, nil
+}
--- a/internal/cli/user_parser.go
+++ b/internal/cli/user_parser.go
@@ -33,12 +33,8 @@ func (p *Parser) parseLoginUser() (*Command, error) {
 	cmd.Params["email"] = email

 	p.nextToken()
-	// Optional: WITH PASSWORD 'password'
-	if p.curToken.Type == TokenWith {
-		p.nextToken()
-		if p.curToken.Type != TokenPassword {
-			return nil, fmt.Errorf("expected PASSWORD after WITH")
-		}
+	// Optional: PASSWORD 'password'
+	if p.curToken.Type == TokenPassword {
 		p.nextToken()
 		password, err := p.parseQuotedString()
 		if err != nil {
@@ -853,6 +849,17 @@ func (p *Parser) parseDeleteCommand() (*Command, error) {
 	}
 }

+func (p *Parser) parseRemoveCommand() (*Command, error) {
+	p.nextToken() // consume RM
+
+	switch p.curToken.Type {
+	case TokenTag:
+		return p.parseRemoveTags()
+	default:
+		return nil, fmt.Errorf("unknown REMOVE target: %s", p.curToken.Value)
+	}
+}
+
 func (p *Parser) parseDropToken() (*Command, error) {
 	p.nextToken() // consume TOKEN

@@ -1574,6 +1581,9 @@ func (p *Parser) parseSetCommand() (*Command, error) {
 	if p.curToken.Type == TokenToken {
 		return p.parseSetToken()
 	}
+	if p.curToken.Type == TokenMetadata {
+		return p.parseSetMeta()
+	}

 	return nil, fmt.Errorf("unknown SET target: %s", p.curToken.Value)
 }
@@ -2229,7 +2239,10 @@ func (p *Parser) parseUserStatement() (*Command, error) {
 		return p.parseInsertCommand()
 	case TokenSearch:
 		return p.parseSearchCommand()
-
+	case TokenUpdate:
+		return p.parseUpdateCommand()
+	case TokenRemove:
+		return p.parseRemoveCommand()
 	default:
 		return nil, fmt.Errorf("invalid user statement: %s", p.curToken.Value)
 	}
@@ -2318,3 +2331,164 @@ func (p *Parser) parseUnsetCommand() (*Command, error) {
 	}
 	return NewCommand("unset_token"), nil
 }
+
+// parseUpdateCommand parses UPDATE CHUNK command
+// UPDATE CHUNK 'chunk_id' OF DATASET 'dataset_name' SET '{"content": "..."}'
+func (p *Parser) parseUpdateCommand() (*Command, error) {
+	p.nextToken() // consume UPDATE
+
+	if p.curToken.Type != TokenChunk {
+		return nil, fmt.Errorf("expected CHUNK after UPDATE")
+	}
+	p.nextToken()
+
+	// Parse chunk_id
+	chunkID, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected chunk_id: %w", err)
+	}
+
+	cmd := NewCommand("update_chunk")
+	cmd.Params["chunk_id"] = chunkID
+
+	p.nextToken()
+	if p.curToken.Type != TokenOf {
+		return nil, fmt.Errorf("expected OF after chunk_id")
+	}
+	p.nextToken()
+
+	if p.curToken.Type != TokenDataset {
+		return nil, fmt.Errorf("expected DATASET after OF")
+	}
+	p.nextToken()
+
+	// Parse dataset_name
+	datasetName, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected dataset_name: %w", err)
+	}
+	cmd.Params["dataset_name"] = datasetName
+
+	p.nextToken()
+	if p.curToken.Type != TokenSet {
+		return nil, fmt.Errorf("expected SET after dataset_name")
+	}
+	p.nextToken()
+
+	// Parse JSON body
+	jsonBody, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected JSON body: %w", err)
+	}
+	cmd.Params["json_body"] = jsonBody
+
+	p.nextToken()
+	// Semicolon is optional
+	if p.curToken.Type == TokenSemicolon {
+		p.nextToken()
+	}
+
+	return cmd, nil
+}
+
+// parseSetMeta parses: SET METADATA OF DOCUMENT 'doc_id' TO '{"key": "value"}'
+func (p *Parser) parseSetMeta() (*Command, error) {
+	p.nextToken() // consume METADATA
+
+	// Expect OF
+	if p.curToken.Type != TokenOf {
+		return nil, fmt.Errorf("expected OF after SET METADATA")
+	}
+	p.nextToken()
+
+	// Expect DOCUMENT
+	if p.curToken.Type != TokenDocument {
+		return nil, fmt.Errorf("expected DOCUMENT after SET METADATA OF")
+	}
+	p.nextToken()
+
+	// Parse doc_id
+	docID, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected doc_id: %w", err)
+	}
+	cmd := NewCommand("set_meta")
+	cmd.Params["doc_id"] = docID
+
+	p.nextToken()
+	// Expect TO
+	if p.curToken.Type != TokenTo {
+		return nil, fmt.Errorf("expected TO after doc_id")
+	}
+	p.nextToken()
+
+	// Parse meta JSON
+	meta, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected meta JSON: %w", err)
+	}
+	cmd.Params["meta"] = meta
+
+	p.nextToken()
+	// Semicolon is optional
+	if p.curToken.Type == TokenSemicolon {
+		p.nextToken()
+	}
+
+	return cmd, nil
+}
+
+// parseRemoveTags parses: REMOVE TAGS 'tag1', 'tag2' from DATASET 'dataset_name';
+func (p *Parser) parseRemoveTags() (*Command, error) {
+	p.nextToken() // consume TAGS
+
+	// Parse first tag
+	tag, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected tag: %w", err)
+	}
+	tags := []string{tag}
+
+	// Parse additional tags separated by commas
+	for {
+		p.nextToken()
+		if p.curToken.Type == TokenComma {
+			p.nextToken()
+			tag, err := p.parseQuotedString()
+			if err != nil {
+				return nil, fmt.Errorf("expected tag after comma: %w", err)
+			}
+			tags = append(tags, tag)
+		} else {
+			break
+		}
+	}
+	cmd := NewCommand("rm_tags")
+	cmd.Params["tags"] = tags
+
+	// Expect from
+	if p.curToken.Type != TokenFrom {
+		return nil, fmt.Errorf("expected FROM after tags")
+	}
+	p.nextToken()
+
+	// Expect DATASET
+	if p.curToken.Type != TokenDataset {
+		return nil, fmt.Errorf("expected DATASET after FROM")
+	}
+	p.nextToken()
+
+	// Parse dataset_name
+	datasetName, err := p.parseQuotedString()
+	if err != nil {
+		return nil, fmt.Errorf("expected dataset_name: %w", err)
+	}
+	cmd.Params["dataset_name"] = datasetName
+
+	// Semicolon is optional
+	if p.curToken.Type == TokenSemicolon {
+		p.nextToken()
+	}
+
+	return cmd, nil
+}