flux_cli_beta.py

#!/usr/bin/env python3

"""
Flux CLI β版

This script provides a command‑line interface to perform text compression using an
external `compression_cli` module and to record detailed logs of each operation.

Features:
 - Token limit enforcement with a configurable `--limit-tokens` option. If the
   number of tokens in an input file exceeds this limit, the file is skipped and
   noted in `skipped.log`.
 - Invocation of an external compression function from `compression_cli`. If
   `compression_cli` is unavailable, a simple fallback implementation that
   returns the original content and a mocked compression ratio is used.
 - Comprehensive logging of every session to a JSON file (`session_*.json`), a
   separate log of skipped files (`skipped.log`), and a summary CSV
   (`summary.csv`) that accumulates results across runs.
 - Errors are caught and recorded with status "error" in the summary.

Usage:
    python flux_cli_beta.py [OPTIONS] FILE [FILE ...]

Example:
    python flux_cli_beta.py --limit-tokens 15000 data1.txt data2.md

The above command processes `data1.txt` and `data2.md`, compressing them if
their token counts do not exceed 15,000 tokens. The results will be logged
accordingly.
"""

import argparse
import json
import csv
import os
import sys
import datetime
from typing import List, Tuple, Dict, Any


def import_compression_module():
    """Attempt to import the external compression_cli module.

    Returns
    -------
    callable
        A function that takes text input and returns a tuple of (output_text,
        output_token_count). A fallback is provided if the module cannot be
        imported.
    """
    try:
        # Attempt to import a function named `compress_file` from compression_cli.
        from compression_cli import compress_file as external_compress

        def wrapper(content: str) -> Tuple[str, int]:
            # We assume external_compress returns the compressed content and logs
            # metrics internally. For compatibility we return the content and
            # compute output tokens as the length of the content split on
            # whitespace.
            compressed = external_compress(content)
            if isinstance(compressed, tuple) and len(compressed) == 2:
                result, meta = compressed
                # If meta contains token information, use it
                if isinstance(meta, dict) and 'output_tokens' in meta:
                    return result, int(meta['output_tokens'])
                return result, len(result.split())
            return compressed, len(compressed.split())

        return wrapper
    except Exception:
        # Provide a fallback compression function that simulates compression by
        # returning the original text and a mocked reduction of approximately
        # 50% of the token count. This fallback also prints a warning once.
        warned = False

        def fallback(content: str) -> Tuple[str, int]:
            nonlocal warned
            if not warned:
                sys.stderr.write(
                    "Warning: compression_cli module not found. "
                    "Using fallback compression (no actual compression).\n"
                )
                warned = True
            tokens = content.split()
            # simulate a 50% reduction rate
            output_tokens = max(1, len(tokens) // 2)
            return content, output_tokens

        return fallback


def count_tokens(content: str) -> int:
    """Count tokens in the content using a simple whitespace split.

    Parameters
    ----------
    content : str
        Text content to count tokens for.

    Returns
    -------
    int
        The number of tokens found in the content.
    """
    return len(content.split())


def ensure_logs_directory(base_path: str) -> str:
    """Ensure that a `logs` directory exists.

    Parameters
    ----------
    base_path : str
        The base directory in which to create the logs directory.

    Returns
    -------
    str
        The path to the logs directory.
    """
    logs_dir = os.path.join(base_path, "logs")
    os.makedirs(logs_dir, exist_ok=True)
    return logs_dir


def write_session_log(logs_dir: str, session_data: List[Dict[str, Any]]) -> str:
    """Write the session log to a JSON file.

    Each session log is named with the current timestamp down to seconds to
    avoid collisions.

    Parameters
    ----------
    logs_dir : str
        Directory where logs are stored.
    session_data : list of dict
        Data to write into the session log.

    Returns
    -------
    str
        The path to the session log file.
    """
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    session_file = os.path.join(logs_dir, f"session_{timestamp}.json")
    with open(session_file, "w", encoding="utf-8") as f:
        json.dump(session_data, f, ensure_ascii=False, indent=2)
    return session_file


def write_skipped_log(logs_dir: str, skipped_entries: List[Dict[str, Any]]) -> str:
    """Append skipped file entries to a log.

    Parameters
    ----------
    logs_dir : str
        Directory where logs are stored.
    skipped_entries : list of dict
        Each dict should include 'file', 'tokens' and 'reason'.

    Returns
    -------
    str
        The path to the skipped log file.
    """
    skipped_file = os.path.join(logs_dir, "skipped.log")
    if not skipped_entries:
        return skipped_file
    with open(skipped_file, "a", encoding="utf-8") as f:
        for entry in skipped_entries:
            f.write(f"{json.dumps(entry, ensure_ascii=False)}\n")
    return skipped_file


def update_summary_csv(logs_dir: str, summary_rows: List[Dict[str, Any]]) -> str:
    """Update the summary CSV file with new rows.

    Parameters
    ----------
    logs_dir : str
        Directory where logs are stored.
    summary_rows : list of dict
        Rows to append to summary CSV. Keys should match column names.

    Returns
    -------
    str
        The path to the summary CSV file.
    """
    summary_file = os.path.join(logs_dir, "summary.csv")
    file_exists = os.path.exists(summary_file)
    fieldnames = [
        "file",
        "input_tokens",
        "output_tokens",
        "reduction_rate",
        "status",
        "skip_reason",
        "error_message",
        "elapsed_time_sec",
    ]
    with open(summary_file, "a", newline="", encoding="utf-8") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        if not file_exists:
            writer.writeheader()
        for row in summary_rows:
            # Fill missing columns with empty strings
            for field in fieldnames:
                row.setdefault(field, "")
            writer.writerow(row)
    return summary_file


def process_file(
    filepath: str,
    compress_func,
    token_limit: int,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Process a single file for compression.

    Parameters
    ----------
    filepath : str
        Path to the input file.
    compress_func : callable
        Function that compresses text and returns (compressed_text, output_tokens).
    token_limit : int
        Maximum number of tokens allowed before skipping the file.

    Returns
    -------
    (dict, dict)
        A tuple containing the session log entry and the summary row. In case the
        file is skipped or an error occurs, certain fields may be empty or
        omitted.
    """
    import time
    start_time = time.time()
    session_entry: Dict[str, Any] = {
        "file": filepath,
    }
    summary_row: Dict[str, Any] = {
        "file": filepath,
    }
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read()
        input_tokens = count_tokens(content)
        session_entry["input_tokens"] = input_tokens
        summary_row["input_tokens"] = input_tokens

        if input_tokens > token_limit:
            # Skip due to token limit
            session_entry.update(
                {
                    "status": "skipped",
                    "skip_reason": "limit_exceeded",
                    "output_tokens": None,
                    "reduction_rate": None,
                    "elapsed_time_sec": time.time() - start_time,
                }
            )
            summary_row.update(
                {
                    "status": "skipped",
                    "skip_reason": "limit_exceeded",
                }
            )
            return session_entry, summary_row

        # Process compression
        compressed_text, output_tokens = compress_func(content)
        reduction_rate = (
            (input_tokens - output_tokens) / input_tokens if input_tokens else 0
        )

        session_entry.update(
            {
                "status": "success",
                "output_tokens": output_tokens,
                "reduction_rate": reduction_rate,
                "elapsed_time_sec": time.time() - start_time,
            }
        )
        summary_row.update(
            {
                "output_tokens": output_tokens,
                "reduction_rate": reduction_rate,
                "status": "success",
            }
        )
        return session_entry, summary_row
    except Exception as e:
        # Handle any exceptions during processing
        session_entry.update(
            {
                "status": "error",
                "error_message": str(e),
                "elapsed_time_sec": time.time() - start_time,
            }
        )
        summary_row.update(
            {
                "status": "error",
                "error_message": str(e),
            }
        )
        return session_entry, summary_row


def main(argv: List[str] = None) -> int:
    """Entry point for the CLI script."""
    parser = argparse.ArgumentParser(
        description="Flux CLI β版 - Compress text files and log results",
    )
    parser.add_argument(
        "files",
        nargs="+",
        help="Paths to the files to compress",
    )
    parser.add_argument(
        "--limit-tokens",
        dest="limit_tokens",
        type=int,
        default=10000,
        help="Maximum number of tokens allowed before a file is skipped (default: 10000)",
    )
    parser.add_argument(
        "--logs-dir",
        dest="logs_dir",
        default=None,
        help="Directory to store logs (default: ./logs)",
    )

    args = parser.parse_args(argv)

    # Clamp limit_tokens to a maximum of 50000
    if args.limit_tokens > 50000:
        print(
            f"Warning: limit_tokens capped at 50000 (requested {args.limit_tokens}).",
            file=sys.stderr,
        )
        token_limit = 50000
    elif args.limit_tokens < 1:
        print(
            f"Warning: limit_tokens must be positive; using default 10000 instead.",
            file=sys.stderr,
        )
        token_limit = 10000
    else:
        token_limit = args.limit_tokens

    base_path = os.getcwd()
    logs_dir = (
        args.logs_dir
        if args.logs_dir is not None
        else ensure_logs_directory(base_path)
    )
    os.makedirs(logs_dir, exist_ok=True)

    compress_func = import_compression_module()

    session_entries: List[Dict[str, Any]] = []
    summary_rows: List[Dict[str, Any]] = []
    skipped_entries: List[Dict[str, Any]] = []

    for file_path in args.files:
        if not os.path.isfile(file_path):
            # Record missing file as error
            session_entry = {
                "file": file_path,
                "status": "error",
                "error_message": "file not found",
                "elapsed_time_sec": 0,
            }
            summary_row = {
                "file": file_path,
                "status": "error",
                "error_message": "file not found",
            }
            session_entries.append(session_entry)
            summary_rows.append(summary_row)
            continue

        session_entry, summary_row = process_file(
            file_path,
            compress_func,
            token_limit,
        )
        session_entries.append(session_entry)
        summary_rows.append(summary_row)
        if session_entry.get("status") == "skipped":
            skipped_entries.append(
                {
                    "file": file_path,
                    "tokens": session_entry.get("input_tokens"),
                    "reason": session_entry.get("skip_reason"),
                }
            )

    # Write logs
    session_log_file = write_session_log(logs_dir, session_entries)
    skipped_log_file = write_skipped_log(logs_dir, skipped_entries)
    summary_file = update_summary_csv(logs_dir, summary_rows)

    # Print a summary for the user
    print(f"Session log written to: {session_log_file}")
    if skipped_entries:
        print(f"Skipped entries logged in: {skipped_log_file}")
    print(f"Summary CSV updated at: {summary_file}")
    return 0


if __name__ == "__main__":
    sys.exit(main())