interlens.conversation¶

def branch(self) -> "Conversation":
	"""Fork into a new ``Conversation`` that **reuses the same participant objects** (shared weights, zero GPU
	cost) with a copied, independent transcript. The branch can diverge freely; the original is untouched. This
	is exactly a no-op copy-on-write clone (``self.set()``). To fork from a specific point in the history rather
	than the end, use ``branch_from``."""
	return self.set()

def branch_from(self, ref: MessageRef) -> "Conversation":
	"""Fork a new conversation whose history is this one's turns **up to and including** the turn ``ref`` — i.e.
	branch as if the conversation had stopped right after ``ref``, ready for a different continuation. ``ref`` is
	a ``MessageRef``: an ``int`` index (negatives count from the end) or the ``Message`` object itself (matched
	by identity). The original is untouched; the fork shares participants (zero GPU cost)."""
	cut = self.transcript.resolve_index(ref)  # resolve against the original (the fork holds copies, new identities)
	fork = self.branch()
	fork.transcript.rewind(to=cut)
	return fork

@contextmanager
def capture(self, sites=("residual",), layers=None, offload: OffloadLocation = "cpu"):
	"""Context manager that captures activations for every ``step`` inside the block into a fresh
	``ActivationCache``, auto-tagged by the current speaker + turn::

		with conv.capture(sites=["residual"], layers=[8, 12]) as cache:
			conv.step(bob)
		cache.at(participant="bob", layer=12)
	"""
	cache = ActivationCache(offload=offload)
	spec = CaptureSpec(sites=tuple(sites), layers=tuple(layers) if layers is not None else None, offload=offload)
	self._pending_capture = CaptureRequest(cache=cache, spec=spec)
	try:
		yield cache
	finally:
		self._pending_capture = None

def edit(self, ref: MessageRef, content: str | None = None, *, author: str | None = None, **metadata) -> Message:
	"""Edit a committed past turn in place and return it (see ``Transcript.edit``): ``ref`` is a ``MessageRef``
	(int index, negatives allowed, or a ``Message`` object matched by identity), and ``content`` / ``author`` /
	``**metadata`` are the fields to change. Editing the returned ``Message``'s fields directly works too, since
	the transcript holds it by reference."""
	return self.transcript.edit(ref, content, author=author, **metadata)

@classmethod
def from_models(cls, models: tuple[ModelLike, ...], names: tuple[str, ...] = ("a", "b"),
                device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16,
                shared_context: str | None = None, shared_system_prompt: str | None = None,
                prompt: PromptLike = None, **gen_kwargs) -> "Conversation":
	"""Scaffold a conversation directly from a tuple of ``models`` — each an HF id or an already-loaded
	``PreTrainedModel`` (see ``ModelLike``). Convenience wrapper around
	``factories.conversation_from_models``; each model becomes a family-correct participant and ``names`` gives
	them identities. **The order of ``models`` / ``names`` is the speaking order** — the first speaks first
	unless you pass ``first=`` to ``run``. ``**gen_kwargs`` are forwarded to every participant.

	Scenario framing is available here too: ``shared_system_prompt`` (instructions, system role) and
	``shared_context`` (a neutral ``moderator``-voiced opening seen by everyone). ``prompt`` is a separate
	convenience for a *participant*-voiced opener (a ``str`` is attributed to the last participant). See that
	function for the details."""
	from .factories import conversation_from_models  # lazy: factories imports this module

	return conversation_from_models(models, names=names, device=device, dtype=dtype,
	                                shared_context=shared_context, shared_system_prompt=shared_system_prompt,
	                                prompt=prompt, **gen_kwargs)

