Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py (in TxAgent-API)
 import os
 import sys
 import json
@@ -9,9 +10,11 @@ import base64
 from datetime import datetime
 from typing import List, Dict, Optional, Tuple
 from enum import Enum
-from fastapi import FastAPI, HTTPException, UploadFile, File, Query, Form
+from fastapi import FastAPI, HTTPException, UploadFile, File, Query, Form, Depends
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.security import OAuth2PasswordBearer
+from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel
 import asyncio
 from bson import ObjectId
@@ -20,17 +23,19 @@ from gtts import gTTS
 from pydub import AudioSegment
 import PyPDF2
 import mimetypes
+from docx import Document
+from jose import JWTError, jwt
 from txagent.txagent import TxAgent
 from db.mongo import get_mongo_client
-from fastapi.encoders import jsonable_encoder
-from docx import Document 
+
 # Logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("TxAgentAPI")
 
 # App
-app = FastAPI(title="TxAgent API", version="2.6.0")  # Updated version for optional patient_id
+app = FastAPI(title="TxAgent API", version="2.6.0")
 
+# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -39,6 +44,13 @@ app.add_middleware(
     allow_headers=["*"]
 )
 
+# JWT settings (must match CPS-API)
+SECRET_KEY = os.getenv("SECRET_KEY", "your-secret-key")  # Same as CPS-API
+ALGORITHM = "HS256"
+
+# OAuth2 scheme (point to CPS-API's login endpoint)
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="https://rocketfarmstudios-cps-api.hf.space/auth/login")
+
 # Pydantic Models
 class ChatRequest(BaseModel):
     message: str
@@ -55,7 +67,7 @@ class VoiceOutputRequest(BaseModel):
     text: str
     language: str = "en"
     slow: bool = False
-    return_format: str = "mp3"  # mp3 or base64
+    return_format: str = "mp3"
 
 # Enums
 class RiskLevel(str, Enum):
@@ -71,7 +83,26 @@ patients_collection = None
 analysis_collection = None
 alerts_collection = None
 
-# Helpers
+# JWT validation
+async def get_current_user(token: str = Depends(oauth2_scheme)):
+    credentials_exception = HTTPException(
+        status_code=401,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        email: str = payload.get("sub")
+        if email is None:
+            raise credentials_exception
+    except JWTError:
+        raise credentials_exception
+    user = await users_collection.find_one({"email": email})
+    if user is None:
+        raise credentials_exception
+    return user
+
+# Helper functions (unchanged from your original code)
 def clean_text_response(text: str) -> str:
     text = re.sub(r'\n\s*\n', '\n\n', text)
     text = re.sub(r'[ ]+', ' ', text)
@@ -87,7 +118,6 @@ def extract_section(text: str, heading: str) -> str:
         return ""
 
 def structure_medical_response(text: str) -> Dict:
-    """Improved version that handles both markdown and plain text formats"""
     def extract_improved(text: str, heading: str) -> str:
         patterns = [
             rf"{re.escape(heading)}:\s*\n(.*?)(?=\n\s*\n|\Z)",
@@ -95,7 +125,6 @@ def structure_medical_response(text: str) -> Dict:
             rf"{re.escape(heading)}[\s\-]+(.*?)(?=\n\s*\n|\Z)",
             rf"\n{re.escape(heading)}\s*\n(.*?)(?=\n\s*\n|\Z)"
         ]
-        
         for pattern in patterns:
             match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
             if match:
@@ -103,9 +132,8 @@ def structure_medical_response(text: str) -> Dict:
                 content = re.sub(r'^\s*[\-\*]\s*', '', content, flags=re.MULTILINE)
                 return content
         return ""
-
-    text = text.replace('**', '').replace('__', '')
     
+    text = text.replace('**', '').replace('__', '')
     return {
         "summary": extract_improved(text, "Summary of Patient's Medical History") or 
                   extract_improved(text, "Summarize the patient's medical history"),
@@ -118,15 +146,12 @@ def structure_medical_response(text: str) -> Dict:
     }
 
 def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
-    """Analyze text for suicide risk factors and return assessment"""
     suicide_keywords = [
         'suicide', 'suicidal', 'kill myself', 'end my life', 
         'want to die', 'self-harm', 'self harm', 'hopeless',
         'no reason to live', 'plan to die'
     ]
-    
     explicit_mentions = [kw for kw in suicide_keywords if kw in text.lower()]
-    
     if not explicit_mentions:
         return RiskLevel.NONE, 0.0, []
     
@@ -145,7 +170,6 @@ def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
             temperature=0.2,
             max_new_tokens=256
         )
