LORKT_AI_Service/main.py

import os
import shutil
import uuid
from typing import Annotated
from fastapi import FastAPI, Request, UploadFile, Depends, Form, File, Cookie
from fastapi.templating import Jinja2Templates
from fastapi.security import OAuth2PasswordRequestForm
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from starlette.middleware.sessions import SessionMiddleware
from service import auth, reports
from service.db_requests import get_user
from service import structs
from service.models import User

app = FastAPI()
templates = Jinja2Templates(directory="client/templates")
app.mount("/static", StaticFiles(directory="client/static"), name="static")

app.add_middleware(
    SessionMiddleware,
    secret_key="CHANGEME",
    session_cookie="session"
)


@app.get("/", response_class=RedirectResponse)
async def main_page(access_token: Annotated[str | None, Cookie()] = None):
    if access_token:
        try: 
            await auth.get_current_user(access_token)
        except:
            return RedirectResponse("/login")
        return RedirectResponse("/upload-study")
    else:
        return RedirectResponse("/login")


@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    access_token = request.cookies.get("access_token")
    if access_token:
        try:
            await auth.get_current_user(access_token)
            return RedirectResponse("/upload-study", status_code=303)
        except:
            pass
    response = templates.TemplateResponse("login.html", {"request": request})
    response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, private"
    response.headers["Pragma"] = "no-cache"
    response.headers["Expires"] = "0"
    return response


@app.post("/login")
async def login_submit(request: Request, 
                       login_data: OAuth2PasswordRequestForm = Depends()):
    user = get_user(login_data.username)
    if user is None or not auth.verify_password(login_data.password,
                                                user.hashed_password):
        return templates.TemplateResponse("login.html",
                                          {"request": request, "error": True})
    response = RedirectResponse("/upload-study", status_code=303)
    response.set_cookie(
        key = "access_token",
        value = auth.create_access_token(user.username),
        httponly=True
    )
    return response


@app.get("/upload-study", dependencies=[Depends(auth.get_current_user)])
async def upload_page(request: Request):
    return templates.TemplateResponse("upload.html", {"request": request})


@app.post("/upload-study")
async def study_submit(
    request: Request,
    user: User = Depends(auth.get_current_user),
    file: UploadFile = File(None),
    demo_filename: str = Form(None),
    pathology: str = Form(...)
):
    user_dir = os.path.join("data", user.username)
    os.makedirs(user_dir, exist_ok=True)
    fpath = os.path.join(user_dir, str(uuid.uuid4()) + ".dcm")

    if file and file.filename:
        with open(fpath, 'wb') as buffer:
            buffer.write(await file.read())
    else:
        demo_source_path = os.path.join("data/demo", pathology, demo_filename)
        shutil.copy2(demo_source_path, fpath)

    try:
        prediction = reports.make_reports(pathology, fpath, user.username)
    except (structs.TagError, structs.ImagesError) as e:
        error_msg = e.msg + ". Загрузите другой файл."
        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return JSONResponse({"error": error_msg}, status_code=400)
        else:
            return templates.TemplateResponse("upload.html", {"request": request, "error": error_msg})
    except Exception as e:
        error_msg = "Ошибка анализа. Загрузите другой файл."
        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return JSONResponse({"error": error_msg}, status_code=500)
        else:
            return templates.TemplateResponse("upload.html", {"request": request, "error": error_msg})

    request.session["prediction"] = prediction._asdict()
    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
        return JSONResponse({"redirect": "/ai-result"})
    else:
        return RedirectResponse("/ai-result", status_code=303)


@app.get("/ai-result", dependencies=[Depends(auth.get_current_user)])
async def result_page(request: Request):
    prediction = request.session["prediction"]
    return templates.TemplateResponse("result.html", {
        "request": request,
        "prediction": prediction
    })


@app.get("/download/{filename}")
async def download_report(filename: str, 
                          user: User = Depends(auth.get_current_user)):
    return FileResponse(f"data/{user.username}/reports/{filename}")