mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
Go CLI: embed and rerank (#14735)
### 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:
@@ -894,3 +894,156 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) {
|
||||
"answer": response.Answer,
|
||||
})
|
||||
}
|
||||
|
||||
type EmbedTextRequest struct {
|
||||
ProviderName *string `json:"provider_name"`
|
||||
InstanceName *string `json:"instance_name"`
|
||||
ModelName *string `json:"model_name"`
|
||||
Texts []string `json:"texts"`
|
||||
Dimension int `json:"dimension"`
|
||||
}
|
||||
|
||||
func (h *ProviderHandler) EmbedText(c *gin.Context) {
|
||||
var req EmbedTextRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
println("JSON bind error: %v (type: %T)", err, err)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeBadRequest,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.ProviderName == nil || *req.ProviderName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Provider name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.InstanceName == nil || *req.InstanceName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Instance name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.ModelName == nil || *req.ModelName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Model name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
userID := c.GetString("user_id")
|
||||
|
||||
apiConfig := models.APIConfig{
|
||||
ApiKey: nil,
|
||||
Region: nil,
|
||||
}
|
||||
|
||||
embeddingConfig := models.EmbeddingConfig{
|
||||
Dimension: req.Dimension,
|
||||
}
|
||||
|
||||
// Non-stream response
|
||||
var response *models.EmbeddingResponse
|
||||
var errorCode common.ErrorCode
|
||||
var err error
|
||||
|
||||
response, errorCode, err = h.modelProviderService.EmbedText(*req.ProviderName, *req.InstanceName, *req.ModelName, userID, req.Texts, &apiConfig, &embeddingConfig)
|
||||
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": errorCode,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 0,
|
||||
"data": response.Data,
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
|
||||
type RerankDocumentRequest struct {
|
||||
ProviderName *string `json:"provider_name"`
|
||||
InstanceName *string `json:"instance_name"`
|
||||
ModelName *string `json:"model_name"`
|
||||
Query string `json:"query"`
|
||||
Documents []string `json:"documents"`
|
||||
TopN int `json:"top_n"`
|
||||
}
|
||||
|
||||
func (h *ProviderHandler) RerankDocument(c *gin.Context) {
|
||||
var req RerankDocumentRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
println("JSON bind error: %v (type: %T)", err, err)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeBadRequest,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.ProviderName == nil || *req.ProviderName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Provider name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.InstanceName == nil || *req.InstanceName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Instance name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if req.ModelName == nil || *req.ModelName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "Model name is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
userID := c.GetString("user_id")
|
||||
|
||||
apiConfig := models.APIConfig{
|
||||
ApiKey: nil,
|
||||
Region: nil,
|
||||
}
|
||||
|
||||
rerankConfig := models.RerankConfig{
|
||||
TopN: req.TopN,
|
||||
}
|
||||
|
||||
// Non-stream response
|
||||
var response *models.RerankResponse
|
||||
var errorCode common.ErrorCode
|
||||
var err error
|
||||
|
||||
response, errorCode, err = h.modelProviderService.RerankDocument(*req.ProviderName, *req.InstanceName, *req.ModelName, userID, req.Query, req.Documents, &apiConfig, &rerankConfig)
|
||||
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": errorCode,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 0,
|
||||
"data": response.Data,
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user