fix: Detect and handle watsonx.ai model provider rate limit exceptions

Issue

- #1599

Summary

- Fix Ruff linting errors

Author: mpawlow

src/api/provider_health.py

@@ -1,22 +1,23 @@
 """Provider health check endpoint."""
 
 import asyncio
-from typing import Optional
+
 import httpx
 from fastapi import Depends
 from fastapi.responses import JSONResponse
-from utils.logging_config import get_logger
-from utils import provider_health_cache
-from config.settings import get_openrag_config
+
 from api.provider_validation import validate_provider_setup
+from config.settings import get_openrag_config
 from dependencies import get_current_user
 from session_manager import User
+from utils import provider_health_cache
+from utils.logging_config import get_logger
 
 logger = get_logger(__name__)
 
 
 async def check_provider_health(
-    provider: Optional[str] = None,
+    provider: str | None = None,
     test_completion: bool = False,
     user: User = Depends(get_current_user),
 ):

src/api/provider_validation.py

@@ -1,7 +1,9 @@
 """Provider validation utilities for testing API keys and models during onboarding."""
 
 import json
+
 import httpx
+
 from utils.container_utils import transform_localhost_url
 from utils.logging_config import get_logger
 from utils.watsonx_retry import request_with_retry
@@ -270,7 +272,7 @@ async def _test_openai_lightweight_health(api_key: str) -> None:
 
     except httpx.TimeoutException:
         logger.error("OpenAI lightweight health check timed out")
-        raise Exception("OpenAI API request timed out")
+        raise Exception("OpenAI API request timed out") from None
     except Exception as e:
         logger.error(f"OpenAI lightweight health check failed: {str(e)}")
         raise
@@ -340,7 +342,7 @@ async def _test_openai_completion_with_tools(api_key: str, llm_model: str) -> No
 
     except httpx.TimeoutException:
         logger.error("OpenAI completion test timed out")
-        raise Exception("Request timed out")
+        raise Exception("Request timed out") from None
     except Exception as e:
         logger.error(f"OpenAI completion test failed: {str(e)}")
         raise
@@ -382,7 +384,7 @@ async def _test_openai_embedding(api_key: str, embedding_model: str) -> None:
 
     except httpx.TimeoutException:
         logger.error("OpenAI embedding test timed out")
-        raise Exception("Request timed out")
+        raise Exception("Request timed out") from None
     except Exception as e:
         logger.error(f"OpenAI embedding test failed: {str(e)}")
         raise
