db.py

import pandas as pd


"""
Small Database for managing prompts, estimates, and votes in a simple arena-like application.
Les données sont stockées dans des fichiers CSV.
Il faut toujours se référer au fichier et pas à la variable en mémoire.
La mémoire est propre à chaque utilisateur. 
Pour être à jour des mises à jours des autres utilisateurs, il faut recharger les données depuis le fichier.
"""

DATABASE = {
    "prompts": {
        "filename": "database/prompts.csv",
        "columns": ["id", "name", "text"],
    },
    "estimates": {
        "filename": "database/estimates.csv",
        "columns": ["id", "prompt_id", "mu", "sigma"],
    },
    "votes": {
        "filename": "database/votes.csv",
        "columns": ["id", "winner_id", "loser_id"],
    },
}


def load(table_name: str) -> pd.DataFrame:
    try:
        result = pd.read_csv(DATABASE[table_name]["filename"])
    except FileNotFoundError:
        result = pd.DataFrame(columns=DATABASE[table_name]["columns"])
    return result


def insert(table_name: str, data: dict) -> None:
    """Inserts data into the specified table.
    Id is automatically generated .
    example:
        db.insert("prompts", {"name": "Test", "text": "This is a test prompt."})
    """
    df = load(table_name)
    data["id"] = 1 if df.empty else df["id"].max() + 1
    df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)
    df.to_csv(DATABASE[table_name]["filename"], index=False)


def update(table_name: str, row_id: int, data: dict) -> None:
    """
    Met à jour les estimations d'un prompt dans la table spécifiée par son id d'estimation.
    Exemple :
        db.update("estimates", 3, {"mu": 25.0, "sigma": 8.3})
    """
    df = load(table_name)
    idx = df.index[df["id"] == row_id]
    if not idx.empty:
        for key, value in data.items():
            df.loc[idx, key] = value
        df.to_csv(DATABASE[table_name]["filename"], index=False)


def replace(table_name: str, df: pd.DataFrame) -> None:
    """
    Remplace la totalité du contenu de la table spécifiée par le DataFrame fourni.
    Exemple :
        db.replace("prompts", df)
    """
    df.to_csv(DATABASE[table_name]["filename"], index=False)