utils/oneclick.py

import os
import json
import logging
from typing import Optional, Dict, Any
from huggingface_hub import InferenceClient
from utils.meldrx import MeldRxAPI
from utils.pdfutils import PDFGenerator
from datetime import datetime

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize Hugging Face Inference Client
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set. Please set your Hugging Face API token.")
client = InferenceClient(api_key=HF_TOKEN)
MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"  # Model to use for discharge summary generation

def generate_ai_discharge_summary(patient_data: Dict[str, Any]) -> Optional[str]:
    """
    Generate a discharge summary using the Hugging Face Inference Client based on patient data.

    Args:
        patient_data (Dict[str, Any]): Patient data in FHIR JSON format.

    Returns:
        Optional[str]: Generated discharge summary text or None if generation fails.
    """
    try:
        # Extract relevant patient information
        name = patient_data.get("name", [{}])[0]
        full_name = f"{name.get('given', ['Unknown'])[0]} {name.get('family', 'Unknown')}"
        gender = patient_data.get("gender", "Unknown").capitalize()
        birth_date = patient_data.get("birthDate", "Unknown")
        age = calculate_age(birth_date) if birth_date != "Unknown" else "Unknown"

        # Placeholder for additional clinical data (e.g., diagnosis, treatment)
        # In a real scenario, this would come from related FHIR resources like Encounter, Condition, etc.
        patient_info = (
            f"Patient Name: {full_name}\n"
            f"Gender: {gender}\n"
            f"Age: {age}\n\n"
            f"Presentation and Diagnosis:\n[Diagnosis data not provided in this snippet; assumed from related FHIR resources]\n\n"
            f"Hospital Course:\n[Treatment data not provided in this snippet; assumed from related FHIR resources]\n\n"
            f"Outcome:\n[Outcome data not provided in this snippet; assumed from related FHIR resources]"
        )

        # Define the prompt for the AI model
        messages = [
            {"role": "user", "content": ""},
            {
                "role": "assistant",
                "content": (
                    "You are a senior expert medical health practitioner known for producing discharge papers. "
                    "You will receive patient information and treatment details. Produce a complete discharge summary "
                    "based on the information provided."
                )
            },
            {"role": "user", "content": patient_info}
        ]

        # Generate discharge summary using streaming
        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

        return discharge_summary.strip()

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

def calculate_age(birth_date: str) -> str:
    """
    Calculate age from birth date.

    Args:
        birth_date (str): Birth date in YYYY-MM-DD format.

    Returns:
        str: Calculated age or 'Unknown' if calculation fails.
    """
    try:
        birth = datetime.strptime(birth_date, "%Y-%m-%d")
        today = datetime.today()
        age = today.year - birth.year - ((today.month, today.day) < (birth.month, birth.day))
        return str(age)
    except ValueError:
        return "Unknown"

def generate_discharge_paper_one_click(
    meldrx_api: MeldRxAPI,
    patient_id: str = None,
    first_name: str = None,
    last_name: str = None
) -> tuple[Optional[str], str]:
    """
    Generate a discharge paper with AI content in one click.

    Args:
        meldrx_api (MeldRxAPI): Initialized and authenticated MeldRxAPI instance.
        patient_id (str, optional): Patient ID to fetch specific patient data.
        first_name (str, optional): First name for patient lookup if patient_id is not provided.
        last_name (str, optional): Last name for patient lookup if patient_id is not provided.

    Returns:
        tuple[Optional[str], str]: (PDF file path, Status message)
    """
    try:
        # Check if already authenticated
        if not meldrx_api.access_token:
            return None, "Error: Not authenticated. Please authenticate first in the 'Authenticate with MeldRx' tab."

        # Fetch patient data
        if patient_id:
            patient_data = meldrx_api.get_patients()
            if not patient_data or "entry" not in patient_data:
                return None, "Error: Failed to fetch patient data by ID."
            patients = [entry["resource"] for entry in patient_data.get("entry", [])]
            patient = next((p for p in patients if p.get("id") == patient_id), None)
            if not patient:
                return None, f"Error: Patient with ID {patient_id} not found."
        else:
            patient_data = meldrx_api.get_patients()
            if not patient_data or "entry" not in patient_data:
                return None, "Error: Failed to fetch patient data."
            patients = [entry["resource"] for entry in patient_data.get("entry", [])]
            if first_name and last_name:
                patient = next(
                    (p for p in patients if
                     p.get("name", [{}])[0].get("given", [""])[0].lower() == first_name.lower() and
                     p.get("name", [{}])[0].get("family", "").lower() == last_name.lower()),
                    None
                )
                if not patient:
                    return None, f"Error: Patient with name {first_name} {last_name} not found."
            else:
                patient = patients[0] if patients else None
                if not patient:
                    return None, "Error: No patients found in the workspace."

        # Generate AI discharge summary
        ai_content = generate_ai_discharge_summary(patient)
        if not ai_content:
            return None, "Error: Failed to generate AI discharge summary."

        # Generate PDF
        pdf_generator = PDFGenerator()
        pdf_path = pdf_generator.generate_pdf_from_text(
            ai_content,
            f"discharge_summary_{patient.get('id', 'unknown')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        )

        if pdf_path:
            return pdf_path, f"Success: Discharge paper generated for {patient.get('name', [{}])[0].get('given', ['Unknown'])[0]} {patient.get('name', [{}])[0].get('family', 'Unknown')}"
        else:
            return None, "Error: Failed to generate PDF."

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