wehub-resource-sync
e4dcfc49aa
Tests / Import Check (Python 3.13) (push) Has been cancelled
Tests / Import Check (Python 3.14) (push) Has been cancelled
Tests / Python Tests (Python 3.11) (push) Has been cancelled
Tests / Python Tests (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.14) (push) Has been cancelled
Tests / Test Summary (push) Has been cancelled
Tests / Lint and Format (push) Has been cancelled
Tests / Web Node Tests (push) Has been cancelled
Tests / Import Check (Python 3.11) (push) Has been cancelled
Tests / Import Check (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.13) (push) Has been cancelled
96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
"""Token-usage accumulator shared across LLM calls within a single turn."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
|
|
class UsageTracker:
|
|
"""Accumulate prompt/completion tokens across many streaming LLM calls.
|
|
|
|
Two ingestion paths:
|
|
|
|
* :meth:`add_from_response` — read OpenAI ``CompletionUsage`` (or the
|
|
streaming ``usage`` chunk) when the provider returns it.
|
|
* :meth:`add_estimated` — fall back to a coarse ``chars / 3.5`` estimate
|
|
for providers that don't emit ``usage`` (used by chat's answer-now path).
|
|
|
|
Construct with ``model=<name>`` so :meth:`summary` can resolve a
|
|
``total_cost_usd`` via the pricing table in ``deeptutor.logging.stats``.
|
|
"""
|
|
|
|
def __init__(self, *, model: str | None = None) -> None:
|
|
self.prompt_tokens: int = 0
|
|
self.completion_tokens: int = 0
|
|
self.total_tokens: int = 0
|
|
self.calls: int = 0
|
|
self.model: str | None = model
|
|
|
|
def add_from_response(self, response_or_usage: Any) -> None:
|
|
usage = getattr(response_or_usage, "usage", None) or response_or_usage
|
|
prompt = int(getattr(usage, "prompt_tokens", 0) or 0)
|
|
completion = int(getattr(usage, "completion_tokens", 0) or 0)
|
|
total = int(getattr(usage, "total_tokens", prompt + completion) or 0)
|
|
if prompt or completion or total:
|
|
self.prompt_tokens += prompt
|
|
self.completion_tokens += completion
|
|
self.total_tokens += total
|
|
self.calls += 1
|
|
|
|
def add_estimated(self, *, input_chars: int, output_chars: int) -> None:
|
|
est_input = int(input_chars / 3.5)
|
|
est_output = int(output_chars / 3.5)
|
|
self.prompt_tokens += est_input
|
|
self.completion_tokens += est_output
|
|
self.total_tokens += est_input + est_output
|
|
self.calls += 1
|
|
|
|
def add_usage(
|
|
self,
|
|
*,
|
|
agent_name: str = "",
|
|
stage: str = "",
|
|
model: str = "",
|
|
system_prompt: str = "",
|
|
user_prompt: str = "",
|
|
response_text: str = "",
|
|
) -> None:
|
|
"""Adapter for :class:`~deeptutor.agents.base_agent.BaseAgent`.
|
|
|
|
``BaseAgent._track_tokens`` looks for an external tracker exposing
|
|
``add_usage(...)``; this method lets a :class:`UsageTracker` be
|
|
passed as the ``token_tracker`` constructor argument so a
|
|
capability pipeline can aggregate cost across all of its
|
|
BaseAgent-derived sub-agents in one place.
|
|
|
|
We fall back to a character-based estimate because BaseAgent
|
|
only hands us the prompt/response text (the raw provider usage
|
|
object is not available at that layer).
|
|
"""
|
|
if model and not self.model:
|
|
self.model = model
|
|
input_chars = len(system_prompt or "") + len(user_prompt or "")
|
|
output_chars = len(response_text or "")
|
|
if input_chars or output_chars:
|
|
self.add_estimated(input_chars=input_chars, output_chars=output_chars)
|
|
|
|
def summary(self) -> dict[str, Any] | None:
|
|
if self.calls == 0:
|
|
return None
|
|
cost_usd = 0.0
|
|
if self.model:
|
|
# Local import keeps ``core.agentic`` import-light at module load.
|
|
from deeptutor.logging.stats.llm_stats import get_pricing
|
|
|
|
pricing = get_pricing(self.model)
|
|
cost_usd = (self.prompt_tokens / 1000.0) * pricing.get("input", 0.0) + (
|
|
self.completion_tokens / 1000.0
|
|
) * pricing.get("output", 0.0)
|
|
return {
|
|
"total_cost_usd": cost_usd,
|
|
"total_tokens": self.total_tokens,
|
|
"total_calls": self.calls,
|
|
"prompt_tokens": self.prompt_tokens,
|
|
"completion_tokens": self.completion_tokens,
|
|
}
|