Skip to content

Commit 518b26a

Browse files
committed
feat: implement retry logic for file uploads and introduce ValidationError for client input validation.
1 parent 4228bfb commit 518b26a

12 files changed

Lines changed: 500 additions & 126 deletions

File tree

src/knowhere/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
ConflictError,
2323
GatewayTimeoutError,
2424
InternalServerError,
25+
InvalidStateError,
2526
JobFailedError,
2627
KnowhereError,
2728
NotFoundError,
@@ -30,6 +31,7 @@
3031
PollingTimeoutError,
3132
RateLimitError,
3233
ServiceUnavailableError,
34+
ValidationError,
3335
)
3436
from knowhere._types import PollProgressCallback, UploadProgressCallback
3537
from knowhere._version import __version__
@@ -58,6 +60,8 @@
5860
"__version__",
5961
# Exceptions
6062
"KnowhereError",
63+
"ValidationError",
64+
"InvalidStateError",
6165
"APIConnectionError",
6266
"APITimeoutError",
6367
"APIStatusError",

src/knowhere/_base_client.py

Lines changed: 96 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from knowhere._exceptions import (
2626
APIConnectionError,
2727
APITimeoutError,
28+
ValidationError,
2829
makeStatusError,
2930
)
3031
from knowhere._logging import getLogger, redactSensitiveHeaders
@@ -35,17 +36,23 @@
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

5158
class 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",

src/knowhere/_client.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
from knowhere._base_client import AsyncAPIClient, SyncAPIClient
1515
from knowhere._constants import DEFAULT_POLL_INTERVAL, DEFAULT_POLL_TIMEOUT
16+
from knowhere._exceptions import ValidationError
1617
from knowhere._logging import getLogger
1718
from knowhere._types import (
1819
PollProgressCallback,
@@ -94,9 +95,9 @@ def parse(
9495
Provide exactly one of *url* or *file*.
9596
"""
9697
if url and file:
97-
raise ValueError("Provide either 'url' or 'file', not both.")
98+
raise ValidationError("Provide either 'url' or 'file', not both.")
9899
if not url and file is None:
99-
raise ValueError("Provide either 'url' or 'file'.")
100+
raise ValidationError("Provide either 'url' or 'file'.")
100101

101102
# Determine source type and create job
102103
if url:
@@ -196,9 +197,9 @@ async def parse(
196197
) -> ParseResult:
197198
"""Parse a document end-to-end (async version)."""
198199
if url and file:
199-
raise ValueError("Provide either 'url' or 'file', not both.")
200+
raise ValidationError("Provide either 'url' or 'file', not both.")
200201
if not url and file is None:
201-
raise ValueError("Provide either 'url' or 'file'.")
202+
raise ValidationError("Provide either 'url' or 'file'.")
202203

203204
if url:
204205
job: Job = await self.jobs.create(

src/knowhere/_constants.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
# Retry configuration
2020
DEFAULT_MAX_RETRIES: int = 5
21+
DEFAULT_UPLOAD_MAX_RETRIES: int = 2
2122

2223
# Polling configuration
2324
MAX_POLL_INTERVAL: float = 30.0

0 commit comments

Comments
 (0)