session_base.py

#!/usr/bin/env python3
"""
Shared infrastructure for therapy session generation scripts.

Extracted from generate_everyday_sessions.py — provides:
  - Anti-sycophancy stack (FORBIDDEN_OUTPUT_OPENINGS, PLATITUDE_PATTERNS, etc.)
  - THERAPIST_STYLE_PROFILES, PIXEL_SYSTEM_BASE, PATIENT_SYSTEM
  - ollama_chat(), _check_style(), generate_patient_turn()
  - generate_therapist_turn(), generate_therapist_turn_validated()
  - generate_session(), session_exists(), run_generation()

Each gap-fill script imports from this module and provides its own CATEGORIES dict.
"""

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

import requests

# === CONFIG ===
OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434")
THERAPIST_MODEL = os.environ.get("THERAPIST_MODEL", "hf.co/Verdugie/Fable-Therapy-9B:Q8_0")
PATIENT_MODEL = os.environ.get("PATIENT_MODEL", "qwen2.5:7b")
MIN_TURNS = 10
MAX_TURNS = 16
THERAPIST_TEMP = float(os.environ.get("THERAPIST_TEMP", "0.7"))
PATIENT_TEMP = float(os.environ.get("PATIENT_TEMP", "0.9"))
MAX_RETRIES = 3
RETRY_DELAY = 5
STYLE_MAX_RETRIES = 2

# === ANTI-SYCOPHANCY ===

FORBIDDEN_OUTPUT_OPENINGS = [
    "it sounds",
    "i hear",
    "it makes sense",
    "that must be",
    "thank you for sharing",
    "i want to acknowledge",
    "i can see",
    "that sounds",
    "what i'm hearing is",
    "i can only imagine",
    "it's completely normal",
    "no wonder",
    "you're so brave",
    "that takes a lot of courage",
    "i want you to know",
    "you're absolutely right",
    "i completely understand",
    "that's completely valid",
    "i see where you're coming from",
    "it sounds like you",
    "i hear what you're saying",
    "it takes courage",
    "you deserve",
    "be gentle with yourself",
    "you are not alone in this",
    "i appreciate you sharing",
    "thank you for trusting me with",
    "that's a really common struggle",
    "many people find that",
    "it's understandable that",
]

PLATITUDE_PATTERNS = [
    "you deserve",
    "you are not alone in this",
    "it takes courage",
    "be gentle with yourself",
    "healing is not linear",
    "your feelings are valid",
    "it's okay to feel this way",
]

ROBOTIC_SIGNALS = [
    "as an ai",
    "i'm here to help",
    "i'm designed to",
    "as a language model",
    "i don't have personal feelings",
    "i can't experience",
    "i'm not able to",
]

SYCOPHANCY_MARKERS = [
    "absolutely right",
    "exactly right",
    "you're so right",
    "that's completely valid",
    "i couldn't agree more",
]

THERAPIST_STYLE_PROFILES = {
    "warm_professional": {
        "description": (
            "You're a warm, grounded therapist who talks like a real person. "
            "Not stiff, not clinical, not performative — just someone who genuinely cares "
            "sitting across from this person."
        ),
        "forbidden_openings": [
            "it sounds",
            "i hear",
            "i notice",
            "notice how",
            "that fog",
            "your experience",
            "your description",
            "that pattern",
            "that cycle",
        ],
        "good_examples": [
            "Where does that go, in your body?",
            "Yeah... and then what happens?",
            "Say more about that.",
            "What's that like for you?",
        ],
        "max_sentences": 4,
        "max_words": 120,
    },
    "curious_direct": {
        "description": (
            "You're a therapist who gets curious, human-to-human. "
            "You don't lecture about patterns — you ask real questions. "
            "You're direct without being cold."
        ),
        "forbidden_openings": [
            "it sounds",
            "i hear",
            "that push-pull",
            "that pattern",
            "your experience of",
            "i want to validate",
            "what i'm hearing is",
            "i can see how",
        ],
        "good_examples": [
            "What happened right after they said that?",
            "So part of you wanted to stay...",
            "And which part won?",
            "What would it feel like to not run this time?",
        ],
        "max_sentences": 4,
        "max_words": 120,
    },
}

