File size: 4,734 Bytes
454f1d5 cc826a1 454f1d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import logging
import re
import tempfile
import zipfile
from pathlib import Path
from typing import List
from fastapi import APIRouter, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
logger = logging.getLogger(__name__)
router = APIRouter()
plugin = None
def set_plugin_instance(plugin_instance):
global plugin
plugin = plugin_instance
@router.get("/status")
async def get_status():
if plugin is None:
return {
"name": "doc",
"enabled": False,
"message": "插件未加载",
}
return plugin.get_status()
@router.post("/merge")
async def merge_documents(
files: List[UploadFile] = File(...),
order: str = Form(None),
output_format: str = Form("md"),
):
if plugin is None or not plugin.enabled:
raise HTTPException(status_code=400, detail="插件未启用")
if output_format not in {"md", "txt"}:
raise HTTPException(status_code=400, detail="格式参数错误，仅支持 md 或 txt")
if not files:
raise HTTPException(status_code=400, detail="至少需要上传一个文件")
allowed_extensions = {".md", ".txt"}
contents = []
for file in files:
ext = Path(file.filename).suffix.lower()
if ext not in allowed_extensions:
raise HTTPException(
status_code=400, detail=f"不支持的文件格式: {file.filename}"
)
content = await file.read()
try:
decoded_content = content.decode("utf-8")
contents.append(decoded_content)
except UnicodeDecodeError:
raise HTTPException(
status_code=400, detail=f"文件编码不支持: {file.filename}"
)
if order:
try:
order_list = [int(x) for x in order.split(",")]
reordered_contents = [
contents[i] for i in order_list if 0 <= i < len(contents)
]
if not reordered_contents:
raise HTTPException(status_code=400, detail="顺序参数无效")
contents = reordered_contents
except (ValueError, IndexError):
raise HTTPException(
status_code=400, detail="顺序参数格式错误，应为逗号分隔的数字"
)
merged_content = "\n\n".join([item.strip() for item in contents])
file_suffix = f".{output_format}"
temp_file = tempfile.NamedTemporaryFile(
mode="w", delete=False, suffix=file_suffix, encoding="utf-8"
)
temp_file.write(merged_content)
temp_file.close()
media_type = "text/markdown" if output_format == "md" else "text/plain"
filename = f"merged_document.{output_format}"
return FileResponse(temp_file.name, media_type=media_type, filename=filename)
@router.post("/split")
async def split_document(file: UploadFile = File(...), output_format: str = Form("md")):
if plugin is None or not plugin.enabled:
raise HTTPException(status_code=400, detail="插件未启用")
if output_format not in {"md", "txt"}:
raise HTTPException(status_code=400, detail="格式参数错误，仅支持 md 或 txt")
ext = Path(file.filename).suffix.lower()
if ext not in {".md", ".txt"}:
raise HTTPException(
status_code=400, detail="不支持的文件格式，仅支持 .md 和 .txt"
)
content_bytes = await file.read()
try:
content = content_bytes.decode("utf-8")
except UnicodeDecodeError:
raise HTTPException(status_code=400, detail="文件编码不支持，请使用 UTF-8 编码")
pattern = re.compile(r"^第\d+章", re.MULTILINE)
matches = list(pattern.finditer(content))
if not matches:
raise HTTPException(status_code=400, detail="未找到 '第X章' 格式的章节标记")
temp_dir = tempfile.mkdtemp()
zip_path = Path(temp_dir) / "split_documents.zip"
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
for i, match in enumerate(matches):
start_pos = match.start()
end_pos = matches[i + 1].start() if i + 1 < len(matches) else len(content)
chapter_content = content[start_pos:end_pos].strip()
first_line = chapter_content.split("\n", 1)[0].strip()
sanitized = re.sub(r'[<>:"/\\|?*]', "_", first_line)
filename = f"{sanitized}.{output_format}"
temp_file = Path(temp_dir) / filename
with open(temp_file, "w", encoding="utf-8") as f:
f.write(chapter_content)
zipf.write(temp_file, filename)
return FileResponse(
zip_path, media_type="application/zip", filename="split_documents.zip"
)
|