beck chnages

app.py

@@ -11,7 +11,7 @@ from groq_client import GroqClient
 from database import get_db
 from auth import register_user, login_user, token_required
 from crisis_detector import check_for_crisis, get_crisis_response, get_crisis_resources
-from prompts import STAGE_GOALS, SUMMARIES, NATURAL_ENDINGS
+from prompts import NATURAL_ENDINGS
 from exercises import get_exercise_for_group
 from wearable import wearable_bp
 from admin import admin_bp
@@ -159,182 +159,172 @@ def get_sessions():
 @token_required
 def chat():
     """
-    Main chat endpoint.
-    Handles classification, stage tracking, crisis detection, and response generation.
+    Main chat endpoint using Beck's 3-Agent Protocol.
     """
     try:
         user = request.current_user
         db = get_db()
-        
+
         data = request.json
         user_message = data.get("message", "").strip()
         session_id = data.get("session_id")
-        conversation_history = data.get("conversation_history", [])  # Get history from client
-
-        if not user_message:
-            return jsonify({"error": "Message is required"}), 400
-        
-        if not session_id:
-            return jsonify({"error": "session_id is required"}), 400
-        
+        conversation_history = data.get("conversation_history", [])
+
+        if not user_message or not session_id:
+            return jsonify({"error": "Message and session_id required"}), 400
+
         # Get session
         session = db.get_session(session_id)
-        if not session:
+        if not session or session["user_id"] != user["id"]:
             return jsonify({"error": "Session not found"}), 404
-        
-        if session["user_id"] != user["id"]:
-            return jsonify({"error": "Unauthorized"}), 403
-        
+
         # ========== CRISIS DETECTION ==========
         is_crisis, trigger_word = check_for_crisis(user_message)
-        
         if is_crisis:
-            # Flag in database
-            db.flag_crisis(
-                user_id=user["id"],
-                user_name=user["name"],
-                user_email=user["email"],
-                session_id=session_id,
-                message_content=user_message,
-                trigger_word=trigger_word
-            )
-            
-            # Save user message
-            db.add_message(session_id, user["id"], "user", user_message)
-            
-            # Get crisis response
+            db.flag_crisis(user["id"], user["name"], user["email"],
+                          session_id, user_message, trigger_word)
             crisis_response = get_crisis_response(user["name"])
-            db.add_message(session_id, user["id"], "assistant", crisis_response)
-            
             return jsonify({
                 "response": crisis_response,
                 "is_crisis": True,
-                "crisis_resources": get_crisis_resources(),
-                "session_state": session["state"],
-                "current_stage": session["current_stage"]
+                "crisis_resources": get_crisis_resources()
             })
-        
+
         # ========== NATURAL EXIT CHECK ==========
         if is_natural_exit(user_message):
-            response = handle_natural_exit(session, user, db, session_id)
-            # Only save to DB if crisis - skip normal messages
-            return jsonify(response)
-        
-        # ========== CLASSIFICATION ==========
-        locked_group = session["locked_group"]
-        current_stage = session["current_stage"]
-        
-        if session["state"] == "WAITING_FOR_PROBLEM" or locked_group == "G0":
-            # Run classifier
+            return jsonify(handle_natural_exit(session, user, db, session_id))
+
+        # ========== GET OR CREATE BECK SESSION ==========
+        beck_data = db.get_beck_session(session_id)
+
+        if not beck_data:
+            # First message - check if distorted
             classification = classifier.classify(user_message)
-            detected_group = classification["group"]
-            confidence = classification["confidence"]
-            
-            if detected_group == "G0":
-                # No distortion - listen supportively
-                db.update_session(session_id, locked_group="G0")
 
+            if classification["group"] == "G0":
+                # No distortion - supportive listening
                 response_text = groq_client.generate_supportive_response(
                     user_message=user_message,
-                    conversation_history=conversation_history[-6:],  # From client
+                    conversation_history=conversation_history[-6:],
                     user_name=user["name"]
                 )
-
-                # Don't save to DB - only save crisis messages
-
                 return jsonify({
                     "response": response_text,
-                    "detected_group": "G0",
-                    "group_name": "No Distortion Detected",
-                    "confidence": confidence,
-                    "current_stage": None,
-                    "total_stages": None,
-                    "session_state": "LISTENING"
+                    "beck_state": None,
+                    "is_beck_protocol": False
                 })
-            
             else:
