Spaces:
Sleeping
Sleeping
| """ | |
| 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 # | |
| # ------------------------------------------------------------------ # | |
| 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) | |
| 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(), | |
| } | |