2525from knowhere ._exceptions import (
2626 APIConnectionError ,
2727 APITimeoutError ,
28+ ValidationError ,
2829 makeStatusError ,
2930)
3031from knowhere ._logging import getLogger , redactSensitiveHeaders
3536
3637_logger = getLogger ()
3738
38- # Error codes that are safe to retry
39- _RETRYABLE_ERROR_CODES : frozenset [str ] = frozenset ({
40- "rate_limit_exceeded" ,
41- "service_unavailable" ,
42- "gateway_timeout" ,
43- "internal_server_error" ,
44- "timeout" ,
39+ # Error codes that are always safe to retry (matches server ALWAYS_RETRYABLE_ERROR_CODES)
40+ _ALWAYS_RETRYABLE_ERROR_CODES : frozenset [str ] = frozenset ({
41+ "ABORTED" , # 409 - Concurrency conflict
42+ "UNAVAILABLE" , # 503 - Service temporarily down
43+ "DEADLINE_EXCEEDED" , # 504 - Timeout
4544})
4645
47- # Status codes that are safe to retry
48- _RETRYABLE_STATUS_CODES : frozenset [int ] = frozenset ({408 , 429 , 500 , 502 , 503 , 504 })
46+ # RESOURCE_EXHAUSTED (429) is conditionally retryable:
47+ # - Rate limit: details.retry_after present → RETRY
48+ # - Quota exceeded: no retry_after → DO NOT RETRY
49+ _CONDITIONALLY_RETRYABLE_ERROR_CODE : str = "RESOURCE_EXHAUSTED"
50+
51+ # HTTP status codes that are always safe to retry
52+ _ALWAYS_RETRYABLE_STATUS_CODES : frozenset [int ] = frozenset ({409 , 502 , 503 , 504 })
53+
54+ # HTTP status code that is conditionally retryable (only with retry_after)
55+ _CONDITIONALLY_RETRYABLE_STATUS_CODE : int = 429
4956
5057
5158class BaseClient :
@@ -71,7 +78,7 @@ def __init__(
7178 # Resolve: arg > env > default
7279 resolved_key : Optional [str ] = api_key or os .environ .get (ENV_API_KEY )
7380 if not resolved_key :
74- raise ValueError (
81+ raise ValidationError ( # TODO
7582 "An API key must be provided via the 'api_key' argument "
7683 f"or the { ENV_API_KEY } environment variable."
7784 )
@@ -122,12 +129,68 @@ def _shouldRetry(
122129 self ,
123130 status_code : int ,
124131 error_code : Optional [str ] = None ,
125- details : Optional [Any ] = None ,
132+ details : Optional [Dict [ str , Any ] ] = None ,
126133 ) -> bool :
127- """Decide whether a request should be retried."""
128- if error_code and error_code in _RETRYABLE_ERROR_CODES :
134+ """Decide whether a request should be retried.
135+
136+ Follows server-side retry semantics:
137+ - ABORTED, UNAVAILABLE, DEADLINE_EXCEEDED → always retry
138+ - RESOURCE_EXHAUSTED (429) → retry only if details.retry_after present
139+ - All other errors → never retry
140+ """
141+ if error_code :
142+ if error_code in _ALWAYS_RETRYABLE_ERROR_CODES :
143+ return True
144+ if error_code == _CONDITIONALLY_RETRYABLE_ERROR_CODE :
145+ return self ._hasRetryAfter (details )
146+ return False
147+
148+ # Fallback to status code when error_code is unavailable
149+ if status_code in _ALWAYS_RETRYABLE_STATUS_CODES :
129150 return True
130- return status_code in _RETRYABLE_STATUS_CODES
151+ if status_code == _CONDITIONALLY_RETRYABLE_STATUS_CODE :
152+ return self ._hasRetryAfter (details )
153+ return False
154+
155+ @staticmethod
156+ def _hasRetryAfter (details : Optional [Dict [str , Any ]]) -> bool :
157+ """Check if details contains a retry_after hint."""
158+ if not isinstance (details , dict ):
159+ return False
160+ retry_after : Any = details .get ("retry_after" )
161+ return retry_after is not None
162+
163+ @staticmethod
164+ def _extractRetryAfter (
165+ error_body : Optional [Dict [str , Any ]],
166+ response : httpx .Response ,
167+ ) -> Optional [float ]:
168+ """Extract retry_after from the response body or Retry-After header.
169+
170+ The server puts retry_after in ``error.details.retry_after``.
171+ Falls back to the HTTP ``Retry-After`` header.
172+ """
173+ # Prefer body: error.details.retry_after
174+ if isinstance (error_body , dict ):
175+ err_obj : Any = error_body .get ("error" , error_body )
176+ if isinstance (err_obj , dict ):
177+ details : Any = err_obj .get ("details" )
178+ if isinstance (details , dict ):
179+ raw : Any = details .get ("retry_after" )
180+ if raw is not None :
181+ try :
182+ return float (raw )
183+ except (ValueError , TypeError ):
184+ pass
185+
186+ # Fallback: HTTP Retry-After header
187+ header_raw : Optional [str ] = response .headers .get ("retry-after" )
188+ if header_raw is not None :
189+ try :
190+ return float (header_raw )
191+ except (ValueError , TypeError ):
192+ pass
193+ return None
131194
132195 def _calculateRetryDelay (
133196 self ,
@@ -257,24 +320,24 @@ def _request(
257320 response
258321 )
259322 error_code : Optional [str ] = None
323+ error_details : Optional [Dict [str , Any ]] = None
260324 if isinstance (error_body , dict ):
261325 err_obj : Any = error_body .get ("error" , error_body )
262326 if isinstance (err_obj , dict ):
263327 error_code = err_obj .get ("code" )
328+ raw_details : Any = err_obj .get ("details" )
329+ if isinstance (raw_details , dict ):
330+ error_details = raw_details
264331
265332 if (
266333 attempt < self .max_retries
267- and self ._shouldRetry (response .status_code , error_code )
334+ and self ._shouldRetry (
335+ response .status_code , error_code , error_details
336+ )
268337 ):
269- retry_after_raw : Optional [str ] = response . headers . get (
270- "retry-after"
338+ retry_after_val : Optional [float ] = self . _extractRetryAfter (
339+ error_body , response
271340 )
272- retry_after_val : Optional [float ] = None
273- if retry_after_raw :
274- try :
275- retry_after_val = float (retry_after_raw )
276- except (ValueError , TypeError ):
277- pass
278341 delay = self ._calculateRetryDelay (attempt , retry_after_val )
279342 _logger .warning (
280343 "Retryable error %d on attempt %d/%d, retrying in %.1fs" ,
@@ -404,22 +467,24 @@ async def _request(
404467
405468 error_body : Optional [Dict [str , Any ]] = self ._parseErrorResponse (response )
406469 error_code : Optional [str ] = None
470+ error_details : Optional [Dict [str , Any ]] = None
407471 if isinstance (error_body , dict ):
408472 err_obj : Any = error_body .get ("error" , error_body )
409473 if isinstance (err_obj , dict ):
410474 error_code = err_obj .get ("code" )
475+ raw_details : Any = err_obj .get ("details" )
476+ if isinstance (raw_details , dict ):
477+ error_details = raw_details
411478
412479 if (
413480 attempt < self .max_retries
414- and self ._shouldRetry (response .status_code , error_code )
481+ and self ._shouldRetry (
482+ response .status_code , error_code , error_details
483+ )
415484 ):
416- retry_after_raw : Optional [str ] = response .headers .get ("retry-after" )
417- retry_after_val : Optional [float ] = None
418- if retry_after_raw :
419- try :
420- retry_after_val = float (retry_after_raw )
421- except (ValueError , TypeError ):
422- pass
485+ retry_after_val : Optional [float ] = self ._extractRetryAfter (
486+ error_body , response
487+ )
423488 delay = self ._calculateRetryDelay (attempt , retry_after_val )
424489 _logger .warning (
425490 "Retryable error %d on attempt %d/%d, retrying in %.1fs" ,
0 commit comments