yarab

chatbot.py

@@ -242,7 +242,8 @@ class MentalHealthChatbot:
         self.summary_model = pipeline(
             "summarization",
             model="philschmid/bart-large-cnn-samsum",
-            device=0 if self.device == "cuda" else -1
+            device=0 if self.device == "cuda" else -1,
+            model_kwargs={"cache_dir": CACHE_DIR}
         )
         logger.info("Summary model loaded successfully")
         
@@ -295,7 +296,10 @@ Response:"""
         # Setup embeddings for vector search
         self.embeddings = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-MiniLM-L6-v2",
-            model_kwargs={"device": self.device}
+            model_kwargs={
+                "device": self.device,
+                "cache_dir": CACHE_DIR
+            }
         )
         
         # Setup vector database for retrieving relevant past conversations
@@ -321,7 +325,7 @@ Response:"""
                 model="SamLowe/roberta-base-go_emotions",
                 top_k=None,
                 device_map="auto" if torch.cuda.is_available() else None,
-                cache_dir=CACHE_DIR,
+                model_kwargs={"cache_dir": CACHE_DIR},
                 local_files_only=False  # Ensure we download from Hugging Face
             )
         except Exception as e:
@@ -333,7 +337,7 @@ Response:"""
                     model="j-hartmann/emotion-english-distilroberta-base", 
                     return_all_scores=True,
                     device_map="auto" if torch.cuda.is_available() else None,
-                    cache_dir=CACHE_DIR,
+                    model_kwargs={"cache_dir": CACHE_DIR},
                     local_files_only=False  # Ensure we download from Hugging Face
                 )
             except Exception as e:

conversation_flow.py

@@ -1,5 +1,6 @@
 import logging
 import json
+import json5
 import time
 from datetime import datetime
 from typing import List, Dict, Any, Optional
@@ -8,6 +9,21 @@ from pydantic import BaseModel, Field
 # Configure logging
 logger = logging.getLogger(__name__)
 
+class PhaseTransitionResponse(BaseModel):
+    goals_progress: Dict[str, float]
+    should_transition: bool
+    next_phase: str
+    reasoning: str
+
+class SessionCharacteristics(BaseModel):
+    alliance_strength: float = Field(ge=0.0, le=1.0)
+    engagement_level: float = Field(ge=0.0, le=1.0)
+    emotional_pattern: str
+    cognitive_pattern: str
+    coping_mechanisms: List[str] = Field(min_items=2)
+    progress_quality: float = Field(ge=0.0, le=1.0)
+    recommended_focus: str
+
 class ConversationPhase(BaseModel):
     name: str
     description: str
@@ -15,7 +31,7 @@ class ConversationPhase(BaseModel):
     typical_duration: int  # in minutes
     started_at: Optional[str] = None  # ISO timestamp
     ended_at: Optional[str] = None  # ISO timestamp
-    completion_metrics: Dict[str, float] = {}  # e.g., {'goal_progress': 0.8}
+    completion_metrics: Dict[str, float] = Field(default_factory=dict)  # e.g., {'goal_progress': 0.8}
 
 class FlowManager:
     
@@ -230,27 +246,22 @@ class FlowManager:
         
         response = self.llm.invoke(prompt)
         
-        # Extract JSON from response
-        import re
-        json_match = re.search(r'\{.*\}', response, re.DOTALL)
-        if json_match:
-            try:
-                evaluation = json.loads(json_match.group(0))
-                
-                # Update goal progress metrics
-                if 'goals_progress' in evaluation:
-                    for goal, score in evaluation['goals_progress'].items():
-                        if goal in current_phase.goals:
-                            current_phase.completion_metrics[goal] = score
-                
-                # Check if we should transition
-                if evaluation.get('should_transition', False):
-                    next_phase_name = evaluation.get('next_phase')
-                    if next_phase_name in self.PHASES:
-                        self._transition_to_phase(user_id, next_phase_name, evaluation.get('reasoning', ''))
-            except json.JSONDecodeError:
-                self._check_time_based_transition(user_id)
-        else:
+        try:
+            # Parse with json5 for more tolerant parsing
+            evaluation = json5.loads(response)
+            # Validate with Pydantic
+            phase_transition = PhaseTransitionResponse.parse_obj(evaluation)
+            
+            # Update goal progress metrics
+            for goal, score in phase_transition.goals_progress.items():
+                if goal in current_phase.goals:
+                    current_phase.completion_metrics[goal] = score
+            
+            # Check if we should transition
+            if phase_transition.should_transition:
+                if phase_transition.next_phase in self.PHASES:
+                    self._transition_to_phase(user_id, phase_transition.next_phase, phase_transition.reasoning)
+        except (json5.Json5DecodeError, ValueError):
             self._check_time_based_transition(user_id)
     
     def _check_time_based_transition(self, user_id: str):
@@ -366,27 +377,15 @@ class FlowManager:
         
         response = self.llm.invoke(prompt)
         
-        # Extract JSON from response
-        import re
-        json_match = re.search(r'\{.*\}', response, re.DOTALL)
-        if json_match:
-            try:
-                characteristics = json.loads(json_match.group(0))
-                # Validate required fields
-                required_fields = [
-                    'alliance_strength', 'engagement_level', 'emotional_pattern',
-                    'cognitive_pattern', 'coping_mechanisms', 'progress_quality',
-                    'recommended_focus'
-                ]
-                if all(field in characteristics for field in required_fields):
-                    session['llm_context']['session_characteristics'] = characteristics
-                    logger.info(f"Updated session characteristics for user {user_id}")
-                else:
-                    logger.warning("Missing required fields in session characteristics")
-            except json.JSONDecodeError:
-                logger.warning("Failed to parse session characteristics from LLM")
-        else:
-            logger.warning("No JSON object found in LLM response")
+        try:
+            # Parse with json5 for more tolerant parsing
+            characteristics = json5.loads(response)
+            # Validate with Pydantic
+            session_chars = SessionCharacteristics.parse_obj(characteristics)
+            session['llm_context']['session_characteristics'] = session_chars.dict()
+            logger.info(f"Updated session characteristics for user {user_id}")
+        except (json5.Json5DecodeError, ValueError) as e:
+            logger.warning(f"Failed to parse session characteristics: {e}")
     
     def _create_flow_context(self, user_id: str) -> Dict[str, Any]: