zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-26 02:13:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context (#17314)

Browse Source 
### Summary

As title.

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-24 16:47:12 +08:00
committed by GitHub
parent 008fa3e10e
commit bdfc3ada41
78 changed files with 1031 additions and 904 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
internal/handler/agent.go
View File

@@ -48,15 +48,13 @@ import (
// the FileHandler (handler/file.go), not by any agent handler, so
// the interface deliberately does NOT list it. (Code review CR1.)
type agentFileService interface {
DownloadAgentFile(tenantID, location string) ([]byte, error)
DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error)
// UploadInfos stores raw bytes in the per-user downloads bucket and
// returns lightweight descriptors. Mirrors python FileService.upload_info
// (multi-file path) used by the agent upload endpoint.
UploadInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error)
// returns lightweight descriptors.
UploadInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error)
// UploadFromURL downloads a remote file (with SSRF protection) and
// stores it as an info blob. Mirrors python FileService.upload_info
// (single-file path with ?url=) used by the agent upload endpoint.
UploadFromURL(tenantID, rawURL string) (map[string]interface{}, error)
// stores it as an info blob.
UploadFromURL(ctx context.Context, tenantID, rawURL string) (map[string]interface{}, error)
}
// chatAgentService is the subset of AgentService used by the chat-completion

6
internal/handler/agent_attachment.go
View File

