diff --git a/scripts/collectivespec/create_controlled_greedy_window.py b/scripts/collectivespec/create_controlled_greedy_window.py index 148494d..e2d0acf 100644 --- a/scripts/collectivespec/create_controlled_greedy_window.py +++ b/scripts/collectivespec/create_controlled_greedy_window.py @@ -45,6 +45,11 @@ def numeric(value: Any, *, field: str, row_index: int) -> float: return value +def stable_session_key(value: Any) -> str: + """Render a session identifier without conflating JSON types.""" + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--windows-path", type=Path, required=True) @@ -80,12 +85,15 @@ def main() -> int: source_trace = resolve_trace_path(args.windows_path, trace_file) if not source_trace.is_file(): raise SystemExit(f"source trace does not exist: {source_trace}") + source_sampling_strategy = source_window.get("sampling_strategy") + session_coherent_source = source_sampling_strategy == "session_coherent_uniform_score" args.output_dir.mkdir(parents=True, exist_ok=True) output_trace = args.output_dir / f"{args.output_window_id}.jsonl" selected = 0 sampling_values: list[float] = [] timestamps: list[float] = [] + session_counts: dict[str, dict[str, int]] = {} with source_trace.open("r", encoding="utf-8") as src, output_trace.open( "w", encoding="utf-8" ) as dest: @@ -96,6 +104,17 @@ def main() -> int: if not isinstance(row, dict): raise SystemExit(f"row {row_index}: expected an object") sampling_u = numeric(row.get("sampling_u"), field="sampling_u", row_index=row_index) + session_entry: dict[str, int] | None = None + if session_coherent_source: + if "session_root" not in row: + raise SystemExit( + f"row {row_index}: session-coherent source has no session_root" + ) + session_key = stable_session_key(row["session_root"]) + session_entry = session_counts.setdefault( + session_key, {"total_turns": 0, "selected_turns": 0} + ) + session_entry["total_turns"] += 1 if sampling_u > threshold: continue timestamp = numeric(row.get("timestamp"), field="timestamp", row_index=row_index) @@ -104,9 +123,27 @@ def main() -> int: selected += 1 sampling_values.append(sampling_u) timestamps.append(timestamp) + if session_entry is not None: + session_entry["selected_turns"] += 1 if selected == 0: output_trace.unlink(missing_ok=True) raise SystemExit("threshold selected zero requests") + split_sessions = [ + entry + for entry in session_counts.values() + if 0 < entry["selected_turns"] < entry["total_turns"] + ] + if split_sessions: + output_trace.unlink(missing_ok=True) + raise SystemExit( + "session-coherent source produced a partial selected session; " + f"split_session_count={len(split_sessions)}" + ) + selected_session_turn_counts = [ + entry["selected_turns"] + for entry in session_counts.values() + if entry["selected_turns"] > 0 + ] output_window = dict(source_window) output_window.update( @@ -114,12 +151,18 @@ def main() -> int: "window_id": args.output_window_id, "trace_file": str(output_trace.resolve()), "num_requests": selected, - "sampling_strategy": "fixed_uniform_score_then_greedy_temperature", + "sampling_strategy": ( + "session_coherent_uniform_score_then_greedy_temperature" + if session_coherent_source + else "fixed_uniform_score_then_greedy_temperature" + ), "collectivespec_source_window_id": args.window_id, "collectivespec_sampling_u_threshold": threshold, "collectivespec_temperature_override": 0.0, "collectivespec_source_trace_path": str(source_trace), "collectivespec_source_trace_sha256": sha256(source_trace), + "collectivespec_source_sampling_strategy": source_sampling_strategy, + "collectivespec_session_closed_within_window": session_coherent_source, } ) output_windows = { @@ -153,6 +196,24 @@ def main() -> int: ), }, "temperature": {"distinct_value_count": 1, "all_zero": True}, + "session": { + "source_is_session_coherent": session_coherent_source, + "session_root_count": len(session_counts) if session_coherent_source else None, + "selected_session_count": ( + len(selected_session_turn_counts) if session_coherent_source else None + ), + "split_session_count": len(split_sessions) if session_coherent_source else None, + "all_selected_or_dropped": not split_sessions if session_coherent_source else None, + "selected_turns": { + "min": min(selected_session_turn_counts) + if selected_session_turn_counts + else None, + "max": max(selected_session_turn_counts) + if selected_session_turn_counts + else None, + "distinct_value_count": len(set(selected_session_turn_counts)), + }, + }, "source_trace_sha256": sha256(source_trace), "output_trace_sha256": sha256(output_trace), }