social-data-attribution-w2/src/data_attribution/cli/handlers.py

"""Command handlers for the data attribution CLI."""

from __future__ import annotations

import argparse
import logging
from pathlib import Path
from typing import Any, Dict

from data_attribution.cli.config import default_run_id


def _log_only_handler(args: argparse.Namespace, config: Dict[str, Any]) -> int:
    logger = logging.getLogger(args.command)
    run_id = args.run_id or default_run_id()
    logger.info("Starting %s (run_id=%s)", args.command, run_id)
    if config:
        logger.debug("Loaded config: %s", config)
    return 0


def _attribution_handler(args: argparse.Namespace, config: Dict[str, Any]) -> int:
    from data_attribution.attribution import AttributionRunConfig, run_and_store

    logger = logging.getLogger(args.command)
    run_id = args.run_id or default_run_id(prefix="attr")
    logger.info("Starting %s (run_id=%s)", args.command, run_id)
    config_payload = config.get("attribution", {}) if config else {}

    output_dir = args.output_dir or Path("runs/attributions") / run_id

    run_config = AttributionRunConfig(
        run_id=run_id,
        query_gradients=args.query_gradients,
        index_path=args.index_path,
        output_dir=output_dir,
        metadata_path=args.metadata_path,
        output_format=args.output_format,
        top_k=args.top_k,
        query_id_field=args.query_id_field,
        doc_id_field=args.doc_id_field,
        metadata_join_keys=tuple(args.metadata_join_keys or ()),
        device=args.device,
        unit_norm=args.unit_norm,
        config=config_payload,
    )

    output_path = run_and_store(run_config)
    logger.info("Wrote attribution results to %s", output_path)
    return 0


__all__ = ["_attribution_handler", "_log_only_handler"]