From 2f2d1569e6c5a36800c1f44211e985e236620441 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Mon, 11 May 2026 20:19:08 +0800 Subject: [PATCH] Go: fix retrieval test error (#14794) ### 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 --- internal/entity/models/zhipu-ai.go | 2 +- internal/handler/chunk.go | 48 ++---------------------------- internal/service/chunk.go | 34 +++++++-------------- internal/service/model_service.go | 2 +- internal/service/nlp/retrieval.go | 5 +++- 5 files changed, 19 insertions(+), 72 deletions(-) diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index adccae702..e4041614f 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -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 } diff --git a/internal/handler/chunk.go b/internal/handler/chunk.go index 207edfee4..8159ce059 100644 --- a/internal/handler/chunk.go +++ b/internal/handler/chunk.go @@ -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 } diff --git a/internal/service/chunk.go b/internal/service/chunk.go index c2ce08d4e..4930ae5ad 100644 --- a/internal/service/chunk.go +++ b/internal/service/chunk.go @@ -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, diff --git a/internal/service/model_service.go b/internal/service/model_service.go index a32daa7ee..5ac249519 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -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 } diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index a3a2e8deb..4cfd197f8 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -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 }