diff --git a/ai_mouse/data_adapters/__main__.py b/ai_mouse/data_adapters/__main__.py new file mode 100644 index 0000000..bb368bd --- /dev/null +++ b/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()) diff --git a/ai_mouse/data_adapters/balabit.py b/ai_mouse/data_adapters/balabit.py index da10070..1771d6c 100644 --- a/ai_mouse/data_adapters/balabit.py +++ b/ai_mouse/data_adapters/balabit.py @@ -197,5 +197,130 @@ def process_session( output_jsonl: Path, config, ) -> int: - """Stub — implemented in Task 6.""" - raise NotImplementedError + """Convert one Balabit session CSV to JSONL traces, append to output. + + 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()) diff --git a/tests/test_balabit_adapter.py b/tests/test_balabit_adapter.py index d6bf1f0..995dacb 100644 --- a/tests/test_balabit_adapter.py +++ b/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