app.py

"""
app.py — AI Visit Plan Generator (Hugging Face Spaces Version)

This is a single-file application that:
1. Accepts free-text care referral instructions
2. Sends them to an AI model on Hugging Face
3. Extracts a structured Visit Plan (JSON)
4. Displays the result in a web interface

Deployed as a Hugging Face Space with Gradio.
"""

import os
import json
import re
import gradio as gr
from huggingface_hub import InferenceClient
from pydantic import BaseModel, Field
from typing import Optional
from enum import Enum


# ══════════════════════════════════════════════════════════════════
# SECTION 1: CONFIGURATION
# ══════════════════════════════════════════════════════════════════

HUGGINGFACE_API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN", "")

if not HUGGINGFACE_API_TOKEN:
    print("WARNING: HUGGINGFACE_API_TOKEN not found in environment.")
    print("Go to your Space's Settings > Repository secrets to add it.")

DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct"


# ══════════════════════════════════════════════════════════════════
# SECTION 2: DATA MODELS (What a Visit Plan looks like)
# ══════════════════════════════════════════════════════════════════

class Confidence(str, Enum):
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class VisitOccurrence(BaseModel):
    """A single unique visit type within the plan."""
    visit_name: str = Field(..., description="Name of the visit, e.g. 'Morning Call'")
    days: list[str] = Field(
        ..., description="Days this visit occurs, e.g. ['Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','Sunday']"
    )
    start_time: str = Field(..., description="Preferred start time, e.g. '07:30'")
    end_time: str = Field(..., description="Preferred end time, e.g. '08:15'")
    duration_minutes: int = Field(..., description="Duration in minutes")
    care_tasks: list[str] = Field(
        default_factory=list,
        description="Specific tasks for this visit"
    )
    notes: Optional[str] = Field(None, description="Additional notes")
    confidence: Confidence = Field(
        Confidence.MEDIUM,
        description="Confidence level for this extraction"
    )


class VisitPlan(BaseModel):
    """Top-level visit plan generated from referral text."""
    service_user_name: Optional[str] = Field(None, description="Name of the service user")
    visit_frequency: str = Field(..., description="e.g. 'Four times daily, 7 days a week'")
    total_hours_per_week: Optional[str] = Field(None, description="e.g. '14 hours'")
    preferred_times: list[str] = Field(
        default_factory=list, description="e.g. ['07:30', '12:30', '17:00', '21:00']"
    )
    qualifications_required: list[str] = Field(
        default_factory=list, description="e.g. ['insulin-trained', 'moving-and-handling']"
    )
    gender_preference: Optional[str] = Field(
        None, description="e.g. 'Female carer preferred'"
    )
    environmental_requirements: list[str] = Field(
        default_factory=list, description="e.g. ['key-safe fitted', 'grab rails in bathroom']"
    )
    medications: list[str] = Field(
        default_factory=list, description="Medications the carer administers or prompts"
    )
    occurrences: list[VisitOccurrence] = Field(
        default_factory=list, description="Unique visit types (NOT one per day — group by visit type)"
    )
    summary: str = Field(..., description="Brief plain-English summary")
    extraction_notes: list[str] = Field(
        default_factory=list, description="Flags for ambiguous or missing info"
    )


# ══════════════════════════════════════════════════════════════════
# SECTION 3: THE PROMPT (Instructions for the AI)
# ══════════════════════════════════════════════════════════════════

# Build a compact schema description instead of the full JSON Schema
# (saves hundreds of tokens in the prompt)
COMPACT_SCHEMA = """{
  "service_user_name": "string or null",
  "visit_frequency": "string (e.g. 'Four times daily, 7 days a week')",
  "total_hours_per_week": "string or null",
  "preferred_times": ["07:30", "12:30"],
  "qualifications_required": ["insulin-trained"],
  "gender_preference": "string or null",
  "environmental_requirements": ["key-safe fitted", "grab rails in bathroom"],
  "medications": ["Insulin Glargine 20 units morning & evening", "Metformin 500mg BD"],
  "occurrences": [
    {
      "visit_name": "Morning Call",
      "days": ["Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday"],
      "start_time": "07:30",
      "end_time": "08:15",
      "duration_minutes": 45,
      "care_tasks": ["personal care", "insulin administration", "breakfast prep"],
      "notes": "any extra notes",
      "confidence": "high"
    }
  ],
  "summary": "Brief plain-English summary of the care plan",
  "extraction_notes": ["any assumptions or missing info flags"]
}"""