-                # Distortion detected - LOCK the group
-                db.update_session(
-                    session_id,
-                    locked_group=detected_group,
-                    stages_reached=1,
-                    messages_in_current_stage=0
-                )
-                locked_group = detected_group
-                current_stage = 1
-        
-        # ========== THERAPEUTIC RESPONSE ==========
-        
-        # Increment message counter
-        messages_in_stage = db.increment_stage_messages(session_id)
-        
-        # Get stage info
-        stage_info = STAGE_GOALS[locked_group][current_stage]
-        
-        # Generate response (use client-provided history, not DB)
-        llm_response = groq_client.generate_therapeutic_response(
-            user_message=user_message,
-            conversation_history=conversation_history[-6:],  # Last 6 messages from client
-            user_name=user["name"],
-            user_context=user["context"],
-            detected_group=locked_group,
-            current_stage=current_stage,
-            stage_goal=stage_info["goal"],
-            stage_instruction=stage_info["instruction"]
-        )
-        
-        response_text = llm_response["response"]
-        should_advance = llm_response.get("advance_to_next_stage", False)
-        
-        # ========== STAGE ADVANCEMENT LOGIC ==========
-        MIN_MESSAGES_PER_STAGE = 3
-        MAX_MESSAGES_PER_STAGE = 8
-        
-        if messages_in_stage < MIN_MESSAGES_PER_STAGE:
-            should_advance = False
-        elif current_stage < 3 and messages_in_stage >= MAX_MESSAGES_PER_STAGE:
-            should_advance = True
-        # Stage 3: NO auto-advance
-        
-        # Handle advancement
-        if should_advance and current_stage < 3:
-            current_stage += 1
-            db.update_session(session_id, stages_reached=current_stage, messages_in_current_stage=0)
-        
-        # Check if session complete
-        session_complete = should_advance and current_stage >= 3
-        
-        if session_complete:
-            db.update_session(session_id, completed=1)
-            summary = SUMMARIES[locked_group]
-            response_text = response_text + "\n\n---\n\n" + summary
-
-        # Don't save to DB - only crisis messages are saved
-
-        # Group display names
-        group_names = {
-            "G1": "Binary & Absolute Thinking",
-            "G2": "Overgeneralized Beliefs",
-            "G3": "Attention & Focus Bias",
-            "G4": "Emotion-Driven Reasoning"
-        }
-        
+                # Distortion detected - START BECK PROTOCOL
+                db.create_beck_session(session_id)
+                db.update_beck_state(session_id, "VALIDATE",
+                                    original_thought=user_message)
+                beck_data = db.get_beck_session(session_id)
+
+        # ========== BECK PROTOCOL STATE MACHINE ==========
+        from prompts import BECK_STATES, AGENT1_STATES, AGENT3_STATES, get_next_state, get_field_to_save
+
+        current_state = beck_data["beck_state"]
+
+        # Determine which agent handles this state
+        if current_state in AGENT1_STATES:
+            # Agent 1: Warm Questioner
+            response_text = groq_client.agent1_warm_questioner(
+                current_state=current_state,
+                user_message=user_message,
+                beck_data=beck_data,
+                user_name=user["name"],
+                conversation_history=conversation_history[-6:]
+            )
+
+            # Save user's response to appropriate field
+            field_to_save = get_field_to_save(current_state)
+            next_state = get_next_state(current_state)
+
+            # Extract rating if this is a rating state
+            if field_to_save and ('rating' in field_to_save or 'intensity' in field_to_save):
+                rating = extract_rating(user_message)
+                if rating is not None:
+                    db.update_beck_state(session_id, next_state, **{field_to_save: rating})
+                else:
+                    db.update_beck_state(session_id, next_state)
+            elif field_to_save:
+                db.update_beck_state(session_id, next_state, **{field_to_save: user_message})
+            else:
+                db.update_beck_state(session_id, next_state)
+
+        elif current_state == "SUMMARIZING":
+            # Agent 2: Clinical Summarizer (internal)
+            clinical_summary = groq_client.agent2_clinical_summarizer(beck_data)
+
+            # Move to Agent 3
+            db.update_beck_state(session_id, "DELIVER_REFRAME")
+
+            # Get updated beck_data and call Agent 3
+            beck_data = db.get_beck_session(session_id)
+            response_text = groq_client.agent3_treatment_agent(
+                current_state="DELIVER_REFRAME",
+                user_message=user_message,
+                beck_data=beck_data,
+                clinical_summary=clinical_summary,
+                user_name=user["name"],
+                conversation_history=conversation_history[-6:]
+            )
+
+        elif current_state in AGENT3_STATES:
+            # Agent 3: Treatment Agent
+            # Need clinical summary - regenerate if needed
+            clinical_summary = groq_client.agent2_clinical_summarizer(beck_data)
+
+            response_text = groq_client.agent3_treatment_agent(
+                current_state=current_state,
+                user_message=user_message,
+                beck_data=beck_data,
+                clinical_summary=clinical_summary,
+                user_name=user["name"],
+                conversation_history=conversation_history[-6:]
+            )
+
+            # Save response and advance state
+            field_to_save = get_field_to_save(current_state)
+            next_state = get_next_state(current_state)
+
+            if field_to_save and ('rating' in field_to_save or 'belief' in field_to_save or 'intensity' in field_to_save):
+                rating = extract_rating(user_message)
+                if rating is not None:
+                    db.update_beck_state(session_id, next_state, **{field_to_save: rating})
+                else:
+                    db.update_beck_state(session_id, next_state)
+            elif field_to_save:
+                db.update_beck_state(session_id, next_state, **{field_to_save: user_message})
+            else:
+                if next_state:
+                    db.update_beck_state(session_id, next_state)
+                else:
+                    # Protocol complete
+                    db.complete_beck_session(session_id)
+
+        else:
+            # Fallback
+            response_text = groq_client.generate_supportive_response(
+                user_message, conversation_history[-6:], user["name"]
+            )
+
+        # Get updated state for response
+        beck_data = db.get_beck_session(session_id)
+
         return jsonify({
             "response": response_text,
-            "detected_group": locked_group,
-            "group_name": group_names.get(locked_group, locked_group),
-            "current_stage": current_stage,
-            "total_stages": 3,
-            "stage_name": STAGE_GOALS[locked_group][current_stage]["name"],
-            "session_state": "COMPLETED" if session_complete else "IN_PROGRESS",
-            "session_complete": session_complete
+            "beck_state": beck_data["beck_state"] if beck_data else None,
+            "is_beck_protocol": beck_data is not None,
+            "protocol_complete": beck_data["beck_state"] == "COMPLETE" if beck_data else False,
+            "belief_improvement": calculate_improvement(beck_data) if beck_data else None
         })
-        
+
     except Exception as e:
         print(f"Error in /api/chat: {str(e)}")
         import traceback
         traceback.print_exc()
-        return jsonify({"error": "Something went wrong. Please try again."}), 500
+        return jsonify({"error": "Something went wrong"}), 500
 
 
 # ==================== EXERCISE ROUTES ====================
@@ -454,33 +444,59 @@ def is_natural_exit(message: str) -> bool:
 def handle_natural_exit(session: dict, user: dict, db, session_id: str) -> dict:
     """Handle natural conversation exit."""
     import random
-    
+
     endings = NATURAL_ENDINGS
     response = random.choice(endings).replace("{name}", user["name"])
-    
-    locked_group = session.get("locked_group")
-    current_stage = session.get("current_stage")
-    
-    # If exiting from Stage 3, add summary
-    if current_stage == 3 and locked_group and locked_group != "G0":
-        summary = SUMMARIES.get(locked_group, "")
-        if summary:
-            response = summary + "\n\n" + response
-    
-    # Update session and stats
-    db.update_session(session_id, completed=1 if current_stage == 3 else 0)
-    db.update_user_stats_on_session_end(user["id"], locked_group)
-    
+
+    # Check if Beck protocol was completed
+    beck_data = db.get_beck_session(session_id)
+    if beck_data and beck_data.get("beck_state") == "COMPLETE":
+        # Get improvement stats
+        improvement = calculate_improvement(beck_data)
+        if improvement:
+            improvement_msg = "You made real progress today! "
+            if improvement.get('belief_change', 0) >= 10:
+                improvement_msg += f"Your belief shifted by {improvement['belief_change']}%. "
+            if improvement.get('emotion_change', 0) >= 10:
+                improvement_msg += f"Your emotion intensity dropped by {improvement['emotion_change']}%. "
+            response = improvement_msg + "\n\n" + response
+
+    # Update session
+    db.update_session(session_id, completed=1)
+
     return {
         "response": response,
-        "session_state": "COMPLETED" if current_stage == 3 else "ENDED",
-        "detected_group": locked_group,
-        "current_stage": current_stage,
+        "session_state": "COMPLETED",
         "natural_exit": True,
-        "session_complete": current_stage == 3,
-        "suggest_exercise": current_stage == 3 and locked_group != "G0"
+        "session_complete": True
     }
 
 
+def extract_rating(message: str) -> int:
+    """Extract a 0-100 rating from user message."""
+    import re
+    # Look for numbers
+    numbers = re.findall(r'\d+', message)
+    for num in numbers:
+        n = int(num)
+        if 0 <= n <= 100:
+            return n
+    return None
+
+
+def calculate_improvement(beck_data: dict) -> dict:
+    """Calculate belief and emotion improvement."""
+    if not beck_data:
+        return None
+
+    result = {}
+    if beck_data.get('initial_belief_rating') and beck_data.get('final_belief_rating'):
+        result['belief_change'] = beck_data['initial_belief_rating'] - beck_data['final_belief_rating']
+    if beck_data.get('initial_emotion_intensity') and beck_data.get('final_emotion_intensity'):
+        result['emotion_change'] = beck_data['initial_emotion_intensity'] - beck_data['final_emotion_intensity']
+
+    return result if result else None
+
+
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)

database.py

