Add better error message when request times out

njbbaer · Aug 21, 2024 · 0fc92d2 · 0fc92d2
1 parent aadd91f
commit 0fc92d2
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/src/api_client.py b/src/api_client.py
@@ -23,14 +23,17 @@ def create(client_type):
 
     async def request_completion(self, messages, parameters, pricing):
         body = self.prepare_body(messages, parameters)
-        async with httpx.AsyncClient(timeout=self.TIMEOUT) as client:
-            response = await client.post(
-                self.API_URL, headers=self.get_headers(), json=body
-            )
-            response.raise_for_status()
-            completion = self.create_completion(response.json(), pricing)
-            self.logger.log(parameters, messages, completion.content)
-            return completion
+        try:
+            async with httpx.AsyncClient(timeout=self.TIMEOUT) as client:
+                response = await client.post(
+                    self.API_URL, headers=self.get_headers(), json=body
+                )
+                response.raise_for_status()
+                completion = self.create_completion(response.json(), pricing)
+                self.logger.log(parameters, messages, completion.content)
+                return completion
+        except httpx.ReadTimeout:
+            raise Exception("Request timed out")
 
     @abstractmethod
     def get_headers(self):