From b65b1bbc52bea31d7012de30684a0a47261fd03f Mon Sep 17 00:00:00 2001 From: Aaron Pham <29749331+aarnphm@users.noreply.github.com> Date: Fri, 10 Nov 2023 17:50:12 -0500 Subject: [PATCH] fix(client): check for should retry header (#606) Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com> --- openllm-client/src/openllm_client/_shim.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/openllm-client/src/openllm_client/_shim.py b/openllm-client/src/openllm_client/_shim.py index fcb8b7de..b5fe92f2 100644 --- a/openllm-client/src/openllm_client/_shim.py +++ b/openllm-client/src/openllm_client/_shim.py @@ -313,10 +313,11 @@ class BaseClient(t.Generic[InnerClient, StreamType]): def _should_retry(self, response: httpx.Response) -> bool: should_retry_header = response.headers.get('x-should-retry') - if should_retry_header.lower() == 'true': - return True - if should_retry_header.lower() == 'false': - return False + if should_retry_header: + if should_retry_header.lower() == 'true': + return True + if should_retry_header.lower() == 'false': + return False if response.status_code in {408, 409, 429}: return True if response.status_code >= 500: