Add urls property to Prediction (#128)

Signed-off-by: Mattt Zmuda <[email protected]>
replicate · Jul 31, 2023 · 6f35c4f · 6f35c4f
1 parent ea0fe2c
commit 6f35c4f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/replicate/prediction.py b/replicate/prediction.py
@@ -20,6 +20,7 @@ class Prediction(BaseModel):
     started_at: Optional[str]
     created_at: Optional[str]
     completed_at: Optional[str]
+    urls: Optional[Dict[str, str]]
 
     def wait(self) -> None:
         """Wait for prediction to finish."""

diff --git a/tests/test_prediction.py b/tests/test_prediction.py
@@ -156,6 +156,7 @@ def test_async_timings():
     assert prediction.created_at == "2022-04-26T20:00:40.658234Z"
     assert prediction.completed_at is None
     assert prediction.output is None
+    assert prediction.urls["get"] == "https://api.replicate.com/v1/predictions/p1"
     prediction.wait()
     assert prediction.created_at == "2022-04-26T20:00:40.658234Z"
     assert prediction.completed_at == "2022-04-26T20:02:27.648305Z"