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
|
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(
|
async def run_generation(
|
||||||
buf: GenerationBuffer,
|
buf: GenerationBuffer,
|
||||||
history: list[dict],
|
history: list[dict],
|
||||||
@@ -467,7 +502,7 @@ async def run_generation(
|
|||||||
|
|
||||||
if use_tool_path:
|
if use_tool_path:
|
||||||
# Declined write tool — the optimistic stream was cancelled, start fresh.
|
# 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:
|
else:
|
||||||
stream_source = _drain_queue(prefetched, stream_queue)
|
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..."})
|
buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
|
||||||
t_stream = time.monotonic()
|
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():
|
if buf.cancel_event.is_set():
|
||||||
cancelled = True
|
cancelled = True
|
||||||
break
|
break
|
||||||
|
|||||||
Reference in New Issue
Block a user