Spaces:
Paused
Paused
| from AdviceGenerator import AdviceGenerator | |
| from SIDetector import SIDetector | |
| from IllnessClassifier import IllnessClassifier | |
| from typing import List, Dict | |
| class UofTearsBot(object): | |
| def __init__(self, llm, threshold: float = 0.86, max_history_msgs: int = 50): | |
| self.suicidality_detector = SIDetector() | |
| self.illness_classifier = IllnessClassifier() | |
| self.chatbot = AdviceGenerator(llm) | |
| 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 input | |
| self.history.append({"role": "user", "content": user_text}) | |
| # crisis branch | |
| if self.FLAG: | |
| crisis_msg = self.userCrisis() | |
| self.history.append({"role": "assistant", "content": crisis_msg}) | |
| yield crisis_msg | |
| return | |
| # normal branch: stream advice tokens | |
| reply_so_far = "" | |
| for delta in self.chatbot.generate_advice( | |
| disorder=disorder, | |
| user_text=user_text, | |
| history=self._prune_history(), | |
| ): | |
| reply_so_far += delta | |
| yield delta # stream to FastAPI as soon as a token arrives | |
| # once stream is done, save full reply | |
| self.history.append({"role": "assistant", "content": reply_so_far}) |