sezer-muhammed/mri-inference-api

1import os2from datetime import datetime3 4import libsql_client5 6_state: dict = {}7 8 9async def init():10    _state["client"] = libsql_client.create_client(11        url=os.environ["TURSO_DATABASE_URL"],12        auth_token=os.environ["TURSO_AUTH_TOKEN"],13    )14    await _state["client"].execute("""15        CREATE TABLE IF NOT EXISTS inference_results (16            id          INTEGER PRIMARY KEY AUTOINCREMENT,17            filename    TEXT    NOT NULL,18            model_name  TEXT    NOT NULL,19            centiloid   REAL,20            raw_output  REAL,21            label       TEXT,22            fold        INTEGER,23            created_at  TEXT    NOT NULL,24            UNIQUE (filename, model_name)25        )26    """)27    await _ensure_fold_column()28 29 30async def _ensure_fold_column() -> None:31    rs = await _state["client"].execute("PRAGMA table_info(inference_results)")32    cols = [row[1] for row in rs.rows]33    if "fold" not in cols:34        await _state["client"].execute(35            "ALTER TABLE inference_results ADD COLUMN fold INTEGER"36        )37 38 39async def save_result(40    filename: str,41    model_name: str,42    centiloid: float,43    raw_output: float,44    label: str | None = None,45    fold: int | None = None,46) -> int:47    now = datetime.utcnow().isoformat()48    rs = await _state["client"].execute(49        "INSERT INTO inference_results"50        " (filename, model_name, centiloid, raw_output, label, fold, created_at)"51        " VALUES (?, ?, ?, ?, ?, ?, ?)"52        " ON CONFLICT(filename, model_name) DO UPDATE SET"53        "   centiloid  = excluded.centiloid,"54        "   raw_output = excluded.raw_output,"55        "   label      = excluded.label,"56        "   fold       = COALESCE(excluded.fold, inference_results.fold),"57        "   created_at = excluded.created_at",58        [filename, model_name, centiloid, raw_output, label, fold, now],59    )60    return rs.last_insert_rowid61 62 63async def get_results_page(64    limit: int,65    offset: int,66    fold: int | None = None,67) -> tuple[int, list[dict]]:68    where = ""69    params: list = []70    if fold is not None:71        where = " WHERE fold = ?"72        params.append(fold)73 74    count_rs = await _state["client"].execute(75        f"SELECT COUNT(*) AS count FROM inference_results{where}",76        params,77    )78    total = int(count_rs.rows[0][0])79 80    rs = await _state["client"].execute(81        "SELECT id, filename, model_name, centiloid, raw_output, label, fold, created_at"82        f" FROM inference_results{where} ORDER BY created_at DESC LIMIT ? OFFSET ?",83        [*params, limit, offset],84    )85    cols = [c.name if hasattr(c, "name") else c for c in rs.columns]86    return total, [dict(zip(cols, row)) for row in rs.rows]87 88 89async def get_done_pairs() -> list[dict]:90    rs = await _state["client"].execute(91        "SELECT filename, model_name FROM inference_results"92    )93    cols = [c.name if hasattr(c, "name") else c for c in rs.columns]94    return [dict(zip(cols, row)) for row in rs.rows]95 96 97async def get_folds() -> list[int]:98    rs = await _state["client"].execute(99        "SELECT DISTINCT fold FROM inference_results"100        " WHERE fold IS NOT NULL ORDER BY fold ASC"101    )102    return [int(row[0]) for row in rs.rows]103 104 105async def update_result_fold_by_id(row_id: int, fold: int | None) -> bool:106    rs = await _state["client"].execute(107        "UPDATE inference_results SET fold = ? WHERE id = ?",108        [fold, row_id],109    )110    return rs.rows_affected > 0111 112 113async def update_result_fold_by_pair(114    filename: str,115    model_name: str,116    fold: int | None,117) -> bool:118    rs = await _state["client"].execute(119        "UPDATE inference_results SET fold = ? WHERE filename = ? AND model_name = ?",120        [fold, filename, model_name],121    )122    return rs.rows_affected > 0123 124 125async def close():126    if "client" in _state:127        await _state["client"].close()128