Fix API: Improve error handling and add fallbacks for NLTK and model components

api_mental_health.py

@@ -261,46 +261,94 @@ def clean_text(text):
         return ""
     text = str(text).lower()
     text = re.sub(r"[^a-zA-Z']", " ", text)
-    tokens = word_tokenize(text)
+    
+    # Simple tokenization by splitting on whitespace instead of using word_tokenize
+    # This avoids the dependency on punkt_tab
+    tokens = text.split()
+    
+    # Filter out stopwords and short tokens
     tokens = [lemmatizer.lemmatize(tok) for tok in tokens if tok not in STOPWORDS and len(tok) > 2]
     return " ".join(tokens)
 
 # Feature engineering function
 @traceable(run_type="tool", name="Engineer Features")
 def engineer_features(context, response=""):
-    context_clean = clean_text(context)
-    context_len = len(context_clean.split())
-    context_vader = analyzer.polarity_scores(context)['compound']
-    context_questions = context.count('?')
-    crisis_keywords = ['suicide', 'hopeless', 'worthless', 'kill', 'harm', 'desperate', 'overwhelmed', 'alone']
-    context_crisis_score = sum(1 for word in crisis_keywords if word in context.lower())
-    
-    context_tfidf = vectorizer.transform([context_clean]).toarray()
-    tfidf_cols = [f"tfidf_context_{i}" for i in range(context_tfidf.shape[1])]
-    response_tfidf = np.zeros_like(context_tfidf)
-    
-    lda_topics = lda.transform(context_tfidf)
-    
-    feature_cols = ["context_len", "context_vader", "context_questions", "crisis_flag"] + \
-                   [f"topic_{i}" for i in range(10)] + tfidf_cols + \
-                   [f"tfidf_response_{i}" for i in range(response_tfidf.shape[1])]
-    
-    features = pd.DataFrame({
-        "context_len": [context_len],
-        "context_vader": [context_vader],
-        "context_questions": [context_questions],
-        **{f"topic_{i}": [lda_topics[0][i]] for i in range(10)},
-        **{f"tfidf_context_{i}": [context_tfidf[0][i]] for i in range(context_tfidf.shape[1])},
-        **{f"tfidf_response_{i}": [response_tfidf[0][i]] for i in range(response_tfidf.shape[1])},
-    })
-    
-    crisis_features = features[["context_len", "context_vader", "context_questions"] + [f"topic_{i}" for i in range(10)]]
-    crisis_flag = crisis_clf.predict(crisis_features)[0]
-    if context_crisis_score > 0:
-        crisis_flag = 1
-    features["crisis_flag"] = crisis_flag
-    
-    return features, feature_cols
+    try:
+        context_clean = clean_text(context)
+        context_len = len(context_clean.split())
+        context_vader = analyzer.polarity_scores(context)['compound']
+        context_questions = context.count('?')
+        crisis_keywords = ['suicide', 'hopeless', 'worthless', 'kill', 'harm', 'desperate', 'overwhelmed', 'alone']
+        context_crisis_score = sum(1 for word in crisis_keywords if word in context.lower())
+        
+        # Check if vectorizer is properly initialized
+        if vectorizer is None or not hasattr(vectorizer, 'transform'):
+            logger.warning("Vectorizer not properly initialized, using placeholder")
+            # Create a simple placeholder for features
+            features = pd.DataFrame({
+                "context_len": [context_len],
+                "context_vader": [context_vader],
+                "context_questions": [context_questions],
+                "crisis_flag": [1 if context_crisis_score > 0 else 0]
+            })
+            feature_cols = ["context_len", "context_vader", "context_questions", "crisis_flag"]
+            return features, feature_cols
+            
+        # Use vectorizer if available
+        context_tfidf = vectorizer.transform([context_clean]).toarray()
+        tfidf_cols = [f"tfidf_context_{i}" for i in range(context_tfidf.shape[1])]
+        response_tfidf = np.zeros_like(context_tfidf)
+        
+        # Check if LDA model is properly initialized
+        if lda is None or not hasattr(lda, 'transform'):
+            logger.warning("LDA model not properly initialized, using zeros")
+            lda_topics = np.zeros((1, 10))
+        else:
+            lda_topics = lda.transform(context_tfidf)
+        
+        feature_cols = ["context_len", "context_vader", "context_questions", "crisis_flag"] + \
+                      [f"topic_{i}" for i in range(10)] + tfidf_cols + \
+                      [f"tfidf_response_{i}" for i in range(response_tfidf.shape[1])]
+        
+        features = pd.DataFrame({
+            "context_len": [context_len],
+            "context_vader": [context_vader],
+            "context_questions": [context_questions],
+            **{f"topic_{i}": [lda_topics[0][i]] for i in range(10)},
+            **{f"tfidf_context_{i}": [context_tfidf[0][i]] for i in range(context_tfidf.shape[1])},
+            **{f"tfidf_response_{i}": [response_tfidf[0][i]] for i in range(response_tfidf.shape[1])},
+        })
+        
+        # Check if crisis classifier is properly initialized
+        if crisis_clf is None or not hasattr(crisis_clf, 'predict'):
+            logger.warning("Crisis classifier not properly initialized, using keyword detection")
+            crisis_flag = 1 if context_crisis_score > 0 else 0
+        else:
+            crisis_features = features[["context_len", "context_vader", "context_questions"] + [f"topic_{i}" for i in range(10)]]
+            crisis_flag = crisis_clf.predict(crisis_features)[0]
+            if context_crisis_score > 0:
+                crisis_flag = 1
+                
+        features["crisis_flag"] = crisis_flag
+        
+        return features, feature_cols
+        
+    except Exception as e:
+        # Fallback to very basic features if anything goes wrong
+        logger.error(f"Error in engineer_features: {e}")
+        context_len = len(context.split())
+        context_questions = context.count('?')
+        crisis_keywords = ['suicide', 'hopeless', 'worthless', 'kill', 'harm', 'desperate', 'overwhelmed', 'alone']
+        context_crisis_score = sum(1 for word in crisis_keywords if word in context.lower())
+        
+        features = pd.DataFrame({
+            "context_len": [context_len],
+            "context_vader": [0.0],  # Default neutral sentiment
+            "context_questions": [context_questions],
+            "crisis_flag": [1 if context_crisis_score > 0 else 0]
+        })
+        feature_cols = ["context_len", "context_vader", "context_questions", "crisis_flag"]
+        return features, feature_cols
 
 # Prediction function
 @traceable(run_type="chain", name="Predict Response Type")