SYSTEM_PROMPT = f"""You are a care coordination assistant. Read the referral or care plan text and extract a structured visit plan as JSON.

CRITICAL RULES:
1. Return ONLY valid JSON. No markdown, no code fences, no explanation before or after.
2. Group visits by TYPE, not by day. If "Morning Call" happens 7 days a week, that is ONE occurrence with days: ["Monday","Tuesday",...,"Sunday"]. Do NOT create separate entries for each day.
3. Map clinical shorthand: "BD"="twice daily", "OD"="once daily", "TDS"="three times daily", "PRN"="as needed".
4. Set confidence: "high" if explicitly stated, "medium" if inferred, "low" if guessed.
5. Never invent care tasks not mentioned in the referral.
6. If gender preference is not mentioned, set to null.
7. Times in 24-hour format.
8. Keep the JSON concise. Do not add unnecessary whitespace or long descriptions.
9. Ensure all commas and brackets are correct before returning.

OUTPUT SCHEMA:
{COMPACT_SCHEMA}"""

FEW_SHOT_INPUT = """Mrs Joan Smith, 82. Lives alone, ground floor flat. Key-safe code 4521.
Requires personal care morning and evening, 7 days. Medication administration BD
(insulin). Female carer preferred. Moving and handling required - uses zimmer frame."""

FEW_SHOT_OUTPUT = json.dumps({
    "service_user_name": "Joan Smith",
    "visit_frequency": "Twice daily, 7 days a week",
    "total_hours_per_week": "14 hours",
    "preferred_times": ["08:00", "20:00"],
    "qualifications_required": ["insulin-trained", "moving-and-handling"],
    "gender_preference": "Female carer preferred",
    "environmental_requirements": ["Key-safe code: 4521", "Ground floor flat", "Uses zimmer frame"],
    "medications": ["Insulin (BD)"],
    "occurrences": [
        {
            "visit_name": "Morning Call",
            "days": ["Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday"],
            "start_time": "08:00",
            "end_time": "09:00",
            "duration_minutes": 60,
            "care_tasks": ["personal care", "insulin administration"],
            "notes": None,
            "confidence": "high"
        },
        {
            "visit_name": "Evening Call",
            "days": ["Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday"],
            "start_time": "20:00",
            "end_time": "21:00",
            "duration_minutes": 60,
            "care_tasks": ["personal care", "insulin administration"],
            "notes": None,
            "confidence": "high"
        }
    ],
    "summary": "Joan Smith, 82, requires twice-daily personal care and insulin. Female carer preferred. Key-safe access, ground floor, zimmer frame.",
    "extraction_notes": ["Visit times assumed as 08:00 and 20:00 - not explicitly stated"]
}, separators=(',', ':'))  # Compact JSON to save tokens in the example


# ══════════════════════════════════════════════════════════════════
# SECTION 4: LLM CLIENT (Talks to Hugging Face's AI)
# ══════════════════════════════════════════════════════════════════

def call_huggingface_model(referral_text: str, model_id: str = DEFAULT_MODEL) -> str:
    """
    Send the referral text to a Hugging Face model and get the response.
    Uses InferenceClient which handles routing internally.
    """
    client = InferenceClient(token=HUGGINGFACE_API_TOKEN)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Example referral:\n{FEW_SHOT_INPUT}"},
        {"role": "assistant", "content": FEW_SHOT_OUTPUT},
        {"role": "user", "content": f"Extract a visit plan from this referral. Return ONLY valid JSON, no other text:\n\n{referral_text}"},
    ]

    response = client.chat_completion(
        model=model_id,
        messages=messages,
        max_tokens=4096,
        temperature=0.1,
    )

    return response.choices[0].message.content


# ══════════════════════════════════════════════════════════════════
# SECTION 5: JSON REPAIR & EXTRACTION
# ══════════════════════════════════════════════════════════════════

