76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from band_aider.ai_client.llm_context_reducer import LLMContextReducer
|
|
from band_aider.ai_client.llm_context_batch import LLMContextBatch
|
|
from band_aider.ai_client.json_array_prompt_reducer import GeneralPromptReducer
|
|
from band_aider.ai_client.llm_context_reduction import LLMContextReduction
|
|
from band_aider.ai_client.token_budget_estimator import TokenBudgetEstimator
|
|
|
|
class _FixtureLLMContextReducer(LLMContextReducer):
|
|
"""A concrete reducer fixture that implements the complete public protocol."""
|
|
|
|
def can_reduce(self, prompt: str) -> bool:
|
|
return len(prompt) > 10
|
|
|
|
def split(self, prompt: str, token_budget: int, estimator: TokenBudgetEstimator) -> list[LLMContextBatch]:
|
|
del prompt, token_budget, estimator
|
|
return [LLMContextBatch(batch_index=0, total_batches=1, content="fixture")]
|
|
|
|
def prompt_for_batch(
|
|
self,
|
|
original_prompt: str,
|
|
batch: LLMContextBatch,
|
|
previous_result: str,
|
|
) -> str:
|
|
return f"batch prompt for {previous_result}"
|
|
|
|
def final_result(self, reductions: list[LLMContextReduction]) -> str:
|
|
return "final result"
|
|
|
|
@pytest.fixture
|
|
def reducer() -> _FixtureLLMContextReducer:
|
|
"""Provide one explicit reducer implementation."""
|
|
|
|
return _FixtureLLMContextReducer()
|
|
|
|
|
|
def test_llm_context_reducer_can_reduce(reducer: _FixtureLLMContextReducer) -> None:
|
|
"""Verify can_reduce correctly identifies prompts that need reduction."""
|
|
assert reducer.can_reduce("short") is False
|
|
assert reducer.can_reduce("this is a long prompt") is True
|
|
|
|
def test_llm_context_reducer_split(reducer: _FixtureLLMContextReducer) -> None:
|
|
"""Verify split returns a list of batches."""
|
|
estimator = TokenBudgetEstimator()
|
|
batches = reducer.split("some prompt", 100, estimator)
|
|
assert isinstance(batches, list)
|
|
assert len(batches) == 1
|
|
|
|
def test_llm_context_reducer_prompt_for_batch(reducer: _FixtureLLMContextReducer) -> None:
|
|
"""Verify prompt_for_batch generates the expected prompt string."""
|
|
batch = LLMContextBatch(batch_index=0, total_batches=1, content="fixture")
|
|
prompt = reducer.prompt_for_batch("orig", batch, "prev")
|
|
assert prompt == "batch prompt for prev"
|
|
|
|
def test_llm_context_reducer_final_result(reducer: _FixtureLLMContextReducer) -> None:
|
|
"""Verify final_result returns the expected final string."""
|
|
reductions = [
|
|
LLMContextReduction(
|
|
previous_result="",
|
|
current_batch=LLMContextBatch(0, 1, "fixture"),
|
|
current_result="reduced",
|
|
),
|
|
]
|
|
result = reducer.final_result(reductions)
|
|
assert result == "final result"
|
|
|
|
|
|
def test_general_prompt_reducer_explicitly_implements_protocol() -> None:
|
|
reducer = GeneralPromptReducer()
|
|
|
|
assert isinstance(reducer, LLMContextReducer)
|