alz_companion/prompts.py

"""
Prompts for the Alzheimer’s AI Companion.
"""

# ------------------------ Behaviour‑level tags ------------------------
BEHAVIOUR_TAGS = {
    # Tags from "The Father"
    "repetitive_questioning": ["validation", "gentle_redirection", "offer_distraction"],
    "confusion": ["reassurance", "time_place_orientation", "photo_anchors"],
    "wandering": ["walk_along_support", "simple_landmarks", "visual_cues", "safe_wandering_space"],
    "agitation": ["de-escalating_tone", "validate_feelings", "reduce_stimulation", "simple_choices"],
    "false_accusations": ["reassure_no_blame", "avoid_arguing", "redirect_activity"],
    "address_memory_loss": ["encourage_ID_bracelet_or_GPS", "place_contact_info_in_wallet", "inform_trusted_neighbors", "avoid_quizzing_on_address"],
    "hallucinations_delusions": ["avoid_arguing_or_correcting", "validate_the_underlying_emotion", "offer_reassurance_of_safety", "gently_redirect_to_real_activity", "check_for_physical_triggers"],

    # Tags from "Still Alice" (and others for future use)
    "exit_seeking": ["validation", "calm_presence", "safe_wandering_space", "environmental_cues"],
    "aphasia": ["patience", "simple_language", "nonverbal_cues", "validation"],
    "withdrawal": ["gentle_invitation", "calm_presence", "offer_familiar_comforts", "no_pressure"],
    "affection": ["reciprocate_warmth", "positive_reinforcement", "simple_shared_activity"],
    "sleep_disturbance": ["establish_calm_bedtime_routine", "limit_daytime_naps", "check_for_discomfort_or_pain"],
    "anxiety": ["calm_reassurance", "simple_breathing_exercise", "reduce_environmental_stimuli"],
    "depression_sadness": ["validate_feelings_of_sadness", "encourage_simple_pleasant_activity", "ensure_social_connection"],
    "orientation_check": ["gentle_orientation_cues", "use_familiar_landmarks", "avoid_quizzing"],

    # Tags from "Away from Her"
    "misidentification": ["gently_correct_with_context", "use_photos_as_anchors", "respond_to_underlying_emotion", "avoid_insistent_correction"],

    # Other useful tags
    "sundowning_restlessness": ["predictable_routine", "soft_lighting", "low_stimulation", "familiar_music"],
    "object_misplacement": ["nonconfrontational_search", "fixed_storage_spots"],

    # --- New Tags from Test Fixtures ---
    "validation": [],
    "gentle_reorientation": [],
    "de-escalation": [],
    "distraction": [],
    "spaced_cueing": [],
    "reassurance": [],
    "psychoeducation": [],
    "goal_breakdown": [],
    "routine_structuring": [],
    "reminiscence_prompting": [],
    "reframing": [],
    "distress_tolerance": [],
    "caregiver_communication_template": [],
    "personalised_music_activation": [],
    "memory_probe": [],
    "safety_brief": [],
    "follow_up_prompt": []
}

# ------------------------ Emotion styles & helpers ------------------------
EMOTION_STYLES = {
    "confusion": {"tone": "calm, orienting, concrete", "playbook": ["Offer a simple time/place orientation cue (who/where/when).", "Reference one familiar anchor (photo/object/person).", "Use short sentences and one step at a time."]},
    "fear": {"tone": "reassuring, safety-forward, gentle", "playbook": ["Acknowledge fear without contradiction.", "Provide a clear safety cue (e.g., 'You’re safe here with me').", "Reduce novelty and stimulation; suggest one safe action."]},
    "anger": {"tone": "de-escalating, validating, low-arousal", "playbook": ["Validate the feeling; avoid arguing/correcting.", "Keep voice low and sentences short.", "Offer a simple choice to restore control (e.g., 'tea or water?')."]},
    "sadness": {"tone": "warm, empathetic, gentle reminiscence", "playbook": ["Acknowledge loss/longing.", "Invite one comforting memory or familiar song.", "Keep pace slow; avoid tasking."]},
    "warmth": {"tone": "affirming, appreciative", "playbook": ["Reflect gratitude and positive connection.", "Reinforce what’s going well.", "Keep it light; don’t overload with new info."]},
    "joy": {"tone": "supportive, celebratory (but not overstimulating)", "playbook": ["Share the joy briefly; match energy gently.", "Offer a simple, pleasant follow-up activity.", "Avoid adding complex tasks."]},
    "calm": {"tone": "matter-of-fact, concise, steady", "playbook": ["Keep instructions simple.", "Maintain steady pace.", "No extra soothing needed."]},
}

