Update app.py

app.py

@@ -1,131 +1,179 @@
+# app.py
 import os
 import requests
 import json
+import logging
 import pandas as pd
+import faiss
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# Load DSM-5 Dataset
-file_path = "dsm5_final_cleaned.csv"
-df = pd.read_csv(file_path)
+# ✅ Initialize FastAPI
+app = FastAPI()
 
-# OpenRouter API Configuration (DeepSeek Model)
-OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # Use environment variable for security
+# ✅ Set OpenRouter API Key
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 if not OPENROUTER_API_KEY:
-    raise ValueError("OPENROUTER_API_KEY is missing. Set it as an environment variable.")
+    raise ValueError("❌ OPENROUTER_API_KEY is missing. Set it as an environment variable.")
+
 OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"
 
-# Initialize FastAPI
-app = FastAPI()
+# ✅ Load AI Models
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+summarization_model = AutoModelForSeq2SeqLM.from_pretrained("google/long-t5-tglobal-base")
+summarization_tokenizer = AutoTokenizer.from_pretrained("google/long-t5-tglobal-base")
+
+# ✅ Load Datasets
+try:
+    recommendations_df = pd.read_csv("treatment_recommendations .csv")
+    questions_df = pd.read_csv("symptom_questions.csv")
+    print("✅ Datasets Loaded Successfully!")
+except FileNotFoundError as e:
+    logging.error(f"❌ Missing dataset file: {e}")
+    raise HTTPException(status_code=500, detail=f"Dataset file not found: {str(e)}")
+
+# ✅ Create FAISS Indexes
+question_embeddings = embedding_model.encode(questions_df["Questions"].tolist(), convert_to_numpy=True)
+question_index = faiss.IndexFlatL2(question_embeddings.shape[1])
+question_index.add(question_embeddings)
+
+treatment_embeddings = embedding_model.encode(recommendations_df["Disorder"].tolist(), convert_to_numpy=True)
+index = faiss.IndexFlatIP(treatment_embeddings.shape[1])
+index.add(treatment_embeddings)
 
-# Pydantic Models
+# ✅ Chat History Storage
+chat_history = []
+
+# ✅ Request Models
 class ChatRequest(BaseModel):
     message: str
 
 class SummaryRequest(BaseModel):
     chat_history: list
 
+
+# ✅ Function: Call DeepSeek via OpenRouter
 def deepseek_request(prompt, max_tokens=300):
-    """Helper function to send a request to DeepSeek API and handle response."""
-    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
+    """Send a request to OpenRouter's DeepSeek model."""
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "Content-Type": "application/json"
+    }
     payload = {
         "model": "deepseek/deepseek-r1-distill-llama-8b",
         "messages": [{"role": "user", "content": prompt}],
         "max_tokens": max_tokens,
-        "temperature": 0.7
+        "temperature": 0.8
     }
-    response = requests.post(OPENROUTER_BASE_URL, headers=headers, data=json.dumps(payload))
-    if response.status_code == 200:
+    try:
+        response = requests.post(OPENROUTER_BASE_URL, headers=headers, data=json.dumps(payload))
+        response.raise_for_status()
         response_json = response.json()
+        
         if "choices" in response_json and response_json["choices"]:
             return response_json["choices"][0].get("message", {}).get("content", "").strip()
-    return "Error: Unable to process the request."
-
-def match_disorders(chat_history):
-    """Match user symptoms with DSM-5 disorders based on keyword occurrence."""
-    disorder_scores = {}
-    for _, row in df.iterrows():
-        disorder = row["Disorder"]
-        keywords = row["Criteria"].split(", ")
-        match_count = sum(1 for word in keywords if word in chat_history.lower())
-        if match_count > 0:
-            disorder_scores[disorder] = match_count
-    sorted_disorders = sorted(disorder_scores, key=disorder_scores.get, reverse=True)
-    return sorted_disorders[:3] if sorted_disorders else []
+    except Exception as e:
+        logging.error(f"OpenRouter DeepSeek API error: {e}")
+        return "I'm here to support you. Can you share more about what you're feeling?"
 
-@app.post("/detect_disorders")
-def detect_disorders(request: SummaryRequest):
-    """Detect psychiatric disorders using DSM-5 keyword matching + DeepSeek validation."""
-    full_chat = " ".join(request.chat_history)
-    matched_disorders = match_disorders(full_chat)
-    
-    prompt = f"""
-    The following is a psychiatric conversation:
-    {full_chat}
-    
-    Based on DSM-5 diagnostic criteria, analyze the symptoms and determine the most probable psychiatric disorders.
-    Here are possible disorder matches from DSM-5 keyword analysis: {', '.join(matched_disorders) if matched_disorders else 'None found'}.
-    If no clear matches exist, diagnose based purely on symptom patterns and clinical reasoning.
-    Return a **list** of disorders, separated by commas, without extra text.
-    """
-    
-    response = deepseek_request(prompt, max_tokens=150)
-    disorders = [disorder.strip() for disorder in response.split(",")] if response and response.lower() != "unspecified disorder" else matched_disorders
-    return {"disorders": disorders if disorders else ["Unspecified Disorder"]}
+# ✅ Function: Retrieve Relevant Diagnostic Question
+def retrieve_relevant_question(user_input):
+    """Find the most relevant diagnostic question from the dataset using FAISS."""
+    input_embedding = embedding_model.encode([user_input], convert_to_numpy=True)
+    _, indices = question_index.search(input_embedding, 1)
 
