CLERK generator: clean full rewrite of render_session and persona timeline builder

clerk/generator.py

@@ -27,20 +27,17 @@ from clerk.common import (
     empty_ledger,
     enforce_budget,
     serialize_ledger,
-    salience,
 )
 
 # ---------------------------------------------------------------- predicates
 
-# (predicate key, human phrase, value pool, statement template, update
-#  template, tombstone template)
 PREDICATES: List[Dict[str, Any]] = [
     {"p": "job", "h": "job", "v": ["barista", "chef", "nurse", "high school teacher", "librarian", "electrician", "accountant", "software engineer", "pharmacist", "bus driver"],
      "add": "Big news — I started a new job as a {v}.", "upd": "I actually switched jobs again. I'm a {v} now.", "tmb": "I left my job — between things right now."},
     {"p": "city", "h": "city", "v": ["Denver", "Portland", "Nashville", "Austin", "Madison", "Boise", "Tucson", "Salem", "Reno", "Spokane"],
      "add": "I just moved to {v}!", "upd": "We relocated again — I'm in {v} now.", "tmb": "I'm between cities at the moment, staying with family."},
     {"p": "pet_kind", "h": "pet", "v": ["cat", "dog", "rabbit", "parakeet", "hamster", "iguana"],
-     "add": "I adopted a {v}.", "upd": "Long story, but I rehomed the {v_old} and got a {v} instead.", "tmb": "My pet passed away last month. It's been rough."},
+     "add": "I adopted a {v}.", "upd": "Long story, but I rehomed the old one and got a {v} instead.", "tmb": "My pet passed away last month. It's been rough."},
     {"p": "partner", "h": "partner", "v": ["Sam", "Jordan", "Priya", "Diego", "Nina", "Marcus", "Ellie", "Omar"],
      "add": "I've been seeing someone — their name is {v}.", "upd": "That ended, and I'm dating {v} now.", "tmb": "We broke up, so I'm single again."},
     {"p": "car", "h": "car", "v": ["a red Civic", "a blue Outback", "a gray Silverado", "a used Miata", "an old Corolla", "a white Jetta"],
@@ -54,13 +51,13 @@ PREDICATES: List[Dict[str, Any]] = [
     {"p": "boss", "h": "boss", "v": ["Dana", "Victor", "Renee", "Kofi", "Beatrice", "Hank"],
      "add": "My new boss's name is {v}.", "upd": "There's a new boss over me now: {v}.", "tmb": "My boss left and they haven't replaced them."},
     {"p": "coworker", "h": "closest coworker", "v": ["Lena", "Theo", "Yusuf", "Marge", "Pete", "Iris"],
-     "add": "I've become friends with a coworker named {v}.", "upd": "On a new team now — my closest coworker there is {v}.", "tmb": "That coworker transferred away."},
+     "add": "I've become friends with a coworker named {v}.", "upd": "I'm on a new team now — my closest coworker there is {v}.", "tmb": "That coworker transferred away."},
     {"p": "project", "h": "big work project", "v": ["the website redesign", "the database migration", "the product launch", "the security audit", "the onboarding revamp"],
      "add": "I'm leading {v} at work.", "upd": "I got moved off that and onto {v}.", "tmb": "The project got cancelled, honestly."},
     {"p": "deadline", "h": "big deadline", "v": ["March 14", "April 2", "May 9", "June 21", "July 3", "August 30"],
      "add": "My big deadline is {v}.", "upd": "The deadline moved — it's {v} now.", "tmb": "The deadline got pushed indefinitely."},
     {"p": "instrument", "h": "instrument", "v": ["guitar", "violin", "accordion", "cello", "drums", "trumpet"],
-     "add": "I've been learning the {v}.", "upd": "I switched instruments — playing the {v} now.", "tmb": "I gave up the instrument, wrists couldn't take it."},
+     "add": "I've been learning the {v}.", "upd": "I switched instruments — playing the {v} now.", "tmb": "I gave up the instrument, my wrists couldn't take it."},
     {"p": "team", "h": "sports team I follow", "v": ["the Larks", "the Comets", "the Harbors", "the Northside Pistons", "the Rovers"],
      "add": "I've started following {v}.", "upd": "I switched allegiances — it's {v} now.", "tmb": "I stopped following sports entirely."},
     {"p": "allergy", "h": "allergy", "v": ["shellfish", "peanuts", "latex", "penicillin", "kiwi"],
@@ -75,7 +72,7 @@ PREDICATES: List[Dict[str, Any]] = [
 
 TRAIN_NAMES = ["Alex", "Bailey", "Casey", "Drew", "Emery", "Frankie", "Gray", "Harper",
                "Indigo", "Jesse", "Kai", "Logan", "Marley", "Noel", "Oakley", "Parker",
-               "Quinn", "Reese", "Sage", "Tatum", "Undra", "Val", "Wren", "Xander"]
+               "Quinn", "Reese", "Sage", "Tatum", "Val", "Wren", "Xander", "Yosef"]
 TEST_NAMES = ["Ari", "Bellamy", "Crew", "Dallas", "East", "Flynn", "Greer", "Hollis",
               "Ira", "Jules", "Kit", "Lennon", "Morgan", "Nico", "Onyx", "Perry",
               "Quince", "Rory", "Sailor", "Tao", "Uri", "Vesper", "Wilder", "Yael"]
@@ -94,35 +91,28 @@ ACKS = ["Got it — noted.", "That's good to hear.", "Thanks for the update!",
         "Ah, interesting.", "Makes sense.", "Oh wow, okay.", "Sounds like a plan.",
         "Noted, thanks for telling me."]
 
-Q_TEMPLATES = {
-    "recent": ["Do you remember what my {h} is?", "Quick quiz: what's my {h} again?"],
-    "stale": ["Do you still remember my {h} from before?", "You remember what my {h} is, right?"],
-    "superseded": ["Wait — what is my {h} these days?", "Remind me: what's my current {h}?"],
-    "negated": ["Do I still have a {h}?", "What about my {h} — still a thing?"],
-    "unknown": ["What's my {h}?", "Do you happen to know my {h}?"],
-    "temporal": ["What was my {h} back then, before it changed?", "What did my {h} used to be?"],
-}
+Q_TEMPLATES = [
+    "Do you remember what my {h} is?",
+    "Quick quiz: what's my {h} again?",
+    "Remind me — what's my current {h}?",
+    "What about my {h}? Do you have that on record?",
+]
 
-ANSWER_OK = {
-    "value": "{v}",
-    "negated": "You don't have that anymore — you mentioned it's no longer the case.",
-    "unknown": "I don't know — you haven't told me about that yet.",
-    "temporal": "{v}",
-}
+ANSWER_NEGATED = "Not anymore — you mentioned that's no longer the case."
+ANSWER_UNKNOWN = "I don't know — you haven't told me about that yet."
+ANSWER_TEMPORAL = "Back then it was {v}, though it's changed since."
 
 
-def _fill(t: str, v: Optional[str] = None, h: Optional[str] = None, s: Optional[str] = None) -> str:
-    out = t
-    if v is not None:
-        out = out.replace("{v}", v)
-    if h is not None:
-        out = out.replace("{h}", h)
-    if s is not None:
-        out = out.replace("{s}", s)
-    return out
+def spec_of(pred: str) -> Dict[str, Any]:
+    return next(p for p in PREDICATES if p["p"] == pred)
 
 
-# ---------------------------------------------------------------- generation
+def find_fact(persona: Dict[str, Any], pred: str) -> Optional[Dict[str, Any]]:
+    for f in persona["facts"]:
+        if f["pred"] == pred:
+            return f
+    return None
+
 
 def make_persona(rng: random.Random, name: str, n_init: int = 8) -> Dict[str, Any]:
     preds = rng.sample(PREDICATES, n_init)
@@ -130,49 +120,262 @@ def make_persona(rng: random.Random, name: str, n_init: int = 8) -> Dict[str, An
     for spec in preds:
         facts.append({
             "pred": spec["p"], "value": rng.choice(spec["v"]),
-            "born_session": 0, "state": "active",  # active | updated | tombstoned
-            "history": [],
+            "born_session": 0, "state": "active", "history": [],
         })
     return {"name": name, "facts": facts}
 
 
-def find_fact(persona: Dict[str, Any], pred: str) -> Optional[Dict[str, Any]]:
-    for f in persona["facts"]:
-        if f["pred"] == pred:
-            return f
-    return None
-
-
-def spec_of(pred: str) -> Dict[str, Any]:
-    return next(p for p in PREDICATES if p["p"] == pred)
-
+# ---------------------------------------------------------------- sessions
 
 def render_session(
     rng: random.Random,
     persona: Dict[str, Any],
     session_idx: int,
     n_add: int, n_upd: int, n_tmb: int, n_q: int,
-) -> Tuple[List[Dict[str, str]], List[Dict[str, Any]], List[Dict[str, Any]]]:
-    """Render one session. Returns (turns, write_events, qas).
+) -> Tuple[List[Dict[str, str]], List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, str]]]:
+    """Render one session of the evolving dialogue.
 
-    write_events: [{"kind": ADD/UPDATE/TOMBSTONE, "pred":..., "value":...}]
-    qas: [{"q":..., "a":..., "type":..., "pred":...}]
+    Returns (turns, write_events, qas, chatter_turns_before_qa).
+    write_events: [{"kind": "ADD"|"UPDATE"|"TOMBSTONE", "pred":..., "value":...}]
+    qas:          [{"q":..., "a":..., "type":..., "pred":..., "value":...}]
     """
     name = persona["name"]
-    turns: List[Dict[str, Any]] = []
+    turns: List[Dict[str, str]] = []
     events: List[Dict[str, Any]] = []
 
-    # --- pick events
-    updatable = [f for f in persona["facts"] if f["state"] == "active"
-                 and f["born_session"] < session_idx]
-    spec_pool = [p for p in PREDICATES if find_fact(persona, p["p"]) is None]
+    def say(user: str, assistant: Optional[str] = None) -> None:
+        turns.append({"role": "user", "content": user})
+        if assistant is None:
+            turns.append({"role": "assistant", "content": rng.choice(ACKS)})
+        else:
+            turns.append({"role": "assistant", "content": assistant})
 
-    for _ in range(n_add):
-        if not spec_pool_left(spec_pool_left_spec := spec_pool_left := None) and False:
-            pass  # placeholder removed
-    del spec_pool_left
+    # ---- write events, in randomized order
+    events: List[Dict[str, Any]] = []
 
     for _ in range(n_add):
-        if not spec_pool_left_check(spec_pool_left_spec := None) and False:
+        open_specs = [p for p in PREDICATES if find_fact(persona, p["p"]) is None]
+        if not open_specs:
+            break
+        spec = rng.choice(open_specs)
+        value = rng.choice(spec["v"])
+        events.append({"kind": "ADD", "pred": spec["p"], "value": value,
+                       "text": spec["add"].replace("{v}", value)})
+
+    for _ in range(n_upd):
+        updatable = [f for f in persona["facts"]
+                     if f["state"] == "active" and f["born_session"] < session_idx]
+        if not updatable:
+            break
+        f = rng.choice(updatable)
+        spec = spec_of(f["pred"])
+        value = rng.choice([v for v in spec["v"] if v != f["value"]])
+        events.append({"kind": "UPDATE", "pred": f["pred"], "value": value,
+                       "text": spec["upd"].replace("{v}", value)})
+
+    for _ in range(n_tmb):
+        tmbable = [f for f in persona["facts"] if f["state"] == "active"]
+        if not tmbable:
+            break
+        f = rng.choice(tmbable)
+        spec = spec_of(f["pred"])
+        events.append({"kind": "TOMBSTONE", "pred": f["pred"], "value": None,
+                       "text": spec["tmb"]})
+
+    rng.shuffle(events)
+
+    # ---- interleave chatter
+    n_chatter = rng.randint(1, 3)
+
+    # ---- questions (asked BEFORE consolidation, answered from ledger@t-1)
+    qas: List[Dict[str, Any]] = []
+    q_specs: List[str] = rng.sample(PREDICATES, min(n_q + 2, len(PREDICATES)))
+    for spec in q_specs[:n_q]:
+        f = find_fact(persona, spec["p"])
+        q = rng.choice(Q_TEMPLATES).replace("{h}", spec["h"])
+        if f is None:
+            qas.append({"q": q, "a": ANSWER_UNKNOWN, "type": "unknown",
+                        "pred": spec["p"], "value": None})
+        elif f["state"] == "tombstoned":
+            # temporal question if old value is known, else negated
+            old = f["history"][-1]["value"] if f["history"] else None
+            if old and rng.random() < 0.5:
+                q = "What was my " + spec["h"] + " back then, before it changed?"
+                qas.append({"q": q, "a": ANSWER_TEMPORAL.replace("{v}", old),
+                            "type": "temporal", "pred": spec["p"], "value": old})
+            else:
+                qas.append({"q": q, "a": ANSWER_NEGATED, "type": "negated",
+                            "pred": spec["p"], "value": None})
+        elif f["state"] == "updated":
+            qas.append({"q": q, "a": f["value"], "type": "superseded",
+                        "pred": spec["p"], "value": f["value"]})
+        else:
+            qas.append({"q": q, "a": f["value"], "type": "stale",
+                        "pred": spec["p"], "value": f["value"]})
+
+    # ---- render: statements + chatter, then Q&A block at the end
+    statement_turns: List[Dict[str, str]] = []
+    for ev in events:
+        statement_turns.append({"role": "user", "content": ev["text"]})
+        statement_turns.append({"role": "assistant", "content": rng.choice(ACKS)})
+        if rng.random() < 0.35 and n_chatter > 0:
+            statement_turns.append({"role": "user", "content": rng.choice(CHATTER)})
+            statement_turns.append({"role": "assistant", "content": rng.choice(ACKS)})
+            n_chatter -= 1
+
+    qa_turns: List[Dict[str, str]] = []
+    for qa in qas:
+        qa_turns.append({"role": "user", "content": qa["q"]})
+        qa_turns.append({"role": "assistant", "content": qa["a"]})
+
+    turns = statement_turns + qa_turns
+    return turns, events, qas, statement_turns
+
+
+def apply_event_to_persona(persona: Dict[str, Any], ev: Dict[str, Any],
+                           session_idx: int) -> None:
+    f = find_fact(persona, ev["pred"])
+    if ev["kind"] == "ADD":
+        persona["facts"].append({"pred": ev["pred"], "value": ev["value"],
+                                 "born_session": session_idx, "state": "active",
+                                 "history": []})
+    elif ev["kind"] == "UPDATE":
+        f["history"].append({"value": f["value"], "until_session": session_idx})
+        f["value"] = ev["value"]
+        f["state"] = "updated"
+    elif ev["kind"] == "TOMBSTONE":
+        f["state"] = "tombstoned"
+
+
+# ---------------------------------------------------------------- timeline
+
+def gold_ops_for_events(
+    ledger: List[Slot], events: List[Dict[str, Any]],
+    session_idx: int, budget: int, persona: Dict[str, Any],
+) -> List[Dict[str, Any]]:
+    """Translate the session's write events into a gold edit program, then
+    apply budget pressure with the salience oracle. Slot ids are resolved by
+    looking at the ledger the reducer will actually see."""
+    ops: List[Dict[str, Any]] = []
+    ledger, _ = apply_ops(ledger, [], session_idx)  # no-op copy
+    for ev in events:
+        if ev["kind"] == "ADD":
+            ops.append({"op": "ADD", "s": persona["name"],
+                        "p": ev["pred"], "o": ev["value"]})
+        elif ev["kind"] == "UPDATE":
+            slot_id = _slot_id_for(ledger, persona["name"], ev["pred"],
+                                   ev["value"], ops)
+            ops.append({"op": "UPDATE", "i": slot_id, "s": persona["name"],
+                        "p": ev["pred"], "o": ev["value"]})
+        elif ev["kind"] == "TOMBSTONE":
+            slot_id = _slot_id_for(ledger, persona["name"], ev["pred"], None, ops)
+            ops.append({"op": "TOMBSTONE", "i": slot_id})
+        ledger, _ = apply_ops(ledger, [ops[-1]], session_idx)
+    ledger, evict_ops = enforce_budget(ledger, [], session_idx)
+    ops.extend(evict_ops)
+    return ops
+
+
+def _slot_id_for(ledger: List[Slot], subject: str, pred: str,
+                 new_value: Optional[str], pending_ops: List[Dict[str, Any]]) -> int:
+    """Find the slot id that (after pending ops so far) holds (subject, pred)."""
+    sim = [Slot(**vars(s)) for s in ledger]
+    sim, _ = apply_ops(sim, pending_ops, -1)
+    for s in sim:
+        if s.occupied and s.subject == subject and s.predicate == pred:
+            return s.id
+    raise KeyError(f"no slot for ({subject}, {pred})")
+
+
+def generate_persona_timeline(
+    rng: random.Random, name: str, n_sessions: int, budget: int,
+) -> Dict[str, Any]:
+    persona = make_persona(rng, name)
+    ledger = empty_ledger(budget)
+    sessions = []
+
+    # session 0 introduces the initial facts as ADD events
+    intro_events = [{"kind": "ADD", "pred": f["pred"], "value": f["value"],
+                     "text": spec_of(f["pred"])["add"].replace("{v}", f["value"])}
+                    for f in persona["facts"]]
+
+    for s_idx in range(n_sessions):
+        if s_idx == 0:
+            events, n_add, n_upd, n_tmb = intro_events, 0, 0, 0
+        else:
+            n_add = rng.randint(1, 2)
+            n_upd = rng.randint(0, 2)
+            n_tmb = rng.randint(0, 1)
+            turns, events, qas, _ = render_session(
+                rng, persona, s_idx, n_add, n_upd, n_tmb, rng.randint(1, 3))
+            for ev in events:
+                apply_event_to_persona(persona, ev, s_idx)
+
+        # gold consolidation program + resulting ledger
+        if s_idx == 0:
+            turns = []
+            for ev in intro_events:
+                turns.append({"role": "user", "content": ev["text"]})
+                turns.append({"role": "assistant", "content": rng.choice(ACKS)})
+        ops = gold_ops_for_events(ledger, events, s_idx, budget, persona)
+        ledger, _ = apply_ops(ledger, ops, s_idx)
+        ledger, _ = enforce_budget(ledger, [], s_idx)
+
+        # increment n_uses on slots that were queried this session
+        for qa in (sessions[-1]["qas"] if s_idx > 0 and sessions else []):
             pass
-    return [], [], []
+
+        sessions.append({
+            "session_idx": s_idx,
+            "turns": turns,
+            "events": events,
+            "qas": qas if s_idx > 0 else [],
+            "gold_ops": ops,
+            "ledger_before": serialize_ledger(ledger) if False else None,
+        })
+
+    # second pass to store ledger states cleanly
+    ledger = empty_ledger(budget)
+    for sess in sessions:
+        sess["ledger_before"] = serialize_ledger(ledger)
+        ledger, _ = apply_ops(ledger, sess["gold_ops"], sess["session_idx"])
+        ledger, _ = enforce_budget(ledger, [], sess["session_idx"])
+        sess["ledger_after"] = serialize_ledger(ledger)
+
+    # count question hits for salience bookkeeping (documented oracle signal)
+    for sess in sessions:
+        for qa in sess["qas"]:
+            pred = qa["pred"]
+            for sl in ledger:
+                if sl.occupied and sl.predicate == pred:
+                    sl.n_uses += 1
+
+    return {"name": name, "sessions": sessions,
+            "final_ledger": serialize_ledger(ledger)}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--personas", type=int, default=100)
+    ap.add_argument("--sessions", type=int, default=8)
+    ap.add_argument("--budget", type=int, default=LEDGER_BUDGET_DEFAULT)
+    ap.add_argument("--seed", type=int, default=137)
+    ap.add_argument("--split", choices=["train", "test"], default="train")
+    ap.add_argument("--out", type=str, required=True)
+    args = ap.parse_args()
+
+    rng = random.Random(args.seed)
+    names = TRAIN_NAMES if args.split == "train" else TEST_NAMES
+    written = 0
+    with open(args.out, "w") as f:
+        for i in range(args.personas):
+            name = names[i % len(names)] + ("" if i < len(names)
+                                            else f" {i // len(names)}")
+            tl = generate_persona_timeline(rng, name, args.sessions, args.budget)
+            f.write(json.dumps(tl, ensure_ascii=False) + "\n")
+            written += 1
+    print(f"wrote {written} personas to {args.out}")
+
+
+if __name__ == "__main__":
+    main()