@@ -341,109 +389,185 @@ def predict_response_type(context):
 # RAG suggestion function
 @traceable(run_type="chain", name="RAG Suggestion")
 def generate_suggestion_rag(context, response_type, crisis_flag):
-    results = vector_store.similarity_search_with_score(context, k=3)
-    retrieved_contexts = [
-        f"Patient: {res[0].page_content}\nCounselor: {res[0].metadata['response']} (Type: {res[0].metadata['response_type']}, Crisis: {res[0].metadata['crisis_flag']}, Score: {res[1]:.2f})"
-        for res in results
-    ]
+    # Check if essential components are available
+    if vector_store is None or llm is None:
+        logger.warning("Vector store or LLM not available for RAG suggestions, using fallback")
+        risk_level = "High" if crisis_flag else "Low"
+        
+        # Simple fallback suggestions based on response type
+        if response_type == "Empathetic Listening":
+            suggestion = "I can hear that you're going through a difficult time. It sounds really challenging, and I appreciate you sharing this with me."
+        elif response_type == "Question":
+            suggestion = "Could you tell me more about how this has been affecting your daily life?"
+        elif response_type == "Advice":
+            suggestion = "It might be helpful to consider speaking with a mental health professional who can provide personalized support for what you're experiencing."
+        elif response_type == "Validation":
+            suggestion = "It's completely understandable to feel this way given what you're going through. Your feelings are valid."
+        else:
+            suggestion = "Thank you for sharing that with me. Let's explore this further together."
+            
+        # Add crisis resources if needed
+        if crisis_flag:
+            suggestion += " If you're in crisis, please remember help is available 24/7 through the National Suicide Prevention Lifeline at 988."
+            
+        return {
+            "suggested_response": suggestion,
+            "risk_level": risk_level
+        }
     
