zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-06 19:38:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: fix retrieval test error (#14794)

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

1. Add region check in zhipu-ai embed method
2. Fix retrieval test

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-05-11 20:19:08 +08:00
committed by GitHub
parent 3e90d303e0
commit 2f2d1569e6
5 changed files with 19 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/entity/models/zhipu-ai.go
View File

@@ -396,7 +396,7 @@ func (z *ZhipuAIModel) Embed(modelName *string, texts []string, apiConfig *APICo
}
var region = "default"
if apiConfig.Region != nil {
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}

48
internal/handler/chunk.go
View File

@@ -92,7 +92,7 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) {
})
return
}
if req.KbID == nil {
if req.Datasets == nil {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id is required",
@@ -100,52 +100,10 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) {
return
}
// Validate kb_id type: string or []string
switch v := req.KbID.(type) {
case string:
if v == "" {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id cannot be empty string",
})
return
}
case []interface{}:
// Convert to []string
var kbIDs []string
for _, item := range v {
if str, ok := item.(string); ok && str != "" {
kbIDs = append(kbIDs, str)
} else {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id array must contain non-empty strings",
})
return
}
}
if len(kbIDs) == 0 {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id array cannot be empty",
})
return
}
// Convert back to interface{} for service
req.KbID = kbIDs
case []string:
// Already correct type
if len(v) == 0 {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id array cannot be empty",
})
return
}
default:
if len(req.Datasets) == 0 {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "kb_id must be string or array of strings",
"message": "kb_id array cannot be empty",
})
return
}

34
internal/service/chunk.go
View File

@@ -63,7 +63,7 @@ func NewChunkService() *ChunkService {
// RetrievalTestRequest retrieval test request
type RetrievalTestRequest struct {
KbID interface{} `json:"kb_id" binding:"required"` // string or []string
Datasets []string `json:"dataset_ids" binding:"required"` // string or []string
Question string `json:"question" binding:"required"`
Page *int `json:"page,omitempty"`
Size *int `json:"size,omitempty"`
@@ -105,7 +105,7 @@ type RetrievalTestResponse struct {
// 7. knowledge graph retrieval (not implemented)
// 8. Apply retrieval by children to group child chunks under parent chunks
func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) {
common.Info("RetrievalTest started", zap.String("userID", userID), zap.Any("kbID", req.KbID), zap.String("question", req.Question))
common.Info("RetrievalTest started", zap.String("userID", userID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question))
common.Debug(fmt.Sprintf("RetrievalTest request:\n"+
" kbID=%v\n"+
@@ -120,7 +120,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
" rerankID=%v\n"+
" keyword=%v\n"+
" similarityThreshold=%v, vectorSimilarityWeight=%v",
req.KbID, req.Question,
req.Datasets, req.Question,
ptrString(req.Page), ptrString(req.Size), req.DocIDs,
ptrString(req.UseKG), ptrString(req.TopK), req.CrossLanguages, ptrString(req.SearchID),
req.Filter,
@@ -134,20 +134,6 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
ctx := context.Background()
// Determine kb_id list and check permission for each kb_id
var kbIDs []string
switch v := req.KbID.(type) {
case string:
kbIDs = []string{v}
case []string:
kbIDs = v
default:
return nil, fmt.Errorf("kb_id must be string or array of strings")
}
if len(kbIDs) == 0 {
return nil, fmt.Errorf("kb_id cannot be empty")
}
tenants, err := s.userTenantDAO.GetByUserID(userID)
if err != nil {
return nil, fmt.Errorf("failed to get user tenants: %w", err)
@@ -159,13 +145,13 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
var tenantIDs []string
var kbRecords []*entity.Knowledgebase
for _, kbID := range kbIDs {
for _, datasetID := range req.Datasets {
found := false
for _, tenant := range tenants {
kb, err := s.kbDAO.GetByIDAndTenantID(kbID, tenant.TenantID)
kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenant.TenantID)
if err == nil && kb != nil {
common.Debug("Found knowledge base in database",
zap.String("kbID", kbID),
zap.String("datasetID", datasetID),
zap.String("tenantID", tenant.TenantID),
zap.String("kbName", kb.Name),
zap.String("embdID", kb.EmbdID))
@@ -227,7 +213,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
}
}
// If no chatID from search_config, or chatModel not found, use tenant default
// If no chatID from search_config, or chatModel not found, use tenant default
if chatModelForFilter == nil {
tenantSvc := NewTenantService()
modelName, err := tenantSvc.GetDefaultModelName(tenantIDs[0], entity.ModelTypeChat)
@@ -253,7 +239,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
if filter != nil {
// Get flattened metadata
metadataSvc := NewMetadataService()
flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(kbIDs)
flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(req.Datasets)
if err != nil {
common.Warn("Failed to get flatted metadata", zap.Error(err))
} else {
@@ -393,7 +379,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
retrievalReq := &nlp.RetrievalRequest{
TenantIDs: tenantIDs,
Question: modifiedQuestion,
KbIDs: kbIDs,
KbIDs: req.Datasets,
DocIDs: docIDs,
Page: getPageNum(req.Page, 1),
PageSize: getPageSize(req.Size, 30),
@@ -427,7 +413,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (
delete(filteredChunks[i], "vector")
}
common.Info("RetrievalTest completed", zap.String("userID", userID), zap.Any("kbID", req.KbID), zap.String("question", req.Question), zap.Int64("chunkCount", int64(len(filteredChunks))))
common.Info("RetrievalTest completed", zap.String("userID", userID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question), zap.Int64("chunkCount", int64(len(filteredChunks))))
return &RetrievalTestResponse{
Chunks: filteredChunks,

2
internal/service/model_service.go
View File

@@ -101,7 +101,7 @@ func (m *ModelProviderService) AddModelProvider(providerName, userID string) (co
tenantModelProvider.UpdateDate = &nowDate
err = m.modelProviderDAO.Create(tenantModelProvider)
if err != nil {
return common.CodeServerError, errors.New("fail to create model provider")
return common.CodeServerError, fmt.Errorf("fail to create model provider: %s", err.Error())
}
return common.CodeSuccess, nil
}

5
internal/service/nlp/retrieval.go
View File

@@ -607,7 +607,10 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque
// GetVector computes query vector and returns MatchDenseExpr for hybrid search
func (s *RetrievalService) GetVector(txt string, embModel *models.EmbeddingModel, topk int, similarity float64) (*types.MatchDenseExpr, error) {
embeddings, err := embModel.ModelDriver.Embed(embModel.ModelName, []string{txt}, embModel.APIConfig, nil)
embeddingConfig := &models.EmbeddingConfig{
Dimension: 0,
}
embeddings, err := embModel.ModelDriver.Embed(embModel.ModelName, []string{txt}, embModel.APIConfig, embeddingConfig)
if err != nil {
return nil, err
}