Update endpoints.py

endpoints.py

@@ -1,242 +1,246 @@
-from fastapi import Depends, HTTPException, UploadFile, File, Query, Form
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Query, Form
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.encoders import jsonable_encoder
-from config import app, agent, logger
 from models import ChatRequest, VoiceOutputRequest, RiskLevel
 from auth import get_current_user
 from utils import clean_text_response
-from analysis import analyze_patient_report, analyze_all_patients
+from analysis import analyze_patient_report
 from voice import recognize_speech, text_to_speech, extract_text_from_pdf
 from docx import Document
 import re
-import mimetypes
-from bson import ObjectId
+import io
 from datetime import datetime
+from bson import ObjectId
 import asyncio
 
-@app.get("/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(),
-        "version": "2.6.0",
-        "features": ["chat", "voice-input", "voice-output", "patient-analysis", "report-upload"]
-    }
-
-@app.get("/patients/analysis-results")
-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:
-            name_regex = re.compile(name, re.IGNORECASE)
-            matching_patients = await patients_collection.find({"full_name": name_regex}).to_list(length=None)
-            patient_ids = [p["fhir_id"] for p in matching_patients if "fhir_id" in p]
-            if not patient_ids:
-                return []
-            query = {"patient_id": {"$in": patient_ids}}
-
-        analyses = await analysis_collection.find(query).sort("timestamp", -1).to_list(length=100)
-        enriched_results = []
-        for analysis in analyses:
-            patient = await patients_collection.find_one({"fhir_id": analysis.get("patient_id")})
-            if patient:
-                analysis["full_name"] = patient.get("full_name", "Unknown")
-            analysis["_id"] = str(analysis["_id"])
-            enriched_results.append(analysis)
-
-        return enriched_results
-
-    except Exception as e:
-        logger.error(f"Error fetching analysis results: {e}")
-        raise HTTPException(status_code=500, detail="Failed to retrieve analysis results")
-
-@app.post("/chat-stream")
-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():
+def create_router(agent, logger, patients_collection, analysis_collection):
+    router = APIRouter()
+
+    @router.get("/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(),
+            "version": "2.6.0",
+            "features": ["chat", "voice-input", "voice-output", "patient-analysis", "report-upload"]
+        }
+
+    @router.get("/patients/analysis-results")
+    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:
-            conversation = [{"role": "system", "content": agent.chat_prompt}]
-            if request.history:
-                conversation.extend(request.history)
-            conversation.append({"role": "user", "content": request.message})
-
-            input_ids = agent.tokenizer.apply_chat_template(
-                conversation, add_generation_prompt=True, return_tensors="pt"
-            ).to(agent.device)
-
-            output = agent.model.generate(
-                input_ids,
-                do_sample=True,
-                temperature=request.temperature,
-                max_new_tokens=request.max_new_tokens,
-                pad_token_id=agent.tokenizer.eos_token_id,
-                return_dict_in_generate=True
-            )
+            query = {}
+            if name:
+                name_regex = re.compile(name, re.IGNORECASE)
+                matching_patients = await patients_collection.find({"full_name": name_regex}).to_list(length=None)
+                patient_ids = [p["fhir_id"] for p in matching_patients if "fhir_id" in p]
+                if not patient_ids:
+                    return []
+                query = {"patient_id": {"$in": patient_ids}}
+
+            analyses = await analysis_collection.find(query).sort("timestamp", -1).to_list(length=100)
+            enriched_results = []
+            for analysis in analyses:
+                patient = await patients_collection.find_one({"fhir_id": analysis.get("patient_id")})
+                if patient:
+                    analysis["full_name"] = patient.get("full_name", "Unknown")
+                analysis["_id"] = str(analysis["_id"])
+                enriched_results.append(analysis)
+
+            return enriched_results
 
-            text = agent.tokenizer.decode(output["sequences"][0][input_ids.shape[1]:], skip_special_tokens=True)
-            for chunk in text.split():
-                yield chunk + " "
-                await asyncio.sleep(0.05)
         except Exception as e:
