242 lines
7.7 KiB
Python
242 lines
7.7 KiB
Python
"""Tests for the LLM gateway: complete(), classified recovery, and fallback."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.messages import AIMessage, AIMessageChunk
|
|
import pytest
|
|
|
|
from artemis.llm.reliability import LLMExhaustedError, LLMPermanentError
|
|
from artemis.services import llm as llm_service
|
|
from artemis.services.llm import (
|
|
RobustChatModelWrapper,
|
|
acomplete,
|
|
acomplete_structured,
|
|
with_fallback,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _isolate_gateway_state(monkeypatch, tmp_path):
|
|
llm_service._ENDPOINT_BREAKER._states.clear()
|
|
llm_service._NON_STREAMING_ENDPOINTS.clear()
|
|
monkeypatch.setattr(llm_service, "PAUSE_FILE", tmp_path / ".artemis_paused")
|
|
monkeypatch.setattr(
|
|
llm_service,
|
|
"retry_policy_for",
|
|
lambda category: SimpleNamespace(max_attempts=2, delay_for=lambda attempt: 0.0),
|
|
)
|
|
monkeypatch.setattr(llm_service.settings, "LLM_PAUSE_TIMEOUT_SECONDS", 1.0, raising=False)
|
|
yield
|
|
llm_service._ENDPOINT_BREAKER._states.clear()
|
|
llm_service._NON_STREAMING_ENDPOINTS.clear()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mid_stream_failure_retries_without_duplicating_chunks():
|
|
class BreaksMidStream:
|
|
def __init__(self):
|
|
self.astream_calls = 0
|
|
|
|
async def astream(self, *args, **kwargs):
|
|
self.astream_calls += 1
|
|
if self.astream_calls == 1:
|
|
yield AIMessageChunk(content="Hello")
|
|
raise RuntimeError("503 service unavailable")
|
|
yield AIMessageChunk(content="Hello")
|
|
yield AIMessageChunk(content=" world")
|
|
|
|
base = BreaksMidStream()
|
|
wrapper = RobustChatModelWrapper(base)
|
|
result = await wrapper.complete([])
|
|
|
|
assert base.astream_calls == 2
|
|
# The partial first attempt was discarded entirely: no duplicated prefix.
|
|
assert result.content == "Hello world"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auth_error_fails_fast_without_retry_or_pause(tmp_path):
|
|
class AuthFails:
|
|
def __init__(self):
|
|
self.astream_calls = 0
|
|
self.ainvoke_calls = 0
|
|
|
|
async def astream(self, *args, **kwargs):
|
|
self.astream_calls += 1
|
|
raise RuntimeError("401 unauthorized: invalid api key")
|
|
yield # pragma: no cover - makes this an async generator
|
|
|
|
async def ainvoke(self, *args, **kwargs):
|
|
self.ainvoke_calls += 1
|
|
return AIMessage(content="should never happen")
|
|
|
|
base = AuthFails()
|
|
wrapper = RobustChatModelWrapper(base)
|
|
with pytest.raises(LLMPermanentError):
|
|
await wrapper.complete([])
|
|
|
|
assert base.astream_calls == 1
|
|
assert base.ainvoke_calls == 0
|
|
assert not llm_service.PAUSE_FILE.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_unsupported_downgrades_loudly_and_is_remembered():
|
|
class NoStreaming:
|
|
def __init__(self):
|
|
self.astream_calls = 0
|
|
self.ainvoke_calls = 0
|
|
|
|
async def astream(self, *args, **kwargs):
|
|
self.astream_calls += 1
|
|
raise NotImplementedError("streaming is not supported")
|
|
yield # pragma: no cover
|
|
|
|
async def ainvoke(self, *args, **kwargs):
|
|
self.ainvoke_calls += 1
|
|
return AIMessage(content="ok")
|
|
|
|
base = NoStreaming()
|
|
wrapper = RobustChatModelWrapper(base)
|
|
|
|
first = await wrapper.complete([])
|
|
assert first.content == "ok"
|
|
assert wrapper._endpoint_key() in llm_service._NON_STREAMING_ENDPOINTS
|
|
|
|
second = await wrapper.complete([])
|
|
assert second.content == "ok"
|
|
# The capability was cached: streaming is not probed again.
|
|
assert base.astream_calls == 1
|
|
assert base.ainvoke_calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exhausted_retries_pause_then_raise_on_deadline():
|
|
class AlwaysUnavailable:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
async def astream(self, *args, **kwargs):
|
|
self.calls += 1
|
|
raise RuntimeError("503 service unavailable")
|
|
yield # pragma: no cover
|
|
|
|
wrapper = RobustChatModelWrapper(AlwaysUnavailable())
|
|
with pytest.raises(LLMExhaustedError):
|
|
await wrapper.complete([])
|
|
# The pause file was written for the pause window, then cleaned up.
|
|
assert not llm_service.PAUSE_FILE.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_with_fallback_takes_over_without_pausing():
|
|
class AlwaysUnavailable:
|
|
async def astream(self, *args, **kwargs):
|
|
raise RuntimeError("503 service unavailable")
|
|
yield # pragma: no cover
|
|
|
|
wrapper = RobustChatModelWrapper(AlwaysUnavailable())
|
|
|
|
async def fallback_call():
|
|
return AIMessage(content="fallback answer")
|
|
|
|
result = await with_fallback(
|
|
main_call=lambda: wrapper.complete([]),
|
|
fallback_call=fallback_call,
|
|
)
|
|
assert result.content == "fallback answer"
|
|
# Because a fallback existed, exhaustion handed over instead of pausing.
|
|
assert not llm_service.PAUSE_FILE.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_with_fallback_refuses_to_hide_bad_requests():
|
|
class CodedError(Exception):
|
|
code = 400
|
|
|
|
calls = {"fallback": 0}
|
|
|
|
async def main_call():
|
|
raise CodedError("malformed tool schema")
|
|
|
|
async def fallback_call():
|
|
calls["fallback"] += 1
|
|
return "hidden bug"
|
|
|
|
with pytest.raises(CodedError):
|
|
await with_fallback(main_call=main_call, fallback_call=fallback_call)
|
|
assert calls["fallback"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acomplete_structured_corrective_reask():
|
|
class CorrectsItself:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def ainvoke(self, messages, **kwargs):
|
|
self.calls.append(messages)
|
|
if len(self.calls) == 1:
|
|
return AIMessage(content="oops, here is prose instead of JSON")
|
|
return AIMessage(content='{"a": 1}')
|
|
|
|
base = CorrectsItself()
|
|
result = await acomplete_structured(base, [])
|
|
assert result == {"a": 1}
|
|
assert len(base.calls) == 2
|
|
correction_text = str(base.calls[1][-1].content)
|
|
assert "could not be parsed" in correction_text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acomplete_compat_path_for_plain_models():
|
|
class PlainModel:
|
|
async def astream(self, *args, **kwargs):
|
|
yield AIMessageChunk(content="a")
|
|
yield AIMessageChunk(content="b")
|
|
|
|
result = await acomplete(PlainModel(), [])
|
|
assert result.content == "ab"
|
|
|
|
class InvokeOnly:
|
|
async def ainvoke(self, *args, **kwargs):
|
|
return AIMessage(content="direct")
|
|
|
|
result = await acomplete(InvokeOnly(), [])
|
|
assert result.content == "direct"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mid_stream_failure_records_stream_reset_payload(monkeypatch):
|
|
recorded_events = []
|
|
|
|
def mock_record_event(name, payload, *, status="retrying"):
|
|
recorded_events.append((name, payload, status))
|
|
|
|
monkeypatch.setattr(llm_service, "_record_llm_event", mock_record_event)
|
|
|
|
class BreaksMidStream:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
async def astream(self, *args, **kwargs):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
yield AIMessageChunk(content="Partial thoughts...")
|
|
raise RuntimeError("503 high demand spike")
|
|
yield AIMessageChunk(content="Recovered output")
|
|
|
|
base = BreaksMidStream()
|
|
wrapper = RobustChatModelWrapper(base)
|
|
result = await wrapper.complete([])
|
|
|
|
assert result.content == "Recovered output"
|
|
resets = [e for e in recorded_events if e[0] == "llm_stream_reset"]
|
|
assert len(resets) == 1
|
|
event_name, payload, status = resets[0]
|
|
assert event_name == "llm_stream_reset"
|
|
assert payload["action"] == "discard"
|
|
assert payload["reason"] == "mid_stream_failure"
|
|
assert "stream_exec_id" in payload
|
|
assert "lower API priority" in payload["message"]
|