zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-15 05:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go CLI: embed and rerank (#14735)

Browse Source 
### What problem does this PR solve?

```
RAGFlow(user)> embed text 'what is rag' 'who are you' with 'embedding-3@test@zhipu-ai' dimension 16;
+-----------+-------+
| dimension | index |
+-----------+-------+
| 16        | 0     |
| 16        | 1     |
+-----------+-------+

RAGFlow(user)> rerank query 'what is rag' document 'rag is retrieval augment generation' 'rag need llm' 'famous rag project includes ragflow' with 'rerank@test@zhipu-ai' top 2;
+-------+-----------------+
| index | relevance_score |
+-------+-----------------+
| 0     | 1               |
| 2     | 0.99999976      |
+-------+-----------------+
```

### 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-05-09 17:41:54 +08:00
committed by GitHub
parent efe6d23d61
commit 17d71e5d79
31 changed files with 919 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/cli/client.go
View File

@@ -263,6 +263,10 @@ func (c *RAGFlowClient) ExecuteUserCommand(cmd *Command) (ResponseIf, error) {
return c.ChatToModel(cmd)
case "think_chat_to_model":
return c.ChatToModel(cmd)
case "embed_user_text":
return c.EmbedUserText(cmd)
case "rarank_user_document":
return c.RerankUserDocument(cmd)
case "check_provider_connection":
return c.CheckProviderConnection(cmd)
case "use_model":

10
internal/cli/lexer.go
View File

@@ -363,6 +363,16 @@ func (l *Lexer) lookupIdent(ident string) Token {
return Token{Type: TokenASR, Value: ident}
case "TTS":
return Token{Type: TokenTTS, Value: ident}
case "EMBED":
return Token{Type: TokenEmbed, Value: ident}
case "TEXT":
return Token{Type: TokenText, Value: ident}
case "QUERY":
return Token{Type: TokenQuery, Value: ident}
case "TOP":
return Token{Type: TokenTop, Value: ident}
case "DIMENSION":
return Token{Type: TokenDimension, Value: ident}
case "OCR":
return Token{Type: TokenOCR, Value: ident}
case "ASYNC":

64
internal/cli/parser.go
View File

@@ -197,6 +197,10 @@ func (p *Parser) parseUserCommand() (*Command, error) {
return p.parseChatCommand()
case TokenThink:
return p.parseThinkCommand()
case TokenEmbed:
return p.parseEmbedCommand()
case TokenRerank:
return p.parseRerankCommand()
case TokenCheck:
return p.parseCheckCommand()
case TokenLS:
@@ -495,43 +499,43 @@ func (p *Parser) parseCESearchCommand() (*Command, error) {
p.curToken.Type == TokenChats || p.curToken.Type == TokenDatasets {
path = path + "/" + p.curToken.Value
p.nextToken()
} else if p.curToken.Type == TokenNumber {
// Handle version numbers like 1.0.0 (parsed as number . number . number)
// OR filenames starting with numbers like 3_list_compressors.pdf
numberPart := p.curToken.Value
p.nextToken()
// Continue reading .number parts (version number format)
if p.curToken.Type == TokenIllegal && p.curToken.Value == "." {
versionPart := numberPart
for p.curToken.Type == TokenIllegal && p.curToken.Value == "." {
p.nextToken() // consume .
if p.curToken.Type == TokenNumber {
versionPart = versionPart + "." + p.curToken.Value
p.nextToken()
} else {
break
} else if p.curToken.Type == TokenNumber {
// Handle version numbers like 1.0.0 (parsed as number . number . number)
// OR filenames starting with numbers like 3_list_compressors.pdf
numberPart := p.curToken.Value
p.nextToken()
// Continue reading .number parts (version number format)
if p.curToken.Type == TokenIllegal && p.curToken.Value == "." {
versionPart := numberPart
for p.curToken.Type == TokenIllegal && p.curToken.Value == "." {
p.nextToken() // consume .
if p.curToken.Type == TokenNumber {
versionPart = versionPart + "." + p.curToken.Value
p.nextToken()
} else {
break
}
}
path = path + "/" + versionPart
} else if p.curToken.Type == TokenIdentifier {
// Filename starting with number: 3_list_compressors.pdf
path = path + "/" + numberPart + p.curToken.Value
p.nextToken()
} else {
// Just a number
path = path + "/" + numberPart
}
path = path + "/" + versionPart
} else if p.curToken.Type == TokenIdentifier {
// Filename starting with number: 3_list_compressors.pdf
path = path + "/" + numberPart + p.curToken.Value
} else if p.curToken.Type == TokenQuotedString {
path = path + "/" + strings.Trim(p.curToken.Value, "\"'")
p.nextToken()
} else {
// Just a number
path = path + "/" + numberPart
// Trailing slash, just append it
path = path + "/"
break
}
} else if p.curToken.Type == TokenQuotedString {
path = path + "/" + strings.Trim(p.curToken.Value, "\"'")
p.nextToken()
} else {
// Trailing slash, just append it
path = path + "/"
break
}
}
cmd.Params["path"] = path
cmd.Params["path"] = path
} else {
cmd.Params["path"] = "."
}

5
internal/cli/types.go
View File

@@ -102,6 +102,11 @@ const (
TokenASR
TokenTTS
TokenOCR
TokenEmbed
TokenText
TokenQuery
TokenTop
TokenDimension
TokenAsync
TokenSync
TokenBenchmark

148
internal/cli/user_command.go
View File

@@ -1572,7 +1572,6 @@ func (c *RAGFlowClient) ChatToModel(cmd *Command) (ResponseIf, error) {
"text": message,
})
}
}
images, ok := cmd.Params["images"].([]string)
@@ -1783,6 +1782,146 @@ func (c *RAGFlowClient) ChatToModel(cmd *Command) (ResponseIf, error) {
return &result, nil
}
func (c *RAGFlowClient) EmbedUserText(cmd *Command) (ResponseIf, error) {
if c.HTTPClient.APIToken == "" && c.HTTPClient.LoginToken == "" {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.ServerType != "user" {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
var providerName, instanceName, modelName string
// Check if composite_model_name is provided in command
if compositeModelName, ok := cmd.Params["composite_model_name"].(string); ok && compositeModelName != "" {
names := strings.Split(compositeModelName, "@")
if len(names) != 3 {
return nil, fmt.Errorf("model name must be in format 'model@instance@provider'")
}
providerName = names[2]
instanceName = names[1]
modelName = names[0]
} else if c.CurrentModel != nil {
// Use current model if set
providerName = c.CurrentModel.Provider
instanceName = c.CurrentModel.Instance
modelName = c.CurrentModel.Model
} else {
return nil, fmt.Errorf("model name not provided and no current model set. Use 'use model' command first")
}
texts, ok := cmd.Params["texts"].([]string)
if !ok {
return nil, fmt.Errorf("texts not provided")
}
dimension, ok := cmd.Params["dimension"].(int)
if !ok {
dimension = 0
}
payload := map[string]interface{}{
"provider_name": providerName,
"instance_name": instanceName,
"model_name": modelName,
"texts": texts,
"dimension": dimension,
}
url := "/embeddings"
resp, err := c.HTTPClient.Request("POST", url, "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to embed text: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to embed text: 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("embed text failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *RAGFlowClient) RerankUserDocument(cmd *Command) (ResponseIf, error) {
if c.HTTPClient.APIToken == "" && c.HTTPClient.LoginToken == "" {
return nil, fmt.Errorf("API token not set. Please login first")
}
if c.ServerType != "user" {
return nil, fmt.Errorf("this command is only allowed in USER mode")
}
var providerName, instanceName, modelName string
// Check if composite_model_name is provided in command
if compositeModelName, ok := cmd.Params["composite_model_name"].(string); ok && compositeModelName != "" {
names := strings.Split(compositeModelName, "@")
if len(names) != 3 {
return nil, fmt.Errorf("model name must be in format 'model@instance@provider'")
}
providerName = names[2]
instanceName = names[1]
modelName = names[0]
} else if c.CurrentModel != nil {
// Use current model if set
providerName = c.CurrentModel.Provider
instanceName = c.CurrentModel.Instance
modelName = c.CurrentModel.Model
} else {
return nil, fmt.Errorf("model name not provided and no current model set. Use 'use model' command first")
}
query, ok := cmd.Params["query"].(string)
if !ok {
return nil, fmt.Errorf("query not provided")
}
documents, ok := cmd.Params["documents"].([]string)
if !ok {
return nil, fmt.Errorf("documents not provided")
}
topN, ok := cmd.Params["top_n"].(int)
if !ok {
return nil, fmt.Errorf("top n not provided")
}
payload := map[string]interface{}{
"provider_name": providerName,
"instance_name": instanceName,
"model_name": modelName,
"query": query,
"documents": documents,
"top_n": topN,
}
url := "/rerank"
resp, err := c.HTTPClient.Request("POST", url, "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to rerank document: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to rerank document: 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("rerank document failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)
}
result.Duration = resp.Duration
return &result, nil
}
func (c *RAGFlowClient) CheckProviderConnection(cmd *Command) (ResponseIf, error) {
if c.HTTPClient.APIToken == "" && c.HTTPClient.LoginToken == "" {
return nil, fmt.Errorf("API token not set. Please login first")
@@ -1820,7 +1959,6 @@ func (c *RAGFlowClient) CheckProviderConnection(cmd *Command) (ResponseIf, error
}
result.Duration = resp.Duration
return &result, nil
}
// UseModel sets the current model for chat
@@ -1928,14 +2066,14 @@ func (c *RAGFlowClient) AddCustomModel(cmd *Command) (ResponseIf, error) {
resp, err := c.HTTPClient.Request("POST", url, "web", nil, payload)
if err != nil {
return nil, fmt.Errorf("failed to check provider connection: %w", err)
return nil, fmt.Errorf("failed to add custom model: %w", err)
}
if resp.StatusCode != 200 {
return nil, fmt.Errorf("failed to check provider connection: HTTP %d, body: %s", resp.StatusCode, string(resp.Body))
return nil, fmt.Errorf("failed to add custom model: 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("check provider connection failed: invalid JSON (%w)", err)
return nil, fmt.Errorf("add custom model failed: invalid JSON (%w)", err)
}
if result.Code != 0 {
return nil, fmt.Errorf("%s", result.Message)

120
internal/cli/user_parser.go
View File

@@ -2603,6 +2603,126 @@ func (p *Parser) parseStreamCommand() (*Command, error) {
return command, nil
}
func (p *Parser) parseEmbedCommand() (*Command, error) {
p.nextToken() // consume EMBED
if p.curToken.Type != TokenText {
return nil, fmt.Errorf("expected WITH after EMBED")
}
p.nextToken() // consume TEXT
var texts []string
textLoop:
for {
if p.curToken.Type != TokenQuotedString {
break textLoop
}
text, err := p.parseQuotedString()
if err != nil {
return nil, err
}
text = strings.TrimSpace(text)
texts = append(texts, text)
p.nextToken()
}
if p.curToken.Type != TokenWith {
return nil, fmt.Errorf("expected WITH after EMBED")
}
p.nextToken() // consume WITH
compositeModelName, err := p.parseQuotedString()
if err != nil {
return nil, err
}
p.nextToken()
if p.curToken.Type != TokenDimension {
return nil, fmt.Errorf("expected DIMENSION")
}
p.nextToken() // consume WITH
dimension, err := p.parseNumber()
if err != nil {
return nil, err
}
p.nextToken()
cmd := NewCommand("embed_user_text")
cmd.Params["composite_model_name"] = compositeModelName
cmd.Params["texts"] = texts
cmd.Params["dimension"] = dimension
return cmd, nil
}
func (p *Parser) parseRerankCommand() (*Command, error) {
p.nextToken() // consume RERANK
if p.curToken.Type != TokenQuery {
return nil, fmt.Errorf("expected WITH after EMBED")
}
p.nextToken() // consume QUERY
query, err := p.parseQuotedString()
if err != nil {
return nil, err
}
query = strings.TrimSpace(query)
p.nextToken() // consume query
if p.curToken.Type != TokenDocument {
return nil, fmt.Errorf("expected DOCUMENT after query")
}
p.nextToken() // consume DOCUMENT
var documents []string
documentLoop:
for {
if p.curToken.Type != TokenQuotedString {
break documentLoop
}
var document string
document, err = p.parseQuotedString()
if err != nil {
return nil, err
}
document = strings.TrimSpace(document)
documents = append(documents, document)
p.nextToken()
}
if p.curToken.Type != TokenWith {
return nil, fmt.Errorf("expected WITH after EMBED")
}
p.nextToken() // consume WITH
compositeModelName, err := p.parseQuotedString()
if err != nil {
return nil, err
}
p.nextToken()
if p.curToken.Type != TokenTop {
return nil, fmt.Errorf("expected TOP after model")
}
p.nextToken()
topN, err := p.parseNumber()
if err != nil {
return nil, err
}
p.nextToken()
cmd := NewCommand("rarank_user_document")
cmd.Params["composite_model_name"] = compositeModelName
cmd.Params["query"] = query
cmd.Params["documents"] = documents
cmd.Params["top_n"] = topN
return cmd, nil
}
func (p *Parser) parseCheckCommand() (*Command, error) {
p.nextToken() // consume CHECK