feat(rag): retry a failed LLM request before the first token lands — BOR_LLM_RETRIES/BOR_LLM_RETRY_DELAY with a live 'retrying' status
This commit is contained in:
@@ -61,6 +61,8 @@ class FakeRagLLM:
|
||||
stream_error: Exception | None = None,
|
||||
fail_mid_stream: bool = False,
|
||||
tool_script: list[list[StreamPiece | ToolCallPiece]] | None = None,
|
||||
embed_fail_count: int = 0,
|
||||
stream_fail_count: int = 0,
|
||||
) -> None:
|
||||
self.settings = Settings(_env_file=None) # pyright: ignore[reportCallIssue]
|
||||
self.embed_batches = 0
|
||||
@@ -69,6 +71,14 @@ class FakeRagLLM:
|
||||
self.embed_error = embed_error
|
||||
self.stream_error = stream_error
|
||||
self.fail_mid_stream = fail_mid_stream
|
||||
#: Phase 67: the first N ``embed_one`` calls raise an
|
||||
#: ``EmbeddingError`` (then succeed) — a dead-then-recovered
|
||||
#: embeddings endpoint for the retry loop.
|
||||
self.embed_fail_count = embed_fail_count
|
||||
#: Phase 67: the first N ``chat_stream`` requests die with an
|
||||
#: ``LLMError`` BEFORE any piece (then succeed) — a dead-then-
|
||||
#: recovered answer endpoint for the pre-first-piece retry rule.
|
||||
self.stream_fail_count = stream_fail_count
|
||||
self.question_embeds: list[str] = []
|
||||
self.seen_messages: list[list[dict[str, str]]] = []
|
||||
#: Every request's ``tools`` value (phase 37) — ``None`` is the
|
||||
@@ -100,6 +110,10 @@ class FakeRagLLM:
|
||||
async def embed_one(self, text: str) -> list[float]:
|
||||
if self.embed_error is not None:
|
||||
raise self.embed_error
|
||||
if self.embed_fail_count > 0:
|
||||
self.embed_fail_count -= 1
|
||||
self.question_embeds.append(text)
|
||||
raise EmbeddingError("simulated embeddings endpoint failure")
|
||||
self.question_embeds.append(text)
|
||||
return _token_vec(text)
|
||||
|
||||
@@ -118,6 +132,9 @@ class FakeRagLLM:
|
||||
self.seen_tools.append(tools)
|
||||
if self.stream_error is not None:
|
||||
raise self.stream_error
|
||||
if self.stream_fail_count > 0:
|
||||
self.stream_fail_count -= 1
|
||||
raise LLMError("simulated pre-piece endpoint failure")
|
||||
if tools is not None and self.tool_script:
|
||||
for piece in self.tool_script.pop(0):
|
||||
yield piece
|
||||
@@ -388,17 +405,33 @@ def test_chat_empty_kb_streams_empty_sources(client, db) -> None:
|
||||
assert row.sources == ""
|
||||
|
||||
|
||||
def test_chat_embed_failure_yields_error_event(client, db, seeded_kb: FakeRagLLM) -> None:
|
||||
def test_chat_embed_failure_yields_error_event(
|
||||
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Phase 67: a dead embeddings endpoint retries on the configured
|
||||
budget — one ``retry`` frame per restart (the attempt about to be
|
||||
tried, 1-based) — and settles on the existing terminal error frame;
|
||||
no query_log row. Zero delay keeps the exhaustion path fast."""
|
||||
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||
try:
|
||||
_, _, frames = _stream_chat(client, QUESTION)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert len(frames) == 1
|
||||
assert frames[0]["type"] == "error"
|
||||
assert "embedding" in frames[0]["detail"]
|
||||
retries = live.llm_retries
|
||||
assert [f["type"] for f in frames] == ["retry"] * retries + ["error"]
|
||||
assert [f["attempt"] for f in frames if f["type"] == "retry"] == list(
|
||||
range(2, retries + 2)
|
||||
)
|
||||
assert all(
|
||||
f["max_attempts"] == retries + 1 for f in frames if f["type"] == "retry"
|
||||
)
|
||||
assert "embedding" in frames[-1]["detail"]
|
||||
assert db.scalars(select(QueryLog)).all() == []
|
||||
|
||||
|
||||
@@ -416,11 +449,19 @@ def test_chat_mid_stream_failure_yields_error_after_partial_deltas(client, db) -
|
||||
assert db.scalars(select(QueryLog)).all() == []
|
||||
|
||||
|
||||
def test_error_event_matches_contract_shape(client, db, seeded_kb) -> None:
|
||||
def test_error_event_matches_contract_shape(
|
||||
client, db, seeded_kb, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The SSE error event (PLAN §4) is exactly ``{type, detail}`` — the
|
||||
client's loading-feedback state machine (phase 06) keys off this shape
|
||||
to flip to the error state and re-enable the send button."""
|
||||
to flip to the error state and re-enable the send button.
|
||||
``llm_retries=0`` keeps this a single-attempt turn: the contract under
|
||||
test is the error frame itself, not the phase-67 retry loop."""
|
||||
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=0)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||
try:
|
||||
_, _, frames = _stream_chat(client, QUESTION)
|
||||
@@ -718,3 +759,156 @@ def test_tool_execution_db_failure_yields_error_event(
|
||||
assert frames[0]["name"] == "list_documents"
|
||||
assert "offline mid-question" in frames[1]["detail"]
|
||||
assert db.scalars(select(QueryLog)).all() == [] # no row for a failed turn
|
||||
|
||||
|
||||
# ---------- phase 67: LLM retries before the first token ----------
|
||||
|
||||
|
||||
def _retry_settings(live: Settings, **overrides: Any) -> Settings:
|
||||
"""Settings for the retry tests: the live (mock-calibrated) threshold
|
||||
plus the phase-67 knobs, with a ZERO delay so the suite never sleeps.
|
||||
(The 5 s default is unit-pinned in ``tests/unit/test_config.py``.)"""
|
||||
kwargs: dict[str, Any] = {
|
||||
"relevance_threshold": live.relevance_threshold,
|
||||
"llm_retry_delay": 0.0,
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return Settings(_env_file=None, **kwargs) # pyright: ignore[reportCallIssue]
|
||||
|
||||
|
||||
def test_embed_failure_retries_then_turn_completes(
|
||||
client,
|
||||
db,
|
||||
seeded_kb: FakeRagLLM,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""A dead-then-recovered embeddings endpoint: one SSE ``retry`` frame
|
||||
(the attempt about to be tried, 1-based) ahead of the normal answer
|
||||
frames; the turn completes and the per-turn log line counts the
|
||||
retry (``retries=1``)."""
|
||||
flaky = FakeRagLLM(embed_fail_count=1)
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=1)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: flaky
|
||||
try:
|
||||
caplog.set_level(logging.INFO, logger="app.chat")
|
||||
_, _, frames = _stream_chat(client, QUESTION)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert frames[0] == {"type": "retry", "attempt": 2, "max_attempts": 2}
|
||||
assert not any(f["type"] == "error" for f in frames)
|
||||
deltas = [f for f in frames if f["type"] == "delta"]
|
||||
assert len(deltas) >= 2
|
||||
assert "".join(d["text"] for d in deltas) == flaky.answer
|
||||
assert frames[-1]["type"] == "done"
|
||||
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
|
||||
assert lines and "retries=1" in lines[-1]
|
||||
|
||||
|
||||
def test_embed_failure_exhausts_retries_then_terminal_error(
|
||||
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A dead embeddings endpoint (``llm_retries=2`` → 3 attempts): one
|
||||
``retry`` frame per restart (attempts 2 and 3 of 3), then the
|
||||
EXISTING terminal error frame — the copy is unchanged, no query_log
|
||||
row."""
|
||||
dead = FakeRagLLM(embed_fail_count=99) # every attempt fails
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=2)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: dead
|
||||
try:
|
||||
_, _, frames = _stream_chat(client, QUESTION)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert [f["type"] for f in frames] == ["retry", "retry", "error"]
|
||||
assert [f["attempt"] for f in frames if f["type"] == "retry"] == [2, 3]
|
||||
assert all(f["max_attempts"] == 3 for f in frames if f["type"] == "retry")
|
||||
assert "embedding" in frames[-1]["detail"]
|
||||
assert db.scalars(select(QueryLog)).all() == []
|
||||
|
||||
|
||||
def test_deflected_stream_retries_before_the_first_piece(
|
||||
client,
|
||||
db,
|
||||
seeded_kb: FakeRagLLM,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Deflected answer stream: the first attempt dies before any piece,
|
||||
the restart streams — a ``retry`` frame ahead of the deltas, the
|
||||
request restarted with the same messages (no tools key), and the
|
||||
per-turn log line counts the retry."""
|
||||
flaky = FakeRagLLM(stream_fail_count=1)
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=2)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: flaky
|
||||
try:
|
||||
caplog.set_level(logging.INFO, logger="app.chat")
|
||||
_, _, frames = _stream_chat(client, OFF_TOPIC)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert frames[0] == {"type": "retry", "attempt": 2, "max_attempts": 3}
|
||||
rest = frames[1:]
|
||||
assert all(f["type"] in ("delta", "done") for f in rest)
|
||||
assert "".join(f["text"] for f in rest if f["type"] == "delta") == flaky.answer
|
||||
assert rest[-1]["type"] == "done" and rest[-1]["deflected"] is True
|
||||
assert len(flaky.seen_messages) == 2 # the request was restarted
|
||||
assert flaky.seen_tools == [None, None] # …byte-identical (no tools key)
|
||||
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
|
||||
assert lines and "retries=1" in lines[-1]
|
||||
|
||||
|
||||
def test_deflected_stream_failure_after_first_frame_is_terminal(
|
||||
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Locked A2: a stream failure AFTER the first output frame is
|
||||
terminal — no ``retry`` frame, the existing error copy, no row (a
|
||||
partial answer is never redone)."""
|
||||
broken = FakeRagLLM(fail_mid_stream=True)
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=3)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||
try:
|
||||
_, _, frames = _stream_chat(client, OFF_TOPIC)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert [f["type"] for f in frames] == ["delta", "error"]
|
||||
assert not any(f["type"] == "retry" for f in frames)
|
||||
assert "dropped the connection" in frames[1]["detail"]
|
||||
assert db.scalars(select(QueryLog)).all() == []
|
||||
|
||||
|
||||
def test_zero_retries_keep_the_pre_phase_wire_shape(
|
||||
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The ``BOR_LLM_RETRIES=0`` kill switch: one attempt, the existing
|
||||
terminal error frame, no ``retry`` frames — the pre-phase-67
|
||||
byte-identical wire shape."""
|
||||
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
|
||||
live = get_settings()
|
||||
monkeypatch.setattr(
|
||||
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=0)
|
||||
)
|
||||
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||
try:
|
||||
_, _, frames = _stream_chat(client, QUESTION)
|
||||
finally:
|
||||
fastapi_app.dependency_overrides.clear()
|
||||
|
||||
assert len(frames) == 1
|
||||
assert frames[0]["type"] == "error"
|
||||
assert "embedding" in frames[0]["detail"]
|
||||
assert not any(f["type"] == "retry" for f in frames)
|
||||
|
||||
Reference in New Issue
Block a user