@classmethod
def load(cls, directory, devices="cuda") -> "Conversation":
	"""Rebuild a saved conversation: reconstruct (lazy) participants from the recipe and attach the saved
	transcript, resuming from that state (does NOT regenerate messages). Raises on an unsupported schema."""
	import json
	from pathlib import Path
	from .participant.serialize import participant_from_dict
	from .transcript import Transcript, SCHEMA_VERSION

	directory = Path(directory)
	recipe = json.loads((directory / "conversation.json").read_text())
	version = recipe.get("schema_version")
	if version != SCHEMA_VERSION:
		raise ValueError(f"checkpoint schema {version!r} is no longer supported (current {SCHEMA_VERSION}); "
		                 f"re-generate the run with the current interlens.")
	device = devices[0] if isinstance(devices, (list, tuple)) else devices
	participants = tuple(participant_from_dict(p, device=device) for p in recipe["participants"])
	return cls(
		participants=participants,
		transcript=Transcript.load(directory / "transcript.json"),
		shared_context=recipe.get("shared_context"),
		shared_system_prompt=recipe.get("shared_system_prompt"),
		moderator_name=recipe.get("moderator_name", "moderator"),
		context_limit=recipe.get("context_limit"),
		reasoning_visibility=ReasoningVisibility(recipe.get("reasoning_visibility", "strip")),
		execution_mode=ExecutionMode(recipe.get("execution_mode", "throughput")),
	)

def participant(self, which: ParticipantLike) -> Participant:
	"""Return one of this conversation's participants, resolved from a ``Participant`` object, its ``name``
	(str), or its index (int). Raises (via ``_resolve_participant``) if it isn't in this conversation."""
	return self.participants[self._resolve_participant(which)]

def render_templated(self, pov: ParticipantLike, *, extra=(), add_generation_prompt: bool = False,
                     tokenize: bool = False):
	"""The full ``view`` run through ``pov``'s tokenizer chat template — the **exact prompt its model sees**,
	special/control tokens and all (a str, or token ids with ``tokenize=True``). This is the truthful
	counterpart to ``transcript.render_templated``, which templates the transcript turns ONLY (no system /
	private framing, no context-fit). Requires a local ``ModelParticipant`` (needs a tokenizer)."""
	participant = self.participant(pov)
	tokenizer = getattr(participant, "tokenizer", None)
	if tokenizer is None:
		raise TypeError(f"{type(participant).__name__} {participant.name!r} has no tokenizer; "
		                f"render_templated needs a local ModelParticipant")
	return tokenizer.apply_chat_template(self.view(pov, extra=extra), tokenize=tokenize,
	                                     add_generation_prompt=add_generation_prompt)

def reset(self) -> "Conversation":
	"""Return the conversation to its fresh, pre-``run`` state: empty the transcript, then **re-seed** the
	``shared_context`` moderator turn (exactly as ``__post_init__`` does on a fresh conversation). Use this
	rather than ``transcript.clear()`` when you want to rerun the same scenario — the raw ``transcript.clear()``
	drops the ``shared_context`` framing too, whereas ``reset`` restores it."""
	self.transcript.clear()
	if self.shared_context is not None:
		self.transcript.append(self.moderator_name, self.shared_context)
	return self

def rewind(self, *, to: MessageRef) -> "Conversation":
	"""Rewind in place so the turn ``to`` becomes the new last turn, dropping everything after it (returns
	``self`` for chaining). ``to`` is a ``MessageRef`` (int index, negatives allowed, or a ``Message`` object).
	Mutates this conversation — use ``branch_from`` instead to keep the original. See ``Transcript.rewind``."""
	self.transcript.rewind(to=to)
	return self

