Skip to content

Commit 63a1a1a

Browse files
committed
fix: Detect and handle watsonx.ai model provider rate limit exceptions
Issue - #1599 Summary - Added async retry logic with Full Jitter exponential backoff for all watsonx.ai HTTP calls to handle HTTP 429 rate limit responses gracefully - Added a short-TTL in-process cache for provider health check responses to prevent concurrent identical watsonx.ai validation round-trips from amplifying rate limit exposure New Utilities - Added `src/utils/watsonx_retry.py`: async `request_with_retry` helper that retries only on HTTP 429, respects `Retry-After` headers (integer seconds and HTTP-date formats), applies Full Jitter backoff, and enforces both per-attempt and total wall-clock budgets - Added `src/utils/provider_health_cache.py`: TTL-based in-process cache (default 10s, configurable via `OPENRAG_PROVIDER_HEALTH_TTL`) keyed on a SHA-256 fingerprint of provider config; API keys are hashed and never stored in plaintext Retry Integration - Replaced direct `client.post` / `client.get` calls with `request_with_retry` in `src/api/provider_validation.py` for all three watsonx validation paths (IAM token, completion with tools, embedding) - Applied the same replacement in `src/services/models_service.py` for IAM token fetch and both model-listing requests Provider Health Cache - Updated `src/api/provider_health.py` to short-circuit repeated polled health checks using the new cache; only successful 200 responses on the default path are cached — explicit `?provider=` overrides and 503 error paths bypass the cache Tests - Added `tests/unit/utils/test_watsonx_retry.py` covering: immediate return on non-429, retry-then-succeed, `Retry-After` integer/HTTP-date header handling, cap clamping, non-retried status codes, max-attempt exhaustion, and total wall-clock budget short-circuit - Added `tests/unit/utils/test_provider_health_cache.py` covering: cache miss, set/get round-trip, invalidation, key uniqueness per field, key stability, API key non-leakage in plaintext, and TTL expiry
1 parent c2c10ca commit 63a1a1a

8 files changed

Lines changed: 507 additions & 23 deletions

File tree

