database.py

"""Database for conversations"""

import sqlite3
from typing import List, Dict
from config import DATABASE_PATH


class VedaDatabase:
    """Database handler"""
    
    def __init__(self):
        self._init_db()
    
    def _get_conn(self):
        conn = sqlite3.connect(DATABASE_PATH)
        conn.row_factory = sqlite3.Row
        return conn
    
    def _init_db(self):
        conn = self._get_conn()
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                user_input TEXT NOT NULL,
                assistant_response TEXT NOT NULL,
                feedback INTEGER DEFAULT 0
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def save_conversation(self, user_input: str, response: str) -> int:
        conn = self._get_conn()
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO conversations (user_input, assistant_response)
            VALUES (?, ?)
        ''', (user_input, response))
        
        conv_id = cursor.lastrowid
        conn.commit()
        conn.close()
        
        return conv_id
    
    def update_feedback(self, conv_id: int, feedback: int):
        conn = self._get_conn()
        cursor = conn.cursor()
        
        cursor.execute('''
            UPDATE conversations SET feedback = ? WHERE id = ?
        ''', (feedback, conv_id))
        
        conn.commit()
        conn.close()
    
    def get_good_conversations(self, limit: int = 100) -> List[Dict]:
        conn = self._get_conn()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_input, assistant_response
            FROM conversations
            WHERE feedback > 0
            ORDER BY timestamp DESC
            LIMIT ?
        ''', (limit,))
        
        rows = cursor.fetchall()
        conn.close()
        
        return [dict(row) for row in rows]
    
    def get_stats(self) -> Dict:
        conn = self._get_conn()
        cursor = conn.cursor()
        
        cursor.execute('SELECT COUNT(*) FROM conversations')
        total = cursor.fetchone()[0]
        
        cursor.execute('SELECT COUNT(*) FROM conversations WHERE feedback > 0')
        positive = cursor.fetchone()[0]
        
        cursor.execute('SELECT COUNT(*) FROM conversations WHERE feedback < 0')
        negative = cursor.fetchone()[0]
        
        conn.close()
        
        return {'total': total, 'positive': positive, 'negative': negative}


db = VedaDatabase()