feat(adapter): implement process_session and CLI

ai_mouse/data_adapters/__main__.py

@@ -0,0 +1,13 @@
+"""CLI dispatch: `python -m ai_mouse.data_adapters.balabit ...`
+
+Note: This file makes `python -m ai_mouse.data_adapters` invokable but for
+clarity prefer the explicit form `python -m ai_mouse.data_adapters.balabit`.
+"""
+from __future__ import annotations
+
+import sys
+
+from ai_mouse.data_adapters.balabit import main
+
+if __name__ == "__main__":
+    sys.exit(main())

ai_mouse/data_adapters/balabit.py

@@ -197,5 +197,130 @@ def process_session(
     output_jsonl: Path,
     config,
 ) -> int:
-    """Stub — implemented in Task 6."""
+    """Convert one Balabit session CSV to JSONL traces, append to output.
-    raise NotImplementedError
+
+    Args:
+        csv_path: Path to a Balabit session CSV.
+        output_jsonl: Output JSONL file (will be appended to).
+        config: BalabitAdapterConfig with window_ms / min_dist / etc.
+
+    Returns:
+        Number of valid segments written.
+    """
+    import math
+
+    session_id = csv_path.stem  # e.g. "session_42"
+    events = parse_session_csv(csv_path)
+    if not events:
+        return 0
+
+    raw_segments = segment_by_clicks(
+        events, window_ms=config.window_ms, session_id=session_id
+    )
+    valid_segments = filter_segments(
+        raw_segments,
+        min_events=config.min_events,
+        min_dist=config.min_dist,
+        max_span_ms=config.max_span_ms,
+        max_gap_ms=config.max_gap_ms,
+    )
+
+    if not valid_segments:
+        return 0
+
+    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+    n_written = 0
+    with output_jsonl.open("a", encoding="utf-8") as f:
+        for seg in valid_segments:
+            sx, sy = seg.events[0].x, seg.events[0].y
+            ex, ey = seg.click_x, seg.click_y
+            dist = math.hypot(ex - sx, ey - sy)
+            angle = math.degrees(math.atan2(ey - sy, ex - sx))
+            t0 = seg.events[0].t_ms
+            record = {
+                "meta": {
+                    "start": [sx, sy],
+                    "end": [ex, ey],
+                    "dist": int(round(dist)),
+                    "angle": round(angle, 1),
+                    "source": "balabit",
+                    "session_id": seg.session_id,
+                },
+                "events": [
+                    {"type": "move", "x": e.x, "y": e.y, "t": e.t_ms - t0}
+                    for e in seg.events
+                ],
+            }
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+            n_written += 1
+    return n_written
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
+    import argparse
+
+    from ai_mouse.config import BalabitAdapterConfig
+
+    parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
+    parser.add_argument(
+        "--input", type=Path, required=True,
+        help="Directory containing Balabit session CSV files (recursive)",
+    )
+    parser.add_argument(
+        "--output", type=Path, default=Path("data/pretrain_traces.jsonl"),
+        help="Output JSONL path (default: data/pretrain_traces.jsonl)",
+    )
+    parser.add_argument("--window-ms", type=int, default=1200)
+    parser.add_argument("--min-dist", type=int, default=50)
+    parser.add_argument("--min-events", type=int, default=5)
+    parser.add_argument("--max-span-ms", type=int, default=5000)
+    parser.add_argument("--max-gap-ms", type=int, default=200)
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Truncate output file before writing (default: append)",
+    )
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    if not args.input.is_dir():
+        logger.error("Input is not a directory: %s", args.input)
+        return 2
+
+    config = BalabitAdapterConfig(
+        window_ms=args.window_ms,
+        min_dist=args.min_dist,
+        min_events=args.min_events,
+        max_span_ms=args.max_span_ms,
+        max_gap_ms=args.max_gap_ms,
+    )
+
+    if args.overwrite and args.output.exists():
+        args.output.unlink()
+
+    # Recursively walk the input directory looking for session files.
+    # Balabit files have no extension; we accept any regular file.
+    csv_files = sorted([p for p in args.input.rglob("*") if p.is_file() and not p.name.startswith(".")])
+    if not csv_files:
+        logger.error("No session files found under %s", args.input)
+        return 2
+
+    total = 0
+    for i, csv_path in enumerate(csv_files, 1):
+        try:
+            n = process_session(csv_path, args.output, config)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("Skipping %s due to error: %s", csv_path.name, exc)
+            continue
+        total += n
+        if i % 10 == 0 or i == len(csv_files):
+            logger.info("Processed %d/%d sessions, %d segments so far", i, len(csv_files), total)
+
+    logger.info("Done. Wrote %d segments to %s", total, args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())

