database.py

from dataclasses import dataclass
from typing import Any

import psycopg2 as pg


@dataclass
class Entry:
    model: str
    card_file: str
    topic: str
    question: str
    guess: bool
    ground_truth: bool
    reasoning: str
    confidence: int


class PostgreSQL:
    hostname: str
    dbname: str
    user: str
    password: str
    port: int
    #            pg.connection
    _connection: Any

    def __init__(self, hostname: str = "129.153.49.94", port: int = 5432,
                 dbname: str = "llm-eval", user: str = 'dev', password: str = "llmeval"):
        self.hostname = hostname
        self.port = port
        self.dbname = dbname
        self.user = user
        self.password = password
        self._connection = self._connect()

    def insert(self, entry: Entry):
        conn = self.get_connection()
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO responses (model, card_file, topic, question, guess, ground_truth, reasoning, confidence)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """,
                (entry.model, entry.card_file, entry.topic, entry.question, entry.guess, entry.ground_truth,
                 entry.reasoning, entry.confidence)
            )
        conn.commit()

    def get_connection(self) -> Any:
        """Get the connection to the PostgreSQL database server"""
        if self._connection.closed:
            self._connection = self._connect()
        return self._connection

    def _connect(self) -> Any:
        """Connect to the PostgreSQL database server"""
        conn = pg.connect(host=self.hostname, port=self.port, dbname=self.dbname, user=self.user,
                          password=self.password)
        with conn.cursor() as cur:
            cur.execute("""SET search_path TO %s;""", ("website",))
        return conn

    def close(self):
        # self._connection: pg.connection
        self._connection.close()