src/api/provider_health.py

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from fastapi import Depends
77
from fastapi.responses import JSONResponse
88
from utils.logging_config import get_logger
9+
from utils import provider_health_cache
910
from config.settings import get_openrag_config
1011
from api.provider_validation import validate_provider_setup
1112
from dependencies import get_current_user
@@ -94,7 +95,29 @@ async def check_provider_health(
9495
embedding_endpoint = getattr(embedding_provider_config, "endpoint", None)
9596
embedding_project_id = getattr(embedding_provider_config, "project_id", None)
9697
embedding_model = current_config.knowledge.embedding_model
97-
98+
99+
# Short-circuit identical concurrent polls from the provider-health
100+
# banner so we don't fan out N watsonx round-trips per poll cycle.
101+
# Only the polled (no `check_provider`) success path is cached; the
102+
# 503 branch and the specific-provider branch always re-validate.
103+
health_cache_key = provider_health_cache.cache_key(
104+
provider=provider,
105+
embedding_provider=embedding_provider,
106+
test_completion=test_completion,
107+
llm_model=llm_model,
108+
embedding_model=embedding_model,
109+
endpoint=endpoint,
110+
project_id=project_id,
111+
api_key=api_key,
112+
embedding_api_key=embedding_api_key,
113+
embedding_endpoint=embedding_endpoint,
114+
embedding_project_id=embedding_project_id,
115+
)
116+
cached_payload = provider_health_cache.get(health_cache_key)
117+
if cached_payload is not None:
118+
logger.debug("Returning cached provider-health response")
119+
return JSONResponse(cached_payload, status_code=200)
120+
98121
logger.info(f"Checking health for provider: {provider}")
99122

100123
# Validate provider setup
@@ -204,19 +227,18 @@ async def check_provider_health(
204227
status_code=503,
205228
)
206229

207-
return JSONResponse(
208-
{
209-
"status": "healthy",
210-
"message": "Both providers properly configured and validated",
211-
"llm_provider": provider,
212-
"embedding_provider": embedding_provider,
213-
"details": {
214-
"llm_model": llm_model,
215-
"embedding_model": embedding_model,
216-
},
230+
healthy_payload = {
231+
"status": "healthy",
232+
"message": "Both providers properly configured and validated",
233+
"llm_provider": provider,
234+
"embedding_provider": embedding_provider,
235+
"details": {
236+
"llm_model": llm_model,
237+
"embedding_model": embedding_model,
217238
},
218-
status_code=200,
219-
)
239+
}
240+
provider_health_cache.set_(health_cache_key, healthy_payload)
241+
return JSONResponse(healthy_payload, status_code=200)
220242

221243
except Exception as e:
222244
error_message = str(e)

src/api/provider_validation.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import httpx
55
from utils.container_utils import transform_localhost_url
66
from utils.logging_config import get_logger
7+
from utils.watsonx_retry import request_with_retry
78

89
logger = get_logger(__name__)
910
def _parse_json_error_message(error_text: str) -> str:
@@ -392,7 +393,9 @@ async def _test_watsonx_lightweight_health(
392393
try:
393394
# Get bearer token from IBM IAM - this validates the API key without consuming credits
394395
async with httpx.AsyncClient() as client:
395-
token_response = await client.post(
396+
token_response = await request_with_retry(
397+
client,
398+
"POST",
396399
"https://iam.cloud.ibm.com/identity/token",
397400
headers={"Content-Type": "application/x-www-form-urlencoded"},
398401
data={
@@ -428,7 +431,9 @@ async def _test_watsonx_completion_with_tools(
428431
try:
429432
# Get bearer token from IBM IAM
430433
async with httpx.AsyncClient() as client:
431-
token_response = await client.post(
434+
token_response = await request_with_retry(
435+
client,
436+
"POST",
432437
"https://iam.cloud.ibm.com/identity/token",
433438
headers={"Content-Type": "application/x-www-form-urlencoded"},
434439
data={
@@ -484,7 +489,9 @@ async def _test_watsonx_completion_with_tools(
484489
}
485490

486491
async with httpx.AsyncClient() as client:
487-
response = await client.post(
492+
response = await request_with_retry(
493+
client,
494+
"POST",
488495
url,
489496
headers=headers,
490497
params=params,
@@ -523,7 +530,9 @@ async def _test_watsonx_embedding(
523530
try:
524531
# Get bearer token from IBM IAM
525532
async with httpx.AsyncClient() as client:
526-
token_response = await client.post(
533+
token_response = await request_with_retry(
534+
client,
535+
"POST",
527536
"https://iam.cloud.ibm.com/identity/token",
528537
headers={"Content-Type": "application/x-www-form-urlencoded"},
529538
data={
@@ -557,7 +566,9 @@ async def _test_watsonx_embedding(
557566
}
558567

559568
async with httpx.AsyncClient() as client:
560-
response = await client.post(
569+
response = await request_with_retry(
570+
client,
571+
"POST",
561572
url,
562573
headers=headers,
563574
params=params,

src/services/models_service.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from config.embedding_constants import OPENAI_DEFAULT_EMBEDDING_MODEL, OPENAI_EMBEDDING_MODEL_PREFIX
1212
from utils.container_utils import transform_localhost_url
1313
from utils.logging_config import get_logger
14+
from utils.watsonx_retry import request_with_retry
1415

1516
logger = get_logger(__name__)
1617

@@ -456,7 +457,9 @@ async def get_ibm_models(
456457
bearer_token = None
457458
if api_key:
458459
async with httpx.AsyncClient() as client:
459-
token_response = await client.post(
460+
token_response = await request_with_retry(
461+
client,
462+
"POST",
460463
"https://iam.cloud.ibm.com/identity/token",
461464
headers={"Content-Type": "application/x-www-form-urlencoded"},
462465
data={
@@ -501,8 +504,8 @@ async def get_ibm_models(
501504
if project_id:
502505
text_params["project_id"] = project_id
503506

504-
text_response = await client.get(
505-
models_url, params=text_params, headers=headers, timeout=10.0
507+
text_response = await request_with_retry(
508+
client, "GET", models_url, params=text_params, headers=headers, timeout=10.0
506509
)
507510

508511
if text_response.status_code == 200:
@@ -538,8 +541,8 @@ async def get_ibm_models(
538541
if project_id:
539542
embed_params["project_id"] = project_id
540543

541-
embed_response = await client.get(
542-
models_url, params=embed_params, headers=headers, timeout=10.0
544+
embed_response = await request_with_retry(
545+
client, "GET", models_url, params=embed_params, headers=headers, timeout=10.0
543546
)
544547

545548
if embed_response.status_code == 200:

src/utils/provider_health_cache.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
"""Short-TTL cache for ``GET /api/provider/health`` responses.
2+
3+
The provider-health banner mounts on every page and polls the endpoint every
4+
5-30 seconds (per browser tab). With multiple tabs open, this fans out to many
5+
identical watsonx.ai validation calls and significantly raises the chance of
6+
hitting watsonx rate limits. A small in-process cache coalesces concurrent
7+
identical health checks so a single watsonx round-trip serves all of them.
8+
9+
Only successful (200) responses on the default polled path are cached. The
10+
explicit ``?provider=`` query bypass and the 503 error path are not cached, so
11+
real outages and on-demand checks are never masked.
12+
"""
13+
14+
import hashlib
15+
import os
16+
17+
from cachetools import TTLCache
18+
19+
_DEFAULT_TTL_S = 10
20+
21+
22+
def _ttl_seconds() -> int:
23+
raw = os.environ.get("OPENRAG_PROVIDER_HEALTH_TTL")
24+
if not raw:
25+
return _DEFAULT_TTL_S
26+
try:
27+
value = int(raw)
28+
except ValueError:
29+
return _DEFAULT_TTL_S
30+
return value if value > 0 else _DEFAULT_TTL_S
31+
32+
33+
_HEALTH_CACHE: TTLCache[str, dict] = TTLCache(maxsize=64, ttl=_ttl_seconds())
34+
35+
36+
def _fingerprint(value: str | None) -> str:
37+
return hashlib.sha256((value or "").encode()).hexdigest()[:16]
38+
39+
40+
def cache_key(
41+
provider: str | None,
42+
embedding_provider: str | None,
43+
test_completion: bool,
44+
llm_model: str | None,
45+
embedding_model: str | None,
46+
endpoint: str | None,
47+
project_id: str | None,
48+
api_key: str | None,
49+
embedding_api_key: str | None = None,
50+
embedding_endpoint: str | None = None,
51+
embedding_project_id: str | None = None,
52+
) -> str:
53+
"""Build the cache key for a polled health-check call.
54+
55+
The API keys are hashed (never stored in plaintext); rotating a key busts
56+
the cache automatically because the fingerprint changes.
57+
"""
58+
parts = [
59+
provider or "",
60+
embedding_provider or "",
61+
"1" if test_completion else "0",
62+
llm_model or "",
63+
embedding_model or "",
64+
endpoint or "",
65+
project_id or "",
66+
_fingerprint(api_key),
67+
embedding_endpoint or "",
68+
embedding_project_id or "",
69+
_fingerprint(embedding_api_key),
70+
]
71+
return hashlib.sha256("|".join(parts).encode()).hexdigest()
72+
73+
74+
def get(key: str) -> dict | None:
75+
return _HEALTH_CACHE.get(key)
76+
77+
78+
def set_(key: str, value: dict) -> None:
79+
_HEALTH_CACHE[key] = value
80+
81+
82+
def invalidate() -> None:
83+
"""Clear the entire cache. Intended for settings-save flows and tests."""
84+
_HEALTH_CACHE.clear()

src/utils/watsonx_retry.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""Async retry helper for watsonx.ai HTTP calls.
2+
3+
Retries only on HTTP 429 using AWS-style "Full Jitter" exponential backoff,
4+
honors the ``Retry-After`` response header when present, and bounds both the
5+
attempt count and the total wall-clock budget so a caller (e.g. a health-check
6+
poll) cannot hang waiting on a throttled upstream.
7+
8+
See https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/.
9+
"""
10+
11+
import asyncio
12+
import random
13+
import time
14+
from datetime import UTC, datetime
15+
from email.utils import parsedate_to_datetime
16+
17+
import httpx
18+
19+
from utils.logging_config import get_logger
20+
21+
logger = get_logger(__name__)
22+
23+
24+
def _retry_after_seconds(resp: httpx.Response, cap: float) -> float | None:
25+
"""Parse the ``Retry-After`` header (int seconds or HTTP-date), clamped at ``cap``.
26+
27+
Returns ``None`` when the header is absent or unparseable, so the caller
28+
falls back to jittered exponential backoff.
29+
"""
30+
raw = resp.headers.get("Retry-After")
31+
if not raw:
32+
return None
33+
try:
34+
return min(max(0.0, float(raw)), cap)
35+
except (TypeError, ValueError):
36+
pass
37+
try:
38+
dt = parsedate_to_datetime(raw)
39+
if dt is None:
40+
return None
41+
if dt.tzinfo is None:
42+
dt = dt.replace(tzinfo=UTC)
43+
delta = (dt - datetime.now(UTC)).total_seconds()
44+
return max(0.0, min(delta, cap))
45+
except (TypeError, ValueError):
46+
return None
47+
48+
49+
async def request_with_retry(
50+
client: httpx.AsyncClient,
51+
method: str,
52+
url: str,
53+
*,
54+
max_attempts: int = 5,
55+
base: float = 1.0,
56+
cap: float = 15.0,
57+
total_cap_s: float = 30.0,
58+
**kwargs,
59+
) -> httpx.Response:
60+
"""Issue an httpx request, retrying only on HTTP 429.
61+
62+
Non-429 responses (including 4xx auth failures and 5xx server errors) are
63+
returned to the caller on the first attempt so existing error-shaping logic
64+
runs unchanged.
65+
"""
66+
start = time.monotonic()
67+
resp: httpx.Response | None = None
68+
for attempt in range(1, max_attempts + 1):
69+
resp = await client.request(method, url, **kwargs)
70+
if resp.status_code != 429:
71+
return resp
72+
if attempt == max_attempts:
73+
break
74+
elapsed = time.monotonic() - start
75+
remaining = total_cap_s - elapsed
76+
if remaining <= 0:
77+
break
78+
ra = _retry_after_seconds(resp, cap)
79+
if ra is not None:
80+
sleep = ra
81+
else:
82+
# Full Jitter: random.uniform(0, min(cap, base * 2 ** (attempt - 1)))
83+
sleep = random.uniform(0, min(cap, base * (2 ** (attempt - 1))))
84+
sleep = min(sleep, remaining)
85+
logger.warning(
86+
"watsonx 429 on %s %s (attempt %d/%d), sleeping %.2fs",
87+
method,
88+
url,
89+
attempt,
90+
max_attempts,
91+
sleep,
92+
)
93+
await asyncio.sleep(sleep)
94+
return resp # type: ignore[return-value] # loop runs at least once

tests/unit/utils/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)