utils/oneclick.py

# utils/oneclick.py
from typing import Tuple, Optional, Dict
from .meldrx import MeldRxAPI
from .responseparser import PatientDataExtractor
from .pdfutils import PDFGenerator
from .verifier import DischargeVerifier
import logging
import json
from huggingface_hub import InferenceClient
import os

logger = logging.getLogger(__name__)

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set.")
client = InferenceClient(api_key=HF_TOKEN)
MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
verifier = DischargeVerifier()

def generate_ai_discharge_summary(patient_dict: Dict[str, str], client) -> Tuple[Optional[str], Optional[str]]:
    """Generate a discharge summary using AI and verify it for hallucinations."""
    try:
        formatted_summary = format_discharge_summary(patient_dict)
        
        logger.info("Generating AI discharge summary with patient info: %s", formatted_summary)

        messages = [
            {
                "role": "assistant",
                "content": (
                    "You are a senior medical practitioner tasked with creating discharge summaries. "
                    "Generate a complete discharge summary based on the provided patient information."
                )
            },
            {"role": "user", "content": formatted_summary}
        ]

        stream = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            temperature=0.4,
            max_tokens=3584,
            top_p=0.7,
            stream=True
        )

        discharge_summary = ""
        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                discharge_summary += content

        discharge_summary = discharge_summary.strip()
        logger.info("AI discharge summary generated successfully")

        # Verify the summary for hallucinations
        question = "Provide a complete discharge summary based on the patient information."
        verified_summary = verifier.verify_discharge_summary(
            context=formatted_summary,
            question=question,
            answer=discharge_summary
        )
        
        return discharge_summary, verified_summary

    except Exception as e:
        logger.error(f"Error generating AI discharge summary: {str(e)}", exc_info=True)
        return None, None

def generate_discharge_paper_one_click(
    api: MeldRxAPI,
    client,
    patient_id: str = "",
    first_name: str = "",
    last_name: str = ""
) -> Tuple[Optional[str], str, Optional[str], Optional[str], Optional[str]]:
    try:
        patients_data = api.get_patients()
        if not patients_data or "entry" not in patients_data:
            logger.error("No patient data received from MeldRx API")
            return None, "Failed to fetch patient data from MeldRx API", None, None, None

        logger.debug(f"Raw patient data from API: {patients_data}")
        
        extractor = PatientDataExtractor(patients_data, "json")
        
        if not extractor.patients:
            logger.error("No patients found in the parsed data")
            return None, "No patients found in the data", None, None, None

        logger.info(f"Found {len(extractor.patients)} patients in the data")
        
        matching_patients = []
        all_patient_ids = []
        all_patient_names = []
        
        for i in range(len(extractor.patients)):
            extractor.set_patient_by_index(i)
            patient_data = extractor.get_patient_dict()
            
            patient_id_from_data = str(patient_data.get('id', '')).strip().lower()
            first_name_from_data = str(patient_data.get('first_name', '')).strip().lower()
            last_name_from_data = str(patient_data.get('last_name', '')).strip().lower()
            
            all_patient_ids.append(patient_id_from_data)
            all_patient_names.append(f"{first_name_from_data} {last_name_from_data}".strip())
            
            patient_id_input = str(patient_id).strip().lower()
            first_name_input = str(first_name).strip().lower()
            last_name_input = str(last_name).strip().lower()
            
            logger.debug(f"Patient {i}: ID={patient_id_from_data}, Name={first_name_from_data} {last_name_from_data}")
            logger.debug(f"Comparing - Input: ID={patient_id_input}, First={first_name_input}, Last={last_name_input}")
            
            # Match logic: ID takes precedence, then first/last name
            matches = False
            if patient_id_input and patient_id_from_data == patient_id_input:
                matches = True
            elif not patient_id_input and first_name_input and last_name_input:
                if first_name_input == first_name_from_data and last_name_input == last_name_from_data:
                    matches = True
            elif not patient_id_input and not first_name_input and not last_name_input:
                continue  # Skip if no criteria provided
            
            if matches:
                matching_patients.append(patient_data)
                logger.info(f"Found matching patient: ID={patient_id_from_data}, "
                           f"Name={first_name_from_data} {last_name_from_data}")

        if not matching_patients:
            search_criteria = f"ID: {patient_id or 'N/A'}, First: {first_name or 'N/A'}, Last: {last_name or 'N/A'}"
            logger.warning(f"No patients matched criteria: {search_criteria}")
            logger.info(f"Available patient IDs: {all_patient_ids}")
            logger.info(f"Available patient names: {all_patient_names}")
            return None, (f"No patients found matching criteria: {search_criteria}\n"
                         f"Available IDs: {', '.join(all_patient_ids)}\n"
                         f"Available Names: {', '.join(all_patient_names)}"), None, None, None
        
        patient_data = matching_patients[0]  # Take the first match
        logger.info(f"Selected patient data: {patient_data}")
        
        basic_summary = format_discharge_summary(patient_data)
        ai_summary, verified_summary = generate_ai_discharge_summary(patient_data, client)
        
        if not ai_summary or not verified_summary:
            return None, "Failed to generate or verify AI summary", basic_summary, None, None
            
        pdf_gen = PDFGenerator()
        filename = f"discharge_{patient_data.get('id', 'unknown')}_{patient_data.get('last_name', 'patient')}.pdf"
        pdf_path = pdf_gen.generate_pdf_from_text(ai_summary, filename)
        
        if pdf_path:
            return pdf_path, "Discharge summary generated and verified successfully", basic_summary, ai_summary, verified_summary
        return None, "Failed to generate PDF file", basic_summary, ai_summary, verified_summary

    except Exception as e:
        logger.error(f"Error in one-click discharge generation: {str(e)}", exc_info=True)
        return None, f"Error generating discharge summary: {str(e)}", None, None, None

