src/prompt_builder.py

# src/prompt_builder.py
from __future__ import annotations
import re
from typing import Any, Dict, List, Optional

__all__ = [
    "normalize_intake",
    "build_referral_summary",
    "build_soap_prompt",
    "build_guideline_rationale_prompt",
]

# --------------------------- helpers ---------------------------

def _as_str(x: Any) -> str:
    if x is None:
        return ""
    if isinstance(x, str):
        return x.strip()
    if isinstance(x, (int, float, bool)):
        return str(x)
    if isinstance(x, (list, tuple, set)):
        vals = [str(v).strip() for v in x if str(v).strip()]
        return "; ".join(vals)
    if isinstance(x, dict):
        parts = [f"{k}: {v}" for k, v in x.items() if v]
        return "; ".join(parts)
    return str(x)

def normalize_intake(raw: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    raw = raw or {}
    p = raw.get("patient", {}) or {}
    c = raw.get("consult", {}) or {}
    patient = {
        "name": _as_str(p.get("name")),
        "age": _as_str(p.get("age")),
        "sex": _as_str(p.get("sex")),
        "history": _as_str(p.get("history")),
        "medications": _as_str(p.get("medications")),
        "labs": _as_str(p.get("labs")),
        "vitals": _as_str(p.get("vitals")),
        "imaging": _as_str(p.get("imaging")),
        "allergies": _as_str(p.get("allergies")),
    }
    consult = {
        "specialty": _as_str(c.get("specialty")),
        "question": _as_str(c.get("question")),
        "summary": _as_str(c.get("summary") or c.get("history")),
        "context": _as_str(c.get("context")),
        "referrer": _as_str(c.get("referrer")),
        "priority": _as_str(c.get("priority")),
    }
    return {"patient": patient, "consult": consult}

def _auto_tag_referral_text(text: str) -> str:
    if not text:
        return ""
    text = text.replace(";", ". ").replace("..", ".")
    text = re.sub(r"\s+", " ", text.strip())
    return text

# --------------------------- summary builder ---------------------------

def build_referral_summary(intake: Dict[str, Any], max_chars: int = 2200) -> str:
    """Compact but rich summary for SOAP prompting (raised limit to 2200)."""
    data = normalize_intake(intake)
    p, c = data["patient"], data["consult"]

    instruction = (
        "Use this summary to draft a concise narrative SOAP note.\n\n"
    )

    age = p.get("age")
    sex = (p.get("sex") or "").lower()
    if sex in ["m", "male"]:
        subj = f"{age or 'Adult'}-year-old man" if age else "Adult man"
    elif sex in ["f", "female"]:
        subj = f"{age or 'Adult'}-year-old woman" if age else "Adult woman"
    else:
        subj = f"{age or 'Adult'}-year-old patient" if age else "Adult patient"

    summary_text = _auto_tag_referral_text(c.get("summary", ""))

    meds = p.get("medications") or c.get("medications")
    labs = p.get("labs") or c.get("labs")
    question = c.get("question")
    specialty = c.get("specialty")

    parts: List[str] = [instruction, f"Patient: {subj}."]
    if summary_text:
        parts.append(summary_text)
    if meds:
        parts.append(f"\nMedications: {meds}")
    if labs:
        parts.append(f"\nLabs: {labs}")
    if question:
        parts.append(f"\nConsult question: {question}")
    if specialty:
        parts.append(f"\nReferral intended for {specialty}.")

    summary = "\n".join(parts).strip()
    if max_chars and len(summary) > max_chars:
        summary = summary[: max_chars - 3].rstrip() + "..."
    return summary

# --------------------------- SOAP prompt ---------------------------

def build_soap_prompt(summary: str) -> List[Dict[str, str]]:
    """
    System instruction per spec: concise narrative SOAP, 1–3 sentences per section.
    Plan may include up to 3 short actions/sentences. STRICT JSON with string values.
    """
    sys = (
        "You are a specialist clinician drafting a concise narrative SOAP note for an e-consult. "
        "Each section (Subjective, Objective, Assessment, Plan) should be written in natural-language prose, "
        "1–3 sentences per section, concise and clinically relevant. "
        "The Plan may include up to 3 short action items or sentences if needed. "
        "Return a STRICT JSON object with these keys: subjective, objective, assessment, plan. "
        "Each value must be a plain-text string (no lists, markdown, or bullets). "
        "Do not include commentary outside the JSON."
    )
    user = "Referral summary:\n" + summary.strip() + "\n\nReturn ONLY the JSON object."
    return [{"role": "system", "content": sys}, {"role": "user", "content": user}]

# ---------------------- Guideline Rationale prompt --------------------------

def build_guideline_rationale_prompt(
    assessment_text: str,
    plan_text: str,
    guideline_chunks: List[Dict[str, str]],
) -> List[Dict[str, str]]:
    """
    System instruction per spec: up to 3 short bullets (≤25 words each), STRICT JSON:
    {"rationale": ["..."]}. No markdown/numbering outside JSON.
    """
    sys = (
        "You are a clinical summarizer. Given the patient's assessment, plan, and the following guideline excerpts, "
        "list up to 3 concise bullet points (≤ 25 words each) explaining how the assessment and plan align with these guidelines. "
        "Each bullet should reference the relevant guideline name or section when possible. "
        "Return STRICT JSON: {\"rationale\": [\"...\"]}. Do not include markdown, numbering, or text outside the JSON object."
    )
    # Build compact excerpt list (max 3)
    lines = []
    for i, ch in enumerate(guideline_chunks[:3], start=1):
        text = (ch.get("text") or "").strip()
        src = (ch.get("source") or "Guideline").split("/")[-1]
        if text:
            lines.append(f"{i}. ({src}) {text[:600]}")
    user = (
        f"Assessment:\n{assessment_text.strip()}\n\n"
        f"Plan:\n{plan_text.strip()}\n\n"
        "Guideline excerpts:\n" + ("\n".join(lines) if lines else "(none)")
    )
    return [{"role": "system", "content": sys}, {"role": "user", "content": user}]