def repair_json(text: str) -> str:
    """Aggressively repair common JSON issues from LLM output."""
    # Remove any trailing commas before } or ]
    text = re.sub(r',\s*([}\]])', r'\1', text)

    # Add missing commas between adjacent string values on separate lines
    # e.g.  "value"\n  "key"  →  "value",\n  "key"
    text = re.sub(r'(\")\s*\n(\s*\")', r'\1,\n\2', text)

    # Add missing commas after } or ] followed by { or [ or "
    text = re.sub(r'(\})\s*\n(\s*\{)', r'\1,\n\2', text)
    text = re.sub(r'(\])\s*\n(\s*\[)', r'\1,\n\2', text)
    text = re.sub(r'(\})\s*\n(\s*\")', r'\1,\n\2', text)
    text = re.sub(r'(\])\s*\n(\s*\")', r'\1,\n\2', text)

    # Fix missing commas after numbers, true, false, null followed by a new key
    text = re.sub(r'(\d)\s*\n(\s*\")', r'\1,\n\2', text)
    text = re.sub(r'(true|false|null)\s*\n(\s*\")', r'\1,\n\2', text)

    # Add missing commas between adjacent objects in arrays: }\n{
    text = re.sub(r'(\})\s*(\{)', r'\1,\2', text)

    return text


def close_truncated_json(text: str) -> str:
    """
    If the JSON was truncated (e.g. hit token limit), try to close
    all open brackets/braces to make it parseable.
    """
    # Count unclosed brackets
    open_braces = text.count('{') - text.count('}')
    open_brackets = text.count('[') - text.count(']')

    if open_braces <= 0 and open_brackets <= 0:
        return text  # Already balanced

    # Strip any trailing partial content (e.g. a half-written string)
    # Find the last complete value
    # Remove trailing partial string if we're mid-string
    in_string = False
    escaped = False
    last_good_pos = 0
    for i, ch in enumerate(text):
        if escaped:
            escaped = False
            continue
        if ch == '\\':
            escaped = True
            continue
        if ch == '"':
            in_string = not in_string
        if not in_string and ch in ',:{}[]':
            last_good_pos = i

    if in_string:
        # We're mid-string — truncate to last good position and close the string
        text = text[:last_good_pos + 1]

    # Remove any trailing comma
    text = text.rstrip().rstrip(',')

    # Close open brackets and braces
    # We need to close them in reverse order of opening
    closing = ''
    stack = []
    in_str = False
    esc = False
    for ch in text:
        if esc:
            esc = False
            continue
        if ch == '\\':
            esc = True
            continue
        if ch == '"':
            in_str = not in_str
        if not in_str:
            if ch == '{':
                stack.append('}')
            elif ch == '[':
                stack.append(']')
            elif ch == '}':
                if stack and stack[-1] == '}':
                    stack.pop()
            elif ch == ']':
                if stack and stack[-1] == ']':
                    stack.pop()

    # Close everything that's still open
    closing = ''.join(reversed(stack))
    return text + closing


def extract_json_from_text(text: str) -> str | None:
    """Find JSON in the AI's response (handles markdown code fences)."""
    # Try markdown code fence first
    match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
    if match:
        return match.group(1)
    # Try to find the outermost JSON object
    # Find the first { and try to match to the last }
    first_brace = text.find('{')
    if first_brace == -1:
        return None
    last_brace = text.rfind('}')
    if last_brace == -1:
        # JSON might be truncated — take everything from first brace
        return text[first_brace:]
    return text[first_brace:last_brace + 1]


def validate_and_warn(visit_plan: dict) -> list[str]:
    """Check the visit plan for potential issues."""
    warnings = []

    for i, occ in enumerate(visit_plan.get("occurrences", [])):
        if occ.get("confidence") == "low":
            name = occ.get("visit_name", f"Visit {i+1}")
            warnings.append(f"{name} has LOW confidence - please verify.")

    if not visit_plan.get("qualifications_required"):
        warnings.append("No qualifications extracted - check the referral.")

    if not visit_plan.get("occurrences"):
        warnings.append("No visit times generated - the AI may have struggled to parse the schedule.")

    if not visit_plan.get("environmental_requirements"):
        warnings.append("No environmental requirements found - check for key-safe codes, access info, etc.")

    return warnings


