Update src/chatbot.py

src/chatbot.py

@@ -98,19 +98,61 @@ def is_political_stress(prompt: str) -> bool:
     return any(keyword in prompt.lower() for keyword in political_keywords)
 
 ########################################
-# 4. GENERATE RESPONSE
+# 4A. SUMMARIZATION HELPER (to keep conversation shorter)
 ########################################
-def generate_response(prompt: str, country: str) -> str:
-    lower_prompt = prompt.lower()
 
-    # 1. Self-harm detection
+def summarize_text(text: str) -> str:
+    """
+    Very basic summarization approach:
+    1) We feed a short prompt to the same model to summarize the text.
+    2) Return a short summary from the model.
+
+    In a real application, you might use a specialized summarization model.
+    """
+    summary_prompt = (
+        "Summarize this conversation in a concise way, focusing on key points:\n\n"
+        f"{text}\n\nSummary:"
+    )
+    inputs = tokenizer(summary_prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        summary_outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            temperature=0.7,
+            top_p=0.9,
+            top_k=50,
+            repetition_penalty=1.2,
+            no_repeat_ngram_size=2,
+            do_sample=True
+        )
+    summary = tokenizer.decode(summary_outputs[0], skip_special_tokens=True)
+    # Basic cleanup
+    if "Summary:" in summary:
+        summary = summary.split("Summary:")[-1].strip()
+    return summary
+
+########################################
+# 4B. GPT-Style Response with Conversation History
+########################################
+
+def generate_response_with_history(conversation_history, country: str) -> str:
+    """
+    Generates a response based on the entire conversation history.
+    - conversation_history: a list of turns, each is { "role": "user"/"bot", "content": str }
+    - country: used for helpline logic if self-harm is detected in the LAST user message
+    """
+    # Check the last user message for self-harm, etc.
+    last_user_message = conversation_history[-1]["content"]
+    lower_prompt = last_user_message.lower()
+
+    # 1) Self-harm detection
     self_harm_keywords = [
-        "suicide", "kill myself", "end my life", "self-harm", ...
+        "suicide", "kill myself", "end my life", "self-harm",
+        "hurt myself", "want to die", "no reason to live", "tired of living"
     ]
     if any(keyword in lower_prompt for keyword in self_harm_keywords):
-        logging.info(f"Self-harm keyword detected in prompt: {prompt}")
+        logging.info(f"Self-harm keyword detected in prompt: {last_user_message}")
         helpline_str = get_helpline_for_country(country)
-
         return (
             "I’m really sorry you’re feeling like this. It sounds like you’re in a very tough place right now. "
             "If you’re comfortable, could you share more about what’s bringing you to feel this way? "
@@ -119,28 +161,34 @@ def generate_response(prompt: str, country: str) -> str:
             "You’re not alone, and there are caring people who want to help you."
         )
 
-    # 2. Harm-to-others detection...
-    # 3. Hate speech detection...
+    # 2) Harm-to-others / Hate speech can be similarly handled, or you can do a single pass prior to generation.
 
-    # 4. If no major triggers:
+    # 3) Construct a system prompt to direct the style:
     system_prompt = (
         "You are a supportive, empathetic companion. "
         "Your top priority is to listen and help the user feel heard. "
         "You:\n"
         "- Acknowledge the user's feelings.\n"
-        "- Ask gentle follow-up questions to encourage them to share more.\n"
-        "- Suggest small, practical ways they might cope or self-soothe.\n"
+        "- Ask gentle follow-up questions.\n"
+        "- Suggest small, practical coping ideas.\n"
         "- Avoid judgment or lecturing.\n"
-        "- Offer additional support if it seems they are in crisis.\n"
-        "But do NOT prefix your messages with 'User:' or 'Bot:'.\n"
+        "- Provide helpline info only if user is in crisis.\n"
+        "Do NOT prefix lines with 'User:' or 'Bot:'.\n"
     )
 
-    # Political stress addition if needed...
-    # system_prompt += ...
+    # 4) Build a single text prompt from the entire conversation
+    # Start with system instructions:
+    conversation_text = system_prompt
 
-    full_prompt = f"{system_prompt}\nThe user says: {prompt}\nYour response:"
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+    for turn in conversation_history:
+        if turn["role"] == "user":
+            conversation_text += f"\nUser says: {turn['content']}"
+        else:
+            conversation_text += f"\nBot says: {turn['content']}"
+    conversation_text += "\nBot says:"
 
