309 lines
9.7 KiB
Python
309 lines
9.7 KiB
Python
"""mail_summary API — local LLM weekly mail work summary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import shutil
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from datetime import date
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel, Field
|
|
|
|
from config import CORS_ORIGINS, INCOMING_DIR, LLM_MODEL, UPLOAD_DIR
|
|
from db.pool import close_pool, init_pool, pool_enabled
|
|
from job_repository import create_job, ensure_tables, get_job, get_report, list_jobs
|
|
from llm_client import ping_llm
|
|
from pipeline import run_analysis_job
|
|
|
|
logger = logging.getLogger("mail_summary")
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
MailboxScope = Literal["inbox", "sent", "both"]
|
|
Mailbox = Literal["inbox", "sent"]
|
|
STAGING_DIR = UPLOAD_DIR / "staging"
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
|
|
STAGING_DIR.mkdir(parents=True, exist_ok=True)
|
|
INCOMING_DIR.mkdir(parents=True, exist_ok=True)
|
|
await init_pool()
|
|
await ensure_tables()
|
|
yield
|
|
await close_pool()
|
|
|
|
|
|
app = FastAPI(title="Mail Summary API", version="1.0.0", lifespan=lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=CORS_ORIGINS or ["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
class JobCreateResponse(BaseModel):
|
|
job: dict
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
status: str
|
|
database: str
|
|
model: str = Field(default="")
|
|
llm: dict = Field(default_factory=dict)
|
|
|
|
|
|
class UploadResponse(BaseModel):
|
|
upload_id: str
|
|
mailbox: str
|
|
filename: str
|
|
size_bytes: int
|
|
source: str # upload | incoming
|
|
|
|
|
|
def _staging_path(upload_id: str) -> Path:
|
|
path = (STAGING_DIR / upload_id).resolve()
|
|
if not str(path).startswith(str(STAGING_DIR.resolve())):
|
|
raise HTTPException(status_code=400, detail="invalid_upload_id")
|
|
return path
|
|
|
|
|
|
def _write_staging_meta(stage_dir: Path, meta: dict) -> None:
|
|
(stage_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
|
|
|
|
|
|
def _read_staging_meta(upload_id: str) -> dict:
|
|
stage_dir = _staging_path(upload_id)
|
|
meta_path = stage_dir / "meta.json"
|
|
if not meta_path.is_file():
|
|
raise HTTPException(status_code=404, detail="upload_not_found")
|
|
return json.loads(meta_path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def _attach_source(src: Path, stage_dir: Path, mailbox: str, filename: str) -> Path:
|
|
dest = stage_dir / f"{mailbox}_{filename}"
|
|
if dest.exists() or dest.is_symlink():
|
|
dest.unlink()
|
|
try:
|
|
dest.symlink_to(src)
|
|
except OSError:
|
|
shutil.copy2(src, dest)
|
|
return dest
|
|
|
|
|
|
def _find_incoming_match(filename: str, size_bytes: int) -> Path | None:
|
|
candidate = INCOMING_DIR / Path(filename).name
|
|
if not candidate.is_file():
|
|
return None
|
|
try:
|
|
if candidate.stat().st_size != size_bytes:
|
|
return None
|
|
except OSError:
|
|
return None
|
|
return candidate.resolve()
|
|
|
|
|
|
def _attach_staging_to_job(upload_id: str, job_dir: Path, mailbox: str) -> str:
|
|
"""Copy/link staged zip into a job dir, keeping staging for re-analysis."""
|
|
meta = _read_staging_meta(upload_id)
|
|
if meta.get("mailbox") != mailbox:
|
|
raise HTTPException(status_code=400, detail="upload_mailbox_mismatch")
|
|
stage_dir = _staging_path(upload_id)
|
|
filename = meta["filename"]
|
|
src = stage_dir / f"{mailbox}_{filename}"
|
|
if not src.exists() and not src.is_symlink():
|
|
raise HTTPException(status_code=404, detail="upload_file_missing")
|
|
dest = job_dir / f"{mailbox}_{filename}"
|
|
if dest.exists() or dest.is_symlink():
|
|
dest.unlink()
|
|
# Prefer symlink so multi-GB files are not duplicated.
|
|
try:
|
|
target = src.resolve() if src.is_symlink() else src
|
|
dest.symlink_to(target)
|
|
except OSError:
|
|
shutil.copy2(src, dest)
|
|
return filename
|
|
|
|
|
|
@app.get("/api/health", response_model=HealthResponse)
|
|
async def health() -> HealthResponse:
|
|
llm = await ping_llm()
|
|
return HealthResponse(
|
|
status="ok",
|
|
database="up" if pool_enabled() else "down",
|
|
model=LLM_MODEL,
|
|
llm=llm,
|
|
)
|
|
|
|
|
|
@app.post("/api/uploads/resolve", response_model=UploadResponse)
|
|
async def api_resolve_upload(
|
|
mailbox: Mailbox = Form(...),
|
|
filename: str = Form(...),
|
|
size_bytes: int = Form(...),
|
|
) -> UploadResponse:
|
|
"""If the selected file already exists on the incoming mount, stage it without re-upload."""
|
|
match = _find_incoming_match(filename, size_bytes)
|
|
if not match:
|
|
raise HTTPException(status_code=404, detail="incoming_match_not_found")
|
|
|
|
upload_id = uuid.uuid4().hex
|
|
stage_dir = _staging_path(upload_id)
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
safe_name = Path(filename).name
|
|
_attach_source(match, stage_dir, mailbox, safe_name)
|
|
meta = {
|
|
"upload_id": upload_id,
|
|
"mailbox": mailbox,
|
|
"filename": safe_name,
|
|
"size_bytes": size_bytes,
|
|
"source": "incoming",
|
|
}
|
|
_write_staging_meta(stage_dir, meta)
|
|
return UploadResponse(**meta)
|
|
|
|
|
|
@app.post("/api/uploads", response_model=UploadResponse)
|
|
async def api_upload_zip(
|
|
mailbox: Mailbox = Form(...),
|
|
file: UploadFile = File(...),
|
|
) -> UploadResponse:
|
|
name = (file.filename or "").lower()
|
|
if not name.endswith(".zip"):
|
|
raise HTTPException(status_code=400, detail="file_must_be_zip")
|
|
|
|
upload_id = uuid.uuid4().hex
|
|
stage_dir = _staging_path(upload_id)
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
safe_name = Path(file.filename or f"{mailbox}.zip").name
|
|
dest = stage_dir / f"{mailbox}_{safe_name}"
|
|
size = 0
|
|
try:
|
|
with dest.open("wb") as out:
|
|
while True:
|
|
chunk = await file.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
size += len(chunk)
|
|
except Exception:
|
|
shutil.rmtree(stage_dir, ignore_errors=True)
|
|
logger.exception("Upload failed")
|
|
raise HTTPException(status_code=500, detail="upload_failed") from None
|
|
|
|
meta = {
|
|
"upload_id": upload_id,
|
|
"mailbox": mailbox,
|
|
"filename": safe_name,
|
|
"size_bytes": size,
|
|
"source": "upload",
|
|
}
|
|
_write_staging_meta(stage_dir, meta)
|
|
return UploadResponse(**meta)
|
|
|
|
|
|
@app.delete("/api/uploads/{upload_id}")
|
|
async def api_delete_upload(upload_id: str) -> dict:
|
|
stage_dir = _staging_path(upload_id)
|
|
if stage_dir.is_dir():
|
|
shutil.rmtree(stage_dir, ignore_errors=True)
|
|
return {"ok": True, "upload_id": upload_id}
|
|
|
|
|
|
@app.get("/api/jobs")
|
|
async def api_list_jobs(limit: int = 20) -> dict:
|
|
if not pool_enabled():
|
|
raise HTTPException(status_code=503, detail="database_unavailable")
|
|
jobs = await list_jobs(limit=max(1, min(limit, 100)))
|
|
return {"jobs": jobs}
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}")
|
|
async def api_get_job(job_id: int) -> dict:
|
|
if not pool_enabled():
|
|
raise HTTPException(status_code=503, detail="database_unavailable")
|
|
job = await get_job(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job_not_found")
|
|
return {"job": job}
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}/report")
|
|
async def api_get_report(job_id: int) -> dict:
|
|
if not pool_enabled():
|
|
raise HTTPException(status_code=503, detail="database_unavailable")
|
|
job = await get_job(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job_not_found")
|
|
report = await get_report(job_id)
|
|
return {"job": job, "report": report}
|
|
|
|
|
|
@app.post("/api/jobs", response_model=JobCreateResponse)
|
|
async def api_create_job(
|
|
date_from: date = Form(...),
|
|
date_to: date = Form(...),
|
|
mailbox_scope: MailboxScope = Form("both"),
|
|
model: str = Form(""),
|
|
inbox_upload_id: str = Form(""),
|
|
sent_upload_id: str = Form(""),
|
|
) -> JobCreateResponse:
|
|
if not pool_enabled():
|
|
raise HTTPException(status_code=503, detail="database_unavailable")
|
|
if date_from > date_to:
|
|
raise HTTPException(status_code=400, detail="invalid_date_range")
|
|
|
|
inbox_upload_id = inbox_upload_id.strip()
|
|
sent_upload_id = sent_upload_id.strip()
|
|
|
|
if mailbox_scope in ("inbox", "both") and not inbox_upload_id:
|
|
raise HTTPException(status_code=400, detail="inbox_zip_required")
|
|
if mailbox_scope in ("sent", "both") and not sent_upload_id:
|
|
raise HTTPException(status_code=400, detail="sent_zip_required")
|
|
|
|
selected_model = (model or LLM_MODEL).strip() or LLM_MODEL
|
|
job_token = uuid.uuid4().hex
|
|
job_dir = UPLOAD_DIR / job_token
|
|
job_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
inbox_name = (
|
|
_attach_staging_to_job(inbox_upload_id, job_dir, "inbox") if inbox_upload_id else None
|
|
)
|
|
sent_name = (
|
|
_attach_staging_to_job(sent_upload_id, job_dir, "sent") if sent_upload_id else None
|
|
)
|
|
job = await create_job(
|
|
mailbox_scope=mailbox_scope,
|
|
date_from=date_from,
|
|
date_to=date_to,
|
|
model=selected_model,
|
|
inbox_zip_name=inbox_name,
|
|
sent_zip_name=sent_name,
|
|
)
|
|
final_dir = UPLOAD_DIR / str(job["id"])
|
|
if final_dir.exists():
|
|
shutil.rmtree(final_dir)
|
|
job_dir.rename(final_dir)
|
|
except HTTPException:
|
|
shutil.rmtree(job_dir, ignore_errors=True)
|
|
raise
|
|
except Exception:
|
|
shutil.rmtree(job_dir, ignore_errors=True)
|
|
logger.exception("Failed to create analysis job")
|
|
raise HTTPException(status_code=500, detail="job_create_failed") from None
|
|
|
|
asyncio.create_task(run_analysis_job(job["id"]))
|
|
return JobCreateResponse(job=job)
|