@@ -29,6 +29,7 @@
package handler
import (
"context"
"net/http"
"path/filepath"
"strings"
@@ -42,7 +43,7 @@ import (
// agentAttachmentFileService is the subset of FileService used by
// the attachment-download handler.
type agentAttachmentFileService interface {
DownloadAgentFile(tenantID, location string) ([]byte, error)
DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error)
}
// attachmentRequestMetadata holds the parsed query params used when
@@ -81,7 +82,8 @@ func (h *AgentHandler) streamAgentAttachment(c *gin.Context, tenantID, attachmen
return
}
blob, err := h.fileService.DownloadAgentFile(tenantID, attachmentID)
ctx := c.Request.Context()
blob, err := h.fileService.DownloadAgentFile(ctx, tenantID, attachmentID)
if err != nil {
common.ResponseWithCodeData(c, common.CodeDataError, nil, "Attachment not found!")
return

3
internal/handler/agent_download.go
View File

@@ -54,6 +54,7 @@ func (h *AgentHandler) DownloadAgentFile(c *gin.Context) {
return
}
ctx := c.Request.Context()
// IDOR note (security review H1): the Go User struct has no
// TenantID field — the project collapses user and tenant in a
// single-tenant session model. Python's @add_tenant_id_to_kwargs
@@ -62,7 +63,7 @@ func (h *AgentHandler) DownloadAgentFile(c *gin.Context) {
// per-object ownership check, so this port preserves the python
// shape. A future per-object ownership check should be added in
// both the python and Go code paths.
blob, err := h.fileService.DownloadAgentFile(user.ID, fileID)
blob, err := h.fileService.DownloadAgentFile(ctx, user.ID, fileID)
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
return

7
internal/handler/agent_download_test.go
View File

@@ -17,6 +17,7 @@
package handler
import (
"context"
"errors"
"mime/multipart"
"net/http/httptest"
@@ -43,19 +44,19 @@ type fakeAgentFileService struct {
urlUploadErr error
}
func (f *fakeAgentFileService) UploadInfos(_ string, _ []*multipart.FileHeader) ([]map[string]interface{}, error) {
func (f *fakeAgentFileService) UploadInfos(ctx context.Context, _ string, _ []*multipart.FileHeader) ([]map[string]interface{}, error) {
if f.uploadErr != nil {
return nil, f.uploadErr
}
return f.uploadList, nil
}
func (f *fakeAgentFileService) UploadFromURL(_ string, _ string) (map[string]interface{}, error) {
func (f *fakeAgentFileService) UploadFromURL(ctx context.Context, _ string, _ string) (map[string]interface{}, error) {
if f.urlUploadErr != nil {
return nil, f.urlUploadErr
}
return f.urlUpload, nil
}
func (f *fakeAgentFileService) DownloadAgentFile(_ string, _ string) ([]byte, error) {
func (f *fakeAgentFileService) DownloadAgentFile(ctx context.Context, _ string, _ string) ([]byte, error) {
if f.err != nil {
return nil, f.err
}

5
internal/handler/agent_upload.go
View File

@@ -116,8 +116,9 @@ func (h *AgentHandler) UploadAgentFile(c *gin.Context) {
// into the normal UploadInfos path. We replicate that with a
// guard that dispatches to UploadFromURL only when both
// conditions are met.
ctx := c.Request.Context()
if url := c.Query("url"); url != "" && len(files) == 1 {
uploaded, err := h.fileService.UploadFromURL(user.ID, url)
uploaded, err := h.fileService.UploadFromURL(ctx, user.ID, url)
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
return
@@ -132,7 +133,7 @@ func (h *AgentHandler) UploadAgentFile(c *gin.Context) {
return
}
results, err := h.fileService.UploadInfos(user.ID, files)
results, err := h.fileService.UploadInfos(ctx, user.ID, files)
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
return

6
internal/handler/bot_test.go
View File

@@ -711,15 +711,15 @@ type fakeFileService struct {
err error
}
func (f *fakeFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) {
func (f *fakeFileService) DownloadAgentFile(ctx context.Context, tenantID, location string) ([]byte, error) {
return f.blob, f.err
}
func (f *fakeFileService) UploadInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) {
func (f *fakeFileService) UploadInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) {
return nil, nil
}
func (f *fakeFileService) UploadFromURL(tenantID, rawURL string) (map[string]interface{}, error) {
func (f *fakeFileService) UploadFromURL(ctx context.Context, tenantID, rawURL string) (map[string]interface{}, error) {
return nil, nil
}

3
internal/handler/chat.go
View File

@@ -202,7 +202,8 @@ func (h *ChatHandler) MindMap(c *gin.Context) {
return
}
mindMap, err := runMindMap(mindMapRunConfig{
ctx := c.Request.Context()
mindMap, err := runMindMap(ctx, mindMapRunConfig{
Question: req.Question,
KbIDs: kbIDs,
SearchID: req.SearchID,

57
internal/handler/chunk.go
View File

@@ -33,15 +33,15 @@ import (
// 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)
SwitchChunks(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error
UpdateChunk(req *service.UpdateChunkRequest, userID string) error
RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error)
Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error)
RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
Get(ctx context.Context, req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error)
List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
SwitchChunks(ctx context.Context, userID string, datasetID, documentID string, availableInt int, chunkIDs []string) error
UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error
RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error)
Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error)
AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
}
// ChunkHandler chunk handler
@@ -62,8 +62,6 @@ func NewChunkHandler(chunkService chunkService, userService *service.UserService
// @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]
@@ -125,8 +123,9 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) {
return
}
ctx := c.Request.Context()
// Call service with user ID for permission checks
resp, err := h.chunkService.RetrievalTest(&req, user.ID)
resp, err := h.chunkService.RetrievalTest(ctx, &req, user.ID)
if err != nil {
common.Warn("dataset search failed", zap.String("error", err.Error()))
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "dataset search failed")
@@ -140,8 +139,6 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) {
// @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"
@@ -164,7 +161,8 @@ func (h *ChunkHandler) Get(c *gin.Context) {
ChunkID: chunkID,
}
resp, err := h.chunkService.Get(req, user.ID)
ctx := c.Request.Context()
resp, err := h.chunkService.Get(ctx, req, user.ID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error())
return
@@ -198,7 +196,8 @@ func (h *ChunkHandler) Parse(c *gin.Context) {
return
}
data, code, err := h.chunkService.Parse(userID, datasetId, &req)
ctx := c.Request.Context()
data, code, err := h.chunkService.Parse(ctx, userID, datasetId, &req)
if code != common.CodeSuccess {
common.ResponseWithCodeData(c, code, nil, err.Error())
return
@@ -249,7 +248,8 @@ func (h *ChunkHandler) ListChunks(c *gin.Context) {
req.AvailableInt = &available
}
resp, err := h.chunkService.List(&req, user.ID)
ctx := c.Request.Context()
resp, err := h.chunkService.List(ctx, &req, user.ID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, err.Error())
return
@@ -306,7 +306,8 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) {
return
}
resp, code, err := h.chunkService.StopParsing(user.ID, datasetID, req)
ctx := c.Request.Context()
resp, code, err := h.chunkService.StopParsing(ctx, user.ID, datasetID, req)
if err != nil {
var data interface{}
if resp != nil {
@@ -332,8 +333,6 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) {
// @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]
@@ -361,7 +360,8 @@ func (h *ChunkHandler) List(c *gin.Context) {
req.Size = &defaultSize
}
resp, err := h.chunkService.List(&req, user.ID)
ctx := c.Request.Context()
resp, err := h.chunkService.List(ctx, &req, user.ID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error())
return
@@ -420,7 +420,8 @@ func (h *ChunkHandler) SwitchChunks(c *gin.Context) {
return
}
if err = h.chunkService.SwitchChunks(userID, datasetID, documentID, availableInt, chunkIDs); err != nil {
ctx := c.Request.Context()
if err = h.chunkService.SwitchChunks(ctx, userID, datasetID, documentID, availableInt, chunkIDs); err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, err.Error())
return
}
@@ -476,18 +477,16 @@ func parseAvailableBody(rawBody map[string]interface{}) (int, error) {
return 0, fmt.Errorf("available must be a boolean")
}
}
return 0, fmt.Errorf("`available_int` or `available` is required.")
return 0, fmt.Errorf("`available_int` or `available` is required")
}
// UpdateChunk updates a chunk
// @Summary Update Chunk
// @Description Update chunk fields
// @Tags chunks
// @Accept json
// @Produce json
// @Param request body service.UpdateChunkRequest true "update chunk"
// @Success 200 {object} map[string]interface{}
// @Router /v1/chunk/update [post]
// @Router /api/v1/chunk/update [post]
func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
@@ -581,7 +580,8 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
req.DocumentID = documentID
req.ChunkID = chunkID
err := h.chunkService.UpdateChunk(&req, user.ID)
ctx := c.Request.Context()
err := h.chunkService.UpdateChunk(ctx, &req, user.ID)
if err != nil {
var coded interface {
Code() common.ErrorCode
@@ -605,8 +605,6 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
// @Summary Remove Chunks
// @Description Remove chunks from a document
// @Tags chunks
// @Accept json
// @Produce json
// @Param request body service.RemoveChunksRequest true "remove chunks request"
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks [delete]
@@ -637,7 +635,8 @@ func (h *ChunkHandler) RemoveChunks(c *gin.Context) {
return
}
deletedCount, err := h.chunkService.RemoveChunks(&req, user.ID)
ctx := c.Request.Context()
deletedCount, err := h.chunkService.RemoveChunks(ctx, &req, user.ID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error())
return

60
internal/handler/chunk_test.go
View File

@@ -20,12 +20,12 @@ import (
// 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)
retrievalTestFn func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
addChunkFn func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
listFn func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
switchChunksFn func(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error
updateChunkFn func(req *service.UpdateChunkRequest, userID string) error
stopParsingFn func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
listFn func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
switchChunksFn func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error
updateChunkFn func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error
stopParsingFn func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
}
type codedTestError struct {
@@ -41,46 +41,46 @@ func (e codedTestError) Code() common.ErrorCode {
return e.code
}
func (m *mockChunkSvc) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
func (m *mockChunkSvc) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
if m.retrievalTestFn != nil {
return m.retrievalTestFn(req, userID)
return m.retrievalTestFn(ctx, 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) {
func (m *mockChunkSvc) Get(context.Context, *service.GetChunkRequest, string) (*service.GetChunkResponse, error) {
panic("not implemented")
}
func (m *mockChunkSvc) List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
func (m *mockChunkSvc) List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
if m.listFn != nil {
return m.listFn(req, userID)
return m.listFn(ctx, req, userID)
}
panic("not implemented")
}
func (m *mockChunkSvc) SwitchChunks(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
func (m *mockChunkSvc) SwitchChunks(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if m.switchChunksFn != nil {
return m.switchChunksFn(userID, datasetID, documentID, availableInt, chunkIDs)
return m.switchChunksFn(ctx, userID, datasetID, documentID, availableInt, chunkIDs)
}
panic("not implemented")
}
func (m *mockChunkSvc) UpdateChunk(req *service.UpdateChunkRequest, userID string) error {
func (m *mockChunkSvc) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
if m.updateChunkFn != nil {
return m.updateChunkFn(req, userID)
return m.updateChunkFn(ctx, req, userID)
}
panic("not implemented")
}
func (m *mockChunkSvc) RemoveChunks(*service.RemoveChunksRequest, string) (int64, error) {
func (m *mockChunkSvc) RemoveChunks(context.Context, *service.RemoveChunksRequest, string) (int64, error) {
panic("not implemented")
}
func (m *mockChunkSvc) StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
func (m *mockChunkSvc) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if m.stopParsingFn != nil {
return m.stopParsingFn(userID, datasetID, req)
return m.stopParsingFn(ctx, userID, datasetID, req)
}
panic("not implemented")
}
func (m *mockChunkSvc) Parse(string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
func (m *mockChunkSvc) Parse(context.Context, string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
panic("not implemented")
}
func (m *mockChunkSvc) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
@@ -139,7 +139,7 @@ func TestChunkHandlerListChunksMapsPathAndQuery(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks)
mock.listFn = func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
}
@@ -188,7 +188,7 @@ func TestChunkHandlerListChunksMapsAvailableFalse(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks)
mock.listFn = func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
}
@@ -216,7 +216,7 @@ func TestChunkHandlerSwitchChunksCallsService(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.SwitchChunks)
mock.switchChunksFn = func(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
mock.switchChunksFn = func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if userID != "user-1" || datasetID != "kb-1" || documentID != "doc-1" {
t.Fatalf("ids = %q/%q/%q, want user-1/kb-1/doc-1", userID, datasetID, documentID)
}
@@ -267,7 +267,7 @@ func TestChunkHandlerUpdateChunkUsesPathIDs(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(req *service.UpdateChunkRequest, userID string) error {
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
}
@@ -295,7 +295,7 @@ func TestChunkHandlerUpdateChunkValidationErrorIsBadRequest(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(req *service.UpdateChunkRequest, userID string) error {
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
return codedTestError{code: common.CodeArgumentError, msg: "`tag_feas` values must be finite numbers greater than 0"}
}
@@ -343,7 +343,7 @@ func TestChunkRetrieval_EmptyQuestion(t *testing.T) {
func TestChunkStopParsing_Success(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if userID != "user1" {
t.Fatalf("expected user1, got %q", userID)
}
@@ -378,7 +378,7 @@ func TestChunkStopParsing_Success(t *testing.T) {
func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
t.Fatal("service should not be called when document_ids is missing")
return nil, common.CodeSuccess, nil
}
@@ -405,10 +405,10 @@ func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) {
func TestChunkStopParsing_InvalidStateIncludesPythonErrorCode(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
return &service.StopParsingResponse{
Data: map[string]interface{}{"error_code": "DOC_STOP_PARSING_INVALID_STATE"},
}, common.CodeDataError, errors.New("Can't stop parsing document that has not started or already completed")
}, common.CodeDataError, errors.New("can't stop parsing document that has not started or already completed")
}
w := httptest.NewRecorder()
@@ -433,7 +433,7 @@ func TestChunkStopParsing_InvalidStateIncludesPythonErrorCode(t *testing.T) {
if data["error_code"] != "DOC_STOP_PARSING_INVALID_STATE" {
t.Fatalf("unexpected error_code: %v", data["error_code"])
}
if resp["message"] != "Can't stop parsing document that has not started or already completed" {
if resp["message"] != "can't stop parsing document that has not started or already completed" {
t.Fatalf("unexpected message: %v", resp["message"])
}
}
@@ -552,7 +552,7 @@ func TestChunkRetrieval_InvalidJSON(t *testing.T) {
func TestChunkRetrieval_Success(t *testing.T) {
_, mock := setupChunkRetrievalTest("user1")
mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
mock.retrievalTestFn = func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
if userID != "user1" {
t.Errorf("expected userID 'user1', got %q", userID)
}
@@ -598,7 +598,7 @@ func TestChunkRetrieval_Success(t *testing.T) {
func TestChunkRetrieval_ServiceError(t *testing.T) {
_, mock := setupChunkRetrievalTest("user1")
mock.retrievalTestFn = func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
mock.retrievalTestFn = func(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
return nil, errors.New("db connection refused")
}

10
internal/handler/dataset.go
View File

@@ -284,9 +284,9 @@ func (h *DatasetsHandler) GetIngestionSummary(c *gin.Context) {
common.ErrorWithCode(c, errorCode, errorMessage)
return
}
ctx := c.Request.Context()
datasetID := c.Param("dataset_id")
result, code, err := h.datasetsService.GetIngestionSummary(datasetID, user.ID)
result, code, err := h.datasetsService.GetIngestionSummary(ctx, datasetID, user.ID)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -784,8 +784,9 @@ func (h *DatasetsHandler) RunIndex(c *gin.Context) {
return
}
ctx := c.Request.Context()
indexType := strings.ToLower(strings.TrimSpace(c.Query("type")))
data, code, err := h.datasetsService.RunIndex(userID, datasetID, indexType)
data, code, err := h.datasetsService.RunIndex(ctx, userID, datasetID, indexType)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -950,7 +951,8 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) {
return
}
data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(userID, datasetID, documentID, req)
ctx := c.Request.Context()
data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(ctx, userID, datasetID, documentID, req)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return

6
internal/handler/dify_retrieval_handler.go
View File

@@ -21,6 +21,7 @@ import (
"errors"
"fmt"
"net/http"
"ragflow/internal/dao"
"strconv"
"strings"
@@ -66,7 +67,7 @@ type RetrievalServiceIface interface {
// DocumentDAOIface abstracts DocumentDAO for the Dify handler.
type DocumentDAOIface interface {
GetByIDs(ids []string) ([]*entity.Document, error)
GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error)
}
// --- Request / Response types ---
@@ -321,10 +322,11 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) {
allDocIDs = append(allDocIDs, id)
}
ctx := c.Request.Context()
docMap := make(map[string]*entity.Document)
if len(allDocIDs) > 0 {
var docs []*entity.Document
docs, err = h.docDAO.GetByIDs(allDocIDs)
docs, err = h.docDAO.GetByIDs(ctx, dao.DB, allDocIDs)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, fmt.Sprintf("failed to load documents: %v", err))
return

11
internal/handler/dify_retrieval_handler_test.go
View File

@@ -20,12 +20,13 @@ import (
"context"
"encoding/json"
"errors"
"gorm.io/gorm"
"net/http"
"net/http/httptest"
"strings"
"testing"
"gorm.io/gorm"
"ragflow/internal/common"
"ragflow/internal/engine"
"ragflow/internal/engine/types"
@@ -131,12 +132,12 @@ func (m *mockRetrievalService) Retrieval(ctx context.Context, req *nlp.Retrieval
type mockDocDAO struct {
DocumentDAOIface
getByIDsFn func(ids []string) ([]*entity.Document, error)
getByIDsFn func(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error)
}
func (m *mockDocDAO) GetByIDs(ids []string) ([]*entity.Document, error) {
func (m *mockDocDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) {
if m.getByIDsFn != nil {
return m.getByIDsFn(ids)
return m.getByIDsFn(ctx, db, ids)
}
return []*entity.Document{
{ID: "doc1", Name: strPtr("Test Doc")},
@@ -382,7 +383,7 @@ func TestDifyRetrieval_KBDBError(t *testing.T) {
func TestDifyRetrieval_DocLoadError(t *testing.T) {
h, r := setupDifyTest("user1")
h.docDAO = &mockDocDAO{
getByIDsFn: func(ids []string) ([]*entity.Document, error) {
getByIDsFn: func(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Document, error) {
return nil, errors.New("db unavailable")
},
}

201
internal/handler/document.go
View File

@@ -18,6 +18,7 @@ package handler
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
@@ -38,7 +39,7 @@ import (
"ragflow/internal/dao"
"ragflow/internal/service"
dataset "ragflow/internal/service/dataset"
"ragflow/internal/service/dataset"
"ragflow/internal/service/document"
)
@@ -46,46 +47,46 @@ var IMG_BASE64_PREFIX = "data:image/png;base64,"
// documentServiceIface defines the DocumentService methods used by DocumentHandler.
type documentServiceIface interface {
GetDocumentByID(id string) (*document.DocumentResponse, error)
UpdateDocument(id string, req *document.UpdateDocumentRequest) error
DeleteDocument(id string) error
DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error)
ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error)
ListDocuments(page, pageSize int) ([]*document.DocumentResponse, int64, error)
ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error)
ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error)
ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error)
GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error)
GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error)
GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error)
GetThumbnails(userID string, docIDs []string) (map[string]string, error)
GetDocumentImage(imageID string) ([]byte, error)
GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error)
SetDocumentMetadata(docID string, meta map[string]interface{}) error
DeleteDocumentMetadata(docID string, keys []string) error
DeleteDocumentAllMetadata(docID string) error
GetDocumentMetadataByID(docID string) (map[string]interface{}, error)
GetDocumentArtifact(filename, userID string) (*document.ArtifactResponse, error)
GetDocumentPreview(docID string) (*document.DocumentPreview, error)
UploadLocalDocuments(kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string)
UploadWebDocument(kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error)
UploadEmptyDocument(kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error)
DownloadDocument(datasetID, docID string) (*document.DownloadDocumentResp, error)
UpdateDatasetDocument(userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error)
BatchUpdateDocumentMetadatas(datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error)
ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error)
IngestDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error)
Ingest(userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error)
RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error)
BatchUpdateDocumentStatus(userID, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error)
GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error)
UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) error
DeleteDocument(ctx context.Context, id string) error
DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error)
ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error)
ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error)
ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error)
ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error)
ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error)
GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error)
GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error)
GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error)
GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error)
GetDocumentImage(ctx context.Context, imageID string) ([]byte, error)
GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error)
SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error
DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error
DeleteDocumentAllMetadata(ctx context.Context, docID string) error
GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error)
GetDocumentArtifact(ctx context.Context, filename, userID string) (*document.ArtifactResponse, error)
GetDocumentPreview(ctx context.Context, docID string) (*document.DocumentPreview, error)
UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string)
UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error)
UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error)
DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error)
UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error)
BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error)
ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error)
IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error)
Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error)
RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error)
BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error)
}
// fileUploadIface defines the FileService upload methods used by DocumentHandler.
type fileUploadIface interface {
UploadDocumentInfos(userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error)
UploadDocumentInfoByURL(userID, rawURL string) (map[string]interface{}, common.ErrorCode, error)
UploadDocumentInfos(ctx context.Context, userID string, files []*multipart.FileHeader) ([]map[string]interface{}, common.ErrorCode, error)
UploadDocumentInfoByURL(ctx context.Context, userID, rawURL string) (map[string]interface{}, common.ErrorCode, error)
}
// DocumentHandler document handler
@@ -108,8 +109,6 @@ func NewDocumentHandler(documentService documentServiceIface, datasetService *da
// @Summary Get Document Info
// @Description Get document details by ID
// @Tags documents
// @Accept json
// @Produce json
// @Param id path int true "document ID"
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/documents/{id} [get]
@@ -128,7 +127,8 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) {
return
}
document, err := h.documentService.GetDocumentByID(id)
ctx := c.Request.Context()
doc, err := h.documentService.GetDocumentByID(ctx, id)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{
"error": "document not found",
@@ -137,7 +137,7 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) {
}
c.JSON(http.StatusOK, gin.H{
"data": document,
"data": doc,
})
}
@@ -155,7 +155,8 @@ func (h *DocumentHandler) GetThumbnail(c *gin.Context) {
return
}
result, err := h.documentService.GetThumbnails(user.ID, docIDs)
ctx := c.Request.Context()
result, err := h.documentService.GetThumbnails(ctx, user.ID, docIDs)
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
return
@@ -187,7 +188,8 @@ func parseThumbnailDocIDs(c *gin.Context) []string {
// GetDocumentImage returns a document image from object storage.
func (h *DocumentHandler) GetDocumentImage(c *gin.Context) {
imageID := c.Param("image_id")
data, err := h.documentService.GetDocumentImage(imageID)
ctx := c.Request.Context()
data, err := h.documentService.GetDocumentImage(ctx, imageID)
if err != nil {
common.ResponseWithCodeData(c, common.CodeDataError, nil, "Image not found.")
return
@@ -224,7 +226,8 @@ func (h *DocumentHandler) GetDocumentArtifact(c *gin.Context) {
return
}
filename := c.Param("filename")
artifact, err := h.documentService.GetDocumentArtifact(filename, user.ID)
ctx := c.Request.Context()
artifact, err := h.documentService.GetDocumentArtifact(ctx, filename, user.ID)
if err != nil {
switch {
case errors.Is(err, document.ErrArtifactInvalidFilename),
@@ -256,7 +259,8 @@ func (h *DocumentHandler) GetDocumentPreview(c *gin.Context) {
return
}
preview, err := h.documentService.GetDocumentPreview(docID)
ctx := c.Request.Context()
preview, err := h.documentService.GetDocumentPreview(ctx, docID)
if err != nil {
common.ErrorWithCode(c, common.CodeDataError, "Document not found!")
return
@@ -298,7 +302,8 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) {
return
}
doc, err := h.documentService.GetDocumentByID(id)
ctx := c.Request.Context()
doc, err := h.documentService.GetDocumentByID(ctx, id)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!")
return
@@ -316,7 +321,7 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) {
return
}
if err = h.documentService.UpdateDocument(id, &req); err != nil {
if err = h.documentService.UpdateDocument(ctx, id, &req); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{
"error": err.Error(),
})
@@ -350,7 +355,8 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) {
return
}
doc, err := h.documentService.GetDocumentByID(id)
ctx := c.Request.Context()
doc, err := h.documentService.GetDocumentByID(ctx, id)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!")
return
@@ -360,7 +366,7 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) {
return
}
if err = h.documentService.DeleteDocument(id); err != nil {
if err = h.documentService.DeleteDocument(ctx, id); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{
"error": err.Error(),
})
@@ -411,7 +417,8 @@ func (h *DocumentHandler) DeleteDocuments(c *gin.Context) {
}
userID := c.GetString("user_id")
deleted, err := h.documentService.DeleteDocuments(ids, req.DeleteAll, datasetID, userID)
ctx := c.Request.Context()
deleted, err := h.documentService.DeleteDocuments(ctx, ids, req.DeleteAll, datasetID, userID)
if err != nil {
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
return
@@ -473,7 +480,8 @@ func (h *DocumentHandler) BatchUpdateDocumentStatus(c *gin.Context) {
return
}
result, code, err := h.documentService.BatchUpdateDocumentStatus(userID, datasetID, status, documentIDs)
ctx := c.Request.Context()
result, code, err := h.documentService.BatchUpdateDocumentStatus(ctx, userID, datasetID, status, documentIDs)
if err != nil {
message := err.Error()
if code == common.CodeServerError {
@@ -520,8 +528,9 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) {
return
}
ctx := c.Request.Context()
if c.Query("type") == "filter" {
filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(opts)
filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(ctx, opts)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, map[string]interface{}{"total": 0, "filter": map[string]interface{}{}}, "failed to get document filters")
return
@@ -533,7 +542,7 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) {
// Use kbID to filter documents.
// Note: create_time_from / create_time_to are applied post-query (not at DB level)
// so that total reflects the unfiltered count, matching the Python API contract.
documents, total, err := h.documentService.ListDocumentsByDatasetIDWithOptions(opts, page, pageSize)
documents, total, err := h.documentService.ListDocumentsByDatasetIDWithOptions(ctx, opts, page, pageSize)
if err != nil {
common.ResponseWithCodeData(c, 1, map[string]interface{}{"total": 0, "docs": []interface{}{}}, "failed to get documents")
return
@@ -547,7 +556,7 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) {
if opts.CreateTimeTo > 0 && doc.CreateTime != nil && *doc.CreateTime > opts.CreateTimeTo {
continue
}
metaFields, err := h.documentService.GetDocumentMetadataByID(doc.ID)
metaFields, err := h.documentService.GetDocumentMetadataByID(ctx, doc.ID)
if err != nil {
metaFields = make(map[string]interface{})
}
@@ -615,13 +624,15 @@ func (h *DocumentHandler) applyDocumentMetadataFilter(c *gin.Context, opts dao.D
return opts, ""
}
candidateIDs, err := h.documentService.ListDocumentIDsByDatasetIDWithOptions(opts)
ctx := c.Request.Context()
candidateIDs, err := h.documentService.ListDocumentIDsByDatasetIDWithOptions(ctx, opts)
if err != nil {
return opts, "failed to get documents"
}
candidateSet := stringSet(candidateIDs)
metadataByKey, err := h.documentService.GetMetadataByKBs([]string{opts.KbID})
metadataByKey, err := h.documentService.GetMetadataByKBs(ctx, []string{opts.KbID})
if err != nil {
return opts, err.Error()
}
@@ -879,8 +890,8 @@ func (h *DocumentHandler) uploadLocalDocuments(c *gin.Context, kb *entity.Knowle
}
}
}
data, errMsgs := h.documentService.UploadLocalDocuments(kb, tenantID, files, c.PostForm("parent_path"), override)
ctx := c.Request.Context()
data, errMsgs := h.documentService.UploadLocalDocuments(ctx, kb, tenantID, files, c.PostForm("parent_path"), override)
if len(data) == 0 && len(errMsgs) > 0 {
common.ResponseWithCodeData(c, common.CodeServerError, nil, strings.Join(errMsgs, "\n"))
return
@@ -929,7 +940,8 @@ func (h *DocumentHandler) uploadEmptyDocument(c *gin.Context, kb *entity.Knowled
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "File name must be 255 bytes or less.")
return
}
data, code, err := h.documentService.UploadEmptyDocument(kb, tenantID, name)
ctx := c.Request.Context()
data, code, err := h.documentService.UploadEmptyDocument(ctx, kb, tenantID, name)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -956,7 +968,8 @@ func (h *DocumentHandler) uploadWebDocument(c *gin.Context, kb *entity.Knowledge
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "The URL format is invalid")
return
}
data, code, err := h.documentService.UploadWebDocument(kb, tenantID, name, rawURL)
ctx := c.Request.Context()
data, code, err := h.documentService.UploadWebDocument(ctx, kb, tenantID, name, rawURL)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -1005,8 +1018,8 @@ func (h *DocumentHandler) DownloadDocument(c *gin.Context) {
common.ErrorWithCode(c, common.CodeDataError, fmt.Sprintf("The dataset not own the document %s.", docID))
return
}
res, err := h.documentService.DownloadDocument(datasetID, docID)
ctx := c.Request.Context()
res, err := h.documentService.DownloadDocument(ctx, datasetID, docID)
if err != nil {
common.ErrorWithCode(c, common.CodeDataError, err.Error())
@@ -1140,7 +1153,8 @@ func (h *DocumentHandler) MetadataSummary(c *gin.Context) {
return
}
summary, err := h.documentService.GetMetadataSummary(kbID, requestBody.DocIDs)
ctx := c.Request.Context()
summary, err := h.documentService.GetMetadataSummary(ctx, kbID, requestBody.DocIDs)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 1, nil, "Failed to get metadata summary: "+err.Error())
return
@@ -1216,7 +1230,8 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) {
}
// Authorization: user must be able to access the document's dataset.
doc, err := h.documentService.GetDocumentByID(req.DocID)
ctx := c.Request.Context()
doc, err := h.documentService.GetDocumentByID(ctx, req.DocID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found")
return
@@ -1226,7 +1241,7 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) {
return
}
err = h.documentService.SetDocumentMetadata(req.DocID, meta)
err = h.documentService.SetDocumentMetadata(ctx, req.DocID, meta)
if err != nil {
errMsg := err.Error()
if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") {
@@ -1244,8 +1259,6 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) {
// @Summary Ingest Document
// @Description Ingest a document for processing
// @Tags documents
// @Accept json
// @Produce json
// @Security ApiKeyAuth
// @Param request body document.IngestDocumentRequest true "ingestion info"
// @Success 200 {object} map[string]interface{}
@@ -1268,8 +1281,8 @@ func (h *DocumentHandler) Ingest(c *gin.Context) {
common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
return
}
if code, err := h.documentService.Ingest(userID, &req); err != nil {
ctx := c.Request.Context()
if code, err := h.documentService.Ingest(ctx, userID, &req); err != nil {
common.ErrorWithCode(c, code, err.Error())
return
}
@@ -1293,7 +1306,7 @@ type DeleteMetaRequest struct {
// @Security ApiKeyAuth
// @Param request body DeleteMetaRequest true "metadata keys to delete or empty to delete all"
// @Success 200 {object} map[string]interface{}
// @Router /v1/document/delete_meta [post]
// @Router /api/v1/document/delete_meta [post]
func (h *DocumentHandler) DeleteMeta(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
@@ -1312,8 +1325,8 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) {
return
}
// Authorization: user must be able to access the document's dataset.
doc, err := h.documentService.GetDocumentByID(req.DocID)
ctx := c.Request.Context()
doc, err := h.documentService.GetDocumentByID(ctx, req.DocID)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found")
return
@@ -1337,7 +1350,7 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) {
return
}
err = h.documentService.DeleteDocumentMetadata(req.DocID, keys)
err = h.documentService.DeleteDocumentMetadata(ctx, req.DocID, keys)
if err != nil {
errMsg := err.Error()
if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") {
@@ -1349,7 +1362,7 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) {
}
} else {
// Delete entire document metadata
err = h.documentService.DeleteDocumentAllMetadata(req.DocID)
err = h.documentService.DeleteDocumentAllMetadata(ctx, req.DocID)
if err != nil {
errMsg := err.Error()
if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") {
@@ -1385,8 +1398,8 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) {
return
}
}
parseResult, err = h.documentService.ListIngestionTasks(userID, req.DatasetID, 0, 0)
ctx := c.Request.Context()
parseResult, err = h.documentService.ListIngestionTasks(ctx, userID, req.DatasetID, 0, 0)
if err != nil {
common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error())
return
@@ -1414,8 +1427,8 @@ func (h *DocumentHandler) StartIngestionTask(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID))
return
}
parseResult, err := h.documentService.IngestDocuments(datasetID, userID, req.DocumentIDs)
ctx := c.Request.Context()
parseResult, err := h.documentService.IngestDocuments(ctx, datasetID, userID, req.DocumentIDs)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error())
return
@@ -1442,8 +1455,8 @@ func (h *DocumentHandler) StopIngestionTasks(c *gin.Context) {
}
userID := c.GetString("user_id")
parseResult, err := h.documentService.StopIngestionTasks(req.Tasks, userID)
ctx := c.Request.Context()
parseResult, err := h.documentService.StopIngestionTasks(ctx, req.Tasks, userID)
if err != nil {
common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error())
return
@@ -1468,8 +1481,9 @@ func (h *DocumentHandler) RemoveIngestionTasks(c *gin.Context) {
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
deletedTasks, err := h.documentService.RemoveIngestionTasks(req.Tasks, userID)
deletedTasks, err := h.documentService.RemoveIngestionTasks(ctx, req.Tasks, userID)
if err != nil {
common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error())
return
@@ -1496,8 +1510,8 @@ func (h *DocumentHandler) ParseDocuments(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization to access the dataset.")
return
}
parseResult, err := h.documentService.ParseDocuments(datasetID, userID, req.Documents)
ctx := c.Request.Context()
parseResult, err := h.documentService.ParseDocuments(ctx, datasetID, userID, req.Documents)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error())
return
@@ -1529,8 +1543,8 @@ func (h *DocumentHandler) StopParseDocuments(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID))
return
}
result, err := h.documentService.StopParseDocuments(datasetID, req.DocumentIDs)
ctx := c.Request.Context()
result, err := h.documentService.StopParseDocuments(ctx, datasetID, req.DocumentIDs)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error())
return
@@ -1559,8 +1573,8 @@ func (h *DocumentHandler) MetadataSummaryByDataset(c *gin.Context) {
if docIDsParam := c.Query("doc_ids"); docIDsParam != "" {
docIDS = strings.Split(docIDsParam, ",")
}
summary, err := h.documentService.GetMetadataSummary(datasetID, docIDS)
ctx := c.Request.Context()
summary, err := h.documentService.GetMetadataSummary(ctx, datasetID, docIDS)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "Failed to get metadata summary"+err.Error())
return
@@ -1606,8 +1620,8 @@ func (h *DocumentHandler) UpdateDatasetDocument(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
return
}
data, code, err := h.documentService.UpdateDatasetDocument(user.ID, datasetID, documentID, &req, present)
ctx := c.Request.Context()
data, code, err := h.documentService.UpdateDatasetDocument(ctx, user.ID, datasetID, documentID, &req, present)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -1644,8 +1658,9 @@ func (h *DocumentHandler) UploadInfo(c *gin.Context) {
return
}
ctx := c.Request.Context()
if rawURL != "" {
data, code, err := h.fileService.UploadDocumentInfoByURL(user.ID, rawURL)
data, code, err := h.fileService.UploadDocumentInfoByURL(ctx, user.ID, rawURL)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -1654,7 +1669,7 @@ func (h *DocumentHandler) UploadInfo(c *gin.Context) {
return
}
data, code, err := h.fileService.UploadDocumentInfos(user.ID, fileHeaders)
data, code, err := h.fileService.UploadDocumentInfos(ctx, user.ID, fileHeaders)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -1727,8 +1742,8 @@ func (h *DocumentHandler) handleBatchUpdateDocumentMetadatas(c *gin.Context) {
Updates: updates,
Deletes: deletes,
}
resp, code, err := h.documentService.BatchUpdateDocumentMetadatas(datasetID, req.Selector, req.Updates, req.Deletes)
ctx := c.Request.Context()
resp, code, err := h.documentService.BatchUpdateDocumentMetadatas(ctx, datasetID, req.Selector, req.Updates, req.Deletes)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return

69
internal/handler/document_test.go
View File

@@ -18,6 +18,7 @@ package handler
import (
"bytes"
"context"
"encoding/json"
"fmt"
"mime/multipart"
@@ -82,7 +83,7 @@ type fakeDocumentService struct {
metadataByKBs map[string]interface{}
}
func (f *fakeDocumentService) Ingest(userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) {
func (f *fakeDocumentService) Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) {
f.ingestUserID = userID
f.ingestReq = req
if f.ingestCode != 0 || f.ingestErr != nil {
@@ -93,14 +94,14 @@ func (f *fakeDocumentService) Ingest(userID string, req *document.IngestDocument
const uploadTestDatasetID = "123e4567-e89b-12d3-a456-426614174000"
func (f *fakeDocumentService) UpdateDatasetDocument(userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) {
func (f *fakeDocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) {
return nil, common.CodeSuccess, nil
}
func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) {
func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.DocumentMetadataSelector, updates []document.DocumentMetadataUpdate, deletes []document.DocumentMetadataDelete) (*document.BatchUpdateDocumentMetadatasResponse, common.ErrorCode, error) {
return nil, common.CodeSuccess, nil
}
func (f *fakeDocumentService) GetDocumentArtifact(filename, _ string) (*document.ArtifactResponse, error) {
func (f *fakeDocumentService) GetDocumentArtifact(ctx context.Context, filename, _ string) (*document.ArtifactResponse, error) {
if filename == "error.txt" {
return nil, document.ErrArtifactNotFound
}
@@ -114,7 +115,7 @@ func (f *fakeDocumentService) GetDocumentArtifact(filename, _ string) (*document
ForceAttachment: false,
}, nil
}
func (f *fakeDocumentService) GetDocumentPreview(docID string) (*document.DocumentPreview, error) {
func (f *fakeDocumentService) GetDocumentPreview(ctx context.Context, docID string) (*document.DocumentPreview, error) {
if docID == "not-found" {
return nil, fmt.Errorf("not found")
}
@@ -124,7 +125,7 @@ func (f *fakeDocumentService) GetDocumentPreview(docID string) (*document.Docume
FileName: "preview.txt",
}, nil
}
func (f *fakeDocumentService) DownloadDocument(datasetID, docID string) (*document.DownloadDocumentResp, error) {
func (f *fakeDocumentService) DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error) {
if docID == "not-found" {
return nil, fmt.Errorf("not found")
}
@@ -134,7 +135,7 @@ func (f *fakeDocumentService) DownloadDocument(datasetID, docID string) (*docume
FileName: "doc.pdf",
}, nil
}
func (f *fakeDocumentService) GetDocumentByID(id string) (*document.DocumentResponse, error) {
func (f *fakeDocumentService) GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error) {
if f.docErr != nil {
return nil, f.docErr
}
@@ -143,109 +144,109 @@ func (f *fakeDocumentService) GetDocumentByID(id string) (*document.DocumentResp
}
return nil, fmt.Errorf("document not found")
}
func (f *fakeDocumentService) UpdateDocument(id string, req *document.UpdateDocumentRequest) error {
func (f *fakeDocumentService) UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) error {
f.updateCalled = true
f.updatedID = id
return nil
}
func (f *fakeDocumentService) DeleteDocument(id string) error {
func (f *fakeDocumentService) DeleteDocument(ctx context.Context, id string) error {
f.deleteCalled = true
f.deletedID = id
return nil
}
func (f *fakeDocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error) {
func (f *fakeDocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) {
return f.deleted, f.err
}
func (f *fakeDocumentService) ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (f *fakeDocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
return nil, nil
}
func (f *fakeDocumentService) StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) {
func (f *fakeDocumentService) StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) {
return f.stopResult, f.stopErr
}
func (f *fakeDocumentService) ListDocuments(page, pageSize int) ([]*document.DocumentResponse, int64, error) {
func (f *fakeDocumentService) ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error) {
return nil, 0, nil
}
func (f *fakeDocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
func (f *fakeDocumentService) ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
return nil, 0, nil
}
func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
f.listOpts = opts
return nil, 0, nil
}
func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) {
func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error) {
f.listOpts = opts
return f.listIDs, nil
}
func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) {
func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error) {
f.filterOpts = opts
if f.filterResult != nil {
return f.filterResult, f.filterTotal, nil
}
return map[string]interface{}{}, 0, nil
}
func (f *fakeDocumentService) GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) {
func (f *fakeDocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error) {
if f.metadataByKBs != nil {
return f.metadataByKBs, nil
}
return map[string]interface{}{}, nil
}
func (f *fakeDocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
func (f *fakeDocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
return map[string]interface{}{}, common.CodeSuccess, nil
}
func (f *fakeDocumentService) GetThumbnails(userID string, docIDs []string) (map[string]string, error) {
func (f *fakeDocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) {
f.thumbnailUserID = userID
f.thumbnailDocIDs = append([]string(nil), docIDs...)
return f.thumbnails, f.thumbnailErr
}
func (f *fakeDocumentService) GetDocumentImage(imageID string) ([]byte, error) {
func (f *fakeDocumentService) GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) {
return nil, nil
}
func (f *fakeDocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) {
func (f *fakeDocumentService) GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) {
return nil, 0, nil
}
func (f *fakeDocumentService) GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) {
func (f *fakeDocumentService) GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) {
f.metadataKBID = kbID
f.metadataDocIDs = docIDs
return f.metadataSummary, f.metadataErr
}
func (f *fakeDocumentService) SetDocumentMetadata(docID string, meta map[string]interface{}) error {
func (f *fakeDocumentService) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error {
f.setMetaCalled = true
f.setMetaDocID = docID
f.setMetaValue = meta
return nil
}
func (f *fakeDocumentService) DeleteDocumentMetadata(docID string, keys []string) error {
func (f *fakeDocumentService) DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error {
return nil
}
func (f *fakeDocumentService) DeleteDocumentAllMetadata(docID string) error {
func (f *fakeDocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID string) error {
return nil
}
func (f *fakeDocumentService) GetDocumentMetadataByID(docID string) (map[string]interface{}, error) {
func (f *fakeDocumentService) GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error) {
return nil, nil
}
func (f *fakeDocumentService) UploadLocalDocuments(kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) {
func (f *fakeDocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) {
f.uploadLocalKB = kb
f.uploadLocalPath = parentPath
f.uploadOverride = parserConfigOverride
return f.uploadLocalData, f.uploadLocalErrs
}
func (f *fakeDocumentService) UploadWebDocument(kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) {
func (f *fakeDocumentService) UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) {
return nil, common.CodeServerError, fmt.Errorf("not implemented")
}
func (f *fakeDocumentService) UploadEmptyDocument(kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) {
func (f *fakeDocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) {
return nil, common.CodeServerError, fmt.Errorf("not implemented")
}
func (f *fakeDocumentService) ListIngestionTasks(userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
func (f *fakeDocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
return nil, nil
}
func (f *fakeDocumentService) IngestDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (f *fakeDocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
return nil, nil
}
func (f *fakeDocumentService) StopIngestionTasks(tasks []string, userID string) ([]*entity.IngestionTask, error) {
func (f *fakeDocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) {
return f.stopIngestionTasks, f.stopIngestionTaskErr
}
func (f *fakeDocumentService) RemoveIngestionTasks(tasks []string, userID string) ([]map[string]string, error) {
func (f *fakeDocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) {
return f.removeIngestionTasks, f.removeIngestionTaskErr
}

13
internal/handler/file.go
View File

@@ -306,7 +306,8 @@ func (h *FileHandler) UploadFile(c *gin.Context) {
}
}
result, err := h.fileService.UploadFile(userID, parentID, files)
ctx := c.Request.Context()
result, err := h.fileService.UploadFile(ctx, userID, parentID, files)
if err != nil {
common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
return
@@ -427,7 +428,8 @@ func (h *FileHandler) MoveFiles(c *gin.Context) {
return
}
success, message := h.fileService.MoveFiles(user.ID, req.SrcFileIDs, req.DestFileID, req.NewName)
ctx := c.Request.Context()
success, message := h.fileService.MoveFiles(ctx, user.ID, req.SrcFileIDs, req.DestFileID, req.NewName)
if !success {
common.ResponseWithCodeData(c, common.CodeBadRequest, nil, message)
return
@@ -482,7 +484,8 @@ func (h *FileHandler) Download(c *gin.Context) {
// If blob is empty, try fallback via file2document
if len(blob) == 0 {
storageAddr, err := h.fileService.GetStorageAddress(fileID)
ctx := c.Request.Context()
storageAddr, err := h.fileService.GetStorageAddress(ctx, fileID)
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, nil, "Failed to get file storage address: "+err.Error())
return
@@ -564,8 +567,8 @@ func (h *FileHandler) LinkToDatasets(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "mode must be 'add' or 'replace'")
return
}
if err := h.file2DocumentService.LinkToDatasets(user.ID, &req, mode); err != nil {
ctx := c.Request.Context()
if err := h.file2DocumentService.LinkToDatasets(ctx, user.ID, &req, mode); err != nil {
common.ResponseWithCodeData(c, linkToDatasetsErrorCode(err), nil, err.Error())
return
}

4
internal/handler/mindmap.go
View File

@@ -42,7 +42,7 @@ type mindMapRunConfig struct {
TenantSvc *service.TenantService
}
func runMindMap(config mindMapRunConfig) (mindMapNode, error) {
func runMindMap(ctx context.Context, config mindMapRunConfig) (mindMapNode, error) {
if config.ChunkSvc == nil {
return mindMapNode{}, fmt.Errorf("chunk service not configured")
}
@@ -54,7 +54,7 @@ func runMindMap(config mindMapRunConfig) (mindMapNode, error) {
modelTenantID = config.AuthUserID
}
retrievalReq := mindMapRetrievalRequest(config.Question, config.KbIDs, config.SearchID, config.SearchConfig)
ranks, err := config.ChunkSvc.RetrievalTest(retrievalReq, config.AuthUserID)
ranks, err := config.ChunkSvc.RetrievalTest(ctx, retrievalReq, config.AuthUserID)
if err != nil {
return mindMapNode{}, err
}

13
internal/handler/searchbot.go
View File

@@ -162,9 +162,7 @@ func (h *SearchBotHandler) Handle(c *gin.Context) {
// 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
// @Tags searchBots
// @Param request body SearchBotRetrievalTestRequest true "Retrieval test parameters"
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/searchbots/retrieval_test [post]
@@ -203,11 +201,11 @@ func (h *SearchBotHandler) RetrievalTest(c *gin.Context) {
}
svcReq := toRetrievalServiceRequest(&req)
result, err := h.chunkSvc.RetrievalTest(svcReq, user.ID)
ctx := c.Request.Context()
result, err := h.chunkSvc.RetrievalTest(ctx, svcReq, user.ID)
if err != nil {
common.Warn("search bot retrieval test failed", zap.String("error", err.Error()))
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeServerError, nil, "retrieval test failed")
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "retrieval test failed")
return
}
@@ -360,7 +358,8 @@ func (h *SearchBotHandler) MindMap(c *gin.Context) {
searchConfig = searchConfigFromDetail(detail)
}
mindMap, err := runMindMap(mindMapRunConfig{
ctx := c.Request.Context()
mindMap, err := runMindMap(ctx, mindMapRunConfig{
Question: req.Question,
KbIDs: filtered,
SearchID: req.SearchID,

5
internal/handler/searchbot_test.go
View File

@@ -17,6 +17,7 @@
package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
@@ -40,7 +41,7 @@ type mockChunkService struct {
LastUserID string
}
func (m *mockChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
func (m *mockChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
m.LastReq = req
m.LastUserID = userID
if m.retrievalTestFn != nil {
@@ -466,7 +467,7 @@ type fakeChunkRetriever struct {
err error
}
func (f *fakeChunkRetriever) RetrievalTest(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
func (f *fakeChunkRetriever) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) {
if f.err != nil {
return nil, f.err
}