Repository navigation
feat(export): export schema v2 manifests from ClickHouse prompt chains - #36
Conversation
…kHouse chains `batchbench export-plans` (also `python -m batchbench.export_plans` and the `batchbench-export-plans` script) reads clay.prompt_chains joined with clay.http_analytics for a time window and writes a schema version 2 trajectory manifest. Sessions are reconstructed from the chains directly: within a principal a request continues the latest earlier request (within the link window) whose full chain is a strict prefix of its own, and each root with its descendants becomes one trajectory. Per request the exporter emits prompt/output tokens, inter-request delay, reset markers for sessions that began before the window, per-block seeds/roles/tokens with the reply to the previous request marked live, overhead tokens, and the recorded stream and max_tokens. Deterministic (optionally length-stratified) sampling and a time-scale factor shape the output, which is validated against the schema version 2 rules before it is written. The ClickHouse client is confined to one function behind the optional `export` extra; the pure logic works on row dictionaries and is covered by fixture tests, including a dry run of the exported manifest through batchbench-agent when the binary is available. The package now imports without the compiled extension so the exporter can run from a pure-Python checkout. Claude-Session: https://claude.ai/code/session_015Rtq2g2DyKJkqHgbi8qK6f
…review - Read chains through clay.prompt_chains_current by default (the view that warehouse readers are granted); FINAL on the base table is now opt-in. - Bound the analytics side of the join to the export window before joining instead of scanning the whole http_analytics table. - Fail instead of writing an empty manifest that batchbench-agent rejects. - Keep per-seed token definitions across resets so a re-tokenized shared block is normalized rather than redefined. - Reject sample fractions outside [0, 1] or non-finite. Claude-Session: https://claude.ai/code/session_015Rtq2g2DyKJkqHgbi8qK6f
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
🟡 Changes recommended
The current trajectory root resolution and per-root counter logic can produce unstable/incorrect trajectory_id values (especially for pre-window sessions and empty-chain roots).
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds a pure-Python exporter (batchbench export-plans) that converts ClickHouse prompt-chain + analytics rows into schema v2 trajectory manifests consumable by batchbench-agent --agent-plans-jsonl, along with docs, tests, packaging, and CI coverage.
Changes:
- Implement
src/batchbench/export_plans.py(ClickHouse query builder + row parsing + session reconstruction + v2 manifest generation + sampling + validator + CLI). - Make
batchbenchdispatchexport-plansbefore loading the Rust extension, and allow importingbatchbenchwithout_corefor pure-Python exports. - Add comprehensive unit tests for the exporter, document the workflow in README, add an optional
exportdependency, and run Python unit tests in CI.
File summaries
| File | Description |
|---|---|
src/batchbench/export_plans.py |
New ClickHouse-backed exporter producing schema v2 manifests with linking, sampling, validation, and a CLI. |
tests/test_export_plans.py |
Fixture-based unit tests covering linking, shaping, validation, query shape, and CLI paths. |
src/batchbench/__main__.py |
Routes batchbench export-plans to the pure-Python exporter before importing Rust _core. |
src/batchbench/__init__.py |
Makes batchbench importable without _core; API functions now error lazily when _core is missing. |
README.md |
Adds documentation for exporting schema v2 plans from production ClickHouse records. |
pyproject.toml |
Adds optional extra export (clickhouse-connect) and a batchbench-export-plans console script. |
.github/workflows/ci.yaml |
Runs Python unit tests via unittest discover. |
Review details
Suppressed comments (1)
src/batchbench/export_plans.py:471
build_plans()usesstem = trajectory.root.leaf or ""for the per-root counter map. If a root has an empty chain (leaf is None), all such trajectories share the same stem and get counters incremented across unrelated sessions, making thetrajectory_idsuffix depend on other empty-chain trajectories. Use the same stem logic astrajectory_id_for()instead.
for trajectory in trajectories:
stem = trajectory.root.leaf or ""
counter = counters.get(stem, 0)
counters[stem] = counter + 1
- Files reviewed: 7/7 changed files
- Comments generated: 4
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| try: | ||
| from . import _core | ||
| except ImportError: # pragma: no cover - pure-Python checkout without the compiled extension | ||
| _core = None |
| def resolve_root(record: ChainRecord) -> int: | ||
| path = [] | ||
| current = record | ||
| while current.index not in root_of: | ||
| parent_index = parents.get(current.index) | ||
| parent = by_index.get(parent_index) if parent_index is not None else None | ||
| if parent is None or not in_window(parent): | ||
| root_of[current.index] = current.index | ||
| began_before[current.index] = parent is not None | ||
| break | ||
| path.append(current.index) | ||
| current = parent | ||
| root = root_of[current.index] | ||
| for index in path: | ||
| root_of[index] = root | ||
| return root |
| the next request's start, `blocks` zips the chain hashes (as seeds) with the block | ||
| roles and token counts, `overhead_tokens` is `prompt_tokens - sum(block_tokens)`, and |
| Plans are ordered by a seeded hash of their trajectory id and the first ``round(fraction * | ||
| n)`` are kept. With stratification the same rule is applied per session-length bucket with | ||
| ``ceil`` instead of ``round``, so every non-empty bucket keeps at least one trajectory. | ||
| """ |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f4776e0afd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| known = seed_tokens.get(seed) | ||
| if known is None: | ||
| seed_tokens[seed] = tokens | ||
| elif known != tokens: | ||
| stats.normalized_block_tokens += 1 | ||
| tokens = known |
There was a problem hiding this comment.
When a repeated seed's recorded token count decreases—for example, from 100 tokens in the first request to 50 in a later 60-token prompt—this replaces the later count with 100, while the overhead calculation merely clamps the resulting negative remainder to zero. The manifest still passes validate_plans, but build_block_replay_plan replays the oversized blocks as written, so the request can substantially exceed its recorded prompt_tokens and corrupt the workload shape. Conflicting records should be split/dropped or normalized in a way that preserves every request's prompt target.
Useful? React with 👍 / 👎.
Stacked on the schema v2 PR (
feat/manifest-schema-v2); merge that first.What
batchbench export-plans(alsopython -m batchbench.export_plansand thebatchbench-export-plansscript) reads a window ofclay.prompt_chainsrecords joined withclay.http_analyticsand writes a schema version 2 trajectory manifest forbatchbench-agent --agent-plans-jsonl.--clickhouse-url/--clickhouse-user/--clickhouse-password(orCLICKHOUSE_URL/USER/PASSWORD),--start/--end(RFC 3339), optional--principal-id,--model,--served-by,--sample(deterministic with--seed),--stratify-by-session-length(power-of-two request-count buckets, at least one trajectory per non-empty bucket),--time-scale,--output.--dump-rows-jsonl/--rows-jsonlsave and replay fetched rows so an export can be reproduced without the warehouse.--link-window-hours, default 24) whose full chain is a strict prefix of B's chain; cumulative hashing makes that a lookup of A's leaf hash inside B's chain. Each root and its descendants form one trajectory ordered by timestamp. The query starts one link window before--start, so a session that began earlier getsreset_beforeon its first in-window request. Branches (a second child of the same parent) also getreset_beforeand no live block.prompt_tokens,output_tokens = completion_tokens,delay_after_ms = next.ts − (this.ts + duration_ms)(clamped at 0),blocks= chain hashes (hex) as seeds zipped with roles and token counts, the block immediately after the previous request's full chain markedlivewhen it is an assistant block,overhead_tokens = max(prompt_tokens − sum(block_tokens), 0),streamandmax_tokenswhen recorded. Requests without block token counts become a singleuserblock ofprompt_tokenswith zero overhead. Requests without positive prompt or completion counts, or with block arrays that disagree with the chain, are dropped and counted. Token counts for a seed seen earlier in a trajectory are normalized to the first value so the replay's per-seed cache stays consistent.start_after_msand validated against the schema v2 rules (a Python-side validator mirroring the Rust checks) before writing; an empty export is an error.fetch_rowsbehind the new optional extraexport = ["clickhouse-connect>=0.7"]; chains are read throughclay.prompt_chains_currentby default (--chains-table/--chains-finalfor the base table withFINAL), and both join sides are bounded to the window before joining.batchbench/__init__.pynow imports without the compiled extension (the API functions raise a clear error if_coreis missing) so the exporter runs from a pure-Python checkout;batchbench export-plansis dispatched in__main__before the Rust CLI.python -m unittest discover -s tests.Tests
tests/test_export_plans.py(19 tests, fixtures only, no live ClickHouse): strict-prefix linking within a principal, latest-parent choice and link window, pre-window sessions andreset_before, branching, live-block marking, delays/offsets/time scale, single-block fallback, dropped-row accounting, per-seed normalization across resets, deterministic and stratified sampling (including range validation), the validator against each Rust rule, query parameterization, CLI filters and error paths, and a dry run of the exported manifest throughbatchbench-agentin both admission modes (skipped when the binary is not built).Codex review findings (default table/FINAL against granted views, unbounded analytics join, empty manifest reported as success, seed definitions lost across resets, sample range) are addressed in the second commit.
I could not run the query against the real warehouse from here; the SQL is parameterized and covered by tests for its shape, but the first real export should be checked against the column names of
clay.http_analyticsin ClickHouse (--analytics-ts-columndefaults totimestamp).https://claude.ai/code/session_015Rtq2g2DyKJkqHgbi8qK6f
Reopened as a new PR against main: the original #34 was auto-closed when its stacked base branch was deleted on merge of #33.