Skip to content

Commit d662f2f

Browse files
authored
Add metrics field to Prediction (#136)
Signed-off-by: Mattt Zmuda <[email protected]>
1 parent 52f88eb commit d662f2f

File tree

2 files changed

+7
-0
lines changed

2 files changed

+7
-0
lines changed

replicate/prediction.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ class Prediction(BaseModel):
3535
error: Optional[str]
3636
"""The error encountered during the prediction, if any."""
3737

38+
metrics: Optional[Dict[str, Any]]
39+
"""Metrics for the prediction."""
40+
3841
created_at: Optional[str]
3942
"""When the prediction was created."""
4043

tests/test_prediction.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,9 @@ def test_async_timings():
193193
"output": "hello world",
194194
"error": None,
195195
"logs": "",
196+
"metrics": {
197+
"predict_time": 1.2345,
198+
},
196199
},
197200
)
198201

@@ -210,3 +213,4 @@ def test_async_timings():
210213
assert prediction.created_at == "2022-04-26T20:00:40.658234Z"
211214
assert prediction.completed_at == "2022-04-26T20:02:27.648305Z"
212215
assert prediction.output == "hello world"
216+
assert prediction.metrics["predict_time"] == 1.2345

0 commit comments

Comments
 (0)