Update endpoints.py

endpoints.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Query, Form
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.encoders import jsonable_encoder
-from typing import Optional
+from typing import Optional, List
 from models import ChatRequest, VoiceOutputRequest, RiskLevel
 from auth import get_current_user
 from utils import clean_text_response
@@ -87,15 +87,61 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
                 )
 
                 text = agent.tokenizer.decode(output["sequences"][0][input_ids.shape[1]:], skip_special_tokens=True)
-                for chunk in text.split():
+                cleaned_text = clean_text_response(text)
+                full_response = ""
+
+                # Store chat session in database
+                chat_entry = {
+                    "user_id": current_user["email"],
+                    "patient_id": request.patient_id if request.patient_id else None,
+                    "message": request.message,
+                    "response": cleaned_text,
+                    "chat_type": "chat",
+                    "timestamp": datetime.utcnow(),
+                    "temperature": request.temperature,
+                    "max_new_tokens": request.max_new_tokens
+                }
+                result = await analysis_collection.insert_one(chat_entry)
+                chat_entry["_id"] = str(result.inserted_id)
+
+                for chunk in cleaned_text.split():
+                    full_response += chunk + " "
                     yield chunk + " "
                     await asyncio.sleep(0.05)
+
+                # Update chat entry with full response
+                await analysis_collection.update_one(
+                    {"_id": result.inserted_id},
+                    {"$set": {"response": full_response}}
+                )
+
             except Exception as e:
                 logger.error(f"Streaming error: {e}")
                 yield f"⚠️ Error: {e}"
 
         return StreamingResponse(token_stream(), media_type="text/plain")
 
+    @router.get("/chats")
+    async def get_chats(
+        current_user: dict = Depends(get_current_user)
+    ):
+        logger.info(f"Fetching chats for {current_user['email']}")
+        try:
+            chats = await analysis_collection.find({"user_id": current_user["email"], "chat_type": "chat"}).sort("timestamp", -1).to_list(length=100)
+            return [
+                {
+                    "id": str(chat["_id"]),
+                    "title": chat.get("message", "Untitled Chat")[:30],  # First 30 chars of message as title
+                    "timestamp": chat["timestamp"].isoformat(),
+                    "message": chat["message"],
+                    "response": chat["response"]
+                }
+                for chat in chats
+            ]
+        except Exception as e:
+            logger.error(f"Error fetching chats: {e}")
+            raise HTTPException(status_code=500, detail="Failed to retrieve chats")
+
     @router.post("/voice/transcribe")
     async def transcribe_voice(
         audio: UploadFile = File(...),
@@ -163,6 +209,20 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
             
             audio_data = text_to_speech(chat_response, language.split('-')[0])
             
+            # Store voice chat in database
+            chat_entry = {
+                "user_id": current_user["email"],
+                "patient_id": None,
+                "message": user_message,
+                "response": chat_response,
+                "chat_type": "voice_chat",
+                "timestamp": datetime.utcnow(),
+                "temperature": temperature,
+                "max_new_tokens": max_new_tokens
+            }
+            result = await analysis_collection.insert_one(chat_entry)
+            chat_entry["_id"] = str(result.inserted_id)
+
             return StreamingResponse(
                 io.BytesIO(audio_data),
                 media_type="audio/mpeg",
@@ -262,15 +322,15 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
             if patient.get("created_by") != current_user["email"] and not current_user.get("is_admin", False):
                 raise HTTPException(status_code=403, detail="Not authorized to delete this patient")
 
-            # Delete all analyses associated with this patient
+            # Delete all analyses and chats associated with this patient
             await analysis_collection.delete_many({"patient_id": patient_id})
-            logger.info(f"Deleted analyses for patient {patient_id}")
+            logger.info(f"Deleted analyses and chats for patient {patient_id}")
 
             # Delete the patient
             await patients_collection.delete_one({"fhir_id": patient_id})
             logger.info(f"Patient {patient_id} deleted successfully")
 
-            return {"status": "success", "message": f"Patient {patient_id} and associated analyses deleted"}
+            return {"status": "success", "message": f"Patient {patient_id} and associated analyses/chats deleted"}
 
         except HTTPException:
             raise