+    # 5) Convert text to tokens and generate
+    inputs = tokenizer(conversation_text, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -154,15 +202,53 @@ def generate_response(prompt: str, country: str) -> str:
         )
 
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Clean up
-    response = response.replace(system_prompt, "").replace("The user says:", "").replace("Your response:", "")
+
+    # 6) Clean up any repeated system text
+    response = response.replace(system_prompt, "")
+    response = response.replace("User says:", "").replace("Bot says:", "")
+    response = response.strip()
+
+    # Make sure it ends with punctuation
     if response and response[-1] not in ".!?":
-        if "." in response:
-            response = response.rsplit(".", 1)[0].strip() + "."
-        else:
-            response += "."
+        response += "."
 
-    return response.strip()
+    return response
+
+
+########################################
+# 4C. Possibly Summarize Older History
+########################################
+
+def maybe_summarize_history(conversation_history, max_turns=10):
+    """
+    If conversation_history is too long, summarize the earliest turns into a single 'summary so far'
+    turn. This helps keep token usage manageable.
+    """
+    if len(conversation_history) > max_turns:
+        # Separate out the first chunk (e.g. first 5 turns).
+        old_turns = conversation_history[:-5]
+        recent_turns = conversation_history[-5:]
+
+        # Build a text from old_turns
+        old_text = ""
+        for turn in old_turns:
+            role = turn["role"]
+            content = turn["content"]
+            old_text += f"{role.upper()}:\n{content}\n\n"
+
+        summary = summarize_text(old_text)
+        # Now store that summary as a single turn with role "bot" (or "system")
+        summary_turn = {
+            "role": "bot",
+            "content": f"(A summary of earlier conversation: {summary})"
+        }
+
+        # Rebuild the conversation: summary + recent 5 turns
+        new_history = [summary_turn] + recent_turns
+
+        # Clear and replace
+        conversation_history.clear()
+        conversation_history.extend(new_history)
 
 
 ########################################
@@ -186,6 +272,10 @@ def save_user_feedback(user_input, bot_response, rating, extra_comments=""):
 ########################################
 # 6. CONVERSATION LOOP
 ########################################
+
+# Keep a global conversation history for the console-based loop
+conversation_history = []
+
 def chatbot_conversation(country_selection: str):
     """
     Starts the console-based conversation loop, using 'country_selection'
@@ -196,6 +286,7 @@ def chatbot_conversation(country_selection: str):
     print(f"Country set to: {country_selection}\n")
 
     print("Type 'exit' or 'quit' to end the conversation.")
+
     while True:
         user_input = input("\nYou: ")
         if user_input.lower() in ["exit", "quit"]:
@@ -212,8 +303,19 @@ def chatbot_conversation(country_selection: str):
             truncated_input = user_input[:HARD_LIMIT]
             print("(System Note): Your message exceeded 2048 chars and was truncated.")
 
-        # Generate & display response
-        bot_response = generate_response(truncated_input, country_selection)
+        # 1) Append user turn to the conversation
+        conversation_history.append({"role": "user", "content": truncated_input})
+
+        # 2) Summarize older history if we exceed ~10 turns
+        maybe_summarize_history(conversation_history, max_turns=10)
+
+        # 3) Generate response with the entire conversation
+        bot_response = generate_response_with_history(conversation_history, country_selection)
+
+        # 4) Append the bot's response to the conversation
+        conversation_history.append({"role": "bot", "content": bot_response})
+
+        # 5) Display
         print(bot_response)
 
         # Ask for feedback
@@ -225,12 +327,12 @@ def chatbot_conversation(country_selection: str):
                 print("Please enter 'y' or 'n' or just hit Enter to skip.")
 
         if rating in ["y", "n"]:
-            # Optional additional comments
             extra_comments = input("Any additional comments? (press Enter to skip): ")
             rating_symbol = "👍" if rating == "y" else "👎"
             save_user_feedback(truncated_input, bot_response, rating_symbol, extra_comments)
             print("(Feedback saved.)")
 
+
 ########################################
 # 7. TKINTER GUI WITH COUNTRY SELECTION
 ########################################
@@ -243,7 +345,11 @@ def main():
     root.title("Mental Health Chatbot Prototype")
 
     # Disclaimer message
-    disclaimer_label = tk.Label(root, text="Disclaimer: This is just a prototype and not a substitute for professional help.", fg="red")
+    disclaimer_label = tk.Label(
+        root,
+        text="Disclaimer: This is just a prototype and not a substitute for professional help.",
+        fg="red"
+    )
     disclaimer_label.pack(padx=20, pady=10)
 
     # Country selection label
@@ -268,3 +374,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+