feat: Add retry logic for transient LiteLLM errors
linuztx committed
Sep 3, 2025 at 20:28 UTC
143678d9d201cd1e069bc808eceb1dee2ad02d85
1 file changed
+84
-39
models.py
+84
-39
@@ -16,6 +16,7 @@ from typing import (
16
17
from litellm import completion, acompletion, embedding
18
import litellm
19
+import openai
20
21
from python.helpers import dotenv
22
from python.helpers import settings
@@ -220,6 +221,31 @@ def get_rate_limiter(
221
return limiter
222
223
224
+def _is_transient_litellm_error(exc: Exception) -> bool:
225
+ """Uses status_code when available, else falls back to exception types"""
226
+ # Prefer explicit status codes if present
227
+ status_code = getattr(exc, "status_code", None)
228
+ if isinstance(status_code, int):
229
+ if status_code in (408, 429, 500, 502, 503, 504):
230
+ return True
231
+ # Treat other 5xx as retriable
232
+ if status_code >= 500:
233
+ return True
234
+ return False
235
+
236
+ # Fallback to exception classes mapped by LiteLLM/OpenAI
237
+ transient_types = (
238
+ getattr(openai, "APITimeoutError", Exception),
239
+ getattr(openai, "APIConnectionError", Exception),
240
+ getattr(openai, "RateLimitError", Exception),
241
+ getattr(openai, "APIError", Exception),
242
+ getattr(openai, "InternalServerError", Exception),
243
+ # Some providers map overloads to ServiceUnavailable-like errors
244
+ getattr(openai, "APIStatusError", Exception),
245
+ )
246
+ return isinstance(exc, transient_types)
247
+
248
+
249
async def apply_rate_limiter(
250
model_config: ModelConfig | None,
251
input_text: str,
@@ -454,50 +480,69 @@ class LiteLLMChatWrapper(SimpleChatModel):
480
self.a0_model_conf, str(msgs_conv), rate_limiter_callback
481
)
482
457
- # call model
458
- _completion = await acompletion(
459
- model=self.model_name,
460
- messages=msgs_conv,
461
- stream=True,
462
- **{**self.kwargs, **kwargs},
463
- )
483
+ # Prepare call kwargs and retry config (strip A0-only params before calling LiteLLM)
484
+ call_kwargs: dict[str, Any] = {**self.kwargs, **kwargs}
485
+ max_retries: int = int(call_kwargs.pop("a0_retry_attempts", 2))
486
+ retry_delay_s: float = float(call_kwargs.pop("a0_retry_delay_seconds", 1.5))
487
488
# results
489
result = ChatGenerationResult()
490
468
- # iterate over chunks
469
- async for chunk in _completion: # type: ignore
470
- # parse chunk
471
- parsed = _parse_chunk(chunk)
472
- output = result.add_chunk(parsed)
473
-
474
- # collect reasoning delta and call callbacks
475
- if output["reasoning_delta"]:
476
- if reasoning_callback:
477
- await reasoning_callback(output["reasoning_delta"], result.reasoning)
478
- if tokens_callback:
479
- await tokens_callback(
480
- output["reasoning_delta"],
481
- approximate_tokens(output["reasoning_delta"]),
482
- )
483
- # Add output tokens to rate limiter if configured
484
- if limiter:
485
- limiter.add(output=approximate_tokens(output["reasoning_delta"]))
486
- # collect response delta and call callbacks
487
- if output["response_delta"]:
488
- if response_callback:
489
- await response_callback(output["response_delta"], result.response)
490
- if tokens_callback:
491
- await tokens_callback(
492
- output["response_delta"],
493
- approximate_tokens(output["response_delta"]),
494
- )
495
- # Add output tokens to rate limiter if configured
496
- if limiter:
497
- limiter.add(output=approximate_tokens(output["response_delta"]))
491
+ attempt = 0
492
+ while True:
493
+ got_any_chunk = False
494
+ try:
495
+ # call model
496
+ _completion = await acompletion(
497
+ model=self.model_name,
498
+ messages=msgs_conv,
499
+ stream=True,
500
+ **call_kwargs,
501
+ )
502
499
- # return complete results
500
- return result.response, result.reasoning
503
+ # iterate over chunks
504
+ async for chunk in _completion: # type: ignore
505
+ got_any_chunk = True
506
+ # parse chunk
507
+ parsed = _parse_chunk(chunk)
508
+ output = result.add_chunk(parsed)
509
+
510
+ # collect reasoning delta and call callbacks
511
+ if output["reasoning_delta"]:
512
+ if reasoning_callback:
513
+ await reasoning_callback(output["reasoning_delta"], result.reasoning)
514
+ if tokens_callback:
515
+ await tokens_callback(
516
+ output["reasoning_delta"],
517
+ approximate_tokens(output["reasoning_delta"]),
518
+ )
519
+ # Add output tokens to rate limiter if configured
520
+ if limiter:
521
+ limiter.add(output=approximate_tokens(output["reasoning_delta"]))
522
+ # collect response delta and call callbacks
523
+ if output["response_delta"]:
524
+ if response_callback:
525
+ await response_callback(output["response_delta"], result.response)
526
+ if tokens_callback:
527
+ await tokens_callback(
528
+ output["response_delta"],
529
+ approximate_tokens(output["response_delta"]),
530
+ )
531
+ # Add output tokens to rate limiter if configured
532
+ if limiter:
533
+ limiter.add(output=approximate_tokens(output["response_delta"]))
534
+
535
+ # Successful completion of stream
536
+ return result.response, result.reasoning
537
+
538
+ except Exception as e:
539
+ import asyncio
540
+
541
+ # Retry only if no chunks received and error is transient
542
+ if got_any_chunk or not _is_transient_litellm_error(e) or attempt >= max_retries:
543
+ raise
544
+ attempt += 1
545
+ await asyncio.sleep(retry_delay_s)
546
547
548
class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):