53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from .shared import *
|
|
|
|
def _resolve_tenant_id(user: User, fallback_tenant_id: int) -> int:
|
|
try:
|
|
return int(getattr(user, "tenant_id") or fallback_tenant_id or 1)
|
|
except Exception:
|
|
return int(fallback_tenant_id or 1)
|
|
|
|
def _safe_decode_bytes(raw: bytes) -> str:
|
|
for encoding in ("utf-8", "utf-8-sig", "gb18030", "latin-1"):
|
|
try:
|
|
return raw.decode(encoding)
|
|
except Exception:
|
|
continue
|
|
return raw.decode("utf-8", errors="ignore")
|
|
|
|
def _extract_text_from_upload(filename: str, raw: bytes) -> tuple[str, str]:
|
|
name = str(filename or "uploaded.txt")
|
|
suffix = name.rsplit(".", 1)[-1].lower() if "." in name else "txt"
|
|
|
|
if suffix in {"txt", "md", "markdown", "csv", "json", "py", "ts", "tsx", "js", "jsx"}:
|
|
return ("md" if suffix in {"md", "markdown"} else "txt", _safe_decode_bytes(raw))
|
|
|
|
if suffix == "pdf":
|
|
try:
|
|
from pypdf import PdfReader # type: ignore
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=f"PDF 解析依赖缺失: {exc}") from exc
|
|
reader = PdfReader(io.BytesIO(raw))
|
|
parts: list[str] = []
|
|
for idx, page in enumerate(reader.pages, start=1):
|
|
page_text = (page.extract_text() or "").strip()
|
|
if not page_text:
|
|
continue
|
|
# 写入页码标记，向量分块时会保留 page_number 元数据，便于引用来源页码。
|
|
parts.append(f"[[PAGE:{idx}]]\n{page_text}")
|
|
text = "\n\n".join(parts)
|
|
return ("pdf", text)
|
|
|
|
if suffix in {"docx", "doc"}:
|
|
try:
|
|
from docx import Document # type: ignore
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=f"Word 解析依赖缺失: {exc}") from exc
|
|
doc = Document(io.BytesIO(raw))
|
|
text = "\n".join([p.text for p in doc.paragraphs if p.text])
|
|
return ("word", text)
|
|
|
|
raise HTTPException(status_code=400, detail=f"不支持的文件类型: .{suffix}")
|
|
|