Experimental/brain-of-reese
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(rag): retry a failed LLM request before the first token lands — BOR_LLM_RETRIES/BOR_LLM_RETRY_DELAY with a live 'retrying' status

Browse Source
This commit is contained in:
ducoterra
2026-09-02 10:52:38 -04:00
parent f04ddbe1f8
commit 88293ed02f
44 changed files with 2488 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/chat.py
+103 -24
View File
@@ -68,9 +68,27 @@ byte-identical to the pre-phase path (A8):** the LOW prompt never
carries tools, and with ``agent_max_rounds`` at **0** ``run_agent``
makes exactly one ``tools=None`` request, reproducing the pre-phase
behavior (the kill switch).
LLM retries (phase 67, ``BOR_LLM_RETRIES`` / ``BOR_LLM_RETRY_DELAY``,
owner-locked 2026-09-01): when the aipi endpoint dies before a request
has streamed its first output frame (locked A2), the request is
restarted — up to ``llm_retries`` times (default 3), a flat
``llm_retry_delay`` (default 5 s) between attempts. Every restart is
announced with an SSE ``retry`` frame (``{"type": "retry",
"attempt": n, "max_attempts": N}`` — the attempt about to be tried,
1-based, ahead of the pre-retry wait) so the UI can show the transient
"Communication interrupted — retrying (n of N)…" status (locked A4);
exhaustion and any failure after the first frame keep the existing
terminal ``error`` frames. The pre-stream question embedding retries on
the same budget, and the deflected answer stream goes through
``chat_stream_retried`` (the agent loop retries per round — task 03).
The per-turn log line records ``retries=N`` after ``total_ms=N`` (0
when nothing was retried — the field is uniform across all turn
shapes).
"""
from __future__ import annotations
import asyncio
import json
import logging
import time
@@ -91,8 +109,10 @@ from app.rag.llm import (
EmbeddingError,
LLMClient,
LLMError,
RetryPiece, # phase 67: one LLM request restart (an SSE retry frame)
StreamPiece, # type of the answer pieces streamed by the agent loop
ToolCallPiece, # phase 37: one model-requested tool call
chat_stream_retried, # phase 67: the retry-before-first-piece primitive
)
from app.rag.overview import load_kb_overview
from app.rag.prompts import build_deflect_prompt, build_high_prompt
@@ -102,6 +122,7 @@ from app.schemas import (
ChatDoneEvent,
ChatErrorEvent,
ChatRequest,
ChatRetryEvent,
ChatThinkingEvent,
ChatToolEvent,
SourceRef,
@@ -246,34 +267,66 @@ async def chat(
# consumer went away before any terminal frame; it must not
# yield (GeneratorExit handling).
settled = False
retries_used = 0 # phase 67: LLM requests restarted this turn (log line)
try:
settings = get_settings()
# 1. Embed the question.
# Phase 67: a dead embeddings endpoint is retried before any
# frame has left the server — up to ``llm_retries`` restarts,
# a flat ``llm_retry_delay`` between attempts, one SSE
# ``retry`` frame per restart (the UI shows the transient
# "retrying" status, not an error — locked A4). The final
# failure keeps the EXISTING terminal ``error`` frame (the
# copy reads correctly after N tries); ``llm_retries=0`` is
# byte-identical to the pre-phase-67 single attempt.
t0 = time.monotonic()
try:
question_vec = await llm.embed_one(request.message)
except EmbeddingError as e:
embed_ms = int((time.monotonic() - t0) * 1000)
total_ms = int((time.monotonic() - started) * 1000)
logger.error(
"chat: question=%r embed_ms=%d total_ms=%d — embedding failed: %s",
request.message,
embed_ms,
total_ms,
e,
)
settled = True # terminal: the error frame settles the turn
yield sse_event(
ChatErrorEvent(
detail="I couldn't reach the embedding model — please try again."
).model_dump()
)
return
max_attempts = settings.llm_retries + 1
attempt = 1
while True:
try:
question_vec = await llm.embed_one(request.message)
break
except EmbeddingError as e:
embed_ms = int((time.monotonic() - t0) * 1000)
if attempt >= max_attempts:
total_ms = int((time.monotonic() - started) * 1000)
logger.error(
"chat: question=%r embed_ms=%d total_ms=%d — embedding failed: %s",
request.message,
embed_ms,
total_ms,
e,
)
settled = True # terminal: the error frame settles the turn
yield sse_event(
ChatErrorEvent(
detail="I couldn't reach the embedding model — please try again."
).model_dump()
)
return
logger.warning(
"chat: question=%r embedding failed (attempt %d/%d) — "
"retrying in %.1fs: %s",
request.message,
attempt,
max_attempts,
settings.llm_retry_delay,
e,
)
retries_used += 1
yield sse_event(
ChatRetryEvent(
attempt=attempt + 1, max_attempts=max_attempts
).model_dump()
)
await asyncio.sleep(settings.llm_retry_delay)
attempt += 1
embed_ms = int((time.monotonic() - t0) * 1000)
# 2. Retrieve top-K chunks, load the owner's steering notes
# (phase 15), then the honesty gate (A8) picks the HIGH
# (grounded) or LOW (deflected) prompt + context.
settings = get_settings()
try:
steering_notes = load_steering_notes(db)
# KB overview (phase 31): one indexed PK lookup per turn —
@@ -313,9 +366,21 @@ async def chat(
# ``agent_max_rounds=0`` ``run_agent`` is a single
# ``tools=None`` request anyway (the kill switch).
holder = AgentHolder()
answer_stream: AsyncIterator[StreamPiece | ToolCallPiece]
answer_stream: AsyncIterator[StreamPiece | ToolCallPiece | RetryPiece]
if plan.deflected:
answer_stream = llm.chat_stream(messages)
# Phase 67: the deflected stream goes through the retry
# primitive — a dead endpoint is restarted (SSE ``retry``
# frames) only before its first piece (locked A2); the
# grounded path stays a plain ``run_agent`` call (task 03
# makes IT retry internally) — its ``RetryPiece``s flow
# through the shared piece loop below.
answer_stream = chat_stream_retried(
llm,
messages,
tools=None,
retries=settings.llm_retries,
delay=settings.llm_retry_delay,
)
else:
answer_stream = run_agent(
llm,
@@ -328,7 +393,7 @@ async def chat(
)
thinking_chars = 0
try:
async for piece in answer_stream: # StreamPiece | ToolCallPiece
async for piece in answer_stream: # StreamPiece | ToolCallPiece | RetryPiece
if isinstance(piece, ToolCallPiece):
# Phase 37 (PLAN §4 extension): one SSE ``tool``
# frame per model-requested call; ``argument`` is
@@ -345,6 +410,19 @@ async def chat(
).model_dump()
)
continue
if isinstance(piece, RetryPiece):
# Phase 67: the answer stream was restarted before
# its first piece (locked A2) — a transient status
# frame, never an error. No other state changes:
# the thinking/clock/timeout handling is the
# client's job.
retries_used += 1
yield sse_event(
ChatRetryEvent(
attempt=piece.attempt, max_attempts=piece.max_attempts
).model_dump()
)
continue
if piece.kind == "thinking":
thinking_chars += len(piece.text)
if settings.stream_thinking:
@@ -419,7 +497,7 @@ async def chat(
logger.info(
"question=%r embed_ms=%d top_score=%.3f fts_hits=%d summary_hits=%d tuning=%d "
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d "
"tool_calls=%d total_ms=%d",
"tool_calls=%d total_ms=%d retries=%d",
request.message,
embed_ms,
plan.top_score,
@@ -433,6 +511,7 @@ async def chat(
thinking_chars,
holder.tool_calls,
total_ms,
retries_used,
)
settled = True # terminal: the done frame settles the turn
yield sse_event(
app/config.py
+25
View File
@@ -73,6 +73,13 @@ class Settings(BaseSettings):
#: pieces are still counted for the per-turn log line but never
#: emitted — the answer stream itself is unchanged.
stream_thinking: bool = True
#: Retries of a failed LLM request when the endpoint stops responding
#: (phase 67, ``BOR_LLM_RETRIES``); ``0`` = no retries (the turn fails
#: on the first error, pre-phase-67 behavior).
llm_retries: int = 3
#: Flat seconds to wait between attempts (phase 67,
#: ``BOR_LLM_RETRY_DELAY``); the TODO-locked 5 s, no backoff.
llm_retry_delay: float = 5.0
# --- RAG tuning ---
embedding_dim: int = 768 # verified against aipi /v1 (embed model)
@@ -261,6 +268,24 @@ class Settings(BaseSettings):
raise ValueError("agent_max_rounds must be >= 0 (0 = no tools)")
return v
@field_validator("llm_retries")
@classmethod
def _llm_retries_non_negative(cls, v: int) -> int:
"""``0`` is the no-retry kill switch (pre-phase-67 behavior) — a
negative value is a typo (the ``agent_max_rounds`` pattern)."""
if v < 0:
raise ValueError("llm_retries must be >= 0 (0 = no retries)")
return v
@field_validator("llm_retry_delay")
@classmethod
def _llm_retry_delay_non_negative(cls, v: float) -> float:
"""A negative delay is a typo — fail loud at startup (the
``agent_max_rounds`` pattern)."""
if v < 0:
raise ValueError("llm_retry_delay must be >= 0 (seconds)")
return v
@field_validator("upload_max_mb")
@classmethod
def _upload_max_mb_positive(cls, v: int) -> int:
app/rag/agent.py
+60 -16
View File
@@ -43,16 +43,27 @@ task 04):
the assistant tool-call message + the tool result (refusals included),
consumes one round, and the model is called again. At the round cap —
``max_rounds = settings.agent_max_rounds`` (``BOR_AGENT_MAX_ROUNDS``,
default 10) — the loop forces one final ``chat_stream(messages,
tools=None)`` and returns: the cap is the **only** forced exit
(besides "the stream carried no calls"), and it bounds pathological
rejected-call streams.
default 10) — the loop forces one final retried no-tools request
(``chat_stream_retried`` with ``tools=None``) and returns: the cap is
the **only** forced exit (besides "the stream carried no calls"), and
it bounds pathological rejected-call streams.
5. A rare stream that carries both content and a tool call keeps the
content (it was already emitted) **and** still runs the tool.
6. *holder* (an :class:`AgentHolder`) records the read documents and the
number of executed tool calls (re-lists included); the API layer
(task 04) reads it after the stream to extend ``done.sources`` /
``query_log.sources`` and the per-turn log line (``tool_calls=N``).
7. Retries (phase 67, owner-locked A2): every model request — each tool
round and the forced final ``tools=None`` call — goes through
``chat_stream_retried``: a round that dies before its first piece is
restarted with the SAME messages (up to ``settings.llm_retries``
restarts, a flat ``settings.llm_retry_delay`` between attempts, each
preceded by a :class:`app.rag.llm.RetryPiece` the API layer turns into
an SSE ``retry`` frame); a round that already streamed a piece fails
the turn as before (no partial answer is ever redone). Retries are
invisible to the round cap: a round that needed a retry still consumes
exactly one round. With ``settings.llm_retries=0`` every request is a
single plain attempt (the pre-phase-67 path).
The DB accessors (:func:`list_catalog`, :func:`find_document`) are
module-level functions so unit tests can monkeypatch them without a
@@ -71,7 +82,13 @@ from sqlalchemy.orm import Session
from app.config import Settings
from app.models import Document
from app.rag.llm import LLMClient, StreamPiece, ToolCallPiece
from app.rag.llm import (
LLMClient,
RetryPiece,
StreamPiece,
ToolCallPiece,
chat_stream_retried,
)
logger = logging.getLogger("app.agent")
@@ -224,13 +241,20 @@ async def run_agent(
seed_docs: Sequence[Document],
settings: Settings,
holder: AgentHolder,
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
) -> AsyncIterator[StreamPiece | ToolCallPiece | RetryPiece]:
"""Run the grounded-turn tool loop, yielding every stream piece.
Every piece (``thinking`` / ``content`` / tool calls) is yielded as it
arrives; the API layer (task 04) turns tool-call pieces into SSE
``tool`` events. After the loop finishes, *holder* carries the read
documents and the executed tool-call count (re-lists included).
Every piece (``thinking`` / ``content`` / tool calls /
:class:`RetryPiece`) is yielded as it arrives; the API layer (task 04)
turns tool-call pieces into SSE ``tool`` events and retry pieces into
SSE ``retry`` events. After the loop finishes, *holder* carries the
read documents and the executed tool-call count (re-lists included).
Retries (phase 67, owner-locked A2): every model request goes through
:func:`chat_stream_retried` — a failed round is retried **before** its
first piece (same messages, ``settings.llm_retries`` restarts, a flat
``settings.llm_retry_delay``); a round that already streamed pieces
fails the turn as before.
``seed_docs`` are the documents the retrieval already put in context
(they shape the *system_prompt* the caller built); re-reading one of
@@ -252,10 +276,22 @@ async def run_agent(
calls: list[ToolCallPiece] = []
# Phase 48: bind the round's stream so a consumer abandon
# (GeneratorExit into the yield below) tears down the in-flight
# model stream deterministically — not GC-dependent. Awaiting
# ``aclose()`` in the ``finally`` is safe because it does not
# yield; on a fully consumed round it is a quiet no-op.
stream = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=tools)
# model stream deterministically — not GC-dependent. Phase 67:
# the round goes through the retry primitive — a failure before
# the first piece restarts the request (locked A2) after a
# RetryPiece; closing the OUTER generator propagates GeneratorExit
# into ``chat_stream_retried``, whose own ``finally`` closes the
# in-flight inner ``chat_stream``, so teardown stays deterministic
# on consumer abandon. Awaiting ``aclose()`` in the ``finally`` is
# safe because it does not yield; on a fully consumed round it is
# a quiet no-op.
stream = chat_stream_retried(
llm,
cast("list[dict[str, str]]", messages),
tools=tools,
retries=settings.llm_retries,
delay=settings.llm_retry_delay,
)
try:
async for piece in stream:
if isinstance(piece, ToolCallPiece):
@@ -300,8 +336,16 @@ async def run_agent(
)
# Phase 48: the forced final answer gets the same explicit
# teardown as the loop rounds (consumer abandon mid-final
# answer must still close the model's stream).
final = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=None)
# answer must still close the model's stream). Phase 67: the
# forced call retries under the same locked-A2 rule as the
# loop rounds.
final = chat_stream_retried(
llm,
cast("list[dict[str, str]]", messages),
tools=None,
retries=settings.llm_retries,
delay=settings.llm_retry_delay,
)
try:
async for piece in final:
yield piece
app/rag/llm.py
+94
View File
@@ -18,6 +18,7 @@ vectors that pgvector rejects.
"""
from __future__ import annotations
import asyncio
import json
import logging
from collections.abc import AsyncGenerator
@@ -82,6 +83,25 @@ class ToolCallPiece:
arguments: dict[str, Any]
@dataclass(frozen=True)
class RetryPiece:
"""One LLM request retry that is about to start (phase 67, locked A2).
``attempt`` is the 1-based number of the attempt that is about to be
tried — the one AFTER the attempt that just failed (a first-attempt
failure carries ``attempt=2``, so the API's SSE ``retry`` frame reads
"retrying (2 of N)" — the same convention the endpoint's embedding
retry loop uses, phase 67 task 02); ``max_attempts`` is the total
attempt budget (``llm_retries + 1``). One piece per wait: the API
layer turns it into an SSE ``retry`` frame, and it always precedes
the pre-retry sleep so the frame reaches the client before the wait
starts.
"""
attempt: int
max_attempts: int
@dataclass
class _ToolCallSlot:
"""Mutable accumulator for one streamed tool call (phase 37, private).
@@ -442,6 +462,80 @@ class LLMClient:
await stream.close()
async def chat_stream_retried(
llm: LLMClient,
messages: list[dict[str, str]],
*,
tools: list[dict[str, Any]] | None = None,
retries: int = 0,
delay: float = 0.0,
) -> AsyncGenerator[StreamPiece | ToolCallPiece | RetryPiece, None]:
"""Stream a chat turn, retrying a dead endpoint (phase 67).
Wraps :meth:`LLMClient.chat_stream` with the retry-before-first-piece
rule (owner-locked A2): a request is restarted **only** while no output
piece (thinking/tool/delta) has been yielded for it. Once pieces have
flowed, an :class:`LLMError` is re-raised unchanged — a partial answer
is never redone, and the API layer's terminal ``error`` frame applies.
This primitive is the ONLY place that rule lives (the chat endpoint
and the agent loop both build on it).
Up to *retries* restarts after the initial attempt (``retries + 1``
attempts total; ``retries=0`` is exactly one attempt with no
:class:`RetryPiece` — the pre-phase-67 kill-switch path). Each restart
is preceded by one :class:`RetryPiece` (``attempt`` = the 1-based
number of the attempt about to be tried — the failed attempt + 1 —
and ``max_attempts`` = ``retries + 1``) and a flat
``asyncio.sleep(delay)`` — the TODO-locked fixed interval, no
backoff. The ``RetryPiece`` always precedes its sleep: the API frame
must reach the client before the wait starts.
The request is restarted byte-identical: ``chat_stream`` is stateless,
so every attempt is opened with the SAME *messages*/*tools*.
Teardown (phase 48, extended): every attempt's stream is explicitly
closed in a ``finally`` — normal exhaustion, a terminal
:class:`LLMError`, and a consumer abandon (``GeneratorExit`` mid-attempt
or during the pre-retry sleep) all run it, so an abandoned turn never
leaves the endpoint's stream open.
"""
max_attempts = retries + 1
for attempt in range(1, max_attempts + 1):
emitted = False
stream = llm.chat_stream(messages, tools=tools)
try:
async for piece in stream:
emitted = True
yield piece
except LLMError as e:
if emitted:
# Locked A2: tokens already flowed — the failure is
# terminal, never redo a partial answer.
raise
if attempt >= max_attempts:
# Retries exhausted — the API layer turns this into the
# terminal error frame.
raise
logger.warning(
"llm stream failed before the first piece (attempt %d/%d) — "
"retrying in %.1fs: %s",
attempt,
max_attempts,
delay,
e,
)
yield RetryPiece(attempt + 1, max_attempts)
await asyncio.sleep(delay)
else:
# A fully consumed attempt — the turn is done; the loop must
# NOT open another stream for an attempt that never failed.
return
finally:
# Phase 48 teardown for THIS attempt's stream (a no-op once the
# stream already ended; the real close on a consumer abandon).
await stream.aclose()
async def check_models(llm: LLMClient) -> None:
"""Verify the models a sync needs (embed + summary) before any
expensive work; raise ModelUnavailableError naming the model.
app/schemas.py
+24
View File
@@ -108,6 +108,30 @@ class ChatErrorEvent(BaseModel):
detail: str
class ChatRetryEvent(BaseModel):
"""SSE retry event: an LLM request is restarted before the first token
(phase 67, owner-locked 2026-09-01).
Sibling of :class:`ChatErrorEvent`, but transient — the client shows a
live status on the existing ``#send-status`` line (locked A4:
"Communication interrupted — retrying (n of N)…") and the send button
stays the Stop control; it never flips the state machine to error. It
is only ever sent when the failed attempt had NOT streamed a single
output frame yet (locked A2: no thinking/tool/delta emitted) — once
tokens are flowing, a failure is terminal (the ``error`` frame) and
this event cannot appear.
``attempt`` is the 1-based number of the attempt the endpoint is about
to try next (what the endpoint sends — the first failure of a
4-attempt budget carries ``attempt=2``); ``max_attempts`` is the total
attempt budget (``llm_retries + 1``).
"""
type: Literal["retry"] = "retry"
attempt: int
max_attempts: int
class DocSummary(BaseModel):
"""One indexed document as shown on the Sources page / API."""