infer.online¶

def run_online_loop(
    *,
    model_dir: Path,
    models_dir: Path | None = None,
    aw_host: str = DEFAULT_AW_HOST,
    poll_seconds: int = DEFAULT_POLL_SECONDS,
    smooth_window: int = DEFAULT_SMOOTH_WINDOW,
    title_salt: str = DEFAULT_TITLE_SALT,
    out_dir: Path = Path(DEFAULT_OUT_DIR),
    bucket_seconds: int = DEFAULT_BUCKET_SECONDS,
    idle_gap_seconds: float = DEFAULT_IDLE_GAP_SECONDS,
    reject_threshold: float | None = DEFAULT_REJECT_THRESHOLD,
    taxonomy_path: Path | None = None,
    calibrator_path: Path | None = None,
    calibrator_store_path: Path | None = None,
    label_queue_path: Path | None = None,
    label_confidence_threshold: float = DEFAULT_LABEL_CONFIDENCE_THRESHOLD,
) -> None:
    """Poll ActivityWatch, predict, smooth, and write results continuously.

    Runs until interrupted with ``KeyboardInterrupt`` (Ctrl+C).  On
    shutdown, writes final segments and prints a summary.

    Session state is tracked across poll cycles.  A new session starts
    when the gap since the last observed event exceeds
    *idle_gap_seconds*.

    When *models_dir* is provided, the loop watches
    ``models/active.json`` for changes and hot-reloads the model bundle
    without restarting.  The swap only occurs after the new bundle loads
    successfully; on failure the current model is kept.

    Args:
        model_dir: Path to a trained model bundle directory.
        models_dir: Base directory for model bundles.  When provided,
            enables automatic model reload on ``active.json`` changes.
        aw_host: Base URL of the ActivityWatch server.
        poll_seconds: Seconds between polling iterations.
        smooth_window: Rolling majority window size.
        title_salt: Salt for hashing window titles.
        out_dir: Directory for ``predictions.csv`` and ``segments.json``.
        bucket_seconds: Width of each time bucket in seconds.
        idle_gap_seconds: Minimum gap (seconds) that starts a new session.
        reject_threshold: If given, predictions with
            ``max(proba) < reject_threshold`` become ``Mixed/Unknown``.
        taxonomy_path: Optional path to a taxonomy YAML file.  When
            provided, output labels are mapped to user-defined buckets.
        calibrator_path: Optional path to a calibrator JSON file.  When
            provided, raw model probabilities are calibrated before the
            reject decision.
        calibrator_store_path: Optional path to a calibrator store
            directory.  When provided, per-user calibration is applied.
            Takes precedence over *calibrator_path*.
        label_queue_path: Optional path to the labeling queue JSON.
            When provided, low-confidence predictions are auto-enqueued
            for manual labeling.
        label_confidence_threshold: Predictions with confidence below
            this value are enqueued when *label_queue_path* is set.
    """
    from taskclf.adapters.activitywatch.client import (
        AWConnectionError,
        AWTimeoutError,
        fetch_aw_events,
        fetch_aw_input_events,
        find_input_bucket_id,
        find_window_bucket_id,
    )
    from taskclf.features.build import build_features_from_aw_events
    from taskclf.infer.calibration import load_calibrator, load_calibrator_store
    from taskclf.infer.feature_state import OnlineFeatureState
    from taskclf.infer.resolve import (
        InferencePolicyReloader,
        resolve_inference_config,
    )
    from taskclf.infer.taxonomy import load_taxonomy

    taxonomy: TaxonomyConfig | None = None
    if taxonomy_path is not None:
        taxonomy = load_taxonomy(taxonomy_path)
        logger.info(
            "Loaded taxonomy from %s (user=%s)", taxonomy_path, taxonomy.user_id
        )

    effective_threshold: float | None
    per_user_thresholds: dict[str, float] | None = None
    if models_dir is not None:
        resolved = resolve_inference_config(
            models_dir,
            model_dir_override=str(model_dir),
            reject_threshold_override=reject_threshold,
            calibrator_store_override=calibrator_store_path,
            calibrator_path_override=calibrator_path,
        )
        model = resolved.model
        metadata = resolved.metadata
        cat_encoders = resolved.cat_encoders
        effective_threshold = resolved.reject_threshold
        calibrator = resolved.calibrator
        cal_store = resolved.calibrator_store
        per_user_thresholds = resolved.per_user_reject_thresholds
    else:
        calibrator_: Calibrator = IdentityCalibrator()
        if calibrator_path is not None:
            calibrator_ = load_calibrator(calibrator_path)
            logger.info("Loaded calibrator from %s", calibrator_path)

        cal_store_: CalibratorStore | None = None
        if calibrator_store_path is not None:
            cal_store_ = load_calibrator_store(calibrator_store_path)
            logger.info(
                "Loaded calibrator store from %s (%d per-user calibrators)",
                calibrator_store_path,
                len(cal_store_.user_calibrators),
            )

        model, metadata, cat_encoders = load_model_bundle(Path(model_dir))
        logger.info("Loaded model from %s (schema=%s)", model_dir, metadata.schema_hash)
        effective_threshold = reject_threshold
        calibrator = calibrator_
        cal_store = cal_store_

    bucket_id = find_window_bucket_id(aw_host)
    logger.info("Using AW window bucket: %s", bucket_id)

    input_bucket_id = find_input_bucket_id(aw_host)
    if input_bucket_id:
        logger.info("Using AW input bucket: %s", input_bucket_id)
    else:
        logger.info(
            "No aw-watcher-input bucket found; keyboard/mouse features will be None"
        )

    predictor = OnlinePredictor(
        model,
        metadata,
        cat_encoders=cat_encoders,
        smooth_window=smooth_window,
        bucket_seconds=bucket_seconds,
        reject_threshold=effective_threshold,
        taxonomy=taxonomy,
        calibrator=calibrator,
        calibrator_store=cal_store,
        per_user_reject_thresholds=per_user_thresholds,
    )

    feature_state = OnlineFeatureState(
        bucket_seconds=bucket_seconds,
        idle_gap_seconds=idle_gap_seconds,
    )

    policy_reloader: InferencePolicyReloader | None = None
    if models_dir is not None:
        policy_reloader = InferencePolicyReloader(models_dir)

    label_queue = None
    if label_queue_path is not None:
        from taskclf.labels.queue import ActiveLabelingQueue

        label_queue = ActiveLabelingQueue(label_queue_path)
        logger.info(
            "Label queue active: %s (threshold=%.2f)",
            label_queue_path,
            label_confidence_threshold,
        )

    pred_path = out_dir / "predictions.csv"
    seg_path = out_dir / "segments.json"

    session_start: datetime | None = None
    last_event_ts: datetime | None = None
    idle_gap = timedelta(seconds=idle_gap_seconds)
    total_enqueued = 0

    print(
        f"Online inference started (polling every {poll_seconds}s, bucket={bucket_id})"
    )
    if input_bucket_id:
        print(f"Input watcher active: {input_bucket_id}")
    if policy_reloader is not None and models_dir is not None:
        print(f"Policy reload enabled (watching {models_dir})")
    print("Press Ctrl+C to stop.\n")

    try:
        while True:
            if policy_reloader is not None:
                new_config = policy_reloader.check_reload()
                if new_config is not None:
                    predictor = OnlinePredictor(
                        new_config.model,
                        new_config.metadata,
                        cat_encoders=new_config.cat_encoders,
                        smooth_window=smooth_window,
                        bucket_seconds=bucket_seconds,
                        reject_threshold=new_config.reject_threshold,
                        taxonomy=taxonomy,
                        calibrator=new_config.calibrator,
                        calibrator_store=new_config.calibrator_store,
                        per_user_reject_thresholds=new_config.per_user_reject_thresholds,
                    )
                    print(
                        f"Config reloaded (schema={new_config.metadata.schema_hash}, "
                        f"threshold={new_config.reject_threshold:.4f})"
                    )

            now = datetime.now(timezone.utc)
            window_start = now - timedelta(seconds=poll_seconds)

            try:
                events = fetch_aw_events(
                    aw_host,
                    bucket_id,
                    window_start,
                    now,
                    title_salt=title_salt,
                )
            except AWConnectionError:
                logger.warning(
                    "ActivityWatch unreachable, will retry in %ds", poll_seconds
                )
                time.sleep(poll_seconds)
                continue
            except AWTimeoutError:
                logger.warning(
                    "ActivityWatch request timed out, will retry in %ds", poll_seconds
                )
                time.sleep(poll_seconds)
                continue
            except Exception:
                logger.warning("Failed to fetch AW events, will retry", exc_info=True)
                time.sleep(poll_seconds)
                continue

            if not events:
                logger.debug("No events in window [%s, %s)", window_start, now)
                time.sleep(poll_seconds)
                continue

            input_events = None
            if input_bucket_id:
                try:
                    input_events = (
                        fetch_aw_input_events(
                            aw_host,
                            input_bucket_id,
                            window_start,
                            now,
                        )
                        or None
                    )
                except Exception:
                    logger.warning("Failed to fetch input events", exc_info=True)

            earliest_new = min(ev.timestamp for ev in events)
            if last_event_ts is not None and earliest_new - last_event_ts >= idle_gap:
                session_start = earliest_new
                logger.info("New session started at %s", session_start)

            if session_start is None:
                session_start = earliest_new

            rows = build_features_from_aw_events(
                events,
                input_events=input_events,
                bucket_seconds=bucket_seconds,
                session_start=session_start,
                schema_version=predictor.metadata.schema_version,
            )
            if not rows:
                time.sleep(poll_seconds)
                continue

            last_event_ts = max(ev.timestamp for ev in events)

            for row in rows:
                feature_state.push(row)
                context = feature_state.get_context()
                row = row.model_copy(update=context)
                prediction = predictor.predict_bucket(row)
                ts_str = row.bucket_start_ts.strftime("%H:%M")
                conf_str = f"{prediction.confidence:.2f}"
                print(
                    f"[{ts_str}] {prediction.mapped_label_name} (confidence: {conf_str})"
                )
                _append_prediction_csv(pred_path, prediction)

                if (
                    label_queue is not None
                    and prediction.confidence < label_confidence_threshold
                ):
                    import pandas as pd

                    enqueue_df = pd.DataFrame(
                        [
                            {
                                "user_id": prediction.user_id or "default-user",
                                "bucket_start_ts": prediction.bucket_start_ts,
                                "bucket_end_ts": prediction.bucket_start_ts
                                + timedelta(seconds=bucket_seconds),
                                "confidence": prediction.confidence,
                                "predicted_label": prediction.core_label_name,
                            }
                        ]
                    )
                    n = label_queue.enqueue_low_confidence(
                        enqueue_df,
                        threshold=label_confidence_threshold,
                    )
                    if n > 0:
                        total_enqueued += n
                        print(f"  → enqueued for labeling (conf={conf_str})")

            segments = predictor.get_segments()
            write_segments_json(segments, seg_path)

            time.sleep(poll_seconds)

    except KeyboardInterrupt:
        print("\nShutting down...")

    segments = predictor.get_segments()
    if segments:
        write_segments_json(segments, seg_path)
        print(f"Final segments written to {seg_path} ({len(segments)} segments)")

        try:
            from taskclf.report.daily import build_daily_report
            from taskclf.report.export import export_report_json

            report = build_daily_report(
                segments,
                bucket_seconds=bucket_seconds,
                raw_labels=predictor._all_raw or None,
                smoothed_labels=predictor._all_smoothed or None,
            )
            report_path = export_report_json(
                report, out_dir / f"report_{report.date}.json"
            )
            print(f"Daily report written to {report_path}")
            if report.flap_rate_raw is not None:
                print(
                    f"Flap rate: raw={report.flap_rate_raw:.4f}  "
                    f"smoothed={report.flap_rate_smoothed:.4f}"
                )
        except Exception:
            logger.warning("Could not generate daily report", exc_info=True)

        try:
            from taskclf.core.telemetry import TelemetryStore, compute_telemetry

            import pandas as pd

            ts_records = [{"bucket_start_ts": ts} for ts in predictor._all_bucket_ts]
            ts_df = pd.DataFrame(ts_records) if ts_records else pd.DataFrame()

            if not ts_df.empty:
                snapshot = compute_telemetry(
                    ts_df,
                    labels=predictor._all_smoothed or None,
                )
                store = TelemetryStore(out_dir / "telemetry")
                store_path = store.append(snapshot)
                print(f"Telemetry snapshot written to {store_path}")
        except Exception:
            logger.warning("Could not write telemetry snapshot", exc_info=True)
    else:
        print("No predictions were made.")

    if total_enqueued > 0:
        print(f"Enqueued {total_enqueued} low-confidence bucket(s) for labeling.")

