188 lines
7.2 KiB
Python
188 lines
7.2 KiB
Python
"""Runaway click arguments must never reach the executor or endless retries."""
|
|
|
|
from copy import deepcopy
|
|
import json
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
from langchain_core.messages import AIMessageChunk
|
|
import pytest
|
|
|
|
from artemis.llm.reliability import FailureCategory, LLMExhaustedError, retry_policy_for
|
|
from artemis.llm.router import ModelEndpoint
|
|
from artemis.llm.stream_guard import (
|
|
SMALL_ACTION_ARGUMENT_LIMIT,
|
|
ToolArgumentGuard,
|
|
ToolArgumentOverflow,
|
|
)
|
|
from artemis.services import llm
|
|
|
|
|
|
def delta(arguments, *, name=None, index=0):
|
|
return AIMessageChunk(
|
|
content=[],
|
|
tool_call_chunks=[{"name": name, "args": arguments, "index": index}],
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def isolate(monkeypatch, tmp_path):
|
|
llm._ENDPOINT_BREAKER._states.clear()
|
|
llm._NON_STREAMING_ENDPOINTS.clear()
|
|
monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
|
|
monkeypatch.setattr(
|
|
llm, "retry_policy_for", lambda _: SimpleNamespace(max_attempts=2, delay_for=lambda _: 0)
|
|
)
|
|
monkeypatch.setattr(llm, "_wait_for_resume", AsyncMock(return_value=False))
|
|
monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", None)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["save_note", "input_text", "report_task_status", "plugin_tool"])
|
|
def test_free_text_tools_and_regular_text_are_not_limited(name):
|
|
guard = ToolArgumentGuard()
|
|
text = "报告内容" * SMALL_ACTION_ARGUMENT_LIMIT
|
|
guard.observe(AIMessageChunk(content=text))
|
|
guard.observe(delta(json.dumps({"text": text}, ensure_ascii=False), name=name))
|
|
|
|
|
|
def test_split_and_interleaved_arguments_are_counted_per_call_without_storing_body():
|
|
guard = ToolArgumentGuard()
|
|
for index in (0, 1):
|
|
guard.observe(delta('{"target_description":"', name="click", index=index))
|
|
for index in (0, 1):
|
|
guard.observe(delta("private" * 1200, index=index))
|
|
with pytest.raises(ToolArgumentOverflow) as caught:
|
|
guard.observe(delta("private" * 1200, index=1))
|
|
assert caught.value.details["tool"] == "click"
|
|
assert caught.value.details["argument_chars"] > SMALL_ACTION_ARGUMENT_LIMIT
|
|
assert "private" not in str(caught.value)
|
|
assert "private" not in repr(guard.calls)
|
|
|
|
|
|
def test_name_arriving_after_arguments_still_triggers_guard():
|
|
guard = ToolArgumentGuard()
|
|
guard.observe(delta("x" * (SMALL_ACTION_ARGUMENT_LIMIT + 1)))
|
|
with pytest.raises(ToolArgumentOverflow):
|
|
guard.observe(delta("", name="click"))
|
|
|
|
|
|
def test_invalid_output_retry_policy_is_bounded_and_honors_zero_retries(monkeypatch):
|
|
monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 5)
|
|
assert retry_policy_for(FailureCategory.INVALID_OUTPUT).max_attempts == 2
|
|
monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 0)
|
|
assert retry_policy_for(FailureCategory.INVALID_OUTPUT).max_attempts == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("recover", [False, True])
|
|
@pytest.mark.parametrize("first_chunk_overflows", [False, True])
|
|
async def test_gateway_closes_bad_stream_discards_partial_and_bounds_retry(
|
|
monkeypatch, recover, first_chunk_overflows
|
|
):
|
|
records, resets, closed = [], [], []
|
|
|
|
class Engine:
|
|
current_session_id = "test-session"
|
|
current_step_id = None
|
|
|
|
def record_trace(self, **record):
|
|
records.append(deepcopy(record))
|
|
|
|
def _publish(self, *_):
|
|
pass
|
|
|
|
class Model:
|
|
calls = 0
|
|
|
|
async def astream(self, *_args, **_kwargs):
|
|
self.calls += 1
|
|
try:
|
|
if recover and self.calls == 2:
|
|
yield delta('{"target":3}', name="click")
|
|
return
|
|
if not first_chunk_overflows:
|
|
yield delta('{"target_description":"', name="click")
|
|
yield delta(
|
|
"private" * (SMALL_ACTION_ARGUMENT_LIMIT // 7 + 1),
|
|
name="click" if first_chunk_overflows else None,
|
|
)
|
|
pytest.fail("Runaway stream was consumed after the safety bound")
|
|
finally:
|
|
closed.append(self.calls)
|
|
|
|
monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", Engine())
|
|
monkeypatch.setattr(llm, "_record_llm_event", lambda *a, **_: resets.append(a))
|
|
model = Model()
|
|
wrapper = llm.RobustChatModelWrapper(model, endpoint=ModelEndpoint(max_retries=5))
|
|
if recover:
|
|
response = await wrapper.complete([])
|
|
assert response.tool_calls == [
|
|
{"name": "click", "args": {"target": 3}, "id": None, "type": "tool_call"}
|
|
]
|
|
llm._wait_for_resume.assert_not_awaited()
|
|
else:
|
|
with pytest.raises(LLMExhaustedError) as caught:
|
|
await wrapper.complete([])
|
|
assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
llm._wait_for_resume.assert_not_awaited()
|
|
assert not llm.PAUSE_FILE.exists()
|
|
assert any(name == "llm_gave_up" and p["handover"] == "task_failed" for name, p in resets)
|
|
assert model.calls == 2
|
|
assert closed == [1, 2]
|
|
failure = next(
|
|
r["payload"]
|
|
for r in records
|
|
if r["type"] == "llm_diagnostic" and r["payload"]["phase"] == "failed"
|
|
)
|
|
assert failure["stream_guard"]["tool"] == "click"
|
|
assert failure["stream_guard"]["argument_chars"] > SMALL_ACTION_ARGUMENT_LIMIT
|
|
assert failure["failure_phase"] == "generating_tool_arguments"
|
|
assert failure["category"] == "invalid_output"
|
|
assert failure["usage"] is None
|
|
assert "private" not in json.dumps(records + resets, default=str)
|
|
assert any(name == "llm_stream_reset" and p["action"] == "discard" for name, p in resets)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("fallback_kind", ["same", "failing", "working"])
|
|
@pytest.mark.parametrize("already_paused", [False, True])
|
|
async def test_invalid_output_fallback_does_not_pause_or_clear_another_request(
|
|
fallback_kind, already_paused
|
|
):
|
|
calls = []
|
|
if already_paused:
|
|
llm.PAUSE_FILE.write_text("another request or user pause", encoding="utf-8")
|
|
|
|
async def main():
|
|
calls.append("main")
|
|
raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 1)
|
|
|
|
async def fallback():
|
|
calls.append("fallback")
|
|
if fallback_kind == "working":
|
|
return "valid response"
|
|
raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 2)
|
|
|
|
async def run():
|
|
return await llm.with_fallback(
|
|
lambda: llm._run_with_recovery(main, endpoint_key="main", max_retries=5),
|
|
lambda: llm._run_with_recovery(
|
|
fallback,
|
|
endpoint_key="main" if fallback_kind == "same" else "fallback",
|
|
max_retries=5,
|
|
),
|
|
)
|
|
|
|
if fallback_kind == "working":
|
|
assert await run() == "valid response"
|
|
else:
|
|
with pytest.raises(LLMExhaustedError) as caught:
|
|
await run()
|
|
assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT
|
|
assert calls.count("main") == 2
|
|
assert calls.count("fallback") == {"same": 0, "failing": 2, "working": 1}[fallback_kind]
|
|
llm._wait_for_resume.assert_not_awaited()
|
|
assert llm.PAUSE_FILE.exists() is already_paused
|
|
if already_paused:
|
|
assert llm.PAUSE_FILE.read_text(encoding="utf-8") == "another request or user pause"
|