| import sqlite3 |
| import threading |
| from datetime import datetime |
| import json |
| from .config import SQLITE_DB_PATH |
|
|
| class TrainingBuffer: |
| def __init__(self): |
| self.db_path = SQLITE_DB_PATH |
| self.lock = threading.Lock() |
| self._init_db() |
|
|
| def _init_db(self): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| cursor = conn.cursor() |
| cursor.execute(''' |
| CREATE TABLE IF NOT EXISTS utterances ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| timestamp TEXT NOT NULL, |
| pictogram_seq TEXT NOT NULL, |
| confirmed_sentence TEXT NOT NULL, |
| persona_id TEXT NOT NULL, |
| context TEXT, |
| was_first_alternative INTEGER DEFAULT 1, |
| session_id TEXT |
| ) |
| ''') |
| conn.commit() |
| conn.close() |
|
|
| def log(self, pictogram_seq, confirmed_sentence, persona_id, context=None, was_first_alternative=True, session_id=None): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| cursor = conn.cursor() |
| timestamp = datetime.now().isoformat() |
| pictogram_seq_str = json.dumps(pictogram_seq) |
| cursor.execute(''' |
| INSERT INTO utterances (timestamp, pictogram_seq, confirmed_sentence, persona_id, context, was_first_alternative, session_id) |
| VALUES (?, ?, ?, ?, ?, ?, ?) |
| ''', (timestamp, pictogram_seq_str, confirmed_sentence, persona_id, context, 1 if was_first_alternative else 0, session_id)) |
| conn.commit() |
| conn.close() |
|
|
| def get_for_persona(self, persona_id): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| conn.row_factory = sqlite3.Row |
| cursor = conn.cursor() |
| cursor.execute('SELECT * FROM utterances WHERE persona_id = ?', (persona_id,)) |
| rows = cursor.fetchall() |
| conn.close() |
| return [dict(row) for row in rows] |
|
|
| def count(self, persona_id=None): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| cursor = conn.cursor() |
| if persona_id: |
| cursor.execute('SELECT COUNT(*) FROM utterances WHERE persona_id = ?', (persona_id,)) |
| else: |
| cursor.execute('SELECT COUNT(*) FROM utterances') |
| count = cursor.fetchone()[0] |
| conn.close() |
| return count |
|
|
| def export_jsonl(self, persona_id): |
| utterances = self.get_for_persona(persona_id) |
| return "\n".join([json.dumps(u) for u in utterances]) |
|
|
| def clear(self, persona_id): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| cursor = conn.cursor() |
| cursor.execute('DELETE FROM utterances WHERE persona_id = ?', (persona_id,)) |
| conn.commit() |
| count = conn.total_changes |
| conn.close() |
| return count |
|
|
| def get_stats(self): |
| with self.lock: |
| conn = sqlite3.connect(self.db_path) |
| cursor = conn.cursor() |
| cursor.execute('SELECT persona_id, COUNT(*) FROM utterances GROUP BY persona_id') |
| stats = dict(cursor.fetchall()) |
| conn.close() |
| return stats |
|
|
| training_buffer = TrainingBuffer() |
|
|
