{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 1 — Train a Dialogue Model with StrataSynth\n",
    "\n",
    "**What you'll learn:**\n",
    "- Generate a structured conversation dataset via the StrataSynth API\n",
    "- Explore the unique fields: `intent`, `communication_act`, `belief_state`, `relationship_state`\n",
    "- Prepare data for fine-tuning (OpenAI, Hugging Face, or raw transformers)\n",
    "\n",
    "**Requirements:**\n",
    "```\n",
    "pip install requests pandas datasets transformers\n",
    "```\n",
    "\n",
    "---\n",
    "> **API endpoint:** `https://api.stratasynth.com`  \n",
    "> Get your API key at [app.stratasynth.com](https://app.stratasynth.com)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "import time\n",
    "import pandas as pd\n",
    "\n",
    "# --- Configuration ---\n",
    "API_KEY = \"ss_live_...\"       # Replace with your API key\n",
    "BASE_URL = \"https://api.stratasynth.com\"\n",
    "\n",
    "# Authenticate: exchange API key for JWT\n",
    "resp = requests.post(\n",
    "    f\"{BASE_URL}/auth/token\",\n",
    "    json={\"api_key\": API_KEY}\n",
    ")\n",
    "resp.raise_for_status()\n",
    "TOKEN = resp.json()[\"data\"][\"token\"]\n",
    "HEADERS = {\"Authorization\": f\"Bearer {TOKEN}\"}\n",
    "print(\"Authenticated. Token valid 24h.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 — Browse available scenarios\n",
    "\n",
    "StrataSynth includes 13 scenarios across 4 categories: family (FAM), romantic (ROM), professional (PRO), and life transitions (VIT)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scenarios = requests.get(f\"{BASE_URL}/scenarios\", headers=HEADERS).json()[\"data\"]\n",
    "\n",
    "df_scenarios = pd.DataFrame([\n",
    "    {\"id\": s[\"id\"], \"name\": s[\"name\"], \"category\": s[\"category\"], \"description\": s[\"description\"]}\n",
    "    for s in scenarios\n",
    "])\n",
    "df_scenarios"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 — Generate a dataset\n",
    "\n",
    "We'll generate 20 conversations from scenario `FAM-01` (Family Caregiver).  \n",
    "Generation is **async** — we create a job and poll until completed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create generation job\n",
    "job_resp = requests.post(\n",
    "    f\"{BASE_URL}/jobs/dataset\",\n",
    "    headers=HEADERS,\n",
    "    json={\n",
    "        \"scenario_id\": \"FAM-01\",\n",
    "        \"conversation_count\": 20,\n",
    "        \"complexity\": 3,        # 1-5: depth of psychological profile\n",
    "        \"language\": \"en\",\n",
    "        \"adapter\": \"flat_jsonl\", # output format\n",
    "        \"noise_pct\": 0.1,       # 10% deliberate noise (lies, retractions)\n",
    "        \"seed\": 42              # reproducibility\n",
    "    }\n",
    ").json()\n",
    "\n",
    "job_id = job_resp[\"data\"][\"jobId\"]\n",
    "print(f\"Job created: {job_id}\")\n",
    "\n",
    "# Poll until COMPLETED\n",
    "while True:\n",
    "    status_resp = requests.get(f\"{BASE_URL}/jobs/{job_id}\", headers=HEADERS).json()[\"data\"]\n",
    "    status = status_resp[\"status\"]\n",
    "    progress = status_resp.get(\"progress\", 0)\n",
    "    print(f\"  {status} — {progress}%\", end=\"\\r\")\n",
    "    if status == \"COMPLETED\":\n",
    "        print(f\"\\nDone! Download URL ready.\")\n",
    "        download_url = status_resp[\"download_url\"]\n",
    "        break\n",
    "    elif status == \"FAILED\":\n",
    "        raise RuntimeError(f\"Job failed: {status_resp.get('error')}\")\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3 — Load and explore the dataset\n",
    "\n",
    "Each record in a StrataSynth dataset is a full conversation with turn-level ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "\n",
    "# Download the JSONL file\n",
    "urllib.request.urlretrieve(download_url, \"social_reasoning.jsonl\")\n",
    "\n",
    "# Load conversations\n",
    "conversations = []\n",
    "with open(\"social_reasoning.jsonl\") as f:\n",
    "    for line in f:\n",
    "        conversations.append(json.loads(line))\n",
    "\n",
    "print(f\"Loaded {len(conversations)} conversations\")\n",
    "print(f\"Keys in each conversation: {list(conversations[0].keys())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect the first turn of the first conversation\n",
    "turn = conversations[0][\"turns\"][0]\n",
    "print(json.dumps(turn, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What makes StrataSynth data different\n",
    "\n",
    "Each turn includes:\n",
    "\n",
    "| Field | What it captures |\n",
    "|---|---|\n",
    "| `text` | The actual utterance |\n",
    "| `speaker` | A or B |\n",
    "| `intent` | Why the speaker said it (e.g. `express_frustration`) |\n",
    "| `goal` | What they want from the exchange (e.g. `seek_validation`) |\n",
    "| `communication_act` | Pragmatic move (e.g. `accusation`, `reassurance`, `deflection`) |\n",
    "| `emotional_state` | Emotion + intensity |\n",
    "| `relationship_state` | trust, tension, connection, dominance_balance |\n",
    "| `belief_state` | 12 beliefs × (value, confidence) |\n",
    "| `belief_delta` | How beliefs shifted this turn |\n",
    "| `is_noise` | Whether this turn contains deliberate inaccuracy |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flatten to a DataFrame — one row per turn\n",
    "rows = []\n",
    "for conv in conversations:\n",
    "    conv_id = conv[\"conversation_id\"]\n",
    "    for turn in conv[\"turns\"]:\n",
    "        rows.append({\n",
    "            \"conversation_id\": conv_id,\n",
    "            \"turn_index\": turn[\"turn_index\"],\n",
    "            \"speaker\": turn[\"speaker\"],\n",
    "            \"text\": turn[\"text\"],\n",
    "            \"intent\": turn.get(\"intent\"),\n",
    "            \"communication_act\": turn.get(\"communication_act\"),\n",
    "            \"trust\": turn.get(\"relationship_state\", {}).get(\"trust\"),\n",
    "            \"tension\": turn.get(\"relationship_state\", {}).get(\"tension\"),\n",
    "            \"is_noise\": turn.get(\"is_noise\", False),\n",
    "        })\n",
    "\n",
    "df = pd.DataFrame(rows)\n",
    "print(f\"Total turns: {len(df)}\")\n",
    "df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Distribution of communication acts\n",
    "df[\"communication_act\"].value_counts().head(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relationship tension progression across a single conversation\n",
    "conv_df = df[df[\"conversation_id\"] == df[\"conversation_id\"].iloc[0]]\n",
    "conv_df[[\"turn_index\", \"speaker\", \"trust\", \"tension\"]].plot(\n",
    "    x=\"turn_index\", y=[\"trust\", \"tension\"], title=\"Relationship state over turns\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 — Prepare for fine-tuning\n",
    "\n",
    "### Option A: OpenAI fine-tuning format\n",
    "\n",
    "Re-generate with `adapter=\"openai_finetuning\"` to get the messages format directly. Or convert from flat_jsonl:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_openai_messages(conversation):\n",
    "    \"\"\"Convert a StrataSynth conversation to OpenAI fine-tuning format.\n",
    "    Injects psychological context into the system prompt.\n",
    "    \"\"\"\n",
    "    persona_a = conversation.get(\"persona_a\", {})\n",
    "    system_prompt = (\n",
    "        f\"You are persona A: {persona_a.get('archetype', 'person')}, \"\n",
    "        f\"attachment style: {persona_a.get('psychological_core', {}).get('attachment_style', 'unknown')}. \"\n",
    "        f\"Respond authentically to the conversation.\"\n",
    "    )\n",
    "    messages = [{\"role\": \"system\", \"content\": system_prompt}]\n",
    "    for turn in conversation[\"turns\"]:\n",
    "        role = \"assistant\" if turn[\"speaker\"] == \"A\" else \"user\"\n",
    "        messages.append({\"role\": role, \"content\": turn[\"text\"]})\n",
    "    return {\"messages\": messages}\n",
    "\n",
    "openai_records = [to_openai_messages(c) for c in conversations]\n",
    "\n",
    "with open(\"openai_finetuning.jsonl\", \"w\") as f:\n",
    "    for r in openai_records:\n",
    "        f.write(json.dumps(r) + \"\\n\")\n",
    "\n",
    "print(f\"Saved {len(openai_records)} records for OpenAI fine-tuning\")\n",
    "print(json.dumps(openai_records[0][\"messages\"][:2], indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option B: Hugging Face Hub\n",
    "\n",
    "Push directly to your HF organization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "HF_TOKEN = \"hf_...\"          # Your Hugging Face token\n",
    "HF_REPO  = \"stratasynth/social-reasoning-fam-01\"  # Your org/repo\n",
    "\n",
    "hf_rows = [\n",
    "    {\n",
    "        \"conversation_id\": conv[\"conversation_id\"],\n",
    "        \"turns\": json.dumps(conv[\"turns\"]),   # serialize nested structure\n",
    "        \"scenario_id\": conv[\"scenario_id\"],\n",
    "        \"turn_count\": len(conv[\"turns\"]),\n",
    "    }\n",
    "    for conv in conversations\n",
    "]\n",
    "\n",
    "dataset = Dataset.from_list(hf_rows)\n",
    "dataset.push_to_hub(HF_REPO, token=HF_TOKEN)\n",
    "print(f\"Pushed to https://huggingface.co/datasets/{HF_REPO}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option C: Preview using GET /preview (no download needed)\n",
    "\n",
    "Before committing to a full download, inspect the first 2 conversations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preview = requests.get(\n",
    "    f\"{BASE_URL}/jobs/{job_id}/preview\",\n",
    "    headers=HEADERS\n",
    ").json()[\"data\"]\n",
    "\n",
    "print(f\"Preview: {preview['conversation_count']} conversations available\")\n",
    "first_turn = preview[\"conversations\"][0][\"turns\"][0]\n",
    "print(json.dumps(first_turn, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Summary\n",
    "\n",
    "You now have a structured dataset ready for fine-tuning. The key advantage of StrataSynth data over generic conversation datasets:\n",
    "\n",
    "- Every turn has **intent** and **communication act** — not just text\n",
    "- **Belief state** and **relationship state** evolve per turn — the data tells a psychological story\n",
    "- **Noise is labelled** — you can train systems that distinguish honest from dishonest speech\n",
    "- **Reproducible** via seed — same seed = same dataset\n",
    "\n",
    "**Next:** See Notebook 2 to evaluate a system trained on this data using StrataSynth's deterministic metrics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus — V2 dimensional mode (country-conditioned personas)\n",
    "\n",
    "Since June 2026 the Dataset Engine supports an opt-in **dimensional mode**: personas are\n",
    "sampled from published psychometric norms (Big Five calibrated to Schmitt 2007 per country,\n",
    "Schwartz values, attachment) instead of fixed archetypes, and conversations carry a full\n",
    "`psychometric` layer per persona. Same API, three extra fields:\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Same job as above, but with country-conditioned dimensional personas (GB)\n",
    "job_resp = requests.post(\n",
    "    f\"{BASE_URL}/jobs/dataset\",\n",
    "    headers=HEADERS,\n",
    "    json={\n",
    "        \"scenario_id\": \"PRO-04\",     # B2B procurement negotiation (new)\n",
    "        \"conversation_count\": 10,\n",
    "        \"complexity\": 3,\n",
    "        \"language\": \"en\",\n",
    "        \"adapter\": \"flat_jsonl\",\n",
    "        \"seed\": 42,\n",
    "        \"mode\": \"dimensional\",      # opt-in V2\n",
    "        \"country\": \"GB\",            # ES US GB DE FR IT MX BR\n",
    "        # \"country_b\": \"US\",        # optional: cross-cultural pair\n",
    "    }\n",
    ").json()\n",
    "print(job_resp[\"data\"][\"jobId\"])\n",
    "# The public Cross-Cultural Negotiation Benchmark was generated exactly this way:\n",
    "# https://huggingface.co/datasets/StrataSynth/stratasynth-cross-cultural-negotiation\n"
   ]
  }
 ]
}