-            logger.error(f"Streaming error: {e}")
-            yield f"⚠️ Error: {e}"
-
-    return StreamingResponse(token_stream(), media_type="text/plain")
-
-@app.post("/voice/transcribe")
-async def transcribe_voice(
-    audio: UploadFile = File(...),
-    language: str = Query("en-US", description="Language code for speech recognition"),
-    current_user: dict = Depends(get_current_user)
-):
-    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')):
-            raise HTTPException(status_code=400, detail="Unsupported audio format")
+            logger.error(f"Error fetching analysis results: {e}")
+            raise HTTPException(status_code=500, detail="Failed to retrieve analysis results")
+
+    @router.post("/chat-stream")
+    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}]
+                if request.history:
+                    conversation.extend(request.history)
+                conversation.append({"role": "user", "content": request.message})
+
+                input_ids = agent.tokenizer.apply_chat_template(
+                    conversation, add_generation_prompt=True, return_tensors="pt"
+                ).to(agent.device)
+
+                output = agent.model.generate(
+                    input_ids,
+                    do_sample=True,
+                    temperature=request.temperature,
+                    max_new_tokens=request.max_new_tokens,
+                    pad_token_id=agent.tokenizer.eos_token_id,
+                    return_dict_in_generate=True
+                )
+
+                text = agent.tokenizer.decode(output["sequences"][0][input_ids.shape[1]:], skip_special_tokens=True)
+                for chunk in text.split():
+                    yield chunk + " "
+                    await asyncio.sleep(0.05)
+            except Exception as e:
+                logger.error(f"Streaming error: {e}")
+                yield f"⚠️ Error: {e}"
+
+        return StreamingResponse(token_stream(), media_type="text/plain")
+
+    @router.post("/voice/transcribe")
+    async def transcribe_voice(
+        audio: UploadFile = File(...),
+        language: str = Query("en-US", description="Language code for speech recognition"),
+        current_user: dict = Depends(get_current_user)
+    ):
+        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')):
+                raise HTTPException(status_code=400, detail="Unsupported audio format")
+            
+            text = recognize_speech(audio_data, language)
+            return {"text": text}
         
-        text = recognize_speech(audio_data, language)
-        return {"text": text}
-    
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error in voice transcription: {e}")
-        raise HTTPException(status_code=500, detail="Error processing voice input")
-
-@app.post("/voice/synthesize")
-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)
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Error in voice transcription: {e}")
+            raise HTTPException(status_code=500, detail="Error processing voice input")
+
+    @router.post("/voice/synthesize")
+    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)
+            
+            if request.return_format == "base64":
+                return {"audio": base64.b64encode(audio_data).decode('utf-8')}
+            else:
+                return StreamingResponse(
+                    io.BytesIO(audio_data),
+                    media_type="audio/mpeg",
+                    headers={"Content-Disposition": "attachment; filename=speech.mp3"}
+                )
         
-        if request.return_format == "base64":
-            return {"audio": base64.b64encode(audio_data).decode('utf-8')}
-        else:
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Error in voice synthesis: {e}")
+            raise HTTPException(status_code=500, detail="Error generating voice output")
+
+    @router.post("/voice/chat")
+    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),
+        current_user: dict = Depends(get_current_user)
+    ):
+        logger.info(f"Voice chat initiated by {current_user['email']}")
+        try:
+            audio_data = await audio.read()
+            user_message = recognize_speech(audio_data, language)
+            
+            chat_response = agent.chat(
+                message=user_message,
+                history=[],
+                temperature=temperature,
+                max_new_tokens=max_new_tokens
+            )
+            
+            audio_data = text_to_speech(chat_response, language.split('-')[0])
+            
             return StreamingResponse(
                 io.BytesIO(audio_data),
                 media_type="audio/mpeg",
-                headers={"Content-Disposition": "attachment; filename=speech.mp3"}
+                headers={"Content-Disposition": "attachment; filename=response.mp3"}
             )
