report_generator.py

#!/usr/bin/env python3
"""
Pre-Visit Report Generator

Creates a concise, one-page summary for healthcare providers including:
- Patient's main concerns
- How they're feeling
- Relevant medical context
- Attachments (voice recordings, charts)
"""

import os
import json
from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass, field
import httpx

LLAMA_SERVER_URL = os.getenv("LLAMA_SERVER_URL", "http://localhost:8081")

LLM_HEADERS = {
    "Content-Type": "application/json",
    "ngrok-skip-browser-warning": "true"
}


@dataclass
class ReportAttachment:
    """Attachment for the report (voice recording, chart, etc.)"""
    type: str  # "audio", "chart", "image"
    title: str
    data: Optional[str] = None  # base64 or URL
    summary: Optional[str] = None  # Text summary of the attachment


@dataclass
class PreVisitReport:
    """Structured pre-visit report for healthcare provider."""
    patient_name: str
    patient_age: int
    patient_gender: str
    generated_at: str
    
    # Core content
    chief_concerns: List[str] = field(default_factory=list)
    patient_feelings: str = ""
    symptom_summary: str = ""
    
    # Medical context
    relevant_conditions: List[str] = field(default_factory=list)
    current_medications: List[str] = field(default_factory=list)
    recent_vitals: Dict[str, str] = field(default_factory=dict)
    
    # NEW: Additional medical history
    immunizations: List[Dict] = field(default_factory=list)
    procedures: List[Dict] = field(default_factory=list)
    recent_encounters: List[Dict] = field(default_factory=list)
    allergies: List[Dict] = field(default_factory=list)
    
    # Attachments
    attachments: List[Dict] = field(default_factory=list)
    
    # Conversation references
    key_quotes: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict:
        return {
            "patient_name": self.patient_name,
            "patient_age": self.patient_age,
            "patient_gender": self.patient_gender,
            "generated_at": self.generated_at,
            "chief_concerns": self.chief_concerns,
            "patient_feelings": self.patient_feelings,
            "symptom_summary": self.symptom_summary,
            "relevant_conditions": self.relevant_conditions,
            "current_medications": self.current_medications,
            "recent_vitals": self.recent_vitals,
            "immunizations": self.immunizations,
            "procedures": self.procedures,
            "recent_encounters": self.recent_encounters,
            "allergies": self.allergies,
            "attachments": self.attachments,
            "key_quotes": self.key_quotes
        }


async def call_llm(prompt: str, max_tokens: int = 1024) -> str:
    """Call LLM for report generation."""
    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            f"{LLAMA_SERVER_URL}/completion",
            headers=LLM_HEADERS,
            json={
                "prompt": prompt,
                "n_predict": max_tokens,
                "temperature": 0.3,
                "stop": ["<end_of_turn>", "</s>", "<|im_end|>"],
                "stream": False
            }
        )
        response.raise_for_status()
        result = response.json()
        return result.get("content", "").strip()


