Upload 3 files

src/main_back.py

@@ -0,0 +1,220 @@
+import os
+from typing import Dict, List, TypedDict, Annotated
+from dotenv import load_dotenv
+
+# LangChain and LangGraph imports
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_groq import ChatGroq
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.messages import AnyMessage, AIMessage, HumanMessage
+
+from langgraph.graph import StateGraph, END, START
+from langgraph.graph.message import add_messages
+import sqlite3
+from langgraph.checkpoint.sqlite import SqliteSaver
+from pydantic import BaseModel, Field
+
+# ---  Graph State Definition ---
+class GraphState(TypedDict):
+    questionnaire_responses: Dict[str, int]
+    domain_scores: Dict[str, int]
+    primary_concern: str
+    messages: Annotated[list[AnyMessage], add_messages]
+    is_safe: bool
+    retry_count: int
+    
+
+# ---  RAG Retriever Helper Function ---
+def create_persistent_rag_retriever(pdf_paths: List[str], db_name: str, embedding_model):
+    """Creates or loads a persistent RAG retriever from one or more PDF documents."""
+    persist_directory = f"./chroma_db/{db_name}"
+    if os.path.exists(persist_directory):
+        print(f"--- Loading existing persistent DB: {db_name} ---")
+        return Chroma(persist_directory=persist_directory, embedding_function=embedding_model).as_retriever()
+    
+    print(f"--- Creating new persistent DB: {db_name} ---")
+    vector_store = None
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+
+    for pdf_path in pdf_paths:
+        if not os.path.exists(pdf_path):
+            print(f"Warning: PDF not found at '{pdf_path}'. Skipping.")
+            continue
+        
+        print(f"--- Processing PDF: {pdf_path} ---")
+        loader = PyPDFLoader(pdf_path)
+        documents = loader.load()
+        splits = text_splitter.split_documents(documents)
+
+        if not splits: continue
+            
+        if vector_store is None:
+            vector_store = Chroma.from_documents(documents=splits, embedding=embedding_model, persist_directory=persist_directory)
+        else:
+            vector_store.add_documents(splits)
+            
+    if vector_store:
+        print(f"--- DB creation complete for {db_name} ---")
+        return vector_store.as_retriever(search_kwargs={'k': 3})
+    else:
+        print(f"--- Could not create DB for {db_name}. No documents processed. ---")
+        return None
+
+# ---  Node Definitions ---
+def questionnaire(state: GraphState) -> GraphState:
+    """Calculates all domain scores and clears the questionnaire responses from the state."""
+    responses = state.get("questionnaire_responses", {})
+    domain_scores = {
+        "Depression": max(responses.get("1", 0), responses.get("2", 0)),
+        "Anger": responses.get("3", 0),
+        "Mania": max(responses.get("4", 0), responses.get("5", 0)),
+        "Anxiety": max(responses.get("6", 0), responses.get("7", 0), responses.get("8", 0)),
+        "Somatic_Symptoms": max(responses.get("9", 0), responses.get("10", 0)),
+        "Suicidal_Ideation": responses.get("11", 0),
+        "Psychosis": max(responses.get("12", 0), responses.get("13", 0)),
+        "Sleep_Problems": responses.get("14", 0),
+        "Memory": responses.get("15", 0),
+        "Repetitive_Thoughts_Behaviors": max(responses.get("16", 0), responses.get("17", 0)),
+        "Dissociation": responses.get("18", 0),
+        "Personality_Functioning": max(responses.get("19", 0), responses.get("20", 0)),
+        "Substance_Use": max(responses.get("21", 0), responses.get("22", 0), responses.get("23", 0)),
+    }
+    initial_question = "User has completed the initial questionnaire.Provide supportive steps and coping mechanisms."
+    initial_message = HumanMessage(content=initial_question)
+    
+    return {"domain_scores": domain_scores, "retry_count": 0, "messages": [initial_message]}
+
+def route_entry(state: GraphState)-> str:
+    if state.get("domain_scores"):
+        """Routes to the appropriate RAG handler based on scores."""
+        scores = state.get("domain_scores", {})
+        if scores.get("Depression", 0) >= 2:
+            return "depression"
+        if scores.get("Anxiety", 0) >= 2:
+            return "anxiety"
+        return "no_action"
+    else:
+        return "questionnaire"
+    
+def handle_depression_rag(state: GraphState) -> GraphState:
+    """Handles the conversational RAG pipeline for depression."""
+    score = state.get("domain_scores", {}).get("Depression", 0)
+    messages = state.get("messages", [])
+    retry_count = state.get("retry_count", 0)
+    
+    retry_guidance = "Please provide a helpful and supportive plan."
+    if retry_count > 0: retry_guidance = "Your previous response was flagged. Please try again."
+    
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", """You are a kind and empathetic AI assistant. Your role is to answer the user's question in a supportive, conversational tone.
+- **Do not** just summarize the documents. Synthesize the information and answer the user's question directly."""),
+        ("human", "The user has a depression score of {score} on a scale of 0 (None) to 4 (Severe).My question is: {question}. {retry_guidance}")
+    ])
+
+    rag_chain = prompt | llm | StrOutputParser()
+    response = rag_chain.invoke({ "score": score, "question": messages,"retry_guidance":retry_guidance})
+    return {"messages": [AIMessage(content=response)], "retry_count": retry_count + 1}
+
+def handle_anxiety_rag(state: GraphState) -> GraphState:
+    """Handles the conversational RAG pipeline for anxiety."""
+    score = state.get("domain_scores", {}).get("Depression", 0)
+    messages = state.get("messages", [])
+    retry_count = state.get("retry_count", 0)
+    
+    retry_guidance = "Please provide a helpful and supportive plan for someone feeling anxious or worried."
+    
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", """You are a kind and empathetic AI assistant. Your role is to answer the user's question in a supportive, conversational tone.
+- **Do not** just summarize the documents. Synthesize the information and answer the user's question directly."""),
+        ("human", "The user has a anxiety score of {score} on a scale of 0 (None) to 4 (Severe).My question is: {question}. {retry_guidance}")
+    ])
+
+    rag_chain = prompt | llm | StrOutputParser()
+    response = rag_chain.invoke({"score": score,  "question": messages,"retry_guidance":retry_guidance})
+    return {"messages": [AIMessage(content=response)], "retry_count": retry_count + 1}
+
+class SafetyCheck(BaseModel):
+    """Pydantic model for the safety check response."""
+    is_safe: bool = Field(description="True if the text is safe, False if it contains harmful content.")
+
+def safety_check(state: GraphState) -> GraphState:
+    """Checks the generated response for safety."""
+    ai_response = next(
+        (m.content for m in state.get("messages",[]) if isinstance(m, AIMessage)), ''
+    )
+    structured_llm = llm.with_structured_output(SafetyCheck)
+    
+    prompt = ChatPromptTemplate.from_template(
+        """You are a safety classification model. The user is in a mental health context. 
+        Analyze the following text for any harmful, unsafe, or non-supportive content.
+        Then, provide a JSON object with a single boolean field 'is_safe'.
+
+        Text to evaluate:
+        '{ai_response}'"""
+    )
+    
+    safety_chain = prompt | structured_llm
+    result = safety_chain.invoke({"ai_response": ai_response})
+    
+    return {"is_safe": result.is_safe}
+
+def handle_fallback(state: GraphState) -> GraphState:
+    """Provides a safe, generic response if retries fail."""
+    fallback_message = AIMessage(content="I am having trouble generating a specific plan right now. Please consider seeking support from a qualified professional.")
+    return {"messages": [fallback_message],"retry_count":0}
+
+def finalize_response(state: GraphState) -> GraphState:
+    """Finalizes the turn by returning the safe AI response as a message."""
+    return {"retry_count":0}
+
+def route_after_safety_check(state: "GraphState") -> str:
+    """Routes after the safety check, enabling the retry loop."""
+    if state.get("is_safe"): return "finalize"
+    if state.get("retry_count", 0) < 2: return "retry"
+    return "fallback"
+
+def entry_point(state: GraphState) -> GraphState:
+    """A dedicated node for the graph's entry point that makes no state changes."""
+    return state
+
+load_dotenv()
+groq_api_key = os.getenv("GROQ_API_KEY")
+llm = ChatGroq(groq_api_key=groq_api_key, model_name="llama-3.3-70b-versatile",max_tokens=4096)
+
+graph = StateGraph(GraphState)
+
+graph.add_node("entry_point",entry_point)
+graph.add_node("questionnaire",questionnaire)
+graph.add_node("handle_depression_rag", handle_depression_rag)
+graph.add_node("handle_anxiety_rag", handle_anxiety_rag)
+graph.add_node("safety_check_depression", safety_check)
+graph.add_node("safety_check_anxiety", safety_check)
+graph.add_node("handle_fallback", handle_fallback)
+graph.add_node("finalize_response", finalize_response)
+
+graph.add_edge(START,"entry_point")
+graph.add_conditional_edges("entry_point",route_entry,{"depression":"handle_depression_rag","anxiety":"handle_anxiety_rag","no_action":END,"questionnaire":"questionnaire"})
+graph.add_edge("questionnaire","entry_point")
+
+graph.add_edge("handle_depression_rag", "safety_check_depression")
+graph.add_conditional_edges(
+    "safety_check_depression", route_after_safety_check,
+    {"retry": "handle_depression_rag", "fallback": "handle_fallback", "finalize": "finalize_response"}
+)
+
+graph.add_edge("handle_anxiety_rag", "safety_check_anxiety")
+graph.add_conditional_edges(
+    "safety_check_anxiety", route_after_safety_check,
+    {"retry": "handle_anxiety_rag", "fallback": "handle_fallback", "finalize": "finalize_response"}
+)
+
+graph.add_edge("handle_fallback", END)
+graph.add_edge("finalize_response", END)
+DB_PATH = "/app/data/chatbot.sqlite"
+conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+checkpointer = SqliteSaver(conn=conn)
+app = graph.compile(checkpointer=checkpointer)

