UofTearsBot.py

from AdviceGenerator import AdviceGenerator
from SIDetector import SIDetector
from IllnessClassifier import IllnessClassifier
from typing import List, Dict

class UofTearsBot(object):
    def __init__(self, threshold: float = 0.86, max_history_msgs: int = 50):
        self.suicidality_detector = SIDetector()
        self.illness_classifier = IllnessClassifier()
        self.chatbot = AdviceGenerator()
        self.history: List[Dict[str, str]] = []
        self.FLAG = False   # suicidal crisis flag
        self.threshold = threshold
        self.max_history_msgs = max_history_msgs

    def safety_check(self, user_text: str):
        suicidal, confidence = self.suicidality_detector.forward(user_text)
        self.FLAG = (suicidal and confidence >= self.threshold)
        if self.FLAG:
            return "suicidal"
        disorder, conf = self.illness_classifier.forward(user_text)
        return disorder

    def userCrisis(self) -> str:
        return (
            "I'm really sorry you're feeling this way. Your safety matters.\n"
            "If you are in immediate danger, please call your local emergency number now.\n"
            "US & Canada: dial **988** (24/7)\n"
            "International: call your local emergency services.\n"
            "If you can, reach out to someone you trust and let them know what’s going on.\n"
            "We can also talk through coping strategies together."
        )

    def _prune_history(self) -> List[Dict[str, str]]:
        """Keep only the last `max_history_msgs` turns."""
        return self.history[-self.max_history_msgs :]

    def converse(self, user_text: str) -> str:
        disorder = self.safety_check(user_text)

        # store user text into history
        self.history.append({"role": "user", "content": user_text})

        if self.FLAG:
            # crisis flow: respond with fixed crisis message only
            crisis_msg = self.userCrisis()
            self.history.append({"role": "assistant", "content": crisis_msg})
            return crisis_msg

        # normal advice generation
        pruned_history = self._prune_history()
        advice = self.chatbot.generate_advice(
            disorder=disorder,
            user_text=user_text,
            history=pruned_history,
        )['text']

        # add bot response to history
        self.history.append({"role": "assistant", "content": advice})
        return advice