def format_discharge_summary(patient_data: dict) -> str:
    """Format patient data into a discharge summary text."""
    patient_data.setdefault('id', 'Unknown')
    patient_data.setdefault('first_name', '')
    patient_data.setdefault('last_name', '')
    patient_data.setdefault('dob', 'Unknown')
    patient_data.setdefault('age', 'Unknown')
    patient_data.setdefault('sex', 'Unknown')
    patient_data.setdefault('address', 'Unknown')
    patient_data.setdefault('city', 'Unknown')
    patient_data.setdefault('state', 'Unknown')
    patient_data.setdefault('zip_code', 'Unknown')
    patient_data.setdefault('phone', 'Unknown')
    patient_data.setdefault('admission_date', 'Unknown')
    patient_data.setdefault('discharge_date', 'Unknown')
    patient_data.setdefault('diagnosis', 'Unknown')
    patient_data.setdefault('medications', 'None specified')
    patient_data.setdefault('doctor_first_name', 'Unknown')
    patient_data.setdefault('doctor_last_name', 'Unknown')
    patient_data.setdefault('hospital_name', 'Unknown')
    patient_data.setdefault('doctor_address', 'Unknown')
    patient_data.setdefault('doctor_city', 'Unknown')
    patient_data.setdefault('doctor_state', 'Unknown')
    patient_data.setdefault('doctor_zip', 'Unknown')

    summary = [
        "DISCHARGE SUMMARY",
        "",
        "PATIENT INFORMATION",
        f"Name: {patient_data['first_name']} {patient_data['last_name']}".strip(),
        f"Patient ID: {patient_data['id']}",
        f"Date of Birth: {patient_data['dob']}",
        f"Age: {patient_data['age']}",
        f"Gender: {patient_data['sex']}",
        "",
        "CONTACT INFORMATION",
        f"Address: {patient_data['address']}",
        f"City: {patient_data['city']}, {patient_data['state']} {patient_data['zip_code']}",
        f"Phone: {patient_data['phone']}",
        "",
        "ADMISSION INFORMATION",
        f"Admission Date: {patient_data['admission_date']}",
        f"Discharge Date: {patient_data['discharge_date']}",
        f"Diagnosis: {patient_data['diagnosis']}",
        "",
        "MEDICATIONS",
        f"{patient_data['medications']}",
        "",
        "PHYSICIAN INFORMATION",
        f"Physician: Dr. {patient_data['doctor_first_name']} {patient_data['doctor_last_name']}".strip(),
        f"Hospital: {patient_data['hospital_name']}",
        f"Address: {patient_data['doctor_address']}",
        f"City: {patient_data['doctor_city']}, {patient_data['doctor_state']} {patient_data['doctor_zip']}",
    ]
    
    return "\n".join(line for line in summary if line.strip() or line == "")