async def generate_report(
    patient_info: Dict,
    conversation_history: List[Dict],
    tool_results: List[Dict],
    attachments: List[Dict] = None,
    immunizations: List[Dict] = None,
    procedures: List[Dict] = None,
    encounters: List[Dict] = None,
    allergies: List[Dict] = None
) -> PreVisitReport:
    """
    Generate a pre-visit report from conversation and medical data.
    
    Args:
        patient_info: Patient demographics
        conversation_history: List of {"role": "user"|"assistant", "content": "..."}
        tool_results: List of {"tool": "...", "facts": "..."}
        attachments: List of {"type": "audio"|"chart", "title": "...", "summary": "..."}
        immunizations: List of immunization records from DB
        procedures: List of procedure records from DB
        encounters: List of encounter records from DB
        allergies: List of allergy records from DB
    
    Returns:
        PreVisitReport object
    """
    
    # Format conversation for analysis
    conversation_text = ""
    for msg in conversation_history:
        role = "Patient" if msg.get("role") == "user" else "Assistant"
        conversation_text += f"{role}: {msg.get('content', '')}\n\n"
    
    # Format tool results
    medical_context = ""
    for result in tool_results:
        tool = result.get("tool", "").replace("get_", "").replace("_", " ")
        facts = result.get("facts", "")
        if facts:
            medical_context += f"[{tool.upper()}]\n{facts}\n\n"
    
    # Format attachments
    attachment_text = ""
    if attachments:
        for att in attachments:
            attachment_text += f"- {att.get('type', 'file').upper()}: {att.get('title', 'Untitled')}\n"
            if att.get('summary'):
                attachment_text += f"  Summary: {att.get('summary')}\n"
    
    # Generate report using LLM
    prompt = f"""<start_of_turn>user
You are a medical assistant creating a concise PRE-VISIT SUMMARY for a healthcare provider.

PATIENT INFO:
- Name: {patient_info.get('name', 'Unknown')}
- Age: {patient_info.get('age', 'Unknown')}
- Gender: {patient_info.get('gender', 'Unknown')}

CONVERSATION WITH PATIENT:
{conversation_text}

MEDICAL RECORD DATA REFERENCED:
{medical_context if medical_context else "No medical records were referenced."}

ATTACHMENTS:
{attachment_text if attachment_text else "None"}

Create a structured pre-visit report in JSON format with these fields:
{{
    "chief_concerns": ["list of 1-3 main reasons for visit"],
    "patient_feelings": "Brief description of how the patient is feeling emotionally/physically",
    "symptom_summary": "Concise summary of any symptoms mentioned (2-3 sentences max)",
    "relevant_conditions": ["list of relevant existing conditions mentioned"],
    "current_medications": ["list of relevant medications mentioned"],
    "recent_vitals": {{"vital_name": "value"}},
    "key_quotes": ["1-2 important direct quotes from patient"]
}}

Keep it CONCISE - this should fit on one page. Focus on what the doctor needs to know.
Output ONLY the JSON, no other text.
<end_of_turn>
<start_of_turn>model
"""
    
    response = await call_llm(prompt, max_tokens=1000)
    
    # Parse the response
    try:
        # Clean up response
        response = response.strip()
        if response.startswith("```json"):
            response = response[7:]
        if response.startswith("```"):
            response = response[3:]
        if response.endswith("```"):
            response = response[:-3]
        
        report_data = json.loads(response.strip())
    except json.JSONDecodeError:
        # Fallback if JSON parsing fails
        report_data = {
            "chief_concerns": ["Unable to parse - please review conversation"],
            "patient_feelings": "See conversation history",
            "symptom_summary": "See conversation history",
            "relevant_conditions": [],
            "current_medications": [],
            "recent_vitals": {},
            "key_quotes": []
        }
    
    # Create report object
    report = PreVisitReport(
        patient_name=patient_info.get('name', 'Unknown'),
        patient_age=patient_info.get('age', 0),
        patient_gender=patient_info.get('gender', 'Unknown'),
        generated_at=datetime.now().strftime("%Y-%m-%d %H:%M"),
        chief_concerns=report_data.get('chief_concerns', []),
        patient_feelings=report_data.get('patient_feelings', ''),
        symptom_summary=report_data.get('symptom_summary', ''),
        relevant_conditions=report_data.get('relevant_conditions', []),
        current_medications=report_data.get('current_medications', []),
        recent_vitals=report_data.get('recent_vitals', {}),
        immunizations=immunizations or [],
        procedures=procedures or [],
        recent_encounters=encounters or [],
        allergies=allergies or [],
        attachments=attachments or [],
        key_quotes=report_data.get('key_quotes', [])
    )
    
    return report