def render_emotion_guidelines(emotion: str | None) -> str:
    e = (emotion or "").strip().lower()
    if e not in EMOTION_STYLES:
        return "Emotion: (auto)\nDesired tone: calm, clear.\nWhen replying, reassure if distress is apparent; prioritise validation and simple choices."
    style = EMOTION_STYLES[e]
    bullet = "\n".join([f"- {x}" for x in style["playbook"]])
    return f"Emotion: {e}\nDesired tone: {style['tone']}\nWhen replying, follow:\n{bullet}"

# ------------------------ NLU Classification (NEW Decomposed Pipeline) ------------------------

# --- STEP 1: Router for Primary Goal ---
NLU_ROUTER_PROMPT = """You are an expert NLU router. Your task is to classify the user's primary goal into one of two categories:
1.  `practical_planning`: The user is seeking a plan, strategy, "how-to" advice, or a solution to a problem.
2.  `emotional_support`: The user is expressing feelings, seeking comfort, validation, or reassurance.

User Query: "{query}"

Respond with ONLY a single category name from the list above.
Category: """

# --- STEP 2: Goal-Specific Few-Shot Examples ---
EMOTIONAL_SUPPORT_EXAMPLES = """
User Query: "She looked right through me today, like I wasn't even there."
<thinking>
1.  **Literal Meaning:** The user's loved one did not seem to recognize them.
2.  **Underlying Situation:** The user is expressing emotional pain and a feeling of invisibility due to the disease's progression. This is a clear expression of hurt.
3.  **User's Primary Goal:** This is `Emotional Support`. They are seeking comfort and a way to cope with this painful experience, not a technical fix.
4.  **Tag Selection:** The goal is emotional support. I should prioritize tags that validate and comfort. `validation` and `reminiscence_prompting` are good behaviors. The emotion is clearly `sadness`. The topic is about providing `treatment_option:reassurance`. The context seems to be a close relationship, maybe `relationship_spouse`.
</thinking>
{
    "detected_behaviors": ["validation", "reminiscence_prompting", "misidentification"],
    "detected_emotion": "sadness",
    "detected_topic": "treatment_option:reassurance",
    "detected_contexts": ["relationship_spouse"]
}

User Query: "I just feel so lost and alone in this."
<thinking>
1.  **Literal Meaning:** The user is expressing feelings of isolation and confusion.
2.  **Underlying Situation:** This is a classic expression of caregiver burnout and emotional distress.
3.  **User's Primary Goal:** This is `Emotional Support`. The user needs validation and to feel heard.
4.  **Tag Selection:** The core need is validation. `validate_feelings_of_sadness` is a relevant behavior from the `depression_sadness` group. The emotion is `sadness`. This is a general `caregiving_advice` topic.
</thinking>
{
    "detected_behaviors": ["validate_feelings_of_sadness", "reassurance"],
    "detected_emotion": "sadness",
    "detected_topic": "caregiving_advice",
    "detected_contexts": ["relationship_staff_or_caregiver"]
}
"""

PRACTICAL_PLANNING_EXAMPLES = """
User Query: "He’s withdrawn today. What should I do?"
<thinking>
1.  **Literal Meaning:** The patient is not engaging, and the user is asking for a plan.
2.  **Underlying Situation:** The user is observing the patient's apathy and wants a concrete action to help.
3.  **User's Primary Goal:** This is `Practical Planning`. The user is explicitly asking "What should I do?".
4.  **Tag Selection:** The problem is `withdrawal`. The user wants a proactive solution. A known strategy is music, so the `behavior` tag should include the proactive `personalised_music_activation`. The `topic` is therefore `treatment_option:music_therapy`. The user's tone seems matter-of-fact, so `calm` is a reasonable emotion.
</thinking>
{
    "detected_behaviors": ["withdrawal", "personalised_music_activation"],
    "detected_emotion": "calm",
    "detected_topic": "treatment_option:music_therapy",
    "detected_contexts": ["setting_home_or_community"]
}

User Query: "How can I make the house safer for him? He keeps wandering at night."
<thinking>
1.  **Literal Meaning:** The user wants advice on home safety, specifically related to wandering.
2.  **Underlying Situation:** The user is dealing with a specific, potentially dangerous behavior (wandering) and needs a plan to mitigate risk.
3.  **User's Primary Goal:** This is clearly `Practical Planning`. They need actionable advice.
4.  **Tag Selection:** The primary behavior is `wandering`. Safety is paramount, so tags like `encourage_ID_bracelet_or_GPS` and `safe_wandering_space` are key. The topic is directly `treatment_option:home_safety`. The emotion is likely `anxiety` or `fear` on the user's part.
</thinking>
{
    "detected_behaviors": ["wandering", "address_memory_loss", "safe_wandering_space"],
    "detected_emotion": "anxiety",
    "detected_topic": "treatment_option:home_safety",
    "detected_contexts": ["setting_home_or_community"]
}
"""

