chore: import upstream snapshot with attribution

cognee/eval_framework/token_usage_analysis/cli.py

@@ -0,0 +1,81 @@
+"""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)")