402 lines
14 KiB
Python
402 lines
14 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
from band_aider.ai_client.json_array_prompt_reducer import GeneralPromptReducer
|
|
from band_aider.ai_client.llm_context_batch import LLMContextBatch
|
|
from band_aider.ai_client.llm_context_reduction import LLMContextReduction
|
|
from band_aider.ai_client.token_budget_estimator import TokenBudgetEstimator
|
|
from band_aider.ai_client.token_miser import TokenMiser
|
|
from band_aider.context.class_meta import ClassMeta
|
|
from band_aider.context.token_budget_state import (
|
|
DEFAULT_MAX_OUTPUT_TOKENS,
|
|
TOKEN_BUCKET_HISTORY_COUNT,
|
|
TOKEN_BUCKET_WINDOW_SECONDS,
|
|
TokenBudgetBucket,
|
|
TokenBudgetSettings,
|
|
configure_token_budget_settings,
|
|
get_token_budget_settings,
|
|
persist_token_budget_stats_json,
|
|
reset_token_budget_state_for_tests,
|
|
restore_token_budget_stats_json,
|
|
update_token_budget_stats,
|
|
)
|
|
from band_aider.judgement.intent_check import IntentCheck
|
|
from band_aider.judgement.prompt_formatter import MarkdownPromptFormatterBuilder
|
|
from tests.support.memory_repo_exec import MemoryRepoExec
|
|
|
|
|
|
class _RecordingClient(CompletionClient):
|
|
"""_RecordingClient records prompts and returns configured responses."""
|
|
|
|
model = "default"
|
|
api_base = "https://example.invalid"
|
|
api_key = "key"
|
|
temperature = 0
|
|
max_tokens = None
|
|
max_attempts = 8
|
|
retry_temperature_step = 0.2
|
|
max_retry_temperature = 0.8
|
|
|
|
def __init__(self, responses: list[str] | None = None) -> None:
|
|
self.prompts: list[str] = []
|
|
self.responses = responses or ["ok"]
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Record the prompt and return the next stationary response."""
|
|
|
|
del logger
|
|
self.prompts.append(prompt)
|
|
index = min(len(self.prompts) - 1, len(self.responses) - 1)
|
|
return self.responses[index]
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
"""Expose no external provider configuration for this in-memory client."""
|
|
|
|
return ()
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient | None:
|
|
"""Declare this stationary client as the root completion boundary."""
|
|
|
|
return None
|
|
|
|
|
|
class _BudgetRecordingReducer(GeneralPromptReducer):
|
|
"""Reducer that records the calculated batch budget without performing reduction."""
|
|
|
|
def __init__(self) -> None:
|
|
self.token_budgets: list[int] = []
|
|
|
|
def can_reduce(self, prompt: str) -> bool:
|
|
"""Accept every prompt to observe the proposed budget."""
|
|
|
|
return True
|
|
|
|
def split(
|
|
self,
|
|
prompt: str,
|
|
token_budget: int,
|
|
estimator: TokenBudgetEstimator,
|
|
) -> list[LLMContextBatch]:
|
|
"""Record the budget and produce no batches for this focused boundary test."""
|
|
|
|
del prompt, estimator
|
|
self.token_budgets.append(token_budget)
|
|
return []
|
|
|
|
|
|
def _intent_prompt(count: int) -> str:
|
|
"""Return an IntentCheck-style prompt with many relevant candidates."""
|
|
|
|
candidates = [
|
|
{
|
|
"location": f"module_{index}.py",
|
|
"summary": " ".join(["candidate summary"] * 8),
|
|
"relevance": [" ".join(["important reason"] * 8)],
|
|
"relevant_symbols": {"Symbol": " ".join(["symbol reason"] * 8)},
|
|
}
|
|
for index in range(count)
|
|
]
|
|
return f"""
|
|
You are IntentCheck.
|
|
|
|
Return strict JSON only:
|
|
{{
|
|
"intent": "new_feature",
|
|
"new_feature": true,
|
|
"reason": "one sentence"
|
|
}}
|
|
|
|
User request:
|
|
Do work.
|
|
|
|
Relevant candidates:
|
|
{json.dumps(candidates, indent=2)}
|
|
""".strip()
|
|
|
|
|
|
def test_token_budget_estimator_estimates_nonzero_and_respects_tokens_per_word() -> None:
|
|
"""Verify the heuristic estimator scales with configured tokens_per_word."""
|
|
|
|
low = TokenBudgetEstimator(tokens_per_word=1.0).estimate("alpha beta")
|
|
high = TokenBudgetEstimator(tokens_per_word=3.0).estimate("alpha beta")
|
|
|
|
assert low > 0
|
|
assert high > low
|
|
|
|
|
|
def test_token_miser_forwards_unchanged_when_under_budget() -> None:
|
|
"""Verify under-budget prompts pass through without reducer changes."""
|
|
|
|
inner = _RecordingClient(["unchanged"])
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(max_prompt_tokens=10_000),
|
|
reducers=[GeneralPromptReducer()],
|
|
)
|
|
|
|
assert client.complete("small prompt") == "unchanged"
|
|
assert inner.prompts == ["small prompt"]
|
|
|
|
|
|
def test_generic_reducer_matches_plain_unmarked_text() -> None:
|
|
"""Verify the reducer can handle unmarked text prompts generically."""
|
|
|
|
reducer = GeneralPromptReducer()
|
|
prompt = "\n".join(["plain"] * 200)
|
|
|
|
assert reducer.can_reduce(prompt) is True
|
|
batches = reducer.split(prompt, token_budget=120, estimator=TokenBudgetEstimator(tokens_per_word=1.0))
|
|
|
|
assert len(batches) > 1
|
|
assert "".join(batch.content for batch in batches).strip()
|
|
|
|
|
|
def test_generic_reducer_prompt_carries_structure_preservation_guidance() -> None:
|
|
"""Verify reduction prompts tell the LLM to preserve structure best effort."""
|
|
|
|
reducer = GeneralPromptReducer()
|
|
batches = reducer.split(_intent_prompt(6), token_budget=280, estimator=TokenBudgetEstimator(tokens_per_word=1.0))
|
|
|
|
second_prompt = reducer.prompt_for_batch(_intent_prompt(6), batches[1], "[Reduced previous batch]")
|
|
|
|
assert "Preserve structure the best you can." in second_prompt
|
|
assert "Keep merge JSON, arrays, headings, sections" in second_prompt
|
|
assert "Keep keys, labels, file paths, test names, class names, member names, and output schema when visible." in second_prompt
|
|
assert "Previous reduced prompt:" in second_prompt
|
|
assert "Merge the previous reduced prompt with the current batch into one reduced prompt." in second_prompt
|
|
|
|
|
|
def test_token_miser_uses_prompt_budget_minus_reserved_output_budget() -> None:
|
|
"""Verify reduction batches reserve output budget from the overall prompt budget."""
|
|
|
|
inner = _RecordingClient(["forwarded"])
|
|
reducer = _BudgetRecordingReducer()
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(
|
|
max_prompt_tokens=100,
|
|
max_output_tokens=20,
|
|
tokens_per_word=1.0,
|
|
batch_target_ratio=0.5,
|
|
),
|
|
reducers=[reducer],
|
|
estimator=TokenBudgetEstimator(tokens_per_word=1.0),
|
|
)
|
|
|
|
client.complete(_intent_prompt(8))
|
|
|
|
assert reducer.token_budgets == [40]
|
|
|
|
|
|
def test_token_miser_reduces_then_runs_final_completion() -> None:
|
|
"""Verify oversized prompts are reduced first and completed once on the reduced prompt."""
|
|
|
|
prompt = _intent_prompt(8)
|
|
inner = _RecordingClient(
|
|
[
|
|
"[Reduced batch 1]",
|
|
"[Reduced batch 2]",
|
|
"[Reduced final prompt]",
|
|
'{"intent": "new_feature", "new_feature": true, "reason": "final"}',
|
|
]
|
|
)
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(max_prompt_tokens=450, max_output_tokens=225, tokens_per_word=1.0, batch_target_ratio=0.85),
|
|
reducers=[GeneralPromptReducer()],
|
|
estimator=TokenBudgetEstimator(tokens_per_word=1.0),
|
|
)
|
|
|
|
result = client.complete(prompt)
|
|
|
|
assert result == '{"intent": "new_feature", "new_feature": true, "reason": "final"}'
|
|
assert len(inner.prompts) > 4
|
|
assert "Reduce this oversized prompt so a later pass can continue." in inner.prompts[0]
|
|
assert inner.prompts[-1] == '{"intent": "new_feature", "new_feature": true, "reason": "final"}'
|
|
|
|
|
|
def test_token_miser_logs_batch_progress_without_batch_index_attribute(caplog) -> None:
|
|
"""Verify reduction logging uses loop enumeration rather than a non-existent batch.index field."""
|
|
|
|
class _Reducer(GeneralPromptReducer):
|
|
"""Stationary reducer that exposes explicit recurrent batch records."""
|
|
|
|
def can_reduce(self, prompt: str) -> bool:
|
|
return True
|
|
|
|
def split(self, prompt: str, token_budget: int, estimator: TokenBudgetEstimator) -> list[LLMContextBatch]:
|
|
return [
|
|
LLMContextBatch(batch_index=0, total_batches=2, content="first"),
|
|
LLMContextBatch(batch_index=1, total_batches=2, content="second"),
|
|
]
|
|
|
|
def prompt_for_batch(
|
|
self,
|
|
original_prompt: str,
|
|
batch: LLMContextBatch,
|
|
previous_result: str,
|
|
) -> str:
|
|
return f"{batch.content}:{previous_result}"
|
|
|
|
def final_result(self, reductions: list[LLMContextReduction]) -> str:
|
|
return "reduced prompt"
|
|
|
|
inner = _RecordingClient(["r1", "r2", "done"])
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(max_prompt_tokens=3, max_output_tokens=1, tokens_per_word=1.0),
|
|
reducers=[_Reducer()],
|
|
estimator=TokenBudgetEstimator(tokens_per_word=1.0),
|
|
)
|
|
|
|
logger = logging.getLogger("test_token_miser")
|
|
with caplog.at_level(logging.INFO, logger="test_token_miser"):
|
|
assert client.complete("this prompt is oversized", logger=logger) == "done"
|
|
|
|
assert "TokenMiser reducing batch 1/2 with _Reducer." in caplog.text
|
|
assert "TokenMiser reducing batch 2/2 with _Reducer." in caplog.text
|
|
|
|
|
|
def test_oversized_prompt_without_any_reducer_logs_and_forwards() -> None:
|
|
"""Verify oversized prompts are forwarded unchanged when no reducer is configured."""
|
|
|
|
prompt = " ".join(["plain"] * 100)
|
|
inner = _RecordingClient(["forwarded"])
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(max_prompt_tokens=10, tokens_per_word=1.0),
|
|
reducers=[],
|
|
estimator=TokenBudgetEstimator(tokens_per_word=1.0),
|
|
)
|
|
|
|
assert client.complete(prompt) == "forwarded"
|
|
assert inner.prompts == [prompt]
|
|
|
|
|
|
def test_intent_check_accepts_decorated_client_without_call_semantic_changes() -> None:
|
|
"""Verify IntentCheck can call a decorated completion client normally."""
|
|
|
|
inner = _RecordingClient(['{"intent": "existing_files", "new_feature": false, "reason": "covered"}'])
|
|
client = TokenMiser(
|
|
inner=inner,
|
|
settings=TokenBudgetSettings(max_prompt_tokens=10_000),
|
|
reducers=[GeneralPromptReducer()],
|
|
)
|
|
check = IntentCheck(
|
|
changes=[],
|
|
prompt="Change existing file.",
|
|
client=client,
|
|
prompt_formatter_builder=MarkdownPromptFormatterBuilder(),
|
|
)
|
|
candidate = ClassMeta("existing.py", "Existing file.", ["relevant"], {"Existing": "reason"})
|
|
|
|
assert check(candidate) is False
|
|
assert check.reason == "covered"
|
|
assert len(inner.prompts) == 1
|
|
|
|
|
|
def test_token_budget_settings_can_be_reconfigured_for_tests() -> None:
|
|
"""An explicit test configuration replaces the preceding process-local settings."""
|
|
|
|
reset_token_budget_state_for_tests()
|
|
configure_token_budget_settings(max_prompt_tokens=123, max_output_tokens=61)
|
|
assert get_token_budget_settings().max_prompt_tokens == 123
|
|
assert get_token_budget_settings().max_output_tokens == 61
|
|
|
|
configure_token_budget_settings(max_prompt_tokens=456, max_output_tokens=228)
|
|
|
|
assert get_token_budget_settings().max_prompt_tokens == 456
|
|
assert get_token_budget_settings().max_output_tokens == 228
|
|
reset_token_budget_state_for_tests()
|
|
|
|
|
|
def test_token_budget_settings_default_output_tokens() -> None:
|
|
"""Verify the default output budget matches the configured default constant."""
|
|
|
|
reset_token_budget_state_for_tests()
|
|
|
|
assert get_token_budget_settings().max_output_tokens == DEFAULT_MAX_OUTPUT_TOKENS
|
|
reset_token_budget_state_for_tests()
|
|
|
|
|
|
def test_update_token_budget_stats_keeps_configured_five_minute_bucket_history() -> None:
|
|
reset_token_budget_state_for_tests()
|
|
|
|
for index in range(TOKEN_BUCKET_HISTORY_COUNT + 1):
|
|
update_token_budget_stats(
|
|
(index + 1) * 100,
|
|
when=TOKEN_BUCKET_WINDOW_SECONDS * index,
|
|
)
|
|
|
|
settings = get_token_budget_settings()
|
|
|
|
assert settings.last_ai_token_used == (TOKEN_BUCKET_HISTORY_COUNT + 1) * 100
|
|
assert settings.total_tokens == sum(
|
|
(index + 1) * 100
|
|
for index in range(TOKEN_BUCKET_HISTORY_COUNT + 1)
|
|
)
|
|
assert settings.five_minute_buckets == tuple(
|
|
TokenBudgetBucket(
|
|
started_at=TOKEN_BUCKET_WINDOW_SECONDS * index,
|
|
tokens_used=(index + 1) * 100,
|
|
)
|
|
for index in range(1, TOKEN_BUCKET_HISTORY_COUNT + 1)
|
|
)
|
|
reset_token_budget_state_for_tests()
|
|
|
|
|
|
def test_update_token_budget_stats_accumulates_within_same_bucket() -> None:
|
|
reset_token_budget_state_for_tests()
|
|
|
|
update_token_budget_stats(120, when=TOKEN_BUCKET_WINDOW_SECONDS * 2)
|
|
update_token_budget_stats(30, when=TOKEN_BUCKET_WINDOW_SECONDS * 2 + 5)
|
|
|
|
assert get_token_budget_settings().five_minute_buckets[-1] == TokenBudgetBucket(
|
|
started_at=TOKEN_BUCKET_WINDOW_SECONDS * 2,
|
|
tokens_used=150,
|
|
)
|
|
reset_token_budget_state_for_tests()
|
|
|
|
|
|
def test_token_budget_stats_json_round_trips_bucket_history() -> None:
|
|
reset_token_budget_state_for_tests()
|
|
configure_token_budget_settings(
|
|
last_ai_token_used=222,
|
|
total_tokens=999,
|
|
five_minute_buckets=(
|
|
TokenBudgetBucket(started_at=300, tokens_used=100),
|
|
TokenBudgetBucket(started_at=600, tokens_used=200),
|
|
TokenBudgetBucket(started_at=900, tokens_used=300),
|
|
TokenBudgetBucket(started_at=1200, tokens_used=400),
|
|
TokenBudgetBucket(started_at=1500, tokens_used=500),
|
|
TokenBudgetBucket(started_at=1800, tokens_used=600),
|
|
),
|
|
)
|
|
|
|
repo_exec = MemoryRepoExec()
|
|
checkpoint_dir = repo_exec.root / "checkpoints"
|
|
persist_token_budget_stats_json(checkpoint_dir, repo_exec=repo_exec)
|
|
reset_token_budget_state_for_tests()
|
|
restore_token_budget_stats_json(checkpoint_dir, repo_exec=repo_exec)
|
|
|
|
settings = get_token_budget_settings()
|
|
assert settings.last_ai_token_used == 222
|
|
assert settings.total_tokens == 999
|
|
assert settings.five_minute_buckets == (
|
|
TokenBudgetBucket(started_at=300, tokens_used=100),
|
|
TokenBudgetBucket(started_at=600, tokens_used=200),
|
|
TokenBudgetBucket(started_at=900, tokens_used=300),
|
|
TokenBudgetBucket(started_at=1200, tokens_used=400),
|
|
TokenBudgetBucket(started_at=1500, tokens_used=500),
|
|
TokenBudgetBucket(started_at=1800, tokens_used=600),
|
|
)
|
|
reset_token_budget_state_for_tests()
|