246 lines
8.5 KiB
Python
246 lines
8.5 KiB
Python
import asyncio
|
|
from copy import deepcopy
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.messages import AIMessageChunk, HumanMessage, ToolMessage
|
|
from langchain_core.tools import tool
|
|
import pytest
|
|
|
|
from artemis.agents.checker import checker
|
|
from artemis.config.llm import LLM
|
|
from artemis.llm import final_review
|
|
from artemis.llm.diagnostics import ACTIVE_ATTEMPT
|
|
from artemis.llm.final_review import (
|
|
ACTIVE_FINAL_REVIEW,
|
|
FinalReviewBudget,
|
|
FinalReviewBudgetExceeded,
|
|
)
|
|
from artemis.llm.reliability import LLMExhaustedError
|
|
from artemis.llm.router import ModelEndpoint
|
|
from artemis.services import llm
|
|
|
|
|
|
class RecordingEngine:
|
|
current_session_id = "review-test"
|
|
current_step_id = None
|
|
|
|
def __init__(self):
|
|
self.records = []
|
|
|
|
def record_trace(self, **record):
|
|
self.records.append(deepcopy(record))
|
|
|
|
def _publish(self, *_):
|
|
pass
|
|
|
|
def latest(self, kind):
|
|
return next(r["payload"] for r in reversed(self.records) if r["type"] == kind)
|
|
|
|
|
|
@pytest.fixture
|
|
def engine(monkeypatch, tmp_path):
|
|
engine = RecordingEngine()
|
|
monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
|
|
monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path))
|
|
monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
|
|
monkeypatch.setattr(
|
|
llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
|
|
)
|
|
llm._ENDPOINT_BREAKER._states.clear()
|
|
yield engine
|
|
llm._ENDPOINT_BREAKER._states.clear()
|
|
assert ACTIVE_FINAL_REVIEW.get() is None
|
|
assert ACTIVE_ATTEMPT.get() is None
|
|
|
|
|
|
def budget(engine, *, timeout=1, retries=1):
|
|
return FinalReviewBudget(
|
|
engine, attempt_id="final#1", request_timeout=timeout, max_retries=retries, retry_delay=0
|
|
)
|
|
|
|
|
|
def test_budget_uses_effective_checker_config_and_respects_overall_cap(monkeypatch):
|
|
monkeypatch.setattr(llm.settings, "OPENAI_TIMEOUT_SECONDS", 300)
|
|
monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 5)
|
|
monkeypatch.setattr(llm.settings, "OPENAI_RETRY_INTERVAL_SECONDS", 10)
|
|
monkeypatch.setattr(llm.settings, "OPENAI_RATE_LIMIT_RETRY_SECONDS", 60)
|
|
cfg = LLM(model="test", timeout=20, max_retries=2)
|
|
ctx = SimpleNamespace(llm_config=SimpleNamespace(get_agent=lambda name: cfg), data_engine=None)
|
|
review = llm.final_review_budget(ctx, attempt_id="final#1")
|
|
assert review.total_seconds == 280 # three normal calls + two retries + waits/tool reads
|
|
assert review.data["request_timeout_seconds"] == 20
|
|
assert review.data["max_retries"] == 2
|
|
|
|
cfg.timeout = 300
|
|
cfg.max_retries = 100
|
|
assert llm.final_review_budget(ctx, attempt_id="final#2").total_seconds == 3600
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_checker_retries_only_failed_request_with_existing_tool_evidence(engine, monkeypatch):
|
|
calls = []
|
|
evidence_reads = 0
|
|
stream_closed = asyncio.Event()
|
|
|
|
@tool
|
|
async def read_evidence() -> str:
|
|
"""Read the already-recorded page evidence."""
|
|
nonlocal evidence_reads
|
|
evidence_reads += 1
|
|
return "实际为中文搜索，预期为英文搜索"
|
|
|
|
class ScriptedModel:
|
|
def bind_tools(self, *args, **_):
|
|
return self
|
|
|
|
async def astream(self, messages, **_):
|
|
calls.append(deepcopy(messages))
|
|
if len(calls) == 1:
|
|
yield AIMessageChunk(
|
|
content="",
|
|
tool_calls=[{"name": "read_evidence", "args": {}, "id": "read-1"}],
|
|
)
|
|
elif len(calls) == 2:
|
|
try:
|
|
yield AIMessageChunk(content='{"test_result":{"status":"passed","reason":"')
|
|
await asyncio.Event().wait() # the *request* timer triggers a retry
|
|
finally:
|
|
stream_closed.set()
|
|
else:
|
|
yield AIMessageChunk(
|
|
content=json.dumps(
|
|
{
|
|
"test_result": {
|
|
"status": "failed",
|
|
"reason": "中文搜索不满足英文预期",
|
|
"evidence": ["read_evidence：实际为中文搜索"],
|
|
}
|
|
}
|
|
)
|
|
)
|
|
|
|
wrapper = llm.RobustChatModelWrapper(
|
|
ScriptedModel(),
|
|
endpoint=ModelEndpoint(model_name="test", timeout_seconds=0.02, max_retries=1),
|
|
)
|
|
monkeypatch.setattr(checker, "get_llm", lambda **_: wrapper)
|
|
monkeypatch.setattr(checker, "trace_langchain_tool", lambda tool, ctx: tool)
|
|
events = []
|
|
monkeypatch.setattr(llm, "_record_llm_event", lambda *args, **kwargs: events.append(args))
|
|
ctx = SimpleNamespace(execution_setup=SimpleNamespace(checker_max_iterations=4))
|
|
review = budget(engine)
|
|
report = await review.run(
|
|
lambda: checker._run_check_loop(
|
|
ctx, [HumanMessage(content="核实英文搜索")], [read_evidence], []
|
|
)
|
|
)
|
|
|
|
assert report.test_result.status == "failed" # partial `passed` was never a verdict
|
|
assert evidence_reads == 1
|
|
assert len(calls) == 3
|
|
assert calls[1] == calls[2]
|
|
assert any(isinstance(m, ToolMessage) and "中文搜索" in m.content for m in calls[2])
|
|
assert stream_closed.is_set()
|
|
assert any(e[0] == "llm_stream_reset" and e[1]["action"] == "discard" for e in events)
|
|
stored = engine.latest("final_review_diagnostic")
|
|
assert stored["status"] == "completed"
|
|
assert stored["request_count"] == 2
|
|
assert stored["retry_count"] == 1
|
|
assert stored["last_failure_category"] == "timeout"
|
|
assert engine.latest("llm_diagnostic")["final_review"]["review_id"] == stored["review_id"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("phase", ["request", "retry_wait", "evidence"])
|
|
async def test_total_budget_expiry_is_not_user_cancellation_or_a_new_retry(
|
|
engine, monkeypatch, tmp_path, phase
|
|
):
|
|
monkeypatch.setattr(final_review, "MAX_FINAL_REVIEW_SECONDS", 0.02)
|
|
monkeypatch.setattr(
|
|
llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 60, max_attempts=6)
|
|
)
|
|
cancelled = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def request():
|
|
nonlocal calls
|
|
calls += 1
|
|
if phase == "retry_wait":
|
|
raise ConnectionError("connection reset")
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
cancelled.set()
|
|
|
|
review = budget(engine, retries=5)
|
|
|
|
async def operation():
|
|
if phase == "evidence":
|
|
return await request()
|
|
return await llm._run_with_recovery(request, request_timeout=1, max_retries=5)
|
|
|
|
with pytest.raises(FinalReviewBudgetExceeded, match="总预算已用尽"):
|
|
await review.run(operation)
|
|
assert calls == 1
|
|
assert not (tmp_path / "pause").exists()
|
|
assert engine.latest("final_review_diagnostic")["status"] == "budget_exhausted"
|
|
if phase != "retry_wait":
|
|
assert cancelled.is_set()
|
|
if phase != "evidence":
|
|
stored = engine.latest("llm_diagnostic")
|
|
assert stored["phase"] == "failed"
|
|
assert stored["category"] == "final_review_budget_exhausted"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exhausted_final_review_request_ends_without_hidden_pause(
|
|
engine, monkeypatch, tmp_path
|
|
):
|
|
calls = 0
|
|
|
|
async def fails():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise TimeoutError()
|
|
|
|
pause = []
|
|
monkeypatch.setattr(llm, "_handle_llm_pause_and_resume", lambda *_: pause.append(True))
|
|
with pytest.raises(LLMExhaustedError):
|
|
await budget(engine).run(
|
|
lambda: llm._run_with_recovery(fails, request_timeout=1, max_retries=1)
|
|
)
|
|
assert calls == 2
|
|
assert not pause
|
|
assert not (tmp_path / "pause").exists()
|
|
assert engine.latest("final_review_diagnostic")["status"] == "request_retries_exhausted"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_cancellation_closes_review_and_child_without_retry(engine):
|
|
started = asyncio.Event()
|
|
cancelled = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def request():
|
|
nonlocal calls
|
|
calls += 1
|
|
started.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
cancelled.set()
|
|
|
|
task = asyncio.create_task(
|
|
budget(engine).run(
|
|
lambda: llm._run_with_recovery(request, request_timeout=1, max_retries=5)
|
|
)
|
|
)
|
|
await started.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert calls == 1 and cancelled.is_set()
|
|
assert engine.latest("final_review_diagnostic")["status"] == "cancelled"
|
|
assert engine.latest("llm_diagnostic")["category"] == "cancelled"
|