-        
         json_match = re.search(r'\{.*\}', response, re.DOTALL)
         if json_match:
             assessment = json.loads(json_match.group())
@@ -165,7 +189,6 @@ def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
     return RiskLevel.LOW, risk_score, explicit_mentions
 
 async def create_alert(patient_id: str, risk_data: dict):
-    """Create an alert document in the database"""
     alert_doc = {
         "patient_id": patient_id,
         "type": "suicide_risk",
@@ -185,16 +208,13 @@ def serialize_patient(patient: dict) -> dict:
     return patient_copy
 
 def compute_patient_data_hash(data: dict) -> str:
-    """Compute SHA-256 hash of patient data or report."""
     serialized = json.dumps(data, sort_keys=True)
     return hashlib.sha256(serialized.encode()).hexdigest()
 
 def compute_file_content_hash(file_content: bytes) -> str:
-    """Compute SHA-256 hash of file content."""
     return hashlib.sha256(file_content).hexdigest()
 
 def extract_text_from_pdf(pdf_data: bytes) -> str:
-    """Extract text from a PDF file."""
     try:
         pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_data))
         text = ""
@@ -206,85 +226,70 @@ def extract_text_from_pdf(pdf_data: bytes) -> str:
         raise HTTPException(status_code=400, detail="Failed to extract text from PDF")
 
 async def analyze_patient_report(patient_id: Optional[str], report_content: str, file_type: str, file_content: bytes):
-    """Analyze a patient report and store results."""
-    try:
-        # Use file content hash as identifier if no patient_id is provided
-        identifier = patient_id if patient_id else compute_file_content_hash(file_content)
-        report_data = {"identifier": identifier, "content": report_content, "file_type": file_type}
-        report_hash = compute_patient_data_hash(report_data)
-        logger.info(f"🧾 Analyzing report for identifier: {identifier}")
-
-        # Check if analysis exists and hash matches
-        existing_analysis = await analysis_collection.find_one({"identifier": identifier, "report_hash": report_hash})
-        if existing_analysis:
-            logger.info(f"✅ No changes in report data for {identifier}, skipping analysis")
-            return existing_analysis
-
-        # Construct analysis prompt
-        prompt = (
-            "You are a clinical decision support AI. Analyze the following patient report:\n"
-            "1. Summarize the patient's medical history.\n"
-            "2. Identify risks or red flags (including mental health and suicide risk).\n"
-            "3. Highlight missed diagnoses or treatments.\n"
-            "4. Suggest next clinical steps.\n"
-            f"\nPatient Report ({file_type}):\n{'-'*40}\n{report_content[:10000]}"
-        )
-
-        # Perform analysis
-        raw_response = agent.chat(
-            message=prompt,
-            history=[],
-            temperature=0.7,
-            max_new_tokens=1024
-        )
-        structured_response = structure_medical_response(raw_response)
+    identifier = patient_id if patient_id else compute_file_content_hash(file_content)
+    report_data = {"identifier": identifier, "content": report_content, "file_type": file_type}
+    report_hash = compute_patient_data_hash(report_data)
+    logger.info(f"🧾 Analyzing report for identifier: {identifier}")
+
+    existing_analysis = await analysis_collection.find_one({"identifier": identifier, "report_hash": report_hash})
+    if existing_analysis:
+        logger.info(f"✅ No changes in report data for {identifier}, skipping analysis")
+        return existing_analysis
+
+    prompt = (
+        "You are a clinical decision support AI. Analyze the following patient report:\n"
+        "1. Summarize the patient's medical history.\n"
+        "2. Identify risks or red flags (including mental health and suicide risk).\n"
+        "3. Highlight missed diagnoses or treatments.\n"
+        "4. Suggest next clinical steps.\n"
+        f"\nPatient Report ({file_type}):\n{'-'*40}\n{report_content[:10000]}"
+    )
 
-        # Suicide risk assessment
-        risk_level, risk_score, risk_factors = detect_suicide_risk(raw_response)
-        suicide_risk = {
-            "level": risk_level.value,
-            "score": risk_score,
-            "factors": risk_factors
-        }
+    raw_response = agent.chat(
+        message=prompt,
+        history=[],
+        temperature=0.7,
+        max_new_tokens=1024
+    )
+    structured_response = structure_medical_response(raw_response)
 
-        # Store analysis
-        analysis_doc = {
-            "identifier": identifier,
-            "patient_id": patient_id,  # May be None
-            "timestamp": datetime.utcnow(),
-            "summary": structured_response,
-            "suicide_risk": suicide_risk,
-            "raw": raw_response,
-            "report_hash": report_hash,
-            "file_type": file_type
-        }
+    risk_level, risk_score, risk_factors = detect_suicide_risk(raw_response)
+    suicide_risk = {
+        "level": risk_level.value,
+        "score": risk_score,
+        "factors": risk_factors
+    }
 
