session_base_multi.py

#!/usr/bin/env python3
"""
Shared infrastructure for multi-patient therapy session generation.

Extends session_base.py with support for:
  - Multiple patients in a single session (group therapy, couples/family)
  - Name-tagged ChatML messages: {"role": "user", "name": "partner_a", "content": "..."}
  - Custom turn-taking patterns (round-robin, therapist-picks, naturalistic)

Each multi-patient script imports from this module.
"""

import json
import os
import random
import re
import sys
import time
import argparse
from pathlib import Path

from session_base import (
    OLLAMA_BASE,
    THERAPIST_MODEL,
    PATIENT_MODEL,
    THERAPIST_TEMP,
    PATIENT_TEMP,
    MAX_RETRIES,
    RETRY_DELAY,
    STYLE_MAX_RETRIES,
    FORBIDDEN_OUTPUT_OPENINGS,
    PLATITUDE_PATTERNS,
    ROBOTIC_SIGNALS,
    SYCOPHANCY_MARKERS,
    THERAPIST_STYLE_PROFILES,
    PIXEL_SYSTEM_BASE,
    PATIENT_SYSTEM,
    ollama_chat,
    _check_style,
    run_generation,
)

MULTI_PATIENT_SYSTEM = (
    "You are a method actor playing a therapy participant in a simulated session. "
    "You are ONE specific person in a multi-person session (group therapy or couples/family therapy). "
    "Generate realistic, emotionally authentic dialogue for YOUR character only. "
    "Speak only as your character — no narration, no stage directions, no labels. "
    "Keep responses concise (2-5 sentences) and emotionally genuine. "
    "React to what the therapist AND other participants say. "
    "Do NOT start with 'This sounds crazy but' or similar phrases. "
    "Vary your openings — start mid-thought, with a question, or with a specific moment."
)


def _strip_name_for_ollama(messages):
    """Strip 'name' field and merge consecutive same-role messages for Ollama compat."""
    result = []
    for msg in messages:
        clean = {"role": msg["role"], "content": msg["content"]}
        if result and result[-1]["role"] == clean["role"]:
            # Merge consecutive same-role messages
            name_prefix = msg.get("name", "")
            if name_prefix:
                result[-1]["content"] += f"\n\n[{name_prefix}]: {clean['content']}"
            else:
                result[-1]["content"] += f"\n\n{clean['content']}"
        else:
            if msg.get("name"):
                clean["content"] = f"[{msg['name']}]: {clean['content']}"
            result.append(clean)
    return result


def generate_multi_patient_turn(
    persona,
    patient_name,
    presentation,
    category_name,
    conversation,
    turn_num,
    total_patient_turns,
    all_personas,
):
    """Generate a single patient's response in a multi-patient session."""
    if turn_num == 1:
        direction = f"The patient is arriving for a session. Their presenting concern: {presentation}. They're nervous but willing to talk."
    elif turn_num <= 3:
        direction = (
            "The patient is opening up, sharing more details. Starting to trust the therapist and other participants."
        )
    elif turn_num <= 5:
        direction = "The patient is going deeper — revealing the emotional impact, not just the surface problem. Becoming more vulnerable."
    elif turn_num == total_patient_turns:
        direction = "Final turn. The patient is reflecting on what they've discussed, maybe feeling a shift or maybe just sitting with it."
    else:
        direction = "The patient is processing what the therapist and other participants said. May push back, have a realization, or share something they held back."

    conv_text = ""
    for msg in conversation:
        if msg["role"] == "assistant":
            role = "Therapist"
        elif msg.get("name"):
            role = msg["name"]
        else:
            role = "Patient"
        conv_text += f"{role}: {msg['content']}\n"

    other_participants = [p for p in all_personas if p["name"] != patient_name]
    others_desc = "\n".join(
        f"- {p['name']} (age {p['age']}, {p['gender']}, {p['occupation']})" for p in other_participants
    )

    prompt = f"""You are playing {patient_name} in a multi-person therapy session. Stay completely in character.

YOUR CHARACTER:
Name: {patient_name}
Age: {persona["age"]}, Gender: {persona["gender"]}, Occupation: {persona["occupation"]}
Presenting concern: {persona["presenting"]}

OTHER PARTICIPANTS IN THE ROOM:
{others_desc}

SESSION FOCUS: {category_name} — {presentation}

DIRECTION FOR THIS TURN:
{direction}

This is your turn {turn_num} of {total_patient_turns}.

CONVERSATION SO FAR:
{conv_text if conv_text else "(First turn — session is just starting.)"}

What does {patient_name} say next? Generate ONLY {patient_name}'s spoken words — no labels, no narration, no other participant's dialogue. 2-5 sentences."""

    messages = [
        {"role": "system", "content": MULTI_PATIENT_SYSTEM},
        {"role": "user", "content": prompt},
    ]
    return ollama_chat(messages, model=PATIENT_MODEL, temperature=PATIENT_TEMP, num_predict=250)