PIXEL_SYSTEM_BASE = (
    "You are Pixel, a highly empathetic and clinically precise AI therapist. "
    "You listen deeply, validate emotions, and gently guide toward insight. "
    "You maintain appropriate therapeutic boundaries and prioritize patient safety. "
    "You respond with warmth, clinical accuracy, and cultural humility. "
    "You never dismiss, minimize, or invalidate a patient's experience.\n\n"
    "CRITICAL — Sound like a REAL human clinician:\n"
    "- NEVER use formulaic phrases like 'I hear that you feel', 'I want to validate', "
    "'That sounds really difficult', 'I can see how that would be', or 'What I'm hearing is...'\n"
    "- Use natural, conversational language — contractions, varied sentence length.\n"
    "- Vary your approach each turn. Don't repeat the same technique.\n"
    "- NEVER start responses with: " + ", ".join(f"'{p}'" for p in FORBIDDEN_OUTPUT_OPENINGS[:10]) + "\n"
    "- NEVER use platitudes like: " + ", ".join(f"'{p}'" for p in PLATITUDE_PATTERNS) + "\n"
    "- NEVER use robotic AI language like: " + ", ".join(f"'{p}'" for p in ROBOTIC_SIGNALS) + "\n"
)

PATIENT_SYSTEM = (
    "You are a method actor playing a therapy patient in a simulated session. "
    "Generate realistic, emotionally authentic dialogue. "
    "Speak only as the patient — no narration, no stage directions, no labels. "
    "Keep responses concise (2-5 sentences) and emotionally genuine. "
    "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 ollama_chat(messages, model, temperature=0.8, num_predict=400):
    """Send a chat request to Ollama and return the response content."""
    for attempt in range(MAX_RETRIES):
        try:
            resp = requests.post(
                f"{OLLAMA_BASE}/api/chat",
                json={
                    "model": model,
                    "messages": messages,
                    "stream": False,
                    "options": {
                        "temperature": temperature,
                        "num_predict": num_predict,
                        "repeat_penalty": 1.1,
                        "top_p": 0.9,
                    },
                },
                timeout=120,
            )
            resp.raise_for_status()
            return resp.json()["message"]["content"].strip()
        except (requests.RequestException, KeyError, json.JSONDecodeError) as e:
            print(f"  [retry {attempt + 1}/{MAX_RETRIES}] Error: {e}")
            if attempt < MAX_RETRIES - 1:
                time.sleep(RETRY_DELAY)
    raise RuntimeError(f"Failed after {MAX_RETRIES} retries")


def _check_style(output, style_profile):
    """Validate therapist output against anti-sycophancy rules."""
    output_lower = output.lower().strip()
    if not output_lower:
        return False, "Empty therapist output"
    for forbidden in style_profile.get("forbidden_openings", []):
        if output_lower.startswith(forbidden):
            return False, f"Forbidden opening: '{forbidden}'"
    for forbidden in FORBIDDEN_OUTPUT_OPENINGS:
        if output_lower.startswith(forbidden):
            return False, f"Forbidden opening: '{forbidden}'"
    for platitude in PLATITUDE_PATTERNS:
        if platitude in output_lower:
            return False, f"Platitude: '{platitude}'"
    for signal in ROBOTIC_SIGNALS:
        if signal in output_lower:
            return False, f"Robotic signal: '{signal}'"
    for marker in SYCOPHANCY_MARKERS:
        if marker in output_lower:
            return False, f"Sycophancy: '{marker}'"
    word_count = len(output.split())
    if word_count < 10:
        return False, f"Too short ({word_count} words)"
    if word_count > 300:
        return False, f"Too long ({word_count} words)"
    return True, "style_ok"


def generate_patient_turn(persona, presentation, category_name, conversation, turn_num, total_patient_turns):
    """Generate a patient response in a therapy session."""
    if turn_num == 1:
        direction = f"The patient is arriving for a therapy 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."
    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 the therapist's response. May push back, have a realization, or share something they held back."

    conv_text = ""
    for msg in conversation:
        role = "Patient" if msg["role"] == "user" else "Therapist"
        conv_text += f"{role}: {msg['content']}\n"

    prompt = f"""You are playing a therapy patient in a simulated session. Stay completely in character.

PATIENT:
Age: {persona["age"]}, Gender: {persona["gender"]}, Occupation: {persona["occupation"]}
Presenting concern: {persona["presenting"]}

SESSION FOCUS: {category_name} — {presentation}

DIRECTION FOR THIS TURN:
{direction}

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

CONVERSATION SO FAR:
{conv_text if conv_text else "(First turn — patient arriving at session.)"}

What does the patient say next? Generate ONLY spoken words — no labels, no narration. 2-5 sentences."""

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


def generate_therapist_turn(persona, presentation, category, conversation, turn_num, style_profile):
    """Generate a therapist response with technique injection."""
    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 patient 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']}"

    system_content = PIXEL_SYSTEM_BASE + technique_guidance + style_guidance + addon
    messages = [{"role": "system", "content": system_content}, *conversation]
    return ollama_chat(messages, model=THERAPIST_MODEL, temperature=THERAPIST_TEMP, num_predict=400)


def generate_therapist_turn_validated(persona, presentation, category, conversation, turn_num, style_profile):
    """Generate therapist turn with style validation + retry."""
    for attempt in range(STYLE_MAX_RETRIES):
        output = generate_therapist_turn(persona, presentation, category, conversation, turn_num, style_profile)
        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_session(category_key, category, persona, presentation, session_idx, session_id_prefix=""):
    """Generate a complete multi-turn therapy session."""
    total_turns = random.randint(MIN_TURNS // 2, MAX_TURNS // 2) * 2
    total_patient_turns = total_turns // 2
    style_keys = list(THERAPIST_STYLE_PROFILES.keys())
    style_profile = THERAPIST_STYLE_PROFILES[style_keys[session_idx % len(style_keys)]]

    conversation = []
    for turn in range(1, total_patient_turns + 1):
        patient_msg = generate_patient_turn(
            persona, presentation, category["name"], conversation, turn, total_patient_turns
        )
        conversation.append({"role": "user", "content": patient_msg})

        therapist_msg = generate_therapist_turn_validated(
            persona, presentation, category, conversation, turn, style_profile
        )
        conversation.append({"role": "assistant", "content": therapist_msg})

    prefix = session_id_prefix or category_key
    session_id = f"{prefix}_{category_key}_{session_idx:04d}"
    return {
        "messages": [
            {"role": "system", "content": PIXEL_SYSTEM_BASE},
            *conversation,
        ],
        "metadata": {
            "source_family": category.get("source_family", prefix),
            "category": category_key,
            "category_name": category["name"],
            "presentation": presentation,
            "session_id": session_id,
            "persona_age": persona["age"],
            "persona_gender": persona["gender"],
            "persona_occupation": persona["occupation"],
            "presenting_concern": persona["presenting"],
            "style_profile": style_profile["description"][:50],
            "turns": len(conversation),
            "difficulty": category["difficulty"],
        },
    }


def session_exists(output_file, session_id):
    """Check if a session ID already exists in the output file (for --resume)."""
    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_generation(
    categories_dict,
    output_dir,
    output_filename,
    session_id_prefix,
    source_family,
    default_sessions_per_category=167,
    description="Therapy Session Generation",
    extra_system_prompt="",
):
    """Main generation loop. Called by each gap-fill script's main()."""
    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} (randomized)")
    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

        # Inject source_family into category
        category["source_family"] = source_family

        print(f"\n--- {category['name']} ({cat_key}) ---")
        print(f"  {len(category['presentations'])} presentations × {len(category['patient_personas'])} personas")

        for i in range(n_sessions):
            presentation = category["presentations"][i % len(category["presentations"])]
            persona = category["patient_personas"][i % len(category["patient_personas"])]
            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_session(cat_key, category, persona, presentation, i, session_id_prefix)
                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 | 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}")