from __future__ import annotations import argparse import json import sys from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any CURRENT_DIR = Path(__file__).resolve().parent PROJECT_ROOT = CURRENT_DIR.parents[1] if str(CURRENT_DIR) not in sys.path: sys.path.append(str(CURRENT_DIR)) if str(PROJECT_ROOT) not in sys.path: sys.path.append(str(PROJECT_ROOT)) from emotion_classifier import EmotionClassifier from intent_classifier import IntentClassifier from language_classifier import LanguageDetector from response_generator import ResponseGenerator from safety_router import crisis_reply, detect_crisis from src.retrieval.retrieval_engine import RetrievalEngine class ChatbotPipeline: def __init__(self, retrieval_source: str = "both", top_k: int = 5, retrieval_collection: str | None = None) -> None: self.retrieval_source = retrieval_source self.top_k = top_k self.retrieval_collection = retrieval_collection self.language_detector = LanguageDetector() self.language_detector.load_model() self.emotion_classifier = EmotionClassifier() self.intent_classifier = IntentClassifier() self.retrieval_engine = None self.response_generator = ResponseGenerator() def run(self, user_message: str, history: list[dict[str, str]] | None = None) -> dict[str, Any]: clean_message = user_message.strip() if not clean_message: return {"response": "Please enter a message.", "state": {}} conversation_history = self._prepare_history(clean_message, history or []) state = self._analyze(clean_message, conversation_history) state["conversation_history"] = conversation_history language_code = state["language"].get("language_code", "en") if state["safety"]["is_crisis"]: state["route"] = "crisis" return {"response": crisis_reply(language_code), "state": state} intent = state["intent"]["intent"] use_retrieval = intent == "asking_mental_health_question" state["route"] = "rag" if use_retrieval else "direct_response" state["retrieval"] = { "enabled": use_retrieval, "source": self.retrieval_source, "top_k": self.top_k, "results": [], } if use_retrieval: retrieval_query = state["intent"].get("retrieval_query") or clean_message state["retrieval"]["query"] = retrieval_query try: state["retrieval"]["results"] = self._retrieve(retrieval_query) except Exception as error: state["retrieval"]["error"] = f"{type(error).__name__}" try: generated = self.response_generator.generate(state) state["llm_review"] = { "language": generated.get("language_review", {}), "emotion": generated.get("emotion_review", {}), "intent": generated.get("intent_review", {}), } corrected_intent = state["llm_review"]["intent"].get("corrected_intent") state["final_intent"] = corrected_intent or state["intent"].get("intent") state["final_route"] = "rag" if state["final_intent"] == "asking_mental_health_question" else "direct_response" if state["final_intent"] == "asking_mental_health_question": state["suggested_questions"] = generated.get("suggested_questions", []) else: state["suggested_questions"] = [] response = generated.get("answer") or "I am here with you, but I could not generate a complete response." except RuntimeError as error: state["llm_review"] = {"language": {}, "emotion": {}, "intent": {}} state["suggested_questions"] = [] state["generation_error"] = str(error) response = ( "I'm here with you \u2764\ufe0f. I can listen, help you slow things down, and support you with mental-health questions. " "The advanced response generator is not available right now, so please try again shortly. " "If this feels urgent or unsafe, contact local emergency support or someone you trust right away." ) except Exception as error: state["llm_review"] = {"language": {}, "emotion": {}, "intent": {}} state["suggested_questions"] = [] state["generation_error"] = f"{type(error).__name__}: {error}" response = ( "I am here with you, but I could not complete a full answer at the moment. " "Try again shortly, or contact a trusted person or professional support if you need help now." ) return {"response": response, "suggested_questions": state.get("suggested_questions", []), "state": state} def _analyze(self, message: str, history: list[dict[str, str]]) -> dict[str, Any]: safety = detect_crisis(message) with ThreadPoolExecutor(max_workers=3) as executor: language_future = executor.submit(self._safe_language, message) emotion_future = executor.submit(self._safe_emotion, message) if safety["is_crisis"]: intent_future = None else: intent_future = executor.submit(self._safe_intent, message, history) language = language_future.result() emotion = emotion_future.result() intent = self._crisis_intent(message) if intent_future is None else intent_future.result() return { "user_message": message, "language": language, "emotion": emotion, "intent": intent, "safety": safety, "retrieval": {"enabled": False, "results": []}, } def _crisis_intent(self, message: str) -> dict[str, Any]: return { "intent": "asking_mental_health_question", "confidence": 0.0, "confidence_margin": 0.0, "intent_scores": {}, "reason": "Crisis guardrail matched before live intent classification.", "retrieval_query": message, "contextual_follow_up": False, "interaction_type": "standalone", "classification_skipped": True, } def _safe_intent(self, message: str, history: list[dict[str, str]]) -> dict[str, Any]: try: return self.intent_classifier.classify(message, history=history) except Exception as error: return self._fallback_intent(message, error) def set_retrieval_collection(self, collection_name: str | None) -> None: if collection_name != self.retrieval_collection: self.retrieval_collection = collection_name self.retrieval_engine = None def _safe_language(self, message: str) -> dict[str, Any]: try: return self.language_detector.predict_with_confidence(message) except Exception as error: return { "language_code": "en", "language_name": "English", "confidence": 0.0, "is_confident": False, "message": f"Language detection unavailable: {type(error).__name__}", } def _safe_emotion(self, message: str) -> dict[str, Any]: try: return self.emotion_classifier.predict_with_confidence(message) except Exception as error: return { "emotion": "unknown", "confidence": 0.0, "is_confident": False, "scores": {}, "message": f"Emotion detection unavailable: {type(error).__name__}", } def _fallback_intent(self, message: str, error: Exception) -> dict[str, Any]: return { "intent": "out_of_scope", "confidence": 0.0, "confidence_margin": 0.0, "intent_scores": {}, "reason": f"Intent classification unavailable: {type(error).__name__}.", "retrieval_query": message, "contextual_follow_up": False, "interaction_type": "standalone", } @staticmethod def _prepare_history(message: str, history: list[dict[str, str]]) -> list[dict[str, str]]: clean_history = [ {"role": item.get("role", ""), "content": str(item.get("content", "")).strip()} for item in history if item.get("role") in {"user", "assistant"} and str(item.get("content", "")).strip() ] if clean_history and clean_history[-1]["role"] == "user" and clean_history[-1]["content"] == message: clean_history.pop() return clean_history[-8:] def _retrieve(self, message: str) -> list[dict[str, Any]]: if self.retrieval_engine is None: self.retrieval_engine = RetrievalEngine(collection_name=self.retrieval_collection) return self.retrieval_engine.search(message, source=self.retrieval_source, top_k=self.top_k) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run the integrated mental-health chatbot pipeline.") parser.add_argument("message", nargs="?", default="I feel anxious and cannot sleep.") parser.add_argument("--source", choices=["both", "cci", "amod"], default="both") parser.add_argument("--top-k", type=int, default=5) return parser.parse_args() if __name__ == "__main__": args = parse_args() pipeline = ChatbotPipeline(retrieval_source=args.source, top_k=args.top_k) output = pipeline.run(args.message) print(json.dumps(output, indent=2, ensure_ascii=False))