@@ -176,6 +176,52 @@ class Database:
             ON device_keys(api_key)
         """)
 
+        # Beck sessions table for cognitive restructuring protocol
+        self.conn.execute("""
+            CREATE TABLE IF NOT EXISTS beck_sessions (
+                session_id TEXT PRIMARY KEY,
+
+                -- Current state in the protocol
+                beck_state TEXT DEFAULT 'VALIDATE',
+
+                -- Phase 1: Capture
+                original_thought TEXT,
+                initial_belief_rating INTEGER,
+                emotion TEXT,
+                initial_emotion_intensity INTEGER,
+
+                -- Phase 2: Discovery (6 Questions)
+                q1_evidence_for TEXT,
+                q1_evidence_against TEXT,
+                q2_alternative TEXT,
+                q3_worst TEXT,
+                q3_best TEXT,
+                q3_realistic TEXT,
+                q4_effect TEXT,
+                q5_friend TEXT,
+                q6_action TEXT,
+
+                -- Phase 3: Reframe
+                adaptive_thought TEXT,
+                new_thought_belief INTEGER,
+
+                -- Phase 4: Measure
+                final_belief_rating INTEGER,
+                final_emotion_intensity INTEGER,
+
+                -- Phase 5: Action
+                action_plan TEXT,
+
+                -- Metadata
+                started_at TEXT DEFAULT CURRENT_TIMESTAMP,
+                completed_at TEXT,
+                belief_improvement INTEGER,
+                emotion_improvement INTEGER,
+
+                FOREIGN KEY (session_id) REFERENCES sessions(id)
+            )
+        """)
+
         self.conn.commit()
     
     # ==================== USER OPERATIONS ====================
@@ -986,6 +1032,107 @@ class Database:
             "G4": counts.get("G4", 0)
         }
 
+    # ==================== BECK SESSION OPERATIONS ====================
+
+    def create_beck_session(self, session_id: str) -> dict:
+        """Initialize a Beck session when distortion is detected."""
+        try:
+            self.conn.execute(
+                """INSERT INTO beck_sessions (session_id, beck_state)
+                   VALUES (?, 'VALIDATE')""",
+                (session_id,)
+            )
+            self.conn.commit()
+            return {"session_id": session_id, "beck_state": "VALIDATE"}
+        except Exception as e:
+            print(f"Error creating Beck session: {e}")
+            return None
+
+    def get_beck_session(self, session_id: str) -> dict:
+        """Get current Beck session state and all data."""
+        result = self.conn.execute(
+            """SELECT session_id, beck_state, original_thought, initial_belief_rating,
+                      emotion, initial_emotion_intensity, q1_evidence_for, q1_evidence_against,
+                      q2_alternative, q3_worst, q3_best, q3_realistic, q4_effect, q5_friend,
+                      q6_action, adaptive_thought, new_thought_belief, final_belief_rating,
+                      final_emotion_intensity, action_plan, started_at, completed_at,
+                      belief_improvement, emotion_improvement
+               FROM beck_sessions WHERE session_id = ?""",
+            (session_id,)
+        ).fetchone()
+
+        if result:
+            return {
+                "session_id": result[0],
+                "beck_state": result[1],
+                "original_thought": result[2],
+                "initial_belief_rating": result[3],
+                "emotion": result[4],
+                "initial_emotion_intensity": result[5],
+                "q1_evidence_for": result[6],
+                "q1_evidence_against": result[7],
+                "q2_alternative": result[8],
+                "q3_worst": result[9],
+                "q3_best": result[10],
+                "q3_realistic": result[11],
+                "q4_effect": result[12],
+                "q5_friend": result[13],
+                "q6_action": result[14],
+                "adaptive_thought": result[15],
+                "new_thought_belief": result[16],
+                "final_belief_rating": result[17],
+                "final_emotion_intensity": result[18],
+                "action_plan": result[19],
+                "started_at": result[20],
+                "completed_at": result[21],
+                "belief_improvement": result[22],
+                "emotion_improvement": result[23]
+            }
+        return None
+
+    def update_beck_state(self, session_id: str, new_state: str, **fields):
+        """Update state and save any new field values."""
+        # Start with state update
+        updates = {"beck_state": new_state}
+        updates.update(fields)
+
+        set_clause = ", ".join(f"{k} = ?" for k in updates.keys())
+        values = tuple(list(updates.values()) + [session_id])
+
+        self.conn.execute(
+            f"UPDATE beck_sessions SET {set_clause} WHERE session_id = ?",
+            values
+        )
+        self.conn.commit()
+
+    def complete_beck_session(self, session_id: str):
+        """Mark session complete, calculate improvements."""
+        # Get current data
+        beck_data = self.get_beck_session(session_id)
+        if not beck_data:
+            return
+
+        # Calculate improvements
+        belief_improvement = None
+        emotion_improvement = None
+
+        if beck_data.get('initial_belief_rating') and beck_data.get('final_belief_rating'):
+            belief_improvement = beck_data['initial_belief_rating'] - beck_data['final_belief_rating']
+
+        if beck_data.get('initial_emotion_intensity') and beck_data.get('final_emotion_intensity'):
+            emotion_improvement = beck_data['initial_emotion_intensity'] - beck_data['final_emotion_intensity']
+
+        # Update completion status
+        self.conn.execute(
+            """UPDATE beck_sessions SET
+               completed_at = CURRENT_TIMESTAMP,
+               belief_improvement = ?,
+               emotion_improvement = ?
+               WHERE session_id = ?""",
+            (belief_improvement, emotion_improvement, session_id)
+        )
+        self.conn.commit()
+
 
 # Singleton instance
 _db_instance = None

groq_client.py

@@ -1,6 +1,5 @@
 """
-Groq API Client
-Generates therapeutic responses using LLaMA 3
+Groq API Client - 3 Agent System for Beck CBT Protocol
 """
 
 import json
@@ -9,82 +8,96 @@ from groq import Groq
 
 class GroqClient:
     """
-    Client for Groq API to generate therapeutic responses.
-    Uses LLaMA 3.3 70B for warm, empathetic responses.
+    3-Agent system for Beck's Cognitive Restructuring Protocol.
+    - Agent 1: Warm Questioner (validation + 6 questions)
+    - Agent 2: Clinical Summarizer (internal analysis)
+    - Agent 3: Treatment Agent (reframe + measurement + action)
     """
-    
+
     MODEL = "llama-3.3-70b-versatile"
-    
+
     def __init__(self, api_key: str):
         if not api_key:
             raise ValueError("GROQ_API_KEY is required")
         self.client = Groq(api_key=api_key)
         print(f"✅ Groq client initialized with model: {self.MODEL}")
-    
-    def generate_therapeutic_response(
+
+    # ==================== AGENT 1: WARM QUESTIONER ====================
+
+    def agent1_warm_questioner(
         self,
+        current_state: str,
         user_message: str,
-        conversation_history: list,
+        beck_data: dict,
         user_name: str,
-        user_context: str,
-        detected_group: str,
-        current_stage: int,
-        stage_goal: str,
-        stage_instruction: str
-    ) -> dict:
+        conversation_history: list = None
+    ) -> str:
         """
-        Generate a therapeutic response based on the current stage.
-        
+        Agent 1: Asks Beck's questions one at a time in a warm, supportive way.
+
+        Args:
+            current_state: Current Beck protocol state (e.g., "Q1_EVIDENCE_FOR")
+            user_message: User's latest response
+            beck_data: All collected data so far
+            user_name: User's name for personalization
+            conversation_history: Recent messages for context
+
         Returns:
