From b1e301d95667de7a0998908bee18711541f3c2c7 Mon Sep 17 00:00:00 2001 From: isarasua Date: Sat, 10 Oct 2026 16:12:39 +0200 Subject: [PATCH 1/2] feat: opt-in replay of the assistant's own earlier reasoning Before every request, User Sim drops each earlier turn's reasoning. A model trained to receive its own earlier reasoning ("preserved thinking") is then evaluated on a history it would never see in deployment. `replay_reasoning = true` on `assistant_model` in the models file now sends the assistant the reasoning of its own earlier messages, tool steps included. Off by default, so requests, results and trajectory ids are unchanged. - Only the assistant takes the switch: its requests are the only ones that hold its own earlier messages. During a simulation every other model, the simulated user above all, sees the conversation as text and never receives the reasoning. - It becomes `ConversationSimulatorConfig.replay_assistant_reasoning`, read through a context variable that `core.probes.dispatch_episode` installs around every episode. The generator, the hosted runtime and the probes' `simulate_` shortcuts all run their episodes through it, so every path replays alike, and the compression decision reads the same setting. - With replay on, the assistant's history is not compressed, so it sees what it produced rather than a summary beside the original reasoning. - A setup that cannot replay is refused before any model call: the switch on another alias, replay without stored reasoning (also refused by the config itself), or a provider that is not OpenAI-compatible (Data Designer's Anthropic adapter drops reasoning). - A run that replays gets its own trajectory ids, like any other model change (so tool_calling draws its own tools), and records the switch in its manifest, from the CLI and the simulate notebook. Closes #23 Co-Authored-By: Claude Opus 5.5 Signed-off-by: isarasua --- architecture/cli.md | 16 ++ docs/engine/AUTHORING_A_PROBE.md | 6 +- docs/engine/EXTERNAL_PROBE_RUNTIME.md | 8 + docs/engine/README.md | 5 +- docs/outputs.md | 2 +- notebooks/01_simulate.ipynb | 3 + src/usersim/cli/_models.py | 89 +++++- src/usersim/cli/_pipeline.py | 9 +- src/usersim/cli/models_default.toml | 8 + src/usersim/cli/simulate.py | 13 +- src/usersim/engine/config.py | 25 +- src/usersim/engine/core/context.py | 10 +- src/usersim/engine/core/episode_input.py | 11 +- src/usersim/engine/core/episode_runtime.py | 10 +- src/usersim/engine/core/llm.py | 52 +++- src/usersim/engine/core/manifest.py | 3 + src/usersim/engine/core/probes.py | 23 ++ src/usersim/engine/core/simulation.py | 6 +- src/usersim/engine/generator.py | 4 +- .../probes/financial_services/generator.py | 3 +- .../probes/general_educational/generator.py | 4 +- .../probes/general_open_ended/generator.py | 4 +- .../engine/probes/safety_agentic/generator.py | 3 +- .../probes/safety_chat_pressure/generator.py | 3 +- .../engine/probes/sov_ai_dynamic/generator.py | 3 +- .../engine/probes/sov_ai_facts/generator.py | 3 +- .../sov_ai_multilingual_parity/generator.py | 3 +- .../engine/probes/tool_calling/generator.py | 3 +- src/usersim/reporting/capability_report.py | 9 +- src/usersim/reporting/dashboard.py | 8 +- src/usersim/testing/__init__.py | 4 + tests/engine/core/test_llm.py | 72 ++++- tests/engine/core/test_reasoning_replay.py | 249 ++++++++++++++++ tests/test_replay_reasoning_config.py | 272 ++++++++++++++++++ 34 files changed, 894 insertions(+), 52 deletions(-) create mode 100644 tests/engine/core/test_reasoning_replay.py create mode 100644 tests/test_replay_reasoning_config.py diff --git a/architecture/cli.md b/architecture/cli.md index ff3bd47..b221909 100644 --- a/architecture/cli.md +++ b/architecture/cli.md @@ -49,6 +49,22 @@ top_p, max_tokens, and any model-family-specific `extra_body`. A TOML row overrides the catalogue per field, and the row is self-describing for anything it sets. +`replay_reasoning = true` on the `assistant_model` row sends the assistant +the reasoning of its own earlier messages, for a model trained to receive it. +It becomes `ConversationSimulatorConfig.replay_assistant_reasoning`, because +Data Designer's model configs take no extra fields. A setup that cannot +replay fails at config construction: +- the switch on another alias +- `--no-store-reasoning` +- a provider that is not OpenAI-compatible + +A model override that changes the assistant's model turns it off unless the +override sets it. + +It only helps where the endpoint passes `reasoning_content` on earlier +assistant messages through to the model; a server or chat template that +drops earlier thinking makes it a no-op. + ## Failure handling `usersim.cli._errors.ConfigError` is a `ValueError` subclass, never a diff --git a/docs/engine/AUTHORING_A_PROBE.md b/docs/engine/AUTHORING_A_PROBE.md index 0e2fc61..5a0fd66 100644 --- a/docs/engine/AUTHORING_A_PROBE.md +++ b/docs/engine/AUTHORING_A_PROBE.md @@ -580,7 +580,7 @@ from usersim.engine.core.cooking_bank import ( ) from usersim.engine.core.outcomes import WarningKind from usersim.engine.core.probes import ( - BankBackedProbe, BankVerbatimMixin, register_probe, + BankBackedProbe, BankVerbatimMixin, dispatch_episode, register_probe, ) from usersim.engine.probes.cooking_advisor.prompts import ( COOKING_GATE_PROMPT, COOKING_USER_PROMPTS, @@ -672,7 +672,9 @@ async def simulate_cooking_advisor( cfg=cfg, provenance=provenance, profile=profile, data=data, outcome_builder=outcome_builder, ) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + # dispatch_episode installs the run's per-episode settings (reasoning + # replay) around run_dispatch, as the generator does. + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) ``` ### File 6: built-in probe list diff --git a/docs/engine/EXTERNAL_PROBE_RUNTIME.md b/docs/engine/EXTERNAL_PROBE_RUNTIME.md index 7cfe954..b816234 100644 --- a/docs/engine/EXTERNAL_PROBE_RUNTIME.md +++ b/docs/engine/EXTERNAL_PROBE_RUNTIME.md @@ -83,6 +83,14 @@ exposing `agenerate_text_embeddings()`); without one, dense retrieval falls back to lexical search with a warning, so results can differ from a standalone run that has one configured. +### Replaying the assistant's reasoning + +When the row's `usersim_config` sets `replay_assistant_reasoning` (from +`replay_reasoning = true` on `assistant_model` in the models file), each +assistant request's earlier assistant messages carry their +`reasoning_content`. Forward them to the model as they are. No other role's +messages ever carry the assistant's reasoning. Replay only helps where the endpoint passes `reasoning_content` on earlier assistant messages through to the model; a server or chat template that drops earlier thinking makes it a no-op. + ## Driving the episode `advance()` is the whole loop. The first call takes no result; every later call diff --git a/docs/engine/README.md b/docs/engine/README.md index 0195115..6b48774 100644 --- a/docs/engine/README.md +++ b/docs/engine/README.md @@ -450,7 +450,8 @@ diagnostics in the simulator-health report quantify that for a given run, and | `persona_grounding_ratio` | float | `1.0` | Fraction of trajectories whose first user turn incorporates persona-specific details. Lower it for a generic-query control arm | | `context_compression` | bool | `True` | Summarize older assistant responses to bound context growth | | `compression_window` | int | `1` | Keep last N assistant responses verbatim before compressing | -| `store_reasoning` | bool | `True` | Keep the assistant's thinking trace on the assistant message when the model emits one (`--no-store-reasoning` to turn off). Gated at capture, so a suppressed trace is never attached. Independent of whether the model *produces* a trace: that is a model-side setting (`chat_template_kwargs` in the models TOML, plus a server-side reasoning parser). Never affects what a model receives: traces are stored, never replayed | +| `store_reasoning` | bool | `True` | Keep the assistant's thinking trace on the assistant message when the model emits one (`--no-store-reasoning` to turn off). Gated at capture, so a suppressed trace is never attached. Independent of whether the model *produces* a trace: that is a model-side setting (`chat_template_kwargs` in the models TOML, plus a server-side reasoning parser). A trace reaches a model again only through `replay_assistant_reasoning` | +| `replay_assistant_reasoning` | bool | `False` | Send the assistant model the reasoning of its own earlier messages (`replay_reasoning = true` on `assistant_model` in the models TOML). Needs `store_reasoning`; the assistant's history is then not compressed. No other model receives the assistant's reasoning during a simulation | | `verbosity` | int | `1` | Logging detail: 0=quiet, 1=normal, 2=detailed | | `random_seed` | int? | `None` | Seed for reproducible tool sampling | @@ -469,7 +470,7 @@ what each one carries: | `simulation_traces` | str? (JSON) | Per-turn / per-call telemetry (keyed `turn_idx` x `call_idx`) with `kind` discriminator: `user_query_gate`, `user_followup_gate`, `fourth_wall_prefilter`, `assistant_quality`, `early_stop`, `api_response_reroll`, `tool_call_verifier`, etc. The debugging surface: **not** the assistant-evaluation surface. | | `probe_family` | str | Probe module's `PROBE_FAMILY` constant (e.g., `"general_open_ended"`, `"sov_ai_facts"`). | | `probe_variant` | str | Probe-specific discriminator (first entry of `PROBE_VARIANTS`, or an explicit row value when provided). Part of `trajectory_id` and the evaluator's stratification key. | -| `conversation_messages` | str (JSON) | Full message list. Each entry is `{role, content}`; assistant entries add `tool_calls` when the turn called a tool, `tool` entries add `tool_call_id`, and an assistant entry adds **`reasoning_content`** when that model call emitted a thinking trace **and `store_reasoning` is on** (the default; `--no-store-reasoning` suppresses it). The key is absent (not empty) whenever either condition fails; the run manifest's `simulation_config.store_reasoning` records which was requested. Traces are attributed per call, so a tool-calling turn that appends several assistant messages carries each trace on the message that produced it. Stored for analysis only: `core/llm.py::_dicts_to_chat_messages` hard-codes the assistant's `reasoning_content` to `None` on the way back into a model, so a later turn never sees an earlier turn's thinking. | +| `conversation_messages` | str (JSON) | Full message list. Each entry is `{role, content}`; assistant entries add `tool_calls` when the turn called a tool, `tool` entries add `tool_call_id`, and an assistant entry adds **`reasoning_content`** when that model call emitted a thinking trace **and `store_reasoning` is on** (the default; `--no-store-reasoning` suppresses it). The key is absent (not empty) whenever either condition fails; the run manifest's `simulation_config.store_reasoning` records which was requested. Traces are attributed per call, so a tool-calling turn that appends several assistant messages carries each trace on the message that produced it. Stored for analysis: `core/llm.py::_dicts_to_chat_messages` drops it on the way back into a model, so a later turn never sees an earlier turn's thinking, unless the run sets `replay_assistant_reasoning`, which sends the assistant its own earlier traces. | | `conversation_metadata` | str (JSON) | Loose per-turn judge / event log surfaced to interactive previews and the deep-dive analysis script. | | `conversation_status` | bool | True if the trajectory ran to completion (no gate exhaustion / no infrastructure failure) AND any optional `should_succeed(state)` hook returned True. | | `behavioral_profile` | str (JSON) | Computed behavioral parameters | diff --git a/docs/outputs.md b/docs/outputs.md index 9c96688..7d55f35 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -9,7 +9,7 @@ report files. `output/trajectories/run=/locale=/probe_family=/`. One row per simulated conversation, each carrying: -- `conversation_messages` (JSON): the full multi-turn message list. Each entry is `{role, content}`, plus `tool_calls` / `tool_call_id` on tool-using turns and `reasoning_content` on an assistant turn whose model emitted a thinking trace, unless `--no-store-reasoning` was passed. The trace is stored for analysis only; it is never replayed into a later turn (see [`docs/engine/README.md`](engine/README.md)). +- `conversation_messages` (JSON): the full multi-turn message list. Each entry is `{role, content}`, plus `tool_calls` / `tool_call_id` on tool-using turns and `reasoning_content` on an assistant turn whose model emitted a thinking trace, unless `--no-store-reasoning` was passed. The trace is stored for analysis. It reaches a model again only when the models file sets `replay_reasoning = true` on `assistant_model`, which sends the assistant the reasoning of its own earlier messages (see [`docs/engine/README.md`](engine/README.md)). - `simulation_outcome` (JSON) holds a row-level structured record: status, failure-class taxonomy, per-actor attribution, per-model token counts, wall-clock, provenance (code SHA + asset bank versions + prompt versions). - `simulation_traces` (JSON): per-turn / per-call telemetry (judge-gate ratings, fourth-wall pre-filter triggers, capitulation checks, tool-call verification, etc.). - A content-hashed `persona_uuid` (stable across runs) and a deterministic `trajectory_id` (hash of `persona + probe + resolved model identities + prompt version`). These are the simulator's idempotency keys. diff --git a/notebooks/01_simulate.ipynb b/notebooks/01_simulate.ipynb index 776b740..bd43774 100644 --- a/notebooks/01_simulate.ipynb +++ b/notebooks/01_simulate.ipynb @@ -80,6 +80,7 @@ "from usersim.cli._models import (\n", " REQUIRED_SIMULATOR_ALIASES,\n", " apply_model_overrides,\n", + " assistant_replays_reasoning,\n", " bundled_models_path,\n", " default_models_path,\n", " describe_models_source,\n", @@ -712,6 +713,8 @@ " context_compression=True,\n", " compression_window=1,\n", " random_seed=RANDOM_SEED,\n", + " # As build_simulator_config_builder sets it for the run itself.\n", + " replay_assistant_reasoning=assistant_replays_reasoning(MODELS, store_reasoning=True),\n", ")\n", "# Describe the RUN, not just this pass. On a resumed run LOCALES holds only the\n", "# locales this pass filled in, so recording it verbatim would understate a run that\n", diff --git a/src/usersim/cli/_models.py b/src/usersim/cli/_models.py index 0d881fc..8cf77a5 100644 --- a/src/usersim/cli/_models.py +++ b/src/usersim/cli/_models.py @@ -97,6 +97,10 @@ class ModelSpec: top_p: float | None = 0.95 timeout: float | None = None extra_body: dict[str, Any] | None = None + #: Send this model the reasoning of its own earlier messages, for a model + #: trained to receive it. Honoured on ``assistant_model`` only; see + #: :func:`assistant_replays_reasoning`. + replay_reasoning: bool = False @dataclass(frozen=True) @@ -255,9 +259,73 @@ def _maybe_float(v: Any) -> float | None: top_p=_maybe_float(pick("top_p", 0.95)), timeout=pick("timeout"), extra_body=pick("extra_body"), + replay_reasoning=_replay_flag(pick("replay_reasoning", False), where=f"{source.name}: alias={alias!r}"), ) +def _replay_flag(value: Any, *, where: str) -> bool: + """``replay_reasoning`` as given, which must be a real boolean. + + A string such as ``"false"`` is refused rather than read as true. + """ + if not isinstance(value, bool): + raise ConfigError(f"{where}: replay_reasoning must be true or false, got {value!r}") + return value + + +def _provider_type(config: ModelsConfig, name: str) -> str | None: + """The ``provider_type`` Data Designer serves provider ``name`` with, if known. + + The config's own providers first, then the defaults Data Designer loads + for itself, as :func:`to_data_designer_kwargs` merges them. + """ + for provider in config.providers: + if provider.name == name: + return provider.provider_type + from data_designer.config.default_model_settings import get_builtin_model_providers, get_default_providers + + try: + defaults = get_default_providers() or get_builtin_model_providers() + except FileNotFoundError: + defaults = get_builtin_model_providers() + return next((provider.provider_type for provider in defaults if provider.name == name), None) + + +def assistant_replays_reasoning(config: ModelsConfig, *, store_reasoning: bool) -> bool: + """Whether a run replays the assistant's own reasoning; refuses a setup that cannot. + + ``replay_reasoning`` is honoured on ``assistant_model`` only: the + assistant's requests are the only ones that hold its own earlier messages, + and a stored message does not record which model produced it. Raises + ``ConfigError`` when the switch is set on another alias, when reasoning is + not stored (there would be nothing to replay), or when the assistant's + provider is not OpenAI-compatible, since only that Data Designer adapter + sends ``reasoning_content`` back (the Anthropic one drops it). + """ + others = sorted(spec.alias for spec in config.models if spec.replay_reasoning and spec.alias != "assistant_model") + if others: + raise ConfigError( + f"replay_reasoning is set on {others}; it applies to assistant_model only, the one model whose " + "requests hold its own earlier messages" + ) + assistant = next((spec for spec in config.models if spec.alias == "assistant_model"), None) + if assistant is None or not assistant.replay_reasoning: + return False + if not store_reasoning: + raise ConfigError( + "replay_reasoning on assistant_model needs the reasoning kept: drop --no-store-reasoning " + "(store_reasoning=False), or turn replay off" + ) + provider_type = _provider_type(config, assistant.provider) + # Data Designer normalizes the type's case before choosing an adapter. + if provider_type is not None and provider_type.lower() != "openai": + raise ConfigError( + f"replay_reasoning on assistant_model needs an OpenAI-compatible provider; {assistant.provider!r} is " + f"{provider_type!r}, and Data Designer does not send reasoning back through it" + ) + return True + + def require_aliases(config: ModelsConfig, *, required: tuple[str, ...], context: str) -> None: """Raise ``ConfigError`` if any required alias is absent.""" have = set(config.aliases()) @@ -275,6 +343,7 @@ def require_aliases(config: ModelsConfig, *, required: tuple[str, ...], context: "top_p", "max_tokens", "extra_body", + "replay_reasoning", } ) @@ -351,7 +420,9 @@ def apply_model_overrides( - **dict** (extended form): per-knob override on top of the catalog. The dict MUST include ``"model"`` and may optionally include any - of ``"temperature"``, ``"top_p"``, ``"max_tokens"``, ``"extra_body"``:: + of ``"temperature"``, ``"top_p"``, ``"max_tokens"``, ``"extra_body"``, + ``"replay_reasoning"``. Replay is kept when the model stays the same and + turned off when it changes, unless the override sets it:: MODEL_OVERRIDES = { "assistant_model": { @@ -575,6 +646,19 @@ def apply_model_overrides( else: new_extra_body = None + # Replay belongs to the model it was set for, not to the alias: a new + # model starts with it off unless the override turns it on. + if "replay_reasoning" in ov: + new_replay = _replay_flag(ov["replay_reasoning"], where=f"model override for alias {spec.alias!r}") + else: + new_replay = spec.replay_reasoning and new_model == spec.model + if spec.replay_reasoning and not new_replay: + logger.warning( + "model override: alias=%r switched to model=%r, so replay_reasoning is off; set it on the " + "override to keep it", + spec.alias, + new_model, + ) new_specs.append( dataclasses.replace( spec, @@ -584,6 +668,7 @@ def apply_model_overrides( top_p=new_top_p, max_tokens=new_max_tokens, extra_body=new_extra_body, + replay_reasoning=new_replay, ) ) return dataclasses.replace(config, models=tuple(new_specs)) @@ -1050,6 +1135,8 @@ def print_resolved_models(config: ModelsConfig) -> None: extras += f" extra_body={spec.extra_body}" if spec.timeout: extras += f" timeout={spec.timeout}" + if spec.replay_reasoning: + extras += " replay_reasoning=True" print( f" {spec.alias:20s} {spec.model:30s} provider={spec.provider:25s} " f"temp={spec.temperature} top_p={spec.top_p} " diff --git a/src/usersim/cli/_pipeline.py b/src/usersim/cli/_pipeline.py index 49bec8d..d6dba11 100644 --- a/src/usersim/cli/_pipeline.py +++ b/src/usersim/cli/_pipeline.py @@ -36,7 +36,12 @@ from typing import Any from usersim.cli._errors import ConfigError -from usersim.cli._models import ModelsConfig, to_data_designer_kwargs, to_model_configs +from usersim.cli._models import ( + ModelsConfig, + assistant_replays_reasoning, + to_data_designer_kwargs, + to_model_configs, +) from usersim.cli._persona_language import ( DEFAULT_PERSONA_DATASETS_DIR, persona_sampler_params, @@ -417,6 +422,7 @@ def build_simulator_config_builder( max_turns=max_turns, max_assistant_attempts=max_assistant_attempts, store_reasoning=store_reasoning, + replay_assistant_reasoning=assistant_replays_reasoning(models, store_reasoning=store_reasoning), context_compression=True, compression_window=1, verbosity=verbosity, @@ -539,6 +545,7 @@ def materialize_episode_inputs( max_turns=max_turns, max_assistant_attempts=max_assistant_attempts, store_reasoning=store_reasoning, + replay_assistant_reasoning=assistant_replays_reasoning(models, store_reasoning=store_reasoning), finance_tier_mix=finance_tier_mix, finance_retrieval_mode=finance_retrieval_mode, **toolset_kwargs, diff --git a/src/usersim/cli/models_default.toml b/src/usersim/cli/models_default.toml index 1d59896..8d5ac8c 100644 --- a/src/usersim/cli/models_default.toml +++ b/src/usersim/cli/models_default.toml @@ -9,6 +9,14 @@ # Per-model sampling parameters live in cli/model_catalog.py. Per-alias knobs # (max_parallel_requests, timeout, max_tokens) go here and win over the catalog. # +# `replay_reasoning = true` on `assistant_model` sends it the reasoning of its +# own earlier messages, for a model trained to receive it ("preserved +# thinking"). Off by default. It needs stored reasoning and an OpenAI-compatible +# provider, and turns off compression of the assistant's history. It only +# helps where the endpoint passes `reasoning_content` on earlier assistant +# messages through to the model; a chat template that drops earlier thinking +# makes it a no-op. +# # Judges are a different model from `assistant_model` on purpose, so the model # under test is not scoring its own output. diff --git a/src/usersim/cli/simulate.py b/src/usersim/cli/simulate.py index 71a144b..8c37bed 100644 --- a/src/usersim/cli/simulate.py +++ b/src/usersim/cli/simulate.py @@ -195,8 +195,9 @@ def register(subparsers: argparse._SubParsersAction) -> None: "Drop the assistant's thinking trace instead of storing it on " "the message. Default is to store it when the model emits one " "(reasoning must also be enabled model-side -- see the models " - "TOML). Traces are never replayed to a model either way; this " - "only controls whether the trajectory carries them." + "TOML). This controls whether the trajectory carries them; a " + "trace reaches a model again only when the models file sets " + "replay_reasoning on assistant_model, which needs them stored." ), ) p.add_argument( @@ -276,6 +277,7 @@ def _parse_probe_mix(spec: str) -> dict: def run(args: argparse.Namespace) -> int: from usersim.cli._models import ( REQUIRED_SIMULATOR_ALIASES, + assistant_replays_reasoning, default_models_path, load_models_config, require_aliases, @@ -285,6 +287,8 @@ def run(args: argparse.Namespace) -> int: models_path = args.models or default_models_path() models = load_models_config(models_path) require_aliases(models, required=REQUIRED_SIMULATOR_ALIASES, context="`usersim simulate`") + # Refused here, before anything runs, rather than at the first locale. + replay_reasoning = assistant_replays_reasoning(models, store_reasoning=not args.no_store_reasoning) probe_mix = _parse_probe_mix(args.probe_mix) @@ -338,6 +342,7 @@ def run(args: argparse.Namespace) -> int: print(f" output: {args.out}") print(f" resumable: {not args.no_skip_existing}") print(f" store_reasoning: {not args.no_store_reasoning}") + print(f" replay_reasoning: {replay_reasoning}") return 0 if (msg := warn_if_missing_api_key(models)) is not None: @@ -454,6 +459,7 @@ def _simulate_to_parquet( max_assistant_attempts=1, store_reasoning=True, ) -> int: + from usersim.cli._models import assistant_replays_reasoning from usersim.cli._persona_language import ( DEFAULT_PERSONA_DATASETS_DIR, log_persona_language_population, @@ -603,6 +609,7 @@ def _simulate_to_parquet( max_turns=max_turns, random_seed=random_seed, store_reasoning=store_reasoning, + replay_assistant_reasoning=assistant_replays_reasoning(models, store_reasoning=store_reasoning), ) try: manifest_path_written = write_simulator_manifest( @@ -633,6 +640,7 @@ def _build_sim_config_for_manifest( max_turns: int, random_seed, store_reasoning: bool = True, + replay_assistant_reasoning: bool = False, ): """Snapshot the sim-config knobs the manifest captures. @@ -651,6 +659,7 @@ def _build_sim_config_for_manifest( compression_window=1, random_seed=random_seed, store_reasoning=store_reasoning, + replay_assistant_reasoning=replay_assistant_reasoning, ) diff --git a/src/usersim/engine/config.py b/src/usersim/engine/config.py index d221523..155146f 100644 --- a/src/usersim/engine/config.py +++ b/src/usersim/engine/config.py @@ -8,7 +8,7 @@ from typing import Literal from data_designer.config.base import SingleColumnConfig -from pydantic import Field +from pydantic import Field, model_validator MODEL_USER = "user_model" MODEL_ASSISTANT = "assistant_model" @@ -125,10 +125,21 @@ class ConversationSimulatorConfig(SingleColumnConfig): # not stored cannot be recovered without re-simulating. Turn it off to # keep trajectories narrow: traces dominate the stored text rather than # merely adding to it, so an eval-only run that will never read them - # pays a large size cost for nothing. Never affects what a model - # receives: traces are stored, never replayed. + # pays a large size cost for nothing. Replaying traces needs them kept + # (``replay_assistant_reasoning``). store_reasoning: bool = True + # Send the assistant model the reasoning of its own earlier messages, for + # a model trained to receive it ("preserved thinking"); set from + # ``replay_reasoning`` on ``assistant_model`` in the models file. Off, + # every earlier turn's reasoning is dropped before a request. During a + # simulation no other model receives the assistant's reasoning: their + # requests show the conversation as text. Needs ``store_reasoning``. The + # assistant's history is not compressed while it is on, so it sees what + # it produced rather than a summary beside the original reasoning. A run + # with it on has its own trajectory ids, like any other model change. + replay_assistant_reasoning: bool = False + # Logging verbosity: 0=quiet, 1=normal (default), 2=detailed per-call timing verbosity: int = 1 @@ -171,6 +182,14 @@ def get_model_aliases(self) -> list[str]: """ return list(REQUIRED_MODEL_ALIASES) + @model_validator(mode="after") + def _replay_needs_stored_reasoning(self) -> "ConversationSimulatorConfig": + # A trace that is not stored is gone before the next request, so + # there would be nothing to replay. + if self.replay_assistant_reasoning and not self.store_reasoning: + raise ValueError("replay_assistant_reasoning needs store_reasoning: an unstored trace cannot be replayed") + return self + @property def required_columns(self) -> list[str]: return [ diff --git a/src/usersim/engine/core/context.py b/src/usersim/engine/core/context.py index ebdb83a..a630621 100644 --- a/src/usersim/engine/core/context.py +++ b/src/usersim/engine/core/context.py @@ -14,7 +14,7 @@ import time from typing import Any -from usersim.engine.core.llm import acall_llm, append_debug_record +from usersim.engine.core.llm import acall_llm, append_debug_record, replaying_assistant_reasoning logger = logging.getLogger("usersim.engine") @@ -111,9 +111,15 @@ def prepare_assistant_history( Compression runs before the probe transform because summary keys refer to canonical message positions. The finance probe uses the transform to bound retrieved document bodies; every other probe inherits the identity view. + + An episode that replays the assistant's reasoning sends its history + uncompressed, so the assistant sees what it produced: a summary beside + the full original reasoning is a history it never produced. Read from the + same setting the request code reads, so the two never disagree. """ messages: list[dict[str, Any]] - if getattr(cfg, "context_compression", False) and getattr(state, "conv_summaries", None): + compress = getattr(cfg, "context_compression", False) and not replaying_assistant_reasoning() + if compress and getattr(state, "conv_summaries", None): messages = compress_history( state.messages, state.conv_summaries, diff --git a/src/usersim/engine/core/episode_input.py b/src/usersim/engine/core/episode_input.py index 4590ccf..dee9ee8 100644 --- a/src/usersim/engine/core/episode_input.py +++ b/src/usersim/engine/core/episode_input.py @@ -179,7 +179,16 @@ def construct_episode_preamble( user_model=resolve_model_name(models.get(MODEL_USER), MODEL_USER), assistant_model=resolve_model_name(models.get(MODEL_ASSISTANT), MODEL_ASSISTANT), prompt_version=prompt_version, - extra_keys=[("locale", locale)], + # A run that replays the assistant's reasoning is a different + # conversation; without replay the id is what it always was. + extra_keys=[ + ("locale", locale), + *( + [("replay_assistant_reasoning", "1")] + if getattr(config, "replay_assistant_reasoning", False) + else [] + ), + ], ) ) language = get_conversation_language(locale) diff --git a/src/usersim/engine/core/episode_runtime.py b/src/usersim/engine/core/episode_runtime.py index 02add9c..d8f63ec 100644 --- a/src/usersim/engine/core/episode_runtime.py +++ b/src/usersim/engine/core/episode_runtime.py @@ -31,7 +31,7 @@ from usersim.engine.core.identity import resolve_model_name from usersim.engine.core.llm import get_current_outcome_builder, set_current_outcome_builder from usersim.engine.core.outcomes import Provenance -from usersim.engine.core.probes import BaseProbe +from usersim.engine.core.probes import BaseProbe, dispatch_episode from usersim.engine.core.simulation import ConversationState if TYPE_CHECKING: @@ -579,7 +579,8 @@ async def _drive_native_lifecycle(self) -> None: set_current_outcome_builder(self.outcome) self.probe.set_tool_call_observer(self._record_executed_tool_call) try: - result = await self.probe.run_dispatch( + result = await dispatch_episode( + self.probe, models=self._activation_models, data=self.probe._data, cfg=self.config, @@ -698,6 +699,11 @@ def _chat_message_to_dict(message: Any) -> dict[str, Any]: "role": str(role_value), "content": getattr(message, "content", "") or "", } + # Present only on the assistant's earlier messages, when its run replays + # reasoning, so a host forwards exactly what a configured model receives. + reasoning = getattr(message, "reasoning_content", None) + if reasoning: + normalized["reasoning_content"] = str(reasoning) tool_calls = getattr(message, "tool_calls", None) if tool_calls: normalized["tool_calls"] = [ diff --git a/src/usersim/engine/core/llm.py b/src/usersim/engine/core/llm.py index e0d81ae..7a47b36 100644 --- a/src/usersim/engine/core/llm.py +++ b/src/usersim/engine/core/llm.py @@ -12,10 +12,14 @@ import random as _random import threading import time +from collections.abc import Iterator +from contextlib import contextmanager from contextvars import ContextVar from pathlib import Path from typing import Any +from usersim.engine.config import MODEL_ASSISTANT + _MAX_RETRIES = 3 _BASE_DELAY = 1.0 _MAX_DELAY = 30.0 @@ -54,6 +58,31 @@ "usersim_pending_debug_records", default=None, ) +#: Whether the assistant's requests carry the reasoning of its own earlier +#: messages (``ConversationSimulatorConfig.replay_assistant_reasoning``). +#: Installed per episode by ``replay_reasoning_scope``. +_REPLAY_ASSISTANT_REASONING: ContextVar[bool] = ContextVar("usersim_replay_assistant_reasoning", default=False) + + +@contextmanager +def replay_reasoning_scope(config: Any) -> Iterator[None]: + """Install ``config``'s reasoning-replay setting for the episode run inside the block. + + ``core.probes.dispatch_episode`` runs every episode inside one, in the + task that runs the episode, so a hosted episode replays as a local one + does and no episode sees another's setting. The previous setting is + restored on the way out, also when the episode fails. + """ + token = _REPLAY_ASSISTANT_REASONING.set(bool(getattr(config, "replay_assistant_reasoning", False))) + try: + yield + finally: + _REPLAY_ASSISTANT_REASONING.reset(token) + + +def replaying_assistant_reasoning() -> bool: + """Whether the episode running in this context replays the assistant's reasoning.""" + return _REPLAY_ASSISTANT_REASONING.get() class ContextWindowError(RuntimeError): @@ -351,17 +380,19 @@ def flush_debug_log() -> None: from data_designer.engine.models.utils import ChatMessage -def _dicts_to_chat_messages(messages: list[dict[str, Any]]) -> list[ChatMessage]: +def _dicts_to_chat_messages(messages: list[dict[str, Any]], *, keep_reasoning: bool = False) -> list[ChatMessage]: """Convert a list of plain dicts to ChatMessage objects for ModelFacade. - A prior turn's ``reasoning_content`` is never forwarded: the - assistant's ``reasoning_content`` is hard-coded to ``None`` below, - matching ``ChatMessage``'s own default for the field. A trace is an - artefact of the turn that produced it, not conversation state. + A prior turn's ``reasoning_content`` is dropped unless ``keep_reasoning``, + which ``_begin_call`` sets only for the assistant model's own requests + when its run replays reasoning: a model trained on its own earlier + reasoning then sees the history it would see in deployment. Every other + request carries none, so during a simulation no model receives another's + reasoning. - Stored trajectories DO keep the trace on the assistant message (see + Stored trajectories keep the trace on the assistant message (see ``assistant_message`` in ``core/probes.py``), so it stays available - for analysis -- it is dropped only on the way back into a model. + for analysis either way. """ out: list[ChatMessage] = [] for msg in messages: @@ -373,7 +404,7 @@ def _dicts_to_chat_messages(messages: list[dict[str, Any]]) -> list[ChatMessage] out.append( ChatMessage.as_assistant( content=content, - reasoning_content=None, + reasoning_content=(msg.get("reasoning_content") or None) if keep_reasoning else None, tool_calls=msg.get("tool_calls") or None, ) ) @@ -437,7 +468,10 @@ def _begin_call( has_tools = "tools" in kwargs and kwargs["tools"] tool_tag = f" + {len(kwargs['tools'])} tools" if has_tools else "" - chat_messages = _dicts_to_chat_messages(messages) + # Only the assistant's requests hold its own earlier messages; every other + # model sees the conversation as text, so only the assistant may replay. + replay = alias == MODEL_ASSISTANT and _REPLAY_ASSISTANT_REASONING.get() + chat_messages = _dicts_to_chat_messages(messages, keep_reasoning=replay) logger.debug(f" |-- LLM call: {alias} ({len(messages)} msgs{tool_tag}) ...") return facade, chat_messages, time.monotonic() diff --git a/src/usersim/engine/core/manifest.py b/src/usersim/engine/core/manifest.py index 4165aea..cbc96e5 100644 --- a/src/usersim/engine/core/manifest.py +++ b/src/usersim/engine/core/manifest.py @@ -172,6 +172,7 @@ class SimulationConfigSnapshot: compression_window: int random_seed: int | None = None store_reasoning: bool = True + replay_assistant_reasoning: bool = False @dataclass(frozen=True) @@ -337,6 +338,7 @@ def _sim_config_from_dict(d: Mapping[str, Any]) -> SimulationConfigSnapshot: max_query_attempts=int(d.get("max_query_attempts", 3)), max_assistant_attempts=int(d.get("max_assistant_attempts", 1)), store_reasoning=bool(d.get("store_reasoning", True)), + replay_assistant_reasoning=bool(d.get("replay_assistant_reasoning", False)), incremental_disclosure_ratio=float(d.get("incremental_disclosure_ratio", 0.6)), persona_grounding_ratio=float(d.get("persona_grounding_ratio", 1.0)), context_compression=bool(d.get("context_compression", True)), @@ -777,6 +779,7 @@ def build_run_manifest( max_query_attempts=int(getattr(sim_config, "max_query_attempts", 3)), max_assistant_attempts=int(getattr(sim_config, "max_assistant_attempts", 1)), store_reasoning=bool(getattr(sim_config, "store_reasoning", True)), + replay_assistant_reasoning=bool(getattr(sim_config, "replay_assistant_reasoning", False)), incremental_disclosure_ratio=float(getattr(sim_config, "incremental_disclosure_ratio", 0.6)), persona_grounding_ratio=float(getattr(sim_config, "persona_grounding_ratio", 1.0)), context_compression=bool(getattr(sim_config, "context_compression", True)), diff --git a/src/usersim/engine/core/probes.py b/src/usersim/engine/core/probes.py index 956f5e4..e064066 100644 --- a/src/usersim/engine/core/probes.py +++ b/src/usersim/engine/core/probes.py @@ -1495,3 +1495,26 @@ async def _baseprobe_run_dispatch( BaseProbe.run_dispatch = _baseprobe_run_dispatch # type: ignore[attr-defined] + + +async def dispatch_episode( + probe: BaseProbe, + *, + models: dict[str, Any], + data: dict[str, Any], + cfg: Any, + **dispatch_kwargs: Any, +) -> dict: + """Run one episode through ``probe.run_dispatch`` with ``cfg``'s per-episode settings installed. + + The generator, the hosted runtime and every ``simulate_`` shortcut + run their episodes through here, so a setting that holds for one episode + (today: reasoning replay) is installed in this one place. The settings are + context variables, so call it in the task that runs the episode. + ``dispatch_kwargs`` (``state``, ``seed_state``) pass through to + ``run_dispatch``. + """ + from usersim.engine.core.llm import replay_reasoning_scope + + with replay_reasoning_scope(cfg): + return await probe.run_dispatch(models=models, data=data, cfg=cfg, **dispatch_kwargs) diff --git a/src/usersim/engine/core/simulation.py b/src/usersim/engine/core/simulation.py index 5e2e5e9..b4cbde9 100644 --- a/src/usersim/engine/core/simulation.py +++ b/src/usersim/engine/core/simulation.py @@ -270,9 +270,9 @@ def make_result( Assistant thinking traces, when the model emits them, ride on the assistant message itself as ``reasoning_content`` -- the OpenAI-style shape ``{role, content, reasoning_content, tool_calls}``. They are - stored for analysis but NOT replayed to a model: ``acall_llm`` converts - messages with ``include_reasoning=False``, so a later turn never sees - an earlier turn's thinking. + stored for analysis, and sent back to the assistant model only when its + run replays reasoning (``replay_assistant_reasoning``); otherwise a later + turn never sees an earlier turn's thinking. """ return { "conversation_messages": json.dumps(messages, ensure_ascii=False, default=str), diff --git a/src/usersim/engine/generator.py b/src/usersim/engine/generator.py index de920f5..77d0cbf 100644 --- a/src/usersim/engine/generator.py +++ b/src/usersim/engine/generator.py @@ -37,7 +37,7 @@ set_current_outcome_builder, ) from usersim.engine.core.outcomes import Provenance -from usersim.engine.core.probes import BankLoadError, load_builtin_probes +from usersim.engine.core.probes import BankLoadError, dispatch_episode, load_builtin_probes from usersim.engine.core.provenance import get_code_sha logger = logging.getLogger("usersim.engine") @@ -205,7 +205,7 @@ async def _generate_row(self, data: dict) -> dict: if construction_error is not None: raise construction_error.cause assert constructed is not None - result = await constructed.probe.run_dispatch(models=models, data=data, cfg=cfg) + result = await dispatch_episode(constructed.probe, models=models, data=data, cfg=cfg) except ContextWindowError as e: logger.warning( " |-- Context window failure for %s: %s", diff --git a/src/usersim/engine/probes/financial_services/generator.py b/src/usersim/engine/probes/financial_services/generator.py index 4db7017..d474ef2 100644 --- a/src/usersim/engine/probes/financial_services/generator.py +++ b/src/usersim/engine/probes/financial_services/generator.py @@ -68,6 +68,7 @@ from usersim.engine.core.probes import ( BankBackedProbe, ToolExecutionMixin, + dispatch_episode, register_probe, ) from usersim.engine.core.simulation import ( @@ -1092,7 +1093,7 @@ async def simulate_financial_services( return _aborted(f"financial_services bank load failed: {e}", provenance) except FinancialServicesProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/general_educational/generator.py b/src/usersim/engine/probes/general_educational/generator.py index b57cce9..6e9c17d 100644 --- a/src/usersim/engine/probes/general_educational/generator.py +++ b/src/usersim/engine/probes/general_educational/generator.py @@ -20,7 +20,7 @@ ) from usersim.engine.core.messages import _parse_theme from usersim.engine.core.persona import format_persona_for_prompt -from usersim.engine.core.probes import BaseProbe, register_probe +from usersim.engine.core.probes import BaseProbe, dispatch_episode, register_probe from usersim.engine.core.prompt_loader import render_prompt from usersim.engine.core.simulation import language_instruction from usersim.engine.probes.general_educational.prompts import ( @@ -122,4 +122,4 @@ async def simulate_general_educational( data=data, outcome_builder=kwargs.get("outcome_builder"), ) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) diff --git a/src/usersim/engine/probes/general_open_ended/generator.py b/src/usersim/engine/probes/general_open_ended/generator.py index 647052f..52f076c 100644 --- a/src/usersim/engine/probes/general_open_ended/generator.py +++ b/src/usersim/engine/probes/general_open_ended/generator.py @@ -20,7 +20,7 @@ ) from usersim.engine.core.messages import _parse_theme from usersim.engine.core.persona import format_persona_for_prompt -from usersim.engine.core.probes import BaseProbe, register_probe +from usersim.engine.core.probes import BaseProbe, dispatch_episode, register_probe from usersim.engine.core.prompt_loader import render_prompt from usersim.engine.core.simulation import language_instruction from usersim.engine.probes.general_open_ended.prompts import ( @@ -124,4 +124,4 @@ async def simulate_general_open_ended( data=data, outcome_builder=kwargs.get("outcome_builder"), ) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) diff --git a/src/usersim/engine/probes/safety_agentic/generator.py b/src/usersim/engine/probes/safety_agentic/generator.py index 8655fc0..d890fc5 100644 --- a/src/usersim/engine/probes/safety_agentic/generator.py +++ b/src/usersim/engine/probes/safety_agentic/generator.py @@ -106,6 +106,7 @@ AgenticMixin, BankBackedProbe, assistant_message, + dispatch_episode, register_probe, ) from usersim.engine.core.simulation import ( @@ -585,7 +586,7 @@ async def simulate_safety_agentic( return _aborted(f"agentic-bank load failed: {e}", provenance) except SafetyAgenticProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/safety_chat_pressure/generator.py b/src/usersim/engine/probes/safety_chat_pressure/generator.py index eecb34b..1403a78 100644 --- a/src/usersim/engine/probes/safety_chat_pressure/generator.py +++ b/src/usersim/engine/probes/safety_chat_pressure/generator.py @@ -74,6 +74,7 @@ from usersim.engine.core.probes import ( BankBackedProbe, BankReframingMixin, + dispatch_episode, register_probe, ) from usersim.engine.core.simulation import ConversationState, make_failed @@ -529,7 +530,7 @@ async def simulate_safety_chat_pressure( return _aborted(f"pressure-bank load failed: {e}", provenance) except SafetyChatPressureProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/sov_ai_dynamic/generator.py b/src/usersim/engine/probes/sov_ai_dynamic/generator.py index 2a42e07..3368726 100644 --- a/src/usersim/engine/probes/sov_ai_dynamic/generator.py +++ b/src/usersim/engine/probes/sov_ai_dynamic/generator.py @@ -58,6 +58,7 @@ from usersim.engine.core.persona import format_persona_for_prompt from usersim.engine.core.probes import ( BankBackedProbe, + dispatch_episode, register_probe, ) from usersim.engine.core.probing_taxonomy import ( @@ -367,7 +368,7 @@ async def simulate_sov_ai_dynamic( ) except SovAiDynamicProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/sov_ai_facts/generator.py b/src/usersim/engine/probes/sov_ai_facts/generator.py index 8fb2775..69631a2 100644 --- a/src/usersim/engine/probes/sov_ai_facts/generator.py +++ b/src/usersim/engine/probes/sov_ai_facts/generator.py @@ -60,6 +60,7 @@ from usersim.engine.core.probes import ( BankBackedProbe, BankVerbatimMixin, + dispatch_episode, register_probe, ) from usersim.engine.core.simulation import ConversationState, make_failed @@ -316,7 +317,7 @@ async def simulate_sov_ai_facts( ) except SovAiFactsProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/sov_ai_multilingual_parity/generator.py b/src/usersim/engine/probes/sov_ai_multilingual_parity/generator.py index 9cdd070..4c3794e 100644 --- a/src/usersim/engine/probes/sov_ai_multilingual_parity/generator.py +++ b/src/usersim/engine/probes/sov_ai_multilingual_parity/generator.py @@ -45,6 +45,7 @@ from usersim.engine.core.probes import ( BankBackedProbe, BankVerbatimMixin, + dispatch_episode, register_probe, ) from usersim.engine.core.query_bank import ( @@ -352,7 +353,7 @@ async def simulate_sov_ai_multilingual_parity( return _aborted(f"query-bank load failed: {e}", provenance) except SovAiMultilingualParityProbeError as e: return _aborted(str(e), provenance) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) def _aborted(reason: str, provenance: Any) -> dict[str, Any]: diff --git a/src/usersim/engine/probes/tool_calling/generator.py b/src/usersim/engine/probes/tool_calling/generator.py index f9ee6a5..f83d6b0 100644 --- a/src/usersim/engine/probes/tool_calling/generator.py +++ b/src/usersim/engine/probes/tool_calling/generator.py @@ -45,6 +45,7 @@ BaseProbe, ToolCallingMixin, ToolExecutionMixin, + dispatch_episode, register_probe, ) from usersim.engine.core.prompt_loader import render_prompt @@ -443,4 +444,4 @@ async def simulate_tool_calling( data=data, outcome_builder=kwargs.get("outcome_builder"), ) - return await probe.run_dispatch(models=models, data=data, cfg=cfg) + return await dispatch_episode(probe, models=models, data=data, cfg=cfg) diff --git a/src/usersim/reporting/capability_report.py b/src/usersim/reporting/capability_report.py index a8b3c33..74861a5 100644 --- a/src/usersim/reporting/capability_report.py +++ b/src/usersim/reporting/capability_report.py @@ -1658,7 +1658,7 @@ def _render_resources_section(report: CapabilityReport) -> str: - **Row 1 (Overall)**: total / input / output / reasoning across every alias. The ``input`` card explains "where did the gap go?" when reasoning looks small relative to total -- input tokens - never carry reasoning content. + carry earlier reasoning only when the run replays it. - **Row 2 (User Tokens)**: total / output / reasoning / conversation for ``user_model`` alone (the simulated-user agent). - **Row 3 (Assistant Tokens)**: same decomposition for @@ -1716,7 +1716,8 @@ def _render_resources_section(report: CapabilityReport) -> str: ) reasoning_tooltip = ( "Reasoning output tokens (invisible thinking content for reasoning models).\n" - "Subset of OUTPUT tokens only -- input tokens never carry reasoning.\n" + "Newly generated reasoning, a subset of OUTPUT tokens. Input tokens include\n" + "earlier reasoning only when the run replays assistant reasoning (replay_reasoning).\n" "Estimated post-hoc as: output_tokens - tiktoken(visible_content).\n" "\n" "Per-alias breakdown:\n" @@ -1755,9 +1756,9 @@ def _render_resources_section(report: CapabilityReport) -> str: # Row 1 -- Overall "

Overall

\n" '
\n' - f'
' + f'
' f"{total_tokens:,}total tokens
\n" - f'
' + f'
' f"{total_input_tokens:,}input tokens
\n" f'
' f"{total_output_tokens:,}output tokens
\n" diff --git a/src/usersim/reporting/dashboard.py b/src/usersim/reporting/dashboard.py index 22efd77..b2c7bdb 100644 --- a/src/usersim/reporting/dashboard.py +++ b/src/usersim/reporting/dashboard.py @@ -1155,16 +1155,16 @@ def _write_no_build_react_dashboard(path: Path, report: CapabilityReport) -> Non // into input + output, with the reasoning subset of output called // out. The "where is the gap?" answer is the input column: input // tokens dominate when conversations are long and judge / summary - // see big contexts, and they NEVER carry reasoning -- so reasoning - // looks tiny against total tokens but is a real share of output. + // see big contexts, and they carry earlier reasoning only when the run + // replays it -- so newly generated reasoning looks tiny against total tokens but is a real share of output. h("div", {{className:"sim-health-status-row"}}, h("div", {{ className:"sim-stat", - title: `Gross billable in-sim tokens across every alias (user / assistant / api_response / judge / summary). Reasoning is a subset of output tokens only -- input tokens never carry reasoning content. Out-of-sim evaluator-judge tokens are not included.` + title: `Gross billable in-sim tokens across every alias (user / assistant / api_response / judge / summary). The reasoning card counts newly generated reasoning, a subset of output tokens; input tokens include earlier reasoning only when the run replays assistant reasoning (replay_reasoning). Out-of-sim evaluator-judge tokens are not included.` }}, h("strong", null, formatTokens(totalTokens)), h("span", null, "total tokens")), h("div", {{ className:"sim-stat", - title: "Input tokens across every alias. Input tokens never carry reasoning content -- this is the bulk of total tokens for long conversations / large summary contexts." + title: "Input tokens across every alias. They include earlier reasoning only when the run replays assistant reasoning (replay_reasoning) -- this is the bulk of total tokens for long conversations / large summary contexts." }}, h("strong", null, formatTokens(totalInputTokens)), h("span", null, "input tokens")), h("div", {{ className:"sim-stat", diff --git a/src/usersim/testing/__init__.py b/src/usersim/testing/__init__.py index 2865f06..b27e7d1 100644 --- a/src/usersim/testing/__init__.py +++ b/src/usersim/testing/__init__.py @@ -499,6 +499,10 @@ def _plain_message(message: Any) -> dict[str, Any]: "role": str(getattr(role, "value", role)), "content": getattr(message, "content", "") or "", } + # Kept so parity compares replayed reasoning too. + reasoning = getattr(message, "reasoning_content", None) + if reasoning: + value["reasoning_content"] = str(reasoning) tool_calls = getattr(message, "tool_calls", None) if tool_calls: value["tool_calls"] = [ diff --git a/tests/engine/core/test_llm.py b/tests/engine/core/test_llm.py index 08ecace..68705eb 100644 --- a/tests/engine/core/test_llm.py +++ b/tests/engine/core/test_llm.py @@ -29,6 +29,7 @@ ContextWindowError, _dicts_to_chat_messages, acall_llm, + replay_reasoning_scope, set_current_outcome_builder, ) from usersim.engine.core.outcomes import OutcomeBuilder, OutcomeStatus @@ -546,12 +547,12 @@ async def acompletion(self, messages, **kwargs): # --------------------------------------------------------------------------- -# Thinking traces are stored, never replayed +# Thinking traces are stored, and replayed only to the assistant, on request # --------------------------------------------------------------------------- def test_trace_does_not_reach_the_chat_message() -> None: - """A stored thinking trace must not travel back into a model call. + """By default, a stored thinking trace must not travel back into a model call. Trajectories keep ``reasoning_content`` on the assistant message for analysis, so the guarantee lives at the conversion boundary: whatever @@ -579,6 +580,73 @@ def test_trace_does_not_reach_the_chat_message() -> None: assert trace not in str(payload) +class _RecordingFacade: + """Records the wire form of every request it receives.""" + + def __init__(self) -> None: + self.payloads: list[list[dict]] = [] + + async def acompletion(self, messages, **kwargs): + self.payloads.append([m.to_dict() for m in messages]) + return SimpleNamespace( + message=SimpleNamespace(content="ok", reasoning_content=None, tool_calls=None), + usage=None, + ) + + +_REPLAY_ON = SimpleNamespace(replay_assistant_reasoning=True) +_REPLAY_OFF = SimpleNamespace(replay_assistant_reasoning=False) +_TRACE = "Plan: ask about sleep first." +_HISTORY = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer", "reasoning_content": _TRACE}, + {"role": "user", "content": "second question"}, +] + + +async def _sent(alias: str, config) -> str: + facade = _RecordingFacade() + with replay_reasoning_scope(config): + await acall_llm({alias: facade}, alias, _HISTORY) + return str(facade.payloads) + + +async def test_replay_sends_the_assistant_its_earlier_reasoning() -> None: + assert _TRACE in await _sent("assistant_model", _REPLAY_ON) + assert _TRACE not in await _sent("assistant_model", _REPLAY_OFF) + + +@pytest.mark.parametrize("alias", ["user_model", "judge_model", "summary_model", "api_response_model"]) +async def test_replay_never_sends_reasoning_to_another_model(alias: str) -> None: + """Even a request that holds an assistant message with reasoning: only the assistant's own history replays.""" + assert _TRACE not in await _sent(alias, _REPLAY_ON) + + +async def test_the_replay_setting_is_restored_after_a_failure() -> None: + with pytest.raises(RuntimeError), replay_reasoning_scope(_REPLAY_ON): + raise RuntimeError("the episode failed") + + # A call outside any episode, in the same context, must not replay. + facade = _RecordingFacade() + await acall_llm({"assistant_model": facade}, "assistant_model", _HISTORY) + assert _TRACE not in str(facade.payloads) + + +async def test_concurrent_episodes_each_keep_their_own_setting() -> None: + import asyncio + + async def episode(config) -> str: + facade = _RecordingFacade() + with replay_reasoning_scope(config): + await asyncio.sleep(0) + await acall_llm({"assistant_model": facade}, "assistant_model", _HISTORY) + return str(facade.payloads) + + on, off = await asyncio.gather(episode(_REPLAY_ON), episode(_REPLAY_OFF)) + + assert _TRACE in on and _TRACE not in off + + # --------------------------------------------------------------------------- # Settings Data Designer's facade would drop # --------------------------------------------------------------------------- diff --git a/tests/engine/core/test_reasoning_replay.py b/tests/engine/core/test_reasoning_replay.py new file mode 100644 index 0000000..0653c1d --- /dev/null +++ b/tests/engine/core/test_reasoning_replay.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in replay of the assistant's own earlier reasoning. + +The invariants: with ``replay_assistant_reasoning`` on, the assistant's +requests carry the reasoning of its own earlier messages, on a hosted run as on +a local one; with it off, requests are unchanged; and no other model, the +simulated user above all, ever receives the assistant's reasoning. + +These run whole episodes through the real model-call path, against scripted +models that record every request; every scripted reply carries reasoning +(``"assistant reasoning"`` for the assistant). +""" + +from __future__ import annotations + +import ast +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from test_simulator_pipeline_e2e import _cli_built_simulator_config, _synthetic_row_data + +import usersim +from usersim.engine.config import ConversationSimulatorConfig +from usersim.engine.core.probes import known_probes, resolve_probe +from usersim.engine.generator import ConversationSimulatorGenerator +from usersim.testing import ScriptedModel, hosted_parity_report + +_ALIASES = { + "api_response_model": "api", + "assistant_model": "assistant", + "judge_model": "judge", + "summary_model": "summary", + "user_model": "user", +} +_TRACE = "assistant reasoning" + + +@pytest.fixture(autouse=True) +def _no_moves_override(monkeypatch): + monkeypatch.delenv("USERSIM_DISCLOSURE_MOVES", raising=False) + + +class _Continuing(ScriptedModel): + """Scripted models whose summary model answers the early-stop check with "no", so turn 2 happens.""" + + CONTENT = {**ScriptedModel.CONTENT, "summary": "no"} + + +def _config(probe_type: str, **changes: Any) -> ConversationSimulatorConfig: + return _cli_built_simulator_config(probe_type, "en_US").model_copy(update={"max_turns": 2, **changes}) + + +async def _requests(cfg: ConversationSimulatorConfig, probe_type: str, model=_Continuing, variant=None): + """Run one episode and return every request, in order, as ``(role, messages)``.""" + prompts: list[dict[str, Any]] = [] + models = {alias: model(role, prompts) for alias, role in _ALIASES.items()} + data = _synthetic_row_data(probe_type, "en_US", cfg) + if variant is not None: + data["probe_variant"] = variant + await ConversationSimulatorGenerator.with_models(cfg, models).agenerate(data) + return [(prompt["role"], prompt["messages"]) for prompt in prompts] + + +def _replayed(messages: list[dict[str, Any]]) -> list[str]: + return [m.get("reasoning_content") for m in messages if m.get("role") == "assistant" and m.get("reasoning_content")] + + +async def test_the_assistant_receives_its_earlier_reasoning_when_replay_is_on(): + requests = await _requests(_config("general_open_ended", replay_assistant_reasoning=True), "general_open_ended") + + later = [messages for role, messages in requests if role == "assistant" and _replayed(messages)] + assert later, "the assistant's second-turn request should carry its first-turn reasoning" + assert all(trace == _TRACE for messages in later for trace in _replayed(messages)) + + +async def test_requests_are_unchanged_when_replay_is_off(): + requests = await _requests(_config("general_open_ended"), "general_open_ended") + + assistant = [messages for role, messages in requests if role == "assistant"] + assert len(assistant) >= 2, "the assistant needs a second turn for this to check anything" + assert not any(_replayed(messages) for _, messages in requests) + + +@pytest.mark.parametrize("replay", [False, True]) +async def test_no_other_model_ever_receives_the_assistants_reasoning(replay): + requests = await _requests(_config("general_open_ended", replay_assistant_reasoning=replay), "general_open_ended") + + others = [messages for role, messages in requests if role != "assistant"] + assert others and _TRACE not in json.dumps(others) + + +async def test_a_tool_step_carries_the_reasoning_of_the_step_before_it(): + requests = await _requests(_config("tool_calling", replay_assistant_reasoning=True), "tool_calling") + + tool_steps = [ + messages + for role, messages in requests + if role == "assistant" and any(m.get("role") == "tool" for m in messages) + ] + assert tool_steps + assert all(_replayed(messages) for messages in tool_steps) + + +class _LongAnswers(_Continuing): + """Long enough answers that the simulator summarizes them for later requests.""" + + CONTENT = {**_Continuing.CONTENT, "assistant": "A long scripted answer. " * 40} + + +@pytest.mark.parametrize(("replay", "compressed"), [(False, True), (True, False)]) +async def test_the_assistants_history_is_compressed_only_without_replay(replay, compressed): + """A provider such as DeepSeek rejects a tool request whose earlier turns lack their reasoning.""" + cfg = _config("general_open_ended", replay_assistant_reasoning=replay, context_compression=True, max_turns=3) + + requests = await _requests(cfg, "general_open_ended", model=_LongAnswers) + + last_assistant = [messages for role, messages in requests if role == "assistant"][-1] + assert ("[Summary of previous response" in json.dumps(last_assistant)) is compressed + + +def _probe_variant_pairs() -> list[tuple[str, str]]: + pairs = [] + for probe in sorted(known_probes()): + module = sys.modules[resolve_probe(probe).__module__] + pairs.extend((probe, str(v)) for v in module.PROBE_VARIANTS) + return pairs + + +@pytest.mark.parametrize(("probe_type", "variant"), _probe_variant_pairs()) +async def test_every_probe_runs_with_replay_and_keeps_the_reasoning_from_others(probe_type, variant): + requests = await _requests( + _config(probe_type, replay_assistant_reasoning=True), + probe_type, + variant=variant if variant != "default" else None, + ) + + assert requests + others = [messages for role, messages in requests if role != "assistant"] + assert _TRACE not in json.dumps(others) + + +@pytest.mark.parametrize("probe_type", ["general_open_ended", "tool_calling"]) +async def test_a_hosted_episode_replays_as_a_local_one_does(probe_type): + cfg = _config(probe_type, replay_assistant_reasoning=True) + data = _synthetic_row_data(probe_type, "en_US", cfg) + persona = data.pop(cfg.persona_column) + + report = await hosted_parity_report(probe_type, config=cfg, persona=persona, data=data, model_factory=_Continuing) + + assert report.differences == [] + hosted = [p["messages"] for p in report.hosted_prompts if p["role"] == "assistant"] + assert any(_replayed(messages) for messages in hosted) + + +async def test_a_probe_shortcut_replays_as_the_generator_does(): + """The public ``simulate_`` functions promise the dispatcher path's behaviour.""" + from usersim.engine.core.episode_input import construct_episode_preamble + from usersim.engine.probes.general_open_ended.generator import simulate_general_open_ended + + cfg = _config("general_open_ended", replay_assistant_reasoning=True) + prompts: list[dict[str, Any]] = [] + models = {alias: _Continuing(role, prompts) for alias, role in _ALIASES.items()} + preamble = construct_episode_preamble( + _synthetic_row_data("general_open_ended", "en_US", cfg), config=cfg, models=models + ) + + await simulate_general_open_ended( + models=models, + data=preamble.data, + persona=preamble.persona, + profile=preamble.behavioral_profile, + locale=preamble.locale, + language=preamble.language, + cfg=cfg, + ) + + assert any(_replayed(p["messages"]) for p in prompts if p["role"] == "assistant") + + +class _RunDispatchCalls(ast.NodeVisitor): + """Each ``run_dispatch`` call, other than ``super().run_dispatch(...)``, by its enclosing function.""" + + def __init__(self) -> None: + self.functions: list[str] = [] + self.calls: list[str] = [] + + def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + self.functions.append(node.name) + self.generic_visit(node) + self.functions.pop() + + visit_AsyncFunctionDef = visit_FunctionDef + + def visit_Call(self, node: ast.Call) -> None: + func = node.func + if isinstance(func, ast.Attribute) and func.attr == "run_dispatch": + receiver = func.value + via_super = ( + isinstance(receiver, ast.Call) and isinstance(receiver.func, ast.Name) and receiver.func.id == "super" + ) + if not via_super: + self.calls.append(self.functions[-1] if self.functions else "") + self.generic_visit(node) + + +def test_every_episode_in_the_package_runs_through_dispatch_episode(): + """Every ``run_dispatch`` call in the package, other than ``super().run_dispatch(...)``, is the one in ``dispatch_episode``. + + Per-episode settings such as replay are installed there and nowhere else, + so a path that called ``run_dispatch`` itself would run without them. + """ + package = Path(usersim.__file__).parent + callers = [] + for path in sorted(package.rglob("*.py")): + visitor = _RunDispatchCalls() + visitor.visit(ast.parse(path.read_text(encoding="utf-8"))) + callers += [(path.relative_to(package).as_posix(), function) for function in visitor.calls] + + assert callers == [("engine/core/probes.py", "dispatch_episode")] + + +def test_compression_follows_the_replay_in_force_not_the_config_alone(): + """A path that runs without the replay switch compresses as usual, so it is never half on.""" + from types import SimpleNamespace + + from usersim.engine.core.context import prepare_assistant_history + from usersim.engine.core.llm import replay_reasoning_scope + + cfg = SimpleNamespace(context_compression=True, compression_window=1, replay_assistant_reasoning=True) + state = SimpleNamespace( + messages=[ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a long answer", "reasoning_content": _TRACE}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "the latest answer", "reasoning_content": _TRACE}, + {"role": "user", "content": "q3"}, + ], + conv_summaries={1: "short", 3: "latest, summarized"}, + ) + probe = SimpleNamespace(transform_assistant_history=None) + + assert "[Summary" in json.dumps(prepare_assistant_history(state, probe, cfg)) + with replay_reasoning_scope(cfg): + assert "[Summary" not in json.dumps(prepare_assistant_history(state, probe, cfg)) diff --git a/tests/test_replay_reasoning_config.py b/tests/test_replay_reasoning_config.py new file mode 100644 index 0000000..a61819b --- /dev/null +++ b/tests/test_replay_reasoning_config.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``replay_reasoning`` in the models file, from the TOML row to the simulator config. + +The invariants: replay is off unless the models file turns it on for +``assistant_model``; a setup that cannot replay (the switch on another alias, +reasoning not stored, a provider whose requests drop reasoning) is refused +before any model call; and a run that replays never shares trajectory ids with +one that does not, while ids without replay stay what they were. +""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +import pytest + +from usersim.cli._errors import ConfigError +from usersim.cli._models import ( + ModelsConfig, + ProviderSpec, + apply_model_overrides, + assistant_replays_reasoning, + default_models_path, + load_models_config, +) +from usersim.cli._pipeline import build_simulator_config_builder +from usersim.engine.config import ConversationSimulatorConfig +from usersim.engine.core._assets import packaged_assets_dir +from usersim.engine.core.episode_input import construct_episode_preamble +from usersim.engine.core.manifest import SimulationConfigSnapshot, build_run_manifest + + +def _models(providers=(), **by_alias): + """The default models, with per-alias changes and extra providers.""" + base = load_models_config(default_models_path()) + specs = tuple(replace(spec, **by_alias.get(spec.alias, {})) for spec in base.models) + return ModelsConfig(providers=base.providers + tuple(providers), models=specs) + + +def _assistant(config): + return next(spec for spec in config.models if spec.alias == "assistant_model") + + +_REPLAY = {"assistant_model": {"replay_reasoning": True}} + + +def test_replay_is_off_unless_the_models_file_turns_it_on(tmp_path): + path = tmp_path / "models.toml" + path.write_text( + '[[models]]\nalias = "assistant_model"\nmodel = "x"\nprovider = "nvidia"\nreplay_reasoning = true\n' + '[[models]]\nalias = "user_model"\nmodel = "x"\nprovider = "nvidia"\n' + ) + config = load_models_config(path) + + assert _assistant(load_models_config(default_models_path())).replay_reasoning is False + assert [spec.replay_reasoning for spec in config.models] == [True, False] + + +def test_only_the_assistant_takes_the_switch(): + with pytest.raises(ConfigError, match="assistant_model"): + assistant_replays_reasoning(_models(user_model={"replay_reasoning": True}), store_reasoning=True) + + +def test_replay_needs_stored_reasoning(): + assert assistant_replays_reasoning(_models(**_REPLAY), store_reasoning=True) is True + with pytest.raises(ConfigError, match="store"): + assistant_replays_reasoning(_models(**_REPLAY), store_reasoning=False) + + +def test_replay_needs_a_provider_that_sends_reasoning(): + """Data Designer's Anthropic adapter drops reasoning from requests.""" + claude = ProviderSpec(name="claude", endpoint="https://example.invalid", provider_type="anthropic", api_key="K") + config = _models([claude], assistant_model={"replay_reasoning": True, "provider": "claude"}) + + with pytest.raises(ConfigError, match="anthropic"): + assistant_replays_reasoning(config, store_reasoning=True) + + +def test_the_run_config_carries_the_switch(tmp_path): + def column(models, **kwargs) -> ConversationSimulatorConfig: + _, builder = build_simulator_config_builder( + locale="en_US", + models=models, + assets_dir=packaged_assets_dir(), + probe_mix={"general_open_ended": 1.0}, + **kwargs, + ) + return builder.get_column_config("conversation_messages") + + assert column(_models()).replay_assistant_reasoning is False + assert column(_models(**_REPLAY)).replay_assistant_reasoning is True + with pytest.raises(ConfigError, match="store"): + column(_models(**_REPLAY), store_reasoning=False) + + +def test_swapping_the_assistant_model_turns_replay_off_unless_the_override_says_so(): + """Replay is a property of the model: a different model has not been trained on it.""" + config = _models(**_REPLAY) + model = "openai/gpt-oss-20b" # one the default provider serves + + assert _assistant(apply_model_overrides(config, {"assistant_model": model})).replay_reasoning is False + kept = apply_model_overrides(config, {"assistant_model": {"model": model, "replay_reasoning": True}}) + assert _assistant(kept).replay_reasoning is True + + +def _trajectory_id(**changes) -> str: + config = ConversationSimulatorConfig(name="conversation_messages", locale="en_US", random_seed=1, **changes) + data = {"persona": {"first_name": "Ana", "age": 30}, "probe_type": "general_open_ended", "theme": "cooking"} + return construct_episode_preamble(data, config=config, models={}).trajectory_id + + +def test_ids_change_only_when_replay_is_on(): + assert _trajectory_id() == _trajectory_id(replay_assistant_reasoning=False) + assert _trajectory_id(replay_assistant_reasoning=True) != _trajectory_id() + + +def test_the_manifest_records_the_switch(): + config = ConversationSimulatorConfig(name="conversation_messages", replay_assistant_reasoning=True) + manifest = build_run_manifest( + run_id="1700000000", + started_at_epoch=1700000000, + models_config=load_models_config(default_models_path()), + sim_config=config, + locales_requested=["en_US"], + probe_mix={"general_open_ended": 1.0}, + ) + + assert manifest.simulation_config.replay_assistant_reasoning is True + assert SimulationConfigSnapshot.__dataclass_fields__["replay_assistant_reasoning"].default is False + + +def test_simulate_refuses_replay_without_stored_reasoning_before_running(tmp_path, caplog): + import logging + + from usersim import cli + + models = tmp_path / "models.toml" + models.write_text( + Path(default_models_path()) + .read_text() + .replace('alias = "assistant_model",', 'alias = "assistant_model", replay_reasoning = true,', 1) + ) + argv = ["simulate", "--locale", "en_US", "--num-rows", "1", "--out", str(tmp_path / "o"), "--models", str(models)] + + with caplog.at_level(logging.ERROR): + assert cli.main([*argv, "--no-store-reasoning", "--dry-run"]) == 2 + assert "replay_reasoning" in caplog.text + assert cli.main([*argv, "--dry-run"]) == 0 + + +def test_a_manifest_keeps_the_switch_when_read_back(): + from usersim.engine.core.manifest import RunManifest + + config = ConversationSimulatorConfig(name="conversation_messages", replay_assistant_reasoning=True) + manifest = build_run_manifest( + run_id="1700000000", + started_at_epoch=1700000000, + models_config=load_models_config(default_models_path()), + sim_config=config, + locales_requested=["en_US"], + probe_mix={"general_open_ended": 1.0}, + ) + + assert RunManifest.from_dict(manifest.to_dict()).simulation_config.replay_assistant_reasoning is True + + +def test_a_plain_run_records_the_switch_in_its_manifest(tmp_path): + from unittest.mock import patch + + import pandas as pd + + from usersim.cli.simulate import _simulate_to_parquet + from usersim.engine.core.manifest import load_run_manifest + from usersim.engine.core.storage import latest_run_id + + rows = pd.DataFrame( + { + "trajectory_id": ["t1"], + "locale": ["en_US"], + "probe_family": ["general"], + "probe_type": ["general_open_ended"], + } + ) + + class _Designer: + def create(self, builder, num_records=None): + return type("Result", (), {"load_dataset": lambda self: rows})() + + with patch("usersim.cli._pipeline.build_simulator_config_builder", lambda **kwargs: (_Designer(), object())): + _simulate_to_parquet( + models=_models(**_REPLAY), + models_path=default_models_path(), + panel=None, + locales=["en_US"], + num_rows=1, + out=tmp_path, + assets_dir=packaged_assets_dir(), + max_turns=1, + probe_mix={"general_open_ended": 1.0}, + random_seed=1, + skip_existing=False, + ) + + manifest = load_run_manifest(latest_run_id(tmp_path), tmp_path) + assert manifest.simulation_config.replay_assistant_reasoning is True + + +def test_materialized_rows_carry_the_switch_to_a_host(monkeypatch): + """A host learns the setting from the row's ``usersim_config``.""" + from usersim.cli._pipeline import materialize_episode_inputs + + monkeypatch.setenv("NVIDIA_API_KEY", "not-used") + (row,) = materialize_episode_inputs( + locale="en_US", + num_rows=1, + models=_models(**_REPLAY), + assets_dir=packaged_assets_dir(), + probe_mix={"general_open_ended": 1.0}, + ) + + assert row["usersim_config"]["replay_assistant_reasoning"] is True + + +def test_the_config_itself_refuses_replay_without_stored_reasoning(): + with pytest.raises(ValueError, match="store_reasoning"): + ConversationSimulatorConfig( + name="conversation_messages", replay_assistant_reasoning=True, store_reasoning=False + ) + + +def test_the_switch_must_be_a_real_boolean(tmp_path): + path = tmp_path / "models.toml" + path.write_text( + '[[models]]\nalias = "assistant_model"\nmodel = "x"\nprovider = "nvidia"\nreplay_reasoning = "false"\n' + ) + + with pytest.raises(ConfigError, match="true or false"): + load_models_config(path) + + +def test_the_provider_type_is_matched_whatever_its_case(): + """Data Designer normalizes the case before choosing an adapter.""" + hub = ProviderSpec(name="hub", endpoint="https://example.invalid", provider_type="OpenAI", api_key="K") + config = _models([hub], assistant_model={"replay_reasoning": True, "provider": "hub"}) + + assert assistant_replays_reasoning(config, store_reasoning=True) is True + + +def test_an_override_that_keeps_the_model_keeps_replay(caplog): + import logging + + # Models the default provider serves, so an override can route them. + config = _models(assistant_model={"replay_reasoning": True, "model": "openai/gpt-oss-20b"}) + + kept = apply_model_overrides(config, {"assistant_model": {"model": "openai/gpt-oss-20b", "temperature": 0.2}}) + assert _assistant(kept).replay_reasoning is True + with caplog.at_level(logging.WARNING): + swapped = apply_model_overrides(config, {"assistant_model": "openai/gpt-oss-120b"}) + assert _assistant(swapped).replay_reasoning is False + assert "replay_reasoning is off" in caplog.text + + +def test_the_simulate_notebook_records_the_switch_in_its_manifest(): + import json + + notebook = Path(__file__).resolve().parents[1] / "notebooks" / "01_simulate.ipynb" + source = "\n".join("".join(cell["source"]) for cell in json.loads(notebook.read_text())["cells"]) + + assert "replay_assistant_reasoning=assistant_replays_reasoning(MODELS" in source From e23b0c396c33453aa38292533a4be66d9e4fe289 Mon Sep 17 00:00:00 2001 From: isarasua Date: Tue, 13 Oct 2026 09:11:25 +0200 Subject: [PATCH 2/2] fix(cli): take replay from the run's models file in `--inputs` runs Follows #30, as agreed in review: whichever of the two landed second adds replay to the setup fingerprint and keeps one provider lookup. - An `--inputs` run replays the assistant's reasoning when its own models file sets `replay_reasoning` on `assistant_model`, whatever the rows were stored with, since the assistant under test comes from that file too. Rows stored with different replay settings can run in one call. - The setup fingerprint includes it, so runs that differ only in replay get runs and ids of their own. - The provider check uses #30's `resolved_model_providers`, and `_provider_type` is gone. An empty providers file now reads as empty, as Data Designer reads it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: isarasua --- docs/engine/EXTERNAL_PROBE_RUNTIME.md | 8 +++-- src/usersim/cli/_inputs.py | 47 +++++++++++++++++++++------ src/usersim/cli/_models.py | 31 ++++++------------ tests/test_cli_simulate_inputs.py | 41 +++++++++++++++++++++-- tests/test_replay_reasoning_config.py | 12 +++++++ 5 files changed, 104 insertions(+), 35 deletions(-) diff --git a/docs/engine/EXTERNAL_PROBE_RUNTIME.md b/docs/engine/EXTERNAL_PROBE_RUNTIME.md index a35ac53..0818a13 100644 --- a/docs/engine/EXTERNAL_PROBE_RUNTIME.md +++ b/docs/engine/EXTERNAL_PROBE_RUNTIME.md @@ -72,6 +72,7 @@ print(result.run_id, result.written, result.skipped) conversation: - each simulator model as Data Designer sends it, provider settings included + - whether the assistant's earlier reasoning is replayed - the rows' shared settings - the asset contents - the `USERSIM_*` variables @@ -102,8 +103,11 @@ print(result.run_id, result.written, result.skipped) saved row's `input_id`. - **Settings come from the rows.** Each row runs with the settings it was stored with, so a sampling flag (`--max-turns`, `--probe-mix`, ...) set to - anything but its default is refused. The rows in one call must share their - settings apart from locale and asset folder. + anything but its default is refused. The one exception is replaying the + assistant's reasoning: like the assistant itself, it comes from this run's + models file (`replay_reasoning` on `assistant_model`), whatever the rows were + stored with. The rows in one call must share their settings apart from + locale, asset folder and that replay setting. `--assets-dir` replaces the asset folder the rows were stored with, and is required when that folder doesn't exist on this machine. - **Concurrent writers.** Writes into `--out` take a lock, so two calls can diff --git a/src/usersim/cli/_inputs.py b/src/usersim/cli/_inputs.py index ed15cf0..ac655fe 100644 --- a/src/usersim/cli/_inputs.py +++ b/src/usersim/cli/_inputs.py @@ -60,6 +60,9 @@ #: Settings that may differ between the rows of one invocation. _PER_ROW_SETTINGS = ("locale", "assets_dir") +#: Settings the run's models file decides, whatever the rows were stored with: +#: replaying the assistant's reasoning belongs to the assistant under test. +_RUN_SETTINGS = ("replay_assistant_reasoning",) @dataclass(frozen=True) @@ -67,7 +70,8 @@ class EpisodeInputs: """Stored rows checked and ready to run.""" rows: list[dict[str, Any]] - #: The ``usersim_config`` every row shares, apart from ``_PER_ROW_SETTINGS``. + #: The ``usersim_config`` every row shares, apart from ``_PER_ROW_SETTINGS`` + #: and ``_RUN_SETTINGS``, which the run settles (see :func:`_run_config`). config: dict[str, Any] #: The asset root the rows run against. assets_dir: Path @@ -127,8 +131,9 @@ def prepare_episode_inputs( """Check stored rows before any model call, and settle what they run against. Raises ``ValueError`` naming the row and the problem. The rows must share - one setup: their ``usersim_config`` may differ only in locale and asset - root, so that one run holds one setup. + one setup: their ``usersim_config`` may differ only in locale, asset + root and the settings the run's models file decides, so that one run + holds one setup. """ from pydantic import ValidationError @@ -185,7 +190,11 @@ def _shared_config(rows: list[dict[str, Any]]) -> dict[str, Any]: """The setup every row shares; ``ValueError`` naming the settings that differ.""" def setup(row: dict[str, Any]) -> dict[str, Any]: - return {key: value for key, value in row["usersim_config"].items() if key not in _PER_ROW_SETTINGS} + return { + key: value + for key, value in row["usersim_config"].items() + if key not in _PER_ROW_SETTINGS and key not in _RUN_SETTINGS + } first = setup(rows[0]) for row in rows[1:]: @@ -363,7 +372,8 @@ def setup_fingerprint(models: Any, inputs: EpisodeInputs) -> tuple[str, dict[str - every simulator model alias as Data Designer sends it, provider settings included (``api_key``, a provider's name and ``max_parallel_requests`` are left out: they don't change what a model says); - - the shared ``usersim_config`` as loaded, defaults applied; + - the shared ``usersim_config`` as loaded, defaults applied, with whether + the assistant's reasoning is replayed as the run's models file sets it; - the contents of every asset root, and every non-operational ``USERSIM_*`` variable; - every probe's ``PROMPT_VERSION``, every extension's package release, and @@ -375,7 +385,7 @@ def setup_fingerprint(models: Any, inputs: EpisodeInputs) -> tuple[str, dict[str from usersim.engine.core.manifest import _own_version from usersim.engine.core.storage import RUN_KIND_STORED - simulation = ConversationSimulatorConfig.model_validate(dict(inputs.config)).model_dump(mode="json") + simulation = ConversationSimulatorConfig.model_validate(_run_config(inputs, models)).model_dump(mode="json") for key in _PER_ROW_SETTINGS: simulation.pop(key, None) aliases = [*MODEL_ALIASES, simulation.get("finance_embedding_model_alias") or ""] @@ -487,12 +497,29 @@ def plan_inputs_run(inputs: EpisodeInputs, models: Any, out: str | Path, *, skip return InputsPlan(fingerprint, record, matching_run=matching, pending=pending, saved=saved) -def _column_config(inputs: EpisodeInputs, locale: str): +def _run_config(inputs: EpisodeInputs, models: Any) -> dict[str, Any]: + """The rows' shared ``usersim_config``, with ``_RUN_SETTINGS`` as this run sets them. + + The assistant under test comes from the run's models file, so whether its + earlier reasoning is replayed does too, as in a plain run. Raises + ``ConfigError`` for a models file that cannot replay. + """ + from usersim.cli._models import assistant_replays_reasoning + from usersim.engine.config import ConversationSimulatorConfig + + store_reasoning = ConversationSimulatorConfig.model_validate(dict(inputs.config)).store_reasoning + return { + **inputs.config, + "replay_assistant_reasoning": assistant_replays_reasoning(models, store_reasoning=store_reasoning), + } + + +def _column_config(inputs: EpisodeInputs, locale: str, models: Any): """The simulator column for ``locale``'s rows, on the run's asset root.""" from usersim.engine.config import ConversationSimulatorConfig return ConversationSimulatorConfig.model_validate( - {**inputs.config, "locale": locale, "assets_dir": str(inputs.assets_dir)} + {**_run_config(inputs, models), "locale": locale, "assets_dir": str(inputs.assets_dir)} ) @@ -615,7 +642,7 @@ def simulate_inputs_to_dataset( for locale, group in by_locale.items(): requested[locale] += len(group) data_designer, builder = build_inputs_config_builder( - models=models, config=_column_config(inputs, locale), rows=[row for row, _ in group.values()] + models=models, config=_column_config(inputs, locale, models), rows=[row for row, _ in group.values()] ) # Data Designer names a batch's artifacts by the second it started, so # two invocations in one directory could share them; each batch gets @@ -652,7 +679,7 @@ def simulate_inputs_to_dataset( trajectory_root=out, models_config=models, models_config_path=Path(models_path) if models_path else None, - sim_config=_column_config(inputs, next(iter(requested))), + sim_config=_column_config(inputs, next(iter(requested)), models), locales_requested=sorted(requested), probe_mix={probe: count / len(stored) for probe, count in sorted(probes.items())}, panel_path=inputs.source, diff --git a/src/usersim/cli/_models.py b/src/usersim/cli/_models.py index 7d07616..9a79c86 100644 --- a/src/usersim/cli/_models.py +++ b/src/usersim/cli/_models.py @@ -278,24 +278,6 @@ def _replay_flag(value: Any, *, where: str) -> bool: return value -def _provider_type(config: ModelsConfig, name: str) -> str | None: - """The ``provider_type`` Data Designer serves provider ``name`` with, if known. - - The config's own providers first, then the defaults Data Designer loads - for itself, as :func:`to_data_designer_kwargs` merges them. - """ - for provider in config.providers: - if provider.name == name: - return provider.provider_type - from data_designer.config.default_model_settings import get_builtin_model_providers, get_default_providers - - try: - defaults = get_default_providers() or get_builtin_model_providers() - except FileNotFoundError: - defaults = get_builtin_model_providers() - return next((provider.provider_type for provider in defaults if provider.name == name), None) - - def assistant_replays_reasoning(config: ModelsConfig, *, store_reasoning: bool) -> bool: """Whether a run replays the assistant's own reasoning; refuses a setup that cannot. @@ -318,10 +300,17 @@ def assistant_replays_reasoning(config: ModelsConfig, *, store_reasoning: bool) return False if not store_reasoning: raise ConfigError( - "replay_reasoning on assistant_model needs the reasoning kept: drop --no-store-reasoning " - "(store_reasoning=False), or turn replay off" + "replay_reasoning on assistant_model needs the reasoning kept, but store_reasoning is off: " + "drop --no-store-reasoning, or for --inputs, materialize the rows without it; or turn replay off" ) - provider_type = _provider_type(config, assistant.provider) + provider_type = next( + ( + provider.provider_type + for provider in resolved_model_providers(config) + if provider.name == assistant.provider + ), + None, + ) # Data Designer normalizes the type's case before choosing an adapter. if provider_type is not None and provider_type.lower() != "openai": raise ConfigError( diff --git a/tests/test_cli_simulate_inputs.py b/tests/test_cli_simulate_inputs.py index a353c39..0d39259 100644 --- a/tests/test_cli_simulate_inputs.py +++ b/tests/test_cli_simulate_inputs.py @@ -38,6 +38,7 @@ from usersim import cli from usersim.cli import _inputs +from usersim.cli._errors import ConfigError from usersim.cli._inputs import ( INPUT_ID_COLUMN, load_episode_inputs, @@ -95,7 +96,7 @@ def _models(provider_api_key: str = "NVIDIA_API_KEY", **by_alias: dict[str, Any] @contextmanager -def _offline(): +def _offline(side_effect=_mock_call_llm_side_effect): """Stub every simulator-side model call and skip Data Designer's health checks.""" def no_health_checks(models): @@ -104,7 +105,7 @@ def no_health_checks(models): with ExitStack() as stack: stack.enter_context(patch("usersim.cli._pipeline.to_model_configs", no_health_checks)) for site in _LLM_SITES: - stack.enter_context(patch(f"{site}.acall_llm", side_effect=_mock_call_llm_side_effect)) + stack.enter_context(patch(f"{site}.acall_llm", side_effect=side_effect)) yield @@ -240,12 +241,48 @@ def test_debug_logging_keeps_the_setup(rows, monkeypatch): pytest.param(lambda: _models(assistant_model={"temperature": 0.1}), id="temperature"), pytest.param(lambda: _models(judge_model={"extra_body": {"reasoning_effort": "low"}}), id="judge"), pytest.param(lambda: _models(assistant_model={"model": "another/model"}), id="model"), + pytest.param(lambda: _models(assistant_model={"replay_reasoning": True}), id="replay"), ], ) def test_a_model_setting_that_can_change_a_conversation_is_a_new_setup(rows, models): assert _fingerprint(rows, models()) != _fingerprint(rows) +def test_the_runs_models_file_decides_replay_whatever_the_rows_were_stored_with(rows, tmp_path): + """The assistant under test comes from the run's models file, and so does replaying its reasoning.""" + from usersim.engine.core.llm import replaying_assistant_reasoning + + seen: list[bool] = [] + + def recording(models, alias, msgs, **kwargs): + if alias == "assistant_model": + seen.append(replaying_assistant_reasoning()) + return _mock_call_llm_side_effect(models, alias, msgs, **kwargs) + + mixed = _edited(rows, lambda r: r[0]["usersim_config"].update(replay_assistant_reasoning=True)) + for stored, models, replayed in [ + (mixed, _models(), False), + (rows, _models(assistant_model={"replay_reasoning": True}), True), + ]: + seen.clear() + out = tmp_path / f"replay-{replayed}" + with _offline(recording): + result = simulate_inputs_to_dataset(inputs=prepare_episode_inputs(stored), models=models, out=out) + assert seen and set(seen) == {replayed} + assert load_run_manifest(result.run_id, out).simulation_config.replay_assistant_reasoning is replayed + assert _fingerprint(mixed) == _fingerprint(rows) + + +def test_a_models_file_that_cannot_replay_is_refused_before_any_model_call(rows, tmp_path): + unstored = _edited(rows, lambda r: [row["usersim_config"].update(store_reasoning=False) for row in r]) + replay = _models(assistant_model={"replay_reasoning": True}) + + with patch("usersim.cli._pipeline.build_inputs_config_builder", side_effect=AssertionError("ran")): + with pytest.raises(ConfigError, match="store"): + _run(unstored, tmp_path / "out", replay) + assert list_runs(tmp_path / "out") == [] + + def test_a_providers_own_request_settings_are_part_of_the_setup(rows, monkeypatch): """Data Designer merges a provider's extra_body and extra_headers into every request it serves.""" from data_designer.config.models import ModelProvider diff --git a/tests/test_replay_reasoning_config.py b/tests/test_replay_reasoning_config.py index 7d92b03..a3ed915 100644 --- a/tests/test_replay_reasoning_config.py +++ b/tests/test_replay_reasoning_config.py @@ -249,6 +249,18 @@ def test_the_provider_type_is_matched_whatever_its_case(): assert assistant_replays_reasoning(config, store_reasoning=True) is True +def test_without_providers_of_its_own_the_check_reads_the_ones_data_designer_loads(monkeypatch): + """A models file that declares no provider runs on Data Designer's own, so their type decides.""" + from data_designer.config.models import ModelProvider + + claude = ModelProvider(name="nvidia", endpoint="https://example.invalid", provider_type="anthropic", api_key="K") + monkeypatch.setattr("data_designer.config.default_model_settings.get_default_providers", lambda: [claude]) + config = replace(_models(**_REPLAY), providers=()) + + with pytest.raises(ConfigError, match="OpenAI-compatible"): + assistant_replays_reasoning(config, store_reasoning=True) + + def test_an_override_that_keeps_the_model_keeps_replay(caplog): import logging