diff --git a/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/DatasetController.java b/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/DatasetController.java index ebb0d855a..7c1b2a0e7 100644 --- a/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/DatasetController.java +++ b/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/DatasetController.java @@ -6,27 +6,32 @@ import cn.iocoder.yudao.framework.common.pojo.PageParam; import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.common.util.object.BeanUtils; import cn.iocoder.yudao.framework.excel.core.util.ExcelUtils; -import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetPageReqVO; -import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetRespVO; -import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetSaveReqVO; -import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetTreeNode; +import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.*; import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetDO; import cn.iocoder.yudao.module.llm.service.dataset.DatasetService; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.Parameter; import io.swagger.v3.oas.annotations.tags.Tag; import org.apache.commons.io.IOUtils; +import org.apache.tomcat.jni.FileInfo; import org.springframework.core.io.ClassPathResource; +import org.springframework.core.io.InputStreamResource; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; import org.springframework.validation.annotation.Validated; import org.springframework.web.bind.annotation.*; import javax.annotation.Resource; +import javax.servlet.ServletOutputStream; import javax.servlet.http.HttpServletResponse; import javax.validation.Valid; -import java.io.FileNotFoundException; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; +import java.io.*; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Base64; import java.util.List; import static cn.iocoder.yudao.framework.apilog.core.enums.OperateTypeEnum.EXPORT; @@ -104,50 +109,37 @@ public class DatasetController { @GetMapping("/download-example") @Operation(summary = "下载示例文件") - public void downloadExampleFile (@RequestParam("type") int type, HttpServletResponse response) throws IOException { - String fileName; - String contentType; + public ResponseEntity downloadExampleFile(@RequestParam("type") int type, HttpServletResponse response) throws IOException { + FileInfoVO fileInfo = getFileInfo(type); - // 根据 type 参数确定文件名和 Content-Type + // 从 resources/file/dataset_example 目录加载文件 + ClassPathResource resource = new ClassPathResource("file/dataset_example/" + fileInfo.getFileName()); + if (!resource.exists()) { + throw new FileNotFoundException("文件未找到: " + fileInfo.getFileName()); + } + + InputStream inputStream = resource.getInputStream(); + HttpHeaders headers = new HttpHeaders(); + headers.add("Content-Disposition", "attachment; filename=\"" + URLEncoder.encode(fileInfo.getFileName(), "UTF-8") + "\""); + + return ResponseEntity.ok() + .headers(headers) + .contentType(MediaType.APPLICATION_OCTET_STREAM) + .body(new InputStreamResource(inputStream)); + } + + private FileInfoVO getFileInfo(int type) { switch (type) { - // txt 文件 case 1: - fileName = "dataset_example_txt.txt"; - contentType = "text/plain"; - break; - // xlsx + return new FileInfoVO("dataset_example_txt.txt", "text/plain"); case 2: - fileName = "dataset_example_xlsx.xlsx"; - contentType = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"; - break; - // csv + return new FileInfoVO("dataset_example_xlsx.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"); case 3: - fileName = "dataset_example_csv.csv"; - contentType = "text/csv"; - break; - // json + return new FileInfoVO("dataset_example_csv.csv", "text/csv"); case 4: - fileName = "dataset_example_json.json"; - contentType = "application/json"; - break; + return new FileInfoVO("dataset_example_json.json", "application/json"); default: throw new IllegalArgumentException("无效的 type 参数"); } - - // 从 resources/file/dataset_example 目录加载文件 - ClassPathResource resource = new ClassPathResource("file/dataset_example/" + fileName); - if (!resource.exists()) { - throw new FileNotFoundException("文件未找到: " + fileName); - } - - // 设置响应头 - response.setContentType(contentType); - response.setHeader("Content-Disposition", "attachment; filename=" + fileName); - - // 将文件内容写入响应输出流 - try (InputStream inputStream = resource.getInputStream(); - OutputStream outputStream = response.getOutputStream()) { - IOUtils.copy(inputStream, outputStream); - } } } diff --git a/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/vo/FileInfoVO.java b/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/vo/FileInfoVO.java new file mode 100644 index 000000000..daf0505ac --- /dev/null +++ b/yudao-module-llm/yudao-module-llm-biz/src/main/java/cn/iocoder/yudao/module/llm/controller/admin/dataset/vo/FileInfoVO.java @@ -0,0 +1,14 @@ +package cn.iocoder.yudao.module.llm.controller.admin.dataset.vo; + +import lombok.AllArgsConstructor; +import lombok.Data; + +/** + * @Description + */ +@Data +@AllArgsConstructor +public class FileInfoVO { + private String fileName; + private String contentType; +}