Spaces:
Sleeping
Sleeping
File size: 10,747 Bytes
73c6377 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 |
"""
Medical Query Router for RAG AI Advisor
"""
import asyncio
import logging
from fastapi import APIRouter, HTTPException, status
from fastapi.responses import StreamingResponse
import sys
import os
import json
# Add src to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
from core.agent import safe_run_agent, safe_run_agent_streaming, clear_session_memory, get_active_sessions
from api.models import ChatRequest, ChatResponse, HBVPatientInput, HBVAssessmentResponse
from typing import Optional
logger = logging.getLogger(__name__)
router = APIRouter(tags=["medical"])
def _build_contextual_query(
query: str,
patient_context: Optional[HBVPatientInput] = None,
assessment_result: Optional[HBVAssessmentResponse] = None
) -> str:
"""
Build an enhanced query that includes patient context and assessment results.
This helps the agent provide more relevant answers by understanding the specific
patient case being discussed.
Args:
query: The doctor's original question
patient_context: Optional patient data from assessment
assessment_result: Optional assessment result with eligibility and recommendations
Returns:
Enhanced query string with context
"""
if not patient_context and not assessment_result:
# No context, return original query
return query
context_parts = [query]
# Add patient context if available
if patient_context:
context_parts.append("\n\n[PATIENT CONTEXT FOR THIS QUESTION]")
context_parts.append(f"- Age: {patient_context.age}, Sex: {patient_context.sex}")
context_parts.append(f"- HBsAg: {patient_context.hbsag_status}, HBeAg: {patient_context.hbeag_status}")
context_parts.append(f"- HBV DNA: {patient_context.hbv_dna_level:,.0f} IU/mL")
context_parts.append(f"- ALT: {patient_context.alt_level} U/L")
context_parts.append(f"- Fibrosis: {patient_context.fibrosis_stage}")
if patient_context.pregnancy_status == "Pregnant":
context_parts.append(f"- Pregnancy: {patient_context.pregnancy_status}")
if patient_context.immunosuppression_status and patient_context.immunosuppression_status != "None":
context_parts.append(f"- Immunosuppression: {patient_context.immunosuppression_status}")
if patient_context.coinfections:
context_parts.append(f"- Coinfections: {', '.join(patient_context.coinfections)}")
# Add assessment result if available
if assessment_result:
context_parts.append("\n[PRIOR ASSESSMENT RESULT]")
context_parts.append(f"- Eligible for treatment: {assessment_result.eligible}")
# Include brief summary of recommendations (first 200 chars)
rec_summary = assessment_result.recommendations[:200] + "..." if len(assessment_result.recommendations) > 200 else assessment_result.recommendations
context_parts.append(f"- Assessment summary: {rec_summary}")
return "\n".join(context_parts)
@router.post("/ask", response_model=ChatResponse)
async def ask(request: ChatRequest):
"""
Interactive chat endpoint for doctors to ask questions about HBV guidelines.
This endpoint:
1. Accepts doctor's questions about HBV treatment guidelines
2. Maintains conversation context via session_id
3. Optionally includes patient context from prior assessment
4. Uses the same SASLT 2021 guidelines vector store as /assess
5. Returns evidence-based answers with guideline citations
Args:
request: ChatRequest containing query, session_id, and optional patient/assessment context
Returns:
ChatResponse with AI answer and session_id
"""
try:
# Validate input
if not request.query or not request.query.strip():
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Query cannot be empty"
)
if len(request.query) > 2000:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Query is too long. Maximum length is 2000 characters."
)
logger.info(f"Processing chat request - Session: {request.session_id}, Query length: {len(request.query)}")
# Build enhanced query with context if provided
enhanced_query = _build_contextual_query(
query=request.query,
patient_context=request.patient_context,
assessment_result=request.assessment_result
)
# Process through agent with session context
response = await safe_run_agent(
user_input=enhanced_query,
session_id=request.session_id
)
if not response or not response.strip():
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Received empty response from AI agent"
)
logger.info(f"Chat request completed - Session: {request.session_id}")
return ChatResponse(
response=response,
session_id=request.session_id
)
except HTTPException:
# Re-raise HTTP exceptions as-is
raise
except Exception as e:
logger.error(f"Error processing chat request: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error processing medical query: {str(e)}"
)
@router.post("/ask/stream")
async def ask_stream(request: ChatRequest):
"""
Interactive streaming chat endpoint for doctors to ask questions about HBV guidelines.
This endpoint:
1. Streams AI responses in real-time for better UX
2. Accepts doctor's questions about HBV treatment guidelines
3. Maintains conversation context via session_id
4. Optionally includes patient context from prior assessment
5. Uses the same SASLT 2021 guidelines vector store as /assess
6. Returns evidence-based answers with guideline citations
Args:
request: ChatRequest containing query, session_id, and optional patient/assessment context
Returns:
StreamingResponse with markdown-formatted AI answer
"""
# Validate input before starting stream
try:
if not request.query or not request.query.strip():
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Query cannot be empty"
)
if len(request.query) > 2000:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Query is too long. Maximum length is 2000 characters."
)
logger.info(f"Processing streaming chat request - Session: {request.session_id}, Query length: {len(request.query)}")
except HTTPException:
raise
except Exception as e:
logger.error(f"Validation error in streaming chat: {str(e)}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Invalid request: {str(e)}"
)
async def event_stream():
try:
# Build enhanced query with context if provided
enhanced_query = _build_contextual_query(
query=request.query,
patient_context=request.patient_context,
assessment_result=request.assessment_result
)
chunk_buffer = ""
chunk_count = 0
async for chunk in safe_run_agent_streaming(
user_input=enhanced_query,
session_id=request.session_id
):
chunk_buffer += chunk
chunk_count += 1
# Send chunks in reasonable sizes for smoother streaming
if len(chunk_buffer) >= 10:
yield chunk_buffer
chunk_buffer = ""
await asyncio.sleep(0.01)
# Send any remaining content
if chunk_buffer:
yield chunk_buffer
logger.info(f"Streaming chat completed - Session: {request.session_id}, Chunks: {chunk_count}")
except Exception as e:
error_msg = f"\n\n**Error**: An error occurred while processing your request. Please try again or contact support if the issue persists."
logger.error(f"Error in streaming chat: {str(e)}", exc_info=True)
yield error_msg
return StreamingResponse(event_stream(), media_type="text/markdown")
@router.delete("/session/{session_id}")
async def clear_session(session_id: str):
"""
Clear conversation history for a specific session.
This is useful when:
- Starting a new patient case
- Switching between different patient discussions
- Resetting the conversation context
Args:
session_id: The session identifier to clear
Returns:
Success message with session status
"""
try:
logger.info(f"Clearing session: {session_id}")
success = clear_session_memory(session_id)
if success:
return {
"status": "success",
"message": f"Session '{session_id}' cleared successfully",
"session_id": session_id
}
else:
return {
"status": "not_found",
"message": f"Session '{session_id}' not found or already cleared",
"session_id": session_id
}
except Exception as e:
logger.error(f"Error clearing session {session_id}: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error clearing session: {str(e)}"
)
@router.get("/sessions")
async def list_sessions():
"""
List all active chat sessions.
Returns:
List of active session IDs
"""
try:
sessions = get_active_sessions()
return {
"status": "success",
"active_sessions": sessions,
"count": len(sessions)
}
except Exception as e:
logger.error(f"Error listing sessions: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error listing sessions: {str(e)}"
)
|