218 lines
8.3 KiB
Python
218 lines
8.3 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, field
|
|
import hashlib
|
|
import logging
|
|
import time
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient, HumanResponse, OpenAIConfig
|
|
from band_aider.ai_client.token_budget_estimator import TokenBudgetEstimator
|
|
from band_aider.context.band_aider_facts import get_band_aider_facts
|
|
from band_aider.ui.status import publish_supreme_status
|
|
from band_aider.context.token_budget_state import token_budget_stats_dto
|
|
from band_aider.ui.dto.model_escalation import ModelEscalationDTO
|
|
from band_aider.ui.status_publisher import DeferredHumanCommandReceiver, HumanCommandReceiver, NullStatusPublisher, StatusPublisher
|
|
|
|
|
|
@dataclass
|
|
class _TokenTrackerCallContext:
|
|
estimated_tokens: int
|
|
wall_time: float
|
|
started_at: float
|
|
level: str
|
|
parent: "_TokenTrackerCallContext | None" = None
|
|
suppress_outer_record: bool = False
|
|
|
|
|
|
_CURRENT_CALL: ContextVar[_TokenTrackerCallContext | None] = ContextVar("band_aider_token_tracker_call", default=None)
|
|
|
|
|
|
def record_provider_llm_call(
|
|
*,
|
|
usage: object,
|
|
input_field: str,
|
|
model_name: str,
|
|
endpoint: str,
|
|
) -> None:
|
|
"""Record one provider usage report or its estimator fallback."""
|
|
|
|
current = _CURRENT_CALL.get()
|
|
# Provider clients remain usable outside a tracked runtime without mutating global facts.
|
|
if current is None:
|
|
return
|
|
values = usage if isinstance(usage, dict) else {}
|
|
input_tokens = values.get(input_field)
|
|
details = values.get(f"{input_field}_details")
|
|
cache_hits = details.get("cached_tokens", 0) if isinstance(details, dict) else 0
|
|
# Provider totals replace estimation only when the provider supplied a usable integer count.
|
|
if isinstance(input_tokens, bool) or not isinstance(input_tokens, int) or input_tokens < 0:
|
|
input_tokens, cache_hits = current.estimated_tokens, 0
|
|
# Invalid cache detail means no confirmed cache hit, not an invalid prompt total.
|
|
elif isinstance(cache_hits, bool) or not isinstance(cache_hits, int) or cache_hits < 0:
|
|
cache_hits = 0
|
|
prompt_tokens = input_tokens
|
|
spent = max(0, prompt_tokens - cache_hits)
|
|
get_band_aider_facts().record_llm_call(
|
|
tokens_spent=spent,
|
|
cache_hit_tokens=cache_hits,
|
|
wait_seconds=time.monotonic() - current.started_at,
|
|
wall_time=current.wall_time,
|
|
model_name=model_name,
|
|
endpoint=endpoint,
|
|
level=current.level,
|
|
)
|
|
current.suppress_outer_record = True
|
|
|
|
|
|
def record_completion_cache_hit(
|
|
*,
|
|
model_name: str,
|
|
endpoint: str,
|
|
) -> None:
|
|
"""Record avoided prompt usage for one in-process completion-cache hit."""
|
|
|
|
current = _CURRENT_CALL.get()
|
|
# A cache decorator outside TokenTracker has no active accounting record to replace.
|
|
if current is None:
|
|
return
|
|
get_band_aider_facts().record_llm_call(
|
|
tokens_spent=0,
|
|
cache_hit_tokens=current.estimated_tokens,
|
|
wait_seconds=0.0,
|
|
wall_time=current.wall_time,
|
|
model_name=model_name,
|
|
endpoint=endpoint,
|
|
level=current.level,
|
|
provider_call=False,
|
|
)
|
|
current.suppress_outer_record = True
|
|
|
|
|
|
@dataclass(init=False)
|
|
class TokenTracker(CompletionClient):
|
|
"""TokenTracker records prompt token usage across all completion calls."""
|
|
|
|
estimator: TokenBudgetEstimator
|
|
config_index: int
|
|
config_count: int
|
|
event_publisher: StatusPublisher
|
|
_inner: CompletionClient = field(init=False, repr=False)
|
|
|
|
def __init__(
|
|
self,
|
|
inner: CompletionClient,
|
|
estimator: TokenBudgetEstimator | None = None,
|
|
config_index: int = 0,
|
|
config_count: int | None = None,
|
|
event_publisher: StatusPublisher | None = None,
|
|
command_receiver: HumanCommandReceiver | None = None,
|
|
) -> None:
|
|
self._inner = inner
|
|
self.estimator = estimator or TokenBudgetEstimator()
|
|
self.config_index = max(0, int(config_index))
|
|
self.config_count = max(1, int(config_count if config_count is not None else len(inner.openai_configs) or 1))
|
|
self.event_publisher = event_publisher or NullStatusPublisher()
|
|
self.command_receiver = command_receiver or DeferredHumanCommandReceiver()
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient:
|
|
return self._inner
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
return self.inner.openai_configs
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
estimated_tokens = max(0, int(self.estimator.estimate(prompt)))
|
|
wall_time = time.time()
|
|
started = time.monotonic()
|
|
level = "escalated" if self.config_index > 0 else "baseline"
|
|
context = _TokenTrackerCallContext(
|
|
estimated_tokens=estimated_tokens,
|
|
wall_time=wall_time,
|
|
started_at=started,
|
|
level=level,
|
|
parent=_CURRENT_CALL.get(),
|
|
)
|
|
token = _CURRENT_CALL.set(context)
|
|
try:
|
|
response = self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
# Human dictation is not a provider call and must not create estimated token spend.
|
|
if isinstance(response, HumanResponse):
|
|
context.suppress_outer_record = True
|
|
return response
|
|
finally:
|
|
if not context.suppress_outer_record:
|
|
config = self._active_config()
|
|
get_band_aider_facts().record_llm_call(
|
|
tokens_spent=estimated_tokens,
|
|
cache_hit_tokens=0,
|
|
wait_seconds=time.monotonic() - started,
|
|
wall_time=wall_time,
|
|
model_name=config.model,
|
|
endpoint=config.api_base,
|
|
level=level,
|
|
)
|
|
context.suppress_outer_record = True
|
|
if self.config_index > 0:
|
|
self._publish_escalation(prompt, estimated_tokens, wall_time, config)
|
|
# A nested tracker accounts the same physical call on behalf of its outer decorator stack.
|
|
if context.parent is not None and context.suppress_outer_record:
|
|
context.parent.suppress_outer_record = True
|
|
_CURRENT_CALL.reset(token)
|
|
publish_supreme_status(token_budget=token_budget_stats_dto())
|
|
|
|
def _active_config(self) -> OpenAIConfig:
|
|
configs = self.openai_configs
|
|
if configs:
|
|
return configs[0]
|
|
model = str(getattr(self.inner, "model", "unknown") or "unknown")
|
|
endpoint = str(getattr(self.inner, "api_base", "unknown") or "unknown")
|
|
return OpenAIConfig(model=model, api_base=endpoint)
|
|
|
|
def _publish_escalation(
|
|
self,
|
|
prompt: str,
|
|
tokens_escalated: int,
|
|
occurred_at: float,
|
|
config: OpenAIConfig,
|
|
) -> None:
|
|
request_md5 = hashlib.md5(prompt.encode("utf-8")).hexdigest()
|
|
self.event_publisher.publish(
|
|
ModelEscalationDTO(
|
|
trigger_name=type(self).__name__,
|
|
occurred_at=occurred_at,
|
|
prompt_preview=_prompt_preview(prompt),
|
|
reason="Completion used an escalated model tier.",
|
|
request_md5=request_md5,
|
|
local_model=self._local_model_name(),
|
|
active_model=config.model,
|
|
config_index=self.config_index,
|
|
config_count=self.config_count,
|
|
state="spent",
|
|
tokens_escalated=tokens_escalated,
|
|
active_endpoint=config.api_base,
|
|
)
|
|
)
|
|
|
|
def _local_model_name(self) -> str:
|
|
configs = self.openai_configs
|
|
if self.config_index == 0 and configs:
|
|
return configs[0].model
|
|
return "baseline"
|
|
|
|
|
|
def _prompt_preview(prompt: str, max_lines: int = 3, max_width: int = 88) -> str:
|
|
# TODO: Deliberate with a human whether to remove these arbitrary prompt-preview caps.
|
|
preview_lines: list[str] = []
|
|
for line in prompt.splitlines()[:max_lines]:
|
|
cleaned = " ".join(line.strip().split())
|
|
if len(cleaned) > max_width:
|
|
cleaned = f"{cleaned[: max_width - 3].rstrip()}..."
|
|
preview_lines.append(cleaned)
|
|
return "\n".join(preview_lines).strip() or prompt[:max_width].strip()
|