refactor

Author: wemoveon2

tests/public_models/test_public_models_predicts.py

@@ -350,35 +350,76 @@ def test_multimodal_predict_on_public_models(channel):
 
 # Helper functions for the new OpenAI compatible endpoint test
 def _call_openai_model(model_id):
-    """Attempts to call a model using OpenAI's chat completions and image generation APIs."""
+    """
+    Attempts to call a model using OpenAI's chat completions and image generation APIs,
+    with an integrated retry mechanism for transient network errors.
+    """
     client = OpenAI(api_key=API_KEY, base_url="https://api.clarifai.com/v2/ext/openai/v1")
-    try:
-        response = client.chat.completions.create(
-            model=model_id,
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Who are you?"},
-            ],
-            max_tokens=50,
-        )
-        assert hasattr(response, "choices") and len(response.choices) > 0, "No choices in response"
-        return response, None
-    except Exception as e1:
+    last_err_chat = None
+    last_err_image = None
+
+    # --- Attempt 1: Chat Completions with Retry ---
+    for attempt in range(MAX_RETRY_ATTEMPTS):
+        try:
+            response = client.chat.completions.create(
+                model=model_id,
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Who are you?"},
+                ],
+                max_tokens=50,
+            )
+            if hasattr(response, 'choices') and response.choices:
+                return response, None  # Success
+            else:
+                last_err_chat = ValueError(
+                    f"Chat completions returned no choices. Response: {response}"
+                )
+                break  # Successful call but no data, so don't retry.
+        except (APIConnectionError, APITimeoutError, RateLimitError) as e:
+            last_err_chat = e
+            if attempt == MAX_RETRY_ATTEMPTS - 1:
+                break  # Last attempt failed, break to move on
+            print(
+                f"Retrying chat predict for '{model_id}' after error: {e}. Attempt #{attempt + 1}"
+            )
+            time.sleep(attempt + 1)  # Exponential backoff
+        except Exception as e:
+            last_err_chat = e
+            break  # Non-retriable error, break immediately
+
+    # --- Attempt 2: Image Generation with Retry ---
+    for attempt in range(MAX_RETRY_ATTEMPTS):
         try:
             response = client.images.generate(
                 model=model_id,
                 prompt="A cat and a dog sitting together in a park",
             )
-            assert hasattr(response, "data") and len(response.data) > 0, (
-                "No image data in response"
+            if hasattr(response, 'data') and response.data:
+                return response, None  # Success
+            else:
+                last_err_image = ValueError(
+                    f"Image generation returned no data. Response: {response}"
+                )
+                break  # Successful call but no data, so don't retry.
+        except (APIConnectionError, APITimeoutError, RateLimitError) as e:
+            last_err_image = e
+            if attempt == MAX_RETRY_ATTEMPTS - 1:
+                break  # Last attempt failed, break
+            print(
+                f"Retrying image predict for '{model_id}' after error: {e}. Attempt #{attempt + 1}"
             )
-            return response, None
-        except Exception as e2:
-            return None, f"chat.completions error: {e1}; image.generate error: {e2}"
+            time.sleep(attempt + 1)  # Exponential backoff
+        except Exception as e:
+            last_err_image = e
+            break  # Non-retriable error, break immediately
+
+    return None, f"chat.completions error: {last_err_chat}; image.generate error: {last_err_image}"
 
 
 def _list_featured_models(per_page=50):
     """Lists featured models from the Clarifai platform."""
+    # This function remains unchanged
     channel = ClarifaiChannel.get_grpc_channel()
     stub = service_pb2_grpc.V2Stub(channel)
     auth_metadata = (("authorization", f"Key {API_KEY}"),)
@@ -389,9 +430,16 @@ def _list_featured_models(per_page=50):
     return response.models
 
 
+# The test functions below remain unchanged as the retry logic
+# is now encapsulated within the _call_openai_model helper.
+
+
 # New integrated test
 def test_openai_compatible_endpoint_on_featured_models():
     """Tests the OpenAI compatible endpoint with featured models."""
+    if not API_KEY:
+        pytest.skip("Skipping test: CLARIFAI_PAT environment variable not set.")
+
     featured_models = _list_featured_models()
     failed_models = []
 
@@ -411,17 +459,16 @@ def test_openai_compatible_endpoint_on_featured_models():
 # New integrated async test
 async def _call_openai_model_async(model_identifier, session):
     """Async helper to call a single model."""
-    # Note: The OpenAI library doesn't natively support asyncio for this type of call.
-    # To make this truly async, one would typically use an async-compatible HTTP client
-    # like `httpx` or `aiohttp`. For simplicity and to match the provided script's
-    # library, we run the synchronous `_call_openai_model` in a thread pool executor.
     loop = asyncio.get_running_loop()
     return await loop.run_in_executor(None, _call_openai_model, model_identifier)
 
 
 @pytest.mark.asyncio
 async def test_openai_compatible_endpoint_on_featured_models_async():
     """Tests the OpenAI compatible endpoint concurrently with featured models."""
+    if not API_KEY:
+        pytest.skip("Skipping test: CLARIFAI_PAT environment variable not set.")
+
     featured_models = _list_featured_models()
     tasks = []
     model_identifiers = []