-    
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error in voice synthesis: {e}")
-        raise HTTPException(status_code=500, detail="Error generating voice output")
-
-@app.post("/voice/chat")
-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),
-    current_user: dict = Depends(get_current_user)
-):
-    logger.info(f"Voice chat initiated by {current_user['email']}")
-    try:
-        audio_data = await audio.read()
-        user_message = recognize_speech(audio_data, language)
         
-        chat_response = agent.chat(
-            message=user_message,
-            history=[],
-            temperature=temperature,
-            max_new_tokens=max_new_tokens
-        )
-        
-        audio_data = text_to_speech(chat_response, language.split('-')[0])
-        
-        return StreamingResponse(
-            io.BytesIO(audio_data),
-            media_type="audio/mpeg",
-            headers={"Content-Disposition": "attachment; filename=response.mp3"}
-        )
-    
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error in voice chat: {e}")
-        raise HTTPException(status_code=500, detail="Error processing voice chat")
-
-@app.post("/analyze-report")
-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),
-    current_user: dict = Depends(get_current_user)
-):
-    logger.info(f"Report analysis initiated by {current_user['email']}")
-    try:
-        content_type = file.content_type
-        allowed_types = [
-            'application/pdf',
-            'text/plain',
-            'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
-        ]
-
-        if content_type not in allowed_types:
-            raise HTTPException(
-                status_code=400,
-                detail=f"Unsupported file type: {content_type}. Supported types: PDF, TXT, DOCX"
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Error in voice chat: {e}")
+            raise HTTPException(status_code=500, detail="Error processing voice chat")
+
+    @router.post("/analyze-report")
+    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),
+        current_user: dict = Depends(get_current_user)
+    ):
+        logger.info(f"Report analysis initiated by {current_user['email']}")
+        try:
+            content_type = file.content_type
+            allowed_types = [
+                'application/pdf',
+                'text/plain',
+                'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
+            ]
+
+            if content_type not in allowed_types:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Unsupported file type: {content_type}. Supported types: PDF, TXT, DOCX"
+                )
+
+            file_content = await file.read()
+
+            if content_type == 'application/pdf':
+                text = extract_text_from_pdf(file_content)
+            elif content_type == 'text/plain':
+                text = file_content.decode('utf-8')
+            elif content_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
+                doc = Document(io.BytesIO(file_content))
+                text = "\n".join([para.text for para in doc.paragraphs])
+            else:
+                raise HTTPException(status_code=400, detail="Unsupported file type")
+
+            text = clean_text_response(text)
+            if len(text.strip()) < 50:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Extracted text is too short (minimum 50 characters required)"
+                )
+
+            analysis = await analyze_patient_report(
+                patient_id=patient_id,
+                report_content=text,
+                file_type=content_type,
+                file_content=file_content
             )
 
-        file_content = await file.read()
-
-        if content_type == 'application/pdf':
-            text = extract_text_from_pdf(file_content)
-        elif content_type == 'text/plain':
-            text = file_content.decode('utf-8')
-        elif content_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
-            doc = Document(io.BytesIO(file_content))
-            text = "\n".join([para.text for para in doc.paragraphs])
-        else:
-            raise HTTPException(status_code=400, detail="Unsupported file type")
-
-        text = clean_text_response(text)
-        if len(text.strip()) < 50:
+            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 JSONResponse(content=jsonable_encoder({
+                "status": "success",
+                "analysis": analysis,
+                "patient_id": patient_id,
+                "file_type": content_type,
+                "file_size": len(file_content)
+            }))
+
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Error in report analysis: {str(e)}")
             raise HTTPException(
-                status_code=400,
-                detail="Extracted text is too short (minimum 50 characters required)"
+                status_code=500,
+                detail=f"Failed to analyze report: {str(e)}"
             )
 
-        analysis = await analyze_patient_report(
-            patient_id=patient_id,
-            report_content=text,
-            file_type=content_type,
-            file_content=file_content
-        )
-
-        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 JSONResponse(content=jsonable_encoder({
-            "status": "success",
-            "analysis": analysis,
-            "patient_id": patient_id,
-            "file_type": content_type,
-            "file_size": len(file_content)
-        }))
-
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error in report analysis: {str(e)}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Failed to analyze report: {str(e)}"
-        )
+    return router