Mirrowel commited on
Commit
6aa3124
·
1 Parent(s): c8b2334

fix(rotator_library): ensure classify_error receives unwrapped exception

Browse files

The error classification logic was inconsistently unwrapping StreamedAPIError before passing to classify_error. This fix aligns both call sites to extract the original exception from the wrapper's 'data' attribute, preventing misclassification of streamed errors.

Files changed (1) hide show
  1. src/rotator_library/client.py +4 -3
src/rotator_library/client.py CHANGED
@@ -675,7 +675,9 @@ class RotatingClient:
675
  raise e
676
 
677
  last_exception = e
678
- classified_error = classify_error(e)
 
 
679
  await self.usage_manager.record_failure(current_cred, model, classified_error)
680
  lib_logger.warning(f"Credential ...{current_cred[-6:]} encountered a recoverable error ({classified_error.error_type}) during custom provider stream. Rotating key.")
681
  break
@@ -777,13 +779,12 @@ class RotatingClient:
777
  except (StreamedAPIError, litellm.RateLimitError) as e:
778
  last_exception = e
779
 
780
- classified_error = classify_error(e)
781
-
782
  # This is the final, robust handler for streamed errors.
783
  error_payload = {}
784
  cleaned_str = None
785
  # The actual exception might be wrapped in our StreamedAPIError.
786
  original_exc = getattr(e, 'data', e)
 
787
 
788
  try:
789
  # The full error JSON is in the string representation of the exception.
 
675
  raise e
676
 
677
  last_exception = e
678
+ # If the exception is our custom wrapper, unwrap the original error
679
+ original_exc = getattr(e, 'data', e)
680
+ classified_error = classify_error(original_exc)
681
  await self.usage_manager.record_failure(current_cred, model, classified_error)
682
  lib_logger.warning(f"Credential ...{current_cred[-6:]} encountered a recoverable error ({classified_error.error_type}) during custom provider stream. Rotating key.")
683
  break
 
779
  except (StreamedAPIError, litellm.RateLimitError) as e:
780
  last_exception = e
781
 
 
 
782
  # This is the final, robust handler for streamed errors.
783
  error_payload = {}
784
  cleaned_str = None
785
  # The actual exception might be wrapped in our StreamedAPIError.
786
  original_exc = getattr(e, 'data', e)
787
+ classified_error = classify_error(original_exc)
788
 
789
  try:
790
  # The full error JSON is in the string representation of the exception.