def process_referral(referral_text: str, model_id: str) -> dict:
    """Main pipeline: referral text -> validated VisitPlan."""
    # Call the AI
    try:
        raw_response = call_huggingface_model(referral_text, model_id)
    except Exception as e:
        return {
            "success": False,
            "visit_plan": None,
            "raw_response": "",
            "errors": [f"Error calling AI model: {str(e)}"],
        }

    # Extract JSON from the response
    json_str = extract_json_from_text(raw_response)
    if not json_str:
        return {
            "success": False,
            "visit_plan": None,
            "raw_response": raw_response,
            "errors": ["AI did not return valid JSON. Try again or use a different model."],
        }

    # Try to parse the JSON (with repair attempts)
    data = None
    parse_errors = []

    # Attempt 1: Parse as-is
    try:
        data = json.loads(json_str)
    except json.JSONDecodeError as e:
        parse_errors.append(f"Direct parse: {str(e)}")

    # Attempt 2: Repair common issues
    if data is None:
        try:
            repaired = repair_json(json_str)
            data = json.loads(repaired)
        except json.JSONDecodeError as e:
            parse_errors.append(f"After repair: {str(e)}")

    # Attempt 3: Close truncated JSON
    if data is None:
        try:
            closed = close_truncated_json(json_str)
            data = json.loads(closed)
        except json.JSONDecodeError as e:
            parse_errors.append(f"After closing truncation: {str(e)}")

    # Attempt 4: Repair + close
    if data is None:
        try:
            repaired = repair_json(json_str)
            closed = close_truncated_json(repaired)
            data = json.loads(closed)
        except json.JSONDecodeError as e:
            parse_errors.append(f"After repair+close: {str(e)}")

    if data is None:
        return {
            "success": False,
            "visit_plan": None,
            "raw_response": raw_response,
            "errors": [
                "Could not parse JSON from AI response after multiple repair attempts.",
                "Parse attempts: " + " | ".join(parse_errors),
                "Try clicking Generate again, or try a different model.",
            ],
        }

    # Validate against schema
    try:
        plan = VisitPlan(**data)
        return {
            "success": True,
            "visit_plan": plan.model_dump(),
            "raw_response": raw_response,
            "errors": [],
        }
    except Exception as e:
        # Even if schema validation fails, return what we have
        # (it's useful for the user to see the partial extraction)
        return {
            "success": True,
            "visit_plan": data,
            "raw_response": raw_response,
            "errors": [],
        }


# ══════════════════════════════════════════════════════════════════
# SECTION 6: SAMPLE REFERRALS (for testing)
# ══════════════════════════════════════════════════════════════════

SAMPLE_REFERRAL_1 = """
Mr Arthur Davies, 78 years old. Referred by Dr Patel, Elm Surgery.
Diagnosis: Type 2 diabetes, early-stage dementia, limited mobility (uses walking stick).
Lives with wife (Margaret, 81) in a two-storey house - bedroom is upstairs.

Care required:
- Morning visit (approx 8am): assist with washing, dressing, breakfast prep, and
  morning medication (metformin 500mg, donepezil 10mg).
- Lunchtime check (approx 12:30): ensure lunch is prepared, prompt medication if
  needed, welfare check.
- Evening visit (approx 6pm): evening meal support, help preparing for bed,
  evening medication (metformin 500mg).

Frequency: 7 days a week. Carer must be medication-trained.
No gender preference expressed. Wife is sometimes out at day centre (Tues/Thurs)
so key-safe access needed - code 7890. Please note: Mr Davies can become confused
in the late afternoon (sundowning).
"""

SAMPLE_REFERRAL_2 = """
DISCHARGE SUMMARY - NHS Trust Reference: DIS-2024-8891
Patient: Mr James Thompson, DOB 15/03/1945 (79 years)
Diagnosis: Right hip replacement (total), Type 1 diabetes (insulin-dependent),
mild cognitive impairment.

Discharge to home. Lives with wife (Betty, 77) in semi-detached house.
Stairs present - bedroom moved to ground floor temporarily.

Post-operative care plan:
AM visit 0730-0830: Personal care, wound dressing check, insulin
administration (NovoRapid 8 units pre-breakfast), compression stockings.
Must be moving and handling trained - patient cannot weight-bear fully.
PM visit 1900-2000: Evening personal care, insulin administration
(Lantus 20 units), wound check, help into bed.

Duration: 6 weeks initially, review at 4 weeks.
Frequency: BD, 7/7. Female carer preferred (patient request).
Key-safe code: 5566. OT assessment pending for stair lift.
District nurse visiting alternate days for wound management.
"""