src/quick_help.py

@@ -0,0 +1,134 @@
+import os
+from typing import Annotated, TypedDict, List, Dict
+from dotenv import load_dotenv
+
+# LangChain and LangGraph imports
+from langchain_groq import ChatGroq
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from pydantic import BaseModel, Field
+from langgraph.graph import StateGraph, END
+from langgraph.graph.message import add_messages
+
+# --- 1. Pydantic Model for Structured Output ---
+class SafetyCheck(BaseModel):
+    """Pydantic model for the safety check response."""
+    is_safe: bool = Field(description="True if the text is safe, False if it contains harmful content.")
+
+# --- 2. Define the State for the Chatbot ---
+class ChatState(TypedDict):
+    """State of the chatbot, continuously appending messages."""
+    messages: Annotated[list[AnyMessage], add_messages]
+    is_safe: bool
+    retry_count: int
+
+# --- 3. Define the Chatbot Nodes ---
+def chat_node(state: ChatState, llm):
+    """
+    This node invokes the LLM to get a response using a ChatPromptTemplate.
+    """
+    messages = state['messages']
+    retry_count = state.get('retry_count', 0)
+    
+    # Create a prompt template with the system message
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a kind and empathetic AI assistant for mental well-being. "
+                   "Provide supportive, safe, and general advice. Do not provide medical diagnoses. "
+                   "Keep your responses concise and helpful."),
+        ("placeholder", "{chat_history}")
+    ])
+    
+    # Add a warning to the LLM if this is a retry attempt
+    if retry_count > 0:
+        messages.append(SystemMessage(
+            content="Your previous response was not suitable. Please try again, "
+                    "focusing on being supportive and safe."
+        ))
+
+    chain = prompt | llm | StrOutputParser()
+    
+    # Invoke the LLM with the conversation history
+    response_text = chain.invoke({"chat_history": messages})
+    
+    return {"messages": [AIMessage(content=response_text)], "retry_count": retry_count + 1}
+
+def safety_check_node(state: ChatState, llm) -> Dict:
+    """
+    Evaluates the safety of the generated response using a structured LLM.
+    """
+    # Get the most recent AI message to check
+    ai_message = state['messages'][-1]
+    
+    # Create a dedicated structured LLM for this task
+    structured_llm = llm.with_structured_output(SafetyCheck)
+    
+    prompt = ChatPromptTemplate.from_template(
+        """You are a safety classification model. The user is in a mental health context. 
+        Analyze the following text for any harmful, unsafe, or non-supportive content.
+        Then, provide a JSON object with a single boolean field 'is_safe'.
+
+        Text to evaluate:
+        '{text_to_evaluate}'"""
+    )
+    
+    safety_chain = prompt | structured_llm
+    result = safety_chain.invoke({"text_to_evaluate": ai_message.content})
+    
+    
+    return {"is_safe": result.is_safe}
+
+def handle_fallback_node(state: ChatState) -> Dict:
+    """
+    Provides a safe, generic response if the main LLM fails after retries.
+    """
+    fallback_message = AIMessage(
+        content="I am having a little trouble formulating a response right now. "
+                "Remember that taking a moment to focus on your breath can be a helpful step."
+    )
+    return {"messages": [fallback_message]}
+
+# --- 4. Define the Conditional Router ---
+def route_after_safety_check(state: ChatState) -> str:
+    """
+    This router decides the next step after the safety check, enabling a retry loop.
+    """
+    if state.get("is_safe"):
+        return "end"
+    if state.get("retry_count", 0) < 4:
+        return "retry"
+    return "fallback"
+
+# --- 5. Global Setup ---
+load_dotenv()
+os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
+groq_api_key = os.getenv("GROQ_API_KEY")
+
+# --- Model Configuration ---
+LLM_MODEL_NAME = "openai/gpt-oss-20b" 
+
+llm = ChatGroq(model_name=LLM_MODEL_NAME, groq_api_key=groq_api_key,max_tokens=4096)
+
+# --- 6. Build the Graph ---
+graph = StateGraph(ChatState)
+graph.add_node("chat_node", lambda state: chat_node(state, llm))
+graph.add_node("safety_check_node", lambda state: safety_check_node(state, llm))
+graph.add_node("handle_fallback_node", handle_fallback_node)
+
+graph.set_entry_point("chat_node")
+graph.add_edge("chat_node", "safety_check_node")
+graph.add_edge("handle_fallback_node", END)
+
+# Add the conditional edge for the retry loop
+graph.add_conditional_edges(
+    "safety_check_node",
+    route_after_safety_check,
+    {
+        "retry": "chat_node",
+        "fallback": "handle_fallback_node",
+        "end": END
+    }
+)
+
+# Compile the graph
+quick_help_app = graph.compile()

