Create distillation.py

distillation.py

@@ -0,0 +1,170 @@
+"""Knowledge Distillation Engine - Student learns from Teacher"""
+
+from typing import Optional, Tuple, Dict
+from teacher import teacher
+from database import db
+from config import (
+    DISTILLATION_ENABLED,
+    AUTO_LEARN_FROM_TEACHER,
+    MIN_RESPONSE_LENGTH,
+    MIN_SAMPLES_FOR_DISTILL_TRAINING,
+)
+
+
+class DistillationEngine:
+    """
+    Manages knowledge distillation from teacher to student.
+    
+    Modes:
+    1. AUTO: Always ask teacher, save responses, use student for speed
+    2. FALLBACK: Only ask teacher if student response is poor
+    3. COMPARE: Show both responses for comparison
+    """
+
+    def __init__(self):
+        self.teacher = teacher
+        self.mode = "fallback"  # "auto", "fallback", "compare"
+        self.teacher_call_count = 0
+        self.student_call_count = 0
+
+    def should_ask_teacher(self, student_response: str) -> bool:
+        """Decide if we should ask the teacher based on student response quality"""
+        if not DISTILLATION_ENABLED:
+            return False
+
+        if not self.teacher.is_available():
+            return False
+
+        # Heuristics for low-quality response
+        if not student_response:
+            return True
+
+        if len(student_response.strip()) < MIN_RESPONSE_LENGTH:
+            return True
+
+        # Check for error messages
+        low_quality_indicators = [
+            "I'm not sure",
+            "I don't know",
+            "Could you try rephrasing",
+            "Error:",
+            "not sure how to respond",
+        ]
+        for indicator in low_quality_indicators:
+            if indicator.lower() in student_response.lower():
+                return True
+
+        return False
+
+    def get_teacher_response(
+        self,
+        user_input: str,
+        conversation_history: list = None,
+        student_response: str = None,
+    ) -> Optional[str]:
+        """Get response from teacher and optionally save for training"""
+        
+        teacher_response = self.teacher.ask(
+            user_message=user_input,
+            conversation_history=conversation_history,
+        )
+
+        if teacher_response and AUTO_LEARN_FROM_TEACHER:
+            # Save for future training
+            db.save_distillation_data(
+                user_input=user_input,
+                teacher_response=teacher_response,
+                student_response=student_response,
+                quality_score=1.0,  # Teacher responses are high quality
+            )
+
+        if teacher_response:
+            self.teacher_call_count += 1
+
+        return teacher_response
+
+    def process_with_distillation(
+        self,
+        user_input: str,
+        student_response: str,
+        conversation_history: list = None,
+    ) -> Tuple[str, str]:
+        """
+        Process a response with potential teacher assistance.
+        
+        Returns:
+            Tuple of (final_response, source) where source is "student", "teacher", or "both"
+        """
+        self.student_call_count += 1
+
+        if self.mode == "auto":
+            # Always get teacher response for learning, but return student for speed
+            teacher_resp = self.get_teacher_response(
+                user_input, conversation_history, student_response
+            )
+            return student_response, "student"
+
+        elif self.mode == "fallback":
+            # Only ask teacher if student response is poor
+            if self.should_ask_teacher(student_response):
+                teacher_resp = self.get_teacher_response(
+                    user_input, conversation_history, student_response
+                )
+                if teacher_resp:
+                    return teacher_resp, "teacher"
+            return student_response, "student"
+
+        elif self.mode == "compare":
+            # Return both for comparison (useful for debugging/evaluation)
+            teacher_resp = self.get_teacher_response(
+                user_input, conversation_history, student_response
+            )
+            if teacher_resp:
+                combined = f"**🎓 Teacher (Dolphin):**\n{teacher_resp}\n\n---\n\n**🧠 Student (Veda):**\n{student_response}"
+                return combined, "both"
+            return student_response, "student"
+
+        return student_response, "student"
+
+    def set_mode(self, mode: str):
+        """Set distillation mode: 'auto', 'fallback', or 'compare'"""
+        if mode in ["auto", "fallback", "compare", "disabled"]:
+            self.mode = mode
+            return True
+        return False
+
+    def get_stats(self) -> Dict:
+        """Get distillation statistics"""
+        distill_data = db.get_distillation_count()
+        return {
+            "mode": self.mode,
+            "teacher_calls": self.teacher_call_count,
+            "student_calls": self.student_call_count,
+            "teacher_available": self.teacher.is_available(),
+            "distillation_samples": distill_data["total"],
+            "unused_samples": distill_data["unused"],
+            "ready_for_training": distill_data["unused"] >= MIN_SAMPLES_FOR_DISTILL_TRAINING,
+        }
+
+    def get_training_data(self) -> str:
+        """Get accumulated teacher responses as training data"""
+        unused = db.get_unused_distillation_data()
+
+        if not unused:
+            return ""
+
+        training_text = ""
+        for item in unused:
+            training_text += f"<USER> {item['user_input']}\n"
+            training_text += f"<ASSISTANT> {item['teacher_response']}\n\n"
+
+        return training_text
+
+    def mark_training_complete(self, ids: list):
+        """Mark distillation data as used after training"""
+        if ids:
+            db.mark_distillation_used(ids)
+
+
+# Global engine instance
+distillation_engine = DistillationEngine()