Update chain_problems.py

chain_problems.py

@@ -1,4 +1,3 @@
-# chain_problems.py
 import json
 import logging
 from typing import Dict
@@ -7,6 +6,7 @@ from models import chat_model
 
 logger = logging.getLogger(__name__)
 
+# Updated prompt template to include eight new themes
 problem_prompt_template = PromptTemplate(
     input_variables=["responses", "internal_report"],
     template=(
@@ -15,30 +15,50 @@ problem_prompt_template = PromptTemplate(
         "You also have an internal analysis report:\n"
         "{internal_report}\n\n"
         "From these inputs, determine a 'problem severity percentage' for the user in the following areas: "
-        "sleep, exercise, stress, and diet. "
-        "Return your answer in JSON format with keys: sleep_problem, exercise_problem, stress_problem, diet_problem.\n"
+        "stress_management, low_therapy, balanced_weight, restless_night, lack_of_motivation, gut_health, anxiety, burnout. "
+        "Return your answer in JSON format with keys: stress_management, low_therapy, balanced_weight, restless_night, "
+        "lack_of_motivation, gut_health, anxiety, burnout.\n"
         "Ensure severity percentages are numbers from 0 to 100.\n\n"
         "JSON Output:"
     )
 )
+
 problem_chain = LLMChain(llm=chat_model, prompt=problem_prompt_template)
 
 def analyze_problems_with_chain(responses: Dict[str, str], internal_report: str) -> Dict[str, float]:
     responses_str = "\n".join(f"{q}: {a}" for q, a in responses.items())
     raw_text = problem_chain.run(responses=responses_str, internal_report=internal_report)
     try:
+        # Extract JSON from the LLM output
         start_idx = raw_text.find('{')
         end_idx = raw_text.rfind('}') + 1
         json_str = raw_text[start_idx:end_idx]
         problems = json.loads(json_str)
-        for key in ["sleep_problem", "exercise_problem", "stress_problem", "diet_problem"]:
+        
+        # Ensure all eight keys are present with default values
+        for key in [
+            "stress_management",
+            "low_therapy",
+            "balanced_weight",
+            "restless_night",
+            "lack_of_motivation",
+            "gut_health",
+            "anxiety",
+            "burnout"
+        ]:
             problems.setdefault(key, 0.0)
+        
         return {k: float(v) for k, v in problems.items()}
     except Exception as e:
         logger.error(f"Error parsing problem percentages from LLM: {e}")
+        # Return default values for all eight themes in case of an error
         return {
-            "sleep_problem": 0.0,
-            "exercise_problem": 0.0,
-            "stress_problem": 0.0,
-            "diet_problem": 0.0
+            "stress_management": 0.0,
+            "low_therapy": 0.0,
+            "balanced_weight": 0.0,
+            "restless_night": 0.0,
+            "lack_of_motivation": 0.0,
+            "gut_health": 0.0,
+            "anxiety": 0.0,
+            "burnout": 0.0
         }