topoteretes--cognee/cognee/eval_framework/token_usage_analysis/cli.py

"""Command-line interface: argument definitions and default resolution.

Kept separate so analyze.py reads as pure orchestration.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from corpus import DEFAULT_MAX_CHUNK_SIZE

DESCRIPTION = "Estimate the token cost of cognee memory vs. full-context prompting."


def parse_args() -> argparse.Namespace:
    parser = _build_parser()
    args = parser.parse_args()
    _resolve_llm_models(args, parser)
    _require_corpus_tokens_with_text(args, parser)
    return args


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--file", help="text file to chunk")
    source.add_argument("--dir", help="directory of .txt files to pool and chunk")
    source.add_argument("--text", help="a single representative chunk")

    parser.add_argument("--samples", type=int, default=3, help="chunks to measure")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-chunk-size", type=int, default=DEFAULT_MAX_CHUNK_SIZE)
    parser.add_argument(
        "--llm-models",
        type=_comma_list,
        default=None,
        help="comma list; default = the model cognee is configured with in .env",
    )
    parser.add_argument("--reduction-factors", type=_comma_factors, default=[1, 2, 7, 10])
    parser.add_argument(
        "--corpus-tokens",
        type=int,
        default=None,
        help="corpus size; defaults to the input's token count. Required with --text.",
    )
    parser.add_argument("--retrieved-context", type=int, default=1118)
    parser.add_argument("--query-overhead", type=int, default=32)
    parser.add_argument("--out", type=Path, default=Path("token_usage_report.json"))
    parser.add_argument("--plot", action="store_true")
    parser.add_argument("--plot-dir", type=Path, default=Path("."))
    return parser


def _comma_list(value: str) -> list[str]:
    return [item.strip() for item in value.split(",") if item.strip()]


def _comma_factors(value: str) -> list[float]:
    factors = [float(item) for item in _comma_list(value)]
    return [int(factor) if factor.is_integer() else factor for factor in factors]


def _resolve_llm_models(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
    """Default to the single llm_model cognee is configured with in .env."""
    if args.llm_models:
        return
    from cognee.infrastructure.llm.config import get_llm_config

    configured = get_llm_config().llm_model
    if not configured:
        parser.error("no --llm-models given and no LLM_MODEL configured in .env")
    args.llm_models = [configured]


def _require_corpus_tokens_with_text(
    args: argparse.Namespace, parser: argparse.ArgumentParser
) -> None:
    if args.text is not None and args.corpus_tokens is None:
        parser.error("--corpus-tokens is required with --text (a lone chunk is not a corpus)")