harness/simulate.py

"""Build simulated users and ordered sessions from BFCL tasks.

The persistence/lifecycle claim (ToolSpec's own flagged gap) is about a memory
that must keep absorbing tool-use patterns *introduced after* the datastore was
first built. A frozen datastore cannot contain tools/workflows it never saw; an
online personal memory incorporates them. To test this we introduce each user's
signature tasks PROGRESSIVELY across sessions: roughly half are "known" (present
at warmup, so the frozen static datastore has them) and half are "novel" (first
appear only in later sessions, so the frozen store has never seen them). Novel
tasks recur after their introduction -- the recurring, evolving per-user usage a
persistent memory is meant to exploit.
"""
from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Any

from .data import Task, perturb_numeric


@dataclass
class Instance:
    user_id: str
    session: int          # 0-indexed session number
    query: str            # (possibly perturbed) natural-language request
    functions: list[dict[str, Any]]
    signature_id: str     # which signature task this instance came from
    novel: bool           # True if this task was introduced after warmup


def build_users(tasks: list[Task], n_users: int, tasks_per_user: int,
                n_sessions: int, queries_per_session: int,
                warmup_frac: float = 0.5, seed: int = 0,
                perturb_prob: float = 1.0, arrival: str = "spread",
                novel_weight: float = 3.0, overlap_frac: float = 0.0,
                user_consistent: bool = False) -> list[Instance]:
    """Assign disjoint signature tasks to users and roll out their sessions.

    Each signature task j gets an introduction session: the first
    ``n_known = round(tasks_per_user * warmup_frac)`` are introduced at session
    0 (warmup); the remaining "novel" tasks are introduced one-by-one across the
    later sessions. A session issues queries drawn only from tasks introduced so
    far, biased toward the most recently introduced (novel) task so the novel
    load actually shows up in the metric.
    """
    rng = random.Random(seed)
    # Separate rng for the perturb coin so perturb_prob=1.0 reproduces the
    # legacy stream bit-for-bit (no extra draws on the main rng).
    coin = random.Random(seed + 991)
    pool = tasks[:]
    rng.shuffle(pool)

    # Phase 4.3: shared-user-task overlap. overlap_frac=0.0 is the legacy
    # fully-disjoint assignment (byte-identical path below); >0 makes
    # n_shared = round(tasks_per_user*overlap_frac) of every user's signature
    # tasks come from ONE common pool that all users reuse (same templates,
    # per-occurrence perturbation), the rest from disjoint private blocks. This
    # stress-tests personalization: at high overlap a global store sees many
    # users' calls for the same template.
    overlap_mode = user_consistent or overlap_frac > 0.0
    if overlap_mode:
        n_shared = min(tasks_per_user, round(tasks_per_user * overlap_frac))
        n_private = tasks_per_user - n_shared
        need = tasks_per_user + n_users * n_private   # shared pool + privates
        if len(pool) < need:
            raise ValueError(f"need {need} tasks, have {len(pool)}")
        shared_pool = pool[:tasks_per_user][:n_shared]
        priv_pool = pool[tasks_per_user:]
    else:
        need = n_users * tasks_per_user
        if len(pool) < need:
            raise ValueError(f"need {need} tasks, have {len(pool)}")

    n_known = max(1, round(tasks_per_user * warmup_frac))
    n_novel = tasks_per_user - n_known
    # introduction session for each novel task, spread over sessions 1..n-1
    if n_novel > 0:
        if arrival == "burst":            # all novel tasks arrive at session 2
            novel_intro = [2] * n_novel
        elif arrival == "late":           # all arrive late (session n-3)
            novel_intro = [max(1, n_sessions - 3)] * n_novel
        else:                              # "spread" (legacy): one per session
            step = max(1, (n_sessions - 1) // (n_novel + 1))
            novel_intro = [min(n_sessions - 1, 1 + step * (k + 1))
                           for k in range(n_novel)]
    else:
        novel_intro = []
    intro = [0] * n_known + novel_intro     # per signature-task intro session

    instances: list[Instance] = []
    cursor = 0
    for u in range(n_users):
        uid = f"user_{u:02d}"
        user_fixed_q = None
        if overlap_mode:
            priv = priv_pool[u * n_private:(u + 1) * n_private]
            sig = list(shared_pool) + list(priv)   # shared slots first
            # Each user gets a FIXED argument realization per signature task
            # (user A always NYC, user B always Boston), consistent across the
            # user's sessions but differing across users — the premise
            # personalization exploits. Shared templates thus map to different
            # concrete queries per user; a global store mixes them.
            urng = random.Random(seed * 100003 + u)
            user_fixed_q = [perturb_numeric(t.query, urng) for t in sig]
        else:
            sig = pool[cursor:cursor + tasks_per_user]
            cursor += tasks_per_user
        novel_flags = [False] * n_known + [True] * n_novel
        for s in range(n_sessions):
            active = [(sig[j], novel_flags[j], intro[j])
                      for j in range(tasks_per_user) if intro[j] <= s]
            # map each active signature task back to its slot index (for the
            # per-user fixed-query lookup in the overlap experiment)
            active_idx = [j for j in range(tasks_per_user) if intro[j] <= s]
            if s == 0:
                chosen = [(j, novel_flags[j]) for j in active_idx]  # once each
            else:
                weights = [novel_weight if intro[j] == s else 1.0
                           for j in active_idx]
                chosen = []
                for _ in range(queries_per_session):
                    j = rng.choices(active_idx, weights=weights, k=1)[0]
                    chosen.append((j, novel_flags[j]))
            for j, nv in chosen:
                t = sig[j]
                if user_fixed_q is not None:
                    q = user_fixed_q[j]                 # user-consistent args
                elif s == 0:
                    q = t.query
                elif perturb_prob >= 1.0:
                    q = perturb_numeric(t.query, rng)   # legacy path, exact
                elif perturb_prob <= 0.0:
                    q = t.query
                else:
                    q = (perturb_numeric(t.query, rng)
                         if coin.random() < perturb_prob else t.query)
                instances.append(Instance(uid, s, q, t.functions, t.id, nv))
    return instances