wavespeed 0.0.1__tar.gz → 0.0.2__tar.gz

{wavespeed-0.0.1 → wavespeed-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wavespeed
-Version: 0.0.1
+Version: 0.0.2
 Summary: Python client for Wavespeed
 Author: Wavespeed, Inc.
 License: MIT License

{wavespeed-0.0.1 → wavespeed-0.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "wavespeed"
-version = "0.0.1"
+version = "0.0.2"
 description = "Python client for Wavespeed"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -37,7 +37,7 @@ asyncio_mode = "auto"
 testpaths = "tests/"
 
 [tool.setuptools]
-packages = ["wavespeed"]
+packages = ["wavespeed", "wavespeed.schemas"]
 
 [tool.setuptools.package-data]
 "wavespeed" = ["py.typed"]

wavespeed-0.0.2/wavespeed/schemas/__init__.py

@@ -0,0 +1,7 @@
+"""
+Schemas for the Wavespeed API.
+"""
+
+from .prediction import PredictionUrls, Prediction, PredictionResponse
+
+__all__ = [ "PredictionUrls", "Prediction", "PredictionResponse"]

wavespeed-0.0.2/wavespeed/schemas/prediction.py

@@ -0,0 +1,76 @@
+import asyncio
+import logging
+import time
+from typing import List, Dict, Any, Optional
+from datetime import datetime
+from pydantic import BaseModel, Field
+import pydantic
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from wavespeed.client import Wavespeed
+
+class PredictionUrls(BaseModel):
+    """URLs associated with a prediction."""
+    get: str
+
+
+class Prediction(BaseModel):
+    """Response from a prediction API call."""
+    _client: "Wavespeed" = pydantic.PrivateAttr()
+
+    id: str
+    model: str
+    input: Dict[str, Any] | None = None
+    outputs: List[str]
+    urls: PredictionUrls
+    has_nsfw_contents: List[bool]
+    status: str
+    created_at: datetime
+    error: str = ""
+    executionTime: int
+
+    def wait(self) -> "Prediction":
+        while self.status not in ['completed', 'failed']:
+            time.sleep(self._client.poll_interval)
+            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)
+        return self
+    
+    async def async_wait(self) -> "Prediction":
+        while self.status not in ['completed', 'failed']:
+            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)
+        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)
+        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)
+        return self
+    
+    def _update_from_dict(self, data: Dict[str, Any]) -> None:
+        """Update the object from a dictionary, handling nested objects properly."""
+        for key, value in data.items():
+            if key == 'urls' and isinstance(value, dict):
+                self.urls = PredictionUrls(**value)
+            elif hasattr(self, key):
+                setattr(self, key, value)
+
+class PredictionResponse(BaseModel):
+    code: int
+    message: str
+    data: Prediction

{wavespeed-0.0.1 → wavespeed-0.0.2}/wavespeed.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wavespeed
-Version: 0.0.1
+Version: 0.0.2
 Summary: Python client for Wavespeed
 Author: Wavespeed, Inc.
 License: MIT License

{wavespeed-0.0.1 → wavespeed-0.0.2}/wavespeed.egg-info/SOURCES.txt

@@ -8,4 +8,6 @@ wavespeed.egg-info/PKG-INFO
 wavespeed.egg-info/SOURCES.txt
 wavespeed.egg-info/dependency_links.txt
 wavespeed.egg-info/requires.txt
-wavespeed.egg-info/top_level.txt
+wavespeed.egg-info/top_level.txt
+wavespeed/schemas/__init__.py
+wavespeed/schemas/prediction.py