-        await analysis_collection.update_one(
-            {"identifier": identifier, "report_hash": report_hash},
-            {"$set": analysis_doc},
-            upsert=True
-        )
+    analysis_doc = {
+        "identifier": identifier,
+        "patient_id": patient_id,
+        "timestamp": datetime.utcnow(),
+        "summary": structured_response,
+        "suicide_risk": suicide_risk,
+        "raw": raw_response,
+        "report_hash": report_hash,
+        "file_type": file_type
+    }
 
-        # Create alert for high-risk cases only if patient_id is provided
-        if patient_id and risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
-            await create_alert(patient_id, suicide_risk)
+    await analysis_collection.update_one(
+        {"identifier": identifier, "report_hash": report_hash},
+        {"$set": analysis_doc},
+        upsert=True
+    )
 
-        logger.info(f"✅ Stored analysis for identifier {identifier}")
-        return analysis_doc
+    if patient_id and risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
+        await create_alert(patient_id, suicide_risk)
 
-    except Exception as e:
-        logger.error(f"Error analyzing patient report: {e}")
-        raise HTTPException(status_code=500, detail="Failed to analyze patient report")
+    logger.info(f"✅ Stored analysis for identifier {identifier}")
+    return analysis_doc
 
 async def analyze_all_patients():
-    """Analyze all patients in the database."""
     patients = await patients_collection.find({}).to_list(length=None)
     for patient in patients:
         await analyze_patient(patient)
         await asyncio.sleep(0.1)
 
 async def analyze_patient(patient: dict):
-    """Analyze patient data (existing logic for patient records)."""
     try:
         serialized = serialize_patient(patient)
         patient_id = serialized.get("fhir_id")
@@ -342,9 +347,7 @@ async def analyze_patient(patient: dict):
         logger.error(f"Error analyzing patient: {e}")
 
 def recognize_speech(audio_data: bytes, language: str = "en-US") -> str:
-    """Convert speech to text using Google's speech recognition."""
     recognizer = sr.Recognizer()
-    
     try:
         with io.BytesIO(audio_data) as audio_file:
             with sr.AudioFile(audio_file) as source:
@@ -362,7 +365,6 @@ def recognize_speech(audio_data: bytes, language: str = "en-US") -> str:
         raise HTTPException(status_code=500, detail="Error processing speech")
 
 def text_to_speech(text: str, language: str = "en", slow: bool = False) -> bytes:
-    """Convert text to speech using gTTS and return as MP3 bytes."""
     try:
         tts = gTTS(text=text, lang=language, slow=slow)
         mp3_fp = io.BytesIO()
@@ -394,6 +396,8 @@ async def startup_event():
     logger.info("✅ TxAgent initialized")
 
     db = get_mongo_client()["cps_db"]
+    global users_collection  # Add this to access users_collection for authentication
+    users_collection = db["users"]
     patients_collection = db["patients"]
     analysis_collection = db["patient_analysis_results"]
     alerts_collection = db["clinical_alerts"]
@@ -401,8 +405,10 @@ async def startup_event():
 
     asyncio.create_task(analyze_all_patients())
 
+# Protected Endpoints (add Depends(get_current_user) to all endpoints)
 @app.get("/status")
-async def status():
+async def status(current_user: dict = Depends(get_current_user)):
+    logger.info(f"Status endpoint accessed by {current_user['email']}")
     return {
         "status": "running",
         "timestamp": datetime.utcnow().isoformat(),
@@ -411,7 +417,11 @@ async def status():
     }
 
 @app.get("/patients/analysis-results")
-async def get_patient_analysis_results(name: Optional[str] = Query(None)):
+async def get_patient_analysis_results(
+    name: Optional[str] = Query(None),
+    current_user: dict = Depends(get_current_user)
+):
+    logger.info(f"Fetching analysis results by {current_user['email']}")
     try:
         query = {}
         if name:
@@ -438,7 +448,11 @@ async def get_patient_analysis_results(name: Optional[str] = Query(None)):
         raise HTTPException(status_code=500, detail="Failed to retrieve analysis results")
 
 @app.post("/chat-stream")
-async def chat_stream_endpoint(request: ChatRequest):
+async def chat_stream_endpoint(
+    request: ChatRequest,
+    current_user: dict = Depends(get_current_user)
+):
+    logger.info(f"Chat stream initiated by {current_user['email']}")
     async def token_stream():
         try:
             conversation = [{"role": "system", "content": agent.chat_prompt}]