-            dict with 'response' and 'advance_to_next_stage' flag
+            Response string to send to user
         """
-        
-        # Build conversation context
-        history_text = self._format_history(conversation_history)
-        
-        system_prompt = f"""You are a warm, caring CBT companion talking to {user_name}, a {user_context}.
-You speak like a supportive friend - not a clinical therapist. Use casual, warm language.
-
-CURRENT SITUATION:
-- Detected thinking pattern: {detected_group}
-- Current stage: {current_stage} of 3
-- Stage goal: {stage_goal}
-
-YOUR TASK FOR THIS STAGE:
-{stage_instruction}
-
-IMPORTANT RULES:
-1. Use {user_name}'s name occasionally (not every message)
-2. Keep responses SHORT (1-2 sentences max) - be concise and warm
-3. Be validating first, then gently guide
-4. Use 1-2 emojis per message to feel warm (like 💙 🌟 ✨ 🤗) - not more
-5. Don't lecture or be preachy - keep it light
-6. Ask ONE short question at a time, if any
-
-STAGE ADVANCEMENT GUIDELINES:
-{"STAGE 3 IS SPECIAL: This is the final stage. Do NOT try to end or wrap up. Just keep being supportive and helpful until " + user_name + " naturally wants to stop. Set advance_to_next_stage to FALSE always." if current_stage == 3 else '''Set "advance_to_next_stage" to TRUE when ''' + user_name + ''' shows reasonable progress:
-- They're engaging with the idea (not just "yeah" but some reflection)
-- They've acknowledged or explored the concept
-- They seem ready to move forward
-- You've made your key therapeutic point for this stage
-
-Set "advance_to_next_stage" to FALSE when:
-- ''' + user_name + ''' is still venting or needs to be heard
-- They seem confused or resistant
-- You haven't yet introduced the core idea of this stage
-- They gave a very short/dismissive response
-
-Don't demand perfection. Progress over perfection.'''}
-
-Respond in JSON format:
-{{"response": "your message here", "advance_to_next_stage": {"false" if current_stage == 3 else "true/false"}}}"""
-
-        user_prompt = f"""CONVERSATION SO FAR:
+        from prompts import BECK_STATES
+
+        state_info = BECK_STATES.get(current_state, {})
+        instruction = state_info.get("instruction", "")
+        example = state_info.get("example", "")
+
+        # Build context of what we know so far
+        context_parts = []
+        if beck_data.get("original_thought"):
+            context_parts.append(f"Their thought: \"{beck_data['original_thought']}\"")
+        if beck_data.get("initial_belief_rating"):
+            context_parts.append(f"Belief rating: {beck_data['initial_belief_rating']}%")
+        if beck_data.get("emotion"):
+            context_parts.append(f"Emotion: {beck_data['emotion']}")
+        if beck_data.get("initial_emotion_intensity"):
+            context_parts.append(f"Emotion intensity: {beck_data['initial_emotion_intensity']}%")
+
+        context_text = "\n".join(context_parts) if context_parts else "Just starting the conversation."
+
+        history_text = self._format_history(conversation_history) if conversation_history else ""
+
+        system_prompt = f"""You are a warm, caring CBT companion named Aria talking to {user_name}.
+You are NOT a clinical therapist - you're like a supportive friend who knows CBT techniques.
+
+CURRENT STATE: {current_state}
+YOUR TASK: {instruction}
+EXAMPLE RESPONSE: "{example}"
+
+WHAT WE KNOW SO FAR:
+{context_text}
+
+CRITICAL RULES:
+1. Be WARM and GENTLE - like a caring friend, not a clinician
+2. Keep responses SHORT (1-3 sentences max) - be concise and warm
+3. Ask only ONE thing at a time
+4. Use {user_name}'s name occasionally (not every message)
+5. Use 1-2 emojis to feel warm (💙 🌟 ✨)
+6. VALIDATE before asking - acknowledge what they shared
+7. Don't lecture or be preachy - keep it light
+8. If they give a number (like "80"), acknowledge it warmly before moving on
+9. Don't explain what you're doing - just do it naturally
+
+IMPORTANT:
+- If they share something painful, validate it first
+- Match their emotional tone
+- Be conversational, not scripted"""
+
+        user_prompt = f"""RECENT CONVERSATION:
 {history_text}
 
 {user_name}'s latest message: "{user_message}"
 
-Respond as the CBT companion. Remember: warm, friendly, like a supportive friend."""
+Respond as Aria, the warm CBT companion. Remember: short, warm, one question at a time."""
 
         try:
             response = self.client.chat.completions.create(
@@ -94,45 +107,191 @@ Respond as the CBT companion. Remember: warm, friendly, like a supportive friend
                     {"role": "user", "content": user_prompt}
                 ],
                 temperature=0.7,
+                max_tokens=300
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            print(f"Agent 1 error: {e}")
+            return f"I hear you, {user_name}. 💙 Tell me more about that."
+
+    # ==================== AGENT 2: CLINICAL SUMMARIZER ====================
+
+    def agent2_clinical_summarizer(self, beck_data: dict) -> dict:
+        """
+        Agent 2: Analyzes all responses and prepares summary for Agent 3.
+        This is INTERNAL - user never sees this output.
+
+        Args:
+            beck_data: All collected data from Agent 1
+
+        Returns:
+            dict with analysis for Agent 3
+        """
+        system_prompt = """You are a clinical analyzer for a CBT system.
+Analyze the patient's responses and extract key patterns for the treatment agent.
+
+Your task:
+1. Identify CONTRADICTIONS (where their evidence contradicts their belief)
+2. Extract their OWN WISDOM (what they'd tell a friend)
+3. Note the COST of their belief (how it hurts them)
+4. Pull their REALISTIC prediction (usually more balanced than their fear)
+5. Suggest elements for a reframe using THEIR OWN WORDS
+
+Respond ONLY in JSON format."""
+
+        user_prompt = f"""Analyze these responses:
+
+ORIGINAL THOUGHT: {beck_data.get('original_thought', 'Not captured')}
+INITIAL BELIEF: {beck_data.get('initial_belief_rating', '?')}%
+EMOTION: {beck_data.get('emotion', '?')} at {beck_data.get('initial_emotion_intensity', '?')}%
+
+EVIDENCE FOR THE THOUGHT: {beck_data.get('q1_evidence_for', 'Not answered')}
+EVIDENCE AGAINST: {beck_data.get('q1_evidence_against', 'Not answered')}
+ALTERNATIVE EXPLANATION: {beck_data.get('q2_alternative', 'Not answered')}
+WORST CASE: {beck_data.get('q3_worst', 'Not answered')}
+BEST CASE: {beck_data.get('q3_best', 'Not answered')}
+REALISTIC CASE: {beck_data.get('q3_realistic', 'Not answered')}
+EFFECT OF BELIEVING: {beck_data.get('q4_effect', 'Not answered')}
+WHAT THEY'D TELL A FRIEND: {beck_data.get('q5_friend', 'Not answered')}
+WHAT THEY THINK THEY SHOULD DO: {beck_data.get('q6_action', 'Not answered')}
+
+Return JSON with:
+{{
+  "contradictions": ["list of contradictions between thought and evidence"],
+  "patient_wisdom": "what they said they'd tell a friend",
+  "cost_of_belief": "how believing this thought hurts them",
+  "realistic_prediction": "their realistic outcome prediction",
+  "reframe_elements": ["key phrases from their answers to use in reframe"],
+  "suggested_balanced_thought": "a balanced thought using their own words"
+}}"""
+
+        try:
+            response = self.client.chat.completions.create(
+                model=self.MODEL,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ],
+                temperature=0.3,
                 max_tokens=500,
                 response_format={"type": "json_object"}
             )
-            
-            result_text = response.choices[0].message.content
-            result = json.loads(result_text)
-            
-            return {
-                "response": result.get("response", "I hear you. Tell me more."),
-                "advance_to_next_stage": result.get("advance_to_next_stage", False)
-            }
-            
-        except json.JSONDecodeError:
-            # If JSON parsing fails, extract response text
-            return {
-                "response": response.choices[0].message.content,
-                "advance_to_next_stage": False
-            }
+            return json.loads(response.choices[0].message.content)
         except Exception as e:
-            print(f"Groq API error: {str(e)}")
+            print(f"Agent 2 error: {e}")
             return {
-                "response": f"I hear you, {user_name}. That sounds tough. Can you tell me more about what you're feeling?",
-                "advance_to_next_stage": False
+                "contradictions": [],
+                "patient_wisdom": beck_data.get('q5_friend', ''),
+                "cost_of_belief": beck_data.get('q4_effect', ''),
+                "realistic_prediction": beck_data.get('q3_realistic', ''),
+                "reframe_elements": [],
+                "suggested_balanced_thought": "Things are hard right now, but one moment doesn't define everything."
             }
-    
+
+    # ==================== AGENT 3: TREATMENT AGENT ====================
+
+    def agent3_treatment_agent(
+        self,
+        current_state: str,
+        user_message: str,
+        beck_data: dict,
+        clinical_summary: dict,
+        user_name: str,
+        conversation_history: list = None
+    ) -> str:
+        """
+        Agent 3: Delivers the reframe, measures change, creates action plan.
+
+        Args:
+            current_state: Current state (DELIVER_REFRAME, RATE_NEW_THOUGHT, etc.)
+            user_message: User's latest response
+            beck_data: All collected data
+            clinical_summary: Analysis from Agent 2
+            user_name: User's name
+            conversation_history: Recent messages
+
+        Returns:
+            Response string to send to user
+        """
+        from prompts import BECK_STATES
+
+        state_info = BECK_STATES.get(current_state, {})
+        instruction = state_info.get("instruction", "")
+
+        # Build improvement stats if available
+        improvement_text = ""
+        if beck_data.get('initial_belief_rating') and beck_data.get('final_belief_rating'):
+            belief_change = beck_data['initial_belief_rating'] - beck_data['final_belief_rating']
+            improvement_text += f"Belief dropped from {beck_data['initial_belief_rating']}% to {beck_data['final_belief_rating']}% ({belief_change}% improvement). "
+        if beck_data.get('initial_emotion_intensity') and beck_data.get('final_emotion_intensity'):
+            emotion_change = beck_data['initial_emotion_intensity'] - beck_data['final_emotion_intensity']
+            improvement_text += f"Emotion dropped from {beck_data['initial_emotion_intensity']}% to {beck_data['final_emotion_intensity']}% ({emotion_change}% improvement)."
+
+        history_text = self._format_history(conversation_history) if conversation_history else ""
+
+        system_prompt = f"""You are Aria, a warm CBT companion helping {user_name} complete cognitive restructuring.
+
+CURRENT STATE: {current_state}
+YOUR TASK: {instruction}
+
+ORIGINAL THOUGHT: "{beck_data.get('original_thought', '')}"
+INITIAL BELIEF: {beck_data.get('initial_belief_rating', '?')}%
+EMOTION: {beck_data.get('emotion', '?')} at {beck_data.get('initial_emotion_intensity', '?')}%
+
+CLINICAL ANALYSIS:
+- Contradictions found: {clinical_summary.get('contradictions', [])}
+- Their own wisdom: "{clinical_summary.get('patient_wisdom', '')}"
+- Cost of belief: "{clinical_summary.get('cost_of_belief', '')}"
+- Realistic prediction: "{clinical_summary.get('realistic_prediction', '')}"
+- Suggested reframe: "{clinical_summary.get('suggested_balanced_thought', '')}"
+
+{f"IMPROVEMENT SO FAR: {improvement_text}" if improvement_text else ""}
+
+CRITICAL RULES:
+1. Use THEIR OWN WORDS when possible - this makes the reframe feel like theirs
+2. Be warm and celebratory about any progress
+3. Keep responses SHORT but meaningful
+4. If delivering reframe, list what THEY discovered first
+5. For re-ratings, acknowledge the number warmly
+6. 10%+ improvement = SUCCESS - celebrate it!
+7. Use 1-2 emojis (💙 🌟 ✨)
+8. Don't be clinical or lecture-y"""
+
+        user_prompt = f"""RECENT CONVERSATION:
+{history_text}
+
+{user_name}'s latest message: "{user_message}"
+
+Respond as Aria for the {current_state} state."""
+
+        try:
+            response = self.client.chat.completions.create(
+                model=self.MODEL,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ],
+                temperature=0.7,
+                max_tokens=400
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            print(f"Agent 3 error: {e}")
+            return f"You've done really good work here, {user_name}. 💙"
+
+    # ==================== SUPPORTIVE RESPONSE (G0 - No Distortion) ====================
+
     def generate_supportive_response(
         self,
         user_message: str,
         conversation_history: list,
         user_name: str
     ) -> str:
-        """
-        Generate a supportive response for G0 (no distortion) cases.
-        Just listen and be supportive, gently ask if anything else is bothering them.
-        """
-        
+        """For G0 cases - just listen supportively, no intervention."""
+
         history_text = self._format_history(conversation_history)
-        
-        system_prompt = f"""You are a warm, caring companion talking to {user_name}.
+
+        system_prompt = f"""You are Aria, a warm companion talking to {user_name}.
 They're sharing something with you, and right now they don't seem caught in negative thinking.
 
 YOUR TASK:
@@ -140,13 +299,13 @@ YOUR TASK:
 2. Gently ask if there's anything else on their mind
 3. Keep it friendly and brief
 
-IMPORTANT:
+RULES:
 - Use {user_name}'s name occasionally
 - Keep responses SHORT (1-2 sentences max)
-- Use 1-2 emojis to feel warm (like 💙 🌟 ✨ 🤗)
+- Use 1-2 emojis (💙 🌟 ✨)
 - Don't lecture - just be a supportive friend"""
 
-        user_prompt = f"""CONVERSATION SO FAR:
+        user_prompt = f"""CONVERSATION:
 {history_text}
 
 {user_name}'s latest message: "{user_message}"
@@ -161,48 +320,23 @@ Respond warmly as their supportive companion."""
                     {"role": "user", "content": user_prompt}
                 ],
                 temperature=0.7,
-                max_tokens=300
+                max_tokens=200
             )
-            
             return response.choices[0].message.content
-            
         except Exception as e:
-            print(f"Groq API error: {str(e)}")
-            return f"I hear you, {user_name}. Thanks for sharing that with me. Is there anything else on your mind?"
-    
+            print(f"Supportive response error: {e}")
+            return f"I hear you, {user_name}. Thanks for sharing that with me. 💙"
+
+    # ==================== HELPER METHODS ====================
+
     def _format_history(self, history: list) -> str:
-        """Format conversation history for the prompt."""
+        """Format conversation history for prompts."""
         if not history:
             return "(This is the start of the conversation)"
-        
-        formatted = []
-        for msg in history:
-            role = "User" if msg["role"] == "user" else "Companion"
-            formatted.append(f"{role}: {msg['content']}")
-        
-        return "\n".join(formatted)
 
+        formatted = []
+        for msg in history[-6:]:  # Last 6 messages
+            role = "User" if msg.get("role") == "user" else "Aria"
+            formatted.append(f"{role}: {msg.get('content', '')}")
 
-# Test if run directly
-if __name__ == "__main__":
-    import os
-    
-    api_key = os.environ.get("GROQ_API_KEY")
-    if api_key:
-        client = GroqClient(api_key)
-        
-        result = client.generate_therapeutic_response(
-            user_message="I failed my exam and I feel like I'll never succeed at anything",
-            conversation_history=[],
-            user_name="Max",
-            user_context="college student",
-            detected_group="G2",
-            current_stage=1,
-            stage_goal="Make underlying assumption explicit",
-            stage_instruction="Help Max identify the prediction or assumption he's making. Gently point out the overgeneralization."
-        )
-        
-        print("Response:", result["response"])
-        print("Advance:", result["advance_to_next_stage"])
-    else:
-        print("Set GROQ_API_KEY to test")
+        return "\n".join(formatted)

prompts.py

@@ -1,108 +1,183 @@
 """
-Prompts and Content
-Stage goals, summaries, and other static content
+Beck CBT Protocol - States and Prompts
+Based on Beck Institute Thought Record Worksheet (J. Beck, 2020)
 """
 
-# Stage goals for each distortion group
-STAGE_GOALS = {
-    "G1": {
-        1: {
-            "name": "Recognize the Pattern",
-            "goal": "Help them see the all-or-nothing thinking",
-            "instruction": "Gently point out that they're using extreme/absolute language (always, never, completely, totally). Help them notice the black-and-white framing without being preachy."
-        },
-        2: {
-            "name": "Find the Gray",
-            "goal": "Explore the middle ground",
-            "instruction": "Help them see that reality usually exists on a spectrum. Ask about exceptions or times when things weren't all bad/all good. Guide them to find nuance."
-        },
-        3: {
-            "name": "Balanced View",
-            "goal": "Develop a more balanced perspective",
-            "instruction": "Help them form a more balanced statement that acknowledges both struggles AND strengths/exceptions. Celebrate any shift toward nuanced thinking."
-        }
-    },
-    "G2": {
-        1: {
-            "name": "Spot the Assumption",
-            "goal": "Identify the overgeneralization or prediction",
-            "instruction": "Help them notice they're making a broad conclusion from limited evidence, or predicting the future/reading minds. Gently highlight the assumption."
-        },
-        2: {
-            "name": "Examine the Evidence",
-            "goal": "Look at actual evidence",
-            "instruction": "Ask what evidence supports or contradicts their belief. Help them see they might be jumping to conclusions. Explore alternative explanations."
-        },
-        3: {
-            "name": "Reality Check",
-            "goal": "Form a more realistic view",
-            "instruction": "Help them develop a more evidence-based perspective. Acknowledge uncertainty is okay. Celebrate willingness to consider other possibilities."
-        }
-    },
-    "G3": {
-        1: {
-            "name": "Notice the Focus",
-            "goal": "Recognize the negative filter",
-            "instruction": "Help them see they might be focusing heavily on negatives while filtering out positives. Gently point out what they might be overlooking."
-        },
-        2: {
-            "name": "Widen the Lens",
-            "goal": "Bring in the full picture",
-            "instruction": "Ask about positives they might be discounting. Help them recall good things or strengths. Balance the perspective without dismissing their struggles."
-        },
-        3: {
-            "name": "Full Picture",
-            "goal": "See both positives and negatives",
-            "instruction": "Help them hold both the difficulties AND the good things. Develop a more complete view. Celebrate their ability to see the bigger picture."
-        }
-    },
-    "G4": {
-        1: {
-            "name": "Feelings vs Facts",
-            "goal": "Separate emotions from reality",
-            "instruction": "Help them see they might be treating feelings as facts, or taking things too personally, or using 'should' rigidly. Validate feelings while questioning conclusions."
-        },
-        2: {
-            "name": "Challenge the Logic",
-            "goal": "Examine the reasoning",
-            "instruction": "Explore whether their conclusion logically follows. Just because they feel something doesn't make it true. Help them see other interpretations."
-        },
-        3: {
-            "name": "New Perspective",
-            "goal": "Develop healthier self-talk",
-            "instruction": "Help them reframe with more compassion and flexibility. Replace rigid 'shoulds' with preferences. Celebrate growth in self-compassion."
-        }
-    }
-}
-
-# Session summaries (shown at end)
-SUMMARIES = {
-    "G1": """✨ **Great work today!**
-
-You practiced noticing all-or-nothing thinking and finding the middle ground. Remember: life rarely fits into neat boxes of "always" or "never." The truth usually lives somewhere in between.
-
-Keep noticing when you use extreme words, and gently remind yourself to look for the gray areas. 💙""",
-
-    "G2": """✨ **Great work today!**
-
-You practiced catching those assumptions and predictions your mind makes, and checking them against actual evidence. Our brains love to jump to conclusions, but we don't have to believe every thought!
+# All possible states in the Beck protocol
+BECK_STATES = {
+    # Agent 1: Warm Questioner States
+    "VALIDATE": {
+        "agent": 1,
+        "instruction": "Acknowledge their pain warmly. Don't ask questions yet. Just validate.",
+        "example": "That sounds really heavy. Thank you for sharing that with me. 💙",
+        "next_state": "RATE_BELIEF",
+        "saves_field": None
+    },
+    "RATE_BELIEF": {
+        "agent": 1,
+        "instruction": "Ask how much they believe the thought, 0-100%. Be warm about it.",
+        "example": "When that thought hits you — how much does it feel true? Like 0 being 'not at all' and 100 being 'completely certain'?",
+        "next_state": "CAPTURE_EMOTION",
+        "saves_field": "initial_belief_rating"
+    },
+    "CAPTURE_EMOTION": {
+        "agent": 1,
+        "instruction": "Ask what emotion comes up with this thought.",
+        "example": "And when you think this thought, what emotion shows up? Sadness? Anxiety? Anger? Something else?",
+        "next_state": "RATE_EMOTION",
+        "saves_field": "emotion"
+    },
+    "RATE_EMOTION": {
+        "agent": 1,
+        "instruction": "Ask how intense the emotion is, 0-100%.",
+        "example": "How intense is that feeling right now, 0 to 100?",
+        "next_state": "Q1_EVIDENCE_FOR",
+        "saves_field": "initial_emotion_intensity"
+    },
+    "Q1_EVIDENCE_FOR": {
+        "agent": 1,
+        "instruction": "Ask what happened that made them feel this way. Get the evidence FOR the thought.",
+        "example": "I'd like to understand better. What happened that made this thought feel so true?",
+        "next_state": "Q1_EVIDENCE_AGAINST",
+        "saves_field": "q1_evidence_for"
+    },
+    "Q1_EVIDENCE_AGAINST": {
+        "agent": 1,
+        "instruction": "Gently ask if there's any evidence against the thought. Any exceptions.",
+        "example": "That makes sense. Can I ask — has there been any time, even small, when things went differently? When you didn't fail, or it worked out okay?",
+        "next_state": "Q2_ALTERNATIVE",
+        "saves_field": "q1_evidence_against"
+    },
+    "Q2_ALTERNATIVE": {
+        "agent": 1,
+        "instruction": "Ask if there could be another explanation for what happened.",
+        "example": "If your best friend was in this exact situation, what other explanation might you suggest to them?",
+        "next_state": "Q3_WORST",
+        "saves_field": "q2_alternative"
+    },
+    "Q3_WORST": {
+        "agent": 1,
+        "instruction": "Ask what the worst case scenario would be.",
+        "example": "Let's imagine together for a moment. What's the worst that could happen here?",
+        "next_state": "Q3_BEST",
+        "saves_field": "q3_worst"
+    },
+    "Q3_BEST": {
+        "agent": 1,
+        "instruction": "Ask what the best case scenario would be.",
+        "example": "And what's the best that could happen?",
+        "next_state": "Q3_REALISTIC",
+        "saves_field": "q3_best"
+    },
+    "Q3_REALISTIC": {
+        "agent": 1,
+        "instruction": "Ask what will realistically/probably happen.",
+        "example": "And honestly — what do you think will probably happen? The realistic middle ground?",
+        "next_state": "Q4_EFFECT",
+        "saves_field": "q3_realistic"
+    },
+    "Q4_EFFECT": {
+        "agent": 1,
+        "instruction": "Ask how believing this thought affects them. Does it help or hurt?",
+        "example": "When you carry this thought around with you, how does it affect you? Does believing it help you in any way, or does it mostly hurt?",
+        "next_state": "Q5_FRIEND",
+        "saves_field": "q4_effect"
+    },
+    "Q5_FRIEND": {
+        "agent": 1,
+        "instruction": "Ask what they would tell a friend who had this exact thought.",
+        "example": "Last question in this part. If someone you really loved — a close friend, a sibling — came to you feeling exactly this way, with this exact thought... what would you say to them?",
+        "next_state": "Q6_ACTION",
+        "saves_field": "q5_friend"
+    },
+    "Q6_ACTION": {
+        "agent": 1,
+        "instruction": "Ask what would be good to do about the situation.",
+        "example": "And what do you think would be good to do about this situation?",
+        "next_state": "SUMMARIZING",
+        "saves_field": "q6_action"
+    },
 
-Keep asking yourself: "What's the evidence? Is there another way to see this?" 💙""",
+    # Agent 2: Internal summarization (no user interaction)
+    "SUMMARIZING": {
+        "agent": 2,
+        "instruction": "INTERNAL: Analyze all responses, extract contradictions, prepare for Agent 3.",
+        "next_state": "DELIVER_REFRAME",
+        "saves_field": None
+    },
 
-    "G3": """✨ **Great work today!**
+    # Agent 3: Treatment Agent States
+    "DELIVER_REFRAME": {
+        "agent": 3,
+        "instruction": "Summarize what they discovered and present a balanced thought using THEIR words.",
+        "example": "I really appreciate you walking through this with me. Here's what I noticed from YOUR words: [list their insights]. So maybe a thought that fits better might be: '[balanced thought]'. How does that land?",
+        "next_state": "RATE_NEW_THOUGHT",
+        "saves_field": "adaptive_thought"
+    },
+    "RATE_NEW_THOUGHT": {
+        "agent": 3,
+        "instruction": "Ask how much they believe the new balanced thought, 0-100%.",
+        "example": "How much do you believe that new thought? 0 to 100?",
+        "next_state": "RERATE_ORIGINAL",
+        "saves_field": "new_thought_belief"
+    },
+    "RERATE_ORIGINAL": {
+        "agent": 3,
+        "instruction": "Ask how much they NOW believe the original thought. Should be lower.",
+        "example": "Now thinking back to the original thought — '[original]' — how much do you believe that one now?",
+        "next_state": "RERATE_EMOTION",
+        "saves_field": "final_belief_rating"
+    },
+    "RERATE_EMOTION": {
+        "agent": 3,
+        "instruction": "Ask how intense the emotion is now. Should be lower.",
+        "example": "And that [emotion] you mentioned — how intense is it now, 0 to 100?",
+        "next_state": "ACTION_PLAN",
+        "saves_field": "final_emotion_intensity"
+    },
+    "ACTION_PLAN": {
+        "agent": 3,
+        "instruction": "Help them create a small behavioral experiment or action plan.",
+        "example": "You've made some real shifts today. Is there one small thing you could do this week to test this new perspective? Even something tiny counts.",
+        "next_state": "COMPLETE",
+        "saves_field": "action_plan"
+    },
+    "COMPLETE": {
+        "agent": 3,
+        "instruction": "Celebrate their progress. Summarize the improvement. End warmly.",
+        "example": "You started at [X]% and you're at [Y]% now — that's real movement. 💙 You did good work today. Remember, you can always come back when you need to talk.",
+        "next_state": None,
+        "saves_field": None
+    }
+}
 
-You practiced widening your focus to see the full picture, not just the negatives. It's natural to notice what's wrong, but there's usually more to the story.
+# States handled by each agent
+AGENT1_STATES = [
+    "VALIDATE", "RATE_BELIEF", "CAPTURE_EMOTION", "RATE_EMOTION",
+    "Q1_EVIDENCE_FOR", "Q1_EVIDENCE_AGAINST", "Q2_ALTERNATIVE",
+    "Q3_WORST", "Q3_BEST", "Q3_REALISTIC", "Q4_EFFECT", "Q5_FRIEND", "Q6_ACTION"
+]
 
-Keep practicing noticing the good stuff too, even the small things. They count! 💙""",
+AGENT3_STATES = [
+    "DELIVER_REFRAME", "RATE_NEW_THOUGHT", "RERATE_ORIGINAL",
+    "RERATE_EMOTION", "ACTION_PLAN", "COMPLETE"
+]
 
-    "G4": """✨ **Great work today!**
+def get_state_info(state: str) -> dict:
+    """Get info about a Beck state."""
+    return BECK_STATES.get(state, {})
 
-You practiced separating feelings from facts and being more flexible with yourself. Emotions are real, but they're not always accurate reporters of reality.
+def get_next_state(current_state: str) -> str:
+    """Get the next state in the protocol."""
+    state_info = BECK_STATES.get(current_state, {})
+    return state_info.get("next_state")
 
-Keep practicing self-compassion and replacing rigid "shoulds" with gentler preferences. You deserve kindness - especially from yourself! 💙"""
-}
+def get_field_to_save(state: str) -> str:
+    """Get which database field this state saves to."""
+    state_info = BECK_STATES.get(state, {})
+    return state_info.get("saves_field")
 
-# Natural conversation endings
+# Natural conversation endings (kept for compatibility)
 NATURAL_ENDINGS = [
     "Take care of yourself, {name}! I'm always here when you need to talk. 💙",
     "Glad I could help, {name}! Remember, you're doing better than you think. See you next time! 💙",
@@ -110,11 +185,3 @@ NATURAL_ENDINGS = [
     "Take it easy, {name}! You've got this. See you soon! 💙",
     "Catch you later, {name}! Remember, one step at a time. 💙"
 ]
-
-# Greeting messages
-GREETINGS = [
-    "Hey {name}! 👋 What's on your mind today?",
-    "Hi {name}! How are you doing? I'm here to listen.",
-    "Hey there, {name}! 💙 What would you like to talk about?",
-    "Hi {name}! Good to see you. What's going on?"
-]

test_beck_protocol.py

@@ -0,0 +1,253 @@
+"""
+Quick test script for Beck Protocol implementation
+Run this to verify the 3-agent system works
+"""
+
+import os
+from groq_client import GroqClient
+from classifier import DistortionClassifier
+
+# Test the components
+def test_classifier():
+    """Test that classifier still works for binary detection"""
+    print("=" * 60)
+    print("TEST 1: CLASSIFIER (Binary Detection)")
+    print("=" * 60)
+
+    classifier = DistortionClassifier()
+
+    test_cases = [
+        ("I had a nice day today", "G0"),
+        ("I'm such a failure, I always mess up", "G1 or G2"),
+        ("Nobody likes me, everyone thinks I'm weird", "G2 or G3"),
+    ]
+
+    for text, expected in test_cases:
+        result = classifier.classify(text)
+        print(f"\nText: {text[:50]}...")
+        print(f"  Predicted: {result['group']} ({result['confidence']:.2%})")
+        print(f"  Expected: {expected}")
+
+    print("\n✅ Classifier test complete\n")
+
+
+def test_agent1():
+    """Test Agent 1 - Warm Questioner"""
+    print("=" * 60)
+    print("TEST 2: AGENT 1 - Warm Questioner")
+    print("=" * 60)
+
+    api_key = os.environ.get("GROQ_API_KEY")
+    if not api_key:
+        print("❌ GROQ_API_KEY not set - skipping LLM tests")
+        return
+
+    client = GroqClient(api_key)
+
+    # Simulate VALIDATE state
+    beck_data = {
+        "original_thought": "I'm a complete failure",
+        "beck_state": "VALIDATE"
+    }
+
+    response = client.agent1_warm_questioner(
+        current_state="VALIDATE",
+        user_message="I'm a complete failure",
+        beck_data=beck_data,
+        user_name="TestUser",
+        conversation_history=[]
+    )
+
+    print(f"\nAgent 1 response (VALIDATE state):")
+    print(f"  {response}")
+
+    # Test Q1_EVIDENCE_FOR
+    beck_data["beck_state"] = "Q1_EVIDENCE_FOR"
+    beck_data["initial_belief_rating"] = 85
+    beck_data["emotion"] = "sadness"
+    beck_data["initial_emotion_intensity"] = 75
+
+    response = client.agent1_warm_questioner(
+        current_state="Q1_EVIDENCE_FOR",
+        user_message="85",
+        beck_data=beck_data,
+        user_name="TestUser",
+        conversation_history=[
+            {"role": "user", "content": "I'm a complete failure"},
+            {"role": "assistant", "content": "That sounds really heavy..."}
+        ]
+    )
+
+    print(f"\nAgent 1 response (Q1_EVIDENCE_FOR state):")
+    print(f"  {response}")
+
+    print("\n✅ Agent 1 test complete\n")
+
+
+def test_agent2():
+    """Test Agent 2 - Clinical Summarizer"""
+    print("=" * 60)
+    print("TEST 3: AGENT 2 - Clinical Summarizer")
+    print("=" * 60)
+
+    api_key = os.environ.get("GROQ_API_KEY")
+    if not api_key:
+        print("❌ GROQ_API_KEY not set - skipping LLM tests")
+        return
+
+    client = GroqClient(api_key)
+
+    # Simulate collected Beck data
+    beck_data = {
+        "original_thought": "I always fail at everything",
+        "initial_belief_rating": 85,
+        "emotion": "sadness",
+        "initial_emotion_intensity": 70,
+        "q1_evidence_for": "I failed my exam yesterday",
+        "q1_evidence_against": "Well, I did get a good grade on my report last week",
+        "q2_alternative": "Maybe I was just tired and didn't prepare enough",
+        "q3_worst": "I'll fail out of school",
+        "q3_best": "I'll do better next time",
+        "q3_realistic": "I'll probably pass the class even if I don't ace everything",
+        "q4_effect": "It makes me feel awful and want to give up",
+        "q5_friend": "I'd tell them one test doesn't define them",
+        "q6_action": "Maybe study more next time"
+    }
+
+    summary = client.agent2_clinical_summarizer(beck_data)
+
+    print("\nAgent 2 Clinical Summary:")
+    print(f"  Contradictions: {summary.get('contradictions', [])}")
+    print(f"  Patient Wisdom: {summary.get('patient_wisdom', 'N/A')}")
+    print(f"  Cost of Belief: {summary.get('cost_of_belief', 'N/A')}")
+    print(f"  Realistic Prediction: {summary.get('realistic_prediction', 'N/A')}")
+    print(f"  Suggested Reframe: {summary.get('suggested_balanced_thought', 'N/A')}")
+
+    print("\n✅ Agent 2 test complete\n")
+
+
+def test_agent3():
+    """Test Agent 3 - Treatment Agent"""
+    print("=" * 60)
+    print("TEST 4: AGENT 3 - Treatment Agent")
+    print("=" * 60)
+
+    api_key = os.environ.get("GROQ_API_KEY")
+    if not api_key:
+        print("❌ GROQ_API_KEY not set - skipping LLM tests")
+        return
+
+    client = GroqClient(api_key)
+
+    beck_data = {
+        "original_thought": "I always fail at everything",
+        "initial_belief_rating": 85,
+        "emotion": "sadness",
+        "initial_emotion_intensity": 70,
+    }
+
+    clinical_summary = {
+        "contradictions": ["Said 'always fails' but got good grade last week"],
+        "patient_wisdom": "One test doesn't define them",
+        "cost_of_belief": "Makes them feel awful and want to give up",
+        "realistic_prediction": "Will probably pass the class",
+        "suggested_balanced_thought": "I failed this test, but I've succeeded before and can improve"
+    }
+
+    response = client.agent3_treatment_agent(
+        current_state="DELIVER_REFRAME",
+        user_message="I'm ready to hear it",
+        beck_data=beck_data,
+        clinical_summary=clinical_summary,
+        user_name="TestUser",
+        conversation_history=[]
+    )
+
+    print(f"\nAgent 3 response (DELIVER_REFRAME):")
+    print(f"  {response}")
+
+    print("\n✅ Agent 3 test complete\n")
+
+
+def test_state_flow():
+    """Test state machine flow"""
+    print("=" * 60)
+    print("TEST 5: STATE MACHINE FLOW")
+    print("=" * 60)
+
+    from prompts import BECK_STATES, get_next_state, get_field_to_save
+
+    states_to_test = [
+        "VALIDATE",
+        "RATE_BELIEF",
+        "CAPTURE_EMOTION",
+        "Q1_EVIDENCE_FOR",
+        "SUMMARIZING",
+        "DELIVER_REFRAME",
+        "COMPLETE"
+    ]
+
+    for state in states_to_test:
+        next_state = get_next_state(state)
+        field = get_field_to_save(state)
+        print(f"\n{state}")
+        print(f"  → Next: {next_state}")
+        print(f"  → Saves to: {field}")
+
+    print("\n✅ State flow test complete\n")
+
+
+def test_rating_extraction():
+    """Test rating extraction from messages"""
+    print("=" * 60)
+    print("TEST 6: RATING EXTRACTION")
+    print("=" * 60)
+
+    import re
+
+    def extract_rating(message: str) -> int:
+        numbers = re.findall(r'\d+', message)
+        for num in numbers:
+            n = int(num)
+            if 0 <= n <= 100:
+                return n
+        return None
+
+    test_cases = [
+        ("85", 85),
+        ("I'd say about 70", 70),
+        ("100%", 100),
+        ("Maybe like... 50?", 50),
+        ("Not sure, somewhere around 0", 0),
+        ("I don't know", None),
+    ]
+
+    for text, expected in test_cases:
+        result = extract_rating(text)
+        status = "✅" if result == expected else "❌"
+        print(f"{status} '{text}' → {result} (expected {expected})")
+
+    print("\n✅ Rating extraction test complete\n")
+
+
+if __name__ == "__main__":
+    print("\n" + "=" * 60)
+    print("BECK PROTOCOL IMPLEMENTATION - TEST SUITE")
+    print("=" * 60 + "\n")
+
+    # Run tests
+    test_classifier()
+    test_agent1()
+    test_agent2()
+    test_agent3()
+    test_state_flow()
+    test_rating_extraction()
+
+    print("\n" + "=" * 60)
+    print("ALL TESTS COMPLETE")
+    print("=" * 60)
+    print("\nNext steps:")
+    print("1. Start the Flask server: python app.py")
+    print("2. Test via API with a real session")
+    print("3. Monitor database for beck_sessions table")
+    print("\n")