def format_report_html(report: PreVisitReport) -> str:
    """Format report as HTML for display."""
    
    concerns_html = "".join([f"<li>{c}</li>" for c in report.chief_concerns]) or "<li>None specified</li>"
    conditions_html = "".join([f"<li>{c}</li>" for c in report.relevant_conditions]) or "<li>None noted</li>"
    medications_html = "".join([f"<li>{m}</li>" for m in report.current_medications]) or "<li>None noted</li>"
    
    vitals_html = ""
    for vital, value in report.recent_vitals.items():
        vitals_html += f"<div class='vital-item'><span class='vital-name'>{vital}:</span> <span class='vital-value'>{value}</span></div>"
    if not vitals_html:
        vitals_html = "<div class='vital-item'>No vitals referenced</div>"
    
    quotes_html = ""
    for quote in report.key_quotes:
        quotes_html += f'<blockquote>"{quote}"</blockquote>'
    
    attachments_html = ""
    for att in report.attachments:
        icon = "🎤" if att.get('type') == 'audio' else "📊" if att.get('type') == 'chart' else "📎"
        attachments_html += f"""
            <div class='attachment-item'>
                <span class='attachment-icon'>{icon}</span>
                <span class='attachment-title'>{att.get('title', 'Attachment')}</span>
                {f"<div class='attachment-summary'>{att.get('summary', '')}</div>" if att.get('summary') else ""}
            </div>
        """
    
    # NEW: Format immunizations
    immunizations_html = ""
    if report.immunizations:
        for imm in report.immunizations[:8]:  # Limit to 8 most recent
            vaccine_name = imm.get('vaccine_display', 'Unknown vaccine')
            # Truncate long vaccine names
            if len(vaccine_name) > 40:
                vaccine_name = vaccine_name[:37] + "..."
            date = imm.get('occurrence_date', '')[:10] if imm.get('occurrence_date') else 'Unknown'
            immunizations_html += f"<li>{vaccine_name} <span class='date-badge'>({date})</span></li>"
    else:
        immunizations_html = "<li>No immunization records</li>"
    
    # NEW: Format procedures (surgical history)
    procedures_html = ""
    if report.procedures:
        for proc in report.procedures[:6]:  # Limit to 6 most recent
            proc_name = proc.get('display', 'Unknown procedure')
            if len(proc_name) > 45:
                proc_name = proc_name[:42] + "..."
            date = proc.get('performed_date', '')[:10] if proc.get('performed_date') else 'Unknown'
            procedures_html += f"<li>{proc_name} <span class='date-badge'>({date})</span></li>"
    else:
        procedures_html = "<li>No surgical history</li>"
    
    # NEW: Format recent encounters
    encounters_html = ""
    if report.recent_encounters:
        for enc in report.recent_encounters[:5]:  # Limit to 5 most recent
            enc_type = enc.get('type_display', enc.get('class_display', 'Visit'))
            if len(enc_type) > 35:
                enc_type = enc_type[:32] + "..."
            reason = enc.get('reason_display', '')
            if reason and len(reason) > 30:
                reason = reason[:27] + "..."
            date = enc.get('period_start', '')[:10] if enc.get('period_start') else 'Unknown'
            reason_text = f" - {reason}" if reason else ""
            encounters_html += f"<li>{enc_type}{reason_text} <span class='date-badge'>({date})</span></li>"
    else:
        encounters_html = "<li>No recent encounters</li>"
    
    # NEW: Format allergies with severity
    allergies_html = ""
    if report.allergies:
        for allergy in report.allergies:
            substance = allergy.get('substance', 'Unknown')
            reaction = allergy.get('reaction', allergy.get('reaction_display', ''))
            criticality = allergy.get('criticality', '')
            severity_class = 'allergy-high' if criticality == 'high' else 'allergy-low'
            reaction_text = f" → {reaction}" if reaction else ""
            allergies_html += f"<li class='{severity_class}'><strong>{substance}</strong>{reaction_text}</li>"
    else:
        allergies_html = "<li>No known allergies</li>"
    
    html = f"""
    <div class="previsit-report">
        <div class="report-header">
            <h2>📋 Pre-Visit Summary</h2>
            <div class="report-meta">
                <div class="patient-info">
                    <strong>{report.patient_name}</strong> · {report.patient_age}y · {report.patient_gender}
                </div>
                <div class="report-date">Generated: {report.generated_at}</div>
            </div>
        </div>
        
        <div class="report-section">
            <h3>🎯 Chief Concerns</h3>
            <ul class="concerns-list">{concerns_html}</ul>
        </div>
        
        <div class="report-section">
            <h3>💭 How Patient Feels</h3>
            <p>{report.patient_feelings or "Not specified"}</p>
        </div>
        
        {f'''<div class="report-section">
            <h3>🩺 Symptoms</h3>
            <p>{report.symptom_summary}</p>
        </div>''' if report.symptom_summary else ""}
        
        <div class="report-columns">
            <div class="report-section half">
                <h3>📁 Active Conditions</h3>
                <ul>{conditions_html}</ul>
            </div>
            <div class="report-section half">
                <h3>⚠️ Allergies</h3>
                <ul class="allergies-list">{allergies_html}</ul>
            </div>
        </div>
        
        <div class="report-section">
            <h3>💊 Current Medications</h3>
            <ul>{medications_html}</ul>
        </div>
        
        {f'''<div class="report-section">
            <h3>📊 Recent Vitals</h3>
            <div class="vitals-grid">{vitals_html}</div>
        </div>''' if report.recent_vitals else ""}
        
        <div class="report-columns">
            <div class="report-section half">
                <h3>💉 Immunizations</h3>
                <ul class="compact-list">{immunizations_html}</ul>
            </div>
            <div class="report-section half">
                <h3>🏥 Surgical History</h3>
                <ul class="compact-list">{procedures_html}</ul>
            </div>
        </div>
        
        <div class="report-section">
            <h3>📅 Recent Encounters</h3>
            <ul class="compact-list">{encounters_html}</ul>
        </div>
        
        {f'''<div class="report-section">
            <h3>💬 Patient Quotes</h3>
            {quotes_html}
        </div>''' if report.key_quotes else ""}
        
        {f'''<div class="report-section">
            <h3>📎 Attachments</h3>
            <div class="attachments-list">{attachments_html}</div>
        </div>''' if report.attachments else ""}
        
        <div class="report-footer">
            <p><em>This summary was auto-generated from a patient conversation. Please verify all information.</em></p>
        </div>
    </div>
    """
    
    return html