def generate_therapist_turn_multi(
    category,
    conversation,
    turn_num,
    style_profile,
    all_patient_names,
):
    """Generate therapist response in a multi-patient session."""
    technique = random.choice(category["therapist_techniques"])
    technique_guidance = (
        f"\n\n[INTERNAL CLINICAL GUIDANCE — embody, never state explicitly]: "
        f"Use this technique naturally: {technique}. "
        f"Weave it into the conversation — don't announce it. "
        f"Respond to what the participants actually said, don't pivot to a technique if it doesn't fit."
    )

    style_guidance = (
        f"\n\nSTYLE: {style_profile['description']}\n"
        f"NEVER start with: {', '.join(style_profile['forbidden_openings'])}\n"
        f"Good examples: {'; '.join(style_profile['good_examples'][:3])}\n"
        f"MAX {style_profile['max_sentences']} sentences, {style_profile['max_words']} words."
    )

    addon = f"\n\n{category['therapist_prompt_addon']}"

    participants_desc = f"Participants in this session: {', '.join(all_patient_names)}."

    system_content = PIXEL_SYSTEM_BASE + technique_guidance + style_guidance + addon + f"\n\n{participants_desc}"
    # Strip 'name' field — Ollama doesn't support it
    clean_conversation = _strip_name_for_ollama(conversation)
    messages = [{"role": "system", "content": system_content}, *clean_conversation]
    return ollama_chat(messages, model=THERAPIST_MODEL, temperature=THERAPIST_TEMP, num_predict=400)


def generate_therapist_turn_multi_validated(
    category,
    conversation,
    turn_num,
    style_profile,
    all_patient_names,
):
    """Generate therapist turn with style validation + retry."""
    for attempt in range(STYLE_MAX_RETRIES):
        output = generate_therapist_turn_multi(category, conversation, turn_num, style_profile, all_patient_names)
        passed, reason = _check_style(output, style_profile)
        if passed:
            return output
        print(f"    [style retry {attempt + 1}/{STYLE_MAX_RETRIES}] {reason}")
    return output