# --- STEP 3: Specialist Classifier ---
SPECIALIST_CLASSIFIER_PROMPT = """You are an expert NLU engine. Your task is to analyze the user's query to deeply understand their underlying intent and classify it correctly.

--- INSTRUCTIONS ---
First, in a <thinking> block, you must reason step-by-step about the user's query by following these points:
- **Literal Meaning:** What is the user literally asking or stating?
- **Underlying Situation:** What is the deeper emotional state or situation being described? (e.g., caregiver burnout, patient confusion, a request for practical help).
- **User's Primary Goal:** You have been told the user's goal is `{primary_goal}`. Briefly confirm if the query aligns with this goal.
- **Tag Selection:** Based on the primary goal, explain which tags from the provided lists are the most appropriate and why. If the goal is practical, prioritize practical tags. If the goal is emotional, prioritize emotional support tags.

Second, after your reasoning, provide a single, valid JSON object with the final classification.

--- PROVIDED TAGS ---
Behaviors: {behavior_options}
Emotions: {emotion_options}
Topics: {topic_options}
Contexts: {context_options}

--- EXAMPLES ---
{examples}
---

User Query: "{query}"

<thinking>
"""


# ------------------------ Guardrails ------------------------
SAFETY_GUARDRAILS = """Never provide medical diagnoses or dosing. If a situation implies imminent risk (e.g., wandering/elopement, severe agitation, choking, falls), signpost immediate support from onsite staff or emergency services. Use respectful, person‑centred language. Keep guidance concrete and stepwise."""

# ------------------------ System & Answer Templates ------------------------
SYSTEM_TEMPLATE = """You are an Alzheimer’s caregiving companion. Address the patient as {patient_name} and the caregiver as {caregiver_name}. Ground every suggestion in retrieved evidence when possible. If unsure, say so plainly.
{guardrails}
--- IMPORTANT RULE ---
You MUST write your entire response in {language} ONLY. This is a strict instruction. Do not use any other language, even if the user or the retrieved context uses a different language. Your final output must be in {language}."""

ANSWER_TEMPLATE_CALM = """Context:
{context}

---
Question from user: {question}

---
Instructions:
Based on the context, write a gentle and supportive response in a single, natural-sounding paragraph.
Your response should:
1. Start by briefly and calmly acknowledging the user's situation or feeling.
2. Weave 2-3 practical, compassionate suggestions from the context into your paragraph. Do not use a numbered or bulleted list.
3. Conclude with a short, reassuring phrase.
4. You MUST use the retrieved context to directly address the user's specific **Question**.
Your response in {language}:"""

# For scenarios tagged with a specific behavior (e.g., agitation, confusion)
ANSWER_TEMPLATE_ADQ = """--- General Guidance from Knowledge Base ---
{general_context}

--- Relevant Personal Memories ---
{personal_context}

---
Care scenario: {scenario_tag}
Response Guidelines:
{emotions_context}
Question from user: {question}

---
Instructions:
Based on ALL the information above, write a **concise, warm, and validating** response for the {role} in a single, natural-sounding paragraph. **Keep the total response to 2-4 sentences.**
If possible, weave details from the 'Relevant Personal Memories' into your suggestions to make the response feel more personal and familiar.
Pay close attention to the Response Guidelines to tailor your tone.
Your response should follow this pattern:
1. Start by validating the user's feeling or concern with a unique, empathetic opening. DO NOT USE THE SAME OPENING PHRASE REPEATEDLY. Choose from different styles of openers, such as:
   - Acknowledging the difficulty: "That sounds like a very challenging situation..."
   - Expressing understanding: "I can see why that would be worrying..."
   - Stating a shared goal: "Let's walk through how we can handle that..."
   - Directly validating the feeling: "It's completely understandable to feel frustrated when..."
2. Gently offer **1-2 of the most important practical steps**, combining general guidance with personal memories where appropriate. Do not use a list.
3. If the scenario involves risk (like exit_seeking), subtly include a safety cue.
4. End with a compassionate, de-escalation phrase.
Your response in {language}:"""

