Wraps a stream generator with resiliency logic. Handles QuotaExceededError by triggering auth rotation and retrying.
(
req_id: str,
model_name: str,
generator_factory: Callable[[Event], AsyncGenerator[str, None]],
completion_event: Event,
)
| 46 | |
| 47 | |
| 48 | async def resilient_stream_generator( |
| 49 | req_id: str, |
| 50 | model_name: str, |
| 51 | generator_factory: Callable[[Event], AsyncGenerator[str, None]], |
| 52 | completion_event: Event, |
| 53 | ) -> AsyncGenerator[str, None]: |
| 54 | """ |
| 55 | Wraps a stream generator with resiliency logic. |
| 56 | Handles QuotaExceededError by triggering auth rotation and retrying. |
| 57 | """ |
| 58 | from api_utils.server_state import state |
| 59 | |
| 60 | logger = state.logger |
| 61 | from browser_utils.auth_rotation import perform_auth_rotation |
| 62 | |
| 63 | max_retries = 3 |
| 64 | retry_count = 0 |
| 65 | |
| 66 | inner_event = Event() |
| 67 | |
| 68 | try: |
| 69 | while retry_count <= max_retries: |
| 70 | try: |
| 71 | if inner_event.is_set(): |
| 72 | inner_event.clear() |
| 73 | |
| 74 | async for chunk in generator_factory(inner_event): |
| 75 | yield chunk |
| 76 | |
| 77 | return |
| 78 | |
| 79 | except (QuotaExceededError, QuotaExceededRetry) as e: |
| 80 | retry_count += 1 |
| 81 | if retry_count > max_retries: |
| 82 | logger.error( |
| 83 | f"[{req_id}] Max retries ({max_retries}) exhausted for quota recovery." |
| 84 | ) |
| 85 | yield f"data: {json.dumps({'error': 'Max retries exhausted for quota recovery.'}, ensure_ascii=False)}\n\n" |
| 86 | return |
| 87 | |
| 88 | logger.warning( |
| 89 | f"[{req_id}] Quota limit hit during stream: {str(e)}. Initiating rotation (Attempt {retry_count}/{max_retries})..." |
| 90 | ) |
| 91 | yield f": processing auth rotation (attempt {retry_count})...\n\n" |
| 92 | |
| 93 | rotation_task = asyncio.create_task( |
| 94 | perform_auth_rotation(target_model_id=model_name) |
| 95 | ) |
| 96 | |
| 97 | rotation_start = time.time() |
| 98 | while not rotation_task.done(): |
| 99 | if time.time() - rotation_start > 120: |
| 100 | logger.error(f"[{req_id}] Rotation timed out.") |
| 101 | yield f"data: {json.dumps({'error': 'Auth rotation timed out.'}, ensure_ascii=False)}\n\n" |
| 102 | return |
| 103 | |
| 104 | yield ": processing auth rotation...\n\n" |
| 105 | await asyncio.sleep(2) |