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/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.