Spaces:
Running on Zero
Running on Zero
| """ | |
| UI-agnostic retry / failure-state logic. | |
| `stream_with_retries` turns a stream (or its failure) into a sequence of | |
| Events. A UI only has to render events; it never catches exceptions. | |
| """ | |
| import math | |
| import random | |
| import time | |
| from dataclasses import dataclass | |
| from typing import Callable, Iterator, Optional | |
| from providers import (ProviderError, RateLimitError, RefusalError, | |
| RequestTimeoutError, StallError, with_watchdog) | |
| class Event: | |
| kind: str # start | chunk | retry_wait | done | stalled | refused | error | |
| text: str = "" # accumulated response text so far | |
| message: str = "" | |
| seconds: int = 0 # for retry_wait: seconds remaining | |
| attempt: int = 0 | |
| status_code: Optional[int] = None | |
| def backoff_seconds(attempt: int, retry_after: Optional[float] = None, | |
| base: float = 2.0, cap: float = 30.0) -> int: | |
| """Honor the server's Retry-After if present; otherwise exponential | |
| backoff with jitter (so many clients don't retry in lockstep).""" | |
| if retry_after: | |
| return math.ceil(retry_after) | |
| return math.ceil(min(cap, base * 2 ** (attempt - 1)) * random.uniform(0.75, 1.25)) | |
| def stream_with_retries(make_stream: Callable[[], Iterator[str]], *, | |
| max_retries: int = 3, | |
| stall_seconds: Optional[float] = None, | |
| sleep: Callable[[float], None] = time.sleep) -> Iterator[Event]: | |
| attempt = 0 | |
| while True: | |
| attempt += 1 | |
| text = "" | |
| yield Event("start", attempt=attempt) | |
| try: | |
| stream = make_stream() | |
| if stall_seconds: | |
| stream = with_watchdog(stream, stall_seconds) | |
| for chunk in stream: | |
| text += chunk | |
| yield Event("chunk", text=text, attempt=attempt) | |
| yield Event("done", text=text, attempt=attempt) | |
| return | |
| except RateLimitError as e: | |
| if attempt > max_retries: | |
| yield Event("error", text=text, attempt=attempt, status_code=429, | |
| message=f"Still rate limited after {attempt} attempts. Giving up.") | |
| return | |
| wait = backoff_seconds(attempt, e.retry_after) | |
| source = "Retry-After header" if e.retry_after else "exponential backoff + jitter" | |
| yield Event("retry_wait", text=text, attempt=attempt, seconds=wait, status_code=429, | |
| message=f"429 on attempt {attempt}. Waiting {wait}s ({source}).") | |
| for left in range(wait, 0, -1): | |
| yield Event("retry_wait", text=text, attempt=attempt, seconds=left, | |
| status_code=429, message="") | |
| sleep(1) | |
| except StallError as e: | |
| yield Event("stalled", text=text, attempt=attempt, message=str(e)) | |
| return | |
| except RefusalError as e: | |
| yield Event("refused", text=text, attempt=attempt, message=str(e)) | |
| return | |
| except RequestTimeoutError as e: | |
| yield Event("error", text=text, attempt=attempt, | |
| message=f"No response from the API: {e}") | |
| return | |
| except ProviderError as e: | |
| yield Event("error", text=text, attempt=attempt, | |
| status_code=e.status_code, message=str(e)) | |
| return | |