""" utils/token_counter.py ----------------------- Groq token usage tracker for AutoDevAgent. Tracks cumulative token consumption across all LLM calls in a single pipeline run and warns the user when they are approaching the Groq free tier rate limit. Design: - One TokenCounter instance is created per pipeline run and passed through the pipeline so every LLM call can update it. - Rate limit tracking is based on requests per minute (RPM), not tokens per minute, since Groq's free tier limit is RPM-based. - The warning threshold is configurable in config.py — default is 80% of the 30 RPM limit (i.e. warn at 24+ requests/minute). - Token counts are extracted from LangChain response objects using the helper in langsmith_tracer.py. - Provides a formatted summary string for display in the Gradio UI. Usage: from utils.token_counter import TokenCounter counter = TokenCounter() counter.record(prompt_tokens=150, completion_tokens=80, llm_calls=1) print(counter.summary()) # "230 tokens · 1 call · OK" print(counter.is_near_limit()) # False print(counter.warning_message()) # "" (no warning) """ import logging import time from dataclasses import dataclass, field from config import settings logger = logging.getLogger(__name__) # ------------------------------------------------------------------ # # Call record # # ------------------------------------------------------------------ # @dataclass class LLMCallRecord: """ Record of a single LLM call's token usage and timing. Stored in TokenCounter.call_history for detailed observability. Attributes: agent_name: Which agent made the call. prompt_tokens: Input tokens consumed. completion_tokens: Output tokens produced. timestamp: Unix timestamp when the call was made. """ agent_name: str = "" prompt_tokens: int = 0 completion_tokens: int = 0 timestamp: float = field(default_factory=time.time) @property def total_tokens(self) -> int: """Total tokens for this call.""" return self.prompt_tokens + self.completion_tokens # ------------------------------------------------------------------ # # Counter # # ------------------------------------------------------------------ # class TokenCounter: """ Tracks cumulative token usage and rate limit status for a pipeline run. Maintains a running total of all token counts and a per-minute request count to detect when the Groq free tier rate limit is being approached. Attributes: prompt_tokens: Total input tokens used across all calls. completion_tokens: Total output tokens produced. total_tokens: Sum of prompt + completion tokens. total_calls: Total LLM calls made this run. call_history: Detailed record of each individual call. _run_start: Unix timestamp when this counter was created. """ def __init__(self) -> None: """Initialise all counters to zero.""" self.prompt_tokens: int = 0 self.completion_tokens: int = 0 self.total_tokens: int = 0 self.total_calls: int = 0 self.call_history: list[LLMCallRecord] = [] self._run_start: float = time.time() def record( self, prompt_tokens: int, completion_tokens: int, agent_name: str = "", llm_calls: int = 1, ) -> None: """ Record token usage from one LLM call. Args: prompt_tokens: Input tokens consumed by this call. completion_tokens: Output tokens produced by this call. agent_name: Name of the agent that made the call. llm_calls: Number of LLM calls this represents (default 1). """ self.prompt_tokens += prompt_tokens self.completion_tokens += completion_tokens self.total_tokens += prompt_tokens + completion_tokens self.total_calls += llm_calls record = LLMCallRecord( agent_name = agent_name, prompt_tokens = prompt_tokens, completion_tokens = completion_tokens, ) self.call_history.append(record) logger.debug( "TokenCounter: %s used %d tokens (%d prompt + %d completion). " "Running total: %d tokens across %d calls.", agent_name or "unknown", prompt_tokens + completion_tokens, prompt_tokens, completion_tokens, self.total_tokens, self.total_calls, ) def calls_in_last_minute(self) -> int: """ Count how many LLM calls were made in the last 60 seconds. Used to estimate whether we are approaching the Groq RPM limit. Returns: Integer count of calls in the last 60 seconds. """ now = time.time() cutoff = now - 60.0 recent_calls = [r for r in self.call_history if r.timestamp >= cutoff] return len(recent_calls) def is_near_limit(self) -> bool: """ Return True if we are approaching the Groq rate limit. Compares calls in the last minute against the configured warning threshold (default: 80% of 30 RPM = 24 calls/min). Returns: True if a warning should be shown to the user. """ calls_per_minute = self.calls_in_last_minute() threshold = int( settings.groq_rate_limit_rpm * settings.groq_rate_limit_warning_threshold ) near = calls_per_minute >= threshold if near: logger.warning( "TokenCounter: approaching rate limit — %d calls in last minute " "(threshold: %d, limit: %d)", calls_per_minute, threshold, settings.groq_rate_limit_rpm, ) return near def warning_message(self) -> str: """ Return a warning string if approaching the rate limit, else empty. Returns: Warning string for the UI, or "" if no warning is needed. """ if not self.is_near_limit(): return "" calls = self.calls_in_last_minute() return ( f"⚠ Approaching Groq rate limit: {calls} calls in the last minute " f"(limit: {settings.groq_rate_limit_rpm} RPM). " "Consider switching to llama3-8b to reduce call overhead." ) def summary(self) -> str: """ Return a compact summary string for the UI stats strip. Format: "1,250 tokens · 6 calls · OK" or "1,250 tokens · 6 calls · ⚠ near limit" Returns: Formatted summary string. """ status = "⚠ near limit" if self.is_near_limit() else "OK" return ( f"{self.total_tokens:,} tokens · " f"{self.total_calls} call{'s' if self.total_calls != 1 else ''} · " f"{status}" ) def elapsed_seconds(self) -> float: """ Return seconds elapsed since this counter was created. Used to compute total pipeline run time for the stats strip. Returns: Float seconds elapsed. """ return round(time.time() - self._run_start, 2) def reset(self) -> None: """ Reset all counters to zero for a new pipeline run. Call this at the start of each new run to prevent state bleeding from one task into the next. """ self.prompt_tokens = 0 self.completion_tokens = 0 self.total_tokens = 0 self.total_calls = 0 self.call_history = [] self._run_start = time.time() logger.debug("TokenCounter reset for new run") def to_dict(self) -> dict: """ Serialize counter state to a dict for W&B logging or session history. Returns: Dict with token counts and call breakdown. """ return { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, "total_tokens": self.total_tokens, "total_calls": self.total_calls, "elapsed_seconds": self.elapsed_seconds(), "is_near_limit": self.is_near_limit(), }