271 lines
9.2 KiB
Python
271 lines
9.2 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
from band_aider.ai_client.cached_completion_client import CachedCompletionClient, CompletionCache
|
|
from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
from band_aider.ai_client.token_budget_estimator import TokenBudgetEstimator
|
|
from band_aider.ai_client.token_tracker import TokenTracker, record_provider_llm_call
|
|
from band_aider.context.band_aider_facts import get_band_aider_facts, reset_band_aider_facts_for_tests
|
|
from band_aider.context.token_budget_state import token_budget_stats_dto
|
|
from band_aider.ui.dto.model_escalation import ModelEscalationDTO
|
|
|
|
|
|
class _Client(CompletionClient):
|
|
@property
|
|
def inner(self) -> CompletionClient | None:
|
|
return None
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
return (OpenAIConfig(model="gpt-5.5", api_base="https://cloud.example", api_key="key"),)
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Return a stationary completion without a provider request."""
|
|
|
|
del logger
|
|
return f"done: {prompt}"
|
|
|
|
|
|
class _Estimator(TokenBudgetEstimator):
|
|
"""Use a fixed estimator result for precise accounting assertions."""
|
|
|
|
def estimate(self, text: str) -> int:
|
|
"""Return the stationary prompt estimate."""
|
|
|
|
del text
|
|
return 50
|
|
|
|
|
|
class _Publisher:
|
|
def __init__(self) -> None:
|
|
self.events: list[object] = []
|
|
|
|
def publish(self, event: object) -> None:
|
|
self.events.append(event)
|
|
|
|
|
|
class _ProviderUsageClient(CompletionClient):
|
|
"""Nominal provider boundary that records one stationary usage response."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
response: str,
|
|
usage: object,
|
|
input_field: str,
|
|
config: OpenAIConfig,
|
|
) -> None:
|
|
self._response = response
|
|
self._usage = usage
|
|
self._input_field = input_field
|
|
self._config = config
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient | None:
|
|
"""Declare this direct fake as a root provider client."""
|
|
|
|
return None
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
"""Expose the one stable provider configuration."""
|
|
|
|
return (self._config,)
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Record the fixed provider usage inside the active tracker scope."""
|
|
|
|
del prompt, logger
|
|
record_provider_llm_call(
|
|
usage=self._usage,
|
|
input_field=self._input_field,
|
|
model_name=self._config.model,
|
|
endpoint=self._config.api_base,
|
|
)
|
|
return self._response
|
|
|
|
|
|
def test_token_tracker_records_escalated_model_endpoint_and_publishes_spend_event() -> None:
|
|
"""Verify escalated runtime clients populate granular token facts."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
publisher = _Publisher()
|
|
tracker = TokenTracker(
|
|
inner=_Client(),
|
|
estimator=_Estimator(),
|
|
config_index=1,
|
|
config_count=2,
|
|
event_publisher=publisher,
|
|
)
|
|
|
|
assert tracker.complete("prompt") == "done: prompt"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == 50
|
|
assert snapshot.baseline_tokens_spent == 0
|
|
assert snapshot.escalated_tokens_spent == 50
|
|
assert [
|
|
(item.level, item.model_name, item.endpoint, item.tokens_spent)
|
|
for item in snapshot.token_breakdown
|
|
] == [("escalated", "gpt-5.5", "https://cloud.example", 50)]
|
|
events = [event for event in publisher.events if isinstance(event, ModelEscalationDTO)]
|
|
assert len(events) == 1
|
|
assert events[0].tokens_escalated == 50
|
|
assert events[0].active_endpoint == "https://cloud.example"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("usage", "spent", "cache_hits"),
|
|
(
|
|
({"prompt_tokens": 100}, 100, 0),
|
|
({"prompt_tokens": 100, "prompt_tokens_details": {"cached_tokens": 40}}, 60, 40),
|
|
({"prompt_tokens": 100, "prompt_tokens_details": {"cached_tokens": 100}}, 0, 100),
|
|
),
|
|
ids=("no_cache", "partial_cache", "full_cache"),
|
|
)
|
|
def test_provider_usage_separates_spent_from_cached_input(usage, spent, cache_hits) -> None:
|
|
"""Verify provider prompt usage assigns each token to spend or cache hits, never both."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
tracker = TokenTracker(
|
|
_ProviderUsageClient(
|
|
response="done",
|
|
usage=usage,
|
|
input_field="prompt_tokens",
|
|
config=OpenAIConfig(model="chat-model", api_base="https://chat.invalid"),
|
|
),
|
|
estimator=_Estimator(),
|
|
)
|
|
|
|
assert tracker.complete("prompt") == "done"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == spent
|
|
assert snapshot.total_cache_hit_tokens == cache_hits
|
|
assert snapshot.llm_call_count == 1
|
|
assert snapshot.token_breakdown[0].model_name == "chat-model"
|
|
assert snapshot.token_breakdown[0].cache_hit_tokens == cache_hits
|
|
dto = token_budget_stats_dto()
|
|
assert (dto.total_tokens, dto.total_cache_hit_tokens) == (spent, cache_hits)
|
|
|
|
|
|
def test_responses_provider_usage_separates_cached_input() -> None:
|
|
"""Verify Responses-compatible input usage follows the same accounting invariant."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
tracker = TokenTracker(
|
|
_ProviderUsageClient(
|
|
response="done",
|
|
usage={"input_tokens": 90, "input_tokens_details": {"cached_tokens": 30}},
|
|
input_field="input_tokens",
|
|
config=OpenAIConfig(model="responses-model", api_base="https://responses.invalid"),
|
|
),
|
|
estimator=_Estimator(),
|
|
)
|
|
|
|
assert tracker.complete("prompt") == "done"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == 60
|
|
assert snapshot.total_cache_hit_tokens == 30
|
|
|
|
|
|
def test_provider_without_usable_usage_falls_back_to_estimator() -> None:
|
|
"""Verify absent provider usage retains the existing prompt estimator fallback."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
tracker = TokenTracker(
|
|
_ProviderUsageClient(
|
|
response="done",
|
|
usage=None,
|
|
input_field="prompt_tokens",
|
|
config=OpenAIConfig(model="chat-model", api_base="https://chat.invalid"),
|
|
),
|
|
estimator=_Estimator(),
|
|
)
|
|
|
|
assert tracker.complete("prompt") == "done"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == 50
|
|
assert snapshot.total_cache_hit_tokens == 0
|
|
|
|
|
|
def test_in_process_completion_cache_hit_records_zero_spent_for_wrapped_config() -> None:
|
|
"""Verify a Band-Aider cache hit records avoided input without a provider call."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
cache = CompletionCache()
|
|
cached = CachedCompletionClient(_Client(), cache=cache)
|
|
tracker = TokenTracker(cached, estimator=_Estimator())
|
|
assert tracker.complete("prompt") == "done: prompt"
|
|
reset_band_aider_facts_for_tests()
|
|
|
|
assert tracker.complete("prompt") == "done: prompt"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == 0
|
|
assert snapshot.total_cache_hit_tokens == 50
|
|
assert snapshot.llm_call_count == 0
|
|
assert [
|
|
(item.model_name, item.endpoint, item.tokens_spent, item.cache_hit_tokens)
|
|
for item in snapshot.token_breakdown
|
|
] == [("gpt-5.5", "https://cloud.example", 0, 50)]
|
|
|
|
|
|
def test_baseline_fallback_and_escalated_usage_each_record_once() -> None:
|
|
"""Verify later escalated usage cannot erase a baseline provider fallback record."""
|
|
|
|
reset_band_aider_facts_for_tests()
|
|
|
|
baseline = _ProviderUsageClient(
|
|
response="baseline",
|
|
usage=None,
|
|
input_field="prompt_tokens",
|
|
config=OpenAIConfig(model="baseline-model", api_base="https://baseline.invalid"),
|
|
)
|
|
escalated = TokenTracker(
|
|
_ProviderUsageClient(
|
|
response="escalated",
|
|
usage={"prompt_tokens": 80, "prompt_tokens_details": {"cached_tokens": 20}},
|
|
input_field="prompt_tokens",
|
|
config=OpenAIConfig(model="escalated-model", api_base="https://escalated.invalid"),
|
|
),
|
|
estimator=_Estimator(),
|
|
config_index=1,
|
|
config_count=2,
|
|
)
|
|
|
|
class _TwoProviderCalls(CompletionClient):
|
|
"""_TwoProviderCalls makes one baseline call followed by one escalated call."""
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient:
|
|
return baseline
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
return baseline.openai_configs
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
baseline.complete(prompt, logger=logger)
|
|
return escalated.complete(prompt, logger=logger)
|
|
|
|
tracker = TokenTracker(_TwoProviderCalls(), estimator=_Estimator(), config_count=2)
|
|
|
|
assert tracker.complete("prompt") == "escalated"
|
|
|
|
snapshot = get_band_aider_facts().snapshot()
|
|
assert snapshot.total_tokens_spent == 110
|
|
assert snapshot.total_cache_hit_tokens == 20
|
|
assert snapshot.baseline_tokens_spent == 50
|
|
assert snapshot.escalated_tokens_spent == 60
|
|
assert snapshot.llm_call_count == 2
|
|
assert len(snapshot.token_breakdown) == 2
|