@@ -425,7 +427,7 @@ async def _test_watsonx_lightweight_health(api_key: str, endpoint: str, project_
 
     except httpx.TimeoutException:
         logger.error("WatsonX lightweight health check timed out")
-        raise Exception("WatsonX API request timed out")
+        raise Exception("WatsonX API request timed out") from None
     except Exception as e:
         logger.error(f"WatsonX lightweight health check failed: {str(e)}")
         raise
@@ -516,7 +518,7 @@ async def _test_watsonx_completion_with_tools(
 
     except httpx.TimeoutException:
         logger.error("IBM Watson completion test timed out")
-        raise Exception("Request timed out")
+        raise Exception("Request timed out") from None
     except Exception as e:
         logger.error(f"IBM Watson completion test failed: {str(e)}")
         # If the error message contains JSON, parse it to extract just the message
@@ -525,7 +527,7 @@ async def _test_watsonx_completion_with_tools(
             json_part = error_str.split("IBM Watson API error: ", 1)[1]
             parsed_message = _parse_json_error_message(json_part)
             if parsed_message != json_part:
-                raise Exception(f"IBM Watson API error: {parsed_message}")
+                raise Exception(f"IBM Watson API error: {parsed_message}") from e
         raise
 
 
@@ -601,7 +603,7 @@ async def _test_watsonx_embedding(
 
     except httpx.TimeoutException:
         logger.error("IBM Watson embedding test timed out")
-        raise Exception("Request timed out")
+        raise Exception("Request timed out") from None
     except Exception as e:
         logger.error(f"IBM Watson embedding test failed: {str(e)}")
         # If the error message contains JSON, parse it to extract just the message
@@ -610,7 +612,7 @@ async def _test_watsonx_embedding(
             json_part = error_str.split("IBM Watson API error: ", 1)[1]
             parsed_message = _parse_json_error_message(json_part)
             if parsed_message != json_part:
-                raise Exception(f"IBM Watson API error: {parsed_message}")
+                raise Exception(f"IBM Watson API error: {parsed_message}") from e
         raise
 
 
@@ -640,7 +642,7 @@ async def _test_ollama_lightweight_health(endpoint: str) -> None:
 
     except httpx.TimeoutException:
         logger.error("Ollama lightweight health check timed out")
-        raise Exception("Ollama endpoint timed out")
+        raise Exception("Ollama endpoint timed out") from None
     except Exception as e:
         logger.error(f"Ollama lightweight health check failed: {str(e)}")
         raise
@@ -692,7 +694,7 @@ async def _test_ollama_completion_with_tools(llm_model: str, endpoint: str) -> N
 
     except httpx.TimeoutException:
         logger.error("Ollama completion test timed out")
-        raise httpx.TimeoutException("Ollama is busy or model inference timed out")
+        raise httpx.TimeoutException("Ollama is busy or model inference timed out") from None
     except Exception as e:
         logger.error(f"Ollama completion test failed: {str(e)}")
         raise
@@ -731,7 +733,7 @@ async def _test_ollama_embedding(embedding_model: str, endpoint: str) -> None:
 
     except httpx.TimeoutException:
         logger.error("Ollama embedding test timed out")
-        raise httpx.TimeoutException("Ollama is busy or embedding generation timed out")
+        raise httpx.TimeoutException("Ollama is busy or embedding generation timed out") from None
     except Exception as e:
         logger.error(f"Ollama embedding test failed: {str(e)}")
         raise
@@ -769,7 +771,7 @@ async def _test_anthropic_lightweight_health(api_key: str) -> None:
 
     except httpx.TimeoutException:
         logger.error("Anthropic lightweight health check timed out")
-        raise Exception("Anthropic API request timed out")
+        raise Exception("Anthropic API request timed out") from None
     except Exception as e:
         logger.error(f"Anthropic lightweight health check failed: {str(e)}")
         raise
@@ -823,7 +825,7 @@ async def _test_anthropic_completion_with_tools(api_key: str, llm_model: str) ->
 
     except httpx.TimeoutException:
         logger.error("Anthropic completion test timed out")
-        raise Exception("Request timed out")
+        raise Exception("Request timed out") from None
     except Exception as e:
         logger.error(f"Anthropic completion test failed: {str(e)}")
         raise

src/services/models_service.py

@@ -1,14 +1,15 @@
-import httpx
 import asyncio
-from typing import Dict, List, Optional
+
+import httpx
+
+from config.embedding_constants import OPENAI_DEFAULT_EMBEDDING_MODEL, OPENAI_EMBEDDING_MODEL_PREFIX
 from config.model_constants import (
     ANTHROPIC_DEFAULT_LANGUAGE_MODEL,
     ANTHROPIC_VALIDATION_MODELS,
     OLLAMA_DEFAULT_LANGUAGE_MODEL_PATTERN,
     OPENAI_DEFAULT_LANGUAGE_MODEL,
     OPENAI_VALIDATION_MODELS,
 )
-from config.embedding_constants import OPENAI_DEFAULT_EMBEDDING_MODEL, OPENAI_EMBEDDING_MODEL_PREFIX
 from utils.container_utils import transform_localhost_url
 from utils.logging_config import get_logger
 from utils.watsonx_retry import request_with_retry
@@ -32,7 +33,7 @@ class ModelsService:
     """Service for fetching available models from different AI providers and managing a model registry."""
 
     # Registry for caching model-to-provider mapping
-    _model_provider_registry: Dict[str, str] = {}
+    _model_provider_registry: dict[str, str] = {}
     _registry_lock = asyncio.Lock()
 
     def __init__(self):
@@ -125,7 +126,7 @@ async def update_model_registry(self):
     async def get_litellm_model_name(
         self,
         model_name: str,
-        provider: Optional[str] = None,
+        provider: str | None = None,
         strict: bool = False,
     ) -> str:
         """Resolve ``model_name`` to a LiteLLM-routable string.
@@ -170,7 +171,7 @@ async def get_litellm_model_name(
 
     async def get_openai_models(
         self, api_key: str, update_index: bool = True
-    ) -> Dict[str, List[Dict[str, str]]]:
+    ) -> dict[str, list[dict[str, str]]]:
         """Fetch available models from OpenAI API with lightweight validation"""
         try:
             headers = {
@@ -256,7 +257,7 @@ async def get_openai_models(
 
     async def get_anthropic_models(
         self, api_key: str, update_index: bool = True
-    ) -> Dict[str, List[Dict[str, str]]]:
+    ) -> dict[str, list[dict[str, str]]]:
         """Fetch available models from Anthropic API"""
         try:
             headers = {
@@ -324,7 +325,7 @@ async def get_anthropic_models(
 
     async def get_ollama_models(
         self, endpoint: str = None, update_index: bool = True
-    ) -> Dict[str, List[Dict[str, str]]]:
+    ) -> dict[str, list[dict[str, str]]]:
         """Fetch available models from Ollama API with tool calling capabilities for language models"""
         try:
             ollama_url = transform_localhost_url(endpoint)
@@ -446,7 +447,7 @@ async def get_ibm_models(
         api_key: str = None,
         project_id: str = None,
         update_index: bool = True,
-    ) -> Dict[str, List[Dict[str, str]]]:
+    ) -> dict[str, list[dict[str, str]]]:
         """Fetch available models from IBM Watson API"""
         try:
             # Use provided endpoint or default