def rollout(self, n: int | None = None, *, devices=None, out_dir=None, resume: bool = False,
            batched: bool = True, max_batch_size: int | None = None, seed: int | None = None) -> "RunReport":
	"""Run many conversations from this one recipe and return a ``RunReport`` — the benchmark/scale entry point.

	**Rollout does NOT mutate this conversation.** It makes an independent copy-on-write clone per job and runs
	*those*; ``self`` stays the unrun recipe. Finished conversations live in ``report.results[job_id].conversation``
	(and their transcripts / analyses on the same ``RunResult``) — that is where to ``sample()`` or inspect
	afterwards, NOT on the original.

	Two modes: with ``data()`` set, expands to ONE conversation per row (job ids ``row_00000``…), resolving each
	templated field (``dataset_field``) against the row; otherwise expands to ``n`` seeded copies of one scenario
	(job ids ``rollout_0000``…). ``n`` defaults to ``len(data)``. Sample *i* gets ``seed + i`` (``seed`` defaults
	to the conversation's ``seed`` field). Any ambient ``with StopCondition(...)`` blocks and the ``run_until``
	field are captured now and attached to every job (so they survive the spawn boundary).

	Parallel by default on two axes (see :func:`interlens.run`): one worker process per device, and batched
	co-stepping within each device. ``batched=False`` gives the token-identical DETERMINISTIC path.
	"""
	from .runner.pool import run_jobs
	seed = self._seed if seed is None else seed
	stop = self._combined_stop_for_jobs()
	jobs = self._expand(n, seed, stop)
	return run_jobs(jobs, devices=devices, out_dir=out_dir, resume=resume, batched=batched,
	                max_batch_size=max_batch_size)

def run(self, turns: int | None = None, until: StopCondition | list | None = None,
        first: ParticipantLike | None = None, prompt: PromptLike = None) -> Transcript:
	"""Alternate speakers until ``turns`` elapse and/or a ``StopCondition`` fires (whichever comes first).

	Speakers are taken in ``participants`` order (that tuple's order IS the turn order); ``first`` sets who
	starts (default: ``participants[0]``) and the rest follow round-robin. ``first`` may be a ``Participant``,
	its ``name`` (str), or its index (int) — resolved via ``_resolve_participant`` (raises if it isn't in this
	conversation).

	Stop conditions are combined from three sources (any of which stops the run): ``until`` (this call), the
	conversation's own ``run_until`` field, and any ambient ``with StopCondition(...):`` blocks. A condition may
	also cap each turn's tokens (``turn_cap``) — e.g. a ``TokenBudget`` shrinks the final turn to land on budget.
	Conditions are ``reset()`` at the start so a reused instance doesn't leak state. ``turns`` defaults to the
	conversation's ``turns`` field; at least one of turns / a stop condition must be present.

	``prompt`` (optional) is appended to the transcript **before** the run: a ``str`` is attributed to the LAST
	participant (so ``first`` replies to it), a ``Message`` is appended as-is, ``None`` appends nothing. It
	appends to the current end and never resets/reseeds.
	"""
	self._check_resolved("run")
	if turns is None:
		turns = self._turns
	stop = self._resolve_stop(until)
	if turns is None and stop is None:
		raise ValueError("run requires at least one of `turns`, a `run_until`/`until` stop condition, or an "
		                 "ambient `with StopCondition(...)` block")
	self._append_prompt(prompt)
	if stop is not None:
		stop.reset()

	start = self._resolve_participant(first) if first is not None else 0
	n = len(self.participants)
	i = 0
	while turns is None or i < turns:
		cap = stop.turn_cap(self) if stop is not None else None
		message = self.step(self.participants[(start + i) % n], max_new_tokens=self._turn_budget(start, i, cap))
		i += 1
		# A hook may have denied the turn (message is None); only a committed message is checked for stopping.
		if message is not None and stop is not None and stop.should_stop(self, message):
			break
	return self.transcript

def sample(self, speaker, message=None, *, as_author=None, steering=None, capture=None, patch=None,
           return_logprobs: bool = False, max_new_tokens: int | None = None):
	"""Ephemerally sample a response **without mutating the transcript** — a pure read of current state, safe to
	call repeatedly / in a loop. ``speaker`` is a ``ParticipantLike`` (name / index / ``Participant``); pass a
	**list or tuple** of them to sample each and get back ``{name: Message}`` instead of a single ``Message``
	(see also ``sample_all``). ``message`` is an optional temporary incoming turn; it defaults to being
	attributed to the **moderator** (a neutral, external voice — so "What do you think of Bob?" reads as an
	interviewer asking, NOT as Bob speaking). Pass ``as_author="bob"`` to make it read as that participant's
	turn instead (e.g. "what would you say if Bob had just said X?"). The same interp options as ``step`` are
	honored on each ephemeral generation."""
	self._check_resolved("sample")
	if isinstance(speaker, (list, tuple)):
		return {self.participant(s).name: self.sample(s, message, as_author=as_author, steering=steering,
		                                              capture=capture, patch=patch, return_logprobs=return_logprobs,
		                                              max_new_tokens=max_new_tokens) for s in speaker}
	speaker = self.participant(speaker)
	extra = []
	if message is not None:
		extra = [Message(author=as_author or self.moderator_name, content=message)]
	return speaker.generate(self._view(speaker, extra=extra), steering=steering, capture=capture,
	                        patch=patch, return_logprobs=return_logprobs, turn=len(self.transcript),
	                        max_new_tokens=max_new_tokens)

