chenmingxuan
e1ea099592
Install and update artemis-use with task diagnostics and plugin guidance. Preserve cross-device file receipts, review canonical reports, guard stage transitions and reduce small history summaries. Improve exploration lifecycle, request recovery and browser live connections.
404 lines
13 KiB
Python
404 lines
13 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_consumer_error_closes_live_model_stream_before_returning():
|
|
closed = []
|
|
|
|
class StreamingModel:
|
|
async def astream(self, *args, **kwargs):
|
|
try:
|
|
yield AIMessageChunk(content="partial")
|
|
finally:
|
|
closed.append(True)
|
|
|
|
wrapper = RobustChatModelWrapper(StreamingModel())
|
|
|
|
def failing_consumer(*_):
|
|
raise RuntimeError("output sink closed")
|
|
|
|
wrapper._emit_stream_delta = failing_consumer
|
|
with pytest.raises(RuntimeError, match="output sink closed"):
|
|
await wrapper._complete_attempt(True, [])
|
|
assert closed == [True]
|
|
|
|
|
|
@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
|
|
@pytest.mark.parametrize("method", ["ainvoke", "complete"])
|
|
async def test_streamed_tool_pairing_error_does_not_retry_pause_or_fallback(monkeypatch, method):
|
|
import httpx
|
|
from openai import APIError
|
|
from unittest.mock import AsyncMock
|
|
|
|
error = APIError(
|
|
"No tool call found for function call output with call_id call_test.",
|
|
request=httpx.Request("POST", "https://example.test/responses"),
|
|
body=None,
|
|
)
|
|
calls = []
|
|
|
|
class BrokenHistory:
|
|
async def ainvoke(self, *_args, **_kwargs):
|
|
calls.append("invoke")
|
|
raise error
|
|
|
|
async def astream(self, *_args, **_kwargs):
|
|
calls.append("stream")
|
|
raise error
|
|
yield # Make this a stream whose error arrives before any chunks.
|
|
|
|
pause, fallback = AsyncMock(), AsyncMock()
|
|
monkeypatch.setattr(llm_service, "_wait_for_resume", pause)
|
|
wrapper = RobustChatModelWrapper(BrokenHistory())
|
|
with pytest.raises(LLMPermanentError):
|
|
await with_fallback(lambda: getattr(wrapper, method)([]), fallback)
|
|
assert len(calls) == 1
|
|
pause.assert_not_called()
|
|
fallback.assert_not_called()
|
|
|
|
|
|
@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 "未完成的输出已丢弃" in payload["message"]
|
|
assert "lower API priority" not 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()
|