wavespeed 0.0.3__tar.gz → 0.0.5__tar.gz

{wavespeed-0.0.3 → wavespeed-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wavespeed
-Version: 0.0.3
+Version: 0.0.5
 Summary: Python client for WaveSpeed AI
 Author: WaveSpeed AI
 License: MIT License
@@ -264,8 +264,7 @@ await prediction.async_reload() -> Prediction
 ## Environment Variables
 
 - `WAVESPEED_API_KEY`: Your WaveSpeed API key
-- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 1)
-- `WAVESPEED_TIMEOUT`: Timeout in seconds for API requests (default: 60)
+- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 0.5)
 
 ## License
 

{wavespeed-0.0.3 → wavespeed-0.0.5}/README.md

@@ -227,8 +227,7 @@ await prediction.async_reload() -> Prediction
 ## Environment Variables
 
 - `WAVESPEED_API_KEY`: Your WaveSpeed API key
-- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 1)
-- `WAVESPEED_TIMEOUT`: Timeout in seconds for API requests (default: 60)
+- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 0.5)
 
 ## License
 

{wavespeed-0.0.3 → wavespeed-0.0.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "wavespeed"
-version = "0.0.3"
+version = "0.0.5"
 description = "Python client for WaveSpeed AI"
 readme = "README.md"
 license = { file = "LICENSE" }

{wavespeed-0.0.3 → wavespeed-0.0.5}/tests/test_client.py

@@ -2,13 +2,19 @@
 Tests for the Wavespeed client.
 """
 
+import unittest
 import pytest
 import respx
 from httpx import Response
+import httpx
 from datetime import datetime
+import unittest.mock
 
-from wavespeed.client import WaveSpeed
+from wavespeed.client import WaveSpeed, RetryTransport
 from wavespeed.schemas.prediction import Prediction, PredictionUrls
+import logging
+
+logging.basicConfig(level=logging.INFO)
 
 
 @pytest.fixture
@@ -112,6 +118,86 @@ def mock_prediction_completed_response():
     }
 
 
+@pytest.fixture
+def mock_transport():
+    """Create a mock transport for testing retry functionality."""
+    class MockTransport(httpx.BaseTransport):
+        def __init__(self):
+            self.request_count = 0
+            self.requests = []
+            
+        def handle_request(self, request):
+            self.request_count += 1
+            self.requests.append(request)
+            
+            # Return 503 for the first two requests, then 200
+            if self.request_count <= 2:
+                response = Response(503, json={"error": "Service Unavailable"})
+                response._request = request  # Set the request on the response
+                return response
+            else:
+                response = Response(200, json={
+                    "code": 200,
+                    "message": "Success",
+                    "data": {
+                        "id": "test_prediction_id",
+                        "model": "wavespeed-ai/flux-dev",
+                        "input": {"prompt": "A test prompt"},
+                        "outputs": ["https://example.com/generated_image.jpg"],
+                        "urls": {"get": "https://api.wavespeed.ai/api/v2/predictions/test_prediction_id/result"},
+                        "has_nsfw_contents": [False],
+                        "status": "completed",
+                        "created_at": datetime.now().isoformat(),
+                        "error": "",
+                        "executionTime": 1000
+                    }
+                })
+                response._request = request  # Set the request on the response
+                return response
+    
+    return MockTransport()
+
+
+@pytest.fixture
+def mock_async_transport():
+    """Create a mock async transport for testing retry functionality."""
+    class MockAsyncTransport(httpx.AsyncBaseTransport):
+        def __init__(self):
+            self.request_count = 0
+            self.requests = []
+            
+        async def handle_async_request(self, request):
+            self.request_count += 1
+            self.requests.append(request)
+            
+            # Return 503 for the first two requests, then 200
+            if self.request_count <= 2:
+                response = Response(503, json={"error": "Service Unavailable"})
+                response._request = request  # Set the request on the response
+                return response
+            else:
+                response = Response(200, json={
+                    "code": 200,
+                    "message": "Success",
+                    "data": {
+                        "id": "test_prediction_id",
+                        "model": "wavespeed-ai/flux-dev",
+                        "input": {"prompt": "A test prompt"},
+                        "outputs": ["https://example.com/generated_image.jpg"],
+                        "urls": {"get": "https://api.wavespeed.ai/api/v2/predictions/test_prediction_id/result"},
+                        "has_nsfw_contents": [False],
+                        "status": "completed",
+                        "created_at": datetime.now().isoformat(),
+                        "error": "",
+                        "executionTime": 1000
+                    }
+                })
+                response._request = request  # Set the request on the response
+                return response
+    
+    return MockAsyncTransport()
+
+
 @respx.mock
 def test_run(client, mock_prediction_response, mock_prediction_completed_response):
     """Test the run method."""
@@ -274,6 +360,7 @@ def test_prediction_wait(client, mock_prediction_in_progress_response, mock_pred
     assert result.status == "completed"
     assert len(result.outputs) == 1
     assert result.outputs[0] == "https://example.com/generated_image.jpg"
+    assert result.has_nsfw_contents == [False]
 
 
 @respx.mock
@@ -313,3 +400,98 @@ async def test_prediction_async_wait(async_client, mock_prediction_in_progress_r
     assert result.status == "completed"
     assert len(result.outputs) == 1
     assert result.outputs[0] == "https://example.com/generated_image.jpg"
+    assert result.has_nsfw_contents == [False]
+
+@unittest.mock.patch("httpx.HTTPTransport.handle_request")
+def test_client_retry_integration(mock_send):
+    """Test that the client's retry logic works correctly when integrated with httpx."""
+    # Create a counter to track the number of requests
+    request_count = {"count": 0}
+    
+    # Configure the mock to return 502 for the first request, 503 for the second, then 200
+    def side_effect(request, **kwargs):
+        request_count["count"] += 1
+        if request_count["count"] == 1:
+            response = Response(502, json={"error": "Bad Gateway"})
+        elif request_count["count"] == 2:
+            response = Response(503, json={"error": "Service Unavailable"})
+        else:
+            response = Response(200, json={
+                "code": 200,
+                "message": "Success",
+                "data": {
+                    "id": "test_prediction_id",
+                    "model": "wavespeed-ai/flux-dev",
+                    "input": {"prompt": "A test prompt"},
+                    "outputs": ["https://example.com/generated_image.jpg"],
+                    "urls": {"get": "https://api.wavespeed.ai/api/v2/predictions/test_prediction_id/result"},
+                    "has_nsfw_contents": [False],
+                    "status": "completed",
+                    "created_at": datetime.now().isoformat(),
+                    "error": "",
+                    "executionTime": 1000
+                }
+            })
+        response._request = request
+        return response
+    mock_send.side_effect = side_effect
+    
+    # Create a test client
+    client = WaveSpeed(api_key="test_api_key")
+    
+    prediction = client.get_prediction("test_prediction_id")
+    assert prediction.status == "completed"
+    # Verify that the transport was called the expected number of times (3 total: 2 failures + 1 success)
+    assert request_count["count"] == 3
+    assert mock_send.call_count == 3
+
+
+@pytest.mark.asyncio
+@unittest.mock.patch("httpx.AsyncHTTPTransport.handle_async_request")
+async def test_client_async_retry_integration(mock_send):
+    """Test that the client's async retry logic works correctly when integrated with httpx."""
+    # Create a counter to track the number of requests
+    request_count = {"count": 0}
+    
+    # Configure the mock to return 502 for the first request, 503 for the second, then 200
+    async def side_effect(request, **kwargs):
+        request_count["count"] += 1
+        if request_count["count"] == 1:
+            response = Response(502, json={"error": "Bad Gateway"})
+        elif request_count["count"] == 2:
+            response = Response(503, json={"error": "Service Unavailable"})
+        else:
+            response = Response(200, json={
+                "code": 200,
+                "message": "Success",
+                "data": {
+                    "id": "test_prediction_id",
+                    "model": "wavespeed-ai/flux-dev",
+                    "input": {"prompt": "A test prompt"},
+                    "outputs": ["https://example.com/generated_image.jpg"],
+                    "urls": {"get": "https://api.wavespeed.ai/api/v2/predictions/test_prediction_id/result"},
+                    "has_nsfw_contents": [False],
+                    "status": "completed",
+                    "created_at": datetime.now().isoformat(),
+                    "error": "",
+                    "executionTime": 1000
+                }
+            })
+        response._request = request
+        return response
+    mock_send.side_effect = side_effect
+    
+    # Create a test client
+    client = WaveSpeed(api_key="test_api_key")
+    
+    # Send the request through the client
+    prediction = await client.async_get_prediction(
+        predictionId="test_prediction_id",
+    )
+    
+    # Verify the response is successful
+    assert prediction.status == "completed"
+    
+    # Verify that the transport was called the expected number of times (3 total: 2 failures + 1 success)
+    assert request_count["count"] == 3
+    assert mock_send.call_count == 3

wavespeed-0.0.5/wavespeed/client.py

@@ -0,0 +1,334 @@
+import os
+import time
+import httpx
+from typing import Dict, Any
+from urllib.parse import urljoin
+import asyncio
+import random
+from datetime import datetime
+from typing import Iterable, Optional, Mapping, Union
+import logging
+
+from wavespeed.schemas.prediction import Prediction
+
+class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
+    """A custom HTTP transport that automatically retries requests using an exponential backoff strategy
+    for specific HTTP status codes and request methods.
+    """
+
+    RETRYABLE_METHODS = frozenset(["HEAD", "GET", "PUT", "DELETE", "OPTIONS", "TRACE"])
+    RETRYABLE_STATUS_CODES = frozenset(
+        [
+            429,  # Too Many Requests
+            502,  # Bad Gateway
+            503,  # Service Unavailable
+            504,  # Gateway Timeout
+        ]
+    )
+    MAX_BACKOFF_WAIT = 60
+
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        wrapped_transport: Union[httpx.BaseTransport, httpx.AsyncBaseTransport],
+        *,
+        max_attempts: int = 10,
+        max_backoff_wait: float = MAX_BACKOFF_WAIT,
+        backoff_factor: float = 0.1,
+        jitter_ratio: float = 0.1,
+        retryable_methods: Optional[Iterable[str]] = None,
+        retry_status_codes: Optional[Iterable[int]] = None,
+    ) -> None:
+        self._wrapped_transport = wrapped_transport
+
+        if jitter_ratio < 0 or jitter_ratio > 0.5:
+            raise ValueError(
+                f"jitter ratio should be between 0 and 0.5, actual {jitter_ratio}"
+            )
+
+        self.max_attempts = max_attempts
+        self.backoff_factor = backoff_factor
+        self.retryable_methods = (
+            frozenset(retryable_methods)
+            if retryable_methods
+            else self.RETRYABLE_METHODS
+        )
+        self.retry_status_codes = (
+            frozenset(retry_status_codes)
+            if retry_status_codes
+            else self.RETRYABLE_STATUS_CODES
+        )
+        self.jitter_ratio = jitter_ratio
+        self.max_backoff_wait = max_backoff_wait
+
+    def _calculate_sleep(
+        self, attempts_made: int, headers: Union[httpx.Headers, Mapping[str, str]]
+    ) -> float:
+        retry_after_header = (headers.get("Retry-After") or "").strip()
+        if retry_after_header:
+            if retry_after_header.isdigit():
+                return float(retry_after_header)
+
+            try:
+                parsed_date = datetime.fromisoformat(retry_after_header).astimezone()
+                diff = (parsed_date - datetime.now().astimezone()).total_seconds()
+                if diff > 0:
+                    return min(diff, self.max_backoff_wait)
+            except ValueError:
+                pass
+
+        backoff = self.backoff_factor * (2 ** (attempts_made - 1))
+        jitter = (backoff * self.jitter_ratio) * random.choice([1, -1])  # noqa: S311
+        total_backoff = backoff + jitter
+        return min(total_backoff, self.max_backoff_wait)
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        response = self._wrapped_transport.handle_request(request)  # type: ignore
+
+        if request.method not in self.retryable_methods:
+            return response
+
+        remaining_attempts = self.max_attempts - 1
+        attempts_made = 1
+
+        while True:
+            if (
+                remaining_attempts < 1
+                or response.status_code not in self.retry_status_codes
+            ):
+                return response
+
+            response.close()
+
+            sleep_for = self._calculate_sleep(attempts_made, response.headers)
+            logging.info("Got %s, Retrying request after %s seconds", response.status_code, sleep_for)
+            time.sleep(sleep_for)
+
+            response = self._wrapped_transport.handle_request(request)  # type: ignore
+
+            attempts_made += 1
+            remaining_attempts -= 1
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        response = await self._wrapped_transport.handle_async_request(request)  # type: ignore
+
+        if request.method not in self.retryable_methods:
+            return response
+
+        remaining_attempts = self.max_attempts - 1
+        attempts_made = 1
+
+        while True:
+            if (
+                remaining_attempts < 1
+                or response.status_code not in self.retry_status_codes
+            ):
+                return response
+
+            await response.aclose()
+
+            sleep_for = self._calculate_sleep(attempts_made, response.headers)
+            logging.info("Got %s, Retrying request after %s seconds", response.status_code, sleep_for)
+            await asyncio.sleep(sleep_for)
+
+            response = await self._wrapped_transport.handle_async_request(request)  # type: ignore
+
+            attempts_made += 1
+            remaining_attempts -= 1
+
+    async def aclose(self) -> None:
+        await self._wrapped_transport.aclose()  # type: ignore
+
+    def close(self) -> None:
+        self._wrapped_transport.close()  # type: ignore
+
+class WaveSpeed:
+    """
+    A client for interacting with the Wavespeed AI API.
+    """
+    _async_client: Optional[httpx.AsyncClient] = None
+    _client: Optional[httpx.Client] = None
+    
+    def __init__(self, api_key: str="", base_url: str = "https://api.wavespeed.ai/api/v2/", timeout: int | None = 120, **kwargs):
+        """
+        Initialize the WaveSpeed client.
+        
+        Args:
+            api_key: Your WaveSpeed API key
+            base_url: Base URL for the API
+            timeout: Timeout in seconds for http client send request
+        """
+        self.api_key = api_key
+        if not self.api_key:
+            self.api_key = os.environ.get("WAVESPEED_API_KEY", '')
+        if not self.api_key:
+            raise ValueError("API key is required.")
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        self.base_url = base_url
+        self.timeout = timeout
+        self.poll_interval = float(os.environ.get("WAVESPEED_POLL_INTERVAL", 0.5))
+        self._client_kwargs = kwargs
+
+    @property
+    def async_client(self) -> httpx.AsyncClient:
+        if self._async_client:
+            return self._async_client
+        transport = self._client_kwargs.get("async_transport", None) or httpx.AsyncHTTPTransport()
+        self._async_client = httpx.AsyncClient(headers=self.headers, timeout=self.timeout, transport=RetryTransport(transport))
+        return self._async_client
+    
+    @property
+    def client(self) -> httpx.Client:
+        if self._client:
+            return self._client
+        transport = self._client_kwargs.get("transport", None) or httpx.HTTPTransport()
+        self._client = httpx.Client(headers=self.headers, timeout=self.timeout, transport=RetryTransport(transport))
+        return self._client
+
+    async def async_run(
+        self,
+        modelId: str,
+        input: Dict[str, Any],
+        **kwargs
+    ) -> Prediction:
+        """
+        Generate an image using the Wavespeed AI API.
+        
+        Args:
+            modelId: The ID of the model to use
+            input: Input parameters for the model
+            **kwargs: Additional parameters to pass to the API
+            
+        Returns:
+            The API response as a dictionary
+        """
+        url = urljoin(self.base_url, modelId)
+        
+        payload = input
+        
+        response = await self.async_client.post(
+            url,
+            headers=self.headers,
+            json=payload,
+        )
+        
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return await prediction.async_wait()
+    
+    def run(
+        self,
+        modelId: str,
+        input: Dict[str, Any],
+        **kwargs
+    ) -> Prediction:
+        """
+        Generate an image using the Wavespeed AI API.
+        
+        Args:
+            modelId: The ID of the model to use
+            input: Input parameters for the model
+            **kwargs: Additional parameters to pass to the API
+            
+        Returns:
+            The API response as a dictionary
+        """
+        url = urljoin(self.base_url, modelId)
+        
+        payload = input
+        
+        response = self.client.post(
+            url,
+            headers=self.headers,
+            json=payload,
+        )
+
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return prediction.wait()
+    
+    async def async_create(self, modelId: str, input: Dict[str, Any], **kwargs) -> Prediction:
+        url = urljoin(self.base_url, modelId)
+        payload = input
+        response = await self.async_client.post(
+            url,
+            headers=self.headers,
+            json=payload,
+        )
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return prediction
+    
+    def create(self, modelId: str, input: Dict[str, Any], **kwargs) -> Prediction:
+        url = urljoin(self.base_url, modelId)
+        payload = input
+        response = self.client.post(
+            url,
+            headers=self.headers,
+            json=payload,
+        )
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return prediction
+    
+    def get_prediction(self, predictionId: str) -> Prediction:
+        url = urljoin(self.base_url, f"predictions/{predictionId}/result")
+        response = self.client.get(
+            url,
+            headers=self.headers,
+        )
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return prediction
+    
+    async def async_get_prediction(self, predictionId: str) -> Prediction:
+        url = urljoin(self.base_url, f"predictions/{predictionId}/result")
+        response = await self.async_client.get(
+            url,
+            headers=self.headers,
+        )
+        # Raise an exception for HTTP errors
+        response.raise_for_status()
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        prediction = Prediction(**data['data'])
+        prediction._client = self
+        return prediction
+        
+    async def close(self):
+        """Close the httpx client session."""
+        self.client.close()
+        await self.async_client.aclose()
+        
+    def __str__(self) -> str:
+        """String representation of the client."""
+        return f"WaveSpeed()"

