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'`
2026-07-21 23:21:04 +08:00 · 2026-06-25 13:32:22 +08:00
parent 824c88423c
commit ced51114f4
6 changed files with 517 additions and 28 deletions
--- 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:
--- a/internal/handler/datasets_search_test.go
+++ 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
+}