tracer-cloud--opensre/core/agent_harness/accounting/token_usage.py

"""Per-session token accounting, extracted from the session state object.

Groups the token-cost bookkeeping (``/cost``) into one cohesive value object so
the session state class does not carry accounting fields and methods directly.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class TokenUsage:
    """Accumulated token counts and LLM call count for one session.

    ``totals`` keeps running sums under ``input`` / ``output`` plus the
    ``*_measured`` / ``*_estimated`` breakdown buckets. ``call_count`` is the
    number of recorded LLM calls (for ``/cost``).
    """

    totals: dict[str, int] = field(default_factory=dict)
    call_count: int = 0

    @property
    def has_estimates(self) -> bool:
        """True when any recorded tokens were estimated rather than measured."""
        return bool(self.totals.get("input_estimated") or self.totals.get("output_estimated"))

    def record(
        self,
        *,
        input_tokens: int = 0,
        output_tokens: int = 0,
        estimated: bool = False,
    ) -> None:
        """Accumulate one LLM call's token counts (input/output + breakdown)."""
        if not input_tokens and not output_tokens:
            return
        suffix = "estimated" if estimated else "measured"
        for direction, count in (("input", input_tokens), ("output", output_tokens)):
            if not count:
                continue
            self.totals[direction] = self.totals.get(direction, 0) + count
            bucket = f"{direction}_{suffix}"
            self.totals[bucket] = self.totals.get(bucket, 0) + count
        self.call_count += 1

    def reset(self) -> None:
        """Clear all accumulated counts (used by ``/new``)."""
        self.totals.clear()
        self.call_count = 0


__all__ = ["TokenUsage"]