{wavespeed-0.0.3 → wavespeed-0.0.5}/wavespeed/schemas/prediction.py

@@ -35,8 +35,10 @@ class Prediction(BaseModel):
             print('Waiting for prediction to complete: ', self.urls.get, type(self.urls.get))
             response = self._client.client.get(self.urls.get)
             response.raise_for_status()
-            data = response.json()['data']
-            self._update_from_dict(data)
+            data = response.json()
+            if data.get('code') != 200:
+                raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+            self._update_from_dict(data['data'])
         return self
     
     async def async_wait(self) -> "Prediction":
@@ -44,22 +46,28 @@ class Prediction(BaseModel):
             await asyncio.sleep(self._client.poll_interval)
             response = await self._client.async_client.get(self.urls.get)
             response.raise_for_status()
-            data = response.json()['data']
-            self._update_from_dict(data)
+            data = response.json()
+            if data.get('code') != 200:
+                raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+            self._update_from_dict(data['data'])
         return self
     
     async def async_reload(self) -> "Prediction":
         response = await self._client.async_client.get(self.urls.get)
         response.raise_for_status()
-        data = response.json()['data']
-        self._update_from_dict(data)
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        self._update_from_dict(data['data'])
         return self
     
     def reload(self) -> "Prediction":
         response = self._client.client.get(self.urls.get)
         response.raise_for_status()
