db/db.py

import sqlite3
from contextlib import contextmanager

DATABASE_PATH = "database.db"  # Chemin de la base de données

# Gestionnaire de contexte pour gérer les connexions
@contextmanager
def connect_db():
    conn = sqlite3.connect(DATABASE_PATH)
    try:
        yield conn
    finally:
        conn.commit()
        conn.close()

# Classe responsable de la base de données
class DatabaseHandler:
    def __init__(self, db_path=DATABASE_PATH):
        self.db_path = db_path
        self._initialize_db()

    def _initialize_db(self):
        with connect_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''CREATE TABLE IF NOT EXISTS prompts (
                                id INTEGER PRIMARY KEY AUTOINCREMENT,
                                type TEXT NOT NULL CHECK(type IN ('installation', 'difficult')),
                                num INTEGER NOT NULL,
                                title TEXT NOT NULL,
                                prompt TEXT NOT NULL
                              );''')
            if not self._is_data_present(cursor):
                self._load_default_data(cursor)

    def _is_data_present(self, cursor):
        cursor.execute("SELECT COUNT(*) FROM prompts;")
        return cursor.fetchone()[0] > 0

    def _load_default_data(self, cursor):
        data = [
            {"type": "installation", "num": 1, "title": "Contexte et objectifs", "prompt": "context_objectives"},
            {"type": "installation", "num": 2, "title": "Statut social, juridique et fiscal", "prompt": "social_legal_fiscal_status"},
            {"type": "installation", "num": 3, "title": "Moyens humains", "prompt": "human_resources"},
            {"type": "installation", "num": 4, "title": "Moyens de production", "prompt": "production_resources"},
            {"type": "installation", "num": 5, "title": "Production par atelier", "prompt": "workshop_production"},
            {"type": "difficult", "num": 11, "title": "Contexte et objectifs", "prompt": "difficult  context_objectives"},
            {"type": "difficult", "num": 12, "title": "Statut social, juridique et fiscal", "prompt": "difficult  social_legal_fiscal_status"},
            {"type": "difficult", "num": 13, "title": "Moyens humains", "prompt": "difficult  human_resources"},
            {"type": "difficult", "num": 14, "title": "Moyens de production", "prompt": "difficult  production_resources"},
            {"type": "difficult", "num": 15, "title": "Production par atelier", "prompt": "difficult  workshop_production"}
        ]
        for record in data:
            cursor.execute(
                "INSERT INTO prompts (type, num, title, prompt) VALUES (:type, :num, :title, :prompt);",
                record
            )
        cursor.connection.commit()

    def _rows_to_dicts(self, rows):
        result = []
        for r in rows:
            result.append({
                "id": r[0],
                "type": r[1],
                "num": r[2],
                "title": r[3],
                "prompt_system": r[4],
            })
        return result

    def get_prompts(self):
        with connect_db() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM prompts;")
            rows = cursor.fetchall()
            return self._rows_to_dicts(rows)

    def get_prompt_by_filters(self, type_=None, num=None):
        with connect_db() as conn:
            cursor = conn.cursor()
            query = "SELECT * FROM prompts"
            conditions = []
            params = []
            
            if type_:
                conditions.append("type = ?")
                params.append(type_)
            if num:
                conditions.append("num = ?")
                params.append(num)
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            
            cursor.execute(query, params)
            rows = cursor.fetchall()
            return self._rows_to_dicts(rows)

    def add_prompt(self, type_, num, title, prompt):
        with connect_db() as conn:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO prompts (type, num, title, prompt) VALUES (?, ?, ?);", (type_, num, title, prompt))

    def update_prompt(self, prompt_id, type_=None, num=None, title=None, prompt=None):
        with connect_db() as conn:
            cursor = conn.cursor()
            query = "UPDATE prompts SET "
            fields = []
            params = []
            if type_:
                fields.append("type = ?")
                params.append(type_)
            if title:
                fields.append("title = ?")
                params.append(title)
            if prompt is not None :
                fields.append("prompt = ?")
                params.append(prompt)
            query += ", ".join(fields) + " WHERE num = ?;"
            params.append(prompt_id)
            cursor.execute(query, params)