diff --git a/codeframe/cli/stats_commands.py b/codeframe/cli/stats_commands.py index 9263a031..b28c4bfc 100644 --- a/codeframe/cli/stats_commands.py +++ b/codeframe/cli/stats_commands.py @@ -112,32 +112,14 @@ def tokens( console.print(table) else: - # Workspace-wide summary - records = db.get_workspace_token_usage() - - total_input = 0 - total_output = 0 - total_cost = 0.0 - model_stats: dict[str, dict] = {} - - for record in records: - total_input += record["input_tokens"] - total_output += record["output_tokens"] - total_cost += record["estimated_cost_usd"] - - model = record["model_name"] - if model not in model_stats: - model_stats[model] = { - "input_tokens": 0, - "output_tokens": 0, - "cost_usd": 0.0, - "calls": 0, - } - model_stats[model]["input_tokens"] += record["input_tokens"] - model_stats[model]["output_tokens"] += record["output_tokens"] - model_stats[model]["cost_usd"] += record["estimated_cost_usd"] - model_stats[model]["calls"] += 1 - + # Workspace-wide summary — aggregated in SQL, then totals + # derived from the handful of per-model rows. + by_model = db.get_costs_by_model() + + total_input = sum(m["input_tokens"] for m in by_model) + total_output = sum(m["output_tokens"] for m in by_model) + total_cost = sum(m["total_cost_usd"] for m in by_model) + total_calls = sum(m["call_count"] for m in by_model) total_tokens = total_input + total_output console.print("\n[bold]Workspace Token Usage Summary[/bold]\n") @@ -150,11 +132,11 @@ def tokens( summary_table.add_row("Input Tokens", _format_number(total_input)) summary_table.add_row("Output Tokens", _format_number(total_output)) summary_table.add_row("Total Cost", f"${total_cost:.4f}") - summary_table.add_row("LLM Calls", str(len(records))) + summary_table.add_row("LLM Calls", str(total_calls)) console.print(summary_table) - if model_stats: + if by_model: console.print("\n[bold]By Model:[/bold]") model_table = Table(show_header=True) model_table.add_column("Model", style="cyan") @@ -162,12 +144,12 @@ def tokens( model_table.add_column("Cost", justify="right") model_table.add_column("Calls", justify="right") - for model_name, stats in model_stats.items(): + for m in by_model: model_table.add_row( - model_name, - _format_number(stats["input_tokens"] + stats["output_tokens"]), - f"${stats['cost_usd']:.4f}", - str(stats["calls"]), + m["model_name"], + _format_number(m["input_tokens"] + m["output_tokens"]), + f"${m['total_cost_usd']:.4f}", + str(m["call_count"]), ) console.print(model_table) @@ -214,24 +196,13 @@ def costs( ) raise typer.Exit(1) - # Single fetch: get raw records and compute summary + per-model breakdown in one pass - records = db.get_workspace_token_usage(start_date=start_date, end_date=end_date) - - total_cost = 0.0 - total_tokens = 0 - model_costs: dict[str, dict] = {} - for record in records: - cost = record["estimated_cost_usd"] - tokens = record["input_tokens"] + record["output_tokens"] - total_cost += cost - total_tokens += tokens - - model = record["model_name"] - if model not in model_costs: - model_costs[model] = {"cost_usd": 0.0, "tokens": 0, "calls": 0} - model_costs[model]["cost_usd"] += cost - model_costs[model]["tokens"] += tokens - model_costs[model]["calls"] += 1 + # Per-model rollup aggregated in SQL; totals summed over the + # small per-model result set. + by_model = db.get_costs_by_model(start_date=start_date, end_date=end_date) + + total_cost = sum(m["total_cost_usd"] for m in by_model) + total_tokens = sum(m["input_tokens"] + m["output_tokens"] for m in by_model) + total_calls = sum(m["call_count"] for m in by_model) period_label = f" ({period})" if period else " (all time)" console.print(f"\n[bold]Cost Report{period_label}[/bold]\n") @@ -242,11 +213,11 @@ def costs( table.add_row("Total Cost", f"${total_cost:.4f}") table.add_row("Total Tokens", _format_number(total_tokens)) - table.add_row("LLM Calls", str(len(records))) + table.add_row("LLM Calls", str(total_calls)) console.print(table) - if model_costs: + if by_model: console.print("\n[bold]By Model:[/bold]") model_table = Table(show_header=True) model_table.add_column("Model", style="cyan") @@ -254,12 +225,12 @@ def costs( model_table.add_column("Tokens", justify="right") model_table.add_column("Calls", justify="right") - for model_name, stats in model_costs.items(): + for m in by_model: model_table.add_row( - model_name, - f"${stats['cost_usd']:.4f}", - _format_number(stats["tokens"]), - str(stats["calls"]), + m["model_name"], + f"${m['total_cost_usd']:.4f}", + _format_number(m["input_tokens"] + m["output_tokens"]), + str(m["call_count"]), ) console.print(model_table) @@ -293,19 +264,21 @@ def export_data( db = _get_db() try: + if format not in ("csv", "json"): + console.print(f"[red]Error:[/red] Unknown format '{format}'. Use 'csv' or 'json'.") + raise typer.Exit(1) + if task is not None: records = db.get_batch_token_usage(task_ids=[task]) else: - records = db.get_workspace_token_usage() + # Stream the workspace table — never buffered into a list. + records = db.get_token_usage_iter() if format == "csv": - MetricsTracker.export_to_csv(records, output) - elif format == "json": - MetricsTracker.export_to_json(records, output) + n = MetricsTracker.export_to_csv(records, output) else: - console.print(f"[red]Error:[/red] Unknown format '{format}'. Use 'csv' or 'json'.") - raise typer.Exit(1) + n = MetricsTracker.export_to_json(records, output) - console.print(f"Exported {len(records)} records to {output}") + console.print(f"Exported {n} records to {output}") finally: db.close() diff --git a/codeframe/lib/metrics_tracker.py b/codeframe/lib/metrics_tracker.py index 822e9f04..e6f01fa4 100644 --- a/codeframe/lib/metrics_tracker.py +++ b/codeframe/lib/metrics_tracker.py @@ -38,9 +38,11 @@ import csv import json import logging +import os import re +import tempfile from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, Iterable, Optional, TextIO, Union from codeframe.core.models import CallType, TokenUsage from codeframe.platform_store.database import Database @@ -334,30 +336,70 @@ def get_workspace_costs( "total_calls": int, } """ - records = self.db.get_workspace_token_usage( - start_date=start_date, end_date=end_date - ) - - result: Dict[str, Any] = { - "total_cost_usd": 0.0, - "total_tokens": 0, - "total_calls": len(records), + # Aggregation is pushed into SQL: the per-model rollup returns a + # handful of rows; summing those in Python is O(models), not O(records). + by_model = self.db.get_costs_by_model(start_date=start_date, end_date=end_date) + + total_cost = sum(m["total_cost_usd"] for m in by_model) + total_tokens = sum(m["input_tokens"] + m["output_tokens"] for m in by_model) + total_calls = sum(m["call_count"] for m in by_model) + + return { + "total_cost_usd": round(total_cost, 6), + "total_tokens": total_tokens, + "total_calls": total_calls, } - for record in records: - result["total_cost_usd"] += record["estimated_cost_usd"] - result["total_tokens"] += record["input_tokens"] + record["output_tokens"] + @staticmethod + def _atomic_stream_write( + output_path: str, write_fn: Callable[[TextIO], int] + ) -> int: + """Stream through a temp file in the same dir, then ``os.replace``. - result["total_cost_usd"] = round(result["total_cost_usd"], 6) - return result + The exporters write incrementally, so a mid-stream failure (source + iterator raises, disk fills) must not leave a truncated file at + ``output_path``. Writing to a sibling temp file and atomically renaming + on success means readers only ever see a complete export; the temp file + is unlinked on any failure. + + Args: + output_path: Final destination path. + write_fn: Callback that writes to the open file and returns a count. + + Returns: + Whatever ``write_fn`` returns (the record count). + """ + directory = os.path.dirname(os.path.abspath(output_path)) + # mkstemp creates the temp file 0600 and os.replace preserves that mode, + # so exports land owner-only (not umask-derived 0644). Intentional: token + # spend data is mildly sensitive and this file is the user's named output. + fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".export-", suffix=".tmp") + try: + with os.fdopen(fd, "w", newline="") as f: + count = write_fn(f) + os.replace(tmp_path, output_path) + return count + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise @staticmethod - def export_to_csv(records: List[Dict[str, Any]], output_path: str) -> None: - """Export token usage records to a CSV file. + def export_to_csv(records: Iterable[Dict[str, Any]], output_path: str) -> int: + """Stream token usage records to a CSV file. + + Consumes ``records`` lazily (accepts the ``get_token_usage_iter`` + generator) so a large table is never buffered into a list. Written + atomically — a partial write never lands at ``output_path``. Args: - records: List of token usage record dictionaries - output_path: Path to write the CSV file + records: Iterable of token usage record dictionaries. + output_path: Path to write the CSV file. + + Returns: + Number of rows written. """ fieldnames = [ "id", "task_id", "agent_id", "project_id", "model_name", @@ -365,40 +407,54 @@ def export_to_csv(records: List[Dict[str, Any]], output_path: str) -> None: "actual_cost_usd", "call_type", "session_id", "timestamp", ] - with open(output_path, "w", newline="") as f: + def _write(f: TextIO) -> int: writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") writer.writeheader() + count = 0 for record in records: writer.writerow(record) + count += 1 + return count + + return MetricsTracker._atomic_stream_write(output_path, _write) @staticmethod - def export_to_json(records: List[Dict[str, Any]], output_path: str) -> None: - """Export token usage records to a JSON file with metadata. + def export_to_json(records: Iterable[Dict[str, Any]], output_path: str) -> int: + """Stream token usage records to a JSON file with metadata. + + Writes the ``records`` array incrementally as it consumes the iterator, + so the whole table is never held in memory at once. ``metadata`` (with + ``record_count``) is written last; JSON object key order is not + semantically meaningful, so consumers using ``json.load`` are unaffected. + Written atomically — a partial write never lands at ``output_path``. Args: - records: List of token usage record dictionaries - output_path: Path to write the JSON file + records: Iterable of token usage record dictionaries. + output_path: Path to write the JSON file. + + Returns: + Number of records written. """ - # Convert sqlite3.Row objects to plain dicts if needed - serializable_records = [] - for record in records: - row = dict(record) - # Ensure all values are JSON-serializable - for key, value in row.items(): - if isinstance(value, datetime): - row[key] = value.isoformat() - serializable_records.append(row) - - data = { - "metadata": { + def _write(f: TextIO) -> int: + count = 0 + f.write('{\n "records": [') + for record in records: + row = { + k: (v.isoformat() if isinstance(v, datetime) else v) + for k, v in dict(record).items() + } + f.write("\n " if count == 0 else ",\n ") + f.write(json.dumps(row, default=str)) + count += 1 + f.write("\n ],\n") + metadata = { "exported_at": datetime.now(timezone.utc).isoformat(), - "record_count": len(serializable_records), - }, - "records": serializable_records, - } + "record_count": count, + } + f.write(' "metadata": ' + json.dumps(metadata) + "\n}\n") + return count - with open(output_path, "w") as f: - json.dump(data, f, indent=2, default=str) + return MetricsTracker._atomic_stream_write(output_path, _write) async def get_project_costs( self, diff --git a/codeframe/platform_store/database.py b/codeframe/platform_store/database.py index c1a9bca6..8e8e1635 100644 --- a/codeframe/platform_store/database.py +++ b/codeframe/platform_store/database.py @@ -271,6 +271,14 @@ def get_workspace_token_usage(self, *args, **kwargs): """Delegate to token_usage.get_workspace_token_usage().""" return self.token_usage.get_workspace_token_usage(*args, **kwargs) + def get_token_usage_iter(self, *args, **kwargs): + """Delegate to token_usage.get_token_usage_iter().""" + return self.token_usage.get_token_usage_iter(*args, **kwargs) + + def get_costs_by_model(self, *args, **kwargs): + """Delegate to token_usage.get_costs_by_model().""" + return self.token_usage.get_costs_by_model(*args, **kwargs) + # ----- Audit log ----- def create_audit_log(self, *args, **kwargs): """Delegate to audit_logs.create_audit_log().""" diff --git a/codeframe/platform_store/repositories/token_repository.py b/codeframe/platform_store/repositories/token_repository.py index 0aa88df1..b48008f2 100644 --- a/codeframe/platform_store/repositories/token_repository.py +++ b/codeframe/platform_store/repositories/token_repository.py @@ -4,7 +4,7 @@ """ from datetime import datetime, timedelta, timezone -from typing import List, Optional, Dict, Any, Union, TYPE_CHECKING +from typing import List, Optional, Dict, Any, Union, Iterator, TYPE_CHECKING import logging @@ -274,6 +274,110 @@ def get_workspace_token_usage( cursor.execute(query, params) return [dict(row) for row in cursor.fetchall()] + def get_token_usage_iter( + self, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + batch_size: int = 1000, + ) -> "Iterator[Dict[str, Any]]": + """Stream workspace token_usage rows without materialising the whole table. + + Yields records one at a time, pulling from SQLite in ``batch_size`` + chunks via ``fetchmany`` so an export of a large table never loads the + entire result set into a Python list. + + Args: + start_date: Optional start of date range (inclusive). + end_date: Optional end of date range (inclusive). + batch_size: Rows fetched per round-trip. + + Yields: + Token usage records as dictionaries, newest first. + """ + # ponytail: uses a dedicated cursor on the shared connection; the export + # path is single-threaded, so no other query interleaves mid-iteration. + query = "SELECT * FROM token_usage WHERE 1=1" + params: list = [] + if start_date is not None: + query += " AND timestamp >= ?" + params.append(start_date.isoformat()) + if end_date is not None: + query += " AND timestamp <= ?" + params.append(end_date.isoformat()) + query += " ORDER BY timestamp DESC" + + cursor = self.conn.cursor() + cursor.execute(query, params) + # try/finally closes the cursor even if the consumer breaks or raises + # mid-stream (GeneratorExit), rather than leaving it open until GC. + try: + while True: + batch = cursor.fetchmany(batch_size) + if not batch: + break + for row in batch: + yield dict(row) + finally: + cursor.close() + + def get_costs_by_model( + self, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + ) -> List[Dict[str, Any]]: + """Aggregate spend per model in SQL, honouring an optional date window. + + Replaces the old ``SELECT *`` + Python for-loop rollup used by + ``cf stats``: the SUM/COUNT/GROUP BY runs in SQLite and only a + handful of per-model rows come back. + + Args: + start_date: Optional start of date range (inclusive). + end_date: Optional end of date range (inclusive). + + Returns: + List of dicts sorted by total_cost_usd DESC:: + + { + "model_name": str, + "input_tokens": int, + "output_tokens": int, + "total_cost_usd": float, + "call_count": int, + } + """ + query = """ + SELECT + model_name, + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COALESCE(SUM(estimated_cost_usd), 0.0) AS total_cost_usd, + COUNT(*) AS call_count + FROM token_usage + WHERE 1=1 + """ + params: list = [] + if start_date is not None: + query += " AND timestamp >= ?" + params.append(start_date.isoformat()) + if end_date is not None: + query += " AND timestamp <= ?" + params.append(end_date.isoformat()) + query += " GROUP BY model_name ORDER BY total_cost_usd DESC" + + cursor = self.conn.cursor() + cursor.execute(query, params) + return [ + { + "model_name": row["model_name"], + "input_tokens": int(row["input_tokens"] or 0), + "output_tokens": int(row["output_tokens"] or 0), + "total_cost_usd": float(row["total_cost_usd"] or 0.0), + "call_count": int(row["call_count"] or 0), + } + for row in cursor.fetchall() + ] + def get_costs_summary(self, days: int) -> Dict[str, Any]: """Aggregate token_usage costs into daily buckets for analytics. diff --git a/tests/core/test_token_repository_aggregation.py b/tests/core/test_token_repository_aggregation.py new file mode 100644 index 00000000..fce4246b --- /dev/null +++ b/tests/core/test_token_repository_aggregation.py @@ -0,0 +1,163 @@ +"""Tests for SQL-side token_usage aggregation + streaming export (issue #752 / P2.3). + +Before the fix, `cf stats` pulled the whole `token_usage` table into Python and +aggregated with for-loops (`get_workspace_token_usage` → SELECT * → sum in +Python). These tests lock in that: + +1. `get_costs_by_model()` pushes the per-model rollup into SQL (SUM/GROUP BY), + honours the [start, end] window, and orders by cost DESC. +2. `get_token_usage_iter()` is a lazy generator (streams rows) and honours + the same window. +3. `MetricsTracker.get_workspace_costs()` derives its totals from the SQL rollup. +4. `export_to_csv` / `export_to_json` consume an iterator (no full-table list), + return the row count, and JSON round-trips. +""" + +import json +from datetime import datetime, timedelta, timezone +from pathlib import Path +from types import GeneratorType + +import pytest + +from codeframe.core.models import CallType, TokenUsage +from codeframe.core.workspace import create_or_load_workspace +from codeframe.lib.metrics_tracker import MetricsTracker +from codeframe.platform_store.database import Database + +pytestmark = pytest.mark.v2 + + +def _mk(db: Database, *, model: str, inp: int, out: int, cost: float, ts: datetime, + task_id: str = "t1", agent_id: str = "a1") -> None: + db.save_token_usage( + TokenUsage( + task_id=task_id, + agent_id=agent_id, + project_id=0, + model_name=model, + input_tokens=inp, + output_tokens=out, + estimated_cost_usd=cost, + call_type=CallType.TASK_EXECUTION, + timestamp=ts, + ) + ) + + +@pytest.fixture +def db(tmp_path: Path): + repo = tmp_path / "repo" + repo.mkdir() + ws = create_or_load_workspace(repo) + database = Database(str(ws.db_path)) + database.initialize() + now = datetime.now(timezone.utc) + # Two models, three records; one record 10 days old for window tests. + _mk(database, model="claude-sonnet-4-5", inp=100, out=50, cost=0.01, ts=now) + _mk(database, model="claude-sonnet-4-5", inp=200, out=100, cost=0.02, ts=now) + _mk(database, model="gpt-4o", inp=1000, out=500, cost=0.50, ts=now - timedelta(days=10)) + yield database + database.close() + + +class TestGetCostsByModel: + def test_groups_and_sums_per_model(self, db: Database): + rows = db.get_costs_by_model() + by_model = {r["model_name"]: r for r in rows} + + assert by_model["claude-sonnet-4-5"]["input_tokens"] == 300 + assert by_model["claude-sonnet-4-5"]["output_tokens"] == 150 + assert by_model["claude-sonnet-4-5"]["call_count"] == 2 + assert by_model["claude-sonnet-4-5"]["total_cost_usd"] == pytest.approx(0.03) + + assert by_model["gpt-4o"]["input_tokens"] == 1000 + assert by_model["gpt-4o"]["call_count"] == 1 + + def test_ordered_by_cost_desc(self, db: Database): + rows = db.get_costs_by_model() + costs = [r["total_cost_usd"] for r in rows] + assert costs == sorted(costs, reverse=True) + assert rows[0]["model_name"] == "gpt-4o" # 0.50 > 0.03 + + def test_window_excludes_old_records(self, db: Database): + start = datetime.now(timezone.utc) - timedelta(days=1) + rows = db.get_costs_by_model(start_date=start) + models = {r["model_name"] for r in rows} + assert models == {"claude-sonnet-4-5"} # the 10-day-old gpt-4o row is out + + def test_empty_window_returns_empty(self, db: Database): + future = datetime.now(timezone.utc) + timedelta(days=1) + assert db.get_costs_by_model(start_date=future) == [] + + +class TestStreamingIterator: + def test_is_a_generator(self, db: Database): + it = db.get_token_usage_iter() + assert isinstance(it, GeneratorType) + + def test_yields_all_rows(self, db: Database): + rows = list(db.get_token_usage_iter()) + assert len(rows) == 3 + assert all("model_name" in r for r in rows) + + def test_window_filters(self, db: Database): + start = datetime.now(timezone.utc) - timedelta(days=1) + rows = list(db.get_token_usage_iter(start_date=start)) + assert len(rows) == 2 # old gpt-4o excluded + + +class TestWorkspaceCostsTotals: + def test_totals_match_sql_rollup(self, db: Database): + tracker = MetricsTracker(db=db) + result = tracker.get_workspace_costs() + # 0.01 + 0.02 + 0.50 + assert result["total_cost_usd"] == pytest.approx(0.53) + # (100+50)+(200+100)+(1000+500) + assert result["total_tokens"] == 1950 + assert result["total_calls"] == 3 + + def test_totals_honour_window(self, db: Database): + tracker = MetricsTracker(db=db) + start = datetime.now(timezone.utc) - timedelta(days=1) + result = tracker.get_workspace_costs(start_date=start) + assert result["total_cost_usd"] == pytest.approx(0.03) + assert result["total_calls"] == 2 + + +class TestStreamingExport: + def test_csv_streams_iterator_and_counts(self, db: Database, tmp_path: Path): + out = tmp_path / "out.csv" + n = MetricsTracker.export_to_csv(db.get_token_usage_iter(), str(out)) + assert n == 3 + lines = out.read_text().strip().splitlines() + assert len(lines) == 4 # header + 3 rows + + def test_json_streams_and_roundtrips(self, db: Database, tmp_path: Path): + out = tmp_path / "out.json" + n = MetricsTracker.export_to_json(db.get_token_usage_iter(), str(out)) + assert n == 3 + data = json.loads(out.read_text()) + assert data["metadata"]["record_count"] == 3 + assert len(data["records"]) == 3 + + def test_json_empty_records_is_valid_json(self, tmp_path: Path): + out = tmp_path / "empty.json" + n = MetricsTracker.export_to_json(iter([]), str(out)) + assert n == 0 + data = json.loads(out.read_text()) # must not raise + assert data["records"] == [] + assert data["metadata"]["record_count"] == 0 + + def test_mid_stream_failure_leaves_no_partial_file(self, tmp_path: Path): + out = tmp_path / "partial.csv" + + def boom(): + yield {"id": 1, "input_tokens": 1, "output_tokens": 1} + raise RuntimeError("source blew up mid-stream") + + with pytest.raises(RuntimeError): + MetricsTracker.export_to_csv(boom(), str(out)) + # Atomic write: the destination must not exist, and no temp junk remains. + assert not out.exists() + assert list(tmp_path.iterdir()) == []