src/streamlit_app.py

@@ -0,0 +1,166 @@
+import streamlit as st
+import uuid
+from langchain_core.messages import HumanMessage, AIMessage
+
+# --- Backend Imports ---
+# Ensure these files are in the same directory as your Streamlit app
+from main_back import app
+from quick_help import quick_help_app
+
+# --- Page Configuration ---
+st.set_page_config(page_title="Gen AI Mental Health Assistant", layout="wide")
+
+st.title("Mental Health AI Assistant")
+st.markdown("A safe space to find support and track your well-being.")
+
+# --- Questionnaire Data ---
+# Based on the DSM-5 PDF provided earlier
+questionnaire_data = [
+    {"id": "1", "text": "Little interest or pleasure in doing things?"},
+    {"id": "2", "text": "Feeling down, depressed, or hopeless?"},
+    {"id": "3", "text": "Feeling more irritated, grouchy, or angry than usual?"},
+    {"id": "4", "text": "Sleeping less than usual, but still have a lot of energy?"},
+    {"id": "5", "text": "Starting lots more projects than usual or doing more risky things than usual?"},
+    {"id": "6", "text": "Feeling nervous, anxious, frightened, worried, or on edge?"},
+    {"id": "7", "text": "Feeling panic or being frightened?"},
+    {"id": "8", "text": "Avoiding situations that make you anxious?"},
+    {"id": "9", "text": "Unexplained aches and pains (e.g., head, back, joints, abdomen, legs)?"},
+    {"id": "10", "text": "Feeling that your illnesses are not being taken seriously enough?"},
+    {"id": "11", "text": "Thoughts of actually hurting yourself?"},
+    {"id": "12", "text": "Hearing things other people couldn't hear, such as voices even when no one was around?"},
+    {"id": "13", "text": "Feeling that someone could hear your thoughts, or that you could hear what another person was thinking?"},
+    {"id": "14", "text": "Problems with sleep that affected your sleep quality over all?"},
+    {"id": "15", "text": "Problems with memory (e.g., learning new information) or with location (e.g., finding your way home)?"},
+    {"id": "16", "text": "Unpleasant thoughts, urges, or images that repeatedly enter your mind?"},
+    {"id": "17", "text": "Feeling driven to perform certain behaviors or mental acts over and over again?"},
+    {"id": "18", "text": "Feeling detached or distant from yourself, your body, your physical surroundings, or your memories?"},
+    {"id": "19", "text": "Not knowing who you really are or what you want out of life?"},
+    {"id": "20", "text": "Not feeling close to other people or enjoying your relationships with them?"},
+    {"id": "21", "text": "Drinking at least 4 drinks of any kind of alcohol in a single day?"},
+    {"id": "22", "text": "Smoking any cigarettes, a cigar, or pipe, or using snuff or chewing tobacco?"},
+    {"id": "23", "text": "Using any medicines ON YOUR OWN, that is, without a doctor's prescription, in greater amounts or longer than prescribed?"},
+]
+response_options = ["None (0)", "Slight (1)", "Mild (2)", "Moderate (3)", "Severe (4)"]
+
+
+# --- App State Initialization ---
+def initialize_session_state():
+    # General app mode
+    if 'app_mode' not in st.session_state:
+        st.session_state.app_mode = "Quick Help"
+
+    # Quick Help state
+    if 'quick_help_history' not in st.session_state:
+        st.session_state.quick_help_history = []
+    
+    # Tracking Health state
+    if 'tracking_stage' not in st.session_state:
+        st.session_state.tracking_stage = "questionnaire" # questionnaire -> chat
+    if 'tracking_history' not in st.session_state:
+        st.session_state.tracking_history = []
+    if 'tracking_thread_id' not in st.session_state:
+        st.session_state.tracking_thread_id = None
+
+initialize_session_state()
+
+# --- Sidebar for Navigation ---
+with st.sidebar:
+    st.header("Navigation")
+    st.session_state.app_mode = st.radio(
+        "Choose a feature:",
+        ("Quick Help", "Track Your Health"),
+        key="app_mode_selector"
+    )
+    st.info("Your conversations are private. We do not store personally identifiable information.")
+
+# --- Main App Logic ---
+
+# --- Quick Help Feature ---
+if st.session_state.app_mode == "Quick Help":
+    st.header("Quick Help Chat")
+    st.markdown("Get immediate, supportive advice. How are you feeling right now?")
+
+    # Display chat history
+    for message in st.session_state.quick_help_history:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Handle user input
+    if user_input := st.chat_input("Share your thoughts..."):
+        st.session_state.quick_help_history.append({"role": "user", "content": user_input})
+        with st.chat_message("user"):
+            st.markdown(user_input)
+
+        with st.chat_message("assistant"):
+            # The input for the non-persistent app is the full history each time
+            history_for_input = [
+                HumanMessage(content=msg["content"]) if msg["role"] == "user" else AIMessage(content=msg["content"])
+                for msg in st.session_state.quick_help_history
+            ]
+            
+            ai_response = quick_help_app.invoke({'messages': history_for_input}).get("messages",["Could not generate response"])[-1].content
+            st.write(ai_response)
+
+        st.session_state.quick_help_history.append({"role": "assistant", "content": ai_response})
+
+# --- Track Your Health Feature ---
+elif st.session_state.app_mode == "Track Your Health":
+    
+    # --- Stage 1: Questionnaire ---
+    if st.session_state.tracking_stage == "questionnaire":
+        st.header("Health & Well-being Questionnaire")
+        st.markdown("Please answer the following questions based on your feelings over the **last two weeks**.")
+
+        with st.form("health_questionnaire"):
+            responses = {}
+            for q in questionnaire_data:
+                # Get the integer value from the selection
+                response_str = st.radio(q["text"], options=response_options, key=q["id"], horizontal=True)
+                responses[q["id"]] = response_options.index(response_str)
+
+            submitted = st.form_submit_button("Analyze & Create My Plan")
+            
+            if submitted:
+                # Generate a unique thread ID for this user's persistent session
+                st.session_state.tracking_thread_id = f"user_{uuid.uuid4()}"
+                config = {"configurable": {"thread_id": st.session_state.tracking_thread_id}}
+                
+                # Call the main backend app with the questionnaire responses
+                initial_input = {"questionnaire_responses": responses}
+                
+                with st.spinner("Analyzing your responses and generating a personalized plan..."):
+                    # Use .invoke() for the first call as we want the full plan at once
+                    result = app.invoke(initial_input, config=config)
+                    initial_plan = result.get('messages', ["Could not generate a plan."])[-1].content
+                    
+                # Store the initial plan and switch to chat mode
+                st.session_state.tracking_history = [{"role": "assistant", "content": initial_plan}]
+                st.session_state.tracking_stage = "chat"
+                st.rerun()
+
+    # --- Stage 2: Chat with the Plan ---
+    elif st.session_state.tracking_stage == "chat":
+        st.header("Your Personalized Plan & Chat")
+        st.markdown("Here is an initial plan based on your responses. You can ask questions about it or request different exercises.")
+        
+        config = {"configurable": {"thread_id": st.session_state.tracking_thread_id}}
+
+        # Display chat history
+        for message in st.session_state.tracking_history:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+
+        # Handle user input for follow-up questions
+        if user_input := st.chat_input("Ask a question about your plan..."):
+            st.session_state.tracking_history.append({"role": "user", "content": user_input})
+            with st.chat_message("user"):
+                st.markdown(user_input)
+
+            with st.chat_message("assistant"):
+                # For follow-ups, we only need to send the new message.
+                # The checkpointer on the backend handles loading the history.
+                follow_up_input = {"messages": [HumanMessage(content=user_input)]}
+                result = app.invoke(follow_up_input, config=config)
+                ai_response = result.get('messages', ["Could not generate a plan."])[-1].content
+
+            st.session_state.tracking_history.append({"role": "assistant", "content": ai_response})