aituner/scripts/collectivespec/create_window_session_closed_windows.py

#!/usr/bin/env python3
"""Build a session-closed replay window from an already materialized trace.

This is deliberately a *window-scoped fallback* for when the raw trace span
needed by ``prepare_trace_windows.py`` is unavailable.  It closes every known
parent/child edge within the supplied window and groups siblings whose common
parent lies outside the window.  It cannot recover turns outside the window;
the manifest records that residual rather than calling the result fully
session-coherent.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from collections import Counter
from pathlib import Path
from typing import Any


SCORE_ALGORITHM = "collectivespec-window-cc-v1"


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def typed_id(value: Any) -> str:
    """Encode IDs without conflating values such as ``1`` and ``"1"``."""
    return json.dumps(
        {"type": type(value).__name__, "value": value},
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )


def session_uniform(*, seed: int, window_id: str, component_key: str) -> float:
    payload = json.dumps(
        {
            "algorithm": SCORE_ALGORITHM,
            "seed": seed,
            "window_id": window_id,
            "component_key": component_key,
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    digest = hashlib.blake2b(payload, digest_size=8).digest()
    return int.from_bytes(digest, byteorder="big", signed=False) / float(1 << 64)


def resolve_trace_path(windows_path: Path, trace_file: str) -> Path:
    candidate = Path(trace_file)
    if candidate.is_absolute():
        return candidate
    return (windows_path.parent / candidate).resolve()


def is_root_parent(value: Any) -> bool:
    return value is None or (
        isinstance(value, (int, float)) and not isinstance(value, bool) and int(value) < 0
    )


class UnionFind:
    def __init__(self) -> None:
        self.parent: dict[str, str] = {}

    def add(self, item: str) -> None:
        self.parent.setdefault(item, item)

    def find(self, item: str) -> str:
        parent = self.parent[item]
        if parent != item:
            self.parent[item] = self.find(parent)
        return self.parent[item]

    def union(self, left: str, right: str) -> None:
        left_root, right_root = self.find(left), self.find(right)
        if left_root != right_root:
            # A canonical representative makes the component key independent
            # of trace order and union order.
            self.parent[max(left_root, right_root)] = min(left_root, right_root)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--windows-path", type=Path, required=True)
    parser.add_argument("--window-id", required=True)
    parser.add_argument("--sample-seed", type=int, default=20260325)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--output-window-id", required=True)
    args = parser.parse_args()

    source_payload = json.loads(args.windows_path.read_text(encoding="utf-8"))
    source_windows = source_payload.get("windows") if isinstance(source_payload, dict) else source_payload
    if not isinstance(source_windows, list):
        raise SystemExit("windows payload must contain a list")
    source_window = next(
        (item for item in source_windows if item.get("window_id") == args.window_id),
        None,
    )
    if not isinstance(source_window, dict):
        raise SystemExit(f"window id not found: {args.window_id}")
    trace_file = source_window.get("trace_file")
    if not isinstance(trace_file, str) or not trace_file:
        raise SystemExit(f"window {args.window_id} has no trace_file")
    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}")

    rows: list[dict[str, Any]] = []
    id_to_index: dict[str, int] = {}
    with source_trace.open("r", encoding="utf-8") as source_handle:
        for row_index, raw in enumerate(source_handle):
            if not raw.strip():
                continue
            row = json.loads(raw)
            if not isinstance(row, dict):
                raise SystemExit(f"row {row_index}: expected an object")
            if "chat_id" not in row or "parent_chat_id" not in row:
                raise SystemExit(f"row {row_index}: chat_id and parent_chat_id are required")
            if "timestamp" not in row:
                raise SystemExit(f"row {row_index}: timestamp is required")
            try:
                timestamp = float(row["timestamp"])
            except (TypeError, ValueError) as error:
                raise SystemExit(f"row {row_index}: timestamp must be numeric") from error
            if not math.isfinite(timestamp):
                raise SystemExit(f"row {row_index}: timestamp must be finite")
            key = typed_id(row["chat_id"])
            if key in id_to_index:
                raise SystemExit(f"row {row_index}: duplicate chat_id {row['chat_id']!r}")
            id_to_index[key] = len(rows)
            rows.append(row)
    if not rows:
        raise SystemExit("source trace has no rows")

    union_find = UnionFind()
    for chat_key in id_to_index:
        union_find.add(chat_key)
    known_parent_edges = 0
    absent_parent_edges = 0
    root_request_count = 0
    for row in rows:
        chat_key = typed_id(row["chat_id"])
        parent = row["parent_chat_id"]
        if is_root_parent(parent):
            root_request_count += 1
            continue
        parent_key = typed_id(parent)
        union_find.add(parent_key)
        if parent_key in id_to_index:
            union_find.union(chat_key, parent_key)
            known_parent_edges += 1
        else:
            # Keep the absent parent as a virtual node. This groups siblings
            # with a common boundary parent without pretending that the parent
            # turn itself is in this window.
            union_find.union(chat_key, parent_key)
            absent_parent_edges += 1

    component_members: dict[str, list[str]] = {}
    for node in union_find.parent:
        component_members.setdefault(union_find.find(node), []).append(node)
    component_key_by_chat: dict[str, str] = {}
    for members in component_members.values():
        # Virtual boundary parents intentionally participate: they ensure all
        # children of an unavailable parent share a deterministic component.
        component_key = min(members)
        for chat_key in (member for member in members if member in id_to_index):
            component_key_by_chat[chat_key] = component_key
    for row in rows:
        parent = row["parent_chat_id"]
        if is_root_parent(parent):
            continue
        parent_key = typed_id(parent)
        if parent_key in id_to_index and union_find.find(typed_id(row["chat_id"])) != union_find.find(parent_key):
            raise SystemExit("same-window parent edge did not close into one component")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    trace_dir = args.output_dir / "traces"
    trace_dir.mkdir(exist_ok=True)
    output_trace = trace_dir / f"{args.output_window_id}.jsonl"
    score_values: list[float] = []
    timestamps: list[float] = []
    component_counts: Counter[str] = Counter()
    with output_trace.open("w", encoding="utf-8") as handle:
        for row in rows:
            chat_key = typed_id(row["chat_id"])
            component_key = component_key_by_chat[chat_key]
            component_id = f"window-component:{component_key}"
            out = dict(row)
            out["collectivespec_original_sampling_u"] = out.get("sampling_u")
            out["session_root"] = component_id
            out["sampling_u"] = session_uniform(
                seed=args.sample_seed,
                window_id=args.window_id,
                component_key=component_key,
            )
            handle.write(json.dumps(out, ensure_ascii=False, separators=(",", ":")) + "\n")
            score_values.append(float(out["sampling_u"]))
            timestamps.append(float(out["timestamp"]))
            component_counts[component_id] += 1

    component_sizes = list(component_counts.values())
    output_window = dict(source_window)
    output_window.update(
        {
            "window_id": args.output_window_id,
            "trace_file": str(output_trace.resolve()),
            "num_requests": len(rows),
            "sampling_strategy": "window_session_closed_uniform_score",
            "sampling_seed": args.sample_seed,
            "collectivespec_sampling_algorithm": SCORE_ALGORITHM,
            "collectivespec_source_window_id": args.window_id,
            "collectivespec_source_trace_path": str(source_trace),
            "collectivespec_source_trace_sha256": sha256(source_trace),
            "collectivespec_session_closure_scope": "window_only",
            "collectivespec_known_parent_edges": known_parent_edges,
            "collectivespec_absent_parent_edges": absent_parent_edges,
        }
    )
    output_windows = {
        "kind": "collectivespec_window_session_closed_fallback",
        "source_windows_path": str(args.windows_path.resolve()),
        "source_window_id": args.window_id,
        "sample_seed": args.sample_seed,
        "windows": [output_window],
    }
    output_manifest = args.output_dir / "windows.json"
    output_manifest.write_text(
        json.dumps(output_windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    sanity = {
        "n": len(rows),
        "timestamp": {
            "min": min(timestamps),
            "max": max(timestamps),
            "non_negative": all(value >= 0 for value in timestamps),
            "nondecreasing": all(
                timestamps[index] >= timestamps[index - 1]
                for index in range(1, len(timestamps))
            ),
        },
        "sampling_u": {
            "min": min(score_values),
            "max": max(score_values),
            "distinct_value_count": len(set(score_values)),
            "within_0_1": all(0 <= value < 1 for value in score_values),
        },
        "window_session_closure": {
            "component_count": len(component_sizes),
            "component_size_min": min(component_sizes),
            "component_size_max": max(component_sizes),
            "component_size_distinct_value_count": len(set(component_sizes)),
            "known_parent_edges": known_parent_edges,
            "absent_parent_edges": absent_parent_edges,
            "root_request_count": root_request_count,
            "all_known_parent_edges_closed": True,
            "scope": "window_only",
        },
        "source_trace_sha256": sha256(source_trace),
        "output_trace_sha256": sha256(output_trace),
    }
    print(
        json.dumps(
            {
                "windows_path": str(output_manifest),
                "trace_path": str(output_trace),
                "data_sanity": sanity,
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())