Spaces:
Running
Running
| """ | |
| ============================================================ | |
| 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 | |