-@app.post("/get_treatment")
-def get_treatment(request: SummaryRequest):
-    """Retrieve structured treatment recommendations based on detected disorders."""
-    detected_disorders = detect_disorders(request)["disorders"]
-    disorders_text = ", ".join(detected_disorders)
-    prompt = f"""
-    The user has been diagnosed with: {disorders_text}.
-    Provide a structured, evidence-based treatment plan including:
-    - Therapy recommendations (e.g., CBT, DBT, psychotherapy).
-    - Possible medications if applicable (e.g., SSRIs, anxiolytics, sleep aids).
-    - Lifestyle and self-care strategies (e.g., sleep hygiene, mindfulness, exercise).
-    If the user has suicidal thoughts, emphasize **immediate crisis intervention and emergency medical support.**
-    """
-    treatment_response = deepseek_request(prompt, max_tokens=200)
-    return {"treatments": treatment_response}
+    if indices[0][0] == -1:
+        return "I'm here to listen. Can you tell me more about your symptoms?"
 
-@app.post("/summarize_chat")
-def summarize_chat(request: SummaryRequest):
-    """Generate a structured summary of the psychiatric consultation."""
-    full_chat = " ".join(request.chat_history)
-    detected_disorders = detect_disorders(request)["disorders"]
-    treatment_response = get_treatment(request)["treatments"]
-    prompt = f"""
-    Summarize the following psychiatric conversation:
-    {full_chat}
-    
-    - **Detected Disorders:** {', '.join(detected_disorders)}
-    - **Suggested Treatments:** {treatment_response}
-    
-    The summary should include:
-    - Main concerns reported by the user.
-    - Key symptoms observed.
-    - Possible underlying psychological conditions.
-    - Recommended next steps, including professional consultation and self-care strategies.
-    If suicidal thoughts were mentioned, highlight the **need for immediate crisis intervention and professional support.**
-    """
-    summary = deepseek_request(prompt, max_tokens=300)
-    return {"summary": summary}
+    return questions_df["Questions"].iloc[indices[0][0]]
 
-@app.post("/chat")
+# ✅ API Endpoint: Chat Interaction
+@app.post("/get_questions")
 def chat(request: ChatRequest):
-    """Generate AI psychiatric response for user input."""
+    """Patient enters data, AI responds and stores conversation."""
+    user_message = request.message
+    chat_history.append(user_message)
+
+    # Retrieve relevant diagnostic question
+    relevant_question = retrieve_relevant_question(user_message)
+
+    # Constructing the DeepSeek prompt
     prompt = f"""
     You are an AI psychiatrist conducting a mental health consultation.
-    The user is discussing their concerns and symptoms. Engage in a supportive conversation,
-    ask relevant follow-up questions, and maintain an empathetic tone.
-    
+    Engage in a supportive, natural conversation, maintaining an empathetic tone.
+
+    - Always provide a thoughtful and compassionate response.
+    - If a user shares distressing emotions, acknowledge their feelings and ask relevant follow-up questions.
+    - Ask a symptom-related question to explore their concerns in depth.
+
+    Previous conversation:
+    {chat_history}
+
     User input:
-    {request.message}
-    
-    Provide a meaningful response and a follow-up question if necessary.
-    If the user mentions suicidal thoughts, respond with an urgent and compassionate tone,
-    suggesting that they seek immediate professional help while providing emotional support.
+    "{user_message}"
+
+    Generate:
+    - An empathetic response.
+    - A related follow-up question.
+    - The most relevant diagnostic question: "{relevant_question}".
+
+    Ensure your response is always meaningful and non-empty.
     """
-    response = deepseek_request(prompt, max_tokens=200)
-    return {"response": response}
+
+    ai_response = deepseek_request(prompt, max_tokens=250)
+
+    chat_history.append(ai_response)
+
+    return {"response": ai_response}
+
+# ✅ API Endpoint: Detect Disorders from Chat History
+@app.post("/detect_disorders")
+def detect_disorders():
+    """Detect psychiatric disorders based on full chat history."""
+    full_chat_text = " ".join(chat_history)
+    text_embedding = embedding_model.encode([full_chat_text], convert_to_numpy=True)
+    distances, indices = index.search(text_embedding, 3)
+
+    if indices[0][0] == -1:
+        return {"disorders": ["No matching disorder found."]}
+
+    disorders = [recommendations_df["Disorder"].iloc[i] for i in indices[0]]
+    return {"disorders": disorders}
+
+# ✅ API Endpoint: Get Treatment Recommendations
+@app.post("/get_treatment")
+def get_treatment():
+    """Retrieve treatment recommendations based on detected disorders."""
+    detected_disorders = detect_disorders()["disorders"]
+    treatments = {}
+
+    for disorder in detected_disorders:
+        if disorder in recommendations_df["Disorder"].values:
+            treatments[disorder] = recommendations_df[recommendations_df["Disorder"] == disorder]["Treatment Recommendation"].values[0]
+        else:
+            # Generate treatment if not in dataset
+            treatment_prompt = f"""
+            The user has been diagnosed with {disorder}. Provide a structured treatment plan including:
+
+            - **Therapy options** (CBT, psychotherapy, etc.).
+            - **Medications** (if applicable).
+            - **Lifestyle strategies** (exercise, mindfulness, etc.).
+            - **When to seek professional help**.
+            - **Encouragement**.
+
+            Ensure your response is clear and medically sound.
+            """
+            treatments[disorder] = deepseek_request(treatment_prompt, max_tokens=250)
+
+    return {"treatments": treatments}
+
+# ✅ API Endpoint: Summarize Chat
+@app.post("/summarize_chat")
+def summarize_chat():
+    """Summarize full chat session using DeepSeek."""
+    chat_text = " ".join(chat_history)
+    summary_prompt = f"The following is a conversation between a patient and an AI psychiatrist. Summarize it clearly:\n{chat_text}"
+    summary = deepseek_request(summary_prompt, max_tokens=500)
+    return {"summary": summary}