Spaces:
Sleeping
Sleeping
File size: 8,739 Bytes
8edee29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 | """
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(),
}
|