112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
from band_aider.ai_client.json_guarding_client import JsonGuardingClient
|
|
from band_aider.judgement.prompt_formatter import MarkdownPromptFormatterBuilder
|
|
|
|
|
|
class _Client(CompletionClient):
|
|
"""_Client returns configured responses in sequence."""
|
|
|
|
def __init__(self, responses: list[str]) -> None:
|
|
self.responses = list(responses)
|
|
self.prompts: list[str] = []
|
|
self._openai_configs: tuple[OpenAIConfig, ...] = ()
|
|
|
|
def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
|
|
"""Return the next stationary local 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, ...]:
|
|
return self._openai_configs
|
|
|
|
@property
|
|
def inner(self) -> CompletionClient | None:
|
|
return None
|
|
|
|
|
|
def test_json_guarding_client_repairs_malformed_json_for_json_prompt() -> None:
|
|
"""Verify malformed JSON triggers one repair retry with the parse error fed back."""
|
|
|
|
inner = _Client(["{bad json", '{"ok": true}'])
|
|
client = JsonGuardingClient(inner=inner, prompt_formatter_builder=MarkdownPromptFormatterBuilder())
|
|
|
|
result = client.complete('Return strict JSON only: {"ok": true}')
|
|
|
|
assert result == '{"ok": true}'
|
|
assert len(inner.prompts) == 2
|
|
assert "malformed json" in inner.prompts[1].lower()
|
|
assert "json parse error" in inner.prompts[1].lower()
|
|
|
|
|
|
def test_json_guarding_client_leaves_non_json_prompt_unchanged() -> None:
|
|
"""Verify non-JSON prompts do not trigger repair behavior."""
|
|
|
|
inner = _Client(["plain text"])
|
|
client = JsonGuardingClient(inner=inner, prompt_formatter_builder=MarkdownPromptFormatterBuilder())
|
|
|
|
assert client.complete("Reply exactly: hello") == "plain text"
|
|
assert inner.prompts == ["Reply exactly: hello"]
|
|
|
|
|
|
def test_json_guarding_client_marks_exhausted_malformed_repair_invalid() -> None:
|
|
"""Verify an exhausted local malformed-repair sequence is explicitly invalid."""
|
|
|
|
inner = _Client(["{bad json", "{still bad json"])
|
|
client = JsonGuardingClient(
|
|
inner=inner,
|
|
prompt_formatter_builder=MarkdownPromptFormatterBuilder(),
|
|
max_repair_attempts=1,
|
|
)
|
|
|
|
response, valid = client._complete_with_repairs(inner, 'Return strict JSON only: {"ok": true}')
|
|
|
|
assert response == "{still bad json"
|
|
assert valid is False
|
|
assert len(inner.prompts) == 3
|
|
|
|
|
|
def test_json_guarding_client_marks_exhausted_null_repair_invalid() -> None:
|
|
"""Verify an exhausted local null-repair sequence is explicitly invalid."""
|
|
|
|
inner = _Client(["null", "null"])
|
|
client = JsonGuardingClient(
|
|
inner=inner,
|
|
prompt_formatter_builder=MarkdownPromptFormatterBuilder(),
|
|
max_repair_attempts=1,
|
|
)
|
|
|
|
response, valid = client._complete_with_repairs(inner, 'Return strict JSON only: {"ok": true}')
|
|
|
|
assert response == "null"
|
|
assert valid is False
|
|
assert len(inner.prompts) == 3
|
|
|
|
|
|
def test_json_guarding_client_preserves_multiple_escalation_config_order() -> None:
|
|
"""Verify the wrapper retains its configured escalation order for an owner to execute."""
|
|
|
|
inner = _Client(["{bad json", "{still bad json"])
|
|
inner._openai_configs = (
|
|
OpenAIConfig(model="local", api_base="http://local"),
|
|
OpenAIConfig(model="cloud-1", api_base="http://cloud-1"),
|
|
OpenAIConfig(model="cloud-2", api_base="http://cloud-2"),
|
|
)
|
|
client = JsonGuardingClient(
|
|
inner=inner,
|
|
prompt_formatter_builder=MarkdownPromptFormatterBuilder(),
|
|
max_repair_attempts=0,
|
|
)
|
|
|
|
assert [config.model for config in client.openai_configs] == ["local", "cloud-1", "cloud-2"]
|