104 lines
4 KiB
Python
104 lines
4 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from band_aider.ai_client.llm_context_batch import LLMContextBatch
|
|
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.token_budget_estimator import TokenBudgetEstimator
|
|
|
|
|
|
_REDUCTION_GUIDANCE = """Reduce this oversized prompt so a later pass can continue.
|
|
|
|
Preserve structure the best you can. Your priorities should be as follow:
|
|
- Infer the intent of the prompt you're summarizing and ensure the goal survives compression.
|
|
- Infer schema and preserve structure as much as possible whether it's json, yaml, kvp, array, stack trace, etc.
|
|
- Keep stack trace, especially the innermost stack trace and first throw location that is inside the src/project root.
|
|
- Keep merge JSON, arrays, headings, sections, but you may drop some free-formed text members of an array or summarize text fields.
|
|
- Keep keys, labels, file paths, test names, class names, member names, and output schema when visible.
|
|
- Replace only oversized parts with concise reduced equivalents and clear "..." markers.
|
|
- Return only the reduced prompt text.
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GeneralPromptReducer(LLMContextReducer):
|
|
"""Reduce oversized prompts generically with best-effort structure preservation."""
|
|
|
|
marker: str | None = None
|
|
|
|
def _content(self, prompt: str) -> str:
|
|
"""Return the prompt content to reduce."""
|
|
|
|
if self.marker and self.marker in prompt:
|
|
return prompt.split(self.marker, 1)[1]
|
|
return prompt
|
|
|
|
def can_reduce(self, prompt: str) -> bool:
|
|
"""Return whether the prompt has any reducible content."""
|
|
|
|
return bool(self._content(prompt).strip())
|
|
|
|
def split(self, prompt: str, token_budget: int, estimator: TokenBudgetEstimator) -> list[LLMContextBatch]:
|
|
"""Split a prompt into line-preserving batches under the token budget."""
|
|
|
|
lines = self._content(prompt).splitlines(keepends=True) or [self._content(prompt)]
|
|
overhead = estimator.estimate(self.prompt_for_batch(prompt, LLMContextBatch(0, 2, ""), ""))
|
|
content_budget = max(1, token_budget - overhead)
|
|
|
|
chunks: list[str] = []
|
|
current = ""
|
|
current_tokens = 0
|
|
|
|
for line in lines:
|
|
line_tokens = estimator.estimate(line)
|
|
if current and current_tokens + line_tokens > content_budget:
|
|
chunks.append(current)
|
|
current = ""
|
|
current_tokens = 0
|
|
current += line
|
|
current_tokens += line_tokens
|
|
|
|
if current or not chunks:
|
|
chunks.append(current)
|
|
|
|
return [
|
|
LLMContextBatch(batch_index=index, total_batches=len(chunks), content=chunk)
|
|
for index, chunk in enumerate(chunks)
|
|
]
|
|
|
|
def prompt_for_batch(
|
|
self,
|
|
original_prompt: str,
|
|
batch: LLMContextBatch,
|
|
previous_result: str,
|
|
) -> str:
|
|
"""Return one reduction prompt for the current batch."""
|
|
|
|
metadata = [
|
|
"This original prompt was reduced and split because it exceeded the configured token budget.",
|
|
f"Batch: {batch.batch_index + 1}/{batch.total_batches}",
|
|
_REDUCTION_GUIDANCE.strip(),
|
|
]
|
|
if batch.batch_index > 0:
|
|
metadata.extend(
|
|
[
|
|
"Previous reduced prompt:",
|
|
previous_result,
|
|
"Merge the previous reduced prompt with the current batch into one reduced prompt.",
|
|
]
|
|
)
|
|
|
|
metadata.extend(["Current batch:", batch.content, "Reduced prompt:"])
|
|
return "\n\n".join(metadata)
|
|
|
|
def final_result(self, reductions: list[LLMContextReduction]) -> str:
|
|
"""Return the final reduced prompt string."""
|
|
|
|
if not reductions:
|
|
return ""
|
|
return reductions[-1].current_result
|