tests/test_balabit_adapter.py

@@ -1,6 +1,7 @@
 """Tests for Balabit Mouse Dynamics Challenge data adapter."""
 from __future__ import annotations
 
+import json
 from pathlib import Path
 
 import pytest
@@ -261,3 +262,75 @@ class TestFilterSegments:
         result = filter_segments([seg], min_events=5, min_dist=50,
                                   max_span_ms=5000, max_gap_ms=200)
         assert len(result) == 1
+
+
+class TestProcessSession:
+    def test_writes_jsonl_in_expected_format(self, tmp_path):
+        from ai_mouse.config import BalabitAdapterConfig
+        from ai_mouse.data_adapters.balabit import process_session
+
+        # Construct a Balabit-format CSV with one valid segment
+        csv_path = tmp_path / "user_session_42"
+        rows = []
+        # 10 Move events going from (100,100) to (500,200) over 500ms
+        for i in range(10):
+            t = i * 0.05
+            x = 100 + i * 40
+            y = 100 + i * 10
+            rows.append(f"0,{t:.3f},NoButton,Move,{x},{y}")
+        rows.append(f"0,0.550,Left,Pressed,510,210")
+        _write_csv(csv_path, rows)
+
+        out = tmp_path / "out.jsonl"
+        config = BalabitAdapterConfig(
+            window_ms=1200, min_dist=50, min_events=5,
+            max_span_ms=5000, max_gap_ms=200,
+        )
+        n = process_session(csv_path, out, config)
+        assert n == 1
+        assert out.exists()
+
+        line = out.read_text(encoding="utf-8").strip()
+        record = json.loads(line)
+        assert record["meta"]["start"] == [100, 100]
+        assert record["meta"]["end"] == [510, 210]
+        assert record["meta"]["dist"] > 0
+        assert record["meta"]["source"] == "balabit"
+        assert record["meta"]["session_id"] == "user_session_42"
+        # Events: only move events, no down/up
+        assert all(e["type"] == "move" for e in record["events"])
+        # Timestamps relative (start at 0)
+        assert record["events"][0]["t"] == 0
+
+    def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
+        from ai_mouse.config import BalabitAdapterConfig
+        from ai_mouse.data_adapters.balabit import process_session
+
+        csv_path = tmp_path / "empty_session"
+        _write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"])  # no clicks
+
+        out = tmp_path / "out.jsonl"
+        config = BalabitAdapterConfig()
+        n = process_session(csv_path, out, config)
+        assert n == 0
+        # File may not exist or may be empty — both acceptable
+        if out.exists():
+            assert out.read_text() == ""
+
+    def test_appends_to_existing_jsonl(self, tmp_path):
+        from ai_mouse.config import BalabitAdapterConfig
+        from ai_mouse.data_adapters.balabit import process_session
+
+        out = tmp_path / "out.jsonl"
+        out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")
+
+        csv_path = tmp_path / "session_x"
+        rows = [f"0,{i*0.05:.3f},NoButton,Move,{100+i*40},100" for i in range(10)]
+        rows.append("0,0.550,Left,Pressed,510,100")
+        _write_csv(csv_path, rows)
+
+        config = BalabitAdapterConfig()
+        process_session(csv_path, out, config)
+
+        lines = out.read_text(encoding="utf-8").strip().split("\n")
+        assert len(lines) == 2  # original + appended