Update the app connection to MongoDB user storage + chat history LTM/STM. Distinguish patient/doctor roles on chat session

src/api/routes/chat.py

@@ -11,6 +11,8 @@ from src.services.medical_response import generate_medical_response
 from src.services.summariser import summarise_title_with_nvidia
 from src.utils.logger import get_logger
 
+from src.data.mongodb import ensure_session
+
 logger = get_logger("CHAT_ROUTES", __name__)
 router = APIRouter()
 
@@ -19,14 +21,16 @@ async def chat_endpoint(
 	request: ChatRequest,
 	state: MedicalState = Depends(get_state)
 ):
-	"""Handle chat messages and generate medical responses"""
+	"""
+	Process a chat message, generate response, and persist short-term cache + long-term Mongo.
+	"""
 	start_time = time.time()
 
 	try:
 		logger.info(f"Chat request from user {request.user_id} in session {request.session_id}")
 		logger.info(f"Message: {request.message[:100]}...")  # Log first 100 chars of message
 
-		# Get or create user profile
+		# Get or create user profile (doctor as current user profile)
 		user_profile = state.memory_system.get_user(request.user_id)
 		if not user_profile:
 			state.memory_system.create_user(request.user_id, request.user_role or "Anonymous")
@@ -37,7 +41,7 @@ async def chat_endpoint(
 				else:
 					logger.warning("Failued to retrieve user")
 
-		# Get or create session
+		# Get or create session (cache)
 		session = state.memory_system.get_session(request.session_id)
 		if not session:
 			session_id = state.memory_system.create_session(request.user_id, request.title or "New Chat")
@@ -45,11 +49,15 @@ async def chat_endpoint(
 			session = state.memory_system.get_session(session_id)
 			logger.info(f"Created new session: {session_id}")
 
-		# Get medical context from memory
+		# Ensure session exists in Mongo with patient/doctor context
+		ensure_session(session_id=request.session_id, patient_id=request.patient_id, doctor_id=request.doctor_id, title=request.title or "New Chat", last_activity=datetime.now(timezone.utc))
+
+		# Get medical context from memory (short-term) + Mongo long-term
 		medical_context = state.history_manager.get_conversation_context(
 			request.user_id,
 			request.session_id,
-			request.message
+			request.message,
+			patient_id=request.patient_id
 		)
 
 		# Generate response using Gemini AI
@@ -70,7 +78,10 @@ async def chat_endpoint(
 				request.message,
 				response,
 				state.gemini_rotator,
-				state.nvidia_rotator
+				state.nvidia_rotator,
+				patient_id=request.patient_id,
+				doctor_id=request.doctor_id,
+				session_title=request.title or "New Chat"
 			)
 		except Exception as e:
 			logger.warning(f"Failed to process medical exchange: {e}")

src/api/routes/session.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from src.core.state import MedicalState, get_state
 from src.models.chat import SessionRequest
 from src.utils.logger import get_logger
+from src.data.mongodb import list_patient_sessions, list_session_messages, ensure_session
 
 logger = get_logger("SESSION_ROUTES", __name__)
 router = APIRouter()
@@ -15,9 +16,11 @@ async def create_chat_session(
 	request: SessionRequest,
 	state: MedicalState = Depends(get_state)
 ):
-	"""Create a new chat session"""
+	"""Create a new chat session (cache + Mongo)"""
 	try:
 		session_id = state.memory_system.create_session(request.user_id, request.title or "New Chat")
+		# Also ensure in Mongo with patient/doctor
+		ensure_session(session_id=session_id, patient_id=request.patient_id, doctor_id=request.doctor_id, title=request.title or "New Chat")
 		return {"session_id": session_id, "message": "Session created successfully"}
 	except Exception as e:
 		logger.error(f"Error creating session: {e}")
@@ -28,7 +31,7 @@ async def get_chat_session(
 	session_id: str,
 	state: MedicalState = Depends(get_state)
 ):
-	"""Get chat session details and messages"""
+	"""Get session from cache (for quick preview)"""
 	try:
 		session = state.memory_system.get_session(session_id)
 		if not session:
@@ -52,6 +55,30 @@ async def get_chat_session(
 		logger.error(f"Error getting session: {e}")
 		raise HTTPException(status_code=500, detail=str(e))
 
+@router.get("/patients/{patient_id}/sessions")
+async def list_sessions_for_patient(patient_id: str):
+	"""List sessions for a patient from Mongo"""
+	try:
+		return {"sessions": list_patient_sessions(patient_id)}
+	except Exception as e:
+		logger.error(f"Error listing sessions: {e}")
+		raise HTTPException(status_code=500, detail=str(e))
+
+@router.get("/sessions/{session_id}/messages")
+async def list_messages_for_session(session_id: str, limit: int | None = None):
+	"""List messages for a session from Mongo"""
+	try:
+		msgs = list_session_messages(session_id, limit=limit)
+		# ensure JSON-friendly timestamps
+		for m in msgs:
+			if isinstance(m.get("timestamp"), datetime):
+				m["timestamp"] = m["timestamp"].isoformat()
+			m["_id"] = str(m["_id"]) if "_id" in m else None
+		return {"messages": msgs}
+	except Exception as e:
+		logger.error(f"Error listing messages: {e}")
+		raise HTTPException(status_code=500, detail=str(e))
+
 @router.delete("/sessions/{session_id}")
 async def delete_chat_session(
 	session_id: str,

src/core/memory/history.py

@@ -4,14 +4,16 @@ import json
 from typing import Any
 
 import numpy as np
+from datetime import datetime, timezone
 
 from src.services.nvidia import nvidia_chat
 from src.services.summariser import (summarise_qa_with_gemini,
                                      summarise_qa_with_nvidia)
 from src.utils.embeddings import EmbeddingClient
 from src.utils.logger import get_logger
+from src.data.mongodb import save_memory_summary, save_chat_message, ensure_session, get_recent_memory_summaries
 
-logger = get_logger("RAG", __name__)
+logger = get_logger("MED_HISTORY")
 
 def _safe_json(s: str) -> Any:
 	try:
@@ -91,9 +93,9 @@ class MedicalHistoryManager:
 		self.memory = memory
 		self.embedder = embedder
 
-	async def process_medical_exchange(self, user_id: str, session_id: str, question: str, answer: str, gemini_rotator, nvidia_rotator=None) -> str:
+	async def process_medical_exchange(self, user_id: str, session_id: str, question: str, answer: str, gemini_rotator, nvidia_rotator=None, *, patient_id: str | None = None, doctor_id: str | None = None, session_title: str | None = None) -> str:
 		"""
-		Process a medical Q&A exchange and store it in memory
+		Process a medical Q&A exchange and store it in memory and MongoDB
 		"""
 		try:
 			# Check if we have valid API keys
@@ -117,13 +119,21 @@ class MedicalHistoryManager:
 					logger.warning(f"Failed to create AI summary: {e}")
 					summary = f"q: {question}\na: {answer}"
 
-			# Store in memory
-			self.memory.add(user_id, summary)
+			# Short-term cache under patient_id when available
+			cache_key = patient_id or user_id
+			self.memory.add(cache_key, summary)
 
-			# Add to session history
+			# Add to session history in cache
 			self.memory.add_message_to_session(session_id, "user", question)
 			self.memory.add_message_to_session(session_id, "assistant", answer)
 
+			# Persist to MongoDB with patient/doctor context
+			if patient_id and doctor_id:
+				ensure_session(session_id=session_id, patient_id=patient_id, doctor_id=doctor_id, title=session_title or "New Chat", last_activity=datetime.now(timezone.utc))
+				save_chat_message(session_id=session_id, patient_id=patient_id, doctor_id=doctor_id, role="user", content=question)
+				save_chat_message(session_id=session_id, patient_id=patient_id, doctor_id=doctor_id, role="assistant", content=answer)
+				save_memory_summary(patient_id=patient_id, doctor_id=doctor_id, summary=summary)
+
 			# Update session title if it's the first message
 			session = self.memory.get_session(session_id)
 			if session and len(session.messages) == 2:  # Just user + assistant
@@ -137,18 +147,50 @@ class MedicalHistoryManager:
 			logger.error(f"Error processing medical exchange: {e}")
 			# Fallback: store without summary
 			summary = f"q: {question}\na: {answer}"
-			self.memory.add(user_id, summary)
+			cache_key = patient_id or user_id
+			self.memory.add(cache_key, summary)
 			self.memory.add_message_to_session(session_id, "user", question)
 			self.memory.add_message_to_session(session_id, "assistant", answer)
 			return summary
 
-	def get_conversation_context(self, user_id: str, session_id: str, question: str) -> str:
+	def get_conversation_context(self, user_id: str, session_id: str, question: str, *, patient_id: str | None = None) -> str:
 		"""
-		Get relevant conversation context for a new question
+		Get relevant conversation context combining short-term cache (3) and long-term Mongo (20)
 		"""
-		return self.memory.get_medical_context(user_id, session_id, question)
+		# Short-term summaries
+		cache_key = patient_id or user_id
+		recent_qa = self.memory.recent(cache_key, 3)
 
-	def get_user_medical_history(self, user_id: str, limit: int = 10) -> list[str]:
+		# Long-term summaries from Mongo (exclude ones already likely in cache by time order)
+		long_term = []
+		if patient_id:
+			try:
+				long_term = get_recent_memory_summaries(patient_id, limit=20)
+			except Exception as e:
+				logger.warning(f"Failed to fetch long-term memory: {e}")
+
+		# Get current session messages for context
+		session = self.memory.get_session(session_id)
+		session_context = ""
+		if session:
+			recent_messages = session.get_messages(10)
+			session_context = "\n".join([f"{msg['role']}: {msg['content']}" for msg in recent_messages])
+
+		# Combine context
+		context_parts = []
+		combined = []
+		if long_term:
+			combined.extend(long_term[::-1])  # oldest to newest within limit
+		if recent_qa:
+			combined.extend(recent_qa[::-1])
+		if combined:
+			context_parts.append("Recent medical context:\n" + "\n".join(combined[-20:]))
+		if session_context:
+			context_parts.append("Current conversation:\n" + session_context)
+
+		return "\n\n".join(context_parts) if context_parts else ""
+
+	def get_user_medical_history(self, user_id: str, limit: int = 20) -> list[str]:
 		"""
 		Get user's medical history (QA summaries)
 		"""

src/core/state.py

@@ -19,7 +19,8 @@ class MedicalState:
 
 	def initialize(self):
 		"""Initialize all core components"""
-		self.memory_system = MemoryLRU(capacity=50, max_sessions_per_user=20)
+		# Keep only 3 short-term summaries/messages in cache
+		self.memory_system = MemoryLRU(capacity=3, max_sessions_per_user=20)
 		self.embedding_client = create_embedding_client("all-MiniLM-L6-v2", dimension=384)
 		self.history_manager = MedicalHistoryManager(self.memory_system, self.embedding_client)
 		self.gemini_rotator = APIKeyRotator("GEMINI_API_", max_slots=5)
@@ -31,6 +32,7 @@ class MedicalState:
 			cls._instance = MedicalState()
 		return cls._instance
 
+
 def get_state() -> MedicalState:
 	"""FastAPI dependency for getting application state"""
 	return MedicalState.get_instance()

src/data/medical_kb.py

@@ -1,6 +1,7 @@
 # data/medical_kb.py
-# Medical Knowledge Base for the Medical AI Assistant
 
+# TODO: This should be replaced with a more robust knowledge base system that can be updated by the user.
+# Medical Knowledge Base for the Medical AI Assistant
 MEDICAL_KB = {
 	"symptoms": {
 		"fever": "Fever is a temporary increase in body temperature, often due to illness. Normal body temperature is around 98.6°F (37°C).",

src/data/mongodb.py

@@ -16,6 +16,7 @@ from pymongo.database import Database
 from pymongo.errors import DuplicateKeyError
 
 from src.utils.logger import get_logger
+import os
 
 logger = get_logger("MONGO")
 
@@ -32,8 +33,7 @@ def get_database() -> Database:
 	"""Get database instance with connection management"""
 	global _mongo_client
 	if _mongo_client is None:
-		# TODO This needs to use an environment variable when deployed
-		CONNECTION_STRING = "mongodb://127.0.0.1:27017/"
+		CONNECTION_STRING = os.getenv("MONGO_URI", "mongodb://127.0.0.1:27017/")
 		try:
 			logger.info("Initializing MongoDB connection")
 			_mongo_client = MongoClient(CONNECTION_STRING)
@@ -41,13 +41,14 @@ def get_database() -> Database:
 			logger.error(f"Failed to connect to MongoDB: {str(e)}")
 			# Pass the error down, code that calls this function should handle it
 			raise e
-	return _mongo_client['medicaldiagnosissystem']
+	db_name = os.getenv("MONGO_DB", "medicaldiagnosissystem")
+	return _mongo_client[db_name]
 
 def close_connection():
 	"""Close MongoDB connection"""
 	global _mongo_client
 	if _mongo_client is not None:
-		logger.info("Closing MongoDB connection")
+		# Close the connection and reset the client
 		_mongo_client.close()
 		_mongo_client = None
 
@@ -251,3 +252,114 @@ def backup_collection(collection_name: str) -> str:
 
 	logger.info(f"Created backup {backup_name} with {doc_count} documents")
 	return backup_name
+
+# New: Chat and Medical Memory Persistence Helpers
+
+CHAT_MESSAGES_COLLECTION = "chat_messages"
+MEDICAL_MEMORY_COLLECTION = "medical_memory"
+PATIENTS_COLLECTION = "patients"
+
+
+def ensure_session(
+	*,
+	session_id: str,
+	patient_id: str,
+	doctor_id: str,
+	title: str,
+	last_activity: datetime | None = None,
+	collection_name: str = CHAT_SESSIONS_COLLECTION
+) -> None:
+	collection = get_collection(collection_name)
+	now = datetime.now(timezone.utc)
+	collection.update_one(
+		{"session_id": session_id},
+		{"$set": {
+			"session_id": session_id,
+			"patient_id": patient_id,
+			"doctor_id": doctor_id,
+			"title": title,
+			"last_activity": (last_activity or now),
+			"updated_at": now
+		}, "$setOnInsert": {"created_at": now}},
+		upsert=True
+	)
+
+
+def save_chat_message(
+	*,
+	session_id: str,
+	patient_id: str,
+	doctor_id: str,
+	role: str,
+	content: str,
+	timestamp: datetime | None = None,
+	collection_name: str = CHAT_MESSAGES_COLLECTION
+) -> ObjectId:
+	collection = get_collection(collection_name)
+	ts = timestamp or datetime.now(timezone.utc)
+	doc = {
+		"session_id": session_id,
+		"patient_id": patient_id,
+		"doctor_id": doctor_id,
+		"role": role,
+		"content": content,
+		"timestamp": ts,
+		"created_at": ts
+	}
+	result = collection.insert_one(doc)
+	return result.inserted_id
+
+
+def list_session_messages(
+	session_id: str,
+	/,
+	*,
+	limit: int | None = None,
+	collection_name: str = CHAT_MESSAGES_COLLECTION
+) -> list[dict[str, Any]]:
+	collection = get_collection(collection_name)
+	cursor = collection.find({"session_id": session_id}).sort("timestamp", ASCENDING)
+	if limit is not None:
+		cursor = cursor.limit(limit)
+	return list(cursor)
+
+
+def save_memory_summary(
+	*,
+	patient_id: str,
+	doctor_id: str,
+	summary: str,
+	created_at: datetime | None = None,
+	collection_name: str = MEDICAL_MEMORY_COLLECTION
+) -> ObjectId:
+	collection = get_collection(collection_name)
+	ts = created_at or datetime.now(timezone.utc)
+	result = collection.insert_one({
+		"patient_id": patient_id,
+		"doctor_id": doctor_id,
+		"summary": summary,
+		"created_at": ts
+	})
+	return result.inserted_id
+
+
+def get_recent_memory_summaries(
+	patient_id: str,
+	/,
+	*,
+	limit: int = 20,
+	collection_name: str = MEDICAL_MEMORY_COLLECTION
+) -> list[str]:
+	collection = get_collection(collection_name)
+	docs = list(collection.find({"patient_id": patient_id}).sort("created_at", DESCENDING).limit(limit))
+	return [d.get("summary", "") for d in docs]
+
+
+def list_patient_sessions(
+	patient_id: str,
+	/,
+	*,
+	collection_name: str = CHAT_SESSIONS_COLLECTION
+) -> list[dict[str, Any]]:
+	collection = get_collection(collection_name)
+	return list(collection.find({"patient_id": patient_id}).sort("last_activity", DESCENDING))

src/models/chat.py

@@ -4,6 +4,8 @@ from pydantic import BaseModel
 
 class ChatRequest(BaseModel):
 	user_id: str
+	patient_id: str
+	doctor_id: str
 	session_id: str
 	message: str
 	user_role: str | None = "Medical Professional"
@@ -18,6 +20,8 @@ class ChatResponse(BaseModel):
 
 class SessionRequest(BaseModel):
 	user_id: str
+	patient_id: str
+	doctor_id: str
 	title: str | None = "New Chat"
 
 class SummariseRequest(BaseModel):

static/css/styles.css

@@ -858,3 +858,48 @@ body {
     color: var(--primary-color);
     margin-right: var(--spacing-sm);
 }
+
+.patient-section {
+    padding: var(--spacing-lg);
+    border-bottom: 1px solid var(--border-color);
+}
+
+.patient-header {
+    font-weight: 600;
+    margin-bottom: var(--spacing-sm);
+    color: var(--text-primary);
+}
+
+.patient-input-group {
+    display: flex;
+    gap: var(--spacing-sm);
+    align-items: center;
+}
+
+.patient-input {
+    flex: 1;
+    padding: 8px 10px;
+    border: 1px solid var(--border-color);
+    border-radius: 6px;
+    background: var(--bg-primary);
+    color: var(--text-primary);
+}
+
+.patient-load-btn {
+    padding: 8px 10px;
+    background: var(--primary-color);
+    color: #fff;
+    border: none;
+    border-radius: 6px;
+    cursor: pointer;
+}
+
+.patient-load-btn:hover {
+    background: var(--primary-hover);
+}
+
+.patient-status {
+    margin-top: var(--spacing-sm);
+    font-size: 0.8rem;
+    color: var(--text-secondary);
+}

static/index.html

@@ -26,7 +26,7 @@
                     </div>
                     <div class="user-info">
                         <div class="user-name" id="userName">Anonymous</div>
-                        <div class="user-status">Medical Professional</div>
+                        <div class="user-status" id="userStatus">Medical Professional</div>
                     </div>
                     <button class="user-menu-btn" id="userMenuBtn">
                         <i class="fas fa-ellipsis-v"></i>
@@ -34,6 +34,18 @@
                 </div>
             </div>
 
+            <!-- Patient Login Section -->
+            <div class="patient-section">
+                <div class="patient-header">Patient</div>
+                <div class="patient-input-group">
+                    <input type="text" id="patientIdInput" class="patient-input" placeholder="Enter 8-digit Patient ID" maxlength="8" inputmode="numeric" pattern="\\d{8}">
+                    <button class="patient-load-btn" id="loadPatientBtn" title="Load Patient">
+                        <i class="fas fa-user-injured"></i>
+                    </button>
+                </div>
+                <div class="patient-status" id="patientStatus">No patient selected</div>
+            </div>
+
             <div class="chat-sessions" id="chatSessions">
                 <!-- Chat sessions will be populated here -->
             </div>

static/js/app.js

@@ -2,7 +2,8 @@
 
 class MedicalChatbotApp {
     constructor() {
-        this.currentUser = null;
+        this.currentUser = null; // doctor
+        this.currentPatientId = null;
         this.currentSession = null;
         this.memory = new Map(); // In-memory storage for demo
         this.isLoading = false;
@@ -31,6 +32,20 @@ class MedicalChatbotApp {
             this.startNewChat();
         });
 
+        // Patient load button
+        const loadBtn = document.getElementById('loadPatientBtn');
+        if (loadBtn) {
+            loadBtn.addEventListener('click', () => this.loadPatient());
+        }
+        const patientInput = document.getElementById('patientIdInput');
+        if (patientInput) {
+            patientInput.addEventListener('keydown', (e) => {
+                if (e.key === 'Enter') {
+                    this.loadPatient();
+                }
+            });
+        }
+
         // Send button and input
         document.getElementById('sendBtn').addEventListener('click', () => {
             this.sendMessage();
@@ -267,11 +282,48 @@ I'm here to help you with medical questions, diagnosis assistance, and healthcar
 How can I assist you today?`;
     }
 
+    async loadPatient() {
+        const input = document.getElementById('patientIdInput');
+        const status = document.getElementById('patientStatus');
+        const id = (input?.value || '').trim();
+        if (!/^\d{8}$/.test(id)) {
+            status.textContent = 'Invalid patient ID. Use 8 digits.';
+            status.style.color = 'var(--warning-color)';
+            return;
+        }
+        // For now we accept ID and load sessions from backend
+        this.currentPatientId = id;
+        status.textContent = `Patient: ${id}`;
+        status.style.color = 'var(--text-secondary)';
+        await this.fetchAndRenderPatientSessions();
+    }
+
+    async fetchAndRenderPatientSessions() {
+        if (!this.currentPatientId) return;
+        try {
+            const resp = await fetch(`/sessions/patients/${this.currentPatientId}/sessions`.replace('/sessions/', '/sessions/'));
+            if (resp.ok) {
+                const data = await resp.json();
+                // Map to sidebar session cards if needed. For now, rely on local sessions until full backend sync is added.
+                // Future: hydrate local UI from data.sessions
+            }
+        } catch (e) {
+            console.error('Failed to load patient sessions', e);
+        }
+        this.loadChatSessions();
+    }
+
     async sendMessage() {
         const input = document.getElementById('chatInput');
         const message = input.value.trim();
 
         if (!message || this.isLoading) return;
+        if (!this.currentPatientId) {
+            const status = document.getElementById('patientStatus');
+            status.textContent = 'Select a patient before chatting.';
+            status.style.color = 'var(--warning-color)';
+            return;
+        }
 
         // Clear input
         input.value = '';
@@ -322,6 +374,8 @@ How can I assist you today?`;
                 },
                 body: JSON.stringify({
                 user_id: this.currentUser.id,
+                patient_id: this.currentPatientId,
+                doctor_id: this.currentUser.id,
                 session_id: this.currentSession?.id || 'default',
                 message: message,
                 user_role: this.currentUser.role,
@@ -344,7 +398,8 @@ How can I assist you today?`;
                 message: error.message,
                 stack: error.stack,
                 user: this.currentUser,
-                session: this.currentSession
+                session: this.currentSession,
+                patientId: this.currentPatientId
             });
 
             // Only return mock response if it's a network error, not a server error