def generate_multi_patient_session(
    category_key,
    category,
    personas,
    presentation,
    session_idx,
    session_id_prefix,
    min_turns=16,
    max_turns=24,
    turn_pattern="naturalistic",
):
    """Generate a complete multi-patient therapy session.

    Args:
        personas: list of persona dicts, each with "name" key
        turn_pattern: "round_robin", "therapist_picks", or "naturalistic"
    """
    total_turns = random.randint(min_turns // 2, max_turns // 2) * 2
    total_therapist_turns = total_turns // 2
    style_keys = list(THERAPIST_STYLE_PROFILES.keys())
    style_profile = THERAPIST_STYLE_PROFILES[style_keys[session_idx % len(style_keys)]]

    patient_names = [p["name"] for p in personas]
    conversation = []

    # Track how many turns each patient has had
    patient_turn_counts = {name: 0 for name in patient_names}
    total_patient_turns_target = total_therapist_turns  # each therapist turn followed by 1-2 patient turns

    # First patient opens the session, then therapist responds (matching single-patient pattern)
    first_persona = personas[0]
    patient_turn_counts[first_persona["name"]] += 1
    first_patient_msg = generate_multi_patient_turn(
        first_persona,
        first_persona["name"],
        presentation,
        category["name"],
        conversation,
        1,
        total_patient_turns_target,
        personas,
    )
    conversation.append({"role": "user", "name": first_persona["name"], "content": first_patient_msg})

    for round_num in range(1, total_therapist_turns + 1):
        # Therapist turn
        therapist_msg = generate_therapist_turn_multi_validated(
            category, conversation, round_num, style_profile, patient_names
        )
        conversation.append({"role": "assistant", "content": therapist_msg})

        # Determine which patient(s) speak after this therapist turn
        if turn_pattern == "round_robin":
            # Each round, the next patient in rotation speaks
            patient_idx = (round_num - 1) % len(personas)
            selected_patients = [personas[patient_idx]]
        elif turn_pattern == "therapist_picks":
            # Simulate therapist directing to a specific patient
            patient_idx = (round_num - 1) % len(personas)
            selected_patients = [personas[patient_idx]]
        else:  # naturalistic
            # 1-2 patients respond, weighted by who's spoken least
            n_responders = random.choice([1, 1, 1, 2])  # mostly 1, sometimes 2
            sorted_by_turns = sorted(personas, key=lambda p: patient_turn_counts[p["name"]])
            selected_patients = sorted_by_turns[:n_responders]

        for persona in selected_patients:
            patient_turn_counts[persona["name"]] += 1
            patient_turn_num = patient_turn_counts[persona["name"]]
            patient_msg = generate_multi_patient_turn(
                persona,
                persona["name"],
                presentation,
                category["name"],
                conversation,
                patient_turn_num,
                total_patient_turns_target,
                personas,
            )
            conversation.append({"role": "user", "name": persona["name"], "content": patient_msg})

    session_id = f"{session_id_prefix}_{category_key}_{session_idx:04d}"
    return {
        "messages": [
            {"role": "system", "content": PIXEL_SYSTEM_BASE},
            *conversation,
        ],
        "metadata": {
            "source_family": category.get("source_family", session_id_prefix),
            "category": category_key,
            "category_name": category["name"],
            "presentation": presentation,
            "session_id": session_id,
            "participants": [
                {"name": p["name"], "age": p["age"], "gender": p["gender"], "occupation": p["occupation"]}
                for p in personas
            ],
            "n_participants": len(personas),
            "style_profile": style_profile["description"][:50],
            "turns": len(conversation),
            "difficulty": category["difficulty"],
            "turn_pattern": turn_pattern,
        },
    }


def session_exists(output_file, session_id):
    if not output_file.exists():
        return False
    with open(output_file) as f:
        for line in f:
            if session_id in line:
                return True
    return False


def run_multi_generation(
    categories_dict,
    output_dir,
    output_filename,
    session_id_prefix,
    source_family,
    default_sessions_per_category=50,
    description="Multi-Patient Therapy Session Generation",
    min_turns=16,
    max_turns=24,
    turn_pattern="naturalistic",
):
    """Main generation loop for multi-patient sessions."""
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--categories", default="all", help="Comma-separated category keys or 'all'")
    parser.add_argument("--sessions-per-category", type=int, default=default_sessions_per_category)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--spot-check", type=int, default=None, help="Generate N sessions from first category only")
    args = parser.parse_args()

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / output_filename

    if args.categories == "all":
        cats = list(categories_dict.keys())
    else:
        cats = [c.strip() for c in args.categories.split(",")]

    if args.spot_check:
        cats = cats[:1]
        total_sessions = args.spot_check
    else:
        total_sessions = len(cats) * args.sessions_per_category

    print(f"\n=== {description.upper()} ===")
    print(f"Categories: {len(cats)} ({', '.join(cats)})")
    print(f"Sessions per category: {args.spot_check or args.sessions_per_category}")
    print(f"Total sessions: {total_sessions}")
    print(f"Output: {output_file}")
    print(f"Therapist: {THERAPIST_MODEL}")
    print(f"Patient: {PATIENT_MODEL}")
    print(f"Turns: {min_turns}-{max_turns} (multi-patient)")
    print(f"Turn pattern: {turn_pattern}")
    print()

    completed = 0
    skipped = 0
    failed = 0
    start_time = time.time()

    for cat_key in cats:
        category = categories_dict[cat_key]
        n_sessions = args.spot_check or args.sessions_per_category

        category["source_family"] = source_family

        print(f"\n--- {category['name']} ({cat_key}) ---")
        print(f"  {len(category['presentations'])} presentations × {len(category['persona_sets'])} persona sets")

        for i in range(n_sessions):
            presentation = category["presentations"][i % len(category["presentations"])]
            persona_set = category["persona_sets"][i % len(category["persona_sets"])]
            session_id = f"{session_id_prefix}_{cat_key}_{i:04d}"

            if args.resume and session_exists(output_file, session_id):
                skipped += 1
                continue

            try:
                session = generate_multi_patient_session(
                    cat_key,
                    category,
                    persona_set,
                    presentation,
                    i,
                    session_id_prefix,
                    min_turns,
                    max_turns,
                    turn_pattern,
                )
                with open(output_file, "a") as f:
                    f.write(json.dumps(session) + "\n")
                completed += 1
                elapsed = time.time() - start_time
                rate = completed / (elapsed / 3600) if elapsed > 0 else 0
                remaining = (total_sessions - completed - skipped) / rate if rate > 0 else 0
                print(
                    f"  ✓ {session_id} {len(session['messages'])} msgs, {len(persona_set)} patients | done: {completed}/{total_sessions} | ~{remaining:.1f}h left"
                )
            except Exception as e:
                failed += 1
                print(f"  ✗ {session_id} FAILED: {e}")
                with open(output_dir / "errors.log", "a") as f:
                    f.write(f"{session_id}: {e}\n")

    elapsed = time.time() - start_time
    print(f"\n=== COMPLETE ===")
    print(f"Generated: {completed}")
    print(f"Skipped: {skipped}")
    print(f"Failed: {failed}")
    print(f"Elapsed: {elapsed / 3600:.1f}h")
    print(f"Output: {output_file}")