Skip to content

feat(arms): multi-turn session measurement with per-turn usage - #174

Open
napetrov wants to merge 4 commits into
intel:mainfrom
napetrov:fix/skill-gap-W7
Open

napetrov wants to merge 4 commits into
intel:mainfrom
napetrov:fix/skill-gap-W7

Conversation

@napetrov

Copy link
Copy Markdown
Contributor

Part of #4.

Summary

dokimi arms run can now ask each question as a multi-turn session, with the history kept and the token usage recorded per turn.

 dokimi arms run ... --session-turns 3 [--turn-reminder TEXT]
   question: "How do I parallelize a loop?"            # turn 0: the plain single-turn prompt
   follow_ups: ["Now with a reduction?", "What does it cost?"]
   turn 0  send [user]                                  -> a1   usage: in/out/cache/cost
   turn 1  send [user, asst, user]                      -> a2   usage: in/out/cache/cost
   turn 2  send [user, asst, user, asst, user]          -> a3   usage: in/out/cache/cost
+  arm record: answer = a3, token_usage / metrics = sum over turns,
+              session.turns = per-turn metrics, stopped_reason
+  <arm>+reminder variant (only with --turn-reminder): same arm, reminder text
+              appended to every user turn, not to the assistant history

Off by default. Without --session-turns, the runner and its output are unchanged.

  • Sessions are measurement only. --judge is refused with --session-turns, because the judge would score the last answer against the first question.
  • Agentic arms (mcp:, skill-agent:, skill-catalog:) are refused in a session with session_unsupported.
  • A question with fewer than N-1 follow_ups fails before any answer is paid for. Every record of a run has the same number of turns.
  • A turn that raises ends the session and keeps the completed turns. An empty reply ends it with stopped_reason: empty_reply.
  • The new logic is in dokimi/eval/session.py, so arm_runner.py grows by one hook (_answer_session_arm) and the output shape gains a session block only in session mode.

Evidence

Tests (new file tests/test_session_runner.py, 13 tests) use the fake_llm adapter only. Nothing reaches a provider.

Fail-on-old: dokimi/eval/arm_runner.py and dokimi/commands/arms.py replaced with the intel/main versions, the rest of the branch kept:

FAILED test_the_cli_refuses_judge_with_session_before_any_answer
FAILED test_an_empty_reply_ends_the_session
FAILED test_session_options_are_validated_up_front
FAILED test_a_failing_turn_keeps_the_completed_turns
FAILED test_reminder_arm_repeats_the_reminder_on_every_user_turn
FAILED test_per_turn_usage_with_cache_tokens_is_recorded
FAILED test_an_agentic_arm_is_refused_in_a_session
FAILED test_a_question_short_of_follow_ups_fails_before_any_answer_is_paid
FAILED test_sessions_are_not_judged
FAILED test_skill_context_is_in_the_opening_turn_only
FAILED test_the_artifact_names_the_session_settings
FAILED test_session_sends_the_full_history_each_turn
12 failed, 1 passed in 1.62s   (test_plain_runs_are_unchanged_by_the_session_fields passes on both: it is the guard)

Local checks on the rebased head (rebased onto intel/main caae46a):

uv run ruff check .                              All checks passed!
uv run mypy dokimi                               Success: no issues found in 223 source files
uv run python -m dokimi.schema_check             exit 0
uv run python -m dokimi.config_check             exit 0
uv run pytest -n auto -q                         7081 passed, 22 skipped

Merge Danger

Door: two-way. The flags are opt-in and additive. Reverting the commit restores the old runner.

Blast Radius: low. The code runs only when --session-turns is passed.

Behaviour changes for callers:

  • Two new arms run flags. --session-turns N (N >= 2) and --turn-reminder TEXT. --turn-reminder without --session-turns is an error.
  • follow_ups is a new optional question field. Questions without it run as before. A session run needs N-1 entries on every question.
  • Session records are not judged. --judge with --session-turns fails up front, before any call is made.
  • Cost reading. token_usage and metrics of a session arm are the sums over all turns, so a session arm's cost is not comparable with a single-turn arm's cost for the same question.
  • Matrix and subjects parsers do not expose the new flags. They forward the two fields only when present, so no existing config changes meaning.

Checklist

  • Tests added/updated and pytest passes
  • ruff, mypy, and schema checks pass
  • Docs updated: docs/benchmarking-and-comparison.md (new "Multi-turn sessions" section)
  • changelog fragment added (changelog.d/)
  • No ad-hoc run artifacts committed
  • New code and test files carry the SPDX license header

🤖 Generated with Claude Code