SAMPLE_REFERRAL_3 = """
Mrs Margaret Williams, 75. Lives alone in a bungalow. Key-safe: 3344.
Needs help with personal care every morning, 7 days a week.
Medication: paracetamol 500mg BD. No gender preference.
Mobile with walking frame. Contact: son David 07700 111222.
"""


# ══════════════════════════════════════════════════════════════════
# SECTION 7: GRADIO USER INTERFACE
# ══════════════════════════════════════════════════════════════════

def generate_visit_plan(referral_text: str, model_choice: str):
    """Called when the user clicks 'Generate Visit Plan'."""
    if not referral_text.strip():
        return "Please paste referral text in the box on the left.", "", ""

    if not HUGGINGFACE_API_TOKEN:
        return (
            "ERROR: No API token found.\n\n"
            "Go to your Space's Settings > Repository secrets\n"
            "and add a secret called HUGGINGFACE_API_TOKEN\n"
            "with your Hugging Face token (starts with hf_...).",
            "",
            ""
        )

    result = process_referral(referral_text, model_choice)

    if result["success"]:
        plan_json = json.dumps(result["visit_plan"], indent=2)
        warnings = validate_and_warn(result["visit_plan"])
        all_notes = []
        if isinstance(result["visit_plan"], dict):
            all_notes = result["visit_plan"].get("extraction_notes", [])
        all_notes = all_notes + warnings
        if all_notes:
            notes_text = "\n".join(f">> {note}" for note in all_notes)
        else:
            notes_text = "No issues detected. Please still review against the original referral."
        return plan_json, notes_text, result["raw_response"]
    else:
        error_text = "ERRORS:\n" + "\n".join(result["errors"])
        return error_text, "", result.get("raw_response", "")


# Build the Gradio interface
with gr.Blocks(
    title="AI Visit Plan Generator",
) as demo:

    gr.Markdown("# AI Visit Plan Generator")
    gr.Markdown(
        "Paste a free-text care referral below. The AI will extract a structured "
        "visit plan for review. **Always verify the output against the original referral.**"
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            model_dropdown = gr.Dropdown(
                choices=[
                    "Qwen/Qwen2.5-72B-Instruct",
                    "meta-llama/Meta-Llama-3.1-8B-Instruct",
                    "mistralai/Mistral-Nemo-Instruct-2407",
                    "microsoft/Phi-3.5-mini-instruct",
                ],
                value="Qwen/Qwen2.5-72B-Instruct",
                label="AI Model",
                info="Try different models to compare results"
            )
            referral_input = gr.Textbox(
                label="Referral Text",
                placeholder="Paste the GP referral, hospital discharge summary, "
                           "or social worker assessment here...",
                lines=15,
            )
            with gr.Row():
                generate_btn = gr.Button("Generate Visit Plan", variant="primary")
            with gr.Row():
                example1_btn = gr.Button("Example 1: GP Referral", size="sm")
                example2_btn = gr.Button("Example 2: Hospital Discharge", size="sm")
                example3_btn = gr.Button("Example 3: Simple", size="sm")

        with gr.Column(scale=1):
            gr.Markdown("### Output")
            plan_output = gr.Textbox(
                label="Generated Visit Plan (JSON)",
                lines=20,
                interactive=False,
            )
            notes_output = gr.Textbox(
                label="Notes & Warnings",
                lines=5,
                interactive=False,
            )

    with gr.Accordion("Debug: Raw Model Response", open=False):
        gr.Markdown(
            "*Shows the raw text from the AI model. Useful for troubleshooting.*"
        )
        raw_output = gr.Textbox(label="Raw Response", lines=10, interactive=False)

    # Connect buttons
    generate_btn.click(
        fn=generate_visit_plan,
        inputs=[referral_input, model_dropdown],
        outputs=[plan_output, notes_output, raw_output],
    )
    example1_btn.click(fn=lambda: SAMPLE_REFERRAL_1, outputs=referral_input)
    example2_btn.click(fn=lambda: SAMPLE_REFERRAL_2, outputs=referral_input)
    example3_btn.click(fn=lambda: SAMPLE_REFERRAL_3, outputs=referral_input)


# Launch
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)