Pickle exceptions with custom reduce

Author: srilman

src/openai/_exceptions.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Optional, cast
-from typing_extensions import Literal
+from typing_extensions import Literal, Self
 
 import httpx
 
@@ -66,6 +66,13 @@ def __init__(self, message: str, request: httpx.Request, *, body: object | None)
             self.param = None
             self.type = None
 
+    @classmethod
+    def _reconstruct(cls, message: str, request: httpx.Request, body: object | None) -> Self:
+        return cls(message, request, body=body)
+
+    def __reduce__(self) -> tuple[type[APIError], tuple[str, httpx.Request, object | None]]:
+        return (self.__class__._reconstruct, (self.message, self.request, self.body))
+
 
 class APIResponseValidationError(APIError):
     response: httpx.Response
@@ -76,6 +83,13 @@ def __init__(self, response: httpx.Response, body: object | None, *, message: st
         self.response = response
         self.status_code = response.status_code
 
+    @classmethod
+    def _reconstruct(cls, response: httpx.Response, body: object | None, message: str | None) -> Self:
+        return cls(response, body, message=message)
+
+    def __reduce__(self) -> tuple[type[APIResponseValidationError], tuple[httpx.Response, object | None, str | None]]:
+        return (self.__class__._reconstruct, (self.response, self.body, self.message))
+
 
 class APIStatusError(APIError):
     """Raised when an API response has a status code of 4xx or 5xx."""
@@ -90,11 +104,25 @@ def __init__(self, message: str, *, response: httpx.Response, body: object | Non
         self.status_code = response.status_code
         self.request_id = response.headers.get("x-request-id")
 
+    @classmethod
+    def _reconstruct(cls, message: str, response: httpx.Response, body: object | None) -> Self:
+        return cls(message, response=response, body=body)
+
+    def __reduce__(self) -> tuple[type[APIStatusError], tuple[str, httpx.Response, object | None]]:
+        return (self.__class__._reconstruct, (self.message, self.response, self.body))
+
 
 class APIConnectionError(APIError):
     def __init__(self, *, message: str = "Connection error.", request: httpx.Request) -> None:
         super().__init__(message, request, body=None)
 
+    @classmethod
+    def _reconstruct(cls, message: str, request: httpx.Request) -> Self:
+        return cls(message=message, request=request)
+
+    def __reduce__(self) -> tuple[type[APIConnectionError], tuple[str, httpx.Request]]:
+        return (self.__class__._reconstruct, (self.message, self.request))
+
 
 class APITimeoutError(APIConnectionError):
     def __init__(self, request: httpx.Request) -> None:
@@ -149,6 +177,13 @@ def __init__(self, *, completion: ChatCompletion) -> None:
         super().__init__(msg)
         self.completion = completion
 
+    @classmethod
+    def _reconstruct(cls, completion: ChatCompletion) -> Self:
+        return cls(completion=completion)
+
+    def __reduce__(self) -> tuple[type[LengthFinishReasonError], tuple[ChatCompletion]]:
+        return (self.__class__._reconstruct, (self.completion,))
+
 
 class ContentFilterFinishReasonError(OpenAIError):
     def __init__(self) -> None: