REFACTOR: Improve chat pipeline.

Refactoring the code pipeline for sending a message to prepare for implementing RAG.

src/api/routes/session.py

@@ -1,14 +1,11 @@
 # src/api/routes/chat.py
 
-from datetime import datetime, timezone
-
 from fastapi import APIRouter, Depends, HTTPException, status
 
+from src.core.response_pipeline import generate_chat_response
 from src.core.state import AppState, get_state
 from src.models.session import (ChatRequest, ChatResponse, Message, Session,
                                 SessionCreateRequest)
-from src.services.medical_response import generate_medical_response
-from src.services.guard import SafetyGuard
 from src.utils.logger import logger
 
 router = APIRouter(prefix="/session", tags=["Session & Chat"])
@@ -85,80 +82,21 @@ async def post_chat_message(
 	and persists the full exchange to long-term memory.
 	"""
 	logger().info(f"POST /session/{session_id}/messages")
-
-	# 0. Safety Guard: Validate user query
-	try:
-		safety_guard = SafetyGuard(state.nvidia_rotator)
-		is_safe, safety_reason = safety_guard.check_user_query(req.message)
-		if not is_safe:
-			logger().warning(f"Safety guard blocked user query: {safety_reason}")
-			raise HTTPException(
-				status_code=status.HTTP_400_BAD_REQUEST, 
-				detail=f"Query blocked for safety reasons: {safety_reason}"
-			)
-		logger().info(f"User query passed safety validation: {safety_reason}")
-	except Exception as e:
-		logger().error(f"Safety guard error: {e}")
-		# Fail open for now - allow query through if guard fails
-		logger().warning("Safety guard failed, allowing query through")
-
-	# 1. Get Enhanced Context
 	try:
-		medical_context = await state.memory_manager.get_enhanced_context(
+		response = await generate_chat_response(
+			state=state,
+			message=req.message,
 			session_id=session_id,
 			patient_id=req.patient_id,
-			question=req.message,
-			nvidia_rotator=state.nvidia_rotator
+			account_id=req.account_id
 		)
+		return ChatResponse(response=response)
+	except HTTPException as e:
+		# Re-raise HTTPException to let FastAPI handle it
+		raise e
 	except Exception as e:
-		logger().error(f"Error getting medical context: {e}")
-		raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to build medical context.")
-
-	# 2. Generate AI Response
-	try:
-		# In a real app, user role/specialty would come from the authenticated user
-		response_text = await generate_medical_response(
-			user_message=req.message,
-			user_role="Medical Professional",
-			user_specialty="",
-			rotator=state.gemini_rotator,
-			medical_context=medical_context,
-			nvidia_rotator=state.nvidia_rotator
+		logger().error(f"Unhandled error in chat pipeline: {e}")
+		raise HTTPException(
+			status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+			detail="An unexpected error occurred."
 		)
-	except Exception as e:
-		logger().error(f"Error generating medical response: {e}")
-		raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to generate AI response.")
-
-	# 2.5. Safety Guard: Validate AI response
-	try:
-		is_safe, safety_reason = safety_guard.check_model_answer(req.message, response_text)
-		if not is_safe:
-			logger().warning(f"Safety guard blocked AI response: {safety_reason}")
-			# Replace with safe fallback response
-			response_text = "I apologize, but I cannot provide a response to that query as it may contain unsafe content. Please consult with a qualified healthcare professional for medical advice."
-		else:
-			logger().info(f"AI response passed safety validation: {safety_reason}")
-	except Exception as e:
-		logger().error(f"Safety guard error for response: {e}")
-		# Fail open for now - allow response through if guard fails
-		logger().warning("Safety guard failed for response, allowing through")
-
-	# 3. Process and Store the Exchange
-	summary = await state.memory_manager.process_medical_exchange(
-		session_id=session_id,
-		patient_id=req.patient_id,
-		doctor_id=req.account_id,
-		question=req.message,
-		answer=response_text,
-		gemini_rotator=state.gemini_rotator,
-		nvidia_rotator=state.nvidia_rotator
-	)
-	if not summary:
-		logger().warning(f"Failed to process and store medical exchange for session {session_id}")
-
-	return ChatResponse(
-		response=response_text,
-		session_id=session_id,
-		timestamp=datetime.now(timezone.utc),
-		medical_context=medical_context
-	)

src/core/prompt_builder.py

@@ -2,15 +2,20 @@
 
 import json
 
+from src.models.account import Account
 from src.models.medical import MedicalMemory
 
 
-def medical_response_prompt(user_role: str, user_specialty: str, medical_context: str, user_message: str) -> str:
+def medical_response_prompt(
+	account: Account,
+	user_message: str,
+	medical_context: str | None = None
+) -> str:
 	"""Generates the prompt for creating a medical response."""
 	return f"""You are a knowledgeable medical AI assistant. Provide a comprehensive, accurate, and helpful response to this medical question.
-**User Role:** {user_role}
-**User Specialty:** {user_specialty if user_specialty else 'General'}
-**Medical Context:** {medical_context if medical_context else 'No previous context'}
+**User Role:** {account.role}
+**User Specialty:** {account.specialty or "No specialty"}
+**Medical Context:** {medical_context or 'No previous context'}
 **Question:** {user_message}
 **Instructions:**
 1. Provide a detailed, medically accurate response.

src/{services/medical_response.py → core/response_pipeline.py}

@@ -1,23 +1,116 @@
-# src/services/medical_response.py
+# src/core/response_pipeline.py
+
+from fastapi import HTTPException, status
 
 from src.core import prompt_builder
+from src.core.state import AppState
 from src.data.medical_kb import search_medical_kb
+from src.models.account import Account
 from src.services.gemini import gemini_chat
 from src.services.guard import SafetyGuard
 from src.utils.logger import logger
 from src.utils.rotator import APIKeyRotator
 
 
+async def generate_chat_response(
+	state: AppState,
+	message: str,
+	session_id: str,
+	patient_id: str,
+	account_id: str
+) -> str:
+	"""
+	Handles the pipeline for generating a chat response, including safety checks,
+	context retrieval, response generation, and memory persistence.
+	"""
+	logger().info(f"Starting response pipeline for session {session_id}")
+
+	# 0. Safety Guard: Validate user query
+	try:
+		safety_guard = SafetyGuard(state.nvidia_rotator)
+		is_safe, safety_reason = safety_guard.check_user_query(message)
+		if not is_safe:
+			logger().warning(f"Safety guard blocked user query: {safety_reason}")
+			raise HTTPException(
+				status_code=status.HTTP_400_BAD_REQUEST,
+				detail=f"Query blocked for safety reasons: {safety_reason}"
+			)
+		logger().info(f"User query passed safety validation: {safety_reason}")
+	except Exception as e:
+		logger().error(f"Safety guard error: {e}")
+		raise e
+
+	# 1. Get Enhanced Context
+	# TODO Implement RAG
+	try:
+		medical_context = await state.memory_manager.get_enhanced_context(
+			session_id=session_id,
+			patient_id=patient_id,
+			question=message,
+			nvidia_rotator=state.nvidia_rotator
+		)
+	except Exception as e:
+		logger().error(f"Error getting medical context: {e}")
+		raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to build medical context.")
+
+	account = state.memory_manager.get_account(account_id)
+	if not account:
+		raise Exception("Account not found")
+
+	# 2. Generate AI Response
+	try:
+		response_text = await generate_medical_response(
+			message=message,
+			account=account,
+			rotator=state.gemini_rotator,
+			medical_context=medical_context,
+			nvidia_rotator=state.nvidia_rotator
+		)
+	except Exception as e:
+		logger().error(f"Error generating medical response: {e}")
+		raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to generate AI response.")
+
+	# TODO Safety guard is applied twice
+	# 2.5. Safety Guard: Validate AI response
+	try:
+		is_safe, safety_reason = safety_guard.check_model_answer(message, response_text) # type: ignore
+		if not is_safe:
+			logger().warning(f"Safety guard blocked AI response: {safety_reason}")
+			response_text = "I apologize, but I cannot provide a response to that query as it may contain unsafe content. Please consult with a qualified healthcare professional for medical advice."
+		else:
+			logger().info(f"AI response passed safety validation: {safety_reason}")
+	except Exception as e:
+		logger().error(f"Safety guard error for response: {e}")
+		logger().warning("Safety guard failed for response, allowing through")
+
+	# 3. Process and Store the Exchange
+	summary = await state.memory_manager.process_medical_exchange(
+		session_id=session_id,
+		patient_id=patient_id,
+		doctor_id=account_id,
+		question=message,
+		answer=response_text,
+		gemini_rotator=state.gemini_rotator,
+		nvidia_rotator=state.nvidia_rotator
+	)
+	if not summary:
+		logger().warning(f"Failed to process and store medical exchange for session {session_id}")
+
+	return response_text
+
 async def generate_medical_response(
-	user_message: str,
-	user_role: str,
-	user_specialty: str,
+	account: Account,
+	message: str,
 	rotator: APIKeyRotator,
 	medical_context: str = "",
-	nvidia_rotator: APIKeyRotator = None
+	nvidia_rotator: APIKeyRotator | None = None
 ) -> str:
 	"""Generates an intelligent, contextual medical response using Gemini AI."""
-	prompt = prompt_builder.medical_response_prompt(user_role, user_specialty, medical_context, user_message)
+	prompt = prompt_builder.medical_response_prompt(
+		account=account,
+		user_message=message,
+		medical_context=medical_context
+	)
 
 	# Generate response using Gemini
 	response_text = await gemini_chat(prompt, rotator)
@@ -27,11 +120,12 @@ async def generate_medical_response(
 		if "disclaimer" not in response_text.lower() and "consult" not in response_text.lower():
 			response_text += "\n\n⚠️ **Important Disclaimer:** This information is for educational purposes only and should not replace professional medical advice, diagnosis, or treatment. Always consult with qualified healthcare professionals."
 
+		# TODO Safety guard is applied to the response twice
 		# Safety Guard: Validate the generated response
 		if nvidia_rotator:
 			try:
 				safety_guard = SafetyGuard(nvidia_rotator)
-				is_safe, safety_reason = safety_guard.check_model_answer(user_message, response_text)
+				is_safe, safety_reason = safety_guard.check_model_answer(message, response_text)
 				if not is_safe:
 					logger().warning(f"Safety guard blocked generated response: {safety_reason}")
 					# Return safe fallback response
@@ -47,16 +141,14 @@ async def generate_medical_response(
 		return response_text
 
 	logger().warning("Gemini response failed, using fallback.")
-	return _generate_fallback_response(user_message, user_role, user_specialty)
+	return _generate_fallback_response(message=message, account=account)
 
 def _generate_fallback_response(
-	user_message: str,
-	user_role: str,
-	user_specialty: str,
-	medical_context: str = ""
+	message: str,
+	account: Account
 ) -> str:
 	"""Generates a fallback response using a local knowledge base."""
-	kb_info = search_medical_kb(user_message)
+	kb_info = search_medical_kb(message)
 
 	logger().info("Generating backup response")
 
@@ -64,7 +156,7 @@ def _generate_fallback_response(
 	response_parts = []
 
 	# Analyze the question to provide more specific responses
-	question_lower = user_message.lower()
+	question_lower = message.lower()
 
 	if kb_info:
 		response_parts.append(f"Based on your question about medical topics, here's what I found:\n\n{kb_info}")
@@ -117,21 +209,21 @@ def _generate_fallback_response(
 			response_parts.append("Thank you for your medical question. While I can provide general information, it's important to consult with healthcare professionals for personalized medical advice.")
 
 	# Add role-specific guidance
-	if user_role.lower() in ["physician", "doctor", "nurse"]:
+	if account.role.lower() in ["physician", "doctor", "nurse"]:
 		response_parts.append("\n\n**Professional Context:** As a healthcare professional, you're likely familiar with these concepts. Remember to always follow your institution's protocols and guidelines, and consider the latest clinical evidence in your practice.")
-	elif user_role.lower() in ["medical student", "student"]:
+	elif account.role.lower() in ["medical student", "student"]:
 		response_parts.append("\n\n**Educational Context:** As a medical student, this information can help with your studies. Always verify information with your professors and clinical supervisors, and use this as a starting point for further research.")
-	elif user_role.lower() in ["patient"]:
+	elif account.role.lower() in ["patient"]:
 		response_parts.append("\n\n**Patient Context:** As a patient, this information is for educational purposes only. Please discuss any concerns with your healthcare provider, and don't make treatment decisions based solely on this information.")
 	else:
 		response_parts.append("\n\n**General Context:** This information is provided for educational purposes. Always consult with qualified healthcare professionals for medical advice.")
 
 	# Add specialty-specific information if available
-	if user_specialty and user_specialty.lower() in ["cardiology", "cardiac"]:
+	if account.specialty and account.specialty.lower() in ["cardiology", "cardiac"]:
 		response_parts.append("\n\n**Cardiology Perspective:** Given your interest in cardiology, consider how this information relates to cardiovascular health and patient care. Many conditions can have cardiac implications.")
-	elif user_specialty and user_specialty.lower() in ["pediatrics", "pediatric"]:
+	elif account.specialty and account.specialty.lower() in ["pediatrics", "pediatric"]:
 		response_parts.append("\n\n**Pediatric Perspective:** In pediatric care, remember that children may present differently than adults and may require specialized approaches. Consider age-appropriate considerations.")
-	elif user_specialty and user_specialty.lower() in ["emergency", "er"]:
+	elif account.specialty and account.specialty.lower() in ["emergency", "er"]:
 		response_parts.append("\n\n**Emergency Medicine Perspective:** In emergency settings, rapid assessment and intervention are crucial. Consider the urgency and severity of presenting symptoms.")
 
 	# Add medical disclaimer

src/models/session.py

@@ -44,13 +44,9 @@ class ChatRequest(BaseModel):
 	account_id: str # For context, though session_id implies this
 	patient_id: str # For context, though session_id implies this
 	message: str
-	session_id: str | None = None # Optional session ID for continuing existing sessions
 
 # --- API Response Models ---
 
 class ChatResponse(BaseModel):
 	"""Response model for a chat interaction."""
 	response: str
-	session_id: str
-	timestamp: datetime
-	medical_context: str | None = None