wavespeed 0.1.0__tar.gz → 0.1.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/python-publish.yml +2 -2
  2. {wavespeed-0.1.0 → wavespeed-0.1.2}/PKG-INFO +56 -7
  3. {wavespeed-0.1.0 → wavespeed-0.1.2}/README.md +55 -6
  4. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/_version.py +1 -1
  5. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/api/__init__.py +20 -0
  6. wavespeed-0.1.2/src/wavespeed/api/client.py +393 -0
  7. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/config.py +9 -0
  8. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/PKG-INFO +56 -7
  9. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_api.py +13 -1
  10. wavespeed-0.1.0/src/wavespeed/api/client.py +0 -238
  11. {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/claude-code-review.yml +0 -0
  12. {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/claude.yml +0 -0
  13. {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/pre-commit.yml +0 -0
  14. {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/python-packages.yml +0 -0
  15. {wavespeed-0.1.0 → wavespeed-0.1.2}/.gitignore +0 -0
  16. {wavespeed-0.1.0 → wavespeed-0.1.2}/.pre-commit-config.yaml +0 -0
  17. {wavespeed-0.1.0 → wavespeed-0.1.2}/CLAUDE.md +0 -0
  18. {wavespeed-0.1.0 → wavespeed-0.1.2}/LICENSE +0 -0
  19. {wavespeed-0.1.0 → wavespeed-0.1.2}/MANIFEST.in +0 -0
  20. {wavespeed-0.1.0 → wavespeed-0.1.2}/VERSIONING.md +0 -0
  21. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/build_and_push.sh +0 -0
  22. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/Dockerfile +0 -0
  23. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/build_and_test.sh +0 -0
  24. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/handler.py +0 -0
  25. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/requirements.txt +0 -0
  26. {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/test_input.json +0 -0
  27. {wavespeed-0.1.0 → wavespeed-0.1.2}/pyproject.toml +0 -0
  28. {wavespeed-0.1.0 → wavespeed-0.1.2}/pytest.ini +0 -0
  29. {wavespeed-0.1.0 → wavespeed-0.1.2}/setup.cfg +0 -0
  30. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/__init__.py +0 -0
  31. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/_config_module.py +0 -0
  32. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/__init__.py +0 -0
  33. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/__init__.py +0 -0
  34. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/fastapi.py +0 -0
  35. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/handler.py +0 -0
  36. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/heartbeat.py +0 -0
  37. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/http.py +0 -0
  38. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/job.py +0 -0
  39. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/local.py +0 -0
  40. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/logger.py +0 -0
  41. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/progress.py +0 -0
  42. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/scaler.py +0 -0
  43. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/state.py +0 -0
  44. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/__init__.py +0 -0
  45. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/boto3_utils.py +0 -0
  46. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/validator.py +0 -0
  47. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/worker.py +0 -0
  48. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/SOURCES.txt +0 -0
  49. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/dependency_links.txt +0 -0
  50. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/not-zip-safe +0 -0
  51. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/requires.txt +0 -0
  52. {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/top_level.txt +0 -0
  53. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/__init__.py +0 -0
  54. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/conftest.py +0 -0
  55. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_config.py +0 -0
  56. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/__init__.py +0 -0
  57. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/__init__.py +0 -0
  58. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_fastapi.py +0 -0
  59. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_handler.py +0 -0
  60. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_http.py +0 -0
  61. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_job.py +0 -0
  62. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_local.py +0 -0
  63. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_logger.py +0 -0
  64. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_progress.py +0 -0
  65. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_scaler.py +0 -0
  66. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_state.py +0 -0
  67. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_utils/__init__.py +0 -0
  68. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_utils/test_validator.py +0 -0
  69. {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_worker.py +0 -0
@@ -39,8 +39,8 @@ jobs:
39
39
  - name: Install dependencies
40
40
  run: |
41
41
  python -m pip install --upgrade pip
42
- pip install build 'setuptools>=64' 'setuptools_scm>=8'
42
+ pip install build wheel 'setuptools>=64' 'setuptools_scm>=8'
43
43
  - name: Build package
44
- run: python -m build --sdist -n
44
+ run: python -m build
45
45
  - name: Publish package
46
46
  uses: pypa/gh-action-pypi-publish@release/v1
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: wavespeed
3
- Version: 0.1.0
3
+ Version: 0.1.2
4
4
  Summary: WaveSpeedAI Python Client — Official Python SDK for WaveSpeedAI inference platform
5
5
  Author-email: WaveSpeedAI <support@wavespeed.ai>
6
6
  License: MIT
@@ -62,7 +62,7 @@ import wavespeed
62
62
 
63
63
  output = wavespeed.run(
64
64
  "wavespeed-ai/z-image/turbo",
65
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
65
+ input={"prompt": "Cat"},
66
66
  )
67
67
 
68
68
  print(output["outputs"][0]) # Output URL
@@ -70,7 +70,7 @@ print(output["outputs"][0]) # Output URL
70
70
 
71
71
  ### Authentication
72
72
 
73
- Set your API key via environment variable:
73
+ Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
74
74
 
75
75
  ```bash
76
76
  export WAVESPEED_API_KEY="your-api-key"
@@ -82,7 +82,7 @@ Or pass it directly:
82
82
  from wavespeed import Client
83
83
 
84
84
  client = Client(api_key="your-api-key")
85
- output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds a sign that says 'Hello World'"})
85
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
86
86
  ```
87
87
 
88
88
  ### Options
@@ -90,9 +90,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
90
90
  ```python
91
91
  output = wavespeed.run(
92
92
  "wavespeed-ai/z-image/turbo",
93
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
94
- timeout=36000.0 # Max wait time in seconds (default: 36000.0)
95
- poll_interval=1.0, # Status check interval (default: 1.0)
93
+ input={"prompt": "Cat"},
94
+ timeout=36000.0, # Max wait time in seconds (default: 36000.0)
95
+ poll_interval=1.0, # Status check interval (default: 1.0)
96
+ enable_sync_mode=False, # Single request mode, no polling (default: False)
97
+ )
98
+ ```
99
+
100
+ ### Sync Mode
101
+
102
+ Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
103
+
104
+ > **Note:** Not all models support sync mode. Check the model documentation for availability.
105
+
106
+ ```python
107
+ output = wavespeed.run(
108
+ "wavespeed-ai/z-image/turbo",
109
+ input={"prompt": "Cat"},
110
+ enable_sync_mode=True,
111
+ )
112
+ ```
113
+
114
+ ### Retry Configuration
115
+
116
+ Configure retries at the client level:
117
+
118
+ ```python
119
+ from wavespeed import Client
120
+
121
+ client = Client(
122
+ api_key="your-api-key",
123
+ max_retries=0, # Task-level retries (default: 0)
124
+ max_connection_retries=5, # HTTP connection retries (default: 5)
125
+ retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
96
126
  )
97
127
  ```
98
128
 
@@ -175,6 +205,25 @@ def handler(job):
175
205
  return {"output": "done"}
176
206
  ```
177
207
 
208
+ ### Concurrent Execution
209
+
210
+ Enable concurrent job processing with `concurrency_modifier`:
211
+
212
+ ```python
213
+ import wavespeed.serverless as serverless
214
+
215
+ def handler(job):
216
+ return {"output": job["input"]["data"]}
217
+
218
+ def concurrency_modifier(current_concurrency):
219
+ return 2 # Process 2 jobs concurrently
220
+
221
+ serverless.start({
222
+ "handler": handler,
223
+ "concurrency_modifier": concurrency_modifier
224
+ })
225
+ ```
226
+
178
227
  ## Local Development
179
228
 
180
229
  ### Test with JSON Input
@@ -33,7 +33,7 @@ import wavespeed
33
33
 
34
34
  output = wavespeed.run(
35
35
  "wavespeed-ai/z-image/turbo",
36
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
36
+ input={"prompt": "Cat"},
37
37
  )
38
38
 
39
39
  print(output["outputs"][0]) # Output URL
@@ -41,7 +41,7 @@ print(output["outputs"][0]) # Output URL
41
41
 
42
42
  ### Authentication
43
43
 
44
- Set your API key via environment variable:
44
+ Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
45
45
 
46
46
  ```bash
47
47
  export WAVESPEED_API_KEY="your-api-key"
@@ -53,7 +53,7 @@ Or pass it directly:
53
53
  from wavespeed import Client
54
54
 
55
55
  client = Client(api_key="your-api-key")
56
- output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds a sign that says 'Hello World'"})
56
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
57
57
  ```
58
58
 
59
59
  ### Options
@@ -61,9 +61,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
61
61
  ```python
62
62
  output = wavespeed.run(
63
63
  "wavespeed-ai/z-image/turbo",
64
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
65
- timeout=36000.0 # Max wait time in seconds (default: 36000.0)
66
- poll_interval=1.0, # Status check interval (default: 1.0)
64
+ input={"prompt": "Cat"},
65
+ timeout=36000.0, # Max wait time in seconds (default: 36000.0)
66
+ poll_interval=1.0, # Status check interval (default: 1.0)
67
+ enable_sync_mode=False, # Single request mode, no polling (default: False)
68
+ )
69
+ ```
70
+
71
+ ### Sync Mode
72
+
73
+ Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
74
+
75
+ > **Note:** Not all models support sync mode. Check the model documentation for availability.
76
+
77
+ ```python
78
+ output = wavespeed.run(
79
+ "wavespeed-ai/z-image/turbo",
80
+ input={"prompt": "Cat"},
81
+ enable_sync_mode=True,
82
+ )
83
+ ```
84
+
85
+ ### Retry Configuration
86
+
87
+ Configure retries at the client level:
88
+
89
+ ```python
90
+ from wavespeed import Client
91
+
92
+ client = Client(
93
+ api_key="your-api-key",
94
+ max_retries=0, # Task-level retries (default: 0)
95
+ max_connection_retries=5, # HTTP connection retries (default: 5)
96
+ retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
67
97
  )
68
98
  ```
69
99
 
@@ -146,6 +176,25 @@ def handler(job):
146
176
  return {"output": "done"}
147
177
  ```
148
178
 
179
+ ### Concurrent Execution
180
+
181
+ Enable concurrent job processing with `concurrency_modifier`:
182
+
183
+ ```python
184
+ import wavespeed.serverless as serverless
185
+
186
+ def handler(job):
187
+ return {"output": job["input"]["data"]}
188
+
189
+ def concurrency_modifier(current_concurrency):
190
+ return 2 # Process 2 jobs concurrently
191
+
192
+ serverless.start({
193
+ "handler": handler,
194
+ "concurrency_modifier": concurrency_modifier
195
+ })
196
+ ```
197
+
149
198
  ## Local Development
150
199
 
151
200
  ### Test with JSON Input
@@ -1,2 +1,2 @@
1
1
  # This file is automatically generated by setuptools_scm
2
- __version__ = "0.1.0"
2
+ __version__ = "0.1.2"
@@ -41,6 +41,8 @@ def run(
41
41
  *,
42
42
  timeout: float | None = None,
43
43
  poll_interval: float = 1.0,
44
+ enable_sync_mode: bool = False,
45
+ max_retries: int | None = None,
44
46
  ) -> dict:
45
47
  """Run a model and wait for the output.
46
48
 
@@ -49,6 +51,8 @@ def run(
49
51
  input: Input parameters for the model.
50
52
  timeout: Maximum time to wait for completion (None = no timeout).
51
53
  poll_interval: Interval between status checks in seconds.
54
+ enable_sync_mode: If True, use synchronous mode (single request).
55
+ max_retries: Maximum retries for this request (overrides default setting).
52
56
 
53
57
  Returns:
54
58
  Dict containing "outputs" array with model outputs.
@@ -64,12 +68,28 @@ def run(
64
68
  input={"prompt": "A cat sitting on a windowsill"}
65
69
  )
66
70
  print(output["outputs"][0]) # First output URL
71
+
72
+ # With sync mode
73
+ output = wavespeed.run(
74
+ "wavespeed-ai/z-image/turbo",
75
+ input={"prompt": "A cat"},
76
+ enable_sync_mode=True
77
+ )
78
+
79
+ # With retry
80
+ output = wavespeed.run(
81
+ "wavespeed-ai/z-image/turbo",
82
+ input={"prompt": "A cat"},
83
+ max_retries=3
84
+ )
67
85
  """
68
86
  return _get_default_client().run(
69
87
  model,
70
88
  input=input,
71
89
  timeout=timeout,
72
90
  poll_interval=poll_interval,
91
+ enable_sync_mode=enable_sync_mode,
92
+ max_retries=max_retries,
73
93
  )
74
94
 
75
95
 
@@ -0,0 +1,393 @@
1
+ """WaveSpeed API client implementation."""
2
+
3
+ import os
4
+ import time
5
+ import traceback
6
+ from typing import Any, BinaryIO
7
+
8
+ import requests
9
+
10
+ from wavespeed.config import api as api_config
11
+
12
+
13
+ class Client:
14
+ """WaveSpeed API client.
15
+
16
+ Args:
17
+ api_key: WaveSpeed API key. If not provided, uses wavespeed.config.api.api_key.
18
+ base_url: Base URL for the API. If not provided, uses wavespeed.config.api.base_url.
19
+ connection_timeout: Timeout for HTTP requests in seconds.
20
+ max_retries: Maximum number of retries for the entire operation.
21
+ max_connection_retries: Maximum retries for individual HTTP requests.
22
+ retry_interval: Base interval between retries in seconds.
23
+
24
+ Example:
25
+ client = Client(api_key="your-api-key")
26
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
27
+
28
+ # With sync mode (single request, waits for result)
29
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"}, enable_sync_mode=True)
30
+
31
+ # With retry
32
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"}, max_retries=3)
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ api_key: str | None = None,
38
+ base_url: str | None = None,
39
+ connection_timeout: float | None = None,
40
+ max_retries: int | None = None,
41
+ max_connection_retries: int | None = None,
42
+ retry_interval: float | None = None,
43
+ ) -> None:
44
+ """Initialize the client."""
45
+ self.api_key = api_key or api_config.api_key
46
+ self.base_url = (base_url or api_config.base_url).rstrip("/")
47
+ self.connection_timeout = connection_timeout or api_config.connection_timeout
48
+ self.max_retries = (
49
+ max_retries if max_retries is not None else api_config.max_retries
50
+ )
51
+ self.max_connection_retries = (
52
+ max_connection_retries
53
+ if max_connection_retries is not None
54
+ else api_config.max_connection_retries
55
+ )
56
+ self.retry_interval = (
57
+ retry_interval if retry_interval is not None else api_config.retry_interval
58
+ )
59
+
60
+ def _get_headers(self) -> dict[str, str]:
61
+ """Get request headers with authentication."""
62
+ if not self.api_key:
63
+ raise ValueError(
64
+ "API key is required. Set WAVESPEED_API_KEY environment variable "
65
+ "or pass api_key to Client()."
66
+ )
67
+ return {
68
+ "Content-Type": "application/json",
69
+ "Authorization": f"Bearer {self.api_key}",
70
+ }
71
+
72
+ def _submit(
73
+ self,
74
+ model: str,
75
+ input: dict[str, Any] | None,
76
+ enable_sync_mode: bool = False,
77
+ timeout: float | None = None,
78
+ ) -> tuple[str | None, dict[str, Any] | None]:
79
+ """Submit a prediction request.
80
+
81
+ Args:
82
+ model: Model identifier.
83
+ input: Input parameters.
84
+ enable_sync_mode: If True, wait for result in single request.
85
+ timeout: Request timeout in seconds.
86
+
87
+ Returns:
88
+ Tuple of (request_id, result). In async mode, result is None.
89
+ In sync mode, request_id is None and result contains the response.
90
+
91
+ Raises:
92
+ RuntimeError: If submission fails after retries.
93
+ """
94
+ url = f"{self.base_url}/api/v3/{model}"
95
+ body = dict(input) if input else {}
96
+
97
+ if enable_sync_mode:
98
+ body["enable_sync_mode"] = True
99
+
100
+ request_timeout = timeout if timeout is not None else api_config.timeout
101
+ # Use connection timeout for connect, request_timeout for read
102
+ connect_timeout = (
103
+ min(self.connection_timeout, request_timeout)
104
+ if request_timeout
105
+ else self.connection_timeout
106
+ )
107
+ timeouts = (connect_timeout, request_timeout)
108
+
109
+ for retry in range(self.max_connection_retries + 1):
110
+ try:
111
+ response = requests.post(
112
+ url, json=body, headers=self._get_headers(), timeout=timeouts
113
+ )
114
+
115
+ if response.status_code != 200:
116
+ raise RuntimeError(
117
+ f"Failed to submit prediction: HTTP {response.status_code}: "
118
+ f"{response.text}"
119
+ )
120
+
121
+ result = response.json()
122
+
123
+ if enable_sync_mode:
124
+ return None, result
125
+
126
+ request_id = result.get("data", {}).get("id")
127
+ if not request_id:
128
+ raise RuntimeError(f"No request ID in response: {result}")
129
+
130
+ return request_id, None
131
+
132
+ except (
133
+ requests.exceptions.ConnectionError,
134
+ requests.exceptions.Timeout,
135
+ ) as e:
136
+ print(
137
+ f"Connection error on attempt {retry + 1}/{self.max_connection_retries + 1}:"
138
+ )
139
+ traceback.print_exc()
140
+
141
+ if retry < self.max_connection_retries:
142
+ delay = self.retry_interval * (retry + 1)
143
+ print(f"Retrying in {delay} seconds...")
144
+ time.sleep(delay)
145
+ else:
146
+ raise RuntimeError(
147
+ f"Failed to submit prediction after {self.max_connection_retries + 1} attempts"
148
+ ) from e
149
+
150
+ def _get_result(
151
+ self, request_id: str, timeout: float | None = None
152
+ ) -> dict[str, Any]:
153
+ """Get prediction result.
154
+
155
+ Args:
156
+ request_id: The prediction request ID.
157
+ timeout: Request timeout in seconds.
158
+
159
+ Returns:
160
+ Full API response.
161
+
162
+ Raises:
163
+ RuntimeError: If fetching result fails after retries.
164
+ """
165
+ url = f"{self.base_url}/api/v3/predictions/{request_id}/result"
166
+ request_timeout = timeout if timeout is not None else api_config.timeout
167
+ connect_timeout = (
168
+ min(self.connection_timeout, request_timeout)
169
+ if request_timeout
170
+ else self.connection_timeout
171
+ )
172
+ timeouts = (connect_timeout, request_timeout)
173
+
174
+ for retry in range(self.max_connection_retries + 1):
175
+ try:
176
+ response = requests.get(
177
+ url, headers=self._get_headers(), timeout=timeouts
178
+ )
179
+
180
+ if response.status_code != 200:
181
+ raise RuntimeError(
182
+ f"Failed to get result: HTTP {response.status_code}: {response.text}"
183
+ )
184
+
185
+ return response.json()
186
+
187
+ except (
188
+ requests.exceptions.ConnectionError,
189
+ requests.exceptions.Timeout,
190
+ ) as e:
191
+ print(
192
+ f"Connection error getting result on attempt {retry + 1}/{self.max_connection_retries + 1}:"
193
+ )
194
+ traceback.print_exc()
195
+
196
+ if retry < self.max_connection_retries:
197
+ delay = self.retry_interval * (retry + 1)
198
+ print(f"Retrying in {delay} seconds...")
199
+ time.sleep(delay)
200
+ else:
201
+ raise RuntimeError(
202
+ f"Failed to get result after {self.max_connection_retries + 1} attempts"
203
+ ) from e
204
+
205
+ def _wait(
206
+ self,
207
+ request_id: str,
208
+ timeout: float | None,
209
+ poll_interval: float,
210
+ ) -> dict[str, Any]:
211
+ """Wait for prediction to complete.
212
+
213
+ Args:
214
+ request_id: The prediction request ID.
215
+ timeout: Maximum wait time in seconds (None = no timeout).
216
+ poll_interval: Time between polls in seconds.
217
+
218
+ Returns:
219
+ Dict with "outputs" array.
220
+
221
+ Raises:
222
+ RuntimeError: If prediction fails.
223
+ TimeoutError: If prediction times out.
224
+ """
225
+ start_time = time.time()
226
+
227
+ while True:
228
+ # Check timeout
229
+ if timeout is not None:
230
+ elapsed = time.time() - start_time
231
+ if elapsed >= timeout:
232
+ raise TimeoutError(f"Prediction timed out after {timeout} seconds")
233
+
234
+ result = self._get_result(request_id, timeout=timeout)
235
+ data = result.get("data", {})
236
+ status = data.get("status")
237
+
238
+ if status == "completed":
239
+ return {"outputs": data.get("outputs", [])}
240
+
241
+ if status == "failed":
242
+ error = data.get("error") or "Unknown error"
243
+ raise RuntimeError(f"Prediction failed: {error}")
244
+
245
+ time.sleep(poll_interval)
246
+
247
+ def _is_retryable_error(self, error: Exception) -> bool:
248
+ """Determine if an error is worth retrying at the task level.
249
+
250
+ Args:
251
+ error: The exception to check.
252
+
253
+ Returns:
254
+ True if the error is retryable.
255
+ """
256
+ # Always retry timeout and connection errors
257
+ if isinstance(
258
+ error, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)
259
+ ):
260
+ return True
261
+
262
+ # Retry server errors (5xx) and rate limiting (429)
263
+ if isinstance(error, RuntimeError):
264
+ error_str = str(error)
265
+ if "HTTP 5" in error_str or "HTTP 429" in error_str:
266
+ return True
267
+
268
+ return False
269
+
270
+ def run(
271
+ self,
272
+ model: str,
273
+ input: dict[str, Any] | None = None,
274
+ *,
275
+ timeout: float | None = None,
276
+ poll_interval: float = 1.0,
277
+ enable_sync_mode: bool = False,
278
+ max_retries: int | None = None,
279
+ ) -> dict[str, Any]:
280
+ """Run a model and wait for the output.
281
+
282
+ Args:
283
+ model: Model identifier (e.g., "wavespeed-ai/flux-dev").
284
+ input: Input parameters for the model.
285
+ timeout: Maximum time to wait for completion (None = no timeout).
286
+ poll_interval: Interval between status checks in seconds.
287
+ enable_sync_mode: If True, use synchronous mode (single request).
288
+ max_retries: Maximum task-level retries (overrides client setting).
289
+
290
+ Returns:
291
+ Dict containing "outputs" array with model outputs.
292
+
293
+ Raises:
294
+ ValueError: If API key is not configured.
295
+ RuntimeError: If the prediction fails.
296
+ TimeoutError: If the prediction times out.
297
+ """
298
+ task_retries = max_retries if max_retries is not None else self.max_retries
299
+ last_error = None
300
+
301
+ for attempt in range(task_retries + 1):
302
+ try:
303
+ request_id, sync_result = self._submit(
304
+ model, input, enable_sync_mode=enable_sync_mode, timeout=timeout
305
+ )
306
+
307
+ if enable_sync_mode:
308
+ # In sync mode, extract outputs from the result
309
+ data = sync_result.get("data", {})
310
+ return {"outputs": data.get("outputs", [])}
311
+
312
+ return self._wait(request_id, timeout, poll_interval)
313
+
314
+ except Exception as e:
315
+ last_error = e
316
+ is_retryable = self._is_retryable_error(e)
317
+
318
+ if not is_retryable or attempt >= task_retries:
319
+ raise
320
+
321
+ print(f"Task attempt {attempt + 1}/{task_retries + 1} failed: {e}")
322
+ delay = self.retry_interval * (attempt + 1)
323
+ print(f"Retrying in {delay} seconds...")
324
+ time.sleep(delay)
325
+
326
+ # Should not reach here, but just in case
327
+ if last_error:
328
+ raise last_error
329
+ raise RuntimeError(f"All {task_retries + 1} attempts failed")
330
+
331
+ def upload(self, file: str | BinaryIO, *, timeout: float | None = None) -> str:
332
+ """Upload a file to WaveSpeed.
333
+
334
+ Args:
335
+ file: File path string or file-like object to upload.
336
+ timeout: Total API call timeout in seconds.
337
+
338
+ Returns:
339
+ URL of the uploaded file.
340
+
341
+ Raises:
342
+ ValueError: If API key is not configured.
343
+ FileNotFoundError: If file path does not exist.
344
+ RuntimeError: If upload fails.
345
+
346
+ Example:
347
+ url = client.upload("/path/to/image.png")
348
+ print(url)
349
+ """
350
+ if not self.api_key:
351
+ raise ValueError(
352
+ "API key is required. Set WAVESPEED_API_KEY environment variable "
353
+ "or pass api_key to Client()."
354
+ )
355
+
356
+ url = f"{self.base_url}/api/v3/media/upload/binary"
357
+ headers = {"Authorization": f"Bearer {self.api_key}"}
358
+ timeout = timeout or api_config.timeout
359
+ request_timeout = (min(self.connection_timeout, timeout), timeout)
360
+
361
+ if isinstance(file, str):
362
+ if not os.path.exists(file):
363
+ raise FileNotFoundError(f"File not found: {file}")
364
+ with open(file, "rb") as f:
365
+ files = {"file": (os.path.basename(file), f)}
366
+ response = requests.post(
367
+ url, headers=headers, files=files, timeout=request_timeout
368
+ )
369
+ else:
370
+ filename = getattr(file, "name", "upload")
371
+ if isinstance(filename, str) and os.path.sep in filename:
372
+ filename = os.path.basename(filename)
373
+ files = {"file": (filename, file)}
374
+ response = requests.post(
375
+ url, headers=headers, files=files, timeout=request_timeout
376
+ )
377
+
378
+ if response.status_code != 200:
379
+ raise RuntimeError(
380
+ f"Failed to upload file: HTTP {response.status_code}: {response.text}"
381
+ )
382
+
383
+ result = response.json()
384
+ if result.get("code") != 200:
385
+ raise RuntimeError(
386
+ f"Upload failed: {result.get('message', 'Unknown error')}"
387
+ )
388
+
389
+ download_url = result.get("data", {}).get("download_url")
390
+ if not download_url:
391
+ raise RuntimeError("Upload failed: no download_url in response")
392
+
393
+ return download_url
@@ -27,6 +27,15 @@ class api:
27
27
  # Total API call timeout in seconds
28
28
  timeout: float = 36000.0
29
29
 
30
+ # Maximum number of retries for the entire operation (task-level retries)
31
+ max_retries: int = 0
32
+
33
+ # Maximum number of retries for individual HTTP requests (connection errors, timeouts)
34
+ max_connection_retries: int = 5
35
+
36
+ # Base interval between retries in seconds (actual delay = retry_interval * attempt)
37
+ retry_interval: float = 1.0
38
+
30
39
 
31
40
  class serverless:
32
41
  """Serverless configuration options.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: wavespeed
3
- Version: 0.1.0
3
+ Version: 0.1.2
4
4
  Summary: WaveSpeedAI Python Client — Official Python SDK for WaveSpeedAI inference platform
5
5
  Author-email: WaveSpeedAI <support@wavespeed.ai>
6
6
  License: MIT
@@ -62,7 +62,7 @@ import wavespeed
62
62
 
63
63
  output = wavespeed.run(
64
64
  "wavespeed-ai/z-image/turbo",
65
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
65
+ input={"prompt": "Cat"},
66
66
  )
67
67
 
68
68
  print(output["outputs"][0]) # Output URL
@@ -70,7 +70,7 @@ print(output["outputs"][0]) # Output URL
70
70
 
71
71
  ### Authentication
72
72
 
73
- Set your API key via environment variable:
73
+ Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
74
74
 
75
75
  ```bash
76
76
  export WAVESPEED_API_KEY="your-api-key"
@@ -82,7 +82,7 @@ Or pass it directly:
82
82
  from wavespeed import Client
83
83
 
84
84
  client = Client(api_key="your-api-key")
85
- output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds a sign that says 'Hello World'"})
85
+ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
86
86
  ```
87
87
 
88
88
  ### Options
@@ -90,9 +90,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
90
90
  ```python
91
91
  output = wavespeed.run(
92
92
  "wavespeed-ai/z-image/turbo",
93
- input={"prompt": "A cat holds a sign that says 'Hello World'"},
94
- timeout=36000.0 # Max wait time in seconds (default: 36000.0)
95
- poll_interval=1.0, # Status check interval (default: 1.0)
93
+ input={"prompt": "Cat"},
94
+ timeout=36000.0, # Max wait time in seconds (default: 36000.0)
95
+ poll_interval=1.0, # Status check interval (default: 1.0)
96
+ enable_sync_mode=False, # Single request mode, no polling (default: False)
97
+ )
98
+ ```
99
+
100
+ ### Sync Mode
101
+
102
+ Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
103
+
104
+ > **Note:** Not all models support sync mode. Check the model documentation for availability.
105
+
106
+ ```python
107
+ output = wavespeed.run(
108
+ "wavespeed-ai/z-image/turbo",
109
+ input={"prompt": "Cat"},
110
+ enable_sync_mode=True,
111
+ )
112
+ ```
113
+
114
+ ### Retry Configuration
115
+
116
+ Configure retries at the client level:
117
+
118
+ ```python
119
+ from wavespeed import Client
120
+
121
+ client = Client(
122
+ api_key="your-api-key",
123
+ max_retries=0, # Task-level retries (default: 0)
124
+ max_connection_retries=5, # HTTP connection retries (default: 5)
125
+ retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
96
126
  )
97
127
  ```
98
128
 
@@ -175,6 +205,25 @@ def handler(job):
175
205
  return {"output": "done"}
176
206
  ```
177
207
 
208
+ ### Concurrent Execution
209
+
210
+ Enable concurrent job processing with `concurrency_modifier`:
211
+
212
+ ```python
213
+ import wavespeed.serverless as serverless
214
+
215
+ def handler(job):
216
+ return {"output": job["input"]["data"]}
217
+
218
+ def concurrency_modifier(current_concurrency):
219
+ return 2 # Process 2 jobs concurrently
220
+
221
+ serverless.start({
222
+ "handler": handler,
223
+ "concurrency_modifier": concurrency_modifier
224
+ })
225
+ ```
226
+
178
227
  ## Local Development
179
228
 
180
229
  ### Test with JSON Input
@@ -29,6 +29,10 @@ class TestClient(unittest.TestCase):
29
29
  """Test client initialization from config."""
30
30
  mock_config.api_key = "config-key"
31
31
  mock_config.base_url = "https://api.wavespeed.ai"
32
+ mock_config.connection_timeout = 10.0
33
+ mock_config.max_retries = 0
34
+ mock_config.max_connection_retries = 5
35
+ mock_config.retry_interval = 1.0
32
36
  client = Client()
33
37
  self.assertEqual(client.api_key, "config-key")
34
38
 
@@ -56,9 +60,12 @@ class TestClient(unittest.TestCase):
56
60
  mock_post.return_value = mock_response
57
61
 
58
62
  client = Client(api_key="test-key")
59
- request_id = client._submit("wavespeed-ai/z-image/turbo", {"prompt": "test"})
63
+ request_id, result = client._submit(
64
+ "wavespeed-ai/z-image/turbo", {"prompt": "test"}
65
+ )
60
66
 
61
67
  self.assertEqual(request_id, "req-123")
68
+ self.assertIsNone(result)
62
69
  mock_post.assert_called_once()
63
70
  call_args = mock_post.call_args
64
71
  self.assertIn("wavespeed-ai/z-image/turbo", call_args[0][0])
@@ -178,6 +185,11 @@ class TestModuleLevelRun(unittest.TestCase):
178
185
  # Mock config
179
186
  mock_config.api_key = "config-key"
180
187
  mock_config.base_url = "https://api.wavespeed.ai"
188
+ mock_config.connection_timeout = 10.0
189
+ mock_config.timeout = 36000.0
190
+ mock_config.max_retries = 0
191
+ mock_config.max_connection_retries = 5
192
+ mock_config.retry_interval = 1.0
181
193
 
182
194
  # Reset default client
183
195
  wavespeed.api._default_client = None
@@ -1,238 +0,0 @@
1
- """WaveSpeed API client implementation."""
2
-
3
- import os
4
- import time
5
- from typing import Any, BinaryIO
6
-
7
- import requests
8
-
9
- from wavespeed.config import api as api_config
10
-
11
-
12
- class Client:
13
- """WaveSpeed API client.
14
-
15
- Args:
16
- api_key: WaveSpeed API key. If not provided, uses wavespeed.config.api.api_key.
17
- base_url: Base URL for the API. If not provided, uses wavespeed.config.api.base_url.
18
- connection_timeout: Timeout for HTTP requests in seconds.
19
-
20
- Example:
21
- client = Client(api_key="your-api-key")
22
- output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Hello"})
23
- """
24
-
25
- def __init__(
26
- self,
27
- api_key: str | None = None,
28
- base_url: str | None = None,
29
- connection_timeout: float | None = None,
30
- ) -> None:
31
- """Initialize the client."""
32
- self.api_key = api_key or api_config.api_key
33
- self.base_url = (base_url or api_config.base_url).rstrip("/")
34
- self.connection_timeout = connection_timeout or api_config.connection_timeout
35
-
36
- def _get_headers(self) -> dict[str, str]:
37
- """Get request headers with authentication."""
38
- if not self.api_key:
39
- raise ValueError(
40
- "API key is required. Set WAVESPEED_API_KEY environment variable "
41
- "or pass api_key to Client()."
42
- )
43
- return {
44
- "Content-Type": "application/json",
45
- "Authorization": f"Bearer {self.api_key}",
46
- }
47
-
48
- def _submit(self, model: str, input: dict[str, Any] | None) -> str:
49
- """Submit a prediction request.
50
-
51
- Args:
52
- model: Model identifier.
53
- input: Input parameters.
54
-
55
- Returns:
56
- Request ID for polling.
57
-
58
- Raises:
59
- RuntimeError: If submission fails.
60
- """
61
- url = f"{self.base_url}/api/v3/{model}"
62
- body = input or {}
63
-
64
- response = requests.post(
65
- url, json=body, headers=self._get_headers(), timeout=self.connection_timeout
66
- )
67
-
68
- if response.status_code != 200:
69
- raise RuntimeError(
70
- f"Failed to submit prediction: HTTP {response.status_code}: "
71
- f"{response.text}"
72
- )
73
-
74
- result = response.json()
75
- request_id = result.get("data", {}).get("id")
76
-
77
- if not request_id:
78
- raise RuntimeError(f"No request ID in response: {result}")
79
-
80
- return request_id
81
-
82
- def _get_result(self, request_id: str) -> dict[str, Any]:
83
- """Get prediction result.
84
-
85
- Args:
86
- request_id: The prediction request ID.
87
-
88
- Returns:
89
- Full API response.
90
-
91
- Raises:
92
- RuntimeError: If fetching result fails.
93
- """
94
- url = f"{self.base_url}/api/v3/predictions/{request_id}/result"
95
-
96
- response = requests.get(
97
- url, headers=self._get_headers(), timeout=self.connection_timeout
98
- )
99
-
100
- if response.status_code != 200:
101
- raise RuntimeError(
102
- f"Failed to get result: HTTP {response.status_code}: {response.text}"
103
- )
104
-
105
- return response.json()
106
-
107
- def _wait(
108
- self,
109
- request_id: str,
110
- timeout: float | None,
111
- poll_interval: float,
112
- ) -> dict[str, Any]:
113
- """Wait for prediction to complete.
114
-
115
- Args:
116
- request_id: The prediction request ID.
117
- timeout: Maximum wait time in seconds (None = no timeout).
118
- poll_interval: Time between polls in seconds.
119
-
120
- Returns:
121
- Dict with "outputs" array.
122
-
123
- Raises:
124
- RuntimeError: If prediction fails.
125
- TimeoutError: If prediction times out.
126
- """
127
- start_time = time.time()
128
-
129
- while True:
130
- # Check timeout
131
- if timeout is not None:
132
- elapsed = time.time() - start_time
133
- if elapsed >= timeout:
134
- raise TimeoutError(f"Prediction timed out after {timeout} seconds")
135
-
136
- result = self._get_result(request_id)
137
- data = result.get("data", {})
138
- status = data.get("status")
139
-
140
- if status == "completed":
141
- return {"outputs": data.get("outputs", [])}
142
-
143
- if status == "failed":
144
- error = data.get("error") or "Unknown error"
145
- raise RuntimeError(f"Prediction failed: {error}")
146
-
147
- time.sleep(poll_interval)
148
-
149
- def run(
150
- self,
151
- model: str,
152
- input: dict[str, Any] | None = None,
153
- *,
154
- timeout: float | None = None,
155
- poll_interval: float = 1.0,
156
- ) -> dict[str, Any]:
157
- """Run a model and wait for the output.
158
-
159
- Args:
160
- model: Model identifier (e.g., "wavespeed-ai/flux-dev").
161
- input: Input parameters for the model.
162
- timeout: Maximum time to wait for completion (None = no timeout).
163
- poll_interval: Interval between status checks in seconds.
164
-
165
- Returns:
166
- Dict containing "outputs" array with model outputs.
167
-
168
- Raises:
169
- ValueError: If API key is not configured.
170
- RuntimeError: If the prediction fails.
171
- TimeoutError: If the prediction times out.
172
- """
173
- request_id = self._submit(model, input)
174
- return self._wait(request_id, timeout, poll_interval)
175
-
176
- def upload(self, file: str | BinaryIO, *, timeout: float | None = None) -> str:
177
- """Upload a file to WaveSpeed.
178
-
179
- Args:
180
- file: File path string or file-like object to upload.
181
- timeout: Total API call timeout in seconds.
182
-
183
- Returns:
184
- URL of the uploaded file.
185
-
186
- Raises:
187
- ValueError: If API key is not configured.
188
- FileNotFoundError: If file path does not exist.
189
- RuntimeError: If upload fails.
190
-
191
- Example:
192
- url = client.upload("/path/to/image.png")
193
- print(url)
194
- """
195
- if not self.api_key:
196
- raise ValueError(
197
- "API key is required. Set WAVESPEED_API_KEY environment variable "
198
- "or pass api_key to Client()."
199
- )
200
-
201
- url = f"{self.base_url}/api/v3/media/upload/binary"
202
- headers = {"Authorization": f"Bearer {self.api_key}"}
203
- timeout = timeout or api_config.timeout
204
- request_timeout = (min(self.connection_timeout, timeout), timeout)
205
-
206
- if isinstance(file, str):
207
- if not os.path.exists(file):
208
- raise FileNotFoundError(f"File not found: {file}")
209
- with open(file, "rb") as f:
210
- files = {"file": (os.path.basename(file), f)}
211
- response = requests.post(
212
- url, headers=headers, files=files, timeout=request_timeout
213
- )
214
- else:
215
- filename = getattr(file, "name", "upload")
216
- if isinstance(filename, str) and os.path.sep in filename:
217
- filename = os.path.basename(filename)
218
- files = {"file": (filename, file)}
219
- response = requests.post(
220
- url, headers=headers, files=files, timeout=request_timeout
221
- )
222
-
223
- if response.status_code != 200:
224
- raise RuntimeError(
225
- f"Failed to upload file: HTTP {response.status_code}: {response.text}"
226
- )
227
-
228
- result = response.json()
229
- if result.get("code") != 200:
230
- raise RuntimeError(
231
- f"Upload failed: {result.get('message', 'Unknown error')}"
232
- )
233
-
234
- download_url = result.get("data", {}).get("download_url")
235
- if not download_url:
236
- raise RuntimeError("Upload failed: no download_url in response")
237
-
238
- return download_url
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes