347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""Tests for the LLM gateway: complete(), classified recovery, and fallback."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage
|
|
import pytest
|
|
|
|
from artemis.llm.reliability import LLMExhaustedError, LLMPermanentError
|
|
from artemis.llm.router import ModelEndpoint
|
|
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_gpt_stream_uses_adapted_messages_without_mutating_history():
|
|
class RecordingStream:
|
|
messages = None
|
|
|
|
async def astream(self, messages, **kwargs):
|
|
self.messages = messages
|
|
yield AIMessageChunk(content="OK")
|
|
|
|
base = RecordingStream()
|
|
wrapper = RobustChatModelWrapper(
|
|
base, endpoint=ModelEndpoint(model_name="gpt-6-luna", max_retries=0)
|
|
)
|
|
original = [SystemMessage(content="Use tools"), HumanMessage(content="Inspect screen")]
|
|
result = await wrapper.complete(original)
|
|
|
|
assert result.content == "OK"
|
|
assert len(base.messages) == 1
|
|
assert isinstance(base.messages[0], HumanMessage)
|
|
assert base.messages[0].content[0]["text"] == "Use tools"
|
|
assert base.messages[0].content[1]["text"] == "Inspect screen"
|
|
assert isinstance(original[0], SystemMessage)
|
|
|
|
|
|
@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"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_configured_retry_budget_is_shared_by_identical_fallback(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from artemis.llm.reliability import retry_policy_for
|
|
from artemis.llm.router import ModelEndpoint
|
|
|
|
monkeypatch.setattr(llm_service, "retry_policy_for", retry_policy_for)
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_MAX_RETRIES", 5)
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_RETRY_INTERVAL_SECONDS", 10)
|
|
monkeypatch.setattr(llm_service, "_wait_for_resume", AsyncMock(return_value=False))
|
|
sleep = AsyncMock()
|
|
monkeypatch.setattr(llm_service.asyncio, "sleep", sleep)
|
|
|
|
class Fails:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
async def ainvoke(self, *args, **kwargs):
|
|
self.calls += 1
|
|
if self.calls % 2:
|
|
raise ConnectionError("connection reset")
|
|
raise TimeoutError("timeout")
|
|
|
|
main, fallback = Fails(), Fails()
|
|
endpoint = ModelEndpoint(model_name="test", api_base="https://example.test/v1", max_retries=5)
|
|
primary = RobustChatModelWrapper(main, endpoint=endpoint)
|
|
secondary = RobustChatModelWrapper(fallback, endpoint=endpoint.model_copy())
|
|
with pytest.raises(LLMExhaustedError):
|
|
await with_fallback(lambda: primary.ainvoke([]), lambda: secondary.ainvoke([]))
|
|
assert main.calls == 6
|
|
assert fallback.calls == 0
|
|
assert [call.args[0] for call in sleep.call_args_list] == [10] * 5
|
|
assert llm_service._FAILED_ENDPOINTS.get() is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_429_waits_one_minute_then_retries(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from artemis.llm.reliability import retry_policy_for
|
|
|
|
class RateLimited(Exception):
|
|
status_code = 429
|
|
|
|
calls = 0
|
|
|
|
async def call_once():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise RateLimited("Too Many Requests")
|
|
return "recovered"
|
|
|
|
sleep = AsyncMock()
|
|
monkeypatch.setattr(llm_service, "retry_policy_for", retry_policy_for)
|
|
monkeypatch.setattr(llm_service.asyncio, "sleep", sleep)
|
|
monkeypatch.setattr(llm_service.settings, "OPENAI_RATE_LIMIT_RETRY_SECONDS", 60)
|
|
assert await llm_service._run_with_recovery(call_once, max_retries=1) == "recovered"
|
|
assert calls == 2
|
|
sleep.assert_awaited_once_with(60)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_deadline_cancels_attempt_and_respects_zero_retries(monkeypatch):
|
|
import asyncio
|
|
from unittest.mock import AsyncMock
|
|
|
|
monkeypatch.setattr(llm_service, "_wait_for_resume", AsyncMock(return_value=False))
|
|
cancelled = asyncio.Event()
|
|
|
|
async def pending():
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
cancelled.set()
|
|
|
|
with pytest.raises(LLMExhaustedError):
|
|
await llm_service._run_with_recovery(pending, request_timeout=0.005, max_retries=0)
|
|
assert cancelled.is_set()
|