WalkXR-AI / Paul-Game-Jam-Phase-2.5 / src / nodes / reflection.py
reflection.py
Raw
"""Node that generates reflective feedback after conversation ends."""

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI

from src.resources.empathetic_reflection_lexicon import LEXICON
from src.schemas.state import WalkState


def _get_llm():
    return ChatOpenAI(model="gpt-4o-mini", temperature=0.5)


REFLECTION_PROMPT = """You are an empathetic communication coach specializing in trauma-informed conversations.
The user just finished a role-play session where they practiced talking to a veteran struggling with PTSD.
Based on the rich data below, provide warm, specific, and actionable feedback.

## Scenario
{scenario_title}: {scenario_description}

## Session Statistics
- Total turns: {turn_count}
- Final confidence (James's trust): {final_confidence:.1%}
- Final rapport score: {rapport_score:.1%}
- NPC final mood: {npc_mood}
- NPC final openness: {npc_openness:.1%}

## Empathy Performance
- Empathetic moments (hits): {empathy_hits}
- Insensitive moments (misses): {empathy_misses}
- Empathy ratio: {empathy_ratio:.0%}

## Conversation Stage Progression
{stage_progression}

## Emotion Changes (User)
{emotion_summary}
- Dominant emotion throughout: {dominant_emotion}

## Observed Behaviors
{observations}

## End Condition
{end_reason}

## Empathetic Expression Guide
Reference these principles:
{lexicon_sample}

## Feedback Structure
Write your feedback with these sections:

**What You Did Well**: Highlight 2-3 specific moments where the user showed genuine empathy. Reference their actual words or approach if visible from observations.

**How James Responded**: Describe James's emotional arc — did he open up? Stay guarded? Withdraw? Connect this to the user's approach. Use the stage progression and NPC mood data.

**Your Empathy Score**: Based on the {empathy_ratio:.0%} empathy ratio and the rapport score of {rapport_score:.1%}. Frame this encouragingly regardless of score.

**Tips for Next Time**: Give 2-3 concrete, trauma-informed tips. Tailor these to what the user struggled with (based on observations and misses).

**Final Thought**: A brief, warm closing that affirms the value of showing up and trying.

Keep it under 250 words. Be direct but warm. No bullet points — use flowing paragraphs.
"""


def _determine_end_reason(state: WalkState) -> str:
    """Infer why the conversation ended for richer feedback."""
    memory = state.get("memory", {})
    agent_output = state.get("agent_output", {})
    scenario = state.get("scenario", {})
    turn_index = state.get("turn_index", 0)

    if agent_output.get("should_end_turn"):
        if memory.get("npc_mood") == "withdrawn":
            return "James withdrew from the conversation due to feeling pushed or misunderstood."
        return "The conversation reached a natural ending — James chose to wrap up."

    if turn_index >= scenario.get("max_turns", 10):
        rapport = memory.get("rapport_score", 0)
        if rapport >= 0.7:
            return "Time ran out, but a meaningful connection was established."
        elif rapport >= 0.4:
            return "Time ran out. James was starting to open up but needed more time."
        else:
            return "Time ran out. James remained guarded throughout."

    if memory.get("npc_mood") == "withdrawn":
        return "James shut down after feeling the conversation became insensitive or pushy."

    if memory.get("silence_count", 0) >= 3:
        return "The conversation ended due to user disengagement (repeated short responses)."

    if memory.get("rapport_score", 0) >= 0.85:
        return "A deep connection was achieved — the conversation concluded on a high note."

    return "The conversation ended."


def reflection(state: WalkState) -> dict:
    """Generate reflection feedback based on full memory data."""

    memory = state.get("memory", {})
    scenario = state.get("scenario", {})
    trajectory = memory.get("emotion_trajectory", [])

    # Emotion summary
    if trajectory:
        emotion_lines = [
            f"  Turn {e['turn']}: {e.get('label', '?')} "
            f"(valence={e.get('valence', 0):.2f}, arousal={e.get('arousal', 0):.2f})"
            for e in trajectory
        ]
        emotion_summary = "\n".join(emotion_lines)
    else:
        emotion_summary = "(No emotion data recorded)"

    # Observations
    observations = "\n".join(
        f"- Turn {i + 1}: {obs}"
        for i, obs in enumerate(memory.get("key_observations", []))
    ) or "(No observations recorded)"

    # Stage progression
    stage_history = memory.get("stage_history", [])
    stage_progression = "".join(stage_history) if stage_history else "opening"

    # Lexicon samples
    lexicon_sample = "\n".join(f"- {expr}" for expr in LEXICON[:6])

    # End reason
    end_reason = _determine_end_reason(state)

    prompt = REFLECTION_PROMPT.format(
        scenario_title=scenario.get("title", ""),
        scenario_description=scenario.get("description", ""),
        turn_count=state.get("turn_index", 0),
        final_confidence=memory.get("confidence_level", 0.5),
        rapport_score=memory.get("rapport_score", 0.0),
        npc_mood=memory.get("npc_mood", "guarded"),
        npc_openness=memory.get("npc_openness", 0.3),
        empathy_hits=memory.get("empathy_hits", 0),
        empathy_misses=memory.get("empathy_misses", 0),
        empathy_ratio=memory.get("empathy_ratio", 0.5),
        stage_progression=stage_progression,
        emotion_summary=emotion_summary,
        dominant_emotion=memory.get("dominant_emotion", "neutral"),
        observations=observations,
        end_reason=end_reason,
        lexicon_sample=lexicon_sample,
    )

    messages = [
        SystemMessage(content=prompt),
        HumanMessage(content="Please write the feedback."),
    ]

    result = _get_llm().invoke(messages)

    return {"reflection_result": result.content}