Add --session-turns N to `dokimi arms run`. Each question is asked as its
opening question followed by its `follow_ups` list, with the conversation
kept across turns. Every turn records its token usage (input, output, cache
read and write, cost), and the arm totals are the sum over turns.

--turn-reminder TEXT adds a `<arm>+reminder` variant of each arm that
appends the reminder to every user turn, so the cost of a reminder can be
read against the same arm without it.

Sessions are measurement only: --judge is refused with --session-turns, and
agentic arms are refused in a session. A question with fewer follow-ups than
N-1 fails before any answer is paid for.

Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com>

@antonwolfy antonwolfy left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review at c2ccd81. The session design is sound: the history is copied per call, the reminder stays out of the assistant turns, and the judge and agentic refusals fire before any answer is paid for. Four things are worth fixing before merge, all inline: the cost of a failed session, the empty-reply path, the pairing commands ignoring session runs, and untested CLI wiring and options.

What I ran, on a Linux review host at c2ccd81:

  • make ci-local: the test job passes (7082 passed, 21 skipped) and so do the changed-line coverage, changelog, schema-check, schema-compat, package-smoke and gitleaks jobs. The ruff, mypy, zizmor, pre-commit, bandit, deptry, mkdocs and pip-audit jobs did not run in my venv: make dev stops at the tools install on this base, which #179 fixed on main. Static checks, mypy and the strict docs build pass on GitHub's CI at this head.
  • Reproductions for the code findings, with the LLM transport mocked (fake_llm).
  • A mutation run: 66 mutants of the new code against tests/test_session_runner.py, 8 of which pass the full file and are findings, plus 1 equivalent mutant. The 13 tests fail on the base modules except the guard test, as the description says.

Checked, no issues

  • Schema versioning: arms.v1.json has no additionalProperties: false on the record or arm, so session, session_turns and the arm session block are additive. A success run and a mid-session failure both validate.
  • Matrix and subjects: those parsers do not define the new flags, so forwarded_run_config forwards nothing and getattr(..., None) applies.
  • Usage sums: n_calls, cache read and write and cost add correctly over turns; cost_summary reads a successful session like any other arm.
  • Docs and fragment: the python cli.py arms run form matches the other examples on that page; the fragment name and format match the others. No new page, so no nav: change.

Comment thread dokimi/dokimi/eval/arm_runner.py
Comment thread dokimi/dokimi/eval/session.py
Comment thread dokimi/dokimi/eval/arm_runner.py
Comment thread dokimi/tests/test_session_runner.py
…sees session settings

A session that failed mid-way dropped token_usage/metrics so cost_summary lost
the completed turns; an empty reply stopped the session but was recorded as a
success. run_config now records session_turns/turn_reminder and plugin-delta
and the matrix pair key refuse to mix session and single-turn runs. Adds tests
for the arms run wiring and per-call session options.

Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com>
@codecov-commenter

codecov-commenter commented Oct 22, 2026 •

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 96.03175% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.70%. Comparing base (caae46a) to head (1a20b18).
⚠️ Report is 11 commits behind head on main.

Files with missing lines Patch % Lines
dokimi/dokimi/eval/arm_runner.py 95.58% 3 Missing ⚠️
dokimi/dokimi/eval/session.py 95.83% 1 Missing and 1 partial ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #174      +/-   ##
==========================================
+ Coverage   91.51%   91.70%   +0.18%     
==========================================
  Files         216      224       +8     
  Lines       27074    27629     +555     
  Branches     4457     4551      +94     
==========================================
+ Hits        24777    25337     +560     
+ Misses       1632     1624       -8     
- Partials      665      668       +3     
Files with missing lines Coverage Δ
dokimi/dokimi/commands/arms.py 89.18% <100.00%> (+0.14%) ⬆️
dokimi/dokimi/eval/matrix_rollup.py 82.00% <ø> (ø)
dokimi/dokimi/eval/plugin_delta.py 86.15% <100.00%> (+0.55%) ⬆️
dokimi/dokimi/eval/session.py 95.83% <95.83%> (ø)
dokimi/dokimi/eval/arm_runner.py 94.94% <95.58%> (+1.26%) ⬆️

... and 27 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@antonwolfy antonwolfy left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review at 41fc167. The four fixes from my earlier review are in and hold up: the error arms keep their spend, an empty reply is an errored arm, run_config records the session settings, and the CLI wiring and per-call options are tested. One bug comes from the spend fix itself, and two test gaps are left. Two nits on the docs.