-    prompt_template = ChatPromptTemplate.from_template(
-        """
-        You are an expert mental health counseling assistant. A counselor has provided the following patient situation:
-        
-        Patient Situation: {context}
-        
-        Predicted Response Type: {response_type}
-        Crisis Flag: {crisis_flag}
-        
-        Based on the predicted response type and crisis flag, provide a suggested response for the counselor to use with the patient. The response should align with the response type ({response_type}) and be sensitive to the crisis level.
+    # If vector store is available, proceed with RAG
+    try:
+        results = vector_store.similarity_search_with_score(context, k=3)
+        retrieved_contexts = [
+            f"Patient: {res[0].page_content}\nCounselor: {res[0].metadata['response']} (Type: {res[0].metadata['response_type']}, Crisis: {res[0].metadata['crisis_flag']}, Score: {res[1]:.2f})"
+            for res in results
+        ]
         
-        For reference, here are similar cases from past conversations:
-        {retrieved_contexts}
+        prompt_template = ChatPromptTemplate.from_template(
+            """
+            You are an expert mental health counseling assistant. A counselor has provided the following patient situation:
+            
+            Patient Situation: {context}
+            
+            Predicted Response Type: {response_type}
+            Crisis Flag: {crisis_flag}
+            
+            Based on the predicted response type and crisis flag, provide a suggested response for the counselor to use with the patient. The response should align with the response type ({response_type}) and be sensitive to the crisis level.
+            
+            For reference, here are similar cases from past conversations:
+            {retrieved_contexts}
+            
+            Guidelines:
+            - If Crisis Flag is True, prioritize safety, empathy, and suggest immediate resources (e.g., National Suicide Prevention Lifeline at 988).
+            - For 'Empathetic Listening', focus on validating feelings without giving direct advice or questions.
+            - For 'Advice', provide practical, actionable suggestions.
+            - For 'Question', pose an open-ended question to encourage further discussion.
+            - For 'Validation', affirm the patient's efforts or feelings.
+            
+            Output in the following format:
+            ```json
+            {{
+                "suggested_response": "Your suggested response here",
+                "risk_level": "Low/Moderate/High"
+            }}
+            ```
+            """
+        )
         
-        Guidelines:
-        - If Crisis Flag is True, prioritize safety, empathy, and suggest immediate resources (e.g., National Suicide Prevention Lifeline at 988).
-        - For 'Empathetic Listening', focus on validating feelings without giving direct advice or questions.
-        - For 'Advice', provide practical, actionable suggestions.
-        - For 'Question', pose an open-ended question to encourage further discussion.
-        - For 'Validation', affirm the patient's efforts or feelings.
+        rag_chain = (
+            {
+                "context": RunnablePassthrough(),
+                "response_type": lambda x: response_type,
+                "crisis_flag": lambda x: "Crisis" if crisis_flag else "No Crisis",
+                "retrieved_contexts": lambda x: "\n".join(retrieved_contexts)
+            }
+            | prompt_template
+            | llm
+        )
         
-        Output in the following format:
-        ```json
-        {{
-            "suggested_response": "Your suggested response here",
-            "risk_level": "Low/Moderate/High"
-        }}
-        ```
-        """
-    )
-    
-    rag_chain = (
-        {
-            "context": RunnablePassthrough(),
-            "response_type": lambda x: response_type,
-            "crisis_flag": lambda x: "Crisis" if crisis_flag else "No Crisis",
-            "retrieved_contexts": lambda x: "\n".join(retrieved_contexts)
-        }
-        | prompt_template
-        | llm
-    )
-    
-    try:
         response = rag_chain.invoke(context)
         return eval(response.content.strip("```json\n").strip("\n```"))
     except Exception as e:
         logger.error(f"Error generating RAG suggestion: {e}")
-        raise HTTPException(status_code=500, detail=f"Error generating RAG suggestion: {str(e)}")
+        risk_level = "High" if crisis_flag else "Low"
+        
+        # Fallback suggestion if RAG fails
+        if crisis_flag:
+            suggestion = "I'm hearing that you're going through a very difficult time. Your safety is the most important thing right now. Would it be helpful to talk about resources that are available to support you, like the National Suicide Prevention Lifeline at 988?"
+        else:
+            suggestion = "Thank you for sharing that with me. I want to understand more about your experience and how I can best support you right now."
+            
+        return {
+            "suggested_response": suggestion,
+            "risk_level": risk_level
+        }
 
 # Direct suggestion function
 @traceable(run_type="chain", name="Direct Suggestion")
 def generate_suggestion_direct(context, response_type, crisis_flag):
