mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-24 17:10:12 +08:00
feat: implement POST /api/v1/searchbots/retrieval_test (#15710)
## What problem does this PR solve? Implements `POST /api/v1/searchbots/retrieval_test` in the Go API server, aligning with the Python `bot_api.py` counterpart. Also applies security hardening and consistency fixes discovered during CTO-level code review: - **Missing endpoint**: `retrieval_test` was not available in Go, requiring Python fallback - **Security**: Both `chunkHandler` and `searchBotHandler` leaked `err.Error()` to API consumers - **Python alignment**: Default values, empty question handling, and `top_k <= 0` validation differed from Python behavior - **Test gaps**: `chunkHandler.RetrievalTest` had zero unit tests; several edge cases uncovered ## Type of change - [x] New Feature (non-breaking change which adds functionality) - [x] Bug Fix (non-breaking change which fixes an issue) - [x] Refactoring ## Summary ### New Endpoint - `POST /api/v1/searchbots/retrieval_test` — retrieval test with full field support (page, size, top_k, use_kg, cross_languages, keyword, similarity_threshold, vector_similarity_weight) ### New Type - `common.StringSlice` — JSON type that accepts both `"kb1"` and `["kb1", "kb2"]`, matching Python API flexibility ### Security - Both `searchBotHandler` and `chunkHandler` now use `common.Warn()` + generic error messages instead of leaking `err.Error()` to API consumers - All error responses include consistent `"data": nil` shape - `chunkHandler.RetrievalTest` uses interface-based DI (`chunkService`) to enable testability ### Python Alignment - Handler-level defaults align with Python `bot_api.py` (page=1, size=30, top_k=1024, similarity_threshold=0.0, vector_similarity_weight=0.3) - `top_k <= 0` validation matching Python behavior - Empty/whitespace question returns 200 + empty result (matches `chunk_api.py`) - `chunkHandler` `Datasets` field uses `common.StringSlice` for string-or-array flexibility ### Refactoring - `ChunkServiceIface` → `ChunkRetriever`, `chunkSvcIface` → `chunkService` (Go-conventional naming) - Extracted `applyRetrievalDefaults`, `toRetrievalServiceRequest` from handler body - Regex moved to package-level var in `parseRelatedQuestions` - `service.RetrievalTestRequest.Datasets` type changed to `common.StringSlice` - `chunkHandler` now uses consumer-side interface for DI ### Tests - 37 unit tests across both handlers: auth, validation, defaults, StringSlice edge cases, empty/whitespace KbID, service errors, JSON format, `top_k <= 0`, field mapping verification ## Files Changed | File | Change | |------|--------| | `cmd/server_main.go` | Wire new handler + chunkService + difyRetrievalHandler | | `internal/common/json_types.go` | New StringSlice type | | `internal/common/json_types_test.go` | StringSlice tests | | `internal/handler/chunk.go` | Interface-based DI, security, Python alignment, defaults | | `internal/handler/chunk_test.go` | New — 9 comprehensive tests | | `internal/handler/searchbot.go` | New endpoint + refactoring + `top_k <= 0` validation | | `internal/handler/searchbot_test.go` | 18 tests covering all edge cases | | `internal/router/router.go` | Register new route + difyRetrievalHandler | | `internal/service/chunk.go` | Datasets type → StringSlice, Question binding relaxed | 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -19,28 +19,155 @@ package handler
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
"ragflow/internal/common"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"ragflow/internal/service"
|
||||
)
|
||||
|
||||
// chunkService is the consumer-side interface for ChunkHandler's service dependency.
|
||||
type chunkService interface {
|
||||
RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
|
||||
Get(req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error)
|
||||
List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
|
||||
UpdateChunk(req *service.UpdateChunkRequest, userID string) error
|
||||
RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error)
|
||||
}
|
||||
|
||||
// ChunkHandler chunk handler
|
||||
type ChunkHandler struct {
|
||||
chunkService *service.ChunkService
|
||||
chunkService chunkService
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
// NewChunkHandler create chunk handler
|
||||
func NewChunkHandler(chunkService *service.ChunkService, userService *service.UserService) *ChunkHandler {
|
||||
func NewChunkHandler(chunkService chunkService, userService *service.UserService) *ChunkHandler {
|
||||
return &ChunkHandler{
|
||||
chunkService: chunkService,
|
||||
userService: userService,
|
||||
}
|
||||
}
|
||||
|
||||
// Get retrieves a chunk by ID
|
||||
// RetrievalTest performs retrieval test for chunks
|
||||
// @Summary Retrieval Test
|
||||
// @Description Test retrieval of chunks based on question and knowledge base
|
||||
// @Tags chunks
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param request body service.RetrievalTestRequest true "retrieval test parameters"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/datasets/search [post]
|
||||
func (h *ChunkHandler) RetrievalTest(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
jsonError(c, errorCode, errorMessage)
|
||||
return
|
||||
}
|
||||
|
||||
// Bind JSON request
|
||||
var req service.RetrievalTestRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": common.CodeArgumentError,
|
||||
"data": nil,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Set default values for optional parameters
|
||||
if req.Page == nil {
|
||||
defaultPage := 1
|
||||
req.Page = &defaultPage
|
||||
}
|
||||
if req.Size == nil {
|
||||
defaultSize := 30
|
||||
req.Size = &defaultSize
|
||||
}
|
||||
if req.TopK == nil {
|
||||
defaultTopK := 1024
|
||||
req.TopK = &defaultTopK
|
||||
}
|
||||
if req.UseKG == nil {
|
||||
defaultUseKG := false
|
||||
req.UseKG = &defaultUseKG
|
||||
}
|
||||
|
||||
// Strip and validate question. Matching Python chunk_api.py which returns
|
||||
// an empty result for blank questions rather than an error.
|
||||
if strings.TrimSpace(req.Question) == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": int(common.CodeSuccess),
|
||||
"data": &service.RetrievalTestResponse{
|
||||
Chunks: []map[string]interface{}{},
|
||||
DocAggs: []map[string]interface{}{},
|
||||
Total: 0,
|
||||
},
|
||||
"message": "success",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Validate required fields
|
||||
if req.Datasets == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": common.CodeArgumentError,
|
||||
"data": nil,
|
||||
"message": "kb_id is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.Datasets) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": common.CodeArgumentError,
|
||||
"data": nil,
|
||||
"message": "kb_id array cannot be empty",
|
||||
})
|
||||
return
|
||||
}
|
||||
if req.TopK != nil && *req.TopK <= 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": common.CodeArgumentError,
|
||||
"data": nil,
|
||||
"message": "top_k must be greater than 0",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Call service with user ID for permission checks
|
||||
resp, err := h.chunkService.RetrievalTest(&req, user.ID)
|
||||
if err != nil {
|
||||
common.Warn("dataset search failed", zap.String("error", err.Error()))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"code": common.CodeServerError,
|
||||
"data": nil,
|
||||
"message": "dataset search failed",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": int(common.CodeSuccess),
|
||||
"data": resp,
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
|
||||
// Get retrieves a chunk by ID.
|
||||
// @Summary Get Chunk
|
||||
// @Description Retrieve a single chunk by its ID.
|
||||
// @Tags chunks
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param dataset_id path string true "Dataset ID"
|
||||
// @Param document_id path string true "Document ID"
|
||||
// @Param chunk_id path string true "Chunk ID"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id} [get]
|
||||
func (h *ChunkHandler) Get(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
@@ -49,20 +176,16 @@ func (h *ChunkHandler) Get(c *gin.Context) {
|
||||
}
|
||||
|
||||
chunkID := c.Param("chunk_id")
|
||||
datasetID := c.Param("dataset_id")
|
||||
documentID := c.Param("document_id")
|
||||
if chunkID == "" || datasetID == "" || documentID == "" {
|
||||
if chunkID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": "dataset_id, document_id and chunk_id are required",
|
||||
"message": "chunk_id is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
req := &service.GetChunkRequest{
|
||||
ChunkID: chunkID,
|
||||
DocumentID: documentID,
|
||||
DatasetID: datasetID,
|
||||
ChunkID: chunkID,
|
||||
}
|
||||
|
||||
resp, err := h.chunkService.Get(req, user.ID)
|
||||
@@ -81,7 +204,15 @@ func (h *ChunkHandler) Get(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// List retrieves chunks for a document
|
||||
// List retrieves chunks for a document.
|
||||
// @Summary List Chunks
|
||||
// @Description Retrieve paginated chunks for a document with optional filtering.
|
||||
// @Tags chunks
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param request body service.ListChunksRequest true "List chunks parameters"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/chunk/list [post]
|
||||
func (h *ChunkHandler) List(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
|
||||
292
internal/handler/chunk_test.go
Normal file
292
internal/handler/chunk_test.go
Normal file
@@ -0,0 +1,292 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// mockChunkSvc implements chunkSvcIface for testing ChunkHandler.
|
||||
// Only the methods actually called by the test are set; others panic.
|
||||
type mockChunkSvc struct {
|
||||
retrievalTestFn func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
|
||||
}
|
||||
|
||||
func (m *mockChunkSvc) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
|
||||
if m.retrievalTestFn != nil {
|
||||
return m.retrievalTestFn(req, userID)
|
||||
}
|
||||
return &service.RetrievalTestResponse{
|
||||
Chunks: []map[string]interface{}{{"docnm_kwd": "test", "content_with_weight": "content"}},
|
||||
Total: 1,
|
||||
}, nil
|
||||
}
|
||||
func (m *mockChunkSvc) Get(*service.GetChunkRequest, string) (*service.GetChunkResponse, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *mockChunkSvc) List(*service.ListChunksRequest, string) (*service.ListChunksResponse, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *mockChunkSvc) UpdateChunk(*service.UpdateChunkRequest, string) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *mockChunkSvc) RemoveChunks(*service.RemoveChunksRequest, string) (int64, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
|
||||
func setupChunkRetrievalTest(userID string) (*gin.Engine, *mockChunkSvc) {
|
||||
mock := &mockChunkSvc{}
|
||||
h := &ChunkHandler{chunkService: mock}
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set("user", &entity.User{ID: userID})
|
||||
})
|
||||
r.POST("/api/v1/datasets/search", h.RetrievalTest)
|
||||
return r, mock
|
||||
}
|
||||
|
||||
func setupChunkRetrievalTestNoAuth() *gin.Engine {
|
||||
// Returns a router without the user middleware — used for error-path
|
||||
// tests that don't call the service.
|
||||
h := &ChunkHandler{}
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.POST("/api/v1/datasets/search", h.RetrievalTest)
|
||||
return r
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_EmptyQuestion(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": ""}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Errorf("expected code 0, got %v: %q", resp["code"], resp["message"])
|
||||
}
|
||||
data, ok := resp["data"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected data to be object, got %T", resp["data"])
|
||||
}
|
||||
chunks, _ := data["chunks"].([]interface{})
|
||||
if chunks == nil || len(chunks) != 0 {
|
||||
t.Errorf("expected empty chunks array, got %v", chunks)
|
||||
}
|
||||
if total, _ := data["total"].(float64); total != 0 {
|
||||
t.Errorf("expected total 0, got %v", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_WhitespaceQuestion(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": " "}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Errorf("expected code 0, got %v", resp["code"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_TopKZero(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": "test", "top_k": 0}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "top_k must be greater than 0" {
|
||||
t.Errorf("expected 'top_k must be greater than 0', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_MissingDatasetIDs(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
body := `{"question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_EmptyDatasetIDs(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
body := `{"dataset_ids": [], "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "kb_id array cannot be empty" {
|
||||
t.Errorf("expected 'kb_id array cannot be empty', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_NoAuth(t *testing.T) {
|
||||
r := setupChunkRetrievalTestNoAuth()
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("expected 200, got %d", w.Code)
|
||||
}
|
||||
// jsonError returns HTTP 200 with error code in body
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp["code"] == float64(common.CodeSuccess) {
|
||||
t.Errorf("expected error code, got %v", resp["code"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_InvalidJSON(t *testing.T) {
|
||||
r, _ := setupChunkRetrievalTest("user1")
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader("{invalid}"))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_Success(t *testing.T) {
|
||||
_, mock := setupChunkRetrievalTest("user1")
|
||||
mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
|
||||
if userID != "user1" {
|
||||
t.Errorf("expected userID 'user1', got %q", userID)
|
||||
}
|
||||
return &service.RetrievalTestResponse{
|
||||
Chunks: []map[string]interface{}{{"docnm_kwd": "result"}},
|
||||
DocAggs: []map[string]interface{}{{"doc_id": "1", "count": float64(1)}},
|
||||
Total: 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set("user", &entity.User{ID: "user1"})
|
||||
})
|
||||
h := &ChunkHandler{chunkService: mock}
|
||||
r.POST("/api/v1/datasets/search", h.RetrievalTest)
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Fatalf("expected code 0, got %v: %q", resp["code"], resp["message"])
|
||||
}
|
||||
data, ok := resp["data"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected data object, got %T", resp["data"])
|
||||
}
|
||||
if total, _ := data["total"].(float64); total != 1 {
|
||||
t.Errorf("expected total 1, got %v", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRetrieval_ServiceError(t *testing.T) {
|
||||
_, mock := setupChunkRetrievalTest("user1")
|
||||
mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
|
||||
return nil, errors.New("db connection refused")
|
||||
}
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set("user", &entity.User{ID: "user1"})
|
||||
})
|
||||
h := &ChunkHandler{chunkService: mock}
|
||||
r.POST("/api/v1/datasets/search", h.RetrievalTest)
|
||||
|
||||
body := `{"dataset_ids": ["kb1"], "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/datasets/search", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Fatalf("expected 500, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msg, _ := resp["message"].(string)
|
||||
if msg != "dataset search failed" {
|
||||
t.Errorf("expected generic error message, got %q", msg)
|
||||
}
|
||||
if strings.Contains(msg, "db connection refused") {
|
||||
t.Errorf("internal error details leaked to response: %q", msg)
|
||||
}
|
||||
}
|
||||
@@ -31,41 +31,72 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// searchbotLLM is the interface for LLM calls used by SearchbotHandler.
|
||||
// searchbotLLM is the interface for LLM calls used by SearchBotHandler.
|
||||
type searchbotLLM interface {
|
||||
Chat(tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (*modelModule.ChatResponse, error)
|
||||
}
|
||||
|
||||
// SearchbotRealLLM wraps ModelProviderService to implement searchbotLLM.
|
||||
type SearchbotRealLLM struct {
|
||||
// ChunkRetriever abstracts chunk retrieval for the searchbots handler.
|
||||
type ChunkRetriever interface {
|
||||
RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
|
||||
}
|
||||
|
||||
// SearchBotRealLLM wraps ModelProviderService to implement searchbotLLM.
|
||||
type SearchBotRealLLM struct {
|
||||
Svc *service.ModelProviderService
|
||||
}
|
||||
|
||||
func (r *SearchbotRealLLM) Chat(tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (*modelModule.ChatResponse, error) {
|
||||
driver, modelName, apiConfig, _, err := r.Svc.GetModelConfigFromProviderInstance(tenantID, entity.ModelTypeChat, modelID)
|
||||
func (r *SearchBotRealLLM) Chat(tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (*modelModule.ChatResponse, error) {
|
||||
chatModel, err := r.Svc.GetChatModel(tenantID, modelID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatModel := modelModule.NewChatModel(driver, &modelName, apiConfig)
|
||||
return chatModel.ModelDriver.ChatWithMessages(*chatModel.ModelName, messages, chatModel.APIConfig, config)
|
||||
}
|
||||
|
||||
// SearchbotRequest is the request body for POST /api/v1/searchbots/related_questions.
|
||||
type SearchbotRequest struct {
|
||||
// SearchBotRetrievalTestRequest is the request body for POST /api/v1/searchbots/retrieval_test.
|
||||
type SearchBotRetrievalTestRequest struct {
|
||||
KbIDs common.StringSlice `json:"kb_id" binding:"required"`
|
||||
Question string `json:"question" binding:"required"`
|
||||
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"`
|
||||
TenantRerankID *string `json:"tenant_rerank_id,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"`
|
||||
// TODO: wire highlight to nlp Retrieval when engine supports highlightFields
|
||||
// Python: bot_api.py → retrieval(highlight=req.get("highlight"))
|
||||
// → search.py highlightFields → ES get_highlight()
|
||||
// Issue: https://github.com/infiniflow/ragflow/issues/15712
|
||||
// Highlight *bool `json:"highlight,omitempty"`
|
||||
}
|
||||
|
||||
// SearchBotRequest is the request body for POST /api/v1/searchbots/related_questions.
|
||||
type SearchBotRequest struct {
|
||||
Question string `json:"question" binding:"required"`
|
||||
SearchID string `json:"search_id,omitempty"`
|
||||
}
|
||||
|
||||
// SearchbotHandler handles POST /api/v1/searchbots/related_questions.
|
||||
type SearchbotHandler struct {
|
||||
// SearchBotHandler handles searchbot endpoints:
|
||||
// POST /api/v1/searchbots/related_questions
|
||||
// POST /api/v1/searchbots/retrieval_test
|
||||
type SearchBotHandler struct {
|
||||
searchSvc *service.SearchService
|
||||
tenantSvc *service.TenantService
|
||||
llm searchbotLLM
|
||||
chunkSvc ChunkRetriever
|
||||
}
|
||||
|
||||
// NewSearchbotHandler creates a new SearchbotHandler.
|
||||
func NewSearchbotHandler(searchSvc *service.SearchService, tenantSvc *service.TenantService, llm searchbotLLM) *SearchbotHandler {
|
||||
return &SearchbotHandler{searchSvc: searchSvc, tenantSvc: tenantSvc, llm: llm}
|
||||
// NewSearchBotHandler creates a new SearchBotHandler.
|
||||
func NewSearchBotHandler(searchSvc *service.SearchService, tenantSvc *service.TenantService, llm searchbotLLM, chunkSvc ChunkRetriever) *SearchBotHandler {
|
||||
return &SearchBotHandler{searchSvc: searchSvc, tenantSvc: tenantSvc, llm: llm, chunkSvc: chunkSvc}
|
||||
}
|
||||
|
||||
// Handle generates related search questions based on a user query.
|
||||
@@ -74,17 +105,17 @@ func NewSearchbotHandler(searchSvc *service.SearchService, tenantSvc *service.Te
|
||||
// @Tags searchbots
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param request body SearchbotRequest true "Request body"
|
||||
// @Param request body SearchBotRequest true "Request body"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/searchbots/related_questions [post]
|
||||
func (h *SearchbotHandler) Handle(c *gin.Context) {
|
||||
func (h *SearchBotHandler) Handle(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
jsonError(c, errorCode, errorMessage)
|
||||
return
|
||||
}
|
||||
|
||||
var req SearchbotRequest
|
||||
var req SearchBotRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeArgumentError,
|
||||
@@ -148,21 +179,133 @@ func (h *SearchbotHandler) Handle(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeSuccess,
|
||||
"data": questions,
|
||||
"message": "",
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
|
||||
// RetrievalTest performs a retrieval test against specified knowledge bases.
|
||||
// @Summary Retrieval Test
|
||||
// @Description Test document retrieval across knowledge bases with optional filters, reranking, and KG search.
|
||||
// @Tags searchbots
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param request body SearchBotRetrievalTestRequest true "Retrieval test parameters"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/searchbots/retrieval_test [post]
|
||||
func (h *SearchBotHandler) RetrievalTest(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"code": errorCode, "data": nil, "message": errorMessage})
|
||||
return
|
||||
}
|
||||
|
||||
var req SearchBotRetrievalTestRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"code": common.CodeArgumentError, "data": nil, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Filter out empty strings from KbIDs before validation.
|
||||
filtered := make(common.StringSlice, 0, len(req.KbIDs))
|
||||
for _, id := range req.KbIDs {
|
||||
if strings.TrimSpace(id) != "" {
|
||||
filtered = append(filtered, id)
|
||||
}
|
||||
}
|
||||
req.KbIDs = filtered
|
||||
|
||||
if len(req.KbIDs) == 0 || req.Question == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"code": common.CodeArgumentError, "data": nil, "message": "kb_id and question are required"})
|
||||
return
|
||||
}
|
||||
|
||||
applyRetrievalDefaults(&req)
|
||||
|
||||
if req.TopK != nil && *req.TopK <= 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"code": common.CodeArgumentError, "data": nil, "message": "top_k must be greater than 0"})
|
||||
return
|
||||
}
|
||||
|
||||
svcReq := toRetrievalServiceRequest(&req)
|
||||
|
||||
result, err := h.chunkSvc.RetrievalTest(svcReq, user.ID)
|
||||
if err != nil {
|
||||
common.Warn("searchbot retrieval test failed", zap.String("error", err.Error()))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"code": common.CodeServerError, "data": nil, "message": "retrieval test failed"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"code": int(common.CodeSuccess), "data": result, "message": "success"})
|
||||
}
|
||||
|
||||
// toRetrievalServiceRequest maps the handler DTO to the service DTO.
|
||||
// The two structs differ in KbIDs (StringSlice → []string) and
|
||||
// MetaDataFilter (→ Filter) to maintain Python API compatibility.
|
||||
func toRetrievalServiceRequest(h *SearchBotRetrievalTestRequest) *service.RetrievalTestRequest {
|
||||
return &service.RetrievalTestRequest{
|
||||
Datasets: common.StringSlice(h.KbIDs),
|
||||
Question: h.Question,
|
||||
Page: h.Page,
|
||||
Size: h.Size,
|
||||
DocIDs: h.DocIDs,
|
||||
UseKG: h.UseKG,
|
||||
TopK: h.TopK,
|
||||
CrossLanguages: h.CrossLanguages,
|
||||
SearchID: h.SearchID,
|
||||
Filter: h.MetaDataFilter,
|
||||
TenantRerankID: h.TenantRerankID,
|
||||
RerankID: h.RerankID,
|
||||
Keyword: h.Keyword,
|
||||
SimilarityThreshold: h.SimilarityThreshold,
|
||||
VectorSimilarityWeight: h.VectorSimilarityWeight,
|
||||
}
|
||||
}
|
||||
|
||||
// ptrFloat64 returns a pointer to a float64 value.
|
||||
func ptrFloat64(v float64) *float64 { return &v }
|
||||
|
||||
// applyRetrievalDefaults fills in default values for optional fields,
|
||||
// matching Python bot_api.py retrieval_test endpoint.
|
||||
func applyRetrievalDefaults(req *SearchBotRetrievalTestRequest) {
|
||||
if req.Page == nil {
|
||||
v := 1
|
||||
req.Page = &v
|
||||
}
|
||||
if req.Size == nil {
|
||||
v := 30
|
||||
req.Size = &v
|
||||
}
|
||||
if req.TopK == nil {
|
||||
v := 1024
|
||||
req.TopK = &v
|
||||
}
|
||||
if req.UseKG == nil {
|
||||
v := false
|
||||
req.UseKG = &v
|
||||
}
|
||||
if req.Keyword == nil {
|
||||
v := false
|
||||
req.Keyword = &v
|
||||
}
|
||||
if req.SimilarityThreshold == nil {
|
||||
v := 0.0
|
||||
req.SimilarityThreshold = &v
|
||||
}
|
||||
if req.VectorSimilarityWeight == nil {
|
||||
v := 0.3
|
||||
req.VectorSimilarityWeight = &v
|
||||
}
|
||||
}
|
||||
|
||||
var relatedQuestionLineRe = regexp.MustCompile(`^\d+\.\s`)
|
||||
|
||||
// parseRelatedQuestions extracts numbered list items from an LLM response.
|
||||
// Lines matching "^N. " are extracted and the number prefix is stripped.
|
||||
func parseRelatedQuestions(text string) []string {
|
||||
lineRe := regexp.MustCompile(`^\d+\.\s`)
|
||||
var result []string
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
if lineRe.MatchString(line) {
|
||||
result = append(result, lineRe.ReplaceAllString(line, ""))
|
||||
if relatedQuestionLineRe.MatchString(line) {
|
||||
result = append(result, relatedQuestionLineRe.ReplaceAllString(line, ""))
|
||||
}
|
||||
}
|
||||
if result == nil {
|
||||
|
||||
@@ -18,17 +18,390 @@ package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/entity"
|
||||
modelModule "ragflow/internal/entity/models"
|
||||
"ragflow/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// mockChunkService implements ChunkRetriever for testing.
|
||||
// It captures the last request received so tests can verify field mapping.
|
||||
type mockChunkService struct {
|
||||
retrievalTestFn func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
|
||||
LastReq *service.RetrievalTestRequest
|
||||
LastUserID string
|
||||
}
|
||||
|
||||
func (m *mockChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
|
||||
m.LastReq = req
|
||||
m.LastUserID = userID
|
||||
if m.retrievalTestFn != nil {
|
||||
return m.retrievalTestFn(req, userID)
|
||||
}
|
||||
return &service.RetrievalTestResponse{
|
||||
Chunks: []map[string]interface{}{{"docnm_kwd": "test", "content_with_weight": "content"}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func setupSearchbotsTest(userID string) (*SearchBotHandler, *mockChunkService, *gin.Engine) {
|
||||
mockSvc := &mockChunkService{}
|
||||
h := &SearchBotHandler{
|
||||
chunkSvc: mockSvc,
|
||||
}
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set("user", &entity.User{ID: userID})
|
||||
})
|
||||
r.POST("/api/v1/searchbots/retrieval_test", h.RetrievalTest)
|
||||
return h, mockSvc, r
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_Basic(t *testing.T) {
|
||||
_, mockSvc, r := setupSearchbotsTest("user1")
|
||||
|
||||
body := `{"kb_id": ["kb1"], "question": "test question"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Errorf("expected code 0, got %v", resp["code"])
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "success" {
|
||||
t.Errorf("expected message 'success', got %q", msg)
|
||||
}
|
||||
// Verify field mapping: handler → service request
|
||||
if mockSvc.LastReq == nil {
|
||||
t.Fatal("RetrievalTest was not called")
|
||||
}
|
||||
if len(mockSvc.LastReq.Datasets) != 1 || mockSvc.LastReq.Datasets[0] != "kb1" {
|
||||
t.Errorf("Datasets = %v, want [\"kb1\"]", mockSvc.LastReq.Datasets)
|
||||
}
|
||||
if mockSvc.LastReq.Question != "test question" {
|
||||
t.Errorf("Question = %q, want \"test question\"", mockSvc.LastReq.Question)
|
||||
}
|
||||
if mockSvc.LastUserID != "user1" {
|
||||
t.Errorf("userID = %q, want \"user1\"", mockSvc.LastUserID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_MissingKbID(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
msg, _ := resp["message"].(string)
|
||||
if msg == "" || msg == "success" {
|
||||
t.Errorf("expected validation error message, got %q", msg)
|
||||
}
|
||||
if !strings.Contains(msg, "KbIDs") || !strings.Contains(msg, "required") {
|
||||
t.Errorf("expected message to mention 'KbIDs' and 'required', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_MissingQuestion(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": ["kb1"]}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
msg, _ := resp["message"].(string)
|
||||
if msg == "" || msg == "success" {
|
||||
t.Errorf("expected validation error message, got %q", msg)
|
||||
}
|
||||
if !strings.Contains(msg, "Question") || !strings.Contains(msg, "required") {
|
||||
t.Errorf("expected message to mention 'Question' and 'required', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_NoAuth(t *testing.T) {
|
||||
h := NewSearchBotHandler(nil, nil, nil, &mockChunkService{})
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.POST("/api/v1/searchbots/retrieval_test", h.RetrievalTest)
|
||||
w := httptest.NewRecorder()
|
||||
body := `{"kb_id": ["kb1"], "question": "test"}`
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Errorf("expected 401, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_ServiceError(t *testing.T) {
|
||||
h, _, r := setupSearchbotsTest("user1")
|
||||
h.chunkSvc = &mockChunkService{
|
||||
retrievalTestFn: func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
|
||||
return nil, errors.New("db error")
|
||||
},
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
body := `{"kb_id": ["kb1"], "question": "test"}`
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("expected 500, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
code, _ := resp["code"].(float64)
|
||||
if code == 0 {
|
||||
t.Errorf("expected non-zero error code, got %v", code)
|
||||
}
|
||||
msg, _ := resp["message"].(string)
|
||||
if msg == "" || msg == "success" {
|
||||
t.Errorf("expected error message, got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_KbIDSingleString(t *testing.T) {
|
||||
// Verify "kb1" (string) is accepted and converted to []string{"kb1"}
|
||||
_, mockSvc, r := setupSearchbotsTest("user1")
|
||||
|
||||
body := `{"kb_id": "kb1", "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if mockSvc.LastReq == nil {
|
||||
t.Fatal("RetrievalTest was not called")
|
||||
}
|
||||
if len(mockSvc.LastReq.Datasets) != 1 || mockSvc.LastReq.Datasets[0] != "kb1" {
|
||||
t.Errorf("Datasets = %v, want [\"kb1\"]", mockSvc.LastReq.Datasets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_KbIDArray(t *testing.T) {
|
||||
// Verify ["a","b"] (array) still works
|
||||
_, mockSvc, r := setupSearchbotsTest("user1")
|
||||
|
||||
body := `{"kb_id": ["a","b"], "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if mockSvc.LastReq == nil {
|
||||
t.Fatal("RetrievalTest was not called")
|
||||
}
|
||||
if len(mockSvc.LastReq.Datasets) != 2 || mockSvc.LastReq.Datasets[0] != "a" || mockSvc.LastReq.Datasets[1] != "b" {
|
||||
t.Errorf("Datasets = %v, want [\"a\",\"b\"]", mockSvc.LastReq.Datasets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_InvalidJSON(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader("{invalid}"))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_EmptyStringKbID(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": "", "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "kb_id and question are required" {
|
||||
t.Errorf("expected message 'kb_id and question are required', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_WhitespaceOnlyKbID(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": " ", "question": "test"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "kb_id and question are required" {
|
||||
t.Errorf("expected message 'kb_id and question are required', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_DefaultsApplied(t *testing.T) {
|
||||
// Verify that when optional fields are omitted, the handler applies
|
||||
// defaults matching Python bot_api.py retrieval_test endpoint.
|
||||
_, mockSvc, r := setupSearchbotsTest("user1")
|
||||
|
||||
body := `{"kb_id": ["kb1"], "question": "does this default?"}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if mockSvc.LastReq == nil {
|
||||
t.Fatal("RetrievalTest was not called")
|
||||
}
|
||||
|
||||
svcReq := mockSvc.LastReq
|
||||
if svcReq.Page == nil || *svcReq.Page != 1 {
|
||||
t.Errorf("Page = %v, want 1", nullableInt(svcReq.Page))
|
||||
}
|
||||
if svcReq.Size == nil || *svcReq.Size != 30 {
|
||||
t.Errorf("Size = %v, want 30", nullableInt(svcReq.Size))
|
||||
}
|
||||
if svcReq.TopK == nil || *svcReq.TopK != 1024 {
|
||||
t.Errorf("TopK = %v, want 1024", nullableInt(svcReq.TopK))
|
||||
}
|
||||
if svcReq.UseKG == nil || *svcReq.UseKG != false {
|
||||
t.Errorf("UseKG = %v, want false", nullableBool(svcReq.UseKG))
|
||||
}
|
||||
if svcReq.Keyword == nil || *svcReq.Keyword != false {
|
||||
t.Errorf("Keyword = %v, want false", nullableBool(svcReq.Keyword))
|
||||
}
|
||||
if svcReq.SimilarityThreshold == nil || *svcReq.SimilarityThreshold != 0.0 {
|
||||
t.Errorf("SimilarityThreshold = %v, want 0.0", nullableFloat(svcReq.SimilarityThreshold))
|
||||
}
|
||||
if svcReq.VectorSimilarityWeight == nil || *svcReq.VectorSimilarityWeight != 0.3 {
|
||||
t.Errorf("VectorSimilarityWeight = %v, want 0.3", nullableFloat(svcReq.VectorSimilarityWeight))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_TopKZero(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": ["kb1"], "question": "test", "top_k": 0}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "top_k must be greater than 0" {
|
||||
t.Errorf("expected message 'top_k must be greater than 0', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBotsRetrieval_TopKNegative(t *testing.T) {
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": ["kb1"], "question": "test", "top_k": -1}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d", w.Code)
|
||||
}
|
||||
if msg := jsonDecodeMessage(t, w.Body.Bytes()); msg != "top_k must be greater than 0" {
|
||||
t.Errorf("expected message 'top_k must be greater than 0', got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func jsonDecodeMessage(t *testing.T, body []byte) string {
|
||||
t.Helper()
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
msg, _ := resp["message"].(string)
|
||||
return msg
|
||||
}
|
||||
|
||||
func nullableInt(p *int) string {
|
||||
if p == nil { return "nil" }
|
||||
return fmt.Sprintf("%d", *p)
|
||||
}
|
||||
func nullableBool(p *bool) string {
|
||||
if p == nil { return "nil" }
|
||||
return fmt.Sprintf("%v", *p)
|
||||
}
|
||||
func nullableFloat(p *float64) string {
|
||||
if p == nil { return "nil" }
|
||||
return fmt.Sprintf("%v", *p)
|
||||
}
|
||||
|
||||
|
||||
func TestSearchBotsRetrieval_EmptyQuestion(t *testing.T) {
|
||||
// Send kb_id but empty question — caught by binding:"required" on the DTO.
|
||||
_, _, r := setupSearchbotsTest("user1")
|
||||
body := `{"kb_id": ["kb1"], "question": ""}`
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest("POST", "/api/v1/searchbots/retrieval_test", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
msg := jsonDecodeMessage(t, w.Body.Bytes())
|
||||
if !strings.Contains(msg, "Question") || !strings.Contains(msg, "required") {
|
||||
t.Errorf("expected validation error mentioning Question and required, got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// fakeSearchbotLLM implements searchbotLLM for testing.
|
||||
type fakeSearchbotLLM struct {
|
||||
response string
|
||||
@@ -42,7 +415,7 @@ func (f *fakeSearchbotLLM) Chat(tenantID, modelID string, messages []modelModule
|
||||
return &modelModule.ChatResponse{Answer: &f.response}, nil
|
||||
}
|
||||
|
||||
func setupSearchbotRequest(body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
func setupSearchBotRequest(body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
@@ -54,17 +427,17 @@ func setupSearchbotRequest(body string) (*gin.Context, *httptest.ResponseRecorde
|
||||
return c, w
|
||||
}
|
||||
|
||||
// TestSearchbotHandler_Success verifies the happy path.
|
||||
func TestSearchbotHandler_Success(t *testing.T) {
|
||||
// TestSearchBotHandler_Success verifies the happy path.
|
||||
func TestSearchBotHandler_Success(t *testing.T) {
|
||||
llm := &fakeSearchbotLLM{
|
||||
response: `Here are some related questions:
|
||||
1. How do EV impact environment?
|
||||
2. What are advantages of EV?
|
||||
3. Cost of EV?`,
|
||||
}
|
||||
h := NewSearchbotHandler(nil, nil, llm)
|
||||
h := NewSearchBotHandler(nil, nil, llm, nil)
|
||||
|
||||
c, w := setupSearchbotRequest(`{"question": "EV benefits"}`)
|
||||
c, w := setupSearchBotRequest(`{"question": "EV benefits"}`)
|
||||
h.Handle(c)
|
||||
|
||||
var resp map[string]interface{}
|
||||
@@ -72,6 +445,9 @@ func TestSearchbotHandler_Success(t *testing.T) {
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Fatalf("expected code 0, got %v: %v", resp["code"], resp["message"])
|
||||
}
|
||||
if msg, _ := resp["message"].(string); msg != "success" {
|
||||
t.Errorf("expected message 'success', got %q", msg)
|
||||
}
|
||||
|
||||
questions, ok := resp["data"].([]interface{})
|
||||
if !ok {
|
||||
@@ -85,14 +461,14 @@ func TestSearchbotHandler_Success(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchbotHandler_EmptyResponse verifies empty LLM response returns empty list.
|
||||
func TestSearchbotHandler_EmptyResponse(t *testing.T) {
|
||||
// TestSearchBotHandler_EmptyResponse verifies empty LLM response returns empty list.
|
||||
func TestSearchBotHandler_EmptyResponse(t *testing.T) {
|
||||
llm := &fakeSearchbotLLM{
|
||||
response: "No related questions found.",
|
||||
}
|
||||
h := NewSearchbotHandler(nil, nil, llm)
|
||||
h := NewSearchBotHandler(nil, nil, llm, nil)
|
||||
|
||||
c, w := setupSearchbotRequest(`{"question": "EV benefits"}`)
|
||||
c, w := setupSearchBotRequest(`{"question": "EV benefits"}`)
|
||||
h.Handle(c)
|
||||
|
||||
var resp map[string]interface{}
|
||||
@@ -109,14 +485,14 @@ func TestSearchbotHandler_EmptyResponse(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchbotHandler_LLMFailure verifies error handling on LLM failure.
|
||||
func TestSearchbotHandler_LLMFailure(t *testing.T) {
|
||||
// TestSearchBotHandler_LLMFailure verifies error handling on LLM failure.
|
||||
func TestSearchBotHandler_LLMFailure(t *testing.T) {
|
||||
llm := &fakeSearchbotLLM{
|
||||
err: errFake{msg: "LLM unavailable"},
|
||||
}
|
||||
h := NewSearchbotHandler(nil, nil, llm)
|
||||
h := NewSearchBotHandler(nil, nil, llm, nil)
|
||||
|
||||
c, w := setupSearchbotRequest(`{"question": "EV benefits"}`)
|
||||
c, w := setupSearchBotRequest(`{"question": "EV benefits"}`)
|
||||
h.Handle(c)
|
||||
|
||||
var resp map[string]interface{}
|
||||
@@ -127,12 +503,12 @@ func TestSearchbotHandler_LLMFailure(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchbotHandler_MissingQuestion verifies validation.
|
||||
func TestSearchbotHandler_MissingQuestion(t *testing.T) {
|
||||
// TestSearchBotHandler_MissingQuestion verifies validation.
|
||||
func TestSearchBotHandler_MissingQuestion(t *testing.T) {
|
||||
llm := &fakeSearchbotLLM{response: "dummy"}
|
||||
h := NewSearchbotHandler(nil, nil, llm)
|
||||
h := NewSearchBotHandler(nil, nil, llm, nil)
|
||||
|
||||
c, w := setupSearchbotRequest(`{}`)
|
||||
c, w := setupSearchBotRequest(`{}`)
|
||||
h.Handle(c)
|
||||
|
||||
var resp map[string]interface{}
|
||||
|
||||
Reference in New Issue
Block a user