diff --git a/internal/dao/document.go b/internal/dao/document.go index 4c291b7cf0..8c73f1074a 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -116,3 +116,10 @@ func (dao *DocumentDAO) GetAllDocIDsByKBIDs(kbIDs []string) ([]map[string]string } return result, nil } + +// CountByTenantID counts documents by tenant ID +func (dao *DocumentDAO) CountByTenantID(tenantID string) (int64, error) { + var count int64 + err := DB.Model(&entity.Document{}).Where("created_by = ?", tenantID).Count(&count).Error + return count, err +} diff --git a/internal/dao/file.go b/internal/dao/file.go index b8d27e613b..f8806e21d4 100644 --- a/internal/dao/file.go +++ b/internal/dao/file.go @@ -221,6 +221,82 @@ func (dao *FileDAO) GetAllIDsByTenantID(tenantID string) ([]string, error) { return ids, err } +// GetByParentIDAndName gets file by parent folder ID and name +func (dao *FileDAO) GetByParentIDAndName(parentID, name string) (*entity.File, error) { + var file entity.File + err := DB.Where("parent_id = ? AND name = ?", parentID, name).First(&file).Error + if err != nil { + return nil, err + } + return &file, nil +} + +// GetIDListByID recursively gets list of file IDs by traversing folder structure +func (dao *FileDAO) GetIDListByID(id string, names []string, count int, res []string) ([]string, error) { + if count < len(names) { + file, err := dao.GetByParentIDAndName(id, names[count]) + if err != nil { + return res, nil + } + res = append(res, file.ID) + return dao.GetIDListByID(file.ID, names, count+1, res) + } + return res, nil +} + +// CreateFolder creates a folder in the database +func (dao *FileDAO) CreateFolder(parentID, tenantID, name, fileType string) (*entity.File, error) { + file := &entity.File{ + ID: generateUUID(), + ParentID: parentID, + TenantID: tenantID, + CreatedBy: tenantID, + Name: name, + Type: fileType, + Size: 0, + SourceType: "", + } + if err := DB.Create(file).Error; err != nil { + return nil, err + } + return file, nil +} + +// Insert inserts a new file record +func (dao *FileDAO) Insert(file *entity.File) error { + return DB.Create(file).Error +} + +// IsParentFolderExist checks if parent folder exists +func (dao *FileDAO) IsParentFolderExist(parentID string) bool { + var count int64 + err := DB.Model(&entity.File{}).Where("id = ?", parentID).Count(&count).Error + if err != nil || count == 0 { + return false + } + return true +} + +// Query retrieves files by conditions +func (dao *FileDAO) Query(name string, parentID string) []*entity.File { + var files []*entity.File + query := DB.Model(&entity.File{}) + if name != "" { + query = query.Where("name = ?", name) + } + if parentID != "" { + query = query.Where("parent_id = ?", parentID) + } + query.Find(&files) + return files +} + +// UpdateByID updates file by ID +func (dao *FileDAO) UpdateByID(id string, updates map[string]interface{}) bool { + result := DB.Model(&entity.File{}).Where("id = ?", id).Updates(updates) + return result.RowsAffected > 0 +} + // generateUUID generates a UUID func generateUUID() string { id := uuid.New().String() diff --git a/internal/handler/file.go b/internal/handler/file.go index cae393ffc8..93828277fa 100644 --- a/internal/handler/file.go +++ b/internal/handler/file.go @@ -20,6 +20,7 @@ import ( "net/http" "ragflow/internal/common" "strconv" + "strings" "github.com/gin-gonic/gin" @@ -226,3 +227,147 @@ func (h *FileHandler) GetAllParentFolders(c *gin.Context) { "message": "success", }) } + +type CreateFolderRequest struct { + Name string `json:"name" binding:"required"` + ParentID string `json:"parent_id"` + Type string `json:"type"` +} + +// UploadFile handles file upload and folder creation +// @Summary Upload Files or Create Folder +// @Description Upload files or create a folder based on content type +// @Tags file +// @Accept multipart/form-data, application/json +// @Produce json +// @Param parent_id query string false "parent folder ID (for multipart/form-data)" +// @Param file formData file false "file to upload (for multipart/form-data)" +// @Success 200 {object} map[string]interface{} +// @Failure 400 {object} map[string]interface{} +// @Router /v1/file/upload [post] +func (h *FileHandler) UploadFile(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + c.JSON(http.StatusBadRequest, gin.H{ + "code": errorCode, + "message": errorMessage, + }) + return + } + + userID := user.ID + + contentType := c.ContentType() + + if strings.Contains(contentType, "multipart/form-data") { + if err := c.Request.ParseMultipartForm(32 << 20); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": "Failed to parse multipart form: " + err.Error(), + }) + return + } + + form := c.Request.MultipartForm + if form == nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": "No file part!", + }) + return + } + parentID := c.PostForm("parent_id") + if parentID == "" { + rootFolder, err := h.fileService.GetRootFolder(userID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "code": 500, + "message": err.Error(), + }) + return + } + parentID = rootFolder["id"].(string) + } + + files := form.File["file"] + if len(files) == 0 { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": "No file selected!", + }) + return + } + + for _, fileHeader := range files { + if fileHeader.Filename == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": "No file selected!", + }) + return + } + } + + result, err := h.fileService.UploadFile(userID, parentID, files) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": 0, + "data": result, + "message": "success", + }) + return + } + + if strings.Contains(contentType, "application/json") { + var req CreateFolderRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": err.Error(), + }) + return + } + + parentID := req.ParentID + if parentID == "" { + rootFolder, err := h.fileService.GetRootFolder(userID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "code": 500, + "message": err.Error(), + }) + return + } + parentID = rootFolder["id"].(string) + } + + result, err := h.fileService.CreateFolder(userID, req.Name, parentID, req.Type) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": 0, + "data": result, + "message": "success", + }) + return + } + + c.JSON(http.StatusBadRequest, gin.H{ + "code": 400, + "message": "Unsupported content type", + }) + return +} diff --git a/internal/handler/memory.go b/internal/handler/memory.go index a1d9e5423c..b8e04d06d8 100644 --- a/internal/handler/memory.go +++ b/internal/handler/memory.go @@ -194,7 +194,7 @@ func (h *MemoryHandler) CreateMemory(c *gin.Context) { // Return success response c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "message": true, + "message": "success", "data": result, }) } @@ -293,7 +293,7 @@ func (h *MemoryHandler) UpdateMemory(c *gin.Context) { // Return success response c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "message": true, + "message": "success", "data": result, }) } @@ -347,7 +347,7 @@ func (h *MemoryHandler) DeleteMemory(c *gin.Context) { // Return success response c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "message": true, + "message": "success", "data": nil, }) } @@ -436,7 +436,7 @@ func (h *MemoryHandler) ListMemories(c *gin.Context) { // Return success response c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "message": true, + "message": "success", "data": result, }) } @@ -490,7 +490,7 @@ func (h *MemoryHandler) GetMemoryConfig(c *gin.Context) { // Return success response c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "message": true, + "message": "success", "data": result, }) } diff --git a/internal/router/router.go b/internal/router/router.go index 72f485d25a..3668effe06 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -189,6 +189,12 @@ func (r *Router) Setup(engine *gin.Engine) { // message.GET("", r.memoryHandler.GetMessages) // message.GET("/:memory_id/:message_id/content", r.memoryHandler.GetMessageContent) // } + + file := v1.Group("/files") + { + file.POST("", r.fileHandler.UploadFile) + } + // provider pool route group provider := v1.Group("/providers") { diff --git a/internal/service/file.go b/internal/service/file.go index 38b88f5907..1c8b89adee 100644 --- a/internal/service/file.go +++ b/internal/service/file.go @@ -17,8 +17,17 @@ package service import ( + "fmt" + "mime/multipart" + "os" + "path/filepath" "ragflow/internal/dao" "ragflow/internal/entity" + "ragflow/internal/storage" + "ragflow/internal/util" + "strings" + + "github.com/google/uuid" ) // FileService file service @@ -218,3 +227,217 @@ func (s *FileService) GetAllParentFolders(fileID string) ([]map[string]interface return result, nil } + +const ( + FileTypeFolder = "folder" + FileTypeVirtual = "virtual" +) + +// GetDocCount gets document count for a tenant +func (s *FileService) GetDocCount(tenantID string) (int64, error) { + documentDAO := dao.NewDocumentDAO() + return documentDAO.CountByTenantID(tenantID) +} + +// UploadFile uploads files to a folder +func (s *FileService) UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { + if parentID == "" { + rootFolder, err := s.fileDAO.GetRootFolder(tenantID) + if err != nil { + return nil, fmt.Errorf("failed to get root folder: %w", err) + } + parentID = rootFolder.ID + } + + _, err := s.fileDAO.GetByID(parentID) + if err != nil { + return nil, fmt.Errorf("Can't find this folder!") + } + + maxFileNumPerUser := os.Getenv("MAX_FILE_NUM_PER_USER") + if maxFileNumPerUser != "" { + var maxNum int64 + if _, err := fmt.Sscanf(maxFileNumPerUser, "%d", &maxNum); err == nil && maxNum > 0 { + docCount, err := s.GetDocCount(tenantID) + if err != nil { + return nil, fmt.Errorf("failed to get document count: %w", err) + } + if docCount >= maxNum { + return nil, fmt.Errorf("Exceed the maximum file number of a free user!") + } + } + } + + storageImpl := storage.GetStorageFactory().GetStorage() + if storageImpl == nil { + return nil, fmt.Errorf("storage not initialized") + } + + var result []map[string]interface{} + + for _, fileHeader := range files { + filename := fileHeader.Filename + if filename == "" { + return nil, fmt.Errorf("No file selected!") + } + + fileType := util.FilenameType(filename) + + fileObjNames := s.parseFilePath(filename) + + idList, err := s.fileDAO.GetIDListByID(parentID, fileObjNames, 1, []string{parentID}) + if err != nil { + return nil, fmt.Errorf("failed to get file ID list: %w", err) + } + + var lastFolder *entity.File + if len(fileObjNames) != len(idList)-1 { + lastID := idList[len(idList)-1] + lastFolder, err = s.fileDAO.GetByID(lastID) + if err != nil { + return nil, fmt.Errorf("Folder not found!") + } + createdFolder, err := s.createFolderRecursive(lastFolder, fileObjNames, len(idList), tenantID) + if err != nil { + return nil, fmt.Errorf("failed to create folder: %w", err) + } + lastFolder = createdFolder + } else { + lastID := idList[len(idList)-2] + lastFolder, err = s.fileDAO.GetByID(lastID) + if err != nil { + return nil, fmt.Errorf("Folder not found!") + } + } + + location := fileObjNames[len(fileObjNames)-1] + for storageImpl.ObjExist(lastFolder.ID, location) { + location += "_" + } + + src, err := fileHeader.Open() + if err != nil { + return nil, fmt.Errorf("failed to open uploaded file: %w", err) + } + defer src.Close() + + data := make([]byte, fileHeader.Size) + if _, err := src.Read(data); err != nil { + return nil, fmt.Errorf("failed to read file data: %w", err) + } + + if err := storageImpl.Put(lastFolder.ID, location, data); err != nil { + return nil, fmt.Errorf("failed to store file: %w", err) + } + + uniqueName := s.getUniqueFilename(fileObjNames[len(fileObjNames)-1], lastFolder.ID) + + fileRecord := &entity.File{ + ID: s.generateUUID(), + ParentID: lastFolder.ID, + TenantID: tenantID, + CreatedBy: tenantID, + Name: uniqueName, + Location: &location, + Size: int64(len(data)), + Type: fileType, + SourceType: "", + } + + if err := s.fileDAO.Insert(fileRecord); err != nil { + return nil, fmt.Errorf("failed to insert file record: %w", err) + } + + result = append(result, s.toFileResponse(fileRecord)) + } + + return result, nil +} + +func (s *FileService) parseFilePath(filename string) []string { + filename = strings.TrimPrefix(filename, "/") + parts := strings.Split(filename, "/") + var result []string + for _, part := range parts { + if part != "" { + result = append(result, part) + } + } + return result +} + +func (s *FileService) createFolderRecursive(parentFolder *entity.File, names []string, count int, tenantID string) (*entity.File, error) { + if count > len(names)-2 { + return parentFolder, nil + } + + newFolder, err := s.fileDAO.CreateFolder(parentFolder.ID, tenantID, names[count], FileTypeFolder) + if err != nil { + return nil, err + } + + return s.createFolderRecursive(newFolder, names, count+1, tenantID) +} + +func (s *FileService) getUniqueFilename(name, parentID string) string { + existingFiles := s.fileDAO.Query(name, parentID) + if len(existingFiles) == 0 { + return name + } + + base := filepath.Base(name) + ext := filepath.Ext(name) + nameWithoutExt := strings.TrimSuffix(base, ext) + + counter := 1 + for { + newName := fmt.Sprintf("%s_%d%s", nameWithoutExt, counter, ext) + existingFiles = s.fileDAO.Query(newName, parentID) + if len(existingFiles) == 0 { + return newName + } + counter++ + } +} + +func (s *FileService) generateUUID() string { + id := uuid.New().String() + return strings.ReplaceAll(id, "-", "") +} + +// CreateFolder creates a new folder or virtual file +func (s *FileService) CreateFolder(tenantID, name, parentID, fileType string) (map[string]interface{}, error) { + if parentID == "" { + rootFolder, err := s.fileDAO.GetRootFolder(tenantID) + if err != nil { + return nil, fmt.Errorf("failed to get root folder: %w", err) + } + parentID = rootFolder.ID + } + + if !s.fileDAO.IsParentFolderExist(parentID) { + return nil, fmt.Errorf("Parent Folder Doesn't Exist!") + } + + existingFiles := s.fileDAO.Query(name, parentID) + if len(existingFiles) > 0 { + return nil, fmt.Errorf("Duplicated folder name in the same folder.") + } + + if fileType == "" { + fileType = FileTypeVirtual + } + + if fileType == FileTypeFolder { + fileType = FileTypeFolder + } else { + fileType = FileTypeVirtual + } + + folder, err := s.fileDAO.CreateFolder(parentID, tenantID, name, fileType) + if err != nil { + return nil, fmt.Errorf("failed to create folder: %w", err) + } + + return s.toFileResponse(folder), nil +} diff --git a/internal/util/file.go b/internal/util/file.go new file mode 100644 index 0000000000..2e1628b54b --- /dev/null +++ b/internal/util/file.go @@ -0,0 +1,134 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package util + +import ( + "path/filepath" + "regexp" + "strings" +) + +const ( + FileTypePDF = "pdf" + FileTypeDOC = "doc" + FileTypeVISUAL = "visual" + FileTypeAURAL = "aural" + FileTypeFOLDER = "folder" + FileTypeOTHER = "other" +) + +var ( + filenameLenLimit = 255 +) + +func init() { +} + +func normalizeFilename(filename string) (string, bool) { + if filename == "" { + return "", false + } + base := filepath.Base(filename) + base = strings.TrimSpace(base) + if base == "" || len(base) > filenameLenLimit { + return "", false + } + return strings.ToLower(base), true +} + +func FilenameType(filename string) string { + normalized, ok := normalizeFilename(filename) + if !ok { + return FileTypeOTHER + } + + if matched, _ := regexp.MatchString(`.*\.pdf$`, normalized); matched { + return FileTypePDF + } + + docExtensions := []string{ + "msg", "eml", "doc", "docx", "ppt", "pptx", "yml", "xml", "htm", "json", "jsonl", "ldjson", + "csv", "txt", "ini", "xls", "xlsx", "wps", "rtf", "hlp", "pages", "numbers", "key", + "md", "mdx", "py", "js", "java", "c", "cpp", "h", "php", "go", "ts", "sh", "cs", "kt", + "html", "sql", "epub", + } + for _, ext := range docExtensions { + if strings.HasSuffix(normalized, "."+ext) { + return FileTypeDOC + } + } + + audioExtensions := []string{ + "wav", "flac", "ape", "alac", "wv", "mp3", "aac", "ogg", "vorbis", "opus", + } + for _, ext := range audioExtensions { + if strings.HasSuffix(normalized, "."+ext) { + return FileTypeAURAL + } + } + + visualExtensions := []string{ + "jpg", "jpeg", "png", "tif", "gif", "pcx", "tga", "exif", "fpx", "svg", "psd", "cdr", + "pcd", "dxf", "ufo", "eps", "ai", "raw", "WMF", "webp", "avif", "apng", "icon", "ico", + "mpg", "mpeg", "avi", "rm", "rmvb", "mov", "wmv", "asf", "dat", "asx", "wvx", "mpe", + "mpa", "mp4", "mkv", + } + for _, ext := range visualExtensions { + if strings.HasSuffix(normalized, "."+ext) { + return FileTypeVISUAL + } + } + + return FileTypeOTHER +} + +func SanitizeFilename(filename string) string { + if filename == "" { + return "" + } + filename = strings.TrimSpace(filename) + if filename == "" { + return "" + } + + filename = strings.ReplaceAll(filename, "\\", "/") + filename = strings.Trim(filename, "/") + + parts := strings.Split(filename, "/") + var sanitizedParts []string + for _, part := range parts { + if part != "" && part != "." && part != ".." { + sanitizedParts = append(sanitizedParts, part) + } + } + + unsafeRegex := regexp.MustCompile(`[^A-Za-z0-9_\-/]`) + for i, part := range sanitizedParts { + sanitizedParts[i] = unsafeRegex.ReplaceAllString(part, "") + } + + result := strings.Join(sanitizedParts, "/") + return result +} + +func GetFileExtension(filename string) string { + ext := filepath.Ext(filename) + if len(ext) > 0 && ext[0] == '.' { + return strings.ToLower(ext[1:]) + } + return strings.ToLower(ext) +}