vent_agent_2.0/api/algorithm_routes.py

"""独立算法的 HTTP 适配。"""
from typing import Literal
from contextlib import asynccontextmanager

from fastapi import APIRouter, File, HTTPException, Query, Request, UploadFile
from fastapi.responses import Response, JSONResponse, StreamingResponse
from starlette.background import BackgroundTask
from starlette.concurrency import run_in_threadpool

import config
from algorithms.characteristic_curve_recognition.batch import BatchManager, MAX_BATCH_FILES, MAX_BATCH_BYTES
from algorithms.characteristic_curve_recognition.service import AlgorithmError, MAX_UPLOAD_BYTES, analyze_image


@asynccontextmanager
async def lifespan(app):
    manager = await run_in_threadpool(BatchManager, config.resolve_path(config.ALGORITHM_JOB_DIR))
    app.state.characteristic_batches = manager
    try:
        yield
    finally:
        await run_in_threadpool(manager.close)


router = APIRouter(prefix="/algorithms", tags=["algorithms"], lifespan=lifespan)
_BATCH_PATH = "/characteristic-curve-recognition/batches"


async def _call(func, *args, **kwargs):
    try:
        return await run_in_threadpool(func, *args, **kwargs)
    except AlgorithmError as exc:
        raise HTTPException(status_code=exc.status_code,
                            detail={"code": exc.code, "message": str(exc)},
                            headers={"Retry-After": "5"} if exc.status_code == 429 else None) from exc


def _links(job: dict) -> dict:
    base = f"/api/algorithms{_BATCH_PATH}/{job['id']}"
    job["status_url"] = base
    job["download_url"] = base + "/download"
    for item in job["items"]:
        item["result_url"] = f"{base}/items/{item['id']}/result" if item["status"] == "completed" else None
    return job


@router.post("/characteristic-curve-recognition", response_model=None, summary="特性曲线识别",
             responses={200: {"content": {"application/json": {}, "application/zip": {}}}})
async def characteristic_curve_recognition(
    file: UploadFile = File(..., description="PNG/JPEG，最多 20 MiB、1600 万像素"),
    detailed: bool = Query(False, description="生成各阶段检查图片，仅影响 ZIP 内可视化"),
    output: Literal["json", "zip"] = Query("json", description="JSON 结果或完整 ZIP 结果包"),
):
    try:
        content = await file.read(MAX_UPLOAD_BYTES + 1)
    finally:
        await file.close()
    result = await _call(analyze_image, content, detailed=detailed, output=output)
    if output == "zip":
        return Response(result, media_type="application/zip",
                        headers={"Content-Disposition": 'attachment; filename="characteristic-curve-recognition.zip"'})
    return result


@router.post(_BATCH_PATH, status_code=202, summary="提交特性曲线识别批量任务")
async def submit_batch(request: Request, files: list[UploadFile] = File(..., description="1～20 张图片，合计不超过 100 MiB"),
                       detailed: bool = Query(False)):
    uploads = []
    total = 0
    try:
        if not 1 <= len(files) <= MAX_BATCH_FILES:
            raise HTTPException(422, detail={"code": "invalid_batch_size", "message": "每批需上传 1～20 张图片"})
        for file in files:
            content = await file.read(MAX_UPLOAD_BYTES + 1)
            if len(content) > MAX_UPLOAD_BYTES:
                raise HTTPException(413, detail={"code": "image_too_large", "message": "单张图片不能超过 20 MiB"})
            total += len(content)
            if total > MAX_BATCH_BYTES:
                raise HTTPException(413, detail={"code": "batch_too_large", "message": "每批图片合计不能超过 100 MiB"})
            uploads.append((file.filename or "image", content))
    finally:
        for file in files:
            await file.close()
    job = _links(await _call(request.app.state.characteristic_batches.submit, uploads, detailed=detailed))
    return JSONResponse(job, status_code=202, headers={"Location": job["status_url"]})


@router.get(_BATCH_PATH + "/{job_id}", summary="查询批量任务和各图片进度")
async def batch_status(request: Request, job_id: str):
    return _links(await _call(request.app.state.characteristic_batches.status, job_id))


@router.get(_BATCH_PATH + "/{job_id}/items/{item_id}/result", summary="读取单张图片的 JSON 结果")
async def batch_item_result(request: Request, job_id: str, item_id: str):
    return await _call(request.app.state.characteristic_batches.result, job_id, item_id)


@router.get(_BATCH_PATH + "/{job_id}/download", summary="下载批量结果 ZIP")
async def batch_download(request: Request, job_id: str):
    archive = await _call(request.app.state.characteristic_batches.open_archive, job_id)

    def chunks():
        try:
            while chunk := archive.read(1024 * 1024):
                yield chunk
        finally:
            archive.close()

    return StreamingResponse(chunks(), media_type="application/zip", background=BackgroundTask(archive.close),
                             headers={"Content-Disposition": f'attachment; filename="characteristic-curves-{job_id}.zip"'})