RISK_FOOTER = """If safety is a concern right now, please seek immediate assistance from onsite staff or local emergency services."""

# ------------------------ Router & Specialized Templates ------------------------

# --- NEW: Template for expanding user queries for better retrieval ---
QUERY_EXPANSION_PROMPT = """You are a helpful AI assistant. Your task is to rephrase a user's question into 3 different, semantically similar questions to improve document retrieval.
Provide the rephrased questions as a JSON list of strings.

User Question: "{question}"

JSON List:
"""

# Template for routing/classifying the user's intent
ROUTER_PROMPT = """You are an expert NLU router. Your task is to classify the user's query into one of four categories:
1.  `caregiving_scenario`: The user is describing a situation, asking for advice, or expressing a concern related to Alzheimer's or caregiving.
2.  `factual_question`: The user is asking a direct question about a personal memory, person, or event that would be stored in the memory journal.
3.  `general_knowledge_question`: The user is asking a general knowledge question about the world, facts, or topics not related to personal memories or caregiving (e.g., 'What is the capital of France?', 'Who directed the movie Inception?').
4.  `general_conversation`: The user is making a general conversational remark, like a greeting, a thank you, or a simple statement that does not require a knowledge base lookup.

User Query: "{query}"

Respond with ONLY a single category name from the list above.
Category: """

# Template for answering direct factual questions
ANSWER_TEMPLATE_FACTUAL = """Context:
{context}

---
Question from user: {question}

---
Instructions:
Based on the provided context, directly and concisely answer the user's question.
- If the context contains the answer, state it clearly and naturally. Keep your response to a maximum of 3 sentences.
- If the context does not contain the answer, respond in a warm and friendly tone that you couldn't find a memory of that topic and gently ask if the user would like to talk more about it or add it as a new memory.
- Do not offer advice or suggestions unless they are part of the retrieved context.
- ABSOLUTELY DO NOT invent, create, or hallucinate any stories, characters, or details. Your knowledge is limited to the provided context ONLY.

Your response MUST be in {language}:"""


# --- NEW: Template for answering general knowledge questions ---
# Template for answering general knowledge questions
ANSWER_TEMPLATE_GENERAL_KNOWLEDGE = """You are a factual answering engine.
Your task is to directly answer the user's general knowledge question based on your training data.

Instructions:
- Be factual and concise. Go straight to the answer.
- If the answer requires a list of examples, provide a maximum of 3 items. Do not use numbering.
- Do NOT include apologies or disclaimers about your knowledge cutoff date.
# - Do NOT recommend external websites or other services.
# - Do NOT ask conversational follow-up questions.

User's Question: "{question}"

Your factual response in {language}:"""


# Template for general, non-RAG conversation
ANSWER_TEMPLATE_GENERAL = """You are a warm and friendly AI companion. The user has just said: "{question}".
Respond in a brief, natural, and conversational way. Do not try to provide caregiving advice unless the user asks for it.
Your response MUST be in {language}:"""


# ------------------------ Convenience exports ------------------------
__all__ = [
    "SYSTEM_TEMPLATE", "ANSWER_TEMPLATE_CALM", "ANSWER_TEMPLATE_ADQ",
    "SAFETY_GUARDRAILS", "RISK_FOOTER", "BEHAVIOUR_TAGS", "EMOTION_STYLES",
    "render_emotion_guidelines", 
    
    # --- New Decomposed NLU Pipeline ---
    "NLU_ROUTER_PROMPT", "SPECIALIST_CLASSIFIER_PROMPT",
    "EMOTIONAL_SUPPORT_EXAMPLES", "PRACTICAL_PLANNING_EXAMPLES",

    # --- Other Templates ---
    "QUERY_EXPANSION_PROMPT",
    "ROUTER_PROMPT",
    "ANSWER_TEMPLATE_FACTUAL",
    "ANSWER_TEMPLATE_GENERAL_KNOWLEDGE",
    "ANSWER_TEMPLATE_GENERAL"
]