Update database.py

database.py

@@ -1,25 +1,27 @@
-"""Database for conversations"""
+"""Database for conversations and distillation data"""
 
 import sqlite3
+from datetime import datetime
 from typing import List, Dict
 from config import DATABASE_PATH
 
 
 class VedaDatabase:
-    """Database handler"""
-    
+    """Database handler with distillation support"""
+
     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()
-        
+
+        # Regular conversations table
         cursor.execute('''
             CREATE TABLE IF NOT EXISTS conversations (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -29,40 +31,66 @@ class VedaDatabase:
                 feedback INTEGER DEFAULT 0
             )
         ''')
-        
+
+        # Distillation data table (teacher responses)
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS distillation_data (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+                user_input TEXT NOT NULL,
+                teacher_response TEXT NOT NULL,
+                student_response TEXT,
+                used_for_training BOOLEAN DEFAULT 0,
+                quality_score REAL DEFAULT 0
+            )
+        ''')
+
+        # Training history
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS training_history (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+                training_type TEXT,
+                samples_used INTEGER,
+                epochs INTEGER,
+                final_loss REAL
+            )
+        ''')
+
         conn.commit()
         conn.close()
-    
+
+    # ===== Conversations =====
     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
@@ -70,28 +98,129 @@ class VedaDatabase:
             ORDER BY timestamp DESC
             LIMIT ?
         ''', (limit,))
-        
+
         rows = cursor.fetchall()
         conn.close()
-        
+
         return [dict(row) for row in rows]
-    
+
+    # ===== Distillation =====
+    def save_distillation_data(
+        self,
+        user_input: str,
+        teacher_response: str,
+        student_response: str = None,
+        quality_score: float = 0.0
+    ) -> int:
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute('''
+            INSERT INTO distillation_data 
+            (user_input, teacher_response, student_response, quality_score)
+            VALUES (?, ?, ?, ?)
+        ''', (user_input, teacher_response, student_response, quality_score))
+
+        data_id = cursor.lastrowid
+        conn.commit()
+        conn.close()
+
+        return data_id
+
+    def get_unused_distillation_data(self, limit: int = 500) -> List[Dict]:
+        """Get teacher responses not yet used for training"""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute('''
+            SELECT id, user_input, teacher_response
+            FROM distillation_data
+            WHERE used_for_training = 0
+            ORDER BY timestamp DESC
+            LIMIT ?
+        ''', (limit,))
+
+        rows = cursor.fetchall()
+        conn.close()
+
+        return [dict(row) for row in rows]
+
+    def mark_distillation_used(self, ids: List[int]):
+        """Mark distillation data as used for training"""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        placeholders = ",".join("?" * len(ids))
+        cursor.execute(f'''
+            UPDATE distillation_data
+            SET used_for_training = 1
+            WHERE id IN ({placeholders})
+        ''', ids)
+
+        conn.commit()
+        conn.close()
+
+    def get_distillation_count(self) -> Dict:
+        """Get count of distillation data"""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute('SELECT COUNT(*) FROM distillation_data')
+        total = cursor.fetchone()[0]
+
+        cursor.execute('SELECT COUNT(*) FROM distillation_data WHERE used_for_training = 0')
+        unused = cursor.fetchone()[0]
+
+        cursor.execute('SELECT COUNT(*) FROM distillation_data WHERE used_for_training = 1')
+        used = cursor.fetchone()[0]
+
+        conn.close()
+
+        return {"total": total, "unused": unused, "used": used}
+
+    # ===== Stats =====
     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]
-        
+
+        distill = self.get_distillation_count()
+
+        conn.close()
+
+        return {
+            "total": total,
+            "positive": positive,
+            "negative": negative,
+            "distillation_total": distill["total"],
+            "distillation_unused": distill["unused"],
+        }
+
+    def save_training_history(
+        self,
+        training_type: str,
+        samples_used: int,
+        epochs: int,
+        final_loss: float
+    ):
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute('''
+            INSERT INTO training_history (training_type, samples_used, epochs, final_loss)
+            VALUES (?, ?, ?, ?)
+        ''', (training_type, samples_used, epochs, final_loss))
+
+        conn.commit()
         conn.close()
-        
-        return {'total': total, 'positive': positive, 'negative': negative}
 
 
 db = VedaDatabase()