-    prompt_template = ChatPromptTemplate.from_template(
-        """
-        You are an expert mental health counseling assistant. A counselor has provided the following patient situation:
-        
-        Patient Situation: {context}
-        
-        Predicted Response Type: {response_type}
-        Crisis Flag: {crisis_flag}
-        
-        Provide a suggested response for the counselor to use with the patient, aligned with the response type ({response_type}) and sensitive to the crisis level.
-        
-        Guidelines:
-        - If Crisis Flag is True, prioritize safety, empathy, and suggest immediate resources (e.g., National Suicide Prevention Lifeline at 988).
-        - For 'Empathetic Listening', focus on validating feelings without giving direct advice or questions.
-        - For 'Advice', provide practical, actionable suggestions.
-        - For 'Question', pose an open-ended question to encourage further discussion.
-        - For 'Validation', affirm the patient's efforts or feelings.
-        - Strictly adhere to the response type. For 'Empathetic Listening', do not include questions or advice.
-        
-        Output in the following format:
-        ```json
-        {{
-            "suggested_response": "Your suggested response here",
-            "risk_level": "Low/Moderate/High"
-        }}
-        ```
-        """
-    )
-    
-    direct_chain = (
-        {
-            "context": RunnablePassthrough(),
-            "response_type": lambda x: response_type,
-            "crisis_flag": lambda x: "Crisis" if crisis_flag else "No Crisis"
+    # Check if essential components are available
+    if llm is None:
+        logger.warning("LLM not available for direct suggestions, using fallback")
+        risk_level = "High" if crisis_flag else "Low"
+        
+        # Simple fallback suggestions based on response type
+        if response_type == "Empathetic Listening":
+            suggestion = "It sounds like this has been really difficult for you. I'm here to listen and support you."
+        elif response_type == "Question":
+            suggestion = "How have you been coping with these feelings recently?"
+        elif response_type == "Advice":
+            suggestion = "One thing that might help is establishing a simple morning routine with small, achievable steps."
+        elif response_type == "Validation":
+            suggestion = "What you're experiencing is a normal response to a difficult situation. Your feelings are valid."
+        else:
+            suggestion = "I appreciate you sharing this with me. Let's work through this together."
+            
+        # Add crisis resources if needed
+        if crisis_flag:
+            suggestion += " Given what you've shared, I want to make sure you know about resources like the National Suicide Prevention Lifeline at 988, which is available 24/7."
+            
+        return {
+            "suggested_response": suggestion,
+            "risk_level": risk_level
         }
-        | prompt_template
-        | llm
-    )
     
+    # If LLM is available, proceed with direct suggestion
     try:
+        prompt_template = ChatPromptTemplate.from_template(
+            """
+            You are an expert mental health counseling assistant. A counselor has provided the following patient situation:
+            
+            Patient Situation: {context}
+            
+            Predicted Response Type: {response_type}
+            Crisis Flag: {crisis_flag}
+            
+            Provide a suggested response for the counselor to use with the patient, aligned with the response type ({response_type}) and sensitive to the crisis level.
+            
+            Guidelines:
+            - If Crisis Flag is True, prioritize safety, empathy, and suggest immediate resources (e.g., National Suicide Prevention Lifeline at 988).
+            - For 'Empathetic Listening', focus on validating feelings without giving direct advice or questions.
+            - For 'Advice', provide practical, actionable suggestions.
+            - For 'Question', pose an open-ended question to encourage further discussion.
+            - For 'Validation', affirm the patient's efforts or feelings.
+            - Strictly adhere to the response type. For 'Empathetic Listening', do not include questions or advice.
+            
+            Output in the following format:
+            ```json
+            {{
+                "suggested_response": "Your suggested response here",
+                "risk_level": "Low/Moderate/High"
+            }}
+            ```
+            """
+        )
+        
+        direct_chain = (
+            {
+                "context": RunnablePassthrough(),
+                "response_type": lambda x: response_type,
+                "crisis_flag": lambda x: "Crisis" if crisis_flag else "No Crisis"
+            }
+            | prompt_template
+            | llm
+        )
+        
         response = direct_chain.invoke(context)
         return eval(response.content.strip("```json\n").strip("\n```"))
     except Exception as e:
         logger.error(f"Error generating direct suggestion: {e}")
-        raise HTTPException(status_code=500, detail=f"Error generating direct suggestion: {str(e)}")
+        risk_level = "High" if crisis_flag else "Low"
+        
+        # Fallback suggestion if direct generation fails
+        if crisis_flag:
+            suggestion = "I'm concerned about what you're sharing. Your wellbeing is important, and I want to make sure you have support. The National Suicide Prevention Lifeline (988) has trained counselors available 24/7."
+        else:
+            suggestion = "I hear you're having a difficult time. Would you like to talk more about how these feelings have been affecting you?"
+            
+        return {
+            "suggested_response": suggestion,
+            "risk_level": risk_level
+        }
 
 # User Profile Endpoints
 @app.post("/users/create", response_model=UserProfile)