260 lines
7.5 KiB
Python
260 lines
7.5 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from app.utils.llm_client import LLMClient, LLMResponseError
|
|
|
|
|
|
class CompletionSequence:
|
|
def __init__(self, *results):
|
|
self.results = list(results)
|
|
self.calls = []
|
|
|
|
def create(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
result = self.results.pop(0)
|
|
if isinstance(result, Exception):
|
|
raise result
|
|
return result
|
|
|
|
|
|
class ProviderError(RuntimeError):
|
|
def __init__(self, *, status_code, body):
|
|
super().__init__(body.get("error", {}).get("message", "provider error"))
|
|
self.status_code = status_code
|
|
self.body = body
|
|
|
|
|
|
def _response(content, *, finish_reason="stop", include_choice=True):
|
|
choices = []
|
|
if include_choice:
|
|
choices.append(
|
|
SimpleNamespace(
|
|
finish_reason=finish_reason,
|
|
message=SimpleNamespace(content=content),
|
|
)
|
|
)
|
|
return SimpleNamespace(choices=choices)
|
|
|
|
|
|
def _client_for(sequence):
|
|
client = object.__new__(LLMClient)
|
|
client.model = "compatible-model"
|
|
client.client = SimpleNamespace(
|
|
chat=SimpleNamespace(completions=sequence)
|
|
)
|
|
return client
|
|
|
|
|
|
def test_chat_json_retries_truncated_completion_without_token_cap():
|
|
sequence = CompletionSequence(
|
|
_response('{"items": [', finish_reason="length"),
|
|
_response('{"items": [1, 2]}'),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
result = client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_tokens=4096,
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert result == {"items": [1, 2]}
|
|
assert sequence.calls[0]["max_tokens"] == 4096
|
|
assert "max_tokens" not in sequence.calls[1]
|
|
|
|
|
|
def test_chat_json_omits_token_cap_when_requested():
|
|
sequence = CompletionSequence(_response('{"ok": true}'))
|
|
client = _client_for(sequence)
|
|
|
|
assert client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_tokens=None,
|
|
) == {"ok": True}
|
|
assert "max_tokens" not in sequence.calls[0]
|
|
|
|
|
|
def test_chat_json_retries_empty_content_once():
|
|
sequence = CompletionSequence(
|
|
_response(None),
|
|
_response('{"ok": true}'),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
result = client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert result == {"ok": True}
|
|
assert len(sequence.calls) == 2
|
|
|
|
|
|
def test_chat_json_accepts_complete_object_before_trailing_text():
|
|
sequence = CompletionSequence(
|
|
_response('{"ok": true}\nThis object is ready to use.')
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
assert client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
) == {"ok": True}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"content",
|
|
[
|
|
'{"status": "draft"}\n{"status": "complete"}',
|
|
'{"status": "draft"}\n```json\n{"status": "complete"}\n```',
|
|
'{"status": "draft"}\nExplanation first.\n{"status": "complete"}',
|
|
],
|
|
)
|
|
def test_chat_json_rejects_a_second_json_document(content):
|
|
sequence = CompletionSequence(_response(content))
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(LLMResponseError, match="multiple JSON"):
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
)
|
|
|
|
|
|
def test_chat_json_rejects_top_level_array():
|
|
sequence = CompletionSequence(_response('[{"ok": true}]'))
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(LLMResponseError, match="JSON object"):
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=1,
|
|
)
|
|
|
|
|
|
def test_chat_json_error_does_not_echo_partial_model_output():
|
|
partial = '{"private_source_text": "SENTINEL-SHOULD-NOT-LEAK"'
|
|
sequence = CompletionSequence(
|
|
_response(partial, finish_reason="length"),
|
|
_response(partial, finish_reason="length"),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(LLMResponseError) as captured:
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert "SENTINEL-SHOULD-NOT-LEAK" not in str(captured.value)
|
|
assert captured.value.finish_reason == "length"
|
|
|
|
|
|
def test_chat_json_falls_back_only_for_explicit_response_format_rejection():
|
|
unsupported = ProviderError(
|
|
status_code=400,
|
|
body={
|
|
"error": {
|
|
"param": "response_format",
|
|
"code": "unsupported_parameter",
|
|
"message": "response_format is not supported by this model",
|
|
}
|
|
},
|
|
)
|
|
sequence = CompletionSequence(unsupported, _response('{"ok": true}'))
|
|
client = _client_for(sequence)
|
|
|
|
result = client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=1,
|
|
)
|
|
|
|
assert result == {"ok": True}
|
|
assert sequence.calls[0]["response_format"] == {"type": "json_object"}
|
|
assert "response_format" not in sequence.calls[1]
|
|
|
|
|
|
def test_response_format_fallback_keeps_content_retry_available():
|
|
unsupported = ProviderError(
|
|
status_code=400,
|
|
body={
|
|
"error": {
|
|
"param": "response_format",
|
|
"code": "unsupported_parameter",
|
|
"message": "response_format is not supported by this model",
|
|
}
|
|
},
|
|
)
|
|
sequence = CompletionSequence(
|
|
unsupported,
|
|
_response('{"items": [', finish_reason="length"),
|
|
_response('{"items": [1]}'),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
result = client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert result == {"items": [1]}
|
|
assert sequence.calls[0]["response_format"] == {"type": "json_object"}
|
|
assert "response_format" not in sequence.calls[1]
|
|
assert "response_format" not in sequence.calls[2]
|
|
assert "max_tokens" not in sequence.calls[2]
|
|
|
|
|
|
def test_chat_json_defaults_to_one_content_attempt():
|
|
sequence = CompletionSequence(
|
|
_response(None),
|
|
_response('{"ok": true}'),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(LLMResponseError, match="empty JSON"):
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
)
|
|
|
|
assert len(sequence.calls) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("status_code", [400, 401, 429, 500])
|
|
def test_chat_json_does_not_retry_unrelated_provider_errors(status_code):
|
|
provider_error = ProviderError(
|
|
status_code=status_code,
|
|
body={
|
|
"error": {
|
|
"param": "messages",
|
|
"code": "invalid_request",
|
|
"message": "request failed for an unrelated reason",
|
|
}
|
|
},
|
|
)
|
|
sequence = CompletionSequence(provider_error, _response('{"ok": true}'))
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(ProviderError) as captured:
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert captured.value is provider_error
|
|
assert len(sequence.calls) == 1
|
|
|
|
|
|
def test_chat_json_reports_missing_choices_without_retrying_forever():
|
|
sequence = CompletionSequence(
|
|
_response(None, include_choice=False),
|
|
_response(None, include_choice=False),
|
|
)
|
|
client = _client_for(sequence)
|
|
|
|
with pytest.raises(LLMResponseError, match="no choices"):
|
|
client.chat_json(
|
|
messages=[{"role": "user", "content": "Return JSON"}],
|
|
max_attempts=2,
|
|
)
|
|
|
|
assert len(sequence.calls) == 2
|