@@ -472,9 +486,10 @@ async def chat_stream_endpoint(request: ChatRequest):
 @app.post("/voice/transcribe")
 async def transcribe_voice(
     audio: UploadFile = File(...),
-    language: str = Query("en-US", description="Language code for speech recognition")
+    language: str = Query("en-US", description="Language code for speech recognition"),
+    current_user: dict = Depends(get_current_user)
 ):
-    """Convert speech to text."""
+    logger.info(f"Voice transcription initiated by {current_user['email']}")
     try:
         audio_data = await audio.read()
         if not audio.filename.lower().endswith(('.wav', '.mp3', '.ogg', '.flac')):
@@ -490,8 +505,11 @@ async def transcribe_voice(
         raise HTTPException(status_code=500, detail="Error processing voice input")
 
 @app.post("/voice/synthesize")
-async def synthesize_voice(request: VoiceOutputRequest):
-    """Convert text to speech."""
+async def synthesize_voice(
+    request: VoiceOutputRequest,
+    current_user: dict = Depends(get_current_user)
+):
+    logger.info(f"Voice synthesis initiated by {current_user['email']}")
     try:
         audio_data = text_to_speech(request.text, request.language, request.slow)
         
@@ -515,9 +533,10 @@ async def voice_chat_endpoint(
     audio: UploadFile = File(...),
     language: str = Query("en-US", description="Language code for speech recognition"),
     temperature: float = Query(0.7, ge=0.1, le=1.0),
-    max_new_tokens: int = Query(512, ge=50, le=1024)
+    max_new_tokens: int = Query(512, ge=50, le=1024),
+    current_user: dict = Depends(get_current_user)
 ):
-    """Complete voice chat interaction (speech-to-text -> AI -> text-to-speech)."""
+    logger.info(f"Voice chat initiated by {current_user['email']}")
     try:
         audio_data = await audio.read()
         user_message = recognize_speech(audio_data, language)
@@ -548,18 +567,11 @@ async def analyze_clinical_report(
     file: UploadFile = File(...),
     patient_id: Optional[str] = Form(None),
     temperature: float = Form(0.5),
-    max_new_tokens: int = Form(1024)
+    max_new_tokens: int = Form(1024),
+    current_user: dict = Depends(get_current_user)
 ):
-    """
-    Analyze a clinical patient report from an uploaded file.
-    Parameters:
-    - file: Uploaded clinical report file (PDF, TXT, DOCX)
-    - patient_id: Optional patient ID to associate with this report
-    - temperature: Controls randomness of response (0.1-1.0)
-    - max_new_tokens: Maximum length of response
-    """
+    logger.info(f"Report analysis initiated by {current_user['email']}")
     try:
-        # Validate file type
         content_type = file.content_type
         allowed_types = [
             'application/pdf',
@@ -573,10 +585,8 @@ async def analyze_clinical_report(
                 detail=f"Unsupported file type: {content_type}. Supported types: PDF, TXT, DOCX"
             )
 
-        # Read file content
         file_content = await file.read()
 
-        # Extract text from file
         if content_type == 'application/pdf':
             text = extract_text_from_pdf(file_content)
         elif content_type == 'text/plain':
@@ -587,7 +597,6 @@ async def analyze_clinical_report(
         else:
             raise HTTPException(status_code=400, detail="Unsupported file type")
 
-        # Clean and validate text
         text = clean_text_response(text)
         if len(text.strip()) < 50:
             raise HTTPException(
@@ -595,7 +604,6 @@ async def analyze_clinical_report(
                 detail="Extracted text is too short (minimum 50 characters required)"
             )
 
-        # Analyze the report
         analysis = await analyze_patient_report(
             patient_id=patient_id,
             report_content=text,
@@ -603,13 +611,11 @@ async def analyze_clinical_report(
             file_content=file_content
         )
 
-        # Manually convert ObjectId and timestamp if needed
         if "_id" in analysis and isinstance(analysis["_id"], ObjectId):
             analysis["_id"] = str(analysis["_id"])
         if "timestamp" in analysis and isinstance(analysis["timestamp"], datetime):
             analysis["timestamp"] = analysis["timestamp"].isoformat()
 
-        # Return response using jsonable_encoder
         return JSONResponse(content=jsonable_encoder({
             "status": "success",
             "analysis": analysis,
@@ -627,7 +633,6 @@ async def analyze_clinical_report(
             detail=f"Failed to analyze report: {str(e)}"
         )
 
-        
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)