137 lines
6.6 KiB
Python
137 lines
6.6 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
import logging
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
from band_aider.ai_client.llm_context_reducer import LLMContextReducer
|
|
from band_aider.ai_client.llm_context_reduction import LLMContextReduction
|
|
from band_aider.ai_client.llm_logging import log_llm_context_warning
|
|
from band_aider.ai_client.token_budget_estimator import TokenBudgetEstimator
|
|
from band_aider.context.token_budget_state import TokenBudgetSettings
|
|
|
|
|
|
@dataclass(init=False)
|
|
class TokenMiser(CompletionClient):
|
|
"""TokenMiser decorates a completion client with prompt budget reduction."""
|
|
|
|
settings: TokenBudgetSettings
|
|
reducers: list[LLMContextReducer] = field(default_factory=list)
|
|
estimator: TokenBudgetEstimator | None = None
|
|
_inner: CompletionClient = field(init=False, repr=False)
|
|
|
|
def __init__(
|
|
self,
|
|
inner: CompletionClient,
|
|
settings: TokenBudgetSettings,
|
|
reducers: list[LLMContextReducer] | None = None,
|
|
estimator: TokenBudgetEstimator | None = None,
|
|
) -> None:
|
|
"""Initialize the wrapper with one inner completion client and reduction settings."""
|
|
|
|
self._inner = inner
|
|
self.settings = settings
|
|
self.reducers = list(reducers or [])
|
|
self.estimator = estimator
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient:
|
|
"""Return the wrapped client."""
|
|
|
|
return self._inner
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
"""Return the delegated OpenAI config chain."""
|
|
|
|
return self.inner.openai_configs
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Complete one prompt, reducing oversized context when a reducer supports it."""
|
|
|
|
estimator = self.estimator or TokenBudgetEstimator(tokens_per_word=self.settings.tokens_per_word)
|
|
estimated_tokens = estimator.estimate(prompt)
|
|
if estimated_tokens <= self.settings.max_prompt_tokens:
|
|
if logger is not None:
|
|
logger.info("TokenMiser forwarded prompt unchanged: estimated_tokens=%s within budget=%s.", estimated_tokens, self.settings.max_prompt_tokens)
|
|
return self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message="Prompt exceeds configured token budget; looking for a reducer.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenMiser chose reducer search branch: estimated_tokens=%s exceeds budget=%s.", estimated_tokens, self.settings.max_prompt_tokens)
|
|
|
|
reducer = next((candidate for candidate in self.reducers if candidate.can_reduce(prompt)), None)
|
|
if reducer is None:
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message="No reducer matched oversized prompt; forwarding unchanged.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenMiser found no reducer; forwarding original prompt unchanged.")
|
|
return self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
|
|
# Reserve output capacity from the overall prompt budget, then target a fraction of the
|
|
# remaining input budget for each reduction batch.
|
|
available_prompt_budget = max(1, self.settings.max_prompt_tokens - self.settings.max_output_tokens)
|
|
batch_budget = max(1, int(available_prompt_budget * self.settings.batch_target_ratio))
|
|
batches = reducer.split(prompt, batch_budget, estimator)
|
|
if not batches:
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message=f"{type(reducer).__name__} produced no batches; forwarding unchanged.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenMiser reducer %s produced no batches; forwarding original prompt.", type(reducer).__name__)
|
|
return self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message=f"Reducing oversized prompt with {type(reducer).__name__} in {len(batches)} batches.",
|
|
)
|
|
|
|
reductions: list[LLMContextReduction] = []
|
|
previous_result = ""
|
|
for batch_number, batch in enumerate(batches, start=1):
|
|
batch_prompt = reducer.prompt_for_batch(prompt, batch, previous_result)
|
|
if logger is not None:
|
|
logger.info("TokenMiser reducing batch %s/%s with %s.", batch_number, len(batches), type(reducer).__name__)
|
|
current_result = self.inner.complete(batch_prompt) if logger is None else self.inner.complete(batch_prompt, logger=logger)
|
|
reductions.append(
|
|
LLMContextReduction(
|
|
previous_result=previous_result,
|
|
current_batch=batch,
|
|
current_result=current_result,
|
|
)
|
|
)
|
|
previous_result = current_result
|
|
|
|
reduced_prompt = reducer.final_result(reductions).strip()
|
|
if not reduced_prompt:
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message=f"{type(reducer).__name__} produced an empty reduced prompt; forwarding unchanged.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenMiser reducer %s produced empty final prompt; forwarding original prompt.", type(reducer).__name__)
|
|
return self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
|
|
if logger is not None:
|
|
logger.info("TokenMiser forwarding reduced prompt produced by %s.", type(reducer).__name__)
|
|
return self.inner.complete(reduced_prompt) if logger is None else self.inner.complete(reduced_prompt, logger=logger)
|