api/routers/hbv_assessment.py

"""
HBV Patient Assessment Router
API endpoint for HBV treatment eligibility assessment
"""
from fastapi import APIRouter, HTTPException
from api.models import HBVPatientInput, HBVAssessmentResponse, TextAssessmentInput
import logging
from contextlib import contextmanager
from typing import List
from core.assessment_chain import (
    run_assessment_chain,
    run_assessment_chain_from_prompt,
    build_prompt_from_raw_text,
)

logger = logging.getLogger(__name__)


VERBOSE_LOGGER_NAMES = [
    "",
    "api",
    "api.app",
    "api.middleware",
    "api.routers",
    "api.routers.hbv_assessment",
    "core",
    "core.assessment_chain",
    "uvicorn",
    "uvicorn.error",
    "uvicorn.access",
]


TEXT_ASSESS_VERBOSE = True


@contextmanager
def temporary_verbose_logging(enabled: bool):
    """
    Temporarily raise logging levels to DEBUG so the request pipeline is visible.
    """
    if not enabled:
        yield
        return

    target_loggers: List[logging.Logger] = [logging.getLogger(name) for name in VERBOSE_LOGGER_NAMES]
    original_levels = [logger.level for logger in target_loggers]

    for log_obj in target_loggers:
        log_obj.setLevel(logging.DEBUG)

    logger.info("Verbose logging enabled for this request")

    try:
        yield
    finally:
        for log_obj, level in zip(target_loggers, original_levels):
            log_obj.setLevel(level)

router = APIRouter(
    prefix="",
    tags=["HBV Assessment"]
)


@router.post("/assess", response_model=HBVAssessmentResponse)
async def assess_patient(patient: HBVPatientInput) -> HBVAssessmentResponse:
    """
    Assess HBV patient eligibility for treatment according to SASLT 2021 guidelines
    
    This endpoint uses a LangChain Chain with hybrid approach:
    1. Phase 1 (Deterministic): Validates patient data, computes eligibility deterministically
    2. Phase 2 (LLM): Generates narrative recommendations with citations
    3. Returns structured assessment with eligibility and comprehensive recommendations
    
    Returns:
        HBVAssessmentResponse containing:
        - eligible: Whether patient is eligible for treatment
        - recommendations: Comprehensive narrative including eligibility determination, 
          specific criteria met, treatment options (ETV, TDF, TAF), and special considerations,
          with inline citations in format [SASLT 2021, Page X]
    """
    try:
        logger.info(f"Assessing HBV patient: Age {patient.age}, Sex {patient.sex}, HBV DNA {patient.hbv_dna_level}")
        
        # Convert Pydantic model to dict for chain
        patient_data = patient.dict()
        
        # Run assessment chain (hybrid: deterministic + LLM)
        result = run_assessment_chain(patient_data)
        
        # Convert dict result back to Pydantic response model
        response = HBVAssessmentResponse(**result)
        
        logger.info(f"Assessment complete: Eligible={response.eligible}")
        return response
        
    except Exception as e:
        logger.error(f"Error assessing patient: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error assessing patient: {str(e)}")


@router.post("/assess/text", response_model=HBVAssessmentResponse)
async def assess_patient_from_text(text_input: TextAssessmentInput) -> HBVAssessmentResponse:
    """
    Assess HBV patient eligibility from free-form text input
    
    This endpoint:
    1. Uses the free-form text directly as the LLM prompt (no intermediate structuring)
    2. Starts the assessment chain from the LLM invocation step
    3. Runs the same Phase 2 post-processing as the structured endpoint
    4. Returns structured assessment with eligibility and recommendations
    
    Example text input:
    "45-year-old male patient
    HBsAg: Positive for 12 months
    HBV DNA: 5000 IU/mL
    HBeAg: Positive
    ALT: 80 U/L
    Fibrosis stage: F2
    Necroinflammatory activity: A2
    No extrahepatic manifestations
    No immunosuppression
    No coinfections (HIV, HCV, HDV)
    No family history of cirrhosis or HCC"
    
    Returns:
        HBVAssessmentResponse containing:
        - eligible: Whether patient is eligible for treatment
        - recommendations: Comprehensive narrative with inline citations
    """
    try:
        raw_text = text_input.text_input
        logger.info(f"Received text input for assessment (length: {len(raw_text)} characters)")
        logger.info(f"Text input preview: {raw_text[:200]}...")

        # Enable verbose logging for observability when configured.
        with temporary_verbose_logging(TEXT_ASSESS_VERBOSE):
            prompt = build_prompt_from_raw_text(raw_text)

            # Run assessment chain starting from the LLM invocation step,
            # using the synthesized prompt that forces JSON output.
            logger.info("Performing HBV eligibility assessment from raw text prompt using LangChain chain...")
            result = run_assessment_chain_from_prompt(prompt)
        
        # Convert dict result to Pydantic response model
        response = HBVAssessmentResponse(**result)
        
        logger.info(f"Text-based assessment complete: Eligible={response.eligible}")
        return response
        
    except ValueError as e:
        logger.error(f"Validation error in text assessment: {str(e)}")
        raise HTTPException(status_code=400, detail=f"Invalid patient data: {str(e)}")
    except Exception as e:
        logger.error(f"Error in text-based assessment: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing text input: {str(e)}")