Update endpoints.py

endpoints.py

@@ -22,7 +22,7 @@ class ChatRequest(BaseModel):
     format: Optional[str] = "clean"
     temperature: Optional[float] = 0.7
     max_new_tokens: Optional[int] = 512
-    patient_id: Optional[str] = None  # Added optional patient_id field
+    patient_id: Optional[str] = None
 
 class VoiceOutputRequest(BaseModel):
     text: str
@@ -35,7 +35,7 @@ class RiskLevel(BaseModel):
     score: float
     factors: Optional[List[str]] = None
 
-def create_router(agent, logger, patients_collection, analysis_collection, users_collection):
+def create_router(agent, logger, patients_collection, analysis_collection, users_collection, chats_collection):
     router = APIRouter()
 
     @router.get("/status")
@@ -110,10 +110,10 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
                 cleaned_text = clean_text_response(text)
                 full_response = ""
 
-                # Store chat session in database
+                # Store chat session in the chats_collection
                 chat_entry = {
                     "user_id": current_user["email"],
-                    "patient_id": request.patient_id,  # Now safely optional, defaults to None
+                    "patient_id": request.patient_id,
                     "message": request.message,
                     "response": cleaned_text,
                     "chat_type": "chat",
@@ -121,8 +121,15 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
                     "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)
+                logger.info(f"Attempting to insert chat entry into chats_collection: {chat_entry}")
+                try:
+                    result = await chats_collection.insert_one(chat_entry)
+                    chat_entry["_id"] = str(result.inserted_id)
+                    logger.info(f"Successfully inserted chat entry with ID: {chat_entry['_id']}")
+                except Exception as db_error:
+                    logger.error(f"Failed to insert chat entry into chats_collection: {str(db_error)}")
+                    yield f"⚠️ Error: Failed to store chat in database: {str(db_error)}"
+                    return
 
                 for chunk in cleaned_text.split():
                     full_response += chunk + " "
@@ -130,10 +137,15 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
                     await asyncio.sleep(0.05)
 
                 # Update chat entry with full response
-                await analysis_collection.update_one(
-                    {"_id": result.inserted_id},
-                    {"$set": {"response": full_response}}
-                )
+                try:
+                    update_result = await chats_collection.update_one(
+                        {"_id": result.inserted_id},
+                        {"$set": {"response": full_response.strip()}}
+                    )
+                    logger.info(f"Updated chat entry {chat_entry['_id']}: matched {update_result.matched_count}, modified {update_result.modified_count}")
+                except Exception as update_error:
+                    logger.error(f"Failed to update chat entry {chat_entry['_id']}: {str(update_error)}")
+                    yield f"⚠️ Warning: Chat streamed successfully, but failed to update in database: {str(update_error)}"
 
             except Exception as e:
                 logger.error(f"Streaming error: {e}")
@@ -147,11 +159,12 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
     ):
         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)
+            chats = await chats_collection.find({"user_id": current_user["email"], "chat_type": "chat"}).sort("timestamp", -1).to_list(length=100)
+            logger.info(f"Retrieved {len(chats)} chats for {current_user['email']}")
             return [
                 {
                     "id": str(chat["_id"]),
-                    "title": chat.get("message", "Untitled Chat")[:30],  # First 30 chars of message as title
+                    "title": chat.get("message", "Untitled Chat")[:30],
                     "timestamp": chat["timestamp"].isoformat(),
                     "message": chat["message"],
                     "response": chat["response"]
@@ -229,7 +242,7 @@ 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
+            # Store voice chat in the chats_collection
             chat_entry = {
                 "user_id": current_user["email"],
                 "patient_id": None,
@@ -240,8 +253,14 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
                 "temperature": temperature,
                 "max_new_tokens": max_new_tokens
             }
-            result = await analysis_collection.insert_one(chat_entry)
-            chat_entry["_id"] = str(result.inserted_id)
+            logger.info(f"Attempting to insert voice chat entry into chats_collection: {chat_entry}")
+            try:
+                result = await chats_collection.insert_one(chat_entry)
+                chat_entry["_id"] = str(result.inserted_id)
+                logger.info(f"Successfully inserted voice chat entry with ID: {chat_entry['_id']}")
+            except Exception as db_error:
+                logger.error(f"Failed to insert voice chat entry into chats_collection: {str(db_error)}")
+                raise HTTPException(status_code=500, detail=f"Failed to store voice chat: {str(db_error)}")
 
             return StreamingResponse(
                 io.BytesIO(audio_data),
@@ -344,6 +363,7 @@ def create_router(agent, logger, patients_collection, analysis_collection, users
 
             # Delete all analyses and chats associated with this patient
             await analysis_collection.delete_many({"patient_id": patient_id})
+            await chats_collection.delete_many({"patient_id": patient_id})
             logger.info(f"Deleted analyses and chats for patient {patient_id}")
 
             # Delete the patient