Spaces:
Sleeping
Sleeping
| from dataclasses import dataclass | |
| from typing import Any | |
| import psycopg2 as pg | |
| 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() | |