ankahi_backend/training_buffer.py

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()