-        data = response.json()['data']
-        self._update_from_dict(data)
+        data = response.json()
+        if data.get('code') != 200:
+            raise ValueError(f"Unexpected code: {data.get('code')}, data: {data}")
+        self._update_from_dict(data['data'])
         return self
     
     def _update_from_dict(self, data: Dict[str, Any]) -> None:

{wavespeed-0.0.3 → wavespeed-0.0.5}/wavespeed.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wavespeed
-Version: 0.0.3
+Version: 0.0.5
 Summary: Python client for WaveSpeed AI
 Author: WaveSpeed AI
 License: MIT License
@@ -264,8 +264,7 @@ await prediction.async_reload() -> Prediction
 ## Environment Variables
 
 - `WAVESPEED_API_KEY`: Your WaveSpeed API key
-- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 1)
-- `WAVESPEED_TIMEOUT`: Timeout in seconds for API requests (default: 60)
+- `WAVESPEED_POLL_INTERVAL`: Interval in seconds for polling prediction status (default: 0.5)
 
 ## License
 

wavespeed-0.0.3/wavespeed/client.py

@@ -1,150 +0,0 @@
-import os
-import httpx
-from typing import Dict, Any
-from urllib.parse import urljoin
-
-from wavespeed.schemas.prediction import Prediction
-
-
-class WaveSpeed:
-    """
-    A client for interacting with the Wavespeed AI API.
-    """
-    
-    BASE_URL = "https://api.wavespeed.ai/api/v2/"
-    
-    def __init__(self, api_key: str):
-        """
-        Initialize the Wavespeed client.
-        
-        Args:
-            api_key: Your Wavespeed API key
-        """
-        self.api_key = api_key
-        if not api_key:
-            api_key = os.environ.get("WAVESPEED_API_KEY", '')
-        if not api_key:
-            raise ValueError("API key is required.")
-        self.headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {api_key}"
-        }
-        self.async_client = httpx.AsyncClient(headers=self.headers)
-        self.client = httpx.Client(headers=self.headers)
-        self.poll_interval = float(os.environ.get("WAVESPEED_POLL_INTERVAL", 1)) # seconds
-        self.timeout = int(os.environ.get("WAVESPEED_TIMEOUT", 60)) # seconds
-    
-    async def async_run(
-        self,
-        modelId: str,
-        input: Dict[str, Any],
-        **kwargs
-    ) -> Prediction:
-        """
-        Generate an image using the Wavespeed AI API.
-        
-        Args:
-            modelId: The ID of the model to use
-            input: Input parameters for the model
-            **kwargs: Additional parameters to pass to the API
-            
-        Returns:
-            The API response as a dictionary
-        """
-        url = urljoin(self.BASE_URL, modelId)
-        
-        payload = input
-        
-        # Reset client if it's closed
-        if self.async_client.is_closed:
-            self.async_client = httpx.AsyncClient()
-            
-        response = await self.async_client.post(
-            url,
-            headers=self.headers,
-            json=payload,
-            timeout=self.timeout
-        )
-        
-        # Raise an exception for HTTP errors
-        response.raise_for_status()
-        data = response.json()
-        prediction = Prediction(**data['data'])
-        prediction._client = self
-        return await prediction.async_wait()
-    
-    def run(
-        self,
-        modelId: str,
-        input: Dict[str, Any],
-        **kwargs
-    ) -> Prediction:
-        """
-        Generate an image using the Wavespeed AI API.
-        
-        Args:
-            modelId: The ID of the model to use
-            input: Input parameters for the model
-            **kwargs: Additional parameters to pass to the API
-            
-        Returns:
-            The API response as a dictionary
-        """
-        url = urljoin(self.BASE_URL, modelId)
-        
-        payload = input
-        
-        response = self.client.post(
-            url,
-            headers=self.headers,
-            json=payload,
-            timeout=self.timeout
-        )
-
-        # Raise an exception for HTTP errors
-        response.raise_for_status()
-        data = response.json()
-        prediction = Prediction(**data['data'])
-        prediction._client = self
-        return prediction.wait()
-    
-    async def async_create(self, modelId: str, input: Dict[str, Any], **kwargs) -> Prediction:
-        url = urljoin(self.BASE_URL, modelId)
-        payload = input
-        response = await self.async_client.post(
-            url,
-            headers=self.headers,
-            json=payload,
-            timeout=self.timeout
-        )
-        # Raise an exception for HTTP errors
-        response.raise_for_status()
-        data = response.json()
-        prediction = Prediction(**data['data'])
-        prediction._client = self
-        return prediction
-    
-    def create(self, modelId: str, input: Dict[str, Any], **kwargs) -> Prediction:
-        url = urljoin(self.BASE_URL, modelId)
-        payload = input
-        response = self.client.post(
-            url,
-            headers=self.headers,
-            json=payload,
-            timeout=self.timeout
-        )
-        # Raise an exception for HTTP errors
-        response.raise_for_status()
-        data = response.json()
-        prediction = Prediction(**data['data'])
-        prediction._client = self
-        return prediction
-        
-    async def close(self):
-        """Close the httpx client session."""
-        self.client.close()
-        await self.async_client.aclose()
-        
-    def __str__(self) -> str:
-        """String representation of the client."""
-        return f"WaveSpeed()"