210 lines
9.2 KiB
Python
210 lines
9.2 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 hashlib
|
|
import logging
|
|
import time
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
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
|
|
from band_aider.ui.dto.model_escalation import ModelEscalationDTO
|
|
from band_aider.ui.status_publisher import NullStatusPublisher, StatusPublisher
|
|
|
|
|
|
@dataclass(init=False)
|
|
class TokenJacker(CompletionClient):
|
|
"""TokenJacker routes oversized prompts to a designated overflow client."""
|
|
|
|
settings: TokenBudgetSettings
|
|
event_publisher: StatusPublisher
|
|
estimator: TokenBudgetEstimator | None = None
|
|
_inner: CompletionClient = field(init=False, repr=False)
|
|
|
|
def __init__(
|
|
self,
|
|
inner: CompletionClient,
|
|
settings: TokenBudgetSettings,
|
|
event_publisher: StatusPublisher | None = None,
|
|
estimator: TokenBudgetEstimator | None = None,
|
|
escalation_clients: tuple[CompletionClient, ...] = (),
|
|
) -> None:
|
|
"""Initialize the wrapper with a primary completion client."""
|
|
|
|
self._inner = inner
|
|
self.settings = settings
|
|
self.event_publisher = event_publisher or NullStatusPublisher()
|
|
self.estimator = estimator
|
|
self.escalation_clients = escalation_clients
|
|
self.last_config_index = 0
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient:
|
|
"""Return the primary wrapped client."""
|
|
|
|
return self._inner
|
|
|
|
@property
|
|
def openai_configs(self) -> tuple[OpenAIConfig, ...]:
|
|
"""Return the delegated OpenAI config chain from the primary client."""
|
|
|
|
return self.inner.openai_configs
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Complete one prompt, routing oversized input to the overflow client."""
|
|
|
|
# A retry owner must count the provider that actually answered an oversized prompt.
|
|
self.last_config_index = 0
|
|
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("TokenJacker kept prompt on primary client: 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; routing to overflow client.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenJacker chose overflow branch: estimated_tokens=%s exceeds budget=%s.", estimated_tokens, self.settings.max_prompt_tokens)
|
|
overflows = self.escalation_clients
|
|
if not overflows:
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message="No overflow client is configured; forwarding oversized prompt to the primary client.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenJacker found no overflow clients; forwarding to primary client.")
|
|
response = self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
self._publish_escalation(
|
|
prompt=prompt,
|
|
response=response,
|
|
estimated_tokens=estimated_tokens,
|
|
local_model=self._model_name(0),
|
|
active_model=self._model_name(0),
|
|
config_index=0,
|
|
state="fallback_primary",
|
|
reason=f"Estimated tokens {estimated_tokens} exceeded budget {self.settings.max_prompt_tokens}; no overflow client is configured.",
|
|
)
|
|
return response
|
|
|
|
last_error: Exception | None = None
|
|
for index, overflow in enumerate(overflows, start=1):
|
|
try:
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message=f"Routing oversized prompt to overflow config index {index}.",
|
|
)
|
|
if logger is not None:
|
|
logger.info("TokenJacker routed prompt to overflow config index %s.", index)
|
|
response = overflow.complete(prompt) if logger is None else overflow.complete(prompt, logger=logger)
|
|
self.last_config_index = index
|
|
self._publish_escalation(
|
|
prompt=prompt,
|
|
response=response,
|
|
estimated_tokens=estimated_tokens,
|
|
local_model=self._model_name(0),
|
|
active_model=self._model_name(index),
|
|
config_index=index,
|
|
state="routed",
|
|
reason=f"Estimated tokens {estimated_tokens} exceeded budget {self.settings.max_prompt_tokens}.",
|
|
)
|
|
return response
|
|
except Exception as error:
|
|
last_error = error
|
|
log_llm_context_warning(
|
|
source=type(self).__name__,
|
|
estimated_tokens=estimated_tokens,
|
|
max_prompt_tokens=self.settings.max_prompt_tokens,
|
|
message=f"Overflow config index {index} failed; escalating to the next overflow tier.",
|
|
)
|
|
if logger is not None:
|
|
logger.warning("TokenJacker overflow config index %s failed; trying next overflow tier.", index)
|
|
self._publish_escalation(
|
|
prompt=prompt,
|
|
response="",
|
|
estimated_tokens=estimated_tokens,
|
|
local_model=self._model_name(0),
|
|
active_model=self._model_name(index),
|
|
config_index=index,
|
|
state="overflow_failed",
|
|
reason=f"Estimated tokens {estimated_tokens} exceeded budget {self.settings.max_prompt_tokens}; overflow tier {index} failed.",
|
|
)
|
|
|
|
if last_error is not None:
|
|
raise last_error
|
|
return self.inner.complete(prompt) if logger is None else self.inner.complete(prompt, logger=logger)
|
|
|
|
def _publish_escalation(
|
|
self,
|
|
prompt: str,
|
|
response: str,
|
|
estimated_tokens: int,
|
|
local_model: str,
|
|
active_model: str,
|
|
config_index: int,
|
|
state: str,
|
|
reason: str,
|
|
) -> None:
|
|
request_md5 = hashlib.md5(prompt.encode("utf-8")).hexdigest()
|
|
response_md5 = hashlib.md5(response.encode("utf-8")).hexdigest() if response else ""
|
|
# TODO: Deliberate with a human whether to remove this arbitrary response-history cap.
|
|
response_tail = response[-1000:] if response else ""
|
|
endpoint = self._endpoint(config_index)
|
|
tokens_escalated = max(0, int(estimated_tokens)) if config_index > 0 else 0
|
|
self.event_publisher.publish(
|
|
ModelEscalationDTO(
|
|
trigger_name=type(self).__name__,
|
|
occurred_at=time.time(),
|
|
prompt_preview=_prompt_preview(prompt),
|
|
reason=reason,
|
|
request_md5=request_md5,
|
|
local_model=local_model,
|
|
active_model=active_model,
|
|
config_index=config_index,
|
|
config_count=max(1, len(self.openai_configs)),
|
|
state=state,
|
|
response_md5=response_md5,
|
|
response_tail=response_tail,
|
|
repeat_count=0,
|
|
tokens_escalated=tokens_escalated,
|
|
active_endpoint=endpoint,
|
|
)
|
|
)
|
|
|
|
def _model_name(self, index: int) -> str:
|
|
configs = self.openai_configs
|
|
if configs and 0 <= index < len(configs):
|
|
return configs[index].model
|
|
return str(getattr(self.inner, "model", "unknown"))
|
|
|
|
def _endpoint(self, index: int) -> str:
|
|
configs = self.openai_configs
|
|
if configs and 0 <= index < len(configs):
|
|
return configs[index].api_base
|
|
return str(getattr(self.inner, "api_base", "unknown"))
|
|
|
|
|
|
def _prompt_preview(prompt: str, max_lines: int = 3, max_width: int = 88) -> str:
|
|
"""Return a bounded preview of the first prompt lines for UI display."""
|
|
|
|
# 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()
|