diff --git a/dokimi/changelog.d/20261021_130000_napetrov_skill-gap-W7.md b/dokimi/changelog.d/20261021_130000_napetrov_skill-gap-W7.md new file mode 100644 index 00000000..711c5568 --- /dev/null +++ b/dokimi/changelog.d/20261021_130000_napetrov_skill-gap-W7.md @@ -0,0 +1,8 @@ + + +### Added + +- **Multi-turn session measurement for `arms run`** — `--session-turns N` asks each question as its own text followed by the first N-1 entries of its `follow_ups` list, keeps the conversation history across the turns, and records each turn's input, output, cache read and write tokens and cost under `session.turns` of the arm. The arm's `token_usage` and `metrics` are the sums over turns, so cost rollups read them as before. `--turn-reminder TEXT` also runs each arm as `+reminder`, with the text appended to every user turn in a `` block, to test whether a per-turn reminder keeps an instruction in force in later turns. Measurement only: `--judge` is refused with `--session-turns`, a question with too few `follow_ups` fails the run before any answer is paid for, and agentic arms (`mcp:`, `skill-agent:`, `skill-catalog:`) are refused per arm. Without the flags, `arms run` behaves as before. diff --git a/dokimi/docs/benchmarking-and-comparison.md b/dokimi/docs/benchmarking-and-comparison.md index 64367f53..94852bf7 100644 --- a/dokimi/docs/benchmarking-and-comparison.md +++ b/dokimi/docs/benchmarking-and-comparison.md @@ -383,6 +383,38 @@ that needed a re-draw. token, but some endpoints compute slightly different scores from one call to the next, so reproducibility has to be measured per endpoint. +### Multi-turn sessions + +A single answer does not show what an arm costs once the conversation grows. +`--session-turns N` asks each question as its own text followed by the first +N-1 entries of its `follow_ups` list (a list of strings in the question record), +keeps the history, and records each turn's usage: + +```bash +python cli.py arms run ... --arms baseline,skill:/path/to/skill --session-turns 4 +python cli.py arms run ... --arms baseline,skill:/path/to/skill --session-turns 4 \ + --turn-reminder "Answer tersely." # also runs each arm as +reminder +``` + +Every turn is recorded under the arm's `session.turns`, with `metrics` for that +call (prompt, completion, cache read and write tokens, cost). The arm's +`token_usage` and `metrics` are the sums over turns, so `cost_summary` reads +them like any other run. A question with fewer than N-1 follow-ups fails the +run before any answer is generated. Agentic arms (`mcp:`, `skill-agent:`, +`skill-catalog:`) are refused per arm, and a session cannot be judged: +`--judge` is refused with `--session-turns`, because the judge would score the +last answer against the first question. The reminder is appended to every user +turn in a `` block; the assistant turns are not changed. + +A session that fails part-way (a turn raises, or a turn comes back empty and the +session stops) is an errored arm: `error` is `generation_failed: turn N: ...` or +`empty_reply: turn N` (turns count from 0), and `answer` is `None`. The turns already paid for stay in +`session.turns`, and the arm's `token_usage` and `metrics` cover them, so the cost +rollup counts what was spent. The artifact's `run_config` records `session_turns` +and `turn_reminder`. `plugin-delta` refuses to pair a session run with a +single-turn run or one with different session settings, and the matrix rollup +does not pair them. + ### Trusting the judge A significance table says the delta is unlikely under *this judge's* noise. It diff --git a/dokimi/dokimi/commands/arms.py b/dokimi/dokimi/commands/arms.py index e07da1f8..d462e12c 100644 --- a/dokimi/dokimi/commands/arms.py +++ b/dokimi/dokimi/commands/arms.py @@ -46,6 +46,8 @@ "judge_panel", "concurrency", "stats_correction", + "session_turns", + "turn_reminder", ) @@ -289,6 +291,14 @@ def register(sub, positive_int) -> None: "'technical_expert,developer_advocate,doc_reviewer'). When " "set, --judge scores each answer with a panel and records " "grader spread and agreement instead of one opinion.") + arms_run_p.add_argument("--session-turns", type=positive_int, default=None, dest="session_turns", + help="Measure multi-turn sessions: each question is asked as its question " + "followed by the first N-1 entries of its 'follow_ups' list, with the " + "history kept, and each turn's token and cache usage recorded. Non-agentic " + "arms only. Measurement only: refused with --judge") + arms_run_p.add_argument("--turn-reminder", default=None, dest="turn_reminder", + help="With --session-turns, also run each arm as '+reminder', with this " + "text appended to every user turn in a block") arms_run_p.add_argument("--concurrency", type=positive_int, default=5) arms_run_p.add_argument("--out-json", default=None, dest="out_json", help="Output JSON path (default: results/arms/{product}.json)") @@ -562,6 +572,12 @@ def _probe_sampling(model: str, provider: str, temperature, seed) -> None: ) from exc +def _refuse_judged_session(args: argparse.Namespace) -> None: + """Refused before any answer is paid for: a session is measured, not judged.""" + if getattr(args, "session_turns", None) is not None and args.judge: + raise ValueError("--session-turns is measurement-only; drop --judge") + + def _build_arms_output(args: argparse.Namespace, specs: list[str], questions: list[dict]) -> dict: from dokimi.eval import Judge, JudgePanel from dokimi.eval.panel import JUDGE_ROLES, JudgeConfig @@ -574,6 +590,7 @@ def _build_arms_output(args: argparse.Namespace, specs: list[str], questions: li ) from dokimi.treatments import create_treatments + _refuse_judged_session(args) treatments = create_treatments(specs, top_k=args.top_k, rerank_threshold=args.rerank_threshold) _validate_expected_skills(questions, treatments) try: @@ -657,6 +674,8 @@ def _build_arms_output(args: argparse.Namespace, specs: list[str], questions: li stats_correction=getattr(args, "stats_correction", "holm"), temperature=getattr(args, "temperature", None), seed=getattr(args, "seed", None), + session_turns=getattr(args, "session_turns", None), + turn_reminder=getattr(args, "turn_reminder", None), ) records = runner.run( library_name=args.product, diff --git a/dokimi/dokimi/eval/arm_runner.py b/dokimi/dokimi/eval/arm_runner.py index 0b930559..1c63ecd6 100644 --- a/dokimi/dokimi/eval/arm_runner.py +++ b/dokimi/dokimi/eval/arm_runner.py @@ -37,6 +37,7 @@ from dokimi.eval.judge_context import build as build_judge_context from dokimi.eval.judge_context import cite as cite_judge_materials from dokimi.eval.judge_context import join_injected +from dokimi.eval.session import REMINDER_ARM_SUFFIX, run_session, session_turns_of from dokimi.metrics.usage import to_record from dokimi.treatments.base import Treatment @@ -109,6 +110,16 @@ def _static_arm_kind(treatment: Treatment) -> Dict[str, Any]: return {"arm_kind": getattr(t, "arm_kind", None)} + +def _refuse_session_records(recs: Sequence[Record]) -> None: + """Refuse to judge a session run: its last answer would be scored against its first question.""" + if any(r.raw.get("session_turns") for r in recs): + raise ValueError( + "Session records are measurement-only: the judge would score a session's last " + "answer against its first question. Judge a single-turn run instead." + ) + + class ArmRunner: """Generate answers for each treatment arm and (optionally) judge them.""" @@ -133,9 +144,17 @@ def __init__( stats_correction: str = "holm", temperature: Optional[float] = None, seed: Optional[int] = None, + session_turns: Optional[int] = None, + turn_reminder: Optional[str] = None, ): if not treatments: raise ValueError("ArmRunner requires at least one treatment arm") + if session_turns is not None and session_turns < 2: + raise ValueError( + f"session_turns must be at least 2 (one turn is the plain run), got {session_turns}" + ) + if turn_reminder is not None and (session_turns is None or not turn_reminder.strip()): + raise ValueError("turn_reminder needs session_turns and non-blank text") # Rejected here rather than at summarize time: the correction is only # applied after every arm has been answered and judged, so a typo would # otherwise surface after the whole run has been paid for. @@ -153,6 +172,8 @@ def __init__( self.max_tokens = max_tokens self.temperature = temperature self.seed = seed + self.session_turns = session_turns + self.turn_reminder = turn_reminder self.grounded_from_topics = grounded_from_topics self.harness = harness self.matrix_cell = matrix_cell @@ -168,7 +189,10 @@ def __init__( @property def arm_names(self) -> List[str]: - return [t.name for t in self.treatments] + names = [t.name for t in self.treatments] + if self.turn_reminder is not None: + names += [name + REMINDER_ARM_SUFFIX for name in names] + return names def _run_config(self) -> Dict[str, Any]: """Request-shaping settings, recorded so a run stays interpretable later. @@ -183,13 +207,19 @@ def _run_config(self) -> Dict[str, Any]: non-deterministic by default, so a pinned run and a free-sampling one are differently reproducible and should not be silently pooled as replicates. """ - return { + config: Dict[str, Any] = { "max_tokens": self.max_tokens, "max_iterations": self.max_iterations, "grounded_from_topics": self.grounded_from_topics, "temperature": self.temperature, "seed": self.seed, } + # A session run and a single-turn run answer different questions of the model (the + # history grows), so the pairing checks must be able to tell them apart. + if self.session_turns is not None: + config["session_turns"] = self.session_turns + config["turn_reminder"] = self.turn_reminder + return config def run( self, @@ -200,6 +230,12 @@ def run( ) -> List[Dict[str, Any]]: """Answer every question under every arm. Returns one record per question.""" n = len(questions) + if self.session_turns is not None: + # Before any answer is paid for: a question short of follow-ups fails the run. + for i, q in enumerate(questions): + qid = question_id(q) + session_turns_of(q, question_text_of(q), self.session_turns, + label=f"#{i}" if qid is None else str(qid)) results: Dict[int, Dict[str, Any]] = {} lock = threading.Lock() done = [0] @@ -230,11 +266,25 @@ def _answer_question( if qid is None: qid = "unknown_q" + turns = None + if self.session_turns is not None: + turns = session_turns_of(q, question_text, self.session_turns, label=qid) + arms: Dict[str, Any] = {} for treatment in self.treatments: - arms[treatment.name] = self._answer_one_arm( - treatment, question_text, library_name, library_id, qid + if turns is None: + arms[treatment.name] = self._answer_one_arm( + treatment, question_text, library_name, library_id, qid + ) + continue + arms[treatment.name] = self._answer_session_arm( + treatment, turns, library_name, library_id, qid ) + if self.turn_reminder is not None: + arms[treatment.name + REMINDER_ARM_SUFFIX] = self._answer_session_arm( + treatment, turns, library_name, library_id, qid, + reminder=self.turn_reminder, + ) rec = { "question_id": qid, @@ -245,6 +295,8 @@ def _answer_question( "persona": q.get("persona"), "arms": arms, } + if turns is not None: + rec["session_turns"] = len(turns) # Selection label: present-and-None means "no skill should load", which is # why presence is tested rather than truthiness. if "expected_skill" in q: @@ -365,6 +417,109 @@ def _answer_one_arm( "elapsed_sec": round(time.time() - t0, 2), } + def _answer_session_arm( + self, + treatment: Treatment, + turns: List[str], + library_name: str, + library_id: Optional[str], + question_id: str, + reminder: Optional[str] = None, + ) -> Dict[str, Any]: + """Answer one arm's multi-turn session; see :mod:`dokimi.eval.session`.""" + t0 = time.time() + try: + cfg = treatment.prepare(turns[0], library_name, library_id) + except Exception as exc: + logger.exception("Arm %s prepare() failed", treatment.name) + return {"error": f"prepare_failed: {exc}", "answer": None, + "metadata": _static_arm_kind(treatment)} + + from dokimi.plugins import take_output_shapers + shapers, cfg_metadata = take_output_shapers(cfg) + if cfg.is_agentic: + return {"error": "session_unsupported: an agentic arm (tools) is answered in one " + "turn only", "answer": None, "metadata": cfg_metadata} + + # The opening message is the single-turn prompt, context included, so turn 1 of + # a session is the answer the plain run would have given. + if cfg.has_context: + context = join_injected(cfg.injected_context) + opening = (CONTEXT_PROMPT_PREFIX.format(context=context[:MAX_CONTEXT_CHARS]) + + CONTEXT_PROMPT_SUFFIX.format(question=turns[0])) + else: + opening = ANSWER_PROMPT_PLAIN.format(question=turns[0]) + + session = run_session( + model=self.model, provider=self.provider, api_key=self.api_key, + opening=opening, follow_ups=turns[1:], system=cfg.system_prompt, + max_tokens=self.max_tokens, temperature=self.temperature, seed=self.seed, + reminder=reminder, + ) + session_block = { + "turns_requested": len(turns), + "turns": session["turns"], + "stopped_reason": session["stopped_reason"], + "reminder": reminder is not None, + } + if session["error"] is not None: + logger.error("Arm %s session failed for %s: %s", treatment.name, + question_id, session["error"]) + return {"error": f"generation_failed: {session['error']}", "answer": None, + **self._session_spend(session), + "session": session_block, "metadata": cfg_metadata} + + if session["stopped_reason"] == "empty_reply": + # The session ended short of the turns asked for. Recording it as a success + # would mix 2-turn and 3-turn sessions in one arm's totals; the usage stays. + error = f"empty_reply: turn {len(session['turns']) - 1}" + logger.error("Arm %s session stopped for %s: %s", treatment.name, question_id, error) + return {"error": error, "answer": None, **self._session_spend(session), + "session": session_block, "metadata": cfg_metadata} + + raw_answer = session["answer"] + final_text = self._apply_shapers(shapers, raw_answer) + usage = session["usage"] + return { + "answer": final_text, + "raw_answer": raw_answer if shapers else None, + "model": self.model, + "harness": self.harness, + "plugin_set": self.plugin_set["plugin_set"], + "plugin_set_id": self.plugin_set["plugin_set_id"], + "used_system_prompt": bool(cfg.system_prompt), + "context_chunks": [ + { + "source": c.get("source", "unknown"), + "snippet": (c["content"][:200] + "…") if len(c["content"]) > 200 else c["content"], + "relevance_score": c.get("relevance_score"), + } + for c in cfg.injected_context + ], + # Whole-session totals, so the cost and token rollups read this arm like any other. + "token_usage": usage.as_token_usage_dict(), + "metrics": usage.as_metrics_dict( + answer_chars=len(raw_answer or ""), final_chars=len(final_text or "")), + "session": session_block, + "metadata": cfg_metadata, + "elapsed_sec": round(time.time() - t0, 2), + } + + @staticmethod + def _session_spend(session: Dict[str, Any]) -> Dict[str, Any]: + """Token usage and metrics of the turns a failed session did complete. + + The calls were paid for, so the cost rollup must count them; an arm without + ``metrics`` is skipped by it. No answer-length keys are written: a failed + session has no answer, and plugin-delta averages those keys over every arm + that has a ``metrics`` block. + """ + if not session["turns"]: + return {} + usage = session["usage"] + return {"token_usage": usage.as_token_usage_dict(), + "metrics": usage.as_metrics_dict()} + @staticmethod def _apply_shapers(shapers, answer_text): """Run output shapers over a raw answer (no-op when there are none).""" @@ -497,6 +652,7 @@ def _score(rec: Record) -> Dict[str, Any]: recs = as_records(records, source="answers", kind="answers") # Before any judge call: refuse, do not grade against "". require_question_texts(r.raw for r in recs) + _refuse_session_records(recs) n = len(recs) out: Dict[int, Dict[str, Any]] = {} with ThreadPoolExecutor(max_workers=concurrency) as pool: @@ -539,6 +695,8 @@ def build_output( "agentic_usage_summary": self._summarize_agentic_usage(records), "skill_selection_summary": summarize_skill_selection(records), } + if self.session_turns is not None: + out["session"] = {"turns": self.session_turns, "turn_reminder": self.turn_reminder} # One copy of each judge context the arm records cite, written whether or # not the run was judged, so an answers-only artifact can be judged later. materials = cite_judge_materials(records, self.judge_context) diff --git a/dokimi/dokimi/eval/matrix_rollup.py b/dokimi/dokimi/eval/matrix_rollup.py index d56639c0..4eb4fef1 100644 --- a/dokimi/dokimi/eval/matrix_rollup.py +++ b/dokimi/dokimi/eval/matrix_rollup.py @@ -124,6 +124,8 @@ def _pair_key(data: dict[str, Any]) -> tuple[Any, ...]: tuple(data.get("arms") or []), data.get("baseline_arm"), data.get("total_questions"), + (data.get("run_config") or {}).get("session_turns"), + (data.get("run_config") or {}).get("turn_reminder"), ) diff --git a/dokimi/dokimi/eval/plugin_delta.py b/dokimi/dokimi/eval/plugin_delta.py index e6a48c67..9399eba9 100644 --- a/dokimi/dokimi/eval/plugin_delta.py +++ b/dokimi/dokimi/eval/plugin_delta.py @@ -58,6 +58,18 @@ def compare_plugin_runs(control: Dict[str, Any], plugin: Dict[str, Any]) -> Dict } +def _validate_session_config(control: Dict[str, Any], plugin: Dict[str, Any]) -> None: + """Request-shaping settings of a multi-turn session run; absent on a single-turn run.""" + for field in ("session_turns", "turn_reminder"): + before = (control.get("run_config") or {}).get(field) + after = (plugin.get("run_config") or {}).get(field) + if before != after: + raise PluginDeltaError( + f"cannot compute plugin delta across different {field}: " + f"{before!r} vs {after!r}" + ) + + def _validate_pair(control: Dict[str, Any], plugin: Dict[str, Any]) -> None: for label, artifact in (("baseline", control), ("plugin", plugin)): if artifact.get("schema_version") != "arms.v1": @@ -73,6 +85,7 @@ def _validate_pair(control: Dict[str, Any], plugin: Dict[str, Any]) -> None: f"cannot compute plugin delta across different {field}: " f"{control.get(field)!r} vs {plugin.get(field)!r}" ) + _validate_session_config(control, plugin) if control.get("plugin_set") == plugin.get("plugin_set"): raise PluginDeltaError( "cannot compute plugin delta for identical plugin_set values " diff --git a/dokimi/dokimi/eval/session.py b/dokimi/dokimi/eval/session.py new file mode 100644 index 00000000..113df64e --- /dev/null +++ b/dokimi/dokimi/eval/session.py @@ -0,0 +1,133 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-turn session measurement for one arm: history kept, usage per turn. + +:class:`~dokimi.eval.arm_runner.ArmRunner` answers each question in one call. +With ``session_turns`` set, a question is asked as its opening question followed +by its ``follow_ups``; the conversation so far is sent with every turn, and the +token usage (input, output, cache read and write, cost) of each turn is recorded. +That is what a per-turn cost or style claim needs: a single-shot answer cannot +show what a skill or instruction costs once the context has grown. + +Measurement only. Sessions are not judged: the judge scores one answer against one +question, and a session's last answer is not that. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional, Sequence + +from dokimi.llm import LLM +from dokimi.metrics.usage import UsageRecord + +logger = logging.getLogger(__name__) + +#: Appended to an arm's name for its reminder variant (the same arm, with the +#: reminder text added to every user turn). +REMINDER_ARM_SUFFIX = "+reminder" + +#: Question field holding the follow-up user turns, after the question itself. +FOLLOW_UP_FIELD = "follow_ups" + + +def session_turns_of(question: Dict[str, Any], question_text: str, + n_turns: int, label: str) -> List[str]: + """The user turns of one question's session: its text, then ``n_turns - 1`` follow-ups. + + Raises ``ValueError`` rather than running a shorter session, so every record of a + run has the same number of turns. + """ + follow_ups = question.get(FOLLOW_UP_FIELD) or [] + if not isinstance(follow_ups, list) or not all( + isinstance(t, str) and t.strip() for t in follow_ups): + raise ValueError(f"question {label}: '{FOLLOW_UP_FIELD}' must be a list of non-blank strings") + if len(follow_ups) < n_turns - 1: + raise ValueError( + f"question {label} has {len(follow_ups)} '{FOLLOW_UP_FIELD}' entr" + f"{'y' if len(follow_ups) == 1 else 'ies'}; --session-turns {n_turns} needs {n_turns - 1}" + ) + return [question_text, *follow_ups[: n_turns - 1]] + + +def _with_reminder(text: str, reminder: Optional[str]) -> str: + if reminder is None: + return text + return f"{text}\n\n{reminder}" + + +def run_session( + *, + model: str, + provider: str, + opening: str, + follow_ups: Sequence[str], + system: Optional[str] = None, + api_key: Optional[str] = None, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + seed: Optional[int] = None, + reminder: Optional[str] = None, +) -> Dict[str, Any]: + """Answer a conversation turn by turn, sending the full history each time. + + ``opening`` is the first user message (the prompt the single-turn path would + send, context included); each ``follow_ups`` entry is a later user message. + ``reminder``, when set, is appended to every user message. + + Returns ``answer`` (the last assistant text, ``None`` if no turn completed), + ``turns`` (per turn: ``user_chars``, ``answer``, ``finish_reason``, + ``truncated`` and ``metrics``, the :meth:`UsageRecord.as_metrics_dict` of that + call), ``stopped_reason`` (``completed``, ``empty_reply`` or ``error``), + ``error`` (``None`` unless a turn raised), and ``usage`` (the sum over turns). + A turn that raises ends the session; the completed turns are kept. + """ + llm = LLM(model=model, provider=provider, api_key=api_key, max_tokens=max_tokens, + temperature=temperature, seed=seed) + messages: List[Dict[str, Any]] = [] + if system: + messages.append({"role": "system", "content": system}) + + total = UsageRecord(model=model, provider=provider, n_calls=0) + turns: List[Dict[str, Any]] = [] + answer: Optional[str] = None + stopped_reason = "completed" + error: Optional[str] = None + + for i, text in enumerate([opening, *follow_ups]): + content = _with_reminder(text, reminder) + messages.append({"role": "user", "content": content}) + try: + # A copy: the history grows after this call, and no adapter may see it grow. + reply = llm.chat(list(messages)) + except Exception as exc: + logger.exception("Session turn %d failed", i) + error = f"turn {i}: {exc}" + stopped_reason = "error" + break + text_out = reply.text or "" + messages.append({"role": "assistant", "content": text_out}) + total = total + reply.usage + answer = text_out + turns.append({ + "turn": i, + "user_chars": len(content), + "answer": text_out, + "finish_reason": reply.finish_reason, + "truncated": reply.truncated, + "metrics": reply.usage.as_metrics_dict(answer_chars=len(text_out)), + }) + if not text_out.strip(): + # An empty assistant message cannot be sent back as history, and a + # provider rejects it; the rest of the session is not asked. + stopped_reason = "empty_reply" + break + + return { + "answer": answer if turns else None, + "turns": turns, + "stopped_reason": stopped_reason, + "error": error, + "usage": total, + } diff --git a/dokimi/tests/test_session_runner.py b/dokimi/tests/test_session_runner.py new file mode 100644 index 00000000..80eb53ac --- /dev/null +++ b/dokimi/tests/test_session_runner.py @@ -0,0 +1,446 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-turn session measurement: history kept across turns, usage per turn. + +Every test runs against the ``fake_llm`` adapter, so nothing reaches a provider. +""" + +from argparse import Namespace +from unittest.mock import MagicMock + +import pytest + +from conftest import DEMO_SKILL_MARKER +from dokimi.eval.arm_runner import ANSWER_PROMPT_PLAIN, ArmRunner +from dokimi.eval.arm_summary import cost_summary +from dokimi.llm import Reply +from dokimi.metrics.usage import UsageRecord +from dokimi.skills import load_skill +from dokimi.treatments import BaselineTreatment, SkillTreatment +from dokimi.treatments.base import AgentConfig, Tool, Treatment + +SESSION_Q = { + "id": "q1", + "question": "How do I parallelize a loop?", + "follow_ups": ["Now with a reduction?", "What does it cost?"], +} + + +def _roles(messages): + return [m["role"] for m in messages] + + +def test_session_sends_the_full_history_each_turn(fake_llm): + fake_llm.push("a1", "a2", "a3") + runner = ArmRunner([BaselineTreatment()], session_turns=3) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + reqs = fake_llm.requests + assert [_roles(r.messages) for r in reqs] == [ + ["user"], + ["user", "assistant", "user"], + ["user", "assistant", "user", "assistant", "user"], + ] + # The opening is the plain single-turn prompt; follow-ups are sent as written. + assert reqs[0].messages[0]["content"] == ANSWER_PROMPT_PLAIN.format( + question="How do I parallelize a loop?") + assert reqs[1].messages[1]["content"] == "a1" + assert reqs[2].messages[2]["content"] == "Now with a reduction?" + assert reqs[2].messages[3]["content"] == "a2" + assert reqs[2].messages[4]["content"] == "What does it cost?" + + arm = records[0]["arms"]["baseline"] + assert arm["answer"] == "a3" + assert [t["answer"] for t in arm["session"]["turns"]] == ["a1", "a2", "a3"] + assert arm["session"]["stopped_reason"] == "completed" + assert records[0]["session_turns"] == 3 + + +def test_skill_context_is_in_the_opening_turn_only(fake_llm, demo_skill): + fake_llm.push("a1", "a2") + runner = ArmRunner([SkillTreatment(load_skill(demo_skill))], session_turns=2) + + runner.run("oneTBB", [SESSION_Q], concurrency=1) + + first, second = fake_llm.requests + assert DEMO_SKILL_MARKER in first.messages[-1]["content"] + assert DEMO_SKILL_MARKER not in second.messages[-1]["content"] + + +def test_per_turn_usage_with_cache_tokens_is_recorded(fake_llm): + fake_llm.push( + Reply(text="a1", usage=UsageRecord(prompt_tokens=100, completion_tokens=10, + total_tokens=110, cache_write_tokens=60, + cost_usd=0.001)), + Reply(text="a2", usage=UsageRecord(prompt_tokens=180, completion_tokens=12, + total_tokens=192, cache_read_tokens=60, + cost_usd=0.002)), + Reply(text="a3", usage=UsageRecord(prompt_tokens=200, completion_tokens=9, + total_tokens=209, cache_read_tokens=160, + cost_usd=0.003)), + ) + runner = ArmRunner([BaselineTreatment()], session_turns=3) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + arm = records[0]["arms"]["baseline"] + turns = arm["session"]["turns"] + assert [t["metrics"]["prompt_tokens"] for t in turns] == [100, 180, 200] + assert [t["metrics"]["completion_tokens"] for t in turns] == [10, 12, 9] + assert [t["metrics"]["cache_read_tokens"] for t in turns] == [0, 60, 160] + assert [t["metrics"]["cache_write_tokens"] for t in turns] == [60, 0, 0] + assert turns[1]["metrics"]["cost_usd"] == pytest.approx(0.002) + + # The arm-level totals are the sum over turns, read by the usual rollups. + assert arm["token_usage"] == {"prompt_tokens": 480, "completion_tokens": 31, + "total_tokens": 511} + assert arm["metrics"]["n_calls"] == 3 + assert arm["metrics"]["cache_read_tokens"] == 220 + summary = cost_summary(records)["baseline"] + assert summary["total_llm_calls"] == 3 + assert summary["total_cost_usd"] == pytest.approx(0.006) + assert summary["cache_hit_ratio"] == pytest.approx(220 / 480, abs=1e-4) + + +def test_reminder_arm_repeats_the_reminder_on_every_user_turn(fake_llm): + fake_llm.push("a1", "a2", "b1", "b2") + runner = ArmRunner([BaselineTreatment()], session_turns=2, turn_reminder="BE TERSE") + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + assert set(records[0]["arms"]) == {"baseline", "baseline+reminder"} + assert runner.arm_names == ["baseline", "baseline+reminder"] + plain, reminded = fake_llm.requests[:2], fake_llm.requests[2:] + for req in plain: + for m in req.messages: + if m["role"] == "user": + assert "system-reminder" not in m["content"] + for req in reminded: + users = [m["content"] for m in req.messages if m["role"] == "user"] + assert all(u.endswith("\n\nBE TERSE") for u in users) + # Only user turns carry the reminder; the history is what the model said. + assert all("system-reminder" not in m["content"] + for m in req.messages if m["role"] == "assistant") + assert records[0]["arms"]["baseline+reminder"]["session"]["reminder"] is True + assert records[0]["arms"]["baseline"]["session"]["reminder"] is False + + +def test_sessions_are_not_judged(fake_llm): + fake_llm.push("a1", "a2") + runner = ArmRunner([BaselineTreatment()], session_turns=2) + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + with pytest.raises(ValueError, match="measurement-only"): + runner.judge(MagicMock(), "oneTBB", records) + + +def test_a_question_short_of_follow_ups_fails_before_any_answer_is_paid(fake_llm): + runner = ArmRunner([BaselineTreatment()], session_turns=3) + short = {"id": "q9", "question": "Short?", "follow_ups": ["only one"]} + + with pytest.raises(ValueError, match="q9.*needs 2"): + runner.run("oneTBB", [SESSION_Q, short], concurrency=1) + assert fake_llm.requests == [] + + +def test_an_empty_reply_ends_the_session_and_is_a_failed_arm_that_keeps_its_spend(fake_llm): + fake_llm.push("a1", "", "never sent") + runner = ArmRunner([BaselineTreatment()], session_turns=3) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + arm = records[0]["arms"]["baseline"] + assert len(fake_llm.requests) == 2 + assert arm["session"]["stopped_reason"] == "empty_reply" + assert len(arm["session"]["turns"]) == 2 + # Not a success: its totals would cover fewer turns than the other questions' rows. + assert arm["error"] == "empty_reply: turn 1" + assert arm["answer"] is None + assert arm["metrics"]["n_calls"] == 2 + assert cost_summary(records)["baseline"]["total_llm_calls"] == 2 + + +def test_a_failing_turn_keeps_the_completed_turns(fake_llm): + fake_llm.push("a1", ValueError("boom")) + runner = ArmRunner([BaselineTreatment()], session_turns=3) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + arm = records[0]["arms"]["baseline"] + assert arm["answer"] is None + assert arm["error"] == "generation_failed: turn 1: boom" + assert arm["session"]["stopped_reason"] == "error" + assert [t["answer"] for t in arm["session"]["turns"]] == ["a1"] + assert len(fake_llm.requests) == 2 + + +def test_a_failed_session_still_counts_the_turns_already_paid_for(fake_llm): + fake_llm.push(Reply(text="a1", usage=UsageRecord(prompt_tokens=100, completion_tokens=10, + total_tokens=110, cost_usd=0.01)), + ValueError("boom")) + runner = ArmRunner([BaselineTreatment()], session_turns=3) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + arm = records[0]["arms"]["baseline"] + assert arm["token_usage"]["total_tokens"] == 110 + summary = cost_summary(records)["baseline"] + assert summary["total_llm_calls"] == 1 + assert summary["total_cost_usd"] == pytest.approx(0.01) + + +def test_a_failed_session_does_not_skew_the_answer_length_comparison(fake_llm): + # Control: q1 answers 100 chars at turn 0 and then returns an empty reply, q2 succeeds. + # Plugin: both succeed. Every real answer is 100 chars on both sides. + questions = [SESSION_Q, {**SESSION_Q, "id": "q2"}] + fake_llm.push("x" * 100, "", "x" * 100, "x" * 100, "x" * 100) + control = ArmRunner([BaselineTreatment()], session_turns=2) + control_records = control.run("oneTBB", questions, concurrency=1) + failed = control_records[0]["arms"]["baseline"] + assert failed["error"] == "empty_reply: turn 1" + assert "raw_answer_chars" not in failed["metrics"] and "final_answer_chars" not in failed["metrics"] + assert failed["metrics"]["n_calls"] == 2 # the spend stays visible to the cost rollup + + fake_llm.push("x" * 100, "x" * 100, "x" * 100, "x" * 100) + plugin_records = ArmRunner([BaselineTreatment()], session_turns=2).run( + "oneTBB", questions, concurrency=1) + + def artifact(records, plugin_set): + out = _artifact(session_turns=2, plugin_set=plugin_set) + out.update(answers=records, total_questions=2) + return out + + from dokimi.eval.plugin_delta import _answer_text_deltas + + rows = _answer_text_deltas(artifact(control_records, "none"), artifact(plugin_records, "caveman"), + ["baseline"])["baseline"] + for key in ("raw_answer_chars", "final_answer_chars"): + assert rows[key] == {"baseline_avg": 100.0, "plugin_avg": 100.0, "delta": 0.0} + + +def test_a_session_that_fails_on_its_first_turn_has_no_spend_to_report(fake_llm): + fake_llm.push(ValueError("boom")) + records = ArmRunner([BaselineTreatment()], session_turns=2).run( + "oneTBB", [SESSION_Q], concurrency=1) + arm = records[0]["arms"]["baseline"] + assert "metrics" not in arm and "token_usage" not in arm + + +class _Noop(Tool): + name = "noop" + description = "Does nothing." + + def call(self, **kwargs) -> str: + return "ok" + + +class _AgenticTreatment(Treatment): + name = "agentic" + + def prepare(self, question_text, library_name, library_id=None): + return AgentConfig(tools=[_Noop()]) + + +def test_an_agentic_arm_is_refused_in_a_session(fake_llm): + runner = ArmRunner([_AgenticTreatment()], session_turns=2) + + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + arm = records[0]["arms"]["agentic"] + assert arm["error"].startswith("session_unsupported") + assert arm["answer"] is None + assert fake_llm.requests == [] + + +def test_session_options_are_validated_up_front(): + with pytest.raises(ValueError, match="at least 2"): + ArmRunner([BaselineTreatment()], session_turns=1) + with pytest.raises(ValueError, match="needs session_turns"): + ArmRunner([BaselineTreatment()], turn_reminder="BE TERSE") + with pytest.raises(ValueError, match="needs session_turns"): + ArmRunner([BaselineTreatment()], session_turns=2, turn_reminder=" ") + + +def test_the_artifact_names_the_session_settings(fake_llm): + fake_llm.push("a1", "a2") + runner = ArmRunner([BaselineTreatment()], session_turns=2, turn_reminder="BE TERSE") + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + out = runner.build_output("oneTBB", records) + + assert out["session"] == {"turns": 2, "turn_reminder": "BE TERSE"} + assert out["arms"] == ["baseline", "baseline+reminder"] + + +def test_plain_runs_are_unchanged_by_the_session_fields(fake_llm): + fake_llm.push("answer") + runner = ArmRunner([BaselineTreatment()]) + records = runner.run("oneTBB", [SESSION_Q], concurrency=1) + + out = runner.build_output("oneTBB", records) + + assert "session" not in out + assert "session_turns" not in records[0] + assert "session" not in records[0]["arms"]["baseline"] + + +def test_the_cli_refuses_judge_with_session_before_any_answer(fake_llm): + from dokimi.commands.arms import _build_arms_output + + args = Namespace(judge=True, session_turns=2, turn_reminder=None) + with pytest.raises(ValueError, match="measurement-only"): + _build_arms_output(args, ["baseline"], [SESSION_Q]) + assert fake_llm.requests == [] + + +class _SystemPromptTreatment(Treatment): + name = "sys" + + def prepare(self, question_text, library_name, library_id=None): + return AgentConfig(system_prompt="You are terse.") + + +class _BigContextTreatment(Treatment): + name = "big" + + def prepare(self, question_text, library_name, library_id=None): + return AgentConfig(injected_context=[{"content": "\u00a7" * 40_000, "source": "s"}]) + + +def test_every_request_of_a_session_carries_the_sampling_options_and_system_prompt(fake_llm): + fake_llm.push("a1", "a2", "a3") + runner = ArmRunner([_SystemPromptTreatment()], session_turns=3, temperature=0.3, seed=7, + max_tokens=50) + + runner.run("oneTBB", [SESSION_Q], concurrency=1) + + assert len(fake_llm.requests) == 3 + for req in fake_llm.requests: + assert (req.temperature, req.seed, req.max_tokens) == (0.3, 7, 50) + assert req.messages[0] == {"role": "system", "content": "You are terse."} + + +def test_the_opening_turn_caps_injected_context_as_the_plain_run_does(fake_llm): + from dokimi.eval.arm_runner import MAX_CONTEXT_CHARS + + fake_llm.push("a1", "a2") + ArmRunner([_BigContextTreatment()], session_turns=2).run( + "oneTBB", [SESSION_Q], concurrency=1) + + opening = fake_llm.requests[0].messages[-1]["content"] + assert opening.count("\u00a7") == MAX_CONTEXT_CHARS + # The context arm opens with the first question, not a follow-up. + assert SESSION_Q["question"] in opening + assert not any(follow_up in opening for follow_up in SESSION_Q["follow_ups"]) + + +def test_opening_is_prepared_from_the_first_turn_not_the_last(fake_llm): + seen = [] + + class _Spy(Treatment): + name = "spy" + + def prepare(self, question_text, library_name, library_id=None): + seen.append(question_text) + return AgentConfig() + + fake_llm.push("a1", "a2", "a3") + ArmRunner([_Spy()], session_turns=3).run("oneTBB", [SESSION_Q], concurrency=1) + assert seen == [SESSION_Q["question"]] + + +def test_output_shapers_apply_to_the_session_answer_but_not_the_raw_one(fake_llm): + from dokimi.plugins import create_plugins, plugin_set_metadata, wrap_treatments + + plugins = create_plugins(["truncate:5"]) + treatments = wrap_treatments([BaselineTreatment()], plugins) + fake_llm.push("a1", "a long final answer") + runner = ArmRunner(treatments, session_turns=2, plugin_set=plugin_set_metadata(plugins)) + + arm = runner.run("oneTBB", [SESSION_Q], concurrency=1)[0]["arms"]["baseline"] + + assert arm["raw_answer"] == "a long final answer" + assert arm["answer"] != arm["raw_answer"] and len(arm["answer"]) == 5 + + +def _parse_arms_run(argv): + from dokimi.cli import build_parser + + return build_parser().parse_args( + ["arms", "run", "--product", "oneTBB", "--questions", "q.json", + "--arms", "baseline"] + argv) + + +def test_arms_run_wires_the_session_flags_into_the_run(fake_llm, monkeypatch): + import dokimi.commands.arms as arms_mod + from dokimi.commands.arms import _build_arms_output + + monkeypatch.setattr(arms_mod, "_check_sampling_support", lambda args: None) + fake_llm.push("a1", "a2", "r1", "r2") + args = _parse_arms_run(["--session-turns", "2", "--turn-reminder", "BE TERSE"]) + + out = _build_arms_output(args, ["baseline"], [SESSION_Q]) + + assert out["session"] == {"turns": 2, "turn_reminder": "BE TERSE"} + assert out["arms"] == ["baseline", "baseline+reminder"] + assert out["run_config"]["session_turns"] == 2 + assert len(fake_llm.requests) == 4 + + +def test_arms_run_without_session_flags_runs_single_turn(fake_llm, monkeypatch): + import dokimi.commands.arms as arms_mod + from dokimi.commands.arms import _build_arms_output + + monkeypatch.setattr(arms_mod, "_check_sampling_support", lambda args: None) + fake_llm.push("a1") + out = _build_arms_output(_parse_arms_run([]), ["baseline"], [SESSION_Q]) + assert "session" not in out and "session_turns" not in out["run_config"] + + +def test_session_flags_parse_to_ints_and_text(): + args = _parse_arms_run(["--session-turns", "3", "--turn-reminder", "X"]) + assert (args.session_turns, args.turn_reminder) == (3, "X") + assert _parse_arms_run([]).session_turns is None + + +def _artifact(session_turns=None, plugin_set="none", reminder=None): + out = {"schema_version": "arms.v1", "library_name": "L", "model": "m", "provider": "p", + "harness": "h", "arms": ["baseline"], "baseline_arm": "baseline", + "total_questions": 1, "plugin_set": plugin_set, "answers": [], + "run_config": {"max_tokens": None}} + if session_turns is not None: + out["run_config"]["session_turns"] = session_turns + out["run_config"]["turn_reminder"] = reminder + return out + + +def test_a_session_run_is_not_paired_with_a_single_turn_run(): + from dokimi.eval.plugin_delta import PluginDeltaError, compare_plugin_runs + + with pytest.raises(PluginDeltaError, match="different session_turns: None vs 3"): + compare_plugin_runs(_artifact(), _artifact(session_turns=3, plugin_set="caveman")) + + +def test_a_run_with_a_turn_reminder_is_not_paired_with_one_without(): + from dokimi.eval.matrix_rollup import _pair_key + from dokimi.eval.plugin_delta import PluginDeltaError, compare_plugin_runs + + # Same turn count: only the request text differs, on every user turn. + with pytest.raises(PluginDeltaError, match="different turn_reminder: None vs 'R'"): + compare_plugin_runs(_artifact(3, reminder=None), _artifact(3, reminder="R", plugin_set="caveman")) + assert _pair_key(_artifact(3, reminder=None)) != _pair_key(_artifact(3, reminder="R")) + + +def test_the_matrix_pairing_key_separates_session_runs(): + from dokimi.eval.matrix_rollup import _pair_key + + assert _pair_key(_artifact()) != _pair_key(_artifact(session_turns=3)) + assert _pair_key(_artifact(session_turns=3)) == _pair_key(_artifact(session_turns=3)) + + +def test_the_run_config_records_the_session_settings_only_when_set(fake_llm): + assert "session_turns" not in ArmRunner([BaselineTreatment()])._run_config() + config = ArmRunner([BaselineTreatment()], session_turns=2, turn_reminder="R")._run_config() + assert (config["session_turns"], config["turn_reminder"]) == (2, "R")