From ced51114f484e015975a553d7d843ef77af6f3b6 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 25 Jun 2026 13:32:22 +0800 Subject: [PATCH] feat(go-api): add dataset search endpoint (#16304) ### What problem does this PR solve? - added the new dataset search route and handler - reused the existing shared SearchDatasets service by adapting single-dataset requests into dataset_ids=[dataset_id] - aligned handler error responses with Python behavior for argument/data errors - aligned key service error messages such as invalid search_id and mixed embedding models - added focused handler and service tests for request mapping and error behavior ### Tests: `/usr/local/go/bin/go test ./internal/service -run 'TestSearchDatasetRequestToSearchDatasetsRequest|TestDatasetServiceSearchDatasets'` `/usr/local/go/bin/go test ./internal/handler -run 'TestDatasetsHandlerSearchDataset'` --- internal/handler/datasets.go | 148 +++++++++-- internal/handler/datasets_search_test.go | 251 ++++++++++++++++++ internal/router/datasets_search_route_test.go | 33 +++ internal/router/router.go | 1 + internal/service/dataset.go | 62 ++++- internal/service/dataset_search_test.go | 50 ++++ 6 files changed, 517 insertions(+), 28 deletions(-) create mode 100644 internal/handler/datasets_search_test.go create mode 100644 internal/router/datasets_search_route_test.go create mode 100644 internal/service/dataset_search_test.go diff --git a/internal/handler/datasets.go b/internal/handler/datasets.go index 95a0447a49..64bbb417ba 100644 --- a/internal/handler/datasets.go +++ b/internal/handler/datasets.go @@ -34,8 +34,18 @@ import ( // DatasetsHandler handles the RESTful dataset endpoints. type DatasetsHandler struct { - datasetsService *service.DatasetService - metadataService *service.MetadataService + datasetsService *service.DatasetService + metadataService *service.MetadataService + searchDatasetsService searchDatasetsService + searchDatasetService searchDatasetService +} + +type searchDatasetsService interface { + SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) +} + +type searchDatasetService interface { + SearchDataset(datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) } type listDatasetsExt struct { @@ -46,10 +56,15 @@ type listDatasetsExt struct { // NewDatasetsHandler creates a new datasets handler. func NewDatasetsHandler(datasetsService *service.DatasetService, metadataService *service.MetadataService) *DatasetsHandler { - return &DatasetsHandler{ + h := &DatasetsHandler{ datasetsService: datasetsService, metadataService: metadataService, } + if datasetsService != nil { + h.searchDatasetsService = datasetsService + h.searchDatasetService = datasetsService + } + return h } // ListDatasets handles GET /api/v1/datasets. @@ -995,10 +1010,7 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) { var req service.SearchDatasetsRequest if err := c.ShouldBindJSON(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{ - "code": 400, - "message": err.Error(), - }) + jsonError(c, common.CodeArgumentError, err.Error()) return } @@ -1019,35 +1031,37 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) { req.UseKG = &defaultUseKG } + req.Question = strings.TrimSpace(req.Question) if req.Question == "" { - c.JSON(http.StatusBadRequest, gin.H{ - "code": 400, - "message": "question is required", - }) + jsonError(c, common.CodeArgumentError, "question is required") return } if req.DatasetIDs == nil { - c.JSON(http.StatusBadRequest, gin.H{ - "code": 400, - "message": "kb_id is required", - }) + jsonError(c, common.CodeArgumentError, "kb_id is required") return } if len(req.DatasetIDs) == 0 { - c.JSON(http.StatusBadRequest, gin.H{ - "code": 400, - "message": "kb_id array cannot be empty", - }) + jsonError(c, common.CodeArgumentError, "kb_id array cannot be empty") + return + } + if err := validateSearchDatasetsRequest(&req); err != nil { + jsonError(c, common.CodeArgumentError, err.Error()) return } - resp, err := h.datasetsService.SearchDatasets(&req, user.ID) + searchService := h.searchDatasetsService + if searchService == nil { + searchService = h.datasetsService + } + if searchService == nil { + jsonError(c, common.CodeDataError, "dataset service is not initialized") + return + } + + resp, err := searchService.SearchDatasets(&req, user.ID) if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ - "code": 500, - "message": err.Error(), - }) + jsonError(c, common.CodeDataError, err.Error()) return } @@ -1057,6 +1071,92 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) { }) } +// SearchDataset searches chunks within a single dataset based on a question. +// @Summary Search Dataset +// @Description Search for relevant chunks within one dataset based on a question +// @Tags datasets +// @Accept json +// @Produce json +// @Param dataset_id path string true "dataset id" +// @Param request body service.SearchDatasetRequest true "search parameters" +// @Success 200 {object} map[string]interface{} +// @Router /api/v1/datasets/{dataset_id}/search [post] +func (h *DatasetsHandler) SearchDataset(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + datasetID := c.Param("dataset_id") + if datasetID == "" { + jsonError(c, common.CodeDataError, "dataset_id is required") + return + } + + var req service.SearchDatasetRequest + if err := c.ShouldBindJSON(&req); err != nil { + jsonError(c, common.CodeArgumentError, err.Error()) + return + } + req.Question = strings.TrimSpace(req.Question) + if req.Question == "" { + jsonError(c, common.CodeArgumentError, "question is required") + return + } + if err := validateSearchDatasetRequest(&req); err != nil { + jsonError(c, common.CodeArgumentError, err.Error()) + return + } + + searchService := h.searchDatasetService + if searchService == nil { + searchService = h.datasetsService + } + if searchService == nil { + jsonError(c, common.CodeDataError, "dataset service is not initialized") + return + } + + resp, err := searchService.SearchDataset(datasetID, user.ID, &req) + if err != nil { + jsonError(c, common.CodeDataError, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": 0, + "data": resp, + }) +} + +func validateSearchDatasetsRequest(req *service.SearchDatasetsRequest) error { + return validateSearchParams(req.Page, req.Size, req.TopK, req.SimilarityThreshold, req.VectorSimilarityWeight) +} + +func validateSearchDatasetRequest(req *service.SearchDatasetRequest) error { + return validateSearchParams(req.Page, req.Size, req.TopK, req.SimilarityThreshold, req.VectorSimilarityWeight) +} + +func validateSearchParams(page, size, topK *int, similarityThreshold, vectorSimilarityWeight *float64) error { + if page != nil && *page < 1 { + return fmt.Errorf("page must be greater than or equal to 1") + } + if size != nil && *size < 1 { + return fmt.Errorf("size must be greater than or equal to 1") + } + if topK != nil && *topK < 1 { + return fmt.Errorf("top_k must be greater than or equal to 1") + } + if similarityThreshold != nil && (*similarityThreshold < 0 || *similarityThreshold > 1) { + return fmt.Errorf("similarity_threshold must be between 0 and 1") + } + if vectorSimilarityWeight != nil && (*vectorSimilarityWeight < 0 || *vectorSimilarityWeight > 1) { + return fmt.Errorf("vector_similarity_weight must be between 0 and 1") + } + return nil +} + func firstStringValue(value interface{}) string { switch v := value.(type) { case string: diff --git a/internal/handler/datasets_search_test.go b/internal/handler/datasets_search_test.go new file mode 100644 index 0000000000..da430ec298 --- /dev/null +++ b/internal/handler/datasets_search_test.go @@ -0,0 +1,251 @@ +package handler + +import ( + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "ragflow/internal/common" + "ragflow/internal/entity" + "ragflow/internal/service" +) + +type fakeSearchDatasetService struct { + datasetID string + userID string + req *service.SearchDatasetRequest + resp *service.SearchDatasetsResponse + err error +} + +func (f *fakeSearchDatasetService) SearchDataset(datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) { + f.datasetID = datasetID + f.userID = userID + f.req = req + return f.resp, f.err +} + +type fakeSearchDatasetsService struct { + userID string + req *service.SearchDatasetsRequest + resp *service.SearchDatasetsResponse + err error +} + +func (f *fakeSearchDatasetsService) SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) { + f.userID = userID + f.req = req + return f.resp, f.err +} + +func TestDatasetsHandlerSearchDataset(t *testing.T) { + gin.SetMode(gin.TestMode) + fake := &fakeSearchDatasetService{resp: &service.SearchDatasetsResponse{Total: 1}} + h := &DatasetsHandler{searchDatasetService: fake} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/ds-1/search", strings.NewReader(`{"question":"hello","doc_ids":["doc-1"],"top_k":7}`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.SearchDataset(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if fake.datasetID != "ds-1" || fake.userID != "user-1" { + t.Fatalf("call args = (%q,%q), want (ds-1,user-1)", fake.datasetID, fake.userID) + } + if fake.req == nil || fake.req.Question != "hello" || len(fake.req.DocIDs) != 1 || fake.req.DocIDs[0] != "doc-1" { + t.Fatalf("request = %#v", fake.req) + } + if len(fake.req.ToSearchDatasetsRequest("ds-1").DatasetIDs) != 1 { + t.Fatal("request conversion failed") + } + + var body map[string]interface{} + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v body=%s", err, rec.Body.String()) + } + if body["code"] != float64(common.CodeSuccess) { + t.Fatalf("code=%v want=%d", body["code"], common.CodeSuccess) + } +} + +func TestDatasetsHandlerSearchDatasetValidatesQuestion(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &DatasetsHandler{} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/ds-1/search", strings.NewReader(`{"question":" "}`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.SearchDataset(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeArgumentError) { + t.Fatalf("code=%v want=%d body=%s", body["code"], common.CodeArgumentError, rec.Body.String()) + } +} + +func TestDatasetsHandlerSearchDatasetPropagatesServiceError(t *testing.T) { + gin.SetMode(gin.TestMode) + fake := &fakeSearchDatasetService{err: errors.New("boom")} + h := &DatasetsHandler{searchDatasetService: fake} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/ds-1/search", strings.NewReader(`{"question":"hello"}`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.SearchDataset(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeDataError) || body["message"] != "boom" { + t.Fatalf("response=%v want code=%d message=boom", body, common.CodeDataError) + } +} + +func TestDatasetsHandlerSearchDatasetParseErrorUsesArgumentEnvelope(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &DatasetsHandler{} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/ds-1/search", strings.NewReader(`{"question"`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.SearchDataset(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeArgumentError) { + t.Fatalf("code=%v want=%d body=%s", body["code"], common.CodeArgumentError, rec.Body.String()) + } +} + +func TestDatasetsHandlerSearchDatasetsSuccess(t *testing.T) { + gin.SetMode(gin.TestMode) + fake := &fakeSearchDatasetsService{resp: &service.SearchDatasetsResponse{Total: 2}} + h := &DatasetsHandler{searchDatasetsService: fake} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/search", strings.NewReader(`{"question":" hello ","dataset_ids":["ds-1"],"top_k":7}`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + + h.SearchDatasets(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if fake.userID != "user-1" || fake.req == nil || fake.req.Question != "hello" || len(fake.req.DatasetIDs) != 1 || fake.req.DatasetIDs[0] != "ds-1" { + t.Fatalf("call args userID=%q req=%#v", fake.userID, fake.req) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeSuccess) { + t.Fatalf("code=%v want=%d body=%s", body["code"], common.CodeSuccess, rec.Body.String()) + } +} + +func TestDatasetsHandlerSearchDatasetsValidationErrorsUseArgumentEnvelope(t *testing.T) { + tests := []struct { + name string + body string + }{ + {name: "parse error", body: `{"question"`}, + {name: "missing question", body: `{"dataset_ids":["ds-1"]}`}, + {name: "blank question", body: `{"question":" ","dataset_ids":["ds-1"]}`}, + {name: "missing dataset ids", body: `{"question":"hello"}`}, + {name: "empty dataset ids", body: `{"question":"hello","dataset_ids":[]}`}, + {name: "invalid top k", body: `{"question":"hello","dataset_ids":["ds-1"],"top_k":0}`}, + {name: "invalid threshold", body: `{"question":"hello","dataset_ids":["ds-1"],"similarity_threshold":1.1}`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &DatasetsHandler{} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/search", strings.NewReader(tt.body)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + + h.SearchDatasets(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeArgumentError) { + t.Fatalf("code=%v want=%d body=%s", body["code"], common.CodeArgumentError, rec.Body.String()) + } + }) + } +} + +func TestDatasetsHandlerSearchDatasetsPropagatesServiceError(t *testing.T) { + gin.SetMode(gin.TestMode) + fake := &fakeSearchDatasetsService{err: errors.New("boom")} + h := &DatasetsHandler{searchDatasetsService: fake} + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/search", strings.NewReader(`{"question":"hello","dataset_ids":["ds-1"]}`)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(rec) + c.Request = req + c.Set("user", &entity.User{ID: "user-1"}) + + h.SearchDatasets(c) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + body := decodeSearchResponse(t, rec) + if body["code"] != float64(common.CodeDataError) || body["message"] != "boom" { + t.Fatalf("response=%v want code=%d message=boom", body, common.CodeDataError) + } +} + +func decodeSearchResponse(t *testing.T, rec *httptest.ResponseRecorder) map[string]interface{} { + t.Helper() + if !json.Valid(rec.Body.Bytes()) { + t.Fatalf("response is not valid json: %s", rec.Body.String()) + } + var body map[string]interface{} + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v body=%s", err, rec.Body.String()) + } + return body +} diff --git a/internal/router/datasets_search_route_test.go b/internal/router/datasets_search_route_test.go new file mode 100644 index 0000000000..7920b27ba1 --- /dev/null +++ b/internal/router/datasets_search_route_test.go @@ -0,0 +1,33 @@ +package router + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "ragflow/internal/handler" +) + +func TestRouterSetupRegistersSearchDatasetRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + engine := gin.New() + r := &Router{ + authHandler: handler.NewAuthHandler(), + datasetsHandler: handler.NewDatasetsHandler(nil, nil), + } + r.Setup(engine) + + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/dataset-1/search", nil) + engine.ServeHTTP(resp, req) + + if resp.Code == http.StatusNotFound { + t.Fatalf("POST /api/v1/datasets/:dataset_id/search returned 404; SearchDataset route is not registered") + } + if resp.Code != http.StatusUnauthorized { + t.Fatalf("status=%d body=%s; want auth middleware to handle registered SearchDataset route", resp.Code, resp.Body.String()) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index d9fd3974f8..7c4c9b6678 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -291,6 +291,7 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.POST("", r.datasetsHandler.CreateDataset) datasets.DELETE("", r.datasetsHandler.DeleteDatasets) datasets.POST("/search", r.datasetsHandler.SearchDatasets) + datasets.POST("/:dataset_id/search", r.datasetsHandler.SearchDataset) datasets.GET("/metadata/flattened", r.datasetsHandler.ListMetadataFlattened) datasets.GET("/:dataset_id/metadata/summary", r.documentHandler.MetadataSummaryByDataset) diff --git a/internal/service/dataset.go b/internal/service/dataset.go index 736f77ca24..3304be203d 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -1704,6 +1704,54 @@ type SearchDatasetsResponse struct { Total int64 `json:"total"` } +// SearchDatasetRequest is the request structure for searching chunks within one dataset. +type SearchDatasetRequest struct { + Question string `json:"question"` + Page *int `json:"page,omitempty"` + Size *int `json:"size,omitempty"` + DocIDs []string `json:"doc_ids,omitempty"` + UseKG *bool `json:"use_kg,omitempty"` + TopK *int `json:"top_k,omitempty"` + CrossLanguages []string `json:"cross_languages,omitempty"` + SearchID *string `json:"search_id,omitempty"` + MetadataFilter map[string]interface{} `json:"meta_data_filter,omitempty"` + RerankID *string `json:"rerank_id,omitempty"` + Keyword *bool `json:"keyword,omitempty"` + SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"` + VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"` +} + +// ToSearchDatasetsRequest converts a single-dataset search request into the multi-dataset form. +func (req *SearchDatasetRequest) ToSearchDatasetsRequest(datasetID string) *SearchDatasetsRequest { + if req == nil { + return &SearchDatasetsRequest{DatasetIDs: []string{datasetID}} + } + return &SearchDatasetsRequest{ + DatasetIDs: []string{datasetID}, + Question: req.Question, + Page: req.Page, + Size: req.Size, + DocIDs: req.DocIDs, + UseKG: req.UseKG, + TopK: req.TopK, + CrossLanguages: req.CrossLanguages, + SearchID: req.SearchID, + MetadataFilter: req.MetadataFilter, + RerankID: req.RerankID, + Keyword: req.Keyword, + SimilarityThreshold: req.SimilarityThreshold, + VectorSimilarityWeight: req.VectorSimilarityWeight, + } +} + +// SearchDataset searches chunks within one knowledge base based on a question. +func (s *DatasetService) SearchDataset(datasetID, userID string, req *SearchDatasetRequest) (*SearchDatasetsResponse, error) { + if datasetID == "" { + return nil, fmt.Errorf("dataset_id is required") + } + return s.SearchDatasets(req.ToSearchDatasetsRequest(datasetID), userID) +} + // SearchDatasets searches chunks across one or more knowledge bases based on a question. // It retrieves relevant chunks using embedding and optional reranking, applying filters, // cross-language translation, and keyword extraction as configured. @@ -1813,7 +1861,7 @@ func (s *DatasetService) SearchDatasets(req *SearchDatasetsRequest, userID strin firstEmbdID := kbRecords[0].EmbdID for i := 1; i < len(kbRecords); i++ { if kbRecords[i].EmbdID != firstEmbdID { - return nil, fmt.Errorf("cannot retrieve across datasets with different embedding models") + return nil, fmt.Errorf("Datasets use different embedding models.") } } } @@ -1821,9 +1869,14 @@ func (s *DatasetService) SearchDatasets(req *SearchDatasetsRequest, userID strin // Override request fields with values from saved search config (if search_id is provided) var chatID string if searchID != "" { + if s.searchService == nil { + common.Warn("Search service is not initialized for search_id", zap.String("searchID", searchID)) + return nil, fmt.Errorf("Invalid search_id") + } searchDetail, err := s.searchService.GetDetail(searchID) - if err != nil { - common.Warn("Failed to get search detail for search_id, proceeding without it", zap.String("searchID", searchID), zap.Error(err)) + if err != nil || searchDetail == nil || len(searchDetail) == 0 { + common.Warn("Invalid search_id", zap.String("searchID", searchID), zap.Error(err)) + return nil, fmt.Errorf("Invalid search_id") } else if searchConfig, ok := searchDetail["search_config"].(map[string]interface{}); ok && searchConfig != nil { if scMetadataFilter, ok := searchConfig["meta_data_filter"].(map[string]interface{}); ok { metadataFilter = scMetadataFilter @@ -1874,7 +1927,8 @@ func (s *DatasetService) SearchDatasets(req *SearchDatasetsRequest, userID strin zap.String("chatID", chatID), zap.Bool("useKG", useKG)) } else { - common.Warn("No search_config found in search detail", zap.String("searchID", searchID)) + common.Warn("Invalid search_id: search_config missing or invalid", zap.String("searchID", searchID)) + return nil, fmt.Errorf("Invalid search_id") } } diff --git a/internal/service/dataset_search_test.go b/internal/service/dataset_search_test.go new file mode 100644 index 0000000000..524ca366c1 --- /dev/null +++ b/internal/service/dataset_search_test.go @@ -0,0 +1,50 @@ +package service + +import "testing" + +func TestSearchDatasetRequestToSearchDatasetsRequest(t *testing.T) { + page := 2 + size := 15 + topK := 128 + useKG := true + keyword := true + similarityThreshold := 0.42 + vectorSimilarityWeight := 0.8 + searchID := "search-1" + rerankID := "rerank-1" + req := &SearchDatasetRequest{ + Question: "hello world", + Page: &page, + Size: &size, + DocIDs: []string{"doc-1", "doc-2"}, + UseKG: &useKG, + TopK: &topK, + CrossLanguages: []string{"en", "zh"}, + SearchID: &searchID, + MetadataFilter: map[string]interface{}{"method": "manual"}, + RerankID: &rerankID, + Keyword: &keyword, + SimilarityThreshold: &similarityThreshold, + VectorSimilarityWeight: &vectorSimilarityWeight, + } + + converted := req.ToSearchDatasetsRequest("dataset-1") + if len(converted.DatasetIDs) != 1 || converted.DatasetIDs[0] != "dataset-1" { + t.Fatalf("dataset_ids=%v want [dataset-1]", converted.DatasetIDs) + } + if converted.Question != req.Question || converted.Page != req.Page || converted.Size != req.Size { + t.Fatalf("converted request did not preserve pagination/question fields: %#v", converted) + } + if len(converted.DocIDs) != 2 || converted.DocIDs[0] != "doc-1" || converted.DocIDs[1] != "doc-2" { + t.Fatalf("doc_ids=%v want [doc-1 doc-2]", converted.DocIDs) + } + if converted.UseKG != req.UseKG || converted.TopK != req.TopK || converted.SearchID != req.SearchID { + t.Fatalf("converted request did not preserve optional fields: %#v", converted) + } + if converted.MetadataFilter["method"] != "manual" || converted.RerankID != req.RerankID || converted.Keyword != req.Keyword { + t.Fatalf("converted request did not preserve search config fields: %#v", converted) + } + if converted.SimilarityThreshold != req.SimilarityThreshold || converted.VectorSimilarityWeight != req.VectorSimilarityWeight { + t.Fatalf("converted request did not preserve thresholds: %#v", converted) + } +}