""" FastAPI REST API layer wrapping the existing medical chatbot logic. Converts the Gradio interface to a REST API for React frontend integration. """ import hashlib import logging from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Optional, Dict, Any import json from app.main import ( chat_fn, _load_state, _save_state, _parse_response, _get_question_hint, ) from app.services.llm_extractor import extract_features_from_text from app.services.feature_builder import count_collected_features, is_ready_for_prediction, prepare_feature_vector from app.services.predictor import get_predictor from app.memory import initialize_state, update_state, get_missing_features from app.config import DEFAULT_MODEL_FEATURES, MIN_FEATURES_FOR_PREDICTION, CLASS_NAMES from app.utils.helpers import generate_question, prioritize_features from app.services.session_manager import get_session_manager logger = logging.getLogger(__name__) # Initialize FastAPI app = FastAPI(title="Medical Diagnosis AI", description="REST API for medical health prediction") # CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Session manager session_manager = get_session_manager() # ===== Request/Response Models ===== class ChatRequest(BaseModel): """Chat message request""" session_id: Optional[str] = None message: str history: list = [] class ChatResponse(BaseModel): """Chat response with features and state""" session_id: str message: str features: Dict[str, Any] collected_count: int total_features: int = 16 is_complete: bool prediction: Optional[Dict[str, Any]] = None hint: str = "" class ResetRequest(BaseModel): """Reset session request""" session_id: str class SessionStateResponse(BaseModel): """Session state response""" session_id: str features: Dict[str, Any] collected_count: int total_features: int = 16 # ===== Helper Functions ===== def _build_acknowledgment(extracted: dict) -> str: """Build acknowledgment message from extracted features""" extracted_items = [] for feature, value in extracted.items(): if value is not None and feature in DEFAULT_MODEL_FEATURES: extracted_items.append(f"{feature}: {value}") if extracted_items: return f"✓ Got your {', '.join(extracted_items[:2])}" return "" # ===== API Endpoints ===== @app.get("/health") def health_check(): """Health check endpoint""" return {"status": "ok", "service": "Medical Diagnosis AI"} @app.post("/api/chat", response_model=ChatResponse) def chat_endpoint(req: ChatRequest): """ Send a message and get AI response with updated features. Handles: - Session ID generation if not provided - Feature extraction from user message - State persistence - Prediction when all 16 features collected """ try: # Generate or use session ID if req.session_id: session_id = req.session_id else: # Generate from first message hash session_id = "sess_" + hashlib.md5(req.message.encode()).hexdigest()[:8] logger.info(f"🔐 Created new session: {session_id}") # Load persisted state state = _load_state(session_id) # Extract features from user message extracted = extract_features_from_text(req.message) # Update state with extracted features state = update_state(state, extracted) # Count collected features collected = count_collected_features(state) missing = get_missing_features(state) # Check if ready for prediction if is_ready_for_prediction(state, MIN_FEATURES_FOR_PREDICTION): # All 16 features collected - make prediction feature_vector = prepare_feature_vector(state) predictor = get_predictor() pred_result = predictor.predict(feature_vector) pred_data = { "prediction_class": int(pred_result.prediction), "prediction_name": CLASS_NAMES[int(pred_result.prediction)], "confidence": float(pred_result.probability), "risk_level": pred_result.risk_level, "explanation": pred_result.explanation, "features": state } # Simple completion message - full details now in DiagnosisCard component response_msg = "✅ Assessment Complete! Your diagnosis is ready below." _save_state(session_id, state) return ChatResponse( session_id=session_id, message=response_msg, features=state, collected_count=collected, is_complete=True, prediction=pred_data, hint="" ) # Not complete yet - ask for next missing feature prioritized_missing = prioritize_features(missing) next_question = generate_question(prioritized_missing[:1]) ack = _build_acknowledgment(extracted) remaining = 16 - collected response_msg = f"""{ack} {next_question} **{remaining} more pieces of information needed.**""" if ack else f"""{next_question} **{remaining} more pieces of information needed.**""" # Get hint for the next question hint = _get_question_hint(next_question) _save_state(session_id, state) return ChatResponse( session_id=session_id, message=response_msg, features=state, collected_count=collected, is_complete=False, hint=hint ) except Exception as e: logger.error(f"❌ Error in chat endpoint: {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/reset") def reset_endpoint(req: ResetRequest): """Reset a session - clear all features and start fresh""" try: session_manager = get_session_manager() success = session_manager.reset_session(req.session_id) if success: logger.info(f"✅ Reset session {req.session_id}") return { "success": True, "message": "Session reset successfully", "session_id": req.session_id } else: raise HTTPException(status_code=404, detail="Session not found") except Exception as e: logger.error(f"❌ Error resetting session: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/api/session/{session_id}", response_model=SessionStateResponse) def get_session_endpoint(session_id: str): """Get current session state""" try: state = _load_state(session_id) collected = count_collected_features(state) return SessionStateResponse( session_id=session_id, features=state, collected_count=collected ) except Exception as e: logger.error(f"❌ Error getting session: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/api/features") def get_features_list(): """Get list of all 16 features with their metadata""" from app.config import FEATURE_RANGES features_info = {} for feature in DEFAULT_MODEL_FEATURES: if feature in FEATURE_RANGES: min_val, max_val, _ = FEATURE_RANGES[feature] features_info[feature] = { "min": min_val, "max": max_val, "type": "numeric" if feature not in ["Smoking", "Alcohol", "Family History"] else "binary" } else: features_info[feature] = {"min": None, "max": None, "type": "unknown"} return { "total": len(DEFAULT_MODEL_FEATURES), "features": DEFAULT_MODEL_FEATURES, "metadata": features_info } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000, reload=True)