feat: Abort API retries on client disconnection

Pass the FastAPI `request` object to the rotating client's `acompletion` and `aembedding` methods.
Add a check within the retry loops to detect client disconnections. If the client disconnects, immediately abort further retries to prevent unnecessary resource usage and provide faster feedback.

src/proxy_app/main.py

@@ -239,7 +239,7 @@ async def chat_completions(
         request_data = await request.json()
         is_streaming = request_data.get("stream", False)
 
-        response = await client.acompletion(**request_data)
+        response = await client.acompletion(request=request, **request_data)
 
         if is_streaming:
             # Wrap the streaming response to enable logging after it's complete
@@ -339,7 +339,7 @@ async def embeddings(
             if isinstance(request_data.get("input"), str):
                 request_data["input"] = [request_data["input"]]
             
-            response = await client.aembedding(**request_data)
+            response = await client.aembedding(request=request, **request_data)
 
         if ENABLE_REQUEST_LOGGING:
             response_summary = {

src/rotator_library/client.py

@@ -176,7 +176,7 @@ class RotatingClient:
                 yield "data: [DONE]\n\n"
 
 
-    async def acompletion(self, pre_request_callback: Optional[callable] = None, **kwargs) -> Union[Any, AsyncGenerator[str, None]]:
+    async def acompletion(self, request: Optional[Any] = None, pre_request_callback: Optional[callable] = None, **kwargs) -> Union[Any, AsyncGenerator[str, None]]:
         """
         Performs a completion call with smart key rotation and retry logic.
         It will try each available key in sequence if the previous one fails.
@@ -275,11 +275,17 @@ class RotatingClient:
                         
                         classified_error = classify_error(e)
 
+                        classified_error = classify_error(e)
+
                         if classified_error.error_type in ['invalid_request', 'context_window_exceeded']:
                             # These errors are not recoverable by rotating keys, so fail fast.
                             lib_logger.error(f"Unrecoverable error '{classified_error.error_type}' with key ...{current_key[-4:]}. Failing request.")
                             raise last_exception
 
+                        if request and await request.is_disconnected():
+                            lib_logger.warning(f"Client disconnected. Aborting retries for key ...{current_key[-4:]}.")
+                            raise last_exception
+
                         if classified_error.error_type in ['server_error', 'api_connection']:
                             # These are temporary, so record the failure and retry with backoff.
                             await self.usage_manager.record_failure(current_key, model, classified_error)
@@ -310,7 +316,7 @@ class RotatingClient:
         
         raise Exception("Failed to complete the request: No available API keys for the provider or all keys failed.")
 
-    async def aembedding(self, **kwargs) -> Any:
+    async def aembedding(self, request: Optional[Any] = None, **kwargs) -> Any:
         """
         Performs an embedding call with smart key rotation and retry logic.
         """
@@ -366,6 +372,10 @@ class RotatingClient:
                             lib_logger.error(f"Unrecoverable error '{classified_error.error_type}' with key ...{current_key[-4:]}. Failing request.")
                             raise last_exception
 
+                        if request and await request.is_disconnected():
+                            lib_logger.warning(f"Client disconnected during embedding. Aborting retries for key ...{current_key[-4:]}.")
+                            raise last_exception
+
                         if classified_error.error_type in ['server_error', 'api_connection']:
                             await self.usage_manager.record_failure(current_key, model, classified_error)