240 lines
8.1 KiB
Python
240 lines
8.1 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
"""Token budget state owns process-wide LLM prompt budget settings."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from band_aider.judgement.repo_exec import DEFAULT_REPO_EXEC, RepoExec, resolve_repo_exec
|
|
|
|
from band_aider.context.band_aider_facts import (
|
|
FACT_BUCKET_WINDOW_SECONDS,
|
|
get_band_aider_facts,
|
|
reset_band_aider_facts_for_tests,
|
|
)
|
|
from band_aider.ui.dto.token_budget_stats import (
|
|
TokenBudgetBucketDTO,
|
|
TokenBudgetStatsDTO,
|
|
TokenUsageBreakdownDTO,
|
|
)
|
|
|
|
DEFAULT_MAX_PROMPT_TOKENS = 128000
|
|
DEFAULT_MAX_OUTPUT_TOKENS = 16000
|
|
DEFAULT_TOKENS_PER_WORD = 2.8
|
|
DEFAULT_BATCH_TARGET_RATIO = 0.85
|
|
TOKEN_BUCKET_WINDOW_SECONDS = FACT_BUCKET_WINDOW_SECONDS
|
|
# TODO: Deliberate with a human whether to remove this arbitrary token-history cap.
|
|
TOKEN_BUCKET_HISTORY_COUNT = 10
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenBudgetBucket:
|
|
"""TokenBudgetBucket stores one five-minute token usage bucket."""
|
|
|
|
started_at: int
|
|
tokens_used: int = 0
|
|
cache_hit_tokens: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenBudgetSettings:
|
|
"""TokenBudgetSettings records process-wide LLM prompt budget configuration."""
|
|
|
|
max_prompt_tokens: int = DEFAULT_MAX_PROMPT_TOKENS
|
|
max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS
|
|
tokens_per_word: float = DEFAULT_TOKENS_PER_WORD
|
|
batch_target_ratio: float = DEFAULT_BATCH_TARGET_RATIO
|
|
|
|
@property
|
|
def last_ai_token_used(self) -> int:
|
|
return get_band_aider_facts().last_query_tokens
|
|
|
|
@property
|
|
def total_tokens(self) -> int:
|
|
return get_band_aider_facts().total_tokens_spent
|
|
|
|
@property
|
|
def max_ai_token_used(self) -> int:
|
|
return get_band_aider_facts().largest_query_tokens
|
|
|
|
@property
|
|
def five_minute_buckets(self) -> tuple[TokenBudgetBucket, ...]:
|
|
return tuple(
|
|
TokenBudgetBucket(
|
|
started_at=bucket.started_at,
|
|
tokens_used=bucket.tokens_spent,
|
|
cache_hit_tokens=bucket.cache_hit_tokens,
|
|
)
|
|
for bucket in get_band_aider_facts().buckets[-TOKEN_BUCKET_HISTORY_COUNT:]
|
|
)
|
|
|
|
|
|
_settings: TokenBudgetSettings | None = None
|
|
|
|
|
|
def configure_token_budget_settings(
|
|
*,
|
|
max_prompt_tokens: int | None = None,
|
|
max_output_tokens: int | None = None,
|
|
tokens_per_word: float | None = None,
|
|
batch_target_ratio: float | None = None,
|
|
last_ai_token_used: int | None = None,
|
|
total_tokens: int | None = None,
|
|
max_ai_token_used: int | None = None,
|
|
five_minute_buckets: tuple[TokenBudgetBucket, ...] | None = None,
|
|
) -> TokenBudgetSettings:
|
|
"""Configure immutable token budget settings for this process."""
|
|
|
|
global _settings
|
|
if _settings is not None:
|
|
_settings = TokenBudgetSettings(
|
|
max_prompt_tokens=_settings.max_prompt_tokens if max_prompt_tokens is None else max_prompt_tokens,
|
|
max_output_tokens=_settings.max_output_tokens if max_output_tokens is None else max_output_tokens,
|
|
tokens_per_word=_settings.tokens_per_word if tokens_per_word is None else tokens_per_word,
|
|
batch_target_ratio=_settings.batch_target_ratio if batch_target_ratio is None else batch_target_ratio,
|
|
)
|
|
_apply_legacy_fact_overrides(
|
|
last_ai_token_used=last_ai_token_used,
|
|
total_tokens=total_tokens,
|
|
max_ai_token_used=max_ai_token_used,
|
|
five_minute_buckets=five_minute_buckets,
|
|
)
|
|
return _settings
|
|
|
|
_settings = TokenBudgetSettings(
|
|
max_prompt_tokens=max_prompt_tokens
|
|
if max_prompt_tokens is not None
|
|
else int(os.getenv("BAND_AIDER_MAX_PROMPT_TOKENS", str(DEFAULT_MAX_PROMPT_TOKENS))),
|
|
max_output_tokens=max_output_tokens
|
|
if max_output_tokens is not None
|
|
else int(os.getenv("BAND_AIDER_MAX_OUTPUT_TOKENS", str(DEFAULT_MAX_OUTPUT_TOKENS))),
|
|
tokens_per_word=tokens_per_word
|
|
if tokens_per_word is not None
|
|
else float(os.getenv("BAND_AIDER_TOKENS_PER_WORD", str(DEFAULT_TOKENS_PER_WORD))),
|
|
batch_target_ratio=batch_target_ratio
|
|
if batch_target_ratio is not None
|
|
else float(os.getenv("BAND_AIDER_CONTEXT_BATCH_TARGET_RATIO", str(DEFAULT_BATCH_TARGET_RATIO))),
|
|
)
|
|
_apply_legacy_fact_overrides(
|
|
last_ai_token_used=last_ai_token_used,
|
|
total_tokens=total_tokens,
|
|
max_ai_token_used=max_ai_token_used,
|
|
five_minute_buckets=five_minute_buckets,
|
|
)
|
|
return _settings
|
|
|
|
|
|
def get_token_budget_settings() -> TokenBudgetSettings:
|
|
"""Return token budget settings initialized from environment defaults when needed."""
|
|
|
|
return _settings or configure_token_budget_settings()
|
|
|
|
|
|
def reset_token_budget_state_for_tests() -> None:
|
|
"""Reset process token budget settings for focused tests."""
|
|
|
|
global _settings
|
|
_settings = None
|
|
reset_band_aider_facts_for_tests()
|
|
|
|
|
|
def token_budget_stats_dto() -> TokenBudgetStatsDTO:
|
|
settings = get_token_budget_settings()
|
|
facts = get_band_aider_facts().snapshot()
|
|
return TokenBudgetStatsDTO(
|
|
max_prompt_tokens=settings.max_prompt_tokens,
|
|
last_ai_token_used=facts.last_query_tokens,
|
|
total_tokens=facts.total_tokens_spent,
|
|
max_ai_token_used=facts.largest_query_tokens,
|
|
bucket_window_seconds=TOKEN_BUCKET_WINDOW_SECONDS,
|
|
buckets=tuple(
|
|
TokenBudgetBucketDTO(
|
|
started_at=bucket.started_at,
|
|
tokens_used=bucket.tokens_spent,
|
|
cache_hit_tokens=bucket.cache_hit_tokens,
|
|
)
|
|
for bucket in facts.buckets[-TOKEN_BUCKET_HISTORY_COUNT:]
|
|
),
|
|
baseline_tokens=facts.baseline_tokens_spent,
|
|
escalated_tokens=facts.escalated_tokens_spent,
|
|
token_breakdown=token_breakdown_dtos(facts.token_breakdown),
|
|
total_cache_hit_tokens=facts.total_cache_hit_tokens,
|
|
)
|
|
|
|
|
|
def token_breakdown_dtos(items) -> tuple[TokenUsageBreakdownDTO, ...]:
|
|
return tuple(
|
|
TokenUsageBreakdownDTO(
|
|
started_at=item.started_at,
|
|
model_name=item.model_name,
|
|
endpoint=item.endpoint,
|
|
level=item.level,
|
|
tokens_spent=item.tokens_spent,
|
|
cache_hit_tokens=item.cache_hit_tokens,
|
|
)
|
|
for item in items
|
|
)
|
|
|
|
|
|
def update_token_budget_stats(tokens_used: int, when: float | None = None) -> TokenBudgetSettings:
|
|
get_band_aider_facts().record_llm_call(
|
|
tokens_spent=max(0, int(tokens_used)),
|
|
wait_seconds=0.0,
|
|
wall_time=when if when is not None else time.time(),
|
|
)
|
|
return get_token_budget_settings()
|
|
|
|
|
|
def persist_token_budget_stats_json(
|
|
checkpoint_dir: Path,
|
|
*,
|
|
repo_exec: RepoExec = DEFAULT_REPO_EXEC,
|
|
) -> None:
|
|
get_band_aider_facts().persist(checkpoint_dir, repo_exec=resolve_repo_exec(repo_exec))
|
|
|
|
|
|
def restore_token_budget_stats_json(
|
|
checkpoint_dir: Path,
|
|
*,
|
|
repo_exec: RepoExec = DEFAULT_REPO_EXEC,
|
|
) -> None:
|
|
get_band_aider_facts().restore(checkpoint_dir, repo_exec=resolve_repo_exec(repo_exec))
|
|
|
|
|
|
def _apply_legacy_fact_overrides(
|
|
*,
|
|
last_ai_token_used: int | None,
|
|
total_tokens: int | None,
|
|
max_ai_token_used: int | None,
|
|
five_minute_buckets: tuple[TokenBudgetBucket, ...] | None,
|
|
) -> None:
|
|
"""Support old tests/callers that seeded token metrics through settings."""
|
|
|
|
if (
|
|
last_ai_token_used is None
|
|
and total_tokens is None
|
|
and max_ai_token_used is None
|
|
and five_minute_buckets is None
|
|
):
|
|
return
|
|
facts = get_band_aider_facts()
|
|
facts.last_query_tokens = max(0, int(last_ai_token_used or 0))
|
|
facts.total_tokens_spent = max(0, int(total_tokens or 0))
|
|
facts.largest_query_tokens = max(0, int(max_ai_token_used or facts.last_query_tokens))
|
|
if five_minute_buckets is not None:
|
|
from band_aider.context.band_aider_facts import BandAiderFactBucket
|
|
|
|
facts.buckets = tuple(
|
|
BandAiderFactBucket(
|
|
started_at=bucket.started_at,
|
|
tokens_spent=bucket.tokens_used,
|
|
cache_hit_tokens=bucket.cache_hit_tokens,
|
|
)
|
|
for bucket in five_minute_buckets
|
|
if bucket.started_at > 0
|
|
)
|