Add stream retry to all generation paths, not just round 0
Adds _stream_with_retry() async generator (wraps stream_chat_with_tools with up to 2 retries on Ollama 500, 3s/6s delay). Previously only the optimistic round 0 _fill_queue had retry logic. Two paths were still bare: the declined-write-tool fresh stream, and the round 1+ stream. Round 1 500s occur when tag suggestions (fire-and-forget inside execute_tool) race the follow-up stream to the same model. The retry waits for tag suggestions to complete before succeeding. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -195,6 +195,41 @@ async def _drain_queue(
|
||||
yield item
|
||||
|
||||
|
||||
async def _stream_with_retry(
|
||||
messages: list[dict],
|
||||
model: str,
|
||||
tools: list[dict],
|
||||
think: bool,
|
||||
) -> AsyncGenerator[ChatChunk, None]:
|
||||
"""stream_chat_with_tools with automatic retry on Ollama 500 errors.
|
||||
|
||||
500s occur when Ollama is still loading a model or handling a concurrent
|
||||
request (e.g. tag suggestions racing with round 1). Retries up to 2 times
|
||||
with a short delay — by which point the model is warm and other calls done.
|
||||
"""
|
||||
last_exc: BaseException | None = None
|
||||
for attempt in range(3):
|
||||
if attempt > 0:
|
||||
delay = 3.0 * attempt
|
||||
logger.warning(
|
||||
"Ollama stream 500 (attempt %d/3), retrying in %.0fs", attempt, delay
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
try:
|
||||
async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think):
|
||||
yield chunk
|
||||
return
|
||||
except httpx.HTTPStatusError as exc:
|
||||
last_exc = exc
|
||||
if exc.response.status_code != 500:
|
||||
break # non-500 is not retryable
|
||||
except BaseException as exc:
|
||||
last_exc = exc
|
||||
break
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
|
||||
|
||||
async def run_generation(
|
||||
buf: GenerationBuffer,
|
||||
history: list[dict],
|
||||
@@ -467,7 +502,7 @@ async def run_generation(
|
||||
|
||||
if use_tool_path:
|
||||
# Declined write tool — the optimistic stream was cancelled, start fresh.
|
||||
stream_source: AsyncGenerator = stream_chat_with_tools(messages, model, tools=tools, think=think)
|
||||
stream_source: AsyncGenerator = _stream_with_retry(messages, model, tools, think)
|
||||
else:
|
||||
stream_source = _drain_queue(prefetched, stream_queue)
|
||||
|
||||
@@ -546,7 +581,7 @@ async def run_generation(
|
||||
|
||||
buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
|
||||
t_stream = time.monotonic()
|
||||
async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think):
|
||||
async for chunk in _stream_with_retry(messages, model, tools, think):
|
||||
if buf.cancel_event.is_set():
|
||||
cancelled = True
|
||||
break
|
||||
|
||||
Reference in New Issue
Block a user