What I ran, on a Linux review host at 41fc167:

  • make ci-local: 27 passed, 1 failed, 7 skipped. The test job passes (7114 passed, 1 skipped) and so do changed-line coverage, mypy, deptry, schema-compat, package-smoke, windows-smoke, mkdocs and gitleaks. pip-audit fails on gitpython 3.2.0 (CVE-2026-114145, fix in 3.2.1), which comes from the lock files this PR does not touch.
  • Reproductions for the code findings, with the LLM transport mocked (fake_llm).
  • A mutation run: 35 mutants of the new code against tests/test_session_runner.py, 4 of which pass the full file (3 are findings below, 1 is a harmless raw_answer duplicate).

Checked, no issues

  • Empty reply: empty_reply: turn N is the index of the empty turn, and its spend is kept. A failing or empty first turn gives a valid arms.v1 artifact and the report renders it.
  • Old artifacts: plugin_delta, matrix_rollup and model_compare.replicate_key treat a missing session_turns and turn_reminder as equal to None, so two old runs still pair.
  • Reminder arm: --turn-reminder gives [baseline, baseline+reminder], validates against the schema and renders in the report.

Comment on lines +506 to +507
usage = session["usage"]
return {"token_usage": usage.as_token_usage_dict(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A failed session records raw_answer_chars and final_answer_chars as 0, which skews the plugin-delta answer-length comparison.

_session_spend calls as_metrics_dict(answer_chars=0, final_chars=0), but the turns that completed did produce text. plugin_delta._answer_metrics (plugin_delta.py:225-241) averages those two fields over every arm that has a metrics dict. It does not look at error, and a plain failed arm has no metrics, so this is new with the spend fix.

Repro (fake_llm, --session-turns 3, every good reply is "x" * 100): in the control run q1 answers 100 chars at turn 0 and then returns an empty reply, q2 succeeds. In the plugin run both succeed. compare_plugin_runs returns raw_answer_chars: {baseline_avg: 50.0, plugin_avg: 100.0, delta: 50.0} (same for final_answer_chars). Every real answer is 100 chars on both sides, so the +50 is made up by the failed arm.

Suggestion: usage.as_metrics_dict() with no char arguments. The keys are then omitted (usage.py:126), _answer_metrics skips them, and cost_summary still sees metrics. Add a test that a failed session's metrics has no raw_answer_chars.

Comment on lines +387 to +398
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_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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

turn_reminder in the pairing checks is untested.

_artifact() always sets turn_reminder=None, and only session_turns differs between the two runs. Mutation check: both of these pass the full file (25 passed):

  • plugin_delta.py:76 with the tuple cut to ("session_turns",)
  • matrix_rollup.py:128 with the turn_reminder line removed

Then a run with a reminder and one without, at the same turn count, are paired as if comparable. Their request text differs on every user turn.

Suggestion: give _artifact a turn_reminder argument, and add compare_plugin_runs(_artifact(3, reminder=None), _artifact(3, reminder="R", plugin_set="caveman")) raising different turn_reminder, and _pair_key differing for the same case.

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]))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The opening turn of a context-bearing arm is not tested against the first question.

Mutation check: CONTEXT_PROMPT_SUFFIX.format(question=turns[-1]) here passes all 25 tests. test_opening_is_prepared_from_the_first_turn_not_the_last covers prepare(turns[0]) and the plain ANSWER_PROMPT_PLAIN branch (both mutants fail), but not this branch. With the mutant, a skill or context arm opens a 3-turn session with the last follow-up and the history is in the wrong order, for those arms only.

Suggestion: in the context-cap test (test_the_opening_turn_caps_injected_context_as_the_plain_run_does), also assert that the opening ends with the first question and does not contain the last follow-up.

Comment thread dokimi/docs/benchmarking-and-comparison.md Outdated
Comment thread dokimi/docs/benchmarking-and-comparison.md Outdated
Signed-off-by: Anton Volkov <antonwolfy@gmail.com>
Signed-off-by: Anton Volkov <antonwolfy@gmail.com>
@antonwolfy

Copy link
Copy Markdown
Contributor

I pushed fixes for the two doc nits from my review, one commit each:

  • 5ddf3c5: only plugin-delta refuses to pair a session run with a single-turn run; the matrix rollup does not pair them.
  • 1a20b18: the turn numbers in generation_failed: turn N and empty_reply: turn N count from 0.

GitHub's Tests now fail at 1a20b18 on every Python version, in test_growth_ratchet.py. commands/arms.py::_build_arms_output has complexity 14 against its baseline of 13, and eval/plugin_delta.py::_validate_pair has 16 against 14. The two commits change only benchmarking-and-comparison.md. The baselines came to main with the growth ratchet (#93) after this PR's earlier green run, and the PR's changes to those two functions take them past it. Splitting the two functions makes Tests green again.

Still open: the bug and the two risks from my review.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants