""" ============================================================ RAG Engine v2 - Integrated with Diagnosis ============================================================ الترقية الرئيسية: - بدلاً من 2 prompts منفصلين (chat + validation) - عندنا 4 modes: 1. chat: رد عادي على رسالة في وسط المحادثة 2. probing: أسئلة استكشافية لو لسه ما يقدرش يشخّص 3. final_diagnosis: تشخيص نهائي + نصايح 4. validate_pdf: التحقق من ملف الرفع - بياخد الـ classifier prediction كـ context قوي - بيستخدم RAG retrieval عشان النصايح مدعومة بمعلومات طبية """ import os from operator import itemgetter from typing import List, Optional, Dict from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import Chroma from langchain_groq import ChatGroq from langchain_core.prompts import ChatPromptTemplate from langchain_core.output_parsers import StrOutputParser from langchain_community.document_loaders import PyPDFLoader from langchain_text_splitters import RecursiveCharacterTextSplitter from dotenv import load_dotenv from config import RAG_CONFIG, CHROMA_DB_PATH, GROQ_API_KEY load_dotenv() class MentallicaRAG: """RAG محسّن للصحة النفسية""" def __init__(self): if not GROQ_API_KEY: raise EnvironmentError( "GROQ_API_KEY not found! Set it in .env file or environment." ) os.environ["GROQ_API_KEY"] = GROQ_API_KEY # Embeddings + Vector DB self.embeddings = HuggingFaceEmbeddings(model_name=RAG_CONFIG['embedding_model']) if os.path.exists(CHROMA_DB_PATH): self.vector_db = Chroma( persist_directory=CHROMA_DB_PATH, embedding_function=self.embeddings ) self.retriever = self.vector_db.as_retriever( search_kwargs={"k": RAG_CONFIG['retrieval_k']} ) print("✓ Loaded existing ChromaDB") else: print(f"⚠ ChromaDB not found at {CHROMA_DB_PATH}") print(" Run ingest.py first to build the knowledge base") self.vector_db = None self.retriever = None # LLM self.llm = ChatGroq( model=RAG_CONFIG['llm_model'], temperature=RAG_CONFIG['temperature'], max_tokens=RAG_CONFIG['max_tokens'], ) # Build prompts self._build_prompts() def _build_prompts(self): """4 modes مختلفة""" # ===== Chat Mode: رد عادي خلال المحادثة ===== self.chat_prompt = ChatPromptTemplate.from_messages([ ("system", "You are Mentallico, an empathetic, professional AI mental health assistant. " "You are having a conversation with a patient. Listen carefully, validate their feelings, " "and ask gentle follow-up questions to understand their situation better. " "DO NOT provide a diagnosis yet — you need more information.\n\n" "🔴 CRITICAL LANGUAGE RULE: You MUST automatically detect the language of the patient's input. " "Your entire response MUST be in the EXACT SAME LANGUAGE. If the patient writes in Arabic, reply in Arabic. If English, reply in English.\n\n" "Use this medical context to inform your responses:\n{context}\n\n" "Conversation history:\n{history}"), ("human", "Patient: {input}"), ]) # ===== Probing Mode: أسئلة استكشافية محددة ===== self.probing_prompt = ChatPromptTemplate.from_messages([ ("system", "You are Mentallico, a mental health assistant. The classifier suggests the patient " "might be experiencing '{predicted_label}' (confidence: {confidence:.0%}), but you need " "MORE information to confirm. Ask 1-2 specific, empathetic questions that would help " "differentiate this condition from similar ones. Be warm, not clinical.\n\n" "🔴 CRITICAL LANGUAGE RULE: You MUST automatically detect the language of the patient's input. " "Your entire response MUST be in the EXACT SAME LANGUAGE. If the patient writes in Arabic, reply in Arabic. If English, reply in English.\n\n" "Medical reference:\n{context}\n\n" "What you know so far:\n{history}"), ("human", "Patient's latest: {input}"), ]) # ===== Final Diagnosis Mode ===== self.diagnosis_prompt = ChatPromptTemplate.from_messages([ ("system", "You are Mentallico, a professional AI mental health assistant. Based on the full " "conversation, you have determined the patient is most likely experiencing: " "**{diagnosis}** (confidence: {confidence:.0%}).\n\n" "Your response MUST follow this structure:\n" "1. Acknowledgment: Validate their experience compassionately (2 sentences).\n" "2. Diagnosis Explanation: Briefly explain what {diagnosis} is and why it matches their symptoms.\n" "3. Coping Strategies: Practical, evidence-based techniques they can start TODAY.\n" "4. Lifestyle Recommendations: Sleep, exercise, nutrition, social.\n" "5. When to Seek Professional Help: Clear signs they should see a specialist.\n" "6. Disclaimer: Remind them this is AI guidance, not a clinical diagnosis.\n\n" "🔴 CRITICAL LANGUAGE RULE: You MUST write this ENTIRE report, including all headings and bullet points, " "in the EXACT SAME LANGUAGE the patient used during the chat (e.g., fully in Arabic if they spoke Arabic).\n\n" "Use the medical context below to ground your suggestions:\n{context}\n\n" "Full conversation:\n{history}\n\n" "If urgency is 'critical', start with crisis resources and helplines."), ("human", "Provide the comprehensive diagnosis report and recommendations. " "Urgency level: {urgency}"), ]) # ===== Validation Mode (للـ PDF upload) ===== # ده هنسيبه بالانجليزي لانه لوجيك داخلي مش بيظهر لليوزر self.validation_prompt = ( "You are a medical document classifier. Read the following text and determine if it is " "related to mental health, psychiatry, psychology, or medicine. " "Respond strictly with a single word: YES or NO.\n\n" "Text: {text}" ) def _format_docs(self, docs): if not docs: return "No specific medical context retrieved." return "\n\n".join(doc.page_content for doc in docs) def _format_history(self, history: List[dict]) -> str: if not history: return "(start of conversation)" lines = [] for msg in history[-10:]: # last 10 messages role = "Patient" if msg['role'] == 'user' else "Mentallico" lines.append(f"{role}: {msg['content']}") return "\n".join(lines) def _retrieve(self, query: str) -> str: """يجيب context relevant""" if self.retriever is None: return "No knowledge base available." try: docs = self.retriever.invoke(query) return self._format_docs(docs) except Exception as e: print(f"⚠ Retrieval error: {e}") return "Could not retrieve context." # ============= Public Methods ============= def chat_response(self, user_input: str, history: List[dict]) -> str: """رد محادثة عادي - بدون تشخيص""" try: context = self._retrieve(user_input) history_text = self._format_history(history) chain = self.chat_prompt | self.llm | StrOutputParser() response = chain.invoke({ "input": user_input, "context": context, "history": history_text }) return response except Exception as e: return f"I'm having trouble responding right now. Could you tell me more about how you've been feeling? (Error: {str(e)[:80]})" def probing_response(self, user_input: str, history: List[dict], predicted_label: str, confidence: float) -> str: """أسئلة استكشافية - النموذج عنده توقع بس مش متأكد""" try: # نستخدم الـ predicted label كـ retrieval query لجيب معلومات أكتر retrieval_query = f"{predicted_label} symptoms diagnosis criteria" context = self._retrieve(retrieval_query) history_text = self._format_history(history) chain = self.probing_prompt | self.llm | StrOutputParser() response = chain.invoke({ "input": user_input, "context": context, "history": history_text, "predicted_label": predicted_label, "confidence": confidence, }) return response except Exception as e: return self.chat_response(user_input, history) def final_diagnosis_response(self, diagnosis: str, confidence: float, history: List[dict], urgency: str = 'normal') -> str: """التشخيص النهائي + النصايح المفصلة""" try: # Retrieval شامل عن المرض ده retrieval_query = f"{diagnosis} treatment coping strategies recommendations therapy" context = self._retrieve(retrieval_query) history_text = self._format_history(history) chain = self.diagnosis_prompt | self.llm | StrOutputParser() response = chain.invoke({ "diagnosis": diagnosis, "confidence": confidence, "context": context, "history": history_text, "urgency": urgency, }) return response except Exception as e: return (f"Based on our conversation, you may be experiencing **{diagnosis}**. " f"I strongly recommend consulting with a mental health professional. " f"(System error in detail generation: {str(e)[:80]})") def validate_pdf_content(self, text: str) -> bool: """التحقق من ان الـ PDF متعلق بصحة نفسية""" try: prompt = self.validation_prompt.format(text=text[:1500]) response = self.llm.invoke(prompt).content.strip().upper() return "YES" in response except Exception as e: print(f"⚠ Validation error: {e}") return False def add_pdf_to_knowledge_base(self, pdf_path: str) -> dict: """يضيف PDF جديد للـ vector DB بعد التحقق منه""" try: loader = PyPDFLoader(pdf_path) docs = loader.load() if not docs: return {"status": "error", "message": "PDF empty or unreadable"} sample = docs[0].page_content[:1500] if not self.validate_pdf_content(sample): return { "status": "error", "message": "Document not related to mental health/psychiatry" } splitter = RecursiveCharacterTextSplitter( chunk_size=RAG_CONFIG['chunk_size'], chunk_overlap=RAG_CONFIG['chunk_overlap'] ) chunks = splitter.split_documents(docs) if self.vector_db is None: # أول مرة - ننشئ الـ DB self.vector_db = Chroma.from_documents( documents=chunks, embedding=self.embeddings, persist_directory=CHROMA_DB_PATH ) else: self.vector_db.add_documents(chunks) self.retriever = self.vector_db.as_retriever( search_kwargs={"k": RAG_CONFIG['retrieval_k']} ) return { "status": "success", "message": f"Added {len(chunks)} chunks from {os.path.basename(pdf_path)}" } except Exception as e: return {"status": "error", "message": f"Processing error: {str(e)}"} # Global instance (lazy init) _rag_instance = None def get_rag() -> MentallicaRAG: global _rag_instance if _rag_instance is None: _rag_instance = MentallicaRAG() return _rag_instance