def sample_all(self, message: str | None = None, *, as_author: str | None = None, steering=None,
               capture=None, patch=None, return_logprobs: bool = False,
               max_new_tokens: int | None = None) -> "dict[str, Message]":
	"""Ephemerally sample **every** participant's response — a convenience for
	``sample(list(self.participants), ...)``. Returns ``{name: Message}``; the transcript is untouched. Handy
	for "what does each model say to this right now?"."""
	return self.sample(list(self.participants), message, as_author=as_author, steering=steering,
	                   capture=capture, patch=patch, return_logprobs=return_logprobs, max_new_tokens=max_new_tokens)

def save(self, directory) -> None:
	"""Persist the conversation: the transcript (``transcript.json``) + the recipe (``conversation.json`` — the
	participants' own constructor kwargs + scenario framing/policies, no weights)."""
	import json
	from pathlib import Path
	from .participant.serialize import participant_to_dict
	from .transcript import SCHEMA_VERSION

	directory = Path(directory)
	directory.mkdir(parents=True, exist_ok=True)
	recipe = {
		"schema_version": SCHEMA_VERSION,
		"participants": [participant_to_dict(p) for p in self.participants],
		"shared_context": self.shared_context if not has_fields(self.shared_context) else None,
		"shared_system_prompt": self.shared_system_prompt if not has_fields(self.shared_system_prompt) else None,
		"moderator_name": self.moderator_name,
		"context_limit": self.context_limit,
		"reasoning_visibility": self.reasoning_visibility.value,
		"execution_mode": self.execution_mode.value,
	}
	(directory / "conversation.json").write_text(json.dumps(recipe, indent=2))
	self.transcript.save(directory / "transcript.json")

def step(self, speaker: Participant, *, steering=None, capture=None, patch=None,
         return_logprobs: bool = False, max_new_tokens: int | None = None) -> Message:
	"""Have ``speaker`` produce and commit one turn. Interp options flow to ``generate``; if no ``capture`` is
	passed but a ``conv.capture(...)`` block is active, its pending request is used (auto-tagged by turn)."""
	if capture is None:
		capture = self._pending_capture
	turn = len(self.transcript)
	message = speaker.generate(self._view(speaker), steering=steering, capture=capture,
	                           patch=patch, return_logprobs=return_logprobs, turn=turn,
	                           max_new_tokens=max_new_tokens)
	message = self._apply_hooks(message)
	if message is None:
		return None  # a hook denied this turn; nothing is committed
	self.transcript.messages.append(message)
	return message

def view(self, pov: ParticipantLike, extra=()) -> list[dict]:
	"""The full ``[{role, content}]`` view ``pov``'s model is conditioned on — the **real generation input**:
	system block (its private ``system_prompt`` + ``shared_system_prompt``) → ``private_context`` → the
	transcript role-swapped to ``pov``, then **context-fit** and **family-flattened**. This is exactly what
	``step`` / ``sample`` feed to ``generate``. Unlike ``transcript.render_roles`` (transcript turns only, no
	framing or fitting), it reflects everything the model actually sees. ``pov`` is a ``ParticipantLike`` (name
	/ index / participant); ``extra`` renders temporary, uncommitted messages (as ``sample`` does)."""
	return self._view(self.participant(pov), extra=extra)