fix(avatar): 异步知识库索引并修复大文件上传 #16
@@ -17,7 +17,8 @@ UPLOAD_DIR = os.path.abspath(os.getenv("UPLOAD_DIR", os.path.join(BASE_DIR, "upl
|
|||||||
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
||||||
|
|
||||||
ALLOWED_EXT = {".md", ".txt", ".pdf", ".doc", ".docx", ".xlsx"}
|
ALLOWED_EXT = {".md", ".txt", ".pdf", ".doc", ".docx", ".xlsx"}
|
||||||
MAX_UPLOAD_BYTES = 10 * 1024 * 1024
|
MAX_UPLOAD_BYTES = 50 * 1024 * 1024
|
||||||
|
UPLOAD_CHUNK_BYTES = 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
class QAIn(BaseModel):
|
class QAIn(BaseModel):
|
||||||
@@ -80,17 +81,25 @@ async def upload_doc(avatar_id: str, file: UploadFile = File(...), authorization
|
|||||||
os.makedirs(avatar_dir, exist_ok=True)
|
os.makedirs(avatar_dir, exist_ok=True)
|
||||||
stored = f"{uuid.uuid4().hex}{ext}"
|
stored = f"{uuid.uuid4().hex}{ext}"
|
||||||
path = os.path.join(avatar_dir, stored)
|
path = os.path.join(avatar_dir, stored)
|
||||||
content = await file.read()
|
file_size = 0
|
||||||
if len(content) > MAX_UPLOAD_BYTES:
|
try:
|
||||||
return fail("文件不能超过 10MB", code=400)
|
# Stream large files to disk so a 100MB upload does not occupy 100MB RAM.
|
||||||
with open(path, "wb") as f:
|
with open(path, "wb") as f:
|
||||||
f.write(content)
|
while chunk := await file.read(UPLOAD_CHUNK_BYTES):
|
||||||
|
file_size += len(chunk)
|
||||||
|
if file_size > MAX_UPLOAD_BYTES:
|
||||||
|
raise ValueError("文件不能超过 50MB")
|
||||||
|
f.write(chunk)
|
||||||
|
except ValueError as exc:
|
||||||
|
if os.path.exists(path):
|
||||||
|
os.remove(path)
|
||||||
|
return fail(str(exc), code=400)
|
||||||
doc = KnowledgeDoc(
|
doc = KnowledgeDoc(
|
||||||
id=uuid.uuid4().hex,
|
id=uuid.uuid4().hex,
|
||||||
avatar_id=avatar_id,
|
avatar_id=avatar_id,
|
||||||
filename=file.filename,
|
filename=file.filename,
|
||||||
file_type=ext.lstrip("."),
|
file_type=ext.lstrip("."),
|
||||||
file_size=len(content),
|
file_size=file_size,
|
||||||
file_url=f"/api/files/{avatar_id}/{stored}",
|
file_url=f"/api/files/{avatar_id}/{stored}",
|
||||||
status="parsing",
|
status="parsing",
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -64,6 +64,29 @@ def test_upload_returns_before_background_vectorization(
|
|||||||
db.close()
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_upload_rejects_oversize_file_before_queuing_indexing(
|
||||||
|
tmp_path: Path,
|
||||||
|
authorization_context,
|
||||||
|
):
|
||||||
|
context = authorization_context
|
||||||
|
with (
|
||||||
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
||||||
|
patch("routers.knowledge.MAX_UPLOAD_BYTES", 4),
|
||||||
|
patch("routers.knowledge.knowledge_vectorizer.enqueue") as enqueue,
|
||||||
|
):
|
||||||
|
response = client.post(
|
||||||
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
||||||
|
headers=context["owner_headers"],
|
||||||
|
files={"file": ("oversize.md", b"12345", "text/markdown")},
|
||||||
|
)
|
||||||
|
|
||||||
|
payload = response.json()
|
||||||
|
assert payload["code"] == 400
|
||||||
|
assert payload["message"] == "文件不能超过 50MB"
|
||||||
|
enqueue.assert_not_called()
|
||||||
|
assert not list((tmp_path / context["avatar"].id).glob("*"))
|
||||||
|
|
||||||
|
|
||||||
def test_background_vectorizer_commits_ready_document_and_chunks_together(
|
def test_background_vectorizer_commits_ready_document_and_chunks_together(
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
authorization_context,
|
authorization_context,
|
||||||
|
|||||||
@@ -117,7 +117,7 @@ location /api/ {
|
|||||||
proxy_set_header X-Forwarded-Proto $scheme;
|
proxy_set_header X-Forwarded-Proto $scheme;
|
||||||
proxy_buffering off;
|
proxy_buffering off;
|
||||||
proxy_read_timeout 300s;
|
proxy_read_timeout 300s;
|
||||||
client_max_body_size 20m;
|
client_max_body_size 100m;
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -23,6 +23,10 @@ http {
|
|||||||
|
|
||||||
root /usr/share/nginx/html;
|
root /usr/share/nginx/html;
|
||||||
index index.html;
|
index index.html;
|
||||||
|
# Keep the application gateway aligned with the production edge gateway.
|
||||||
|
# Without this Nginx rejects ordinary PDF uploads with HTTP 413 before
|
||||||
|
# FastAPI can return its user-facing file-size validation message.
|
||||||
|
client_max_body_size 100m;
|
||||||
|
|
||||||
# SPA 兜底(hash 路由下深链接也可正常加载)
|
# SPA 兜底(hash 路由下深链接也可正常加载)
|
||||||
location / {
|
location / {
|
||||||
|
|||||||
Reference in New Issue
Block a user