def resolve_model_dir(
    model_dir: str | None,
    models_dir: Path,
    policy: SelectionPolicy | None = None,
) -> Path:
    """Resolve the model directory for inference.

    .. deprecated::
        Use :func:`resolve_inference_config` instead.  This function
        only resolves the model bundle path; it does not load the
        calibrator store or reject threshold from the inference policy.

    Resolution precedence:

    1. If *model_dir* is provided, validate that it exists and return it.
    2. Otherwise, delegate to :func:`~taskclf.model_registry.resolve_active_model`
       which reads ``active.json`` or falls back to best-model selection.
    3. If no eligible model is found, raise :class:`ModelResolutionError`
       with a descriptive message including exclusion reasons.

    Args:
        model_dir: Explicit ``--model-dir`` value from CLI, or ``None``.
        models_dir: Base directory containing promoted model bundles.
        policy: Selection policy override (defaults to policy v1).

    Returns:
        Path to the resolved model bundle directory.

    Raises:
        ModelResolutionError: When no model can be resolved.
    """
    if model_dir is not None:
        path = Path(model_dir)
        if not path.is_dir():
            raise ModelResolutionError(
                f"Explicit --model-dir does not exist: {model_dir}"
            )
        return path

    if not models_dir.is_dir():
        raise ModelResolutionError(
            f"Models directory does not exist: {models_dir}. "
            "Provide --model-dir explicitly or train a model first."
        )

    bundle, report = resolve_active_model(models_dir, policy)

    if bundle is not None:
        logger.info("Resolved model: %s", bundle.path)
        return bundle.path

    lines = [
        f"No eligible model found in {models_dir}.",
        "Provide --model-dir explicitly or train a compatible model.",
    ]
    if report is not None and report.excluded:
        lines.append("Excluded bundles:")
        for rec in report.excluded:
            lines.append(f"  - {rec.model_id}: {rec.reason}")
    raise ModelResolutionError("\n".join(lines), report=report)